`nemo_rl.models.megatron.data`#

Module Contents#

Classes#

`ProcessedInputs`	Processed microbatch inputs used for model forward pass.
`ProcessedMicrobatch`	Container for a processed microbatch ready for model forward pass.

Functions#

`make_processed_microbatch_iterator`	Wrap a raw microbatch iterator to yield processed microbatches.
`get_microbatch_iterator`	Create a processed microbatch iterator from a batch of data.
`get_ltor_masks_and_position_ids`	Lazy proxy for `megatron.training.utils.get_ltor_masks_and_position_ids`.
`process_microbatch`	Process a microbatch for Megatron model forward pass.
`_make_r3_trace_token_identity`	Build debug-only `[batch_idx, token_pos, valid]` token identities.
`_verify_r3_trace_cp_token_alignment`	Verify debug identities line up with CP-local tokens and routed experts.
`_fill_routed_experts_padding`	Replace materialized jagged padding with a valid dummy top-k route.
`process_global_batch`	Process a global batch and compute normalization factors.
`_prepare_vlm_batch_for_megatron`	Prepare a [B, max_seq] batch for a model that does its own packing + CP sharding.
`_pack_sequences_for_megatron`	Pack sequences for Megatron model processing with optional context parallelism.
`_shard_routed_experts_for_cp`	CP-shard routed_experts / token_identity onto main’s per-seq packed layout.
`_get_pack_sequence_parameters_for_megatron`	Get pack sequence parameters for Megatron model processing with optional context parallelism.
`_unpack_sequences_from_megatron`	Unpack sequences from Megatron output format.
`get_and_validate_seqlen`

API#

class nemo_rl.models.megatron.data.ProcessedInputs#

Processed microbatch inputs used for model forward pass.

input_ids: torch.Tensor#: None

input_ids_cp_sharded: torch.Tensor#: None

attention_mask: Optional[torch.Tensor]#: None

position_ids: Optional[torch.Tensor]#: None

packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams]#: None

cu_seqlens_padded: Optional[torch.Tensor]#: None

mtp_loss_mask: Optional[torch.Tensor]#: None

routed_experts: Optional[torch.Tensor]#: None

routed_experts_cp_sharded: Optional[torch.Tensor]#: None

class nemo_rl.models.megatron.data.ProcessedMicrobatch#

Container for a processed microbatch ready for model forward pass.

This dataclass holds both the original data dictionary and the processed tensors needed for the Megatron model forward pass.

.. attribute:: data_dict

The original BatchedDataDict containing raw batch data

.. attribute:: input_ids

Processed input token IDs (may be packed for sequence packing)

.. attribute:: input_ids_cp_sharded

Context-parallel sharded input token IDs

.. attribute:: attention_mask

Attention mask tensor (None for packed sequences)

.. attribute:: position_ids

Position IDs tensor (None for packed sequences)

.. attribute:: packed_seq_params

PackedSeqParams for sequence packing (None if not packing)

.. attribute:: cu_seqlens_padded

Padded cumulative sequence lengths (None if not packing)

.. attribute:: mtp_loss_mask

Pre-computed MTP loss mask (token_mask × sample_mask). None when MTP is disabled or token/sample masks are absent.

.. attribute:: routed_experts

Optional token-aligned routed expert ids

.. attribute:: routed_experts_cp_sharded

Context-parallel sharded routed expert ids

data_dict: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]#: None

input_ids: torch.Tensor#: None

input_ids_cp_sharded: torch.Tensor#: None

attention_mask: Optional[torch.Tensor]#: None

position_ids: Optional[torch.Tensor]#: None

packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams]#: None

cu_seqlens_padded: Optional[torch.Tensor]#: None

mtp_loss_mask: Optional[torch.Tensor]#: None

routed_experts: Optional[torch.Tensor]#: None

routed_experts_cp_sharded: Optional[torch.Tensor]#: None

nemo_rl.models.megatron.data.make_processed_microbatch_iterator( raw_iterator: Iterator[nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]], cfg: dict[str, Any], seq_length_key: Optional[str], pad_individual_seqs_to_multiple_of: int, pad_packed_seq_to_multiple_of: int, straggler_timer: megatron.core.utils.StragglerDetector, pad_full_seq_to: Optional[int], delegate_pack_to_model: bool = False, ) → Iterator[nemo_rl.models.megatron.data.ProcessedMicrobatch]#

Wrap a raw microbatch iterator to yield processed microbatches.

This function takes a raw iterator that yields BatchedDataDict objects and wraps it to yield ProcessedMicrobatch objects that contain both the original data and the processed tensors ready for model forward pass.

Parameters:

raw_iterator – Iterator yielding raw BatchedDataDict microbatches
cfg – Configuration dictionary containing sequence_packing settings
seq_length_key – Key for sequence length in data dict (required for packing)
pad_individual_seqs_to_multiple_of – Padding multiple for individual sequences
pad_packed_seq_to_multiple_of – Padding multiple for packed sequences
pad_full_seq_to – Target length for full sequence padding (optional)

Yields:

ProcessedMicrobatch objects containing processed tensors ready for model forward

nemo_rl.models.megatron.data.get_microbatch_iterator( data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], cfg: dict[str, Any], mbs: int, straggler_timer: megatron.core.utils.StragglerDetector, seq_length_key: Optional[str] = None, delegate_pack_to_model: bool = False, ) → Tuple[Iterator[nemo_rl.models.megatron.data.ProcessedMicrobatch], int, int, int, int]#

Create a processed microbatch iterator from a batch of data.

This function creates an iterator that yields ProcessedMicrobatch objects, which contain both the original data dictionary and the processed tensors ready for model forward pass.

Parameters:

data – The batch data to create microbatches from
cfg – Configuration dictionary
mbs – Microbatch size
seq_length_key – Key for sequence lengths in data dict (auto-detected if None)

Returns:

Tuple containing the iterator and metadata

iterator: Iterator yielding ProcessedMicrobatch objects
data_iterator_len: Number of microbatches in the iterator
micro_batch_size: Size of each microbatch
seq_dim_size: Sequence length dimension size
padded_seq_length: Padded sequence length for pipeline parallelism (may differ from seq_length)

nemo_rl.models.megatron.data.get_ltor_masks_and_position_ids(

*args: Any,

**kwargs: Any,

) → Any#

Lazy proxy for megatron.training.utils.get_ltor_masks_and_position_ids.

The underlying import is deferred to call time so that importing this module does not pull in megatron.training -> modelopt -> transformers -> torchvision, which can crash on a duplicate torchvision ``roi_align` meta-kernel registration in the mcore venv.

nemo_rl.models.megatron.data.process_microbatch( data_dict: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], seq_length_key: Optional[str] = None, pad_individual_seqs_to_multiple_of: int = 1, pad_packed_seq_to_multiple_of: int = 1, pad_full_seq_to: Optional[int] = None, pack_sequences: bool = False, delegate_pack_to_model: bool = False, straggler_timer: Optional[megatron.core.utils.StragglerDetector] = None, ) → nemo_rl.models.megatron.data.ProcessedInputs#: Process a microbatch for Megatron model forward pass.

nemo_rl.models.megatron.data._make_r3_trace_token_identity( input_ids: torch.Tensor, seq_lengths: Optional[torch.Tensor] = None, ) → torch.Tensor#: Build debug-only [batch_idx, token_pos, valid] token identities.

nemo_rl.models.megatron.data._verify_r3_trace_cp_token_alignment( *, source_input_ids: torch.Tensor, source_routed_experts: Optional[torch.Tensor], input_ids_cp_sharded: torch.Tensor, routed_experts_cp_sharded: Optional[torch.Tensor], token_identity_cp_sharded: Optional[torch.Tensor], ) → Optional[int]#: Verify debug identities line up with CP-local tokens and routed experts.

nemo_rl.models.megatron.data._fill_routed_experts_padding( routed_experts: torch.Tensor, seq_lengths: torch.Tensor, ) → torch.Tensor#: Replace materialized jagged padding with a valid dummy top-k route.

nemo_rl.models.megatron.data.process_global_batch( data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction, dp_group: torch.distributed.ProcessGroup, *, batch_idx: int, batch_size: int, ) → dict[str, Any]#

Process a global batch and compute normalization factors.

Parameters:

data – Full dataset to extract a batch from
loss_fn – Loss function (used to check loss type for token-level validation)
dp_group – Data parallel process group for all-reduce
batch_idx – Index of batch to extract
batch_size – Size of batch to extract

Returns:

batch: The extracted batch
global_valid_seqs: Number of valid sequences across all ranks
global_valid_toks: Number of valid tokens across all ranks

Return type:

Dictionary containing

nemo_rl.models.megatron.data._prepare_vlm_batch_for_megatron( input_ids: torch.Tensor, seq_lengths: torch.Tensor, pad_individual_seqs_to_multiple_of: int, pad_full_seq_to: Optional[int] = None, ) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, megatron.core.packed_seq_params.PackedSeqParams, Optional[torch.Tensor], torch.Tensor]#

Prepare a [B, max_seq] batch for a model that does its own packing + CP sharding.

Used with mbridge VLM wrappers (e.g. Qwen3VL). The model’s forward calls preprocess_packed_seqs internally, which re-packs + CP-shards from attention_mask. So NeMo-RL must NOT pre-pack / CP-shard; it only:

pads each sequence (along dim 1) to pad_individual_seqs_to_multiple_of,
builds a bool attention_mask describing real token validity,
builds cu_seqlens_padded describing full (pre-shard) packed layout,
hands everything to the model as [B, max_seq].

When pad_full_seq_to is set (PP>1 requires a constant total packed length across microbatches), the last sequence’s effective length is extended so sum(padded_lens) == pad_full_seq_to. These extra positions are treated as “valid” by the model (so mbridge’s internal packing stays consistent) but should be masked out at the loss layer via token_mask.

Returns:

packed [1, T] view for downstream logprob/loss target slicing

input_ids_cp_sharded: [B, padded_max_seq] for the model forward
attention_mask: [B, padded_max_seq] bool (True for valid tokens)
packed_seq_params: PackedSeqParams(qkv_format=”thd”, cu_seqlens_*=padded)
cu_seqlens: None (unpadded cu_seqlens unused in this path)
cu_seqlens_padded: [B+1] int32 matching packed_seq_params

Return type:

input_ids

nemo_rl.models.megatron.data._pack_sequences_for_megatron( input_ids: torch.Tensor, seq_lengths: torch.Tensor, pad_individual_seqs_to_multiple_of: int = 1, pad_packed_seq_to_multiple_of: int = 1, pad_packed_seq_to: Optional[int] = None, cp_rank: int = 0, cp_size: int = 1, ) → tuple[torch.Tensor, megatron.core.packed_seq_params.PackedSeqParams, torch.Tensor, Optional[torch.Tensor]]#

Pack sequences for Megatron model processing with optional context parallelism.

Parameters:

input_ids – Input token IDs [batch_size, seq_length]
seq_lengths – Actual sequence lengths for each sample [batch_size]
pad_individual_seqs_to_multiple_of – Pad individual sequences to a multiple of this value
pad_packed_seq_to_multiple_of – Pad packed sequences to a multiple of this value
pad_packed_seq_to –
Pad packed sequences to this value (before CP)
- The three parameters above can be calculated using _get_pack_sequence_parameters_for_megatron, we do not recommend users to set these parameters manually.
cp_size – Context parallelism size

Returns:

packed_input_ids: Packed input tensor [1, T]
input_ids_cp_sharded: Sharded input tensor [cp_size, T // cp_size]
packed_seq_params: PackedSeqParams object
cu_seqlens: Cumulative sequence lengths
cu_seqlens_padded: Padded cumulative sequence lengths

Return type:

Tuple of

nemo_rl.models.megatron.data._shard_routed_experts_for_cp( routed_experts: Optional[torch.Tensor], token_identity: Optional[torch.Tensor], seq_lengths: torch.Tensor, cu_seqlens: torch.Tensor, cu_seqlens_padded: Optional[torch.Tensor], cp_rank: int, cp_size: int, ) → tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]#

CP-shard routed_experts / token_identity onto main’s per-seq packed layout.

Mirrors _pack_sequences_for_megatron’s per-sequence zigzag for input_ids: each sequence is padded to its padded length (from cu_seqlens_padded, so boundaries are IDENTICAL to input_ids) then sharded with _get_tokens_on_this_cp_rank(seq_dim=0). routed_experts pad rows use arange(topk) (a valid top-k route; mcore _validate_replay_tensor rejects 0/dup/-1). token_identity pads with 0 (verifier skips). No roll (these are per-token, not next-token targets). Returns (routed_packed, routed_cp_sharded, identity_packed, identity_cp_sharded), each [1, T(/cp), …] or None.

This additive helper is the only routed_experts-specific CP code path.

nemo_rl.models.megatron.data._get_pack_sequence_parameters_for_megatron( megatron_cfg: dict, pad_individual_seqs_to_multiple_of: int, max_seq_len_in_batch: int, )#

Get pack sequence parameters for Megatron model processing with optional context parallelism.

Parameters:

megatron_cfg – Megatron configuration
pad_individual_seqs_to_multiple_of – Pad individual sequences to a multiple of this value
max_seq_len_in_batch – Maximum sequence length in batch

Returns:

pad_individual_seqs_to_multiple_of: Pad individual sequences to a multiple of this value
pad_packed_seq_to_multiple_of: Pad packed sequences to a multiple of this value
pad_packed_seq_to: Pad packed sequences to this value (before CP)

Return type:

Tuple of

nemo_rl.models.megatron.data._unpack_sequences_from_megatron( output_tensor: torch.Tensor, seq_lengths: torch.Tensor, cu_seqlens: torch.Tensor, cu_seqlens_padded: Optional[torch.Tensor], original_batch_size: int, original_seq_length: int, ) → torch.Tensor#

Unpack sequences from Megatron output format.

Parameters:

output_tensor – Packed output tensor [1, T, vocab_size]
seq_lengths – Actual sequence lengths for each sample
cu_seqlens – Cumulative sequence lengths
cu_seqlens_padded – Padded cumulative sequence lengths (if CP was used)
original_batch_size – Original batch size
original_seq_length – Original maximum sequence length

Returns:

Unpacked output tensor [batch_size, seq_length, vocab_size]

nemo_rl.models.megatron.data.get_and_validate_seqlen( data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], )#

nemo_rl.models.megatron.data#

Module Contents#

Classes#

Functions#

API#

`nemo_rl.models.megatron.data`#