`nemo_rl.models.megatron.train`#

Module Contents#

Classes#

`LossPostProcessor`
`LogprobsPostProcessor`
`TeacherFullPayloadPostProcessor`	Emit sampled-token logprobs plus the full-vocabulary teacher payload.
`TopkLogitsPostProcessor`

Functions#

`_prepare_padding_mask_for_model`	Match a CP-local padding mask to the model’s sequence-parallel layout.
`suspend_activation_offload_for_forward_only`	Keep inference-only RL phases from consuming MCore’s training warmup.
`model_forward`	Perform a single forward pass through the model.
`apply_temperature_scaling`	Apply temperature scaling to logits.
`forward_with_post_processing_fn`	Perform forward pass with pre-processed microbatch and return output tensor and post-processing function.
`megatron_forward_backward`	Execute forward and backward passes using Megatron’s utilities.
`should_reduce_loss_across_context_parallel`	Return whether scalar loss reporting must include context-parallel ranks.
`strip_context_parallel_local_loss_metric`	Remove CP-local loss after it contributes to the global reduction.
`aggregate_training_statistics`	Aggregate training statistics across microbatches and data-parallel ranks.

Data#

PostProcessingFunction

API#

nemo_rl.models.megatron.train.PostProcessingFunction#: None

nemo_rl.models.megatron.train._prepare_padding_mask_for_model( model: megatron.core.models.gpt.GPTModel, padding_mask: Optional[torch.Tensor], model_slices_context_parallel_inputs: bool = False, ) → Optional[torch.Tensor]#: Match a CP-local padding mask to the model’s sequence-parallel layout.

nemo_rl.models.megatron.train.suspend_activation_offload_for_forward_only( model: Union[megatron.core.models.gpt.GPTModel, List[megatron.core.models.gpt.GPTModel]], forward_only: bool, ) → Iterator[None]#: Keep inference-only RL phases from consuming MCore’s training warmup.

nemo_rl.models.megatron.train.model_forward( model: megatron.core.models.gpt.GPTModel, data_dict: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], input_ids_cp_sharded: torch.Tensor, position_ids: Optional[torch.Tensor], attention_mask: Optional[torch.Tensor], packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None, defer_fp32_logits: Optional[bool] = False, mtp_loss_mask: Optional[torch.Tensor] = None, padding_mask: Optional[torch.Tensor] = None, straggler_timer: Optional[megatron.core.utils.StragglerDetector] = None, use_fused_linear_logprobs: bool = False, labels_cp_sharded: Optional[torch.Tensor] = None, loss_mask_cp_sharded: Optional[torch.Tensor] = None, media_token_validity_mask: Optional[torch.Tensor] = None, model_slices_context_parallel_inputs: bool = False, ) → torch.Tensor#

Perform a single forward pass through the model.

Parameters:

model – The model to run forward pass on
data_dict – Dictionary containing batch data
input_ids_cp_sharded – Model-forward token IDs. Usually CP-sharded; models that insert media before CP selection receive the full packed THD row.
position_ids – Position IDs for tokens
attention_mask – Attention mask for the sequence
packed_seq_params – Parameters for packed sequences (optional)
defer_fp32_logits – Whether to skip the conversion of logits to fp32
mtp_loss_mask – MTP loss mask to exclude prompt tokens from MTP loss (optional)
padding_mask – Packed-sequence padding mask for MoE routing (optional)
straggler_timer – Straggler detector for profiling the forward pass
use_fused_linear_logprobs – Whether to compute logprobs with the fused chunked linear cross-entropy kernel (directly from hidden states)
labels_cp_sharded – Target-aligned labels for direct MCore loss computation
loss_mask_cp_sharded – Target-aligned mask for direct MCore loss computation
media_token_validity_mask – Which media-token positions actually anchor a projected feature, already in this model’s token layout. Only passed when the model accepts it; otherwise the model derives its own.
model_slices_context_parallel_inputs – Whether the model CP-slices its own inputs.

Returns:

Output tensor from the model (logits)

Return type:

torch.Tensor

nemo_rl.models.megatron.train.apply_temperature_scaling( logits: torch.Tensor, sampling_params: Optional[nemo_rl.algorithms.logits_sampling_utils.TrainingSamplingParams], ) → torch.Tensor#

Apply temperature scaling to logits.

Parameters:

logits – Logits tensor to scale
sampling_params – Sampling parameters

Returns:

Temperature-scaled logits

Return type:

torch.Tensor

nemo_rl.models.megatron.train.forward_with_post_processing_fn( data_iterator: Iterator[nemo_rl.models.megatron.data.ProcessedMicrobatch], model: megatron.core.models.gpt.GPTModel, post_processing_fn: nemo_rl.models.megatron.train.PostProcessingFunction, defer_fp32_logits: Optional[bool] = False, global_valid_seqs: Optional[torch.Tensor] = None, global_valid_toks: Optional[torch.Tensor] = None, sampling_params: Optional[nemo_rl.algorithms.logits_sampling_utils.TrainingSamplingParams] = None, straggler_timer: Optional[megatron.core.utils.StragglerDetector] = None, draft_model: Optional[nemo_rl.models.megatron.config.MegatronModule] = None, enable_hidden_capture: Optional[bool] = False, enable_opd_full_capture: bool = False, use_fused_linear_logprobs: bool = False, use_router_replay: bool = False, router_replay_train: bool = False, model_slices_context_parallel_inputs: bool = False, ) → Tuple[torch.Tensor, Callable]#

Perform forward pass with pre-processed microbatch and return output tensor and post-processing function.

This function takes a pre-processed microbatch (with sequence packing already handled), runs the forward step through the model, and prepares a post-processing function for post-processing the outputs.

Parameters:

data_iterator – Iterator yielding ProcessedMicrobatch objects (already processed)
model – The model to run forward pass on
post_processing_fn – Post-processing function to post-process the logits
defer_fp32_logits – Whether to defer FP32 conversion of logits
global_valid_seqs – Global valid sequence count for loss normalization
global_valid_toks – Global valid token count for loss normalization
sampling_params – Sampling parameters (top-k, top-p, temperature)
straggler_timer – Straggler detector for profiling the forward pass
draft_model – Draft model for online draft model training
enable_hidden_capture – Whether to enable hidden state capture for draft model training
enable_opd_full_capture – Whether to capture pre-LM-head hidden states for the full-vocabulary MOPD teacher payload
model_slices_context_parallel_inputs – Whether the model CP-slices its own inputs.

Returns:

(output_tensor, post_processing_fn_wrapped) - output_tensor: Raw model outputs (logits) - post_processing_fn_wrapped: Function to create output post-processing function when called

Return type:

tuple

nemo_rl.models.megatron.train.megatron_forward_backward( model: megatron.core.models.gpt.GPTModel, data_iterator: Iterator[nemo_rl.models.megatron.data.ProcessedMicrobatch], num_microbatches: int, seq_length: int, mbs: int, post_processing_fn: nemo_rl.models.megatron.train.PostProcessingFunction, forward_only: bool = False, defer_fp32_logits: Optional[bool] = False, global_valid_seqs: Optional[torch.Tensor] = None, global_valid_toks: Optional[torch.Tensor] = None, sampling_params: Optional[nemo_rl.algorithms.logits_sampling_utils.TrainingSamplingParams] = None, straggler_timer: Optional[megatron.core.utils.StragglerDetector] = None, draft_model: Optional[nemo_rl.models.megatron.config.MegatronModule] = None, enable_hidden_capture: Optional[bool] = False, enable_opd_full_capture: bool = False, use_fused_linear_logprobs: bool = False, use_router_replay: bool = False, router_replay_train: bool = False, model_slices_context_parallel_inputs: bool = False, ) → Any#

Execute forward and backward passes using Megatron’s utilities.

This is the main training loop function that coordinates forward and backward passes across multiple microbatches using Megatron’s pipeline parallel execution framework.

Parameters:

model – The model to train
data_iterator – Iterator yielding ProcessedMicrobatch objects (already processed)
num_microbatches – Number of microbatches to process
seq_length – Sequence length
mbs – Micro batch size
post_processing_fn – Post-processing function to post-process the logits
forward_only – If True, skip backward pass
defer_fp32_logits – Whether to skip the conversion of logits to fp32
global_valid_seqs – Global valid sequence count for loss normalization
global_valid_toks – Global valid token count for loss normalization
sampling_params – Sampling parameters (top-k, top-p, temperature)
straggler_timer – Straggler detector for profiling the forward pass
draft_model – Draft model for online draft model training
enable_hidden_capture – Whether to enable hidden state capture for draft model training
enable_opd_full_capture – Whether to capture pre-LM-head hidden states for the full-vocabulary MOPD teacher payload
model_slices_context_parallel_inputs – Whether the model CP-slices its own inputs.

Returns:

Results from the forward/backward execution

class nemo_rl.models.megatron.train.LossPostProcessor( loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction, cfg: nemo_rl.models.policy.PolicyConfig, num_microbatches: int = 1, cp_normalize: bool = True, sampling_params: Optional[nemo_rl.algorithms.logits_sampling_utils.TrainingSamplingParams] = None, draft_model: Optional[nemo_rl.models.megatron.config.MegatronModule] = None, prepare_fn: Optional[Callable[..., Any]] = None, defer_draft_normalization: bool = False, teacher_output_layer_weight_by_index: Optional[dict[int, torch.Tensor]] = None, )#

Initialization

Build a per-microbatch loss post-processor for the Megatron train loop.

Parameters:

loss_fn – Loss function to wrap.
cfg – Policy(-like) config; supplies sequence_packing / logprob_chunk_size.
num_microbatches – Microbatch count, used to counteract Megatron’s per-microbatch loss averaging.
cp_normalize – Whether to divide the loss by the context-parallel size.
sampling_params – Optional temperature / top-k/p for logprob losses.
draft_model – Optional EAGLE draft model for distillation.
prepare_fn – Optional override for the default prepare_loss_input. Must accept (logits, data, loss_fn, vocab_parallel_rank, vocab_parallel_group, context_parallel_group) and return (loss_input, data); value models pass one that right-shifts and CP-all-gathers the scalar value-head output.
defer_draft_normalization – Return raw draft loss statistics for split optimizer-step finalization instead of normalizing per microbatch.
teacher_output_layer_weight_by_index – This rank’s teacher LM-head shards for every configured teacher, keyed by the stable per-checkpoint index rows are tagged with (see OPD_FULL_TEACHER_INDEX_FIELD). The full-vocabulary MOPD loss projects each row’s teacher payload with them. They ride this argument rather than the data dict because the sequence-packing wrapper batch-slices every data entry, and rather than the loss object because that is pickled to workers.

__call__( data_dict: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None, global_valid_seqs: Optional[torch.Tensor] = None, global_valid_toks: Optional[torch.Tensor] = None, prepacked_loss_mask: Optional[torch.Tensor] = None, ) → Callable[[torch.Tensor], Tuple[torch.Tensor, Dict[str, Any]]]#

Create a loss post-processing function for training.

This function wraps a loss function with the necessary context and parameters to compute loss and metrics from model outputs. It handles sequence packing and context parallelism normalization.

Parameters:

data_dict – Batched data dictionary for the current microbatch
packed_seq_params – Parameters for packed sequences (optional)
global_valid_seqs – Global valid sequence count for loss normalization
global_valid_toks – Global valid token count for loss normalization
prepacked_loss_mask – CP-local target-aligned mask when MCore returned per-token losses for a direct packed row

Returns:

Function that takes output tensor and returns (loss, metrics) tuple

Return type:

Callable

nemo_rl.models.megatron.train.should_reduce_loss_across_context_parallel( data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any] | dict[str, Any], ) → bool#: Return whether scalar loss reporting must include context-parallel ranks.

nemo_rl.models.megatron.train.strip_context_parallel_local_loss_metric( metrics: Dict[str, List[Any]], enabled: bool, ) → Dict[str, List[Any]]#: Remove CP-local loss after it contributes to the global reduction.

class nemo_rl.models.megatron.train.LogprobsPostProcessor( cfg: nemo_rl.models.policy.PolicyConfig, sampling_params: Optional[nemo_rl.algorithms.logits_sampling_utils.TrainingSamplingParams] = None, use_fused_linear_logprobs: bool = False, )#

Initialization

__call__( data_dict: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], input_ids: torch.Tensor, cu_seqlens_padded: Optional[torch.Tensor | nemo_rl.utils.sequence_lengths.CpuIntTuple], original_seq_length: int, ) → Callable[[torch.Tensor], Tuple[torch.Tensor, Dict[str, torch.Tensor]]]#

Create a post-processing function that computes token log probabilities.

This function returns a processor that takes model logits and converts them to token-level log probabilities, handling both packed and unpacked sequences.

Parameters:

data_dict – Batched data dictionary containing input sequences
input_ids – Processed input token IDs
cu_seqlens_padded – Cumulative sequence lengths for packed sequences
original_seq_length – Sequence width before dense padding was applied

Returns:

Function that takes output tensor and returns (dummy_loss, {“logprobs”: token_logprobs})

Return type:

Callable

class nemo_rl.models.megatron.train.TeacherFullPayloadPostProcessor( cfg: nemo_rl.models.policy.PolicyConfig, payload: str, payload_dtype: torch.dtype, sampling_params: Optional[nemo_rl.algorithms.logits_sampling_utils.TrainingSamplingParams] = None, )#

Emit sampled-token logprobs plus the full-vocabulary teacher payload.

Full-vocabulary MOPD needs the teacher’s whole next-token distribution, not just the sampled token’s log-probability. Both come out of the same forward pass: this wraps :class:LogprobsPostProcessor for the scalar column and adds the payload the student reconstructs the distribution from.

The payload is returned in canonical [B, S, D] layout with tensor, context, and sequence parallelism already undone, so the student’s own parallelism can differ from the teacher’s.

Initialization

__call__( data_dict: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], input_ids: torch.Tensor, cu_seqlens_padded: Optional[torch.Tensor | nemo_rl.utils.sequence_lengths.CpuIntTuple], original_seq_length: int, hidden_states: Optional[torch.Tensor] = None, ) → Callable[[torch.Tensor], Tuple[torch.Tensor, Dict[str, torch.Tensor]]]#

Create the post-processing function for a teacher full-payload forward.

Parameters:

data_dict – Batched data dictionary containing input sequences.
input_ids – Processed input token IDs.
cu_seqlens_padded – Cumulative sequence lengths for packed sequences.
original_seq_length – Sequence width before dense padding was applied.
hidden_states – Captured pre-LM-head hidden states [S, B, H], required for the hidden_states payload.

Returns:

[B, S], “teacher_full_payload”: [B, S, D]})``.

Return type:

Callable mapping the model output to ``(dummy_loss, {“logprobs”

class nemo_rl.models.megatron.train.TopkLogitsPostProcessor( cfg: nemo_rl.models.policy.PolicyConfig, k: int, )#

Initialization

__call__( data_dict: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], cu_seqlens_padded: Optional[torch.Tensor | nemo_rl.utils.sequence_lengths.CpuIntTuple], original_seq_length: int, ) → Callable[[torch.Tensor], Tuple[torch.Tensor, Dict[str, torch.Tensor]]]#

Create a post-processing function that computes top-k logits and indices.

This function returns a processor that extracts the top-k highest logits and their corresponding vocabulary indices from model outputs. It handles tensor parallelism, context parallelism, and sequence packing.

Parameters:

data_dict – Batched data dictionary
cu_seqlens_padded – Cumulative sequence lengths for packed sequences
original_seq_length – Sequence width before dense padding was applied

Returns:

Function that takes output tensor and returns (dummy_loss, {“topk_logits”: values, “topk_indices”: indices})

Return type:

Callable

nemo_rl.models.megatron.train.aggregate_training_statistics( all_mb_metrics: List[Dict[str, Any]], losses: List[Union[float, torch.Tensor]], data_parallel_group: torch.distributed.ProcessGroup, ) → Tuple[Dict[str, List[Any]], torch.Tensor]#

Aggregate training statistics across microbatches and data-parallel ranks.

Computes a global loss by all-reducing per-gradient-buffer losses across the data-parallel group, then collects per-microbatch metrics into lists keyed by metric name.

Parameters:

all_mb_metrics – List of metric dicts from each microbatch.
losses – List of per-gradient-buffer scalar losses on this rank. Tensor values remain on device until the final reduction.
data_parallel_group – The data-parallel process group for all-reduce.

Returns:

mb_metrics: Dict mapping metric names to lists of values across microbatches.
global_loss: Tensor of losses summed across all data-parallel ranks.

Return type:

Tuple of

nemo_rl.models.megatron.train#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.models.megatron.train`#