diff --git a/slime/__init__.py b/slime/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a4d636226b9a4b5130dfd8e94f2cece5fd290e53 --- /dev/null +++ b/slime/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + diff --git a/slime/__pycache__/__init__.cpython-312.pyc b/slime/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c59d5368a89e85fea98fb14db737ec9f070224f Binary files /dev/null and b/slime/__pycache__/__init__.cpython-312.pyc differ diff --git a/slime/backends/__init__.py b/slime/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0944b65fdbc1b6db59071d12271239c029c212a2 --- /dev/null +++ b/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 + + diff --git a/slime/backends/__pycache__/__init__.cpython-312.pyc b/slime/backends/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..88d39e5afda24334a786a128b4525931ab0c3eda Binary files /dev/null and b/slime/backends/__pycache__/__init__.cpython-312.pyc differ diff --git a/slime/backends/fsdp_utils/__init__.py b/slime/backends/fsdp_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cc577d2fa54903b8db12a85c4f9b8489e93f8085 --- /dev/null +++ b/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) diff --git a/slime/backends/fsdp_utils/actor.py b/slime/backends/fsdp_utils/actor.py new file mode 100644 index 0000000000000000000000000000000000000000..09700ca0578b84354436900c3e70b1119514f30d --- /dev/null +++ b/slime/backends/fsdp_utils/actor.py @@ -0,0 +1,1145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +import os +import random +from argparse import Namespace +from itertools import accumulate + +import ray +import torch +import torch.distributed as dist +import torch.nn.functional as F +from ring_flash_attn import substitute_hf_flash_attn, update_ring_flash_attn_params +from tqdm import tqdm +from transformers import AutoConfig + +from slime.ray.train_actor import TrainRayActor +from slime.utils import train_dump_utils, train_metric_utils +from slime.utils.context_utils import with_defer +from slime.utils.data import get_minimum_num_micro_batch_size, process_rollout_data +from slime.utils.distributed_utils import get_gloo_group +from slime.utils.memory_utils import clear_memory, print_memory +from slime.utils.metric_utils import compute_rollout_step +from slime.utils.misc import load_function +from slime.utils.ppo_utils import ( + compute_approx_kl, + compute_gspo_kl, + compute_opsm_mask, + compute_policy_loss, + vanilla_tis_function, +) +from slime.utils.processing_utils import load_processor, load_tokenizer +from slime.utils.ray_utils import Box +from slime.utils.timer import Timer, inverse_timer, timer +from slime.utils.tracking_utils import init_tracking + +from ...utils import tracking_utils +from ...utils.profile_utils import TrainProfiler +from . import checkpoint +from .data_packing import pack_sequences, pad_packed_sequence_with_cp, unpack_sequences +from .lr_scheduler import get_lr_scheduler +from .update_weight_utils import UpdateWeightFromDistributed, UpdateWeightFromTensor + +logger = logging.getLogger(__name__) + + +class FSDPTrainRayActor(TrainRayActor): + """Simplified TrainRayActor for pure HF+FSDP training. + + Responsibilities: + * Initialize model/tokenizer on rank0 sequentially to avoid race on cache + * Wrap model with FSDP + * Provide minimal train / save / update_weights hooks compatible with existing RayTrainGroup + + Weight update strategy: + * Rank0 gathers state_dict (full) and broadcasts tensor-by-tensor. + * For small models this is fine; for larger models consider sharded state_dict type. + """ + + @with_defer(lambda: Timer().start("train_wait")) + def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # type: ignore[override] + super().init(args, role, with_ref) + + # Setup device mesh for parallelism (handles both CP and non-CP cases) + self._setup_device_mesh() + torch.manual_seed(args.seed) + + self.train_parallel_config = { + "dp_size": self.dp_size, + } + + if self.args.debug_rollout_only: + return 0 + + self.fsdp_cpu_offload = getattr(self.args, "fsdp_cpu_offload", False) + # Offload train and fsdp cpu offload cannot be used together, fsdp_cpu_offload is more aggressive + if self.args.offload_train and self.fsdp_cpu_offload: + self.args.offload_train = False + + self._enable_true_on_policy_optimizations(args) + if dist.get_rank() == 0: + init_tracking(args, primary=False) + + if getattr(self.args, "start_rollout_id", None) is None: + self.args.start_rollout_id = 0 + + self.prof = TrainProfiler(args) + + for i in range(dist.get_world_size()): + if i == dist.get_rank(): + self.hf_config = AutoConfig.from_pretrained(self.args.hf_checkpoint, trust_remote_code=True) + self.tokenizer = load_tokenizer(self.args.hf_checkpoint, trust_remote_code=True) + # Vision models have `vision_config` in the config + if hasattr(self.hf_config, "vision_config"): + self.processor = load_processor(self.args.hf_checkpoint, trust_remote_code=True) + dist.barrier(group=get_gloo_group()) + + init_context = self._get_init_weight_context_manager() + + with init_context(): + model = self.get_model_cls().from_pretrained( + self.args.hf_checkpoint, + trust_remote_code=True, + attn_implementation=self.args.attn_implementation, + ) + + model.train() + + full_state = model.state_dict() + + model = apply_fsdp2(model, mesh=self.dp_mesh, cpu_offload=self.fsdp_cpu_offload, args=self.args) + + model = self._fsdp2_load_full_state_dict( + model, full_state, self.dp_mesh, cpu_offload=True if self.fsdp_cpu_offload else None + ) + + self.model = model + + if args.gradient_checkpointing: + self.model.gradient_checkpointing_enable() + + if args.optimizer == "adam": + self.optimizer = torch.optim.AdamW( + self.model.parameters(), + lr=args.lr, + betas=(args.adam_beta1, args.adam_beta2), + eps=args.adam_eps, + weight_decay=args.weight_decay, + ) + else: + raise ValueError(f"Unsupported optimizer: {args.optimizer}. Supported options: 'adam'") + + # Initialize LR scheduler + self.lr_scheduler = get_lr_scheduler(args, self.optimizer) + + self.global_step = 0 + self.micro_step = 0 + + checkpoint_payload = checkpoint.load(self) + + # Create separate ref model if needed (kept in CPU until needed) + self.ref_model = None + if with_ref: + self.ref_model = self._create_ref_model(args.ref_load) + + self.weight_updater = ( + UpdateWeightFromTensor(self.args, self.model) + if self.args.colocate + else UpdateWeightFromDistributed(self.args, self.model) + ) + + checkpoint.finalize_load(self, checkpoint_payload) + + # Initialize data packing parameters + self.max_tokens_per_gpu = args.max_tokens_per_gpu # From main arguments + + if self.args.offload_train: + self.sleep() + + self.prof.on_init_end() + + return int(getattr(self.args, "start_rollout_id", 0)) + + def get_model_cls(self): + # Vision models have `vision_config` in the config + if hasattr(self.hf_config, "vision_config"): + from transformers import AutoModelForImageTextToText + + return AutoModelForImageTextToText + else: + from transformers import AutoModelForCausalLM + + return AutoModelForCausalLM + + def _enable_true_on_policy_optimizations(self, args): + if args.true_on_policy_mode: + from sglang.srt.batch_invariant_ops import enable_batch_invariant_mode + + from .models.qwen3_moe import apply_true_on_policy_patch_for_qwen3_moe + + logger.info("FSDPTrainRayActor call enable_batch_invariant_mode for true-on-policy") + enable_batch_invariant_mode( + # In Qwen3, rope `inv_freq_expanded.float() @ position_ids_expanded.float()` uses bmm + # and disabling it will make it aligned + enable_bmm=False, + ) + + apply_true_on_policy_patch_for_qwen3_moe() + else: + from .models.qwen3_moe_hf import apply_fsdp_moe_patch + + apply_fsdp_moe_patch() + + def _setup_device_mesh(self) -> None: + """Setup device mesh for parallelism (always called, handles both CP and non-CP cases). + + Creates 2D mesh (dp_size, cp_size) for all cases: + - When context_parallel_size > 1: hybrid CP + DP + - When context_parallel_size = 1: pure DP (equivalent to 1D mesh) + + This ensures consistent group management across all parallelism modes. + """ + from torch.distributed.device_mesh import init_device_mesh + + world_size = dist.get_world_size() + rank = dist.get_rank() + + # Use context_parallel_size directly (defaults to 1 for pure DP) + self.cp_size = self.args.context_parallel_size + self.dp_size = world_size // self.cp_size + + # Create 2D device mesh: (dp_size, cp_size) + # Ranks laid out in row-major: mesh[dp_idx, cp_idx] = dp_idx * cp_size + cp_idx + # - CP groups: consecutive ranks along dim 1, e.g., [0,1], [2,3], [4,5], [6,7] + # - DP groups: striped ranks along dim 0, e.g., [0,2,4,6], [1,3,5,7] + # When cp_size=1, this degenerates to pure DP + self.mesh = init_device_mesh("cuda", mesh_shape=(self.dp_size, self.cp_size), mesh_dim_names=("dp", "cp")) + + # Extract process groups from mesh + self.dp_group = self.mesh.get_group("dp") # For FSDP gradient sync, metric reduction + self.cp_group = self.mesh.get_group("cp") # For Ring Flash Attention, logit gathering + self.dp_mesh = self.mesh["dp"] # For FSDP + + # Compute local ranks within each dimension + self.dp_rank = rank // self.cp_size + self.cp_rank = rank % self.cp_size + + logger.info( + f"[Rank {rank}] Device mesh (2D): world_size={world_size}, " + f"cp_size={self.cp_size}, dp_size={self.dp_size}" + ) + logger.info(f"[Rank {rank}] Mesh shape: {self.mesh.shape}, " f"dp_rank={self.dp_rank}, cp_rank={self.cp_rank}") + + # Setup Ring Flash Attention with CP group from mesh (only when cp_size > 1) + if self.cp_size > 1: + substitute_hf_flash_attn(self.cp_group, heads_k_stride=1) + logger.info(f"[Rank {rank}] CP initialized via device mesh") + else: + logger.info(f"[Rank {rank}] Pure DP mode (cp_size=1)") + + def _get_init_weight_context_manager(self): + """Get context manager for model initialization. + + Returns a callable that creates a context manager. + Uses meta device (no memory allocation) for non-rank-0 processes, + UNLESS tie_word_embeddings=True (which causes hangs with meta tensors). + + Ref: verl/utils/fsdp_utils.py::get_init_weight_context_manager + NOTE: tie_word_embedding causes meta_tensor init to hang + """ + from accelerate import init_empty_weights + + # Check if model uses tied word embeddings (which doesn't work with meta tensors) + use_meta_tensor = not self.hf_config.tie_word_embeddings + + def cpu_init_weights(): + return torch.device("cpu") + + if use_meta_tensor: + # Rank 0: CPU, others: meta device (memory efficient for large models) + return init_empty_weights if dist.get_rank() != 0 else cpu_init_weights + else: + logger.info(f"[Rank {dist.get_rank()}] tie_word_embeddings=True, loading full model to CPU on all ranks") + return cpu_init_weights + + def _fsdp2_load_full_state_dict(self, model, full_state, device_mesh, cpu_offload): + """Load full state dict into FSDP2 model with efficient broadcast from rank 0. + + This function loads weights from rank 0 and broadcasts to all other ranks, + avoiding the need for each rank to load the full model from disk. + + Args: + model: FSDP2-wrapped model + full_state: State dict (only rank 0 has real weights, others have empty dict) + device_mesh: Device mesh for FSDP + cpu_offload: If not None, enables StateDictOptions cpu_offload + + Ref:verl/utils/fsdp_utils.py::fsdp2_load_full_state_dict + """ + from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict + + # Rank 0: move with weights, others: allocate empty tensors on device + if dist.get_rank() == 0: + model = model.to(device=torch.cuda.current_device(), non_blocking=True) + else: + # to_empty creates tensors on device without initializing memory + model = model.to_empty(device=torch.cuda.current_device()) + + is_cpu_offload = cpu_offload is not None + options = StateDictOptions(full_state_dict=True, cpu_offload=is_cpu_offload, broadcast_from_rank0=True) + + set_model_state_dict(model, full_state, options=options) + + # set_model_state_dict will not broadcast buffers, so we need to broadcast them manually. + for _name, buf in model.named_buffers(): + dist.broadcast(buf, src=0) + + if is_cpu_offload: + model.to("cpu", non_blocking=True) + for buf in model.buffers(): + buf.data = buf.data.to(torch.cuda.current_device()) + + return model + + @timer + def sleep(self) -> None: + """Pause CUDA memory for all tracked tensors.""" + if not self.args.offload_train: + return + + print_memory("before offload model") + + self.model.cpu() + move_torch_optimizer(self.optimizer, "cpu") + clear_memory() + dist.barrier(group=get_gloo_group()) + print_memory("after offload model") + + @timer + def wake_up(self) -> None: + """Resume CUDA memory for all tracked tensors.""" + if not self.args.offload_train: + return + + self.model.cuda() + move_torch_optimizer(self.optimizer, "cuda") + dist.barrier(group=get_gloo_group()) + print_memory("after wake_up model") + + def save_model(self, rollout_id: int, force_sync: bool = False) -> None: + """Delegate checkpoint saving to the shared checkpoint utilities.""" + if self.args.debug_rollout_only or self.args.save is None: + return + + assert not self.args.async_save, "FSDPTrainRayActor does not support async_save yet." + checkpoint.save(self, rollout_id) + + def _compute_log_prob( + self, + model_tag: str, + packed_batches: list[dict[str, torch.Tensor]], + store_prefix: str = "", + ) -> dict[str, list[torch.Tensor]]: + """Compute token log-probabilities for a list of packed batches. + + Parameters: + model_tag: Which parameters to use, e.g. "actor" or "ref". + packed_batches: A list of packed batch dictionaries produced by + `pack_sequences`, each containing at least `tokens` and + `position_ids`; may also include multimodal keys like `pixel_values`. + store_prefix: Prefix to use for keys in outputs (e.g., "ref_"). + + Returns: + A lightweight dictionary keyed by f"{store_prefix}log_probs". The + actual per-sequence results are written in-place into each element of + `packed_batches` under the same key and can be read back by callers. + + Note: + Uses separate ref model when model_tag == "ref". The ref model is + loaded from CPU to GPU on-demand and offloaded back after use. + """ + # Select which model to use + if model_tag == "ref" and self.ref_model is not None: + if not self.fsdp_cpu_offload: + self.model.cpu() + torch.cuda.empty_cache() + dist.barrier(group=get_gloo_group()) + + active_model = self.ref_model + active_model.eval() + else: + active_model = self.model + + try: + rollout_data = {f"{store_prefix}log_probs": []} + with timer(f"{store_prefix}log_probs"), torch.no_grad(): + for batch in self.prof.iterate_train_log_probs( + tqdm(packed_batches, desc=f"{store_prefix}log_probs", disable=dist.get_rank() != 0) + ): + model_args = self._get_model_inputs_args(batch) + logits = active_model(**model_args).logits.squeeze(0).float() + log_probs_result, entropy_result = get_logprob_and_entropy_with_cp( + logits=logits, + target_tokens=batch["tokens"], + cp_rank=self.cp_rank, + cp_size=self.cp_size, + cp_group=self.cp_group, + model_input_ids=model_args["input_ids"], + allow_compile=not self.args.true_on_policy_mode, + temperature=self.args.rollout_temperature, + ) + batch[f"{store_prefix}log_probs"] = log_probs_result + if store_prefix == "": + batch["entropy"] = entropy_result + return rollout_data + + finally: + # Restore actor model if it was offloaded + if model_tag == "ref" and self.ref_model is not None: + torch.cuda.empty_cache() + dist.barrier(group=get_gloo_group()) + + if not self.fsdp_cpu_offload: + self.model.cuda() + dist.barrier(group=get_gloo_group()) + + def _packed_data( + self, rollout_data: dict[str, list[torch.Tensor]] + ) -> tuple[list[dict[str, torch.Tensor]], list[int]]: + """Pack variable-length sequences for efficient processing. + + Parameters: + rollout_data: Dictionary of lists containing sequence-level tensors + such as `tokens`, `loss_masks`, `rewards`, `response_lengths`, + `advantages`, `returns`, and optional `rollout_log_probs`. + + Returns: + A pair `(packed_batches, grad_accum)` where `packed_batches` is a list + of packed batch dictionaries and `grad_accum` lists the micro-batch + indices at which to perform optimizer steps. + """ + # Pack sequences efficiently + tokens = rollout_data["tokens"] + + packed_batches = [] + mbs_size_list = [] + local_batch_size = self.args.global_batch_size // self.dp_size + assert ( + self.args.global_batch_size % self.dp_size == 0 + ), f"global_batch_size {self.args.global_batch_size} is not divisible by dp_world_size {self.dp_size}" + # Use global_batch_size for splitting when max_tokens_per_gpu is enabled + if self.args.use_dynamic_batch_size: + # In CP mode, CP group shares sequences, so total capacity is max_tokens_per_gpu * cp_size + max_tokens = self.args.max_tokens_per_gpu + if self.cp_size > 1: + max_tokens = max_tokens * self.cp_size + + for i in range(0, len(tokens), local_batch_size): + mbs_size_list.append( + get_minimum_num_micro_batch_size( + [len(t) for t in rollout_data["tokens"][i : i + local_batch_size]], + max_tokens, + ) + ) + num_microbatches = torch.tensor(mbs_size_list, dtype=torch.int, device=torch.cuda.current_device()) + dist.all_reduce(num_microbatches, op=dist.ReduceOp.MAX, group=self.dp_group) + num_microbatches = num_microbatches.tolist() + else: + num_microbatches = [self.args.global_batch_size // (self.args.micro_batch_size * self.dp_size)] * ( + len(tokens) // local_batch_size + ) + + start = 0 + for mbs_size in num_microbatches: + end = start + local_batch_size + packed_batches.extend( + pack_sequences( + rollout_data["tokens"][start:end], + rollout_data["loss_masks"][start:end], + rollout_data["rewards"][start:end], + rollout_data["raw_reward"][start:end], + rollout_data["response_lengths"][start:end], + rollout_data["advantages"][start:end], + rollout_data["returns"][start:end], + rollout_log_probs=( + rollout_data["rollout_log_probs"][start:end] if "rollout_log_probs" in rollout_data else None + ), + multimodal_train_inputs=( + rollout_data["multimodal_train_inputs"][start:end] + if "multimodal_train_inputs" in rollout_data + else None + ), + num_packs=mbs_size, + ) + ) + start = end + grad_accum = list(accumulate(num_microbatches)) + + return packed_batches, grad_accum + + def train(self, rollout_id: int, rollout_data_ref: Box) -> None: + """Run one training update over a rollout batch. + + Parameters: + rollout_id: Monotonic id for logging. + rollout_data_ref: A Box handle wrapping a Ray object reference to a + dictionary with rollout tensors and metadata (e.g., `tokens`, + `loss_masks`, `rewards`, `response_lengths`, optional + `rollout_log_probs`, etc.). It will be fetched and partitioned + by `process_rollout_data` based on data-parallel rank/size. + """ + if self.args.offload_train: + self.wake_up() + + with inverse_timer("train_wait"), timer("train"): + rollout_data = process_rollout_data(self.args, rollout_data_ref, self.dp_rank, self.dp_size) + if self.args.debug_rollout_only: + return + self._train_core(rollout_id=rollout_id, rollout_data=rollout_data) + + train_metric_utils.log_perf_data_raw( + rollout_id=rollout_id, + args=self.args, + is_primary_rank=dist.get_rank() == 0, + compute_total_fwd_flops=None, + ) + + def _log_rollout_data(self, rollout_id: int, rollout_data, packed_batches): + log_dict = {} + if "raw_reward" in rollout_data and dist.get_rank() == 0: + raw_reward_list = rollout_data["raw_reward"] + if raw_reward_list: + log_dict["rollout/raw_reward"] = sum(raw_reward_list) / len(raw_reward_list) + + for metric_key in ["log_probs", "rollout_log_probs", "ref_log_probs", "advantages", "returns"]: + if metric_key not in packed_batches[0]: + continue + val = torch.tensor([0.0], device=torch.cuda.current_device()) + for _mbs_id, batches in enumerate(packed_batches): + unpacked_batches = unpack_sequences(batches) + for unpacked_batch in unpacked_batches: + if isinstance(unpacked_batch[metric_key], torch.Tensor): + loss_masks_tensor = unpacked_batch["loss_masks"].to(device=torch.cuda.current_device()) + metric_tensor = unpacked_batch[metric_key].to(device=torch.cuda.current_device()) + val += (metric_tensor * loss_masks_tensor).sum() / loss_masks_tensor.sum().clamp_min(1) + else: + val += unpacked_batch[metric_key] + dist.all_reduce(val, op=dist.ReduceOp.SUM, group=self.dp_group) + log_dict[f"rollout/{metric_key}"] = ( + val / (self.args.n_samples_per_prompt * self.args.rollout_batch_size) + ).item() + if dist.get_rank() == 0: + logger.info(f"rollout {rollout_id}: {log_dict}") + log_dict["rollout/step"] = compute_rollout_step(self.args, rollout_id) + tracking_utils.log(self.args, log_dict, step_key="rollout/step") + + if self.args.ci_test and self.args.true_on_policy_mode: + assert log_dict["rollout/log_probs"] == log_dict["rollout/rollout_log_probs"], ( + f"CI check failed: true_on_policy_mode is enabled, but log_probs " + f"({log_dict['rollout/log_probs']}) != rollout_log_probs " + f"({log_dict['rollout/rollout_log_probs']})" + ) + + def _train_core(self, rollout_id: int, rollout_data) -> None: + if self.args.advantage_estimator in ["grpo", "gspo"]: + rollout_data["advantages"] = rollout_data["returns"] = [ + torch.tensor([rollout_data["rewards"][i]] * rollout_data["response_lengths"][i]) + for i in range(len(rollout_data["rewards"])) + ] + else: + raise NotImplementedError(f"Unsupported advantage_estimator {self.args.advantage_estimator}") + + packed_batches, grad_accum = self._packed_data(rollout_data) + + assert ( + len(grad_accum) > 0 + ), f"Invalid grad_accum {grad_accum} for micro_batch_size {self.args.micro_batch_size} and global_batch_size {self.args.global_batch_size}" + + if self.ref_model is not None: + self._compute_log_prob("ref", packed_batches, store_prefix="ref_") + + self._compute_log_prob("actor", packed_batches) + self._log_rollout_data(rollout_id, rollout_data, packed_batches) + + with timer("actor_train"): + reported_accum: dict[str, list[torch.Tensor]] = {} + self.optimizer.zero_grad(set_to_none=True) + for mbs_id, packed_batch in self.prof.iterate_train_actor( + enumerate(tqdm(packed_batches, desc="actor_train", disable=dist.get_rank() != 0)) + ): + self._train_step( + packed_batch=packed_batch, + reported_accum=reported_accum, + mbs_id=mbs_id, + grad_accum=grad_accum, + ) + + self.prof.step(rollout_id=rollout_id) + + train_dump_utils.save_debug_train_data(self.args, rollout_id=rollout_id, rollout_data=rollout_data) + + # Update ref model if needed (copy actor weights to ref) + if ( + self.args.ref_update_interval is not None + and (rollout_id + 1) % self.args.ref_update_interval == 0 + and self.ref_model is not None + ): + if dist.get_rank() == 0: + logger.info(f"Updating ref model at rollout_id {rollout_id}") + # Copy actor model state to ref model + actor_state = self.model.state_dict() + self.ref_model.load_state_dict(actor_state) + self.ref_model.cpu() + + def _train_step(self, packed_batch, reported_accum, mbs_id, grad_accum): + # Prepare model inputs + model_args = self._get_model_inputs_args(packed_batch) + logits = self.model(**model_args).logits.squeeze(0).float() + + # Compute log probs and entropy (unified for both CP and non-CP modes) + log_probs, entropy_result = get_logprob_and_entropy_with_cp( + logits=logits, + target_tokens=packed_batch["tokens"], + cp_rank=self.cp_rank, + cp_size=self.cp_size, + cp_group=self.cp_group, + model_input_ids=model_args["input_ids"], + allow_compile=not self.args.true_on_policy_mode, + temperature=self.args.rollout_temperature, + ) + packed_batch["cur_log_probs"] = log_probs + packed_batch["entropy"] = entropy_result + + unpacked_batches = unpack_sequences(packed_batch) + + old_log_prob_key = "rollout_log_probs" if self.args.use_rollout_logprobs else "log_probs" + missing_old_log_probs = [ + idx + for idx, batch in enumerate(unpacked_batches) + if old_log_prob_key not in batch or not isinstance(batch[old_log_prob_key], torch.Tensor) + ] + if missing_old_log_probs: + raise KeyError( + f"{old_log_prob_key} must be provided as torch.Tensor for all microbatches when " + f"use_rollout_logprobs is set to {self.args.use_rollout_logprobs}. Missing in batches: {missing_old_log_probs}" + ) + old_log_probs = torch.cat([batch[old_log_prob_key] for batch in unpacked_batches], dim=0) + log_probs = torch.cat([batch["cur_log_probs"] for batch in unpacked_batches], dim=0) + advantages = torch.cat([batch["advantages"] for batch in unpacked_batches], dim=0) + loss_masks = [batch["loss_masks"].to(device=log_probs.device) for batch in unpacked_batches] + response_lengths = [batch["response_lengths"] for batch in unpacked_batches] + + advantages = advantages.to(device=log_probs.device) + old_log_probs = old_log_probs.to(device=log_probs.device) + ppo_kl = old_log_probs - log_probs + + if self.args.use_opsm: + opsm_mask, opsm_clipfrac = compute_opsm_mask( + args=self.args, + full_log_probs=[batch["cur_log_probs"] for batch in unpacked_batches], + full_old_log_probs=[batch[old_log_prob_key] for batch in unpacked_batches], + advantages=[batch["advantages"] for batch in unpacked_batches], + loss_masks=loss_masks, + ) + + if self.args.advantage_estimator == "gspo": + ppo_kl = compute_gspo_kl( + full_log_probs=[batch["cur_log_probs"] for batch in unpacked_batches], + full_old_log_probs=[batch[old_log_prob_key] for batch in unpacked_batches], + local_log_probs=[batch["cur_log_probs"] for batch in unpacked_batches], + loss_masks=loss_masks, + ) + + pg_loss, pg_clipfrac = compute_policy_loss(ppo_kl, advantages, self.args.eps_clip, self.args.eps_clip_high) + + if self.args.use_opsm: + pg_loss = pg_loss * opsm_mask + + def _has_rollout_log_probs(batch) -> bool: + rollout_tensor = batch.get("rollout_log_probs") + return isinstance(rollout_tensor, torch.Tensor) and rollout_tensor.numel() > 0 + + has_rollout_log_probs = all(_has_rollout_log_probs(batch) for batch in unpacked_batches) + rollout_log_probs = ( + torch.cat([batch["rollout_log_probs"] for batch in unpacked_batches], dim=0) + if has_rollout_log_probs + else None + ) + + # Apply off-policy correction using importance sampling if enabled + if self.args.use_tis: + assert ( + has_rollout_log_probs and rollout_log_probs is not None + ), "rollout_log_probs must be provided as non-empty torch.Tensor for TIS/MIS" + + train_log_probs_list = list(log_probs.split(response_lengths, dim=0)) + rollout_log_probs_list = list(rollout_log_probs.split(response_lengths, dim=0)) + ois = (-ppo_kl).exp() + tis_kwargs = { + "args": self.args, + "pg_loss": pg_loss, + "train_log_probs": train_log_probs_list, + "rollout_log_probs": rollout_log_probs_list, + "loss_masks": loss_masks, + "response_lengths": response_lengths, + "cp_rank": self.cp_rank, + "cp_size": self.cp_size, + "cp_group": self.cp_group, + } + + if self.args.custom_tis_function_path is not None: + tis_func = load_function(self.args.custom_tis_function_path) + else: + tis_func = vanilla_tis_function + pg_loss, loss_masks, tis_metrics = tis_func(**tis_kwargs) + + if self.args.calculate_per_token_loss: + pg_loss = sum_of_token(pg_loss, response_lengths, loss_masks) + pg_clipfrac = sum_of_token(pg_clipfrac, response_lengths, loss_masks) + ppo_kl = sum_of_token(ppo_kl.abs(), response_lengths, loss_masks) + else: + pg_loss = sum_of_sample_mean(pg_loss, response_lengths, loss_masks) + pg_clipfrac = sum_of_sample_mean(pg_clipfrac, response_lengths, loss_masks) + ppo_kl = sum_of_sample_mean(ppo_kl.abs(), response_lengths, loss_masks) + + # Only compare rollout vs. train log probs when they originate from different stages. + train_rollout_logprob_abs_diff = None + if not self.args.use_rollout_logprobs and rollout_log_probs is not None: + train_rollout_logprob_abs_diff = (old_log_probs - rollout_log_probs).abs() + train_rollout_logprob_abs_diff = sum_of_sample_mean( + train_rollout_logprob_abs_diff, response_lengths, loss_masks + ).detach() + + entropy = torch.cat([batch["entropy"] for batch in unpacked_batches], dim=0) + entropy_loss = sum_of_sample_mean(entropy, response_lengths, loss_masks) + + loss = pg_loss - self.args.entropy_coef * entropy_loss + + if self.args.use_kl_loss: + ref_log_probs = torch.cat([batch["ref_log_probs"] for batch in unpacked_batches], dim=0) + importance_ratio = None + if self.args.use_unbiased_kl: + importance_ratio = torch.exp(log_probs - old_log_probs) + kl = compute_approx_kl( + log_probs, + ref_log_probs, + kl_loss_type=self.args.kl_loss_type, + importance_ratio=importance_ratio, + ) + kl_loss = sum_of_sample_mean(kl, response_lengths, loss_masks) + + loss = loss + self.args.kl_loss_coef * kl_loss + + reported = { + "loss": loss.detach(), + "pg_loss": pg_loss.detach(), + "pg_clipfrac": pg_clipfrac.detach(), + "ppo_kl": ppo_kl.detach(), + "entropy_loss": entropy_loss.detach(), + } + + if train_rollout_logprob_abs_diff is not None: + reported["train_rollout_logprob_abs_diff"] = train_rollout_logprob_abs_diff + + if self.args.use_kl_loss: + reported["kl_loss"] = kl_loss.detach() + + if self.args.use_opsm: + reported["opsm_clipfrac"] = opsm_clipfrac + + if self.args.use_tis and tis_metrics: + reported["ois"] = sum_of_sample_mean(ois, response_lengths, loss_masks).detach() + for k, v in tis_metrics.items(): + if self.args.calculate_per_token_loss: + reported[k] = sum_of_token(v, response_lengths, loss_masks).detach() + else: + reported[k] = sum_of_sample_mean(v, response_lengths, loss_masks).detach() + + # Scale loss for gradient accumulation + loss = loss * self.dp_size / self.args.global_batch_size + loss.backward() + + # Accumulate reported metrics (store tensors for later mean) + for k, v in reported.items(): + reported_accum.setdefault(k, []).append(v) + + if (mbs_id + 1) in grad_accum: + # TODO: check if the grad norm is global grad norm. + grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip_grad) + # the grad norm used to be of DTensor + grad_norm = float(grad_norm) + + self.optimizer.step() + # Update learning rate + self.lr_scheduler.step() + self.optimizer.zero_grad(set_to_none=True) + # Aggregate logs + aggregated = {k: torch.stack(v).sum().item() for k, v in reported_accum.items()} + # TODO: change this, this is slow. + reduced_aggregated = [None] * self.dp_size + dist.all_gather_object(reduced_aggregated, aggregated, group=self.dp_group) + aggregated = {} + for k in reported_accum.keys(): + aggregated[k] = sum([r[k] for r in reduced_aggregated]) / (self.args.global_batch_size) + reported_accum.clear() + if dist.get_rank() == 0: + log_dict = { + f"train/{k}": (val.item() if torch.is_tensor(val) else val) for k, val in aggregated.items() + } + log_dict["train/grad_norm"] = grad_norm + + # Log learning rate per parameter group; use scheduler's last computed LRs + lr_values = self.lr_scheduler.get_last_lr() + for gid, _group in enumerate(self.optimizer.param_groups): + log_dict[f"train/lr-pg_{gid}"] = lr_values[gid] + + kl_info = "" + if self.args.use_kl_loss and "kl_loss" in aggregated: + kl_info = f", kl_loss: {aggregated['kl_loss']:.4f}, kl_penalty: {aggregated['kl_loss'] * self.args.kl_loss_coef:.4f}" + logger.info(kl_info) + logger.info(f"step {self.global_step}: {log_dict}") + + log_dict["train/step"] = self.global_step + tracking_utils.log(self.args, log_dict, step_key="train/step") + self.global_step += 1 + + @timer + def update_weights(self) -> None: # type: ignore[override] + """Synchronize actor weights to rollout engines. + + Handles both colocated and distributed update modes. In offload mode, + wakes up parameters as needed to perform the update. + """ + if self.args.debug_train_only or self.args.debug_rollout_only: + return + + rollout_engines, rollout_engine_lock, num_new_engines = ray.get( + self.rollout_manager.get_rollout_engines_and_lock.remote() + ) + if num_new_engines > 0: + self.weight_updater.connect_rollout_engines(rollout_engines, rollout_engine_lock) + dist.barrier(group=get_gloo_group()) + + self.weight_updater.update_weights() + + if self.args.ci_test and len(rollout_engines) > 0: + engine = random.choice(rollout_engines) + engine_version = ray.get(engine.get_weight_version.remote()) + if str(engine_version) != str(self.weight_updater.weight_version): + raise RuntimeError( + f"Weight version mismatch! Engine: {engine_version}, Updater: {self.weight_updater.weight_version}" + ) + + clear_memory() + + def _create_ref_model(self, ref_load_path: str | None): + """Create and initialize a separate reference model with FSDP2 CPUOffloadPolicy. + + Parameters: + ref_load_path: Path to a directory containing a HF checkpoint. If + None, a ValueError is raised. + + Returns: + FSDP2-wrapped ref model with CPU offload enabled + + Note: + Creates a separate FSDP2 model instance for the reference model. + ALWAYS uses CPUOffloadPolicy for the reference model to save memory, + regardless of the actor model's CPU offload setting. + """ + if ref_load_path is None: + raise ValueError("ref_load_path must be provided when loading reference model") + + if os.path.isdir(ref_load_path): + logger.info(f"[Rank {dist.get_rank()}] Creating separate ref model from {ref_load_path}") + + init_context = self._get_init_weight_context_manager() + + with init_context(): + ref_model = self.get_model_cls().from_pretrained( + ref_load_path, + trust_remote_code=True, + attn_implementation=self.args.attn_implementation, + ) + + full_state = ref_model.state_dict() + + # Always use CPUOffloadPolicy for reference, let FSDP2 handle the offload. It is faster than model.cpu(). + ref_model = apply_fsdp2(ref_model, mesh=self.dp_mesh, cpu_offload=True, args=self.args) + ref_model = self._fsdp2_load_full_state_dict(ref_model, full_state, self.dp_mesh, cpu_offload=True) + + logger.info(f"[Rank {dist.get_rank()}] Reference model created with FSDP2 CPUOffloadPolicy") + return ref_model + else: + raise NotImplementedError(f"Loading from checkpoint file {ref_load_path} not yet implemented") + + def _get_model_inputs_args(self, packed_sequence: dict) -> dict: + input_ids = packed_sequence["tokens"].unsqueeze(0) + position_ids = packed_sequence["position_ids"].unsqueeze(0) + if self.cp_size > 1: + + packed_sequence = pad_packed_sequence_with_cp(packed_sequence, self.cp_size) + + if not packed_sequence["cu_seqlens"].is_cuda: + packed_sequence["cu_seqlens"] = packed_sequence["cu_seqlens"].cuda() + cu_seqlens = packed_sequence["cu_seqlens"] + update_ring_flash_attn_params(cu_seqlens, self.cp_group) + + input_ids = torch.chunk(packed_sequence["tokens"].unsqueeze(0), self.cp_size, dim=1)[self.cp_rank] + position_ids = torch.chunk(packed_sequence["position_ids"].unsqueeze(0), self.cp_size, dim=1)[self.cp_rank] + + model_args = { + "input_ids": input_ids, + "position_ids": position_ids, + "attention_mask": None, + } + if packed_sequence.get("multimodal_train_inputs"): + model_args.update(packed_sequence["multimodal_train_inputs"]) + return model_args + + +def selective_log_softmax_raw(logits: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: + """Fused version of the common `log_softmax -> gather` operation. + + The fused version of this operation avoids the (potentially large) memory overhead + of allocating a new tensor to store the full logprobs. + + Parameters: + logits: Tensor of shape [..., V] containing model logits. + input_ids: Tensor of shape [...] of token indices whose log-probabilities are gathered. + + Returns: + Tensor of shape [...] containing the log-probabilities corresponding to `input_ids`. + """ + logprobs = logits.log_softmax(dim=-1) + return torch.gather(logprobs, dim=-1, index=input_ids.unsqueeze(-1)).squeeze(-1) + + +selective_log_softmax_compiled = torch.compile(dynamic=True)(selective_log_softmax_raw) + + +def gather_log_probs_packed( + shifted_logits: torch.Tensor, + input_ids: torch.Tensor, + allow_compile: bool, + cu_seqlens: torch.Tensor | float | None = None, + temperature: torch.Tensor | None = None, +) -> torch.Tensor: + """Gather next-token log probabilities for packed sequences. + + Parameters: + logits: Model logits of shape [B, T, V] or [T, V]. + input_ids: Token ids of shape [B, T] or [T]. + cu_seqlens: Optional cumulative sequence lengths (unused here). Present + for API compatibility with callers. + + Returns: + A tensor of shape [T-1] (or [B, T-1]) with log-probabilities of targets. + """ + # Handle batch dimension - logits should be [batch_size, seq_len, vocab_size] + if shifted_logits.dim() == 3: + # Remove batch dimension for packed sequences + shifted_logits = shifted_logits.squeeze(0) + input_ids = input_ids.squeeze(0) + + if temperature is not None: + shifted_logits = shifted_logits.div(temperature) + + targets = input_ids[1:].to(device=shifted_logits.device) + + # Gather log probs for targets + selective_log_softmax = selective_log_softmax_compiled if allow_compile else selective_log_softmax_raw + return selective_log_softmax(shifted_logits, targets) + + +def get_logprob_and_entropy_with_cp( + logits: torch.Tensor, + target_tokens: torch.Tensor, + cp_rank: int, + cp_size: int, + cp_group, + model_input_ids: torch.Tensor, + allow_compile: bool, + temperature: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute log probabilities and entropy in Context Parallel mode. + + Parameters: + logits: Model output logits with shape [chunk_size, vocab_size] + target_tokens: Target tokens with shape [total_seq_len] + cp_rank: Current CP rank + cp_size: CP world size + cp_group: CP communication group + model_input_ids: Model input_ids (used for the last rank) + allow_compile: Whether to allow compilation + temperature: Temperature parameter (optional) + + Returns: + log_probs: Aggregated log probabilities with shape [total_seq_len - 1] + entropy: Aggregated entropy with shape [total_seq_len - 1] + """ + # Fast path for non-CP mode (cp_size=1): avoid unnecessary communication + if cp_size == 1: + shifted_logits = logits[:-1, :] + local_log_probs = gather_log_probs_packed( + shifted_logits, target_tokens, allow_compile=allow_compile, temperature=temperature + ) + log_probs_full = torch.log_softmax(shifted_logits, dim=-1) + probs = torch.softmax(shifted_logits, dim=-1) + entropy = -(probs * log_probs_full).sum(dim=-1) + return local_log_probs, entropy + + chunk_size = logits.shape[0] + tokens_start_index = chunk_size * cp_rank + tokens_end_index = ( + tokens_start_index + chunk_size + 1 if cp_rank < cp_size - 1 else tokens_start_index + chunk_size + ) + + # For the last rank, remove the last logit + logits = logits if cp_rank < cp_size - 1 else logits[:-1, :] + + # Get local tokens for current rank + local_tokens = ( + target_tokens[tokens_start_index:tokens_end_index] if cp_rank < cp_size - 1 else model_input_ids.squeeze(0) + ) + + # Compute local log probs + local_log_probs = gather_log_probs_packed( + logits, local_tokens, allow_compile=allow_compile, temperature=temperature + ) + + # Pad for the last rank + if cp_rank == cp_size - 1: + local_log_probs = F.pad(local_log_probs, (0, chunk_size - local_log_probs.shape[0]), value=0) + + # Compute entropy + shifted_logits = logits[:-1, :] if cp_rank == cp_size - 1 else logits + log_probs_full = torch.log_softmax(shifted_logits, dim=-1) + probs = torch.softmax(shifted_logits, dim=-1) + entropy = -(probs * log_probs_full).sum(dim=-1) + + # Pad entropy for the last rank + if cp_rank == cp_size - 1: + entropy = F.pad(entropy, (0, chunk_size - entropy.shape[0]), value=0) + + # Merge with a single all_gather: stack as [2, chunk_size] + stacked_local = torch.stack([local_log_probs, entropy], dim=0) + gathered_stacked = torch.distributed.nn.functional.all_gather(stacked_local, group=cp_group) + + # Concatenate by effective length (non-last rank=chunk_size, last rank=chunk_size-1) + lp_parts, ent_parts = [], [] + for r in range(cp_size): + eff_len = chunk_size if r < cp_size - 1 else max(0, chunk_size - 1) + if eff_len > 0: + lp_parts.append(gathered_stacked[r][0][:eff_len]) + ent_parts.append(gathered_stacked[r][1][:eff_len]) + + log_probs = torch.cat(lp_parts, dim=0) if lp_parts else local_log_probs.new_zeros((0,)) + entropy_result = torch.cat(ent_parts, dim=0) if ent_parts else entropy.new_zeros((0,)) + + # Truncate to global effective length T-1 (packed tokens length is T) + log_probs = log_probs[: len(target_tokens) - 1] + entropy_result = entropy_result[: len(target_tokens) - 1] + + return log_probs, entropy_result + + +def sum_of_sample_mean(x: torch.Tensor, response_lengths: list[int], loss_masks: list[torch.Tensor]) -> torch.Tensor: + """Compute sum of per-sample means across variable-length responses. + + Parameters: + x: Flat tensor containing concatenated per-token values across samples. + response_lengths: Lengths of each sample's response segment in `x`. + loss_masks: Per-sample masks aligned with `response_lengths`. + + Returns: + A scalar tensor equal to the sum over samples of the mean value within + each sample's response segment. + """ + return sum( + [ + (x_i * loss_mask_i).sum() / torch.clamp_min(loss_mask_i.sum(), 1) + for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=False) + ] + ) + + +@torch.no_grad() +def move_torch_optimizer(optimizer, device): + """ref: https://github.com/volcengine/verl/blob/main/verl/utils/fsdp_utils.py""" + if not optimizer.state: + return + + for param_group in optimizer.param_groups: + for param in param_group["params"]: + state = optimizer.state[param] + for key, value in state.items(): + if isinstance(value, torch.Tensor): + state[key] = value.to(device, non_blocking=True) + + torch.cuda.synchronize() + + +def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None): + """Apply FSDP v2 to the model. + + Args: + model: The model to wrap with FSDP + mesh: Optional DeviceMesh for FSDP. If None, uses all ranks. + cpu_offload: If True, offload parameters, gradients, and optimizer states + to CPU. The optimizer step will run on CPU. (Default: False) + args: Arguments containing precision settings (fp16/bf16) + + Ref: https://github.com/volcengine/verl/blob/main/verl/utils/fsdp_utils.py + """ + from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard + + offload_policy = CPUOffloadPolicy() if cpu_offload else None + + layer_cls_to_wrap = model._no_split_modules + assert len(layer_cls_to_wrap) > 0 and layer_cls_to_wrap[0] is not None + + modules = [ + module + for name, module in model.named_modules() + if module.__class__.__name__ in layer_cls_to_wrap + or (isinstance(module, torch.nn.Embedding) and not model.config.tie_word_embeddings) + ] + + # Determine precision policy based on args + param_dtype = torch.bfloat16 # Default to bf16 as before + reduce_dtype = torch.float32 + + if args.fp16: + param_dtype = torch.float16 + + logger.info(f"FSDP MixedPrecision Policy: param_dtype={param_dtype}, reduce_dtype={reduce_dtype}") + + fsdp_kwargs = { + "mp_policy": MixedPrecisionPolicy( + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + ), + "offload_policy": offload_policy, + "mesh": mesh, + } + + # Apply FSDP to each module (offload_policy=None is equivalent to not passing it) + for module in modules: + fully_shard(module, **fsdp_kwargs) + + # Apply FSDP to the top-level model + fully_shard(model, **fsdp_kwargs) + + return model + + +def sum_of_token(x: torch.Tensor, response_lengths: list[int], loss_masks: list[torch.Tensor]) -> torch.Tensor: + return sum( + [ + (x_i * loss_mask_i).sum() + for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=False) + ] + ) diff --git a/slime/backends/fsdp_utils/arguments.py b/slime/backends/fsdp_utils/arguments.py new file mode 100644 index 0000000000000000000000000000000000000000..7664c946268a4cf7e3630ab89d2fe474356e2434 --- /dev/null +++ b/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 diff --git a/slime/backends/fsdp_utils/checkpoint.py b/slime/backends/fsdp_utils/checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..5b883c0e279bc64dcff74dd58be0f86b167ed6ce --- /dev/null +++ b/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() diff --git a/slime/backends/fsdp_utils/data_packing.py b/slime/backends/fsdp_utils/data_packing.py new file mode 100644 index 0000000000000000000000000000000000000000..02c5d967f1611de7752563ce94a75e3633a2b69a --- /dev/null +++ b/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 diff --git a/slime/backends/fsdp_utils/kernels/__init__.py b/slime/backends/fsdp_utils/kernels/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a4d636226b9a4b5130dfd8e94f2cece5fd290e53 --- /dev/null +++ b/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 + diff --git a/slime/backends/fsdp_utils/kernels/fused_experts.py b/slime/backends/fsdp_utils/kernels/fused_experts.py new file mode 100644 index 0000000000000000000000000000000000000000..0b8bb783032bc5a8a47da7330973e28a2b6d9012 --- /dev/null +++ b/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 diff --git a/slime/backends/fsdp_utils/kernels/fused_moe_triton_backward_kernels.py b/slime/backends/fsdp_utils/kernels/fused_moe_triton_backward_kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..863f79ae846b3fc118e1b7f31f4d11288ff85ae9 --- /dev/null +++ b/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"], + ) diff --git a/slime/backends/fsdp_utils/lr_scheduler.py b/slime/backends/fsdp_utils/lr_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..5b1cf5be87082a61b6478acc26b8d5ef858d49db --- /dev/null +++ b/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 diff --git a/slime/backends/fsdp_utils/models/__init__.py b/slime/backends/fsdp_utils/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a4d636226b9a4b5130dfd8e94f2cece5fd290e53 --- /dev/null +++ b/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 + diff --git a/slime/backends/fsdp_utils/models/qwen3_moe.py b/slime/backends/fsdp_utils/models/qwen3_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..7ce9cdcfe5b62a50800fe64fca54f21be2d4c551 --- /dev/null +++ b/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 diff --git a/slime/backends/fsdp_utils/models/qwen3_moe_hf.py b/slime/backends/fsdp_utils/models/qwen3_moe_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..ef3041883c6e0933ad34887357924c909b7b1f1e --- /dev/null +++ b/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 diff --git a/slime/backends/fsdp_utils/update_weight_utils.py b/slime/backends/fsdp_utils/update_weight_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..97a50e677c55cefd9ef6254a57d850a65e9c0091 --- /dev/null +++ b/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) diff --git a/slime/backends/megatron_utils/__init__.py b/slime/backends/megatron_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..22682230828571bd4711df0f00d1b23ea4e143d9 --- /dev/null +++ b/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) diff --git a/slime/backends/megatron_utils/__pycache__/__init__.cpython-312.pyc b/slime/backends/megatron_utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a4e41bf9e64acda0e593f017c77f1ac0261c7606 Binary files /dev/null and b/slime/backends/megatron_utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/__pycache__/actor.cpython-312.pyc b/slime/backends/megatron_utils/__pycache__/actor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d221c9a87bc7a6b839c242c1f76f1ab69d4cc8a4 Binary files /dev/null and b/slime/backends/megatron_utils/__pycache__/actor.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/__pycache__/arguments.cpython-312.pyc b/slime/backends/megatron_utils/__pycache__/arguments.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..16426891d0594004c1786368533bc22faaab224d Binary files /dev/null and b/slime/backends/megatron_utils/__pycache__/arguments.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/__pycache__/checkpoint.cpython-312.pyc b/slime/backends/megatron_utils/__pycache__/checkpoint.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0152b56c99895fb591813191d5fe28c8d77cf881 Binary files /dev/null and b/slime/backends/megatron_utils/__pycache__/checkpoint.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/__pycache__/cp_utils.cpython-312.pyc b/slime/backends/megatron_utils/__pycache__/cp_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..18c30f46f25fa75ab0781e9d65f0d5cb54982681 Binary files /dev/null and b/slime/backends/megatron_utils/__pycache__/cp_utils.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/__pycache__/data.cpython-312.pyc b/slime/backends/megatron_utils/__pycache__/data.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c86f1a47215095a68ef552059f9e9fd5bfdc6111 Binary files /dev/null and b/slime/backends/megatron_utils/__pycache__/data.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/__pycache__/initialize.cpython-312.pyc b/slime/backends/megatron_utils/__pycache__/initialize.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cce021e068b0d223f51827b4694c6db347989875 Binary files /dev/null and b/slime/backends/megatron_utils/__pycache__/initialize.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/__pycache__/loss.cpython-312.pyc b/slime/backends/megatron_utils/__pycache__/loss.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d7737590b1107c52f0afd7630a514751b02f32ef Binary files /dev/null and b/slime/backends/megatron_utils/__pycache__/loss.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/__pycache__/misc_utils.cpython-312.pyc b/slime/backends/megatron_utils/__pycache__/misc_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5b63724b7cd4809339aef3f2fade59c5dddb4181 Binary files /dev/null and b/slime/backends/megatron_utils/__pycache__/misc_utils.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/__pycache__/model.cpython-312.pyc b/slime/backends/megatron_utils/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..74f313824e88c0002ff98327afa8853fe2df1121 Binary files /dev/null and b/slime/backends/megatron_utils/__pycache__/model.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/__pycache__/model_provider.cpython-312.pyc b/slime/backends/megatron_utils/__pycache__/model_provider.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..21c9daaebc378b32ce10bd4b57fd2702e55d67d0 Binary files /dev/null and b/slime/backends/megatron_utils/__pycache__/model_provider.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/__pycache__/sglang.cpython-312.pyc b/slime/backends/megatron_utils/__pycache__/sglang.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..97d50b839adf2711d5e051d1778e8716e0c232c6 Binary files /dev/null and b/slime/backends/megatron_utils/__pycache__/sglang.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py new file mode 100644 index 0000000000000000000000000000000000000000..bbe6856d04febc8633592881abf2cdd777e4731c --- /dev/null +++ b/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, + ) diff --git a/slime/backends/megatron_utils/arguments.py b/slime/backends/megatron_utils/arguments.py new file mode 100644 index 0000000000000000000000000000000000000000..7258787da79a12962811c6346d3a9b6657b86045 --- /dev/null +++ b/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 diff --git a/slime/backends/megatron_utils/checkpoint.py b/slime/backends/megatron_utils/checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..65ae078466d8163af3dd6a6d9d4216a9edeabeb4 --- /dev/null +++ b/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) diff --git a/slime/backends/megatron_utils/ci_utils.py b/slime/backends/megatron_utils/ci_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5bf94077b5fb389318baf4d4b96a62081164934a --- /dev/null +++ b/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." + ) diff --git a/slime/backends/megatron_utils/cp_utils.py b/slime/backends/megatron_utils/cp_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..da627809c42a7416e8fa8710c932519c702c4ec6 --- /dev/null +++ b/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) diff --git a/slime/backends/megatron_utils/data.py b/slime/backends/megatron_utils/data.py new file mode 100644 index 0000000000000000000000000000000000000000..3bf844a15205bffb18fd52b7c9f3182c37f39f58 --- /dev/null +++ b/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 + } + ) diff --git a/slime/backends/megatron_utils/initialize.py b/slime/backends/megatron_utils/initialize.py new file mode 100644 index 0000000000000000000000000000000000000000..33cd4b5ae59e90a1ba307cede0332f412f3439bd --- /dev/null +++ b/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 + ) diff --git a/slime/backends/megatron_utils/kernels/int4_qat/fake_int4_quant_cuda.cu b/slime/backends/megatron_utils/kernels/int4_qat/fake_int4_quant_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..a6e9554906cbd401eba93e83b2b11dbb8462baaf --- /dev/null +++ b/slime/backends/megatron_utils/kernels/int4_qat/fake_int4_quant_cuda.cu @@ -0,0 +1,368 @@ +#include +#include + +#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 +__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(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(val); + out_scale[row * stride_osm + item * stride_osn] = static_cast(scale); + if(!sym) { + out_zero[row * stride_ozm + item * stride_ozn] = static_cast(zero); + } + } + } +} + +// for some transpose case, blocksize = [32, 1] +template +__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(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(val); + out_scale[blockIdx.x * stride_osm + item * stride_osn] = static_cast(scale); + if (!sym) { + out_zero[blockIdx.x * stride_ozm + item * stride_ozn] = static_cast(zero); + } + } + } +} + +template +__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(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(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(val); + out_scale[blockIdx.x * stride_osm + item * stride_osn] = static_cast(scale); + if (!sym) { + out_zero[blockIdx.x * stride_ozm + item * stride_ozn] = static_cast(zero); + } + } + } + } +} + +// dispatch +template +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<<>>( + 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<<>>( + 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<<>>( + 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 +fake_int4_quant_cuda( + torch::Tensor& x, + std::vector& 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( + x.const_data_ptr(), + out.data_ptr(), + out_scale.data_ptr(), + out_zero.data_ptr(), + 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"); +} diff --git a/slime/backends/megatron_utils/kernels/int4_qat/setup.py b/slime/backends/megatron_utils/kernels/int4_qat/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..08b9a41169969cbc41cd29ad8ad5e867f7a5e4cd --- /dev/null +++ b/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}, +) diff --git a/slime/backends/megatron_utils/loss.py b/slime/backends/megatron_utils/loss.py new file mode 100644 index 0000000000000000000000000000000000000000..addbd77084e136d57385051483c7e1eb76a08925 --- /dev/null +++ b/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, + ), + }, + ) diff --git a/slime/backends/megatron_utils/megatron_to_hf/__init__.py b/slime/backends/megatron_utils/megatron_to_hf/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..feb6548c1396be69495567ad24cf059aa823d820 --- /dev/null +++ b/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 diff --git a/slime/backends/megatron_utils/megatron_to_hf/__pycache__/__init__.cpython-312.pyc b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0c975bd8c51347fecb683b301dde3704a26f800c Binary files /dev/null and b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/__init__.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/megatron_to_hf/__pycache__/deepseekv3.cpython-312.pyc b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/deepseekv3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f402326fe054c88c430a7e169257b71939c8d470 Binary files /dev/null and b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/deepseekv3.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/megatron_to_hf/__pycache__/glm4.cpython-312.pyc b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/glm4.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1d46547685ee44345b50db9b253c709777924afc Binary files /dev/null and b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/glm4.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/megatron_to_hf/__pycache__/glm4moe.cpython-312.pyc b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/glm4moe.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..96a741f24caa7476e0a4a9c6581c44a0937ed649 Binary files /dev/null and b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/glm4moe.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/megatron_to_hf/__pycache__/llama.cpython-312.pyc b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/llama.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..79787ea214b3ef89e3c91bce28d55b031806abd7 Binary files /dev/null and b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/llama.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/megatron_to_hf/__pycache__/mimo.cpython-312.pyc b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/mimo.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a60333ee3e6ccd9fbe19eb5430b6227297c66de3 Binary files /dev/null and b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/mimo.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/megatron_to_hf/__pycache__/qwen2.cpython-312.pyc b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/qwen2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a227bd2a5c5937d37383c7b3c90a6258543b516a Binary files /dev/null and b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/qwen2.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/megatron_to_hf/__pycache__/qwen3_next.cpython-312.pyc b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/qwen3_next.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..899b6492302c738ade61c8866babb58fb0a64293 Binary files /dev/null and b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/qwen3_next.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/megatron_to_hf/__pycache__/qwen3moe.cpython-312.pyc b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/qwen3moe.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..218d8a19a1d2c350906739ea97cc4f19e719e01a Binary files /dev/null and b/slime/backends/megatron_utils/megatron_to_hf/__pycache__/qwen3moe.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/megatron_to_hf/deepseekv3.py b/slime/backends/megatron_utils/megatron_to_hf/deepseekv3.py new file mode 100644 index 0000000000000000000000000000000000000000..d9af816853e82a118ec01a73629da012446c1704 --- /dev/null +++ b/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}") diff --git a/slime/backends/megatron_utils/megatron_to_hf/glm4.py b/slime/backends/megatron_utils/megatron_to_hf/glm4.py new file mode 100644 index 0000000000000000000000000000000000000000..196a3f6dd35687ca8bf91ffe3e2630beee9e9082 --- /dev/null +++ b/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}") diff --git a/slime/backends/megatron_utils/megatron_to_hf/glm4moe.py b/slime/backends/megatron_utils/megatron_to_hf/glm4moe.py new file mode 100644 index 0000000000000000000000000000000000000000..a93a0942657d552099c1082e5cbf9f09ab3ccffe --- /dev/null +++ b/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}") diff --git a/slime/backends/megatron_utils/megatron_to_hf/llama.py b/slime/backends/megatron_utils/megatron_to_hf/llama.py new file mode 100644 index 0000000000000000000000000000000000000000..13d2fcdf57b0b72ae7a2fc29776d46d4b6fc398f --- /dev/null +++ b/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}") diff --git a/slime/backends/megatron_utils/megatron_to_hf/mimo.py b/slime/backends/megatron_utils/megatron_to_hf/mimo.py new file mode 100644 index 0000000000000000000000000000000000000000..59954b91e09250cc184308ffd172c7f1d92cdd54 --- /dev/null +++ b/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}") diff --git a/slime/backends/megatron_utils/megatron_to_hf/processors/__init__.py b/slime/backends/megatron_utils/megatron_to_hf/processors/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a4d636226b9a4b5130dfd8e94f2cece5fd290e53 --- /dev/null +++ b/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 + diff --git a/slime/backends/megatron_utils/megatron_to_hf/processors/__pycache__/__init__.cpython-312.pyc b/slime/backends/megatron_utils/megatron_to_hf/processors/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cc25068d1773ecc4529e2a193dde53d3bfbf480b Binary files /dev/null and b/slime/backends/megatron_utils/megatron_to_hf/processors/__pycache__/__init__.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/megatron_to_hf/processors/__pycache__/padding_remover.cpython-312.pyc b/slime/backends/megatron_utils/megatron_to_hf/processors/__pycache__/padding_remover.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d0b01ef0f263b8f5648be340ef72292f6dcfa36 Binary files /dev/null and b/slime/backends/megatron_utils/megatron_to_hf/processors/__pycache__/padding_remover.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/megatron_to_hf/processors/__pycache__/quantizer.cpython-312.pyc b/slime/backends/megatron_utils/megatron_to_hf/processors/__pycache__/quantizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5dedd2bcab627b9668c7a13781dbd41efe337eba Binary files /dev/null and b/slime/backends/megatron_utils/megatron_to_hf/processors/__pycache__/quantizer.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/megatron_to_hf/processors/padding_remover.py b/slime/backends/megatron_utils/megatron_to_hf/processors/padding_remover.py new file mode 100644 index 0000000000000000000000000000000000000000..e7c49c9a51aeb5c0f2a66ddb477d6e433d847224 --- /dev/null +++ b/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 diff --git a/slime/backends/megatron_utils/megatron_to_hf/processors/quantizer.py b/slime/backends/megatron_utils/megatron_to_hf/processors/quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..341bc45c5e87b8e3ca7ebab2638d3776f8e2a8bc --- /dev/null +++ b/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)] diff --git a/slime/backends/megatron_utils/megatron_to_hf/qwen2.py b/slime/backends/megatron_utils/megatron_to_hf/qwen2.py new file mode 100644 index 0000000000000000000000000000000000000000..aed81eec86276e5c61291af1b6f810ada9d7ca89 --- /dev/null +++ b/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}") diff --git a/slime/backends/megatron_utils/megatron_to_hf/qwen3_next.py b/slime/backends/megatron_utils/megatron_to_hf/qwen3_next.py new file mode 100644 index 0000000000000000000000000000000000000000..69a3981ee9e81a24660a9466aea30cb91b4a0799 --- /dev/null +++ b/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}") diff --git a/slime/backends/megatron_utils/megatron_to_hf/qwen3moe.py b/slime/backends/megatron_utils/megatron_to_hf/qwen3moe.py new file mode 100644 index 0000000000000000000000000000000000000000..4c24b6ad2faa5d4d9e53f2220214922d25a8b863 --- /dev/null +++ b/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}") diff --git a/slime/backends/megatron_utils/misc_utils.py b/slime/backends/megatron_utils/misc_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..85977768ade73a52cf42c1a1a799788d127a2cb0 --- /dev/null +++ b/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 diff --git a/slime/backends/megatron_utils/model.py b/slime/backends/megatron_utils/model.py new file mode 100644 index 0000000000000000000000000000000000000000..72f190719d6696ba11dd248c8314d4de8219f00a --- /dev/null +++ b/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 diff --git a/slime/backends/megatron_utils/model_provider.py b/slime/backends/megatron_utils/model_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..371b3a17203a47a73a0c1136b34b63af86b7cab3 --- /dev/null +++ b/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 diff --git a/slime/backends/megatron_utils/sglang.py b/slime/backends/megatron_utils/sglang.py new file mode 100644 index 0000000000000000000000000000000000000000..f2c03291c7a2e4b1c7e2adcffe7f2872c3a6477c --- /dev/null +++ b/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", +] diff --git a/slime/backends/megatron_utils/update_weight/__init__.py b/slime/backends/megatron_utils/update_weight/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a4d636226b9a4b5130dfd8e94f2cece5fd290e53 --- /dev/null +++ b/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 + diff --git a/slime/backends/megatron_utils/update_weight/__pycache__/__init__.cpython-312.pyc b/slime/backends/megatron_utils/update_weight/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..726b83c06be5d276336bf6e38bccbcf15172a5d5 Binary files /dev/null and b/slime/backends/megatron_utils/update_weight/__pycache__/__init__.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/update_weight/__pycache__/common.cpython-312.pyc b/slime/backends/megatron_utils/update_weight/__pycache__/common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f37d708d099b8f1812e35027f9661b0559a2a44a Binary files /dev/null and b/slime/backends/megatron_utils/update_weight/__pycache__/common.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/update_weight/__pycache__/hf_weight_iterator_base.cpython-312.pyc b/slime/backends/megatron_utils/update_weight/__pycache__/hf_weight_iterator_base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7f1da7bd603e50dc6f87199a39f9892811cde40d Binary files /dev/null and b/slime/backends/megatron_utils/update_weight/__pycache__/hf_weight_iterator_base.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/update_weight/__pycache__/update_weight_from_distributed.cpython-312.pyc b/slime/backends/megatron_utils/update_weight/__pycache__/update_weight_from_distributed.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d67d42242cce6cf43bb26ce0edd0e6d9b69e97a8 Binary files /dev/null and b/slime/backends/megatron_utils/update_weight/__pycache__/update_weight_from_distributed.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/update_weight/__pycache__/update_weight_from_tensor.cpython-312.pyc b/slime/backends/megatron_utils/update_weight/__pycache__/update_weight_from_tensor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8d79d6dc94c75f2d51040b26ca14fcdae0f19e7 Binary files /dev/null and b/slime/backends/megatron_utils/update_weight/__pycache__/update_weight_from_tensor.cpython-312.pyc differ diff --git a/slime/backends/megatron_utils/update_weight/common.py b/slime/backends/megatron_utils/update_weight/common.py new file mode 100644 index 0000000000000000000000000000000000000000..5422ffa10543e73a7e5edf0b53655000a4b5d57a --- /dev/null +++ b/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 diff --git a/slime/backends/megatron_utils/update_weight/hf_weight_iterator_base.py b/slime/backends/megatron_utils/update_weight/hf_weight_iterator_base.py new file mode 100644 index 0000000000000000000000000000000000000000..a08c30e0b033040542f5b5d99631b4228ecf618c --- /dev/null +++ b/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 diff --git a/slime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py b/slime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py new file mode 100644 index 0000000000000000000000000000000000000000..209edcd8fb9d0341477817a59e6138d21da59c8d --- /dev/null +++ b/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) diff --git a/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py b/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py new file mode 100644 index 0000000000000000000000000000000000000000..fb0da94ad2b0851f57cc97f9df5815a27df9f46a --- /dev/null +++ b/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 diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py new file mode 100644 index 0000000000000000000000000000000000000000..452372bda491bfe396f061b3a16d1c508734e7c9 --- /dev/null +++ b/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 diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..aa5bc314427d2c6cdb16b56b029fea795485486d --- /dev/null +++ b/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 diff --git a/slime/backends/sglang_utils/__init__.py b/slime/backends/sglang_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a4d636226b9a4b5130dfd8e94f2cece5fd290e53 --- /dev/null +++ b/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 + diff --git a/slime/backends/sglang_utils/__pycache__/__init__.cpython-312.pyc b/slime/backends/sglang_utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ccaf6579f338360c3621318d4cd4e203a4319ce8 Binary files /dev/null and b/slime/backends/sglang_utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/slime/backends/sglang_utils/__pycache__/arguments.cpython-312.pyc b/slime/backends/sglang_utils/__pycache__/arguments.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84da3bd682ab06acc016c9b7e73db0fb7721977e Binary files /dev/null and b/slime/backends/sglang_utils/__pycache__/arguments.cpython-312.pyc differ diff --git a/slime/backends/sglang_utils/__pycache__/sglang_engine.cpython-312.pyc b/slime/backends/sglang_utils/__pycache__/sglang_engine.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..41a9df7364b59647668e2b36b7ec5cf1738a1a51 Binary files /dev/null and b/slime/backends/sglang_utils/__pycache__/sglang_engine.cpython-312.pyc differ diff --git a/slime/backends/sglang_utils/arguments.py b/slime/backends/sglang_utils/arguments.py new file mode 100644 index 0000000000000000000000000000000000000000..eb2b82453be4050432b2da40c8acdd2054542876 --- /dev/null +++ b/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) diff --git a/slime/backends/sglang_utils/sglang_engine.py b/slime/backends/sglang_utils/sglang_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..27cec8d615de010a897186b3f255eeb4e7f44a6f --- /dev/null +++ b/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", +] diff --git a/slime/ray/__init__.py b/slime/ray/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a4d636226b9a4b5130dfd8e94f2cece5fd290e53 --- /dev/null +++ b/slime/ray/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + diff --git a/slime/ray/__pycache__/__init__.cpython-312.pyc b/slime/ray/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..559aa6d7337000ec19e2dd992b865d54558e4031 Binary files /dev/null and b/slime/ray/__pycache__/__init__.cpython-312.pyc differ diff --git a/slime/ray/__pycache__/actor_group.cpython-312.pyc b/slime/ray/__pycache__/actor_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7da842e0df9043f5e6d4d9bae698c17b822fc331 Binary files /dev/null and b/slime/ray/__pycache__/actor_group.cpython-312.pyc differ diff --git a/slime/ray/__pycache__/placement_group.cpython-312.pyc b/slime/ray/__pycache__/placement_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7652c7802d71bf8af5edd09262f25be9409db536 Binary files /dev/null and b/slime/ray/__pycache__/placement_group.cpython-312.pyc differ diff --git a/slime/ray/__pycache__/ray_actor.cpython-312.pyc b/slime/ray/__pycache__/ray_actor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5ea8e03ea8e544268cb17d9be08d7b300a08dfc1 Binary files /dev/null and b/slime/ray/__pycache__/ray_actor.cpython-312.pyc differ diff --git a/slime/ray/__pycache__/rollout.cpython-312.pyc b/slime/ray/__pycache__/rollout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d63e5e58de2267a393351e79bc8dd201cc02de6a Binary files /dev/null and b/slime/ray/__pycache__/rollout.cpython-312.pyc differ diff --git a/slime/ray/__pycache__/train_actor.cpython-312.pyc b/slime/ray/__pycache__/train_actor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24b945b31c704c4ae8fb807c0ea57b71fb78d455 Binary files /dev/null and b/slime/ray/__pycache__/train_actor.cpython-312.pyc differ diff --git a/slime/ray/__pycache__/utils.cpython-312.pyc b/slime/ray/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50d5c70a8b181dd7024ab9a88e909f8e8499df54 Binary files /dev/null and b/slime/ray/__pycache__/utils.cpython-312.pyc differ diff --git a/slime/ray/actor_group.py b/slime/ray/actor_group.py new file mode 100644 index 0000000000000000000000000000000000000000..7c96a89ab0ec558c65165c96ae970173a9895fab --- /dev/null +++ b/slime/ray/actor_group.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os + +import ray +from ray.util.placement_group import PlacementGroup +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +from slime.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST + + +class RayTrainGroup: + """ + A group of ray actors + Functions start with 'async' should return list of object refs + + Args: + args (Namespace): Arguments for the actor group. + num_nodes (int): Number of nodes for this actor group. + num_gpus_per_node (int): Number of gpus for this actor group. + pg (PlacementGroup, optional): Placement group to schedule actor on. + If none, create new placement group automatically. Defaults to None. + num_gpus_per_actor (float, optional): Number of gpus allocated for each actor. + If < 1.0, multiple models can share same gpu. Defaults to 1. + resources (Dict[str, float], optional): Custom resources to allocate for each actor. + See https://docs.ray.io/en/latest/ray-core/scheduling/resources.html + num_resources_per_node (int, optional): Number of custom resources to allocate for each node. + See https://docs.ray.io/en/latest/ray-core/scheduling/resources.html + """ + + def __init__( + self, + args, + num_nodes, + num_gpus_per_node, + pg: tuple[PlacementGroup, list[int]], + num_gpus_per_actor: float = 1, + role: str = "actor", + ) -> None: + self.args = args + self._num_nodes = num_nodes + self._num_gpus_per_node = num_gpus_per_node + self.role = role + + # Allocate the GPUs for actors w/o instantiating them + self._allocate_gpus_for_actor(pg, num_gpus_per_actor) + + def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): + world_size = self._num_nodes * self._num_gpus_per_node + + # Use placement group to lock resources for models of same type + assert pg is not None + pg, reordered_bundle_indices = pg + + env_vars = { + # because sglang will always set NCCL_CUMEM_ENABLE to 0 + # we need also set it to 0 to prevent nccl error. + "NCCL_CUMEM_ENABLE": os.environ.get("NCCL_CUMEM_ENABLE", "0"), + "NVTE_FP8_BLOCK_SCALING_FP32_SCALES": "1", + **{name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST}, + **self.args.train_env_vars, + } + + if self.args.offload_train and self.args.train_backend == "megatron": + import torch_memory_saver + + dynlib_path = os.path.join( + os.path.dirname(os.path.dirname(torch_memory_saver.__file__)), + "torch_memory_saver_hook_mode_preload.abi3.so", + ) + assert os.path.exists(dynlib_path), f"LD_PRELOAD so file {dynlib_path} does not exist." + + env_vars["LD_PRELOAD"] = dynlib_path + env_vars["TMS_INIT_ENABLE"] = "1" + env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1" + + # We cannot do routing replay for critic. + if self.args.use_routing_replay and self.role == "actor": + env_vars["ENABLE_ROUTING_REPLAY"] = "1" + + backend = self.args.train_backend + if backend == "megatron": + from slime.backends.megatron_utils.actor import MegatronTrainRayActor + + actor_impl = MegatronTrainRayActor + + else: + from slime.backends.fsdp_utils import FSDPTrainRayActor + + actor_impl = FSDPTrainRayActor + + TrainRayActor = ray.remote(num_gpus=1, runtime_env={"env_vars": env_vars})(actor_impl) + + # Create worker actors + self._actor_handlers = [] + master_addr, master_port = None, None + for rank in range(world_size): + actor = TrainRayActor.options( + num_cpus=num_gpus_per_actor, + num_gpus=num_gpus_per_actor, + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=reordered_bundle_indices[rank], + ), + ).remote(world_size, rank, master_addr, master_port) + if rank == 0: + master_addr, master_port = ray.get(actor.get_master_addr_and_port.remote()) + self._actor_handlers.append(actor) + + def async_init(self, args, role, with_ref=False): + """ + Allocate GPU resourced and initialize model, optimzier, local ckpt, etc. + """ + self.args = args + return [actor.init.remote(args, role, with_ref=with_ref) for actor in self._actor_handlers] + + def async_train(self, rollout_id, rollout_data_ref): + """Do one rollout training""" + return [actor.train.remote(rollout_id, rollout_data_ref) for actor in self._actor_handlers] + + def save_model(self, rollout_id, force_sync=False): + """Save actor model""" + return ray.get([actor.save_model.remote(rollout_id, force_sync=force_sync) for actor in self._actor_handlers]) + + def update_weights(self): + """Broadcast weights from rank 0 to all other ranks.""" + return ray.get([actor.update_weights.remote() for actor in self._actor_handlers]) + + def onload(self): + return ray.get([actor.wake_up.remote() for actor in self._actor_handlers]) + + def offload(self): + return ray.get([actor.sleep.remote() for actor in self._actor_handlers]) + + def clear_memory(self): + return ray.get([actor.clear_memory.remote() for actor in self._actor_handlers]) + + def connect(self, critic_group): + return ray.get( + [ + actor.connect_actor_critic.remote(critic) + for actor, critic in zip(self._actor_handlers, critic_group._actor_handlers, strict=False) + ] + ) + + def set_rollout_manager(self, rollout_manager): + return ray.get([actor.set_rollout_manager.remote(rollout_manager) for actor in self._actor_handlers]) diff --git a/slime/ray/placement_group.py b/slime/ray/placement_group.py new file mode 100644 index 0000000000000000000000000000000000000000..be552392aabb0e49737f03c9b606eff0390c2ae0 --- /dev/null +++ b/slime/ray/placement_group.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +import socket +import ray +from ray.util.placement_group import placement_group +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +from .actor_group import RayTrainGroup +from .rollout import RolloutManager + +logger = logging.getLogger(__name__) + + +@ray.remote(num_gpus=1) +class InfoActor: + def get_ip_and_gpu_id(self): + return ray.util.get_node_ip_address(), ray.get_gpu_ids()[0] + + +def sort_key(x): + index, node_identifier, gpu_id = x + # Sort by node IP number and then by GPU ID + try: + # try to parse it as an IP address. + ip_address = node_identifier + node_ip_parts = list(map(int, ip_address.split("."))) + except ValueError: + # Try to resolve the hostname to an IP address. + try: + ip_address = socket.gethostbyname(node_identifier) + node_ip_parts = list(map(int, ip_address.split("."))) + except (socket.gaierror, TypeError): + # Instead, we convert each character of the original identifier string + # to its ASCII value. This provides a stable and consistent numerical + # representation that allows for sorting. + node_ip_parts = [ord(c) for c in node_identifier] + + return (node_ip_parts, gpu_id) + + +def _create_placement_group(num_gpus): + """Create a placement group with the specified number of GPUs.""" + bundles = [{"GPU": 1, "CPU": 1} for _ in range(num_gpus)] + pg = placement_group(bundles, strategy="PACK") + num_bundles = len(bundles) + + ray.get(pg.ready()) + # use info actor to get the GPU id + info_actors = [] + for i in range(num_bundles): + info_actors.append( + InfoActor.options( + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=i, + ) + ).remote() + ) + gpu_ids = ray.get([actor.get_ip_and_gpu_id.remote() for actor in info_actors]) + for actor in info_actors: + ray.kill(actor) + + bundle_infos = [(i, gpu_ids[i][0], gpu_ids[i][1]) for i in range(num_bundles)] + pg_reordered_bundle_indices = [bundle_info[0] for bundle_info in sorted(bundle_infos, key=sort_key)] + for i in range(num_bundles): + actual_bundle_index = pg_reordered_bundle_indices[i] + logger.info( + f" bundle {i:4}, actual_bundle_index: {actual_bundle_index:4}, " + f"node: {gpu_ids[actual_bundle_index][0]}, gpu: {gpu_ids[actual_bundle_index][1]}" + ) + + return pg, pg_reordered_bundle_indices + + +def create_placement_groups(args): + """Create placement groups for actor and rollout engines.""" + + num_gpus = 0 + if args.debug_train_only: + num_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node + rollout_offset = 0 + if args.use_critic: + num_gpus += args.critic_num_nodes * args.critic_num_gpus_per_node + critic_offset = args.actor_num_nodes * args.actor_num_gpus_per_node + elif args.debug_rollout_only: + num_gpus = args.rollout_num_gpus + rollout_offset = 0 + elif args.colocate: + num_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node + rollout_offset = 0 + if args.use_critic: + num_gpus += args.critic_num_nodes * args.critic_num_gpus_per_node + critic_offset = args.actor_num_nodes * args.actor_num_gpus_per_node + else: + num_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node + args.rollout_num_gpus + rollout_offset = args.actor_num_nodes * args.actor_num_gpus_per_node + if args.use_critic: + num_gpus += args.critic_num_nodes * args.critic_num_gpus_per_node + critic_offset = args.actor_num_nodes * args.actor_num_gpus_per_node + rollout_offset += args.critic_num_nodes * args.critic_num_gpus_per_node + + logger.info(f"Creating placement group with {num_gpus} GPUs...") + pg, actor_pg_reordered_bundle_indices = _create_placement_group(num_gpus) + + rollout_pg_reordered_bundle_indices = actor_pg_reordered_bundle_indices[rollout_offset:] + if args.use_critic: + critic_pg_reordered_bundle_indices = actor_pg_reordered_bundle_indices[critic_offset:] + + return { + "actor": (pg, actor_pg_reordered_bundle_indices), + "critic": (pg, critic_pg_reordered_bundle_indices) if args.use_critic else None, + "rollout": (pg, rollout_pg_reordered_bundle_indices), + } + + +def allocate_train_group(args, num_nodes, num_gpus_per_node, pg): + return RayTrainGroup( + args=args, + num_nodes=num_nodes, + num_gpus_per_node=num_gpus_per_node, + pg=pg, + num_gpus_per_actor=0.4, + ) + + +def create_training_models(args, pgs, rollout_manager): + actor_model = allocate_train_group( + args=args, + num_nodes=args.actor_num_nodes, + num_gpus_per_node=args.actor_num_gpus_per_node, + pg=pgs["actor"], + ) + if args.use_critic: + critic_model = allocate_train_group( + args=args, + num_nodes=args.critic_num_nodes, + num_gpus_per_node=args.critic_num_gpus_per_node, + pg=pgs["critic"], + ) + critic_init_handle = critic_model.async_init(args, role="critic", with_ref=False) + else: + critic_model = None + + start_rollout_ids = ray.get( + actor_model.async_init(args, role="actor", with_ref=args.kl_coef != 0 or args.use_kl_loss) + ) + + assert len(set(start_rollout_ids)) == 1 + if args.start_rollout_id is None: + args.start_rollout_id = start_rollout_ids[0] + + if args.use_critic: + ray.get(critic_init_handle) + actor_model.connect(critic_model) + + actor_model.set_rollout_manager(rollout_manager) + if args.rollout_global_dataset: + ray.get(rollout_manager.load.remote(args.start_rollout_id - 1)) + + return actor_model, critic_model + + +def create_rollout_manager(args, pg): + rollout_manager = RolloutManager.options( + num_cpus=1, + num_gpus=0, + ).remote(args, pg) + + # calculate num_rollout from num_epoch + num_rollout_per_epoch = None + if args.num_rollout is None: + num_rollout_per_epoch = ray.get(rollout_manager.get_num_rollout_per_epoch.remote()) + args.num_rollout = num_rollout_per_epoch * args.num_epoch + assert args.num_rollout > 0 + + if args.check_weight_update_equal: + ray.get(rollout_manager.check_weights.remote(action="snapshot")) + ray.get(rollout_manager.check_weights.remote(action="reset_tensors")) + + if args.offload_rollout: + ray.get(rollout_manager.offload.remote()) + + return rollout_manager, num_rollout_per_epoch diff --git a/slime/ray/ray_actor.py b/slime/ray/ray_actor.py new file mode 100644 index 0000000000000000000000000000000000000000..d5f1b7248ec6c029339504f87840aa62af4a80bd --- /dev/null +++ b/slime/ray/ray_actor.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from slime.utils.misc import get_current_node_ip, get_free_port + + +class RayActor: + @staticmethod + def _get_current_node_ip_and_free_port(start_port=10000, consecutive=1): + return get_current_node_ip(), get_free_port(start_port=start_port, consecutive=consecutive) + + def get_master_addr_and_port(self): + return self.master_addr, self.master_port diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..6cfc5238e4fccc30542de5bea7c8dde71bff6b2f --- /dev/null +++ b/slime/ray/rollout.py @@ -0,0 +1,688 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +import multiprocessing +import random +import time +from pathlib import Path +from typing import Any + +import numpy as np +import ray +import torch +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +from slime.backends.sglang_utils.sglang_engine import SGLangEngine +from slime.rollout.base_types import call_rollout_fn +from slime.utils import tracking_utils +from slime.utils.health_monitor import RolloutHealthMonitor +from slime.utils.http_utils import _wrap_ipv6, find_available_port, get_host_info, init_http_client +from slime.utils.iter_utils import group_by +from slime.utils.logging_utils import configure_logger +from slime.utils.metric_checker import MetricChecker +from slime.utils.metric_utils import compute_pass_rate, compute_rollout_step, compute_statistics, dict_add_prefix +from slime.utils.misc import load_function +from slime.utils.ray_utils import Box +from slime.utils.seqlen_balancing import get_seqlen_balanced_partitions +from slime.utils.tracking_utils import init_tracking +from slime.utils.types import Sample + +from ..utils.metric_utils import has_repetition +from .utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, Lock + +logging.getLogger("httpx").setLevel(logging.WARNING) +logging.getLogger("httpcore").setLevel(logging.WARNING) + +logger = logging.getLogger(__name__) + + +@ray.remote +class RolloutManager: + """The class to run rollout and convert rollout data to training data.""" + + def __init__(self, args, pg): + configure_logger() + + self.args = args + self.pg = pg + _start_router(args) + # TODO make args immutable + init_tracking(args, primary=False, router_addr=f"http://{args.sglang_router_ip}:{args.sglang_router_port}") + init_http_client(args) + + data_source_cls = load_function(self.args.data_source_path) + self.data_source = data_source_cls(args) + + self.generate_rollout = load_function(self.args.rollout_function_path) + self.eval_generate_rollout = load_function(self.args.eval_function_path) + self.custom_reward_post_process_func = None + if self.args.custom_reward_post_process_path is not None: + self.custom_reward_post_process_func = load_function(self.args.custom_reward_post_process_path) + self.custom_convert_samples_to_train_data_func = None + if self.args.custom_convert_samples_to_train_data_path is not None: + self.custom_convert_samples_to_train_data_func = load_function( + self.args.custom_convert_samples_to_train_data_path + ) + logger.info(f"import {self.args.rollout_function_path} as generate_rollout function.") + logger.info(f"import {self.args.eval_function_path} as eval_generate_rollout function.") + + if self.args.debug_train_only: + self.all_rollout_engines = [] + else: + num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node) + num_engines = args.rollout_num_gpus // num_gpu_per_engine + self.all_rollout_engines = [None] * num_engines + self.num_new_engines = init_rollout_engines(args, pg, self.all_rollout_engines) + self.nodes_per_engine = max(1, args.rollout_num_gpus_per_engine // args.num_gpus_per_node) + self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0).remote() + + self._metric_checker = MetricChecker.maybe_create(args) + if self.args.use_fault_tolerance: + self._health_monitor = RolloutHealthMonitor(self, args) + + def dispose(self): + if self._metric_checker is not None: + self._metric_checker.dispose() + + # TODO maybe rename "rollout_engines" and "all_rollout_engines" later + @property + def rollout_engines(self): + # when doing multi-node serving, we will only send request to node-0 for each engine. + return self.all_rollout_engines[:: self.nodes_per_engine] + + def get_rollout_engines_and_lock(self): + return self.rollout_engines, self.rollout_engine_lock, self.num_new_engines + + def get_num_rollout_per_epoch(self): + assert self.args.rollout_global_dataset + return len(self.data_source.dataset) // self.args.rollout_batch_size + + def generate(self, rollout_id): + monitor_started = self.args.use_fault_tolerance and self._health_monitor.start() + start_time = time.time() + try: + data, metrics = self._get_rollout_data(rollout_id=rollout_id) + self._save_debug_rollout_data(data, rollout_id=rollout_id, evaluation=False) + _log_rollout_data(rollout_id, self.args, data, metrics, time.time() - start_time) + data = self._convert_samples_to_train_data(data) + return self._split_train_data_by_dp(data, self.train_parallel_config["dp_size"]) + finally: + if monitor_started: + self._health_monitor.stop() + self.num_new_engines = init_rollout_engines(self.args, self.pg, self.all_rollout_engines) + else: + self.num_new_engines = 0 + + def eval(self, rollout_id): + if self.args.debug_train_only: + # if debug train only, we don't generate evaluation data + return + + # TODO: add fault tolerance to eval + result = call_rollout_fn(self.eval_generate_rollout, self.args, rollout_id, self.data_source, evaluation=True) + data = result.data + self._save_debug_rollout_data(data, rollout_id=rollout_id, evaluation=True) + metrics = _log_eval_rollout_data(rollout_id, self.args, data, result.metrics) + if self._metric_checker is not None: + self._metric_checker.on_eval(metrics) + + def save(self, rollout_id): + self.data_source.save(rollout_id) + + def load(self, rollout_id=None): + self.data_source.load(rollout_id) + + def offload(self): + return ray.get([engine.release_memory_occupation.remote() for engine in self.rollout_engines]) + + def onload(self, tags: list[str] = None): + return ray.get([engine.resume_memory_occupation.remote(tags=tags) for engine in self.rollout_engines]) + + def check_weights(self, action: str): + return ray.get([engine.check_weights.remote(action=action) for engine in self.rollout_engines]) + + def _get_rollout_data(self, rollout_id): + if self.args.load_debug_rollout_data: + data = torch.load( + open(self.args.load_debug_rollout_data.format(rollout_id=rollout_id), "rb"), + weights_only=False, + )["samples"] + data = [Sample.from_dict(sample) for sample in data] + if (ratio := self.args.load_debug_rollout_data_subsample) is not None: + original_num_rows = len(data) + rough_subsample_num_rows = int(original_num_rows * ratio) + data = data[: rough_subsample_num_rows // 2] + data[-rough_subsample_num_rows // 2 :] + logger.info( + f"Subsample loaded debug rollout data using {ratio=} and change num rows {original_num_rows} -> {len(data)}" + ) + metrics = None + else: + data = call_rollout_fn(self.generate_rollout, self.args, rollout_id, self.data_source, evaluation=False) + metrics = data.metrics + data = data.samples + # flatten the data if it is a list of lists + while isinstance(data[0], list): + data = sum(data, []) + + if self.args.disable_rollout_trim_samples: + logger.info(f"Collectd {len(data)} samples from rollout to train") + elif len(data) % self.args.global_batch_size != 0: + trim_len = (len(data) // self.args.global_batch_size) * self.args.global_batch_size + origin_data_length = len(data) + data = data[:trim_len] + logger.info(f"trim number of samples from {origin_data_length} to {trim_len}") + return data, metrics + + def _save_debug_rollout_data(self, data, rollout_id, evaluation: bool): + # TODO to be refactored (originally Buffer._set_data) + if (path_template := self.args.save_debug_rollout_data) is not None: + path = Path(path_template.format(rollout_id=("eval_" if evaluation else "") + str(rollout_id))) + logger.info(f"Save debug rollout data to {path}") + path.parent.mkdir(parents=True, exist_ok=True) + + # TODO may improve the format + if evaluation: + dump_data = dict( + samples=[sample.to_dict() for dataset_name, info in data.items() for sample in info["samples"]] + ) + else: + dump_data = dict( + samples=[sample.to_dict() for sample in data], + ) + + torch.save(dict(rollout_id=rollout_id, **dump_data), path) + + def _post_process_rewards(self, samples: list[Sample] | list[list[Sample]]): + if self.custom_reward_post_process_func is not None: + return self.custom_reward_post_process_func(self.args, samples) + + raw_rewards = [sample.get_reward_value(self.args) for sample in samples] + if ( + self.args.advantage_estimator in ["grpo", "gspo", "reinforce_plus_plus_baseline"] + and self.args.rewards_normalization + ): + # group norm + rewards = torch.tensor(raw_rewards, dtype=torch.float) + if rewards.shape[-1] == self.args.n_samples_per_prompt * self.args.rollout_batch_size: + rewards = rewards.reshape(-1, self.args.n_samples_per_prompt) + else: + # when samples count are not equal in each group + rewards = rewards.view(-1, rewards.shape[-1]) + mean = rewards.mean(dim=-1, keepdim=True) + rewards = rewards - mean + + if self.args.advantage_estimator in ["grpo", "gspo"] and self.args.grpo_std_normalization: + std = rewards.std(dim=-1, keepdim=True) + rewards = rewards / (std + 1e-6) + + return raw_rewards, rewards.flatten().tolist() + + return raw_rewards, raw_rewards + + def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sample]]): + """ + Convert inference generated samples to training data. + """ + if self.custom_convert_samples_to_train_data_func is not None: + return self.custom_convert_samples_to_train_data_func(self.args, samples) + + raw_rewards, rewards = self._post_process_rewards(samples) + + assert len(raw_rewards) == len(samples) + assert len(rewards) == len(samples) + + train_data = { + "tokens": [sample.tokens for sample in samples], + "response_lengths": [sample.response_length for sample in samples], + # some reward model, e.g. remote rm, may return multiple rewards, + # we could use key to select the reward. + "rewards": rewards, + "raw_reward": raw_rewards, + "truncated": [1 if sample.status == Sample.Status.TRUNCATED else 0 for sample in samples], + "sample_indices": [sample.index for sample in samples], + } + + # loss mask + # TODO: compress the loss mask + loss_masks = [] + for sample in samples: + # always instantiate loss_mask if not provided + if sample.loss_mask is None: + sample.loss_mask = [1] * sample.response_length + + assert ( + len(sample.loss_mask) == sample.response_length + ), f"loss mask length {len(sample.loss_mask)} != response length {sample.response_length}" + if sample.remove_sample: + sample.loss_mask = [0] * sample.response_length + loss_masks.append(sample.loss_mask) + train_data["loss_masks"] = loss_masks + + # overwriting the raw reward + if samples[0].metadata and "raw_reward" in samples[0].metadata: + train_data["raw_reward"] = [sample.metadata["raw_reward"] for sample in samples] + + # For rollout buffer + if samples[0].metadata and "round_number" in samples[0].metadata: + train_data["round_number"] = [sample.metadata["round_number"] for sample in samples] + + # Add rollout log probabilities for off-policy correction. + if all(s.rollout_log_probs is not None for s in samples): + train_data["rollout_log_probs"] = [sample.rollout_log_probs for sample in samples] + + if all(s.rollout_routed_experts is not None for s in samples): + train_data["rollout_routed_experts"] = [sample.rollout_routed_experts for sample in samples] + + if all(s.train_metadata is not None for s in samples): + train_data["metadata"] = [sample.train_metadata for sample in samples] + + if all(s.multimodal_train_inputs is not None for s in samples): + train_data["multimodal_train_inputs"] = [sample.multimodal_train_inputs for sample in samples] + + if "teacher_log_probs" in samples[0].__dict__: + train_data["teacher_log_probs"] = [sample.teacher_log_probs for sample in samples] + + if "verifiable_rewards" in samples[0].__dict__: + train_data["verifiable_rewards"] = [ + getattr(sample, "verifiable_rewards", None) for sample in samples + ] + + return train_data + + def set_train_parallel_config(self, config: dict): + self.train_parallel_config = config + + def _split_train_data_by_dp(self, data, dp_size): + """Split the train data by data parallel size.""" + rollout_data = {} + + if "prompt" in data: + rollout_data["prompt"] = data["prompt"] + + total_lengths = [len(t) for t in data["tokens"]] + data["total_lengths"] = total_lengths + + if self.args.balance_data: + partitions = get_seqlen_balanced_partitions(total_lengths, dp_size, equal_size=True) + else: + partitions = [range(i, len(total_lengths), dp_size) for i in range(dp_size)] + + rollout_data_refs = [] + + for i in range(dp_size): + rollout_data = {} + partition = partitions[i] + rollout_data["partition"] = partition + for key in [ + "tokens", + "multimodal_train_inputs", + "response_lengths", + "rewards", + "truncated", + "loss_masks", + "round_number", + "sample_indices", + "rollout_log_probs", + "rollout_routed_experts", + "prompt", + "teacher_log_probs", + "verifiable_rewards", + ]: + if key not in data: + continue + val = [data[key][j] for j in partition] + rollout_data[key] = val + # keys that need to be splited at train side + for key in [ + "raw_reward", + "total_lengths", + ]: + if key not in data: + continue + rollout_data[key] = data[key] + rollout_data_refs.append(Box(ray.put(rollout_data))) + return rollout_data_refs + + +def init_rollout_engines(args, pg, all_rollout_engines): + if args.debug_train_only: + return 0 + + num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node) + num_engines = args.rollout_num_gpus // num_gpu_per_engine + assert len(all_rollout_engines) == num_engines + if args.prefill_num_servers is not None: + prefill_num_servers = args.prefill_num_servers * args.rollout_num_gpus_per_engine // num_gpu_per_engine + assert ( + num_engines > prefill_num_servers + ), f"num_engines {num_engines} should be larger than prefill_num_servers {prefill_num_servers}" + + pg, reordered_bundle_indices = pg + + RolloutRayActor = ray.remote(SGLangEngine) + + rollout_engines = [] + for i in range(num_engines): + if all_rollout_engines[i] is not None: + continue + + num_gpus = 0.2 + num_cpus = num_gpus + + scheduling_strategy = PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_capture_child_tasks=True, + placement_group_bundle_index=reordered_bundle_indices[i * num_gpu_per_engine], + ) + + env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} | { + "SGL_JIT_DEEPGEMM_PRECOMPILE": "false", + "SGLANG_JIT_DEEPGEMM_PRECOMPILE": "false", + "SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK": "true", + "SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK": "true", + "SGLANG_MEMORY_SAVER_CUDA_GRAPH": "true", + "SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_FALLBACK_VARIANT": "true", + "SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION": "false", + } + + worker_type = "regular" + if args.prefill_num_servers is not None: + if i < prefill_num_servers: + worker_type = "prefill" + else: + worker_type = "decode" + + rollout_engine = RolloutRayActor.options( + num_cpus=num_cpus, + num_gpus=num_gpus, + scheduling_strategy=scheduling_strategy, + runtime_env={ + "env_vars": env_vars, + }, + ).remote(args, rank=i, worker_type=worker_type) + + rollout_engines.append((i, rollout_engine)) + all_rollout_engines[i] = rollout_engine + + num_new_engines = len(rollout_engines) + + if num_new_engines == 0: + return num_new_engines + + if args.rollout_external: + addr_and_ports = _allocate_rollout_engine_addr_and_ports_external(args=args, rollout_engines=rollout_engines) + else: + addr_and_ports = _allocate_rollout_engine_addr_and_ports_normal( + args=args, num_engines=num_engines, rollout_engines=rollout_engines + ) + + # TODO: don't ray.get here to overlap train actor init with rollout engine init. + # somehow if we don't sync here, the --debug-rollout-only mode will crash. + init_handles = [engine.init.remote(**(addr_and_ports[rank])) for rank, engine in rollout_engines] + ray.get(init_handles) + + return num_new_engines + + +def _allocate_rollout_engine_addr_and_ports_external(args, rollout_engines): + addr_and_ports = [] + for rank, _ in rollout_engines: + [host, port] = args.rollout_external_engine_addrs[rank].split(":") + addr_and_ports.append( + dict( + dist_init_addr=None, + nccl_port=None, + host=host, + port=int(port), + ) + ) + return addr_and_ports + + +def _allocate_rollout_engine_addr_and_ports_normal(*, args, num_engines, rollout_engines): + # get ports + # there are 4 ports we need to allocate + # 1. server port + # 2. nccl port + # 3. dist_init_addr port + # 4. other ports for dp_attention, which is of size 4 + dp_size + num_engines_per_node = max( + 1, min(args.num_gpus_per_node, args.rollout_num_gpus) // args.rollout_num_gpus_per_engine + ) + addr_and_ports = [{} for _ in range(num_engines)] + + # Calculate prefill limit to identify prefill engines + prefill_limit = 0 + if args.prefill_num_servers is not None: + num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node) + prefill_limit = args.prefill_num_servers * args.rollout_num_gpus_per_engine // num_gpu_per_engine + + visited_nodes = set() + for rank, engine in rollout_engines: + if rank // num_engines_per_node in visited_nodes: + continue + visited_nodes.add(rank // num_engines_per_node) + # TODO: currently when restarting engines, we will set port for all engines on this node starting with this rank. + # e.g. for 8 gpus, if we are restarting engine on gpu 3, we will set port for engine 3,4,5,6,7 on this node. + num_engines_on_this_node = num_engines_per_node - (rank % num_engines_per_node) + + def get_addr_and_ports(engine): + # use small ports to prevent ephemeral port between 32768 and 65536. + # also, ray uses port 10002-19999, thus we avoid near-10002 to avoid racing condition + start_port = 15000 + + def port(consecutive=1): + nonlocal start_port + _, port = ray.get( + engine._get_current_node_ip_and_free_port.remote( + start_port=start_port, + consecutive=consecutive, + ) + ) + start_port = port + consecutive + return port + + def addr(): + addr, _ = ray.get(engine._get_current_node_ip_and_free_port.remote()) + return addr + + return addr, port + + get_addr, get_port = get_addr_and_ports(engine) + + for i in range(num_engines_on_this_node): + current_rank = rank + i + addr_and_ports[current_rank]["host"] = get_addr() + addr_and_ports[current_rank]["port"] = get_port() + addr_and_ports[current_rank]["nccl_port"] = get_port() + + if args.prefill_num_servers is not None and current_rank < prefill_limit: + addr_and_ports[current_rank]["disaggregation_bootstrap_port"] = get_port() + + if args.rollout_num_gpus_per_engine > args.num_gpus_per_node: + num_node_per_engine = args.rollout_num_gpus_per_engine // args.num_gpus_per_node + if rank % num_node_per_engine == 0: + # this is the first node in the engine, we need to allocate the dist_init_addr port + dist_init_addr = f"{get_addr()}:{get_port(30 + args.sglang_dp_size)}" + for i in range(num_node_per_engine): + addr_and_ports[rank + i]["dist_init_addr"] = dist_init_addr + else: + for i in range(num_engines_on_this_node): + addr_and_ports[rank + i]["dist_init_addr"] = f"{get_addr()}:{get_port(30 + args.sglang_dp_size)}" + + for i, _ in rollout_engines: + for key in ["port", "nccl_port", "dist_init_addr"]: + assert key in addr_and_ports[i], f"Engine {i} {key} is not set." + logger.info(f"Ports for engine {i}: {addr_and_ports[i]}") + + return addr_and_ports + + +def _start_router(args): + """start sgl router and slime router""" + if not args.rollout_num_gpus: + # No rollout engines (e.g. Lightning OPD) — skip router entirely. + args.sglang_router_ip = args.sglang_router_ip or "127.0.0.1" + args.sglang_router_port = args.sglang_router_port or 0 + return + if args.sglang_router_ip is not None: + return + + args.sglang_router_ip = _wrap_ipv6(get_host_info()[1]) + if args.sglang_router_port is None: + args.sglang_router_port = find_available_port(random.randint(3000, 4000)) + + if args.use_slime_router: + assert args.prefill_num_servers is None, "slime router does not support prefill_num_servers." + from slime.router.router import run_router + + router_args = args + + else: + from sglang_router.launch_router import RouterArgs + + from slime.utils.http_utils import run_router + + router_args = RouterArgs.from_cli_args(args, use_router_prefix=True) + router_args.host = args.sglang_router_ip + router_args.port = args.sglang_router_port + router_args.prometheus_port = find_available_port(random.randint(4000, 5000)) + router_args.log_level = "warn" + + if args.prefill_num_servers is not None: + router_args.pd_disaggregation = True + + if hasattr(router_args, "request_timeout_secs"): + router_args.request_timeout_secs = args.sglang_router_request_timeout_secs + + logger.info(f"Launch router with args: {router_args}") + + process = multiprocessing.Process( + target=run_router, + args=(router_args,), + ) + process.daemon = True # Set the process as a daemon + process.start() + # Wait 3 seconds + time.sleep(3) + assert process.is_alive() + logger.info(f"Router launched at {args.sglang_router_ip}:{args.sglang_router_port}") + + +def _log_eval_rollout_data(rollout_id, args, data, extra_metrics: dict[str, Any] | None = None): + if args.custom_eval_rollout_log_function_path is not None: + custom_log_func = load_function(args.custom_eval_rollout_log_function_path) + if custom_log_func(rollout_id, args, data, extra_metrics): + return + + log_dict = extra_metrics or {} + for key in data.keys(): + rewards = data[key]["rewards"] + log_dict[f"eval/{key}"] = sum(rewards) / len(rewards) + if (samples := data[key].get("samples")) is not None: + log_dict |= dict_add_prefix(compute_metrics_from_samples(args, samples), f"eval/{key}/") + if "truncated" in data[key]: + truncated = data[key]["truncated"] + log_dict[f"eval/{key}-truncated_ratio"] = sum(truncated) / len(truncated) + if args.log_passrate: + log_dict |= dict_add_prefix( + compute_pass_rate( + flat_rewards=rewards, + group_size=args.n_samples_per_eval_prompt, + ), + f"eval/{key}-", + ) + + logger.info(f"eval {rollout_id}: {log_dict}") + + step = compute_rollout_step(args, rollout_id) + log_dict["eval/step"] = step + tracking_utils.log(args, log_dict, step_key="eval/step") + + return log_dict + + +def _log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time): + if args.custom_rollout_log_function_path is not None: + custom_log_func = load_function(args.custom_rollout_log_function_path) + if custom_log_func(rollout_id, args, samples, rollout_extra_metrics, rollout_time): + return + + if args.load_debug_rollout_data: + return + + log_dict = {**(rollout_extra_metrics or {})} + response_lengths = [sample.effective_response_length for sample in samples] + log_dict["perf/rollout_time"] = rollout_time + if args.rollout_num_gpus: + log_dict["perf/tokens_per_gpu_per_sec"] = sum(response_lengths) / rollout_time / args.rollout_num_gpus + log_dict["perf/longest_sample_tokens_per_sec"] = max(response_lengths) / rollout_time + log_dict |= dict_add_prefix(compute_metrics_from_samples(args, samples), "rollout/") + logger.info(f"perf {rollout_id}: {log_dict}") + step = compute_rollout_step(args, rollout_id) + log_dict["rollout/step"] = step + tracking_utils.log(args, log_dict, step_key="rollout/step") + + +def compute_metrics_from_samples(args, samples): + response_lengths = [sample.effective_response_length for sample in samples] + + log_dict = {} + log_dict |= dict_add_prefix(compute_statistics(response_lengths), "response_len/") + log_dict |= _compute_zero_std_metrics(args, samples) + log_dict |= _compute_spec_metrics(args, samples) + log_dict |= _compute_reward_cat_metrics(args, samples) + log_dict["repetition_frac"] = np.mean([int(has_repetition(s.response)) for s in samples]).item() + log_dict["truncated_ratio"] = np.mean([int(s.status == Sample.Status.TRUNCATED) for s in samples]).item() + return log_dict + + +def _compute_zero_std_metrics(args, all_samples: list[Sample]): + # only compute in GRPO-like algorithms where one prompt has multiple responses + if args.advantage_estimator == "ppo": + return {} + + def _is_zero_std(samples: list[Sample]): + rewards = [sample.get_reward_value(args) for sample in samples] + return len(rewards) == 0 or all(rewards[0] == r for r in rewards) + + all_sample_groups = group_by(all_samples, lambda s: s.group_index) + interesting_sample_groups = [g for g in all_sample_groups.values() if _is_zero_std(g)] + + def _format_reward(reward): + # Handle dict rewards (from RL samples with meta_info) + if isinstance(reward, dict): + return "dict" + try: + return str(round(reward, 1)) + except (TypeError, ValueError): + return str(reward) + + interesting_rewards = [_format_reward(g[0].get_reward_value(args)) for g in interesting_sample_groups] + + return {f"zero_std/count_{reward}": len(items) for reward, items in group_by(interesting_rewards).items()} + + +def _compute_spec_metrics(args, all_samples: list[Sample]): + if args.sglang_speculative_algorithm is None: + return {} + num_samples = len(all_samples) + metrics = {} + metrics["rollout/spec_accept_rate"] = ( + sum(sample.spec_info.spec_accept_rate for sample in all_samples) / num_samples + ) + metrics["rollout/spec_accept_length"] = ( + sum(sample.spec_info.spec_accept_length for sample in all_samples) / num_samples + ) + return metrics + + +def _compute_reward_cat_metrics(args, all_samples: list[Sample]): + reward_cat_key = args.log_reward_category + if reward_cat_key is None: + return {} + + samples_of_reward_cat = group_by(all_samples, lambda s: s.reward[reward_cat_key]) + + return {f"error_cat/{reward_cat}": len(s) / len(all_samples) for reward_cat, s in samples_of_reward_cat.items()} diff --git a/slime/ray/train_actor.py b/slime/ray/train_actor.py new file mode 100644 index 0000000000000000000000000000000000000000..799db4aba3270c2293ca1abe820ae78e660277b9 --- /dev/null +++ b/slime/ray/train_actor.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import abc +import logging +import os +import random +from datetime import timedelta + +import ray +import torch +import torch.distributed as dist + +import slime.utils.eval_config +from slime.ray.ray_actor import RayActor +from slime.utils.distributed_utils import init_gloo_group +from slime.utils.logging_utils import configure_logger +from slime.utils.memory_utils import clear_memory, print_memory + +logger = logging.getLogger(__name__) + + +def get_local_gpu_id(): + cvd = os.environ.get("CUDA_VISIBLE_DEVICES", None) + if cvd is None: + return ray.get_gpu_ids()[0] + else: + return cvd.split(",").index(str(ray.get_gpu_ids()[0])) + + +class TrainRayActor(RayActor): + def __init__(self, world_size, rank, master_addr, master_port): + configure_logger() + + self._world_size = world_size + self._rank = rank + if master_addr: + self.master_addr, self.master_port = master_addr, master_port + else: + self.master_addr, self.master_port = self._get_current_node_ip_and_free_port( + start_port=random.randint(20000, 21000) + ) + + os.environ["MASTER_ADDR"] = self.master_addr + os.environ["MASTER_PORT"] = str(self.master_port) + os.environ["WORLD_SIZE"] = str(self._world_size) + os.environ["RANK"] = str(self._rank) + # TODO: currently this doesn't work as ray has already set torch.cuda.device_count(). + # os.environ.pop("CUDA_VISIBLE_DEVICES", None) + # os.environ["LOCAL_RANK"] = str(ray.get_gpu_ids()[0]) + os.environ["LOCAL_RANK"] = str(get_local_gpu_id()) + + def init(self, args, role, with_ref=False): + self.args = args + self.role = role + self.with_ref = with_ref + + torch.serialization.add_safe_globals([slime.utils.eval_config.EvalDatasetConfig]) + + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(f"cuda:{local_rank}") + + # Use hybrid backend when FSDP CPU offload is enabled with a CPU backend + backend = args.distributed_backend + if getattr(args, "fsdp_cpu_offload", False) and getattr(args, "fsdp_cpu_backend", None): + cpu_backend = args.fsdp_cpu_backend + backend = f"cpu:{cpu_backend},cuda:{args.distributed_backend}" + logger.info(f"FSDP CPU offload enabled, using hybrid backend: {backend}") + + dist.init_process_group( + backend=backend, + timeout=timedelta(minutes=args.distributed_timeout_minutes), + ) + init_gloo_group() + + args.rank = dist.get_rank() + args.world_size = dist.get_world_size() + + try: + if torch.version.hip is not None: + logger.info("Detected ROCm/HIP environment, skipping NUMA affinity setup") + # will find the coresponding API to implement ROCm version as below + else: + import pynvml + + pynvml.nvmlInit() + + local_rank = int(os.environ["RANK"]) % args.num_gpus_per_node + + handle = pynvml.nvmlDeviceGetHandleByIndex(local_rank) + pynvml.nvmlDeviceSetCpuAffinity(handle) + + logger.info(f"Set NUMA affinity for GPU {local_rank}") + pynvml.nvmlShutdown() + + except ImportError: + logger.info("Warning: pynvml not available, skipping NUMA affinity setup") + except Exception as e: + logger.info(f"Warning: Failed to set NUMA affinity: {e}") + + def clear_memory(self): + print_memory("before TrainRayActor.clear_memory") + clear_memory() + print_memory("after TrainRayActor.clear_memory") + + @abc.abstractmethod + def sleep(self, tags): + raise NotImplementedError + + @abc.abstractmethod + def wake_up(self, tags): + raise NotImplementedError + + @abc.abstractmethod + def train(self, rollout_id, rollout_data_ref): + raise NotImplementedError + + @abc.abstractmethod + def save_model(self, rollout_id, force_sync=False): + raise NotImplementedError + + @abc.abstractmethod + def update_weights(self): + raise NotImplementedError + + @abc.abstractmethod + def connect_actor_critic(self, critic_group): + raise NotImplementedError + + @abc.abstractmethod + def _get_parallel_config(self): + raise NotImplementedError + + def set_rollout_manager(self, rollout_manager): + self.rollout_manager = rollout_manager + if self.args.rank == 0: + ray.get(self.rollout_manager.set_train_parallel_config.remote(self.train_parallel_config)) diff --git a/slime/ray/utils.py b/slime/ray/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..39ea8c88abebbeeda80a737d1ad89b3aa96a1a94 --- /dev/null +++ b/slime/ray/utils.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Adapted from https://github.com/OpenRLHF/OpenRLHF/blob/10c733694ed9fbb78a0a2ff6a05efc7401584d46/openrlhf/trainer/ray/utils.py#L1 +import os + +import ray +import torch +from slime.ray.ray_actor import RayActor + + +# Refer to +# https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/nvidia_gpu.py#L95-L96 +# https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/amd_gpu.py#L102-L103 +# https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/npu.py#L94-L95 +# https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/hpu.py#L116-L117 +# https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/neuron.py#L108-L109 +# https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/tpu.py#L171-L172 +# https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/intel_gpu.py#L97-L98 +NOSET_VISIBLE_DEVICES_ENV_VARS_LIST = [ + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES", + "RAY_EXPERIMENTAL_NOSET_NEURON_RT_VISIBLE_CORES", + "RAY_EXPERIMENTAL_NOSET_TPU_VISIBLE_CHIPS", + "RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR", +] + + +def ray_noset_visible_devices(env_vars=os.environ): + return any(env_vars.get(env_var) for env_var in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST) + + +def get_physical_gpu_id(): + device = torch.cuda.current_device() + props = torch.cuda.get_device_properties(device) + return str(props.uuid) + + +@ray.remote +class Lock(RayActor): + def __init__(self): + self._locked = False # False: unlocked, True: locked + + def acquire(self): + """ + Try to acquire the lock. Returns True if acquired, False otherwise. + Caller should retry until it returns True. + """ + if not self._locked: + self._locked = True + return True + return False + + def release(self): + """Release the lock, allowing others to acquire.""" + assert self._locked, "Lock is not acquired, cannot release." + self._locked = False diff --git a/slime/rollout/__init__.py b/slime/rollout/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a4d636226b9a4b5130dfd8e94f2cece5fd290e53 --- /dev/null +++ b/slime/rollout/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + diff --git a/slime/rollout/__pycache__/__init__.cpython-312.pyc b/slime/rollout/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1eb495f2a7e894261541e17f51b3282d890dd571 Binary files /dev/null and b/slime/rollout/__pycache__/__init__.cpython-312.pyc differ diff --git a/slime/rollout/__pycache__/base_types.cpython-312.pyc b/slime/rollout/__pycache__/base_types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..67bd64ac1087efbca086ef83dd7599af0cda0e1b Binary files /dev/null and b/slime/rollout/__pycache__/base_types.cpython-312.pyc differ diff --git a/slime/rollout/__pycache__/data_source.cpython-312.pyc b/slime/rollout/__pycache__/data_source.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6cbaec49412c5a05e45156b9257544f5962d3b2b Binary files /dev/null and b/slime/rollout/__pycache__/data_source.cpython-312.pyc differ diff --git a/slime/rollout/__pycache__/on_policy_distillation.cpython-312.pyc b/slime/rollout/__pycache__/on_policy_distillation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..195c79ddf3d943fea27c5abd284afa4b9296d4c0 Binary files /dev/null and b/slime/rollout/__pycache__/on_policy_distillation.cpython-312.pyc differ diff --git a/slime/rollout/__pycache__/sglang_rollout.cpython-312.pyc b/slime/rollout/__pycache__/sglang_rollout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ef9a5ef3ce9492c9013a71a45424a37cf2981f8 Binary files /dev/null and b/slime/rollout/__pycache__/sglang_rollout.cpython-312.pyc differ diff --git a/slime/rollout/base_types.py b/slime/rollout/base_types.py new file mode 100644 index 0000000000000000000000000000000000000000..ce8c21c0da9a811b0dd59abb2c3b47d20cf08f77 --- /dev/null +++ b/slime/rollout/base_types.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass +from typing import Any + +from slime.utils.types import Sample + + +@dataclass +class RolloutFnTrainOutput: + samples: list[list[Sample]] + metrics: dict[str, Any] = None + + +@dataclass +class RolloutFnEvalOutput: + data: dict[str, dict[str, Any]] + metrics: dict[str, Any] = None + + +def call_rollout_fn(fn, *args, evaluation: bool, **kwargs): + output = fn(*args, **kwargs, evaluation=evaluation) + + # compatibility for legacy version + if not isinstance(output, (RolloutFnTrainOutput, RolloutFnEvalOutput)): + output = RolloutFnEvalOutput(data=output) if evaluation else RolloutFnTrainOutput(samples=output) + + return output diff --git a/slime/rollout/data_source.py b/slime/rollout/data_source.py new file mode 100644 index 0000000000000000000000000000000000000000..fddcc60788b9b72c9984e4ea60c5119e2140c09c --- /dev/null +++ b/slime/rollout/data_source.py @@ -0,0 +1,230 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import abc +import copy +import logging +import os +from pathlib import Path + +import torch + +from slime.utils.data import create_dataset +from slime.utils.misc import load_function +from slime.utils.processing_utils import load_processor, load_tokenizer +from slime.utils.types import Sample + +logger = logging.getLogger(__name__) + + +class DataSource(abc.ABC): + @abc.abstractmethod + def get_samples(self, num_samples: int) -> list[list[Sample]]: + """ + Return num_samples samples + """ + + @abc.abstractmethod + def add_samples(self, samples: list[list[Sample]]): + """ + Add samples to the data source + """ + + @abc.abstractmethod + def save(self, rollout_id): + """ + Save the state of the data source + """ + + @abc.abstractmethod + def load(self, rollout_id=None): + """ + Load the state of the data source + """ + + +# TODO may further refactor data-loading part later +class RolloutDataSource(DataSource): + def __init__(self, args): + self.args = args + + self.epoch_id = 0 + self.sample_group_index = 0 + self.sample_index = 0 + self.sample_offset = 0 + # TODO remove this + self.metadata = {} + + if args.rollout_global_dataset: + tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) + processor = load_processor(args.hf_checkpoint, trust_remote_code=True) + + # TODO move (during the refactor) + if (d := args.dump_details) is not None: + tokenizer.save_pretrained(Path(d) / "tokenizer") + if processor: + processor.save_pretrained(Path(d) / "processor") + + self.dataset = create_dataset( + args.prompt_data, + tokenizer=tokenizer, + processor=processor, + max_length=args.rollout_max_prompt_len, + prompt_key=args.input_key, + multimodal_keys=args.multimodal_keys, + label_key=args.label_key, + metadata_key=args.metadata_key, + tool_key=args.tool_key, + apply_chat_template=args.apply_chat_template, + apply_chat_template_kwargs=args.apply_chat_template_kwargs, + seed=args.rollout_seed, + ) + if self.args.rollout_shuffle: + self.dataset.shuffle(self.epoch_id) + else: + self.dataset = None + + def get_samples(self, num_samples): + # TODO further improve code + if self.dataset is not None: + if self.sample_offset + num_samples <= len(self.dataset): + prompt_samples = self.dataset.samples[self.sample_offset : self.sample_offset + num_samples] + self.sample_offset += num_samples + else: + prompt_samples = self.dataset.samples[self.sample_offset :] + num_samples -= len(prompt_samples) + self.epoch_id += 1 + if self.args.rollout_shuffle: + self.dataset.shuffle(self.epoch_id) + prompt_samples += self.dataset.samples[:num_samples] + self.sample_offset = num_samples + else: + prompt_samples = [Sample() for _ in range(num_samples)] + + samples = [] + for prompt_sample in prompt_samples: + group = [] + for _ in range(self.args.n_samples_per_prompt): + sample = copy.deepcopy(prompt_sample) + sample.group_index = self.sample_group_index + sample.index = self.sample_index + self.sample_index += 1 + group.append(sample) + self.sample_group_index += 1 + samples.append(group) + return samples + + def add_samples(self, samples: list[list[Sample]]): + raise RuntimeError(f"Cannot add samples to {self.__class__.__name__}. This is a read-only data source.") + + def save(self, rollout_id): + if not self.args.rollout_global_dataset: + return + + state_dict = { + "sample_offset": self.sample_offset, + "epoch_id": self.epoch_id, + "sample_group_index": self.sample_group_index, + "sample_index": self.sample_index, + "metadata": self.metadata, + # Save wandb_run_id for resume support + "wandb_run_id": getattr(self.args, "wandb_run_id", None), + } + path = os.path.join(self.args.save, f"rollout/global_dataset_state_dict_{rollout_id}.pt") + os.makedirs(os.path.dirname(path), exist_ok=True) + torch.save(state_dict, path) + + def load(self, rollout_id=None): + if not self.args.rollout_global_dataset: + return + + if self.args.load is None: + return + + path = os.path.join(self.args.load, f"rollout/global_dataset_state_dict_{rollout_id}.pt") + if not os.path.exists(path): + logger.info(f"Checkpoint {path} does not exist.") + return + + logger.info(f"load metadata from {path}") + logger.info(f"load metadata: {self.metadata}") + state_dict = torch.load(path) + self.sample_offset = state_dict.get("sample_offset", 0) + self.epoch_id = state_dict.get("epoch_id", 0) + self.sample_group_index = state_dict.get("sample_group_index", 0) + self.sample_index = state_dict.get("sample_index", 0) + self.metadata = state_dict.get("metadata", {}) + + # Load wandb_run_id for resume support (only if not already set) + if not getattr(self.args, "wandb_run_id", None): + loaded_wandb_run_id = state_dict.get("wandb_run_id") + if loaded_wandb_run_id: + self.args.wandb_run_id = loaded_wandb_run_id + logger.info(f"Loaded wandb_run_id from checkpoint: {loaded_wandb_run_id}") + + if self.args.rollout_global_dataset and self.args.rollout_shuffle: + self.dataset.shuffle(self.epoch_id) + + +class RolloutDataSourceWithBuffer(RolloutDataSource): + def __init__(self, args): + super().__init__(args) + self.buffer = [] + if self.args.buffer_filter_path is None: + self.buffer_filter = pop_first + else: + self.buffer_filter = load_function(self.args.buffer_filter_path) + + def get_samples(self, num_samples: int) -> list[list[Sample]]: + """ + Return num_samples samples + """ + + samples = self._get_samples_from_buffer(num_samples) + num_samples -= len(samples) + + if num_samples == 0: + return samples + + samples += super().get_samples(num_samples=num_samples) + return samples + + def _get_samples_from_buffer(self, num_samples: int) -> list[list[Sample]]: + if len(self.buffer) == 0 or num_samples == 0: + return [] + + samples = self.buffer_filter(self.args, None, self.buffer, num_samples) + return samples + + def add_samples(self, samples: list[list[Sample]]): + """ + Add a sample group to buffer. + """ + if not samples: + return + assert isinstance(samples, list), f"samples must be a list, got {type(samples)}" + assert isinstance(samples[0], list), f"the elements of samples must be list, got {type(samples[0])}" + for i in range(0, len(samples)): + assert ( + len(samples[i]) == self.args.n_samples_per_prompt + ), f"the length of the elements of samples must be equal to n_samples_per_prompt, got {len(samples[i])} != {self.args.n_samples_per_prompt}" + group = samples[i] # type: ignore + self.buffer.append(group) + + # TODO remove + def update_metadata(self, metadata: dict): + self.metadata.update(metadata) + + # TODO remove + def get_metadata(self): + return self.metadata + + def get_buffer_length(self): + return len(self.buffer) + + +def pop_first(args, rollout_id, buffer: list[list[Sample]], num_samples: int) -> list[list[Sample]]: + num_to_pop = min(len(buffer), num_samples) + samples = buffer[:num_to_pop] + del buffer[:num_to_pop] + return samples diff --git a/slime/rollout/filter_hub/__init__.py b/slime/rollout/filter_hub/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a4d636226b9a4b5130dfd8e94f2cece5fd290e53 --- /dev/null +++ b/slime/rollout/filter_hub/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + diff --git a/slime/rollout/filter_hub/__pycache__/__init__.cpython-312.pyc b/slime/rollout/filter_hub/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..98412345a413d7ca85d7fd465101984404ee5418 Binary files /dev/null and b/slime/rollout/filter_hub/__pycache__/__init__.cpython-312.pyc differ diff --git a/slime/rollout/filter_hub/__pycache__/base_types.cpython-312.pyc b/slime/rollout/filter_hub/__pycache__/base_types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b7c6df2189553bb0dce6c7f7ae1112feb51d4b83 Binary files /dev/null and b/slime/rollout/filter_hub/__pycache__/base_types.cpython-312.pyc differ diff --git a/slime/rollout/filter_hub/base_types.py b/slime/rollout/filter_hub/base_types.py new file mode 100644 index 0000000000000000000000000000000000000000..cf49784a30ce653528370ceb1b23b1aa295a2ff4 --- /dev/null +++ b/slime/rollout/filter_hub/base_types.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass + + +@dataclass +class DynamicFilterOutput: + keep: bool + reason: str | None = None diff --git a/slime/rollout/filter_hub/dynamic_sampling_filters.py b/slime/rollout/filter_hub/dynamic_sampling_filters.py new file mode 100644 index 0000000000000000000000000000000000000000..aff4d66c2b232faf73ee5a19bfc217af27732f10 --- /dev/null +++ b/slime/rollout/filter_hub/dynamic_sampling_filters.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from slime.rollout.filter_hub.base_types import DynamicFilterOutput +from slime.utils.types import Sample + +__all__ = ["check_reward_nonzero_std"] + + +def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): + rewards = [sample.get_reward_value(args) for sample in samples] + keep = torch.tensor(rewards, dtype=torch.float).std() > 0.0 + return DynamicFilterOutput( + keep=keep, + reason=None if keep else f"zero_std_{round(rewards[0], 1)}", + ) diff --git a/slime/rollout/on_policy_distillation.py b/slime/rollout/on_policy_distillation.py new file mode 100644 index 0000000000000000000000000000000000000000..3cd92c4aa75a057371abc7b5cf68c837888fba89 --- /dev/null +++ b/slime/rollout/on_policy_distillation.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import aiohttp +import torch + +from slime.utils.processing_utils import encode_image_for_rollout_engine +from slime.utils.types import Sample + + +async def reward_func(args, sample, **kwargs): + # For Lightning OPD: teacher log-probs are pre-computed in metadata, + # no teacher server call needed. Return a sentinel so post_process_rewards knows. + metadata = sample.metadata or {} + if metadata.get("is_lightning_opd", False) or metadata.get("is_offline_opd", False): + return {"lightning_opd": True} + + payload = { + "input_ids": sample.tokens, + "sampling_params": { + "temperature": 0, + "max_new_tokens": 0, + "skip_special_tokens": False, + }, + "return_logprob": True, + "logprob_start_len": 0, + } + + if sample.multimodal_inputs and sample.multimodal_inputs.get("images"): + image_data = sample.multimodal_inputs["images"] + payload["image_data"] = [encode_image_for_rollout_engine(image) for image in image_data] + + session_kwargs = {} + async with aiohttp.ClientSession(**session_kwargs) as session: + async with session.post(args.rm_url, json=payload) as resp: + resp.raise_for_status() + return await resp.json() + + +def post_process_rewards(args, samples: list[Sample], **kwargs): + """Process rewards from teacher model and extract teacher log probabilities. + + This function: + 1. Extracts teacher log-probs from the reward response (which contains sglang's logprob output) + 2. Trims them to match the response length + 3. Stores them in sample.teacher_log_probs for OPD KL penalty computation + 4. Returns scalar rewards (0.0 for pure distillation) compatible with GRPO/PPO + + For Lightning OPD, teacher log-probs are pre-computed in the parquet + metadata instead of being fetched from a teacher server at runtime. + """ + raw_rewards = [sample.get_reward_value(args) for sample in samples] + response_lengths = [sample.response_length for sample in samples] + + for i, (sample, reward) in enumerate(zip(samples, raw_rewards)): + metadata = sample.metadata or {} + if isinstance(reward, dict) and reward.get("lightning_opd"): + # Lightning OPD: teacher log-probs are pre-computed in metadata + pre_teacher_lp = metadata.get("teacher_log_probs", []) + sample.teacher_log_probs = torch.tensor( + [float(x) for x in pre_teacher_lp], dtype=torch.float32 + ) + else: + # Online OPD: extract teacher log-probs from sglang response + t_log_probs = torch.tensor( + [item[0] for item in reward["meta_info"]["input_token_logprobs"][1:]], + dtype=torch.float32, + ) + sample.teacher_log_probs = t_log_probs[-response_lengths[i]:] + + scalar_rewards = [0.0] * len(samples) + return scalar_rewards, scalar_rewards \ No newline at end of file diff --git a/slime/rollout/rm_hub/__init__.py b/slime/rollout/rm_hub/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8235aa29c56a5834c970895c74370496765c29d2 --- /dev/null +++ b/slime/rollout/rm_hub/__init__.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import random + +import aiohttp + +from slime.utils.misc import load_function +from slime.utils.types import Sample + +from .deepscaler import get_deepscaler_rule_based_reward +from .f1 import f1_score +from .gpqa import compute_gpqa_reward +from .math_dapo_utils import compute_score as compute_score_dapo +from .math_utils import extract_answer as extract_boxed_answer +from .math_utils import grade_answer_verl + + +async def remote_rm(args, sample: Sample): + payload = { + "prompt": sample.prompt, + "response": sample.response, + "label": sample.label, + } + session_kwargs = {} + async with aiohttp.ClientSession(**session_kwargs) as session: + async with session.post(args.rm_url, json=payload) as resp: + resp.raise_for_status() + return await resp.json() + + +async def async_rm(args, sample: Sample, **kwargs): + if args.custom_rm_path is not None: + rm_function = load_function(args.custom_rm_path) + return await rm_function(args, sample, **kwargs) + + metadata = sample.metadata if isinstance(sample.metadata, dict) else {} + rm_type = (metadata.get("rm_type") or args.rm_type or "").strip() + response = sample.response + label = sample.label + if rm_type.startswith("boxed_"): + response = extract_boxed_answer(response) or "" + rm_type = rm_type[len("boxed_") :] + + # This function is intended for remote or time-consuming reward model evaluation. + # Implement the actual logic as needed. + if rm_type == "remote_rm": + return await remote_rm(args, sample) + elif rm_type == "deepscaler": + return get_deepscaler_rule_based_reward(response, label) + elif rm_type == "dapo": + return compute_score_dapo(response, label) + elif rm_type == "math": + return 1 if grade_answer_verl(response, label) else 0 + elif rm_type == "f1": + return f1_score(response, label)[0] + elif rm_type == "gpqa": + return compute_gpqa_reward(response, label, metadata=metadata) + elif rm_type == "ifbench": + from .ifbench import compute_ifbench_reward + + return compute_ifbench_reward(response, label, metadata=metadata) + elif rm_type == "random": + return random.randint(0, 1) + elif rm_type: + raise NotImplementedError(f"Rule-based RM for {rm_type} is not implemented.") + else: + raise NotImplementedError("Rule-based RM type is not specified.") + + +async def batched_async_rm( + args, + samples: list[Sample], + **kwargs, +) -> list[int | float]: + if args.custom_rm_path is not None: + # Ensure the custom reward function is implemented in batch mode + rm_function = load_function(args.custom_rm_path) + return await rm_function(args, samples, **kwargs) + tasks = [async_rm(args, sample, **kwargs) for sample in samples] + rewards = await asyncio.gather(*tasks) + return rewards diff --git a/slime/rollout/rm_hub/__pycache__/__init__.cpython-312.pyc b/slime/rollout/rm_hub/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5555c3ec7f47886746ed05f7b8ab62bfeb6760c9 Binary files /dev/null and b/slime/rollout/rm_hub/__pycache__/__init__.cpython-312.pyc differ diff --git a/slime/rollout/rm_hub/__pycache__/deepscaler.cpython-312.pyc b/slime/rollout/rm_hub/__pycache__/deepscaler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46a119e4bbb78caaa979635a23da9cd77fb482f3 Binary files /dev/null and b/slime/rollout/rm_hub/__pycache__/deepscaler.cpython-312.pyc differ diff --git a/slime/rollout/rm_hub/__pycache__/f1.cpython-312.pyc b/slime/rollout/rm_hub/__pycache__/f1.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a13cd3ec781a3c558ec4ad0b6ebd0095d8a39b08 Binary files /dev/null and b/slime/rollout/rm_hub/__pycache__/f1.cpython-312.pyc differ diff --git a/slime/rollout/rm_hub/__pycache__/gpqa.cpython-312.pyc b/slime/rollout/rm_hub/__pycache__/gpqa.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f5a92419a4693e28c6ceb654ca47f6c2949aecc Binary files /dev/null and b/slime/rollout/rm_hub/__pycache__/gpqa.cpython-312.pyc differ diff --git a/slime/rollout/rm_hub/__pycache__/math_dapo_utils.cpython-312.pyc b/slime/rollout/rm_hub/__pycache__/math_dapo_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..19a22f30c47ad97836740aa89041fac07631de47 Binary files /dev/null and b/slime/rollout/rm_hub/__pycache__/math_dapo_utils.cpython-312.pyc differ diff --git a/slime/rollout/rm_hub/__pycache__/math_utils.cpython-312.pyc b/slime/rollout/rm_hub/__pycache__/math_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..848a613b0fca70230a2a28bbc8b7c8fe02fdf871 Binary files /dev/null and b/slime/rollout/rm_hub/__pycache__/math_utils.cpython-312.pyc differ diff --git a/slime/rollout/rm_hub/deepscaler.py b/slime/rollout/rm_hub/deepscaler.py new file mode 100644 index 0000000000000000000000000000000000000000..a6296c731880306b4aea41999732326035641361 --- /dev/null +++ b/slime/rollout/rm_hub/deepscaler.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .math_utils import extract_answer, grade_answer_mathd, grade_answer_sympy + + +def get_deepscaler_rule_based_reward(response, label): + if "" in response: + model_solution = response.split("")[-1] + elif "###Response" in response: + model_solution = response.split("###Response")[1] + else: + return 0 + + model_answer = extract_answer(model_solution) + if model_answer is None: + return 0 + if label == "": + return 0 + + # Convert single answer to list for uniform processing + assert isinstance(label, (str, float, int)) + ground_truths = [label] + + # Process each ground truth + processed_ground_truths = [] + for truth in ground_truths: + truth = str(truth) + if "\\boxed" in truth: + processed_truth = extract_answer(truth) + if processed_truth is not None: + processed_ground_truths.append(processed_truth) + else: + processed_ground_truths.append(truth) + + if not processed_ground_truths: + return 0 + + # Check against all possible correct answers + for ground_truth in processed_ground_truths: + is_correct = grade_answer_mathd(model_answer, ground_truth) or grade_answer_sympy(model_answer, ground_truth) + if is_correct: + return 1 + + return 0 diff --git a/slime/rollout/rm_hub/f1.py b/slime/rollout/rm_hub/f1.py new file mode 100644 index 0000000000000000000000000000000000000000..0fd5c4a096cebfbf78a4496ba87ee93dbeb7e386 --- /dev/null +++ b/slime/rollout/rm_hub/f1.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import re +import string +from collections import Counter + + +def normalize_answer(s): + + def remove_articles(text): + return re.sub(r"\b(a|an|the)\b", " ", text) + + def white_space_fix(text): + return " ".join(text.split()) + + def remove_punc(text): + exclude = set(string.punctuation) + return "".join(ch for ch in text if ch not in exclude) + + def lower(text): + return text.lower() + + return white_space_fix(remove_articles(remove_punc(lower(s)))) + + +def f1_score(prediction, ground_truth): + ZERO_METRIC = (0, 0, 0) + + if prediction is None: + return ZERO_METRIC + + normalized_prediction = normalize_answer(prediction) + normalized_ground_truth = normalize_answer(ground_truth) + + if normalized_prediction in ["yes", "no", "noanswer"] and normalized_prediction != normalized_ground_truth: + return ZERO_METRIC + if normalized_ground_truth in ["yes", "no", "noanswer"] and normalized_prediction != normalized_ground_truth: + return ZERO_METRIC + + prediction_tokens = normalized_prediction.split() + ground_truth_tokens = normalized_ground_truth.split() + common = Counter(prediction_tokens) & Counter(ground_truth_tokens) + num_same = sum(common.values()) + if num_same == 0: + return ZERO_METRIC + precision = 1.0 * num_same / len(prediction_tokens) + recall = 1.0 * num_same / len(ground_truth_tokens) + f1 = (2 * precision * recall) / (precision + recall) + return f1, precision, recall diff --git a/slime/rollout/rm_hub/gpqa.py b/slime/rollout/rm_hub/gpqa.py new file mode 100644 index 0000000000000000000000000000000000000000..93d23e9d11b97e2c132fd83d1257267dbf509f51 --- /dev/null +++ b/slime/rollout/rm_hub/gpqa.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 string +from collections.abc import Iterable + +DEFAULT_VALID_LETTERS = list(string.ascii_uppercase[:8]) + + +def _strip_chain_of_thought(text: str) -> str: + if not text: + return "" + + if "" in text: + return text.rsplit("", 1)[-1] + + return text + + +def _normalize_text(text: str) -> str: + return re.sub(r"[^a-z0-9]+", " ", text.lower()).strip() + + +def _extract_letter_from_response(response: str, valid_letters: Iterable[str]) -> str | None: + """ + Best-effort extraction of the selected option letter from the model response. + """ + if not response: + return None + + text = _strip_chain_of_thought(response) + patterns = [ + r"(?:answer|option|choice)\s*(?:is|:)?\s*([A-Z])", + r"([A-Z])\s*(?:is\s*(?:the)?\s*correct)", + r"final\s*(?:answer|option)\s*(?:is|:)?\s*([A-Z])", + ] + + valid_letters = {letter.upper() for letter in valid_letters} + for pattern in patterns: + match = re.search(pattern, text, flags=re.IGNORECASE) + if match: + letter = match.group(1).upper() + if letter in valid_letters: + return letter + + # Fallback: last standalone capital letter that is valid. + candidates = re.findall(r"\b([A-Z])\b", text) + for letter in reversed(candidates): + letter = letter.upper() + if letter in valid_letters: + return letter + + return None + + +def compute_gpqa_reward(response: str, label, metadata: dict | None = None) -> float: + """Rule-based scorer for GPQA-style multiple-choice evaluation.""" + if response is None: + return 0.0 + + metadata = metadata or {} + + choices = metadata.get("choices") + if isinstance(choices, dict): + choices = list(choices.values()) + elif choices is not None: + choices = list(choices) + + valid_letters = metadata.get("valid_letters") + if valid_letters: + valid_letters = [str(letter).upper() for letter in valid_letters] + elif choices: + valid_letters = list(string.ascii_uppercase[: len(choices)]) + else: + valid_letters = DEFAULT_VALID_LETTERS + + correct_letter = metadata.get("correct_letter") + if isinstance(correct_letter, str): + correct_letter = correct_letter.strip().upper() + else: + correct_letter = None + + label_text = None + if isinstance(label, str): + label_text = label.strip() + if len(label_text) == 1 and label_text.upper() in valid_letters and not correct_letter: + correct_letter = label_text.upper() + elif isinstance(label, (int, float)): + idx = int(label) + if 0 <= idx < len(valid_letters): + correct_letter = valid_letters[idx] + + if not correct_letter and choices and label_text: + normalized_label = _normalize_text(label_text) + for idx, choice in enumerate(choices): + if _normalize_text(str(choice)) == normalized_label: + correct_letter = valid_letters[idx] + metadata.setdefault("correct_answer", choice) + break + + extracted_letter = _extract_letter_from_response(response, valid_letters) + if extracted_letter and correct_letter: + return 1.0 if extracted_letter == correct_letter else 0.0 + + candidate_answers = [] + if correct_letter and choices: + try: + idx = valid_letters.index(correct_letter) + except ValueError: + idx = None + if idx is not None and idx < len(choices): + candidate_answers.append(str(choices[idx])) + + for key in ("correct_answer", "answer_text"): + value = metadata.get(key) + if value: + candidate_answers.append(str(value)) + + if label_text: + candidate_answers.append(label_text) + + normalized_targets = {_normalize_text(text) for text in candidate_answers if text} + normalized_response = _normalize_text(_strip_chain_of_thought(response)) + for target in normalized_targets: + if target and target in normalized_response: + return 1.0 + + if extracted_letter and not correct_letter and label_text: + return 1.0 if extracted_letter == label_text.strip().upper() else 0.0 + + return 0.0 diff --git a/slime/rollout/rm_hub/ifbench.py b/slime/rollout/rm_hub/ifbench.py new file mode 100644 index 0000000000000000000000000000000000000000..be1d8d9570d597933473453839bb3b1d7ca8ab64 --- /dev/null +++ b/slime/rollout/rm_hub/ifbench.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import importlib +import logging +import os +import subprocess +import sys +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + +_WORKSPACE_ROOT = Path(__file__).resolve().parents[3] +_WORKSPACE_PARENT = _WORKSPACE_ROOT.parent +_LOCAL_IFBENCH_REQUIREMENTS = _WORKSPACE_ROOT / "examples" / "eval_multi_task" / "requirements_ifbench.txt" + + +def _ensure_ifbench_repo() -> Path: + """Clone IFBench repo if needed and ensure it is available on sys.path.""" + + repo_path = _WORKSPACE_PARENT / "IFBench" + + if not repo_path.exists(): + clone_cmd = ["git", "clone", "https://github.com/allenai/IFBench.git", str(repo_path)] + try: + subprocess.run(clone_cmd, check=True, capture_output=True) + except Exception as exc: + raise ImportError( + "Unable to automatically clone IFBench. Please clone " + "https://github.com/allenai/IFBench.git into the repo root." + ) from exc + + repo_str = str(repo_path) + if repo_str not in sys.path: + sys.path.insert(0, repo_str) + + current_pythonpath = os.environ.get("PYTHONPATH") + if current_pythonpath is None: + os.environ["PYTHONPATH"] = repo_str + elif repo_str not in current_pythonpath.split(os.pathsep): + os.environ["PYTHONPATH"] = os.pathsep.join([repo_str, current_pythonpath]) + + return repo_path + + +def _ensure_ifbench_dependencies(repo_path: Path) -> None: + """Install IFBench requirements the first time the module is imported.""" + + requirements_file = _LOCAL_IFBENCH_REQUIREMENTS + + if not requirements_file.exists(): + logger.debug("Local IFBench requirements file not found at %s; skipping install.", requirements_file) + return + + sentinel = repo_path / ".deps_installed" + if sentinel.exists(): + return + + install_cmd = [sys.executable, "-m", "pip", "install", "-r", str(requirements_file)] + try: + subprocess.run(install_cmd, check=True) + except Exception as exc: + logger.warning("Failed to install IFBench dependencies automatically: %s", exc) + else: + sentinel.write_text("installed\n") + + +def _load_evaluation_lib(): + repo_path = _ensure_ifbench_repo() + try: + return importlib.import_module("evaluation_lib") + except ImportError: + _ensure_ifbench_dependencies(repo_path) + return importlib.import_module("evaluation_lib") + + +evaluation_lib = _load_evaluation_lib() +InputExample = evaluation_lib.InputExample + + +JsonDict = dict[str, Any] +KwargsDict = dict[str, str | int | float | None] + + +def _normalize_instruction_ids(raw_ids: Sequence[Any]) -> list[str]: + """Ensure instruction identifiers are clean strings.""" + + normalized: list[str] = [] + for entry in raw_ids or []: + if entry is None: + continue + text = str(entry).strip() + if not text: + continue + normalized.append(text) + return normalized + + +def _coerce_kwargs_list( + raw_kwargs: Any, + num_instructions: int, +) -> list[KwargsDict]: + """Convert stored kwargs into the list structure expected by IFBench.""" + + if isinstance(raw_kwargs, list): + processed: list[KwargsDict] = [] + for entry in raw_kwargs: + if isinstance(entry, dict): + processed.append(dict(entry)) + else: + processed.append({}) + elif isinstance(raw_kwargs, dict): + processed = [dict(raw_kwargs) for _ in range(num_instructions)] + else: + processed = [{} for _ in range(num_instructions)] + + if len(processed) < num_instructions: + tail = processed[-1] if processed else {} + processed.extend([dict(tail) for _ in range(num_instructions - len(processed))]) + elif len(processed) > num_instructions: + processed = processed[:num_instructions] + + # Remove explicit None values to match official preprocessing. + sanitized: list[KwargsDict] = [] + for entry in processed: + sanitized.append({k: v for k, v in entry.items() if v is not None}) + return sanitized + + +def _build_input_example(metadata: JsonDict) -> InputExample | None: + instruction_ids = _normalize_instruction_ids(metadata.get("instruction_id_list") or []) + if not instruction_ids: + logger.debug("Missing instruction identifiers in metadata: %s", metadata) + return None + + prompt_text = metadata.get("prompt_text") + if prompt_text is None: + prompt_text = "" + else: + prompt_text = str(prompt_text) + + raw_kwargs = metadata.get("kwargs") + kwargs_list = _coerce_kwargs_list(raw_kwargs, len(instruction_ids)) + + return InputExample( + key=int(metadata.get("record_id") or 0), + instruction_id_list=instruction_ids, + prompt=prompt_text, + kwargs=kwargs_list, + ) + + +def compute_ifbench_reward(response: str, label: Any, metadata: JsonDict | None = None) -> float: + """Score a model response using the official IFBench rules.""" + + if metadata is None: + logger.debug("No metadata provided for IFBench scoring.") + return 0.0 + + if response is None: + return 0.0 + + inp = _build_input_example(metadata) + if inp is None: + return 0.0 + + prompt_to_response = {inp.prompt: str(response or "")} + output = evaluation_lib.test_instruction_following_strict(inp, prompt_to_response) + return 1.0 if output.follow_all_instructions else 0.0 diff --git a/slime/rollout/rm_hub/math_dapo_utils.py b/slime/rollout/rm_hub/math_dapo_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..915c3aa480b1f9812c18fd3255b5c03e774a81f2 --- /dev/null +++ b/slime/rollout/rm_hub/math_dapo_utils.py @@ -0,0 +1,292 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Adapted from https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/hendrycks_math/utils.py + +import re +import signal + + +def last_boxed_only_string(string: str) -> str | None: + """Extract the last LaTeX boxed expression from a string. + + Args: + string: Input string containing LaTeX code + + Returns: + The last boxed expression or None if not found + """ + idx = string.rfind("\\boxed{") + if idx < 0: + return None + + i = idx + right_brace_idx = None + num_left_braces_open = 0 + + while i < len(string): + if string[i] == "{": + num_left_braces_open += 1 + if string[i] == "}": + num_left_braces_open -= 1 + if num_left_braces_open == 0: + right_brace_idx = i + break + i += 1 + + return string[idx : right_brace_idx + 1] if right_brace_idx is not None else None + + +def remove_boxed(s: str) -> str: + """Remove the LaTeX boxed command from a string. + + Args: + s: String with format "\\boxed{content}" + + Returns: + The content inside the boxed command + """ + left = "\\boxed{" + assert s[: len(left)] == left, f"box error: {s}" + assert s[-1] == "}", f"box error: {s}" + return s[len(left) : -1] + + +class timeout: + + def __init__(self, seconds=1, error_message="Timeout"): + self.seconds = seconds + self.error_message = error_message + + def handle_timeout(self, signum, frame): + raise TimeoutError(self.error_message) + + def __enter__(self): + signal.signal(signal.SIGALRM, self.handle_timeout) + signal.alarm(self.seconds) + + def __exit__(self, type, value, traceback): + signal.alarm(0) + + +# Constants for normalization +SUBSTITUTIONS = [ + ("an ", ""), + ("a ", ""), + (".$", "$"), + ("\\$", ""), + (r"\ ", ""), + (" ", ""), + ("mbox", "text"), + (",\\text{and}", ","), + ("\\text{and}", ","), + ("\\text{m}", "\\text{}"), +] + +REMOVED_EXPRESSIONS = [ + "square", + "ways", + "integers", + "dollars", + "mph", + "inches", + "hours", + "km", + "units", + "\\ldots", + "sue", + "points", + "feet", + "minutes", + "digits", + "cents", + "degrees", + "cm", + "gm", + "pounds", + "meters", + "meals", + "edges", + "students", + "childrentickets", + "multiples", + "\\text{s}", + "\\text{.}", + "\\text{\ns}", + "\\text{}^2", + "\\text{}^3", + "\\text{\n}", + "\\text{}", + r"\mathrm{th}", + r"^\circ", + r"^{\circ}", + r"\;", + r",\!", + "{,}", + '"', + "\\dots", + "<|im_end|>", + "<|endoftext|>", +] + + +def normalize_final_answer(final_answer: str) -> str: + """Normalize a final answer to a quantitative reasoning question. + + Args: + final_answer: The answer string to normalize + + Returns: + Normalized answer string + """ + final_answer = str(final_answer) + final_answer = final_answer.split("=")[-1] + + # Apply substitutions and removals + for before, after in SUBSTITUTIONS: + final_answer = final_answer.replace(before, after) + for expr in REMOVED_EXPRESSIONS: + final_answer = final_answer.replace(expr, "") + + # Extract and normalize LaTeX math + final_answer = re.sub(r"(.*?)(\$)(.*?)(\$)(.*)", "$\\3$", final_answer) + final_answer = re.sub(r"(\\text\{)(.*?)(\})", "\\2", final_answer) + final_answer = re.sub(r"(\\textbf\{)(.*?)(\})", "\\2", final_answer) + final_answer = re.sub(r"(\\overline\{)(.*?)(\})", "\\2", final_answer) + final_answer = re.sub(r"(\\boxed\{)(.*)(\})", "\\2", final_answer) + + # Normalize shorthand TeX: + # \fracab -> \frac{a}{b} + # \frac{abc}{bef} -> \frac{abc}{bef} + # \fracabc -> \frac{a}{b}c + # \sqrta -> \sqrt{a} + # \sqrtab -> sqrt{a}b + final_answer = re.sub(r"(frac)([^{])(.)", "frac{\\2}{\\3}", final_answer) + final_answer = re.sub(r"(sqrt)([^{])", "sqrt{\\2}", final_answer) + final_answer = final_answer.replace("$", "") + + # Normalize numbers + if final_answer.replace(",", "").isdigit(): + final_answer = final_answer.replace(",", "") + + return final_answer.strip() + + +def is_correct_minerva( + solution_str: str, gt: str, gt_need_extract: bool = False, answer_pattern: str = r"(?i)Answer\s*:\s*([^\n]+)" +) -> tuple[bool, str]: + """Check if the solution is correct according to Minerva criteria. + + Args: + solution_str: The solution string to check + gt: The ground truth answer + gt_need_extract: Whether the ground truth needs extraction + answer_pattern: Regex pattern to extract the answer + + Returns: + Tuple of (is_correct, normalized_prediction) + """ + # Extract answer from solution + match = re.findall(answer_pattern, solution_str) + extracted_answer = match[-1] if match else "[INVALID]" + pred = normalize_final_answer(extracted_answer) + + # Process ground truth + if gt_need_extract: + gt = normalize_final_answer(remove_boxed(last_boxed_only_string(gt))) + else: + gt = normalize_final_answer(gt) + + gt = str(int(float(gt))) # in dapo, all answers are integers + + return (pred == gt), pred + + +def is_correct_strict_box(pred: str, gt: str, pause_tokens_index: list[int] | None = None) -> tuple[int, str | None]: + """Check if the prediction is correct using strict boxed answer criteria. + + Args: + pred: The prediction string + gt: The ground truth answer + pause_tokens_index: Indices of pause tokens + + Returns: + Tuple of (score, extracted_prediction) + """ + # Extract the relevant part of the prediction + if pause_tokens_index is not None: + assert len(pause_tokens_index) == 4 + pred = pred[pause_tokens_index[-1] - 100 :] + else: + pred = pred[-100:] + + # Extract and check the boxed answer + boxed_pred = last_boxed_only_string(pred) + extracted_pred = remove_boxed(boxed_pred) if boxed_pred is not None else None + + return 1 if (extracted_pred == gt) else -1, extracted_pred + + +def verify( + solution_str: str, answer: str, strict_box_verify: bool = False, pause_tokens_index: list[int] | None = None +) -> bool: + """Verify if the solution is correct. + + Args: + solution_str: The solution string to verify + answer: The ground truth answer + strict_box_verify: Whether to use strict box verification + pause_tokens_index: Indices of pause tokens + + Returns: + True if the solution is correct, False otherwise + """ + if strict_box_verify: + correct, pred = is_correct_strict_box(solution_str, answer, pause_tokens_index) + return correct == 1, pred + + correct, pred = is_correct_minerva(solution_str, answer) + return correct, pred + + +def compute_score( + solution_str: str, + ground_truth: str, + strict_box_verify: bool = False, + pause_tokens_index: list[int] | None = None, +) -> float: + """Compute the reward score for a solution. + + Args: + solution_str: The solution string + ground_truth: The ground truth answer + config: Configuration object containing reward model settings + pause_tokens_index: Indices of pause tokens + + Returns: + Reward score (1.0 for correct, -1.0 for incorrect) + """ + # Limit solution length for efficiency + solution_str = solution_str[-300:] # The longest answer in MATH-500 has 159 characters + + # Verify the solution + correct, pred = verify(solution_str, ground_truth, strict_box_verify, pause_tokens_index) + + reward = 1.0 if correct else -1.0 + acc = correct + + return { + "score": reward, + "acc": acc, + "pred": pred, + } diff --git a/slime/rollout/rm_hub/math_utils.py b/slime/rollout/rm_hub/math_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4e1c47ec133394485fbddc79405d28c20f8fab4c --- /dev/null +++ b/slime/rollout/rm_hub/math_utils.py @@ -0,0 +1,491 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# from https://github.com/agentica-project/deepscaler/blob/e6080ccd974eb64bd3430f0b36108244a6fee330/deepscaler/rewards/math_utils/utils.py +""" +Answer checker API that uses sympy to simplify expressions and check for equality. + +Call grade_answer(given_answer: str, ground_truth: str). +""" +import re + +import sympy +from pylatexenc import latex2text +from sympy.parsing import sympy_parser + + +# Dan Hendrycks' code +def mathd_normalize_answer(answer: str | None) -> str | None: + if answer is None: + return None + answer = answer.strip() + try: + # Remove enclosing `\text{}`. + m = re.search("^\\\\text\{(?P.+?)\}$", answer) + if m is not None: + answer = m.group("text").strip() + return _strip_string(answer) + except Exception: + return answer + + +def _strip_string(string): + def _fix_fracs(string): + substrs = string.split("\\frac") + new_str = substrs[0] + if len(substrs) > 1: + substrs = substrs[1:] + for substr in substrs: + new_str += "\\frac" + if substr[0] == "{": + new_str += substr + else: + try: + assert len(substr) >= 2 + except Exception: + return string + a = substr[0] + b = substr[1] + if b != "{": + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}{" + b + "}" + post_substr + else: + new_str += "{" + a + "}{" + b + "}" + else: + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}" + b + post_substr + else: + new_str += "{" + a + "}" + b + string = new_str + return string + + def _fix_a_slash_b(string): + if len(string.split("/")) != 2: + return string + a = string.split("/")[0] + b = string.split("/")[1] + try: + a = int(a) + b = int(b) + assert string == f"{a}/{b}" + new_string = "\\frac{" + str(a) + "}{" + str(b) + "}" + return new_string + except Exception: + return string + + def _remove_right_units(string): + # "\\text{ " only ever occurs (at least in the val set) when describing units + if "\\text{ " in string: + splits = string.split("\\text{ ") + assert len(splits) == 2 + return splits[0] + else: + return string + + def _fix_sqrt(string): + if "\\sqrt" not in string: + return string + splits = string.split("\\sqrt") + new_string = splits[0] + for split in splits[1:]: + if split[0] != "{": + a = split[0] + new_substr = "\\sqrt{" + a + "}" + split[1:] + else: + new_substr = "\\sqrt" + split + new_string += new_substr + return new_string + + # linebreaks + string = string.replace("\n", "") + + # remove inverse spaces + string = string.replace("\\!", "") + + # replace \\ with \ + string = string.replace("\\\\", "\\") + + # replace tfrac and dfrac with frac + string = string.replace("tfrac", "frac") + string = string.replace("dfrac", "frac") + + # remove \left and \right + string = string.replace("\\left", "") + string = string.replace("\\right", "") + + # Remove circ (degrees) + string = string.replace("^{\\circ}", "") + string = string.replace("^\\circ", "") + + # remove dollar signs + string = string.replace("\\$", "") + + # remove units (on the right) + string = _remove_right_units(string) + + # remove percentage + string = string.replace("\\%", "") + string = string.replace("\%", "") + + # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string + string = string.replace(" .", " 0.") + string = string.replace("{.", "{0.") + # if empty, return empty string + if len(string) == 0: + return string + if string[0] == ".": + string = "0" + string + + # to consider: get rid of e.g. "k = " or "q = " at beginning + if len(string.split("=")) == 2: + if len(string.split("=")[0]) <= 2: + string = string.split("=")[1] + + # fix sqrt3 --> sqrt{3} + string = _fix_sqrt(string) + + # remove spaces + string = string.replace(" ", "") + + # \frac1b or \frac12 --> \frac{1}{b} and \frac{1}{2}, etc. Even works with \frac1{72} (but not \frac{72}1). Also does a/b --> \\frac{a}{b} + string = _fix_fracs(string) + + # manually change 0.5 --> \frac{1}{2} + if string == "0.5": + string = "\\frac{1}{2}" + + # NOTE: X/Y changed to \frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y + string = _fix_a_slash_b(string) + + return string + + +# sympy might hang -- we don't care about trying to be lenient in these cases +BAD_SUBSTRINGS = ["^{", "^("] +BAD_REGEXES = ["\^[0-9]+\^", "\^[0-9][0-9]+"] +TUPLE_CHARS = "()[]" + + +def _sympy_parse(expr: str): + """Parses an expression with sympy.""" + py_expr = expr.replace("^", "**") + return sympy_parser.parse_expr( + py_expr, + transformations=(sympy_parser.standard_transformations + (sympy_parser.implicit_multiplication_application,)), + ) + + +def _parse_latex(expr: str) -> str: + """Attempts to parse latex to an expression sympy can read.""" + expr = expr.replace("\\tfrac", "\\frac") + expr = expr.replace("\\dfrac", "\\frac") + expr = expr.replace("\\frac", " \\frac") # Play nice with mixed numbers. + expr = latex2text.LatexNodes2Text().latex_to_text(expr) + + # Replace the specific characters that this parser uses. + expr = expr.replace("√", "sqrt") + expr = expr.replace("π", "pi") + expr = expr.replace("∞", "inf") + expr = expr.replace("∪", "U") + expr = expr.replace("·", "*") + expr = expr.replace("×", "*") + + return expr.strip() + + +def _is_float(num: str) -> bool: + try: + float(num) + return True + except Exception: + return False + + +def _is_int(x: float) -> bool: + try: + return abs(x - int(round(x))) <= 1e-7 + except Exception: + return False + + +def _is_frac(expr: str) -> bool: + return bool(re.search(r"^-?[0-9]+.?/0*[1-9][0-9]*.?$", expr)) + + +def _str_is_int(x: str) -> bool: + try: + x = _strip_properly_formatted_commas(x) + x = float(x) + return abs(x - int(round(x))) <= 1e-7 + except Exception: + return False + + +def _str_to_int(x: str) -> int: + x = x.replace(",", "") + x = float(x) + return int(x) + + +def _inject_implicit_mixed_number(step: str): + """ + Automatically make a mixed number evalable + e.g. 7 3/4 => 7+3/4 + """ + p1 = re.compile("([0-9]) +([0-9])") + step = p1.sub("\\1+\\2", step) ## implicit mults + return step + + +def _strip_properly_formatted_commas(expr: str): + # We want to be careful because we don't want to strip tuple commas + p1 = re.compile("(\d)(,)(\d\d\d)($|\D)") + while True: + next_expr = p1.sub("\\1\\3\\4", expr) + if next_expr == expr: + break + expr = next_expr + return next_expr + + +def _normalize(expr: str) -> str: + """Normalize answer expressions.""" + if expr is None: + return None + + # Remove enclosing `\text{}`. + m = re.search("^\\\\text\{(?P.+?)\}$", expr) + if m is not None: + expr = m.group("text") + + expr = expr.replace("\\%", "%") + expr = expr.replace("\\$", "$") + expr = expr.replace("$", "") + expr = expr.replace("%", "") + expr = expr.replace(" or ", " , ") + expr = expr.replace(" and ", " , ") + + expr = expr.replace("million", "*10^6") + expr = expr.replace("billion", "*10^9") + expr = expr.replace("trillion", "*10^12") + + for unit in [ + "degree", + "cm", + "centimeter", + "meter", + "mile", + "second", + "minute", + "hour", + "day", + "week", + "month", + "year", + "foot", + "feet", + "inch", + "yard", + ]: + expr = re.sub(f"{unit}(es)?(s)? *(\^[0-9]+)?", "", expr) + expr = re.sub("\^ *\\\\circ", "", expr) + + if len(expr) > 0 and expr[0] == "{" and expr[-1] == "}": + expr = expr[1:-1] + + expr = re.sub(",\\\\! *", "", expr) + if _is_float(expr) and _is_int(float(expr)): + expr = str(int(round(float(expr)))) + if "\\" in expr: + try: + expr = _parse_latex(expr) + except Exception: + pass + + # edge case with mixed numbers and negative signs + expr = re.sub("- *", "-", expr) + + expr = _inject_implicit_mixed_number(expr) + expr = expr.replace(" ", "") + + # if we somehow still have latex braces here, just drop them + expr = expr.replace("{", "") + expr = expr.replace("}", "") + + # don't be case sensitive for text answers + expr = expr.lower() + + if _str_is_int(expr): + expr = str(_str_to_int(expr)) + + return expr + + +def count_unknown_letters_in_expr(expr: str): + expr = expr.replace("sqrt", "") + expr = expr.replace("frac", "") + letters_in_expr = set([x for x in expr if x.isalpha()]) + return len(letters_in_expr) + + +def should_allow_eval(expr: str): + # we don't want to try parsing unknown text or functions of more than two variables + if count_unknown_letters_in_expr(expr) > 2: + return False + + for bad_string in BAD_SUBSTRINGS: + if bad_string in expr: + return False + + for bad_regex in BAD_REGEXES: + if re.search(bad_regex, expr) is not None: + return False + + return True + + +def are_equal_under_sympy(ground_truth_normalized: str, given_normalized: str): + are_equal = False + try: + expr = f"({ground_truth_normalized})-({given_normalized})" + if should_allow_eval(expr): + sympy_diff = _sympy_parse(expr) + simplified = sympy.simplify(sympy_diff) + if simplified == 0: + are_equal = True + except Exception: + pass + return are_equal + + +def split_tuple(expr: str): + """ + Split the elements in a tuple/interval, while handling well-formatted commas in large numbers + """ + expr = _strip_properly_formatted_commas(expr) + if len(expr) == 0: + return [] + if ( + len(expr) > 2 + and expr[0] in TUPLE_CHARS + and expr[-1] in TUPLE_CHARS + and all([ch not in expr[1:-1] for ch in TUPLE_CHARS]) + ): + elems = [elem.strip() for elem in expr[1:-1].split(",")] + else: + elems = [expr] + return elems + + +def last_boxed_only_string(string): + idx = string.rfind("\\boxed") + if idx < 0: + idx = string.rfind("\\fbox") + if idx < 0: + return None + + i = idx + right_brace_idx = None + num_left_braces_open = 0 + while i < len(string): + if string[i] == "{": + num_left_braces_open += 1 + if string[i] == "}": + num_left_braces_open -= 1 + if num_left_braces_open == 0: + right_brace_idx = i + break + i += 1 + + if right_brace_idx is None: + retval = None + else: + retval = string[idx : right_brace_idx + 1] + + return retval + + +def remove_boxed(s): + left = "\\boxed{" + try: + assert s[: len(left)] == left + assert s[-1] == "}" + return s[len(left) : -1] + except Exception: + return None + + +def extract_boxed_answer(solution: str) -> str: + """Extract the answer from inside a LaTeX \\boxed{} command""" + solution = last_boxed_only_string(solution) + solution = remove_boxed(solution) + return solution + + +def grade_answer_sympy(given_answer: str, ground_truth: str) -> bool: + ground_truth_normalized = _normalize(ground_truth) + given_normalized = _normalize(given_answer) + + if ground_truth_normalized is None: + return False + + if ground_truth_normalized == given_normalized: + return True + + if len(given_normalized) == 0: + return False + + ground_truth_elems = split_tuple(ground_truth_normalized) + given_elems = split_tuple(given_normalized) + + if len(ground_truth_elems) > 1 and ( + ground_truth_normalized[0] != given_normalized[0] or ground_truth_normalized[-1] != given_normalized[-1] + ): + is_correct = False + elif len(ground_truth_elems) != len(given_elems): + is_correct = False + else: + for ground_truth_elem, given_elem in zip(ground_truth_elems, given_elems, strict=False): + if _is_frac(ground_truth_elem) and _is_frac(given_elem): + # if fractions aren't reduced, then shouldn't be marked as correct + # so, we don't want to allow sympy.simplify in this case + is_correct = ground_truth_elem == given_elem + elif _str_is_int(ground_truth_elem) != _str_is_int(given_elem): + # if the ground truth answer is an integer, we require the given answer to be a strict match (no sympy.simplify) + is_correct = False + else: + is_correct = are_equal_under_sympy(ground_truth_elem, given_elem) + if not is_correct: + break + + return is_correct + + +def grade_answer_mathd(given_answer: str, ground_truth: str) -> bool: + ground_truth_normalized_mathd = mathd_normalize_answer(ground_truth) + given_answer_normalized_mathd = mathd_normalize_answer(given_answer) + + # be at least as lenient as mathd + if ground_truth_normalized_mathd == given_answer_normalized_mathd: + return True + return False + + +def extract_answer(passage: str) -> str: + if "\\boxed" in passage: + return extract_boxed_answer(passage) + return None + + +def grade_answer_verl(solution_str, ground_truth): + if not ground_truth: + return False + ground_truth = str(ground_truth) + if "\\boxed" in ground_truth: + ground_truth = extract_answer(ground_truth) + given_answer = extract_answer(solution_str) + if given_answer is None: + return False + return grade_answer_mathd(given_answer, ground_truth) or grade_answer_sympy(given_answer, ground_truth) diff --git a/slime/rollout/sglang_rollout.py b/slime/rollout/sglang_rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..dfc9db83a8ff1e2921c76bea4c57cac03c3d494e --- /dev/null +++ b/slime/rollout/sglang_rollout.py @@ -0,0 +1,675 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import copy +import logging +from argparse import Namespace +from collections import defaultdict +from collections.abc import Callable +from typing import Any + +import numpy as np +import pybase64 +import torch +import sglang_router +from packaging.version import parse +from tqdm import tqdm + +from slime.rollout.base_types import RolloutFnEvalOutput, RolloutFnTrainOutput +from slime.rollout.filter_hub.base_types import DynamicFilterOutput +from slime.utils.async_utils import run +from slime.utils.data import Dataset +from slime.utils.eval_config import EvalDatasetConfig +from slime.utils.http_utils import get, post +from slime.utils.mask_utils import get_response_lengths, MultiTurnLossMaskGenerator +from slime.utils.misc import SingletonMeta, load_function +from slime.utils.processing_utils import encode_image_for_rollout_engine, load_processor, load_tokenizer +from slime.utils.types import Sample + +from .rm_hub import async_rm, batched_async_rm + +__all__ = ["generate_rollout"] + +logger = logging.getLogger(__name__) + + +class GenerateState(metaclass=SingletonMeta): + """ + The global state for the generation process. + """ + + def __init__(self, args: Namespace) -> None: + # persistent state for the generation process + self.args = args + self.tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) + self.processor = load_processor(args.hf_checkpoint, trust_remote_code=True) + + num_engines = args.rollout_num_gpus // args.rollout_num_gpus_per_engine + self.semaphore = asyncio.Semaphore( + args.sglang_server_concurrency * num_engines if num_engines > 0 else 1 + ) + self.sampling_params: dict[str, Any] = dict( + temperature=args.rollout_temperature, + top_p=args.rollout_top_p, + top_k=args.rollout_top_k, + max_new_tokens=args.rollout_max_response_len, + stop=args.rollout_stop, + stop_token_ids=args.rollout_stop_token_ids, + skip_special_tokens=args.rollout_skip_special_tokens, + no_stop_trim=False, # Changed to remove stop tokens from rollout output + spaces_between_special_tokens=False, + ) + + if getattr(args, "sglang_enable_deterministic_inference", False): + sampling_seed_base = args.rollout_seed + self.group_sampling_seeds = [sampling_seed_base + i for i in range(args.n_samples_per_prompt)] + + self.reset() + + def reset(self) -> None: + self.remaining_batch_size = 0 + self.pendings = set() + self.aborted = False + self.current_rollout_id = 0 + + def submit_generate_tasks(self, samples: list[list[Sample]]) -> None: + for group in samples: + self.pendings.add( + asyncio.create_task( + # submit a group of samples as a single task. + generate_and_rm_group( + self.args, + group, + sampling_params=self.sampling_params.copy(), + evaluation=False, + ) + ) + ) + self.remaining_batch_size += len(samples) + + + +def _is_lightning_opd(sample: Sample) -> bool: + """Check if sample is a Lightning OPD sample (pre-computed response + teacher logprobs).""" + metadata = sample.metadata or {} + return metadata.get("is_lightning_opd", False) or metadata.get("is_offline_opd", False) + + +def _handle_lightning_opd_sample(sample: Sample, state: "GenerateState") -> Sample: + """Handle Lightning OPD samples: response tokens are pre-computed in parquet metadata. + + Expected metadata fields: + response_tokens: list[int] pre-tokenized response token IDs + loss_mask: list[int] 1 for each response token to compute loss on + response: str decoded response text (used by verifiable reward) + + The sample's prompt tokens are prepended to form the full sequence sent to the teacher + server for logprob computation. The training model then computes student logprobs on + the same sequence during each forward pass, so the OPD advantage + log P_teacher - log P_πt is still computed against the *current* policy. + """ + import numpy as np + + metadata = sample.metadata or {} + + # Prompt may be a raw string (common case) or already tokenized list[int]. + prompt = sample.prompt + if isinstance(prompt, str): + prompt = state.tokenizer.encode(prompt, add_special_tokens=False) + elif isinstance(prompt, np.ndarray): + prompt = prompt.tolist() + + response_tokens = metadata.get("response_tokens", []) + if isinstance(response_tokens, np.ndarray): + response_tokens = response_tokens.tolist() + + loss_mask = metadata.get("loss_mask", [1] * len(response_tokens)) + if isinstance(loss_mask, np.ndarray): + loss_mask = loss_mask.tolist() + + # Full sequence = prompt tokens + response tokens (RM needs full context for logprobs) + # Force Python int conversion: parquet pyarrow backend can produce numpy.int64 elements + # that survive list() but fail JSON serialization when sent to the teacher server. + sample.tokens = [int(x) for x in prompt] + [int(x) for x in response_tokens] + sample.loss_mask = [int(x) for x in loss_mask] + sample.response_length = int(sum(loss_mask)) + sample.response = metadata.get("response", "") + sample.status = Sample.Status.COMPLETED + + # Load pre-computed student (pi_ref) log-probs for importance weight tracking. + # These are produced by data_curation/add_student_logprobs.py and stored as + # metadata["student_log_probs"]. When present they are passed through as + # rollout_log_probs so that loss.py can compute w = pi_theta / pi_ref. + student_log_probs = metadata.get("student_log_probs") + if student_log_probs is not None: + sample.rollout_log_probs = [float(x) for x in student_log_probs] + + return sample + + +async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, Any]) -> Sample: + """Generate using traditional SGLang router with token-based workflow""" + if args.ci_test: + assert isinstance(sample.prompt, str) + + state = GenerateState(args) + url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate" + + assert ( + sample.status == Sample.Status.PENDING or sample.status == Sample.Status.ABORTED + ), f"Sample status is {sample.status}" + + # Handle Lightning OPD samples: response tokens are pre-computed in metadata, skip sglang. + # The RM call (teacher logprob computation) still runs normally after this. + if _is_lightning_opd(sample): + return _handle_lightning_opd_sample(sample, state) + + if state.processor: + processor_output = state.processor(text=sample.prompt, **sample.multimodal_inputs) + prompt_ids = processor_output["input_ids"][0] + sample.multimodal_train_inputs = { + k: v for k, v in processor_output.items() if k not in ["input_ids", "attention_mask"] + } or None + else: + prompt_ids = state.tokenizer.encode(sample.prompt, add_special_tokens=False) + + if len(sample.response) > 0: + sampling_params["max_new_tokens"] -= len(sample.tokens) - len(prompt_ids) + + assert ( + sampling_params["max_new_tokens"] >= 0 + ), f"max_new_tokens: {sampling_params['max_new_tokens']} should not be less than 0" + if sampling_params["max_new_tokens"] == 0: + sample.status = Sample.Status.TRUNCATED + return sample + + # Prepare payload for sglang server + payload = { + "sampling_params": sampling_params, + "return_logprob": True, + } + + if args.use_rollout_routing_replay: + payload["return_routed_experts"] = True + + if sample.multimodal_inputs and sample.multimodal_inputs["images"]: + image_data = sample.multimodal_inputs["images"] + payload["image_data"] = [encode_image_for_rollout_engine(image) for image in image_data] + + # Use existing tokens for multi-turn or tokenize the new prompt + if len(sample.response) > 0: + payload["input_ids"] = sample.tokens + else: + payload["input_ids"] = prompt_ids + if not sample.tokens: # Initialize sample.tokens for the first turn + sample.tokens = prompt_ids + + output = await post(url, payload) + + # Extract new response tokens + + if args.use_slime_router and "RadixTreeMiddleware" in args.slime_router_middleware_paths: + assert not args.partial_rollout, "Currently partial rollout is not supported when using slime router" + retrieve_url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/retrieve_from_text" + retrieve_payload = {"text": sample.prompt + output["text"], "return_logp": True} + retrieve_output = await post(retrieve_url, retrieve_payload) + sample.tokens = retrieve_output["tokens"] + sample.response += output["text"] + sample.loss_mask = retrieve_output["loss_mask"] + sample.response_length = get_response_lengths([sample.loss_mask])[0] + sample.loss_mask = sample.loss_mask[-sample.response_length :] + sample.rollout_log_probs = retrieve_output["rollout_logp"][-sample.response_length :] + # Notice: currently cannot get the spec info from radix router output. + else: + if "output_token_logprobs" in output["meta_info"]: + new_response_tokens = [item[1] for item in output["meta_info"]["output_token_logprobs"]] + new_response_log_probs = [item[0] for item in output["meta_info"]["output_token_logprobs"]] + else: + new_response_tokens, new_response_log_probs = [], [] + + # Update sample with tokens directly - avoiding re-tokenization + sample.tokens = sample.tokens + new_response_tokens + sample.response_length += len(new_response_tokens) + sample.response += output["text"] + + if sample.rollout_log_probs is None: + sample.rollout_log_probs = [] + sample.rollout_log_probs += new_response_log_probs + + if args.sglang_speculative_algorithm: + # cannot directly use spec info from sglang because of partial rollout. + sample.spec_info.add( + meta_info=output["meta_info"], + response_length=sample.response_length, + ) + + if "weight_version" in output["meta_info"]: + sample.weight_versions.append(output["meta_info"]["weight_version"]) + + if "routed_experts" in output["meta_info"]: + sample.rollout_routed_experts = np.frombuffer( + pybase64.b64decode(output["meta_info"]["routed_experts"].encode("ascii")), + dtype=np.int32, + ).reshape( + len(sample.tokens) - 1, + args.num_layers, + args.moe_router_topk, + ) + + match output["meta_info"]["finish_reason"]["type"]: + case "length": + sample.status = Sample.Status.TRUNCATED + case "abort": + sample.status = Sample.Status.ABORTED + case "stop": + sample.status = Sample.Status.COMPLETED + + return sample + + +async def generate_and_rm( + args: Namespace, + sample: Sample | list[Sample], + sampling_params: dict[str, Any], + evaluation: bool = False, +) -> Sample | list[Sample]: + # mask previous off-policy generation for partial rollout + if args.partial_rollout and args.mask_offpolicy_in_partial_rollout and sample.response_length > 0: + sample.loss_mask = [0] * sample.response_length + + # For samples with existing response, check if they're complete + if sample.status == Sample.Status.COMPLETED or sample.status == Sample.Status.TRUNCATED: + assert sample.response is not None + if not args.group_rm: + assert sample.reward is not None + return sample + + state = GenerateState(args) + is_lightning_opd = _is_lightning_opd(sample) + + # generate (skip for Lightning OPD samples whose tokens are pre-computed) + if is_lightning_opd: + sample = _handle_lightning_opd_sample(sample, state) + else: + async with state.semaphore: + if state.aborted: + sample.status = Sample.Status.ABORTED + return sample + + if args.custom_generate_function_path is not None: + custom_generate_func = load_function(args.custom_generate_function_path) + sample = await custom_generate_func(args, sample, sampling_params) + else: + sample = await generate(args, sample, sampling_params) + + # for the rm that need the whole group, we will not do the rm here + if args.group_rm: + return sample + + # multi samples + if isinstance(sample, list): + samples = sample + if any([sample.status == Sample.Status.ABORTED for sample in samples]): + return samples + + # for multi agent system, the reward of some sample is calculated during generation. + samples_need_reward = [sample for sample in samples if sample.reward is None] + rewards = await batched_async_rm(args, samples_need_reward) + for sample, reward in zip(samples_need_reward, rewards, strict=False): + sample.reward = reward + return samples + else: + if sample.status == Sample.Status.ABORTED: + return sample + # for multi-turn environment, a reward could be assigned to the agent. + if sample.reward is None: + sample.reward = await async_rm(args, sample) + + return sample + + +async def generate_and_rm_group( + args: Namespace, group: list[Sample], sampling_params: dict[str, Any], evaluation: bool = False +) -> list[Sample]: + state = GenerateState(args) + + if state.aborted: + return group + + tasks = [] + for idx, sample in enumerate(group): + current_sampling_params = sampling_params.copy() + if getattr(args, "sglang_enable_deterministic_inference", False): + seed = state.group_sampling_seeds[idx] + current_sampling_params["sampling_seed"] = seed + tasks.append( + asyncio.create_task(generate_and_rm(args, sample, current_sampling_params, evaluation=evaluation)) + ) + + group = await asyncio.gather(*tasks) + + # for the rm that need the whole group, we will do the rm here + if not state.aborted and args.group_rm: + rewards = await batched_async_rm(args, group) + for sample, reward in zip(group, rewards, strict=False): + sample.reward = reward + + return group + + +async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: + aborted_samples = [] + + state = GenerateState(args) + assert not state.aborted + state.aborted = True + + # No rollout engines → no router, no pending tasks; nothing to abort. + if not args.rollout_num_gpus: + return aborted_samples + + if parse(sglang_router.__version__) <= parse("0.2.1") or args.use_slime_router: + response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/list_workers") + urls = response["urls"] + else: + response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/workers") + urls = [worker["url"] for worker in response["workers"]] + + logger.info(f"Abort request for {urls}") + await asyncio.gather(*[post(f"{url}/abort_request", {"abort_all": True}) for url in urls]) + + # make sure all the pending tasks are finished + count = 0 + while state.pendings: + done, state.pendings = await asyncio.wait(state.pendings, return_when=asyncio.FIRST_COMPLETED) + + if not args.partial_rollout: + continue + + # for partial rollout, collect the partial samples into the data buffer + for task in done: + group = task.result() + for sample in group: + if sample.response and "start_rollout_id" not in sample.metadata: + sample.metadata["start_rollout_id"] = rollout_id + aborted_samples.append(group) + count += len(group) + + if args.partial_rollout: + logger.info(f"Collected {count} partial samples into the data buffer") + + return aborted_samples + + +async def generate_rollout_async( + args: Namespace, rollout_id: int, data_source: Callable[[int], list[list[Sample]]] +) -> tuple[RolloutFnTrainOutput, list[list[Sample]]]: + """An example to implement the generate_rollout function for an rule based rm rollout generation. + + Args: + args: the whole args + rollout_id: int, the id of the rollout, used for deterministic data generation + data_source: the data source to fetch + + Returns: + tuple[RolloutFnTrainOutput, list[list[Sample]]]: + - data: a list of groups of samples generated by the rollout, length equals `rollout_batch_size` + - aborted_samples: any partial groups collected during abort when partial_rollout is enabled + """ + assert args.rollout_global_dataset + + state = GenerateState(args) + state.current_rollout_id = rollout_id + + # instantiate data filters + dynamic_filter = ( + load_function(args.dynamic_sampling_filter_path) if args.dynamic_sampling_filter_path is not None else None + ) + + metric_gatherer = _MetricGatherer() + + # target_data_size is the total number of valid samples to get + target_data_size = args.rollout_batch_size + + data = [] + all_data = [] + do_print = True + pbar = tqdm(total=target_data_size * args.n_samples_per_prompt, desc="Rollout generation") + while len(data) < target_data_size: + while state.remaining_batch_size < target_data_size: + # get samples from the buffer and submit the generation requests. + samples = data_source(args.over_sampling_batch_size) + state.submit_generate_tasks(samples) + + # wait for the generation to finish + done, state.pendings = await asyncio.wait(state.pendings, return_when=asyncio.FIRST_COMPLETED) + for task in done: + group: list[Sample] = task.result() + + if do_print: + sample = group[0][0] if isinstance(group[0], list) else group[0] + logger.info( + f"First rollout sample: {[str(sample.prompt) + sample.response]}, label: {sample.label}, reward: {sample.reward}", + ) + do_print = False + + assert len(group) == args.n_samples_per_prompt + all_data.append(group) + dynamic_filter_output = _call_dynamic_filter(dynamic_filter, args, group) + if not dynamic_filter_output.keep: + metric_gatherer.on_dynamic_filter_drop(reason=dynamic_filter_output.reason) + state.remaining_batch_size -= 1 + continue + + # add the samples to the data + # NOTE: here we have not stored all the unused samples back to the data buffer. + if len(data) < target_data_size: + data.append(group) + pbar.update(len(group)) + + pbar.close() + sample = data[-1][0][0] if isinstance(data[-1][0], list) else data[-1][0] + logger.info( + f"Finish rollout: {[str(sample.prompt) + sample.response]}, label: {sample.label}, reward: {sample.reward}", + ) + + # there are still some unfinished requests, abort them + aborted_samples = await abort(args, rollout_id) + + assert len(data) == args.rollout_batch_size, f"Got {len(data)} samples, expected {args.rollout_batch_size}" + data = sorted(data, key=lambda group: group[0][0].index if isinstance(group[0], list) else group[0].index) + all_samples = sorted(data, key=lambda group: group[0][0].index if isinstance(group[0], list) else group[0].index) + + # reset the global state to prevent effects on the next rollout or eval. + state.reset() + if args.rollout_sample_filter_path is not None: + filter_func = load_function(args.rollout_sample_filter_path) + filter_func(args, data) + + # There can be circumstances where users want to process all samples including filtered ones. + if args.rollout_all_samples_process_path is not None: + process_func = load_function(args.rollout_all_samples_process_path) + process_func(args, all_samples, data_source) + + return RolloutFnTrainOutput(samples=data, metrics=metric_gatherer.collect()), aborted_samples + + +def _call_dynamic_filter(fn, *args, **kwargs): + if fn is None: + return DynamicFilterOutput(keep=True) + + output = fn(*args, **kwargs) + + # compatibility for legacy version + if not isinstance(output, DynamicFilterOutput): + output = DynamicFilterOutput(keep=output) + + return output + + +class _MetricGatherer: + def __init__(self): + self._dynamic_filter_drop_reason_count = defaultdict(lambda: 0) + + def on_dynamic_filter_drop(self, reason: str | None): + if not reason: + return + self._dynamic_filter_drop_reason_count[reason] += 1 + + def collect(self): + return { + f"rollout/dynamic_filter/drop_{reason}": count + for reason, count in self._dynamic_filter_drop_reason_count.items() + } + + +EVAL_PROMPT_DATASET = {} + + +async def eval_rollout(args: Namespace, rollout_id: int) -> tuple[dict[str, dict[str, list[Any]]], list[list[Sample]]]: + assert not args.group_rm, "Group RM is not supported for eval rollout" + + coros = [] + for dataset_cfg in getattr(args, "eval_datasets", []) or []: + coros.append(eval_rollout_single_dataset(args, rollout_id, dataset_cfg)) + results_list = await asyncio.gather(*coros) + results = {} + for r in results_list: + results.update(r) + return RolloutFnEvalOutput(data=results), [] + + +async def eval_rollout_single_dataset( + args: Namespace, rollout_id: int, dataset_cfg: EvalDatasetConfig +) -> dict[str, dict[str, list[Any]]]: + """An example to implement the eval_rollout function for an rule based rm rollout generation. + + Args: + args: the whole args + rollout_id: int, the id of the rollout, used for deterministic data generation + dataset_cfg: configuration of the dataset + """ + assert not args.group_rm, "Group RM is not supported for eval rollout" + + global EVAL_PROMPT_DATASET + + cache_key = dataset_cfg.cache_key + (args.hf_checkpoint, args.apply_chat_template) + if cache_key not in EVAL_PROMPT_DATASET: + tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) + processor = load_processor(args.hf_checkpoint, trust_remote_code=True) + EVAL_PROMPT_DATASET[cache_key] = Dataset( + path=dataset_cfg.path, + tokenizer=tokenizer, + processor=processor, + max_length=args.eval_max_prompt_len, + prompt_key=dataset_cfg.input_key, + label_key=dataset_cfg.label_key, + multimodal_keys=args.multimodal_keys, + metadata_key=dataset_cfg.metadata_key, + tool_key=dataset_cfg.tool_key, + apply_chat_template=args.apply_chat_template, + apply_chat_template_kwargs=args.apply_chat_template_kwargs, + ) + dataset = EVAL_PROMPT_DATASET[cache_key] + + base_sampling_params = dict( + temperature=dataset_cfg.temperature, + top_p=dataset_cfg.top_p, + top_k=dataset_cfg.top_k, + max_new_tokens=dataset_cfg.max_response_len, + stop=args.rollout_stop, + stop_token_ids=args.rollout_stop_token_ids, + skip_special_tokens=args.rollout_skip_special_tokens, + no_stop_trim=False, # Changed to remove stop tokens from rollout output + spaces_between_special_tokens=False, + ) + + tasks = [] + # do multiple samples for eval prompts + sample_index = 0 + for _i, prompt_sample in enumerate(dataset.samples): + for j in range(dataset_cfg.n_samples_per_eval_prompt): + # use the same prompt for multiple samples + sample = copy.deepcopy(prompt_sample) + sample.index = sample_index + sample_index += 1 + sample.metadata = dataset_cfg.inject_metadata(getattr(sample, "metadata", None)) + sampling_params = base_sampling_params + if getattr(args, "sglang_enable_deterministic_inference", False): + sampling_params = base_sampling_params.copy() + sampling_params["sampling_seed"] = args.rollout_seed + j + tasks.append( + asyncio.create_task( + generate_and_rm( + args, + sample, + sampling_params=sampling_params, + evaluation=True, + ) + ) + ) + + data = [] + do_print = True + pbar = tqdm(total=len(tasks), desc="Rollout generation", disable=not do_print) + for coro in asyncio.as_completed(tasks): + sample = await coro + if do_print: + logger.info( + "eval_rollout_single_dataset example data: " + f"{[str(sample.prompt) + sample.response]} " + f"reward={sample.reward}" + ) + do_print = False + if isinstance(sample, list): + data.extend(sample) + else: + data.append(sample) + pbar.update(1) + pbar.close() + + data.sort(key=lambda sample: sample.index) + + reward_key = args.eval_reward_key or args.reward_key + return { + dataset_cfg.name: { + "rewards": [sample.reward if not reward_key else sample.reward[reward_key] for sample in data], + "truncated": [sample.status == Sample.Status.TRUNCATED for sample in data], + "samples": data, + } + } + + +# TODO remove this temp function +def generate_rollout( + args: Namespace, rollout_id: int, data_buffer: Any, evaluation: bool = False +) -> RolloutFnTrainOutput | RolloutFnEvalOutput: + """An example to implement the generate_rollout function for an rule based rm rollout generation. + + Args: + args: the whole args + rollout_id: int, the id of the rollout, used for deterministic data generation + data_buffer: the data buffer to store the generated samples + evaluation: bool, whether the rollout is for evaluation or not + + Returns: + list[list[Sample]]: a list of list of samples generated by the rollout + """ + output, aborted_samples = generate_abortable_samples( + args, rollout_id, data_buffer.get_samples, evaluation=evaluation + ) + data_buffer.add_samples(aborted_samples) + return output + + +def generate_abortable_samples( + args: Namespace, + rollout_id: int, + data_source: Callable[[int], list[list[Sample]]], + evaluation: bool = False, +) -> tuple[Any, list[list[Sample]]]: + assert args.rollout_global_dataset + if evaluation: + return run(eval_rollout(args, rollout_id)) + return run(generate_rollout_async(args, rollout_id, data_source)) diff --git a/slime/rollout/sleep_rollout.py b/slime/rollout/sleep_rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..a902656c3037ca3d79f68f8671a292014f4f617e --- /dev/null +++ b/slime/rollout/sleep_rollout.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +import time + +logger = logging.getLogger(__name__) + + +def sleep(args, rollout_id, data_source, evaluation=False): + count = 0 + while True: + time.sleep(3600) + count += 1 + logger.info(f"rollout sleep for {count} hours") diff --git a/slime/router/__init__.py b/slime/router/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a4d636226b9a4b5130dfd8e94f2cece5fd290e53 --- /dev/null +++ b/slime/router/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + diff --git a/slime/router/middleware_hub/__init__.py b/slime/router/middleware_hub/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a4d636226b9a4b5130dfd8e94f2cece5fd290e53 --- /dev/null +++ b/slime/router/middleware_hub/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + diff --git a/slime/router/middleware_hub/radix_tree.py b/slime/router/middleware_hub/radix_tree.py new file mode 100644 index 0000000000000000000000000000000000000000..b5b663af53ed3ee619523fca26d0bcf59a33aebc --- /dev/null +++ b/slime/router/middleware_hub/radix_tree.py @@ -0,0 +1,689 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +""" +String-based Radix Trie for efficient prefix matching and token caching. +Optimized for string prefixes with corresponding token IDs. +""" + +import threading +import time +from dataclasses import dataclass +from typing import Any + + +@dataclass +class MatchResult: + """Result of prefix matching operation.""" + + matched_prefix: str + token_ids: list[int] + logp: list[float] + loss_mask: list[int] # Added loss mask for model generation parts + remaining_string: str + last_node: StringTreeNode + + +class StringTreeNode: + """Tree node for string-based radix trie.""" + + counter = 0 + + def __init__(self, node_id: int | None = None): + # Core tree structure + self.children: list[StringTreeNode] = [] # Use list to store children + self.parent: StringTreeNode | None = None + + # Node data + self.string_key: str = "" # The string fragment this node represents + self.token_ids: list[int] | None = None # Token IDs for this node only (not cumulative) + self.logp: list[float] | None = None # Log probabilities for this node's tokens + self.loss_mask: list[int] | None = None # Loss mask for model generation parts + + # Access tracking + self.last_access_time = time.monotonic() + self.access_count = 0 + + # Reference counting for protection from eviction + self.ref_count = 0 + + # Weight version tracking + self.weight_version: int | None = None # Weight version for this node + + # Node identification + self.id = StringTreeNode.counter if node_id is None else node_id + StringTreeNode.counter += 1 + + @property + def is_leaf(self) -> bool: + """Check if this node is a leaf node.""" + return len(self.children) == 0 + + @property + def has_value(self) -> bool: + """Check if this node has token IDs stored.""" + return self.token_ids is not None + + def validate_token_logp_consistency(self) -> bool: + """Validate that token_ids, logp, and loss_mask have consistent lengths.""" + if self.token_ids is None and self.logp is None and self.loss_mask is None: + return True + + # Check if at least one is not None + if self.token_ids is not None and len(self.token_ids) > 0: + token_len = len(self.token_ids) + if self.logp is not None and len(self.logp) != token_len: + return False + if self.loss_mask is not None and len(self.loss_mask) != token_len: + return False + + return True + + @property + def is_evictable(self) -> bool: + """Check if this node can be evicted.""" + return self.ref_count == 0 and self.token_ids is not None + + def touch(self): + """Update access time and count.""" + self.last_access_time = time.monotonic() + self.access_count += 1 + + def __lt__(self, other: StringTreeNode) -> bool: + """For heap operations - least recently used first.""" + return self.last_access_time < other.last_access_time + + +class StringRadixTrie: + """ + String-based Radix Trie for efficient prefix matching and token caching. + Features: + - Efficient string prefix matching + - Token ID caching for matched prefixes + - Thread-safe operations + - Weight version tracking + - Automatic garbage collection based on weight version thresholds + """ + + def __init__(self, max_cache_size: int = 10000, gc_threshold_k: int = 5, tokenizer=None, verbose: bool = False): + """ + Initialize the String Radix Trie. + Args: + max_cache_size: Maximum number of cached token IDs (triggers GC when exceeded) + gc_threshold_k: GC threshold - nodes with weight_version < (current_version - k) will be removed + tokenizer: Optional tokenizer for converting text to tokens when not found in cache + verbose: Whether to print debug information and tree structure + """ + self.max_cache_size = max_cache_size + self.gc_threshold_k = gc_threshold_k + self.tokenizer = tokenizer + self.verbose = verbose + + # Tree structure + self.root = StringTreeNode() + self.root.string_key = "" + self.root.ref_count = 1 # Root is always protected + + # Cache statistics + self.total_entries = 0 + self.cache_hits = 0 + self.cache_misses = 0 + self.cur_cache_size = 0 # Total number of token IDs across all nodes + + # Thread safety + self._lock = threading.RLock() + + def find_longest_prefix(self, text: str) -> MatchResult: + """ + Find the longest cached prefix for the given text. + Args: + text: Input string to find prefix for + Returns: + MatchResult containing matched prefix, token IDs, logp, and remaining string + """ + with self._lock: + if not text: + return MatchResult("", [], [], [], text, self.root) + + matched_tokens = [] + matched_logp = [] + matched_loss_mask = [] + matched_prefix = "" + current_node = self.root + remaining_text = text + + while remaining_text: + # Find the best matching child that completely matches from start + best_child = None + best_key_len = 0 + + for child_node in current_node.children: + # Only consider complete startswith matches using node's string_key + if remaining_text.startswith(child_node.string_key): + if len(child_node.string_key) > best_key_len: + best_child = child_node + best_key_len = len(child_node.string_key) + + if best_child is None: + # No complete startswith match found + break + + # Move to the best matching child + best_child.touch() + current_node = best_child + matched_prefix += best_child.string_key + remaining_text = remaining_text[best_key_len:] + + # Accumulate tokens, logp, and loss_mask from this node + if best_child.has_value: + matched_tokens.extend(best_child.token_ids) + matched_logp.extend(best_child.logp) + if best_child.loss_mask is not None: + matched_loss_mask.extend(best_child.loss_mask) + else: + # If no loss_mask is stored, create default mask same as logp + matched_loss_mask.extend([1] * len(best_child.token_ids)) + self.cache_hits += 1 + + if not matched_tokens: + self.cache_misses += 1 + + result = MatchResult( + matched_prefix, matched_tokens, matched_logp, matched_loss_mask, remaining_text, current_node + ) + + # Print tree structure if verbose is enabled + if self.verbose: + print("Tree structure after find_longest_prefix:") + self.pretty_print() + + return result + + def insert( + self, + text: str, + token_ids: list[int], + logp: list[float] | None = None, + loss_mask: list[int] | None = None, + weight_version: int | None = None, + ) -> bool: + """ + Insert a string and its corresponding token IDs, log probabilities, and loss mask into the trie. + Args: + text: String to insert + token_ids: Corresponding token IDs + logp: Corresponding log probabilities (must match token_ids length) + loss_mask: Corresponding loss mask for model generation parts (must match token_ids length) + weight_version: Optional weight version for this insertion + Returns: + True if insertion was successful + """ + with self._lock: + if not text or not token_ids: + if self.verbose: + print("[RadixTree] Insertion failed: text or token_ids is empty") + return False + + # Use provided weight version + current_weight_version = weight_version + + # Validate logp consistency + if logp is not None and len(logp) != len(token_ids): + if self.verbose: + print( + f"[WARNING] Logp length {len(logp)} does not match token length {len(token_ids)} for text: {text}" + ) + print(f"[WARNING] Logp: {logp}") + print(f"[WARNING] Token IDs: {token_ids}") + return False + + # Validate loss_mask consistency + if loss_mask is not None and len(loss_mask) != len(token_ids): + if self.verbose: + print( + f"[WARNING] Loss mask length {len(loss_mask)} does not match token length {len(token_ids)} for text: {text}" + ) + print(f"[WARNING] Loss mask: {loss_mask}") + print(f"[WARNING] Token IDs: {token_ids}") + return False + + # If logp is not provided, create default values (0.0) + if logp is None: + logp = [0.0] * len(token_ids) + + # If loss_mask is not provided, create default values (1 for model generation parts) + if loss_mask is None: + loss_mask = [0] * len(token_ids) + + result = self._insert(text, token_ids, logp, loss_mask, current_weight_version) + + # Check if GC should be triggered after insert + if self.cur_cache_size > self.max_cache_size and weight_version is not None: + if self.verbose: + print( + f"[RadixTree] Cache size {self.cur_cache_size} exceeds limit {self.max_cache_size}, triggering GC" + ) + gc_removed = self.gc_by_weight_version(weight_version) + if self.verbose: + print(f"[RadixTree] GC removed {gc_removed} nodes, new cache size: {self.cur_cache_size}") + + # Print tree structure if verbose is enabled + if self.verbose: + print("Tree structure after insert:") + self.pretty_print() + + return result + + def _insert( + self, + text: str, + token_ids: list[int], + logp: list[float], + loss_mask: list[int], + weight_version: int | None = None, + ) -> bool: + """Insert tokens - skip tokens for existing nodes just like we skip text.""" + + current_node = self.root + remaining_text = text + remaining_tokens = token_ids[:] # Copy the tokens list + remaining_logp = logp[:] # Copy the logp list + remaining_loss_mask = loss_mask[:] # Copy the loss_mask list + + # Track all nodes traversed during insert for weight version update + traversed_nodes = [current_node] + new_node = None + + while remaining_text: + # Find best startswith match + best_child = None + best_key_len = 0 + + for child_node in current_node.children: + if remaining_text.startswith(child_node.string_key) and len(child_node.string_key) > best_key_len: + best_child = child_node + best_key_len = len(child_node.string_key) + + if best_child is not None: + # Found existing node - skip its text and tokens + current_node = best_child + traversed_nodes.append(current_node) + remaining_text = remaining_text[best_key_len:] + + # Skip the tokens that this existing node covers + if best_child.has_value: + tokens_to_skip = len(best_child.token_ids) + remaining_tokens = remaining_tokens[tokens_to_skip:] + remaining_logp = remaining_logp[tokens_to_skip:] + remaining_loss_mask = remaining_loss_mask[tokens_to_skip:] + else: + # Create new node for remaining text with remaining tokens + new_node = StringTreeNode() + new_node.parent = current_node + new_node.string_key = remaining_text + + if remaining_tokens: # Only assign if there are tokens left + new_node.token_ids = remaining_tokens + new_node.logp = remaining_logp + new_node.loss_mask = remaining_loss_mask + new_node.touch() + # Increment cache size by number of tokens added + self.cur_cache_size += len(remaining_tokens) + + current_node.children.append(new_node) + traversed_nodes.append(new_node) + self.total_entries += 1 + break + + # If we've traversed the entire text and the last node doesn't have tokens, + # assign remaining tokens to it + if remaining_text == "" and not current_node.has_value: + if remaining_tokens: # Only assign if there are tokens left + current_node.token_ids = remaining_tokens + current_node.logp = remaining_logp + current_node.loss_mask = remaining_loss_mask + current_node.touch() + self.cur_cache_size += len(remaining_tokens) + + # Update weight version for all traversed nodes + if weight_version is not None and new_node: + new_node.weight_version = weight_version + + return True + + def remove(self, text: str) -> bool: + """ + Remove a string and all nodes with this text as prefix from the trie. + Args: + text: String to remove (will also remove all strings starting with this text) + Returns: + True if any removal was performed + """ + with self._lock: + node = self._find_node_by_text(text) + if node: + removed_count = self._clean_node_subtree(node) + + # Print tree structure if verbose is enabled + if self.verbose: + print("Tree structure after remove:") + self.pretty_print() + + return removed_count > 0 + return False + + def _find_node_by_text(self, text: str) -> StringTreeNode | None: + """ + Find node by exact text match. + Args: + text: Text to find + Returns: + Node if found, None otherwise + """ + result = self.find_longest_prefix(text) + if result.matched_prefix == text: + return result.last_node + return None + + def _clean_node_subtree(self, node: StringTreeNode) -> int: + """ + Clean a node and all its descendants. + This is the core cleanup function. + Args: + node: Node to clean (including all descendants) + Returns: + Number of nodes removed + """ + if node == self.root: + return 0 + return self._remove_node_and_descendants(node) + + def _remove_node_and_descendants(self, node: StringTreeNode) -> int: + """ + Remove a node and all its descendants from the trie. + Args: + node: The node to remove along with all its descendants + Returns: + Number of nodes removed + """ + if node == self.root: + # Never remove root node + return 0 + + removed_count = 0 + + # First, recursively remove all descendants + for child in list(node.children): # Create a copy to avoid modification during iteration + removed_count += self._remove_node_and_descendants(child) + + # Count this node if it has data and decrement cache size + if node.has_value: + removed_count += 1 + # Decrement cache size by number of tokens removed + self.cur_cache_size -= len(node.token_ids) + + # Remove this node from its parent + if self._remove_node_from_parent(node): + # Update count for the node structure itself + pass # _remove_node_from_parent already decrements total_entries + + return removed_count + + def _remove_node_from_parent(self, node: StringTreeNode) -> bool: + """Remove a node from its parent's children list.""" + if node.parent and node in node.parent.children: + node.parent.children.remove(node) + self.total_entries -= 1 + return True + return False + + def gc_by_weight_version(self, current_weight_version: int | None = None) -> int: + """ + Perform garbage collection based on weight version. + Remove nodes with weight_version < (current_weight_version - gc_threshold_k). + Args: + current_weight_version: Current weight version to use for GC threshold + Returns: + Number of nodes removed + """ + with self._lock: + if current_weight_version is None: + if self.verbose: + print("[RadixTree GC] No weight version provided, skipping GC") + return 0 + + gc_threshold = current_weight_version - self.gc_threshold_k + if self.verbose: + print( + f"[RadixTree GC] Starting GC with threshold: {gc_threshold} (current_version: {current_weight_version}, k: {self.gc_threshold_k})" + ) + + nodes_to_remove = self._find_outdated_nodes(gc_threshold) + removed_count = 0 + + for node in nodes_to_remove: + # Validate that subtree weight versions are <= parent weight version + self._validate_subtree_weight_versions(node) + removed_count += self._clean_node_subtree(node) + + if self.verbose: + print(f"[RadixTree GC] Completed GC, removed {removed_count} nodes") + + return removed_count + + def _find_outdated_nodes(self, gc_threshold: int) -> list[StringTreeNode]: + """ + Find nodes that should be removed based on weight version threshold. + Uses layer-by-layer traversal - if parent is outdated, children are not checked. + Args: + gc_threshold: Weight version threshold (nodes < this value will be removed) + Returns: + List of nodes to remove + """ + outdated_nodes = [] + + def check_node(node): + if node == self.root: + # Root is never removed, check its children + for child in node.children: + check_node(child) + return + + # Check if this node should be removed + if node.weight_version is not None and node.weight_version <= gc_threshold and node.has_value: + outdated_nodes.append(node) + return # Don't check children since entire subtree will be removed + + # Node is not outdated, check its children + for child in node.children: + check_node(child) + + check_node(self.root) + return outdated_nodes + + def _validate_subtree_weight_versions(self, node: StringTreeNode): + """ + Validate that all nodes in subtree have weight_version <= parent weight_version. + Args: + node: Root node of subtree to validate + """ + + def validate_recursive(current_node, parent_weight_version): + if current_node.weight_version is not None and parent_weight_version is not None: + assert current_node.weight_version <= parent_weight_version, ( + f"Child node weight_version {current_node.weight_version} > " + f"parent weight_version {parent_weight_version}" + ) + + # Recursively validate children + for child in current_node.children: + validate_recursive(child, current_node.weight_version) + + # Start validation from the node itself + validate_recursive(node, node.weight_version) + + def get_stats(self) -> dict[str, Any]: + """Get cache statistics.""" + with self._lock: + total_requests = self.cache_hits + self.cache_misses + hit_rate = self.cache_hits / total_requests if total_requests > 0 else 0 + + return { + "total_entries": self.total_entries, + "cache_hits": self.cache_hits, + "cache_misses": self.cache_misses, + "hit_rate": hit_rate, + "max_cache_size": self.max_cache_size, + "cur_cache_size": self.cur_cache_size, + "gc_threshold_k": self.gc_threshold_k, + } + + def clear(self): + """Clear all entries from the trie.""" + with self._lock: + self.root = StringTreeNode() + self.root.string_key = "" + self.root.ref_count = 1 + self.total_entries = 0 + self.cache_hits = 0 + self.cache_misses = 0 + self.cur_cache_size = 0 + + def pretty_print(self): + """Print the trie structure in a readable format.""" + print("String Radix Trie Structure:") + print("=" * 50) + self._print_node(self.root, 0) + print("=" * 50) + stats = self.get_stats() + for key, value in stats.items(): + print(f"{key}: {value}") + + def _print_node(self, node: StringTreeNode, depth: int): + """Recursively print node structure.""" + indent = " " * depth + key_repr = repr(node.string_key) if node.string_key else "" + token_info = "" + if node.has_value: + token_info = f" -> tokens: {node.token_ids}" + if node.logp: + token_info += f", logp: {[round(p, 3) for p in node.logp]}" + if node.loss_mask: + token_info += f", loss_mask: {node.loss_mask}" + access_info = f" (accessed: {node.access_count}, ref: {node.ref_count})" + + print(f"{indent}{key_repr}{token_info}{access_info}") + + for child in node.children: + self._print_node(child, depth + 1) + + def retrieve_from_text(self, text: str, return_logprob: bool = True): + """ + Get tokens from text by looking up in radix tree or using tokenizer. + Also fetches weight version from worker during this operation. + Args: + text: Input text to get tokens for + return_logprob: If True, also return log probabilities + Returns: + List of token IDs corresponding to the input text if return_logprob is False. + Tuple of (token_ids, logp) if return_logprob is True. + """ + # Call find_longest_prefix to get the match result + result = self.find_longest_prefix(text) + + # If we have a match and it covers the entire text, return the tokens + if result.matched_prefix and result.token_ids: + additional_tokens = self.tokenizer(result.remaining_string, add_special_tokens=False)["input_ids"] + return ( + result.token_ids + additional_tokens, + ( + result.logp + len(additional_tokens) * [0.0] + if return_logprob + else [0] * len(result.token_ids + additional_tokens) + ), + result.loss_mask + len(additional_tokens) * [0], + ) + # If result is empty and input text is not empty, tokenize with tokenizer + # This is needed because we cannot get the prompt token id from engine response + # We have to manually insert the text and token into the tree + if self.tokenizer and text: + # Tokenize the text using the provided tokenizer + tokens = self.tokenizer(text, add_special_tokens=False)["input_ids"] + # Insert the text and tokens into the tree + self.insert(text, tokens) + # Return the tokens + return (tokens, [0.0] * len(tokens), [0] * len(tokens)) + else: + raise ValueError("Tokenizer or input text can't be empty") + + +# Example usage and testing +if __name__ == "__main__": + # Create trie instance for testing + trie = StringRadixTrie(max_cache_size=100, verbose=True) + + # Test token retrieval + print("\nTesting token retrieval:") + test_tokens = trie.retrieve_from_text("Hello world") + print(f"Tokens for 'Hello world': {test_tokens}") + + # Example usage with simplified insert + test_cases = [ + ("Hello world", [1, 2, 3], [-0.1, -0.2, -0.3]), + ("Hello", [1, 2], [-0.1, -0.2]), + ("Hi there", [4, 5, 6], [-0.4, -0.5, -0.6]), + ] + + # Insert test data with weight version and loss masks + print("Inserting test data...") + for text, tokens, logp in test_cases: + # Create loss_mask to match tokens length, 1 for model generation parts + loss_mask = [1] * len(tokens) + success = trie.insert(text, tokens, logp, loss_mask, weight_version=1) + print(f"Inserted '{text}' -> {tokens}: {success}") + + print("\nTrie structure:") + trie.pretty_print() + + # Test prefix matching + print("\nTesting prefix matching:") + test_queries = [ + "Hello world!", # Should match "Hello world" completely + "Hello everyone", # Should match "Hello" only + "Hi there", # Should match "Hi" only + "How are you doing?", # Should match "How are you" completely + "Goodbye", # Should not match anything + "Hell", # Should not match anything (not complete startswith) + ] + + for query in test_queries: + result = trie.find_longest_prefix(query) + print(f"Query: '{query}'") + print( + f" Matched: '{result.matched_prefix}' -> tokens: {result.token_ids}, logp: {result.logp}, loss_mask: {result.loss_mask}" + ) + print(f" Remaining: '{result.remaining_string}'") + print() + + # Test removal + print("Testing removal:") + removed = trie.remove("Hello") + print(f"Removed 'Hello': {removed}") + + result = trie.find_longest_prefix("Hello world") + print( + f"After removal - 'Hello world' -> matched: '{result.matched_prefix}', tokens: {result.token_ids}, logp: {result.logp}, loss_mask: {result.loss_mask}" + ) + + # Show final stats + print("\nFinal statistics:") + stats = trie.get_stats() + for key, value in stats.items(): + print(f"{key}: {value}") + + # Test GC with weight version + print("\nTesting GC with weight version 5:") + gc_removed = trie.gc_by_weight_version(5) + print(f"GC removed {gc_removed} nodes") diff --git a/slime/router/middleware_hub/radix_tree_middleware.py b/slime/router/middleware_hub/radix_tree_middleware.py new file mode 100644 index 0000000000000000000000000000000000000000..fbb9ea69a80347657a0c83334e7c2fb19ea19e3c --- /dev/null +++ b/slime/router/middleware_hub/radix_tree_middleware.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +from time import sleep + +from fastapi.responses import JSONResponse +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request +from starlette.responses import Response +from transformers import AutoTokenizer + +from .radix_tree import StringRadixTrie + +# Hop-by-hop headers that should not be forwarded +HOP_BY_HOP = { + "content-length", + "transfer-encoding", + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailers", + "upgrade", +} + + +def _filter_headers(headers): + """Filter out hop-by-hop headers that should not be forwarded.""" + return {k: v for k, v in headers.items() if k.lower() not in HOP_BY_HOP} + + +async def _materialize_response(resp): + """Convert streaming-like Response into a regular Response/JSONResponse safely.""" + # Collect all bytes from the streaming response + body = b"" + async for chunk in resp.body_iterator: + body += chunk + + # Try to parse as JSON based on content-type + ct = resp.headers.get("content-type", "") + headers = _filter_headers(resp.headers) + + if "application/json" in ct: + # If it's JSON, try to parse and return as JSONResponse + try: + data = json.loads(body.decode("utf-8")) + return JSONResponse(content=data, status_code=resp.status_code, headers=headers) + except Exception: + # JSON parsing failed, fall back to raw bytes + pass + + # Other types: return as raw bytes (without content-length) + return Response(content=body, status_code=resp.status_code, headers=headers, media_type=resp.media_type) + + +class RadixTreeMiddleware(BaseHTTPMiddleware): + def __init__(self, app, *, router): + super().__init__(app) + self.router = router + self.args = router.args + self.tokenizer = AutoTokenizer.from_pretrained(self.args.hf_checkpoint, trust_remote_code=True) + self.radix_tree = StringRadixTrie(max_cache_size=10000, tokenizer=self.tokenizer, verbose=False) + self.router.radix_tree = self.radix_tree + + async def dispatch(self, request: Request, call_next): + + path = request.url.path + + if path != "/generate": + return await call_next(request) + + request_json = await request.json() + if "text" in request_json: + input_text = request_json.pop("text", "") + elif "input_ids" in request_json: + input_text = self.tokenizer.decode(request_json["input_ids"]) + else: + input_text = None + if not input_text: + return await call_next(request) + input_tokens, input_logprobs, input_loss_mask = self.radix_tree.retrieve_from_text( + input_text, return_logprob=True + ) + request_json["input_tokens"] = input_tokens + request_json["stream"] = False + request._json = request_json + + response_data = None + for _ in range(5): + response = await call_next(request) + + # If upstream returned a streaming response, materialize it to avoid Content-Length issues + if response.__class__.__name__ == "_StreamingResponse": + response = await _materialize_response(response) + # Try to parse JSON from the current response for meta inspection + try: + if hasattr(response, "body") and isinstance(response.body, (bytes, bytearray)): + response_data = json.loads(response.body.decode("utf-8")) + elif hasattr(response, "content") and isinstance(response.content, (dict, list)): + response_data = response.content # JSONResponse.content is already a dict/list + except Exception: + response_data = None + + if ( + isinstance(response_data, dict) + and "meta_info" in response_data + and "finish_reason" in response_data["meta_info"] + and response_data["meta_info"]["finish_reason"]["type"] != "abort" + ): + break + # await 30 seconds for aborted responses + sleep(30) + + if isinstance(response_data, dict) and "text" in response_data and "output_ids" in response_data: + generated_text = response_data["text"] + + full_text = input_text + generated_text + if full_text: + try: + if "output_token_logprobs" in response_data.get("meta_info", {}): + generated_token_logprobs = [ + item[0] for item in response_data["meta_info"]["output_token_logprobs"] + ] + generated_token_ids = [item[1] for item in response_data["meta_info"]["output_token_logprobs"]] + full_logprobs = input_logprobs + generated_token_logprobs + full_token_ids = input_tokens + generated_token_ids + full_loss_mask = input_loss_mask + [1] * len(generated_token_ids) + self.radix_tree.insert( + full_text, + full_token_ids, + full_logprobs, + full_loss_mask, + weight_version=response_data["meta_info"]["weight_version"], + ) + else: + generated_token_ids = self.tokenizer(generated_text, add_special_tokens=False)["input_ids"] + full_token_ids = input_tokens + generated_token_ids + full_loss_mask = input_loss_mask + [1] * len(generated_token_ids) + self.radix_tree.insert( + full_text, + full_token_ids, + None, + full_loss_mask, + weight_version=response_data["meta_info"]["weight_version"], + ) + + if getattr(self.router, "verbose", False): + print(f"[slime-router] Successfully cached trajectory with {len(full_token_ids)} tokens") + except Exception as e: + if getattr(self.router, "verbose", False): + print(f"[slime-router] Warning: Failed to cache trajectory: {e}") + return response diff --git a/slime/router/router.py b/slime/router/router.py new file mode 100644 index 0000000000000000000000000000000000000000..d79821140bed98f67c161145e7042f80ed5f18e2 --- /dev/null +++ b/slime/router/router.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import json + +import httpx +import uvicorn +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse +from starlette.responses import Response + +from slime.utils.misc import load_function + + +def run_router(args): + """ + Run the Slime router with the specified configuration. + """ + # Initialize the router with tokenizer and lazy worker initialization + slime_router = SlimeRouter(args, verbose=False) + + # Start the server + uvicorn.run(slime_router.app, host=args.sglang_router_ip, port=args.sglang_router_port, log_level="info") + + +class SlimeRouter: + def __init__(self, args, verbose=False): + """Initialize the slime-router with SGLang router address""" + self.args = args + self.verbose = verbose + + self.app = FastAPI() + + # Worker information + self.worker_urls: dict[str, int] = {} + self.max_weight_version = None + + max_connections = getattr(args, "slime_router_max_connections", None) + if max_connections is None: + max_connections = ( + args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine + ) + + timeout = getattr(args, "slime_router_timeout", None) + + self.client = httpx.AsyncClient( + limits=httpx.Limits(max_connections=max_connections), + timeout=httpx.Timeout(timeout), + ) + + self._setup_routes() + + for middleware_path in args.slime_router_middleware_paths or []: + if self.verbose: + print(f"[slime-router] Loading middleware from: {middleware_path}") + middleware = load_function(middleware_path) + self.app.add_middleware(middleware, router=self) + + def _setup_routes(self): + """Setup all the HTTP routes""" + # sglang-router api + self.app.post("/add_worker")(self.add_worker) + self.app.get("/list_workers")(self.list_workers) + self.app.post("/retrieve_from_text")(self.retrieve_from_text) + # Catch-all route for proxying to SGLang - must be registered LAST + self.app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])(self.proxy) + + async def health_check(self, request: Request): + # TODO: do health check in background + pass + + async def proxy(self, request: Request, path: str): + """Proxy all other requests to the SGLang router""" + # Forward all other paths to SGLang router + worker_url = self._use_url() + url = f"{worker_url}/{path}" + + # Get request body and headers + body = await request.body() + headers = dict(request.headers) + + try: + response = await self.client.request(request.method, url, content=body, headers=headers) + # Eagerly read content so we can return JSON (not streaming) + content = await response.aread() + content_type = response.headers.get("content-type", "") + try: + # Prefer parsing JSON if possible + data = json.loads(content) + return JSONResponse( + content=data, + status_code=response.status_code, + headers=dict(response.headers), + ) + except Exception: + # Fall back to raw body with original content type + return Response( + content=content, + status_code=response.status_code, + headers=dict(response.headers), + media_type=content_type or None, + ) + + finally: + self._finish_url(worker_url) + + async def add_worker(self, request: Request): + """Add a new worker to the router. + Supports providing the URL via query string or JSON body. + Examples: + - POST /add_worker?url=http://127.0.0.1:10090 + - POST /add_worker with body {"url": "http://127.0.0.1:10090"} + """ + # 1) Prefer query param + worker_url = request.query_params.get("url") or request.query_params.get("worker_url") + + # 2) Fallback to JSON body + if not worker_url: + body = await request.body() + payload = json.loads(body) if body else {} + worker_url = payload.get("url") or payload.get("worker_url") + + if not worker_url: + return JSONResponse( + status_code=400, content={"error": "worker_url is required (use query ?url=... or JSON body)"} + ) + + # Add if new, keep a simple request count per worker + if worker_url not in self.worker_urls: + self.worker_urls[worker_url] = 0 + if self.verbose: + print(f"[slime-router] Added new worker: {worker_url}") + + return {"status": "success", "worker_urls": self.worker_urls} + + async def list_workers(self, request: Request): + """List all registered workers""" + return {"urls": list(self.worker_urls.keys())} + + async def retrieve_from_text(self, request: Request): + """Get token information from text input""" + body = await request.body() + payload = json.loads(body) if body else {} + + text = payload.get("text", "") + + # Use radix tree's retrieve_from_text method (no need to fetch weight version here) + token_ids, logp, loss_mask = self.radix_tree.retrieve_from_text(text, return_logprob=True) + + # Handle the result based on whether logp was requested + result = { + "tokens": token_ids, # token IDs + "response": text, # The input text + "loss_mask": loss_mask, # Loss mask for the tokens + "token_length": len(token_ids), + "loss_mask_length": len(loss_mask), + "rollout_logp": logp, + } + + return result + + def _use_url(self): + """Select a worker URL using round-robin strategy""" + assert len(self.worker_urls) > 0, "No workers available" + + # get the url with mininal count + url = min(self.worker_urls, key=self.worker_urls.get) + self.worker_urls[url] += 1 + return url + + def _finish_url(self, url): + """Mark the request to the given URL as finished""" + assert url in self.worker_urls, f"URL {url} not recognized" + self.worker_urls[url] -= 1 + assert self.worker_urls[url] >= 0, f"URL {url} count went negative" + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--host", type=str, default="0.0.0.0") + parser.add_argument("--port", type=int, default=30000) + parser.add_argument("--sglang-host", type=str, required=True) + parser.add_argument("--sglang-port", type=int, required=True) + parser.add_argument("--tokenizer-name", type=str, help="Name of the tokenizer to use for tokenization") + parser.add_argument("--verbose", action="store_true", help="Enable verbose output") + + args = parser.parse_args() + + # Run the router + run_router(args) diff --git a/slime/utils/__init__.py b/slime/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..05f902c8f74c6bb5ded1dd5af22cf9d91310c9b9 --- /dev/null +++ b/slime/utils/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Utility package root for Slime.""" diff --git a/slime/utils/__pycache__/__init__.cpython-312.pyc b/slime/utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7fe9258626cc18d6ef3a31542c49d7d4648e01cc Binary files /dev/null and b/slime/utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/arguments.cpython-312.pyc b/slime/utils/__pycache__/arguments.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..79eaa37dfc180736db6f7cc4e3a30cf3d3059ce4 Binary files /dev/null and b/slime/utils/__pycache__/arguments.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/async_utils.cpython-312.pyc b/slime/utils/__pycache__/async_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50409e697501dd4e2f37ad71ea6eadf4b09fbb33 Binary files /dev/null and b/slime/utils/__pycache__/async_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/context_utils.cpython-312.pyc b/slime/utils/__pycache__/context_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3f83a5885679014a8f0dda96c0981c795fbd8d98 Binary files /dev/null and b/slime/utils/__pycache__/context_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/data.cpython-312.pyc b/slime/utils/__pycache__/data.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c49ef40326455c216c911fe7371dca4e47cef751 Binary files /dev/null and b/slime/utils/__pycache__/data.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/distributed_utils.cpython-312.pyc b/slime/utils/__pycache__/distributed_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..12b0728f3f9b625197d5930c02d0418aba8fb3ff Binary files /dev/null and b/slime/utils/__pycache__/distributed_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/eval_config.cpython-312.pyc b/slime/utils/__pycache__/eval_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..05e0b0fe4763be71a4e65a4863c0b5265ceb9bd8 Binary files /dev/null and b/slime/utils/__pycache__/eval_config.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/flops_utils.cpython-312.pyc b/slime/utils/__pycache__/flops_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4cc92d3089caf7ac790d548b4799d7511df896a5 Binary files /dev/null and b/slime/utils/__pycache__/flops_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/fp8_kernel.cpython-312.pyc b/slime/utils/__pycache__/fp8_kernel.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9c585aacfb687ef48bd53a81a3a8971815554dfa Binary files /dev/null and b/slime/utils/__pycache__/fp8_kernel.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/health_monitor.cpython-312.pyc b/slime/utils/__pycache__/health_monitor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..39abd0e9d2a41fb720ded9df5db7fbbaff1b39c7 Binary files /dev/null and b/slime/utils/__pycache__/health_monitor.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/http_utils.cpython-312.pyc b/slime/utils/__pycache__/http_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aa0979a464cc6488c6d8e734132b6990587e332f Binary files /dev/null and b/slime/utils/__pycache__/http_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/iter_utils.cpython-312.pyc b/slime/utils/__pycache__/iter_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c19f7a832a8364b3b8ec21e67db9090547d5dc38 Binary files /dev/null and b/slime/utils/__pycache__/iter_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/logging_utils.cpython-312.pyc b/slime/utils/__pycache__/logging_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a03f6698c4240eaa6cdd93c40b17bfdb4dde607 Binary files /dev/null and b/slime/utils/__pycache__/logging_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/mask_utils.cpython-312.pyc b/slime/utils/__pycache__/mask_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..58f10c36ba34499ae5727fa73cd2c6897ac4c8de Binary files /dev/null and b/slime/utils/__pycache__/mask_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/megatron_bridge_utils.cpython-312.pyc b/slime/utils/__pycache__/megatron_bridge_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0b67a2e39970641cdb5c19d414cc1591cccbdfe4 Binary files /dev/null and b/slime/utils/__pycache__/megatron_bridge_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/memory_utils.cpython-312.pyc b/slime/utils/__pycache__/memory_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c49048c63fd48d03a9b16a62fe25b62813c115d4 Binary files /dev/null and b/slime/utils/__pycache__/memory_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/metric_checker.cpython-312.pyc b/slime/utils/__pycache__/metric_checker.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a127d4d5e88f1b4aad09dfb27cfd860da63c120d Binary files /dev/null and b/slime/utils/__pycache__/metric_checker.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/metric_utils.cpython-312.pyc b/slime/utils/__pycache__/metric_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff9ac0689fff233d8f343152585679676160e69f Binary files /dev/null and b/slime/utils/__pycache__/metric_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/misc.cpython-312.pyc b/slime/utils/__pycache__/misc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..303014f682578bc65f1a18a002e9d75428dad7e3 Binary files /dev/null and b/slime/utils/__pycache__/misc.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/ppo_utils.cpython-312.pyc b/slime/utils/__pycache__/ppo_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c3714acd39f07a234a2dd413db23b1101d5c426 Binary files /dev/null and b/slime/utils/__pycache__/ppo_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/processing_utils.cpython-312.pyc b/slime/utils/__pycache__/processing_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d4766094973122d7649863fa24a1cdc21b367876 Binary files /dev/null and b/slime/utils/__pycache__/processing_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/profile_utils.cpython-312.pyc b/slime/utils/__pycache__/profile_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..deeadc84fd50c40173605841394ef184d7beb260 Binary files /dev/null and b/slime/utils/__pycache__/profile_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/ray_utils.cpython-312.pyc b/slime/utils/__pycache__/ray_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..362f5b976ff239446fbda6f121d8d3466f9e817f Binary files /dev/null and b/slime/utils/__pycache__/ray_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/reloadable_process_group.cpython-312.pyc b/slime/utils/__pycache__/reloadable_process_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..955f377d7d05640bd5e47958b41ff0e862f1789d Binary files /dev/null and b/slime/utils/__pycache__/reloadable_process_group.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/routing_replay.cpython-312.pyc b/slime/utils/__pycache__/routing_replay.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e4fd446d1eea73803d437ff3b45203f3c631fe55 Binary files /dev/null and b/slime/utils/__pycache__/routing_replay.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/seqlen_balancing.cpython-312.pyc b/slime/utils/__pycache__/seqlen_balancing.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce9b4bc66bba1ce7e9b044c439dff707742e40f3 Binary files /dev/null and b/slime/utils/__pycache__/seqlen_balancing.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/tensor_backper.cpython-312.pyc b/slime/utils/__pycache__/tensor_backper.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3542d3d195462651ad5768c890b521f35b3d8b34 Binary files /dev/null and b/slime/utils/__pycache__/tensor_backper.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/tensorboard_utils.cpython-312.pyc b/slime/utils/__pycache__/tensorboard_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1384df193d93d2195381bb7b9e17fffd9c0d231c Binary files /dev/null and b/slime/utils/__pycache__/tensorboard_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/timer.cpython-312.pyc b/slime/utils/__pycache__/timer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b6759923cb9a1780c69f7004d99ce4226c4a6348 Binary files /dev/null and b/slime/utils/__pycache__/timer.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/tracking_utils.cpython-312.pyc b/slime/utils/__pycache__/tracking_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e4aa692192f7bb6b7ac64865673b8b36618bce9f Binary files /dev/null and b/slime/utils/__pycache__/tracking_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/train_dump_utils.cpython-312.pyc b/slime/utils/__pycache__/train_dump_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d87e8b417e59963cccd1bde18533c0a1ff7b828c Binary files /dev/null and b/slime/utils/__pycache__/train_dump_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/train_metric_utils.cpython-312.pyc b/slime/utils/__pycache__/train_metric_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e876bd85026736c383d3318295af5e1267986005 Binary files /dev/null and b/slime/utils/__pycache__/train_metric_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/typer_utils.cpython-312.pyc b/slime/utils/__pycache__/typer_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4207dfc1857766c7296245c8ced468947f0104fa Binary files /dev/null and b/slime/utils/__pycache__/typer_utils.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/types.cpython-312.pyc b/slime/utils/__pycache__/types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..58c7a8fcac5c67c79a1b31a0a6717a401857735a Binary files /dev/null and b/slime/utils/__pycache__/types.cpython-312.pyc differ diff --git a/slime/utils/__pycache__/wandb_utils.cpython-312.pyc b/slime/utils/__pycache__/wandb_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fbf1d12f9620d7b765418c45a1a09ea6eb657727 Binary files /dev/null and b/slime/utils/__pycache__/wandb_utils.cpython-312.pyc differ diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py new file mode 100644 index 0000000000000000000000000000000000000000..327f866ab13137aa98da23c1a861918c2c3666a4 --- /dev/null +++ b/slime/utils/arguments.py @@ -0,0 +1,1737 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import json +import logging +import os +from typing import Any + +import yaml +from sglang_router.launch_router import RouterArgs +from transformers import AutoConfig + +from slime.backends.sglang_utils.arguments import add_sglang_arguments +from slime.backends.sglang_utils.arguments import validate_args as sglang_validate_args +from slime.utils.eval_config import EvalDatasetConfig, build_eval_dataset_configs, ensure_dataset_list +from slime.utils.logging_utils import configure_logger + +logger = logging.getLogger(__name__) + + +def reset_arg(parser, name, **kwargs): + """ + Reset the default value of a Megatron argument. + :param parser: The argument parser. + :param name: The name of the argument to reset. + :param default: The new default value. + """ + for action in parser._actions: + if name in action.option_strings: + if "default" in kwargs: + action.default = kwargs["default"] + break + else: + parser.add_argument(name, **kwargs) + + +def get_slime_extra_args_provider(add_custom_arguments=None): + def add_slime_arguments(parser): + # Ray + def add_cluster_arguments(parser): + parser.add_argument("--actor-num-nodes", type=int, default=1, help="Number of nodes for training actor") + parser.add_argument( + "--actor-num-gpus-per-node", type=int, default=8, help="Number of gpus per node for training actor" + ) + parser.add_argument( + "--critic-num-nodes", type=int, default=None, help="Number of nodes for training actor" + ) + parser.add_argument( + "--critic-num-gpus-per-node", type=int, default=None, help="Number of gpus per node for training actor" + ) + + parser.add_argument( + "--rollout-num-gpus", + type=int, + default=None, + help=( + "Number of GPUs for inference. Note that when using --colocate, " + "i.e. the training and the inference engines are on the same gpus, this param will be ignored and will be set as " + "actor_num_gpus_per_node * actor_num_nodes." + ), + ) + parser.add_argument( + "--rollout-num-gpus-per-engine", + type=int, + default=1, + help="Number of GPUs per inference engine, just like the tp_size in sglang.", + ) + parser.add_argument( + "--num-gpus-per-node", + type=int, + default=8, + help=( + "Number of gpus per node for rollout." + "Notice: If you are going to use less than 8 gpus per node under colocate mode, you should set this number." + ), + ) + parser.add_argument( + "--colocate", + action="store_true", + default=False, + help=( + "Whether to colocate the inference engines and the actor. " + "Turning this on will also set --offload to true." + ), + ) + parser.add_argument( + "--offload", + action="store_true", + default=False, + help=("Equivalent to --offload-train + --offload-rollout. "), + ) + parser.add_argument( + "--offload-train", + action=argparse.BooleanOptionalAction, + help=( + "Whether to offload the training actor to CPU during training. " + "This will always be true when --colocate is set." + ), + ) + parser.add_argument( + "--offload-rollout", + action=argparse.BooleanOptionalAction, + help=( + "Whether to offload the rollout generator to CPU during training. " + "This will always be true when --colocate is set." + ), + ) + + reset_arg(parser, "--distributed-backend", type=str, default="nccl") + reset_arg(parser, "--distributed-timeout-minutes", type=int, default=10) + + return parser + + def add_train_arguments(parser): + parser.add_argument( + "--train-backend", + type=str, + choices=["megatron", "fsdp"], + default="megatron", + help="The backend for training.", + ) + parser.add_argument( + "--true-on-policy-mode", + action="store_true", + default=False, + help="Whether to enable true-on-policy mode.", + ) + parser.add_argument( + "--train-env-vars", + type=json.loads, + default="{}", + help="Extra environment variables for training process, e.g. PyTorch memory management ones.", + ) + parser.add_argument( + "--train-memory-margin-bytes", + type=int, + default=1024**3, + help="Add margin for train memory allocation. By default we will reserve 1GB as margin.", + ) + parser.add_argument( + "--disable-weights-backuper", + action="store_false", + dest="enable_weights_backuper", + help="Whether to disable weights backuper to save host memory.", + ) + parser.add_argument( + "--megatron-to-hf-mode", + choices=["raw", "bridge"], + default="raw", + help="The method to convert megatron weights to hugging face weights for SGLang.", + ) + parser.add_argument( + "--recompute-loss-function", + action="store_true", + help="Whether to disable recompute loss function to save memory during training.", + ) + parser.add_argument( + "--log-probs-chunk-size", type=int, default=-1, help="Chunk size to compute log probs to save memory" + ) + + parser.add_argument( + "--use-ema", + action="store_true", + default=False, + help="Whether to use EMA to save memory during training.", + ) + parser.add_argument( + "--ema-decay", + type=float, + default=0.9999, + help="Decay rate for EMA.", + ) + parser.add_argument( + "--ema-update-interval", + type=int, + default=1, + help="Update EMA every N training steps. Default is 1 (every step).", + ) + + return parser + + # rollout + def add_rollout_arguments(parser): + parser.add_argument( + "--hf-checkpoint", + type=str, + default=None, + help=( + "The huggingface checkpoint of the trained model. " + "This is used to initialize sglang and also provide the tokenizer. " + "Note that, we will always update the parameters in sglang with that of megatron before training, " + "so you only need to provide a huggingface checkpoint that has the same architecture as the model you want to train. " + "It doesn't necessary need to contain the most up-to-date parameters." + ), + ) + parser.add_argument( + "--model-name", + type=str, + default=None, + help=( + "The name of the model, this is used to convert the megatron weights into huggingface format. " + "If not set, we will use `type(AutoConfig.from_pretrained(args.hf_checkpoint)).__name__.lower()` as model_name. " + "Also, sometimes this will help alleviate the bug that transformers cannot find certain model." + ), + ) + parser.add_argument( + "--rollout-function-path", + type=str, + default="slime.rollout.sglang_rollout.generate_rollout", + help=( + "Path to the rollout generation function." + "You should use this model to create your own custom rollout function, " + "and then set this to the path of your custom rollout function. " + "The signature of the function should be " + "`def generate_rollout(args, rollout_id, *, evaluation=False) -> list[list[Sample]]`" + "and within the output sample, you should at least set `tokens`, `response_length`, `reward` " + "and `truncated`." + ), + ) + parser.add_argument( + "--rollout-temperature", + type=float, + default=1.0, + help="the temperature for the inference engine during rollout.", + ) + parser.add_argument( + "--rollout-top-p", type=float, default=1.0, help="the top-p for the inference engine during rollout." + ) + parser.add_argument( + "--rollout-top-k", type=int, default=-1, help="the top-k for the inference engine during rollout." + ) + parser.add_argument( + "--rollout-max-context-len", + type=int, + default=None, + help=( + "The maximum context size for the inference engine during rollout." + "It should no exceed the `max_position_embeddinds` in Huggingface model's `config.json`" + ), + ) + parser.add_argument( + "--rollout-max-prompt-len", + type=int, + default=None, + help=( + "The maximum length of the prompt for the inference engine during rollout. " + "If set, we will filter out the long prompts during initialization of the global dataset. " + "This is not recommended if the dataset is large." + ), + ) + parser.add_argument( + "--rollout-max-response-len", + type=int, + default=None, + help=( + "The maximum length of the response for the inference engine during rollout. " + "It is basically `max_tokens` in sglang." + ), + ) + parser.add_argument( + "--rollout-skip-special-tokens", + action="store_true", + default=False, + help=( + "Whether to skip special tokens in the response during rollout. " + "This is useful when you want to use the response as a prompt for the next rollout." + ), + ) + parser.add_argument( + "--rollout-stop", + type=str, + nargs="+", + default=None, + help=( + "The stop words for the inference engine during rollout. " + "It can be a list of strings or a single string. " + "It may be hard to pass special tokens in command line, in that case rollout_stop_token_ids can be used." + ), + ) + parser.add_argument( + "--rollout-stop-token-ids", + type=int, + nargs="+", + default=None, + help=( + "The stop token ids for the inference engine during rollout. " + "It can be a list of integers or a single integer." + ), + ) + parser.add_argument( + "--rollout-shuffle", + action="store_true", + default=False, + help=("Whether to shuffle the prompts during rollout."), + ) + parser.add_argument( + "--rollout-seed", + type=int, + default=42, + help=( + "The seed for the random number generator during rollout. " + "This is used to shuffle the prompts and also for the random sampling of the prompts." + ), + ) + + # sampling + parser.add_argument( + "--over-sampling-batch-size", + type=int, + default=None, + help=( + "This defines the granularity of the sampling batch in the rollout function. " + "When the number of available samples falls below the target, a sampling " + "operation of size over_sampling_batch_size will be triggered." + "Regardless of whether partial rollout is used or filters are applied, " + "the sampling granularity is always determined by this value. " + "If this value is None, rollout_batch_size will be used as the default over_sampling_batch_size." + ), + ) + parser.add_argument( + "--dynamic-sampling-filter-path", + type=str, + default=None, + help=( + "This is the filter function for dynamic sampling. " + "It should be able to judge whether the result of a prompt should be selected or not." + "We will do dynamic filter for sampling as in DAPO. e.g. not all correct or all wrong samples." + "You could use `slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std` as an example." + ), + ) + + # partial rollout + parser.add_argument( + "--partial-rollout", + action="store_true", + default=False, + help=( + "Whether to use partial rollout. " + "If set, the unfinished samples during dynamic sampling will be recycled back to data buffer. " + "This is useful for long responses." + ), + ) + parser.add_argument( + "--mask-offpolicy-in-partial-rollout", + action="store_true", + default=False, + help=( + "Whether to mask previous generation in partial rollout. " + "If set, only on-policy generated tokens will be used in training" + ), + ) + parser.add_argument( + "--progressive-prefix-start-ratio", + type=float, + default=0.0, + help=( + "Starting prefix ratio for progressive prefix distillation. " + "Fraction of the teacher's response tokens to prepend to the student's prompt " + "at the start of training. Linearly decays to --progressive-prefix-end-ratio " + "over --progressive-prefix-steps rollout steps. Default 0.0 (disabled)." + ), + ) + parser.add_argument( + "--progressive-prefix-end-ratio", + type=float, + default=0.0, + help=( + "Ending prefix ratio for progressive prefix distillation. " + "Target prefix ratio reached at rollout step --progressive-prefix-steps. Default 0.0." + ), + ) + parser.add_argument( + "--progressive-prefix-steps", + type=int, + default=300, + help=( + "Number of rollout steps over which the prefix ratio linearly decays " + "from --progressive-prefix-start-ratio to --progressive-prefix-end-ratio. Default 300." + ), + ) + parser.add_argument( + "--custom-generate-function-path", + type=str, + default=None, + help=( + "Only substitue the `def generate(args, sample, sampling_params)` function within the example rollout function. " + "This should be useful if you need to implement some special rollout logic, e.g. multi-turn, function calling." + ), + ) + parser.add_argument( + "--custom-rollout-log-function-path", + type=str, + default=None, + help=( + "The custom function for logging rollout data. The signature of the functions is: " + "def log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time) -> bool. " + "The return value indicates whether to skip the default logging. " + ), + ) + parser.add_argument( + "--custom-eval-rollout-log-function-path", + type=str, + default=None, + help=( + "The custom function for logging eval rollout data. " + "def log_eval_rollout_data(rollout_id, args, data, extra_metrics) -> bool. " + "The return value indicates whether to skip the default logging. " + ), + ) + + parser.add_argument( + "--buffer-filter-path", + type=str, + default=None, + help=( + "Path to the buffer filter function. " + "It should be able to select the samples in the buffer. " + "The function should take list[list[Sample]] and return list[list[Sample]]." + ), + ) + # update weight + parser.add_argument( + "--update-weight-buffer-size", + type=int, + default=512 * 1024**2, + help=( + "buffer size for update weight, in bytes. " + "This is used for updating weights by chunk and should be useful for MoE models." + ), + ) + parser.add_argument( + "--update-weights-interval", + type=int, + default=1, + help="Interval for updating the weights", + ) + parser.add_argument( + "--keep-old-actor", + action="store_true", + help="Whether to keep the rollout model on training process", + ) + + parser.add_argument( + "--rollout-data-postprocess-path", + type=str, + default=None, + help=( + "The called after we have all the rollout data including log_probs. " + "It may be helpful for updating loss mask." + ), + ) + parser.add_argument( + "--rollout-external", + action="store_true", + default=False, + help="Use external SGLang instances instead of launching them inside the framework.", + ) + parser.add_argument( + "--rollout-external-engine-addrs", + type=str, + default=None, + nargs="+", + help="Address and ports of the external engines.", + ) + return parser + + def add_fault_tolerance_arguments(parser): + parser.add_argument( + "--use-fault-tolerance", + action="store_true", + default=False, + help="Whether to enable the fault tolerance function during rollout.", + ) + parser.add_argument( + "--rollout-health-check-interval", + type=float, + default=30.0, + help="Interval in seconds between rollout engine /health_generate checks during generate/eval.", + ) + parser.add_argument( + "--rollout-health-check-timeout", + type=float, + default=30.0, + help="Timeout in seconds to wait for a rollout engine /health_generate response before killing it.", + ) + parser.add_argument( + "--rollout-health-check-first-wait", + type=float, + default=0, + help="Initial grace period (in seconds) before starting health checks. This allows time for model compilation and initialization. Increase this value significantly when using deepgemm.", + ) + return parser + + # data + def add_data_arguments(parser): + # dataset + # TODO: maybe add an num_epoch and calculate the num_rollout from buffer + parser.add_argument( + "--num-rollout", + type=int, + default=None, + help="Number of rollout steps. If not set, we will calculate the number of rollout steps from the dataset size.", + ) + parser.add_argument( + "--num-epoch", + type=int, + default=None, + help=( + "Number of epochs for the training. " + "This is used to calculate the number of rollout steps from the dataset size. " + "If set, we will calculate the number of rollout steps as `num_rollout = num_epoch * dataset_size // rollout_batch_size`." + "If both `--num-epoch` and `--num-rollout` are set, `--num-epoch` will be ignored." + ), + ) + + parser.add_argument( + "--disable-rollout-global-dataset", + action="store_false", + dest="rollout_global_dataset", + help=( + "Whether to use a global dataset for rollout. " + "If set, the rollout will use the `--prompt-data` as the prompt dataset, " + "and the prompts for rollout will be sampled from the dataset. " + "If not set, you need to manage the data by your self." + ), + ) + + parser.add_argument( + "--data-source-path", + type=str, + default="slime.rollout.data_source.RolloutDataSourceWithBuffer", + help="The data source class for rollout data.", + ) + parser.add_argument( + "--prompt-data", + type=str, + nargs="+", + default=None, + help=( + "The path(s) to the prompt data. Supports files, directories, and multiple datasets for mixture. " + "Format: path1 [weight1] path2 [weight2] ... " + "Each path can be a file (.jsonl or .parquet) or a directory (reads all data files recursively). " + "Supports row slicing with @[start:end] suffix, e.g., 'data.jsonl@[0:1000]'. " + "Examples: " + " --prompt-data /path/to/data.jsonl " + " --prompt-data /path/to/data_folder/ (reads all data files in folder) " + " --prompt-data /path/to/data1.jsonl /path/to/data2.jsonl (equal weights) " + "Supported formats: .jsonl, .parquet." + ), + ) + parser.add_argument("--apply-chat-template", action="store_true", default=False) + # Temporarily be JSON-serialized str, will be a real dict after using Omegaconf + parser.add_argument("--apply-chat-template-kwargs", type=json.loads, default="{}") + parser.add_argument("--input-key", type=str, default="input", help="JSON dataset key") + parser.add_argument("--label-key", type=str, default=None, help="JSON dataset key") + parser.add_argument( + "--multimodal-keys", + type=json.loads, + default=None, + help=( + 'JSON string for multimodal data mapping media types to data keys. Example: \'{"image": "image_file"}\'' + ), + ) + parser.add_argument("--metadata-key", type=str, default="metadata", help="JSON dataset key") + parser.add_argument( + "--tool-key", + type=str, + default=None, + help=( + "When need to add tools during apply_chat_template, you should provide the key for the tools in the prompt dataset." + ), + ) + + parser.add_argument( + "--start-rollout-id", + type=int, + default=None, + help=( + "The starting rollout step, if not set, will try to load the step from --load when doing continue training, " + "otherwise will be set to 0, meaning training from start." + ), + ) + + # batch sizes + parser.add_argument( + "--rollout-batch-size", + type=int, + required=True, + help=( + "The number of prompts in each rollout step. " + "The total data returned should be rollout_batch_size * n_samples_per_prompt. " + ), + ) + parser.add_argument( + "--n-samples-per-prompt", type=int, default=1, help="Number of responses for each prompt in generation" + ) + + # gbs of the training, note that the gbs is of sample, not of prompts, + # so if you hope to train 1 step for each rollout, the global_bach_size should be set as + # `rollout_batch_size * n_samples_per_prompt`. + reset_arg(parser, "--global-batch-size", type=int, default=None) + parser.add_argument( + "--num-steps-per-rollout", + type=int, + default=None, + help=( + "Number of steps per rollout, e.g. It is equivalent to setting gbs as " + "`rollout_batch_size * n_samples_per_prompt // num_steps_per_rollout`." + ), + ) + # mbs for the training, will be ignored if `use_dynamic_batch_size` is set. + reset_arg(parser, "--micro-batch-size", type=int, default=1) + parser.add_argument( + "--balance-data", + action="store_true", + default=False, + help=( + "Balance the number of tokens between data parallel ranks with `karmarkar_karp` for verl. " + "Note that this may allocate the different response of the same prompt into different training steps." + ), + ) + + parser.add_argument( + "--use-dynamic-batch-size", + action="store_true", + default=False, + help=( + "Because the sample length varies, to maximize the GPU utilization, " + "we will use the dynamic batch size to adjust the micro batch size according to the maximum number of tokens each gpu can run. " + "For example, if we have 3 samples, with the length of 100, 200, and 300, and the max_tokens_per_gpu is 300, when enabling " + "dynamic batch size, slime will make 2 micro batches, i.e. [100, 200], [300]." + ), + ) + parser.add_argument( + "--max-tokens-per-gpu", + type=int, + default=None, + help=( + "The maximum number of tokens per GPU for dynamic batch size. " + "Note that when enabling context parallel (CP), the max tokens per gpu should be around " + "`max_response_len // cp_size` instead of `max_response_len`." + ), + ) + parser.add_argument( + "--log-probs-max-tokens-per-gpu", + type=int, + default=None, + help=( + "The maximum number of tokens per GPU for calculating log probs. " + "This is used to calculate the log probs of the responses during rollout, " + "and should be set to a larger value than `max_tokens_per_gpu` if you want better performance. " + ), + ) + return parser + + def add_eval_arguments(parser): + parser.add_argument( + "--eval-function-path", + type=str, + default=None, + help=( + "Path to the eval generation function." + "If not set, we will use rollout_function_path as the default. " + ), + ) + + # change the default value of eval_interval from Megatron to None + reset_arg(parser, "--eval-interval", type=int, default=None) + + parser.add_argument( + "--eval-prompt-data", + type=str, + default=None, + nargs="+", + help=( + "Path to the evaluation prompt data, " + "should first input the name of the eval dataset and then the path, e.g. " + "aime /path/to/aime.jsonl" + ), + ) + parser.add_argument( + "--eval-config", + type=str, + default=None, + help=( + "Path to an OmegaConf YAML/JSON file describing evaluation datasets. " + "When provided, this overrides --eval-prompt-data." + ), + ) + parser.add_argument( + "--skip-eval-before-train", + action="store_true", + default=False, + help="Whether to skip evaluation before training.", + ) + + # The following keys are used to override the rollout version during eval. + parser.add_argument("--eval-input-key", type=str, default=None, help="JSON dataset key") + parser.add_argument("--eval-label-key", type=str, default=None, help="JSON dataset key") + parser.add_argument("--eval-tool-key", type=str, default=None, help="JSON dataset key") + parser.add_argument( + "--n-samples-per-eval-prompt", + type=int, + default=1, + help="number of responses for each prompt in generation", + ) + parser.add_argument("--eval-temperature", type=float, default=None) + parser.add_argument("--eval-top-p", type=float, default=None) + parser.add_argument("--eval-top-k", type=int, default=None) + parser.add_argument("--eval-max-response-len", type=int, default=None) + parser.add_argument("--eval-max-prompt-len", type=int, default=None) + parser.add_argument("--eval-min-new-tokens", type=int, default=None) + parser.add_argument("--eval-max-context-len", type=int, default=None) + + return parser + + def add_algo_arguments(parser): + parser.add_argument( + "--ref-load", + type=str, + default=None, + help=( + "The checkpoint for reference model. " + "When --load is not set, this will be used as the initial checkpoint for training. " + ), + ) + parser.add_argument( + "--ref-ckpt-step", type=int, default=None, help="The checkpoint step for reference model. " + ) + reset_arg(parser, "--load", type=str, default=None) + reset_arg(parser, "--save", type=str, default=None) + reset_arg(parser, "--save-interval", type=int, default=None) + reset_arg(parser, "--async-save", action="store_true") + reset_arg(parser, "--seed", type=int, default=1234) + reset_arg(parser, "--clip-grad", type=float, default=1.0) + reset_arg(parser, "--calculate-per-token-loss", action="store_true") + reset_arg(parser, "--lr", type=float, default=1e-6) + + parser.add_argument("--num-critic-only-steps", type=int, default=0, help="Number of critic only steps") + parser.add_argument("--critic-load", type=str, default=None, help="The checkpoint for critic model.") + parser.add_argument("--critic-save", type=str, default=None, help="The checkpoint for critic model.") + parser.add_argument("--critic-lr", type=float, default=None, help="The lr for critic model") + parser.add_argument( + "--critic-lr-warmup-iters", + type=int, + default=0, + help="number of iterations to linearly warmup for critic model.", + ) + + parser.add_argument("--eps-clip", type=float, default=0.2, help="PPO clip range") + parser.add_argument("--eps-clip-high", type=float, default=None, help="PPO clip upper range") + parser.add_argument( + "--eps-clip-c", + type=float, + default=None, + help="lower bound of the value for Dual-clip PPO from https://arxiv.org/pdf/1912.09729", + ) + parser.add_argument("--value-clip", type=float, default=0.2, help="the clip for value loss") + parser.add_argument( + "--kl-coef", + type=float, + default=0.00, + help="KL penalty coefficient for reward shaping. This is applied to the reward signal before advantage calculation.", + ) + parser.add_argument( + "--loss-type", + type=str, + choices=["policy_loss", "custom_loss"], + default="policy_loss", + help=( + "Choose loss type. Options:\n" + " - policy_loss: PPO/GRPO policy loss for RL training\n" + " - custom_loss: Use custom loss function from `--custom-loss-function-path`\n" + ), + ) + parser.add_argument( + "--custom-loss-function-path", + type=str, + default=None, + help=( + "Path to the custom loss function, if the loss_type is `custom_loss`, " + "we will use this function to calculate the loss. " + ), + ) + parser.add_argument( + "--kl-loss-type", + type=str, + choices=["k1", "k2", "k3", "low_var_kl"], + default="k1", + help="Choose KL loss type: kl, k2, k3, low_var_kl", + ) + parser.add_argument( + "--advantage-estimator", + type=str, + choices=[ + "grpo", + "gspo", + "reinforce_plus_plus", + "reinforce_plus_plus_baseline", + "ppo", + "on_policy_distillation", + ], + default="grpo", + ) + parser.add_argument( + "--disable-compute-advantages-and-returns", + action="store_false", + dest="compute_advantages_and_returns", + help=( + "Whether to disable computing advantages and returns. " + "If set, we will not compute the advantages and returns, " + "This is useful for custom loss function." + ), + ) + parser.add_argument( + "--use-kl-loss", action="store_true", default=False, help="whether to use KL loss from GRPO" + ) + parser.add_argument( + "--kl-loss-coef", + type=float, + default=0.0, + help="KL penalty coefficient for the loss function. This is added to the final PPO loss.", + ) + parser.add_argument( + "--use-unbiased-kl", + action="store_true", + default=False, + help="Whether to enable unbiased KL estimation.", + ) + parser.add_argument( + "--ref-update-interval", + type=int, + default=None, + help="Interval (in rollout steps) to update ref model from actor. If None, ref model is not updated.", + ) + parser.add_argument("--entropy-coef", type=float, default=0.0, help="Entropy loss coef") + parser.add_argument("--gamma", type=float, default=1.0, help="PPO GAE gamma") + parser.add_argument("--lambd", type=float, default=1.0, help="PPO GAE lambd") + parser.add_argument("--normalize-advantages", action="store_true", default=False) + + parser.add_argument( + "--disable-grpo-std-normalization", + action="store_false", + dest="grpo_std_normalization", + help="from Dr.GRPO https://arxiv.org/pdf/2503.20783", + ) + parser.add_argument( + "--disable-rewards-normalization", + action="store_false", + dest="rewards_normalization", + help="Disable rewards normalization", + ) + parser.add_argument( + "--use-rollout-entropy", + action="store_true", + default=False, + help=( + "Whether to calculate the entropy when calculating the logprobs from actor and reference model. " + "This is useful for doing special loss mask." + ), + ) + parser.add_argument( + "--get-mismatch-metrics", + action="store_true", + default=False, + help="Whether to calculate the mismatch metrics.", + ) + parser.add_argument( + "--use-rollout-logprobs", + action="store_true", + default=False, + help=( + "Whether to use the rollout logprobs when calculating the importance sampling ratios. " + "If not set, we will use the logprobs from the actor model." + ), + ) + # Off-Policy Correction using Importance Sampling: https://fengyao.notion.site/off-policy-rl + parser.add_argument( + "--use-tis", + action="store_true", + default=False, + help="Enable TIS from https://fengyao.notion.site/off-policy-rl for off-policy importance sampling.", + ) + parser.add_argument( + "--tis-clip", + type=float, + default=2.0, + help="Clipping threshold C for importance sampling ratios to control variance.", + ) + parser.add_argument( + "--tis-clip-low", + type=float, + default=0, + help="Lower bound clipping threshold C for importance sampling ratios to control variance.", + ) + parser.add_argument( + "--custom-tis-function-path", + type=str, + default=None, + help="Path to the custom TIS/RS function (e.g., examples/train_infer_mismatch_helper/mis.py:compute_mis_weights_with_cp).", + ) + + parser.add_argument( + "--use-routing-replay", + action="store_true", + default=False, + help="The routing replay technique from https://arxiv.org/abs/2507.18071", + ) + parser.add_argument( + "--use-rollout-routing-replay", + action="store_true", + default=False, + help="The rollout routing replay technique from https://arxiv.org/abs/2510.11370", + ) + parser.add_argument( + "--use-opsm", + action="store_true", + default=False, + help="Whether to enable Off-Policy Sequence Masking (OPSM).", + ) + parser.add_argument( + "--opsm-delta", + type=float, + default=1e-4, + help="The threshold for Off-Policy Sequence Masking (OPSM).", + ) + parser.add_argument( + "--include-verifiable-reward", + action="store_true", + default=False, + help="Whether to include the verifiable reward in the log.", + ) + return parser + + def add_router_arguments(parser): + parser.add_argument( + "--use-slime-router", + action="store_true", + default=False, + help="Whether to use SlimeRouter for text-based routing instead of SGLang token-based routing", + ) + parser.add_argument( + "--slime-router-middleware-paths", + type=str, + nargs="+", + default="", + ) + parser.add_argument( + "--slime-router-timeout", + type=float, + default=None, + help="Timeout for SlimeRouter HTTP requests in seconds.", + ) + parser.add_argument( + "--slime-router-max-connections", + type=int, + default=None, + help="Max connections for SlimeRouter HTTP client.", + ) + RouterArgs.add_cli_args(parser, use_router_prefix=True, exclude_host_port=True) + return parser + + # wandb + def add_wandb_arguments(parser): + # wandb parameters + parser.add_argument("--use-wandb", action="store_true", default=False) + parser.add_argument( + "--wandb-mode", + type=str, + default=None, + choices=["online", "offline", "disabled"], + help="W&B mode: online (default), offline (local only), or disabled. Overrides WANDB_MODE env var.", + ) + parser.add_argument( + "--wandb-dir", + type=str, + default=None, + help="Directory to store wandb logs. Default is ./wandb in current directory.", + ) + parser.add_argument("--wandb-key", type=str, default=None) + parser.add_argument("--wandb-host", type=str, default=None) + parser.add_argument("--wandb-team", type=str, default=None) + parser.add_argument("--wandb-group", type=str, default=None) + reset_arg(parser, "--wandb-project", type=str, default=None) + parser.add_argument( + "--disable-wandb-random-suffix", + action="store_false", + dest="wandb_random_suffix", + default=True, + help=( + "Whether to add a random suffix to the wandb run name. " + "By default, we will add a random 6 length string with characters to the run name." + ), + ) + parser.add_argument( + "--wandb-always-use-train-step", + action="store_true", + default=False, + help=( + "Whether to always use train step as the step metric in wandb. " + "If set, we will always use the train steps for wandb logging, " + "otherwise, will use rollout step for most info other than train/*. " + ), + ) + parser.add_argument( + "--log-multi-turn", + action="store_true", + default=False, + help="Whether to log information for multi-turn rollout.", + ) + parser.add_argument( + "--log-passrate", + action="store_true", + default=False, + help="Whether to turn on passrate logging, which will log the pass@n of the responses in the rollout.", + ) + parser.add_argument( + "--log-reward-category", + type=str, + default=None, + help=( + "Log statistics of the category of reward, such as why the reward function considers it as failed. " + "Specify the key in the reward dict using this argument.", + ), + ) + parser.add_argument( + "--log-correct-samples", + action="store_true", + default=False, + help="Whether to turn on passrate logging, which will log the pass@n of the responses in the rollout.", + ) + parser.add_argument("--wandb-run-id", type=str, default=None) + parser.add_argument( + "--wandb-resume-run-id", + type=str, + default=None, + help=( + "Specify the W&B run ID to resume. When set, wandb will resume the specified run " + "instead of creating a new one. This is useful for continuing training from a checkpoint." + ), + ) + return parser + + # tensorboard + def add_tensorboard_arguments(parser): + # tb_project_name, tb_experiment_name + parser.add_argument("--use-tensorboard", action="store_true", default=False) + parser.add_argument( + "--tb-project-name", + type=str, + default=None, + help="Directory to store tensorboard logs. Default is os.environ.get('TENSORBOARD_DIR') directory.", + ) + parser.add_argument("--tb-experiment-name", type=str, default=None) + + return parser + + # debug + def add_debug_arguments(parser): + parser.add_argument( + "--save-debug-rollout-data", + type=str, + default=None, + help=( + "Save the rollout data to this path for debugging. " + "The file will be saved to `save_debug_rollout_data.format(rollout_id)`." + ), + ) + parser.add_argument( + "--load-debug-rollout-data", + type=str, + default=None, + help=( + "Load the rollout data from this path for debugging. " + "The file will be loaded from `load_debug_rollout_data.format(rollout_id)`. " + "When this is enabled, slime will not instantiate sglang servers." + ), + ) + parser.add_argument( + "--load-debug-rollout-data-subsample", + type=float, + default=None, + help="Subsample a portion of the debug rollout data for faster debugging.", + ) + parser.add_argument( + "--debug-rollout-only", + action="store_true", + default=False, + help=( + "Whether to only run the rollout generation without training. " + "This is useful for debugging the rollout generation function." + ), + ) + parser.add_argument( + "--debug-train-only", + action="store_true", + default=False, + help=( + "Whether to only run the training without sglang servers. " + "This is useful for debugging the rollout generation function." + ), + ) + parser.add_argument( + "--save-debug-train-data", + type=str, + default=None, + help=( + "Save the train data to this path for debugging. " + "The file will be saved to `save_debug_train_data.format(rollout_id)`." + ), + ) + parser.add_argument( + "--dump-details", + type=str, + default=None, + help=("Dump all details of training for post-hoc analysis and visualization."), + ) + # use together with --record-memory-history and --memory-snapshot-path (defined in Megatron) + parser.add_argument( + "--memory-snapshot-dir", + type=str, + default=".", + ) + parser.add_argument( + "--memory-snapshot-num-steps", + type=int, + default=None, + ) + parser.add_argument( + "--profile-target", + type=str, + choices=["train_overall", "train_actor", "train_log_probs"], + default=["train_overall"], + nargs="+", + ) + parser.add_argument( + "--memory-recorder", + type=str, + choices=["torch", "memray"], + default="torch", + ) + parser.add_argument("--check-weight-update-equal", action="store_true") + return parser + + def add_network_arguments(parser): + parser.add_argument("--http-proxy", type=str, default=None) + parser.add_argument("--use-distributed-post", action="store_true", default=False) + return parser + + def add_reward_model_arguments(parser): + parser.add_argument( + "--rm-type", + type=str, + default=None, + help="Type of the reward model", + ) + parser.add_argument( + "--reward-key", + type=str, + default=None, + help=( + "Some reward model may return a dict instead of a value, " + "this is the key to extract the reward value from the dict. " + ), + ) + parser.add_argument( + "--eval-reward-key", + type=str, + default=None, + help="The eval variant for --reward-key", + ) + parser.add_argument( + "--group-rm", action="store_true", default=False, help="Whether to do rm on a whole group." + ) + parser.add_argument( + "--rm-url", + type=str, + default=None, + help="URL for the reward model service for --rm-type remote_rm, e.g. http://localhost:8000", + ) + parser.add_argument( + "--custom-rm-path", + type=str, + default=None, + help=( + "Path to the custom reward model function. " + "If set, we will use this function to calculate the reward instead of the default one. " + "The function should have the signature `def custom_rm(args, sample) -> float`." + ), + ) + parser.add_argument( + "--custom-reward-post-process-path", + type=str, + default=None, + help=( + "Path to the custom function that will post process reward, by default it will be the normalization for grpo. " + ), + ) + parser.add_argument( + "--custom-convert-samples-to-train-data-path", + type=str, + default=None, + help=( + "Path to a custom function that converts samples to training data. " + "If set, this function will replace the default _convert_samples_to_train_data. " + "The function should have the signature `def convert_samples_to_train_data(args, samples) -> dict`." + ), + ) + return parser + + def add_rollout_buffer_arguments(parser): + parser.add_argument( + "--rollout-buffer-url", + type=str, + default=None, + help="URL for the rollout buffer", + ) + + parser.add_argument( + "--fetch-trajectory-retry-times", + type=int, + default=-1, + help="Number of times to retry fetching trajectory, -1 means unlimited retry", + ) + parser.add_argument( + "--min-batch-collection-ratio", + type=float, + default=1, + help="Minimum batch collection ratio", + ) + parser.add_argument( + "--rollout-task-type", + type=str, + default="math", + ) + parser.add_argument( + "--loss-mask-type", + type=str, + default="qwen", + choices=["qwen", "qwen3", "distill_qwen"], + help="Loss mask type", + ) + parser.add_argument( + "--data-pad-size-multiplier", + type=int, + default=128, + help="Multiplier for data padding size in data processing.", + ) + parser.add_argument( + "--rollout-sample-filter-path", + type=str, + default=None, + help=( + "Path to the rollout sample filter function. " + "This function determines whether a sample will participate in loss calculation. " + "The function should take args and samples (list[Sample]) as input, and return None. " + "Please directly modify the remove_sample attribute of Sample. " + "Note: This attribute does not determine whether the sample participates in advantage normalization." + ), + ) + parser.add_argument( + "--rollout-all-samples-process-path", + type=str, + default=None, + help=( + "Path to the rollout all samples process function that " + "can process all samples including filtered ones." + ), + ) + parser.add_argument( + "--disable-rollout-trim-samples", + action="store_true", + default=False, + help="disable trim samples in rollout buffer when converting samples to train data", + ) + return parser + + def add_custom_megatron_plugins_arguments(parser): + """ + Add custom Megatron plugins arguments. + This is a placeholder for any additional arguments that might be needed. + """ + # Custom arguments can be added here + parser.add_argument( + "--custom-megatron-init-path", + type=str, + default=None, + ) + parser.add_argument( + "--custom-megatron-before-log-prob-hook-path", + type=str, + default=None, + ) + parser.add_argument( + "--custom-megatron-before-train-step-hook-path", + type=str, + default=None, + ) + return parser + + def add_mtp_training_arguments(parser): + """Add MTP training specific arguments.""" + reset_arg(parser, "--mtp-num-layers", type=int, default=None) + reset_arg(parser, "--mtp-loss-scaling-factor", type=float, default=0.2) + parser.add_argument( + "--enable-mtp-training", + action="store_true", + default=False, + help="Enable MTP layer parameter updates during training", + ) + + return parser + + def add_prefill_decode_disaggregation_arguments(parser): + parser.add_argument( + "--prefill-num-servers", + type=int, + default=None, + help="Number of prefill servers for disaggregation.", + ) + return parser + + def add_ci_arguments(parser): + parser.add_argument( + "--ci-test", + action="store_true", + ) + parser.add_argument( + "--ci-disable-kl-checker", + action="store_true", + ) + parser.add_argument( + "--ci-metric-checker-key", + type=str, + default=None, + ) + parser.add_argument( + "--ci-metric-checker-threshold", + type=float, + default=None, + ) + parser.add_argument( + "--ci-save-grad-norm", + type=str, + default=None, + ) + parser.add_argument( + "--ci-load-grad-norm", + type=str, + default=None, + ) + return parser + + def add_sglang_tp_size(): + temp_parser = argparse.ArgumentParser(add_help=False) + temp_parser.add_argument("--rollout-num-gpus-per-engine", type=int, default=1) + temp_args, _ = temp_parser.parse_known_args() + sglang_tp_size = temp_args.rollout_num_gpus_per_engine + return sglang_tp_size + + # Add custom arguments in front to prevent overwritten some slime arguments. + if add_custom_arguments is not None: + parser = add_custom_arguments(parser) + + parser = add_cluster_arguments(parser) + parser = add_train_arguments(parser) + parser = add_rollout_arguments(parser) + parser = add_fault_tolerance_arguments(parser) + parser = add_data_arguments(parser) + parser = add_eval_arguments(parser) + parser = add_algo_arguments(parser) + parser = add_wandb_arguments(parser) + parser = add_tensorboard_arguments(parser) + parser = add_router_arguments(parser) + parser = add_debug_arguments(parser) + parser = add_sglang_arguments(parser) + parser = add_network_arguments(parser) + parser = add_reward_model_arguments(parser) + parser = add_rollout_buffer_arguments(parser) + parser = add_mtp_training_arguments(parser) + parser = add_prefill_decode_disaggregation_arguments(parser) + parser = add_ci_arguments(parser) + parser = add_custom_megatron_plugins_arguments(parser) + reset_arg( + parser, + "--custom-config-path", + type=str, + default=None, + help="Path to the YAML config for custom function arguments.", + ) + reset_arg(parser, "--padded-vocab-size", type=int, default=None) + + parser.set_defaults(sglang_tensor_parallel_size=add_sglang_tp_size()) + return parser + + return add_slime_arguments + + +def parse_args(add_custom_arguments=None): + # Users may call `parse_args` very early, thus we ensure logger is configured here + configure_logger() + + add_slime_arguments = get_slime_extra_args_provider(add_custom_arguments) + + backend = parse_args_train_backend() + if backend == "megatron": + from slime.backends.megatron_utils.arguments import parse_args as megatron_parse_args + from slime.backends.megatron_utils.arguments import set_default_megatron_args + from slime.backends.megatron_utils.arguments import validate_args as megatron_validate_args + + args = megatron_parse_args(extra_args_provider=add_slime_arguments) + if args.hf_checkpoint: + hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True) + hf_validate_args(args, hf_config) + + args.rank = 0 + args.world_size = args.actor_num_nodes * args.actor_num_gpus_per_node + args = set_default_megatron_args(args) + else: + from slime.backends.fsdp_utils.arguments import load_fsdp_args + + args = load_fsdp_args(extra_args_provider=add_slime_arguments) + args.rank = 0 # Primary process rank for wandb initialization + args.world_size = args.actor_num_nodes * args.actor_num_gpus_per_node + + slime_validate_args(args) + + if backend == "megatron": + megatron_validate_args(args) + + # always use varlen + args.variable_seq_lengths = True + if getattr(args, "moe_token_dispatcher_type", None) == "allgather": + logger.info( + "--moe-token-dispatcher-type allgather does not support variable sequence length, " + "please use alltoall dispatcher instead." + ) + args.moe_token_dispatcher_type = "alltoall" + + sglang_validate_args(args) + + return args + + +def parse_args_train_backend(): + if os.environ.get("SLIME_BACKEND") is not None: + raise Exception("`SLIME_BACKEND` is deprecated, please use --train-backend directly.") + + parser = argparse.ArgumentParser() + get_slime_extra_args_provider()(parser) + args_partial, _ = parser.parse_known_args() + return args_partial.train_backend + + +def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]: + """ + Build evaluation dataset configurations from either --eval-config or --eval-prompt-data. + """ + datasets_config = [] + defaults: dict[str, Any] = {} + + if args.eval_config: + from omegaconf import OmegaConf + + cfg = OmegaConf.load(args.eval_config) + cfg_dict = OmegaConf.to_container(cfg, resolve=True) + if not isinstance(cfg_dict, dict): + raise ValueError("--eval-config must contain a mapping at the root.") + + eval_cfg = cfg_dict.get("eval", cfg_dict) + if not isinstance(eval_cfg, dict): + raise ValueError("--eval-config must define an `eval` mapping or be a mapping itself.") + + defaults = dict(eval_cfg.get("defaults") or {}) + datasets_config = ensure_dataset_list(eval_cfg.get("datasets")) + if not datasets_config: + raise ValueError("--eval-config does not define any datasets under `eval.datasets`.") + elif args.eval_prompt_data: + values = list(args.eval_prompt_data) + if len(values) == 1: + logger.info("[legacy] only one eval_prompt_data detected, will assume it is data for aime") + values = ["aime", values[0]] + if len(values) % 2 != 0: + raise ValueError("eval prompt data must be provided as name/path pairs.") + datasets_config = [{"name": values[i], "path": values[i + 1]} for i in range(0, len(values), 2)] + else: + datasets_config = [] + + eval_datasets = build_eval_dataset_configs(args, datasets_config, defaults) + if eval_datasets: + args.eval_prompt_data = [item for dataset in eval_datasets for item in (dataset.name, dataset.path)] + else: + args.eval_prompt_data = None + + return eval_datasets + + +def slime_validate_args(args): + args.eval_datasets = _resolve_eval_datasets(args) + + if args.kl_coef != 0 or args.use_kl_loss: + if not os.path.exists(args.ref_load): + raise FileNotFoundError(f"ref_load {args.ref_load} does not exist, please check the path.") + + if not os.path.exists(os.path.join(args.ref_load, "latest_checkpointed_iteration.txt")): + logger.info( + f"ref_load {args.ref_load} does not have latest_checkpointed_iteration.txt, " + "please make sure it is a valid megatron checkpoint directory." + ) + + # TODO: During loading, we need to set the start_rollout_id here. + if args.megatron_to_hf_mode == "bridge": + if args.load is None: + args.load = args.ref_load or args.hf_checkpoint + args.start_rollout_id = 0 + else: + if ( + args.load is None + or not os.path.exists(args.load) + or not os.path.exists(os.path.join(args.load, "latest_checkpointed_iteration.txt")) + ): + args.no_load_optim = True + args.no_load_rng = True + args.finetune = True + args.load = args.ref_load + if args.ref_ckpt_step is not None: + args.ckpt_step = args.ref_ckpt_step + args.start_rollout_id = 0 + + if args.eval_interval is not None: + assert args.eval_datasets, "Evaluation datasets must be configured when eval_interval is set." + + if args.save_interval is not None: + assert args.save is not None, "'--save' is required when save_interval is set." + + assert not (args.kl_coef != 0 and args.kl_loss_coef != 0), "Only one of kl_coef and kl_loss_coef can be set" + + if args.advantage_estimator in ["reinforce_plus_plus", "reinforce_plus_plus_baseline"]: + assert args.normalize_advantages, ( + "The 'reinforce_plus_plus' and 'reinforce_plus_plus_baseline' advantage estimators " + "require advantage normalization. Please add `--normalize-advantages` to your command." + ) + + if args.use_rollout_logprobs: + assert not args.use_tis, "use_rollout_logprobs and use_tis cannot be set at the same time." + + if args.get_mismatch_metrics: + assert ( + args.custom_tis_function_path is not None + ), "custom_tis_function_path must be set when get_mismatch_metrics is set" + + if args.use_rollout_logprobs: + logger.info( + "get_mismatch_metrics is set; For metrics calculation, the log probs will still be recomputed by training engine. One more forward pass will be applied." + ) + + if args.use_dynamic_batch_size: + assert args.max_tokens_per_gpu is not None, "max_tokens_per_gpu must be set when use_dynamic_batch_size is set" + if args.log_probs_max_tokens_per_gpu is None: + args.log_probs_max_tokens_per_gpu = args.max_tokens_per_gpu + + if args.eps_clip_high is None: + args.eps_clip_high = args.eps_clip + + if args.eval_reward_key is None: + args.eval_reward_key = args.reward_key + + if args.dump_details is not None: + args.save_debug_rollout_data = f"{args.dump_details}/rollout_data/{{rollout_id}}.pt" + args.save_debug_train_data = f"{args.dump_details}/train_data/{{rollout_id}}_{{rank}}.pt" + + if args.load_debug_rollout_data is not None: + logger.info( + f"load_debug_rollout_data {args.load_debug_rollout_data} is set, " + "will not instantiate sglang servers and will only run the training process." + ) + args.debug_train_only = True + + args.use_critic = args.advantage_estimator == "ppo" + if args.critic_num_gpus_per_node is None: + args.critic_num_gpus_per_node = args.actor_num_gpus_per_node + if args.critic_num_nodes is None: + args.critic_num_nodes = args.actor_num_nodes + if args.critic_load is None: + args.critic_load = args.load + if args.critic_lr is None: + args.critic_lr = args.lr + + if args.offload: + args.offload_train = True + args.offload_rollout = True + del args.offload + + if args.debug_rollout_only: + if args.colocate and (not args.rollout_num_gpus): + args.rollout_num_gpus = args.actor_num_gpus_per_node * args.actor_num_nodes + else: + args.actor_num_gpus_per_node = min(8, args.rollout_num_gpus) + args.actor_num_nodes = args.rollout_num_gpus // args.actor_num_gpus_per_node + args.colocate = False + args.offload_train = args.offload_rollout = False + if args.train_memory_margin_bytes > 0: + logger.warning("Force train_memory_margin_bytes=0 since debug_rollout_only does not support it") + args.train_memory_margin_bytes = 0 + + assert not (args.debug_rollout_only and args.debug_train_only), ( + "debug_rollout_only and debug_train_only cannot be set at the same time, " "please set only one of them." + ) + + # always true on offload for colocate at the moment. + if args.colocate: + if args.offload_train is None: + args.offload_train = True + if args.offload_rollout is None: + args.offload_rollout = True + if args.rollout_num_gpus != args.actor_num_gpus_per_node * args.actor_num_nodes: + logger.info( + f"rollout_num_gpus {args.rollout_num_gpus} != actor_num_gpus_per_node {args.actor_num_gpus_per_node} " + f"* actor_num_nodes {args.actor_num_nodes}, overriding rollout_num_gpus to match actor_num_gpus_per_node * actor_num_nodes." + ) + args.rollout_num_gpus = args.actor_num_gpus_per_node * args.actor_num_nodes + if args.use_critic: + args.rollout_num_gpus += args.critic_num_gpus_per_node * args.critic_num_nodes + + if args.offload_train is None: + args.offload_train = False + if args.offload_rollout is None: + args.offload_rollout = False + + if args.eval_function_path is None: + args.eval_function_path = args.rollout_function_path + + if args.num_steps_per_rollout is not None: + global_batch_size = args.rollout_batch_size * args.n_samples_per_prompt // args.num_steps_per_rollout + if args.global_batch_size is not None: + assert args.global_batch_size == global_batch_size, ( + f"global_batch_size {args.global_batch_size} is not equal to " + f"rollout_batch_size {args.rollout_batch_size} * n_samples_per_prompt {args.n_samples_per_prompt} " + f"// num_steps_per_rollout {args.num_steps_per_rollout}" + ) + args.global_batch_size = global_batch_size + + assert args.rollout_batch_size * args.n_samples_per_prompt % args.global_batch_size == 0, ( + f"rollout_batch_size {args.rollout_batch_size} * n_samples_per_prompt {args.n_samples_per_prompt} " + f"is not a multiple of global_batch_size {args.global_batch_size}" + ) + + if args.n_samples_per_prompt == 1: + args.grpo_std_normalization = False + logger.info("n_samples_per_prompt is set to 1, grpo_std_normalization will be set to False.") + + if args.over_sampling_batch_size is None: + args.over_sampling_batch_size = args.rollout_batch_size + + assert args.over_sampling_batch_size >= args.rollout_batch_size, ( + f"over_sampling_batch_size {args.over_sampling_batch_size} should be greater than or equal to " + f"rollout_batch_size {args.rollout_batch_size}" + ) + + if args.num_epoch is not None: + if args.num_rollout is not None: + logger.info("Both num_epoch and num_rollout are set, num_epoch will be ignored.") + else: + assert args.rollout_global_dataset, ( + "num_epoch is set, but rollout_global_dataset is not set, " + "please remove --disable-rollout-global-dataset to use num_epoch" + ) + else: + # if num_epoch is not set, we should set num_rollout + assert args.num_rollout is not None, ( + "num_epoch is not set, but num_rollout is not set, " "please set --num-rollout or --num-epoch" + ) + + if args.enable_mtp_training: + assert args.mtp_num_layers, "mtp_num_layers must be set when enable_mtp_training is set" + + if args.use_rollout_routing_replay: + args.use_routing_replay = True + + if args.custom_config_path: + with open(args.custom_config_path) as f: + data = yaml.safe_load(f) or {} + for k, v in data.items(): + if hasattr(args, k): + logger.info(f"Warning: Argument {k} is already set to {getattr(args, k)}, will override with {v}.") + setattr(args, k, v) + + if args.eval_max_context_len is None: + logger.info( + f"args.eval_max_context_len is not set. Use args.rollout_max_context_len {args.rollout_max_context_len} as default value." + ) + args.eval_max_context_len = args.rollout_max_context_len + + if args.rollout_max_context_len is not None: + if args.rollout_max_prompt_len is None: + args.rollout_max_prompt_len = args.rollout_max_context_len - 1 + logger.info( + f"args.rollout_max_prompt_len is not set. Use args.rollout_max_context_len - 1 ({args.rollout_max_context_len} - 1) as default value so that there is at least one generated token to compute loss." + ) + assert ( + args.rollout_max_prompt_len <= args.rollout_max_context_len - 1 + ), f"args.rollout_max_prompt_len ({args.rollout_max_prompt_len}) must be smaller than args.rollout_max_context_len ({args.rollout_max_context_len}) so that there is at least one generated token to compute loss." + + assert not ( + args.prefill_num_servers is not None and args.rollout_external + ), "prefill_num_servers cannot be set when rollout_external is set." + + +def hf_validate_args(args, hf_config): + def equal(x, y): + return x == y + + errors = [] + + # multimodal models have different config structure + if hasattr(hf_config, "text_config"): + hf_config = hf_config.text_config + + for hf_config_name, megatron_config_name, compare_fn in [ + ("hidden_size", "hidden_size", equal), + ("num_attention_heads", "num_attention_heads", equal), + ("num_hidden_layers", "num_layers", equal), + ("intermediate_size", "ffn_hidden_size", equal), + ("tie_word_embeddings", "untie_embeddings_and_output_weights", lambda x, y: not x == y), + ("rms_norm_eps", "norm_epsilon", equal), + ("rope_theta", "rotary_base", equal), + ]: + if hasattr(hf_config, hf_config_name): + if not compare_fn(getattr(hf_config, hf_config_name), getattr(args, megatron_config_name)): + errors.append( + f"{hf_config_name} in hf config {getattr(hf_config, hf_config_name)} is not equal to " + f"{megatron_config_name} {getattr(args, megatron_config_name)}, please check the config." + ) + + if len(errors) > 0: + raise AssertionError("hf_validate_args failed: " + "; ".join(errors)) diff --git a/slime/utils/async_utils.py b/slime/utils/async_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8b6d35df8ca50cf6fe6e857ba688fcb2e1efc210 --- /dev/null +++ b/slime/utils/async_utils.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import threading + +__all__ = ["get_async_loop", "run"] + + +# Create a background event loop thread +class AsyncLoopThread: + def __init__(self): + self.loop = asyncio.new_event_loop() + self._thread = threading.Thread(target=self._start_loop, daemon=True) + self._thread.start() + + def _start_loop(self): + asyncio.set_event_loop(self.loop) + self.loop.run_forever() + + def run(self, coro): + # Schedule a coroutine onto the loop and block until it's done + return asyncio.run_coroutine_threadsafe(coro, self.loop).result() + + +# Create one global instance +async_loop = None + + +def get_async_loop(): + global async_loop + if async_loop is None: + async_loop = AsyncLoopThread() + return async_loop + + +def run(coro): + """Run a coroutine in the background event loop.""" + return get_async_loop().run(coro) diff --git a/slime/utils/context_utils.py b/slime/utils/context_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3423f0ea6ef9eb2ced6f44e35376a7cfeff50fd6 --- /dev/null +++ b/slime/utils/context_utils.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from functools import wraps + + +def with_defer(deferred_func): + def decorator(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + try: + return fn(*args, **kwargs) + finally: + deferred_func() + + return wrapper + + return decorator diff --git a/slime/utils/data.py b/slime/utils/data.py new file mode 100644 index 0000000000000000000000000000000000000000..3bce55d56778576993c763229450c759ed416a11 --- /dev/null +++ b/slime/utils/data.py @@ -0,0 +1,342 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import logging +import os +import random +import re + +import numpy as np +import pandas as pd +import ray + +from slime.utils.types import MultimodalTypes, Sample + +from .timer import Timer + +__all__ = ["Dataset", "create_dataset"] + +logger = logging.getLogger(__name__) + + +def _read_single_file(path, row_slice=None): + """Read a single data file (jsonl or parquet).""" + if path.endswith(".jsonl"): + df = pd.read_json(path, lines=True, dtype={"label": str}) + elif path.endswith(".parquet"): + df = pd.read_parquet(path, dtype_backend="pyarrow") + else: + raise ValueError(f"Unsupported file format: {path}. Supported formats are .jsonl and .parquet.") + + if row_slice is not None: + logger.info(f"read_file path={path} slice {len(df)=} rows into {row_slice=}") + df = df.iloc[row_slice] + + for _, row in df.iterrows(): + yield row.to_dict() + + +def _list_data_files(directory): + """List all supported data files in a directory (recursively).""" + supported_extensions = ('.jsonl', '.parquet') + data_files = [] + + for root, _, files in os.walk(directory): + for file in sorted(files): # Sort for deterministic order + if file.endswith(supported_extensions): + data_files.append(os.path.join(root, file)) + + return sorted(data_files) # Sort by full path for deterministic order + + +def read_file(path): + """Read data from a file or directory. + + Args: + path: Path to a data file (.jsonl or .parquet) or a directory containing data files. + If a directory is provided, all .jsonl and .parquet files in it (and subdirectories) + will be read and concatenated. + Supports row slicing with @[start:end] suffix, e.g., "data.jsonl@[0:1000]" + + Yields: + dict: Each row of data as a dictionary. + """ + path, row_slice = _parse_generalized_path(path) + + if not os.path.exists(path): + raise FileNotFoundError(f"Prompt dataset path '{path}' does not exist.") + + # Handle directory: read all data files inside + if os.path.isdir(path): + data_files = _list_data_files(path) + if not data_files: + raise ValueError(f"No .jsonl or .parquet files found in directory: {path}") + + logger.info(f"Found {len(data_files)} data files in directory {path}") + + # For directory, row_slice applies to the combined dataset + if row_slice is not None: + # Collect all data first, then apply slice + all_rows = [] + for file_path in data_files: + for row in _read_single_file(file_path): + all_rows.append(row) + logger.info(f"read_file directory={path} slice {len(all_rows)=} rows into {row_slice=}") + for row in all_rows[row_slice]: + yield row + else: + # Stream data from each file + for file_path in data_files: + for row in _read_single_file(file_path): + yield row + else: + # Handle single file + for row in _read_single_file(path, row_slice): + yield row + + +def _parse_generalized_path(s: str): + if (m := re.match(r"^(?P.*)@\[(?P-?\d*):(?P-?\d*)\]$", s)) is not None: + path = m.group("real_path") + start = int(x) if (x := m.group("start")) != "" else None + end = int(x) if (x := m.group("end")) != "" else None + return path, slice(start, end) + + return s, None + + +def _should_skip_prompt(formatted_prompt: str, tokenizer, processor, max_length, multimodal_inputs=None): + if max_length is None: + return False + + if processor: + processor_output = processor(text=formatted_prompt, **multimodal_inputs) + input_ids = processor_output["input_ids"][0] + else: + input_ids = tokenizer.encode(formatted_prompt, add_special_tokens=False) + + return len(input_ids) > max_length + + +def _build_messages(data: dict, prompt_key: str, as_conversation: bool, multimodal_keys: dict = None): + prompt = data.get(prompt_key) + + if isinstance(prompt, str): + if not as_conversation: + return prompt + else: + prompt = [{"role": "user", "content": prompt}] + + if multimodal_keys: + assert as_conversation, "as_conversation must be True when multimodal_keys is not None" + # Build mapping: placeholder -> (MultimodalType, content_list) + multimodals = {} + for type_name, data_key in multimodal_keys.items(): + mt = MultimodalTypes.get(type_name) + if mt: + multimodals[mt.placeholder] = (mt, list(data.get(data_key))) + + pattern = "(" + "|".join(re.escape(p) for p in multimodals.keys()) + ")" + + for message in prompt: + if isinstance(message["content"], str): + content_list = [] + for segment in re.split(pattern, message["content"]): + if not segment: + continue + if segment in multimodals: + mt, content = multimodals[segment] + content_list.append({"type": mt.name, mt.name: content.pop(0)}) + else: + content_list.append({"type": "text", "text": segment}) + message["content"] = content_list + + elif isinstance(message["content"], list): + # TODO: handle more general cases. where message['content'] is a dict and contains multiple types of content. + # e.g. + # "content": [ + # { + # "type": "image", + # "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg", + # }, + # {"type": "text", "text": "Describe this image."}, + # ], + logger.warning("message['content'] is a list of dicts, no processing will be done.") + continue + else: + raise ValueError( + f"Unsupported content type: {type(message['content'])}, expected str or list of dicts" + ) + + return prompt + + + + + +class Dataset: + def __init__( + self, + path, + tokenizer, + processor, + max_length, + *, + prompt_key="text", + multimodal_keys=None, + label_key=None, + tool_key=None, + metadata_key="metadata", + seed=42, + apply_chat_template=False, + apply_chat_template_kwargs=None, + ): + self.origin_samples = [] + + for data in read_file(path): + metadata = data.get(metadata_key) or {} + + prompt = _build_messages(data, prompt_key, apply_chat_template, multimodal_keys) + + tools = None + if tool_key is not None and tool_key in data: + tools = data[tool_key] + if isinstance(tools, str): + tools = json.loads(tools) + elif isinstance(tools, np.ndarray): + tools = tools.tolist() + assert isinstance(tools, list), f"tools must be a list, got {type(tools)} instead" + metadata["tools"] = tools + + if apply_chat_template: + formatted_prompt = tokenizer.apply_chat_template( + prompt, + tools=tools, + tokenize=False, + add_generation_prompt=True, + **(apply_chat_template_kwargs or {}), + ) + else: + formatted_prompt = prompt + + if processor: + # temporary solution, will write image utils for slime later + from qwen_vl_utils import process_vision_info + + assert isinstance( + prompt, list + ), f"prompt must be a list when processor is not None, got {type(prompt)} instead" + images, videos = process_vision_info(prompt) + multimodal_inputs = {"images": images, "videos": videos} + else: + multimodal_inputs = None + + # TODO: this is slow. + if _should_skip_prompt(formatted_prompt, tokenizer, processor, max_length, multimodal_inputs): + continue + + self.origin_samples.append( + Sample( + prompt=formatted_prompt, + label=data.get(label_key, None) if label_key is not None else None, + metadata=metadata, + multimodal_inputs=multimodal_inputs, + ) + ) + + logger.info(f"Dataset: Loaded {len(self.origin_samples)} samples from {path}") + self.epoch_id = -1 + self.seed = seed + self.samples = self.origin_samples + + def shuffle(self, new_epoch_id): + if self.epoch_id == new_epoch_id: + return + + random.seed(self.seed + new_epoch_id) + permutation = list(range(len(self.samples))) + random.shuffle(permutation) + self.samples = [self.origin_samples[i] for i in permutation] + self.epoch_id = new_epoch_id + + def __getitem__(self, idx): + return self.samples[idx] + + def __len__(self): + return len(self.samples) + + +def get_minimum_num_micro_batch_size(total_lengths, max_tokens_per_gpu): + # use first fit to get the number of micro batches + batches = [] + for length in total_lengths: + for i in range(len(batches)): + if batches[i] + length <= max_tokens_per_gpu: + batches[i] += length + break + else: + batches.append(length) + + return len(batches) + + +def process_rollout_data(args, rollout_data_ref, dp_rank, dp_size): + assert len(rollout_data_ref) == dp_size + rollout_data = ray.get(rollout_data_ref[dp_rank].inner) + + partition = rollout_data.pop("partition") + total_lengths = rollout_data["total_lengths"] + + # save the seqlen of the whole rollout batch + Timer().seq_lens = total_lengths + rollout_data["total_lengths"] = [total_lengths[i] for i in partition] + + return rollout_data + + + +def create_dataset( + paths, + tokenizer, + processor, + max_length, + *, + prompt_key="text", + multimodal_keys=None, + label_key=None, + tool_key=None, + metadata_key="metadata", + seed=42, + apply_chat_template=False, + apply_chat_template_kwargs=None, +): + """Factory function to create a Dataset. + + Args: + paths: A single path string, or a list with one path from --prompt-data. + Other args are the same as Dataset. + + Returns: + Dataset instance. + """ + if isinstance(paths, list): + if len(paths) != 1: + raise ValueError(f"Only single-path datasets are supported, got {len(paths)} paths.") + paths = paths[0] + + return Dataset( + path=paths, + tokenizer=tokenizer, + processor=processor, + max_length=max_length, + prompt_key=prompt_key, + multimodal_keys=multimodal_keys, + label_key=label_key, + tool_key=tool_key, + metadata_key=metadata_key, + seed=seed, + apply_chat_template=apply_chat_template, + apply_chat_template_kwargs=apply_chat_template_kwargs, + ) + diff --git a/slime/utils/debug_utils/__init__.py b/slime/utils/debug_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a4d636226b9a4b5130dfd8e94f2cece5fd290e53 --- /dev/null +++ b/slime/utils/debug_utils/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + diff --git a/slime/utils/debug_utils/display_debug_rollout_data.py b/slime/utils/debug_utils/display_debug_rollout_data.py new file mode 100644 index 0000000000000000000000000000000000000000..a6f9a6f182657a388d0be035a02386e4bfb88742 --- /dev/null +++ b/slime/utils/debug_utils/display_debug_rollout_data.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Annotated + +import torch +import typer + +from slime.ray.rollout import compute_metrics_from_samples +from slime.utils.types import Sample + +_WHITELIST_KEYS = [ + "group_index", + "index", + "prompt", + "response", + "response_length", + "label", + "reward", + "status", + "metadata", +] + + +def main( + # Deliberately make this name consistent with main training arguments + load_debug_rollout_data: Annotated[str, typer.Option()], + show_metrics: bool = True, + show_samples: bool = True, + category: list[str] = None, +): + if category is None: + category = ["train", "eval"] + for rollout_id, path in _get_rollout_dump_paths(load_debug_rollout_data, category): + print("-" * 80) + print(f"{rollout_id=} {path=}") + print("-" * 80) + + pack = torch.load(path) + sample_dicts = pack["samples"] + + if show_metrics: + # TODO read these configs from dumps + args = SimpleNamespace( + advantage_estimator="grpo", + reward_key=None, + log_reward_category=None, + ) + sample_objects = [Sample.from_dict(s) for s in sample_dicts] + metrics = compute_metrics_from_samples(args, sample_objects) + print("metrics", metrics) + + if show_samples: + for sample in sample_dicts: + print(json.dumps({k: v for k, v in sample.items() if k in _WHITELIST_KEYS})) + + +def _get_rollout_dump_paths(load_debug_rollout_data: str, categories: list[str]): + # may improve later + for rollout_id in range(1000): + for category in categories: + prefix = { + "train": "", + "eval": "eval_", + }[category] + path = Path(load_debug_rollout_data.format(rollout_id=f"{prefix}{rollout_id}")) + if path.exists(): + yield rollout_id, path + + +if __name__ == "__main__": + """python -m slime.utils.debug_utils.display_debug_rollout_data --load-debug-rollout-data ...""" + typer.run(main) diff --git a/slime/utils/debug_utils/replay_reward_fn.py b/slime/utils/debug_utils/replay_reward_fn.py new file mode 100644 index 0000000000000000000000000000000000000000..e32a8c9d87658365339eaf5a112ccffb23da32b7 --- /dev/null +++ b/slime/utils/debug_utils/replay_reward_fn.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from typing import Annotated + +import ray +import torch +import typer + +from slime.utils.misc import load_function +from slime.utils.types import Sample + + +def _truncate(text, max_len=200): + """Truncate text and add ellipsis if too long.""" + if text is None: + return None + text = str(text).replace("\n", "\\n") + if len(text) > max_len: + return text[:max_len] + "..." + return text + + +def main( + rollout_data_path: Annotated[str, typer.Option()], + custom_rm_path: Annotated[str, typer.Option()], +): + if not ray.is_initialized(): + ray.init() + + pack = torch.load(rollout_data_path) + samples = [Sample.from_dict(s) for s in pack["samples"]] + asyncio.run(_main_async(samples=samples, custom_rm_path=custom_rm_path)) + + +async def _main_async(samples, custom_rm_path): + rm_function = load_function(custom_rm_path) + rewards = await asyncio.gather(*[rm_function(None, sample) for sample in samples]) + + for i, (sample, reward) in enumerate(zip(samples, rewards, strict=True)): + print("-" * 60) + print(f"Sample {i + 1}/{len(samples)}") + print(f" Index: {sample.index}") + print(f" Status: {sample.status}") + print(f" Reward: {reward}") + print(f" Prompt: {_truncate(sample.prompt, 200)}") + print(f" Response: {_truncate(sample.response, 200)}") + print("-" * 60) + + +if __name__ == "__main__": + typer.run(main) diff --git a/slime/utils/debug_utils/send_to_sglang.py b/slime/utils/debug_utils/send_to_sglang.py new file mode 100644 index 0000000000000000000000000000000000000000..3d482584c6dc9e4fb76455fd8b136c87cf8fd72f --- /dev/null +++ b/slime/utils/debug_utils/send_to_sglang.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import json +from typing import Annotated + +import typer +from openai import AsyncOpenAI + +from slime.utils.data import read_file + + +# can unify w/ sglang_rollout.py later, e.g. add RM, if needed +def main( + prompt_data: Annotated[str, typer.Option()], + url: Annotated[str, typer.Option()] = "http://localhost:30000/v1", + input_key: Annotated[str, typer.Option()] = "input", + n_samples_per_prompt: Annotated[int, typer.Option()] = 1, + rollout_max_response_len: Annotated[int, typer.Option()] = 1024, + rollout_temperature: Annotated[float, typer.Option()] = 1.0, + rollout_top_p: Annotated[float, typer.Option()] = 1.0, +): + """ + Minimally send prompts to SGLang using OpenAI endpoints with arguments in the same format as main Slime. + + Example usage: + python -m slime.utils.debug_utils.send_to_sglang --prompt-data /root/datasets/aime-2024/aime-2024.jsonl --input-key prompt --n-samples-per-prompt 16 --rollout-max-response-len 32768 --rollout-temperature 0.8 --rollout-top-p 0.7 + """ + + async def _main_async(): + tasks = [ + asyncio.create_task(_run_one(row, row_index=row_index, repeat_index=repeat_index)) + for row_index, row in enumerate(read_file(prompt_data)) + for repeat_index in range(n_samples_per_prompt) + ] + outputs = await asyncio.gather(*tasks) + for output in outputs: + print(json.dumps(output)) + + async def _run_one(row, row_index: int, repeat_index: int): + resp = await client.chat.completions.create( + messages=row[input_key], + model="dummy_model", + max_tokens=rollout_max_response_len, + temperature=rollout_temperature, + top_p=rollout_top_p, + ) + return dict( + row_index=row_index, + repeat_index=repeat_index, + **row, + response=resp.choices[0].message.content, + ) + + client = AsyncOpenAI(api_key="dummy_key", base_url=url) + asyncio.run(_main_async()) + + +if __name__ == "__main__": + typer.run(main) diff --git a/slime/utils/distributed_utils.py b/slime/utils/distributed_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..422676a5b63b4d1b43ab2d498360ee441f8fc8e1 --- /dev/null +++ b/slime/utils/distributed_utils.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from datetime import timedelta +from typing import Any + +import torch +import torch.distributed as dist +from torch.distributed.distributed_c10d import ( + Backend, + PrefixStore, + Store, + _new_process_group_helper, + _world, + default_pg_timeout, + rendezvous, +) + + +GLOO_GROUP = None + + +def init_gloo_group(): + """Initialize Gloo group for distributed communication.""" + global GLOO_GROUP + if GLOO_GROUP is None: + GLOO_GROUP = dist.new_group(backend="gloo") + return GLOO_GROUP + + +def get_gloo_group(): + """Get the Gloo group for distributed communication.""" + global GLOO_GROUP + if GLOO_GROUP is None: + raise RuntimeError("Gloo group has not been initialized. Call _init_gloo_group() first.") + return GLOO_GROUP + + +# Copy from pytorch to allow creating multiple main groups. +# https://github.com/pytorch/pytorch/blob/main/torch/distributed/distributed_c10d.py +def init_process_group( + backend: str | Backend = None, + init_method: str | None = None, + timeout: timedelta | None = None, + world_size: int = -1, + rank: int = -1, + store: Store | None = None, + group_name: str = None, + pg_options: Any | None = None, +): + assert (store is None) or (init_method is None), "Cannot specify both init_method and store." + + if store is not None: + assert world_size > 0, "world_size must be positive if using store" + assert rank >= 0, "rank must be non-negative if using store" + elif init_method is None: + init_method = "env://" + + if backend: + backend = Backend(backend) + else: + backend = Backend("undefined") + + if timeout is None: + timeout = default_pg_timeout + + # backward compatible API + if store is None: + rendezvous_iterator = rendezvous(init_method, rank, world_size, timeout=timeout) + store, rank, world_size = next(rendezvous_iterator) + store.set_timeout(timeout) + + # Use a PrefixStore to avoid accidental overrides of keys used by + # different systems (e.g. RPC) in case the store is multi-tenant. + store = PrefixStore(group_name, store) + + # NOTE: The pg_options parameter was renamed into backend_options in PyTorch 2.6.0 + # https://github.com/pytorch/pytorch/commit/a0c7029a75628cd5fa8df83c0de0ea98ee7fd844 + # We need to determine the appropriate parameter name based on PyTorch version + pg_options_param_name = "backend_options" if str(torch.__version__) >= "2.6" else "pg_options" + pg, _ = _new_process_group_helper( + world_size, + rank, + [], + backend, + store, + group_name=group_name, + **{pg_options_param_name: pg_options}, + timeout=timeout, + ) + + _world.pg_group_ranks[pg] = {i: i for i in range(world_size)} + + return pg + + +def distributed_masked_whiten( + values: torch.Tensor, + mask: torch.Tensor, + process_group: dist.ProcessGroup | None = None, + shift_mean: bool = True, + epsilon: float = 1e-8, +): + """ + Performs whitening on a tensor using global statistics from all participating GPUs. + + It calculates the global mean and variance across all ranks in the default + process group (the WORLD) and uses these global statistics to normalize the + local data on each rank. + + Args: + values (torch.Tensor): The local tensor of values to whiten. + mask (torch.Tensor): The local mask corresponding to the values. + process_group: The process group for all_reduce. + If None, uses the default world group. + shift_mean (bool): If True, the output is zero-mean. Defaults to True. + epsilon (float): A small value for numerical stability. + + Returns: + torch.Tensor: The locally whitened tensor using global statistics. + """ + # Calculate local intermediate statistics + local_sum = (values * mask).sum() + local_sum_sq = ((values**2) * mask).sum() + local_mask_sum = mask.sum() + + stats_tensor = torch.tensor( + [local_sum, local_sum_sq, local_mask_sum], + device=values.device, + dtype=torch.float32, + ) + + # Aggregate via all_reduce within the DP group + dist.all_reduce(stats_tensor, group=process_group) + + # Calculate global stats from aggregated results + global_sum, global_sum_sq, global_mask_sum = stats_tensor + + if global_mask_sum.item() == 0: + raise ValueError("The global mask sum across all participating GPUs is zero.") + + global_mean = global_sum / global_mask_sum + global_mean_sq = global_sum_sq / global_mask_sum + global_var = global_mean_sq - global_mean**2 + + # Bessel's correction for unbiased estimate + if global_mask_sum.item() >= 2: + bessel_correction = global_mask_sum / (global_mask_sum - 1) + global_var = global_var * bessel_correction + + # Whiten local data using global stats + whitened_values = (values - global_mean) * torch.rsqrt(global_var + epsilon) + + if not shift_mean: + whitened_values += global_mean + + return whitened_values diff --git a/slime/utils/eval_config.py b/slime/utils/eval_config.py new file mode 100644 index 0000000000000000000000000000000000000000..6738b543b5905afe03a261d29372399b4828d963 --- /dev/null +++ b/slime/utils/eval_config.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass, field +from typing import Any + +_MISSING = object() + +# TODO: This is ugly, temporarily leave this. We should unify all the config name for dataset, default, and args. (advice from Tom.) +DATASET_RUNTIME_SPECS: dict[str, dict[str, tuple[str, ...]]] = { + "n_samples_per_eval_prompt": { + "dataset_keys": ("n_samples_per_eval_prompt",), + "default_keys": ("n_samples_per_eval_prompt",), + "arg_attrs": ("n_samples_per_eval_prompt", "n_samples_per_prompt"), + }, + "temperature": { + "dataset_keys": ("temperature",), + "default_keys": ("temperature",), + "arg_attrs": ("eval_temperature", "rollout_temperature"), + }, + "top_p": { + "dataset_keys": ("top_p",), + "default_keys": ("top_p",), + "arg_attrs": ("eval_top_p", "rollout_top_p"), + }, + "top_k": { + "dataset_keys": ("top_k",), + "default_keys": ("top_k",), + "arg_attrs": ("eval_top_k", "rollout_top_k"), + }, + "max_response_len": { + "dataset_keys": ("max_response_len",), + "default_keys": ("max_response_len",), + "arg_attrs": ("eval_max_response_len", "rollout_max_response_len"), + }, +} + +DATASET_SAMPLE_SPECS: dict[str, dict[str, tuple[str, ...]]] = { + "input_key": { + "dataset_keys": ("input_key",), + "default_keys": ("input_key",), + "arg_attrs": ("eval_input_key", "input_key"), + }, + "label_key": { + "dataset_keys": ("label_key",), + "default_keys": ("label_key",), + "arg_attrs": ("eval_label_key", "label_key"), + }, + "tool_key": { + "dataset_keys": ("tool_key",), + "default_keys": ("tool_key",), + "arg_attrs": ("eval_tool_key", "tool_key"), + }, + "metadata_key": { + "dataset_keys": ("metadata_key",), + "default_keys": ("metadata_key",), + "arg_attrs": ("metadata_key",), + }, +} + + +def _first_not_missing(*values: Any) -> Any: + for value in values: + if value is not _MISSING: + return value + return _MISSING + + +def _pick_from_mapping(data: dict[str, Any], key_names: tuple[str, ...] | None) -> Any: + if key_names is None: + return _MISSING + for key_name in key_names: + if key_name in data: + return data[key_name] + return _MISSING + + +def pick_from_args(args: Any, attrs: tuple[str, ...]) -> Any: + for attr in attrs: + value = getattr(args, attr, None) + if value is not None: + return value + return None + + +def _ensure_metadata_overrides(value: Any) -> dict[str, Any]: + if value is None: + return {} + if not isinstance(value, dict): + raise TypeError("metadata_overrides must be a mapping.") + return value + + +@dataclass +class EvalDatasetConfig: + """Configuration for a single evaluation dataset.""" + + name: str + path: str + rm_type: str | None = None + + # Dataset-specific overrides + input_key: str | None = None + label_key: str | None = None + tool_key: str | None = None + metadata_key: str | None = None + + n_samples_per_eval_prompt: int | None = None + + temperature: float | None = None + top_p: float | None = None + top_k: int | None = None + max_response_len: int | None = None + stop: list[str] | None = None + stop_token_ids: list[int] | None = None + min_new_tokens: int | None = None + + metadata_overrides: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + self.metadata_overrides = _ensure_metadata_overrides(self.metadata_overrides) + + @property + def cache_key(self) -> tuple[Any, ...]: + """Return a tuple uniquely identifying dataset config for caching.""" + return ( + self.name, + self.path, + self.input_key, + self.label_key, + self.tool_key, + self.metadata_key, + ) + + def inject_metadata(self, sample_metadata: Any) -> dict[str, Any]: + """Return updated metadata merging overrides.""" + if not isinstance(sample_metadata, dict): + metadata = {} + else: + metadata = dict(sample_metadata) + + if self.rm_type is not None: + metadata["rm_type"] = self.rm_type + + for key, value in self.metadata_overrides.items(): + metadata[key] = value + + return metadata + + +def ensure_dataset_list(config: Any) -> list[dict[str, Any]]: + """ + Normalize OmegaConf containers into a list of dicts. + Accepts either a list or dictionary keyed by dataset name. + """ + if config is None: + return [] + + if isinstance(config, dict): + datasets = [] + for name, cfg in config.items(): + dataset = dict(cfg or {}) + dataset.setdefault("name", name) + datasets.append(dataset) + return datasets + + if isinstance(config, (list, tuple)): + datasets = [] + for item in config: + dataset = dict(item or {}) + if "name" not in dataset: + raise ValueError("Each evaluation dataset entry must include a `name` field.") + datasets.append(dataset) + return datasets + + raise TypeError("eval.datasets must be either a list or a mapping.") + + +def _apply_dataset_field_overrides( + args: Any, dataset_cfg: dict[str, Any], defaults: dict[str, Any], spec_names: dict[str, Any] +) -> None: + for field_name, spec in spec_names.items(): + dataset_value = _pick_from_mapping(dataset_cfg, spec["dataset_keys"]) + default_value = _pick_from_mapping(defaults, spec["default_keys"]) + resolved_value = _first_not_missing(dataset_value, default_value) + if resolved_value is not _MISSING: + dataset_cfg[field_name] = resolved_value + continue + dataset_cfg[field_name] = pick_from_args(args, spec["arg_attrs"]) + + +def build_eval_dataset_configs( + args: Any, + raw_config: Iterable[dict[str, Any]], + defaults: dict[str, Any], +) -> list[EvalDatasetConfig]: + defaults = defaults or {} + datasets: list[EvalDatasetConfig] = [] + for cfg in raw_config: + cfg_dict = dict(cfg or {}) + combined_specs = {**DATASET_RUNTIME_SPECS, **DATASET_SAMPLE_SPECS} + _apply_dataset_field_overrides(args, cfg_dict, defaults, combined_specs) + dataset = EvalDatasetConfig(**cfg_dict) + datasets.append(dataset) + return datasets diff --git a/slime/utils/external_utils/__init__.py b/slime/utils/external_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a4d636226b9a4b5130dfd8e94f2cece5fd290e53 --- /dev/null +++ b/slime/utils/external_utils/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + diff --git a/slime/utils/external_utils/__pycache__/__init__.cpython-312.pyc b/slime/utils/external_utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..68056db08c40b5470027fc25da78c5818739de4d Binary files /dev/null and b/slime/utils/external_utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/slime/utils/external_utils/__pycache__/command_utils.cpython-312.pyc b/slime/utils/external_utils/__pycache__/command_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0132ef43ebf3fbfe6126f0712fdab1dc40c235ac Binary files /dev/null and b/slime/utils/external_utils/__pycache__/command_utils.cpython-312.pyc differ diff --git a/slime/utils/external_utils/command_utils.py b/slime/utils/external_utils/command_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ae93511121ff80b8c16dc8de2476a686ea5d3057 --- /dev/null +++ b/slime/utils/external_utils/command_utils.py @@ -0,0 +1,275 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +This file is not for slime framework itself, but as an optional utility to easily launch slime jobs and tests. +""" + +import datetime +import json +import os +import random +import time +from dataclasses import dataclass +from pathlib import Path + +from slime.utils.misc import exec_command +from slime.utils.typer_utils import dataclass_cli + +_ = exec_command, dataclass_cli + +repo_base_dir = Path(os.path.abspath(__file__)).resolve().parents[3] + + +def convert_checkpoint( + model_name, + megatron_model_type, + num_gpus_per_node: int, + multinode: bool = False, + extra_args: str = "", + dir_dst: str = "/root/models", + hf_checkpoint: str | None = None, +): + hf_checkpoint = hf_checkpoint or f"/root/models/{model_name}" + + # TODO shall we make it in host-mapped folder and thus can cache it to speedup CI + path_dst = f"{dir_dst}/{model_name}_torch_dist" + if Path(path_dst).exists(): + print(f"convert_checkpoint skip {path_dst} since exists") + return + + multinode_args = "" + if multinode: + # This variable can be provided via: + # `export SLURM_JOB_HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")` + print(f"{os.environ.get('SLURM_JOB_HOSTNAMES')=} {os.environ.get('SLURM_NODEID')=}") + job_hostnames = os.environ["SLURM_JOB_HOSTNAMES"].strip().split("\n") + master_addr = job_hostnames[0] + nnodes = len(job_hostnames) + node_rank = int(os.environ["SLURM_NODEID"]) + + multinode_args = ( + f"--master-addr {master_addr} " "--master-port 23456 " f"--nnodes={nnodes} " f"--node-rank {node_rank} " + ) + + exec_command( + f"source {repo_base_dir}/configs/models/{megatron_model_type}.sh && " + f"PYTHONPATH=/root/Megatron-LM " + f"torchrun " + f"--nproc-per-node {num_gpus_per_node} " + f"{multinode_args}" + f"tools/convert_hf_to_torch_dist.py " + "${MODEL_ARGS[@]} " + f"--hf-checkpoint {hf_checkpoint} " + f"--save {path_dst}" + f"{extra_args}" + ) + + +def rsync_simple(path_src: str, path_dst: str): + exec_command(f"mkdir -p {path_dst} && rsync -a --info=progress2 {path_src}/ {path_dst}") + + +def hf_download_dataset(full_name: str): + _, partial_name = full_name.split("/") + exec_command(f"hf download --repo-type dataset {full_name} --local-dir /root/datasets/{partial_name}") + + +def fp8_cast_bf16(path_src, path_dst): + if Path(path_dst).exists(): + print(f"fp8_cast_bf16 skip {path_dst} since exists") + return + + exec_command( + "python tools/fp8_cast_bf16.py " f"--input-fp8-hf-path {path_src} " f"--output-bf16-hf-path {path_dst} " + ) + + +# This class can be extended by concrete scripts +@dataclass +class ExecuteTrainConfig: + cuda_core_dump: bool = False + num_nodes: int = int(os.environ.get("SLURM_JOB_NUM_NODES", "1")) + extra_env_vars: str = "" + + +def execute_train( + train_args: str, + num_gpus_per_node: int, + megatron_model_type: str | None, + train_script: str = "train.py", + before_ray_job_submit=None, + extra_env_vars=None, + config: ExecuteTrainConfig | None = None, + rerun: bool = False, +): + if extra_env_vars is None: + extra_env_vars = {} + if config is None: + config = ExecuteTrainConfig() + external_ray = get_bool_env_var("SLIME_SCRIPT_EXTERNAL_RAY") + master_addr = os.environ.get("MASTER_ADDR", "127.0.0.1") + + train_backend_fsdp = "--train-backend fsdp" in train_args + assert train_backend_fsdp == (megatron_model_type is None) + + if rerun: + exec_command( + "pkill -9 sglang; " + "sleep 3; " + f"{'' if external_ray else 'ray stop --force; '}" + f"{'' if external_ray else 'pkill -9 ray; '}" + # cannot be run in CI, o/w kill the parent script + # TODO: do we really need this kill? (or can we instead kill slime) + # "pkill -9 python; " + "pkill -9 slime; " + "sleep 3; " + f"{'' if external_ray else 'pkill -9 ray; '}" + # "pkill -9 python; " + "pkill -9 slime; " + "pkill -9 redis; " + "true; " + ) + + if not external_ray: + exec_command( + # will prevent ray from buffering stdout/stderr + f"export PYTHONBUFFERED=16 && " + f"ray start --head --node-ip-address {master_addr} --num-gpus {num_gpus_per_node} --disable-usage-stats" + ) + + if (f := before_ray_job_submit) is not None: + f() + + runtime_env_json = json.dumps( + { + "env_vars": { + "PYTHONPATH": "/root/Megatron-LM/", + # If setting this in FSDP, the computation communication overlapping may have issues + **( + {} + if train_backend_fsdp + else { + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + } + ), + "NCCL_NVLS_ENABLE": str(int(check_has_nvlink())), + "no_proxy": f"127.0.0.1,{master_addr}", + # This is needed by megatron / torch distributed in multi-node setup + "MASTER_ADDR": master_addr, + **( + { + "CUDA_ENABLE_COREDUMP_ON_EXCEPTION": "1", + "CUDA_COREDUMP_SHOW_PROGRESS": "1", + "CUDA_COREDUMP_GENERATION_FLAGS": "skip_nonrelocated_elf_images,skip_global_memory,skip_shared_memory,skip_local_memory,skip_constbank_memory", + "CUDA_COREDUMP_FILE": "/root/shared_data/cuda_coredump_%h.%p.%t", + } + if config.cuda_core_dump + else {} + ), + **extra_env_vars, + **_parse_extra_env_vars(config.extra_env_vars), + } + } + ) + + if get_bool_env_var("SLIME_SCRIPT_ENABLE_RAY_SUBMIT", "1"): + cmd_megatron_model_source = ( + f'source "{repo_base_dir}/configs/models/{megatron_model_type}.sh" && ' + if megatron_model_type is not None + else "" + ) + exec_command( + f"export no_proxy=127.0.0.1 && export PYTHONBUFFERED=16 && " + f"{cmd_megatron_model_source}" + f'ray job submit --address="http://127.0.0.1:8265" ' + f"--runtime-env-json='{runtime_env_json}' " + f"-- python3 {train_script} " + f"{'${MODEL_ARGS[@]}' if megatron_model_type is not None else ''} " + f"{train_args}" + ) + + +def _parse_extra_env_vars(text: str): + try: + return json.loads(text) + except ValueError: + return {kv[0]: kv[1] for item in text.split(" ") if item.strip() != "" if (kv := item.split("=")) or True} + + +def check_has_nvlink(): + output = exec_command("nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l", capture_output=True) + return int(output) > 0 + + +def get_default_wandb_args(test_file: str, run_name_prefix: str | None = None, run_id: str | None = None): + if not os.environ.get("WANDB_API_KEY"): + print("Skip wandb configuration since WANDB_API_KEY is not found") + return "" + + test_file = Path(test_file) + test_name = test_file.stem + if len(test_name) < 6: + test_name = f"{test_file.parent.name}_{test_name}" + + wandb_run_name = run_id or create_run_id() + if (x := os.environ.get("GITHUB_COMMIT_NAME")) is not None: + wandb_run_name += f"_{x}" + if (x := run_name_prefix) is not None: + wandb_run_name = f"{x}_{wandb_run_name}" + + # do not put wandb_api_key value here to avoid leaking to logs explicitly + return ( + "--use-wandb " + f"--wandb-project slime-{test_name} " + f"--wandb-group {wandb_run_name} " + f"--wandb-key ${{WANDB_API_KEY}} " + "--disable-wandb-random-suffix " + ) + + +def create_run_id() -> str: + return datetime.datetime.utcnow().strftime("%y%m%d-%H%M%S") + f"-{random.Random().randint(0, 999):03d}" + + +_warned_bool_env_var_keys = set() + + +# copied from SGLang +def get_bool_env_var(name: str, default: str = "false") -> bool: + value = os.getenv(name, default) + value = value.lower() + + truthy_values = ("true", "1") + falsy_values = ("false", "0") + + if (value not in truthy_values) and (value not in falsy_values): + if value not in _warned_bool_env_var_keys: + print(f"get_bool_env_var({name}) see non-understandable value={value} and treat as false") + _warned_bool_env_var_keys.add(value) + + return value in truthy_values + + +def get_env_enable_infinite_run(): + return get_bool_env_var("SLIME_TEST_ENABLE_INFINITE_RUN", "false") + + +def save_to_temp_file(text: str, ext: str): + path = Path(f"/tmp/slime_temp_file_{time.time()}_{random.randrange(0, 10000000)}.{ext}") + path.write_text(text) + print(f"Write the following content to {path=}: {text=}") + return str(path) + + +NUM_GPUS_OF_HARDWARE = { + "H100": 8, + "GB200": 4, + "GB300": 4, +} + +GENERATION_HARDWARE = { + "H100": "Hopper", + "GB200": "Blackwell", + "GB300": "Blackwell", +} diff --git a/slime/utils/flops_utils.py b/slime/utils/flops_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a6f4b522905f8d9117392eca2b2727b44a2fb6b3 --- /dev/null +++ b/slime/utils/flops_utils.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +def calculate_embedding_flops(seqlen, hidden_size): + return 2 * seqlen * hidden_size + + +def calculate_lm_head_flops(seqlen, hidden_size, vocab_size): + return 2 * seqlen * hidden_size * vocab_size + + +def calculate_qkv_projection_flops(args, seqlen, hidden_size, num_attention_heads, num_query_groups): + if args.q_lora_rank is None: + q_flops = 2 * seqlen * hidden_size * num_attention_heads * args.kv_channels + else: + q_flops = ( + 2 + * seqlen + * args.q_lora_rank + * (args.hidden_size + args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim)) + ) + if args.kv_lora_rank is None: + kv_flops = 2 * 2 * seqlen * hidden_size * num_query_groups * args.kv_channels + else: + kv_flops = ( + 2 + * seqlen + * ( + args.kv_lora_rank + * (args.hidden_size + args.num_attention_heads * (args.qk_head_dim + args.v_head_dim)) + + args.hidden_size * args.qk_pos_emb_head_dim + ) + ) + + return q_flops + kv_flops + + +def calculate_attention_flops(args, seqlen, num_attention_heads): + # QK^T with causal + if args.qk_pos_emb_head_dim: + flops = 2 * num_attention_heads * seqlen * seqlen * (args.qk_head_dim + args.qk_pos_emb_head_dim) / 2 + else: + flops = 2 * num_attention_heads * seqlen * seqlen * args.kv_channels / 2 + # A*V + if args.v_head_dim: + flops += num_attention_heads * seqlen * seqlen * args.v_head_dim + else: + flops += num_attention_heads * seqlen * seqlen * args.kv_channels + return flops + + +def calculate_output_flops(seqlen, hidden_size): + return 2 * seqlen * hidden_size * hidden_size + + +def calculate_mlp_flops(seqlen, hidden_size, ffn_hidden_size): + return 2 * seqlen * hidden_size * ffn_hidden_size * 3 + + +def calculate_layer_flops(args, seqlen, hidden_size, num_attention_heads, num_query_groups, ffn_hidden_size): + return ( + calculate_qkv_projection_flops(args, seqlen, hidden_size, num_attention_heads, num_query_groups) + + calculate_attention_flops(args, seqlen, num_attention_heads) + + calculate_output_flops(seqlen, hidden_size) + + calculate_mlp_flops(seqlen, hidden_size, ffn_hidden_size) + ) + + +def calculate_fwd_flops( + seqlens, + args, +): + hidden_size = args.hidden_size + num_attention_heads = args.num_attention_heads + num_query_groups = args.num_query_groups + vocab_size = args.vocab_size + + total_flops = 0 + + dense_ffn = args.ffn_hidden_size + if args.num_experts is None: + num_dense_layers = args.num_layers + num_moe_layers = 0 + else: + shared_expert_ffn = getattr(args, "moe_shared_expert_intermediate_size", None) + if shared_expert_ffn is None: + shared_expert_ffn = 0 + + moe_ffn = args.moe_ffn_hidden_size * args.moe_router_topk + shared_expert_ffn + if hasattr(args, "moe_layer_freq"): + if isinstance(args.moe_layer_freq, list): + num_dense_layers = sum(1 for freq in args.moe_layer_freq if freq == 0) + num_moe_layers = sum(1 for freq in args.moe_layer_freq if freq > 0) + else: + num_dense_layers = sum(1 for i in range(args.num_layers) if i % args.moe_layer_freq != 0) + num_moe_layers = sum(1 for i in range(args.num_layers) if i % args.moe_layer_freq == 0) + else: + num_dense_layers = 0 + num_moe_layers = args.num_layers + + for seqlen in seqlens: + if num_dense_layers > 0: + total_flops += ( + calculate_layer_flops( + args, + seqlen, + hidden_size, + num_attention_heads, + num_query_groups, + dense_ffn, + ) + * num_dense_layers + ) + + if num_moe_layers > 0: + total_flops += ( + calculate_layer_flops( + args, + seqlen, + hidden_size, + num_attention_heads, + num_query_groups, + moe_ffn, + ) + * num_moe_layers + ) + + total_flops += calculate_lm_head_flops(seqlen, hidden_size, vocab_size) + + return total_flops diff --git a/slime/utils/fp8_kernel.py b/slime/utils/fp8_kernel.py new file mode 100644 index 0000000000000000000000000000000000000000..4d89375b9544700bd7e0299ec92aaa5876a470f6 --- /dev/null +++ b/slime/utils/fp8_kernel.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +import triton +import triton.language as tl + +fp8_dtype = torch.float8_e4m3fn +fp8_max = torch.finfo(fp8_dtype).max +fp8_min = -fp8_max + + +def ceil_div(x: int, y: int) -> int: + """ + Perform ceiling division of two integers. + + Args: + x: the dividend. + y: the divisor. + + Returns: + The result of the ceiling division. + """ + return (x + y - 1) // y + + +@triton.jit +def _blockwise_cast_to_fp8_triton( + X, + Y, + S, + stride_xm, + stride_xn, + stride_ym, + stride_yn, + stride_sm, + stride_sn, + M, + N, + eps, + fp8_min, + fp8_max, + BLOCK_M: tl.constexpr = 32, + BLOCK_N: tl.constexpr = 128, +): + pid_m = tl.cast(tl.program_id(axis=0), tl.int64) + pid_n = tl.cast(tl.program_id(axis=1), tl.int64) + off_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + off_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + mask_m = off_m < M + mask_n = off_n < N + mask = mask_m[:, None] & mask_n[None, :] + + x = tl.load(X + off_m[:, None] * stride_xm + off_n[None, :] * stride_xn, mask=mask, other=0.0).to(tl.float32) + _absmax = tl.maximum(tl.max(tl.abs(x)), eps) + x_s = _absmax / fp8_max + s_inv = 1.0 / x_s + y_q = tl.clamp(x * s_inv, fp8_min, fp8_max).to(Y.dtype.element_ty) + + tl.store(Y + off_m[:, None] * stride_ym + off_n[None, :] * stride_yn, y_q, mask=mask) + tl.store(S + pid_m * stride_sm + pid_n * stride_sn, x_s) + + +def blockwise_cast_to_fp8_triton(x: torch.Tensor, block_size=None) -> tuple[torch.Tensor, torch.Tensor]: + BLOCK_M, BLOCK_N = 128, 128 + if block_size: + BLOCK_M, BLOCK_N = block_size[0], block_size[1] + M, N = x.shape + y = torch.empty(M, N, device=x.device, dtype=torch.float8_e4m3fn) + s = torch.empty(ceil_div(M, BLOCK_M), ceil_div(N, BLOCK_N), dtype=torch.float32, device=x.device) + + def grid(meta): + return (triton.cdiv(M, meta["BLOCK_M"]), triton.cdiv(N, meta["BLOCK_N"])) + + if x.is_contiguous(): + kwargs = {"BLOCK_M": BLOCK_M, "BLOCK_N": BLOCK_N, "num_warps": 8, "num_stages": 2} + else: + kwargs = {"BLOCK_M": BLOCK_M, "BLOCK_N": BLOCK_N, "num_warps": 1, "num_stages": 4} + _blockwise_cast_to_fp8_triton[grid]( + x, y, s, *x.stride(), *y.stride(), *s.stride(), M, N, 1e-10, fp8_min, fp8_max, **kwargs + ) + return y, s diff --git a/slime/utils/health_monitor.py b/slime/utils/health_monitor.py new file mode 100644 index 0000000000000000000000000000000000000000..78e7f0f36fbd19e4da5c8afde07b2daad2eacc70 --- /dev/null +++ b/slime/utils/health_monitor.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +import threading + +import ray + + +logger = logging.getLogger(__name__) + + +class RolloutHealthMonitor: + def __init__(self, rollout_manager, args): + # TODO may remove this dependency after refactoring + self._rollout_manager = rollout_manager + + self._thread = None + self._stop_event = None + self._check_interval = args.rollout_health_check_interval + self._check_timeout = args.rollout_health_check_timeout + self._check_first_wait = args.rollout_health_check_first_wait + + def start(self) -> bool: + if not self._rollout_manager.rollout_engines: + return False + + assert self._thread is None, "Health monitor thread is already running." + + logger.info("Starting RolloutHealthMonitor...") + self._stop_event = threading.Event() + self._thread = threading.Thread( + target=self._health_monitor_loop, + name="RolloutHealthMonitor", + daemon=True, + ) + self._thread.start() + logger.info("RolloutHealthMonitor started.") + return True + + def stop(self) -> None: + if not self._thread: + return + + logger.info("Stopping RolloutHealthMonitor...") + assert self._stop_event is not None + self._stop_event.set() + timeout = self._check_timeout + self._check_interval + 5 + self._thread.join(timeout=timeout) + if self._thread.is_alive(): + logging.warning("Rollout health monitor thread did not terminate within %.1fs", timeout) + else: + logger.info("RolloutHealthMonitor stopped.") + + self._thread = None + self._stop_event = None + + def _health_monitor_loop(self) -> None: + assert self._stop_event is not None + logger.info(f"Health monitor loop started. Waiting for first wait: {self._check_first_wait}s") + # TODO: need to be waiting for the large moe to be ready. this is hacky. + if self._stop_event.wait(self._check_first_wait): + logger.info("Health monitor stopped during first wait.") + return + while not self._stop_event.is_set(): + self._run_health_checks() + if self._stop_event.wait(self._check_interval): + break + + def _run_health_checks(self) -> None: + for rollout_engine_id, engine in enumerate(self._rollout_manager.rollout_engines): + if self._stop_event is not None and self._stop_event.is_set(): + break + self._check_engine_health(rollout_engine_id, engine) + + def _check_engine_health(self, rollout_engine_id, engine) -> None: + if engine is None: + logger.info(f"Skipping health check for engine {rollout_engine_id} (None)") + return + + try: + ray.get(engine.health_generate.remote(timeout=self._check_timeout)) + except Exception as e: + logger.error( + f"Health check failed for rollout engine {rollout_engine_id} (ray timeout or error). Killing actor. Exception: {e}" + ) + self._kill_engine(rollout_engine_id=rollout_engine_id) + + def _kill_engine(self, rollout_engine_id: int): + logger.info(f"Killing engine group {rollout_engine_id}...") + for i in range( + rollout_engine_id * self._rollout_manager.nodes_per_engine, + (rollout_engine_id + 1) * self._rollout_manager.nodes_per_engine, + ): + engine = self._rollout_manager.all_rollout_engines[i] + if engine: + logger.info(f"Shutting down and killing engine at index {i}") + try: + ray.get(engine.shutdown.remote()) + ray.kill(engine) + logger.info(f"Successfully killed engine at index {i}") + except Exception as e: + logger.warning(f"Fail to kill engine at index {i} (e: {e})") + else: + logger.info(f"Engine at index {i} is already None") + self._rollout_manager.all_rollout_engines[i] = None diff --git a/slime/utils/http_utils.py b/slime/utils/http_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..35d6cffc36fee705a3bf205ab5c8b3ccee83e1d3 --- /dev/null +++ b/slime/utils/http_utils.py @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import ipaddress +import json +import logging +import multiprocessing +import os +import random +import socket + +import httpx + +logger = logging.getLogger(__name__) + +SLIME_HOST_IP_ENV = "SLIME_HOST_IP" + + +def find_available_port(base_port: int): + port = base_port + random.randint(100, 1000) + while True: + if is_port_available(port): + return port + if port < 60000: + port += 42 + else: + port -= 43 + + +def is_port_available(port): + """Return whether a port is available.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + try: + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind(("", port)) + s.listen(1) + return True + except OSError: + return False + except OverflowError: + return False + + +def get_host_info(): + hostname = socket.gethostname() + + if env_overwrite_local_ip := os.getenv(SLIME_HOST_IP_ENV, None): + return hostname, env_overwrite_local_ip + + # try DNS + try: + return hostname, socket.gethostbyname(hostname) + except socket.gaierror: + pass + + # try IPv4 + try: + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as udp_sock: + udp_sock.connect(("8.8.8.8", 80)) # Google DNS + return hostname, udp_sock.getsockname()[0] + except OSError: + pass + + # try IPv6 + try: + with socket.socket(socket.AF_INET6, socket.SOCK_DGRAM) as s6: + s6.connect(("2001:4860:4860::8888", 80)) + return hostname, s6.getsockname()[0] + except OSError: + pass + + # hostname -I + try: + local_ip = os.popen("hostname -I | awk '{print $1}'").read().strip() + return hostname, local_ip or "::1" + except Exception: + return hostname, "::1" + + +def _wrap_ipv6(host): + """Wrap IPv6 address in [] if needed.""" + try: + ipaddress.IPv6Address(host.strip("[]")) + return f"[{host.strip('[]')}]" + except ipaddress.AddressValueError: + return host + + +def run_router(args): + try: + from sglang_router.launch_router import launch_router + + router = launch_router(args) + if router is None: + return 1 + return 0 + except Exception as e: + logger.info(e) + return 1 + + +def terminate_process(process: multiprocessing.Process, timeout: float = 1.0) -> None: + """Terminate a process gracefully, with forced kill as fallback. + + Args: + process: The process to terminate + timeout: Seconds to wait for graceful termination before forcing kill + """ + if not process.is_alive(): + return + + process.terminate() + process.join(timeout=timeout) + if process.is_alive(): + process.kill() + process.join() + + +_http_client: httpx.AsyncClient | None = None +_client_concurrency: int = 0 + +# Optional Ray-based distributed POST dispatch +_distributed_post_enabled: bool = False +_post_actors: list[object] = [] +_post_actor_idx: int = 0 + + +def _next_actor(): + global _post_actor_idx + if not _post_actors: + return None + actor = _post_actors[_post_actor_idx % len(_post_actors)] + _post_actor_idx = (_post_actor_idx + 1) % len(_post_actors) + return actor + + +async def _post(client, url, payload, max_retries=60): + retry_count = 0 + while retry_count < max_retries: + try: + response = await client.post(url, json=payload or {}) + response.raise_for_status() + try: + output = response.json() + except json.JSONDecodeError: + output = response.text + except Exception as e: + retry_count += 1 + + if isinstance(e, httpx.HTTPStatusError): + response_text = e.response.text + else: + response_text = None + + logger.info( + f"Error: {e}, retrying... (attempt {retry_count}/{max_retries}, url={url}, response={response_text})" + ) + if retry_count >= max_retries: + logger.info(f"Max retries ({max_retries}) reached, failing... (url={url})") + raise e + await asyncio.sleep(1) + continue + break + + return output + + +def init_http_client(args): + """Initialize HTTP client and optionally enable distributed POST via Ray.""" + global _http_client, _client_concurrency, _distributed_post_enabled + if not args.rollout_num_gpus: + return + + _client_concurrency = args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine + if _http_client is None: + _http_client = httpx.AsyncClient( + limits=httpx.Limits(max_connections=_client_concurrency), + timeout=httpx.Timeout(None), + ) + + # Optionally initialize distributed POST via Ray without changing interfaces + if args.use_distributed_post: + _init_ray_distributed_post(args) + _distributed_post_enabled = True + + +def _init_ray_distributed_post(args): + """Initialize one or more Ray async actors per node for HTTP POST. + + Uses NodeAffinitySchedulingStrategy to place actors on distinct nodes. + Controlled by SLIME_HTTP_POST_ACTORS_PER_NODE. + """ + global _post_actors + if _post_actors: + return # Already initialized + + import ray + from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + + # Discover alive nodes + nodes = [n for n in ray.nodes() if n.get("Alive")] + if not nodes: + raise RuntimeError("No alive Ray nodes to place HTTP POST actors.") + + # Define the async actor + @ray.remote + class _HttpPosterActor: + def __init__(self, concurrency: int): + # Lazy creation to this actor's event loop + self._client = httpx.AsyncClient( + limits=httpx.Limits(max_connections=max(1, concurrency)), + timeout=httpx.Timeout(None), + ) + + async def do_post(self, url, payload, max_retries=60): + return await _post(self._client, url, payload, max_retries) + + # Create actors per node + created = [] + # Distribute client concurrency across actors (at least 1 per actor) + per_actor_conc = (_client_concurrency + len(nodes)) // len(nodes) + + for node in nodes: + node_id = node["NodeID"] + scheduling = NodeAffinitySchedulingStrategy(node_id=node_id, soft=False) + for _ in range(args.num_gpus_per_node): + actor = _HttpPosterActor.options( + name=None, + lifetime="detached", + scheduling_strategy=scheduling, + max_concurrency=per_actor_conc, + # Use tiny CPU to schedule + num_cpus=0.001, + ).remote(per_actor_conc) + created.append(actor) + + _post_actors = created + + +async def post(url, payload, max_retries=60): + # If distributed mode is enabled and actors exist, dispatch via Ray. + if _distributed_post_enabled and _post_actors: + try: + import ray + + actor = _next_actor() + if actor is not None: + # Use a thread to avoid blocking the event loop on ray.get + obj_ref = actor.do_post.remote(url, payload, max_retries) + return await asyncio.to_thread(ray.get, obj_ref) + except Exception as e: + logger.info(f"[http_utils] Distributed POST failed, falling back to local: {e} (url={url})") + # fall through to local + + return await _post(_http_client, url, payload, max_retries) + + +async def get(url): + response = await _http_client.get(url) + response.raise_for_status() + output = response.json() + return output diff --git a/slime/utils/iter_utils.py b/slime/utils/iter_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f802ccf2c6f35de2eb3216da2e6b3e7b8743371c --- /dev/null +++ b/slime/utils/iter_utils.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections import defaultdict +from collections.abc import Callable, Iterable +from typing import Any + +import torch + + +# details: https://stackoverflow.com/questions/773/how-do-i-use-itertools-groupby +def group_by(iterable, key=None): + """Similar to itertools.groupby, but do not require iterable to be sorted""" + ret = defaultdict(list) + for item in iterable: + ret[key(item) if key is not None else item].append(item) + return dict(ret) + + +# TODO fsdp can also use this +def chunk_named_params_by_size(named_params: Iterable[tuple[str, torch.Tensor]], chunk_size: int): + return _chunk_by_size( + named_params, + compute_size=lambda named_weight: named_weight[1].nbytes, + chunk_size=chunk_size, + ) + + +def _chunk_by_size(objects: Iterable[Any], compute_size: Callable[[Any], int], chunk_size: int): + bucket: list[Any] = [] + bucket_size = 0 + + for obj in objects: + obj_size = compute_size(obj) + + if bucket and (bucket_size + obj_size) >= chunk_size: + yield bucket + bucket = [] + bucket_size = 0 + + bucket.append(obj) + bucket_size += obj_size + + if bucket: + yield bucket diff --git a/slime/utils/logging_utils.py b/slime/utils/logging_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a38a039787901f5cb685ef0c09e45a41c072ad97 --- /dev/null +++ b/slime/utils/logging_utils.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging + +_LOGGER_CONFIGURED = False + + +# ref: SGLang +def configure_logger(prefix: str = ""): + global _LOGGER_CONFIGURED + if _LOGGER_CONFIGURED: + return + + _LOGGER_CONFIGURED = True + + logging.basicConfig( + level=logging.INFO, + format=f"[%(asctime)s{prefix}] %(filename)s:%(lineno)d - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + force=True, + ) diff --git a/slime/utils/mask_utils.py b/slime/utils/mask_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..bb3319fb9e99d4c46b14c9f12fb36d3d987a2304 --- /dev/null +++ b/slime/utils/mask_utils.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from transformers import AutoTokenizer + + +def get_response_lengths(loss_masks: list[list[int]]) -> list[int]: + return [mask.count(1) if 1 in mask else 0 for mask in loss_masks] + + +class MultiTurnLossMaskGenerator: + def __init__(self, tokenizer: AutoTokenizer, tokenizer_type: str = "qwen"): + self.tokenizer = tokenizer + self.system_message_length, self.gen_token_length = self.get_system_message_length() + self.tokenizer_type = tokenizer_type + + def get_response_lengths(self, loss_masks: list[list[int]]) -> list[int]: + return get_response_lengths(loss_masks) + + def find_all_sublist_indices(self, main_list, sublist): + sublist_len = len(sublist) + indices = [] + for i in range(len(main_list) - sublist_len + 1): + if main_list[i : i + sublist_len] == sublist: + indices.append(i) + return indices + + def get_system_message_length(self) -> tuple[int, int]: + test_string = "FOR TESTING ONLY" + test_messages = [ + {"role": "user", "content": test_string}, + {"role": "user", "content": test_string}, + ] + raw_token_ids = self.tokenizer(test_string, add_special_tokens=False)["input_ids"] + chat_template_token = self.tokenizer.apply_chat_template( + test_messages, add_special_tokens=False, tokenize=False + ) + chat_template_token_ids = self.tokenizer(chat_template_token, add_special_tokens=False)["input_ids"] + idx_1, idx_2 = self.find_all_sublist_indices(chat_template_token_ids, raw_token_ids) + end_interval = len(chat_template_token_ids) - len(raw_token_ids) - idx_2 + gen_token_length = len( + self.tokenizer.apply_chat_template( + test_messages, add_special_tokens=False, tokenize=True, add_generation_prompt=True + ) + ) - len(chat_template_token_ids) + + system_message_length = idx_1 - ((idx_2 - idx_1) - end_interval - len(raw_token_ids)) + return system_message_length, gen_token_length + + def gen_multi_turn_loss_mask_qwen( + self, messages: list[dict], tools: list[dict] = None + ) -> tuple[list[int], list[int]]: + all_loss_masks = [] + all_token_ids = [] + + for i, message in enumerate(messages): + if i == 0: + message_ids = self.tokenizer.apply_chat_template([message], tokenize=True, tools=tools) + else: + message_ids = self.tokenizer.apply_chat_template([message], tokenize=True) + + if message["role"] != "system" and i > 0: + message_ids = message_ids[self.system_message_length :] + + if message["role"] == "assistant": + loss_mask = [0] * self.gen_token_length + [1] * (len(message_ids) - self.gen_token_length) + else: + loss_mask = [0] * len(message_ids) + + if message.get("step_loss_mask", 1) != 1: + loss_mask = [0] * len(message_ids) + + all_loss_masks.extend(loss_mask) + all_token_ids.extend(message_ids) + + return all_token_ids, all_loss_masks + + def gen_multi_turn_loss_mask_qwen3( + self, messages: list[dict], tools: list[dict] = None + ) -> tuple[list[int], list[int]]: + all_loss_masks = [] + all_token_ids = [] + + prefix_message = {"role": "user", "content": "FOR CALCULATING LOSS MASK ONLY"} + prefix_token_ids = self.tokenizer.apply_chat_template([prefix_message], tokenize=True) + + for i, message in enumerate(messages): + if i == 0: + tailed_message_ids = self.tokenizer.apply_chat_template( + [message, prefix_message], tokenize=True, tools=tools + ) + message_ids = tailed_message_ids[: -len(prefix_token_ids)] + else: + prefixed_message_ids = self.tokenizer.apply_chat_template([prefix_message, message], tokenize=True) + message_ids = prefixed_message_ids[len(prefix_token_ids) :] + + if message["role"] != "system" and i > 0: + message_ids = message_ids[self.system_message_length :] + + if message["role"] == "assistant": + loss_mask = [0] * self.gen_token_length + [1] * (len(message_ids) - self.gen_token_length) + else: + loss_mask = [0] * len(message_ids) + + if message.get("step_loss_mask", 1) != 1: + loss_mask = [0] * len(message_ids) + + all_loss_masks.extend(loss_mask) + all_token_ids.extend(message_ids) + + return all_token_ids, all_loss_masks + + def gen_multi_turn_loss_mask_distill_qwen( + self, messages: list[dict], tools: list[dict] = None + ) -> tuple[list[int], list[int]]: + prompt = self.tokenizer.apply_chat_template( + messages[:1], tokenize=False, add_generation_prompt=True, tools=tools + ) + response = messages[-1]["content"] + prompt_tokens = self.tokenizer(prompt, add_special_tokens=False)["input_ids"] + response_tokens = self.tokenizer(response, add_special_tokens=False)["input_ids"] + + response_length = len(response_tokens) + token_ids = prompt_tokens + response_tokens + loss_mask = [0] * len(prompt_tokens) + [1] * response_length + + if messages[-1].get("step_loss_mask", 1) != 1: + loss_mask = [0] * len(token_ids) + return token_ids, loss_mask + + def get_loss_mask(self, messages: list[dict], tools: list[dict] = None) -> tuple[list[int], list[int]]: + if self.tokenizer_type == "qwen": + if "<｜Assistant｜>" in self.tokenizer.get_added_vocab(): + return self.gen_multi_turn_loss_mask_distill_qwen(messages, tools) + + return self.gen_multi_turn_loss_mask_qwen(messages, tools) + elif self.tokenizer_type == "qwen3": + return self.gen_multi_turn_loss_mask_qwen3(messages, tools) + elif self.tokenizer_type == "distill_qwen": + return self.gen_multi_turn_loss_mask_distill_qwen(messages, tools) + else: + raise ValueError(f"Unsupported tokenizer type: {self.tokenizer_type}") + + def get_loss_mask_with_multimodal_alignment( + self, messages: list[dict], input_ids: list[int], tools: list[dict] = None + ) -> tuple[list[int], list[int]]: + text = [] + for msg in messages: + if isinstance(msg.get("content"), list): + text_parts = [] + for item in msg["content"]: + if isinstance(item, dict) and item.get("type") == "text": + text_parts.append(item.get("text", "")) + elif isinstance(item, str): + text_parts.append(item) + text.append({"role": msg["role"], "content": " ".join(text_parts)}) + else: + text.append(msg) + + _, loss_mask_text = self.get_loss_mask(text, tools=tools) + + diff = len(input_ids) - len(loss_mask_text) + assert diff >= 0, ( + f"input_ids (length={len(input_ids)}) is shorter than text loss_mask (length={len(loss_mask_text)}) " + f"Please check if processor and tokenizer tokenization are consistent." + ) + loss_mask = [0] * diff + loss_mask_text + + return input_ids, loss_mask + + def get_text_from_loss_mask(self, token_ids: list[int], loss_masks: list[int]) -> list[str]: + selected_texts = [] + current_tokens = [] + + for idx, mask in enumerate(loss_masks): + if mask == 1: + current_tokens.append(token_ids[idx]) + elif current_tokens: + selected_texts.append(self.tokenizer.decode(current_tokens)) + current_tokens = [] + + if current_tokens: + selected_texts.append(self.tokenizer.decode(current_tokens)) + + return selected_texts diff --git a/slime/utils/megatron_bridge_utils.py b/slime/utils/megatron_bridge_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..950f13352c3a4b14ee0c028363403abed52f4ac2 --- /dev/null +++ b/slime/utils/megatron_bridge_utils.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from contextlib import contextmanager + +try: + from megatron.core.utils import unwrap_model +except ImportError: + unwrap_model = None + + +@contextmanager +def patch_megatron_model(model): + unwrapped_model = unwrap_model(model)[0] + model_config = unwrapped_model.config + attribute_was_added = False + if not hasattr(model_config, "share_embeddings_and_output_weights"): + model_config.share_embeddings_and_output_weights = unwrapped_model.share_embeddings_and_output_weights + attribute_was_added = True + + try: + yield + finally: + if attribute_was_added: + delattr(model_config, "share_embeddings_and_output_weights") diff --git a/slime/utils/memory_utils.py b/slime/utils/memory_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..14652067dc842da68539b8106836f81e7e0fa5cc --- /dev/null +++ b/slime/utils/memory_utils.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import gc +import logging + +import torch +import torch.distributed as dist + +logger = logging.getLogger(__name__) + + +def clear_memory(clear_host_memory: bool = False): + torch.cuda.synchronize() + gc.collect() + torch.cuda.empty_cache() + if clear_host_memory: + torch._C._host_emptyCache() + + +def available_memory(): + device = torch.cuda.current_device() + free, total = torch.cuda.mem_get_info(device) + return { + "gpu": str(device), + "total_GB": _byte_to_gb(total), + "free_GB": _byte_to_gb(free), + "used_GB": _byte_to_gb(total - free), + "allocated_GB": _byte_to_gb(torch.cuda.memory_allocated(device)), + "reserved_GB": _byte_to_gb(torch.cuda.memory_reserved(device)), + } + + +def _byte_to_gb(n: int): + return round(n / (1024**3), 2) + + +def print_memory(msg, clear_before_print: bool = False): + if clear_before_print: + clear_memory() + + memory_info = available_memory() + # Need to print for all ranks, b/c different rank can have different behaviors + logger.info( + f"[Rank {dist.get_rank()}] Memory-Usage {msg}{' (cleared before print)' if clear_before_print else ''}: {memory_info}" + ) + return memory_info diff --git a/slime/utils/metric_checker.py b/slime/utils/metric_checker.py new file mode 100644 index 0000000000000000000000000000000000000000..17c5d761723636b1232f77410a6fd80a238244af --- /dev/null +++ b/slime/utils/metric_checker.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging + +logger = logging.getLogger(__name__) + + +class MetricChecker: + @staticmethod + def maybe_create(args): + if args.ci_test and (args.ci_metric_checker_key is not None): + return MetricChecker(args) + return None + + def __init__(self, args): + self.args = args + self._exists_check_success = False + + def on_eval(self, metrics: dict[str, float]): + actual_value = metrics.get(self.args.ci_metric_checker_key) + assert actual_value is not None, f"{metrics=} {self.args.ci_metric_checker_key=}" + + check_success = actual_value >= self.args.ci_metric_checker_threshold + logger.info(f"[MetricChecker] {check_success=} {actual_value=} {self.args.ci_metric_checker_threshold=}") + + self._exists_check_success |= check_success + + def dispose(self): + assert self._exists_check_success, "[MetricChecker] accuracy check failed" + logger.info("[MetricChecker] pass dispose check") diff --git a/slime/utils/metric_utils.py b/slime/utils/metric_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..040a6ef0bf574f24982bec034718ff52639b2ab5 --- /dev/null +++ b/slime/utils/metric_utils.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math +from typing import Any, Literal + +import numpy as np + + +def dict_add_prefix(d: dict[str, Any], prefix: str) -> dict[str, Any]: + return {f"{prefix}{k}": v for k, v in d.items()} + + +def compute_pass_rate( + flat_rewards: list[float], + group_size: int, + num_groups: int | None = None, +): + if group_size == 1: + return {} + + if num_groups is None: + num_groups = len(flat_rewards) // group_size + + pass_rate_name_list = [2**i for i in range(int(math.log2(group_size)) + 1)] + + assert len(flat_rewards) == num_groups * group_size, f"{len(flat_rewards)=} {num_groups=} {group_size=}" + rewards_of_group = np.array(flat_rewards).reshape(num_groups, group_size) + + log_dict = {} + for k in pass_rate_name_list: + num_correct = np.sum(rewards_of_group == 1, axis=1) + num_samples = np.full(num_groups, group_size) + + pass_k_estimates = _estimate_pass_at_k(num_samples, num_correct, k) + + pass_k = np.mean(pass_k_estimates) + log_dict[f"pass@{k}"] = pass_k + + return log_dict + + +def _estimate_pass_at_k(num_samples, num_correct, k): + """ + Estimates pass@k of each problem and returns them in an array. + """ + + def estimator(n, c, k): + """ + Calculates 1 - comb(n - c, k) / comb(n, k). + """ + if n - c < k: + return 1.0 + return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1)) + + return np.array([estimator(int(n), int(c), k) for n, c in zip(num_samples, num_correct, strict=False)]) + + +def compute_statistics(values: list[float]) -> dict[str, float]: + values = np.array(values) + return { + "mean": np.mean(values).item(), + "median": np.median(values).item(), + } + + +def compression_ratio( + data: str | bytes, + *, + encoding: str = "utf-8", + algorithm: Literal["zlib", "gzip", "bz2", "lzma"] = "zlib", + level: int = 9, +) -> tuple[float, float]: + if isinstance(data, str): + raw = data.encode(encoding) + else: + raw = data + + original = len(raw) + if original == 0: + return float("inf"), 0.0 + + if algorithm == "zlib": + import zlib + + compressed = zlib.compress(raw, level) + elif algorithm == "gzip": + import gzip + + compressed = gzip.compress(raw, compresslevel=level) + elif algorithm == "bz2": + import bz2 + + compressed = bz2.compress(raw, compresslevel=level) + elif algorithm == "lzma": + import lzma + + compressed = lzma.compress(raw, preset=level) + else: + raise ValueError(f"Unsupported algorithm: {algorithm}") + + comp_len = len(compressed) + if comp_len == 0: + return float("inf"), 100.0 + + ratio = original / comp_len + savings_pct = 100.0 * (1.0 - comp_len / original) + return ratio, savings_pct + + +def has_repetition(text: str = None): + if len(text) > 10000 and compression_ratio(text[-10000:])[0] > 10: + return True + else: + return False + + +def compute_rollout_step(args, rollout_id): + if args.wandb_always_use_train_step: + return rollout_id * args.rollout_batch_size * args.n_samples_per_prompt // args.global_batch_size + return rollout_id diff --git a/slime/utils/misc.py b/slime/utils/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..50101097d289ecdd35b06b42b8e089674e1e6fdd --- /dev/null +++ b/slime/utils/misc.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import importlib +import subprocess + +import ray + +from slime.utils.http_utils import is_port_available + + +def load_function(path): + """ + Load a function from a module. + :param path: The path to the function, e.g. "module.submodule.function". + :return: The function object. + """ + module_path, _, attr = path.rpartition(".") + module = importlib.import_module(module_path) + return getattr(module, attr) + + +class SingletonMeta(type): + """ + A metaclass for creating singleton classes. + """ + + _instances = {} + + def __call__(cls, *args, **kwargs): + if cls not in cls._instances: + instance = super().__call__(*args, **kwargs) + cls._instances[cls] = instance + return cls._instances[cls] + + +def exec_command(cmd: str, capture_output: bool = False) -> str | None: + print(f"EXEC: {cmd}", flush=True) + + try: + result = subprocess.run( + ["bash", "-c", cmd], + shell=False, + check=True, + capture_output=capture_output, + **(dict(text=True) if capture_output else {}), + ) + except subprocess.CalledProcessError as e: + if capture_output: + print(f"{e.stdout=} {e.stderr=}") + raise + + if capture_output: + print(f"Captured stdout={result.stdout} stderr={result.stderr}") + return result.stdout + + +def get_current_node_ip(): + address = ray._private.services.get_node_ip_address() + # strip ipv6 address + address = address.strip("[]") + return address + + +def get_free_port(start_port=10000, consecutive=1): + # find the port where port, port + 1, port + 2, ... port + consecutive - 1 are all available + port = start_port + while not all(is_port_available(port + i) for i in range(consecutive)): + port += 1 + return port + + +def should_run_periodic_action( + rollout_id: int, + interval: int | None, + num_rollout_per_epoch: int | None = None, + num_rollout: int | None = None, +) -> bool: + """ + Return True when a periodic action (eval/save/checkpoint) should run. + + Args: + rollout_id: The current rollout index (0-based). + interval: Desired cadence; disables checks when None. + num_rollout_per_epoch: Optional epoch boundary to treat as a trigger. + """ + if interval is None: + return False + + if num_rollout is not None and rollout_id == num_rollout - 1: + return True + + step = rollout_id + 1 + return (step % interval == 0) or (num_rollout_per_epoch is not None and step % num_rollout_per_epoch == 0) diff --git a/slime/utils/ppo_utils.py b/slime/utils/ppo_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1ab89240ea3255672e73650d656d8714db71c00f --- /dev/null +++ b/slime/utils/ppo_utils.py @@ -0,0 +1,718 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Adapt from https://github.com/OpenRLHF/OpenRLHF/blob/10c733694ed9fbb78a0a2ff6a05efc7401584d46/openrlhf/models/utils.py +# and https://github.com/OpenRLHF/OpenRLHF/blob/10c733694ed9fbb78a0a2ff6a05efc7401584d46/openrlhf/trainer/ppo_utils/experience_maker.py + +from argparse import Namespace + +import torch +import torch.distributed as dist +import torch.nn.functional as F + + +@torch.compile(dynamic=True) +def compute_approx_kl( + log_probs: torch.Tensor, + log_probs_base: torch.Tensor, + kl_loss_type: str, + importance_ratio: torch.Tensor | None = None, +) -> torch.Tensor: + """ + Compute the approximate KL divergence between two distributions. + Schulman blog: http://joschu.net/blog/kl-approx.html + + Args: + log_probs: Log probabilities of the new distribution. + log_probs_base: Log probabilities of the base distribution. + kl_loss_type: Type of KL estimator (k1, k2, k3, low_var_kl). + importance_ratio: Optional IS ratio (π_θ/π_old) for unbiased KL estimation. + """ + log_ratio = log_probs.float() - log_probs_base.float() + + if kl_loss_type == "k1": + kl = log_ratio + elif kl_loss_type == "k2": + kl = log_ratio**2 / 2.0 + elif kl_loss_type in ["k3", "low_var_kl"]: + # The non negative kl approximation in + # http://joschu.net/blog/kl-approx.html + # Besides non negative, it is also unbiased and have lower variance. + log_ratio = -log_ratio + kl = log_ratio.exp() - 1 - log_ratio + else: + raise ValueError(f"Unknown kl_loss_type: {kl_loss_type}") + + # Apply IS ratio for unbiased KL estimation (DeepSeek-V3.2) + if importance_ratio is not None: + kl = importance_ratio * kl + + # Clamp only for low_var_kl for numerical stability + if kl_loss_type == "low_var_kl": + kl = torch.clamp(kl, min=-10, max=10) + + return kl + + +def compute_opsm_mask( + args: Namespace, + full_log_probs: list[torch.Tensor], + full_old_log_probs: list[torch.Tensor], + advantages: list[torch.Tensor], + loss_masks: list[torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute Off-Policy Sequence Masking (OPSM) mask. + + Args: + args: Configuration containing `opsm_delta` threshold. + full_log_probs: Current policy log-probs per sample. + full_old_log_probs: Old policy log-probs per sample. + advantages: Advantage values per sample. + loss_masks: Loss masks per sample. + + Returns: + Tuple of `(opsm_mask, opsm_clipfrac)` where `opsm_mask` is a + concatenated tensor of per-token masks and + `opsm_clipfrac` is the count of masked sequences. + """ + opsm_mask_list = [] + device = advantages[0].device + opsm_clipfrac = torch.tensor(0.0, device=device) + + for full_log_prob, full_old_log_prob, advantage, loss_mask in zip( + full_log_probs, full_old_log_probs, advantages, loss_masks, strict=False + ): + # Calculate sequence-level KL + seq_kl = ((full_old_log_prob - full_log_prob) * loss_mask).sum() / torch.clamp_min(loss_mask.sum(), 1) + + # Create mask: 0 if (advantage < 0 and seq_kl > delta), else 1 + mask = ((advantage < 0) & (seq_kl > args.opsm_delta)).float() + opsm_clipfrac += mask.sum() / torch.clamp_min(loss_mask.sum(), 1) + + opsm_mask_list.append(1 - mask) + + opsm_mask = torch.cat(opsm_mask_list, dim=0) + return opsm_mask, opsm_clipfrac + + +def compute_gspo_kl( + full_log_probs: list[torch.Tensor], + full_old_log_probs: list[torch.Tensor], + local_log_probs: list[torch.Tensor], + loss_masks: list[torch.Tensor], +) -> torch.Tensor: + """Compute GSPO-style per-sequence KL divergence. + + Args: + full_log_probs: Current policy log-probs per sample (full or CP-local). + full_old_log_probs: Old policy log-probs per sample (full or CP-local). + local_log_probs: Local (CP-local) log-probs for expansion shape reference. + loss_masks: Loss masks per sample. + + Returns: + Concatenated tensor of per-token KL values where each token in a + sequence has the same KL value (the sequence-level KL). + """ + # Compute sequence-level KL and expand to per-token + ppo_kl = [ + ((old_logprob - log_prob) * loss_mask).sum() / torch.clamp_min(loss_mask.sum(), 1) + for log_prob, old_logprob, loss_mask in zip(full_log_probs, full_old_log_probs, loss_masks, strict=False) + ] + ppo_kl = [kl.expand_as(log_prob) for kl, log_prob in zip(ppo_kl, local_log_probs, strict=False)] + ppo_kl = torch.cat(ppo_kl, dim=0) + + return ppo_kl + + +@torch.compile(dynamic=True) +def compute_policy_loss( + ppo_kl: torch.Tensor, + advantages: torch.Tensor, + eps_clip: float, + eps_clip_high: float, + eps_clip_c: float | None = None, +): + ratio = (-ppo_kl).exp() + pg_losses1 = -ratio * advantages + pg_losses2 = -ratio.clamp(1 - eps_clip, 1 + eps_clip_high) * advantages + clip_pg_losses1 = torch.maximum(pg_losses1, pg_losses2) + clipfrac = torch.gt(pg_losses2, pg_losses1).float() + + if eps_clip_c is not None: + assert ( + eps_clip_c > 1.0 + ), f"The lower bound of the clip_ratio_c for dual-clip PPO should be greater than 1.0, but get the value: {eps_clip_c}." + pg_losses3 = -eps_clip_c * advantages + clip_pg_losses2 = torch.min(pg_losses3, clip_pg_losses1) + pg_losses = torch.where(advantages < 0, clip_pg_losses2, clip_pg_losses1) + else: + pg_losses = clip_pg_losses1 + + return pg_losses, clipfrac + + +def compute_log_probs(logits: torch.Tensor, tokens: torch.Tensor, process_group: dist.ProcessGroup | None): + from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy + + # convert to [seq_len, batch_size, vocab_size] as expected by fused_vocab_parallel_cross_entropy + logits = logits.unsqueeze(1) + tokens = tokens.unsqueeze(1) + return -fused_vocab_parallel_cross_entropy(logits, tokens, process_group) + + +# from https://github.com/volcengine/verl/blob/0bdf7f469854815177e73dcfe9e420836c952e6e/verl/utils/megatron/tensor_parallel.py#L99 +class _VocabParallelEntropy(torch.autograd.Function): + + @staticmethod + def forward(ctx, vocab_parallel_logits: torch.Tensor, process_group: dist.ProcessGroup) -> torch.Tensor: + + @torch.compile(dynamic=True) + def mul_reduce(a, b): + return (a * b).sum(dim=-1, keepdim=True) + + logits_max = vocab_parallel_logits.max(dim=-1, keepdim=True).values + dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=process_group) + normalized_vocab_parallel_logits = vocab_parallel_logits - logits_max + normalized_exp_logits = normalized_vocab_parallel_logits.exp_() + normalized_sum_exp_logits = normalized_exp_logits.sum(dim=-1, keepdim=True) + dist.all_reduce(normalized_sum_exp_logits, group=process_group) + softmax_logits = normalized_exp_logits.div_(normalized_sum_exp_logits) + sum_softmax_times_logits = mul_reduce(softmax_logits, vocab_parallel_logits) + dist.all_reduce(sum_softmax_times_logits, group=process_group) + entropy = logits_max + normalized_sum_exp_logits.log() - sum_softmax_times_logits + ctx.save_for_backward(vocab_parallel_logits, softmax_logits, sum_softmax_times_logits) + return entropy.squeeze(dim=-1) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: + vocab_parallel_logits, softmax_logits, sum_softmax_times_logits = ctx.saved_tensors + # reuse softmax_logits as grad + vocab_parallel_logits.sub_(sum_softmax_times_logits) + softmax_logits.mul_(vocab_parallel_logits) + softmax_logits.mul_(grad_output.unsqueeze(dim=-1)) + # recover vocab_parallel_logits + vocab_parallel_logits.add_(sum_softmax_times_logits) + softmax_logits.mul_(-1) + return softmax_logits, None + + +def compute_entropy_from_logits(logits: torch.Tensor, process_group) -> torch.Tensor: + return _VocabParallelEntropy.apply(logits, process_group) + + +def get_grpo_returns( + rewards: torch.Tensor, + kl: list[torch.Tensor], +): + returns = [] + for i in range(len(rewards)): + returns.append(torch.ones_like(kl[i]) * rewards[i]) + return returns + + +def get_reinforce_plus_plus_returns( + rewards: torch.Tensor, + kl: list[torch.Tensor], + loss_masks: list[torch.Tensor], + response_lengths: list[int], + total_lengths: list[int], + kl_coef: float, + gamma: float, +) -> list[torch.Tensor]: + """ + Calculates discounted returns for REINFORCE++ (https://arxiv.org/pdf/2501.03262) + + Args: + rewards (Tensor): A tensor of scalar rewards for each sequence. + kl (List[Tensor]): List of per-token KL divergence tensors for sequence chunks. + loss_masks (List[Tensor]): List of response-only loss masks for each full sequence. + response_lengths (List[int]): The full length of each response sequence. + total_lengths (List[int]): The full length of each sequence (prompt + response). + kl_coef (float): Coefficient for the KL penalty. + gamma (float): The discount factor. + + Returns: + List[torch.Tensor]: A list of return (G_t) tensors for the + local sequence chunks owned by the current GPU rank. + """ + from megatron.core import mpu + + cp_size = mpu.get_context_parallel_world_size() + + final_returns_chunks = [] + for i in range(len(rewards)): + local_kl_chunk = kl[i] + total_len, response_len = total_lengths[i], response_lengths[i] + + if cp_size > 1: + # Step 1,2:Gather all chunks and token_offsets from all ranks and reconstruct the full response tensor by splitting and placing each part + from slime.backends.megatron_utils.cp_utils import all_gather_with_cp + + full_kl_response = all_gather_with_cp(local_kl_chunk, total_len, response_len) + else: + full_kl_response = local_kl_chunk + + # Step 3: Compute returns on full response kl tensor. + token_level_rewards = -kl_coef * full_kl_response + full_mask = loss_masks[i] + assert full_mask.sum().item() > 0, f"Sequence at index {i} is fully masked." + last_idx = full_mask.nonzero(as_tuple=True)[0][-1] + token_level_rewards[last_idx] += rewards[i] + + returns_for_seq = torch.zeros_like(token_level_rewards) + running_return = 0.0 + for t in reversed(range(token_level_rewards.size(0))): + # G_t = r_t + gamma * G_{t+1} + running_return = token_level_rewards[t] + gamma * running_return + returns_for_seq[t] = running_return + + # Step 4: Pick up the results corresponding to our local chunk's parts. + if cp_size > 1: + from slime.backends.megatron_utils.cp_utils import slice_log_prob_with_cp + + local_returns_chunk = slice_log_prob_with_cp(returns_for_seq, total_len, response_len) + else: + local_returns_chunk = returns_for_seq + + final_returns_chunks.append(local_returns_chunk) + + return final_returns_chunks + + +def get_reinforce_plus_plus_baseline_advantages( + rewards: torch.Tensor, + kl: list[torch.Tensor], + loss_masks: list[torch.Tensor], + kl_coef: float, +) -> list[torch.Tensor]: + """ + Calculates the unwhitened advantages for the REINFORCE++-baseline algorithm. + Broadcasting the scalar (reward - group_baseline) to each token. + + Args: + rewards (Tensor): A tensor of scalar rewards, where the group-wise + baseline has already been subtracted. + kl (list[Tensor]): A list of per-token KL divergence tensors. Used to + get the shape for broadcasting. + loss_masks (list[Tensor]): A list of per-token loss masks. + kl_coef (float): Coefficient for the KL penalty. + + Returns: + list[Tensor]: A list of tensors containing the unwhitened advantages. + """ + # Broadcast to get unwhitened advantages + unwhitened_advantages = [ + torch.ones_like(kl_tensor) * reward_val - kl_coef * kl_tensor + for kl_tensor, reward_val in zip(kl, rewards, strict=False) + ] + + return unwhitened_advantages + + +def get_advantages_and_returns( + total_len: int, + response_len: int, + values: torch.Tensor, + rewards: torch.Tensor, + gamma: float, + lambd: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Function that computes advantages and returns from rewards and values. + Calculated as in the original PPO paper: https://arxiv.org/abs/1707.06347 + Note that rewards may include a KL divergence loss term. + + Advantages looks like this: + Adv1 = R1 + γ * λ * R2 + γ^2 * λ^2 * R3 + ... + - V1 + γ * (1 - λ) V2 + γ^2 * λ * (1 - λ) V3 + ... + + Returns looks like this: + Ret1 = R1 + γ * λ * R2 + γ^2 * λ^2 * R3 + ... + + γ * (1 - λ) V2 + γ^2 * λ * (1 - λ) V3 + ... + + Input: + - values: Tensor of shape (response_size,) + - rewards: Tensor of shape (response_size,) + + Output: + - advantages: Tensor of shape (response_size,) + - returns: Tensor of shape (response_size,) + """ + from megatron.core import mpu + + cp_size = mpu.get_context_parallel_world_size() + if cp_size > 1: + from slime.backends.megatron_utils.cp_utils import all_gather_with_cp + + full_rewards = all_gather_with_cp(rewards, total_len, response_len) + full_values = all_gather_with_cp(values, total_len, response_len) + else: + full_rewards = rewards + full_values = values + + lastgaelam = 0 + advantages_reversed = [] + + for t in reversed(range(response_len)): + nextvalues = full_values[t + 1] if t < response_len - 1 else 0.0 + delta = full_rewards[t] + gamma * nextvalues - full_values[t] + lastgaelam = delta + gamma * lambd * lastgaelam + advantages_reversed.append(lastgaelam) + full_advantages = torch.tensor(advantages_reversed[::-1], dtype=full_values.dtype, device=full_values.device) + full_returns = full_advantages + full_values + + if cp_size > 1: + from slime.backends.megatron_utils.cp_utils import slice_log_prob_with_cp + + advantages = slice_log_prob_with_cp(full_advantages, total_len, response_len) + returns = slice_log_prob_with_cp(full_returns, total_len, response_len) + else: + advantages = full_advantages + returns = full_returns + + return advantages.detach(), returns + + +def get_advantages_and_returns_batch( + total_lengths, + response_lengths, + values_list, + rewards_list, + gamma, + lambd, + chunked: bool = True, +): + """ + Batched GAE with CP support. + Input: + total_lengths: list[int], each sample's total_len + response_lengths: list[int], each sample's response_len + values_list: list[Tensor], each shape = [resp_len_i] + rewards_list: list[Tensor], same shape + Output: + advantages_list: list[Tensor], each shape = [resp_len_i] + returns_list: list[Tensor], same shape + """ + + from megatron.core import mpu + + with torch.no_grad(): + B = len(response_lengths) + assert B == len(values_list) + assert B == len(rewards_list) + + cp_size = mpu.get_context_parallel_world_size() + device = values_list[0].device + dtype = values_list[0].dtype + + if cp_size > 1: + from slime.backends.megatron_utils.cp_utils import all_gather_with_cp + + full_values_list = [] + full_rewards_list = [] + + for total_len, resp_len, v, r in zip( + total_lengths, response_lengths, values_list, rewards_list, strict=False + ): + full_v = all_gather_with_cp(v, total_len, resp_len) + full_r = all_gather_with_cp(r, total_len, resp_len) + full_values_list.append(full_v) + full_rewards_list.append(full_r) + + # full_values_list[i].shape = [total_len_i] + else: + full_values_list = values_list + full_rewards_list = rewards_list + + # pad to max_len for batched GAE + max_len = max(response_lengths) + + full_values = torch.zeros(B, max_len, device=device, dtype=dtype) + full_rewards = torch.zeros(B, max_len, device=device, dtype=dtype) + + for i in range(B): + L = response_lengths[i] + full_values[i, :L] = full_values_list[i][:L] + full_rewards[i, :L] = full_rewards_list[i][:L] + + if not chunked: + full_advantages, full_returns = vanilla_gae( + rewards=full_rewards, + values=full_values, + gamma=gamma, + lambd=lambd, + ) + else: + full_advantages, full_returns = chunked_gae( + rewards=full_rewards, + values=full_values, + gamma=gamma, + lambd=lambd, + ) + + advantages_list = [] + returns_list = [] + + if cp_size > 1: + from slime.backends.megatron_utils.cp_utils import slice_log_prob_with_cp + + for total_len, resp_len, adv_row, ret_row in zip( + total_lengths, + response_lengths, + full_advantages, + full_returns, + strict=False, + ): + adv_full = adv_row # shape = [resp_len_i padded to max_len] + ret_full = ret_row + + adv_sliced = slice_log_prob_with_cp(adv_full[:resp_len], total_len, resp_len) + ret_sliced = slice_log_prob_with_cp(ret_full[:resp_len], total_len, resp_len) + + advantages_list.append(adv_sliced) + returns_list.append(ret_sliced) + + else: + for i in range(B): + L = response_lengths[i] + advantages_list.append(full_advantages[i, :L]) + returns_list.append(full_returns[i, :L]) + + return advantages_list, returns_list + + +def vanilla_gae( + rewards: torch.Tensor, + values: torch.Tensor, + gamma: float, + lambd: float, +): + B, T = rewards.shape + device = rewards.device + dtype = rewards.dtype + + lastgaelam = torch.zeros(B, device=device, dtype=dtype) + adv_rev = [] + + for t in reversed(range(T)): + next_value = values[:, t + 1] if t < T - 1 else 0.0 + delta = rewards[:, t] + gamma * next_value - values[:, t] + lastgaelam = delta + gamma * lambd * lastgaelam + adv_rev.append(lastgaelam) + + full_advantages = torch.stack(adv_rev[::-1], dim=1) # [B, max_len] + full_returns = full_advantages + values # [B, max_len] + return full_advantages, full_returns + + +def chunked_gae( + rewards: torch.Tensor, + values: torch.Tensor, + gamma: float, + lambd: float, + chunk_size: int = 128, +): + """ + Compute Generalized Advantage Estimation (GAE) using a FlashLinearAttention- + inspired algorithm: parallel prefix scan within chunks and recurrent state + propagation across chunks. + + This reduces the sequential dependency length from O(T) to O(T / chunk_size), + while keeping chunk computations fully parallelizable (O(C^2) per chunk). + + Args: + rewards (Tensor): [B, T] reward sequence. + values (Tensor): [B, T] value predictions. The next-value of the final + step is assumed to be zero (standard PPO convention). + gamma (float): discount factor. + lam (float): GAE lambda. + chunk_size (int): sequence chunk length for parallel scan. + + Returns: + advantages (Tensor): [B, T] computed advantages. + returns (Tensor): [B, T] advantages + values. + """ + + # ------------------------------------------------------------------------- + # Validate inputs + # ------------------------------------------------------------------------- + assert rewards.ndim == 2 and values.ndim == 2 + B, T = rewards.shape + assert values.shape == (B, T) + + device = rewards.device + dtype = rewards.dtype + + # ------------------------------------------------------------------------- + # Build δ_t = r_t + γ * V_{t+1} - V_t with V_{T} = 0 + # ------------------------------------------------------------------------- + next_values = torch.cat( + [values[:, 1:], torch.zeros(B, 1, device=device, dtype=dtype)], + dim=1, + ) + deltas = rewards + gamma * next_values - values + + # Reformulate backward GAE as a forward scan on the reversed sequence: + # S[i] = Δ[i] + w * S[i - 1], w = γλ + w = gamma * lambd + deltas_rev = torch.flip(deltas, dims=[1]) # [B, T] + + # ------------------------------------------------------------------------- + # Pad to a multiple of chunk_size + # ------------------------------------------------------------------------- + if T % chunk_size != 0: + pad = chunk_size - (T % chunk_size) + deltas_rev = F.pad(deltas_rev, (0, pad)) + else: + pad = 0 + + B, T_pad = deltas_rev.shape + n_chunks = T_pad // chunk_size + + deltas_chunks = deltas_rev.view(B, n_chunks, chunk_size) + + # ------------------------------------------------------------------------- + # Construct the intra-chunk parallel scan kernel M + # + # For a chunk Δ[0..C-1], we want: + # S_local[t] = sum_{k=0..t} w^(t-k) * Δ[k] + # + # This is implemented as: + # S_local = Δ @ M + # + # where: + # M[i, j] = w^(j - i) if j >= i + # 0 otherwise + # ------------------------------------------------------------------------- + idx = torch.arange(chunk_size, device=device) + row = idx[:, None] + col = idx[None, :] + diff = col - row + + M = torch.zeros(chunk_size, chunk_size, device=device, dtype=dtype) + mask = diff >= 0 + + if w == 0.0: + M[mask & (diff == 0)] = 1.0 + else: + M[mask] = w ** diff[mask].to(dtype) + + # pow_vec[t] = w^(t+1), used to inject the recurrent state s_prev + if w == 0.0: + pow_vec = torch.zeros(chunk_size, device=device, dtype=dtype) + else: + pow_vec = w ** torch.arange(1, chunk_size + 1, device=device, dtype=dtype) + + # ------------------------------------------------------------------------- + # Parallel compute local chunk results (assuming initial state = 0) + # ------------------------------------------------------------------------- + deltas_flat = deltas_chunks.reshape(B * n_chunks, chunk_size) + S_local_flat = deltas_flat @ M + S_local_chunks = S_local_flat.view(B, n_chunks, chunk_size) + + # Effective length of each chunk (the last chunk may be padded) + lengths = [chunk_size] * n_chunks + if pad > 0: + lengths[-1] = chunk_size - pad + + # ------------------------------------------------------------------------- + # Recurrent propagation between chunks + # + # Each chunk contributes: + # S_global[t] = S_local[t] + w^(t+1) * s_prev + # + # And updates: + # s_prev = S_global[last_t] + # ------------------------------------------------------------------------- + S_rev = deltas_rev.new_zeros(B, T_pad) + s_prev = torch.zeros(B, device=device, dtype=dtype) + + for c in range(n_chunks): + Lc = lengths[c] + start = c * chunk_size + end = start + Lc + + S_local = S_local_chunks[:, c, :Lc] + S_global = S_local + s_prev.unsqueeze(1) * pow_vec[:Lc] + + S_rev[:, start:end] = S_global + s_prev = S_global[:, -1] # state for next chunk + + # Remove padding and flip back to original time order + if pad > 0: + S_rev = S_rev[:, :T] + + advantages = torch.flip(S_rev, dims=[1]) + returns = advantages + values + + return advantages, returns + + +def calculate_log_probs_and_entropy(logits, tokens, tp_group, with_entropy: bool = False, chunk_size: int = -1): + logits = logits.contiguous() + # TODO: not sure why we need to clone the logits here. + # Without the clone, the backward will trigger inplace edit error. + # It seems that the function with tp will modify the logits inplace. + entropy = None + if logits.size(0) != 0: + if chunk_size > 0: + num_chunks = (logits.size(0) - 1) // chunk_size + 1 + tokens_chunks = tokens.chunk(num_chunks, dim=0) + logits_chunks = logits.chunk(num_chunks, dim=0) + log_probs = [] + for tokens_chunk, logits_chunk in zip(tokens_chunks, logits_chunks, strict=True): + log_prob = compute_log_probs(logits_chunk.clone(), tokens_chunk, tp_group) + log_probs.append(log_prob) + log_prob = torch.cat(log_probs, dim=0) + if with_entropy: + entropys = [] + for _, logits_chunk in zip(tokens_chunks, logits_chunks, strict=True): + entropy = compute_entropy_from_logits(logits_chunk.clone(), tp_group) + entropys.append(entropy) + entropy = torch.cat(entropys, dim=0) + else: + log_prob = compute_log_probs(logits.clone(), tokens, tp_group) + if with_entropy: + entropy = compute_entropy_from_logits(logits.clone(), tp_group) + else: + log_prob = logits.new_zeros((0,)) + if with_entropy: + entropy = logits.new_zeros((0,)) + + return log_prob, entropy + + +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, +) -> tuple[torch.Tensor, list[torch.Tensor], dict[str, torch.Tensor]]: + """Apply TIS off-policy correction using importance sampling. + + Parameters: + args: Arguments containing TIS settings. + pg_loss: Policy gradient loss tensor of shape [total_seq_len - 1]. + train_log_probs: List of tensors containing training log-probabilities + for each sequence. + rollout_log_probs: List of tensors containing rollout log-probabilities + for each sequence. + loss_masks: List of tensors containing loss masks for each sequence. + """ + 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 = (tis - 1).abs() + tis_clip_low = args.tis_clip_low if args.tis_clip_low is not None else 0.1 + tis_clip_high = args.tis_clip if args.tis_clip is not None else 2.0 + tis_weights = torch.clamp(tis, min=tis_clip_low, max=tis_clip_high) + 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 diff --git a/slime/utils/processing_utils.py b/slime/utils/processing_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c3709e6d565346ddfb4f31b5f06681acc2e94f60 --- /dev/null +++ b/slime/utils/processing_utils.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import base64 +import io +import logging + +from transformers import AutoProcessor, AutoTokenizer, PreTrainedTokenizerBase, ProcessorMixin + +logger = logging.getLogger(__name__) + + +def load_tokenizer(name_or_path: str, **kwargs): + return AutoTokenizer.from_pretrained(name_or_path, **kwargs) + + +def load_processor(name_or_path: str, **kwargs): + try: + proc = AutoProcessor.from_pretrained(name_or_path, **kwargs) + except (OSError, ValueError) as e: + logger.warning(f"Failed to load processor from {name_or_path}: {e}") + proc = None + + # If HF returned a tokenizer, discard it. + if isinstance(proc, PreTrainedTokenizerBase) or not isinstance(proc, ProcessorMixin): + proc = None + + return proc + + +def encode_image_for_rollout_engine(image) -> str: + """Load an image from path, ensure RGB, encode as PNG base64 string.""" + buffer = io.BytesIO() + if image.mode != "RGB": + image = image.convert("RGB") + image.save(buffer, format="PNG") + return base64.b64encode(buffer.getvalue()).decode("utf-8") diff --git a/slime/utils/profile_utils.py b/slime/utils/profile_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0975f955cbbe88b4d577f2f02bb65b02d2b03e1c --- /dev/null +++ b/slime/utils/profile_utils.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +import time +import traceback +from pathlib import Path + +import torch + +from slime.utils.memory_utils import print_memory + +logger = logging.getLogger(__name__) + + +class TrainProfiler: + def __init__(self, args): + self.args = args + self._torch_profiler_overall = None + self._memory_profiler_overall = None + + if args.use_pytorch_profiler and ("train_overall" in args.profile_target): + self._torch_profiler_overall = _create_torch_profiler(args, name="train_overall") + + if args.record_memory_history and ("train_overall" in args.profile_target): + self._memory_profiler_overall = _BaseMemoryProfiler.create(args) + self._memory_profiler_overall.start() + + def on_init_end(self): + if self._torch_profiler_overall is not None: + self._torch_profiler_overall.start() + + def step(self, rollout_id: int): + if self._torch_profiler_overall is not None: + self._torch_profiler_overall.step() + + if ( + self._memory_profiler_overall is not None + and ((s := self.args.memory_snapshot_num_steps) is not None) + and (rollout_id == s - 1) + ): + self._memory_profiler_overall.stop() + + def iterate_train_actor(self, iterator): + return _profile_simple_loop(iterator, self.args, name="train_actor") + + def iterate_train_log_probs(self, iterator): + return _profile_simple_loop(iterator, self.args, name="train_log_probs") + + +def _profile_simple_loop(iterator, args, name): + if not (args.use_pytorch_profiler and (name in args.profile_target)): + yield from iterator + return + + torch_profiler = _create_torch_profiler(args, name=name) + torch_profiler.start() + for item in iterator: + yield item + torch_profiler.step() + + +def _create_torch_profiler(args, name): + return torch.profiler.profile( + schedule=torch.profiler.schedule( + # TODO the train_actor and train_log_probs ones may need to have different args to control step + wait=max(args.profile_step_start - 1, 0), + warmup=1 if args.profile_step_start > 0 else 0, + active=args.profile_step_end - args.profile_step_start, + repeat=1, + ), + on_trace_ready=torch.profiler.tensorboard_trace_handler( + args.tensorboard_dir, + worker_name=f"{name}_rank_{torch.distributed.get_rank()}", + use_gzip=True, + ), + record_shapes=True, + with_stack=True, + profile_memory=True, + with_flops=True, + ) + + +class _BaseMemoryProfiler: + @staticmethod + def create(args): + c = { + "torch": _TorchMemoryProfiler, + "memray": _MemrayMemoryProfiler, + }[args.memory_recorder] + return c(args) + + def __init__(self, args): + self._path_dump = ( + Path(args.memory_snapshot_dir) + / f"memory_snapshot_time{time.time()}_rank{torch.distributed.get_rank()}_{args.memory_snapshot_path}" + ) + + def start(self): + raise NotImplementedError + + def stop(self): + raise NotImplementedError + + +class _TorchMemoryProfiler(_BaseMemoryProfiler): + def start(self): + logger.info("Attach OOM dump memory history.") + + torch.cuda.memory._record_memory_history( + max_entries=1000000, + # record stack information for the trace events + # trace_alloc_record_context=True, + stacks="all", + ) + + def oom_observer(device, alloc, device_alloc, device_free): + logger.info( + f"Observe OOM, will dump snapshot to {self._path_dump}. ({device=} {alloc=} {device_alloc=} {device_free=}; stacktrace is as follows)" + ) + traceback.print_stack() + torch.cuda.memory._dump_snapshot(self._path_dump) + print_memory("when oom") + + torch._C._cuda_attach_out_of_memory_observer(oom_observer) + + def stop(self): + logger.info(f"Dump memory snapshot to: {self._path_dump}") + torch.cuda.memory._dump_snapshot(self._path_dump) + torch.cuda.memory._record_memory_history(enabled=None) + + +class _MemrayMemoryProfiler(_BaseMemoryProfiler): + def __init__(self, args): + super().__init__(args) + assert args.memory_snapshot_num_steps is not None, "In memray, must provide --memory-snapshot-num-steps" + + def start(self): + logger.info("Memray tracker started.") + import memray + + self._tracker = memray.Tracker( + file_name=self._path_dump, + native_traces=True, + ) + self._tracker.__enter__() + + def stop(self): + logger.info(f"Memray tracker stopped and dump snapshot to: {self._path_dump}") + self._tracker.__exit__(None, None, None) diff --git a/slime/utils/ray_utils.py b/slime/utils/ray_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..45d31e86f85375fe629188b2cce4d72f476d81d0 --- /dev/null +++ b/slime/utils/ray_utils.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +class Box: + def __init__(self, inner): + self._inner = inner + + @property + def inner(self): + return self._inner diff --git a/slime/utils/reloadable_process_group.py b/slime/utils/reloadable_process_group.py new file mode 100644 index 0000000000000000000000000000000000000000..aa92806044d973eb002b87a24e2feda5e9d8404d --- /dev/null +++ b/slime/utils/reloadable_process_group.py @@ -0,0 +1,285 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +import os +from contextlib import contextmanager + +import torch +import torch.distributed as dist + +from slime.utils.memory_utils import print_memory + +logger = logging.getLogger(__name__) + +old_new_group_dict = {} + + +def monkey_patch_torch_dist(): + pid = os.getpid() + if pid in old_new_group_dict: + assert dist.old_new_group == old_new_group_dict[pid] + return + + logger.info("Applying monkey patch to torch.distributed") + + old_new_group = dist.new_group + old_new_group_dict[pid] = old_new_group + dist.old_new_group = old_new_group + + def new_group(*args, **kwargs): + group = old_new_group(*args, **kwargs) + # skip none nccl group. + if len(args) >= 3 and args[2] == "gloo" or "backend" in kwargs and kwargs["backend"] == "gloo": + return group + + # Get ranks from arguments + if len(args) >= 1 and args[0] is not None: + ranks = args[0] + elif "ranks" in kwargs and kwargs["ranks"] is not None: + ranks = kwargs["ranks"] + else: + # If no ranks specified, use all ranks in world + ranks = list(range(dist.get_world_size())) + + if len(ranks) == 1: + return group + + group = ReloadableProcessGroup(group, ranks) + return group + + dist.new_group = new_group + + def get_new_function(func): + def new_function(*args, **kwargs): + args = tuple([arg.group if isinstance(arg, ReloadableProcessGroup) else arg for arg in args]) + kwargs = {k: (v.group if isinstance(v, ReloadableProcessGroup) else v) for k, v in kwargs.items()} + with _wrap_low_level_call(): + return func(*args, **kwargs) + + return new_function + + dist.get_rank = get_new_function(dist.get_rank) + dist.get_world_size = get_new_function(dist.get_world_size) + dist.get_backend = get_new_function(dist.get_backend) + dist.get_global_rank = get_new_function(dist.get_global_rank) + dist.get_group_rank = get_new_function(dist.get_group_rank) + dist.get_process_group_ranks = get_new_function(dist.get_process_group_ranks) + + dist.all_reduce = get_new_function(dist.all_reduce) + dist.all_gather = get_new_function(dist.all_gather) + dist.all_gather_into_tensor = get_new_function(dist.all_gather_into_tensor) + dist.all_gather_object = get_new_function(dist.all_gather_object) + dist.all_to_all = get_new_function(dist.all_to_all) + dist.all_to_all_single = get_new_function(dist.all_to_all_single) + dist.broadcast = get_new_function(dist.broadcast) + dist.reduce = get_new_function(dist.reduce) + dist.reduce_scatter = get_new_function(dist.reduce_scatter) + dist.reduce_scatter_tensor = get_new_function(dist.reduce_scatter_tensor) + dist.scatter = get_new_function(dist.scatter) + dist.gather = get_new_function(dist.gather) + dist.barrier = get_new_function(dist.barrier) + dist.send = get_new_function(dist.send) + dist.recv = get_new_function(dist.recv) + dist._coalescing_manager = get_new_function(dist._coalescing_manager) + + # p2p + old_isend = dist.isend + old_irecv = dist.irecv + + dist.isend = get_new_function(dist.isend) + dist.irecv = get_new_function(dist.irecv) + + def get_new_p2pop_function(func): + def new_function(*args, **kwargs): + def convert(arg): + if isinstance(arg, ReloadableProcessGroup): + return arg.group + elif arg == dist.isend: + arg = old_isend + elif arg == dist.irecv: + arg = old_irecv + return arg + + args = (convert(arg) for arg in args) + kwargs = {k: convert(v) for k, v in kwargs.items()} + return func(*args, **kwargs) + + return new_function + + dist.P2POp.__new__ = get_new_p2pop_function(dist.P2POp.__new__) + dist.P2POp.__init__ = get_new_p2pop_function(dist.P2POp.__init__) + + +class ReloadableProcessGroup(torch.distributed.ProcessGroup): + GROUPS = {} + + def __init__(self, group, ranks): + super().__init__( + rank=dist.get_rank(group), + size=dist.get_world_size(group), + ) + self.group = group + self.group_info = { + "ranks": ranks, + } + pid = os.getpid() + if pid not in ReloadableProcessGroup.GROUPS: + ReloadableProcessGroup.GROUPS[pid] = [] + ReloadableProcessGroup.GROUPS[pid].append(self) + + def __getattr__(self, name): + return getattr(self.group, name) + + @staticmethod + def destroy_process_groups(): + pid = os.getpid() + for reloadable_group in ReloadableProcessGroup.GROUPS.get(pid, []): + if reloadable_group.group is None: + continue + try: + dist.destroy_process_group(reloadable_group.group) + except ValueError as e: + logger.warning( + f"Process group already invalid/destroyed; skipping cleanup. Exception: {e}", + exc_info=True, + ) + + del reloadable_group.group + reloadable_group.group = None + + @staticmethod + def reload_process_groups(): + pid = os.getpid() + reloadable_groups = ReloadableProcessGroup.GROUPS.get(pid, []) + logger.info(f"Reloading {len(reloadable_groups)} process groups in pid {pid}") + old_new_group = old_new_group_dict.get(pid) + for reloadable_group in reloadable_groups: + if reloadable_group.group is not None: + continue + group = old_new_group(ranks=reloadable_group.group_info["ranks"], backend="nccl") + reloadable_group.group = group + + def rank(self) -> int: + return self.group.rank() + + def size(self) -> int: + return self.group.size() + + def name(self) -> str: + return self.group.name() + + def shutdown(self) -> None: + if self.group is not None: + self.group.shutdown() + + def abort(self) -> None: + if self.group is not None: + self.group.abort() + + def _fwd(self, method, *args, **kwargs): + inner = self.group + if inner is None: + raise RuntimeError("ReloadableProcessGroup: inner PG is None, call reload() first.") + with _wrap_low_level_call(): + return getattr(inner, method)(*args, **kwargs) + + def barrier(self, *a, **kw): + return self._fwd("barrier", *a, **kw) + + def broadcast(self, *a, **kw): + return self._fwd("broadcast", *a, **kw) + + def allreduce(self, *a, **kw): + return self._fwd("allreduce", *a, **kw) + + def allreduce_coalesced(self, *a, **kw): + return self._fwd("allreduce_coalesced", *a, **kw) + + def reduce(self, *a, **kw): + return self._fwd("reduce", *a, **kw) + + def allgather(self, *a, **kw): + return self._fwd("allgather", *a, **kw) + + def _allgather_base(self, *a, **kw): + return self._fwd("_allgather_base", *a, **kw) + + def allgather_coalesced(self, *a, **kw): + return self._fwd("allgather_coalesced", *a, **kw) + + def allgather_into_tensor_coalesced(self, *a, **kw): + return self._fwd("allgather_into_tensor_coalesced", *a, **kw) + + def gather(self, *a, **kw): + return self._fwd("gather", *a, **kw) + + def scatter(self, *a, **kw): + return self._fwd("scatter", *a, **kw) + + def reduce_scatter(self, *a, **kw): + return self._fwd("reduce_scatter", *a, **kw) + + def _reduce_scatter_base(self, *a, **kw): + return self._fwd("_reduce_scatter_base", *a, **kw) + + def reduce_scatter_tensor_coalesced(self, *a, **kw): + return self._fwd("reduce_scatter_tensor_coalesced", *a, **kw) + + def alltoall_base(self, *a, **kw): + return self._fwd("alltoall_base", *a, **kw) + + def alltoall(self, *a, **kw): + return self._fwd("alltoall", *a, **kw) + + def send(self, *a, **kw): + return self._fwd("send", *a, **kw) + + def recv(self, *a, **kw): + return self._fwd("recv", *a, **kw) + + def recv_anysource(self, *a, **kw): + return self._fwd("recv_anysource", *a, **kw) + + def _start_coalescing(self, *a, **kw): + return self._fwd("_start_coalescing", *a, **kw) + + def _end_coalescing(self, *a, **kw): + return self._fwd("_end_coalescing", *a, **kw) + + def _get_backend_name(self): + return self._fwd("_get_backend_name") + + def _get_backend(self, *a, **kw): + return self._fwd("_get_backend", *a, **kw) + + def _set_default_backend(self, *a, **kw): + return self._fwd("_set_default_backend", *a, **kw) + + @property + def bound_device_id(self): + return self.group.bound_device_id + + @bound_device_id.setter + def bound_device_id(self, dev): + self.group.bound_device_id = dev + + +def destroy_process_groups(): + """Destroy all reloadable process groups.""" + ReloadableProcessGroup.destroy_process_groups() + + +def reload_process_groups(): + """Reload all reloadable process groups.""" + ReloadableProcessGroup.reload_process_groups() + + +@contextmanager +def _wrap_low_level_call(): + try: + yield + except Exception as e: + mem_info = print_memory("after torch distributed error") + e.add_note(f"{mem_info=}") + raise diff --git a/slime/utils/rocm_checkpoint_writer.py b/slime/utils/rocm_checkpoint_writer.py new file mode 100644 index 0000000000000000000000000000000000000000..c3b7a7efe817bf6fc7f208618c5762631cf1cbdb --- /dev/null +++ b/slime/utils/rocm_checkpoint_writer.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from megatron.core.dist_checkpointing.strategies.filesystem_async import FileSystemWriterAsync + + +class ROCmFileSystemWriterAsync(FileSystemWriterAsync): + """ + FileSystemWriterAsync wrapper for ROCm compatibility. + + On ROCm/HIP, using non_blocking=True causes tensors to be stored in pinned memory, + which triggers segmentation faults when forking subprocesses afterward. + """ + + @staticmethod + def preload_tensors(*args, **kwargs): + # Change argument non_blocking to False on HIP platform + # The tensors will be stored in pinned memory if non_blocking=True + # Currently on the ROCm platform, forking a subprocess afterward + # with pinned_memory=True will trigger segmentation fault + if torch.version.hip: + print("HIP/ROCm detected: setting non_blocking=False in preload_tensors") + if "non_blocking" in kwargs: + kwargs["non_blocking"] = False + elif len(args) > 1 and isinstance(args[-1], bool): + # non_blocking is typically the last argument + args = args[:-1] + (False,) + + return FileSystemWriterAsync.preload_tensors(*args, **kwargs) diff --git a/slime/utils/routing_replay.py b/slime/utils/routing_replay.py new file mode 100644 index 0000000000000000000000000000000000000000..7d17bbc1b6c5e3f47e893c810056f55b48b13298 --- /dev/null +++ b/slime/utils/routing_replay.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +import torch + + +ROUTING_REPLAY = None + + +def set_routing_replay(replay): + global ROUTING_REPLAY + ROUTING_REPLAY = replay + + +class RoutingReplay: + all_routing_replays = [] + + def __init__(self): + self.forward_index = 0 + self.backward_index = 0 + self.top_indices_list = [] + RoutingReplay.all_routing_replays.append(self) + + def record(self, top_indices): + # offload top_indices to CPU pinned memory + buf = torch.empty_like(top_indices, device="cpu", pin_memory=True) + buf.copy_(top_indices) + self.top_indices_list.append(buf) + + def pop_forward(self): + top_indices = self.top_indices_list[self.forward_index] + self.forward_index += 1 + return top_indices.to(torch.cuda.current_device()) + + def pop_backward(self): + top_indices = self.top_indices_list[self.backward_index] + self.backward_index += 1 + return top_indices.to(torch.cuda.current_device()) + + def clear(self): + self.forward_index = 0 + self.backward_index = 0 + self.top_indices_list = [] + + def clear_forward(self): + self.forward_index = 0 + + @staticmethod + def clear_all(): + for replay in RoutingReplay.all_routing_replays: + replay.clear() + + @staticmethod + def clear_all_forward(): + for replay in RoutingReplay.all_routing_replays: + replay.clear_forward() + + +def get_routing_replay_compute_topk(old_compute_topk): + def compute_topk(scores, topk, num_groups=None, group_topk=None): + if os.environ.get("ENABLE_ROUTING_REPLAY", "0") == "1": + routing_replay_stage = os.environ["ROUTING_REPLAY_STAGE"] + if routing_replay_stage == "fallthrough": + return old_compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk) + if routing_replay_stage == "record": + probs, top_indices = old_compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk) + ROUTING_REPLAY.record(top_indices) + elif routing_replay_stage == "replay_forward": + top_indices = ROUTING_REPLAY.pop_forward() + assert ( + top_indices.shape[0] == scores.shape[0] and top_indices.shape[1] == topk + ), f"[{torch.distributed.get_rank()}] top_indices shape {top_indices.shape} does not match scores shape {scores.shape} and topk {topk}" + probs = scores.gather(1, top_indices) + elif routing_replay_stage == "replay_backward": + top_indices = ROUTING_REPLAY.pop_backward() + assert ( + top_indices.shape[0] == scores.shape[0] and top_indices.shape[1] == topk + ), f"top_indices shape {top_indices.shape} does not match scores shape {scores.shape} and topk {topk}" + probs = scores.gather(1, top_indices) + return probs, top_indices + else: + return old_compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk) + + return compute_topk + + +def register_routing_replay(module): + if os.environ.get("ENABLE_ROUTING_REPLAY", "0") == "1": + module.routing_replay = RoutingReplay() + + def pre_forward_hook(*args, **kwargs): + set_routing_replay(module.routing_replay) + + module.register_forward_pre_hook(pre_forward_hook) diff --git a/slime/utils/seqlen_balancing.py b/slime/utils/seqlen_balancing.py new file mode 100644 index 0000000000000000000000000000000000000000..5bee97c6eada67eb63359953b629d838f6629992 --- /dev/null +++ b/slime/utils/seqlen_balancing.py @@ -0,0 +1,186 @@ +# Copied from https://github.com/volcengine/verl/blob/468adf22c43b744348051fccd7a5d830c6c3c36a/verl/utils/seqlen_balancing.py +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import heapq + + +def karmarkar_karp(seqlen_list: list[int], k_partitions: int, equal_size: bool): + # see: https://en.wikipedia.org/wiki/Largest_differencing_method + class Set: + + def __init__(self) -> None: + self.sum = 0 + self.items = [] + + def add(self, idx: int, val: int): + self.items.append((idx, val)) + self.sum += val + + def merge(self, other): + for idx, val in other.items: + self.items.append((idx, val)) + self.sum += val + + def __lt__(self, other): + if self.sum != other.sum: + return self.sum < other.sum + if len(self.items) != len(other.items): + return len(self.items) < len(other.items) + return self.items < other.items + + class State: + + def __init__(self, items: list[tuple[int, int]], k: int) -> None: + self.k = k + # sets should always be decreasing order + self.sets = [Set() for _ in range(k)] + assert len(items) in [1, k], f"{len(items)} not in [1, {k}]" + for i, (idx, seqlen) in enumerate(items): + self.sets[i].add(idx=idx, val=seqlen) + self.sets = sorted(self.sets, reverse=True) + + def get_partitions(self): + partitions = [] + for i in range(len(self.sets)): + cur_partition = [] + for idx, _ in self.sets[i].items: + cur_partition.append(idx) + partitions.append(cur_partition) + return partitions + + def merge(self, other): + for i in range(self.k): + self.sets[i].merge(other.sets[self.k - 1 - i]) + self.sets = sorted(self.sets, reverse=True) + + @property + def spread(self) -> int: + return self.sets[0].sum - self.sets[-1].sum + + def __lt__(self, other): + # least heap, let the state with largest spread to be popped first, + # if the spread is the same, let the state who has the largest set + # to be popped first. + if self.spread != other.spread: + return self.spread > other.spread + return self.sets[0] > other.sets[0] + + def __repr__(self) -> str: + repr_str = "[" + for i in range(self.k): + if i > 0: + repr_str += "," + repr_str += "{" + for j, (_, seqlen) in enumerate(self.sets[i].items): + if j > 0: + repr_str += "," + repr_str += str(seqlen) + repr_str += "}" + repr_str += "]" + return repr_str + + sorted_seqlen_list = sorted([(seqlen, i) for i, seqlen in enumerate(seqlen_list)]) + states_pq = [] + if equal_size: + assert len(seqlen_list) % k_partitions == 0, f"{len(seqlen_list)} % {k_partitions} != 0" + for offset in range(0, len(sorted_seqlen_list), k_partitions): + items = [] + for i in range(k_partitions): + seqlen, idx = sorted_seqlen_list[offset + i] + items.append((idx, seqlen)) + heapq.heappush(states_pq, State(items=items, k=k_partitions)) + else: + for seqlen, idx in sorted_seqlen_list: + heapq.heappush(states_pq, State(items=[(idx, seqlen)], k=k_partitions)) + + while len(states_pq) > 1: + state0 = heapq.heappop(states_pq) + state1 = heapq.heappop(states_pq) + # merge states + state0.merge(state1) + heapq.heappush(states_pq, state0) + + final_state = states_pq[0] + partitions = final_state.get_partitions() + if equal_size: + for _i, partition in enumerate(partitions): + assert len(partition) * k_partitions == len( + seqlen_list + ), f"{len(partition)} * {k_partitions} != {len(seqlen_list)}" + return partitions + + +def greedy_partition(seqlen_list: list[int], k_partitions: int, equal_size: bool): + bias = sum(seqlen_list) + 1 if equal_size else 0 + sorted_seqlen = [(seqlen + bias, i) for i, seqlen in enumerate(seqlen_list)] + partitions = [[] for _ in range(k_partitions)] + partition_sums = [0 for _ in range(k_partitions)] + for seqlen, i in sorted_seqlen: + min_idx = None + for j in range(k_partitions): + if min_idx is None or partition_sums[j] < partition_sums[min_idx]: + min_idx = j + partitions[min_idx].append(i) + partition_sums[min_idx] += seqlen + if equal_size: + for _i, partition in enumerate(partitions): + assert len(partition) * k_partitions == len( + seqlen_list + ), f"{len(partition)} * {k_partitions} != {len(seqlen_list)}" + return partitions + + +def get_seqlen_balanced_partitions(seqlen_list: list[int], k_partitions: int, equal_size: bool): + """get order of seq lengths to make partitions balanced, this is + used in balacing sum of seqlength across dp ranks and microbatches + Parameters: + seqlen_list (List[int]): + seq lengths of each items + k_partitions (int): + resulting number of partitions + equal_size (bool): + if True, number of items in each partitions must be equal. + if False, only consider balancing the sum, each partition can have + variable number of items + Returns: + partitions (List[List[int]]): + return k_partitions list containing the index of items. + """ + assert len(seqlen_list) >= k_partitions, f"number of items:[{len(seqlen_list)}] < k_partitions:[{k_partitions}]" + + def _check_and_sort_partitions(partitions): + assert len(partitions) == k_partitions, f"{len(partitions)} != {k_partitions}" + seen_idx = set() + sorted_partitions = [None] * k_partitions + for _i, partition in enumerate(partitions): + assert len(partition) > 0, f"the {_i}-th partition is empty" + for idx in partition: + seen_idx.add(idx) + sorted_partitions[_i] = sorted(partition) + assert seen_idx == set(range(len(seqlen_list))) + return sorted_partitions + + partitions = karmarkar_karp(seqlen_list=seqlen_list, k_partitions=k_partitions, equal_size=equal_size) + return _check_and_sort_partitions(partitions) + + +def get_reverse_idx(idx_map): + reverse_idx_map = copy.deepcopy(idx_map) + + for i, idx in enumerate(idx_map): + reverse_idx_map[idx] = i + + return reverse_idx_map diff --git a/slime/utils/tensor_backper.py b/slime/utils/tensor_backper.py new file mode 100644 index 0000000000000000000000000000000000000000..a5db68e6ef33704f01edc36d89a6996a56a59611 --- /dev/null +++ b/slime/utils/tensor_backper.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from abc import ABC, abstractmethod +from collections import defaultdict +from collections.abc import Callable, Iterable + +import torch + +_SourceGetter = Callable[[], Iterable[tuple[str, torch.Tensor]]] + + +class TensorBackuper(ABC): + @staticmethod + def create(source_getter, single_tag): + if single_tag is None: + return _TensorBackuperNormal(source_getter=source_getter) + else: + return _TensorBackuperNoop(source_getter=source_getter, single_tag=single_tag) + + def __init__(self, source_getter: _SourceGetter): + self._source_getter = source_getter + + @property + @abstractmethod + def backup_tags(self): + raise NotImplementedError + + @abstractmethod + def get(self, tag: str): + raise NotImplementedError + + @abstractmethod + def backup(self, tag: str): + raise NotImplementedError + + def copy(self, *, src_tag: str, dst_tag: str): + raise NotImplementedError + + @abstractmethod + def restore(self, tag: str): + raise NotImplementedError + + +class _TensorBackuperNormal(TensorBackuper): + def __init__(self, source_getter): + super().__init__(source_getter=source_getter) + self._backups: dict[str, dict[str, torch.Tensor]] = defaultdict(dict) + + @property + def backup_tags(self): + return list(self._backups) + + def get(self, tag: str): + return self._backups[tag] + + @torch.no_grad() + def backup(self, tag: str) -> None: + backup_dict = self._backups[tag] + for name, param in self._source_getter(): + if name not in backup_dict: + backup_dict[name] = torch.empty_like(param, device=torch.device("cpu"), pin_memory=True) + backup_dict[name].copy_(param.detach(), non_blocking=True) + torch.cuda.synchronize() + + @torch.no_grad() + def copy(self, *, src_tag: str, dst_tag: str): + for name in self._backups[dst_tag]: + self._backups[dst_tag][name].copy_(self._backups[src_tag][name]) + + @torch.no_grad() + def restore(self, tag: str) -> None: + backup_dict = self._backups[tag] + for name, param in self._source_getter(): + assert name in backup_dict + param.copy_(backup_dict[name], non_blocking=True) + torch.cuda.synchronize() + + +class _TensorBackuperNoop(TensorBackuper): + def __init__(self, source_getter, single_tag): + super().__init__(source_getter=source_getter) + self._single_tag = single_tag + # Sanity check for safety + self._backup_hash_dict = None + + @property + def backup_tags(self): + return [self._single_tag] + + def get(self, tag: str): + ans = dict(self._source_getter()) + ans = {k: v.detach() for k, v in ans.items()} + assert _compute_hash_dict(ans) == self._backup_hash_dict + return ans + + def backup(self, tag: str) -> None: + assert tag == self._single_tag + self._backup_hash_dict = _compute_hash_dict(dict(self._source_getter())) + torch.cuda.synchronize() + + def restore(self, tag: str) -> None: + assert tag == self._single_tag + assert _compute_hash_dict(dict(self._source_getter())) == self._backup_hash_dict + torch.cuda.synchronize() + + +def _compute_hash_dict(tensors: dict[str, torch.Tensor]): + return {k: _compute_hash_tensor(v) for k, v in tensors.items()} + + +def _compute_hash_tensor(x: torch.Tensor): + # Not a real/good hash, but pretty fast + x = x.contiguous() + x = x.view(-1) + x = x.view(torch.uint32) + x = x.sum() + return x.item() diff --git a/slime/utils/tensorboard_utils.py b/slime/utils/tensorboard_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4a03519e7f7ad25729fcdfdf44d39f1ac1524b7d --- /dev/null +++ b/slime/utils/tensorboard_utils.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import datetime +import logging +import os +from slime.utils.misc import SingletonMeta + +try: + from torch.utils.tensorboard import SummaryWriter +except ImportError: + SummaryWriter = None + +logger = logging.getLogger(__name__) + + +class _TensorboardAdapter(metaclass=SingletonMeta): + _writer = None + + """ + # Usage example: This will return the same instance every rank + # tb = _TensorboardAdapter(args) # Initialize on first call + # tb.log({"Loss": 0.1}, step=1) + + # In other files: + # from tensorboard_utils import _TensorboardAdapter + # tb = _TensorboardAdapter(args) # No parameters needed to get existing instance + # tb.log({"Accuracy": 0.9}, step=1) + """ + + def __init__(self, args): + assert args.use_tensorboard, f"{args.use_tensorboard=}" + tb_project_name = args.tb_project_name + tb_experiment_name = args.tb_experiment_name + if tb_project_name is not None or os.environ.get("TENSORBOARD_DIR", None): + if tb_project_name is not None and tb_experiment_name is None: + tb_experiment_name = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + self._initialize(tb_project_name, tb_experiment_name) + else: + raise ValueError("tb_project_name and tb_experiment_name, or TENSORBOARD_DIR are required") + + def _initialize(self, tb_project_name, tb_experiment_name): + """Actual initialization logic""" + # Get tensorboard directory from environment variable or use default path + tensorboard_dir = os.environ.get("TENSORBOARD_DIR", f"tensorboard_log/{tb_project_name}/{tb_experiment_name}") + os.makedirs(tensorboard_dir, exist_ok=True) + logger.info(f"Saving tensorboard log to {tensorboard_dir}.") + self._writer = SummaryWriter(tensorboard_dir) + + def log(self, data, step): + """Log data to tensorboard + + Args: + data (dict): Dictionary containing metric names and values + step (int): Current step/epoch number + """ + for key in data: + self._writer.add_scalar(key, data[key], step) + + def finish(self): + """Close the tensorboard writer""" + self._writer.close() diff --git a/slime/utils/timer.py b/slime/utils/timer.py new file mode 100644 index 0000000000000000000000000000000000000000..6031b37878e493130bf4922ac32611c444b6323b --- /dev/null +++ b/slime/utils/timer.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +from contextlib import contextmanager +from functools import wraps +from time import time + +import torch.distributed + +from .misc import SingletonMeta + +__all__ = ["Timer", "timer"] + +logger = logging.getLogger(__name__) + + +class Timer(metaclass=SingletonMeta): + def __init__(self): + self.timers = {} + self.start_time = {} + + def start(self, name): + assert name not in self.start_time, f"Timer {name} already started." + self.start_time[name] = time() + if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0: + logger.info(f"Timer {name} start") + + def end(self, name): + assert name in self.start_time, f"Timer {name} not started." + elapsed_time = time() - self.start_time[name] + self.add(name, elapsed_time) + del self.start_time[name] + if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0: + logger.info(f"Timer {name} end (elapsed: {elapsed_time:.1f}s)") + + def reset(self, name=None): + if name is None: + self.timers = {} + elif name in self.timers: + del self.timers[name] + + def add(self, name, elapsed_time): + self.timers[name] = self.timers.get(name, 0) + elapsed_time + + def log_dict(self): + return self.timers + + @contextmanager + def context(self, name): + self.start(name) + try: + yield + finally: + self.end(name) + + +def timer(name_or_func): + """ + Can be used either as a decorator or a context manager: + + @timer + def func(): + ... + + or + + with timer("block_name"): + ... + """ + # When used as a context manager + if isinstance(name_or_func, str): + name = name_or_func + return Timer().context(name) + + func = name_or_func + + @wraps(func) + def wrapper(*args, **kwargs): + with Timer().context(func.__name__): + return func(*args, **kwargs) + + return wrapper + + +@contextmanager +def inverse_timer(name): + Timer().end(name) + try: + yield + finally: + Timer().start(name) diff --git a/slime/utils/tracking_utils.py b/slime/utils/tracking_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2af640d107e33775793a70d9030190f5fab00fab --- /dev/null +++ b/slime/utils/tracking_utils.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import wandb +from slime.utils.tensorboard_utils import _TensorboardAdapter + +from . import wandb_utils + + +def init_tracking(args, primary: bool = True, **kwargs): + if primary: + wandb_utils.init_wandb_primary(args, **kwargs) + else: + wandb_utils.init_wandb_secondary(args, **kwargs) + + +# TODO further refactor, e.g. put TensorBoard init to the "init" part +def log(args, metrics, step_key: str): + if args.use_wandb: + wandb.log(metrics) + + if args.use_tensorboard: + metrics_except_step = {k: v for k, v in metrics.items() if k != step_key} + _TensorboardAdapter(args).log(data=metrics_except_step, step=metrics[step_key]) diff --git a/slime/utils/train_dump_utils.py b/slime/utils/train_dump_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d31ba6498e3b49d51fb0885661d355a077f6835b --- /dev/null +++ b/slime/utils/train_dump_utils.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +from pathlib import Path + +import torch + +logger = logging.getLogger(__name__) + + +def save_debug_train_data(args, *, rollout_id, rollout_data): + if (path_template := args.save_debug_train_data) is not None: + rank = torch.distributed.get_rank() + path = Path(path_template.format(rollout_id=rollout_id, rank=rank)) + logger.info(f"Save debug train data to {path}") + path.parent.mkdir(parents=True, exist_ok=True) + torch.save( + dict( + rollout_id=rollout_id, + rank=rank, + rollout_data=rollout_data, + ), + path, + ) diff --git a/slime/utils/train_metric_utils.py b/slime/utils/train_metric_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a69c19226101ccb7ef4fdee9be7627f556769374 --- /dev/null +++ b/slime/utils/train_metric_utils.py @@ -0,0 +1,51 @@ +# 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 +from copy import deepcopy + +from slime.utils import tracking_utils +from slime.utils.metric_utils import compute_rollout_step +from slime.utils.timer import Timer + +logger = logging.getLogger(__name__) + + +def log_perf_data_raw( + rollout_id: int, args: Namespace, is_primary_rank: bool, compute_total_fwd_flops: Callable +) -> None: + timer_instance = Timer() + log_dict_raw = deepcopy(timer_instance.log_dict()) + timer_instance.reset() + + if not is_primary_rank: + return + + log_dict = {f"perf/{key}_time": val for key, val in log_dict_raw.items()} + + if ("perf/actor_train_time" in log_dict) and (compute_total_fwd_flops is not None): + total_fwd_flops = compute_total_fwd_flops(seq_lens=timer_instance.seq_lens) + + if "perf/log_probs_time" in log_dict: + log_dict["perf/log_probs_tflops"] = total_fwd_flops / log_dict["perf/log_probs_time"] + + if "perf/ref_log_probs_time" in log_dict: + log_dict["perf/ref_log_probs_tflops"] = total_fwd_flops / log_dict["perf/ref_log_probs_time"] + + if log_dict["perf/actor_train_time"] > 0: + log_dict["perf/actor_train_tflops"] = 3 * total_fwd_flops / log_dict["perf/actor_train_time"] + log_dict["perf/actor_train_tok_per_s"] = sum(timer_instance.seq_lens) / log_dict["perf/actor_train_time"] + + if "perf/train_wait_time" in log_dict and "perf/train_time" in log_dict: + total_time = log_dict["perf/train_wait_time"] + log_dict["perf/train_time"] + if total_time > 0: + log_dict["perf/step_time"] = total_time + log_dict["perf/wait_time_ratio"] = log_dict["perf/train_wait_time"] / total_time + + logger.info(f"perf {rollout_id}: {log_dict}") + + step = compute_rollout_step(args, rollout_id) + log_dict["rollout/step"] = step + tracking_utils.log(args, log_dict, step_key="rollout/step") diff --git a/slime/utils/typer_utils.py b/slime/utils/typer_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3177e850081c6944c8c5051eaaec44c5cf2f2357 --- /dev/null +++ b/slime/utils/typer_utils.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses +import inspect +from typing import Annotated + +import typer + + +def dataclass_cli(func, env_var_prefix: str = "SLIME_SCRIPT_"): + """Modified from https://github.com/fastapi/typer/issues/154#issuecomment-1544876144""" + + # The dataclass type is the first argument of the function. + sig = inspect.signature(func) + param = list(sig.parameters.values())[0] + dataclass_cls = param.annotation + assert dataclasses.is_dataclass(dataclass_cls) + + # To construct the signature, we remove the first argument (self) + # from the dataclass __init__ signature. + signature = inspect.signature(dataclass_cls.__init__) + old_parameters = list(signature.parameters.values()) + if len(old_parameters) > 0 and old_parameters[0].name == "self": + del old_parameters[0] + + new_parameters = [] + for param in old_parameters: + env_var_name = f"{env_var_prefix}{param.name.upper()}" + new_annotation = Annotated[param.annotation, typer.Option(envvar=env_var_name)] + new_parameters.append(param.replace(annotation=new_annotation)) + + def wrapped(**kwargs): + data = dataclass_cls(**kwargs) + print(f"Execute command with args: {data}") + return func(data) + + wrapped.__signature__ = signature.replace(parameters=new_parameters) + wrapped.__doc__ = func.__doc__ + wrapped.__name__ = func.__name__ + wrapped.__qualname__ = func.__qualname__ + + return wrapped + + +# unit test +if __name__ == "__main__": + from typer.testing import CliRunner + + @dataclasses.dataclass + class DemoArgs: + name: str + count: int = 1 + + app = typer.Typer() + + @app.command() + @dataclass_cli + def main(args: DemoArgs): + print(f"{args.name}|{args.count}") + + runner = CliRunner() + + res1 = runner.invoke(app, [], env={"SLIME_SCRIPT_NAME": "EnvName", "SLIME_SCRIPT_COUNT": "10"}) + print(f"{res1.stdout=}") + assert res1.exit_code == 0 + assert "EnvName|10" in res1.stdout.strip() + + res2 = runner.invoke(app, ["--count", "999"], env={"SLIME_SCRIPT_NAME": "EnvName"}) + print(f"{res2.stdout=}") + assert res2.exit_code == 0 + assert "EnvName|999" in res2.stdout.strip() + + print("✅ All Tests Passed!") diff --git a/slime/utils/types.py b/slime/utils/types.py new file mode 100644 index 0000000000000000000000000000000000000000..557fa6044aa03904f97756c2448f46ee767f56b3 --- /dev/null +++ b/slime/utils/types.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any + +import torch + + +@dataclass +class Sample: + """The sample generated""" + + group_index: int | None = None + index: int | None = None + # prompt - can be: + # - str: raw text prompt + # - list[dict[str, str]]: chat messages format + prompt: str | list[dict[str, str]] = "" + tokens: list[int] = field(default_factory=list) + multimodal_inputs: dict[str, Any] = None # raw multimodal data, e.g. images, videos, etc. + multimodal_train_inputs: dict[str, Any] = None # processed multimodal data, e.g. pixel_values, etc. + # response + response: str = "" + response_length: int = 0 + label: str | None = None + reward: float | dict[str, Any] | None = None + loss_mask: list[int] | None = None + weight_versions: list[str] = field(default_factory=list) + rollout_log_probs: list[float] | None = None # Log probabilities from rollout engine + rollout_routed_experts: list[list[int]] | None = None # Routed experts from rollout engine + remove_sample: bool = False + + class Status(Enum): + PENDING = "pending" + COMPLETED = "completed" + TRUNCATED = "truncated" + ABORTED = "aborted" + # Indicates a recoverable or non-critical failure during generation (e.g., tool call failure, + # external API error, parsing error). Unlike ABORTED, FAILED samples may still contain partial + # valid output and can be retried or handled gracefully. + FAILED = "failed" + + status: Status = Status.PENDING + + metadata: dict = field(default_factory=dict) + # metadata used during training, e.g., what loss to use for this sample. + train_metadata: dict | None = None + + class SpecInfo: + spec_accept_token_num: int = 0 + spec_draft_token_num: int = 0 + spec_verify_ct: int = 0 + spec_accept_rate: float = 0.0 + spec_accept_length: float = 0.0 + + def add(self, meta_info: dict, response_length: int): + self.spec_accept_token_num += meta_info["spec_accept_token_num"] + self.spec_draft_token_num += meta_info["spec_draft_token_num"] + self.spec_verify_ct += meta_info["spec_verify_ct"] + if self.spec_draft_token_num > 0: + # Notice: this does not iclude the bonus token generated by verify step. + self.spec_accept_rate = self.spec_accept_token_num / self.spec_draft_token_num + # self.spec_accept_rate = meta_info["spec_accept_rate"] # + if self.spec_verify_ct > 0: + self.spec_accept_length = response_length / self.spec_verify_ct + + def to_dict(self): + return { + "spec_accept_token_num": self.spec_accept_token_num, + "spec_draft_token_num": self.spec_draft_token_num, + "spec_verify_ct": self.spec_verify_ct, + "spec_accept_rate": self.spec_accept_rate, + "spec_accept_length": self.spec_accept_length, + } + + @staticmethod + def from_dict(data: dict): + info = Sample.SpecInfo() + info.spec_accept_token_num = data.get("spec_accept_token_num", 0) + info.spec_draft_token_num = data.get("spec_draft_token_num", 0) + info.spec_verify_ct = data.get("spec_verify_ct", 0) + info.spec_accept_rate = data.get("spec_accept_rate", 0.0) + info.spec_accept_length = data.get("spec_accept_length", 0.0) + return info + + spec_info: SpecInfo = field(default_factory=SpecInfo) + + def to_dict(self): + value = self.__dict__.copy() + value["status"] = self.status.value + value["spec_info"] = self.spec_info.to_dict() + return value + + @staticmethod + def from_dict(data: dict): + data["status"] = Sample.Status(data["status"]) + data["spec_info"] = Sample.SpecInfo.from_dict(data.get("spec_info", {})) + return Sample(**data) + + def get_reward_value(self, args) -> float: + return self.reward if not args.reward_key else self.reward[args.reward_key] + + @property + def effective_response_length(self): + return sum(self.loss_mask) if self.loss_mask is not None else self.response_length + + +@dataclass(frozen=True) +class ParamInfo: + name: str + dtype: torch.dtype + shape: torch.Size + attrs: dict + size: int + src_rank: int + + +# A dict-based batch produced along the rollout -> training path +# In Megatron backend, several fields are converted to torch.Tensor lists on GPU +# before being consumed by data iterators (see megatron_utils.actor._get_rollout_data). +RolloutBatch = dict[str, list[torch.Tensor] | list[int] | list[float] | list[str]] + + +@dataclass +class MultimodalType: + name: str # Type identifier used in message content (e.g., "image") + placeholder: str # Placeholder token in conversation messages (e.g., "") + + +class MultimodalTypes: + IMAGE = MultimodalType(name="image", placeholder="") + VIDEO = MultimodalType(name="video", placeholder="