> This page is for version 0.5.0 · 26.06. > For other versions, use one of these documentation indexes: > - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt > - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt > - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt > - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt > For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server. # nemo_automodel.components.distributed.thd_utils ## Module Contents ### Functions | Name | Description | | --------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------- | | [`_thd_padding_mask`](#nemo_automodel-components-distributed-thd_utils-_thd_padding_mask) | Mark padding slots in a packed THD stream from the pack layout. | | [`process_input_for_thd`](#nemo_automodel-components-distributed-thd_utils-process_input_for_thd) | Process inputs for THD (total, hidden, depth) format. | | [`split_batch_into_thd_chunks`](#nemo_automodel-components-distributed-thd_utils-split_batch_into_thd_chunks) | Process inputs for THD format by splitting batch into chunks for context parallelism. | | [`thd_padding_mask_from_token_ids`](#nemo_automodel-components-distributed-thd_utils-thd_padding_mask_from_token_ids) | Mark padding by token value, rejecting a pad id that is also content. | ### API ```python 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 `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. ```python 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. **Example** \>>> batch\_size, seq\_len = 2, 6 \>>> # 2D Token IDs case with packed sequences \>>> batch = \{ ... 'input\_ids': torch.tensor(\[\[1, 2, 3, 99, 4, 5], \[6, 7, 8, 9, 10, 11]]), ... 'labels': torch.tensor(\[\[2, 3, 99, 4, 5, 6], \[7, 8, 9, 10, 11, 12]]), ... 'position\_ids': torch.tensor(\[\[0, 1, 2, 0, 0, 1], \[0, 1, 2, 3, 4, 5]]), ... 'seq\_lens': torch.tensor(\[\[3, 2], \[6, -1000]]), # Second batch has only 1 sequence ... 'seq\_lens\_padded': torch.tensor(\[\[4, 2], \[6, -1000]]) ... } \>>> \>>> result = process\_input\_for\_thd(batch) \>>> # result\['input\_ids'].shape: \[12] (2D input collapsed to 1D) \>>> # result\['labels'].shape: \[12] \>>> # result\['position\_ids'].shape: \[12] \>>> # result\['cu\_seqlens']: tensor(\[0, 3, 5, 11], dtype=torch.int32) \>>> # Breakdown: \[0] + cumsum(\[3, 2, 6]) = \[0, 3, 5, 11] (from seq\_lens — real lengths) \>>> # result\['cu\_seqlens\_padded']: tensor(\[0, 4, 6, 12], dtype=torch.int32) \>>> # Breakdown: \[0] + cumsum(\[4, 2, 6]) = \[0, 4, 6, 12] (from seq\_lens\_padded) \>>> # result\['max\_seqlen']: tensor(6, dtype=torch.int32) # max slot width in cu\_seqlens \>>> # result\['padding\_mask'].shape: \[12] **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`** `int` — default: -1000 Value used to indicate padding in seq\_lens/seq\_lens\_padded tensors that should be filtered out (default: -1000) --- **`padding_token_id`** `int` — default: 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': Cumulative physical sequence boundaries, retained whenever seq\_lens\_padded is provided, including when equal to cu\_seqlens. * 'pad\_between\_seqs': Whether the complete physical and attention boundary arrays differ after trailing-pad absorption. This controls TE dispatch independently of the presence of physical metadata needed by other consumers such as MTP. * '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') ```python 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. **Example** \>>> batch\_size, seq\_len = 4, 6 \>>> batch = \{ ... 'input\_ids': torch.tensor(\[\[1,2,3,4,5,6], \[7,8,9,10,11,12], ... \[13,14,15,16,17,18], \[19,20,21,22,23,24]]), ... 'labels': torch.tensor(\[\[2,3,4,5,6,7], \[8,9,10,11,12,13], ... \[14,15,16,17,18,19], \[20,21,22,23,24,25]]), ... 'position\_ids': torch.tensor(\[\[0,1,2,3,4,5], \[0,1,2,3,4,5], ... \[0,1,2,3,4,5], \[0,1,2,3,4,5]]), ... 'seq\_lens': torch.tensor(\[\[6], \[6], \[6], \[6]]), ... 'seq\_lens\_padded': torch.tensor(\[\[6], \[6], \[6], \[6]]), ... } \>>> \>>> result = split\_batch\_into\_thd\_chunks(batch, num\_chunks=2) \>>> # result\['input\_ids'].shape: \[2, 12] (2 chunks, each with 2\*6=12 tokens) \>>> # result\['cu\_seqlens'].shape: \[2, 3] (2 chunks, each with \[0, 6, 12]) \>>> # result\['cu\_seqlens']\[0]: tensor(\[0, 6, 12], dtype=torch.int32) \>>> # result\['cu\_seqlens']\[1]: tensor(\[0, 6, 12], dtype=torch.int32) **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`** `int` — default: -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`** `int` — default: 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: ```python 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.