`nemo_rl.models.automodel.train`#

Training utilities for automodel (DTensor-based) policy workers.

This module provides post-processor classes and forward/backward functions that follow the same pattern as nemo_rl/models/megatron/train.py.

Key differences from megatron approach:

Post-processors compute results directly (no callable return pattern)
forward_with_post_processing_fn calls post-processor directly
automodel_forward_backward uses PyTorch autograd instead of Megatron’s pipeline

Module Contents#

Classes#

`PreparedModelForward`	Model inputs and optional Automodel CP state for one microbatch.
`LossPostProcessor`	Post-processor for computing training loss from model outputs.
`LogprobsPostProcessor`	Post-processor for computing log probabilities from model outputs.
`TopkLogitsPostProcessor`	Post-processor for computing top-k logits from model outputs.
`FullLogitsPostProcessor`	Export this rank’s raw teacher logits (full vocab, no reduction at the worker).
`ScorePostProcessor`	Post-processor for computing reward model scores from model outputs.

Functions#

`_build_model_batch`	Build a model-facing batch from canonical inputs.
`prepare_model_forward`	Prepare the model batch and resolve Automodel CP state when active.
`model_forward`	Run a model on an Automodel-prepared batch.
`extract_logits`	Extract logits from model outputs.
`apply_temperature_scaling`	Apply temperature scaling to logits.
`apply_top_k_top_p_filtering_for_local_logits`	Apply top-k and top-p filtering to the non-distributed logits.
`_cp_gather_logits`	Restore CP-local logits to canonical full-sequence order.
`forward_with_post_processing_fn`	Perform forward pass with pre-processed microbatch and apply post-processing.
`automodel_forward_backward`	Execute forward and backward passes for automodel.
`aggregate_training_statistics`	Aggregate training statistics across microbatches and ranks.

Data#

PostProcessingFunction

API#

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

class nemo_rl.models.automodel.train.PreparedModelForward#

Model inputs and optional Automodel CP state for one microbatch.

model_batch: dict[str, Any]#: None

cp_size: int#: None

cp_sharder: Optional[nemo_automodel.components.distributed.context_parallel.ContextParallelSharder]#: None

model_context_factory: Callable[[], contextlib.AbstractContextManager[Any]]#: None

nemo_rl.models.automodel.train._build_model_batch( model: torch.nn.Module, processed_inputs: nemo_rl.models.automodel.data.ProcessedInputs, *, is_reward_model: bool, allow_flash_attn_args: bool, clone_model_tensors: bool, ) → dict[str, Any]#: Build a model-facing batch from canonical inputs.

nemo_rl.models.automodel.train.prepare_model_forward( model: torch.nn.Module, processed_inputs: nemo_rl.models.automodel.data.ProcessedInputs, *, device_mesh: Optional[torch.distributed.device_mesh.DeviceMesh], cp_size: int, padding_token_id: int, is_reward_model: bool, allow_flash_attn_args: bool, ) → nemo_rl.models.automodel.train.PreparedModelForward#: Prepare the model batch and resolve Automodel CP state when active.

nemo_rl.models.automodel.train.model_forward( model: torch.nn.Module, model_batch: dict[str, Any], ) → Any#

Run a model on an Automodel-prepared batch.

Parameters:

model – The model to run.
model_batch – Private batch returned by the CP sharder.

Returns:

Model-specific forward output.

nemo_rl.models.automodel.train.extract_logits( model: torch.nn.Module, outputs: Any, ) → torch.Tensor#

Extract logits from model outputs.

Parameters:

model – The model (used for lm_head if needed)
outputs – Model outputs (can be tensor, DTensor, or object with logits attribute)

Returns:

Logits tensor

Return type:

torch.Tensor

nemo_rl.models.automodel.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.automodel.train.apply_top_k_top_p_filtering_for_local_logits( logits: torch.Tensor, sampling_params: Optional[nemo_rl.algorithms.logits_sampling_utils.TrainingSamplingParams], ) → torch.Tensor#

Apply top-k and top-p filtering to the non-distributed logits.

Parameters:

logits – Logits tensor to filter
sampling_params – Sampling parameters

Returns:

Filtered logits

Return type:

torch.Tensor

nemo_rl.models.automodel.train._cp_gather_logits( logits: torch.Tensor | torch.distributed.tensor.DTensor, cp_sharder: nemo_automodel.components.distributed.context_parallel.ContextParallelSharder, seq_dim: int = 1, ) → torch.Tensor | torch.distributed.tensor.DTensor#

Restore CP-local logits to canonical full-sequence order.

Keeps a tensor-parallel DTensor a DTensor on the same vocabulary mesh: only the sequence dimension is reassembled.

nemo_rl.models.automodel.train.forward_with_post_processing_fn( model: torch.nn.Module, prepared: nemo_rl.models.automodel.train.PreparedModelForward, post_processing_fn: nemo_rl.models.automodel.train.PostProcessingFunction, processed_mb: nemo_rl.models.automodel.data.ProcessedMicrobatch, 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, sequence_dim: int = 1, ) → Tuple[Any, dict[str, Any], nemo_rl.models.automodel.data.ProcessedMicrobatch]#

Perform forward pass with pre-processed microbatch and apply post-processing.

This function takes a pre-processed microbatch (with sequence packing already handled), runs the forward step through the model, and applies the post-processing function to compute the result.

Unlike the megatron approach which returns a callable, this directly computes and returns the result since automodel uses PyTorch autograd.

Parameters:

model – The model to run forward pass on
prepared – Per-microbatch model batch, CP layout, and forward context.
post_processing_fn – Post-processing function to apply to the logits
processed_mb – Pre-fetched ProcessedMicrobatch containing data and processed inputs
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)
sequence_dim – Sequence dimension

Returns:

(result, metrics, processed_microbatch) - result: Output from post-processing (loss, logprobs, topk, or scores) - metrics: Dictionary of metrics from post-processing - processed_microbatch: The ProcessedMicrobatch that was processed

Return type:

tuple

nemo_rl.models.automodel.train.automodel_forward_backward( model: torch.nn.Module, data_iterator: Iterator[nemo_rl.models.automodel.data.ProcessedMicrobatch], post_processing_fn: nemo_rl.models.automodel.train.PostProcessingFunction, device_mesh: Optional[torch.distributed.device_mesh.DeviceMesh], padding_token_id: int, autocast_context_factory: Callable[[], contextlib.AbstractContextManager[Any]], forward_only: bool = False, is_reward_model: bool = False, allow_flash_attn_args: bool = True, 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, sequence_dim: int = 1, dp_size: int = 1, cp_size: int = 1, num_global_batches: int = 1, num_valid_microbatches: Optional[int] = None, on_microbatch_start: Optional[Callable[[int], None]] = None, ) → list[Tuple[Any, dict[str, Any]]]#

Execute forward and backward passes for automodel.

This is the main training loop function that coordinates forward and backward passes across multiple microbatches using PyTorch autograd.

Unlike megatron_forward_backward which uses Megatron’s pipeline parallel framework, this uses standard PyTorch operations.

Parameters:

model – The model to train
data_iterator – Iterator yielding ProcessedMicrobatch objects (already processed)
post_processing_fn – Post-processing function to apply to the logits
device_mesh – Worker device mesh used by Automodel CP resolution.
padding_token_id – Token ID used for Automodel sequence padding.
autocast_context_factory – Worker-owned precision context factory.
forward_only – If True, skip backward pass
is_reward_model – Whether this is a reward model
allow_flash_attn_args – Whether to pass flash_attn_kwargs to model
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)
sequence_dim – Sequence dimension
dp_size – Data parallel size
cp_size – Context parallel size
num_global_batches – Number of global batches (for metric scaling)
num_valid_microbatches – Number of valid (non-dummy) microbatches. If provided, microbatches beyond this index are treated as dummy batches (loss *= 0). If None, all microbatches are considered valid.
on_microbatch_start – Optional callback called at the start of each microbatch with the microbatch index. Useful for cache clearing, etc.

Returns:

List of (result, metrics) tuples from each microbatch

class nemo_rl.models.automodel.train.LossPostProcessor( loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction, cfg: nemo_rl.models.policy.PolicyConfig, cp_mesh: Any, cp_size: int, dp_size: int, enable_seq_packing: bool = False, sampling_params: Optional[nemo_rl.algorithms.logits_sampling_utils.TrainingSamplingParams] = None, )#

Post-processor for computing training loss from model outputs.

Initialization

Initialize LossPostProcessor.

Parameters:

loss_fn – Loss function to compute loss
cfg – Configuration dictionary
cp_mesh – Context parallel mesh, used only to resolve the CP process group handed to the loss. The sequence layout itself belongs to the per-microbatch cp_sharder.
cp_size – Context parallel size
dp_size – Data parallel size
enable_seq_packing – Whether sequence packing is enabled
sampling_params – Sampling parameters

property cp_gradient_fanout: int#: Number of CP loss consumers for each local model contribution.

__call__( logits: torch.Tensor, data_dict: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], processed_inputs: nemo_rl.models.automodel.data.ProcessedInputs, global_valid_seqs: torch.Tensor, global_valid_toks: torch.Tensor, *, cp_sharder: Optional[nemo_automodel.components.distributed.context_parallel.ContextParallelSharder], sequence_dim: int = 1, ) → tuple[torch.Tensor, dict[str, Any]]#

Compute loss from logits.

Parameters:

logits – Model output logits
data_dict – Microbatch data
processed_inputs – Processed inputs
global_valid_seqs – Global valid sequence count
global_valid_toks – Global valid token count
cp_sharder – Per-microbatch Automodel sequence-layout owner, or None when context parallelism is inactive.
sequence_dim – Sequence dimension

Returns:

Tuple of (loss, metrics)

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

Post-processor for computing log probabilities from model outputs.

Initialization

Initialize LogprobsPostProcessor.

Parameters:

cfg – Configuration dictionary
enable_seq_packing – Whether sequence packing is enabled
sampling_params – Sampling parameters

__call__( logits: torch.Tensor, data_dict: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], processed_inputs: nemo_rl.models.automodel.data.ProcessedInputs, original_batch_size: int, original_seq_len: int, *, cp_sharder: Optional[nemo_automodel.components.distributed.context_parallel.ContextParallelSharder], sequence_dim: int = 1, ) → torch.Tensor#

Compute token log probabilities from logits.

Parameters:

logits – Model output logits
data_dict – Microbatch data
processed_inputs – Processed inputs
original_batch_size – Original batch size before packing
original_seq_len – Original sequence length before packing
cp_sharder – Per-microbatch Automodel sequence-layout owner, or None when context parallelism is inactive.
sequence_dim – Sequence dimension

Returns:

Token log probabilities tensor [batch_size, seq_length]

_compute_local_logprobs( logits: torch.Tensor, input_ids: torch.Tensor, ) → torch.Tensor#

Compute logprobs locally without distributed processing.

Parameters:

logits – Model output logits
input_ids – Input token IDs

Returns:

Token log probabilities

When logprob_chunk_size is set, log-softmax and gather run per sequence chunk to bound peak memory for long-context runs.

class nemo_rl.models.automodel.train.TopkLogitsPostProcessor( cfg: nemo_rl.models.policy.PolicyConfig, tp_mesh: Any, k: int, enable_seq_packing: bool = False, )#

Post-processor for computing top-k logits from model outputs.

Initialization

Initialize TopkLogitsPostProcessor.

Parameters:

cfg – Configuration dictionary
tp_mesh – Tensor parallel mesh, used for the vocabulary-parallel top-k. The sequence layout belongs to the per-microbatch cp_sharder.
k – Number of top logits to return
enable_seq_packing – Whether sequence packing is enabled

__call__( logits: torch.Tensor, data_dict: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], processed_inputs: nemo_rl.models.automodel.data.ProcessedInputs, original_batch_size: int, original_seq_len: int, *, cp_sharder: Optional[nemo_automodel.components.distributed.context_parallel.ContextParallelSharder], sequence_dim: int = 1, ) → tuple[torch.Tensor, torch.Tensor]#

Compute top-k logits and indices from model outputs.

Parameters:

logits – Model output logits
data_dict – Microbatch data
processed_inputs – Processed inputs
original_batch_size – Original batch size before packing
original_seq_len – Original sequence length before packing
cp_sharder – Per-microbatch Automodel sequence-layout owner, or None when context parallelism is inactive.
sequence_dim – Sequence dimension

Returns:

Tuple of (top-k values, top-k indices) tensors

class nemo_rl.models.automodel.train.FullLogitsPostProcessor( cfg: nemo_rl.models.policy.PolicyConfig, cp_mesh: Any, cp_size: int, enable_seq_packing: bool = False, )#

Export this rank’s raw teacher logits (full vocab, no reduction at the worker).

Used by cross-tokenizer distillation; the loss fn does all vocab reduction (none at the worker) so the distributed result matches the single-GPU PyTorch reference. Teacher TP/CP are supported and may differ from the student’s: under TP each rank emits its vocab shard, under CP it allgathers and re-emits its contiguous seq slice; the IPC consumer reassembles the global [B, T_t, V_t]. Sequence packing raises NotImplementedError.

Initialization

Initialize FullLogitsPostProcessor.

Parameters:

cfg – Configuration dictionary
cp_mesh – Context parallel mesh, used to pick this rank’s contiguous IPC window. The sequence layout of the logits themselves belongs to the per-microbatch cp_sharder.
cp_size – Context parallel size
enable_seq_packing – Whether sequence packing is enabled

__call__( logits: torch.Tensor, data_dict: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], processed_inputs: Any, original_batch_size: int, original_seq_len: int, *, cp_sharder: Optional[nemo_automodel.components.distributed.context_parallel.ContextParallelSharder], sequence_dim: int = 1, ) → torch.Tensor#

class nemo_rl.models.automodel.train.ScorePostProcessor(cfg: nemo_rl.models.policy.PolicyConfig)#

Post-processor for computing reward model scores from model outputs.

Initialization

Initialize ScorePostProcessor.

Parameters:: cfg – Configuration dictionary

__call__(logits: torch.Tensor) → torch.Tensor#

Extract scores from reward model outputs.

Parameters:: logits – Model output logits
Returns:: Scores tensor

nemo_rl.models.automodel.train.aggregate_training_statistics( losses: list[float], all_mb_metrics: list[dict[str, Any]], grad_norm: Optional[torch.Tensor], dp_group: Any, dtype: torch.dtype, ) → dict[str, Any]#

Aggregate training statistics across microbatches and ranks.

Parameters:

losses – List of loss values from each microbatch
all_mb_metrics – List of metrics dictionaries from each microbatch
grad_norm – Gradient norm tensor (or None if eval mode)
dp_group – Data parallel process group for all-reduce
dtype – Model dtype for metrics

Returns:

Dictionary containing aggregated metrics including global_loss, grad_norm, etc.

nemo_rl.models.automodel.train#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.models.automodel.train`#