nemo_automodel.components.models.common.mtp.mtp

Model-agnostic MTP scaffolding: depth iteration, token rolling, and loss.

Module Contents

Classes

Name	Description
`MTPConfig`	Runtime configuration for the MTP block.
`MTPContextParallelInputs`	Globally shifted MTP tensors prepared before context-parallel sharding.
`MTPModule`	Multi-Token Prediction block.

Functions

Name	Description
`_packed_seq_ids_from_batch`	Normalize supported packed-boundary metadata to token-aligned IDs.
`_packed_seq_ids_from_padded_lengths`	Expand padded packed-sequence lengths into token-aligned sequence IDs.
`get_mtp_loss_scaling_factor`	Return the model’s configured MTP auxiliary-loss scaling factor.
`prepare_mtp_context_parallel_inputs`	Prepare global future-token tensors before context-parallel sharding.
`roll_tensor`	Roll a tensor along `dim` by `shifts` and zero the wrapped slice.
`shift_packed_tensor`	Shift a token-aligned tensor left without crossing sequence boundaries.

API

class nemo_automodel.components.models.common.mtp.mtp.MTPConfig(
    num_layers: int = 0,
    layer_pattern: str = '',
    loss_scaling_factor: float = 0.1,
    use_repeated_layer: bool = False
)

Dataclass

Runtime configuration for the MTP block.

enabled

bool

layer_pattern

str = ''

loss_scaling_factor

float = 0.1

num_layers

int = 0

num_physical_depths

int

pattern_length

int

total_sublayers

int

use_repeated_layer

bool = False

class nemo_automodel.components.models.common.mtp.mtp.MTPContextParallelInputs(
    input_ids: tuple[torch.LongTensor, ...],
    position_ids: tuple[torch.LongTensor, ...],
    targets: tuple[torch.LongTensor, ...],
    valid_masks: tuple[torch.BoolTensor, ...],
    position_ids_seq_dim: int
)

Dataclass

Globally shifted MTP tensors prepared before context-parallel sharding.

input_ids

tuple[LongTensor, ...]

position_ids

tuple[LongTensor, ...]

position_ids_seq_dim

int

targets

tuple[LongTensor, ...]

valid_masks

tuple[BoolTensor, ...]

class nemo_automodel.components.models.common.mtp.mtp.MTPModule(
    mtp_config: nemo_automodel.components.models.common.mtp.mtp.MTPConfig,
    block_types_per_sublayer: list[str],
    sublayer_factory: typing.Callable[..., torch.nn.Module]
)

Bases: Module

Multi-Token Prediction block.

Holds a flat :class:nn.ModuleList of sublayers (length num_physical_depths * pattern_length) where the first sublayer of each physical depth carries the fusion modules (enorm, hnorm, eh_proj) and the last sublayer of each physical depth carries final_layernorm. This flat layout matches the HuggingFace export format used by Nemotron-V3 (mtp.layers.{i}.*).

The model-specific sublayer construction (which decoder block to use, how to handle MoE / attention / Mamba) is delegated to the caller via sublayer_factory.

Parameters:

mtp_config

MTPConfig

:class:MTPConfig describing depth and pattern.

block_types_per_sublayer

list[str]

List of block-type strings (one per inner sublayer position), length must equal mtp_config.pattern_length. Caller is responsible for parsing the model-specific symbol convention; this module does not interpret symbols.

sublayer_factory

Callable[..., nn.Module]

Callable factory(global_idx, depth, sublayer_idx, block_type, has_fusion, has_final_norm) -> nn.Module constructing one sublayer. The returned module must be callable as sublayer(hidden_states, **kwargs) -> Tensor and, when has_fusion=True, expose attributes enorm, hnorm, eh_proj. When has_final_norm=True it must expose final_layernorm.

layers

= nn.ModuleList(layers)

num_depths

int

pattern_length

int

nemo_automodel.components.models.common.mtp.mtp.MTPModule.forward(
    hidden_states: torch.Tensor,
    input_ids: torch.LongTensor | None = None,
    input_ids_per_depth: tuple[torch.LongTensor, ...] | None = None,
    embed_fn: typing.Callable[[torch.LongTensor], torch.Tensor] | None = None,
    embed_inputs: tuple[torch.Tensor, ...] | None = None,
    position_ids: torch.LongTensor | None = None,
    position_ids_per_depth: tuple[torch.LongTensor, ...] | None = None,
    block_kwargs = {}
) -> list[torch.Tensor]

Iterate over MTP depths and return per-depth hidden states.

Three mutually-exclusive input modes:

Single-rank / first-stage PP (default): pass input_ids plus embed_fn. The module rolls input_ids cumulatively left by 1 per depth and applies embed_fn to produce the future-token embedding for that depth.
Context parallel: pass input_ids_per_depth plus embed_fn. Each tensor is globally shifted and then sharded into the local CP token layout, so this module embeds it directly without a rank-local roll.
Final-stage PP / multimodal: pass embed_inputs (a tuple of pre-rolled per-depth embeddings, length num_depths). Used when the last PP stage no longer owns embed_tokens, or for multimodal models (e.g. SALM) where some positions carry continuous audio embeddings that cannot be recovered by re-embedding an integer token id — the caller pre-rolls the fused embedding tensor and passes it here.

Parameters:

hidden_states

torch.Tensor

Output of the main model’s final norm (h_0); tensor of shape [batch, sequence, hidden] or [tokens, hidden] for THD.

input_ids

torch.LongTensor | NoneDefaults to None

Token ids of shape [batch, sequence] (or [tokens] in THD). Rolled cumulatively left by 1 per depth. Mutually exclusive with input_ids_per_depth and embed_inputs.

input_ids_per_depth

tuple[torch.LongTensor, ...] | NoneDefaults to None

Optional tuple of num_depths pre-shifted token-ID tensors. Each has local CP shape [batch, sequence] or [tokens] and is embedded directly with embed_fn. Requires position_ids_per_depth.

embed_fn

Callable[[torch.LongTensor], torch.Tensor] | NoneDefaults to None

Callable applied to rolled input_ids to produce the future-token embedding (typically the model’s input embedding layer). Required when input_ids is supplied.

embed_inputs

tuple[torch.Tensor, ...] | NoneDefaults to None

Optional tuple of num_depths pre-computed future-token embeddings, each of shape [batch, sequence, hidden] or [tokens, hidden]. Mutually exclusive with input_ids/ input_ids_per_depth/embed_fn.

position_ids

torch.LongTensor | NoneDefaults to None

Position ids matching input_ids. When supplied, rolled cumulatively per depth in lockstep with input_ids (so slot t carries the original position of the rolled token) and forwarded to each sublayer via block_kwargs. Required for RoPE-using sublayers; ignored by sublayers that don’t consume it.

position_ids_per_depth

tuple[torch.LongTensor, ...] | NoneDefaults to None

Optional tuple of num_depths pre-computed future-token position tensors. Each tensor has shape [batch, sequence], [axes, batch, sequence] for multi-axis RoPE, or [tokens] for THD. When supplied, these tensors are forwarded directly instead of rolling rank-local position_ids. Use this when context parallelism has already sharded the sequence. Required with input_ids_per_depth and incompatible with rank-local input_ids rolling.

**block_kwargs

Defaults to {}

Forwarded to each sublayer’s __call__ (e.g. attention_mask).

Returns: list[torch.Tensor]

List of length num_depths containing hidden states of shape

nemo_automodel.components.models.common.mtp.mtp._packed_seq_ids_from_batch(
    batch: collections.abc.MutableMapping[str, object],
    input_ids: torch.Tensor
) -> torch.LongTensor | None

Normalize supported packed-boundary metadata to token-aligned IDs.

Parameters:

batch

MutableMapping[str, object]

Unsharded batch. Optional seq_idx or _packed_seq_ids tensors have shape [batch, sequence] (or [sequence] when batch is one). seq_lens_padded has shape [batch, num_sequences]. cu_seqlens_padded or cu_seqlens contains flattened cumulative boundaries of shape [num_sequences + 1] with optional negative sentinels.

input_ids

torch.Tensor

Global token-ID tensor of shape [batch, sequence] whose materialized token layout defines the expected output shape.

Returns: torch.LongTensor | None

Sequence-ID tensor of shape [batch, sequence] on the input device,

nemo_automodel.components.models.common.mtp.mtp._packed_seq_ids_from_padded_lengths(
    seq_lens_padded: torch.Tensor,
    batch_size: int,
    seq_len: int,
    device: torch.device
) -> torch.LongTensor

Expand padded packed-sequence lengths into token-aligned sequence IDs.

Parameters:

seq_lens_padded

torch.Tensor

Tensor of shape [batch, num_sequences] or [num_sequences] for a single batch row. Negative entries are unused sentinels; nonnegative entries are materialized sequence lengths including padding.

batch_size

int

Expected batch dimension of the returned tensor.

seq_len

int

Expected materialized token count in each batch row.

device

torch.device

Device for the returned token-aligned IDs.

Returns: torch.LongTensor

Sequence-ID tensor of shape [batch, sequence] on device.

nemo_automodel.components.models.common.mtp.mtp.get_mtp_loss_scaling_factor(
    model: torch.nn.Module,
    default: float = 0.1
) -> float

Return the model’s configured MTP auxiliary-loss scaling factor.

nemo_automodel.components.models.common.mtp.mtp.prepare_mtp_context_parallel_inputs(
    batch: collections.abc.MutableMapping[str, object],
    num_depths: int,
    ignore_index: int = -100
) -> nemo_automodel.components.models.common.mtp.mtp.MTPContextParallelInputs

Prepare global future-token tensors before context-parallel sharding.

Each MTP depth is shifted in global sequence order before CP partitions the token axis. Packed boundaries are preserved, so no future token, position, or target crosses from one document into another. Missing or shared position IDs are materialized in batch so the main model and MTP heads are subsequently sharded from the same global source.

Parameters:

batch

MutableMapping[str, object]

Mutable unsharded batch. input_ids and labels are tensors of shape [batch, sequence]. Optional position_ids has shape [batch, sequence], shared shape [1, sequence], or multi-axis RoPE shape [axes, batch, sequence]. Packed boundaries may use the tensor layouts documented by _packed_seq_ids_from_batch.

num_depths

int

Number of MTP future-token depths; must be positive.

ignore_index

intDefaults to -100

Fill value for invalid targets at trailing and packed boundary positions.

Returns: MTPContextParallelInputs

Per-depth input IDs, position IDs, targets, and validity masks. Token

nemo_automodel.components.models.common.mtp.mtp.roll_tensor(
    t: torch.Tensor,
    shifts: int = -1,
    dim: int = -1
) -> torch.Tensor

Roll a tensor along dim by shifts and zero the wrapped slice.

Used to shift input_ids / position_ids / labels left by one position per MTP depth. Single-GPU path only (no CP / packed-sequence handling).

Parameters:

torch.Tensor

Input tensor.

shifts

intDefaults to -1

Number of positions to shift (negative = left shift).

dim

intDefaults to -1

Dimension to roll along.

Returns: torch.Tensor

New tensor with the trailing |shifts| positions along dim

nemo_automodel.components.models.common.mtp.mtp.shift_packed_tensor(
    tensor: torch.Tensor,
    depth: int,
    seq_idx: torch.Tensor | None = None,
    fill_value: float | int = 0,
    batch_dim: int = 0,
    seq_dim: int = 1
) -> torch.Tensor

Shift a token-aligned tensor left without crossing sequence boundaries.

Parameters:

tensor

torch.Tensor

Token-aligned tensor in global sequence order. Its batch and sequence axes are selected by batch_dim and seq_dim.

depth

int

Number of future-token positions to shift; must be positive.

seq_idx

torch.Tensor | NoneDefaults to None

Optional sequence IDs of shape [batch, sequence]. Tokens whose shifted source has a different ID are filled.

fill_value

float | intDefaults to 0

Scalar used for trailing and cross-sequence positions.

batch_dim

intDefaults to 0

Batch dimension in tensor.

seq_dim

intDefaults to 1

Sequence dimension in tensor.

Returns: torch.Tensor

Tensor with the same shape, dtype, and device as tensor. The output