初始化项目，由ModelHub XC社区提供模型

Model: ayh015/myLightningOPD
Source: Original Platform

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/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

slime/backends/megatron_utils/megatron_to_hf/deepseekv3.py

@@ -0,0 +1,132 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+
+import torch
+
+
+def convert_deepseekv3_to_hf(args, name, param):
+    if name == "module.module.embedding.word_embeddings.weight":
+        return [("model.embed_tokens.weight", param)]
+    if name == "module.module.output_layer.weight":
+        return [("lm_head.weight", param)]
+    if name == "module.module.decoder.final_layernorm.weight":
+        return [("model.norm.weight", param)]
+
+    try:
+        head_dim = args.kv_channels if args.kv_channels is not None else args.hidden_size // args.num_attention_heads
+    except AttributeError:
+        head_dim = args.hidden_size // args.num_attention_heads
+    value_num_per_group = args.num_attention_heads // args.num_query_groups
+
+    decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)"
+    match = re.match(decoder_layers_pattern, name)
+    if match:
+        layer_idx, rest = match.groups()
+
+        # experts
+        expert_pattern = r"mlp.experts\.(.+)\.weight(\d+)"
+        match = re.match(expert_pattern, rest)
+        if match:
+            rest, expert_idx = match.groups()
+            if rest == "linear_fc1":
+                gate_weight, up_weight = param.chunk(2, dim=0)
+                outputs = [
+                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.weight", gate_weight),
+                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.weight", up_weight),
+                ]
+                return outputs
+            elif rest == "linear_fc2":
+                outputs = [
+                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight", param),
+                ]
+                return outputs
+            else:
+                raise ValueError(f"Unknown expert parameter name: {name}")
+
+        # shared expert
+        shared_expert_pattern = r"mlp.shared_experts\.(.+)"
+        match = re.match(shared_expert_pattern, rest)
+        if match:
+            rest = match.groups()[0]
+            if rest == "linear_fc1.weight":
+                gate_weight, up_weight = param.chunk(2, dim=0)
+                return [
+                    (f"model.layers.{layer_idx}.mlp.shared_experts.gate_proj.weight", gate_weight),
+                    (f"model.layers.{layer_idx}.mlp.shared_experts.up_proj.weight", up_weight),
+                ]
+            elif rest == "linear_fc2.weight":
+                return [(f"model.layers.{layer_idx}.mlp.shared_experts.down_proj.weight", param)]
+            else:
+                raise ValueError(f"Unknown shared expert parameter name: {name}")
+
+        if rest == "self_attention.linear_proj.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.o_proj.weight", param)]
+        elif rest == "self_attention.linear_q_proj.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.q_proj.weight", param)]
+        elif rest == "self_attention.linear_q_down_proj.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.q_a_proj.weight", param)]
+        elif rest == "self_attention.linear_q_up_proj.layer_norm_weight":
+            return [(f"model.layers.{layer_idx}.self_attn.q_a_layernorm.weight", param)]
+        elif rest == "self_attention.linear_q_up_proj.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.q_b_proj.weight", param)]
+        elif rest == "self_attention.linear_qkv.bias":
+            param = param.view(args.num_query_groups, -1)
+            q_bias, k_bias, v_bias = torch.split(
+                param,
+                split_size_or_sections=[value_num_per_group * head_dim, head_dim, head_dim],
+                dim=1,
+            )
+            q_bias = q_bias.contiguous().flatten()
+            k_bias = k_bias.contiguous().flatten()
+            v_bias = v_bias.contiguous().flatten()
+            return [
+                (f"model.layers.{layer_idx}.self_attn.q_proj.bias", q_bias),
+                (f"model.layers.{layer_idx}.self_attn.k_proj.bias", k_bias),
+                (f"model.layers.{layer_idx}.self_attn.v_proj.bias", v_bias),
+            ]
+        elif rest == "mlp.linear_fc1.weight":
+            gate_weight, up_weight = param.chunk(2, dim=0)
+            return [
+                (f"model.layers.{layer_idx}.mlp.gate_proj.weight", gate_weight),
+                (f"model.layers.{layer_idx}.mlp.up_proj.weight", up_weight),
+            ]
+        elif rest == "mlp.linear_fc2.weight":
+            return [(f"model.layers.{layer_idx}.mlp.down_proj.weight", param)]
+        elif rest == "self_attention.linear_qkv.layer_norm_weight" or rest == "input_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)]
+        elif rest == "mlp.linear_fc1.layer_norm_weight":
+            return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)]
+        elif rest == "self_attention.linear_kv_down_proj.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.kv_a_proj_with_mqa.weight", param)]
+        elif rest == "self_attention.linear_kv_up_proj.layer_norm_weight":
+            return [(f"model.layers.{layer_idx}.self_attn.kv_a_layernorm.weight", param)]
+        elif rest == "self_attention.linear_kv_up_proj.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.kv_b_proj.weight", param)]
+        elif rest == "pre_mlp_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)]
+        elif rest == "mlp.router.weight":
+            return [(f"model.layers.{layer_idx}.mlp.gate.weight", param)]
+        elif rest == "mlp.router.expert_bias":
+            return [(f"model.layers.{layer_idx}.mlp.gate.e_score_correction_bias", param)]
+
+    mtp_layer_pattern = r"module\.module\.mtp\.layers\.(\d+)\.(.+)"
+    match = re.match(mtp_layer_pattern, name)
+    if match:
+        layer_idx, rest = match.groups()
+        layer_idx = int(layer_idx) + args.num_layers
+        if rest == "eh_proj.weight":
+            return [(f"model.layers.{layer_idx}.eh_proj.weight", param)]
+        elif rest == "enorm.weight":
+            return [(f"model.layers.{layer_idx}.enorm.weight", param)]
+        elif rest == "hnorm.weight":
+            return [(f"model.layers.{layer_idx}.hnorm.weight", param)]
+        elif rest == "final_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.shared_head.norm.weight", param)]
+        else:
+            name = f"module.module.decoder.layers.{layer_idx}.{rest}"
+            name = name.replace("transformer_layer.", "")
+            return convert_deepseekv3_to_hf(args, name, param)
+
+    raise ValueError(f"Unknown parameter name: {name}")

slime/backends/megatron_utils/megatron_to_hf/glm4.py

@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+import torch
+
+
+def convert_glm4_to_hf(args, name, param):
+    if name == "module.module.embedding.word_embeddings.weight":
+        return [("model.embed_tokens.weight", param)]
+    if name == "module.module.output_layer.weight":
+        return [("lm_head.weight", param)]
+    if name == "module.module.decoder.final_layernorm.weight":
+        return [("model.norm.weight", param)]
+
+    try:
+        head_dim = args.kv_channels if args.kv_channels is not None else args.hidden_size // args.num_attention_heads
+    except AttributeError:
+        head_dim = args.hidden_size // args.num_attention_heads
+    value_num_per_group = args.num_attention_heads // args.num_query_groups
+
+    decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)"
+    match = re.match(decoder_layers_pattern, name)
+    if match:
+        layer_idx, rest = match.groups()
+        if rest == "self_attention.linear_proj.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.o_proj.weight", param)]
+        elif rest == "self_attention.linear_qkv.weight":
+
+            param = param.view(args.num_query_groups, -1, head_dim, args.hidden_size)
+            q_param, k_param, v_param = torch.split(param, split_size_or_sections=[value_num_per_group, 1, 1], dim=1)
+            q_param = q_param.reshape(-1, args.hidden_size)
+            k_param = k_param.reshape(-1, args.hidden_size)
+            v_param = v_param.reshape(-1, args.hidden_size)
+            return [
+                (f"model.layers.{layer_idx}.self_attn.q_proj.weight", q_param),
+                (f"model.layers.{layer_idx}.self_attn.k_proj.weight", k_param),
+                (f"model.layers.{layer_idx}.self_attn.v_proj.weight", v_param),
+            ]
+        elif rest == "self_attention.linear_qkv.bias":
+            param = param.view(args.num_query_groups, -1)
+            q_bias, k_bias, v_bias = torch.split(
+                param,
+                split_size_or_sections=[value_num_per_group * head_dim, head_dim, head_dim],
+                dim=1,
+            )
+            q_bias = q_bias.contiguous().flatten()
+            k_bias = k_bias.contiguous().flatten()
+            v_bias = v_bias.contiguous().flatten()
+            return [
+                (f"model.layers.{layer_idx}.self_attn.q_proj.bias", q_bias),
+                (f"model.layers.{layer_idx}.self_attn.k_proj.bias", k_bias),
+                (f"model.layers.{layer_idx}.self_attn.v_proj.bias", v_bias),
+            ]
+        elif rest == "mlp.linear_fc1.weight":
+            return [
+                (f"model.layers.{layer_idx}.mlp.gate_up_proj.weight", param),
+            ]
+        elif rest == "mlp.linear_fc2.weight":
+            return [(f"model.layers.{layer_idx}.mlp.down_proj.weight", param)]
+        elif rest == "self_attention.linear_qkv.layer_norm_weight":
+            return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)]
+        elif rest == "mlp.linear_fc1.layer_norm_weight":
+            return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)]
+
+        # qk norm
+        elif rest == "self_attention.q_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.q_norm.weight", param)]
+        elif rest == "self_attention.k_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.k_norm.weight", param)]
+
+        # sandwitch norm
+        elif rest == "post_self_attn_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.post_self_attn_layernorm.weight", param)]
+        elif rest == "post_mlp_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.post_mlp_layernorm.weight", param)]
+
+    raise ValueError(f"Unknown parameter name: {name}")

slime/backends/megatron_utils/megatron_to_hf/glm4moe.py

@@ -0,0 +1,142 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+
+import torch
+
+
+def convert_glm4moe_to_hf(args, name, param):
+    if name == "module.module.embedding.word_embeddings.weight":
+        return [("model.embed_tokens.weight", param)]
+    if name == "module.module.output_layer.weight":
+        return [("lm_head.weight", param)]
+    if name == "module.module.decoder.final_layernorm.weight":
+        return [("model.norm.weight", param)]
+
+    try:
+        head_dim = args.kv_channels if args.kv_channels is not None else args.hidden_size // args.num_attention_heads
+    except AttributeError:
+        head_dim = args.hidden_size // args.num_attention_heads
+    value_num_per_group = args.num_attention_heads // args.num_query_groups
+
+    decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)"
+    match = re.match(decoder_layers_pattern, name)
+    if match:
+        layer_idx, rest = match.groups()
+
+        # experts
+        expert_pattern = r"mlp.experts\.(.+)\.weight(\d+)"
+        match = re.match(expert_pattern, rest)
+        if match:
+            rest, expert_idx = match.groups()
+            if rest == "linear_fc1":
+                gate_weight, up_weight = param.chunk(2, dim=0)
+                outputs = [
+                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.weight", gate_weight),
+                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.weight", up_weight),
+                ]
+                return outputs
+            elif rest == "linear_fc2":
+                outputs = [
+                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight", param),
+                ]
+                return outputs
+            else:
+                raise ValueError(f"Unknown expert parameter name: {name}")
+
+        # shared expert
+        shared_expert_pattern = r"mlp.shared_experts\.(.+)"
+        match = re.match(shared_expert_pattern, rest)
+        if match:
+            rest = match.groups()[0]
+            if rest == "linear_fc1.weight":
+                gate_weight, up_weight = param.chunk(2, dim=0)
+                return [
+                    (f"model.layers.{layer_idx}.mlp.shared_experts.gate_proj.weight", gate_weight),
+                    (f"model.layers.{layer_idx}.mlp.shared_experts.up_proj.weight", up_weight),
+                ]
+            elif rest == "linear_fc2.weight":
+                return [
+                    (f"model.layers.{layer_idx}.mlp.shared_experts.down_proj.weight", param),
+                ]
+            else:
+                raise ValueError(f"Unknown shared expert parameter name: {name}")
+
+        if rest == "self_attention.linear_proj.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.o_proj.weight", param)]
+        elif rest == "self_attention.linear_qkv.weight":
+
+            param = param.view(args.num_query_groups, -1, head_dim, args.hidden_size)
+            q_param, k_param, v_param = torch.split(param, split_size_or_sections=[value_num_per_group, 1, 1], dim=1)
+            q_param = q_param.reshape(-1, args.hidden_size)
+            k_param = k_param.reshape(-1, args.hidden_size)
+            v_param = v_param.reshape(-1, args.hidden_size)
+            return [
+                (f"model.layers.{layer_idx}.self_attn.q_proj.weight", q_param),
+                (f"model.layers.{layer_idx}.self_attn.k_proj.weight", k_param),
+                (f"model.layers.{layer_idx}.self_attn.v_proj.weight", v_param),
+            ]
+        elif rest == "self_attention.linear_qkv.bias":
+            param = param.view(args.num_query_groups, -1)
+            q_bias, k_bias, v_bias = torch.split(
+                param,
+                split_size_or_sections=[value_num_per_group * head_dim, head_dim, head_dim],
+                dim=1,
+            )
+            q_bias = q_bias.contiguous().flatten()
+            k_bias = k_bias.contiguous().flatten()
+            v_bias = v_bias.contiguous().flatten()
+            return [
+                (f"model.layers.{layer_idx}.self_attn.q_proj.bias", q_bias),
+                (f"model.layers.{layer_idx}.self_attn.k_proj.bias", k_bias),
+                (f"model.layers.{layer_idx}.self_attn.v_proj.bias", v_bias),
+            ]
+        elif rest == "mlp.linear_fc1.weight":
+            gate_weight, up_weight = param.chunk(2, dim=0)
+            return [
+                (f"model.layers.{layer_idx}.mlp.gate_proj.weight", gate_weight),
+                (f"model.layers.{layer_idx}.mlp.up_proj.weight", up_weight),
+            ]
+        elif rest == "mlp.linear_fc2.weight":
+            return [(f"model.layers.{layer_idx}.mlp.down_proj.weight", param)]
+        elif rest == "self_attention.linear_qkv.layer_norm_weight":
+            return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)]
+        elif rest == "mlp.linear_fc1.layer_norm_weight":
+            return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)]
+        elif rest == "post_self_attn_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.post_self_attn_layernorm.weight", param)]
+        elif rest == "post_mlp_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.post_mlp_layernorm.weight", param)]
+        elif rest == "pre_mlp_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)]
+        elif rest == "mlp.router.weight":
+            return [(f"model.layers.{layer_idx}.mlp.gate.weight", param)]
+        elif rest == "mlp.router.expert_bias":
+            return [(f"model.layers.{layer_idx}.mlp.gate.e_score_correction_bias", param)]
+
+        # qk norm
+        elif rest == "self_attention.q_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.q_norm.weight", param)]
+        elif rest == "self_attention.k_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.k_norm.weight", param)]
+
+    mtp_layer_pattern = r"module\.module\.mtp\.layers\.(\d+)\.(.+)"
+    match = re.match(mtp_layer_pattern, name)
+    if match:
+        layer_idx, rest = match.groups()
+        layer_idx = int(layer_idx) + args.num_layers
+        if rest == "eh_proj.weight":
+            return [(f"model.layers.{layer_idx}.eh_proj.weight", param)]
+        elif rest == "enorm.weight":
+            return [(f"model.layers.{layer_idx}.enorm.weight", param)]
+        elif rest == "hnorm.weight":
+            return [(f"model.layers.{layer_idx}.hnorm.weight", param)]
+        elif rest == "final_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.shared_head.norm.weight", param)]
+        else:
+            name = f"module.module.decoder.layers.{layer_idx}.{rest}"
+            name = name.replace("transformer_layer.", "")
+            return convert_glm4moe_to_hf(args, name, param)
+
+    raise ValueError(f"Unknown parameter name: {name}")

slime/backends/megatron_utils/megatron_to_hf/llama.py

@@ -0,0 +1,56 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+import torch
+
+
+def convert_llama_to_hf(args, name, param):
+    if name == "module.module.embedding.word_embeddings.weight":
+        return [("model.embed_tokens.weight", param)]
+    if name == "module.module.output_layer.weight":
+        return [("lm_head.weight", param)]
+    if name == "module.module.decoder.final_layernorm.weight":
+        return [("model.norm.weight", param)]
+
+    try:
+        head_dim = args.kv_channels if args.kv_channels is not None else args.hidden_size // args.num_attention_heads
+    except AttributeError:
+        head_dim = args.hidden_size // args.num_attention_heads
+    value_num_per_group = args.num_attention_heads // args.num_query_groups
+
+    decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)"
+    match = re.match(decoder_layers_pattern, name)
+    if match:
+        layer_idx, rest = match.groups()
+        if rest == "self_attention.linear_proj.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.o_proj.weight", param)]
+        elif rest == "self_attention.linear_qkv.weight":
+            # Split QKV weight for Llama
+            param = param.view(args.num_query_groups, -1, head_dim, args.hidden_size)
+            q_param, k_param, v_param = torch.split(param, split_size_or_sections=[value_num_per_group, 1, 1], dim=1)
+            q_param = q_param.reshape(-1, args.hidden_size)
+            k_param = k_param.reshape(-1, args.hidden_size)
+            v_param = v_param.reshape(-1, args.hidden_size)
+            return [
+                (f"model.layers.{layer_idx}.self_attn.q_proj.weight", q_param),
+                (f"model.layers.{layer_idx}.self_attn.k_proj.weight", k_param),
+                (f"model.layers.{layer_idx}.self_attn.v_proj.weight", v_param),
+            ]
+        elif rest == "mlp.linear_fc1.weight":
+            # Split gate and up projections for SwiGLU
+            gate_weight, up_weight = param.chunk(2, dim=0)
+            return [
+                (f"model.layers.{layer_idx}.mlp.gate_proj.weight", gate_weight),
+                (f"model.layers.{layer_idx}.mlp.up_proj.weight", up_weight),
+            ]
+        elif rest == "mlp.linear_fc2.weight":
+            return [(f"model.layers.{layer_idx}.mlp.down_proj.weight", param)]
+        elif rest == "self_attention.linear_qkv.layer_norm_weight":
+            return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)]
+        elif rest == "mlp.linear_fc1.layer_norm_weight":
+            return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)]
+        elif rest == "pre_mlp_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)]
+
+    raise ValueError(f"Unknown parameter name: {name}")

slime/backends/megatron_utils/megatron_to_hf/mimo.py

@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+import torch
+from .qwen2 import convert_qwen2_to_hf
+
+
+def convert_mimo_to_hf(args, name, param):
+    """
+    Convert MiMo model parameters from Megatron to HuggingFace format.
+
+    MiMo extends Qwen2 with MTP (Multi-Token Prediction) layers.
+    """
+
+    if "mtp" in name:
+        return convert_mimo_mtp_param(args, name, param)
+
+    return convert_qwen2_to_hf(args, name, param)
+
+
+def convert_mimo_mtp_param(args, name, param):
+    """
+    Convert MTP layer parameters from Megatron to HuggingFace format.
+
+    MTP layers in MiMo contain:
+    - LayerNorms (token_layernorm, hidden_layernorm, final_layernorm)
+    - Input projection (input_proj)
+    - Self attention (reuses Qwen2 attention structure)
+    - MLP (reuses Qwen2 MLP structure)
+
+    Based on MimoBridge._convert_mtp_param logic (reverse mapping)
+    """
+    mtp_pattern = r"module\.module\.mtp\.layers\.(\d+)\.(.+)"
+    match = re.match(mtp_pattern, name)
+
+    if not match:
+        raise ValueError(f"Invalid MTP parameter name: {name}")
+
+    layer_idx, component = match.groups()
+
+    # Direct mappings for MTP-specific components (Megatron -> HF)
+    # Based on MimoBridge direct_name_mapping (reversed)
+    direct_mappings = {
+        "enorm.weight": f"model.mtp_layers.{layer_idx}.token_layernorm.weight",
+        "hnorm.weight": f"model.mtp_layers.{layer_idx}.hidden_layernorm.weight",
+        "eh_proj.weight": f"model.mtp_layers.{layer_idx}.input_proj.weight",
+        "final_layernorm.weight": f"model.mtp_layers.{layer_idx}.final_layernorm.weight",
+    }
+    if component == "eh_proj.weight":
+        first_half, second_half = param.chunk(2, dim=1)
+        param = torch.cat([second_half, first_half], dim=1)
+
+    # Check direct mappings first
+    if component in direct_mappings:
+        return [(direct_mappings[component], param)]
+
+    # Handle transformer_layer components
+    if component.startswith("transformer_layer."):
+        # Remove "transformer_layer." prefix
+        transformer_component = component[len("transformer_layer.") :]
+
+        # Create proxy name for reusing existing Qwen2 conversion functions
+        proxy_name = f"module.module.decoder.layers.{layer_idx}.{transformer_component}"
+
+        # Use existing convert_qwen2_to_hf function for transformer components
+        results = convert_qwen2_to_hf(args, proxy_name, param)
+
+        # Replace model.layers with mtp_layers in results
+        converted_results = []
+        for hf_name, hf_param in results:
+            # Replace model.layers.{idx} with mtp_layers.{idx}
+            hf_name = hf_name.replace(f"model.layers.{layer_idx}", f"model.mtp_layers.{layer_idx}")
+            converted_results.append((hf_name, hf_param))
+
+        return converted_results
+
+    raise ValueError(f"Unknown MTP component: {component} in {name}")

slime/backends/megatron_utils/megatron_to_hf/processors/__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/megatron_utils/megatron_to_hf/processors/padding_remover.py

@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from slime.backends.megatron_utils.misc_utils import strip_param_name_prefix
+
+
+def remove_padding(name: str, param: torch.Tensor, vocab_size: int) -> torch.Tensor:
+    """
+    Remove vocab padding: param[:vocab_size] for embedding/output layers, else unchanged.
+    """
+    if strip_param_name_prefix(name) in {"embedding.word_embeddings.weight", "output_layer.weight"}:
+        return param[:vocab_size]
+    return param

slime/backends/megatron_utils/megatron_to_hf/processors/quantizer.py

@@ -0,0 +1,110 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+
+import torch
+
+from slime.utils.fp8_kernel import blockwise_cast_to_fp8_triton
+
+from ...sglang import quant_weight_ue8m0, should_deepgemm_weight_requant_ue8m0, transform_scale_ue8m0
+
+
+def quantize_params(args, megatron_name, converted_named_params, quantization_config):
+    if quantization_config is None:
+        return converted_named_params
+    assert quantization_config["quant_method"] == "fp8"
+    assert quantization_config["fmt"] == "e4m3"
+    assert quantization_config["activation_scheme"] == "dynamic"
+    weight_block_size = quantization_config.get("weight_block_size", None)
+
+    decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)"
+    match = re.match(decoder_layers_pattern, megatron_name)
+
+    if not match:
+        # check mtp layers
+        mtp_layer_pattern = r"module\.module\.mtp\.layers\.(\d+)\.(.+)"
+        match = re.match(mtp_layer_pattern, megatron_name)
+        if not match:
+            return converted_named_params
+        layer_idx, rest = match.groups()
+        rest = rest.replace("transformer_layer.", "")
+    else:
+        layer_idx, rest = match.groups()
+
+    # experts
+    expert_pattern = r"mlp.experts\.(.+)\.weight(\d+)"
+    match = re.match(expert_pattern, rest)
+    if match:
+        rest, expert_idx = match.groups()
+        if rest in [
+            "linear_fc1",
+            "linear_fc2",
+        ]:
+            quantize_named_params = []
+            for converted_name, param in converted_named_params:
+                # skip bf16 weight_scale and input_scale
+                # TODO: find a clearer way.
+                if converted_name.endswith("_scale"):
+                    continue
+                quantize_named_params.extend(_quantize_param(converted_name, param, weight_block_size))
+
+            return quantize_named_params
+
+    # shared expert
+    shared_expert_pattern = r"mlp.shared_experts\.(.+)"
+    match = re.match(shared_expert_pattern, rest)
+    if match:
+        rest = match.groups()[0]
+        if rest in [
+            "linear_fc1.weight",
+            "linear_fc2.weight",
+        ]:
+            quantize_named_params = []
+            for converted_name, param in converted_named_params:
+                quantize_named_params.extend(_quantize_param(converted_name, param, weight_block_size))
+
+            return quantize_named_params
+
+    if rest in [
+        "self_attention.linear_proj.weight",
+        "self_attention.linear_qkv.weight",
+        "mlp.linear_fc1.weight",
+        "mlp.linear_fc2.weight",
+        # mla
+        "self_attention.linear_q_proj.weight",
+        "self_attention.linear_q_down_proj.weight",
+        "self_attention.linear_q_up_proj.weight",
+        "self_attention.linear_kv_down_proj.weight",
+        "self_attention.linear_kv_up_proj.weight",
+    ]:
+        quantize_named_params = []
+        for converted_name, param in converted_named_params:
+            quantize_named_params.extend(_quantize_param(converted_name, param, weight_block_size))
+
+        return quantize_named_params
+
+    # for other parameters, we just return the original converted_named_params
+    return converted_named_params
+
+
+def _quantize_param(name, weight, weight_block_size):
+    assert name.endswith(".weight"), f"Expected weight parameter, got {name}"
+    FP8_MIN = torch.finfo(torch.float8_e4m3fn).min
+    FP8_MAX = torch.finfo(torch.float8_e4m3fn).max
+    if weight_block_size is not None:
+        if should_deepgemm_weight_requant_ue8m0 and should_deepgemm_weight_requant_ue8m0(
+            weight_block_size=weight_block_size
+        ):
+            qweight, scale = quant_weight_ue8m0(weight, weight_block_size=weight_block_size)
+            scale = transform_scale_ue8m0(scale, mn=qweight.shape[-2])
+        else:
+            qweight, scale = blockwise_cast_to_fp8_triton(weight, weight_block_size)
+        scale_name = name.replace(".weight", ".weight_scale_inv")
+    else:
+        # per tensor quant
+        scale = weight.abs().max().clamp(min=1e-12).to(torch.float32) / FP8_MAX
+        qweight = (weight / scale).clamp(min=FP8_MIN, max=FP8_MAX).to(torch.float8_e4m3fn)
+        scale = scale.view(1)
+        scale_name = name.replace(".weight", ".weight_scale")
+    return [(name, qweight), (scale_name, scale)]

slime/backends/megatron_utils/megatron_to_hf/qwen2.py

@@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+import torch
+
+
+def convert_qwen2_to_hf(args, name, param):
+    if name == "module.module.embedding.word_embeddings.weight":
+        return [("model.embed_tokens.weight", param)]
+    if name == "module.module.output_layer.weight":
+        return [("lm_head.weight", param)]
+    if name == "module.module.decoder.final_layernorm.weight":
+        return [("model.norm.weight", param)]
+
+    try:
+        head_dim = args.kv_channels if args.kv_channels is not None else args.hidden_size // args.num_attention_heads
+    except AttributeError:
+        head_dim = args.hidden_size // args.num_attention_heads
+    value_num_per_group = args.num_attention_heads // args.num_query_groups
+
+    decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)"
+    match = re.match(decoder_layers_pattern, name)
+    if match:
+        layer_idx, rest = match.groups()
+        if rest == "self_attention.linear_proj.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.o_proj.weight", param)]
+        elif rest == "self_attention.linear_qkv.weight":
+
+            param = param.view(args.num_query_groups, -1, head_dim, args.hidden_size)
+            q_param, k_param, v_param = torch.split(param, split_size_or_sections=[value_num_per_group, 1, 1], dim=1)
+            q_param = q_param.reshape(-1, args.hidden_size)
+            k_param = k_param.reshape(-1, args.hidden_size)
+            v_param = v_param.reshape(-1, args.hidden_size)
+            return [
+                (f"model.layers.{layer_idx}.self_attn.q_proj.weight", q_param),
+                (f"model.layers.{layer_idx}.self_attn.k_proj.weight", k_param),
+                (f"model.layers.{layer_idx}.self_attn.v_proj.weight", v_param),
+            ]
+        elif rest == "self_attention.linear_qkv.bias":
+            param = param.view(args.num_query_groups, -1)
+            q_bias, k_bias, v_bias = torch.split(
+                param,
+                split_size_or_sections=[value_num_per_group * head_dim, head_dim, head_dim],
+                dim=1,
+            )
+            q_bias = q_bias.contiguous().flatten()
+            k_bias = k_bias.contiguous().flatten()
+            v_bias = v_bias.contiguous().flatten()
+            return [
+                (f"model.layers.{layer_idx}.self_attn.q_proj.bias", q_bias),
+                (f"model.layers.{layer_idx}.self_attn.k_proj.bias", k_bias),
+                (f"model.layers.{layer_idx}.self_attn.v_proj.bias", v_bias),
+            ]
+        elif rest == "mlp.linear_fc1.weight":
+            gate_weight, up_weight = param.chunk(2, dim=0)
+            return [
+                (f"model.layers.{layer_idx}.mlp.gate_proj.weight", gate_weight),
+                (f"model.layers.{layer_idx}.mlp.up_proj.weight", up_weight),
+            ]
+        elif rest == "mlp.linear_fc2.weight":
+            return [(f"model.layers.{layer_idx}.mlp.down_proj.weight", param)]
+        elif rest == "self_attention.linear_qkv.layer_norm_weight":
+            return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)]
+        elif rest == "mlp.linear_fc1.layer_norm_weight":
+            return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)]
+
+        # qk norm
+        elif rest == "self_attention.q_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.q_norm.weight", param)]
+        elif rest == "self_attention.k_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.k_norm.weight", param)]
+
+    raise ValueError(f"Unknown parameter name: {name}")

slime/backends/megatron_utils/megatron_to_hf/qwen3_next.py

@@ -0,0 +1,145 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+
+import torch
+
+
+def convert_qwen3_next_to_hf(args, name, param):
+    if name == "module.module.embedding.word_embeddings.weight":
+        return [("model.embed_tokens.weight", param)]
+    if name == "module.module.output_layer.weight":
+        return [("lm_head.weight", param)]
+    if name == "module.module.decoder.final_layernorm.weight":
+        return [("model.norm.weight", param)]
+
+    try:
+        head_dim = args.kv_channels if args.kv_channels is not None else args.hidden_size // args.num_attention_heads
+    except AttributeError:
+        head_dim = args.hidden_size // args.num_attention_heads
+    value_num_per_group = args.num_attention_heads // args.num_query_groups
+
+    decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)"
+    match = re.match(decoder_layers_pattern, name)
+    if match:
+        layer_idx, rest = match.groups()
+
+        # experts
+        expert_pattern = r"mlp.experts\.(.+)\.weight(\d+)"
+        match = re.match(expert_pattern, rest)
+        if match:
+            rest, expert_idx = match.groups()
+            if rest == "linear_fc1":
+                gate_weight, up_weight = param.chunk(2, dim=0)
+                outputs = [
+                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.weight", gate_weight),
+                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.weight", up_weight),
+                ]
+                return outputs
+            elif rest == "linear_fc2":
+                outputs = [
+                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight", param),
+                ]
+                return outputs
+            else:
+                raise ValueError(f"Unknown expert parameter name: {name}")
+
+        # shared expert
+        shared_expert_pattern = r"mlp.shared_experts\.(.+)"
+        match = re.match(shared_expert_pattern, rest)
+        if match:
+            rest = match.groups()[0]
+            if rest == "linear_fc1.weight":
+                gate_weight, up_weight = param.chunk(2, dim=0)
+                return [
+                    (f"model.layers.{layer_idx}.mlp.shared_expert.gate_proj.weight", gate_weight),
+                    (f"model.layers.{layer_idx}.mlp.shared_expert.up_proj.weight", up_weight),
+                ]
+            elif rest == "linear_fc2.weight":
+                return [(f"model.layers.{layer_idx}.mlp.shared_expert.down_proj.weight", param)]
+            elif rest == "gate_weight":
+                return [(f"model.layers.{layer_idx}.mlp.shared_expert_gate.weight", param)]
+            else:
+                raise ValueError(f"Unknown shared expert parameter name: {name}")
+
+        if rest == "self_attention.linear_proj.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.o_proj.weight", param)]
+        elif rest == "self_attention.linear_qkv.weight":
+            param = param.view(args.num_query_groups, -1, head_dim, args.hidden_size)
+            q_param, k_param, v_param = torch.split(
+                param, split_size_or_sections=[2 * value_num_per_group, 1, 1], dim=1
+            )
+            q_param = (
+                q_param.reshape(args.num_query_groups, 2, value_num_per_group, head_dim, args.hidden_size)
+                .transpose(1, 2)
+                .reshape(-1, args.hidden_size)
+            )
+            k_param = k_param.reshape(-1, args.hidden_size)
+            v_param = v_param.reshape(-1, args.hidden_size)
+            return [
+                (f"model.layers.{layer_idx}.self_attn.q_proj.weight", q_param),
+                (f"model.layers.{layer_idx}.self_attn.k_proj.weight", k_param),
+                (f"model.layers.{layer_idx}.self_attn.v_proj.weight", v_param),
+            ]
+        elif rest == "self_attention.linear_qkv.bias":
+            param = param.view(args.num_query_groups, -1)
+            q_bias, k_bias, v_bias = torch.split(
+                param,
+                split_size_or_sections=[value_num_per_group * head_dim, head_dim, head_dim],
+                dim=1,
+            )
+            q_bias = q_bias.contiguous().flatten()
+            k_bias = k_bias.contiguous().flatten()
+            v_bias = v_bias.contiguous().flatten()
+            return [
+                (f"model.layers.{layer_idx}.self_attn.q_proj.bias", q_bias),
+                (f"model.layers.{layer_idx}.self_attn.k_proj.bias", k_bias),
+                (f"model.layers.{layer_idx}.self_attn.v_proj.bias", v_bias),
+            ]
+        elif rest == "mlp.linear_fc1.weight":
+            gate_weight, up_weight = param.chunk(2, dim=0)
+            return [
+                (f"model.layers.{layer_idx}.mlp.gate_proj.weight", gate_weight),
+                (f"model.layers.{layer_idx}.mlp.up_proj.weight", up_weight),
+            ]
+        elif rest == "mlp.linear_fc2.weight":
+            return [(f"model.layers.{layer_idx}.mlp.down_proj.weight", param)]
+        elif rest == "self_attention.linear_qkv.layer_norm_weight":
+            return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)]
+        elif rest == "mlp.linear_fc1.layer_norm_weight":
+            return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)]
+        elif rest == "pre_mlp_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)]
+        elif rest == "mlp.router.weight":
+            return [(f"model.layers.{layer_idx}.mlp.gate.weight", param)]
+        elif rest == "mlp.router.expert_bias":
+            return [(f"model.layers.{layer_idx}.mlp.gate.e_score_correction_bias", param)]
+
+        # qk norm
+        elif rest == "self_attention.q_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.q_norm.weight", param)]
+        elif rest == "self_attention.k_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.k_norm.weight", param)]
+        elif rest.startswith("self_attention.") and rest[len("self_attention.") :] in [
+            "input_layernorm.weight",
+            # linear attn
+            "linear_attn.A_log",
+            "linear_attn.conv1d.weight",
+            "linear_attn.dt_bias",
+            "linear_attn.in_proj_ba.weight",
+            "linear_attn.in_proj_qkvz.weight",
+            "linear_attn.norm.weight",
+            "linear_attn.out_proj.weight",
+            # gated attn
+            "self_attn.k_norm.weight",
+            "self_attn.k_proj.weight",
+            "self_attn.o_proj.weight",
+            "self_attn.q_norm.weight",
+            "self_attn.q_proj.weight",
+            "self_attn.v_proj.weight",
+        ]:
+            rest = rest[len("self_attention.") :]
+            return [(f"model.layers.{layer_idx}.{rest}", param)]
+
+    raise ValueError(f"Unknown parameter name: {name}")

slime/backends/megatron_utils/megatron_to_hf/qwen3moe.py

@@ -0,0 +1,120 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+
+import torch
+
+
+def convert_qwen3moe_to_hf(args, name, param):
+    if name == "module.module.embedding.word_embeddings.weight":
+        return [("model.embed_tokens.weight", param)]
+    if name == "module.module.output_layer.weight":
+        return [("lm_head.weight", param)]
+    if name == "module.module.decoder.final_layernorm.weight":
+        return [("model.norm.weight", param)]
+
+    try:
+        head_dim = args.kv_channels if args.kv_channels is not None else args.hidden_size // args.num_attention_heads
+    except AttributeError:
+        head_dim = args.hidden_size // args.num_attention_heads
+    value_num_per_group = args.num_attention_heads // args.num_query_groups
+
+    decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)"
+    match = re.match(decoder_layers_pattern, name)
+    if match:
+        layer_idx, rest = match.groups()
+
+        # experts
+        expert_pattern = r"mlp.experts\.(.+)\.weight(\d+)"
+        match = re.match(expert_pattern, rest)
+        if match:
+            rest, expert_idx = match.groups()
+            if rest == "linear_fc1":
+                gate_weight, up_weight = param.chunk(2, dim=0)
+                outputs = [
+                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.weight", gate_weight),
+                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.weight", up_weight),
+                ]
+                return outputs
+            elif rest == "linear_fc2":
+                outputs = [
+                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight", param),
+                ]
+                return outputs
+            else:
+                raise ValueError(f"Unknown expert parameter name: {name}")
+
+        # shared expert
+        shared_expert_pattern = r"mlp.shared_experts\.(.+)"
+        match = re.match(shared_expert_pattern, rest)
+        if match:
+            rest = match.groups()[0]
+            if rest == "linear_fc1.weight":
+                gate_weight, up_weight = param.chunk(2, dim=0)
+                return [
+                    (f"model.layers.{layer_idx}.mlp.shared_expert.gate_proj.weight", gate_weight),
+                    (f"model.layers.{layer_idx}.mlp.shared_expert.up_proj.weight", up_weight),
+                ]
+            elif rest == "linear_fc2.weight":
+                return [(f"model.layers.{layer_idx}.mlp.shared_expert.down_proj.weight", param)]
+            elif rest == "gate_weight":
+                return [(f"model.layers.{layer_idx}.mlp.shared_expert_gate.weight", param)]
+            else:
+                raise ValueError(f"Unknown shared expert parameter name: {name}")
+
+        if rest == "self_attention.linear_proj.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.o_proj.weight", param)]
+        elif rest == "self_attention.linear_qkv.weight":
+
+            param = param.view(args.num_query_groups, -1, head_dim, args.hidden_size)
+            q_param, k_param, v_param = torch.split(param, split_size_or_sections=[value_num_per_group, 1, 1], dim=1)
+            q_param = q_param.reshape(-1, args.hidden_size)
+            k_param = k_param.reshape(-1, args.hidden_size)
+            v_param = v_param.reshape(-1, args.hidden_size)
+            return [
+                (f"model.layers.{layer_idx}.self_attn.q_proj.weight", q_param),
+                (f"model.layers.{layer_idx}.self_attn.k_proj.weight", k_param),
+                (f"model.layers.{layer_idx}.self_attn.v_proj.weight", v_param),
+            ]
+        elif rest == "self_attention.linear_qkv.bias":
+            param = param.view(args.num_query_groups, -1)
+            q_bias, k_bias, v_bias = torch.split(
+                param,
+                split_size_or_sections=[value_num_per_group * head_dim, head_dim, head_dim],
+                dim=1,
+            )
+            q_bias = q_bias.contiguous().flatten()
+            k_bias = k_bias.contiguous().flatten()
+            v_bias = v_bias.contiguous().flatten()
+            return [
+                (f"model.layers.{layer_idx}.self_attn.q_proj.bias", q_bias),
+                (f"model.layers.{layer_idx}.self_attn.k_proj.bias", k_bias),
+                (f"model.layers.{layer_idx}.self_attn.v_proj.bias", v_bias),
+            ]
+        elif rest == "mlp.linear_fc1.weight":
+            gate_weight, up_weight = param.chunk(2, dim=0)
+            return [
+                (f"model.layers.{layer_idx}.mlp.gate_proj.weight", gate_weight),
+                (f"model.layers.{layer_idx}.mlp.up_proj.weight", up_weight),
+            ]
+        elif rest == "mlp.linear_fc2.weight":
+            return [(f"model.layers.{layer_idx}.mlp.down_proj.weight", param)]
+        elif rest == "self_attention.linear_qkv.layer_norm_weight":
+            return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)]
+        elif rest == "mlp.linear_fc1.layer_norm_weight":
+            return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)]
+        elif rest == "pre_mlp_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)]
+        elif rest == "mlp.router.weight":
+            return [(f"model.layers.{layer_idx}.mlp.gate.weight", param)]
+        elif rest == "mlp.router.expert_bias":
+            return [(f"model.layers.{layer_idx}.mlp.gate.e_score_correction_bias", param)]
+
+        # qk norm
+        elif rest == "self_attention.q_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.q_norm.weight", param)]
+        elif rest == "self_attention.k_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.k_norm.weight", param)]
+
+    raise ValueError(f"Unknown parameter name: {name}")

slime/backends/megatron_utils/misc_utils.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+def strip_param_name_prefix(name: str):
+    prefix = "module."
+    while name.startswith(prefix):
+        name = name.removeprefix(prefix)
+    return name

slime/backends/megatron_utils/model.py

@@ -0,0 +1,730 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import dataclasses
+import gc
+import logging
+import math
+import os
+from argparse import Namespace
+from collections.abc import Callable, Sequence
+from functools import partial
+
+import torch
+from megatron.core import mpu
+from megatron.core.distributed import DistributedDataParallel as DDP
+from megatron.core.distributed import finalize_model_grads
+from megatron.core.enums import ModelType
+from megatron.core.models.gpt import GPTModel
+from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer
+from megatron.core.optimizer.optimizer import MegatronOptimizer
+from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
+from megatron.core.pipeline_parallel import get_forward_backward_func
+from megatron.core.utils import get_model_config
+from megatron.training.global_vars import get_args
+from megatron.training.training import get_model
+
+from slime.utils import tracking_utils
+from slime.utils.memory_utils import clear_memory
+
+from .checkpoint import load_checkpoint, save_checkpoint
+from .data import DataIterator, get_batch
+from .loss import loss_function
+from .model_provider import get_model_provider_func
+
+logger = logging.getLogger(__name__)
+
+
+def get_optimizer_param_scheduler(args: Namespace, optimizer: MegatronOptimizer) -> OptimizerParamScheduler:
+    """Create and configure the optimizer learning-rate/weight-decay scheduler.
+
+    This configures iteration-based schedules derived from the global batch size
+    and run-time arguments.
+
+    Args:
+        args (Namespace): Training/runtime arguments (argparse namespace).
+        optimizer (MegatronOptimizer): Megatron optimizer bound to the model.
+
+    Returns:
+        OptimizerParamScheduler: Initialized scheduler bound to ``optimizer``.
+    """
+    # Iteration-based training.
+    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 * args.global_batch_size
+    wd_incr_steps = args.train_iters * args.global_batch_size
+    wsd_decay_steps = None
+    if args.lr_wsd_decay_iters is not None:
+        wsd_decay_steps = args.lr_wsd_decay_iters * args.global_batch_size
+    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 * args.global_batch_size
+
+    opt_param_scheduler = OptimizerParamScheduler(
+        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,
+        start_wd=args.start_weight_decay,
+        end_wd=args.end_weight_decay,
+        wd_incr_steps=wd_incr_steps,
+        wd_incr_style=args.weight_decay_incr_style,
+        use_checkpoint_opt_param_scheduler=args.use_checkpoint_opt_param_scheduler,
+        override_opt_param_scheduler=args.override_opt_param_scheduler,
+        wsd_decay_steps=wsd_decay_steps,
+        lr_wsd_decay_style=args.lr_wsd_decay_style,
+    )
+
+    return opt_param_scheduler
+
+
+def setup_model_and_optimizer(
+    args: Namespace,
+    role: str = "actor",
+) -> tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler]:
+    """Build model(s), wrap with DDP, and construct optimizer and scheduler.
+
+    Args:
+        args (Namespace): Training/runtime arguments (argparse namespace).
+        role (str): Logical role of the model (e.g., "actor", "critic").
+        no_wd_decay_cond (Callable[..., bool] | None): Predicate to exclude
+            parameters from weight decay.
+        scale_lr_cond (Callable[..., bool] | None): Predicate to scale LR for
+            selected parameter groups.
+        lr_mult (float): Global learning-rate multiplier for the optimizer.
+
+    Returns:
+        tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler]:
+            - List of model chunks wrapped by ``DDP``.
+            - The constructed ``MegatronOptimizer`` instance.
+            - The learning-rate/weight-decay scheduler tied to the optimizer.
+    """
+    assert not args.moe_use_upcycling
+    assert args.load is not None or args.pretrained_checkpoint is not None
+
+    model = get_model(get_model_provider_func(args, role), ModelType.encoder_or_decoder)
+
+    # Optimizer
+    kwargs = {}
+    for f in dataclasses.fields(OptimizerConfig):
+        if hasattr(args, f.name):
+            kwargs[f.name] = getattr(args, f.name)
+    config = OptimizerConfig(**kwargs)
+    config.timers = None
+
+    optimizer = get_megatron_optimizer(
+        config=config,
+        model_chunks=model,
+        use_gloo_process_groups=args.enable_gloo_process_groups,
+    )
+    opt_param_scheduler = get_optimizer_param_scheduler(args, optimizer)
+    return model, optimizer, opt_param_scheduler
+
+
+def enable_forward_pre_hook(model_chunks: Sequence[DDP]) -> None:
+    """Enable forward pre-hooks for provided DDP-wrapped model chunks.
+
+    Args:
+        model_chunks (Sequence[DDP]): Sequence of DDP modules to enable hooks on.
+    """
+    for model_chunk in model_chunks:
+        assert isinstance(model_chunk, DDP)
+        model_chunk.enable_forward_pre_hook()
+
+
+def disable_forward_pre_hook(model_chunks: Sequence[DDP], param_sync: bool = True) -> None:
+    """Disable forward pre-hooks for provided DDP-wrapped model chunks.
+
+    Args:
+        model_chunks (Sequence[DDP]): Sequence of DDP modules to disable hooks on.
+        param_sync (bool): Whether to synchronize parameters when disabling.
+    """
+    for model_chunk in model_chunks:
+        assert isinstance(model_chunk, DDP)
+        model_chunk.disable_forward_pre_hook(param_sync=param_sync)
+
+
+@torch.no_grad()
+def forward_only(
+    f: Callable[..., dict[str, list[torch.Tensor]]],
+    args: Namespace,
+    model: Sequence[DDP],
+    data_iterator: Sequence[DataIterator],
+    num_microbatches: Sequence[int],
+    store_prefix: str = "",
+) -> dict[str, list[torch.Tensor]]:
+    """Run forward passes only and collect non-loss outputs (e.g., logprobs).
+
+    The model is put into evaluation mode, a forward-only pipeline pass is
+    executed, and relevant outputs are aggregated and returned.
+
+    Args:
+        f (Callable[..., dict[str, list[torch.Tensor]]]): Post-forward callback used to
+            compute and package outputs to collect. This should accept a logits
+            tensor as its first positional argument and additional keyword-only
+            arguments; see ``get_log_probs_and_entropy``/``get_values`` in
+            ``megatron_utils.loss`` for examples. It will be partially applied
+            so that the callable returned from the internal forward step only
+            requires the logits tensor.
+        args (Namespace): Runtime arguments.
+        model (Sequence[DDP]): Sequence of DDP-wrapped model chunks.
+        data_iterator (Sequence[DataIterator]): Iterable(s) yielding batches for inference.
+        num_microbatches (Sequence[int]): Number of microbatches per rollout step.
+        store_prefix (str): Prefix to prepend to stored output keys.
+
+    Returns:
+        dict[str, list[torch.Tensor]]: Aggregated outputs keyed by ``store_prefix + key``.
+    """
+
+    # reset data iterator
+    for iterator in data_iterator:
+        iterator.reset()
+
+    config = get_model_config(model[0])
+
+    def forward_step(
+        data_iterator: DataIterator, model: GPTModel, return_schedule_plan: bool = False
+    ) -> tuple[torch.Tensor, Callable[[torch.Tensor], dict[str, list[torch.Tensor]]]]:
+        """Forward step used by Megatron's pipeline engine.
+
+        Args:
+            data_iterator (DataIterator): Input data iterator.
+            model (GPTModel): The GPT model chunk to execute.
+
+        Returns:
+            tuple[torch.Tensor, Callable[[torch.Tensor], dict[str, list[torch.Tensor]]]]:
+            Output tensor(s) and a callable that computes and packages results
+            to be collected by the engine.
+        """
+
+        assert not return_schedule_plan, "forward_only step should never return schedule plan"
+
+        # Get the batch.
+        batch = get_batch(
+            data_iterator,
+            [
+                "tokens",
+                "loss_masks",
+                "multimodal_train_inputs",
+                "total_lengths",
+                "response_lengths",
+            ],
+            args.data_pad_size_multiplier,
+        )
+        unconcat_tokens = batch["unconcat_tokens"]
+        tokens = batch["tokens"]
+        packed_seq_params = batch["packed_seq_params"]
+        total_lengths = batch["total_lengths"]
+        response_lengths = batch["response_lengths"]
+        output_tensor = model(
+            input_ids=tokens,
+            position_ids=None,
+            attention_mask=None,
+            labels=None,
+            packed_seq_params=packed_seq_params,
+            loss_mask=batch["full_loss_masks"],
+            **(batch["multimodal_train_inputs"] if batch["multimodal_train_inputs"] is not None else {}),
+        )
+
+        return output_tensor, partial(
+            f,
+            args=args,
+            unconcat_tokens=unconcat_tokens,
+            total_lengths=total_lengths,
+            response_lengths=response_lengths,
+            with_entropy=args.use_rollout_entropy,
+        )
+
+    # Turn on evaluation mode which disables dropout.
+    for model_module in model:
+        model_module.eval()
+
+    if args.custom_megatron_before_log_prob_hook_path:
+        from slime.utils.misc import load_function
+
+        custom_before_log_prob_hook = load_function(args.custom_megatron_before_log_prob_hook_path)
+        custom_before_log_prob_hook(args, model, store_prefix)
+
+    forward_backward_func = get_forward_backward_func()
+    # Don't care about timing during evaluation
+    config.timers = None
+    forward_data_store = []
+    num_steps_per_rollout = len(num_microbatches)
+    for step_id in range(num_steps_per_rollout):
+        # collect_non_loss_data
+        forward_data_store += forward_backward_func(
+            forward_step_func=forward_step,
+            data_iterator=data_iterator,
+            model=model,
+            num_microbatches=num_microbatches[step_id],
+            seq_length=args.seq_length,
+            micro_batch_size=args.micro_batch_size,
+            forward_only=True,
+            collect_non_loss_data=True,
+        )
+
+    # Move model back to the train mode.
+    for model_module in model:
+        model_module.train()
+
+    rollout_data = {}
+    # Store the results on the last stage
+    if mpu.is_pipeline_last_stage():
+        keys = forward_data_store[0].keys()
+        for key in keys:
+            values = []
+            for value in forward_data_store:
+                assert isinstance(value[key], list)
+                values += value[key]
+
+            if args.use_dynamic_batch_size:
+                # TODO: This is ugly... Find a better way to make the data have the same order.
+                # TODO: move this out of the loop.
+                origin_values = [None] * len(values)
+                origin_indices = sum(data_iterator[0].micro_batch_indices, [])
+                for value, origin_index in zip(values, origin_indices, strict=False):
+                    origin_values[origin_index] = value
+                values = origin_values
+            rollout_data[f"{store_prefix}{key}"] = values
+    return rollout_data
+
+
+def train_one_step(
+    args: Namespace,
+    rollout_id: int,
+    step_id: int,
+    data_iterator: Sequence[DataIterator],
+    model: Sequence[DDP],
+    optimizer: MegatronOptimizer,
+    opt_param_scheduler: OptimizerParamScheduler,
+    num_microbatches: int,
+) -> tuple[dict[str, float], float]:
+    """Execute a single pipeline-parallel training step.
+
+    Runs forward/backward over ``num_microbatches``, applies optimizer step and
+    one scheduler step when gradients are valid.
+
+    Args:
+        args (Namespace): Runtime arguments.
+        rollout_id (int): Rollout identifier.
+        step_id (int): Step index within the current rollout.
+        data_iterator (Sequence[DataIterator]): Iterable(s) yielding training batches.
+        model (Sequence[DDP]): Sequence of DDP-wrapped model chunks.
+        optimizer (MegatronOptimizer): Optimizer instance.
+        opt_param_scheduler (OptimizerParamScheduler): LR/WD scheduler.
+        num_microbatches (int): Number of microbatches to process.
+
+    Returns:
+        tuple[dict[str, float], float]: Reduced loss dictionary (last stage only)
+        and gradient norm for logging.
+    """
+    args = get_args()
+
+    # Set grad to zero.
+    for model_chunk in model:
+        model_chunk.zero_grad_buffer()
+    optimizer.zero_grad()
+
+    if args.custom_megatron_before_train_step_hook_path:
+        from slime.utils.misc import load_function
+
+        custom_before_train_step_hook = load_function(args.custom_megatron_before_train_step_hook_path)
+        custom_before_train_step_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler)
+
+    def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_plan: bool = False) -> tuple[
+        torch.Tensor,
+        Callable[[torch.Tensor], tuple[torch.Tensor, int, dict[str, torch.Tensor | list[str]]]],
+    ]:
+        """Forward step used by Megatron's pipeline engine during training.
+
+        Args:
+            data_iterator (DataIterator): Input data iterator.
+            model (GPTModel): The GPT model chunk to execute.
+
+        Returns:
+            tuple[torch.Tensor, Callable[[torch.Tensor], tuple[torch.Tensor, int, dict[str, torch.Tensor | list[str]]]]]:
+            Output tensor(s) and the loss function, which returns
+            (loss, num_elems, {"keys": list[str], "values": torch.Tensor}).
+        """
+
+        # Get the batch.
+        batch = get_batch(
+            data_iterator,
+            [
+                "tokens",
+                "multimodal_train_inputs",
+                "packed_seq_params",
+                "total_lengths",
+                "response_lengths",
+                "loss_masks",
+                "log_probs",
+                "ref_log_probs",
+                "values",
+                "advantages",
+                "returns",
+                "rollout_log_probs",
+                "teacher_log_probs",    # For OPD distillation loss
+            ],
+            args.data_pad_size_multiplier,
+        )
+
+        if os.environ.get("ENABLE_ROUTING_REPLAY", "0") == "1":
+            old_stage = os.environ["ROUTING_REPLAY_STAGE"]
+            os.environ["ROUTING_REPLAY_STAGE"] = "replay_forward"
+
+        if return_schedule_plan:
+            assert not args.enable_mtp_training, "MTP training should not be enabled when using combined 1f1b"
+            output_tensor = model.build_schedule_plan(
+                input_ids=batch["tokens"],
+                position_ids=None,
+                attention_mask=None,
+                labels=None,
+                packed_seq_params=batch["packed_seq_params"],
+                loss_mask=batch["full_loss_masks"],
+            )
+        else:
+            output_tensor = model(
+                input_ids=batch["tokens"],
+                position_ids=None,
+                attention_mask=None,
+                labels=None,
+                packed_seq_params=batch["packed_seq_params"],
+                loss_mask=batch["full_loss_masks"],
+                mtp_kwargs={"mtp_labels": batch["tokens"]} if args.enable_mtp_training else {},
+                **(batch["multimodal_train_inputs"] if batch["multimodal_train_inputs"] is not None else {}),
+            )
+
+        if os.environ.get("ENABLE_ROUTING_REPLAY", "0") == "1":
+            os.environ["ROUTING_REPLAY_STAGE"] = old_stage
+
+        return output_tensor, partial(loss_function, args, batch, num_microbatches)
+
+    # Forward pass.
+    forward_backward_func = get_forward_backward_func()
+    losses_reduced = forward_backward_func(
+        forward_step_func=forward_step,
+        data_iterator=data_iterator,
+        model=model,
+        num_microbatches=num_microbatches,
+        seq_length=args.seq_length,
+        micro_batch_size=args.micro_batch_size,
+        decoder_seq_length=args.decoder_seq_length,
+        forward_only=False,
+    )
+
+    valid_step = True
+    if not getattr(args, "check_for_nan_in_loss_and_grad", True):
+        found_inf_flag = optimizer.prepare_grads()
+        if found_inf_flag:
+            valid_step = False
+        else:
+            grad_norm = optimizer.get_grad_norm()
+            if isinstance(grad_norm, torch.Tensor):
+                valid_step = not (torch.isnan(grad_norm) or torch.isinf(grad_norm))
+            else:
+                valid_step = not (math.isnan(grad_norm) or math.isinf(grad_norm))
+
+    # CI check: verify only MTP parameters have non-zero gradients when truncation happens
+    # This check must happen before optimizer.step() as gradients may be modified during step
+    if args.ci_test and args.enable_mtp_training:
+        from slime.backends.megatron_utils.ci_utils import check_mtp_only_grad
+
+        check_mtp_only_grad(model, step_id)
+
+    if valid_step:
+        # Update parameters.
+        update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
+
+        # Update learning rate.
+        assert update_successful
+        opt_param_scheduler.step(increment=args.global_batch_size)
+
+    # release grad
+    for model_chunk in model:
+        model_chunk.zero_grad_buffer()
+    optimizer.zero_grad()
+
+    if mpu.is_pipeline_last_stage(ignore_virtual=True):
+        # Average loss across microbatches.
+        keys = losses_reduced[0]["keys"]
+        values = None
+        for x in losses_reduced:
+            if values is None:
+                values = x["values"]
+            else:
+                values += x["values"]
+        assert len(keys) + 1 == values.numel()
+        torch.distributed.all_reduce(values, group=mpu.get_data_parallel_group(with_context_parallel=True))
+
+        loss_reduced = {}
+        values = values.tolist()
+        num_samples_or_tokens = values[0]
+        for key, value in zip(keys, values[1:], strict=False):
+            loss_reduced[key] = value * mpu.get_context_parallel_world_size() / num_samples_or_tokens
+        return loss_reduced, grad_norm
+    return {}, grad_norm
+
+
+def should_disable_forward_pre_hook(args: Namespace) -> bool:
+    """Block forward pre-hook for certain configurations."""
+    return args.use_distributed_optimizer and args.overlap_param_gather
+
+
+def finalize_model_grads_with_empty_cache(*args, **kwargs):
+    # trigger empty cache when there are less than 10% free memory before the final reduce scatter.
+    # TODO: this is an ad-hoc method and we should figure out why the oom happens in the first place.
+    device = torch.cuda.current_device()
+    free, total = torch.cuda.mem_get_info(device)
+    if free / total < 0.1:
+        clear_memory()
+    return finalize_model_grads(*args, **kwargs)
+
+
+def train(
+    rollout_id: int,
+    model: Sequence[DDP],
+    optimizer: MegatronOptimizer,
+    opt_param_scheduler: OptimizerParamScheduler,
+    data_iterator: Sequence[DataIterator],
+    num_microbatches: Sequence[int],
+) -> None:
+    """Run training over a rollout consisting of multiple steps.
+
+    The model is switched to train mode, training hooks are configured, and
+    ``train_one_step`` is invoked for each step in the rollout.
+
+    Args:
+        rollout_id (int): Rollout identifier.
+        model (Sequence[DDP]): Sequence of DDP-wrapped model chunks.
+        optimizer (MegatronOptimizer): Optimizer instance.
+        opt_param_scheduler (OptimizerParamScheduler): LR/WD scheduler.
+        data_iterator (Sequence[DataIterator]): Iterable(s) yielding training batches.
+        num_microbatches (Sequence[int]): Microbatches per step in the rollout.
+    """
+    args = get_args()
+
+    for iterator in data_iterator:
+        iterator.reset()
+
+    # Turn on training mode which enables dropout.
+    for model_module in model:
+        model_module.train()
+
+    # Setup some training config params.
+    config = get_model_config(model[0])
+    config.grad_scale_func = optimizer.scale_loss
+    config.timers = None
+    if isinstance(model[0], DDP) and args.overlap_grad_reduce:
+        assert config.no_sync_func is None, (
+            "When overlap_grad_reduce is True, config.no_sync_func must be None; "
+            "a custom no_sync_func is not supported when overlapping grad-reduce"
+        )
+        config.no_sync_func = [model_chunk.no_sync for model_chunk in model]
+        if len(model) == 1:
+            config.no_sync_func = config.no_sync_func[0]
+        if args.align_grad_reduce:
+            config.grad_sync_func = [model_chunk.start_grad_sync for model_chunk in model]
+            if len(model) == 1:
+                config.grad_sync_func = config.grad_sync_func[0]
+    if args.overlap_param_gather and args.align_param_gather:
+        config.param_sync_func = [model_chunk.start_param_sync for model_chunk in model]
+        if len(model) == 1:
+            config.param_sync_func = config.param_sync_func[0]
+    config.finalize_model_grads_func = finalize_model_grads_with_empty_cache
+
+    pre_hook_enabled = False
+
+    if args.manual_gc:
+        # Disable the default garbage collector and perform the collection manually.
+        # This is to align the timing of garbage collection across ranks.
+        assert args.manual_gc_interval >= 0, "Manual garbage collection interval should be larger than or equal to 0"
+        gc.disable()
+        gc.collect()
+
+    # Disable forward pre-hook to start training to ensure that errors in checkpoint loading
+    # or random initialization don't propagate to all ranks in first all-gather (which is a
+    # no-op if things work correctly).
+    if should_disable_forward_pre_hook(args):
+        disable_forward_pre_hook(model, param_sync=False)
+        # Also remove param_sync_func temporarily so that sync calls made in
+        # `forward_backward_func` are no-ops.
+        param_sync_func = config.param_sync_func
+        config.param_sync_func = None
+        pre_hook_enabled = False
+
+    num_steps_per_rollout = len(num_microbatches)
+
+    # Run training iterations till done.
+    for step_id in range(num_steps_per_rollout):
+
+        # Run training step.
+        loss_dict, grad_norm = train_one_step(
+            args,
+            rollout_id,
+            step_id,
+            data_iterator,
+            model,
+            optimizer,
+            opt_param_scheduler,
+            num_microbatches[step_id],
+        )
+
+        if step_id == 0:
+            # Enable forward pre-hook after training step has successfully run. All subsequent
+            # forward passes will use the forward pre-hook / `param_sync_func` in
+            # `forward_backward_func`.
+            if should_disable_forward_pre_hook(args):
+                enable_forward_pre_hook(model)
+                config.param_sync_func = param_sync_func
+                pre_hook_enabled = True
+
+        if args.enable_mtp_training:
+            from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper
+
+            mtp_loss_scale = 1 / num_microbatches[step_id]
+            tracker = MTPLossLoggingHelper.tracker
+            if "values" in tracker:
+                values = tracker["values"]
+                if tracker.get("reduce_group") is not None:
+                    torch.distributed.all_reduce(values, group=tracker.get("reduce_group"))
+                if tracker.get("avg_group") is not None:
+                    torch.distributed.all_reduce(values, group=tracker["avg_group"], op=torch.distributed.ReduceOp.AVG)
+                # here we assume only one mtp layer
+                mtp_losses = (tracker["values"] * mtp_loss_scale).item()
+                MTPLossLoggingHelper.clean_loss_in_tracker()
+
+                # CI check: verify MTP loss is within expected bounds
+                if args.ci_test:
+                    from slime.backends.megatron_utils.ci_utils import check_mtp_loss
+
+                    check_mtp_loss(mtp_losses)
+
+        # per train step log.
+        if (
+            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
+        ):
+            accumulated_step_id = rollout_id * num_steps_per_rollout + step_id
+            role = getattr(model[0], "role", "actor")
+            role_tag = "" if role == "actor" else f"{role}-"
+            log_dict = {
+                f"train/{role_tag}{key}": val.mean().item() if isinstance(val, torch.Tensor) else val
+                for key, val in loss_dict.items()
+            }
+            log_dict[f"train/{role_tag}grad_norm"] = grad_norm
+            if args.enable_mtp_training:
+                log_dict[f"train/{role_tag}mtp_loss"] = mtp_losses
+
+            for param_group_id, param_group in enumerate(optimizer.param_groups):
+                log_dict[f"train/{role_tag}lr-pg_{param_group_id}"] = opt_param_scheduler.get_lr(param_group)
+
+            log_dict["train/step"] = accumulated_step_id
+            tracking_utils.log(args, log_dict, step_key="train/step")
+
+            if args.ci_test and not args.ci_disable_kl_checker:
+                if step_id == 0 and "train/ppo_kl" in log_dict and "train/pg_clipfrac" in log_dict:
+                    if args.multi_latent_attention:
+                        # TODO: mla currently have non-zero kl, need further investigation
+                        assert log_dict["train/ppo_kl"] < 1e-8, f"{log_dict=}"
+                    else:
+                        assert log_dict["train/ppo_kl"] == 0.0 and log_dict["train/pg_clipfrac"] == 0.0, f"{log_dict=}"
+                if accumulated_step_id == 0 and "train/kl_loss" in log_dict:
+                    assert log_dict["train/kl_loss"] == 0.0, f"{log_dict=}"
+
+            logger.info(f"{role_tag}step {accumulated_step_id}: {log_dict}")
+
+            if args.ci_save_grad_norm is not None:
+                ci_save_grad_norm_path = args.ci_save_grad_norm.format(
+                    role=role,
+                    rollout_id=rollout_id,
+                    step_id=step_id,
+                )
+                torch.save(grad_norm, ci_save_grad_norm_path)
+            elif args.ci_load_grad_norm is not None:
+                ci_load_grad_norm_path = args.ci_load_grad_norm.format(
+                    role=role,
+                    rollout_id=rollout_id,
+                    step_id=step_id,
+                )
+                expected_grad_norm = torch.load(ci_load_grad_norm_path)
+                assert math.isclose(
+                    grad_norm,
+                    expected_grad_norm,
+                    rel_tol=0.01,
+                    abs_tol=0.01,
+                ), f"grad norm mismatch: {grad_norm} != {expected_grad_norm}"
+    # Close out pre-hooks if using distributed optimizer and overlapped param gather.
+    if pre_hook_enabled:
+        disable_forward_pre_hook(model)
+
+
+def save(
+    iteration: int, model: Sequence[DDP], optimizer: MegatronOptimizer, opt_param_scheduler: OptimizerParamScheduler
+) -> None:
+    """Persist a training checkpoint safely with forward hooks disabled.
+
+    Args:
+        iteration (int): Current global iteration number.
+        model (Sequence[DDP]): Sequence of DDP-wrapped model chunks.
+        optimizer (MegatronOptimizer): Optimizer instance.
+        opt_param_scheduler (OptimizerParamScheduler): LR/WD scheduler.
+    """
+    args = get_args()
+    if should_disable_forward_pre_hook(args):
+        disable_forward_pre_hook(model)
+    save_checkpoint(
+        iteration,
+        model,
+        optimizer,
+        opt_param_scheduler,
+        num_floating_point_operations_so_far=0,
+        checkpointing_context=None,
+        train_data_iterator=None,
+        preprocess_common_state_dict_fn=None,
+    )
+    if should_disable_forward_pre_hook(args):
+        enable_forward_pre_hook(model)
+
+
+def initialize_model_and_optimizer(
+    args: Namespace, role: str = "actor"
+) -> tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler, int]:
+    """Initialize model(s), optimizer, scheduler, and load from checkpoint.
+
+    Args:
+        args (Namespace): Runtime arguments.
+        role (str): Logical role of the model (e.g., "actor", "critic").
+
+    Returns:
+        tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler, int]:
+            DDP-wrapped model chunks, optimizer, scheduler, and iteration index.
+    """
+
+    if torch.version.hip:
+        import megatron.core.dist_checkpointing.strategies.filesystem_async as filesystem_async_module
+        from slime.utils.rocm_checkpoint_writer import ROCmFileSystemWriterAsync
+
+        filesystem_async_module.FileSystemWriterAsync = ROCmFileSystemWriterAsync
+        print("[ROCm] Applied FileSystemWriterAsync patch for HIP compatibility")
+
+    model, optimizer, opt_param_scheduler = setup_model_and_optimizer(args, role)
+    model[0].role = role
+    clear_memory()
+    iteration, _ = load_checkpoint(
+        model,
+        optimizer,
+        opt_param_scheduler,
+        checkpointing_context={},
+        skip_load_to_model_and_opt=False,
+    )
+    clear_memory()
+
+    opt_param_scheduler.step(increment=iteration * args.global_batch_size)
+
+    return model, optimizer, opt_param_scheduler, iteration

slime/backends/megatron_utils/model_provider.py

@@ -0,0 +1,180 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Adapt from https://github.com/NVIDIA/Megatron-LM/blob/b1efb3c7126ef7615e8c333432d76e08038e17ff/pretrain_gpt.py
+import argparse
+import inspect
+from contextlib import nullcontext
+from typing import Literal
+
+import torch
+from megatron.core import tensor_parallel
+from megatron.core.models.gpt import GPTModel
+from megatron.core.models.gpt.gpt_layer_specs import (
+    get_gpt_decoder_block_spec,
+    get_gpt_layer_local_spec,
+    get_gpt_layer_with_transformer_engine_spec,
+)
+from megatron.core.transformer.spec_utils import import_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.training.arguments import core_transformer_config_from_args
+
+
+# Adapt from https://github.com/volcengine/verl/blob/c3b20575d2bc815fcccd84bddb4c0401fc4b632b/verl/models/llama/megatron/layers/parallel_linear.py#L82
+class LinearForLastLayer(torch.nn.Linear):
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int,
+        *,
+        config: TransformerConfig,
+        bias: bool = True,
+    ) -> None:
+        super().__init__(in_features=input_size, out_features=output_size, bias=bias)
+        self.sequence_parallel = config.sequence_parallel
+        if self.sequence_parallel:
+            self.weight.sequence_parallel = True
+
+        self.weight.data.normal_(mean=0.0, std=0.02)
+        if bias:
+            self.bias.data.zero_()
+
+    def forward(
+        self,
+        input_: torch.Tensor,
+        weight: torch.Tensor | None = None,
+        runtime_gather_output: bool | None = None,
+    ) -> tuple[torch.Tensor, None]:
+        logits = super().forward(input_)
+        logits = logits.float()
+        if self.sequence_parallel:
+            logits = tensor_parallel.gather_from_sequence_parallel_region(logits, tensor_parallel_output_grad=False)
+        return logits, None
+
+
+def get_model_provider_func(
+    args: argparse.Namespace,
+    role: Literal["actor", "critic"] = "actor",
+):
+    if args.megatron_to_hf_mode == "bridge":
+        from megatron.bridge import AutoBridge
+
+        bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)
+        provider = bridge.to_megatron_provider(load_weights=False)
+        # TODO: we should not manually set this...
+        provider.tensor_model_parallel_size = args.tensor_model_parallel_size
+        provider.pipeline_model_parallel_size = args.pipeline_model_parallel_size
+        provider.expert_model_parallel_size = args.expert_model_parallel_size
+        provider.expert_tensor_parallel_size = args.expert_tensor_parallel_size
+        provider.sequence_parallel = args.sequence_parallel
+        provider.finalize()
+        return provider.provide
+
+    def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage: int | None = None) -> GPTModel:
+        """Builds the model.
+
+        If you set the use_legacy_models to True, it will return the legacy GPT model and if not the mcore GPT model.
+
+        Args:
+            pre_process (bool, optional): Set to true if you need to compute embedings. Defaults to True.
+            post_process (bool, optional): Set to true if you need to want to compute output logits/loss. Defaults to True.
+
+
+        Returns:
+            Union[GPTModel, megatron.legacy.model.GPTModel]: The returned model
+        """
+        use_te = args.transformer_impl == "transformer_engine"
+
+        # Experimental loading arguments from yaml
+        config: TransformerConfig = core_transformer_config_from_args(args)
+
+        if args.spec is not None:
+            transformer_layer_spec = import_module(args.spec)
+            # Allow the spec to be a function so that user can use customized Megatron easier.
+            if callable(transformer_layer_spec):
+                transformer_layer_spec = transformer_layer_spec(args, config, vp_stage)
+        else:
+            if args.num_experts:
+                # Define the decoder block spec
+                kwargs = {
+                    "use_transformer_engine": use_te,
+                }
+                if vp_stage is not None:
+                    kwargs["vp_stage"] = vp_stage
+                transformer_layer_spec = get_gpt_decoder_block_spec(config, **kwargs)
+            else:
+                # Define the decoder layer spec
+                if use_te:
+                    transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(
+                        num_experts=args.num_experts,
+                        moe_grouped_gemm=args.moe_grouped_gemm,
+                        qk_layernorm=args.qk_layernorm,
+                        multi_latent_attention=args.multi_latent_attention,
+                        moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm,
+                    )
+                else:
+                    transformer_layer_spec = get_gpt_layer_local_spec(
+                        num_experts=args.num_experts,
+                        moe_grouped_gemm=args.moe_grouped_gemm,
+                        qk_layernorm=args.qk_layernorm,
+                        multi_latent_attention=args.multi_latent_attention,
+                        moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm,
+                    )
+
+        build_model_context = nullcontext
+        build_model_context_args = {}
+        if args.fp8_param_gather:
+            try:
+                from transformer_engine.pytorch import fp8_model_init
+
+                build_model_context = fp8_model_init
+                build_model_context_args["enabled"] = True
+
+                # Check if fp8_model_init supports preserve_high_precision_init_val
+                if "preserve_high_precision_init_val" in inspect.signature(fp8_model_init).parameters:
+                    build_model_context_args["preserve_high_precision_init_val"] = True
+            except Exception as e:
+                raise RuntimeError(
+                    "--fp8-param-gather requires `fp8_model_init` from TransformerEngine, but not found."
+                ) from e
+
+        kwargs = {
+            "config": config,
+            "transformer_layer_spec": transformer_layer_spec,
+            "vocab_size": args.padded_vocab_size,
+            "max_sequence_length": args.max_position_embeddings,
+            "pre_process": pre_process,
+            "post_process": post_process,
+            "fp16_lm_cross_entropy": args.fp16_lm_cross_entropy,
+            "parallel_output": True,
+            "share_embeddings_and_output_weights": not args.untie_embeddings_and_output_weights,
+            "position_embedding_type": args.position_embedding_type,
+            "rotary_percent": args.rotary_percent,
+            "rotary_base": args.rotary_base,
+            "rope_scaling": args.use_rope_scaling,
+        }
+
+        if vp_stage is not None:
+            kwargs["vp_stage"] = vp_stage
+
+        if args.mtp_num_layers:
+            from megatron.core.models.gpt.gpt_layer_specs import get_gpt_mtp_block_spec
+
+            mtp_kwargs = {
+                "use_transformer_engine": use_te,
+            }
+            if vp_stage is not None:
+                mtp_kwargs["vp_stage"] = vp_stage
+
+            mtp_block_spec = get_gpt_mtp_block_spec(config, transformer_layer_spec, **mtp_kwargs)
+            kwargs["mtp_block_spec"] = mtp_block_spec
+
+        with build_model_context(**build_model_context_args):
+            model = GPTModel(**kwargs)
+
+        if post_process and role == "critic":
+            model.output_layer = LinearForLastLayer(input_size=config.hidden_size, output_size=1, config=config)
+
+        return model
+
+    return model_provider

slime/backends/megatron_utils/sglang.py

@@ -0,0 +1,34 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# the file to manage all sglang deps in the megatron actor
+try:
+    from sglang.srt.layers.quantization.fp8_utils import quant_weight_ue8m0, transform_scale_ue8m0
+    from sglang.srt.model_loader.utils import should_deepgemm_weight_requant_ue8m0
+except ImportError:
+    quant_weight_ue8m0 = None
+    transform_scale_ue8m0 = None
+    should_deepgemm_weight_requant_ue8m0 = None
+
+try:
+    from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
+except ImportError:
+    from sglang.srt.patch_torch import monkey_patch_torch_reductions
+
+
+from sglang.srt.utils import MultiprocessingSerializer
+
+
+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]
+
+__all__ = [
+    "quant_weight_ue8m0",
+    "transform_scale_ue8m0",
+    "should_deepgemm_weight_requant_ue8m0",
+    "monkey_patch_torch_reductions",
+    "MultiprocessingSerializer",
+    "FlattenedTensorBucket",
+]

slime/backends/megatron_utils/update_weight/__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/megatron_utils/update_weight/common.py

@@ -0,0 +1,238 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import inspect
+import re
+from argparse import Namespace
+from collections.abc import Iterator, Sequence
+
+import torch
+import torch.distributed as dist
+from megatron.core import mpu
+from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
+
+from slime.backends.megatron_utils.misc_utils import strip_param_name_prefix
+from slime.utils.types import ParamInfo
+
+
+def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor:
+    """
+    All-gather TP-sharded param to full tensor. expert_bias→param, non-TP/duplicated→param.data.
+    Uses expert-TP for ".experts.", else regular-TP. linear_fc1 rechunked (GLU), linear_fc2 dim fix.
+    """
+    if "expert_bias" in name:
+        return param
+
+    assert hasattr(param, "tensor_model_parallel"), f"{name} does not have tensor_model_parallel attribute"
+    if not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated":
+        return param.data
+
+    if ".experts." in name:
+        tp_size = mpu.get_expert_tensor_parallel_world_size()
+        tp_group = mpu.get_expert_tensor_parallel_group()
+    else:
+        tp_size = mpu.get_tensor_model_parallel_world_size()
+        tp_group = mpu.get_tensor_model_parallel_group()
+
+    param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)]
+    dist.all_gather(param_partitions, param.data, group=tp_group)
+    partition_dim = param.partition_dim
+    assert param.partition_stride == 1, "partition_stride != 1 is not supported"
+    # TODO: here we did an extra copy during concat, maybe merge this with convert_to_hf is better?
+    # TODO: check only GLU is used.
+    if "linear_fc1.weight" in name:
+        param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
+        param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
+    # this is bug in megatron's grouped moe.
+    if "linear_fc2.weight" in name:
+        if partition_dim == 0:
+            partition_dim = 1
+    param = torch.cat(param_partitions, dim=partition_dim)
+    return param
+
+
+def all_gather_params_async(
+    param_infos_and_params: list[tuple[ParamInfo, torch.Tensor]],
+) -> list[torch.Tensor]:
+    """
+    Parallel TP all-gather for multiple params. Loop 1: for each TP param, allocate buffers +
+    dist.all_gather(async_op=True) on expert-TP/regular-TP group (skip expert_bias/non-TP/duplicated).
+    Loop 2: wait all NCCL handles (enables overlap). Loop 3: concat partitions + apply GLU rechunk/MoE dim fix.
+    """
+    # Phase 1: Start all async all_gather operations
+    gather_tasks = []
+    handles = []
+
+    for info, param in param_infos_and_params:
+        # Prepare async all_gather
+        if "expert_bias" in info.name:
+            gather_tasks.append((info, param, None, None, None))
+            handles.append(None)
+        elif not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated":
+            gather_tasks.append((info, param.data, None, None, None))
+            handles.append(None)
+        else:
+            # Start async all_gather
+            if ".experts." in info.name:
+                tp_size = mpu.get_expert_tensor_parallel_world_size()
+                tp_group = mpu.get_expert_tensor_parallel_group()
+            else:
+                tp_size = mpu.get_tensor_model_parallel_world_size()
+                tp_group = mpu.get_tensor_model_parallel_group()
+
+            param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)]
+            handle = dist.all_gather(param_partitions, param.data, group=tp_group, async_op=True)
+            gather_tasks.append((info, None, handle, param_partitions, param.partition_dim))
+            handles.append(handle)
+
+    # Phase 2: Wait for ALL async operations to complete at once
+    # This ensures maximum parallelism by not blocking on individual operations
+    for handle in handles:
+        if handle is not None:
+            handle.wait()
+
+    # Phase 3: Process all results after all communications are done
+    gathered_params = []
+    for info, direct_param, handle, param_partitions, partition_dim in gather_tasks:
+        if handle is None:
+            # No all_gather needed
+            param = direct_param
+        else:
+            # Process the gathered partitions (same logic as original all_gather_param)
+            assert partition_dim is not None, "partition_stride != 1 is not supported"
+            # TODO: here we did an extra copy during concat, maybe merge this with convert_to_hf is better?
+            # TODO: check only GLU is used.
+            if "linear_fc1.weight" in info.name:
+                param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
+                param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
+            # this is bug in megatron's grouped moe.
+            if "linear_fc2.weight" in info.name:
+                if partition_dim == 0:
+                    partition_dim = 1
+            param = torch.cat(param_partitions, dim=partition_dim)
+
+        gathered_params.append(param)
+
+    return gathered_params
+
+
+def named_params_and_buffers(
+    args: Namespace,
+    model: Sequence[torch.nn.Module],
+    convert_to_global_name: bool = True,
+    translate_gpu_to_cpu: bool = False,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    if convert_to_global_name:
+        ans = _named_params_and_buffers_global(args, model)
+    else:
+        ans = _named_params_and_buffers_vanilla(model)
+
+    if translate_gpu_to_cpu:
+        ans = ((name, _maybe_get_cpu_backup(tensor)) for name, tensor in ans)
+
+    return ans
+
+
+def _maybe_get_cpu_backup(x: torch.Tensor):
+    from torch_memory_saver import torch_memory_saver
+
+    if (cpu_tensor := torch_memory_saver.get_cpu_backup(x)) is not None:
+        return cpu_tensor
+
+    return x
+
+
+def _named_params_and_buffers_vanilla(model: Sequence[torch.nn.Module]) -> Iterator[tuple[str, torch.Tensor]]:
+    for vp_stage, model_module in enumerate(model):
+
+        def _compute_fqn(name, vp_stage=vp_stage):
+            return f"vp_stages.{vp_stage}.{strip_param_name_prefix(name)}"
+
+        for name, param in model_module.named_parameters():
+            yield _compute_fqn(name), param
+
+        for name, buffer in model_module.named_buffers():
+            # TODO shall we handle (almost) all buffers like Megatron Bridge
+            if "expert_bias" not in name:
+                continue
+            yield _compute_fqn(name), buffer
+
+
+def _named_params_and_buffers_global(
+    args: Namespace, model: Sequence[torch.nn.Module]
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """
+    Yield (global_name, param/buffer) with consistent names across PP/EP. Adjusts indices for
+    virtual PP + EP offsets. Handles decoder.layers, mtp.layers (Multi-Token Prediction), expert_bias.
+    """
+    ep_size = mpu.get_expert_model_parallel_world_size()
+    ep_rank = mpu.get_expert_model_parallel_rank()
+    if args.num_experts:
+        expert_offset = ep_rank * args.num_experts // ep_size
+
+    sig = inspect.signature(get_transformer_layer_offset)
+    need_vp_stage = "vp_stage" in sig.parameters
+
+    for vp_stage, model_module in enumerate(model):
+        if need_vp_stage:
+            layer_offset = get_transformer_layer_offset(model_module.config, vp_stage)
+        else:
+            layer_offset = get_transformer_layer_offset(model_module.config)
+        for name, param in model_module.named_parameters():
+            # for model without ddp wrap
+            if not name.startswith("module.module."):
+                name = "module." + name
+
+            decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)"
+            match = re.match(decoder_layers_pattern, name)
+            if not match:
+                # MTP (Multi-Token Prediction) layers for speculative decoding
+                mtp_layers_pattern = r"module\.module\.mtp\.layers\.(\d+)\.(.+)"
+                match = re.match(mtp_layers_pattern, name)
+                if not match:
+                    yield name, param
+                    continue
+
+                # MTP layer indices start from 0
+                layer_idx, rest = match.groups()
+                expert_pattern = r"transformer_layer.mlp.experts\.(.+)\.weight(\d+)"
+                match = re.match(expert_pattern, rest)
+                if not match:
+                    yield name, param
+                    continue
+
+                rest, expert_idx = match.groups()
+                expert_idx = int(expert_idx) + expert_offset
+                yield f"module.module.mtp.layers.{layer_idx}.transformer_layer.mlp.experts.{rest}.weight{expert_idx}", param
+                continue
+
+            layer_idx, rest = match.groups()
+            layer_idx = int(layer_idx) + layer_offset
+
+            # this is hardcoded for te grouped matmul
+            expert_pattern = r"mlp.experts\.(.+)\.weight(\d+)"
+            match = re.match(expert_pattern, rest)
+            if match:
+                rest, expert_idx = match.groups()
+                expert_idx = int(expert_idx) + expert_offset
+                yield f"module.module.decoder.layers.{layer_idx}.mlp.experts.{rest}.weight{expert_idx}", param
+            else:
+                yield f"module.module.decoder.layers.{layer_idx}.{rest}", param
+
+        # treat expert bias as normal parameters
+        for name, buffer in model_module.named_buffers():
+            # TODO shall we handle (almost) all buffers like Megatron Bridge
+            if "expert_bias" not in name:
+                continue
+            # for model without ddp wrap
+            if not name.startswith("module.module."):
+                name = "module." + name
+
+            decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)"
+            match = re.match(decoder_layers_pattern, name)
+            if not match:
+                yield name, buffer
+            else:
+                layer_idx, rest = match.groups()
+                layer_idx = int(layer_idx) + layer_offset
+                yield f"module.module.decoder.layers.{layer_idx}.{rest}", buffer

slime/backends/megatron_utils/update_weight/hf_weight_iterator_base.py

@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+
+
+class HfWeightIteratorBase(ABC):
+    @staticmethod
+    def create(args, model, **kwargs):
+        from .hf_weight_iterator_bridge import HfWeightIteratorBridge
+        from .hf_weight_iterator_direct import HfWeightIteratorDirect
+
+        c = {
+            "raw": HfWeightIteratorDirect,
+            "bridge": HfWeightIteratorBridge,
+        }[args.megatron_to_hf_mode]
+
+        return c(args, model, **kwargs)
+
+    def __init__(self, args, model, model_name, quantization_config):
+        self.args = args
+        self.model = model
+        self.model_name = model_name
+        self.quantization_config = quantization_config
+
+    @abstractmethod
+    def get_hf_weight_chunks(self, megatron_local_weights):
+        """
+        Mental model of the API:
+        megatron_model.to_hf_magically().named_parameters()
+        """
+        raise NotImplementedError

slime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py

@@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import dataclasses
+
+from slime.utils import megatron_bridge_utils
+from slime.utils.iter_utils import chunk_named_params_by_size
+
+from ..megatron_to_hf import postprocess_hf_param
+from ..misc_utils import strip_param_name_prefix
+from .hf_weight_iterator_base import HfWeightIteratorBase
+
+
+class HfWeightIteratorBridge(HfWeightIteratorBase):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        from megatron.bridge import AutoBridge
+        import slime_plugins.megatron_bridge  # noqa: F401
+
+        self._bridge = AutoBridge.from_hf_pretrained(self.args.hf_checkpoint)
+
+    def get_hf_weight_chunks(self, megatron_local_weights):
+        # TODO support quantization (e.g. modify megatron-bridge to provide megatron param name)
+        renamed_megatron_local_weights = {strip_param_name_prefix(k): v for k, v in megatron_local_weights.items()}
+        with megatron_bridge_utils.patch_megatron_model(self.model):
+            conversion_tasks = self._bridge.get_conversion_tasks(self.model)
+            conversion_tasks = _process_conversion_tasks(conversion_tasks, renamed_megatron_local_weights)
+
+            named_weights = self._bridge.export_hf_weights(self.model, cpu=False, conversion_tasks=conversion_tasks)
+
+            named_weights = (
+                (
+                    hf_param_name,
+                    postprocess_hf_param(
+                        args=self.args,
+                        megatron_param_name=megatron_param_name,
+                        hf_param_name=hf_param_name,
+                        param=weight,
+                    ),
+                )
+                for hf_param_name, weight, megatron_param_name in named_weights
+            )
+
+            yield from chunk_named_params_by_size(named_weights, chunk_size=self.args.update_weight_buffer_size)
+
+
+def _process_conversion_tasks(vanilla_conversion_tasks, new_weight_dict):
+    def _handle_one(task):
+        if task.param_weight is None:
+            return task
+
+        weight_dict_key = f"vp_stages.{task.vp_stage}.{task.param_name}"
+        assert (
+            weight_dict_key in new_weight_dict
+        ), f"{weight_dict_key=} not in new_weight_dict ({task.vp_stage=}, {task.param_name=}, {list(new_weight_dict)=})"
+
+        new_param_weight = new_weight_dict[weight_dict_key]
+        new_param_weight = new_param_weight.cuda()
+        return dataclasses.replace(task, param_weight=new_param_weight)
+
+    return _MapWithLen(_handle_one, vanilla_conversion_tasks)
+
+
+class _MapWithLen:
+    def __init__(self, fn, xs):
+        self.fn = fn
+        self.xs = xs
+
+    def __len__(self):
+        return len(self.xs)
+
+    def __iter__(self):
+        for x in self.xs:
+            yield self.fn(x)

slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py

@@ -0,0 +1,215 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import dataclasses
+from argparse import Namespace
+from collections.abc import Sequence
+
+import torch
+import torch.distributed as dist
+from megatron.core import mpu
+from tqdm import tqdm
+
+from slime.utils.distributed_utils import get_gloo_group
+from slime.utils.types import ParamInfo
+
+from ..megatron_to_hf import convert_to_hf
+from ..sglang import monkey_patch_torch_reductions
+from .common import all_gather_params_async, named_params_and_buffers
+from .hf_weight_iterator_base import HfWeightIteratorBase
+
+
+class HfWeightIteratorDirect(HfWeightIteratorBase):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.megatron_local_param_info_buckets = _get_megatron_local_param_info_buckets(self.args, self.model)
+
+    def get_hf_weight_chunks(self, megatron_local_weights):
+        rank = dist.get_rank()
+
+        for megatron_local_param_infos in tqdm(
+            self.megatron_local_param_info_buckets, disable=rank != 0, desc="Update weights"
+        ):
+            megatron_full_params = _get_megatron_full_params(megatron_local_param_infos, megatron_local_weights)
+            hf_named_tensors = self._convert_to_hf_named_tensors(megatron_full_params, megatron_local_param_infos)
+            yield hf_named_tensors
+            del megatron_full_params
+
+    def _convert_to_hf_named_tensors(self, megatron_full_params: Sequence[torch.Tensor], param_infos: list[ParamInfo]):
+        hf_named_tensors = []
+        for info, param in zip(param_infos, megatron_full_params, strict=False):
+            hf_named_tensors.extend(
+                convert_to_hf(self.args, self.model_name, info.name, param, self.quantization_config)
+            )
+        return hf_named_tensors
+
+
+def _get_megatron_full_params(
+    megatron_local_param_infos: Sequence[ParamInfo],
+    megatron_local_weights,
+) -> Sequence[torch.Tensor]:
+    monkey_patch_torch_reductions()
+    pp_size = mpu.get_pipeline_model_parallel_world_size()
+    ep_size = mpu.get_expert_model_parallel_world_size()
+    rank = dist.get_rank()
+    # init params:
+    params = []
+    for info in megatron_local_param_infos:
+        if dist.get_rank() == info.src_rank:
+            params.append(
+                torch.nn.Parameter(
+                    megatron_local_weights[info.name].to(device=torch.cuda.current_device(), non_blocking=True),
+                    requires_grad=False,
+                )
+            )
+        else:
+            params.append(torch.empty(info.shape, dtype=info.dtype, device=torch.cuda.current_device()))
+    torch.cuda.synchronize()
+
+    # broadcast params across pp ranks
+    if pp_size > 1:
+        handles = []
+        for info, param in zip(megatron_local_param_infos, params, strict=False):
+            if info.src_rank in dist.get_process_group_ranks(mpu.get_pipeline_model_parallel_group()):
+                handles.append(
+                    torch.distributed.broadcast(
+                        param, src=info.src_rank, group=mpu.get_pipeline_model_parallel_group(), async_op=True
+                    )
+                )
+        for handle in handles:
+            handle.wait()
+
+    # broadcast params across ep ranks
+    if ep_size > 1:
+        handles = []
+        for info, param in zip(megatron_local_param_infos, params, strict=False):
+            if ".experts." in info.name:
+                src_rank = (
+                    info.src_rank
+                    if info.src_rank in dist.get_process_group_ranks(mpu.get_expert_model_parallel_group())
+                    else rank
+                )
+                handles.append(
+                    torch.distributed.broadcast(
+                        param, src=src_rank, group=mpu.get_expert_model_parallel_group(), async_op=True
+                    )
+                )
+        for handle in handles:
+            handle.wait()
+
+    # Set tp attrs for all params
+    for info, param in zip(megatron_local_param_infos, params, strict=False):
+        for key, value in info.attrs.items():
+            setattr(param, key, value)
+
+    # Batch async all_gather for all parameters
+    gathered_params = all_gather_params_async(list(zip(megatron_local_param_infos, params, strict=False)))
+
+    return gathered_params
+
+
+def _get_megatron_local_param_info_buckets(args: Namespace, model: Sequence[torch.nn.Module]) -> list[list[ParamInfo]]:
+    """
+    Partition params into buckets ≤ update_weight_buffer_size (with TP replication).
+    """
+    param_infos = _get_megatron_local_param_infos(args, model)
+    param_info_buckets = [[]]  # Start with one empty bucket
+    buffer_size = 0  # Track current bucket size in bytes
+
+    for info in param_infos:
+        # Expert params use expert-TP size, others use regular-TP size
+        if ".experts." in info.name:
+            tp_size = mpu.get_expert_tensor_parallel_world_size()
+        else:
+            tp_size = mpu.get_tensor_model_parallel_world_size()
+
+        # Full param size = shard size × TP replicas (all-gather will reconstruct full param)
+        param_size = info.size * tp_size
+
+        # If adding this param exceeds limit AND current bucket has params: start new bucket
+        if buffer_size + param_size > args.update_weight_buffer_size and len(param_info_buckets[-1]) > 0:
+            param_info_buckets.append([])
+            buffer_size = 0
+
+        # Add param to current bucket and update size
+        param_info_buckets[-1].append(info)
+        buffer_size += param_size
+
+    return param_info_buckets
+
+
+def _get_megatron_local_param_infos(args: Namespace, model: Sequence[torch.nn.Module]) -> list[ParamInfo]:
+    """
+    Build global param metadata: collect → exchange PP/EP → resolve duplicates (MTP virtual PP)
+    by min src_rank → validate. Returns sorted ParamInfo identical across all ranks.
+    """
+    pp_size = mpu.get_pipeline_model_parallel_world_size()
+    ep_size = mpu.get_expert_model_parallel_world_size()
+
+    param_infos = {}
+    rank = dist.get_rank()
+    for name, param in named_params_and_buffers(args, model):
+        param_infos[name] = ParamInfo(
+            name=name,
+            dtype=param.dtype,
+            shape=param.shape,
+            attrs={
+                "tensor_model_parallel": getattr(param, "tensor_model_parallel", False),
+                "partition_dim": getattr(param, "partition_dim", -1),
+                "partition_stride": getattr(param, "partition_stride", 1),
+                "parallel_mode": getattr(param, "parallel_mode", None),
+            },
+            size=param.numel() * param.element_size(),
+            src_rank=rank,
+        )
+
+    if pp_size > 1:
+        param_infos_list = [None] * pp_size
+        dist.all_gather_object(
+            obj=(rank, param_infos), object_list=param_infos_list, group=mpu.get_pipeline_model_parallel_group()
+        )
+        for src_rank, infos in param_infos_list:
+            if src_rank == rank:
+                continue
+            for name, info in infos.items():
+                if name in param_infos:
+                    assert args.mtp_num_layers is not None
+                    old_info = param_infos[name]
+                    if old_info.src_rank > src_rank:
+                        param_infos[name] = info
+                else:
+                    param_infos[name] = info
+
+    if ep_size > 1:
+        param_infos_list = [None] * ep_size
+        dist.all_gather_object(
+            obj=(rank, param_infos), object_list=param_infos_list, group=mpu.get_expert_model_parallel_group()
+        )
+        for src_rank, infos in param_infos_list:
+            for name, info in infos.items():
+                if name not in param_infos:
+                    # here we need to set the src_rank to the rank within the expert model parallel group
+                    info = dataclasses.replace(info, src_rank=src_rank)
+                    param_infos[name] = info
+
+    param_infos = list(param_infos.values())
+    param_infos = sorted(param_infos, key=lambda info: info.name)
+
+    # Check all ranks has the same parameter info
+    all_param_info_list = [None] * dist.get_world_size()
+    dist.all_gather_object(
+        obj=param_infos,
+        object_list=all_param_info_list,
+        group=get_gloo_group(),
+    )
+    for i, param_info in enumerate(param_infos):
+        for infos in all_param_info_list:
+            assert infos[i].name == param_info.name, f"Parameter name mismatch: {infos[i].name} != {param_info.name}"
+            assert (
+                infos[i].shape == param_info.shape
+            ), f"Parameter shape mismatch: {infos[i].shape} != {param_info.shape}"
+            assert (
+                infos[i].dtype == param_info.dtype
+            ), f"Parameter dtype mismatch: {infos[i].dtype} != {param_info.dtype}"
+
+    return param_infos

slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py

@@ -0,0 +1,306 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import socket
+import time
+from argparse import Namespace
+from collections.abc import Callable, Mapping, Sequence
+
+import ray
+import torch
+import torch.distributed as dist
+from megatron.core import mpu
+from ray import ObjectRef
+from ray.actor import ActorHandle
+from tqdm import tqdm
+
+from slime.utils.distributed_utils import get_gloo_group, init_process_group
+
+from ..megatron_to_hf import convert_to_hf
+from .common import all_gather_param, named_params_and_buffers
+
+
+class UpdateWeightFromDistributed:
+    """
+    Update distributed engines via NCCL. Each PP rank: group "slime-pp_{pp_rank}",
+    only DP=TP=0 broadcasts. Non-expert (TP) and expert (EP) params separate.
+    """
+
+    def __init__(
+        self,
+        args: Namespace,
+        model: Sequence[torch.nn.Module],
+        weights_getter: Callable[[], Mapping[str, torch.Tensor]],
+        *,
+        model_name: str,
+        quantization_config: dict[str, int | str | list[str]] | None,
+    ) -> None:
+        """
+        Initialize. Groups created in connect_rollout_engines.
+        """
+        self.args = args
+        self.model = model
+        self.model_name = model_name
+        self.quantization_config = quantization_config
+        self.weight_version = 0
+        self._model_update_groups = None
+        self.rollout_engines = []
+
+    def connect_rollout_engines(
+        self, rollout_engines: Sequence[ActorHandle], rollout_engine_lock: ActorHandle
+    ) -> None:
+        """
+        Create NCCL "slime-pp_{pp_rank}" if PP source (DP=TP=0). Lock prevents concurrent broadcasts.
+        """
+        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_pp_src_rank = (
+            mpu.get_data_parallel_rank(with_context_parallel=True) == 0 and mpu.get_tensor_model_parallel_rank() == 0
+        )
+        pp_rank = mpu.get_pipeline_model_parallel_rank()
+        if self._is_pp_src_rank:
+            self._group_name = f"slime-pp_{pp_rank}"
+
+        if self._is_pp_src_rank:
+            if self._model_update_groups is not None:
+                disconnect_rollout_engines_from_distributed(
+                    self.args, self._group_name, self._model_update_groups, self.rollout_engines
+                )
+            self._model_update_groups = connect_rollout_engines_from_distributed(
+                self.args, self._group_name, rollout_engines
+            )
+
+    @torch.no_grad()
+    def update_weights(self) -> None:
+        """
+        Pause → flush → non-expert (TP) → expert (EP) → continue. Progress on PP source.
+        """
+        if not self.rollout_engines:
+            return
+
+        self.weight_version += 1
+
+        if dist.get_rank() == 0:
+            ray.get([engine.pause_generation.remote() for engine in self.rollout_engines])
+            ray.get([engine.flush_cache.remote() for engine in self.rollout_engines])
+        dist.barrier(group=get_gloo_group())
+
+        buffer_size = 0
+        converted_named_tensors = []
+        # non expert params
+        pbar = tqdm(desc=f"[{self._group_name}] Update weights", total=0) if self._is_pp_src_rank else None
+
+        for name, param in named_params_and_buffers(self.args, self.model):
+            if ".experts." in name:
+                continue
+            buffer_size = self._update_weight_from_distributed(
+                name, param, converted_named_tensors, buffer_size, pbar=pbar
+            )
+
+        if converted_named_tensors:
+            self._update_bucket_weights_from_distributed(converted_named_tensors, pbar=pbar)
+
+        dist.barrier(group=get_gloo_group())
+
+        buffer_size = 0
+        named_tensors = []
+        for name, param in named_params_and_buffers(self.args, self.model):
+            if ".experts." not in name:
+                continue
+            buffer_size = self._update_expert_weight_from_distributed(
+                name, param, named_tensors, buffer_size, pbar=pbar
+            )
+
+        if named_tensors:
+            self._update_expert_bucket_weights_from_distributed(named_tensors, pbar=pbar)
+
+        dist.barrier(group=get_gloo_group())
+        if dist.get_rank() == 0:
+            ray.get([engine.continue_generation.remote() for engine in self.rollout_engines])
+        dist.barrier(group=get_gloo_group())
+
+    def _update_weight_from_distributed(
+        self,
+        name: str,
+        param: torch.nn.Parameter,
+        converted_named_tensors: list[tuple[str, torch.Tensor]],
+        buffer_size: int,
+        pbar: tqdm | None = None,
+    ) -> int | None:
+        """
+        Non-expert: gather TP → rm pad → HF → buffer (flush if full). All gather, PP source buffers.
+        Returns updated bytes on source, None on non-source.
+        """
+        param = all_gather_param(name, param)
+        if not self._is_pp_src_rank:
+            return
+
+        param_size = param.numel() * param.element_size()
+        if buffer_size + param_size > self.args.update_weight_buffer_size:
+            self._update_bucket_weights_from_distributed(converted_named_tensors, pbar=pbar)
+            buffer_size = 0
+        converted_named_tensors += convert_to_hf(self.args, self.model_name, name, param, self.quantization_config)
+        buffer_size += param_size
+        return buffer_size
+
+    def _update_expert_weight_from_distributed(
+        self,
+        name: str,
+        param: torch.nn.Parameter,
+        named_tensors: list[tuple[str, torch.Tensor]],
+        buffer_size: int,
+        pbar: tqdm | None = None,
+    ) -> int:
+        """
+        Expert: gather TP → rm pad → buffer. EP gather + HF deferred. Threshold × EP size.
+        """
+        param = all_gather_param(name, param)
+
+        param_size = param.numel() * param.element_size()
+        if (
+            buffer_size + param_size
+        ) * mpu.get_expert_model_parallel_world_size() > self.args.update_weight_buffer_size:
+            self._update_expert_bucket_weights_from_distributed(named_tensors, pbar=pbar)
+            buffer_size = 0
+
+        named_tensors.append((name, param))
+        buffer_size += param_size
+        return buffer_size
+
+    def _update_expert_bucket_weights_from_distributed(
+        self, named_tensors: list[tuple[str, torch.Tensor]], pbar: tqdm | None = None
+    ) -> None:
+        """
+        Gather EP → HF → broadcast. Clears buffer.
+        """
+        names = [name for name, _ in named_tensors]
+        all_names = [None] * mpu.get_expert_model_parallel_world_size()
+        dist.all_gather_object(all_names, names, group=mpu.get_expert_model_parallel_group())
+
+        for names in all_names:
+            assert len(named_tensors) == len(names), f"mismatch names length: {len(named_tensors)} != {len(names)}"
+
+        all_gathered_params = [[] for _ in range(mpu.get_expert_model_parallel_world_size())]
+        handles = []
+        for i, (_name, param) in enumerate(named_tensors):
+            params = [
+                torch.empty_like(param.data, device=torch.cuda.current_device())
+                for _ in range(mpu.get_expert_model_parallel_world_size())
+            ]
+            handle = dist.all_gather(params, param.data, group=mpu.get_expert_model_parallel_group(), async_op=True)
+            handles.append(handle)
+            for ep_rank, names in enumerate(all_names):
+                all_gathered_params[ep_rank].append((names[i], params[ep_rank]))
+        for handle in handles:
+            handle.wait()
+
+        named_tensors.clear()
+        if not self._is_pp_src_rank:
+            return
+
+        all_gathered_params = sum(all_gathered_params, [])
+        converted_hf_tensors = []
+        for name, param in all_gathered_params:
+            converted_hf_tensors += convert_to_hf(self.args, self.model_name, name, param, self.quantization_config)
+
+        self._update_bucket_weights_from_distributed(converted_hf_tensors, pbar)
+
+    def _update_bucket_weights_from_distributed(
+        self, converted_named_tensors: list[tuple[str, torch.Tensor]], pbar: tqdm | None = None
+    ) -> None:
+        """
+        Lock → broadcast → clear → unlock → pbar++. Lock prevents NCCL deadlock.
+        """
+        # lock the rollout engines to prevent dead lock on broadcast.
+        while not ray.get(self.rollout_engine_lock.acquire.remote()):
+            time.sleep(0.1)
+
+        refs = update_weights_from_distributed(
+            self._group_name,
+            self._model_update_groups,
+            self.weight_version,
+            self.rollout_engines,
+            converted_named_tensors,
+        )
+
+        ray.get(refs)
+        converted_named_tensors.clear()
+        ray.get(self.rollout_engine_lock.release.remote())
+        pbar.update(1)
+
+
+def connect_rollout_engines_from_distributed(
+    args: Namespace, group_name: str, rollout_engines: Sequence[ActorHandle]
+) -> dist.ProcessGroup:
+    """
+    Create NCCL group: training rank 0 + all engine GPUs. Blocks until joined.
+    """
+    master_address = ray._private.services.get_node_ip_address()
+    with socket.socket() as sock:
+        sock.bind(("", 0))
+        master_port = sock.getsockname()[1]
+    world_size = len(rollout_engines) * args.rollout_num_gpus_per_engine + 1
+
+    refs = [
+        engine.init_weights_update_group.remote(
+            master_address,
+            master_port,
+            i * args.rollout_num_gpus_per_engine + 1,
+            world_size,
+            group_name,
+            backend="nccl",
+        )
+        for i, engine in enumerate(rollout_engines)
+    ]
+    model_update_groups = init_process_group(
+        backend="nccl",
+        init_method=f"tcp://{master_address}:{master_port}",
+        world_size=world_size,
+        rank=0,
+        group_name=group_name,
+    )
+    ray.get(refs)
+    return model_update_groups
+
+
+def disconnect_rollout_engines_from_distributed(args, group_name, model_update_groups, rollout_engines):
+    """
+    Destroy NCCL on training and engines.
+    """
+    refs = [engine.destroy_weights_update_group.remote(group_name) for engine in rollout_engines]
+    dist.destroy_process_group(model_update_groups)
+    ray.get(refs)
+
+
+def update_weights_from_distributed(
+    group_name: str,
+    group: dist.ProcessGroup,
+    weight_version: int,
+    rollout_engines: Sequence[ActorHandle],
+    converted_named_tensors: Sequence[tuple[str, torch.Tensor]],
+) -> list[ObjectRef]:
+    """
+    Send metadata (Ray), broadcast tensors (NCCL rank 0 → engines).
+    """
+    refs = [
+        engine.update_weights_from_distributed.remote(
+            names=[name for name, _ in converted_named_tensors],
+            dtypes=[param.dtype for _, param in converted_named_tensors],
+            shapes=[param.shape for _, param in converted_named_tensors],
+            group_name=group_name,
+            weight_version=str(weight_version),
+        )
+        for engine in rollout_engines
+    ]
+
+    handles = []
+    for _, param in converted_named_tensors:
+        handles.append(dist.broadcast(param.data, 0, group=group, async_op=True))
+    for handle in handles:
+        handle.wait()
+
+    return refs

slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py

@@ -0,0 +1,209 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from argparse import Namespace
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any
+
+import ray
+import torch
+import torch.distributed as dist
+from megatron.core import mpu
+from ray import ObjectRef
+from ray.actor import ActorHandle
+
+from slime.utils.distributed_utils import get_gloo_group
+
+from ..sglang import FlattenedTensorBucket, MultiprocessingSerializer
+from .hf_weight_iterator_base import HfWeightIteratorBase
+from .update_weight_from_distributed import (
+    connect_rollout_engines_from_distributed,
+    disconnect_rollout_engines_from_distributed,
+    update_weights_from_distributed,
+)
+
+
+class UpdateWeightFromTensor:
+    """
+    Update rollout engines from tensor dict:
+    load(dict→GPU) → broadcast PP/EP(GPU NCCL) → gather TP(GPU NCCL) → convert HF(GPU) → send.
+    Colocated: GPU→CPU serialize → gather_object(Gloo CPU, collects from rollout_num_gpus_per_engine ranks) → Ray IPC to engine.
+    Distributed: GPU NCCL broadcast to remote engines.
+    """
+
+    def __init__(
+        self,
+        args: Namespace,
+        model: Sequence[torch.nn.Module],
+        weights_getter: Callable[[], Mapping[str, torch.Tensor]],
+        *,
+        model_name: str,
+        quantization_config: dict[str, int | str | list[str]] | None,
+    ) -> None:
+        """
+        Compute param buckets, create IPC Gloo groups (rollout_num_gpus_per_engine ranks/group).
+        """
+        self.args = args
+        self.model = model
+        self.weights_getter = weights_getter
+        self.model_name = model_name
+        self.quantization_config = quantization_config
+        self.weight_version = 0
+
+        self._hf_weight_iterator = HfWeightIteratorBase.create(
+            args=args, model=model, model_name=model_name, quantization_config=quantization_config
+        )
+
+        # create the group within megatron.
+        for start_rank in range(0, dist.get_world_size(), self.args.rollout_num_gpus_per_engine):
+            end_rank = start_rank + 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_group = new_group
+                self._ipc_gather_src = start_rank
+
+        self._model_update_groups = None
+
+    def connect_rollout_engines(
+        self, rollout_engines: Sequence[ActorHandle], rollout_engine_lock: ActorHandle
+    ) -> None:
+        """
+        Split colocated/distributed engines. Global source rank (DP=TP=PP=0) creates NCCL
+        for distributed. Map ranks to colocated IPC engines.
+        """
+        self.rollout_engines = rollout_engines
+        colocate_engine_nums = (
+            self.args.actor_num_nodes * self.args.actor_num_gpus_per_node // self.args.rollout_num_gpus_per_engine
+        )
+        self.use_distribute = len(rollout_engines) > colocate_engine_nums
+
+        if self.use_distribute:
+            self.rollout_engines = rollout_engines[:colocate_engine_nums]
+            self.distributed_rollout_engines = rollout_engines[colocate_engine_nums:]
+            self._is_distributed_src_rank = (
+                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() == 0
+            )
+            self._group_name = "slime"
+            if self._is_distributed_src_rank:
+                if self._model_update_groups is not None:
+                    disconnect_rollout_engines_from_distributed(
+                        self.args, self._group_name, self._model_update_groups, self.distributed_rollout_engines
+                    )
+
+                self._model_update_groups = connect_rollout_engines_from_distributed(
+                    self.args, self._group_name, self.distributed_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))
+            if dist.get_rank() in group_ranks:
+                self._ipc_engine = engine
+
+    @torch.no_grad()
+    def update_weights(self) -> None:
+        """
+        version++, flush caches, process buckets. Progress on rank 0.
+        """
+        self.weight_version += 1
+
+        rank = dist.get_rank()
+        if rank == 0:
+            ray.get([engine.flush_cache.remote() for engine in self.rollout_engines])
+        dist.barrier(group=get_gloo_group())
+
+        megatron_local_weights = self.weights_getter()
+
+        for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights):
+            refs, long_lived_tensors = self._send_hf_params(hf_named_tensors)
+            ray.get(refs)
+            del long_lived_tensors
+
+        dist.barrier(group=get_gloo_group())
+
+    def _send_hf_params(self, hf_named_tensors) -> tuple[list[ObjectRef], Any]:
+        all_refs = []
+
+        refs_colocated, long_lived_tensors = _send_to_colocated_engine(
+            hf_named_tensors,
+            ipc_engine=self._ipc_engine,
+            ipc_gather_src=self._ipc_gather_src,
+            ipc_gather_group=self._ipc_gather_group,
+            weight_version=self.weight_version,
+        )
+        all_refs.extend(refs_colocated)
+
+        if self.use_distribute and self._is_distributed_src_rank:
+            refs_distributed = update_weights_from_distributed(
+                self._group_name,
+                self._model_update_groups,
+                self.weight_version,
+                self.distributed_rollout_engines,
+                hf_named_tensors,
+            )
+            if refs_distributed:
+                all_refs.extend(refs_distributed)
+
+        return all_refs, long_lived_tensors
+
+
+def _send_to_colocated_engine(
+    hf_named_tensors: list[tuple[str, torch.Tensor]],
+    *,
+    ipc_engine,
+    ipc_gather_src,
+    ipc_gather_group,
+    weight_version,
+) -> tuple[list[ObjectRef], Any]:
+    # TODO improve
+    long_live_tensors = []
+
+    if getattr(FlattenedTensorBucket, "supports_multi_dtypes", False):
+        converted_named_tensors_by_dtypes = {"dtype": hf_named_tensors}
+    else:
+        converted_named_tensors_by_dtypes = {}
+        for name, tensor in hf_named_tensors:
+            dtype = tensor.dtype
+            if dtype not in converted_named_tensors_by_dtypes:
+                converted_named_tensors_by_dtypes[dtype] = []
+            converted_named_tensors_by_dtypes[dtype].append((name, tensor))
+
+    serialized_tensors = []
+    for _dtype, named_tensors in converted_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,
+        }
+        long_live_tensors.append(flattened_tensor_data)
+        serialized_tensors.append(MultiprocessingSerializer.serialize(flattened_tensor_data, output_str=True))
+
+    serialized_named_tensors = (
+        [None] * dist.get_world_size(ipc_gather_group) if ipc_gather_src == dist.get_rank() else None
+    )
+    dist.gather_object(
+        serialized_tensors,
+        object_gather_list=serialized_named_tensors,
+        dst=ipc_gather_src,
+        group=ipc_gather_group,
+    )
+
+    refs = []
+    if dist.get_rank() == ipc_gather_src:
+        # TODO: here we assume all ranks have the same number of dtypes, not sure if that is correct.
+        num_dtypes = len(serialized_named_tensors[0])
+        for i in range(num_dtypes):
+            kwargs = {
+                "serialized_named_tensors": [tensors[i] for tensors in serialized_named_tensors],
+                "load_format": "flattened_bucket",
+                "weight_version": str(weight_version),
+            }
+            refs.append(ipc_engine.update_weights_from_tensor.remote(**kwargs))
+
+    return refs, long_live_tensors

slime/backends/sglang_utils/__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/sglang_utils/arguments.py

@@ -0,0 +1,133 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import sglang
+from packaging.version import parse
+from sglang.srt.server_args import ServerArgs
+from slime.utils.http_utils import _wrap_ipv6
+
+
+# TODO: use all sglang router arguments with `--sglang-router` prefix
+def add_sglang_router_arguments(parser):
+    """
+    Add arguments to the parser for the SGLang router.
+    """
+    parser.add_argument(
+        "--sglang-router-ip",
+        type=str,
+        default=None,
+        help="IP address of the SGLang router",
+    )
+    parser.add_argument(
+        "--sglang-router-port",
+        type=int,
+        default=None,
+        help="Port of the SGLang router",
+    )
+    parser.add_argument(
+        "--sglang-router-request-timeout-secs",
+        type=int,
+        default=14400,
+        help="Timeout for requests to the SGLang router in seconds",
+    )
+    return parser
+
+
+def add_sglang_arguments(parser):
+    """
+    Add arguments to the parser for the SGLang server.
+    """
+    parser = add_sglang_router_arguments(parser)
+    parser.add_argument("--sglang-server-concurrency", type=int, default=512)
+
+    old_add_argument = parser.add_argument
+
+    skipped_args = [
+        "model_path",
+        "dtype",
+        "trust_remote_code",
+        "random_seed",
+        # memory
+        "enable_memory_saver",
+        # distributed
+        "tp_size",
+        "port",
+        "nnodes",
+        "node_rank",
+        "dist_init_addr",
+        "gpu_id_step",
+        "base_gpu_id",
+        "nccl_port",
+        "skip_server_warmup",
+        "enable_return_routed_experts",
+    ]
+
+    def new_add_argument_wrapper(*name_or_flags, **kwargs):
+        """
+        Add arguments to the parser, ensuring that the server arguments are prefixed and skippable.
+        """
+        # Determine the canonical name for skip check (e.g., "model_path")
+        canonical_name_for_skip_check = None
+        if "dest" in kwargs:
+            canonical_name_for_skip_check = kwargs["dest"]
+        else:
+            for flag_name_candidate in name_or_flags:
+                if isinstance(flag_name_candidate, str) and flag_name_candidate.startswith("--"):
+                    # Derive from first long flag: --foo-bar -> foo_bar
+                    stem = flag_name_candidate[2:]
+                    canonical_name_for_skip_check = stem.replace("-", "_")
+                    break
+            # If no long flag and no dest, skip logic might not catch it unless short flags imply a dest.
+
+        if canonical_name_for_skip_check and canonical_name_for_skip_check in skipped_args:
+            return  # Skip this entire argument definition
+
+        # If not skipped, proceed to prefix flags and dest
+        new_name_or_flags_list = []
+        for item_flag in name_or_flags:
+            if isinstance(item_flag, str) and item_flag.startswith("-"):
+                original_flag_stem = item_flag.lstrip("-")  # "foo-bar" from "--foo-bar", or "f" from "-f"
+                prefixed_item = f"--sglang-{original_flag_stem}"
+                new_name_or_flags_list.append(prefixed_item)
+            else:
+                # Positional arguments or non-string items
+                new_name_or_flags_list.append(item_flag)
+
+        # Prepare kwargs for the actual add_argument call.
+        # Make a copy to avoid modifying the original kwargs dict.
+        final_kwargs = kwargs.copy()
+
+        # If 'dest' is explicitly provided and is a string, prefix it.
+        # This ensures the attribute on the args namespace becomes, e.g., args.sglang_dest_name.
+        if "dest" in final_kwargs and isinstance(final_kwargs["dest"], str):
+            original_dest = final_kwargs["dest"]
+            # Avoid double prefixing if dest somehow already starts with sglang_
+            if not original_dest.startswith("sglang_"):
+                final_kwargs["dest"] = f"sglang_{original_dest}"
+        # If 'dest' is not explicitly provided (or is None/not a string),
+        # argparse will derive 'dest' from the (now prefixed) flag names.
+        # E.g., if the first flag is "--sglang-foo-bar", argparse sets dest to "sglang_foo_bar".
+
+        old_add_argument(*new_name_or_flags_list, **final_kwargs)
+
+    parser.add_argument = new_add_argument_wrapper
+    ServerArgs.add_cli_args(parser)
+    parser.add_argument = old_add_argument
+
+    return parser
+
+
+def validate_args(args):
+    if parse(sglang.__version__) == parse("0.4.10") and getattr(args, "sglang_enable_ep_moe", False):
+        args.sglang_expert_parallel_size = args.rollout_num_gpus_per_engine
+
+    args.sglang_tp_size = args.rollout_num_gpus_per_engine
+    args.sglang_dp_size = args.sglang_data_parallel_size
+    args.sglang_pp_size = args.sglang_pipeline_parallel_size
+    args.sglang_ep_size = args.sglang_expert_parallel_size
+
+    if args.sglang_dp_size > 1:
+        assert args.sglang_enable_dp_attention
+
+    if getattr(args, "sglang_router_ip", None):
+        args.sglang_router_ip = _wrap_ipv6(args.sglang_router_ip)

slime/backends/sglang_utils/sglang_engine.py

@@ -0,0 +1,491 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import dataclasses
+import logging
+import multiprocessing
+import time
+from urllib.parse import quote
+
+import requests
+import sglang_router
+from packaging.version import parse
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import kill_process_tree
+from urllib3.exceptions import NewConnectionError
+
+from slime.ray.ray_actor import RayActor
+from slime.utils.http_utils import get_host_info
+
+logger = logging.getLogger(__name__)
+
+
+def get_base_gpu_id(args, rank):
+    num_gpus = min(args.num_gpus_per_node, args.rollout_num_gpus_per_engine)
+    if args.colocate:
+        start_index = (rank * num_gpus) % args.num_gpus_per_node
+    else:
+        num_actor_gpus = 0 if args.debug_rollout_only else args.actor_num_gpus_per_node * args.actor_num_nodes
+        start_index = (num_actor_gpus + rank * num_gpus) % args.num_gpus_per_node
+        if args.use_critic:
+            num_critic_gpus = args.critic_num_gpus_per_node * args.critic_num_nodes
+            start_index = (num_actor_gpus + num_critic_gpus + rank * num_gpus) % args.num_gpus_per_node
+    return start_index
+
+
+def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process:
+    from sglang.srt.entrypoints.http_server import launch_server
+
+    multiprocessing.set_start_method("spawn", force=True)
+    server_args.host = server_args.host.strip("[]")
+    p = multiprocessing.Process(target=launch_server, args=(server_args,))
+    p.start()
+
+    if server_args.node_rank != 0:
+        return
+
+    _wait_server_healthy(
+        base_url=server_args.url(),
+        api_key=server_args.api_key,
+        is_process_alive=lambda: p.is_alive(),
+    )
+
+    return p
+
+
+def _wait_server_healthy(base_url, api_key, is_process_alive):
+    headers = {
+        "Content-Type": "application/json; charset=utf-8",
+        "Authorization": f"Bearer {api_key}",
+    }
+
+    with requests.Session() as session:
+        while True:
+            try:
+                response = session.get(f"{base_url}/health_generate", headers=headers)
+                if response.status_code == 200:
+                    break
+            except requests.RequestException:
+                pass
+
+            if not is_process_alive():
+                raise Exception("Server process terminated unexpectedly.")
+
+            time.sleep(2)
+
+        # use flush_cache to make sure the working queue is empty, so that we can do offload
+        while True:
+            try:
+                response = session.get(f"{base_url}/flush_cache", headers=headers)
+                if response.status_code == 200:
+                    break
+
+            except requests.RequestException:
+                pass
+
+            if not is_process_alive():
+                raise Exception("Server process terminated unexpectedly.")
+
+            time.sleep(2)
+
+
+class SGLangEngine(RayActor):
+    def __init__(self, args, rank: int, worker_type: str = "regular"):
+        self.args = args
+        self.rank = rank
+        self.worker_type = worker_type
+
+    def init(self, dist_init_addr, port, nccl_port, host=None, disaggregation_bootstrap_port=None):
+        self.router_ip = self.args.sglang_router_ip
+        self.router_port = self.args.sglang_router_port
+
+        host = host or get_host_info()[1]
+
+        # support ipv6 address
+        if ":" in host and not host.startswith("["):
+            host = f"[{host}]"
+
+        # dist_init_addr may be 2605:...:10163, should split port
+        *addr_parts, port_str = dist_init_addr.split(":")
+        ipv6_addr = ":".join(addr_parts)
+        if ":" in ipv6_addr and not ipv6_addr.startswith("["):
+            dist_init_addr = f"[{ipv6_addr}]:{port_str}"
+
+        server_args_dict, external_engine_need_check_fields = _compute_server_args(
+            self.args,
+            self.rank,
+            dist_init_addr,
+            nccl_port,
+            host,
+            port,
+            self.worker_type,
+            disaggregation_bootstrap_port,
+        )
+
+        self.node_rank = server_args_dict["node_rank"]
+        self.server_host = server_args_dict["host"]
+        self.server_port = server_args_dict["port"]
+
+        if self.args.rollout_external:
+            self._init_external(server_args_dict, external_engine_need_check_fields=external_engine_need_check_fields)
+        else:
+            self._init_normal(server_args_dict)
+
+    def _init_external(self, expect_server_args, external_engine_need_check_fields):
+        logger.info(f"Use external SGLang engine (rank={self.rank}, expect_server_args={expect_server_args})")
+
+        def _get_actual_server_args():
+            response = requests.get(f"http://{self.server_host}:{self.server_port}/get_server_info")
+            response.raise_for_status()
+            return response.json()
+
+        def _sanity_check_server_args(actual_server_args, expect_server_args):
+            for name in external_engine_need_check_fields:
+                expect_value = expect_server_args.get(name)
+                actual_value = actual_server_args.get(name)
+                assert (
+                    actual_value == expect_value
+                ), f"{name=} {expect_value=} {actual_value=} {expect_server_args=} {actual_server_args=}"
+
+        _wait_server_healthy(
+            base_url=f"http://{self.server_host}:{self.server_port}",
+            api_key=None,
+            is_process_alive=lambda: True,
+        )
+        actual_server_args = _get_actual_server_args()
+        _sanity_check_server_args(actual_server_args, expect_server_args)
+
+    def _init_normal(self, server_args_dict):
+        logger.info(f"Launch HttpServerEngineAdapter at: {self.server_host}:{self.server_port}")
+        self.process = launch_server_process(ServerArgs(**server_args_dict))
+
+        if self.node_rank == 0 and self.router_ip and self.router_port:
+            if parse(sglang_router.__version__) <= parse("0.2.1") or self.args.use_slime_router:
+                assert (
+                    self.worker_type == "regular"
+                ), "pd disaggregation is not supported in old router or slime router."
+                response = requests.post(
+                    f"http://{self.router_ip}:{self.router_port}/add_worker?url=http://{self.server_host}:{self.server_port}"
+                )
+            else:
+                payload = {
+                    "url": f"http://{self.server_host}:{self.server_port}",
+                    "worker_type": self.worker_type,
+                }
+                if self.worker_type == "prefill":
+                    payload["bootstrap_port"] = server_args_dict["disaggregation_bootstrap_port"]
+                response = requests.post(
+                    f"http://{self.router_ip}:{self.router_port}/workers",
+                    json=payload,
+                )
+            response.raise_for_status()
+
+    def _make_request(self, endpoint: str, payload: dict | None = None):
+        """Make a POST request to the specified endpoint with the given payload.
+
+        Args:
+            endpoint: The API endpoint to call
+            payload: The JSON payload to send (default: empty dict)
+
+        Returns:
+            The JSON response from the server
+        """
+        if self.node_rank != 0:
+            return
+
+        url = f"http://{self.server_host}:{self.server_port}/{endpoint}"
+        response = requests.post(url, json=payload or {})
+        try:
+            response.raise_for_status()
+        except requests.exceptions.HTTPError as e:
+            e.add_note(f"{response.text=}")
+            raise
+        return response.json()
+
+    def health_generate(self, timeout: float = 5.0) -> bool:
+        """Run /health_generate on the underlying SGLang HTTP server.
+
+        Args:
+            timeout: Timeout for the health request in seconds.
+
+        Returns:
+            True if the server responds with HTTP 200.
+
+        Raises:
+            requests.RequestException: If the request fails for any reason, including timeout.
+        """
+        if self.node_rank != 0:
+            return True
+
+        response = requests.get(
+            f"http://{self.server_host}:{self.server_port}/health_generate",
+            timeout=timeout,
+        )
+        response.raise_for_status()
+        return True
+
+    def update_weights_from_tensor(
+        self,
+        serialized_named_tensors: list[str],
+        load_format: str | None = None,
+        flush_cache: bool = False,
+        weight_version: str | None = None,
+    ):
+        """
+        Update model weights from tensor data. The HTTP server will only post meta data, and the real weights will be copied directly from GPUs.
+
+        Note: The model should be on GPUs rather than CPU for this functionality to work properly.
+        If you encounter issues, ensure your model is loaded on GPU devices rather than CPU.
+        """
+        payload = {
+            "serialized_named_tensors": serialized_named_tensors,
+            "load_format": load_format,
+            "flush_cache": flush_cache,
+        }
+        if weight_version is not None:
+            payload["weight_version"] = weight_version
+        return self._make_request(
+            "update_weights_from_tensor",
+            payload,
+        )
+
+    def flush_cache(self):
+        """Flush the cache of the server."""
+        if self.node_rank != 0:
+            return
+        # flush cache will not return status_code 200 when there are pending requests
+        for _ in range(60):
+            try:
+                response = requests.get(f"http://{self.server_host}:{self.server_port}/flush_cache")
+                if response.status_code == 200:
+                    break
+            except NewConnectionError as e:
+                raise e
+            except Exception as e:
+                logger.info(f"Error flushing cache: {e}")
+                time.sleep(1)
+                continue
+        else:
+            raise TimeoutError("Timeout while flushing cache.")
+
+    def shutdown(self):
+        if self.args.rollout_external:
+            return
+
+        logger.info(f"Shutdown engine {self.server_host}:{self.server_port}...")
+        if self.node_rank == 0:
+            worker_url = f"http://{self.server_host}:{self.server_port}"
+            response = None
+            if parse(sglang_router.__version__) <= parse("0.2.1") or self.args.use_slime_router:
+                response = requests.post(
+                    f"http://{self.router_ip}:{self.router_port}/remove_worker?url=http://{self.server_host}:{self.server_port}"
+                )
+            elif parse(sglang_router.__version__) < parse("0.3.0"):
+                worker_url = quote(worker_url, safe="")
+                response = requests.delete(f"http://{self.router_ip}:{self.router_port}/workers/{worker_url}")
+            else:
+                try:
+                    all_workers = requests.get(f"http://{self.router_ip}:{self.router_port}/workers").json()["workers"]
+                    for worker in all_workers:
+                        if worker["url"] == worker_url:
+                            worker_id = worker["id"]
+                            response = requests.delete(
+                                f"http://{self.router_ip}:{self.router_port}/workers/{worker_id}"
+                            )
+                            break
+                    else:
+                        logger.warning(f"Worker {worker_url} not found in router during shutdown.")
+                except Exception as e:
+                    logger.warning(f"Failed to fetch workers list or remove worker: {e}")
+
+            if response is not None:
+                response.raise_for_status()
+        kill_process_tree(self.process.pid)
+
+    def get_weight_version(self):
+        if self.node_rank != 0:
+            return
+        url = f"http://{self.server_host}:{self.server_port}/get_weight_version"
+        response = requests.get(url)
+        response.raise_for_status()
+        return response.json()["weight_version"]
+
+    def release_memory_occupation(self):
+        self.flush_cache()
+        return self._make_request("release_memory_occupation")
+
+    def resume_memory_occupation(self, tags: list[str] = None):
+        """
+        Available tags for multi-stage resume: weights, kv_cache
+        """
+        return self._make_request(
+            "resume_memory_occupation",
+            {"tags": tags},
+        )
+
+    def check_weights(self, action: str):
+        return self._make_request("weights_checker", {"action": action})
+
+    def init_weights_update_group(self, master_address, master_port, rank_offset, world_size, group_name, backend):
+        return self._make_request(
+            "init_weights_update_group",
+            {
+                "master_address": master_address,
+                "master_port": master_port,
+                "rank_offset": rank_offset,
+                "world_size": world_size,
+                "group_name": group_name,
+                "backend": backend,
+            },
+        )
+
+    def destroy_weights_update_group(self, group_name):
+        try:
+            return self._make_request(
+                "destroy_weights_update_group",
+                {
+                    "group_name": group_name,
+                },
+            )
+        except requests.exceptions.RequestException:
+            # catch the case there the engine is just created and does not have the group.
+            pass
+
+    def update_weights_from_distributed(
+        self, names, dtypes, shapes, group_name, flush_cache=False, weight_version: str | None = None
+    ):
+        payload = {
+            "names": names,
+            "dtypes": [str(dtype).replace("torch.", "") for dtype in dtypes],
+            "shapes": shapes,
+            "group_name": group_name,
+            "flush_cache": flush_cache,
+        }
+        if weight_version is not None:
+            payload["weight_version"] = weight_version
+        return self._make_request(
+            "update_weights_from_distributed",
+            payload,
+        )
+
+    def pause_generation(self):
+        response = requests.post(f"http://{self.server_host}:{self.server_port}/pause_generation", json={})
+        response.raise_for_status()
+        return response
+
+    def continue_generation(self):
+        response = requests.post(f"http://{self.server_host}:{self.server_port}/continue_generation", json={})
+        response.raise_for_status()
+        return response
+
+    def start_profile(
+        self,
+        # The output directory
+        output_dir: str | None = None,
+        # If set, it profile as many as this number of steps.
+        # If it is set, profiling is automatically stopped after this step, and
+        # the caller doesn't need to run stop_profile.
+        start_step: int | None = None,
+        num_steps: int | None = None,
+        activities: list[str] | None = None,
+        profile_by_stage: bool = False,
+        with_stack: bool | None = None,
+        record_shapes: bool | None = None,
+    ):
+        response = requests.post(
+            f"http://{self.server_host}:{self.server_port}/start_profile",
+            json={
+                "output_dir": output_dir,
+                "start_step": start_step,
+                "num_steps": num_steps,
+                "activities": activities,
+                "profile_by_stage": profile_by_stage,
+                "with_stack": with_stack,
+                "record_shapes": record_shapes,
+            },
+        )
+        response.raise_for_status()
+        return response
+
+    def stop_profile(self):
+        response = requests.post(f"http://{self.server_host}:{self.server_port}/stop_profile", json={})
+        response.raise_for_status()
+        return response
+
+
+def _compute_server_args(
+    args,
+    rank,
+    dist_init_addr,
+    nccl_port,
+    host,
+    port,
+    worker_type: str = "regular",
+    disaggregation_bootstrap_port: int | None = None,
+):
+    nnodes = max(1, args.rollout_num_gpus_per_engine // args.num_gpus_per_node)
+    node_rank = rank % nnodes
+    kwargs = {
+        "model_path": args.hf_checkpoint,
+        "trust_remote_code": True,
+        "random_seed": args.seed + rank,
+        # memory
+        "enable_memory_saver": args.offload_rollout,
+        # distributed
+        "host": host,
+        "port": port,
+        "nccl_port": nccl_port,
+        "nnodes": nnodes,
+        "node_rank": node_rank,
+        "dist_init_addr": dist_init_addr,
+        "gpu_id_step": 1,
+        "base_gpu_id": get_base_gpu_id(args, rank),
+        # parallel
+        "tp_size": args.rollout_num_gpus_per_engine,
+        "dp_size": args.sglang_dp_size,
+        "pp_size": args.sglang_pp_size,
+        "ep_size": args.sglang_ep_size,
+        # always skip warmup to prevent warmup timeout.
+        "skip_server_warmup": True,
+    }
+
+    if worker_type == "prefill":
+        kwargs["disaggregation_mode"] = "prefill"
+        kwargs["load_balance_method"] = "round_robin"
+        assert (
+            disaggregation_bootstrap_port is not None
+        ), "disaggregation_bootstrap_port must be set for prefill worker"
+        kwargs["disaggregation_bootstrap_port"] = disaggregation_bootstrap_port
+    elif worker_type == "decode":
+        kwargs["disaggregation_mode"] = "decode"
+        kwargs["prefill_round_robin_balance"] = True
+
+    if args.use_rollout_routing_replay:
+        kwargs["enable_return_routed_experts"] = True
+    if args.fp16:
+        kwargs["dtype"] = "float16"
+    external_engine_need_check_fields = [k for k in kwargs.keys() if k not in _EXTERNAL_ENGINE_SKIP_CHECK_FIELDS]
+
+    unused_keys = set(kwargs.keys())
+    for attr in dataclasses.fields(ServerArgs):
+        if hasattr(args, f"sglang_{attr.name}") and attr.name not in kwargs:
+            kwargs[attr.name] = getattr(args, f"sglang_{attr.name}")
+        unused_keys.discard(attr.name)
+
+    # for compatibility with old args
+    if len(unused_keys) > 0:
+        logger.info(f"Warning: The following arguments is not supported in the current sglang: {unused_keys}.")
+        for key in unused_keys:
+            kwargs.pop(key)
+
+    return kwargs, external_engine_need_check_fields
+
+
+_EXTERNAL_ENGINE_SKIP_CHECK_FIELDS = [
+    "model_path",
+    "trust_remote_code",
+    "random_seed",
+    "nccl_port",
+    "dist_init_addr",
+    "skip_server_warmup",
+]