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

Model: ayh015/myLightningOPD
Source: Original Platform

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)