nemo_automodel.components.distributed.thd_utils

Module Contents

Functions

Name	Description
`_thd_padding_mask`	Mark padding slots in a packed THD stream from the pack layout.
`process_input_for_thd`	Process inputs for THD (total, hidden, depth) format.
`split_batch_into_thd_chunks`	Process inputs for THD format by splitting batch into chunks for context parallelism.
`thd_padding_mask_from_token_ids`	Mark padding by token value, rejecting a pad id that is also content.

API

nemo_automodel.components.distributed.thd_utils._thd_padding_mask(
    total_tokens: int,
    valid_seq_lens: torch.Tensor,
    valid_seq_lens_padded: torch.Tensor | None,
    device: torch.device
) -> torch.Tensor

Mark padding slots in a packed THD stream from the pack layout.

Sequence i occupies [start_i, start_i + padded_i) and only its first real_i slots hold tokens, so everything else is padding regardless of which token id fills it. Deriving this by comparing against a pad token id instead misclassifies real tokens whenever that id is also content — see :func:thd_padding_mask_from_token_ids for the metadata-free fallback.

Parameters:

total_tokens

int

Length of the packed stream.

valid_seq_lens

torch.Tensor

Real (unpadded) length of each packed sequence.

valid_seq_lens_padded

torch.Tensor | None

Slot count reserved for each packed sequence, or None when sequences carry no individual padding.

device

torch.device

Device of the packed stream.

Returns: torch.Tensor

Boolean tensor [total_tokens]; True marks padding.

nemo_automodel.components.distributed.thd_utils.process_input_for_thd(
    batch: dict[str, torch.Tensor],
    seq_lens_padding_value: int = -1000,
    padding_token_id: int = 0
) -> dict[str, torch.Tensor]

Process inputs for THD (total, hidden, depth) format.

This function converts batched inputs from BSHD (batch, sequence, hidden, depth) format to THD format for packed sequence processing. In THD format, the batch dimension is collapsed and all sequences are concatenated along the sequence dimension. This supports both 2D token IDs and 3D embeddings for pipeline parallelism scenarios.

The function filters out padding values in seq_lens and seq_lens_padded (indicated by seq_lens_padding_value) and computes cumulative sequence lengths for efficient attention computation with Transformer Engine or other packed sequence implementations.

Parameters:

batch

dict[str, torch.Tensor]

Dictionary containing:

‘input_ids’: Input tensor of shape [batch_size, seq_len] for token IDs or [batch_size, seq_len, hidden_dim] for embeddings (in pipeline parallel scenarios)
‘labels’: Labels tensor of shape [batch_size, seq_len]
‘position_ids’: Position IDs tensor of shape [batch_size, seq_len] for standard RoPE, or [n_rope, batch_size, seq_len] for mRoPE (e.g. Qwen-VL, n_rope=3). Required.
‘seq_lens’: Sequence lengths tensor of shape [batch_size, num_packs] containing actual sequence lengths (excluding padding/separators). Values matching seq_lens_padding_value indicate padding and are filtered out.
‘seq_lens_padded’: Padded sequence lengths tensor of shape [batch_size, num_packs] containing lengths including separator tokens. Values matching seq_lens_padding_value indicate padding and are filtered out.

seq_lens_padding_value

intDefaults to -1000

Value used to indicate padding in seq_lens/seq_lens_padded tensors that should be filtered out (default: -1000)

padding_token_id

intDefaults to 0

Filler token id. The padding mask is derived from the pack layout, so this is only consulted by callers that have no pack metadata (see thd_padding_mask_from_token_ids).

Returns: dict[str, torch.Tensor]

Dictionary containing:

‘input_ids’: Reshaped tensor of shape [total_tokens] for 2D token IDs or [total_tokens, hidden_dim] for 3D embeddings
‘labels’: Reshaped labels tensor of shape [total_tokens]
‘position_ids’: Reshaped tensor of shape [total_tokens] for 2D input, or [n_rope, 1, total_tokens] for 3D mRoPE input (leading rope axis and a placeholder batch axis of size 1 preserved)
‘cu_seqlens’: Cumulative REAL sequence lengths tensor of shape [num_sequences + 1] (int32) where num_sequences is the total count of non-padded sequences across the batch. Built from seq_lens (the unpadded real lengths). When the trailing pack-pad is purely at the end (cp_size == 1), the last entry is grown to total_tokens to absorb that pad and avoid TE’s pad_between_seqs=True path; see the absorption block in the function body for the gate.
‘cu_seqlens_padded’: (optional) Cumulative PADDED sequence lengths tensor of the same shape as cu_seqlens. Only emitted when it differs from cu_seqlens after absorption (i.e., when padding lives between sub-sequences, which is the CP case). Forwarded to TE as cu_seqlens_q_padded / cu_seqlens_kv_padded with pad_between_seqs=True so the kernel reads memory offsets from the padded variant while attending only over the real-length slots.
‘max_seqlen’: Scalar int32 tensor equal to max(cu_seqlens[i+1] - cu_seqlens[i]) after any absorption. Honors TE’s contract that max_seqlen_q >= max(cu_seqlens_q[i+1] - cu_seqlens_q[i]).
‘padding_mask’: Boolean tensor of shape [total_tokens] indicating padding positions
Non-tensor keys from input batch are preserved (e.g., ‘qkv_format’)

nemo_automodel.components.distributed.thd_utils.split_batch_into_thd_chunks(
    batch: dict[str, torch.Tensor],
    num_chunks: int,
    seq_lens_padding_value: int = -1000,
    padding_token_id: int = 0
) -> dict[str, torch.Tensor]

Process inputs for THD format by splitting batch into chunks for context parallelism.

This function splits the batch along the batch dimension into num_chunks chunks, processes each chunk with process_input_for_thd, and stacks the tensor results. This is useful for context parallelism where different chunks are processed on different devices/ranks.

The cu_seqlens tensors from different chunks may have different lengths depending on the number of sequences in each chunk. These are padded with seq_lens_padding_value to ensure uniform length across chunks for stacking.

Parameters:

batch

dict[str, torch.Tensor]

Dictionary containing input tensors with same structure as process_input_for_thd:

‘input_ids’: [batch_size, seq_len] or [batch_size, seq_len, hidden_dim]
‘labels’: [batch_size, seq_len]
‘position_ids’: [batch_size, seq_len] (required)
‘seq_lens’: [batch_size, num_packs]
‘seq_lens_padded’: [batch_size, num_packs]

num_chunks

int

Number of chunks to split the batch into. Must evenly divide batch_size. If num_chunks <= 1, returns the result from process_input_for_thd directly.

seq_lens_padding_value

intDefaults to -1000

Value used to indicate padding in seq_lens/seq_lens_padded tensors and for padding cu_seqlens to uniform length (default: -1000)

padding_token_id

intDefaults to 0

Filler token id. Only consulted by the metadata-free fallback when a chunk has no pack metadata.

Returns: dict[str, torch.Tensor]

Dictionary containing:

nemo_automodel.components.distributed.thd_utils.thd_padding_mask_from_token_ids(
    input_ids: torch.Tensor,
    padding_token_id: int
) -> torch.Tensor

Mark padding by token value, rejecting a pad id that is also content.

Only for callers with no pack metadata. A value comparison is correct just when the pad id fills a right-padded tail and appears nowhere else, so validate exactly that: a colliding id then fails loudly instead of silently masking content out of the MoE experts. GLM-5.2, for instance, sets pad_token_id to <|endoftext|>, which is also its first eos_token_id.

Parameters:

input_ids

torch.Tensor

Token ids [total_tokens].

padding_token_id

int

Token id used as filler.

Returns: torch.Tensor

Boolean tensor [total_tokens]; True marks padding.

Raises:

ValueError: If padding_token_id also occurs as content.