[content] train_sft.py (python)

    def generate_sft_rollout(self, rollout_id: int) -> list[Sample]:
        """Generate SFT rollout data (tokenize and create loss masks)."""
        samples = self.data_source.get_samples(self.args.rollout_batch_size)

        result = []
        for i, (sample,) in enumerate(samples):
            messages = sample.prompt
            token_ids, loss_mask = self.mask_generator.get_loss_mask(messages)
            response_length = self.mask_generator.get_response_lengths([loss_mask])[0]

            sample.tokens = token_ids
            sample.response_length = response_length
            sample.reward = 0
            sample.loss_mask = loss_mask[-response_length:]
            result.append(sample)

            if i == 0 and rollout_id == 0 and dist.get_rank() == 0:
                logger.info(
                    f"SFT rollout sample: tokens_len={len(token_ids)}, "
                    f"response_length={response_length}"
                )

        return result

    def _convert_samples_to_train_data(self, samples: list[Sample]) -> dict:
        """Convert samples to training data format."""
        train_data = {
            "tokens": [sample.tokens for sample in samples],
            "response_lengths": [sample.response_length for sample in samples],
            "rewards": [0.0 for _ in samples],
            "raw_reward": [0.0 for _ in samples],
            "truncated": [0 for _ in samples],
            "sample_indices": [sample.index for sample in samples],
        }

        loss_masks = []
        for sample in samples:
            if sample.loss_mask is None:
                sample.loss_mask = [1] * sample.response_length
            loss_masks.append(sample.loss_mask)
        train_data["loss_masks"] = loss_masks

        return train_data

    def _split_train_data_by_dp(self, data: dict) -> dict:
        """Split training data for current DP rank."""
        total_lengths = [len(t) for t in data["tokens"]]
        data["total_lengths"] = total_lengths

        # Simple round-robin partitioning
        partition = list(range(self.dp_rank, len(total_lengths), self.dp_size))

        rollout_data = {"partition": partition, "total_lengths": total_lengths}

        for key in [
            "tokens",
            "response_lengths",
            "rewards",
            "raw_reward",
            "truncated",
            "loss_masks",
            "sample_indices",
        ]:
            if key in data:
                rollout_data[key] = [data[key][j] for j in partition]

        return rollout_data

    def _packed_data(self, rollout_data: dict) -> tuple[list[dict], list[int]]:
        """Pack variable-length sequences for efficient processing."""
        tokens = rollout_data["tokens"]

        packed_batches = []
        mbs_size_list = []
        local_batch_size = self.args.global_batch_size // self.dp_size

        if self.args.use_dynamic_batch_size:
            max_tokens = self.args.max_tokens_per_gpu
            if self.cp_size > 1:
                max_tokens = max_tokens * self.cp_size

            for i in range(0, len(tokens), local_batch_size):
                mbs_size_list.append(
                    get_minimum_num_micro_batch_size(
                        [len(t) for t in rollout_data["tokens"][i : i + local_batch_size]],
                        max_tokens,
                    )
                )
            num_microbatches = torch.tensor(
                mbs_size_list, dtype=torch.int, device=torch.cuda.current_device()
            )
            dist.all_reduce(num_microbatches, op=dist.ReduceOp.MAX, group=self.dp_group)
            num_microbatches = num_microbatches.tolist()
        else:
            num_microbatches = [
                self.args.global_batch_size // (self.args.micro_batch_size * self.dp_size)
            ] * (len(tokens) // local_batch_size)

        start = 0
        for mbs_size in num_microbatches:
            end = start + local_batch_size
            # Create dummy advantages/returns for SFT (not used but required by pack_sequences)
            dummy_advantages = [
                torch.zeros(rollout_data["response_lengths"][i])
                for i in range(start, end)
            ]
            packed_batches.extend(
                pack_sequences(
                    rollout_data["tokens"][start:end],
                    rollout_data["loss_masks"][start:end],
                    rollout_data["rewards"][start:end],
                    rollout_data["raw_reward"][start:end],
                    rollout_data["response_lengths"][start:end],
                    dummy_advantages,  # advantages
                    dummy_advantages,  # returns
                    num_packs=mbs_size,
                )
            )
            start = end

        grad_accum = list(accumulate(num_microbatches))
        return packed_batches, grad_accum

    def _get_model_inputs_args(self, packed_sequence: dict) -> dict:
        """Prepare model input arguments from packed sequence."""
        input_ids = packed_sequence["tokens"].unsqueeze(0)
        position_ids = packed_sequence["position_ids"].unsqueeze(0)

        if self.cp_size > 1:
            packed_sequence = pad_packed_sequence_with_cp(packed_sequence, self.cp_size)

            if not packed_sequence["cu_seqlens"].is_cuda:
                packed_sequence["cu_seqlens"] = packed_sequence["cu_seqlens"].cuda()
            cu_seqlens = packed_sequence["cu_seqlens"]
            update_ring_flash_attn_params(cu_seqlens, self.cp_group)

            input_ids = torch.chunk(
                packed_sequence["tokens"].unsqueeze(0), self.cp_size, dim=1
            )[self.cp_rank]
            position_ids = torch.chunk(
                packed_sequence["position_ids"].unsqueeze(0), self.cp_size, dim=1
            )[self.cp_rank]

        model_args = {
            "input_ids": input_ids,
            "position_ids": position_ids,
            "attention_mask": None,
        }

        if packed_sequence.get("multimodal_inputs"):
            model_args.update(packed_sequence["multimodal_inputs"])

        return model_args

    def _compute_sft_loss(self, unpacked_batches: list[dict], logits: torch.Tensor):
        """Compute SFT loss (negative log likelihood)."""
        loss_masks = [
            batch["loss_masks"].to(device=logits.device) for batch in unpacked_batches
        ]
        response_lengths = [batch["response_lengths"] for batch in unpacked_batches]
        log_probs = torch.cat(
            [batch["cur_log_probs"] for batch in unpacked_batches], dim=0
        )
        loss = -sum_of_sample_mean(log_probs, response_lengths, loss_masks)

        if log_probs.numel() == 0:
            loss += 0 * logits.sum()

        return loss, {"loss": loss.detach()}

    def _train_step(
        self,
        packed_batch: dict,
        reported_accum: dict,
        mbs_id: int,
        grad_accum: list[int],
    ):
        """Execute one training step."""
        # Prepare model inputs
        model_args = self._get_model_inputs_args(packed_batch)
        logits = self.model(**model_args).logits.squeeze(0).float()

        # Compute log probs and entropy (unified for both CP and non-CP modes)
        log_probs, entropy_result = get_logprob_and_entropy_with_cp(
            logits=logits,
            target_tokens=packed_batch["tokens"],
            cp_rank=self.cp_rank,
            cp_size=self.cp_size,
            cp_group=self.cp_group,
            model_input_ids=model_args["input_ids"],
            allow_compile=not self.args.true_on_policy_mode,
            temperature=self.args.rollout_temperature,
        )
        packed_batch["cur_log_probs"] = log_probs
        packed_batch["entropy"] = entropy_result

        unpacked_batches = unpack_sequences(packed_batch)
        loss, reported = self._compute_sft_loss(unpacked_batches, logits)

        # Scale loss for gradient accumulation
        loss = loss * self.dp_size / self.args.global_batch_size
        loss.backward()

        # Accumulate reported metrics (store tensors for later mean)
        for k, v in reported.items():
            reported_accum.setdefault(k, []).append(v)

        if (mbs_id + 1) in grad_accum:
            # TODO: check if the grad norm is global grad norm.
            grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip_grad)
            # the grad norm used to be of DTensor
            grad_norm = float(grad_norm)

            self.optimizer.step()
            # Update learning rate
            self.lr_scheduler.step()
            self.optimizer.zero_grad(set_to_none=True)
            # Aggregate logs
            aggregated = {k: torch.stack(v).sum().item() for k, v in reported_accum.items()}
            # TODO: change this, this is slow.
            reduced_aggregated = [None] * self.dp_size
            dist.all_gather_object(reduced_aggregated, aggregated, group=self.dp_group)
            aggregated = {}
            for k in reported_accum.keys():
                aggregated[k] = sum([r[k] for r in reduced_aggregated]) / (self.args.global_batch_size)
            reported_accum.clear()
            if dist.get_rank() == 0:
                log_dict = {
                    f"train/{k}": (val.item() if torch.is_tensor(val) else val) for k, val in aggregated.items()
                }
                log_dict["train/grad_norm"] = grad_norm

                # Log learning rate per parameter group; use scheduler's last computed LRs
                lr_values = self.lr_scheduler.get_last_lr()
                for gid, _group in enumerate(self.optimizer.param_groups):
                    log_dict[f"train/lr_{gid}"] = lr_values[gid]

                logger.info(f"step {self.global_step}: {log_dict}")
                log_dict["train/step"] = self.global_step
                tracking_utils.log(self.args, log_dict, step_key="train/step")
            self.global_step += 1

    def train_one_rollout(self, rollout_id: int):
        """Execute one rollout's worth of training."""
        samples = self.generate_sft_rollout(rollout_id)

        train_data = self._convert_samples_to_train_data(samples)

        rollout_data = self._split_train_data_by_dp(train_data)

        packed_batches, grad_accum = self._packed_data(rollout_data)

        if len(grad_accum) == 0:
            logger.warning(f"[Rank {dist.get_rank()}] No batches to train on rollout {rollout_id}")
            return

        with timer("actor_train"):
            reported_accum = {}
            self.optimizer.zero_grad(set_to_none=True)

            for mbs_id, packed_batch in enumerate(
                tqdm(packed_batches, desc="actor_train", disable=dist.get_rank() != 0)
            ):
                self._train_step(packed_batch, reported_accum, mbs_id, grad_accum)

        self.prof.step(rollout_id=rollout_id)

    def save_model(self, iteration: int):
        """Save model checkpoint."""
        if self.args.save is None:
            return
        checkpoint.save(self, iteration)