# 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