[tab] train_sft.py (python)
#!/usr/bin/env python3
"""
Ray-free SFT Training Script for Miles.

This script provides a simplified training path for Supervised Fine-Tuning (SFT)
that bypasses Ray entirely and uses torchrun for distributed training.

Usage:
    torchrun --nproc_per_node=2 train_sft.py \
        --hf-checkpoint /path/to/model \
        --prompt-data /path/to/data.parquet \
        --input-key messages \
        --apply-chat-template \
        ...

This is equivalent to the Ray-based SFT with --debug-train-only, but without
the Ray overhead.
"""

import copy
import logging
import os
from argparse import Namespace
from datetime import timedelta
from itertools import accumulate
from pathlib import Path

import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from tqdm import tqdm
from transformers import AutoConfig

from miles.backends.fsdp_utils import checkpoint
from miles.backends.fsdp_utils.actor import (
    apply_fsdp2,
    get_logprob_and_entropy_with_cp,
    sum_of_sample_mean,
)
from miles.backends.fsdp_utils.data_packing import pack_sequences, unpack_sequences
from miles.backends.fsdp_utils.lr_scheduler import get_lr_scheduler
from miles.rollout.data_source import RolloutDataSource
from miles.utils import tracking_utils
from miles.utils.arguments import parse_args
from miles.utils.data import Dataset, get_minimum_num_micro_batch_size
from miles.utils.distributed_utils import get_gloo_group, init_gloo_group
from miles.utils.logging_utils import configure_logger
from miles.utils.mask_utils import MultiTurnLossMaskGenerator
from miles.utils.metric_utils import compute_rollout_step
from miles.utils.misc import load_function, should_run_periodic_action
from miles.utils.processing_utils import load_processor, load_tokenizer
from miles.utils.profile_utils import TrainProfiler
from miles.utils.timer import timer
from miles.utils.tracking_utils import init_tracking
from miles.utils.types import Sample

logger = logging.getLogger(__name__)


class SFTTrainer:
    """
    A simplified trainer for SFT that runs without Ray.

    This class combines the functionality of:
    - FSDPTrainRayActor (model initialization, FSDP wrapping, training)
    - RolloutManager (data loading via generate_rollout)
    - The main training loop from train.py
    """

    def __init__(self, args: Namespace):
        self.args = args
        self.device = torch.device("cuda")

        # Setup distributed
        self._init_distributed()

        # Setup device mesh for FSDP
        self._setup_device_mesh()

        torch.manual_seed(args.seed)

        # Initialize tracking on rank 0
        if dist.get_rank() == 0:
            init_tracking(args, primary=True)

        # Load tokenizer and config
        self._load_tokenizer_and_config()

        # Initialize data source
        self._init_data_source()

        # Initialize model
        self._init_model()

        # Initialize optimizer and scheduler
        self._init_optimizer()

        # Load checkpoint if available
        self._load_checkpoint()

        # Initialize profiler
        self.prof = TrainProfiler(args)
        self.prof.on_init_end()

        logger.info(f"[Rank {dist.get_rank()}] SFTTrainer initialized successfully")

    def _init_distributed(self):
        """Initialize distributed training."""
        # torchrun sets these environment variables
        local_rank = int(os.environ.get("LOCAL_RANK", 0))
        torch.cuda.set_device(f"cuda:{local_rank}")

        backend = self.args.distributed_backend
        dist.init_process_group(
            backend=backend,
            timeout=timedelta(minutes=self.args.distributed_timeout_minutes),
        )
        init_gloo_group()

        self.args.rank = dist.get_rank()
        self.args.world_size = dist.get_world_size()

        logger.info(
            f"[Rank {self.args.rank}] Distributed initialized: "
            f"world_size={self.args.world_size}, local_rank={local_rank}"
        )

    def _setup_device_mesh(self):
        """Setup device mesh for FSDP (pure DP mode for SFT)."""
        world_size = dist.get_world_size()
        rank = dist.get_rank()

        # For SFT, we use pure DP (cp_size=1)
        self.cp_size = getattr(self.args, "context_parallel_size", 1)
        self.dp_size = world_size // self.cp_size

        self.mesh = init_device_mesh(
            "cuda",
            mesh_shape=(self.dp_size, self.cp_size),
            mesh_dim_names=("dp", "cp"),
        )

        self.dp_group = self.mesh.get_group("dp")
        self.cp_group = self.mesh.get_group("cp")
        self.dp_mesh = self.mesh["dp"]

        self.dp_rank = rank // self.cp_size
        self.cp_rank = rank % self.cp_size

        logger.info(
            f"[Rank {rank}] Device mesh: dp_size={self.dp_size}, cp_size={self.cp_size}, "
            f"dp_rank={self.dp_rank}, cp_rank={self.cp_rank}"
        )

    def _load_tokenizer_and_config(self):
        """Load tokenizer and model config sequentially to avoid race conditions."""
        for i in range(dist.get_world_size()):
            if i == dist.get_rank():
                self.hf_config = AutoConfig.from_pretrained(
                    self.args.hf_checkpoint, trust_remote_code=True
                )
                self.tokenizer = load_tokenizer(
                    self.args.hf_checkpoint, trust_remote_code=True
                )
                self.processor = None
                if self.args.multimodal_keys:
                    self.processor = load_processor(
                        self.args.hf_checkpoint, trust_remote_code=True
                    )
            dist.barrier(group=get_gloo_group())

        # Initialize loss mask generator for SFT
        self.mask_generator = MultiTurnLossMaskGenerator(
            self.tokenizer,
            tokenizer_type=getattr(self.args, "loss_mask_type", None),
        )

    def _init_data_source(self):
        """Initialize the data source for SFT training."""
        self.data_source = RolloutDataSource(self.args)

        # Calculate num_rollout from dataset size
        if self.args.num_rollout is None:
            num_rollout_per_epoch = len(self.data_source.dataset) // self.args.rollout_batch_size
            self.args.num_rollout = num_rollout_per_epoch * self.args.num_epoch
            self.num_rollout_per_epoch = num_rollout_per_epoch
        else:
            self.num_rollout_per_epoch = None

        if getattr(self.args, "start_rollout_id", None) is None:
            self.args.start_rollout_id = 0

        logger.info(
            f"[Rank {dist.get_rank()}] Data source initialized: "
            f"dataset_size={len(self.data_source.dataset)}, "
            f"num_rollout={self.args.num_rollout}"
        )

    def _get_init_weight_context_manager(self):
        """Get context manager for model initialization."""
        from accelerate import init_empty_weights

        use_meta_tensor = not self.hf_config.tie_word_embeddings

        def cpu_init_weights():
            return torch.device("cpu")

        if use_meta_tensor:
            return init_empty_weights if dist.get_rank() != 0 else cpu_init_weights
        else:
            return cpu_init_weights

    def _fsdp2_load_full_state_dict(self, model, full_state, device_mesh, cpu_offload):
        """Load full state dict into FSDP2 model with broadcast from rank 0."""
        from torch.distributed.checkpoint.state_dict import (
            StateDictOptions,
            set_model_state_dict,
        )

        if dist.get_rank() == 0:
            model = model.to(device=torch.cuda.current_device(), non_blocking=True)
        else:
            model = model.to_empty(device=torch.cuda.current_device())

        is_cpu_offload = cpu_offload is not None
        options = StateDictOptions(
            full_state_dict=True, cpu_offload=is_cpu_offload, broadcast_from_rank0=True
        )

        set_model_state_dict(model, full_state, options=options)

        for _name, buf in model.named_buffers():
            dist.broadcast(buf, src=0)

        if is_cpu_offload:
            model.to("cpu", non_blocking=True)
            for buf in model.buffers():
                buf.data = buf.data.to(torch.cuda.current_device())

        return model

    def _get_model_cls(self):
        """Get the appropriate model class based on config."""
        if hasattr(self.hf_config, "vision_config"):
            from transformers import AutoModelForImageTextToText

            return AutoModelForImageTextToText
        else:
            from transformers import AutoModelForCausalLM

            return AutoModelForCausalLM

    def _init_model(self):
        """Initialize and wrap model with FSDP."""
        self.fsdp_cpu_offload = getattr(self.args, "fsdp_cpu_offload", False)

        init_context = self._get_init_weight_context_manager()

        with init_context():
            model = self._get_model_cls().from_pretrained(
                self.args.hf_checkpoint,
                trust_remote_code=True,
                attn_implementation=self.args.attn_implementation,
            )

        model.train()
        full_state = model.state_dict()

        model = apply_fsdp2(
            model, mesh=self.dp_mesh, cpu_offload=self.fsdp_cpu_offload, args=self.args
        )

        model = self._fsdp2_load_full_state_dict(
            model,
            full_state,
            self.dp_mesh,
            cpu_offload=True if self.fsdp_cpu_offload else None,
        )

        self.model = model

        if self.args.gradient_checkpointing:
            self.model.gradient_checkpointing_enable()

        logger.info(f"[Rank {dist.get_rank()}] Model initialized with FSDP")

    def _init_optimizer(self):
        """Initialize optimizer and learning rate scheduler."""
        if self.args.optimizer == "adam":
            self.optimizer = torch.optim.AdamW(
                self.model.parameters(),
                lr=self.args.lr,
                betas=(self.args.adam_beta1, self.args.adam_beta2),
                eps=self.args.adam_eps,
                weight_decay=self.args.weight_decay,
            )
        else:
            raise ValueError(f"Unsupported optimizer: {self.args.optimizer}")

        self.lr_scheduler = get_lr_scheduler(self.args, self.optimizer)
        self.global_step = 0
        self.micro_step = 0

    def _load_checkpoint(self):
        """Load checkpoint if available."""
        checkpoint_payload = checkpoint.load(self)
        checkpoint.finalize_load(self, checkpoint_payload)

    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)

        return {
            "input_ids": input_ids,
            "position_ids": position_ids,
            "attention_mask": None,
        }

    def _compute_log_prob(self, packed_batches: list[dict]) -> None:
        """Compute log probabilities for SFT loss."""
        with timer("log_probs"), torch.no_grad():
            for batch in tqdm(
                packed_batches,
                desc="log_probs",
                disable=dist.get_rank() != 0,
            ):
                model_args = self._get_model_inputs_args(batch)
                logits = self.model(**model_args).logits.squeeze(0).float()
                log_probs_result, entropy_result = get_logprob_and_entropy_with_cp(
                    logits=logits,
                    target_tokens=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=True,
                    temperature=self.args.rollout_temperature,
                )
                batch["log_probs"] = log_probs_result
                batch["entropy"] = entropy_result

    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."""
        model_args = self._get_model_inputs_args(packed_batch)
        logits = self.model(**model_args).logits.squeeze(0).float()

        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=True,
            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()

        for k, v in reported.items():
            reported_accum.setdefault(k, []).append(v)

        if (mbs_id + 1) in grad_accum:
            grad_norm = torch.nn.utils.clip_grad_norm_(
                self.model.parameters(), self.args.clip_grad
            )
            grad_norm = float(grad_norm)

            self.optimizer.step()
            self.lr_scheduler.step()
            self.optimizer.zero_grad(set_to_none=True)

            # Aggregate and log metrics
            aggregated = {k: torch.stack(v).sum().item() for k, v in reported_accum.items()}
            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

                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."""
        # Generate SFT data
        samples = self.generate_sft_rollout(rollout_id)

        # Convert to training format
        train_data = self._convert_samples_to_train_data(samples)

        # Split by DP rank
        rollout_data = self._split_train_data_by_dp(train_data)

        # Create advantages/returns (dummy for SFT)
        rollout_data["advantages"] = rollout_data["returns"] = [
            torch.tensor([0.0] * rollout_data["response_lengths"][i])
            for i in range(len(rollout_data["rewards"]))
        ]

        # Pack sequences
        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

        # Compute log probs for logging
        self._compute_log_prob(packed_batches)

        # Training loop
        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)

    def train(self):
        """Main training loop."""
        logger.info(
            f"[Rank {dist.get_rank()}] Starting training: "
            f"rollout_id {self.args.start_rollout_id} -> {self.args.num_rollout}"
        )

        for rollout_id in range(self.args.start_rollout_id, self.args.num_rollout):
            self.train_one_rollout(rollout_id)

            # Save checkpoint periodically
            if should_run_periodic_action(
                rollout_id, self.args.save_interval, self.num_rollout_per_epoch
            ):
                self.save_model(rollout_id)

        logger.info(f"[Rank {dist.get_rank()}] Training completed!")


def set_sft_defaults(args: Namespace) -> Namespace:
    """Set default values appropriate for SFT training."""
    # These flags are normally set by debug_train_only in Ray mode
    if not hasattr(args, "debug_train_only"):
        args.debug_train_only = True

    # SFT-specific defaults
    if not hasattr(args, "loss_type") or args.loss_type is None:
        args.loss_type = "sft_loss"

    if not hasattr(args, "advantage_estimator") or args.advantage_estimator is None:
        args.advantage_estimator = "grpo"

    if not hasattr(args, "n_samples_per_prompt") or args.n_samples_per_prompt is None:
        args.n_samples_per_prompt = 1

    if not hasattr(args, "rollout_global_dataset"):
        args.rollout_global_dataset = True

    if not hasattr(args, "context_parallel_size") or args.context_parallel_size is None:
        args.context_parallel_size = 1

    if not hasattr(args, "attn_implementation") or args.attn_implementation is None:
        args.attn_implementation = "flash_attention_2"

    if not hasattr(args, "rollout_temperature") or args.rollout_temperature is None:
        args.rollout_temperature = 1.0

    # Ensure we don't try to use rollout engines
    args.offload_train = False
    args.offload_rollout = False
    args.colocate = False

    return args


def main():
    configure_logger()

    # Parse arguments
    args = parse_args()

    # Apply SFT-specific defaults
    args = set_sft_defaults(args)

    # Create trainer and run
    trainer = SFTTrainer(args)
    trainer.train()


if __name__ == "__main__":
    main()