nemo_automodel.components.models.common.mtp

Shared scaffolding for Multi-Token Prediction (MTP) auxiliary heads.

MTP follows the DeepSeek-V3 design (Liu et al., 2024). Each MTP “depth” predicts one additional future token; per depth the input is rolled left by one position, fused with the previous-depth hidden state, and passed through an inner block before producing logits via the shared LM head.

Components in this package are model-agnostic. Model-specific glue (building the inner block out of the model’s own decoder layers, wiring HF state-dict keys) lives in the model’s own package.

Submodules

nemo_automodel.components.models.common.mtp.mtp

Package 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
`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 = ''

Per-depth inner-block pattern, e.g. "*E" for one attention + one MoE sublayer per depth.

loss_scaling_factor

float = 0.1

Coefficient applied to the summed per-depth CE loss (default 0.1). The effective per-depth weight is loss_scaling_factor / num_layers.

num_layers

int = 0

Number of MTP forward iterations (D). 0 disables MTP. Equivalent to Megatron’s --mtp-num-layers.

num_physical_depths

int

pattern_length

int

total_sublayers

int

use_repeated_layer

bool = False

When True, build a single physical depth’s worth of sublayers and reuse it for all num_layers forward iterations (weight-tied across depths). Equivalent to Megatron’s --mtp-use-repeated-layer.

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, ...]

Per-depth token IDs of shape [batch, sequence] in global sequence order.

position_ids

tuple[LongTensor, ...]

Per-depth position IDs in global sequence order. Standard positions have shape [batch, sequence]; multi-axis RoPE uses [axes, batch, sequence].

position_ids_seq_dim

int

Sequence dimension of each position-ID tensor.

targets

tuple[LongTensor, ...]

Per-depth loss targets of shape [batch, sequence] in global sequence order, with invalid positions set to the loss ignore index.

valid_masks

tuple[BoolTensor, ...]

Per-depth masks of shape [batch, sequence] identifying positions whose future token remains in the same packed sequence.

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 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

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.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