Upload folder using huggingface_hub

slime/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+

slime/backends/__init__.py

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+

slime/backends/fsdp_utils/__init__.py

@@ -0,0 +1,36 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+
+try:
+    _TORCH_MEMORY_SAVER_AVAILABLE = True
+except ImportError:
+    logging.warning("torch_memory_saver is not installed, refer to : https://github.com/fzyzcjy/torch_memory_saver")
+    _TORCH_MEMORY_SAVER_AVAILABLE = False
+
+try:
+    _FSDP_AVAILABLE = True
+except ImportError as e:
+    logging.warning(f"FSDP backend dependencies not available: {e}")
+    _FSDP_AVAILABLE = False
+
+if _FSDP_AVAILABLE:
+    from .actor import FSDPTrainRayActor
+    from .arguments import load_fsdp_args
+else:
+
+    def _raise_import_error(*args, **kwargs):
+        raise ImportError(
+            "FSDP backend is not available. "
+            "Please ensure PyTorch with FSDP2 support is installed. "
+            "For installation instructions, refer to: https://pytorch.org/docs/stable/distributed.fsdp.fully_shard.html"
+        )
+
+    FSDPTrainRayActor = _raise_import_error
+    load_fsdp_args = _raise_import_error
+
+__all__ = ["load_fsdp_args", "FSDPTrainRayActor"]
+
+logging.getLogger().setLevel(logging.WARNING)

slime/backends/fsdp_utils/actor.py

@@ -0,0 +1,1145 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import os
+import random
+from argparse import Namespace
+from itertools import accumulate
+
+import ray
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from ring_flash_attn import substitute_hf_flash_attn, update_ring_flash_attn_params
+from tqdm import tqdm
+from transformers import AutoConfig
+
+from slime.ray.train_actor import TrainRayActor
+from slime.utils import train_dump_utils, train_metric_utils
+from slime.utils.context_utils import with_defer
+from slime.utils.data import get_minimum_num_micro_batch_size, process_rollout_data
+from slime.utils.distributed_utils import get_gloo_group
+from slime.utils.memory_utils import clear_memory, print_memory
+from slime.utils.metric_utils import compute_rollout_step
+from slime.utils.misc import load_function
+from slime.utils.ppo_utils import (
+    compute_approx_kl,
+    compute_gspo_kl,
+    compute_opsm_mask,
+    compute_policy_loss,
+    vanilla_tis_function,
+)
+from slime.utils.processing_utils import load_processor, load_tokenizer
+from slime.utils.ray_utils import Box
+from slime.utils.timer import Timer, inverse_timer, timer
+from slime.utils.tracking_utils import init_tracking
+
+from ...utils import tracking_utils
+from ...utils.profile_utils import TrainProfiler
+from . import checkpoint
+from .data_packing import pack_sequences, pad_packed_sequence_with_cp, unpack_sequences
+from .lr_scheduler import get_lr_scheduler
+from .update_weight_utils import UpdateWeightFromDistributed, UpdateWeightFromTensor
+
+logger = logging.getLogger(__name__)
+
+
+class FSDPTrainRayActor(TrainRayActor):
+    """Simplified TrainRayActor for pure HF+FSDP training.
+
+    Responsibilities:
+      * Initialize model/tokenizer on rank0 sequentially to avoid race on cache
+      * Wrap model with FSDP
+      * Provide minimal train / save / update_weights hooks compatible with existing RayTrainGroup
+
+    Weight update strategy:
+      * Rank0 gathers state_dict (full) and broadcasts tensor-by-tensor.
+      * For small models this is fine; for larger models consider sharded state_dict type.
+    """
+
+    @with_defer(lambda: Timer().start("train_wait"))
+    def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # type: ignore[override]
+        super().init(args, role, with_ref)
+
+        # Setup device mesh for parallelism (handles both CP and non-CP cases)
+        self._setup_device_mesh()
+        torch.manual_seed(args.seed)
+
+        self.train_parallel_config = {
+            "dp_size": self.dp_size,
+        }
+
+        if self.args.debug_rollout_only:
+            return 0
+
+        self.fsdp_cpu_offload = getattr(self.args, "fsdp_cpu_offload", False)
+        # Offload train and fsdp cpu offload cannot be used together, fsdp_cpu_offload is more aggressive
+        if self.args.offload_train and self.fsdp_cpu_offload:
+            self.args.offload_train = False
+
+        self._enable_true_on_policy_optimizations(args)
+        if dist.get_rank() == 0:
+            init_tracking(args, primary=False)
+
+        if getattr(self.args, "start_rollout_id", None) is None:
+            self.args.start_rollout_id = 0
+
+        self.prof = TrainProfiler(args)
+
+        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)
+                # Vision models have `vision_config` in the config
+                if hasattr(self.hf_config, "vision_config"):
+                    self.processor = load_processor(self.args.hf_checkpoint, trust_remote_code=True)
+            dist.barrier(group=get_gloo_group())
+
+        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 args.gradient_checkpointing:
+            self.model.gradient_checkpointing_enable()
+
+        if args.optimizer == "adam":
+            self.optimizer = torch.optim.AdamW(
+                self.model.parameters(),
+                lr=args.lr,
+                betas=(args.adam_beta1, args.adam_beta2),
+                eps=args.adam_eps,
+                weight_decay=args.weight_decay,
+            )
+        else:
+            raise ValueError(f"Unsupported optimizer: {args.optimizer}. Supported options: 'adam'")
+
+        # Initialize LR scheduler
+        self.lr_scheduler = get_lr_scheduler(args, self.optimizer)
+
+        self.global_step = 0
+        self.micro_step = 0
+
+        checkpoint_payload = checkpoint.load(self)
+
+        # Create separate ref model if needed (kept in CPU until needed)
+        self.ref_model = None
+        if with_ref:
+            self.ref_model = self._create_ref_model(args.ref_load)
+
+        self.weight_updater = (
+            UpdateWeightFromTensor(self.args, self.model)
+            if self.args.colocate
+            else UpdateWeightFromDistributed(self.args, self.model)
+        )
+
+        checkpoint.finalize_load(self, checkpoint_payload)
+
+        # Initialize data packing parameters
+        self.max_tokens_per_gpu = args.max_tokens_per_gpu  # From main arguments
+
+        if self.args.offload_train:
+            self.sleep()
+
+        self.prof.on_init_end()
+
+        return int(getattr(self.args, "start_rollout_id", 0))
+
+    def get_model_cls(self):
+        # Vision models have `vision_config` in the config
+        if hasattr(self.hf_config, "vision_config"):
+            from transformers import AutoModelForImageTextToText
+
+            return AutoModelForImageTextToText
+        else:
+            from transformers import AutoModelForCausalLM
+
+            return AutoModelForCausalLM
+
+    def _enable_true_on_policy_optimizations(self, args):
+        if args.true_on_policy_mode:
+            from sglang.srt.batch_invariant_ops import enable_batch_invariant_mode
+
+            from .models.qwen3_moe import apply_true_on_policy_patch_for_qwen3_moe
+
+            logger.info("FSDPTrainRayActor call enable_batch_invariant_mode for true-on-policy")
+            enable_batch_invariant_mode(
+                # In Qwen3, rope `inv_freq_expanded.float() @ position_ids_expanded.float()` uses bmm
+                # and disabling it will make it aligned
+                enable_bmm=False,
+            )
+
+            apply_true_on_policy_patch_for_qwen3_moe()
+        else:
+            from .models.qwen3_moe_hf import apply_fsdp_moe_patch
+
+            apply_fsdp_moe_patch()
+
+    def _setup_device_mesh(self) -> None:
+        """Setup device mesh for parallelism (always called, handles both CP and non-CP cases).
+
+        Creates 2D mesh (dp_size, cp_size) for all cases:
+        - When context_parallel_size > 1: hybrid CP + DP
+        - When context_parallel_size = 1: pure DP (equivalent to 1D mesh)
+
+        This ensures consistent group management across all parallelism modes.
+        """
+        from torch.distributed.device_mesh import init_device_mesh
+
+        world_size = dist.get_world_size()
+        rank = dist.get_rank()
+
+        # Use context_parallel_size directly (defaults to 1 for pure DP)
+        self.cp_size = self.args.context_parallel_size
+        self.dp_size = world_size // self.cp_size
+
+        # Create 2D device mesh: (dp_size, cp_size)
+        # Ranks laid out in row-major: mesh[dp_idx, cp_idx] = dp_idx * cp_size + cp_idx
+        # - CP groups: consecutive ranks along dim 1, e.g., [0,1], [2,3], [4,5], [6,7]
+        # - DP groups: striped ranks along dim 0, e.g., [0,2,4,6], [1,3,5,7]
+        # When cp_size=1, this degenerates to pure DP
+        self.mesh = init_device_mesh("cuda", mesh_shape=(self.dp_size, self.cp_size), mesh_dim_names=("dp", "cp"))
+
+        # Extract process groups from mesh
+        self.dp_group = self.mesh.get_group("dp")  # For FSDP gradient sync, metric reduction
+        self.cp_group = self.mesh.get_group("cp")  # For Ring Flash Attention, logit gathering
+        self.dp_mesh = self.mesh["dp"]  # For FSDP
+
+        # Compute local ranks within each dimension
+        self.dp_rank = rank // self.cp_size
+        self.cp_rank = rank % self.cp_size
+
+        logger.info(
+            f"[Rank {rank}] Device mesh (2D): world_size={world_size}, "
+            f"cp_size={self.cp_size}, dp_size={self.dp_size}"
+        )
+        logger.info(f"[Rank {rank}] Mesh shape: {self.mesh.shape}, " f"dp_rank={self.dp_rank}, cp_rank={self.cp_rank}")
+
+        # Setup Ring Flash Attention with CP group from mesh (only when cp_size > 1)
+        if self.cp_size > 1:
+            substitute_hf_flash_attn(self.cp_group, heads_k_stride=1)
+            logger.info(f"[Rank {rank}] CP initialized via device mesh")
+        else:
+            logger.info(f"[Rank {rank}] Pure DP mode (cp_size=1)")
+
+    def _get_init_weight_context_manager(self):
+        """Get context manager for model initialization.
+
+        Returns a callable that creates a context manager.
+        Uses meta device (no memory allocation) for non-rank-0 processes,
+        UNLESS tie_word_embeddings=True (which causes hangs with meta tensors).
+
+        Ref: verl/utils/fsdp_utils.py::get_init_weight_context_manager
+        NOTE: tie_word_embedding causes meta_tensor init to hang
+        """
+        from accelerate import init_empty_weights
+
+        # Check if model uses tied word embeddings (which doesn't work with meta tensors)
+        use_meta_tensor = not self.hf_config.tie_word_embeddings
+
+        def cpu_init_weights():
+            return torch.device("cpu")
+
+        if use_meta_tensor:
+            # Rank 0: CPU, others: meta device (memory efficient for large models)
+            return init_empty_weights if dist.get_rank() != 0 else cpu_init_weights
+        else:
+            logger.info(f"[Rank {dist.get_rank()}] tie_word_embeddings=True, loading full model to CPU on all ranks")
+            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 efficient broadcast from rank 0.
+
+        This function loads weights from rank 0 and broadcasts to all other ranks,
+        avoiding the need for each rank to load the full model from disk.
+
+        Args:
+            model: FSDP2-wrapped model
+            full_state: State dict (only rank 0 has real weights, others have empty dict)
+            device_mesh: Device mesh for FSDP
+            cpu_offload: If not None, enables StateDictOptions cpu_offload
+
+        Ref:verl/utils/fsdp_utils.py::fsdp2_load_full_state_dict
+        """
+        from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict
+
+        # Rank 0: move with weights, others: allocate empty tensors on device
+        if dist.get_rank() == 0:
+            model = model.to(device=torch.cuda.current_device(), non_blocking=True)
+        else:
+            # to_empty creates tensors on device without initializing memory
+            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)
+
+        # set_model_state_dict will not broadcast buffers, so we need to broadcast them manually.
+        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
+
+    @timer
+    def sleep(self) -> None:
+        """Pause CUDA memory for all tracked tensors."""
+        if not self.args.offload_train:
+            return
+
+        print_memory("before offload model")
+
+        self.model.cpu()
+        move_torch_optimizer(self.optimizer, "cpu")
+        clear_memory()
+        dist.barrier(group=get_gloo_group())
+        print_memory("after offload model")
+
+    @timer
+    def wake_up(self) -> None:
+        """Resume CUDA memory for all tracked tensors."""
+        if not self.args.offload_train:
+            return
+
+        self.model.cuda()
+        move_torch_optimizer(self.optimizer, "cuda")
+        dist.barrier(group=get_gloo_group())
+        print_memory("after wake_up model")
+
+    def save_model(self, rollout_id: int, force_sync: bool = False) -> None:
+        """Delegate checkpoint saving to the shared checkpoint utilities."""
+        if self.args.debug_rollout_only or self.args.save is None:
+            return
+
+        assert not self.args.async_save, "FSDPTrainRayActor does not support async_save yet."
+        checkpoint.save(self, rollout_id)
+
+    def _compute_log_prob(
+        self,
+        model_tag: str,
+        packed_batches: list[dict[str, torch.Tensor]],
+        store_prefix: str = "",
+    ) -> dict[str, list[torch.Tensor]]:
+        """Compute token log-probabilities for a list of packed batches.
+
+        Parameters:
+            model_tag: Which parameters to use, e.g. "actor" or "ref".
+            packed_batches: A list of packed batch dictionaries produced by
+                `pack_sequences`, each containing at least `tokens` and
+                `position_ids`; may also include multimodal keys like `pixel_values`.
+            store_prefix: Prefix to use for keys in outputs (e.g., "ref_").
+
+        Returns:
+            A lightweight dictionary keyed by f"{store_prefix}log_probs". The
+            actual per-sequence results are written in-place into each element of
+            `packed_batches` under the same key and can be read back by callers.
+
+        Note:
+            Uses separate ref model when model_tag == "ref". The ref model is
+            loaded from CPU to GPU on-demand and offloaded back after use.
+        """
+        # Select which model to use
+        if model_tag == "ref" and self.ref_model is not None:
+            if not self.fsdp_cpu_offload:
+                self.model.cpu()
+                torch.cuda.empty_cache()
+                dist.barrier(group=get_gloo_group())
+
+            active_model = self.ref_model
+            active_model.eval()
+        else:
+            active_model = self.model
+
+        try:
+            rollout_data = {f"{store_prefix}log_probs": []}
+            with timer(f"{store_prefix}log_probs"), torch.no_grad():
+                for batch in self.prof.iterate_train_log_probs(
+                    tqdm(packed_batches, desc=f"{store_prefix}log_probs", disable=dist.get_rank() != 0)
+                ):
+                    model_args = self._get_model_inputs_args(batch)
+                    logits = active_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=not self.args.true_on_policy_mode,
+                        temperature=self.args.rollout_temperature,
+                    )
+                    batch[f"{store_prefix}log_probs"] = log_probs_result
+                    if store_prefix == "":
+                        batch["entropy"] = entropy_result
+            return rollout_data
+
+        finally:
+            # Restore actor model if it was offloaded
+            if model_tag == "ref" and self.ref_model is not None:
+                torch.cuda.empty_cache()
+                dist.barrier(group=get_gloo_group())
+
+                if not self.fsdp_cpu_offload:
+                    self.model.cuda()
+                    dist.barrier(group=get_gloo_group())
+
+    def _packed_data(
+        self, rollout_data: dict[str, list[torch.Tensor]]
+    ) -> tuple[list[dict[str, torch.Tensor]], list[int]]:
+        """Pack variable-length sequences for efficient processing.
+
+        Parameters:
+            rollout_data: Dictionary of lists containing sequence-level tensors
+                such as `tokens`, `loss_masks`, `rewards`, `response_lengths`,
+                `advantages`, `returns`, and optional `rollout_log_probs`.
+
+        Returns:
+            A pair `(packed_batches, grad_accum)` where `packed_batches` is a list
+            of packed batch dictionaries and `grad_accum` lists the micro-batch
+            indices at which to perform optimizer steps.
+        """
+        # Pack sequences efficiently
+        tokens = rollout_data["tokens"]
+
+        packed_batches = []
+        mbs_size_list = []
+        local_batch_size = self.args.global_batch_size // self.dp_size
+        assert (
+            self.args.global_batch_size % self.dp_size == 0
+        ), f"global_batch_size {self.args.global_batch_size} is not divisible by dp_world_size {self.dp_size}"
+        # Use global_batch_size for splitting when max_tokens_per_gpu is enabled
+        if self.args.use_dynamic_batch_size:
+            # In CP mode, CP group shares sequences, so total capacity is max_tokens_per_gpu * cp_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
+            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],
+                    rollout_data["advantages"][start:end],
+                    rollout_data["returns"][start:end],
+                    rollout_log_probs=(
+                        rollout_data["rollout_log_probs"][start:end] if "rollout_log_probs" in rollout_data else None
+                    ),
+                    multimodal_train_inputs=(
+                        rollout_data["multimodal_train_inputs"][start:end]
+                        if "multimodal_train_inputs" in rollout_data
+                        else None
+                    ),
+                    num_packs=mbs_size,
+                )
+            )
+            start = end
+        grad_accum = list(accumulate(num_microbatches))
+
+        return packed_batches, grad_accum
+
+    def train(self, rollout_id: int, rollout_data_ref: Box) -> None:
+        """Run one training update over a rollout batch.
+
+        Parameters:
+            rollout_id: Monotonic id for logging.
+            rollout_data_ref: A Box handle wrapping a Ray object reference to a
+                dictionary with rollout tensors and metadata (e.g., `tokens`,
+                `loss_masks`, `rewards`, `response_lengths`, optional
+                `rollout_log_probs`, etc.). It will be fetched and partitioned
+                by `process_rollout_data` based on data-parallel rank/size.
+        """
+        if self.args.offload_train:
+            self.wake_up()
+
+        with inverse_timer("train_wait"), timer("train"):
+            rollout_data = process_rollout_data(self.args, rollout_data_ref, self.dp_rank, self.dp_size)
+            if self.args.debug_rollout_only:
+                return
+            self._train_core(rollout_id=rollout_id, rollout_data=rollout_data)
+
+        train_metric_utils.log_perf_data_raw(
+            rollout_id=rollout_id,
+            args=self.args,
+            is_primary_rank=dist.get_rank() == 0,
+            compute_total_fwd_flops=None,
+        )
+
+    def _log_rollout_data(self, rollout_id: int, rollout_data, packed_batches):
+        log_dict = {}
+        if "raw_reward" in rollout_data and dist.get_rank() == 0:
+            raw_reward_list = rollout_data["raw_reward"]
+            if raw_reward_list:
+                log_dict["rollout/raw_reward"] = sum(raw_reward_list) / len(raw_reward_list)
+
+        for metric_key in ["log_probs", "rollout_log_probs", "ref_log_probs", "advantages", "returns"]:
+            if metric_key not in packed_batches[0]:
+                continue
+            val = torch.tensor([0.0], device=torch.cuda.current_device())
+            for _mbs_id, batches in enumerate(packed_batches):
+                unpacked_batches = unpack_sequences(batches)
+                for unpacked_batch in unpacked_batches:
+                    if isinstance(unpacked_batch[metric_key], torch.Tensor):
+                        loss_masks_tensor = unpacked_batch["loss_masks"].to(device=torch.cuda.current_device())
+                        metric_tensor = unpacked_batch[metric_key].to(device=torch.cuda.current_device())
+                        val += (metric_tensor * loss_masks_tensor).sum() / loss_masks_tensor.sum().clamp_min(1)
+                    else:
+                        val += unpacked_batch[metric_key]
+            dist.all_reduce(val, op=dist.ReduceOp.SUM, group=self.dp_group)
+            log_dict[f"rollout/{metric_key}"] = (
+                val / (self.args.n_samples_per_prompt * self.args.rollout_batch_size)
+            ).item()
+        if dist.get_rank() == 0:
+            logger.info(f"rollout {rollout_id}: {log_dict}")
+            log_dict["rollout/step"] = compute_rollout_step(self.args, rollout_id)
+            tracking_utils.log(self.args, log_dict, step_key="rollout/step")
+
+        if self.args.ci_test and self.args.true_on_policy_mode:
+            assert log_dict["rollout/log_probs"] == log_dict["rollout/rollout_log_probs"], (
+                f"CI check failed: true_on_policy_mode is enabled, but log_probs "
+                f"({log_dict['rollout/log_probs']}) != rollout_log_probs "
+                f"({log_dict['rollout/rollout_log_probs']})"
+            )
+
+    def _train_core(self, rollout_id: int, rollout_data) -> None:
+        if self.args.advantage_estimator in ["grpo", "gspo"]:
+            rollout_data["advantages"] = rollout_data["returns"] = [
+                torch.tensor([rollout_data["rewards"][i]] * rollout_data["response_lengths"][i])
+                for i in range(len(rollout_data["rewards"]))
+            ]
+        else:
+            raise NotImplementedError(f"Unsupported advantage_estimator {self.args.advantage_estimator}")
+
+        packed_batches, grad_accum = self._packed_data(rollout_data)
+
+        assert (
+            len(grad_accum) > 0
+        ), f"Invalid grad_accum {grad_accum} for micro_batch_size {self.args.micro_batch_size} and global_batch_size {self.args.global_batch_size}"
+
+        if self.ref_model is not None:
+            self._compute_log_prob("ref", packed_batches, store_prefix="ref_")
+
+        self._compute_log_prob("actor", packed_batches)
+        self._log_rollout_data(rollout_id, rollout_data, packed_batches)
+
+        with timer("actor_train"):
+            reported_accum: dict[str, list[torch.Tensor]] = {}
+            self.optimizer.zero_grad(set_to_none=True)
+            for mbs_id, packed_batch in self.prof.iterate_train_actor(
+                enumerate(tqdm(packed_batches, desc="actor_train", disable=dist.get_rank() != 0))
+            ):
+                self._train_step(
+                    packed_batch=packed_batch,
+                    reported_accum=reported_accum,
+                    mbs_id=mbs_id,
+                    grad_accum=grad_accum,
+                )
+
+        self.prof.step(rollout_id=rollout_id)
+
+        train_dump_utils.save_debug_train_data(self.args, rollout_id=rollout_id, rollout_data=rollout_data)
+
+        # Update ref model if needed (copy actor weights to ref)
+        if (
+            self.args.ref_update_interval is not None
+            and (rollout_id + 1) % self.args.ref_update_interval == 0
+            and self.ref_model is not None
+        ):
+            if dist.get_rank() == 0:
+                logger.info(f"Updating ref model at rollout_id {rollout_id}")
+            # Copy actor model state to ref model
+            actor_state = self.model.state_dict()
+            self.ref_model.load_state_dict(actor_state)
+            self.ref_model.cpu()
+
+    def _train_step(self, packed_batch, reported_accum, mbs_id, grad_accum):
+        # 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)
+
+        old_log_prob_key = "rollout_log_probs" if self.args.use_rollout_logprobs else "log_probs"
+        missing_old_log_probs = [
+            idx
+            for idx, batch in enumerate(unpacked_batches)
+            if old_log_prob_key not in batch or not isinstance(batch[old_log_prob_key], torch.Tensor)
+        ]
+        if missing_old_log_probs:
+            raise KeyError(
+                f"{old_log_prob_key} must be provided as torch.Tensor for all microbatches when "
+                f"use_rollout_logprobs is set to {self.args.use_rollout_logprobs}. Missing in batches: {missing_old_log_probs}"
+            )
+        old_log_probs = torch.cat([batch[old_log_prob_key] for batch in unpacked_batches], dim=0)
+        log_probs = torch.cat([batch["cur_log_probs"] for batch in unpacked_batches], dim=0)
+        advantages = torch.cat([batch["advantages"] for batch in unpacked_batches], dim=0)
+        loss_masks = [batch["loss_masks"].to(device=log_probs.device) for batch in unpacked_batches]
+        response_lengths = [batch["response_lengths"] for batch in unpacked_batches]
+
+        advantages = advantages.to(device=log_probs.device)
+        old_log_probs = old_log_probs.to(device=log_probs.device)
+        ppo_kl = old_log_probs - log_probs
+
+        if self.args.use_opsm:
+            opsm_mask, opsm_clipfrac = compute_opsm_mask(
+                args=self.args,
+                full_log_probs=[batch["cur_log_probs"] for batch in unpacked_batches],
+                full_old_log_probs=[batch[old_log_prob_key] for batch in unpacked_batches],
+                advantages=[batch["advantages"] for batch in unpacked_batches],
+                loss_masks=loss_masks,
+            )
+
+        if self.args.advantage_estimator == "gspo":
+            ppo_kl = compute_gspo_kl(
+                full_log_probs=[batch["cur_log_probs"] for batch in unpacked_batches],
+                full_old_log_probs=[batch[old_log_prob_key] for batch in unpacked_batches],
+                local_log_probs=[batch["cur_log_probs"] for batch in unpacked_batches],
+                loss_masks=loss_masks,
+            )
+
+        pg_loss, pg_clipfrac = compute_policy_loss(ppo_kl, advantages, self.args.eps_clip, self.args.eps_clip_high)
+
+        if self.args.use_opsm:
+            pg_loss = pg_loss * opsm_mask
+
+        def _has_rollout_log_probs(batch) -> bool:
+            rollout_tensor = batch.get("rollout_log_probs")
+            return isinstance(rollout_tensor, torch.Tensor) and rollout_tensor.numel() > 0
+
+        has_rollout_log_probs = all(_has_rollout_log_probs(batch) for batch in unpacked_batches)
+        rollout_log_probs = (
+            torch.cat([batch["rollout_log_probs"] for batch in unpacked_batches], dim=0)
+            if has_rollout_log_probs
+            else None
+        )
+
+        # Apply off-policy correction using importance sampling if enabled
+        if self.args.use_tis:
+            assert (
+                has_rollout_log_probs and rollout_log_probs is not None
+            ), "rollout_log_probs must be provided as non-empty torch.Tensor for TIS/MIS"
+
+            train_log_probs_list = list(log_probs.split(response_lengths, dim=0))
+            rollout_log_probs_list = list(rollout_log_probs.split(response_lengths, dim=0))
+            ois = (-ppo_kl).exp()
+            tis_kwargs = {
+                "args": self.args,
+                "pg_loss": pg_loss,
+                "train_log_probs": train_log_probs_list,
+                "rollout_log_probs": rollout_log_probs_list,
+                "loss_masks": loss_masks,
+                "response_lengths": response_lengths,
+                "cp_rank": self.cp_rank,
+                "cp_size": self.cp_size,
+                "cp_group": self.cp_group,
+            }
+
+            if self.args.custom_tis_function_path is not None:
+                tis_func = load_function(self.args.custom_tis_function_path)
+            else:
+                tis_func = vanilla_tis_function
+            pg_loss, loss_masks, tis_metrics = tis_func(**tis_kwargs)
+
+        if self.args.calculate_per_token_loss:
+            pg_loss = sum_of_token(pg_loss, response_lengths, loss_masks)
+            pg_clipfrac = sum_of_token(pg_clipfrac, response_lengths, loss_masks)
+            ppo_kl = sum_of_token(ppo_kl.abs(), response_lengths, loss_masks)
+        else:
+            pg_loss = sum_of_sample_mean(pg_loss, response_lengths, loss_masks)
+            pg_clipfrac = sum_of_sample_mean(pg_clipfrac, response_lengths, loss_masks)
+            ppo_kl = sum_of_sample_mean(ppo_kl.abs(), response_lengths, loss_masks)
+
+        # Only compare rollout vs. train log probs when they originate from different stages.
+        train_rollout_logprob_abs_diff = None
+        if not self.args.use_rollout_logprobs and rollout_log_probs is not None:
+            train_rollout_logprob_abs_diff = (old_log_probs - rollout_log_probs).abs()
+            train_rollout_logprob_abs_diff = sum_of_sample_mean(
+                train_rollout_logprob_abs_diff, response_lengths, loss_masks
+            ).detach()
+
+        entropy = torch.cat([batch["entropy"] for batch in unpacked_batches], dim=0)
+        entropy_loss = sum_of_sample_mean(entropy, response_lengths, loss_masks)
+
+        loss = pg_loss - self.args.entropy_coef * entropy_loss
+
+        if self.args.use_kl_loss:
+            ref_log_probs = torch.cat([batch["ref_log_probs"] for batch in unpacked_batches], dim=0)
+            importance_ratio = None
+            if self.args.use_unbiased_kl:
+                importance_ratio = torch.exp(log_probs - old_log_probs)
+            kl = compute_approx_kl(
+                log_probs,
+                ref_log_probs,
+                kl_loss_type=self.args.kl_loss_type,
+                importance_ratio=importance_ratio,
+            )
+            kl_loss = sum_of_sample_mean(kl, response_lengths, loss_masks)
+
+            loss = loss + self.args.kl_loss_coef * kl_loss
+
+        reported = {
+            "loss": loss.detach(),
+            "pg_loss": pg_loss.detach(),
+            "pg_clipfrac": pg_clipfrac.detach(),
+            "ppo_kl": ppo_kl.detach(),
+            "entropy_loss": entropy_loss.detach(),
+        }
+
+        if train_rollout_logprob_abs_diff is not None:
+            reported["train_rollout_logprob_abs_diff"] = train_rollout_logprob_abs_diff
+
+        if self.args.use_kl_loss:
+            reported["kl_loss"] = kl_loss.detach()
+
+        if self.args.use_opsm:
+            reported["opsm_clipfrac"] = opsm_clipfrac
+
+        if self.args.use_tis and tis_metrics:
+            reported["ois"] = sum_of_sample_mean(ois, response_lengths, loss_masks).detach()
+            for k, v in tis_metrics.items():
+                if self.args.calculate_per_token_loss:
+                    reported[k] = sum_of_token(v, response_lengths, loss_masks).detach()
+                else:
+                    reported[k] = sum_of_sample_mean(v, response_lengths, loss_masks).detach()
+
+        # 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-pg_{gid}"] = lr_values[gid]
+
+                kl_info = ""
+                if self.args.use_kl_loss and "kl_loss" in aggregated:
+                    kl_info = f", kl_loss: {aggregated['kl_loss']:.4f}, kl_penalty: {aggregated['kl_loss'] * self.args.kl_loss_coef:.4f}"
+                    logger.info(kl_info)
+                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
+
+    @timer
+    def update_weights(self) -> None:  # type: ignore[override]
+        """Synchronize actor weights to rollout engines.
+
+        Handles both colocated and distributed update modes. In offload mode,
+        wakes up parameters as needed to perform the update.
+        """
+        if self.args.debug_train_only or self.args.debug_rollout_only:
+            return
+
+        rollout_engines, rollout_engine_lock, num_new_engines = ray.get(
+            self.rollout_manager.get_rollout_engines_and_lock.remote()
+        )
+        if num_new_engines > 0:
+            self.weight_updater.connect_rollout_engines(rollout_engines, rollout_engine_lock)
+            dist.barrier(group=get_gloo_group())
+
+        self.weight_updater.update_weights()
+
+        if self.args.ci_test and len(rollout_engines) > 0:
+            engine = random.choice(rollout_engines)
+            engine_version = ray.get(engine.get_weight_version.remote())
+            if str(engine_version) != str(self.weight_updater.weight_version):
+                raise RuntimeError(
+                    f"Weight version mismatch! Engine: {engine_version}, Updater: {self.weight_updater.weight_version}"
+                )
+
+        clear_memory()
+
+    def _create_ref_model(self, ref_load_path: str | None):
+        """Create and initialize a separate reference model with FSDP2 CPUOffloadPolicy.
+
+        Parameters:
+            ref_load_path: Path to a directory containing a HF checkpoint. If
+                None, a ValueError is raised.
+
+        Returns:
+            FSDP2-wrapped ref model with CPU offload enabled
+
+        Note:
+            Creates a separate FSDP2 model instance for the reference model.
+            ALWAYS uses CPUOffloadPolicy for the reference model to save memory,
+            regardless of the actor model's CPU offload setting.
+        """
+        if ref_load_path is None:
+            raise ValueError("ref_load_path must be provided when loading reference model")
+
+        if os.path.isdir(ref_load_path):
+            logger.info(f"[Rank {dist.get_rank()}] Creating separate ref model from {ref_load_path}")
+
+            init_context = self._get_init_weight_context_manager()
+
+            with init_context():
+                ref_model = self.get_model_cls().from_pretrained(
+                    ref_load_path,
+                    trust_remote_code=True,
+                    attn_implementation=self.args.attn_implementation,
+                )
+
+            full_state = ref_model.state_dict()
+
+            # Always use CPUOffloadPolicy for reference, let FSDP2 handle the offload. It is faster than model.cpu().
+            ref_model = apply_fsdp2(ref_model, mesh=self.dp_mesh, cpu_offload=True, args=self.args)
+            ref_model = self._fsdp2_load_full_state_dict(ref_model, full_state, self.dp_mesh, cpu_offload=True)
+
+            logger.info(f"[Rank {dist.get_rank()}] Reference model created with FSDP2 CPUOffloadPolicy")
+            return ref_model
+        else:
+            raise NotImplementedError(f"Loading from checkpoint file {ref_load_path} not yet implemented")
+
+    def _get_model_inputs_args(self, packed_sequence: dict) -> dict:
+        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_train_inputs"):
+            model_args.update(packed_sequence["multimodal_train_inputs"])
+        return model_args
+
+
+def selective_log_softmax_raw(logits: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
+    """Fused version of the common `log_softmax -> gather` operation.
+
+    The fused version of this operation avoids the (potentially large) memory overhead
+    of allocating a new tensor to store the full logprobs.
+
+    Parameters:
+        logits: Tensor of shape [..., V] containing model logits.
+        input_ids: Tensor of shape [...] of token indices whose log-probabilities are gathered.
+
+    Returns:
+        Tensor of shape [...] containing the log-probabilities corresponding to `input_ids`.
+    """
+    logprobs = logits.log_softmax(dim=-1)
+    return torch.gather(logprobs, dim=-1, index=input_ids.unsqueeze(-1)).squeeze(-1)
+
+
+selective_log_softmax_compiled = torch.compile(dynamic=True)(selective_log_softmax_raw)
+
+
+def gather_log_probs_packed(
+    shifted_logits: torch.Tensor,
+    input_ids: torch.Tensor,
+    allow_compile: bool,
+    cu_seqlens: torch.Tensor | float | None = None,
+    temperature: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Gather next-token log probabilities for packed sequences.
+
+    Parameters:
+        logits: Model logits of shape [B, T, V] or [T, V].
+        input_ids: Token ids of shape [B, T] or [T].
+        cu_seqlens: Optional cumulative sequence lengths (unused here). Present
+            for API compatibility with callers.
+
+    Returns:
+        A tensor of shape [T-1] (or [B, T-1]) with log-probabilities of targets.
+    """
+    # Handle batch dimension - logits should be [batch_size, seq_len, vocab_size]
+    if shifted_logits.dim() == 3:
+        # Remove batch dimension for packed sequences
+        shifted_logits = shifted_logits.squeeze(0)
+        input_ids = input_ids.squeeze(0)
+
+    if temperature is not None:
+        shifted_logits = shifted_logits.div(temperature)
+
+    targets = input_ids[1:].to(device=shifted_logits.device)
+
+    # Gather log probs for targets
+    selective_log_softmax = selective_log_softmax_compiled if allow_compile else selective_log_softmax_raw
+    return selective_log_softmax(shifted_logits, targets)
+
+
+def get_logprob_and_entropy_with_cp(
+    logits: torch.Tensor,
+    target_tokens: torch.Tensor,
+    cp_rank: int,
+    cp_size: int,
+    cp_group,
+    model_input_ids: torch.Tensor,
+    allow_compile: bool,
+    temperature: float | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute log probabilities and entropy in Context Parallel mode.
+
+    Parameters:
+        logits: Model output logits with shape [chunk_size, vocab_size]
+        target_tokens: Target tokens with shape [total_seq_len]
+        cp_rank: Current CP rank
+        cp_size: CP world size
+        cp_group: CP communication group
+        model_input_ids: Model input_ids (used for the last rank)
+        allow_compile: Whether to allow compilation
+        temperature: Temperature parameter (optional)
+
+    Returns:
+        log_probs: Aggregated log probabilities with shape [total_seq_len - 1]
+        entropy: Aggregated entropy with shape [total_seq_len - 1]
+    """
+    # Fast path for non-CP mode (cp_size=1): avoid unnecessary communication
+    if cp_size == 1:
+        shifted_logits = logits[:-1, :]
+        local_log_probs = gather_log_probs_packed(
+            shifted_logits, target_tokens, allow_compile=allow_compile, temperature=temperature
+        )
+        log_probs_full = torch.log_softmax(shifted_logits, dim=-1)
+        probs = torch.softmax(shifted_logits, dim=-1)
+        entropy = -(probs * log_probs_full).sum(dim=-1)
+        return local_log_probs, entropy
+
+    chunk_size = logits.shape[0]
+    tokens_start_index = chunk_size * cp_rank
+    tokens_end_index = (
+        tokens_start_index + chunk_size + 1 if cp_rank < cp_size - 1 else tokens_start_index + chunk_size
+    )
+
+    # For the last rank, remove the last logit
+    logits = logits if cp_rank < cp_size - 1 else logits[:-1, :]
+
+    # Get local tokens for current rank
+    local_tokens = (
+        target_tokens[tokens_start_index:tokens_end_index] if cp_rank < cp_size - 1 else model_input_ids.squeeze(0)
+    )
+
+    # Compute local log probs
+    local_log_probs = gather_log_probs_packed(
+        logits, local_tokens, allow_compile=allow_compile, temperature=temperature
+    )
+
+    # Pad for the last rank
+    if cp_rank == cp_size - 1:
+        local_log_probs = F.pad(local_log_probs, (0, chunk_size - local_log_probs.shape[0]), value=0)
+
+    # Compute entropy
+    shifted_logits = logits[:-1, :] if cp_rank == cp_size - 1 else logits
+    log_probs_full = torch.log_softmax(shifted_logits, dim=-1)
+    probs = torch.softmax(shifted_logits, dim=-1)
+    entropy = -(probs * log_probs_full).sum(dim=-1)
+
+    # Pad entropy for the last rank
+    if cp_rank == cp_size - 1:
+        entropy = F.pad(entropy, (0, chunk_size - entropy.shape[0]), value=0)
+
+    # Merge with a single all_gather: stack as [2, chunk_size]
+    stacked_local = torch.stack([local_log_probs, entropy], dim=0)
+    gathered_stacked = torch.distributed.nn.functional.all_gather(stacked_local, group=cp_group)
+
+    # Concatenate by effective length (non-last rank=chunk_size, last rank=chunk_size-1)
+    lp_parts, ent_parts = [], []
+    for r in range(cp_size):
+        eff_len = chunk_size if r < cp_size - 1 else max(0, chunk_size - 1)
+        if eff_len > 0:
+            lp_parts.append(gathered_stacked[r][0][:eff_len])
+            ent_parts.append(gathered_stacked[r][1][:eff_len])
+
+    log_probs = torch.cat(lp_parts, dim=0) if lp_parts else local_log_probs.new_zeros((0,))
+    entropy_result = torch.cat(ent_parts, dim=0) if ent_parts else entropy.new_zeros((0,))
+
+    # Truncate to global effective length T-1 (packed tokens length is T)
+    log_probs = log_probs[: len(target_tokens) - 1]
+    entropy_result = entropy_result[: len(target_tokens) - 1]
+
+    return log_probs, entropy_result
+
+
+def sum_of_sample_mean(x: torch.Tensor, response_lengths: list[int], loss_masks: list[torch.Tensor]) -> torch.Tensor:
+    """Compute sum of per-sample means across variable-length responses.
+
+    Parameters:
+        x: Flat tensor containing concatenated per-token values across samples.
+        response_lengths: Lengths of each sample's response segment in `x`.
+        loss_masks: Per-sample masks aligned with `response_lengths`.
+
+    Returns:
+        A scalar tensor equal to the sum over samples of the mean value within
+        each sample's response segment.
+    """
+    return sum(
+        [
+            (x_i * loss_mask_i).sum() / torch.clamp_min(loss_mask_i.sum(), 1)
+            for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=False)
+        ]
+    )
+
+
+@torch.no_grad()
+def move_torch_optimizer(optimizer, device):
+    """ref: https://github.com/volcengine/verl/blob/main/verl/utils/fsdp_utils.py"""
+    if not optimizer.state:
+        return
+
+    for param_group in optimizer.param_groups:
+        for param in param_group["params"]:
+            state = optimizer.state[param]
+            for key, value in state.items():
+                if isinstance(value, torch.Tensor):
+                    state[key] = value.to(device, non_blocking=True)
+
+    torch.cuda.synchronize()
+
+
+def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None):
+    """Apply FSDP v2 to the model.
+
+    Args:
+        model: The model to wrap with FSDP
+        mesh: Optional DeviceMesh for FSDP. If None, uses all ranks.
+        cpu_offload: If True, offload parameters, gradients, and optimizer states
+            to CPU. The optimizer step will run on CPU. (Default: False)
+        args: Arguments containing precision settings (fp16/bf16)
+
+    Ref: https://github.com/volcengine/verl/blob/main/verl/utils/fsdp_utils.py
+    """
+    from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard
+
+    offload_policy = CPUOffloadPolicy() if cpu_offload else None
+
+    layer_cls_to_wrap = model._no_split_modules
+    assert len(layer_cls_to_wrap) > 0 and layer_cls_to_wrap[0] is not None
+
+    modules = [
+        module
+        for name, module in model.named_modules()
+        if module.__class__.__name__ in layer_cls_to_wrap
+        or (isinstance(module, torch.nn.Embedding) and not model.config.tie_word_embeddings)
+    ]
+
+    # Determine precision policy based on args
+    param_dtype = torch.bfloat16  # Default to bf16 as before
+    reduce_dtype = torch.float32
+
+    if args.fp16:
+        param_dtype = torch.float16
+
+    logger.info(f"FSDP MixedPrecision Policy: param_dtype={param_dtype}, reduce_dtype={reduce_dtype}")
+
+    fsdp_kwargs = {
+        "mp_policy": MixedPrecisionPolicy(
+            param_dtype=param_dtype,
+            reduce_dtype=reduce_dtype,
+        ),
+        "offload_policy": offload_policy,
+        "mesh": mesh,
+    }
+
+    # Apply FSDP to each module (offload_policy=None is equivalent to not passing it)
+    for module in modules:
+        fully_shard(module, **fsdp_kwargs)
+
+    # Apply FSDP to the top-level model
+    fully_shard(model, **fsdp_kwargs)
+
+    return model
+
+
+def sum_of_token(x: torch.Tensor, response_lengths: list[int], loss_masks: list[torch.Tensor]) -> torch.Tensor:
+    return sum(
+        [
+            (x_i * loss_mask_i).sum()
+            for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=False)
+        ]
+    )

slime/backends/fsdp_utils/arguments.py

@@ -0,0 +1,101 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import dataclasses
+from dataclasses import dataclass
+
+import yaml
+
+
+@dataclass
+class FSDPArgs:
+    # Optim
+    optimizer: str = "adam"  # Optimizer type: "adam" (AdamW)
+    lr: float = 2e-5
+    lr_warmup_init: float = 0.0
+    min_lr: float = 0.0
+    lr_decay_style: str = "constant"
+    lr_decay_iters: int | None = None
+    lr_warmup_iters: int = 0
+    lr_warmup_fraction: float | None = None
+    lr_wsd_decay_iters: int | None = None
+    lr_wsd_decay_style: str | None = None
+    use_checkpoint_lr_scheduler: bool = True
+    override_lr_scheduler: bool = False
+    weight_decay: float = 0.0
+    adam_beta1: float = 0.9
+    adam_beta2: float = 0.95
+    adam_eps: float = 1e-8
+    warmup_ratio: float = 0.03
+
+    attn_implementation: str = "flash_attention_2"
+
+    # Logging
+    wandb_project: str = "slime-fsdp"
+    wandb_run_name: str | None = None
+
+    # Precision
+    gradient_checkpointing: bool = False
+    fp16: bool = False
+
+    # FSDP configuration
+    fsdp_state_dict_cpu_offload: bool = True  # If True, offload full state dict to CPU during collection.
+    fsdp_cpu_offload: bool = (
+        False  # If True, offload parameters, gradients, and optimizer states to CPU (optimizer runs on CPU)
+    )
+    fsdp_cpu_backend: str | None = (
+        "gloo"  # CPU backend for FSDP CPU offload (e.g., "gloo"). Set to None to disable hybrid backend.
+    )
+
+    deterministic_mode: bool = False  # This name must be the same as Megatron's
+
+    # Context Parallelism
+    context_parallel_size: int = 1  # Context Parallelism size
+    # Profile
+    record_memory_history: bool = False
+    memory_snapshot_path: str = "snapshot.pickle"
+    use_pytorch_profiler: bool = False
+    profile_step_start: int = 10
+    profile_step_end: int = 12
+    tensorboard_dir: str | None = None
+
+    # YAML bookkeeping
+    config: str | None = None
+
+
+def parse_fsdp_cli(extra_args_provider=None):
+    parser = argparse.ArgumentParser("FSDP Training (slime)")
+    parser.add_argument("--config", type=str, default=None, help="YAML config path")
+    for f in dataclasses.fields(FSDPArgs):
+        if f.name == "config":
+            continue
+
+        # Handle union types like int | None, str | None, etc.
+        if hasattr(f.type, "__args__"):  # Check if it's a Union type
+            # For T | None, use T as the type
+            non_none_types = [t for t in f.type.__args__ if t is not type(None)]
+            arg_type = non_none_types[0] if non_none_types else str
+        else:
+            arg_type = f.type
+
+        if arg_type is bool:
+            parser.add_argument(f"--{f.name.replace('_', '-')}", action="store_true")
+        else:
+            parser.add_argument(f"--{f.name.replace('_', '-')}", type=arg_type, default=f.default)
+
+    if extra_args_provider is not None:
+        parser = extra_args_provider(parser)
+    args = parser.parse_args()
+    return args
+
+
+def load_fsdp_args(extra_args_provider=None):
+    args = parse_fsdp_cli(extra_args_provider)
+    if args.config:
+        with open(args.config) as f:
+            data = yaml.safe_load(f) or {}
+        for k, v in data.items():
+            if not hasattr(args, k):
+                setattr(args, k, v)
+    return args

slime/backends/fsdp_utils/checkpoint.py

@@ -0,0 +1,252 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.distributed as dist
+import torch.distributed.checkpoint as dcp
+from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
+from torch.distributed.checkpoint.stateful import Stateful
+
+logger = logging.getLogger(__name__)
+
+
+class ModelState(Stateful):
+    """Wrapper for model state only."""
+
+    def __init__(self, model):
+        self.model = model
+
+    def state_dict(self):
+        model_state_dict, _ = get_state_dict(self.model, optimizers=[])
+        return {"model": model_state_dict}
+
+    def load_state_dict(self, state_dict):
+        set_state_dict(self.model, optimizers=[], model_state_dict=state_dict["model"], optim_state_dict=None)
+
+
+class OptimizerState(Stateful):
+    """Wrapper for optimizer state only."""
+
+    def __init__(self, model, optimizer):
+        self.model = model
+        self.optimizer = optimizer
+
+    def state_dict(self):
+        _, optimizer_state_dict = get_state_dict(self.model, optimizers=self.optimizer)
+        return {"optim": optimizer_state_dict}
+
+    def load_state_dict(self, state_dict):
+        set_state_dict(
+            self.model, optimizers=self.optimizer, model_state_dict=None, optim_state_dict=state_dict["optim"]
+        )
+
+
+class LRSchedulerState(Stateful):
+    """Wrapper for LR scheduler state only."""
+
+    def __init__(self, lr_scheduler):
+        self.lr_scheduler = lr_scheduler
+
+    def state_dict(self):
+        return {"lr_scheduler": self.lr_scheduler.state_dict()}
+
+    def load_state_dict(self, state_dict):
+        self.lr_scheduler.load_state_dict(state_dict["lr_scheduler"])
+
+
+def _read_checkpoint_metadata(path: Path) -> dict[str, Any]:
+    if not path.exists():
+        return {}
+    try:
+        return json.loads(path.read_text())
+    except json.JSONDecodeError:
+        logger.warning(f"Failed to parse checkpoint metadata at {path}")
+        return {}
+
+
+def _write_checkpoint_metadata(path: Path, metadata: dict[str, Any]) -> None:
+    tmp_path = path.with_suffix(path.suffix + ".tmp")
+    tmp_path.write_text(json.dumps(metadata, indent=2, sort_keys=True))
+    tmp_path.replace(path)
+
+
+def load(actor: Any) -> dict[str, Any] | None:
+    """Load checkpoint from disk.
+
+    Loads model weights and optionally optimizer state from separate directories.
+    This allows loading weights without optimizer or deleting optimizer before loading.
+    """
+    load_root = getattr(actor.args, "load", None)
+    if load_root is None:
+        return None
+
+    root_path = Path(load_root).expanduser()
+    if not root_path.exists():
+        logger.info(f"[FSDP] Checkpoint directory {root_path} not found; skipping load.")
+        return None
+
+    target_step = getattr(actor.args, "ckpt_step", None)
+    if target_step is None:
+        tracker_file = root_path / "latest_checkpointed_iteration.txt"
+        if not tracker_file.exists():
+            logger.info(f"[FSDP] No tracker file at {tracker_file}; skipping load.")
+            return None
+        tracker_text = tracker_file.read_text().strip()
+        target_step = int(tracker_text)
+
+    checkpoint_dir = root_path / f"iter_{target_step:07d}"
+    model_dir = checkpoint_dir / "model"
+    optimizer_dir = checkpoint_dir / "optimizer"
+    lr_scheduler_dir = checkpoint_dir / "lr_scheduler"
+
+    if not model_dir.exists():
+        logger.info(f"[FSDP] Model checkpoint {model_dir} not found; skipping load.")
+        return None
+
+    # Load model weights (always)
+    model_state = ModelState(actor.model)
+    state_dict = {"model_state": model_state}
+
+    try:
+        dcp.load(state_dict=state_dict, checkpoint_id=str(model_dir))
+        logger.info(f"[FSDP] Loaded model from {model_dir}")
+    except Exception as e:
+        logger.error(f"[FSDP] Failed to load model from {model_dir}: {e}")
+        return None
+
+    # Load optimizer state (optional)
+    load_optimizer = not getattr(actor.args, "no_load_optim", False) and hasattr(actor, "optimizer")
+    if load_optimizer and optimizer_dir.exists():
+        optimizer_state = OptimizerState(actor.model, actor.optimizer)
+        optim_state_dict = {"optim_state": optimizer_state}
+        try:
+            dcp.load(state_dict=optim_state_dict, checkpoint_id=str(optimizer_dir))
+            logger.info(f"[FSDP] Loaded optimizer from {optimizer_dir}")
+        except Exception as e:
+            logger.warning(f"[FSDP] Failed to load optimizer from {optimizer_dir}: {e}")
+    elif load_optimizer:
+        logger.info(f"[FSDP] Optimizer checkpoint not found at {optimizer_dir}, skipping optimizer load.")
+
+    # Load LR scheduler state (optional)
+    load_lr_scheduler = hasattr(actor, "lr_scheduler") and lr_scheduler_dir.exists()
+    if load_lr_scheduler:
+        lr_scheduler_state = LRSchedulerState(actor.lr_scheduler)
+        lr_scheduler_state_dict = {"lr_scheduler_state": lr_scheduler_state}
+        try:
+            dcp.load(state_dict=lr_scheduler_state_dict, checkpoint_id=str(lr_scheduler_dir))
+            logger.info(f"[FSDP] Loaded LR scheduler from {lr_scheduler_dir}")
+        except Exception as e:
+            logger.warning(f"[FSDP] Failed to load LR scheduler from {lr_scheduler_dir}: {e}")
+    elif hasattr(actor, "lr_scheduler"):
+        logger.info(f"[FSDP] LR scheduler checkpoint not found at {lr_scheduler_dir}, skipping LR scheduler load.")
+
+    rng_state = None
+    rng_path = checkpoint_dir / "rng.pt"
+    if rng_path.exists():
+        rng_state = torch.load(rng_path, map_location="cpu")
+
+    metadata = _read_checkpoint_metadata(checkpoint_dir / "meta.json")
+
+    return {
+        "rng": rng_state,
+        "metadata": metadata,
+        "iteration": target_step,
+    }
+
+
+def finalize_load(actor: Any, checkpoint_payload: dict[str, Any] | None) -> None:
+    if checkpoint_payload is None:
+        dist.barrier()
+        return
+
+    if checkpoint_payload.get("rng") is not None and not getattr(actor.args, "no_load_rng", False):
+        rng_state = checkpoint_payload["rng"]
+        if "torch" in rng_state:
+            torch.set_rng_state(rng_state["torch"])
+        if torch.cuda.is_available() and "cuda" in rng_state:
+            torch.cuda.set_rng_state_all(rng_state["cuda"])
+
+    metadata = checkpoint_payload.get("metadata") or {}
+    iteration = checkpoint_payload.get("iteration")
+    if metadata:
+        actor.global_step = int(metadata.get("global_step", actor.global_step))
+        actor.micro_step = int(metadata.get("micro_step", actor.micro_step))
+        next_rollout = metadata.get("next_rollout_id")
+        if next_rollout is not None:
+            actor.args.start_rollout_id = next_rollout
+    elif iteration is not None:
+        if getattr(actor.args, "start_rollout_id", None) is None:
+            actor.args.start_rollout_id = iteration
+
+    torch.cuda.synchronize()
+    dist.barrier()
+
+
+def save(actor: Any, iteration: int) -> None:
+    """Save checkpoint to disk.
+
+    Saves model weights and optimizer state to separate directories.
+    This allows loading weights without optimizer or deleting optimizer before loading.
+    """
+    torch.cuda.synchronize()
+
+    base_dir = Path(actor.args.save).expanduser()
+    step_id = iteration + 1
+    checkpoint_dir = base_dir / f"iter_{step_id:07d}"
+    model_dir = checkpoint_dir / "model"
+    optimizer_dir = checkpoint_dir / "optimizer"
+    lr_scheduler_dir = checkpoint_dir / "lr_scheduler"
+
+    if dist.get_rank() == 0:
+        checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        model_dir.mkdir(parents=True, exist_ok=True)
+        optimizer_dir.mkdir(parents=True, exist_ok=True)
+        lr_scheduler_dir.mkdir(parents=True, exist_ok=True)
+    dist.barrier()
+
+    # Save model weights
+    model_state = ModelState(actor.model)
+    state_dict = {"model_state": model_state}
+    dcp.save(state_dict, checkpoint_id=str(model_dir))
+
+    # Save optimizer state
+    if hasattr(actor, "optimizer") and actor.optimizer is not None:
+        optimizer_state = OptimizerState(actor.model, actor.optimizer)
+        optim_state_dict = {"optim_state": optimizer_state}
+        dcp.save(optim_state_dict, checkpoint_id=str(optimizer_dir))
+
+    # Save LR scheduler state
+    if hasattr(actor, "lr_scheduler") and actor.lr_scheduler is not None:
+        lr_scheduler_state = LRSchedulerState(actor.lr_scheduler)
+        lr_scheduler_state_dict = {"lr_scheduler_state": lr_scheduler_state}
+        dcp.save(lr_scheduler_state_dict, checkpoint_id=str(lr_scheduler_dir))
+
+    if dist.get_rank() == 0:
+        rng_state = {"torch": torch.get_rng_state()}
+        rng_state["cuda"] = torch.cuda.get_rng_state_all()
+        torch.save(rng_state, checkpoint_dir / "rng.pt")
+
+        metadata = {
+            "iteration": step_id,
+            "rollout_id": iteration,
+            "next_rollout_id": iteration + 1,
+            "global_step": actor.global_step,
+            "micro_step": actor.micro_step,
+            "world_size": dist.get_world_size(),
+            "timestamp": time.time(),
+        }
+        _write_checkpoint_metadata(checkpoint_dir / "meta.json", metadata)
+
+        tracker_file = base_dir / "latest_checkpointed_iteration.txt"
+        tracker_file.write_text(str(step_id))
+        logger.info(f"[FSDP] Saved checkpoint to {checkpoint_dir}")
+
+    dist.barrier()

slime/backends/fsdp_utils/data_packing.py

@@ -0,0 +1,221 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Data packing utilities for FSDP backend to reduce padding overhead."""
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+from slime.utils.seqlen_balancing import get_seqlen_balanced_partitions
+
+
+def pack_sequences(
+    tokens: list[list[int]],
+    loss_masks: list[list[int]],
+    rewards: list[float],
+    raw_rewards: list,
+    response_lengths: list[int],
+    advantages: list[float],
+    returns: list[float],
+    rollout_log_probs: list[list[float]] | None = None,
+    multimodal_train_inputs: list[dict] | None = None,
+    max_tokens_per_gpu: int | None = None,
+    num_packs: int | None = None,
+) -> list[dict]:
+    """
+    Pack sequences into dense batches with cumulative sequence lengths.
+
+    Args:
+        tokens: List of token sequences
+        loss_masks: List of loss masks
+        rewards: List of rewards per sequence
+        raw_rewards: List of raw rewards per sequence
+        response_lengths: List of response lengths per sequence
+        advantages: List of advantages per sequence
+        returns: List of returns per sequence
+        rollout_log_probs: List of rollout log probabilities per sequence
+        multimodal_train_inputs: List of dict of multimodal tensors for training per sequence
+        max_tokens_per_gpu: Maximum tokens per GPU pack
+        num_packs: Explicit number of packs to create
+
+    Returns:
+        List of packed batches with tokens, masks, cu_seqlens, rewards, raw_rewards, response_lengths, advantages, returns
+    """
+    if not tokens:
+        return []
+
+    seq_lengths = [len(t) for t in tokens]
+
+    # Determine number of packs and use balanced partitioning
+    if num_packs:
+        k_partitions = num_packs
+    elif max_tokens_per_gpu:
+        total_tokens = sum(seq_lengths)
+        k_partitions = max(1, math.ceil(total_tokens / max_tokens_per_gpu))
+    else:
+        k_partitions = 1
+
+    # Use balanced partitioning for optimal load distribution
+    partitions = get_seqlen_balanced_partitions(
+        seq_lengths, k_partitions=k_partitions, equal_size=False  # Allow variable sizes for better balance
+    )
+
+    # Pack each partition
+    result = []
+    for indices in partitions:
+        # Build cumulative sequence lengths
+        cu_seqlens = [0]
+        flat_tokens = []
+        flat_masks = []
+        flat_positionids = []
+        flat_advantages = []
+        flat_returns = []
+        flat_rollout_log_probs = []
+
+        for i in indices:
+            seq_tokens = tokens[i]
+            seq_mask = loss_masks[i]
+            seq_positionids = list(range(len(seq_tokens)))
+
+            flat_tokens.extend(seq_tokens)
+            flat_positionids.extend(seq_positionids)
+            flat_masks.extend(seq_mask)
+            flat_advantages.extend(advantages[i])
+            flat_returns.extend(returns[i])
+            if rollout_log_probs:
+                flat_rollout_log_probs.extend(rollout_log_probs[i])
+            cu_seqlens.append(cu_seqlens[-1] + len(seq_tokens))
+
+        packed_batch = {
+            "tokens": torch.tensor(flat_tokens, dtype=torch.long),
+            "loss_masks": torch.tensor(flat_masks, dtype=torch.int),
+            "position_ids": torch.tensor(flat_positionids, dtype=torch.int),
+            "cu_seqlens": torch.tensor(cu_seqlens, dtype=torch.int32),
+            "rewards": torch.tensor([rewards[i] for i in indices], dtype=torch.float32),
+            "raw_reward": [raw_rewards[i] for i in indices],
+            "response_lengths": [response_lengths[i] for i in indices],
+            "advantages": torch.tensor(flat_advantages, dtype=torch.float32),
+            "returns": torch.tensor(flat_returns, dtype=torch.float32),
+            "rollout_log_probs": torch.tensor(
+                flat_rollout_log_probs, dtype=torch.float32, device=torch.cuda.current_device()
+            ),
+        }
+
+        # Collect and add multimodal training tensors for this partition
+        if multimodal_train_inputs:
+            multimodal_data = {}  # key -> concatenated tensor
+            multimodal_num_items = {}  # key -> list of item counts per sequence
+            for i in indices:
+                for key, mm_tensor in multimodal_train_inputs[i].items():
+                    if key not in multimodal_data:
+                        multimodal_data[key] = mm_tensor
+                        multimodal_num_items[key] = [mm_tensor.size(0)]
+                    else:
+                        multimodal_data[key] = torch.cat([multimodal_data[key], mm_tensor], dim=0)
+                        multimodal_num_items[key].append(mm_tensor.size(0))
+            packed_batch["multimodal_train_inputs"] = multimodal_data
+            packed_batch["multimodal_num_items"] = multimodal_num_items
+
+        result.append(packed_batch)
+
+    return result
+
+
+def unpack_sequences(packed_batch: dict) -> list[dict]:
+    """
+    Unpack sequences from a packed batch.
+
+    Args:
+        packed_batch: Packed batch
+
+    Returns:
+        List of unpacked batches
+    """
+
+    cu_seqlens = packed_batch["cu_seqlens"]
+    num_sequences = len(cu_seqlens) - 1
+    response_lengths = packed_batch["response_lengths"]
+    multimodal_num_items = packed_batch.get("multimodal_num_items", {})
+
+    instances = []
+
+    # Calculate pad_length by counting trailing zeros
+    tokens = packed_batch["tokens"]
+    nonzero_indices = (tokens != 0).nonzero(as_tuple=True)[0]
+    if len(nonzero_indices) > 0:
+        # Last non-zero index, pad_length is everything after it
+        pad_length = len(tokens) - nonzero_indices[-1].item() - 1
+    else:
+        pad_length = 0  # No padding if no non-zero tokens (or all zeros)
+    for i in range(num_sequences):
+        start_idx = cu_seqlens[i].item()
+        end_idx = cu_seqlens[i + 1].item()
+        instance = {}
+
+        # Copy any additional attributes that might exist in the packed batch
+        for key, value in packed_batch.items():
+            if key not in instance:
+                # Skip multimodal_num_items - it's metadata
+                if key == "multimodal_num_items":
+                    continue
+                # Handle multimodal_train_inputs dict: split each tensor using multimodal_num_items
+                elif key == "multimodal_train_inputs" and isinstance(value, dict):
+                    instance[key] = {}
+                    for mm_key, mm_tensor in value.items():
+                        if mm_key in multimodal_num_items:
+                            num_items_list = multimodal_num_items[mm_key]
+                            start_mm_idx = sum(num_items_list[:i])
+                            end_mm_idx = start_mm_idx + num_items_list[i]
+                            if num_items_list[i] > 0:
+                                instance[key][mm_key] = mm_tensor[start_mm_idx:end_mm_idx]
+                # For tensor attributes, we need to slice them appropriately
+                elif isinstance(value, torch.Tensor):
+                    if key in ["log_probs", "ref_log_probs", "cur_log_probs", "entropy"]:
+                        # These are computed from logits[:-1] so they have length seq_len-1
+                        instance[key] = value[
+                            end_idx - 1 - response_lengths[i] - pad_length : end_idx - 1 - pad_length
+                        ]
+                    elif key == "rollout_log_probs":
+                        # rollout_log_probs is packed based on response_lengths, so slice differently
+                        instance[key] = value[sum(response_lengths[:i]) : sum(response_lengths[: i + 1])]
+                    elif key in ["tokens", "position_ids"]:
+                        # For other tensor attributes, try to slice them
+                        if len(value) > start_idx:
+                            instance[key] = value[start_idx:end_idx]
+                        else:
+                            raise ValueError(f"Attribute {key} is not found in the packed batch")
+                    elif key in ["loss_masks", "advantages", "returns"]:
+                        instance[key] = value[sum(response_lengths[:i]) : sum(response_lengths[: i + 1])]
+                elif isinstance(value, list):
+                    instance[key] = value[i]
+                else:
+                    raise ValueError(f"Attribute {key} is not found in the packed batch")
+
+        instances.append(instance)
+
+    return instances
+
+
+def pad_packed_sequence_with_cp(packed_sequence: dict, cp_size: int) -> dict:
+    """Pad packed sequence to make total length divisible by cp_size.
+
+    Args:
+        packed_sequence: Packed sequence dict containing tokens, position_ids, cu_seqlens, etc.
+        cp_size: Context parallelism world size
+
+    Returns:
+        Padded packed sequence
+    """
+    seq_length = len(packed_sequence["tokens"])
+    # Calculate padding needed: (cp_size - seq_length % cp_size) % cp_size
+    remainder = seq_length % cp_size
+    pad_length = (cp_size - remainder) % cp_size
+
+    if pad_length > 0:
+        packed_sequence["tokens"] = F.pad(packed_sequence["tokens"], (0, pad_length), value=0)
+        packed_sequence["position_ids"] = F.pad(packed_sequence["position_ids"], (0, pad_length), value=0)
+        packed_sequence["loss_masks"] = F.pad(packed_sequence["loss_masks"], (0, pad_length), value=0)
+        packed_sequence["cu_seqlens"][-1] += pad_length
+    return packed_sequence

slime/backends/fsdp_utils/kernels/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+

slime/backends/fsdp_utils/kernels/fused_experts.py

@@ -0,0 +1,384 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import torch
+import triton.language as tl
+from sglang.srt.layers.moe.fused_moe_triton.fused_moe import (
+    invoke_fused_moe_kernel,
+    moe_align_block_size,
+    moe_sum_reduce,
+    silu_and_mul,
+)
+
+from .fused_moe_triton_backward_kernels import invoke_fused_moe_backward_kernel
+
+
+class GateUpProjFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ):
+        num_tokens, _ = hidden_states.shape
+        E, N, _ = w1.shape
+        # We execute the fused_moe kernel in chunks to circumvent this issue:
+        # https://github.com/vllm-project/vllm/issues/5938
+        CHUNK_SIZE = 64 * 1024
+
+        # default deterministic config
+        config = {
+            "BLOCK_SIZE_M": 64,
+            "BLOCK_SIZE_N": 64,
+            "BLOCK_SIZE_K": 32,
+            "GROUP_SIZE_M": 8,
+        }
+
+        topk = topk_ids.shape[1]
+
+        intermediate_cache1 = torch.empty(
+            (num_tokens * topk, N),
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+
+        for chunk in range((num_tokens // CHUNK_SIZE) + 1):
+            begin_chunk_idx, end_chunk_idx = (
+                chunk * CHUNK_SIZE,
+                min((chunk + 1) * CHUNK_SIZE, num_tokens),
+            )
+            curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
+            cur_intermediate_cache1 = intermediate_cache1[begin_chunk_idx * topk : end_chunk_idx * topk]
+
+            curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
+            curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
+
+            sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+                curr_topk_ids, config["BLOCK_SIZE_M"], E
+            )
+
+            invoke_fused_moe_kernel(
+                curr_hidden_states,
+                w1,
+                None,
+                cur_intermediate_cache1,
+                None,
+                None,
+                None,
+                curr_topk_weights,
+                curr_topk_ids,
+                sorted_token_ids,
+                expert_ids,
+                num_tokens_post_padded,
+                False,
+                topk_ids.shape[1],
+                config,
+                compute_type=tl.bfloat16,
+                use_fp8_w8a8=False,
+                use_int8_w8a8=False,
+                use_int8_w8a16=False,
+                use_int4_w4a16=False,
+                per_channel_quant=False,
+                block_shape=None,
+                c_sorted=False,
+                filter_expert=True,
+            )
+
+        ctx.save_for_backward(hidden_states, w1, topk_weights, topk_ids)
+        ctx.config = config
+        ctx.num_tokens = num_tokens
+        ctx.topk = topk
+
+        return intermediate_cache1
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """
+        Backward pass for GateUpProjFunction using Triton kernels.
+
+        Args:
+            grad_output: shape (num_tokens * topk, N)
+
+        Returns:
+            (grad_hidden_states, grad_w1, grad_topk_weights, None)
+        """
+
+        hidden_states, w1, topk_weights, topk_ids = ctx.saved_tensors
+        config = ctx.config
+        num_tokens = ctx.num_tokens
+        topk = ctx.topk
+
+        E, N, D_in = w1.shape
+        CHUNK_SIZE = 64 * 1024
+
+        # Initialize gradient tensors
+        grad_hidden_states = torch.zeros_like(hidden_states)
+        grad_w1 = torch.zeros_like(w1)
+        # GateUpProj stage doesn't need topk_weights gradient
+        grad_topk_weights = torch.zeros_like(topk_weights)
+
+        # Process in chunks to match forward pass
+        for chunk in range((num_tokens // CHUNK_SIZE) + 1):
+            begin_chunk_idx, end_chunk_idx = (
+                chunk * CHUNK_SIZE,
+                min((chunk + 1) * CHUNK_SIZE, num_tokens),
+            )
+
+            curr_num_tokens = end_chunk_idx - begin_chunk_idx
+            if curr_num_tokens == 0:
+                continue
+
+            curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
+            curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
+            curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
+            curr_grad_output = grad_output[begin_chunk_idx * topk : end_chunk_idx * topk]
+
+            # Get aligned metadata
+            sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+                curr_topk_ids, config["BLOCK_SIZE_M"], E
+            )
+
+            # Prepare gradient buffer for this chunk
+            curr_grad_hidden_states = torch.zeros_like(curr_hidden_states)
+            curr_grad_w1 = torch.zeros_like(w1)
+
+            # Call Triton backward kernel with MUL_ROUTED_WEIGHT=False
+            # Use chunk of hidden_states to match sorted_token_ids indices
+            invoke_fused_moe_backward_kernel(
+                grad_output=curr_grad_output,
+                input=curr_hidden_states,  # Use chunk of hidden_states to match sorted_token_ids
+                weight=w1,
+                grad_input=curr_grad_hidden_states,
+                grad_weight=curr_grad_w1,
+                grad_topk_weights=None,  # Not needed for GateUpProj
+                topk_weights=curr_topk_weights,
+                topk_ids=curr_topk_ids,
+                sorted_token_ids=sorted_token_ids,
+                expert_ids=expert_ids,
+                num_tokens_post_padded=num_tokens_post_padded,
+                mul_routed_weight=False,
+                top_k=topk,
+                config=config,
+                compute_type=tl.bfloat16,
+            )
+
+            # Accumulate gradients
+            grad_hidden_states[begin_chunk_idx:end_chunk_idx] += curr_grad_hidden_states
+            grad_w1 += curr_grad_w1
+
+        return grad_hidden_states, grad_w1, grad_topk_weights, None
+
+
+class SiluAndMulFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, intermediate_cache1: torch.Tensor):
+        num_tokens, N = intermediate_cache1.shape
+        intermediate_cache2 = torch.empty(
+            (num_tokens, N // 2),
+            device=intermediate_cache1.device,
+            dtype=intermediate_cache1.dtype,
+        )
+        silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
+
+        ctx.save_for_backward(intermediate_cache1)
+        return intermediate_cache2
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (intermediate_cache1,) = ctx.saved_tensors
+        N = intermediate_cache1.shape[-1]
+        x1, x2 = intermediate_cache1.view(-1, N).chunk(2, dim=-1)
+        silu_x1 = torch.nn.functional.silu(x1)
+
+        sig = torch.sigmoid(x1)
+        dsilu_dx1 = sig + x1 * sig * (1 - sig)
+        grad_x1 = grad_output * x2 * dsilu_dx1
+        grad_x2 = grad_output * silu_x1
+        grad_input = torch.cat([grad_x1, grad_x2], dim=-1)
+
+        return grad_input.view_as(intermediate_cache1)
+
+
+class DownProjFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        intermediate_cache2: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ):
+        num_tokens, _ = intermediate_cache2.shape
+        topk = topk_ids.shape[1]
+        num_tokens //= topk
+        E, _, _ = w2.shape
+        # We execute the fused_moe kernel in chunks to circumvent this issue:
+        # https://github.com/vllm-project/vllm/issues/5938
+        CHUNK_SIZE = 64 * 1024
+
+        # default deterministic config
+        config = {
+            "BLOCK_SIZE_M": 64,
+            "BLOCK_SIZE_N": 64,
+            "BLOCK_SIZE_K": 32,
+            "GROUP_SIZE_M": 8,
+        }
+
+        intermediate_cache3 = torch.empty(
+            (num_tokens, topk, w2.shape[1]),
+            device=intermediate_cache2.device,
+            dtype=intermediate_cache2.dtype,
+        )
+
+        for chunk in range((num_tokens // CHUNK_SIZE) + 1):
+            begin_chunk_idx, end_chunk_idx = (
+                chunk * CHUNK_SIZE,
+                min((chunk + 1) * CHUNK_SIZE, num_tokens),
+            )
+            cur_intermediate_cache2 = intermediate_cache2[begin_chunk_idx * topk : end_chunk_idx * topk]
+            cur_intermediate_cache3 = intermediate_cache3[begin_chunk_idx:end_chunk_idx]
+
+            curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
+            curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
+
+            sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+                curr_topk_ids, config["BLOCK_SIZE_M"], E
+            )
+            invoke_fused_moe_kernel(
+                cur_intermediate_cache2,
+                w2,
+                None,
+                cur_intermediate_cache3,
+                None,
+                None,
+                None,
+                curr_topk_weights,
+                curr_topk_ids,
+                sorted_token_ids,
+                expert_ids,
+                num_tokens_post_padded,
+                True,
+                1,
+                config,
+                compute_type=tl.bfloat16,
+                use_fp8_w8a8=False,
+                use_int8_w8a8=False,
+                use_int8_w8a16=False,
+                use_int4_w4a16=False,
+                per_channel_quant=False,
+                block_shape=None,
+                a_use_tma=False,
+                b_use_tma=False,
+            )
+
+        ctx.save_for_backward(intermediate_cache2, w2, topk_weights, topk_ids)
+        ctx.config = config
+        ctx.num_tokens = num_tokens
+        ctx.topk = topk
+
+        return intermediate_cache3
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """
+        Backward pass for DownProjFunction using Triton kernels.
+
+        Args:
+            grad_output: shape (num_tokens, topk, hidden_size)
+
+        Returns:
+            (grad_intermediate_cache2, grad_w2, grad_topk_weights, None)
+        """
+        intermediate_cache2, w2, topk_weights, topk_ids = ctx.saved_tensors
+        config = ctx.config
+        num_tokens = ctx.num_tokens
+        topk = ctx.topk
+
+        E, hidden_size, intermediate_size = w2.shape
+        CHUNK_SIZE = 64 * 1024
+
+        # Initialize gradient tensors
+        grad_intermediate_cache2 = torch.zeros_like(intermediate_cache2)
+        grad_w2 = torch.zeros_like(w2)
+        grad_topk_weights = torch.zeros_like(topk_weights)
+
+        # Process in chunks to match forward pass
+        for chunk in range((num_tokens // CHUNK_SIZE) + 1):
+            begin_chunk_idx, end_chunk_idx = (
+                chunk * CHUNK_SIZE,
+                min((chunk + 1) * CHUNK_SIZE, num_tokens),
+            )
+
+            curr_num_tokens = end_chunk_idx - begin_chunk_idx
+            if curr_num_tokens == 0:
+                continue
+
+            curr_intermediate_cache2 = intermediate_cache2[begin_chunk_idx * topk : end_chunk_idx * topk]
+            curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
+            curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
+            curr_grad_output = grad_output[begin_chunk_idx:end_chunk_idx]
+
+            # Get aligned metadata
+            sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+                curr_topk_ids, config["BLOCK_SIZE_M"], E
+            )
+
+            # Prepare gradient buffers for this chunk
+            curr_grad_intermediate_cache2 = torch.zeros_like(curr_intermediate_cache2)
+            curr_grad_w2 = torch.zeros_like(w2)
+            curr_grad_topk_weights = torch.zeros_like(curr_topk_weights)
+
+            # Call Triton backward kernel with MUL_ROUTED_WEIGHT=True
+            # Note: Use top_k=1 to match forward pass indexing
+            invoke_fused_moe_backward_kernel(
+                grad_output=curr_grad_output,
+                input=curr_intermediate_cache2,
+                weight=w2,
+                grad_input=curr_grad_intermediate_cache2,
+                grad_weight=curr_grad_w2,
+                grad_topk_weights=curr_grad_topk_weights,
+                topk_weights=curr_topk_weights,
+                topk_ids=curr_topk_ids,
+                sorted_token_ids=sorted_token_ids,
+                expert_ids=expert_ids,
+                num_tokens_post_padded=num_tokens_post_padded,
+                mul_routed_weight=True,
+                top_k=1,
+                config=config,
+                compute_type=tl.bfloat16,
+            )
+
+            # Accumulate gradients
+            grad_intermediate_cache2[begin_chunk_idx * topk : end_chunk_idx * topk] = curr_grad_intermediate_cache2
+            grad_w2 += curr_grad_w2
+            grad_topk_weights[begin_chunk_idx:end_chunk_idx] = curr_grad_topk_weights
+
+        return grad_intermediate_cache2, grad_w2, grad_topk_weights, None
+
+
+class MoeSumReduceFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        intermediate_cache3: torch.Tensor,
+        hidden_states_shape,
+    ):
+        out_hidden_states = torch.empty(
+            hidden_states_shape, device=intermediate_cache3.device, dtype=intermediate_cache3.dtype
+        )
+        moe_sum_reduce(
+            intermediate_cache3,
+            out_hidden_states,
+            1.0,
+        )
+        ctx.save_for_backward(intermediate_cache3)
+        return out_hidden_states
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (intermediate_cache3,) = ctx.saved_tensors
+        return grad_output.unsqueeze(1).expand_as(intermediate_cache3), None

slime/backends/fsdp_utils/kernels/fused_moe_triton_backward_kernels.py

@@ -0,0 +1,543 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def fused_moe_backward_input_kernel(
+    # Pointers to matrices
+    grad_output_ptr,
+    weight_ptr,
+    grad_input_ptr,
+    grad_topk_weights_ptr,
+    topk_weights_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    # Matrix dimensions
+    N,
+    K,
+    EM,
+    num_valid_tokens,
+    # Strides
+    stride_gom,
+    stride_gon,
+    stride_we,
+    stride_wn,
+    stride_wk,
+    stride_gim,
+    stride_gik,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    MUL_ROUTED_WEIGHT: tl.constexpr,
+    top_k: tl.constexpr,
+    compute_type: tl.constexpr,
+):
+    """
+    Backward kernel for computing grad_input.
+
+    Forward: output = input @ weight.T (optionally multiplied by topk_weights)
+    Backward: grad_input = grad_output @ weight (optionally multiplied by topk_weights)
+
+    This kernel computes: grad_input[token] = sum_over_N(grad_output[token, n] * weight[expert, n, :])
+    If MUL_ROUTED_WEIGHT: grad_input[token] *= topk_weights[token]
+
+    Parallelization: Similar to forward, parallel over M and N dimensions, loop over K.
+    """
+    # Map program ids to blocks (parallel over M and N, similar to forward)
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    # Check bounds
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+
+    # Only process if this block is valid
+    if pid_m * BLOCK_SIZE_M < num_tokens_post_padded:
+        # Load token information
+        offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+        offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
+        offs_token = offs_token.to(tl.int64)
+        token_mask = offs_token < num_valid_tokens
+
+        # Get expert ID for this block
+        off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
+
+        # Only process if expert is valid
+        if off_experts != -1:
+            # Initialize offsets for N dimension (current block)
+            offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
+            offs_k = tl.arange(0, BLOCK_SIZE_K)
+
+            # Load grad_output block: shape (BLOCK_SIZE_M, BLOCK_SIZE_N)
+            grad_output_ptrs = grad_output_ptr + (offs_token[:, None] * stride_gom + offs_n[None, :] * stride_gon)
+            grad_out = tl.load(
+                grad_output_ptrs,
+                mask=token_mask[:, None] & (offs_n[None, :] < N),
+                other=0.0,
+            )
+
+            # Apply topk_weights to grad_output if needed
+            if MUL_ROUTED_WEIGHT:
+                moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
+                grad_out = grad_out * moe_weight[:, None]
+
+            # Iterate over K dimension
+            for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+                # Current K offsets
+                curr_offs_k = k * BLOCK_SIZE_K + offs_k
+
+                # Load weight block: shape (BLOCK_SIZE_N, BLOCK_SIZE_K)
+                # weight: shape (E, N, K)
+                weight_ptrs = (
+                    weight_ptr
+                    + off_experts * stride_we
+                    + offs_n[:, None] * stride_wn
+                    + curr_offs_k[None, :] * stride_wk
+                )
+                w = tl.load(
+                    weight_ptrs,
+                    mask=(offs_n[:, None] < N) & (curr_offs_k[None, :] < K),
+                    other=0.0,
+                )
+
+                # Compute contribution: grad_out @ weight
+                # grad_out: (BLOCK_SIZE_M, BLOCK_SIZE_N)
+                # w: (BLOCK_SIZE_N, BLOCK_SIZE_K)
+                # result: (BLOCK_SIZE_M, BLOCK_SIZE_K)
+                contribution = tl.dot(grad_out, w)
+
+                # Atomic add to grad_input because different N blocks contribute to same K
+                grad_input_ptrs = grad_input_ptr + (
+                    (offs_token[:, None] // top_k) * stride_gim + curr_offs_k[None, :] * stride_gik
+                )
+                grad_input_mask = token_mask[:, None] & (curr_offs_k[None, :] < K)
+                tl.atomic_add(grad_input_ptrs, contribution.to(compute_type), mask=grad_input_mask)
+
+
+@triton.jit
+def fused_moe_backward_weight_kernel(
+    # Pointers to matrices
+    grad_output_ptr,
+    input_ptr,
+    grad_weight_ptr,
+    topk_weights_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    # Matrix dimensions
+    N,
+    K,
+    EM,
+    num_valid_tokens,
+    # Strides
+    stride_gom,
+    stride_gon,
+    stride_im,
+    stride_ik,
+    stride_gwe,
+    stride_gwn,
+    stride_gwk,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    MUL_ROUTED_WEIGHT: tl.constexpr,
+    top_k: tl.constexpr,
+    compute_type: tl.constexpr,
+):
+    """
+    Backward kernel for computing grad_weight.
+
+    Forward: output = input @ weight.T (optionally multiplied by topk_weights)
+    Backward: grad_weight = input.T @ grad_output (optionally multiplied by topk_weights)
+
+    This kernel computes: grad_weight[expert, n, k] = sum_over_tokens(input[token, k] * grad_output[token, n])
+    If MUL_ROUTED_WEIGHT: the accumulation is weighted by topk_weights[token]
+
+    Parallelization: Parallel over M and N dimensions with grouping, loop over K.
+    """
+    # Map program ids to blocks (parallel over M and N with grouping, similar to forward and backward_input)
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    # Check bounds
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+
+    # Only process if this block is valid
+    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
+        return
+
+    # Get expert ID for this M block
+    expert_id = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
+
+    # Only process if expert is valid
+    if expert_id == -1:
+        return
+
+    # Load token information for this M block
+    offs_m = tl.arange(0, BLOCK_SIZE_M)
+    offs_token_id = pid_m * BLOCK_SIZE_M + offs_m.to(tl.int64)
+    offs_token = tl.load(
+        sorted_token_ids_ptr + offs_token_id, mask=offs_token_id < num_tokens_post_padded, other=num_valid_tokens
+    )
+    offs_token = offs_token.to(tl.int64)
+    token_mask = (offs_token_id < num_tokens_post_padded) & (offs_token < num_valid_tokens)
+
+    # Clamp offs_token to valid range
+    offs_token_clamped = tl.where(token_mask, offs_token, 0)
+
+    # Determine input token indices based on MUL_ROUTED_WEIGHT
+    if MUL_ROUTED_WEIGHT:
+        input_token_idx = offs_token_clamped
+        input_mask = token_mask
+    else:
+        input_token_idx = offs_token_clamped // top_k
+        num_input_tokens = num_valid_tokens // top_k
+        input_mask = token_mask & (input_token_idx < num_input_tokens)
+
+    # Load topk_weights if needed
+    if MUL_ROUTED_WEIGHT:
+        moe_weight = tl.load(topk_weights_ptr + offs_token_clamped, mask=token_mask, other=0.0)
+
+    # Current N offset for this program
+    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
+
+    # Load grad_output for this N block: shape (M, BLOCK_SIZE_N)
+    # grad_output is always indexed by sorted_token_ids (offs_token_clamped)
+    # because it has shape (num_tokens * topk, N)
+    grad_output_ptrs = grad_output_ptr + (offs_token_clamped[:, None] * stride_gom + offs_n[None, :] * stride_gon)
+    grad_out = tl.load(
+        grad_output_ptrs,
+        mask=token_mask[:, None] & (offs_n[None, :] < N),
+        other=0.0,
+    )
+
+    # Apply topk_weights if needed
+    if MUL_ROUTED_WEIGHT:
+        grad_out = grad_out * moe_weight[:, None]
+
+    # Zero out padding tokens
+    token_mask_col = token_mask[:, None]
+    grad_out = grad_out * token_mask_col
+
+    # Iterate over K blocks and accumulate
+    for k_block in range(tl.cdiv(K, BLOCK_SIZE_K)):
+        offs_k = k_block * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64)
+
+        # Load input for this K block
+        input_ptrs = input_ptr + (input_token_idx[:, None] * stride_im + offs_k[None, :] * stride_ik)
+        inp = tl.load(
+            input_ptrs,
+            mask=input_mask[:, None] & (offs_k[None, :] < K),
+            other=0.0,
+        )
+
+        # Zero out padding tokens - use input_mask for input, token_mask for grad_output
+        input_mask_col = input_mask[:, None]
+        inp = inp * input_mask_col
+
+        # Compute grad_weight contribution: grad_out.T @ inp
+        grad_w_contribution = tl.dot(grad_out.T, inp)
+
+        # Write back using atomic add
+        grad_weight_ptrs = (
+            grad_weight_ptr + expert_id * stride_gwe + offs_n[:, None] * stride_gwn + offs_k[None, :] * stride_gwk
+        )
+        grad_weight_mask = (offs_n[:, None] < N) & (offs_k[None, :] < K)
+        tl.atomic_add(grad_weight_ptrs, grad_w_contribution.to(compute_type), mask=grad_weight_mask)
+
+
+@triton.jit
+def fused_moe_backward_topk_weights_kernel(
+    # Pointers to matrices
+    grad_output_ptr,
+    input_ptr,
+    weight_ptr,
+    grad_topk_weights_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    # Matrix dimensions
+    N,
+    K,
+    EM,
+    num_valid_tokens,
+    # Strides
+    stride_gom,
+    stride_gon,
+    stride_im,
+    stride_ik,
+    stride_we,
+    stride_wn,
+    stride_wk,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    top_k: tl.constexpr,
+    compute_type: tl.constexpr,
+):
+    """
+    Backward kernel for computing grad_topk_weights.
+
+    Forward: output = topk_weights * (input @ weight.T)
+    Backward: grad_topk_weights = sum(grad_output * (input @ weight.T))
+
+    This kernel computes the gradient of topk_weights by computing the dot product
+    of grad_output with the forward output before weight multiplication.
+    """
+    # Map program id to token block
+    pid = tl.program_id(axis=0)
+
+    # Check bounds
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+
+    # Only process if this block is valid
+    if pid * BLOCK_SIZE_M < num_tokens_post_padded:
+        # Load token information
+        offs_token_id = pid * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+        offs_token = tl.load(
+            sorted_token_ids_ptr + offs_token_id, mask=offs_token_id < num_tokens_post_padded, other=num_valid_tokens
+        )
+        offs_token = offs_token.to(tl.int64)
+        token_mask = (offs_token_id < num_tokens_post_padded) & (offs_token < num_valid_tokens)
+
+        # Clamp offs_token to valid range for safe pointer arithmetic
+        offs_token_clamped = tl.where(token_mask, offs_token, 0)
+
+        # Get expert ID for this block
+        off_experts = tl.load(expert_ids_ptr + pid).to(tl.int64)
+
+        # Only process if expert is valid
+        if off_experts != -1:
+            # Initialize offsets
+            offs_n = tl.arange(0, BLOCK_SIZE_N)
+            offs_k = tl.arange(0, BLOCK_SIZE_K)
+
+            # Accumulator for grad_topk_weights
+            accumulator = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)
+
+            # Iterate over N and K dimensions to compute forward output and gradient
+            for n in range(0, tl.cdiv(N, BLOCK_SIZE_N)):
+                # Current N offset
+                curr_offs_n = n * BLOCK_SIZE_N + offs_n
+
+                # Load grad_output block: (M, N)
+                grad_output_ptrs = grad_output_ptr + (
+                    offs_token_clamped[:, None] * stride_gom + curr_offs_n[None, :] * stride_gon
+                )
+                grad_out = tl.load(
+                    grad_output_ptrs,
+                    mask=token_mask[:, None] & (curr_offs_n[None, :] < N),
+                    other=0.0,
+                )
+
+                # Compute forward output for this N block: input @ weight[:, n, :].T
+                forward_output_n = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+
+                for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+                    # Current K offset
+                    curr_offs_k = k * BLOCK_SIZE_K + offs_k
+
+                    # Load input block: (M, K)
+                    input_ptrs = input_ptr + (
+                        (offs_token_clamped[:, None] // top_k) * stride_im + curr_offs_k[None, :] * stride_ik
+                    )
+                    inp = tl.load(
+                        input_ptrs,
+                        mask=token_mask[:, None] & (curr_offs_k[None, :] < K),
+                        other=0.0,
+                    )
+
+                    # Load weight block: (N, K)
+                    weight_ptrs = (
+                        weight_ptr
+                        + off_experts * stride_we
+                        + curr_offs_n[:, None] * stride_wn
+                        + curr_offs_k[None, :] * stride_wk
+                    )
+                    w = tl.load(
+                        weight_ptrs,
+                        mask=(curr_offs_n[:, None] < N) & (curr_offs_k[None, :] < K),
+                        other=0.0,
+                    )
+
+                    # Accumulate forward output: input @ weight.T
+                    # inp: (M, K), w.T: (K, N) -> (M, N)
+                    forward_output_n += tl.dot(inp, w.T)
+
+                # Compute contribution to grad_topk_weights: sum(grad_out * forward_output)
+                # Sum over N dimension
+                accumulator += tl.sum(grad_out * forward_output_n, axis=1)
+
+            # Write back grad_topk_weights using atomic add with clamped token indices
+            tl.atomic_add(grad_topk_weights_ptr + offs_token_clamped, accumulator.to(compute_type), mask=token_mask)
+
+
+def invoke_fused_moe_backward_kernel(
+    grad_output: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    grad_input: torch.Tensor,
+    grad_weight: torch.Tensor,
+    grad_topk_weights: torch.Tensor | None,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    mul_routed_weight: bool,
+    top_k: int,
+    config: dict[str, Any],
+    compute_type: tl.dtype,
+) -> None:
+    """
+    Invoke the fused MOE backward kernels to compute gradients.
+
+    Args:
+        grad_output: Gradient of output, shape (num_tokens * topk, N) or (num_tokens, topk, N)
+        input: Input tensor, shape (num_tokens, K)
+        weight: Weight tensor, shape (E, N, K)
+        grad_input: Output gradient for input, shape (num_tokens, K)
+        grad_weight: Output gradient for weight, shape (E, N, K)
+        grad_topk_weights: Output gradient for topk_weights, shape (num_tokens, topk) or None
+        topk_weights: Top-K routing weights, shape (num_tokens, topk)
+        topk_ids: Top-K expert IDs, shape (num_tokens, topk)
+        sorted_token_ids: Sorted token IDs
+        expert_ids: Expert IDs for each block
+        num_tokens_post_padded: Number of tokens after padding
+        mul_routed_weight: Whether to multiply by routing weights
+        top_k: Number of experts per token
+        config: Kernel configuration
+        compute_type: Computation data type
+    """
+    assert topk_weights.stride(1) == 1
+    assert sorted_token_ids.stride(0) == 1
+
+    # Flatten grad_output if needed
+    # Before: (num_tokens, topk, hidden_size)
+    # After: (num_tokens * topk, hidden_size)
+    if grad_output.ndim == 3:
+        grad_output = grad_output.reshape(-1, grad_output.shape[-1])
+
+    E, N, K = weight.shape
+
+    # ===================== Compute grad_input =====================
+    def grid_input(META):
+        return (triton.cdiv(sorted_token_ids.shape[0], META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),)
+
+    fused_moe_backward_input_kernel[grid_input](
+        grad_output,
+        weight,
+        grad_input,
+        grad_topk_weights if grad_topk_weights is not None else grad_input,  # dummy pointer
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        N,
+        K,
+        sorted_token_ids.shape[0],
+        grad_output.shape[0],
+        grad_output.stride(0),
+        grad_output.stride(1),
+        weight.stride(0),
+        weight.stride(1),
+        weight.stride(2),
+        grad_input.stride(0),
+        grad_input.stride(1),
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        top_k=top_k,
+        compute_type=compute_type,
+        **config,
+    )
+
+    # ===================== Compute grad_weight =====================
+    # Initialize grad_weight to zero
+    grad_weight.zero_()
+
+    # Use same grid configuration as forward kernel: encode both M and N dimensions
+    def grid_weight(META):
+        return (triton.cdiv(sorted_token_ids.shape[0], META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),)
+
+    fused_moe_backward_weight_kernel[grid_weight](
+        grad_output,
+        input,
+        grad_weight,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        N,
+        K,
+        sorted_token_ids.shape[0],
+        grad_output.shape[0],
+        grad_output.stride(0),
+        grad_output.stride(1),
+        input.stride(0),
+        input.stride(1),
+        grad_weight.stride(0),
+        grad_weight.stride(1),
+        grad_weight.stride(2),
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        top_k=top_k,
+        compute_type=compute_type,
+        **config,
+    )
+
+    # ===================== Compute grad_topk_weights (if needed) =====================
+    if mul_routed_weight and grad_topk_weights is not None:
+
+        def grid_topk(META):
+            return (triton.cdiv(sorted_token_ids.shape[0], META["BLOCK_SIZE_M"]),)
+
+        fused_moe_backward_topk_weights_kernel[grid_topk](
+            grad_output,
+            input,
+            weight,
+            grad_topk_weights.view(-1),
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            N,
+            K,
+            sorted_token_ids.shape[0],
+            grad_output.shape[0],
+            grad_output.stride(0),
+            grad_output.stride(1),
+            input.stride(0),
+            input.stride(1),
+            weight.stride(0),
+            weight.stride(1),
+            weight.stride(2),
+            top_k=top_k,
+            compute_type=compute_type,
+            BLOCK_SIZE_M=config["BLOCK_SIZE_M"],
+            BLOCK_SIZE_N=config["BLOCK_SIZE_N"],
+            BLOCK_SIZE_K=config["BLOCK_SIZE_K"],
+        )

slime/backends/fsdp_utils/lr_scheduler.py

@@ -0,0 +1,197 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+"""Learning rate scheduler for FSDP training."""
+
+import logging
+import math
+
+import torch
+from torch.optim.lr_scheduler import LRScheduler
+from typing_extensions import override
+
+logger = logging.getLogger(__name__)
+
+
+class FSDPLRScheduler(LRScheduler):
+    """Learning rate scheduler for FSDP training.
+
+    Args:
+        optimizer (torch.optim.Optimizer): The optimizer to be used.
+        init_lr (float): Initial learning rate.
+        max_lr (float): Maximum learning rate.
+        min_lr (float): Minimum learning rate.
+        lr_warmup_steps (int): Number of warmup steps.
+        lr_decay_steps (int): Number of decay steps.
+        lr_decay_style (str): Decay style for learning rate.
+        use_checkpoint_lr_scheduler (bool, optional): Whether to use the checkpoint values
+            for the lr scheduler.
+        override_lr_scheduler (bool, optional): Whether to override the lr scheduler values
+            with the class values.
+        wsd_decay_steps (int, optional): Number of weight decay decay steps.
+        lr_wsd_decay_style (str, optional): Decay style for learning rate during weight decay decay
+            steps.
+        last_epoch (int, optional): The index of last epoch. Default: -1.
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        init_lr: float,
+        max_lr: float,
+        min_lr: float,
+        lr_warmup_steps: int,
+        lr_decay_steps: int,
+        lr_decay_style: str,
+        use_checkpoint_lr_scheduler: bool | None = True,
+        override_lr_scheduler: bool | None = False,
+        wsd_decay_steps: int | None = None,
+        lr_wsd_decay_style: str | None = None,
+        last_epoch: int = -1,
+    ) -> None:
+        # Store our custom parameters
+        self.init_lr = init_lr
+        self.max_lr = float(max_lr)
+        self.min_lr = min_lr
+        assert self.min_lr >= 0.0
+        assert self.max_lr >= self.min_lr
+        assert self.init_lr <= self.max_lr
+
+        self.lr_warmup_steps = lr_warmup_steps
+        self.lr_decay_steps = lr_decay_steps
+        self.wsd_decay_steps = wsd_decay_steps
+        self.lr_wsd_decay_style = lr_wsd_decay_style
+
+        assert self.lr_decay_steps > 0
+        assert self.lr_warmup_steps < self.lr_decay_steps
+
+        self.lr_decay_style = lr_decay_style
+        if self.lr_decay_style == "WSD":
+            assert self.wsd_decay_steps is not None
+
+        self.override_lr_scheduler = override_lr_scheduler
+        self.use_checkpoint_lr_scheduler = use_checkpoint_lr_scheduler
+
+        if self.override_lr_scheduler:
+            assert not self.use_checkpoint_lr_scheduler, "both override and use-checkpoint are set."
+
+        # Initialize parent class
+        super().__init__(optimizer, last_epoch)
+
+        logger.info(f"> learning rate decay style: {self.lr_decay_style}")
+
+    def _get_lr_for_group(self, param_group: dict) -> float:
+        """Compute learning rate for a specific parameter group.
+
+        Args:
+            param_group (dict): parameter group from the optimizer.
+
+        Returns:
+            float: learning rate for this parameter group.
+        """
+        max_lr = param_group.get("max_lr", self.max_lr)
+        min_lr = param_group.get("min_lr", self.min_lr)
+
+        # Use linear warmup for the initial part.
+        if self.lr_warmup_steps > 0 and self.last_epoch <= self.lr_warmup_steps:
+            return self.init_lr + ((max_lr - self.init_lr) * float(self.last_epoch) / float(self.lr_warmup_steps))
+
+        # If the learning rate is constant, just return the initial value.
+        if self.lr_decay_style == "constant":
+            return max_lr
+
+        # For any steps larger than `self.lr_decay_steps`, use `min_lr`.
+        if self.last_epoch > self.lr_decay_steps:
+            return min_lr
+
+        # If we are done with the warmup period, use the decay style.
+        if self.lr_decay_style == "inverse-square-root":
+            warmup_steps = max(self.lr_warmup_steps, 1)
+            num_steps = max(self.last_epoch, 1)
+            lr = max_lr * warmup_steps**0.5 / (num_steps**0.5)
+            return max(min_lr, lr)
+
+        num_steps_ = self.last_epoch - self.lr_warmup_steps
+        decay_steps_ = self.lr_decay_steps - self.lr_warmup_steps
+        decay_ratio = float(num_steps_) / float(decay_steps_)
+        assert decay_ratio >= 0.0
+        assert decay_ratio <= 1.0
+
+        delta_lr = max_lr - min_lr
+        coeff = None
+
+        if self.lr_decay_style == "linear":
+            coeff = 1.0 - decay_ratio
+        elif self.lr_decay_style == "cosine":
+            coeff = 0.5 * (math.cos(math.pi * decay_ratio) + 1.0)
+        elif self.lr_decay_style == "WSD":
+            wsd_anneal_start_ = self.lr_decay_steps - self.wsd_decay_steps
+            if self.last_epoch <= wsd_anneal_start_:
+                coeff = 1.0
+            else:
+                wsd_steps = self.last_epoch - wsd_anneal_start_
+                wsd_decay_ratio = float(wsd_steps) / float(self.wsd_decay_steps)
+                if self.lr_wsd_decay_style == "linear":
+                    coeff = 1.0 - wsd_decay_ratio
+                elif self.lr_wsd_decay_style == "cosine":
+                    coeff = 0.5 * (math.cos(math.pi * wsd_decay_ratio) + 1.0)
+                elif self.lr_wsd_decay_style == "exponential":
+                    coeff = (2.0 * math.pow(0.5, wsd_decay_ratio)) - 1.0
+                elif self.lr_wsd_decay_style == "minus_sqrt":
+                    coeff = 1.0 - math.sqrt(wsd_decay_ratio)
+        else:
+            raise Exception(f"{self.lr_decay_style} decay style is not supported.")
+
+        assert coeff is not None
+        return min_lr + coeff * delta_lr
+
+    @override
+    def get_lr(self) -> list[float]:
+        """Compute the learning rates for each parameter group.
+
+        Returns:
+            list[float]: A list of learning rates, one for each parameter group.
+        """
+        return [self._get_lr_for_group(group) for group in self.optimizer.param_groups]
+
+
+def get_lr_scheduler(args, optimizer: torch.optim.Optimizer) -> FSDPLRScheduler:
+    """Create and configure the learning-rate scheduler.
+
+    This configures iteration-based schedules derived from the global batch size
+    and run-time arguments.
+
+    Args:
+        args: Training/runtime arguments (namespace).
+        optimizer (torch.optim.Optimizer): Optimizer bound to the model.
+
+    Returns:
+        FSDPLRScheduler: Initialized scheduler bound to ``optimizer``.
+    """
+    args.train_iters = args.num_rollout * args.rollout_batch_size * args.n_samples_per_prompt // args.global_batch_size
+    if args.lr_decay_iters is None:
+        args.lr_decay_iters = args.train_iters
+    lr_decay_steps = args.lr_decay_iters
+    wsd_decay_steps = None
+    if args.lr_wsd_decay_iters is not None:
+        wsd_decay_steps = args.lr_wsd_decay_iters
+    if args.lr_warmup_fraction is not None:
+        lr_warmup_steps = args.lr_warmup_fraction * lr_decay_steps
+    else:
+        lr_warmup_steps = args.lr_warmup_iters
+    lr_scheduler = FSDPLRScheduler(
+        optimizer,
+        init_lr=args.lr_warmup_init,
+        max_lr=args.lr,
+        min_lr=args.min_lr,
+        lr_warmup_steps=lr_warmup_steps,
+        lr_decay_steps=lr_decay_steps,
+        lr_decay_style=args.lr_decay_style,
+        use_checkpoint_lr_scheduler=args.use_checkpoint_lr_scheduler,
+        override_lr_scheduler=args.override_lr_scheduler,
+        wsd_decay_steps=wsd_decay_steps,
+        lr_wsd_decay_style=args.lr_wsd_decay_style,
+    )
+
+    return lr_scheduler

slime/backends/fsdp_utils/models/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+

slime/backends/fsdp_utils/models/qwen3_moe.py

@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeMLP
+
+from slime.backends.fsdp_utils.kernels.fused_experts import (
+    DownProjFunction,
+    GateUpProjFunction,
+    MoeSumReduceFunction,
+    SiluAndMulFunction,
+)
+
+
+def fused_experts_impl(
+    hidden_states: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+):
+    assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
+    assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
+    assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
+    assert w1.is_contiguous(), "Expert weights1 must be contiguous"
+    assert w2.is_contiguous(), "Expert weights2 must be contiguous"
+    assert hidden_states.dtype in [torch.bfloat16]
+
+    intermediate_cache1 = GateUpProjFunction.apply(
+        hidden_states,
+        w1,
+        topk_weights,
+        topk_ids,
+    )
+    intermediate_cache2 = SiluAndMulFunction.apply(intermediate_cache1)
+    intermediate_cache3 = DownProjFunction.apply(
+        intermediate_cache2,
+        w2,
+        topk_weights,
+        topk_ids,
+    )
+    output_hidden_states = MoeSumReduceFunction.apply(
+        intermediate_cache3,
+        hidden_states.shape,
+    )
+    return output_hidden_states
+
+
+class StandardDispatcher:
+    def __init__(self, num_experts: int, num_local_experts: int):
+        self.moe_ep_size = 1
+        self.num_experts = num_experts
+        self.num_local_experts = num_local_experts
+        self.moe_ep_rank = 0
+        self.local_expert_mapping = None
+
+        if self.moe_ep_size > 1:
+            self.local_expert_mapping = torch.full((self.num_experts,), -1, dtype=torch.int32, device="cuda")
+            self.local_expert_mapping[
+                self.moe_ep_rank * self.num_local_experts : (self.moe_ep_rank + 1) * self.num_local_experts
+            ] = torch.arange(0, self.num_local_experts, dtype=torch.int32, device="cuda")
+
+    def dispatch(self, topk_ids) -> torch.Tensor:
+        if self.local_expert_mapping is not None:
+            return self.local_expert_mapping[topk_ids]
+        return topk_ids
+
+
+class Qwen3MoeSparseMoeBlock(nn.Module):
+    dispatcher = None
+    runner = None
+
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.top_k = config.num_experts_per_tok
+        self.norm_topk_prob = config.norm_topk_prob
+
+        # gating
+        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
+
+        self.experts = nn.ModuleList(
+            [Qwen3MoeMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)]
+        )
+
+        if Qwen3MoeSparseMoeBlock.dispatcher is None:
+            Qwen3MoeSparseMoeBlock.dispatcher = StandardDispatcher(
+                num_experts=config.num_experts, num_local_experts=config.num_experts
+            )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        # router_logits: (batch * sequence_length, n_experts)
+        router_logits = self.gate(hidden_states)
+
+        routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
+        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+
+        if self.norm_topk_prob:  # only diff with mixtral sparse moe block!
+            routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        # we cast back to the input dtype
+        routing_weights = routing_weights.to(hidden_states.dtype)
+
+        selected_experts = Qwen3MoeSparseMoeBlock.dispatcher.dispatch(selected_experts)
+
+        w13_weight = torch.stack(
+            [torch.cat([layer.gate_proj.weight, layer.up_proj.weight], dim=0) for layer in self.experts]
+        )
+        w2_weight = torch.stack([layer.down_proj.weight for layer in self.experts], dim=0)
+
+        final_hidden_states = fused_experts_impl(
+            hidden_states.to(torch.bfloat16),
+            w13_weight,
+            w2_weight,
+            routing_weights,
+            selected_experts,
+        )
+
+        return final_hidden_states, router_logits
+
+
+def apply_true_on_policy_patch_for_qwen3_moe():
+    from transformers.models.qwen3_moe import modeling_qwen3_moe
+
+    modeling_qwen3_moe.Qwen3MoeSparseMoeBlock = Qwen3MoeSparseMoeBlock

slime/backends/fsdp_utils/models/qwen3_moe_hf.py

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.nn.functional as F
+
+
+def apply_fsdp_moe_patch():
+
+    from transformers.models.qwen3_moe import modeling_qwen3_moe
+
+    def _forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        router_logits = self.gate(hidden_states)
+
+        routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
+        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+        if self.norm_topk_prob:
+            routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights.to(hidden_states.dtype)
+
+        final_hidden_states = torch.zeros(
+            (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
+        )
+
+        expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
+
+        # Loop over all experts
+        for expert_idx in range(self.num_experts):
+            expert_layer = self.experts[expert_idx]
+            idx, top_x = torch.where(expert_mask[expert_idx])
+
+            if top_x.numel() > 0:
+                current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
+                current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]
+                final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
+            else:
+                # force experts to participate in computation graph
+                dummy_output = expert_layer(hidden_states[:1]) * 0.0
+                final_hidden_states[:1] = final_hidden_states[:1] + dummy_output
+
+        final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+        return final_hidden_states, router_logits
+
+    modeling_qwen3_moe.Qwen3MoeSparseMoeBlock.forward = _forward

slime/backends/fsdp_utils/update_weight_utils.py

@@ -0,0 +1,261 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import abc
+import logging
+import socket
+from argparse import Namespace
+from collections.abc import Sequence
+
+import ray
+import torch
+import torch.distributed as dist
+from ray.actor import ActorHandle
+from torch.distributed.tensor import DTensor, Replicate
+
+try:
+    from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions  # type: ignore[import]
+except ImportError:
+    from sglang.srt.patch_torch import monkey_patch_torch_reductions  # type: ignore[import]
+
+from sglang.srt.utils import MultiprocessingSerializer
+
+from slime.utils.distributed_utils import init_process_group
+
+
+try:
+    from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket  # type: ignore[import]
+except ImportError:
+    from sglang.srt.model_executor.model_runner import FlattenedTensorBucket  # type: ignore[import]
+
+
+logger = logging.getLogger(__name__)
+
+
+class UpdateWeight(abc.ABC):
+    def __init__(self, args: Namespace, model: torch.nn.Module) -> None:
+        self.args = args
+        self.model = model
+        self.weight_version = 0
+
+    @abc.abstractmethod
+    def connect_rollout_engines(
+        self,
+        rollout_engines: Sequence[ActorHandle],
+        rollout_engine_lock: ActorHandle | None,
+    ) -> None:
+        pass
+
+    def update_weights(self) -> None:
+        self.weight_version += 1
+        bucket = []
+        bucket_size = 0
+        for name, param in self.model.state_dict().items():
+            param_size = param.numel() * param.element_size()
+            if bucket and bucket_size + param_size >= self.args.update_weight_buffer_size:
+                self.wait_and_update_bucket_weights(bucket)
+                del bucket
+                bucket = []
+                bucket_size = 0
+
+            param = param.cuda()
+            if isinstance(param, DTensor):
+                # async version of param.full_tensor
+                param = param.redistribute(
+                    placements=[Replicate()] * param.device_mesh.ndim,
+                    async_op=True,
+                ).to_local()
+            bucket.append((name, param))
+            bucket_size += param_size
+
+        if bucket:
+            self.wait_and_update_bucket_weights(bucket)
+            del bucket
+            bucket = []
+            bucket_size = 0
+
+    def wait_and_update_bucket_weights(self, bucket):
+        bucket = [(name, param.wait()) if hasattr(param, "wait") else (name, param) for name, param in bucket]
+        self.update_bucket_weights(bucket, weight_version=self.weight_version)
+
+    @abc.abstractmethod
+    def update_bucket_weights(self, named_tensors, weight_version=None) -> None:
+        pass
+
+
+class UpdateWeightFromTensor(UpdateWeight):
+    """Push model weights to rollout engines using tensors.
+
+    Streams parameters in size-bounded buckets; optionally groups tensors by dtype
+    and flattens per dtype, gathers per-rank blobs to the source, and issues one
+    RPC per dtype per bucket (or one per bucket if not flattened).
+    """
+
+    def connect_rollout_engines(
+        self,
+        rollout_engines: Sequence[ActorHandle],
+        rollout_engine_lock: ActorHandle | None,
+    ) -> None:
+        """Attach rollout engines and create per-engine IPC (Gloo) groups.
+
+        Sets the gather source rank, engine handle, and `tp_rank` within the
+        engine's local group.
+        """
+        self.rollout_engines = rollout_engines
+
+        # Here we assume the gpu id of rollout engines and train actors are the same.
+        for i, engine in enumerate(self.rollout_engines):
+            start_rank = i * self.args.rollout_num_gpus_per_engine
+            end_rank = (i + 1) * self.args.rollout_num_gpus_per_engine
+            group_ranks = list(range(start_rank, end_rank))
+            new_group = dist.new_group(
+                ranks=group_ranks,
+                backend="gloo",
+            )
+            if dist.get_rank() in group_ranks:
+                self._ipc_gather_src = start_rank
+                self._ipc_gather_group = new_group
+                self._ipc_engine = engine
+                # Calculate TP rank within this SGLang engine group
+                self.tp_rank = dist.get_rank() - start_rank
+
+    def update_bucket_weights(self, named_tensors, weight_version=None) -> None:
+        monkey_patch_torch_reductions()
+        # Use flattened bucket approach similar to Megatron
+        logger.info("Using flattened tensor bucket")
+        # Group tensors by dtype (same as Megatron)
+        named_tensors_by_dtypes = {}
+        for name, tensor in named_tensors:
+            dtype = tensor.dtype
+            if dtype not in named_tensors_by_dtypes:
+                named_tensors_by_dtypes[dtype] = []
+            named_tensors_by_dtypes[dtype].append((name, tensor))
+
+        # Create flattened bucket for each dtype group
+        serialized_tensors = []
+        for _dtype, named_tensors in named_tensors_by_dtypes.items():
+            flattened_tensor_bucket = FlattenedTensorBucket(named_tensors=named_tensors)
+            metadata = flattened_tensor_bucket.get_metadata()
+            flattened_tensor_data = {
+                "flattened_tensor": flattened_tensor_bucket.get_flattened_tensor(),
+                "metadata": metadata,
+            }
+            serialized_tensors.append(MultiprocessingSerializer.serialize(flattened_tensor_data, output_str=True))
+
+        if self._ipc_gather_src == dist.get_rank():
+            # On rank 0, prepare a list to hold the gathered batches from all ranks.
+            gathered_serialized_batches = [None for _ in range(dist.get_world_size(self._ipc_gather_group))]
+        else:
+            gathered_serialized_batches = None
+
+        # Gather the serialized batches from all ranks to rank 0.
+        dist.gather_object(
+            obj=serialized_tensors,
+            object_gather_list=gathered_serialized_batches,
+            dst=self._ipc_gather_src,
+            group=self._ipc_gather_group,
+        )
+
+        if dist.get_rank() == self._ipc_gather_src:
+            # Handle flattened bucket format (same as Megatron approach)
+            # Each rank may have multiple dtype buckets
+            # TODO: here we assume all ranks have the same number of dtypes
+            num_dtypes = len(gathered_serialized_batches[0])
+            assert num_dtypes > 0
+            for i in range(num_dtypes):
+                kwargs = {
+                    "serialized_named_tensors": [tensors[i] for tensors in gathered_serialized_batches],
+                    "load_format": "flattened_bucket",
+                    "flush_cache": False,
+                    "weight_version": str(weight_version),
+                }
+                ref = self._ipc_engine.update_weights_from_tensor.remote(**kwargs)
+                ray.get(ref)
+
+        if dist.get_rank() == self._ipc_gather_src:
+            ref = self._ipc_engine.flush_cache.remote()
+            ray.get(ref)
+
+
+class UpdateWeightFromDistributed(UpdateWeight):
+    """Broadcast weights via a temporary NCCL group to rollout engines."""
+
+    def connect_rollout_engines(
+        self,
+        rollout_engines: Sequence[ActorHandle],
+        rollout_engine_lock: ActorHandle | None,
+    ) -> None:
+        """On rank 0, initialize a temporary NCCL group for parameter broadcast."""
+        self.rollout_engines = rollout_engines
+        self.rollout_engine_lock = rollout_engine_lock
+
+        # For TP:
+        #   1. AllGather parameters to rank 0
+        #   2. Broadcast parameters from rank 0 to all sglang engines
+        self._is_src_rank = dist.get_rank() == 0
+        if self._is_src_rank:
+            self._group_name = "slime"
+            master_address = ray._private.services.get_node_ip_address()
+            with socket.socket() as sock:
+                sock.bind(("", 0))
+                master_port = sock.getsockname()[1]
+            ## TODO: why +1?
+            world_size = self.args.rollout_num_gpus + 1
+
+            refs = [
+                engine.init_weights_update_group.remote(
+                    master_address,
+                    master_port,
+                    i * self.args.rollout_num_gpus_per_engine + 1,
+                    world_size,
+                    self._group_name,
+                    backend="nccl",
+                )
+                for i, engine in enumerate(self.rollout_engines)
+            ]
+            self._model_update_groups = init_process_group(
+                backend="nccl",
+                init_method=f"tcp://{master_address}:{master_port}",
+                world_size=world_size,
+                rank=0,
+                group_name=self._group_name,
+            )
+            ray.get(refs)
+
+    def update_bucket_weights(self, named_tensors, weight_version=None) -> None:
+        """Send names/dtypes/shapes metadata to engines, then broadcast tensors.
+
+        Ensures tensors are contiguous; when `world_size == 1`, converts DTensors
+        to full tensors prior to `dist.broadcast`.
+        """
+        if not self._is_src_rank or not named_tensors:
+            return
+
+        refs = [
+            engine.update_weights_from_distributed.remote(
+                names=[name for name, _ in named_tensors],
+                dtypes=[param.dtype for _, param in named_tensors],
+                shapes=[param.shape for _, param in named_tensors],
+                group_name=self._group_name,
+                weight_version=str(weight_version),
+            )
+            for engine in self.rollout_engines
+        ]
+
+        handles = []
+        # Broadcast parameters one by one with memory management
+        for _name, param in named_tensors:
+            torch.cuda.empty_cache()
+            # Ensure tensor is contiguous and on the right device
+            param_data = param.data.contiguous()
+
+            # avoid `DTensor._op_dispatcher.dispatch` has `assert compute_mesh is not None` error
+            if dist.get_world_size() == 1 and isinstance(param_data, DTensor):
+                param_data = param_data.full_tensor()
+
+            # Synchronous broadcast to avoid memory buildup
+            handles.append(dist.broadcast(param_data, 0, group=self._model_update_groups, async_op=True))
+
+        for handle in handles:
+            handle.wait()
+        ray.get(refs)

slime/backends/megatron_utils/__init__.py

@@ -0,0 +1,45 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+import torch
+
+try:
+    import deep_ep
+    from torch_memory_saver import torch_memory_saver
+
+    old_init = deep_ep.Buffer.__init__
+
+    def new_init(self, *args, **kwargs):
+        if torch_memory_saver._impl is not None:
+            torch_memory_saver._impl._binary_wrapper.cdll.tms_set_interesting_region(False)
+        old_init(self, *args, **kwargs)
+        torch.cuda.synchronize()
+        if torch_memory_saver._impl is not None:
+            torch_memory_saver._impl._binary_wrapper.cdll.tms_set_interesting_region(True)
+
+    deep_ep.Buffer.__init__ = new_init
+except ImportError:
+    logging.warning("deep_ep is not installed, some functionalities may be limited.")
+
+try:
+    from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import (
+        Qwen3VLMoETextRotaryEmbedding,
+        Qwen3VLTextRotaryEmbedding,
+    )
+
+    def patch_rotary_embedding(cls):
+        _original_forward = cls.forward
+
+        def _patched_forward(self, *args, packed_seq_params=None, **kwargs):
+            return _original_forward(self, *args, **kwargs)
+
+        cls.forward = _patched_forward
+
+    patch_rotary_embedding(Qwen3VLTextRotaryEmbedding)
+    patch_rotary_embedding(Qwen3VLMoETextRotaryEmbedding)
+except ImportError:
+    pass
+
+logging.getLogger().setLevel(logging.WARNING)

slime/backends/megatron_utils/actor.py

@@ -0,0 +1,575 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import os
+import random
+import socket
+from argparse import Namespace
+from contextlib import nullcontext
+
+import ray
+import torch
+import torch.distributed as dist
+from megatron.core import mpu
+from ray.actor import ActorHandle
+from torch_memory_saver import torch_memory_saver
+from transformers import AutoConfig, AutoTokenizer
+
+from slime.ray.train_actor import TrainRayActor
+from slime.utils import train_dump_utils
+from slime.utils.context_utils import with_defer
+from slime.utils.data import process_rollout_data
+from slime.utils.distributed_utils import get_gloo_group, init_process_group
+from slime.utils.memory_utils import clear_memory, print_memory
+from slime.utils.ray_utils import Box
+from slime.utils.reloadable_process_group import destroy_process_groups, monkey_patch_torch_dist, reload_process_groups
+from slime.utils.routing_replay import RoutingReplay
+from slime.utils.timer import Timer, inverse_timer, timer
+from slime.utils.tracking_utils import init_tracking
+from slime.utils.types import RolloutBatch
+
+from ...utils.profile_utils import TrainProfiler
+from ...utils.tensor_backper import TensorBackuper
+from .checkpoint import load_checkpoint
+from .cp_utils import slice_log_prob_with_cp, slice_with_cp
+from .data import DataIterator, get_data_iterator, log_perf_data, log_rollout_data, sync_actor_critic_data
+from .initialize import init, is_megatron_main_rank
+from .loss import compute_advantages_and_returns, get_log_probs_and_entropy, get_values
+from .model import forward_only, initialize_model_and_optimizer, save, train
+from .update_weight.common import named_params_and_buffers
+from .update_weight.update_weight_from_distributed import UpdateWeightFromDistributed
+from .update_weight.update_weight_from_tensor import UpdateWeightFromTensor
+
+logging.getLogger("megatron").setLevel(logging.WARNING)
+
+logger = logging.getLogger(__name__)
+
+
+class MegatronTrainRayActor(TrainRayActor):
+    @with_defer(lambda: Timer().start("train_wait"))
+    def init(
+        self,
+        args: Namespace,
+        role: str,
+        with_ref: bool = False,
+    ) -> int | None:
+        monkey_patch_torch_dist()
+
+        super().init(args, role, with_ref)
+
+        init(args)
+
+        if is_megatron_main_rank():
+            init_tracking(args, primary=False)
+
+        self.prof = TrainProfiler(args)
+
+        # read config and tokenizer serialized to prevent concurrent writing bug.
+        for i in range(dist.get_world_size()):
+            if i == dist.get_rank():
+                self.hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True)
+                self.tokenizer = AutoTokenizer.from_pretrained(self.args.hf_checkpoint, trust_remote_code=True)
+            dist.barrier(group=get_gloo_group())
+
+        self.train_parallel_config = {
+            "dp_size": mpu.get_data_parallel_world_size(with_context_parallel=False),
+        }
+        dist.barrier(group=get_gloo_group())
+
+        if args.offload_train:
+            if (x := args.train_memory_margin_bytes) > 0:
+                logger.info(f"Set torch_memory_saver.memory_margin_bytes to {x}")
+                torch_memory_saver.memory_margin_bytes = x
+
+        if self.args.debug_rollout_only:
+            return 0
+
+        if role == "critic":
+            self.args.load = self.args.critic_load
+            self.args.save = self.args.critic_save
+            self.args.lr = self.args.critic_lr
+            self.args.lr_warmup_iters = self.args.critic_lr_warmup_iters
+
+        (self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id) = initialize_model_and_optimizer(
+            args, role
+        )
+
+        if role == "critic":
+            if self.args.offload_train:
+                self.sleep()
+            return
+
+        start_rollout_id = loaded_rollout_id + 1
+
+        self.weights_backuper = TensorBackuper.create(
+            source_getter=lambda: named_params_and_buffers(
+                self.args,
+                self.model,
+                convert_to_global_name=args.megatron_to_hf_mode == "raw",
+                translate_gpu_to_cpu=not self.args.enable_weights_backuper,
+            ),
+            single_tag=None if args.enable_weights_backuper else "actor",
+        )
+        self._active_model_tag: str | None = "actor"
+        self.weights_backuper.backup("actor")
+
+        if with_ref:
+            self.load_other_checkpoint("ref", args.ref_load)
+
+        if self.args.keep_old_actor:
+            # Load old_actor checkpoint
+            self.load_other_checkpoint("old_actor", args.load)
+            # Create rollout_actor as a copy of current actor
+            if args.update_weights_interval == 1:
+                self.weights_backuper.backup("rollout_actor")
+
+        if self.args.vocab_size is None:
+            self.args.vocab_size = self.tokenizer.vocab_size
+
+        update_weight_cls = UpdateWeightFromTensor if self.args.colocate else UpdateWeightFromDistributed
+        self.weight_updater = update_weight_cls(
+            self.args,
+            self.model,
+            weights_getter=lambda: self.weights_backuper.get("actor"),
+            model_name=type(self.hf_config).__name__.lower() if self.args.model_name is None else self.args.model_name,
+            quantization_config=getattr(self.hf_config, "quantization_config", None),
+        )
+
+        # empty cache after initialization
+        clear_memory()
+
+        if self.args.offload_train:
+            # recover to actor in the end.
+            self._switch_model("actor")
+            self.sleep()
+
+        self.rollout_engines = None
+
+        self.rollout_data_postprocess = None
+        if self.args.rollout_data_postprocess_path is not None:
+            from slime.utils.misc import load_function
+
+            self.rollout_data_postprocess = load_function(self.args.rollout_data_postprocess_path)
+
+        self.prof.on_init_end()
+
+        return start_rollout_id
+
+    @timer
+    def sleep(self) -> None:
+        assert self.args.offload_train
+
+        clear_memory(clear_host_memory=True)
+        print_memory("before offload model")
+        destroy_process_groups()
+
+        torch_memory_saver.pause()
+
+        print_memory("after offload model")
+
+    @timer
+    def wake_up(self) -> None:
+        assert self.args.offload_train
+        print_memory("before wake_up model")
+
+        torch_memory_saver.resume()
+
+        clear_memory()
+        reload_process_groups()
+        print_memory("after wake_up model")
+
+    def _get_rollout_data(self, rollout_data_ref: Box) -> RolloutBatch:
+        # Fetch data through ray on CPU, not sure if this will be performance bottleneck.
+        # Both first pp stage and the last pp stage will receive the data.
+        rollout_data = process_rollout_data(
+            self.args,
+            rollout_data_ref,
+            mpu.get_data_parallel_rank(with_context_parallel=False),
+            mpu.get_data_parallel_world_size(with_context_parallel=False),
+        )
+        # TODO: this is ugly, move to somewhere else?
+        # move tokens to GPU in advance
+        rollout_data["tokens"] = [
+            torch.tensor(t, dtype=torch.long, device=torch.cuda.current_device()) for t in rollout_data["tokens"]
+        ]
+        rollout_data["loss_masks"] = [
+            torch.tensor(t, dtype=torch.int, device=torch.cuda.current_device()) for t in rollout_data["loss_masks"]
+        ]
+        if "multimodal_train_inputs" in rollout_data:
+            # Move multimodal training tensors to GPU in advance
+            rollout_data["multimodal_train_inputs"] = [
+                (
+                    {key: tensor.to(device=torch.cuda.current_device()) for key, tensor in mm_dict.items()}
+                    if mm_dict is not None
+                    else None
+                )
+                for mm_dict in rollout_data["multimodal_train_inputs"]
+            ]
+        if "rollout_log_probs" in rollout_data:
+            rollout_data["rollout_log_probs"] = [
+                torch.tensor(
+                    slice_log_prob_with_cp(log_prob, total_length, response_length),
+                    device=torch.cuda.current_device(),
+                    dtype=torch.float32,
+                )
+                for log_prob, total_length, response_length in zip(
+                    rollout_data["rollout_log_probs"],
+                    rollout_data["total_lengths"],
+                    rollout_data["response_lengths"],
+                    strict=False,
+                )
+            ]
+        if "rollout_routed_experts" in rollout_data:
+            rollout_data["rollout_routed_experts"] = [
+                torch.from_numpy(r) for r in rollout_data["rollout_routed_experts"]
+            ]
+        return rollout_data
+
+    def _switch_model(self, target_tag: str) -> None:
+        if target_tag not in self.weights_backuper.backup_tags:
+            raise ValueError(f"Cannot switch to unknown model tag: {target_tag}")
+        self.weights_backuper.restore(target_tag)
+        self._active_model_tag = target_tag
+
+    def fill_routing_replay(self, data_iterator, num_microbatches, rollout_data):
+        if "rollout_routed_experts" not in rollout_data:
+            raise ValueError(
+                "rollout_routed_experts is required in rollout_data when use_rollout_routing_replay is set."
+            )
+
+        from megatron.core.transformer.transformer_block import get_num_layers_to_build
+        from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
+
+        from slime.utils.routing_replay import RoutingReplay
+
+        for iterator in data_iterator:
+            iterator.reset()
+
+        tp_rank = mpu.get_tensor_model_parallel_rank()
+        tp_size = mpu.get_tensor_model_parallel_world_size()
+
+        def pad_func(experts, pad):
+            _, num_layers, topk = experts.shape
+            pad = (
+                torch.arange(
+                    pad * num_layers * topk,
+                    device=experts.device,
+                    dtype=experts.dtype,
+                ).reshape((pad, num_layers, topk))
+                % self.args.num_experts
+            )
+            return torch.cat([experts, pad], dim=0)
+
+        for _ in range(sum(num_microbatches)):
+            batch = data_iterator[0].get_next(["rollout_routed_experts", "tokens"])
+            rollout_routed_experts = batch["rollout_routed_experts"]
+            tokens = batch["tokens"]
+            assert len(rollout_routed_experts) == len(tokens)
+            for a, b in zip(rollout_routed_experts, tokens, strict=False):
+                assert a.shape[0] == b.shape[0] - 1, f"{a.shape}, {b.shape}"
+
+            # We need to pad the experts to the last token. We won't calculate loss on this token so this should be fine.
+            # TODO: fuse this padding with the following slice_with_cp to reduce memory copy.
+            rollout_routed_experts = [pad_func(r, 1) for r in rollout_routed_experts]
+            # TODO: maybe extract a common process function for here and get_batch?
+            rollout_routed_experts = [slice_with_cp(r, pad_func) for r in rollout_routed_experts]
+            rollout_routed_experts = torch.cat(rollout_routed_experts, dim=0)
+            pad_size = mpu.get_tensor_model_parallel_world_size() * self.args.data_pad_size_multiplier
+            pad = (pad_size - rollout_routed_experts.size(0) % pad_size) % pad_size
+            if pad != 0:
+                rollout_routed_experts = pad_func(rollout_routed_experts, pad)
+
+            if self.args.sequence_parallel:
+                seqlen = rollout_routed_experts.size(0)
+                assert seqlen % tp_size == 0
+                start, end = seqlen // tp_size * tp_rank, seqlen // tp_size * (tp_rank + 1)
+                rollout_routed_experts = rollout_routed_experts[start:end]
+
+            routing_replay_offset = 0
+            for vp_stage, model in enumerate(self.model):
+                config = model.module.config
+                num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage)
+                offset = get_transformer_layer_offset(config, vp_stage=vp_stage)
+                for layer_id in range(offset, offset + num_layers_to_build):
+                    # skip dense layer
+                    if isinstance(config.moe_layer_freq, int):
+                        if layer_id % config.moe_layer_freq != 0:
+                            continue
+                    elif isinstance(config.moe_layer_freq, list):
+                        assert len(config.moe_layer_freq) == config.num_layers
+                        if config.moe_layer_freq[layer_id] == 0:
+                            continue
+                    layer_routed_experts = rollout_routed_experts[:, layer_id]
+                    RoutingReplay.all_routing_replays[routing_replay_offset].record(layer_routed_experts)
+                    routing_replay_offset += 1
+            assert routing_replay_offset == len(RoutingReplay.all_routing_replays)
+
+        del rollout_data["rollout_routed_experts"]
+
+        for iterator in data_iterator:
+            iterator.reset()
+
+    def compute_log_prob(
+        self,
+        data_iterator: list[DataIterator],
+        num_microbatches: list[int],
+        store_prefix: str = "",
+    ) -> dict[str, list[torch.Tensor]]:
+
+        with timer(f"{store_prefix}log_probs"):
+            return forward_only(
+                get_log_probs_and_entropy,
+                self.args,
+                self.model,
+                data_iterator,
+                num_microbatches,
+                store_prefix=store_prefix,
+            )
+
+    def train(self, rollout_id: int, rollout_data_ref: Box) -> None:
+        if self.args.offload_train:
+            self.wake_up()
+
+        with timer("data_preprocess"):
+            rollout_data = self._get_rollout_data(rollout_data_ref)
+            if self.args.debug_rollout_only:
+                log_rollout_data(rollout_id, self.args, rollout_data)
+                return
+
+        if self.role == "critic":
+            return self.train_critic(rollout_id, rollout_data)
+        else:
+            return self.train_actor(rollout_id, rollout_data)
+
+    def train_critic(self, rollout_id: int, rollout_data: RolloutBatch) -> None:
+        # Create data iterator for log_probs and train.
+        data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data)
+        rollout_data.update(
+            forward_only(
+                get_values,
+                self.args,
+                self.model,
+                data_iterator,
+                num_microbatches,
+            )
+        )
+
+        if rollout_id >= self.args.num_critic_only_steps:
+            sync_actor_critic_data(self.args, rollout_data, self._actor_critic_groups)
+
+        compute_advantages_and_returns(self.args, rollout_data)
+
+        self.args.loss_type = "value_loss"
+        train(
+            rollout_id,
+            self.model,
+            self.optimizer,
+            self.opt_param_scheduler,
+            data_iterator,
+            num_microbatches,
+        )
+
+    def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None:
+        # Create data iterator for log_probs and train.
+        data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data)
+
+        if self.args.use_rollout_routing_replay:
+            self.fill_routing_replay(data_iterator, num_microbatches, rollout_data)
+
+        with inverse_timer("train_wait"), timer("train"):
+            if self.args.compute_advantages_and_returns:
+                if "ref" in self.weights_backuper.backup_tags:
+                    if self.args.use_routing_replay:
+                        os.environ["ROUTING_REPLAY_STAGE"] = "fallthrough"
+                    self._switch_model("ref")
+                    rollout_data.update(
+                        self.compute_log_prob(
+                            data_iterator,
+                            num_microbatches,
+                            store_prefix="ref_",
+                        )
+                    )
+                self._switch_model("old_actor" if self.args.keep_old_actor else "actor")
+                if not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics:
+                    if self.args.use_routing_replay:
+                        if self.args.use_rollout_routing_replay:
+                            os.environ["ROUTING_REPLAY_STAGE"] = "replay_forward"
+                        else:
+                            os.environ["ROUTING_REPLAY_STAGE"] = "record"
+                    rollout_data.update(
+                        self.compute_log_prob(
+                            data_iterator,
+                            num_microbatches,
+                            store_prefix="",
+                        )
+                    )
+                    if self.args.use_rollout_routing_replay:
+                        RoutingReplay.clear_all_forward()
+
+                if self.args.use_critic:
+                    sync_actor_critic_data(
+                        self.args,
+                        rollout_data,
+                        self._actor_critic_groups,
+                    )
+                if self._active_model_tag != "actor":
+                    self._switch_model("actor")
+
+                # Calculate adv and returns. Need to performed before training (instead of on the fly),
+                # because we may need normalize the whole rollout.
+                compute_advantages_and_returns(self.args, rollout_data)
+            
+            if self.rollout_data_postprocess is not None:
+                self.rollout_data_postprocess(self.args)
+
+            log_rollout_data(rollout_id, self.args, rollout_data)
+
+            # Train
+            if self.args.use_routing_replay:
+                os.environ["ROUTING_REPLAY_STAGE"] = "replay_backward"
+            with timer("actor_train"):
+                train(
+                    rollout_id,
+                    self.model,
+                    self.optimizer,
+                    self.opt_param_scheduler,
+                    data_iterator,
+                    num_microbatches,
+                )
+
+            self.prof.step(rollout_id=rollout_id)
+
+        train_dump_utils.save_debug_train_data(self.args, rollout_id=rollout_id, rollout_data=rollout_data)
+
+        if self.args.use_routing_replay:
+            RoutingReplay.clear_all()
+
+        # update the cpu actor weight to the latest model
+        self.weights_backuper.backup("actor")
+
+        # Update ref model if needed
+        if (
+            self.args.ref_update_interval is not None
+            and (rollout_id + 1) % self.args.ref_update_interval == 0
+            and "ref" in self.weights_backuper.backup_tags
+        ):
+            with timer("ref_model_update"):
+                if is_megatron_main_rank():
+                    logger.info(f"Updating ref model at rollout_id {rollout_id}")
+                self.weights_backuper.backup("ref")
+
+        log_perf_data(rollout_id, self.args)
+
+    @timer
+    def save_model(self, rollout_id: int, force_sync: bool = False) -> None:
+        if self.args.debug_rollout_only:
+            return
+
+        # torch dist may trigger nccl communication during saving.
+        if self.args.offload_train:
+            reload_process_groups()
+
+        if self.args.async_save:
+            from megatron.training.async_utils import maybe_finalize_async_save
+
+            maybe_finalize_async_save(blocking=True)
+
+        save(rollout_id, self.model, self.optimizer, self.opt_param_scheduler)
+
+        if force_sync and self.args.async_save:
+            maybe_finalize_async_save(blocking=True)
+
+        if self.args.offload_train:
+            destroy_process_groups()
+
+    @timer
+    def update_weights(self) -> None:
+        if self.args.debug_train_only or self.args.debug_rollout_only:
+            return
+
+        if self.args.offload_train:
+            reload_process_groups()
+
+        rollout_engines, rollout_engine_lock, num_new_engines = ray.get(
+            self.rollout_manager.get_rollout_engines_and_lock.remote()
+        )
+        if num_new_engines > 0:
+            self.weight_updater.connect_rollout_engines(rollout_engines, rollout_engine_lock)
+            dist.barrier(group=get_gloo_group())
+
+        with torch_memory_saver.disable() if self.args.offload_train else nullcontext():
+            print_memory("before update_weights")
+            self.weight_updater.update_weights()
+            print_memory("after update_weights")
+
+            if self.args.ci_test and len(rollout_engines) > 0:
+                engine = random.choice(rollout_engines)
+                engine_version = ray.get(engine.get_weight_version.remote())
+                if str(engine_version) != str(self.weight_updater.weight_version):
+                    raise RuntimeError(
+                        f"Weight version mismatch! Engine: {engine_version}, Updater: {self.weight_updater.weight_version}"
+                    )
+
+            if getattr(self.args, "keep_old_actor", False):
+                if self.args.update_weights_interval == 1:
+                    logger.info("updating model queue: rollout_actor -> old_actor, actor -> rollout_actor")
+                    # Queue-style update: rollout_actor params -> old_actor, actor params -> rollout_actor
+                    # First copy rollout_actor to old_actor
+                    self.weights_backuper.copy(src_tag="rollout_actor", dst_tag="old_actor")
+                    # Then copy current actor to rollout_actor
+                    self.weights_backuper.backup("rollout_actor")
+                else:
+                    self.weights_backuper.backup("old_actor")
+
+        if self.args.offload_train:
+            destroy_process_groups()
+
+    def load_other_checkpoint(self, model_tag: str, path: str) -> None:
+        old_args = self.args.load, self.args.no_load_optim, self.args.no_load_rng, self.args.finetune
+        self.args.load = path
+        self.args.no_load_optim = True
+        self.args.no_load_rng = True
+        self.args.finetune = True
+
+        if model_tag == "ref" and self.args.ref_ckpt_step is not None:
+            old_ckpt_step = self.args.ckpt_step
+            self.args.ckpt_step = self.args.ref_ckpt_step
+
+        _, _ = load_checkpoint(
+            self.model,
+            None,
+            None,
+            checkpointing_context={},
+            skip_load_to_model_and_opt=False,
+        )
+        self.args.load, self.args.no_load_optim, self.args.no_load_rng, self.args.finetune = old_args
+
+        if model_tag == "ref" and self.args.ref_ckpt_step is not None:
+            self.args.ckpt_step = old_ckpt_step
+
+        self.weights_backuper.backup(model_tag)
+        self._active_model_tag = model_tag
+
+    def connect_actor_critic(
+        self,
+        actor_handle: ActorHandle | None = None,
+        master_address: str | None = None,
+        master_port: int | None = None,
+    ) -> None:
+        if self.role == "actor":
+            master_address = ray.util.get_node_ip_address()
+            with socket.socket() as sock:
+                sock.bind(("", 0))
+                master_port = sock.getsockname()[1]
+            actor_handle.connect_actor_critic.remote(master_address=master_address, master_port=master_port)
+
+        group_name = "actor_critic"
+        world_size = 2
+        self._actor_critic_groups = init_process_group(
+            backend="nccl",
+            init_method=f"tcp://{master_address}:{master_port}",
+            world_size=world_size,
+            rank=0 if self.role == "actor" else 1,
+            group_name=group_name,
+        )

slime/backends/megatron_utils/arguments.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+from megatron.training.arguments import parse_args, validate_args
+from megatron.training.tokenizer.tokenizer import _vocab_size_with_padding
+
+__all__ = ["validate_args", "parse_args", "set_default_megatron_args"]
+
+logger = logging.getLogger(__name__)
+
+
+def set_default_megatron_args(args):
+    # always use zero optimizer
+    args.use_distributed_optimizer = True
+    # TODO: maybe change this after megatron has good fp8 support
+    args.bf16 = not args.fp16
+    # placeholders
+    args.seq_length = 4096
+    args.max_position_embeddings = args.seq_length
+    # compatible for megatron
+    if hasattr(args, "rope_type") and args.rope_type is None:
+        args.rope_type = "yarn" if args.multi_latent_attention else "rope"
+
+    if args.vocab_size and not args.padded_vocab_size:
+        args.padded_vocab_size = _vocab_size_with_padding(args.vocab_size, args)
+
+    if not args.tokenizer_model and not args.tokenizer_type:
+        logger.info("--tokenizer-model not set, use --hf-checkpoint as tokenizer model.")
+        args.tokenizer_model = args.hf_checkpoint
+        args.tokenizer_type = "HuggingFaceTokenizer"
+    return args

slime/backends/megatron_utils/checkpoint.py

@@ -0,0 +1,79 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import os
+import re
+from pathlib import Path
+
+# TODO: may need to copy those 2 functions and do refactoring.
+from megatron.training.checkpointing import load_checkpoint as _load_checkpoint_megatron
+from megatron.training.checkpointing import save_checkpoint
+from megatron.training.global_vars import get_args
+
+from slime.utils import megatron_bridge_utils
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["save_checkpoint"]
+
+
+def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, checkpointing_context, skip_load_to_model_and_opt):
+    # ref: how megatron `load_checkpoint` gets directory
+    args = get_args()
+    load_path = args.load
+
+    assert Path(load_path).exists() and _is_dir_nonempty(
+        load_path
+    ), f"{args.load=} does not exist or is an empty directory. Did you specify the wrong folder?"
+
+    if _is_megatron_checkpoint(load_path):
+        return _load_checkpoint_megatron(
+            ddp_model=ddp_model,
+            optimizer=optimizer,
+            opt_param_scheduler=opt_param_scheduler,
+            checkpointing_context=checkpointing_context,
+            skip_load_to_model_and_opt=skip_load_to_model_and_opt,
+        )
+    else:
+        return _load_checkpoint_hf(
+            ddp_model=ddp_model,
+            optimizer=optimizer,
+            args=args,
+            load_path=load_path,
+        )
+
+
+def _is_megatron_checkpoint(path: str | Path) -> bool:
+    return (Path(path) / "latest_checkpointed_iteration.txt").is_file() or bool(
+        re.fullmatch(r"iter_\d{7}", Path(path).name)
+    )
+
+
+def _load_checkpoint_hf(ddp_model, optimizer, args, load_path: str):
+    assert args.megatron_to_hf_mode == "bridge", "Only bridge mode is supported for loading HF checkpoint"
+    from megatron.bridge import AutoBridge
+
+    import slime_plugins.megatron_bridge  # noqa: F401
+
+    logger.info(f"Load checkpoint from HuggingFace model into Megatron (path={load_path})")
+
+    with megatron_bridge_utils.patch_megatron_model(ddp_model):
+        bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)
+        bridge.load_hf_weights(ddp_model)
+
+    # Copied from Megatron-core :: load_checkpoint (with simplifications)
+    if (args.fp16 or args.bf16) and optimizer is not None:
+        assert not args.load_main_params_from_ckpt
+        optimizer.reload_model_params()
+
+    # We can see `successfully loaded checkpoint from ... [ t 1/2, p 1/1 ] at iteration 0`
+    # when loading Megatron, thus it is 0
+    iteration = 0
+    num_floating_point_operations_so_far = 0
+    return iteration, num_floating_point_operations_so_far
+
+
+def _is_dir_nonempty(path):
+    with os.scandir(path) as it:
+        return any(it)

slime/backends/megatron_utils/ci_utils.py

@@ -0,0 +1,87 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""CI utilities for Megatron backend testing."""
+
+import logging
+from collections.abc import Sequence
+
+from megatron.core.distributed import DistributedDataParallel as DDP
+
+logger = logging.getLogger(__name__)
+
+
+def check_mtp_only_grad(model: Sequence[DDP], step_id: int) -> None:
+    """Check that only MTP parameters have non-zero gradients.
+
+    This is used for CI testing to verify that when all outputs are truncated,
+    only the MTP layers receive gradients (since only mtp_loss contributes).
+
+    Args:
+        model: Sequence of DDP-wrapped model chunks.
+        step_id: Current step index for logging.
+
+    Raises:
+        AssertionError: If any non-MTP parameter has a non-zero gradient.
+    """
+    non_mtp_nonzero_grads = []
+    mtp_nonzero_grads = []
+
+    for model_chunk in model:
+        for name, param in model_chunk.named_parameters():
+            # Get the main_grad from the distributed optimizer if available
+            grad = getattr(param, "main_grad", None)
+            if grad is None:
+                grad = param.grad
+            if grad is None:
+                continue
+
+            grad_norm = grad.abs().max().item()
+            is_mtp = ".mtp." in name
+
+            if is_mtp:
+                if grad_norm > 0:
+                    mtp_nonzero_grads.append((name, grad_norm))
+            else:
+                if grad_norm > 0:
+                    non_mtp_nonzero_grads.append((name, grad_norm))
+
+    # Log the results
+    logger.info(
+        f"[CI MTP Grad Check] Step {step_id}: "
+        f"MTP params with non-zero grad: {len(mtp_nonzero_grads)}, "
+        f"non-MTP params with non-zero grad: {len(non_mtp_nonzero_grads)}"
+    )
+
+    if non_mtp_nonzero_grads:
+        # Log the first few non-MTP params with non-zero gradients for debugging
+        for name, grad_norm in non_mtp_nonzero_grads[:5]:
+            logger.error(f"[CI MTP Grad Check] Non-MTP param with non-zero grad: {name}, max_grad={grad_norm}")
+
+    assert len(non_mtp_nonzero_grads) == 0, (
+        f"Expected all non-MTP parameters to have zero gradients, "
+        f"but found {len(non_mtp_nonzero_grads)} with non-zero gradients. "
+        f"First few: {non_mtp_nonzero_grads[:5]}"
+    )
+
+    # Also verify that MTP params do have gradients (otherwise the test is not valid)
+    assert len(mtp_nonzero_grads) > 0, (
+        "Expected MTP parameters to have non-zero gradients, but all were zero. "
+        "This may indicate the MTP loss is not being computed."
+    )
+
+
+def check_mtp_loss(mtp_loss: float, max_mtp_loss: float = 1.0) -> None:
+    """Check that MTP loss is within expected bounds.
+
+    Args:
+        mtp_loss: The computed MTP loss value.
+        max_mtp_loss: Maximum allowed MTP loss (default: 1.0).
+
+    Raises:
+        AssertionError: If MTP loss exceeds the maximum allowed value.
+    """
+    assert mtp_loss < max_mtp_loss, (
+        f"MTP loss {mtp_loss} exceeds maximum allowed value {max_mtp_loss}. "
+        "This may indicate an issue with MTP training."
+    )

slime/backends/megatron_utils/cp_utils.py

@@ -0,0 +1,210 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Callable
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from megatron.core import mpu
+
+
+def get_logits_and_tokens_offset_with_cp(
+    total_length: int,
+    response_length: int,
+):
+    """
+    All offsets start from the begining of the prompt.
+    """
+    cp_rank = mpu.get_context_parallel_rank()
+    cp_size = mpu.get_context_parallel_world_size()
+    assert cp_size > 1
+
+    prompt_length = total_length - response_length
+    chunk_size = (total_length + 2 * cp_size - 1) // (2 * cp_size)
+
+    # the offset of 2 chunks
+    chunk_0 = (cp_rank * chunk_size, (cp_rank + 1) * chunk_size)
+    chunk_1 = ((2 * cp_size - cp_rank - 1) * chunk_size, (2 * cp_size - cp_rank) * chunk_size)
+
+    # the offset of 2 logits, note that the logits need a "-1".
+    logits_0 = (max(chunk_0[0], prompt_length - 1), min(chunk_0[1], total_length - 1))
+    logits_1 = (max(chunk_1[0], prompt_length - 1), min(chunk_1[1], total_length - 1))
+
+    # when the sequence is empty, make an empty slice to continue the gradient flow.
+    if logits_0[0] < logits_0[1]:
+        token_0 = (logits_0[0] + 1, logits_0[1] + 1)
+    else:
+        logits_0 = (0, 0)
+        token_0 = (0, 0)
+
+    if logits_1[0] < logits_1[1]:
+        token_1 = (logits_1[0] + 1, logits_1[1] + 1)
+    else:
+        logits_1 = (0, 0)
+        token_1 = (0, 0)
+
+    return chunk_size, (chunk_0, chunk_1), (logits_0, logits_1), (token_0, token_1)
+
+
+def get_sum_of_sample_mean(
+    total_lengths: list[int],
+    response_lengths: list[int],
+    loss_masks: list[torch.Tensor],
+    calculate_per_token_loss: bool = False,
+) -> Callable[[torch.Tensor], torch.Tensor]:
+    """
+    Calculate correct sample mean for CP
+    """
+    cp_size = mpu.get_context_parallel_world_size()
+    if cp_size == 1:
+
+        def sum_of_sample_mean(x: torch.Tensor) -> torch.Tensor:
+            return sum(
+                [
+                    (x_i * loss_mask_i).sum() / torch.clamp_min(loss_mask_i.sum(), 1)
+                    for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=False)
+                ]
+            )
+
+        def sum_of_token(x: torch.Tensor) -> torch.Tensor:
+            return sum(
+                [
+                    (x_i * loss_mask_i).sum()
+                    for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=False)
+                ]
+            )
+
+    else:
+        cp_chunk_lengths = []
+        chunked_loss_masks = []
+        for i, (total_length, response_length, loss_mask) in enumerate(
+            zip(total_lengths, response_lengths, loss_masks, strict=False)
+        ):
+            prompt_length = total_length - response_length
+            _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp(total_length, response_length)
+            loss_mask_0 = loss_mask[tokens_offset[0][0] - prompt_length : tokens_offset[0][1] - prompt_length]
+            loss_mask_1 = loss_mask[tokens_offset[1][0] - prompt_length : tokens_offset[1][1] - prompt_length]
+            chunked_loss_masks.append(torch.cat([loss_mask_0, loss_mask_1], dim=0))
+            cp_chunk_lengths.append(chunked_loss_masks[i].size(0))
+
+        def sum_of_sample_mean(x: torch.Tensor) -> torch.Tensor:
+            return sum(
+                [
+                    (x_i * chunked_loss_mask).sum() / torch.clamp_min(loss_mask.sum(), 1)
+                    for x_i, chunked_loss_mask, loss_mask in zip(
+                        x.split(cp_chunk_lengths, dim=0), chunked_loss_masks, loss_masks, strict=False
+                    )
+                ]
+            )
+
+        def sum_of_token(x: torch.Tensor) -> torch.Tensor:
+            return sum(
+                [
+                    (x_i * chunked_loss_mask).sum()
+                    for x_i, chunked_loss_mask in zip(
+                        x.split(cp_chunk_lengths, dim=0), chunked_loss_masks, strict=False
+                    )
+                ]
+            )
+
+    return sum_of_sample_mean if not calculate_per_token_loss else sum_of_token
+
+
+def all_gather_with_cp(tensor: torch.Tensor, total_length: int, response_length: int) -> torch.Tensor:
+    """
+    Gather tensors across all ranks in the context parallel group.
+    The first dimension of the output tensor will be the `response_length`.
+    """
+    cp_group = mpu.get_context_parallel_group()
+    cp_size = mpu.get_context_parallel_world_size()
+
+    if cp_size == 1:
+        return tensor
+
+    _, _, logits_offset, _ = get_logits_and_tokens_offset_with_cp(total_length, response_length)
+
+    prompt_length = total_length - response_length
+
+    chunk_0 = tensor[: logits_offset[0][1] - logits_offset[0][0]]
+    chunk_1 = tensor[logits_offset[0][1] - logits_offset[0][0] :]
+    assert chunk_1.shape[0] == logits_offset[1][1] - logits_offset[1][0]
+
+    def zero(len: int) -> torch.Tensor:
+        return torch.zeros(
+            [len] + list(tensor.shape[1:]),
+            dtype=tensor.dtype,
+            device=tensor.device,
+            requires_grad=True,
+        )
+
+    # logprob should be within the range of [prompt_length - 1, total_length - 1]
+    if chunk_0.shape[0] == 0 and chunk_1.shape[0] == 0:
+        # all empty
+        full_tensor = zero(response_length)
+    elif chunk_0.shape[0] != 0 and chunk_1.shape[0] == 0:
+        # only first chunk
+        left = zero(logits_offset[0][0] - (prompt_length - 1))
+        right = zero(total_length - 1 - logits_offset[0][1])
+        full_tensor = torch.cat([left, chunk_0, right], dim=0)
+    elif chunk_0.shape[0] == 0 and chunk_1.shape[0] != 0:
+        # only second chunk
+        left = zero(logits_offset[1][0] - (prompt_length - 1))
+        right = zero(total_length - 1 - logits_offset[1][1])
+        full_tensor = torch.cat([left, chunk_1, right], dim=0)
+    else:
+        left = zero(logits_offset[0][0] - (prompt_length - 1))
+        mid = zero(logits_offset[1][0] - logits_offset[0][1])
+        right = zero(total_length - 1 - logits_offset[1][1])
+        full_tensor = torch.cat([left, chunk_0, mid, chunk_1, right], dim=0)
+
+    assert full_tensor.shape[0] == response_length, f"Expected {response_length}, got {full_tensor.shape}"
+    full_tensor = dist.nn.all_reduce(full_tensor, group=cp_group)
+    return full_tensor
+
+
+def slice_with_cp(tokens: torch.Tensor, pad_value: tuple[int, float, Callable]) -> torch.Tensor:
+    cp_rank = mpu.get_context_parallel_rank()
+    cp_size = mpu.get_context_parallel_world_size()
+
+    if cp_size == 1:
+        return tokens
+
+    # pad
+    chunk_size = (len(tokens) + 2 * cp_size - 1) // (2 * cp_size)
+    pad = 2 * cp_size * chunk_size - len(tokens)
+    if isinstance(pad_value, Callable):
+        pad_func = pad_value
+        tokens = pad_func(tokens, pad)
+    else:
+        # pad on the first dimension
+        pad_tuple = (0, 0) * (tokens.dim() - 1) + (0, pad)
+        tokens = F.pad(tokens, pad_tuple, value=pad_value)
+    # get 2 chunk for thd cp
+    start_1, end_1 = chunk_size * cp_rank, chunk_size * (cp_rank + 1)
+    start_2, end_2 = chunk_size * (2 * cp_size - cp_rank - 1), chunk_size * (2 * cp_size - cp_rank)
+    return torch.cat([tokens[start_1:end_1], tokens[start_2:end_2]])
+
+
+def slice_log_prob_with_cp(
+    log_prob: list[float] | torch.Tensor,
+    total_length: int,
+    response_length: int,
+) -> list[float] | torch.Tensor:
+    assert len(log_prob) == response_length
+
+    cp_size = mpu.get_context_parallel_world_size()
+
+    if cp_size == 1:
+        return log_prob
+
+    prompt_length = total_length - response_length
+    _, _, logits_offset, _ = get_logits_and_tokens_offset_with_cp(total_length, response_length)
+
+    chunk_1 = log_prob[logits_offset[0][0] - (prompt_length - 1) : logits_offset[0][1] - (prompt_length - 1)]
+    chunk_2 = log_prob[logits_offset[1][0] - (prompt_length - 1) : logits_offset[1][1] - (prompt_length - 1)]
+
+    if isinstance(log_prob, list):
+        return chunk_1 + chunk_2
+    else:
+        return torch.cat([chunk_1, chunk_2], dim=0)

slime/backends/megatron_utils/data.py

@@ -0,0 +1,599 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from argparse import Namespace
+from collections.abc import Sequence
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from megatron.core import mpu
+from megatron.core.packed_seq_params import PackedSeqParams
+
+from slime.utils import train_metric_utils
+from slime.utils.data import get_minimum_num_micro_batch_size
+from slime.utils.flops_utils import calculate_fwd_flops
+from slime.utils.metric_utils import compute_pass_rate, compute_rollout_step
+from slime.utils.seqlen_balancing import get_seqlen_balanced_partitions
+from slime.utils.types import RolloutBatch
+
+from ...utils import tracking_utils
+from .cp_utils import get_sum_of_sample_mean, slice_with_cp
+
+logger = logging.getLogger(__name__)
+
+
+def get_batch(
+    data_iterator: "DataIterator",
+    keys: Sequence[str],
+    pad_multiplier: int = 128,
+) -> dict[str, torch.Tensor | PackedSeqParams | list[torch.Tensor] | None]:
+    """
+    Generate a CP-ready micro-batch with packed sequence parameters.
+
+    Steps:
+    - Fetch raw fields via iterator.
+    - Save original token tensors under "unconcat_tokens".
+    - Slice tokens into two chunks for Context Parallelism (CP), concatenate, and pad to a configurable multiple.
+    - Build cu_seqlens and `PackedSeqParams` with T-H-D layout (T: sequence length, H: attention heads, D: head dimension).
+
+    Args:
+        data_iterator: Iterator providing micro-batch data.
+        keys: List of keys to fetch from the iterator.
+        pad_multiplier: Multiplier for padding size calculation (default: 128).
+
+    Returns a dict including:
+    - "tokens": torch.LongTensor of shape [1, T_padded] on the current CUDA device
+    - "unconcat_tokens": list[torch.LongTensor] for the micro-batch before CP slicing/concat
+    - "packed_seq_params": PackedSeqParams with T-H-D settings (cu_seqlens on CUDA, dtype=int)
+    Plus any other requested keys forwarded from the iterator.
+    """
+
+    assert "tokens" in keys
+    batch = data_iterator.get_next(keys)
+
+    tokens = batch["tokens"]
+    # use 0 as the pad token id should be fine?
+    pad_token_id = 0
+
+    # for cp, we need all tokens to calculate logprob
+    batch["unconcat_tokens"] = tokens
+
+    cp_size = mpu.get_context_parallel_world_size()
+    tokens = [slice_with_cp(t, pad_token_id) for t in tokens]
+
+    cu_seqlens = [0]
+    for t in tokens:
+        cu_seqlens.append(cu_seqlens[-1] + t.size(0))
+
+    tokens = torch.cat(tokens)
+
+    # Always pad to reduce memory fragmentation and maybe make the computation faster
+    pad_size = mpu.get_tensor_model_parallel_world_size() * pad_multiplier
+    pad = (pad_size - tokens.size(0) % pad_size) % pad_size
+    if pad != 0:
+        tokens = F.pad(tokens, (0, pad), value=pad_token_id)
+        cu_seqlens.append(cu_seqlens[-1] + pad)
+
+    # thd requires the cu_seqlens to be of the origin length
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int).cuda() * cp_size
+    max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
+
+    packed_seq_params = PackedSeqParams(
+        cu_seqlens_q=cu_seqlens,
+        cu_seqlens_kv=cu_seqlens,
+        max_seqlen_q=max_seqlen,
+        max_seqlen_kv=max_seqlen,
+        qkv_format="thd",
+    )
+
+    tokens = tokens.unsqueeze(0)
+    batch["tokens"] = tokens
+    batch["packed_seq_params"] = packed_seq_params
+
+    # loss masks
+    loss_masks = []
+    for loss_mask, total_length, response_length in zip(
+        batch["loss_masks"],
+        batch["total_lengths"],
+        batch["response_lengths"],
+        strict=True,
+    ):
+        prompt_length = total_length - response_length
+        loss_mask = F.pad(loss_mask, (prompt_length - 1, 1), value=0)
+        loss_mask = slice_with_cp(loss_mask, 0)
+        loss_masks.append(loss_mask)
+    loss_masks = torch.cat(loss_masks)
+    loss_masks = F.pad(loss_masks, (0, pad), value=0).unsqueeze(0)
+    assert loss_masks.shape == tokens.shape, f"loss_masks.shape: {loss_masks.shape}, tokens.shape: {tokens.shape}"
+    batch["full_loss_masks"] = loss_masks
+
+    # Process multimodal training tensors if present
+    multimodal_train_inputs = batch.get("multimodal_train_inputs", None)
+    if multimodal_train_inputs is not None:
+        multimodal_data = {}  # key -> concatenated tensor
+        multimodal_num_items = {}  # key -> list of item counts per sequence
+        for mm_input_dict in multimodal_train_inputs:
+            if mm_input_dict is not None:
+                for key, mm_tensor in mm_input_dict.items():
+                    if key not in multimodal_data:
+                        multimodal_data[key] = mm_tensor
+                        multimodal_num_items[key] = [mm_tensor.size(0)]
+                    else:
+                        multimodal_data[key] = torch.cat([multimodal_data[key], mm_tensor], dim=0)
+                        multimodal_num_items[key].append(mm_tensor.size(0))
+        batch["multimodal_train_inputs"] = multimodal_data
+        batch["multimodal_num_items"] = multimodal_num_items
+
+    return batch
+
+
+def gather_log_data(
+    metric_name: str,
+    args: Namespace,
+    rollout_id: int,
+    log_dict: dict[str, float],
+) -> dict[str, float] | None:
+    """
+    Gather per-rank metrics, reduce by mean on the DP source rank, and log.
+
+    Expects `log_dict` to contain plain scalars. The DP source rank prints and
+    optionally logs to WandB/TensorBoard with a step derived from `rollout_id` and
+    batch sizes. Returns the reduced dict on the DP source rank; returns None on others.
+    """
+
+    if mpu.get_data_parallel_rank(with_context_parallel=True) == 0:
+        dp_size = mpu.get_data_parallel_world_size(with_context_parallel=True)
+
+        gathered_log_dict = [None] * dp_size
+        # Not sure if this will be a performance bottleneck.
+        dist.gather_object(
+            log_dict,
+            gathered_log_dict,
+            dst=mpu.get_data_parallel_src_rank(with_context_parallel=True),
+            group=mpu.get_data_parallel_group_gloo(with_context_parallel=True),
+        )
+
+        reduced_log_dict = {
+            f"{metric_name}/{key}": sum([d[key] for d in gathered_log_dict]) / dp_size for key in log_dict
+        }
+        logger.info(f"{metric_name} {rollout_id}: {reduced_log_dict}")
+
+        # Calculate step once to avoid duplication
+        step = compute_rollout_step(args, rollout_id)
+        reduced_log_dict["rollout/step"] = step
+        tracking_utils.log(args, reduced_log_dict, step_key="rollout/step")
+
+        return reduced_log_dict
+    else:
+        dist.gather_object(
+            log_dict,
+            None,
+            dst=mpu.get_data_parallel_src_rank(with_context_parallel=True),
+            group=mpu.get_data_parallel_group_gloo(with_context_parallel=True),
+        )
+        return None
+
+
+class DataIterator:
+    """Micro-batch iterator over rollout dicts.
+
+    Supports either fixed contiguous micro-batches or an explicit per-step
+    index schedule (for dynamic batch sizing / sequence-length balancing).
+    """
+
+    def __init__(
+        self,
+        rollout_data: RolloutBatch,
+        micro_batch_size: int | None = None,
+        micro_batch_indices: list[list[int]] | None = None,
+    ) -> None:
+        """Initialize an iterator over `rollout_data`.
+
+        Args:
+            rollout_data: Dict of per-sample fields for the local step.
+            micro_batch_size: Fixed contiguous slice size when not using dynamic scheduling.
+            micro_batch_indices: Explicit indices per micro-batch when using dynamic balancing.
+                Must be mutually exclusive with `micro_batch_size`.
+        """
+        self.rollout_data = rollout_data
+        self.micro_batch_size = micro_batch_size
+        self.micro_batch_indices = micro_batch_indices
+        assert micro_batch_size is None or micro_batch_indices is None
+        self.offset = 0
+
+    # Keys that are batch-level (not per-sample) and should be passed through as-is
+    BATCH_LEVEL_KEYS = set()
+    
+    def get_next(self, keys: Sequence[str]) -> dict[str, list[object] | None]:
+        """Return the next micro-batch for the requested keys.
+
+        - If `micro_batch_indices` is provided, selects rows according to the current
+          index list for each requested key.
+        - Otherwise, slices a contiguous window of size `micro_batch_size` starting
+          at the current offset.
+
+        Returns a dict mapping each key to a list subset (or None if absent).
+        """
+        batch = {}
+        for key in keys:
+            vals = self.rollout_data.get(key, None)
+            if vals is None:
+                batch[key] = None
+            elif key in self.BATCH_LEVEL_KEYS:
+                # Batch-level keys are not per-sample, pass through as-is
+                batch[key] = vals
+            else:
+                if self.micro_batch_indices is not None:
+                    indices = self.micro_batch_indices[self.offset]
+                    batch[key] = [vals[i] for i in indices]
+                else:
+                    assert self.offset + self.micro_batch_size <= len(
+                        vals
+                    ), f"offset: {self.offset}, micro_batch_size: {self.micro_batch_size}, len(vals): {len(vals)}"
+                    batch[key] = vals[self.offset : self.offset + self.micro_batch_size]
+
+        if self.micro_batch_indices is not None:
+            self.offset += 1
+        else:
+            self.offset += self.micro_batch_size
+        return batch
+
+    def reset(self) -> "DataIterator":
+        """Reset internal offset to the start and return self."""
+        self.offset = 0
+        return self
+
+
+def get_data_iterator(
+    args: Namespace,
+    model: torch.nn.Module | Sequence[torch.nn.Module],
+    rollout_data: RolloutBatch,
+) -> tuple[list[DataIterator], list[int]]:
+    """
+    Create iterators and a micro-batch schedule for a rollout step.
+
+    - If `use_dynamic_batch_size` is False, splits into fixed-size contiguous
+      micro-batches of `micro_batch_size`.
+    - If True, computes the number of micro-batches per local step based on
+      `max_tokens_per_gpu` and per-sample lengths, all-reduces to a DP-wide
+      maximum, optionally enforces divisibility for Virtual Pipeline Parallelism (VPP), and builds a balanced
+      index schedule to equalize token counts across micro-batches.
+
+    Returns `(data_iterators, num_microbatches)` where:
+    - `data_iterators`: list of `DataIterator`, one per VPP stage (size 1 if VPP disabled)
+    - `num_microbatches`: list[int], one per local step in the rollout (length = steps)
+    """
+    dp_size = mpu.get_data_parallel_world_size(with_context_parallel=False)
+    dp_group = mpu.get_data_parallel_group()
+    vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size()
+    if vpp_size is None:
+        vpp_size = 1
+    if vpp_size > 1:
+        from megatron.core.utils import get_model_config
+
+        config = get_model_config(model[0])
+        microbatch_group_size_per_vp_stage = config.microbatch_group_size_per_vp_stage
+    cp_size = mpu.get_context_parallel_world_size()
+
+    num_local_samples = len(rollout_data["total_lengths"])
+    num_local_gbs = args.global_batch_size // dp_size
+    num_steps_per_rollout = num_local_samples // num_local_gbs
+
+    def _generate_data_iterator(rollout_data, micro_batch_size, micro_batch_indices=None):
+        data_iterator = []
+        for _ in range(vpp_size):
+            data_iterator.append(DataIterator(rollout_data, micro_batch_size, micro_batch_indices))
+        return data_iterator
+
+    if not args.use_dynamic_batch_size:
+        num_microbatches = [num_local_gbs // args.micro_batch_size for _ in range(num_steps_per_rollout)]
+        data_iterator = _generate_data_iterator(rollout_data, args.micro_batch_size)
+    else:
+        assert args.max_tokens_per_gpu is not None
+        # calculate the number of mirobatches for each step
+        samples = rollout_data["total_lengths"]
+        assert len(samples) == num_local_samples
+        num_microbatches = []
+        for i in range(num_steps_per_rollout):
+            start, end = i * num_local_gbs, (i + 1) * num_local_gbs
+            num_microbatches.append(
+                get_minimum_num_micro_batch_size(samples[start:end], args.max_tokens_per_gpu * cp_size)
+            )
+
+        num_microbatches = torch.tensor(num_microbatches, dtype=torch.int, device=torch.cuda.current_device())
+        dist.all_reduce(num_microbatches, op=dist.ReduceOp.MAX, group=dp_group)
+
+        if vpp_size > 1:
+            # vpp requies the number of microbatches to be divisible by vpp_size
+            num_microbatches = torch.clamp(
+                num_microbatches // microbatch_group_size_per_vp_stage * microbatch_group_size_per_vp_stage,
+                min=1,
+            )
+
+        num_microbatches = num_microbatches.tolist()
+
+        # balance the each micro batch
+        samples = rollout_data["total_lengths"]
+        # balance the number of mirobatches across steps
+        micro_batch_indices = []
+        for i, num_mbs in enumerate(num_microbatches):
+            start, end = i * num_local_gbs, (i + 1) * num_local_gbs
+            samples = rollout_data["total_lengths"][start:end]
+            partitions = get_seqlen_balanced_partitions(samples, num_mbs, equal_size=False)
+            for j in range(num_mbs):
+                for k in range(len(partitions[j])):
+                    partitions[j][k] += start
+            micro_batch_indices.extend(partitions)
+
+        assert len(set(sum(micro_batch_indices, []))) == num_local_samples
+
+        data_iterator = _generate_data_iterator(rollout_data, None, micro_batch_indices)
+
+    return (
+        data_iterator,
+        num_microbatches,
+    )
+
+
+def log_rollout_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) -> None:
+    """
+    Summarize rollout fields and log reduced metrics on PP last stage, TP rank 0.
+
+    - Tensor-valued lists are concatenated and averaged. For token-level metrics
+      like log-probs/returns/advantages/values, computes a CP-correct sample mean
+      using `loss_masks` and total/response lengths.
+    - Non-tensor lists are averaged elementwise.
+    - Scalars are converted to Python numbers.
+    """
+    if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():
+        cp_size = mpu.get_context_parallel_world_size()
+        log_dict = {}
+        response_lengths = rollout_data["response_lengths"]
+        loss_masks = rollout_data["loss_masks"]
+        total_lengths = rollout_data["total_lengths"]
+
+        for key, val in rollout_data.items():
+            if key in [
+                "tokens",
+                "multimodal_train_inputs",
+                "loss_masks",
+                "sample_indices",
+                "rollout_routed_experts",
+            ]:
+                continue
+            # Skip None values
+            if val is None:
+                continue
+            # Upload per sample mean for each rollout value
+            # There are the following assumptions:
+            # - Each dp rank has the same number of samples
+            if isinstance(val, (list, tuple)):
+                # Filter out None entries before processing.
+                val = [v for v in val if v is not None]
+                if not val:
+                    continue
+                if all(isinstance(v, torch.Tensor) for v in val):
+                    # NOTE: Here we have to do the clone().detach(), otherwise the tensor will be
+                    # modified in place and will cause problem for the next rollout.
+                    val = torch.cat(val).clone().detach()
+                    if key in ["log_probs", "ref_log_probs", "rollout_log_probs", "returns", "advantages", "values"]:
+                        sum_of_sample_mean = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks)
+                        val = cp_size * sum_of_sample_mean(val) / len(loss_masks)
+                    else:
+                        val = val.mean() * cp_size
+                else:
+                    # Mixed Tensor/scalar list.
+                    # Convert everything to float scalar for logging.
+                    val = sum(float(v.mean()) if isinstance(v, torch.Tensor) else float(v) for v in val) / len(val)
+            elif isinstance(val, torch.Tensor):
+                val = val.float().mean()
+            else:
+                raise ValueError(f"Unsupported type: {type(val)} for key: {key}")
+            log_dict[key] = val.item() if isinstance(val, torch.Tensor) else val
+
+        reduced_log_dict = gather_log_data("rollout", args, rollout_id, log_dict)
+        if args.ci_test and reduced_log_dict is not None:
+            if (
+                rollout_id == 0
+                and "rollout/log_probs" in reduced_log_dict
+                and "rollout/ref_log_probs" in reduced_log_dict
+            ):
+                assert reduced_log_dict["rollout/log_probs"] == reduced_log_dict["rollout/ref_log_probs"]
+            if "rollout/log_probs" in reduced_log_dict:
+                assert -0.5 < reduced_log_dict["rollout/log_probs"] < 0
+            if "rollout/entropy" in reduced_log_dict:
+                assert 0 < reduced_log_dict["rollout/entropy"] < 0.5
+
+    if args.log_multi_turn:
+        log_multi_turn_data(rollout_id, args, rollout_data)
+    if args.log_passrate:
+        log_passrate(rollout_id, args, rollout_data)
+
+    if args.log_correct_samples:
+        if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():
+            cp_size = mpu.get_context_parallel_world_size()
+            log_dict = {}
+            response_lengths = rollout_data["response_lengths"]
+            loss_masks = rollout_data["loss_masks"]
+            total_lengths = rollout_data["total_lengths"]
+
+            def quantile(total_value, n_quantiles, data) -> dict:
+                import math
+
+                assert n_quantiles > 1, f"n_quantiles({n_quantiles}) must be greater than 1."
+
+                quantiles = [((i + 1) / n_quantiles) for i in range(n_quantiles)]
+                cut_points = [total_value * q for q in quantiles]
+                cut_points[-1] = total_value
+
+                count = [0] * n_quantiles
+                for d in data:
+                    for i, point in enumerate(cut_points):
+                        if d <= point:
+                            count[i] += 1
+                            break
+
+                total = sum(count) + 1e-9
+                percentile = [c / total for c in count]
+
+                percentile = {f"p{min(math.ceil(q*100),100)}": p for q, p in zip(quantiles, percentile, strict=True)}
+                return percentile
+
+            raw_rewards = rollout_data["raw_reward"]
+            # Additional metrics for correct cases are calculated separately below.
+            correct_response_lengths = []
+            correct_total_lengths = []
+            correct_loss_masks = []
+            correct_entropy = []
+            for i, raw_reward in enumerate(raw_rewards):
+                if raw_reward == 1:
+                    correct_response_lengths.append(response_lengths[i])
+                    correct_total_lengths.append(total_lengths[i])
+                    correct_loss_masks.append(loss_masks[i])
+                    correct_entropy.append(-rollout_data["log_probs"][i])
+            num_correct_responses = len(correct_total_lengths)
+            rollout_data["correct_response_lengths"] = correct_response_lengths
+            correct_response_length_percentile = quantile(
+                args.rollout_max_response_len, 4, rollout_data["correct_response_lengths"]
+            )
+            for p, val in correct_response_length_percentile.items():
+                rollout_data[f"correct_length/{p}"] = [val] * num_correct_responses
+            if len(correct_entropy) > 0:
+                sum_of_sample_mean = get_sum_of_sample_mean(
+                    correct_total_lengths, correct_response_lengths, correct_loss_masks
+                )
+                correct_entropy = sum_of_sample_mean(torch.cat(correct_entropy, dim=0))
+                rollout_data["correct_entropy"] = [correct_entropy.item()] * num_correct_responses
+            else:
+                rollout_data["correct_entropy"] = [0] * num_correct_responses
+
+
+def log_multi_turn_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) -> None:
+    """
+    Log multi-turn auxiliary metrics such as raw/observed response lengths and rounds.
+
+    Operates only on PP last stage and TP rank 0. Uses GPU tensors when available
+    to compute statistics without host transfers.
+    """
+    if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():
+        log_dict = {}
+        for key, val in rollout_data.items():
+            if key == "loss_masks":
+                if val:  # Check if val is not empty
+                    device = val[0].device  # Get device from first tensor
+
+                    # Vectorized length calculation using torch
+                    raw_response_lengths = torch.tensor([v.shape[0] for v in val], dtype=torch.float32, device=device)
+                    log_dict["raw_response_length/response_length_mean"] = raw_response_lengths.mean().item()
+                    log_dict["raw_response_length/response_length_max"] = raw_response_lengths.max().item()
+                    log_dict["raw_response_length/response_length_min"] = raw_response_lengths.min().item()
+                    log_dict["raw_response_length/response_length_clip_ratio"] = (
+                        (raw_response_lengths >= args.rollout_max_response_len).float().mean().item()
+                    )
+
+                    # Vectorized sum calculation using torch - stay on GPU
+                    wo_obs_response_lengths = torch.tensor(
+                        [v.sum().item() for v in val], dtype=torch.float32, device=device
+                    )
+                    log_dict["wo_obs_response_length/response_length_mean"] = wo_obs_response_lengths.mean().item()
+                    log_dict["wo_obs_response_length/response_length_max"] = wo_obs_response_lengths.max().item()
+                    log_dict["wo_obs_response_length/response_length_min"] = wo_obs_response_lengths.min().item()
+            if key == "round_number":
+                # Use numpy for vectorized round number statistics
+                round_number_array = np.array(val)
+                log_dict["multi_turn_metric/round_number_mean"] = np.mean(round_number_array)
+                log_dict["multi_turn_metric/round_number_max"] = np.max(round_number_array)
+                log_dict["multi_turn_metric/round_number_min"] = np.min(round_number_array)
+        gather_log_data("multi_turn", args, rollout_id, log_dict)
+
+
+def log_passrate(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) -> None:
+    """
+    Compute pass@k metrics from `raw_reward` groups and log the results.
+
+    `raw_reward` is reshaped to `[group_number, group_size]`, then pass@k is
+    estimated per problem and averaged.
+    """
+    if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():
+        log_dict = {}
+        for key, val in rollout_data.items():
+            if key != "raw_reward":
+                continue
+
+            log_dict |= compute_pass_rate(
+                flat_rewards=val,
+                group_size=args.n_samples_per_prompt,
+                num_groups=args.rollout_batch_size,
+            )
+
+        gather_log_data("passrate", args, rollout_id, log_dict)
+
+
+def log_perf_data(rollout_id: int, args: Namespace) -> None:
+    train_metric_utils.log_perf_data_raw(
+        rollout_id=rollout_id,
+        args=args,
+        is_primary_rank=(
+            mpu.get_tensor_model_parallel_rank() == 0
+            and mpu.is_pipeline_last_stage()
+            and mpu.get_data_parallel_rank(with_context_parallel=True) == 0
+        ),
+        compute_total_fwd_flops=lambda seq_lens: calculate_fwd_flops(seqlens=seq_lens, args=args)
+        / dist.get_world_size()
+        / 1e12,
+    )
+
+
+def sync_actor_critic_data(
+    args: Namespace,
+    rollout_data: RolloutBatch | None = None,
+    group: dist.ProcessGroup | None = None,
+) -> None:
+    """
+    Broadcast `values` (from critic) and optionally `log_probs`/`ref_log_probs`
+    (from actor) across PP ranks to align data dependencies.
+
+    - Values are broadcast from src=1.
+    - Log-probs and ref-log-probs are broadcast from src=0 when KL is used.
+    Updates `rollout_data` in place with the synchronized tensors.
+    """
+    log_probs_key = "log_probs" if not args.use_rollout_logprobs else "rollout_log_probs"
+    values, log_probs, ref_log_probs = map(rollout_data.get, ("values", log_probs_key, "ref_log_probs"))
+
+    # return when not the pp last stage
+    if not values and not log_probs:
+        return
+
+    handles = []
+
+    if not values:
+        values = [torch.empty_like(log_prob) for log_prob in log_probs]
+    for value in values:
+        handles.append(dist.broadcast(value, src=1, group=group, async_op=True))
+
+    if args.kl_coef != 0 or args.use_kl_loss:
+        if not log_probs:
+            log_probs = [torch.empty_like(value) for value in values]
+        if not ref_log_probs:
+            ref_log_probs = [torch.empty_like(value) for value in values]
+        for ref_log_prob, log_prob in zip(ref_log_probs, log_probs, strict=False):
+            handles.append(dist.broadcast(log_prob, src=0, group=group, async_op=True))
+            handles.append(dist.broadcast(ref_log_prob, src=0, group=group, async_op=True))
+
+    for handle in handles:
+        handle.wait()
+
+    rollout_data.update(
+        {
+            k: v
+            for k, v in {
+                "values": values,
+                log_probs_key: log_probs,
+                "ref_log_probs": ref_log_probs,
+            }.items()
+            if v is not None
+        }
+    )

slime/backends/megatron_utils/initialize.py

@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import random
+
+import numpy as np
+import torch
+from megatron.core import mpu, tensor_parallel
+from megatron.core.config import set_experimental_flag
+from megatron.core.num_microbatches_calculator import init_num_microbatches_calculator
+from megatron.training.global_vars import _build_tokenizer, set_args
+
+logger = logging.getLogger(__name__)
+
+
+def _set_random_seed(
+    seed_: int,
+    data_parallel_random_init: bool = False,
+    te_rng_tracker: bool = False,
+    inference_rng_tracker: bool = False,
+    use_cudagraphable_rng: bool = False,
+):
+    """Set random seed for reproducability."""
+    # Ensure that different pipeline MP stages get different seeds.
+    seed = seed_ + (100 * mpu.get_pipeline_model_parallel_rank())
+    # Ensure different data parallel ranks get different seeds
+    if data_parallel_random_init:
+        seed = seed + (10 * mpu.get_data_parallel_rank(with_context_parallel=False))
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    tensor_parallel.model_parallel_cuda_manual_seed(seed, te_rng_tracker, inference_rng_tracker, use_cudagraphable_rng)
+
+
+def _initialize_distributed(args, get_embedding_ranks=None, get_position_embedding_ranks=None):
+    """Initialize torch.distributed and core model parallel."""
+    # Set the tensor model-parallel, pipeline model-parallel, and
+    # data-parallel communicators.
+    mpu.initialize_model_parallel(
+        args.tensor_model_parallel_size,
+        args.pipeline_model_parallel_size,
+        args.virtual_pipeline_model_parallel_size,
+        pipeline_model_parallel_comm_backend=args.pipeline_model_parallel_comm_backend,
+        context_parallel_size=args.context_parallel_size,
+        hierarchical_context_parallel_sizes=args.hierarchical_context_parallel_sizes,
+        expert_model_parallel_size=args.expert_model_parallel_size,
+        num_distributed_optimizer_instances=args.num_distributed_optimizer_instances,
+        expert_tensor_parallel_size=args.expert_tensor_parallel_size,
+        distributed_timeout_minutes=args.distributed_timeout_minutes,
+        nccl_communicator_config_path=args.nccl_communicator_config_path,
+        order="tp-cp-ep-dp-pp" if not args.use_tp_pp_dp_mapping else "tp-cp-ep-pp-dp",
+        get_embedding_ranks=get_embedding_ranks,
+        get_position_embedding_ranks=get_position_embedding_ranks,
+        create_gloo_process_groups=args.enable_gloo_process_groups,
+    )
+
+
+def init(args):
+    set_args(args)
+    if args.enable_experimental:
+        logger.info("Enable megatron experimental")
+        set_experimental_flag(True)
+
+    # Pytorch distributed.
+    _initialize_distributed(args)
+
+    # https://github.com/NVIDIA/Megatron-LM/issues/1563
+    assert np.__version__.startswith("1."), "Megatron does not support numpy 2.x"
+
+    # Random seeds for reproducibility.
+    if args.rank == 0:
+        logger.info(f"> setting random seeds to {args.seed} ...")
+    _set_random_seed(
+        args.seed,
+        args.data_parallel_random_init,
+        args.te_rng_tracker,
+        args.inference_rng_tracker,
+    )
+    _build_tokenizer(args)
+    # We won't use this. initialize to pass some validation in megatron.
+    init_num_microbatches_calculator(
+        args.rank,
+        args.rampup_batch_size,
+        args.global_batch_size,
+        args.micro_batch_size,
+        args.data_parallel_size,
+        args.decrease_batch_size_if_needed,
+    )
+
+    if args.deterministic_mode:
+        if args.rank == 0:
+            logger.info("> running in deterministic mode")
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+        torch.use_deterministic_algorithms(True, warn_only=False)
+
+    if args.tp_comm_overlap:
+        from megatron.training.initialize import _initialize_tp_communicators
+
+        _initialize_tp_communicators()
+
+    if getattr(args, "custom_megatron_init_path", None):
+        from slime.utils.misc import load_function
+
+        custom_init = load_function(args.custom_megatron_init_path)
+        custom_init(args)
+
+
+# TODO shall we use a simpler method to determine which rank to init wandb?
+def is_megatron_main_rank():
+    return (
+        mpu.get_data_parallel_rank(with_context_parallel=True) == 0
+        and mpu.get_tensor_model_parallel_rank() == 0
+        and mpu.get_pipeline_model_parallel_rank() == mpu.get_pipeline_model_parallel_world_size() - 1
+    )

slime/backends/megatron_utils/kernels/int4_qat/fake_int4_quant_cuda.cu

@@ -0,0 +1,368 @@
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#define FINAL_MASK 0xFFFFFFFF
+
+__device__ __host__ __forceinline__
+int ceil_div(int a, int b) {
+    return (a + b - 1) / b;
+}
+
+__device__ __forceinline__
+float warpReduceMax(float val) {
+#pragma unroll
+    for (int mask = 16; mask > 0; mask >>= 1)
+        val = fmaxf(val, __shfl_xor_sync(FINAL_MASK, val, mask, 32));
+    return val;
+}
+
+
+__device__ __forceinline__
+float warpReduceMin(float val) {
+#pragma unroll
+    for (int mask = 16; mask > 0; mask >>= 1)
+        val = fminf(val, __shfl_xor_sync(FINAL_MASK, val, mask, 32));
+    return val;
+}
+
+// almost all int4 use blocksize = [1, 32]
+template<typename scalar_t>
+__global__
+void int4_quant_1x32_kernel(
+    const scalar_t* __restrict__ x,
+    scalar_t* __restrict__ out,
+    scalar_t* out_scale,
+    scalar_t* out_zero,
+    const int M, const int N,
+    const int stride_xm, const int stride_xn,
+    const int stride_om, const int stride_on,
+    const int stride_osm, const int stride_osn,
+    const int stride_ozm, const int stride_ozn,
+    bool sym
+) {
+    constexpr int WARPS_PER_BLOCK = 8;
+    const int needed_warps = ceil_div(N, 32);
+    
+    const int tid = threadIdx.x;
+    const int warp_id = tid >> 5;
+    const int lane_id = tid & 0x1F;
+    constexpr float SYM_CONS = 1.0f / 7.0f;
+    constexpr float ASYM_CONS = 1.0f / 15.0f;
+    
+    const int row = blockIdx.x;
+
+    for (int item = warp_id; item < needed_warps; item += WARPS_PER_BLOCK) {
+        const int col = item * 32 + lane_id;
+        float val = 0.0f;
+
+        if (col < N) {
+            val = static_cast<float>(x[row * stride_xm + col * stride_xn]);
+        }
+        
+        float scale = 0.0f;
+        float zero = 0.0f;
+        
+        if (sym) {
+            float abs_val = fabsf(val);
+            
+            float block_max = warpReduceMax(abs_val);
+            
+            scale = fmaxf(block_max * SYM_CONS, 1e-5f);
+
+            val = rintf(val / scale);
+        } else {
+            float block_min = warpReduceMin(val);
+            float block_max = warpReduceMax(val);
+
+            scale = fmaxf((block_max - block_min) * ASYM_CONS, 1e-5f);
+            zero = fminf(fmaxf(-rintf(block_min / scale), 0.0f), 15.0f);
+            
+            val = rintf(val / scale) + zero;
+        }
+        
+        if (col < N) {
+            out[row * stride_om + col * stride_on] = static_cast<scalar_t>(val);
+            out_scale[row * stride_osm + item * stride_osn] = static_cast<scalar_t>(scale);
+            if(!sym) {
+                out_zero[row * stride_ozm + item * stride_ozn] = static_cast<scalar_t>(zero);
+            }
+        }
+    }
+}
+
+// for some transpose case, blocksize = [32, 1]
+template<typename scalar_t>
+__global__
+void int4_quant_32x1_kernel(
+    const scalar_t* __restrict__ x,
+    scalar_t* __restrict__ out,
+    scalar_t* out_scale,
+    scalar_t* out_zero,
+    const int M, const int N,
+    const int stride_xm, const int stride_xn,
+    const int stride_om, const int stride_on,
+    const int stride_osm, const int stride_osn,
+    const int stride_ozm, const int stride_ozn,
+    bool sym
+) {
+    constexpr int WARPS_PER_BLOCK = 8;
+    const int start_row = blockIdx.x * 32;
+    const int end_row = min((blockIdx.x + 1) * 32, M);
+    
+    const int tid = threadIdx.x;
+    const int warp_id = tid >> 5;
+    const int lane_id = tid & 0x1F;
+    constexpr float SYM_CONS = 1.0f / 7.0f;
+    constexpr float ASYM_CONS = 1.0f / 15.0f;
+    
+    for (int item = warp_id; item < N; item += WARPS_PER_BLOCK) {
+        const int col = item;
+        const int row = start_row + lane_id;
+
+        float val = 0.0f;
+
+        if (row < end_row) {
+            val = static_cast<float>(x[row * stride_xm + col * stride_xn]);
+        }
+    
+        float scale = 0.0f;
+        float zero = 0.0f;
+        
+        if (sym) {
+            float abs_val = fabsf(val);
+            
+            float block_max = warpReduceMax(abs_val);
+            
+            scale = fmaxf(block_max * SYM_CONS, 1e-5f);
+
+            val = rintf(val / scale);
+        } else {
+            float block_min = warpReduceMin(val);
+            float block_max = warpReduceMax(val);
+
+            scale = fmaxf((block_max - block_min) * ASYM_CONS, 1e-5f);
+            zero = fminf(fmaxf(-rintf(block_min / scale), 0.0f), 15.0f);
+            
+            val = rintf(val / scale) + zero;
+        }
+        
+        if (row < end_row) {
+            out[row * stride_om + col * stride_on] = static_cast<scalar_t>(val);
+            out_scale[blockIdx.x * stride_osm + item * stride_osn] = static_cast<scalar_t>(scale);
+            if (!sym) {
+                out_zero[blockIdx.x * stride_ozm + item * stride_ozn] = static_cast<scalar_t>(zero);
+            }
+        }
+    }
+}
+
+template<typename scalar_t>
+__global__ void int4_quant_common_kernel(
+    const scalar_t* __restrict__ x,
+    scalar_t* __restrict__ out,
+    scalar_t* out_scale,
+    scalar_t* out_zero,
+    const int M, const int N,
+    const int stride_xm, const int stride_xn,
+    const int stride_om, const int stride_on,
+    const int stride_osm, const int stride_osn,
+    const int stride_ozm, const int stride_ozn,
+    const int BLOCK_M, const int BLOCK_N,
+    bool sym
+) {
+    const int start_row = blockIdx.x * BLOCK_M;
+    const int WARPS_PER_BLOCK = blockDim.x >> 5;
+    
+    const int warp_id = threadIdx.x >> 5;
+    const int lane_id = threadIdx.x & 0x1F;
+    constexpr float SYM_CONS = 1.0f / 7.0f;
+    constexpr float ASYM_CONS = 1.0f / 15.0f;
+    constexpr int WARP_SIZE = 32;
+    
+    const int needed_warps = ceil_div(N, BLOCK_N);
+    const int iters = ceil_div(BLOCK_M * BLOCK_N, 32);
+    int warp_rows = 1;
+
+    if (BLOCK_N <= WARP_SIZE) {
+        warp_rows = WARP_SIZE / BLOCK_N;
+    }
+    
+    for (int item = warp_id; item < needed_warps; item += WARPS_PER_BLOCK) {
+        float local_max = -INFINITY;
+        float local_min = INFINITY;
+
+        float val = 0.0f;
+        float scale, zero = 0.0f;
+
+        const int row_off = lane_id / BLOCK_N;
+        const int col_off = lane_id % BLOCK_N;
+        int row, col = 0;
+        
+        for (int i = 0; i < iters; ++i) {
+            if (BLOCK_N <= WARP_SIZE) {
+                row = start_row + i * warp_rows + row_off;
+                col = item * BLOCK_N + col_off;
+            } else {
+                row = start_row;
+                col = item * BLOCK_N + i * WARP_SIZE + col_off;
+            }
+
+            if (row < M && col < N) {
+                val = static_cast<float>(x[row * stride_xm + col * stride_xn]);
+            } else {
+                val = 0.0f;
+            }
+
+            if (sym) {
+                local_max = fmaxf(local_max, fabsf(val));
+            } else {
+                local_max = fmaxf(local_max, val);
+                local_min = fminf(local_min, val);
+            }
+        }
+
+        if (sym) {
+            float block_max = warpReduceMax(local_max);
+            scale = fmaxf(block_max * SYM_CONS, 1e-5f);
+        } else {
+            float block_max = warpReduceMax(local_max);
+            float block_min = warpReduceMin(local_min);
+            scale = fmaxf((block_max - block_min) * ASYM_CONS, 1e-5f);
+            zero = fminf(fmaxf(-rintf(block_min / scale), 0.0f), 15.0f);
+        }
+
+        for (int i = 0; i < iters; ++i) {
+            if (BLOCK_N <= WARP_SIZE) {
+                row = start_row + i * warp_rows + row_off;
+                col = item * BLOCK_N + col_off;
+            } else {
+                row = start_row;
+                col = item * BLOCK_N + i * WARP_SIZE + col_off;
+            }
+
+            if (row < M && col < N) {
+                float val = static_cast<float>(x[row * stride_xm + col * stride_xn]);
+                if (sym) {
+                    val = rintf(val / scale);
+                } else {
+                    val = rintf(val / scale) + zero;
+                }
+                
+                out[row * stride_om + col * stride_on] = static_cast<scalar_t>(val);
+                out_scale[blockIdx.x * stride_osm + item * stride_osn] = static_cast<scalar_t>(scale);
+                if (!sym) {
+                    out_zero[blockIdx.x * stride_ozm + item * stride_ozn] = static_cast<scalar_t>(zero);
+                }
+            }
+        }
+    }
+}
+
+// dispatch
+template<typename scalar_t>
+void launch_int4_quant_kernel(
+    const scalar_t* x,
+    scalar_t* out,
+    scalar_t* out_scale,
+    scalar_t* out_zero,
+    int M, int N,
+    const int stride_xm, const int stride_xn,
+    const int stride_om, const int stride_on,
+    const int stride_osm, const int stride_osn,
+    const int stride_ozm, const int stride_ozn,
+    int block_m, int block_n,
+    bool sym,
+    cudaStream_t stream
+) {
+    constexpr int WARPS_PER_BLOCK = 8;
+    constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * 32;  // 256
+
+    if (block_m == 1 && block_n == 32) {
+        dim3 grid(M);
+        dim3 block(THREADS_PER_BLOCK);
+        
+        int4_quant_1x32_kernel<scalar_t><<<grid, block, 0, stream>>>(
+            x, out, out_scale, out_zero, M, N,
+            stride_xm, stride_xn,
+            stride_om, stride_on,
+            stride_osm, stride_osn,
+            stride_ozm, stride_ozn,
+            sym
+        );
+    } else if (block_m == 32 && block_n == 1) {
+        dim3 grid(ceil_div(M, block_m));
+        dim3 block(THREADS_PER_BLOCK);
+
+        int4_quant_32x1_kernel<scalar_t><<<grid, block, 0, stream>>>(
+            x, out, out_scale, out_zero, M, N,
+            stride_xm, stride_xn,
+            stride_om, stride_on,
+            stride_osm, stride_osn,
+            stride_ozm, stride_ozn,
+            sym
+        );
+    } else {
+        dim3 grid(ceil_div(M, block_m));
+        dim3 block(THREADS_PER_BLOCK);
+        int4_quant_common_kernel<scalar_t><<<grid, block, 0, stream>>>(
+            x, out, out_scale, out_zero, M, N,
+            stride_xm, stride_xn,
+            stride_om, stride_on,
+            stride_osm, stride_osn,
+            stride_ozm, stride_ozn,
+            block_m, block_n,
+            sym
+        );
+    }
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
+fake_int4_quant_cuda(
+    torch::Tensor& x,
+    std::vector<int64_t>& block_size,
+    bool sym
+) {
+    TORCH_CHECK(x.dim() == 2, "Input must be 2D");
+    TORCH_CHECK(x.is_cuda(), "Input must be on CUDA");
+    
+    int M = x.size(0);
+    int N = x.size(1);
+    int block_m = block_size[0];
+    int block_n = block_size[1];
+
+    TORCH_CHECK(block_m > 0 && block_n > 0, "Block sizes must be positive, got block_m=", block_m, ", block_n=", block_n);
+    TORCH_CHECK((block_m * block_n) % 32 == 0,
+        "block_m * block_n (", block_m * block_n, ") must be divisible by 32. "
+        "But got a ", block_m, "x", block_n, " block.");
+
+    auto out = torch::empty_like(x);
+    auto out_scale = torch::empty({ceil_div(M, block_m), ceil_div(N, block_n)}, x.options());
+    auto out_zero = torch::empty_like(out_scale);
+
+    const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    
+    AT_DISPATCH_FLOATING_TYPES_AND(
+        at::ScalarType::BFloat16,
+        x.scalar_type(), "int4_quant_cuda", [&] {
+        launch_int4_quant_kernel<scalar_t>(
+            x.const_data_ptr<scalar_t>(),
+            out.data_ptr<scalar_t>(),
+            out_scale.data_ptr<scalar_t>(),
+            out_zero.data_ptr<scalar_t>(),
+            M, N,
+            x.stride(0), x.stride(1),
+            out.stride(0), out.stride(1),
+            out_scale.stride(0), out_scale.stride(1),
+            out_zero.stride(0), out_zero.stride(1),
+            block_m, block_n,
+            sym,
+            stream
+        );
+    });
+    
+    return std::make_tuple(out, out_scale, out_zero);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("fake_int4_quant_cuda", &fake_int4_quant_cuda, "fake INT4 quantization cuda");
+}

slime/backends/megatron_utils/kernels/int4_qat/setup.py

@@ -0,0 +1,42 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+import torch
+
+# Get CUDA arch list
+arch_list = []
+if torch.cuda.is_available():
+    for i in range(torch.cuda.device_count()):
+        major, minor = torch.cuda.get_device_capability(i)
+        arch_list.append(f"{major}.{minor}")
+    arch_list = sorted(set(arch_list))
+
+setup(
+    name="fake_int4_quant_cuda",
+    ext_modules=[
+        CUDAExtension(
+            name="fake_int4_quant_cuda",
+            sources=["fake_int4_quant_cuda.cu"],
+            extra_compile_args={
+                "cxx": [
+                    "-O3",
+                    "-std=c++17",
+                ],
+                "nvcc": [
+                    "-O3",
+                    "-std=c++17",
+                    "--expt-relaxed-constexpr",
+                    "-Xcompiler",
+                    "-fPIC",
+                ]
+                + [
+                    f'-gencode=arch=compute_{arch.replace(".", "")},code=sm_{arch.replace(".", "")}'
+                    for arch in arch_list
+                ],
+            },
+        )
+    ],
+    cmdclass={"build_ext": BuildExtension},
+)

slime/backends/megatron_utils/loss.py

@@ -0,0 +1,768 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from argparse import Namespace
+from collections.abc import Callable, Iterator
+from typing import Any
+
+import torch
+
+logger = logging.getLogger(__name__)
+from megatron.core import mpu
+from torch.utils.checkpoint import checkpoint
+
+from slime.utils.distributed_utils import distributed_masked_whiten
+from slime.utils.misc import load_function
+from slime.utils.ppo_utils import (
+    calculate_log_probs_and_entropy,
+    compute_approx_kl,
+    compute_gspo_kl,
+    compute_opsm_mask,
+    compute_policy_loss,
+    get_advantages_and_returns_batch,
+    get_grpo_returns,
+    get_reinforce_plus_plus_baseline_advantages,
+    get_reinforce_plus_plus_returns,
+)
+from slime.utils.types import RolloutBatch
+
+from .cp_utils import all_gather_with_cp, get_logits_and_tokens_offset_with_cp, get_sum_of_sample_mean
+
+
+def get_responses(
+    logits: torch.Tensor,
+    *,
+    args: Namespace,
+    unconcat_tokens: list[torch.Tensor],
+    total_lengths: list[int],
+    response_lengths: list[int],
+) -> Iterator[tuple[torch.Tensor, torch.Tensor]]:
+    """Yield response-aligned `(logits_chunk, tokens_chunk)` pairs per sample.
+
+    After squeezing batch dimension and applying temperature scaling, this
+    function extracts the logits and tokens corresponding to response segments
+    for each sample. When context parallelism is disabled, it slices directly
+    from the concatenated sequence. With context parallelism enabled, it
+    handles split sequences across ranks.
+
+    Args:
+        logits: Model outputs with shape `[1, T, V]` (policy) or `[1, T, 1]`
+            (value). Must be float32.
+        args: Configuration containing `rollout_temperature` for scaling.
+        unconcat_tokens: List of token tensors (prompt+response) per sample.
+        total_lengths: Total sequence lengths (prompt+response) per sample.
+        response_lengths: Response segment lengths per sample.
+
+    Yields:
+        Tuple of `(logits_chunk, tokens_chunk)` where `logits_chunk` is shape
+        `[R, V]` (policy) or `[R, 1]` (value) and `tokens_chunk` is shape `[R]`
+        (1D int64), both aligned to response tokens for one sample.
+    """
+    assert logits.size(0) == 1, f"{logits.shape}"
+    assert logits.dtype == torch.float32, f"{logits.dtype}"
+
+    logits = logits.squeeze(0)
+    logits = logits.div(args.rollout_temperature)
+
+    cp_size = mpu.get_context_parallel_world_size()
+    end = 0
+    for tokens, total_length, response_length in zip(unconcat_tokens, total_lengths, response_lengths, strict=False):
+        if cp_size == 1:
+            end += total_length
+            start = end - response_length
+            logits_chunk = logits[start - 1 : end - 1]
+            tokens_chunk = tokens[-response_length:]
+        else:
+            # TODO: this is super ugly... do better abstraction.
+            chunk_size, chunks_offset, logits_offset, tokens_offset = get_logits_and_tokens_offset_with_cp(
+                total_length, response_length
+            )
+
+            logits_0, logits_1 = logits[end : end + chunk_size], logits[end + chunk_size : end + 2 * chunk_size]
+            end += 2 * chunk_size
+
+            logits_0 = logits_0[logits_offset[0][0] - chunks_offset[0][0] : logits_offset[0][1] - chunks_offset[0][0]]
+            tokens_0 = tokens[tokens_offset[0][0] : tokens_offset[0][1]]
+
+            logits_1 = logits_1[logits_offset[1][0] - chunks_offset[1][0] : logits_offset[1][1] - chunks_offset[1][0]]
+            tokens_1 = tokens[tokens_offset[1][0] : tokens_offset[1][1]]
+
+            assert logits_0.size(0) == tokens_0.size(0), f"{logits_0.size(0)} vs {tokens_0.size(0)}"
+            assert logits_1.size(0) == tokens_1.size(0), f"{logits_1.size(0)} vs {tokens_1.size(0)}"
+
+            logits_chunk = torch.cat([logits_0, logits_1], dim=0)
+            tokens_chunk = torch.cat([tokens_0, tokens_1], dim=0)
+
+        yield logits_chunk, tokens_chunk
+
+
+def get_log_probs_and_entropy(
+    logits: torch.Tensor,
+    *,
+    args: Namespace,
+    unconcat_tokens: list[torch.Tensor],
+    total_lengths: list[int],
+    response_lengths: list[int],
+    with_entropy: bool = False,
+    non_loss_data: bool = True,
+) -> dict[str, list[torch.Tensor]]:
+    """Compute per-token log-probabilities (and optionally entropy) on responses.
+
+    For each sample, extracts response-aligned logits and tokens, then computes
+    log-probabilities via softmax across the tensor-parallel group. Log-probs
+    are squeezed from `[R, 1]` to `[R]`. Entropy values are always appended
+    (even when `with_entropy=False`), but only included in the result dict
+    when requested.
+
+    Args:
+        logits: Policy logits with shape `[1, T, V]`.
+        args: Configuration (temperature applied in `get_responses`).
+        unconcat_tokens: List of token tensors per sample.
+        total_lengths: Total sequence lengths per sample.
+        response_lengths: Response segment lengths per sample.
+        with_entropy: If True, include "entropy" key in result.
+        non_loss_data: Unused; kept for API compatibility.
+
+    Returns:
+        Dict with key "log_probs" mapping to a list of `[R]` tensors per
+        sample. If `with_entropy` is True, also includes "entropy" key with
+        a list of `[R]` tensors.
+    """
+    assert non_loss_data
+    log_probs_list = []
+    entropy_list = []
+    for logits_chunk, tokens_chunk in get_responses(
+        logits,
+        args=args,
+        unconcat_tokens=unconcat_tokens,
+        total_lengths=total_lengths,
+        response_lengths=response_lengths,
+    ):
+        log_prob, entropy = calculate_log_probs_and_entropy(
+            logits_chunk,
+            tokens_chunk,
+            mpu.get_tensor_model_parallel_group(),
+            with_entropy=with_entropy,
+            chunk_size=args.log_probs_chunk_size,
+        )
+
+        log_probs_list.append(log_prob.squeeze(-1))
+        entropy_list.append(entropy)
+
+    res = {
+        "log_probs": log_probs_list,
+    }
+    if with_entropy:
+        res["entropy"] = entropy_list
+    return res
+
+
+def get_values(
+    logits: torch.Tensor,
+    *,
+    args: Namespace,
+    unconcat_tokens: list[torch.Tensor],
+    total_lengths: list[int],
+    response_lengths: list[int],
+    with_entropy: bool = False,
+    non_loss_data: bool = True,
+) -> dict[str, list[torch.Tensor]]:
+    """Extract per-token value predictions over response tokens.
+
+    For each sample, extracts response-aligned chunks from the value head
+    output and squeezes the final dimension from `[R, 1]` to `[R]`.
+
+    Args:
+        logits: Value head output with shape `[1, T, 1]`.
+        args: Configuration (passed to `get_responses` which uses
+            `rollout_temperature` even though values don't need temperature).
+        unconcat_tokens: List of token tensors per sample.
+        total_lengths: Total sequence lengths per sample.
+        response_lengths: Response segment lengths per sample.
+        with_entropy: Unused; kept for signature compatibility.
+        non_loss_data: Unused; kept for signature compatibility.
+
+    Returns:
+        Dict with key "values" mapping to a list of `[R]` value tensors
+        per sample.
+    """
+    value_list = []
+    for logits_chunk, _ in get_responses(
+        logits,
+        args=args,
+        unconcat_tokens=unconcat_tokens,
+        total_lengths=total_lengths,
+        response_lengths=response_lengths,
+    ):
+        assert logits_chunk.size(-1) == 1, f"{logits_chunk.shape}"
+        value_list.append(logits_chunk.squeeze(-1))
+
+    return {
+        "values": value_list,
+    }
+
+
+def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) -> None:
+    """Compute advantages and returns in-place based on `args.advantage_estimator`.
+
+    This function extracts rewards, log-probs, values, and masks from
+    `rollout_data`, computes KL divergences, then applies the chosen advantage
+    estimator. Supported methods: "grpo", "gspo", "ppo", "reinforce_plus_plus",
+    and "reinforce_plus_plus_baseline". When `args.normalize_advantages` is
+    True, advantages are whitened across the data-parallel group using masked
+    statistics.
+
+    Early returns if both `log_probs` and `values` are None (intermediate
+    pipeline stages).
+
+    Args:
+        args: Configuration specifying estimator type, KL coefficient,
+            normalization settings, and other hyperparameters.
+        rollout_data: Dict containing input lists ("log_probs", "ref_log_probs",
+            "rewards", "values", "response_lengths", "loss_masks",
+            "total_lengths"). Modified in-place to add "advantages" and
+            "returns" keys, each mapping to lists of tensors per sample.
+    """
+    log_probs: list[torch.Tensor] = rollout_data.get("rollout_log_probs" if args.use_rollout_logprobs else "log_probs")
+    ref_log_probs: list[torch.Tensor] = rollout_data.get("ref_log_probs")
+    rewards: list[float] = rollout_data.get("rewards")
+    values: None | list[torch.Tensor] = rollout_data.get("values")
+    response_lengths: list[int] = rollout_data.get("response_lengths")
+    loss_masks: list[torch.Tensor] = rollout_data.get("loss_masks")
+    total_lengths: list[int] = rollout_data.get("total_lengths")
+
+    # return when not the last pp stage.
+    if log_probs is None and values is None:
+        return
+
+    if args.kl_coef == 0 or not log_probs:
+        # when kl_coef is 0, we won't compute ref_log_prob
+        xs = log_probs if log_probs is not None else values
+        kl = [torch.zeros_like(x, dtype=torch.float32, device=x.device) for x in xs]
+    else:
+        kl = [
+            compute_approx_kl(
+                log_probs[i],
+                ref_log_probs[i],
+                kl_loss_type=args.kl_loss_type,
+            )
+            for i in range(len(log_probs))
+        ]
+
+    if args.advantage_estimator in ["grpo", "gspo"]:
+        rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device)
+        returns = get_grpo_returns(rewards, kl)
+        # TODO: is the copy necessary?
+        advantages = [r for r in returns]
+
+    elif args.advantage_estimator == "ppo":
+        old_rewards = rewards
+        rewards = []
+        kl_coef = -args.kl_coef
+        cp_rank = mpu.get_context_parallel_rank()
+        for reward, k in zip(old_rewards, kl, strict=False):
+            k *= kl_coef
+            if cp_rank == 0:
+                k[-1] += reward
+            rewards.append(k)
+        advantages, returns = get_advantages_and_returns_batch(
+            total_lengths, response_lengths, values, rewards, args.gamma, args.lambd
+        )
+
+    elif args.advantage_estimator == "reinforce_plus_plus":
+        rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device)
+        returns = get_reinforce_plus_plus_returns(
+            rewards=rewards,
+            kl=kl,
+            loss_masks=loss_masks,
+            response_lengths=response_lengths,
+            total_lengths=total_lengths,
+            kl_coef=args.kl_coef,
+            gamma=args.gamma,
+        )
+        advantages = [r for r in returns]
+
+    elif args.advantage_estimator == "reinforce_plus_plus_baseline":
+        rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device)
+        advantages = get_reinforce_plus_plus_baseline_advantages(
+            rewards=rewards,
+            kl=kl,
+            loss_masks=loss_masks,
+            kl_coef=args.kl_coef,
+        )
+        returns = advantages
+
+    elif args.advantage_estimator == "on_policy_distillation":
+        student_log_probs = log_probs
+        teacher_log_probs = rollout_data.get("teacher_log_probs")
+        response_lengths = rollout_data.get("response_lengths")
+
+        device = student_log_probs[0].device
+        teacher_log_probs = [t_log_prob.to(device=device) for t_log_prob in teacher_log_probs]
+        teacher_log_probs = [
+            t_log_prob[-response_length:]
+            for t_log_prob, response_length in zip(teacher_log_probs, response_lengths, strict=False)
+        ]
+
+        advantages = [
+            teacher_log_prob - student_log_prob
+            for teacher_log_prob, student_log_prob in zip(teacher_log_probs, student_log_probs, strict=False)
+        ]
+
+        returns = advantages
+
+    else:
+        raise NotImplementedError(f"advantage_estimator {args.advantage_estimator} is not supported. ")
+
+    # TODO: OpenRLHF always does advantages normalization but veRL doesn't seem to do it.
+    if args.normalize_advantages:
+        all_advs = torch.cat(advantages)
+        cp_size = mpu.get_context_parallel_world_size()
+        if cp_size == 1:
+            all_masks = torch.cat(loss_masks)
+        else:
+            mask_chunks = []
+            for i in range(len(advantages)):
+                total_len = total_lengths[i]
+                response_len = response_lengths[i]
+                prompt_len = total_len - response_len
+
+                _, _, _, token_offsets = get_logits_and_tokens_offset_with_cp(total_len, response_len)
+
+                # Convert global offsets to response-space offsets
+                s0, e0 = token_offsets[0]
+                s1, e1 = token_offsets[1]
+                res_s0, res_e0 = max(0, s0 - prompt_len), max(0, e0 - prompt_len)
+                res_s1, res_e1 = max(0, s1 - prompt_len), max(0, e1 - prompt_len)
+
+                local_mask_parts = []
+                full_mask = loss_masks[i]
+                if res_e0 > res_s0:
+                    local_mask_parts.append(full_mask[res_s0:res_e0])
+                if res_e1 > res_s1:
+                    local_mask_parts.append(full_mask[res_s1:res_e1])
+
+                # Concatenate the parts to form the final mask chunk for this rank and this sequence
+                local_mask_chunk = (
+                    torch.cat(local_mask_parts)
+                    if local_mask_parts
+                    else torch.tensor([], device=all_advs.device, dtype=full_mask.dtype)
+                )
+                mask_chunks.append(local_mask_chunk)
+
+            all_masks = torch.cat(mask_chunks)
+
+        if all_masks.numel() > 0:
+            assert (
+                all_advs.size() == all_masks.size()
+            ), f"Shape mismatch before whitening: advantages {all_advs.size()}, masks {all_masks.size()}"
+            dp_group = mpu.get_data_parallel_group()
+
+            whitened_advs_flat = distributed_masked_whiten(
+                all_advs,
+                all_masks,
+                process_group=dp_group,
+                shift_mean=True,
+            )
+            chunk_lengths = [chunk.size(0) for chunk in advantages]
+            advantages = list(torch.split(whitened_advs_flat, chunk_lengths))
+
+    rollout_data["advantages"] = advantages
+    rollout_data["returns"] = returns
+
+
+def vanilla_tis_function(
+    args,
+    *,
+    pg_loss: torch.Tensor,
+    train_log_probs: list[torch.Tensor],
+    rollout_log_probs: list[torch.Tensor],
+    loss_masks: list[torch.Tensor],
+    **kwargs: Any,
+) -> tuple[torch.Tensor, list[torch.Tensor], dict[str, torch.Tensor]]:
+    rollout_log_probs = torch.cat(rollout_log_probs, dim=0)
+    old_log_probs = torch.cat(train_log_probs, dim=0)
+    tis = torch.exp(old_log_probs - rollout_log_probs)
+    tis_abs = (torch.exp(old_log_probs - rollout_log_probs) - 1).abs()
+    tis_weights = torch.clamp(tis, min=args.tis_clip_low, max=args.tis_clip)
+    tis_clipfrac = (tis_weights != tis).float()
+    metrics = {
+        "tis": tis.clone().detach(),
+        "tis_clipfrac": tis_clipfrac.clone().detach(),
+        "tis_abs": tis_abs.clone().detach(),
+    }
+    pg_loss = pg_loss * tis_weights
+    return pg_loss, loss_masks, metrics
+
+
+def icepop_function(
+    args,
+    *,
+    pg_loss: torch.Tensor,
+    train_log_probs: list[torch.Tensor],
+    rollout_log_probs: list[torch.Tensor],
+    loss_masks: list[torch.Tensor],
+    **kwargs: Any,
+) -> tuple[torch.Tensor, list[torch.Tensor], dict[str, torch.Tensor]]:
+    rollout_log_probs = torch.cat(rollout_log_probs, dim=0)
+    old_log_probs = torch.cat(train_log_probs, dim=0)
+    ice_ratio = torch.exp(old_log_probs - rollout_log_probs)
+    ice_abs = (torch.exp(old_log_probs - rollout_log_probs) - 1).abs()
+    ice_weight = torch.where(
+        (ice_ratio >= args.tis_clip_low) & (ice_ratio <= args.tis_clip), ice_ratio, torch.zeros_like(ice_ratio)
+    )
+    ice_clipfrac = (ice_weight != ice_ratio).float()
+    metrics = {
+        "tis": ice_ratio.clone().detach(),
+        "tis_clipfrac": ice_clipfrac.clone().detach(),
+        "tis_abs": ice_abs.clone().detach(),
+    }
+    pg_loss = pg_loss * ice_weight
+    return pg_loss, loss_masks, metrics
+
+
+def policy_loss_function(
+    args: Namespace,
+    batch: RolloutBatch,
+    logits: torch.Tensor,
+    sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor],
+) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+    """Compute policy loss (PPO/GSPO) and metrics.
+
+    Computes current log-probabilities and entropy from model logits, then
+    calculates PPO-style clipped policy gradient loss. For GSPO, gathers
+    full sequences via context-parallel all-gather before computing per-sample
+    KL. Optionally applies TIS (Truncated Importance Sampling) correction and
+    adds KL loss term if configured.
+
+    Args:
+        args: Configuration controlling advantage estimator, clipping thresholds,
+            entropy/KL coefficients, and TIS settings.
+        batch: Mini-batch containing "advantages", "log_probs" (old policy),
+            "unconcat_tokens", "response_lengths", "total_lengths", "loss_masks",
+            and optionally "ref_log_probs" and "rollout_log_probs".
+        logits: Policy logits with shape `[1, T, V]`.
+        sum_of_sample_mean: Reduction function that averages per-sample values.
+
+    Returns:
+        Tuple of `(loss, metrics)` where `loss` is a scalar tensor and `metrics`
+        is a dict containing detached scalars: "loss", "pg_loss",
+        "entropy_loss", "pg_clipfrac", "ppo_kl". Additional keys "kl_loss",
+        "tis", "ois", "tis_clipfrac" are included when the respective features
+        are enabled.
+    """
+    advantages = torch.cat(batch["advantages"], dim=0)
+    old_log_probs = batch["rollout_log_probs"] if args.use_rollout_logprobs else batch["log_probs"]
+
+    response_lengths = batch["response_lengths"]
+    total_lengths = batch["total_lengths"]
+
+    log_probs_and_entropy = get_log_probs_and_entropy(
+        logits,
+        args=args,
+        unconcat_tokens=batch["unconcat_tokens"],
+        total_lengths=total_lengths,
+        response_lengths=response_lengths,
+        with_entropy=True,
+    )
+
+    log_probs = log_probs_and_entropy["log_probs"]
+
+    # Pre-gather log probs if needed by OPSM or GSPO to avoid duplicate gathering
+    need_full_log_probs = args.use_opsm or args.advantage_estimator == "gspo"
+
+    full_log_probs = None
+    full_old_log_probs = None
+    if need_full_log_probs:
+        full_log_probs = [
+            all_gather_with_cp(log_prob, total_length, response_length)
+            for log_prob, total_length, response_length in zip(
+                log_probs, total_lengths, response_lengths, strict=False
+            )
+        ]
+        full_old_log_probs = [
+            all_gather_with_cp(old_log_prob, total_length, response_length)
+            for old_log_prob, total_length, response_length in zip(
+                old_log_probs, total_lengths, response_lengths, strict=False
+            )
+        ]
+
+    # Compute OPSM mask if enabled
+    if args.use_opsm:
+        opsm_mask, opsm_clipfrac = compute_opsm_mask(
+            args=args,
+            full_log_probs=full_log_probs,
+            full_old_log_probs=full_old_log_probs,
+            advantages=batch["advantages"],
+            loss_masks=batch["loss_masks"],
+        )
+
+    # Compute KL divergence (GSPO uses sequence-level KL, others use per-token KL)
+    if args.advantage_estimator == "gspo":
+        ppo_kl = compute_gspo_kl(
+            full_log_probs=full_log_probs,
+            full_old_log_probs=full_old_log_probs,
+            local_log_probs=log_probs,
+            loss_masks=batch["loss_masks"],
+        )
+        old_log_probs = torch.cat(old_log_probs, dim=0)
+        log_probs = torch.cat(log_probs, dim=0)
+    else:
+        old_log_probs = torch.cat(old_log_probs, dim=0)
+        log_probs = torch.cat(log_probs, dim=0)
+        ppo_kl = old_log_probs - log_probs
+
+    pg_loss, pg_clipfrac = compute_policy_loss(ppo_kl, advantages, args.eps_clip, args.eps_clip_high)
+
+    if args.use_opsm:
+        pg_loss = pg_loss * opsm_mask
+
+    # Apply off-policy correction using importance sampling if enabled
+    if args.get_mismatch_metrics or args.use_tis:
+        # NOTE:
+        # `tis_func` may apply rejection-sampling style masking (RS) and return `modified_response_masks`.
+        # We rebuild `sum_of_sample_mean` with those masks to correct denominators for loss/backprop.
+        #
+        # However, mismatch/TIS/RS metrics (e.g., "truncate_fraction") are often defined over the
+        # *pre-RS* valid tokens. If we aggregate metrics with `modified_response_masks`, the rejected
+        # tokens are excluded from the denominator and the metric can be artificially driven to 0.
+        # Keep a copy of the original reducer (based on `batch["loss_masks"]`) for metric aggregation.
+        sum_of_sample_mean_for_mismatch_metrics = sum_of_sample_mean
+
+        assert "rollout_log_probs" in batch, "rollout_log_probs must be provided for TIS"
+
+        ois = (-ppo_kl).exp()
+        tis_kwargs = {
+            "args": args,
+            "pg_loss": pg_loss,
+            "train_log_probs": batch["log_probs"],
+            "rollout_log_probs": batch["rollout_log_probs"],
+            "loss_masks": batch["loss_masks"],
+            "total_lengths": total_lengths,
+            "response_lengths": response_lengths,
+        }
+
+        if args.custom_tis_function_path is not None:
+            tis_func = load_function(args.custom_tis_function_path)
+        else:
+            tis_func = vanilla_tis_function
+        pg_loss, modified_response_masks, tis_metrics = tis_func(**tis_kwargs)
+
+        # [decouple IS and rejection] Rebuild sum_of_sample_mean with modified_response_masks for denominator correction
+        # modified_response_masks will be sliced with cp in get_sum_of_sample_mean
+        sum_of_sample_mean = get_sum_of_sample_mean(
+            total_lengths, response_lengths, modified_response_masks, args.calculate_per_token_loss
+        )
+
+    pg_loss = sum_of_sample_mean(pg_loss)
+    pg_clipfrac = sum_of_sample_mean(pg_clipfrac)
+    ppo_kl = sum_of_sample_mean(ppo_kl)
+
+    # entropy loss
+    entropy = log_probs_and_entropy["entropy"]
+    entropy = torch.cat(entropy, dim=0)
+    entropy_loss = sum_of_sample_mean(entropy)
+
+    loss = pg_loss - args.entropy_coef * entropy_loss
+
+    if args.use_kl_loss:
+        ref_log_probs = batch["ref_log_probs"]
+        ref_log_probs = torch.cat(ref_log_probs, dim=0)
+        importance_ratio = None
+        if args.use_unbiased_kl:
+            importance_ratio = torch.exp(log_probs - old_log_probs)
+        kl = compute_approx_kl(
+            log_probs,
+            ref_log_probs,
+            kl_loss_type=args.kl_loss_type,
+            importance_ratio=importance_ratio,
+        )
+        kl_loss = sum_of_sample_mean(kl)
+
+        loss = loss + args.kl_loss_coef * kl_loss
+
+    # make sure the gradient could backprop correctly.
+    if log_probs.numel() == 0:
+        loss += 0 * logits.sum()
+
+    train_rollout_logprob_abs_diff = None
+    importance_weight_mean = None
+    importance_weight_std = None
+    if "rollout_log_probs" in batch and batch["rollout_log_probs"]:
+        rollout_log_probs = torch.cat(batch["rollout_log_probs"], dim=0)
+        train_rollout_logprob_abs_diff = sum_of_sample_mean((old_log_probs - rollout_log_probs).abs())
+        iw = torch.exp(log_probs.detach() - rollout_log_probs)
+        importance_weight_mean = sum_of_sample_mean(iw)
+        importance_weight_std = sum_of_sample_mean((iw - 1).pow(2)).sqrt()
+
+    reported_loss = {
+        "loss": loss.clone().detach(),
+        "pg_loss": pg_loss.clone().detach(),
+        "entropy_loss": entropy_loss.clone().detach(),
+        "pg_clipfrac": pg_clipfrac.clone().detach(),
+        "ppo_kl": ppo_kl.clone().detach(),
+    }
+
+    if train_rollout_logprob_abs_diff is not None:
+        reported_loss["train_rollout_logprob_abs_diff"] = train_rollout_logprob_abs_diff.clone().detach()
+    if importance_weight_mean is not None:
+        reported_loss["importance_weight_mean"] = importance_weight_mean.clone().detach()
+        reported_loss["importance_weight_std"] = importance_weight_std.clone().detach()
+
+    if args.use_kl_loss:
+        reported_loss["kl_loss"] = kl_loss.clone().detach()
+
+    if args.get_mismatch_metrics or args.use_tis:
+        # Aggregate mismatch/TIS/RS related metrics with the *pre-RS* masks.
+        # See comment above where `sum_of_sample_mean_for_mismatch_metrics` is defined.
+        reported_loss["ois"] = sum_of_sample_mean_for_mismatch_metrics(ois).clone().detach()
+        # Assume all metrics are already cloned and detached
+        for metric_key, metric_value in tis_metrics.items():
+            key_name = f"{metric_key}"
+            reported_loss[key_name] = sum_of_sample_mean_for_mismatch_metrics(metric_value)
+
+    if args.use_opsm:
+        reported_loss["opsm_clipfrac"] = opsm_clipfrac
+
+    return loss, reported_loss
+
+
+def value_loss_function(
+    args: Namespace,
+    batch: RolloutBatch,
+    logits: torch.Tensor,
+    sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor],
+) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+    """Compute clipped value loss and metrics.
+
+    Extracts current value predictions from `logits`, compares them against
+    stored old values with clipping, and computes the maximum of clipped and
+    unclipped squared errors (PPO-style value clipping).
+
+    Args:
+        args: Configuration containing `value_clip` threshold.
+        batch: Mini-batch with "values" (old predictions), "returns",
+            "unconcat_tokens", "total_lengths", and "response_lengths".
+        logits: Value head output with shape `[1, T, 1]`.
+        sum_of_sample_mean: Reduction function that averages per-sample values.
+
+    Returns:
+        Tuple of `(loss, metrics)` where `loss` is a scalar tensor and
+        `metrics` contains detached scalars "value_loss" and "value_clipfrac".
+    """
+    old_values = torch.cat(batch["values"], dim=0)
+
+    values = get_values(
+        logits,
+        args=args,
+        unconcat_tokens=batch["unconcat_tokens"],
+        total_lengths=batch["total_lengths"],
+        response_lengths=batch["response_lengths"],
+    )
+    values = torch.cat([value.flatten() for value in values["values"]], dim=0)
+
+    returns = torch.cat(batch["returns"], dim=0)
+
+    values_clipfrac = torch.abs(values - old_values) > args.value_clip
+    values_clipped = old_values + (values - old_values).clamp(-args.value_clip, args.value_clip)
+    surr1 = (values_clipped - returns) ** 2
+    surr2 = (values - returns) ** 2
+    loss = torch.max(surr1, surr2)
+
+    loss = sum_of_sample_mean(loss)
+    values_clipfrac = sum_of_sample_mean(values_clipfrac.float())
+
+    # make sure the gradient could backprop correctly.
+    if values.numel() == 0:
+        loss += 0 * values.sum()
+
+    reported_loss = {
+        "value_loss": loss.clone().detach(),
+        "value_clipfrac": values_clipfrac.clone().detach(),
+    }
+
+    return loss, reported_loss
+
+
+def loss_function(
+    args: Namespace,
+    batch: RolloutBatch,
+    num_microbatches: int,
+    logits: torch.Tensor,
+) -> tuple[torch.Tensor, int | torch.Tensor, dict[str, list[str] | torch.Tensor]]:
+    """Dispatch to the configured loss and rescale for Megatron integration.
+
+    Selects one of "policy_loss", "value_loss", or a custom loss
+    function based on `args.loss_type`, computes the loss and metrics, then
+    rescales the loss by micro-batch and parallelism factors to integrate with
+    Megatron's gradient accumulation.
+
+    Args:
+        args: Configuration specifying `loss_type`, `calculate_per_token_loss`,
+            `global_batch_size`, and optionally `custom_loss_function_path`.
+        batch: Mini-batch with "loss_masks", "response_lengths", and other
+            keys required by the selected loss function.
+        num_microbatches: Number of gradient accumulation steps.
+        logits: Model outputs (policy or value head).
+
+    Returns:
+        Tuple of `(scaled_loss, normalizer, logging_dict)` where:
+        - `scaled_loss` is the loss tensor (scalar) rescaled for Megatron.
+        - `normalizer` is `num_tokens` (scalar tensor) if
+          `args.calculate_per_token_loss` is True, else `1` (int).
+        - `logging_dict` has keys "keys" (list of str metric names) and
+          "values" (1D tensor: [count, metric1, metric2, ...]).
+    """
+    num_tokens = sum([torch.clamp_min(loss_mask.sum(), 1) for loss_mask in batch["loss_masks"]])
+    num_samples = len(batch["response_lengths"])
+
+    sum_of_sample_mean = get_sum_of_sample_mean(
+        batch["total_lengths"],
+        batch["response_lengths"],
+        batch["loss_masks"],
+        args.calculate_per_token_loss,
+    )
+
+    loss_type = args.loss_type
+
+    match loss_type:
+        case "policy_loss":
+            func = policy_loss_function
+        case "value_loss":
+            func = value_loss_function
+        case "custom_loss":
+            func = load_function(args.custom_loss_function_path)
+        case _:
+            raise ValueError(f"Unknown loss type: {loss_type}")
+
+    if args.recompute_loss_function:
+        loss, log = checkpoint(func, args, batch, logits, sum_of_sample_mean)
+    else:
+        loss, log = func(args, batch, logits, sum_of_sample_mean)
+
+    # Here we need to divide by cp_size because to cancel the multiply in Megatron.
+    if not args.calculate_per_token_loss:
+        loss = (
+            loss
+            * num_microbatches
+            / args.global_batch_size
+            * mpu.get_data_parallel_world_size(with_context_parallel=True)
+        )
+    else:
+        loss = loss * mpu.get_context_parallel_world_size()
+
+    return (
+        loss,
+        torch.tensor(num_tokens if args.calculate_per_token_loss else 1, device=logits.device),
+        {
+            "keys": list(log.keys()),
+            "values": torch.tensor(
+                [
+                    num_samples if not args.calculate_per_token_loss else num_tokens,
+                ]
+                + list(log.values()),
+                device=logits.device,
+            ),
+        },
+    )

slime/backends/megatron_utils/megatron_to_hf/__init__.py

@@ -0,0 +1,88 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from .deepseekv3 import convert_deepseekv3_to_hf
+from .glm4 import convert_glm4_to_hf
+from .glm4moe import convert_glm4moe_to_hf
+from .llama import convert_llama_to_hf
+from .mimo import convert_mimo_to_hf
+from .processors.padding_remover import remove_padding
+from .processors.quantizer import quantize_params
+from .qwen2 import convert_qwen2_to_hf
+from .qwen3_next import convert_qwen3_next_to_hf
+from .qwen3moe import convert_qwen3moe_to_hf
+
+
+# TODO unify w/ `convert_to_hf`
+def postprocess_hf_param(args, megatron_param_name, hf_param_name, param):
+    param = remove_padding(megatron_param_name, param, args.vocab_size)
+    # TODO support quant
+    return param
+
+
+# TODO optimize code details
+def convert_to_hf(args, model_name, name, param, quantization_config=None):
+    param = remove_padding(name, param, args.vocab_size)
+
+    converted_named_tensors = _convert_to_hf_core(args, model_name, name, param)
+
+    if not quantization_config:
+        return converted_named_tensors
+
+    return quantize_params(args, name, converted_named_tensors, quantization_config)
+
+
+# TODO optimize
+_cached_tensors = {}
+
+
+# TODO optimize code details
+def _convert_to_hf_core(args, model_name, name, param):
+    if "glm4moe" in model_name:
+        converted_named_tensors = convert_glm4moe_to_hf(args, name, param)
+    elif "glm4" in model_name:
+        converted_named_tensors = convert_glm4_to_hf(args, name, param)
+    elif "qwen3moe" in model_name:
+        converted_named_tensors = convert_qwen3moe_to_hf(args, name, param)
+    elif "qwen3next" in model_name:
+        converted_named_tensors = convert_qwen3_next_to_hf(args, name, param)
+    elif "qwen2" in model_name or "qwen3" in model_name:
+        converted_named_tensors = convert_qwen2_to_hf(args, name, param)
+    elif "deepseekv3" in model_name:
+        converted_named_tensors = convert_deepseekv3_to_hf(args, name, param)
+
+    elif "llama" in model_name:
+        converted_named_tensors = convert_llama_to_hf(args, name, param)
+    elif "mimo" in model_name:
+        converted_named_tensors = convert_mimo_to_hf(args, name, param)
+    else:
+        raise ValueError(f"Unsupported model: {model_name}")
+
+    # to compatible with sglang implementation
+    if args.q_lora_rank is not None:
+        old_converted_named_tensors = converted_named_tensors
+        converted_named_tensors = []
+        for converted_name, converted_param in old_converted_named_tensors:
+            if "q_a_proj" in converted_name:
+                pair_name = converted_name.replace("q_a_proj", "kv_a_proj_with_mqa")
+                if pair_name in _cached_tensors:
+                    converted_named_tensors += [
+                        (converted_name, converted_param),
+                        (pair_name, _cached_tensors[pair_name]),
+                    ]
+                    del _cached_tensors[pair_name]
+                else:
+                    _cached_tensors[converted_name] = converted_param
+            elif "kv_a_proj_with_mqa" in converted_name:
+                pair_name = converted_name.replace("kv_a_proj_with_mqa", "q_a_proj")
+                if pair_name in _cached_tensors:
+                    converted_named_tensors += [
+                        (converted_name, converted_param),
+                        (pair_name, _cached_tensors[pair_name]),
+                    ]
+                    del _cached_tensors[pair_name]
+                else:
+                    _cached_tensors[converted_name] = converted_param
+            else:
+                converted_named_tensors.append((converted_name, converted_param))
+    return converted_named_tensors