nemo_automodel.components.datasets.utils

Module Contents

Classes

SFTSingleTurnPreprocessor

Generic single-turn text-to-text SFT (supervised-fine-tuning) pre-processor.

Functions

batchify	Ensures that the input tensor has at least two dimensions by adding an extra batch dimension if necessary.
extract_key_from_dicts	Extracts the value of the given key from each dictionary in a list of dictionaries.
pad_within_micro	Pads each list in a batch of lists to the same length with a specified token.
find_last_non_pad_token
get_pad_token_from_key
make_attention_mask_from_labels
create_causal_mask_mapping	Create causal mask mapping for pipeline parallelism.
add_causal_masks_to_batch	Add precomputed causal masks to an already-batched data dict.
default_collater	Default batch collator that handles padding and batching.
packed_sequence_thd_collater	Collater for packed sequences in THD (total, hidden, depth) format.

API

nemo_automodel.components.datasets.utils.batchify(tensor, default_tensor_cls=torch.LongTensor)

Ensures that the input tensor has at least two dimensions by adding an extra batch dimension if necessary.

Parameters:

tensor (torch.Tensor) – The input tensor to be batchified.

Returns:

The tensor with an extra dimension added if it was originally 1-dimensional. Otherwise, the tensor is returned as-is.

Return type:

torch.Tensor

nemo_automodel.components.datasets.utils.extract_key_from_dicts(batch, key)

Extracts the value of the given key from each dictionary in a list of dictionaries.

Parameters:

• batch (List[dict]) – A list of dictionaries.

• key (str) – The key whose values are to be extracted from each dictionary.

Returns:

A list of values associated with the specified key, in the same order as the dictionaries in the input batch.

Return type:

List

nemo_automodel.components.datasets.utils.pad_within_micro(batch, pad_token_id, pad_seq_len_divisible=None)

Pads each list in a batch of lists to the same length with a specified token.

Parameters:

• batch (List[List[int]]) – A batch of sequences (e.g., token IDs), where each sequence is a list of integers.

• pad_token_id (int) – The token ID to use for padding shorter sequences.

• pad_seq_len_divisible (int) – The value to use for padding sequence length so that it is divisible by pad_seq_len_divisible.

Returns:

A batch of sequences where each inner list has been padded with the pad token to match the length of the longest sequence in the batch.

Return type:

List[List[int]]

nemo_automodel.components.datasets.utils.find_last_non_pad_token(lst: list[int], value: int) → int | None

nemo_automodel.components.datasets.utils.get_pad_token_from_key(

val: str,

pad_token_ids: Optional[dict[str, int]] = None,

) → int | None

nemo_automodel.components.datasets.utils.make_attention_mask_from_labels(

ids: list[int],

ignore_token: int = -100,

) → list[int]

nemo_automodel.components.datasets.utils.create_causal_mask_mapping(

model_config,

batch_size,

seq_len,

position_ids=None,

attention_mask=None,

device=None,

)

Create causal mask mapping for pipeline parallelism.

This is the core mask creation logic that can be reused by different collate functions. Extracts common mask creation logic to avoid duplication between collate functions.

Parameters:

• model_config – HuggingFace model config

• batch_size – Batch size

• seq_len – Sequence length

• position_ids – Optional position IDs tensor [batch_size, seq_len]

• attention_mask – Optional 2D attention mask tensor [batch_size, seq_len] for padding

• device – Device to create tensors on (defaults to cpu)

Returns:

Mapping of mask types to 4D mask tensors - “full_attention”: [batch_size, 1, seq_len, seq_len] - “sliding_attention”: [batch_size, 1, seq_len, seq_len] (if model uses sliding window)

Return type:

dict

nemo_automodel.components.datasets.utils.add_causal_masks_to_batch(batch_dict, model_config)

Add precomputed causal masks to an already-batched data dict.

This function is designed for datasets that yield complete batches (like MockIterableDataset), where we want to add mask precomputation as a separate processing step.

Parameters:

• batch –

  A dict or list containing a single batched dict with tensors:

  • input_ids: [batch_size, seq_length]

  • position_ids: [batch_size, seq_length] (optional)

  • labels: [batch_size, seq_length]

• model_config – HuggingFace model config for creating causal masks

• precompute_masks – If False, skip mask creation (for compatibility with train_ft.py wrapper)

Returns:

Same batch with added causal_mask_mapping field

Return type:

dict

nemo_automodel.components.datasets.utils.default_collater(batch, pad_seq_len_divisible=None)

Default batch collator that handles padding and batching.

Parameters:

• batch – A batch of examples.

• pad_seq_len_divisible – If provided, pad sequence length to be divisible by this value.

Returns:

A dictionary containing batched tensors.

Return type:

dict

nemo_automodel.components.datasets.utils.packed_sequence_thd_collater(batch)

Collater for packed sequences in THD (total, hidden, depth) format.

This collater is designed for THD format, where multiple variable-length sequences are concatenated with/without padding tokens between them. The THD format represents sequences as (total_tokens, hidden_dim, depth) where total_tokens is the sum of all sequence lengths in the batch.

Unlike traditional padding-based approaches (BSHD/SBHD formats), this THD format:

• Concatenates sequences directly: [a a a b b c c c c]

• Uses seq_lens to identify sequence boundaries for attention computation

• Supports optional identifier or padding tokens between sequences via seq_lens_padded

This collater supports both pipeline parallelism (PP) and non-PP use cases by:

• Stacking token-level tensors (input_ids, labels, position_ids) along batch dimension

• Padding and stacking seq_lens and seq_lens_padded with sentinel value -1000

• Including ‘qkv_format’: ‘thd’ in the output to indicate THD format

IMPORTANT: All examples in the batch must have the same token sequence length for input_ids, labels, and position_ids. This is typically ensured by the dataset/packing logic that creates fixed-length packed sequences.

Parameters:

batch (List[dict]) –

A list of dictionaries, where each dictionary represents one packed example. Each dictionary should contain:

• ’input_ids’: List[int] - Token IDs for all packed sequences (must be same length across batch)

• ’labels’: List[int] - Labels for all packed sequences (must be same length across batch)

• ’position_ids’: List[int] - Position IDs for all tokens (must be same length across batch)

• ’seq_lens’: List[int] - Actual sequence lengths for each packed sequence

• ’seq_lens_padded’: List[int] - Sequence lengths including identifier/padding tokens

Example batch with 2 examples, both with 6 total tokens: [ { ‘input_ids’: [1, 2, 3, 99, 4, 5], # Two sequences: [1,2,3] and [4,5] with sep token 99 ‘labels’: [1, 2, 3, -100, 4, 5], ‘position_ids’: [0, 1, 2, 0, 0, 1], ‘seq_lens’: [3, 2], # Actual sequence lengths (excluding separator) ‘seq_lens_padded’: [4, 2] # Including separator token }, { ‘input_ids’: [6, 7, 99, 8, 9, 10], # Two sequences with separator ‘labels’: [6, 7, -100, 8, 9, 10], ‘position_ids’: [0, 1, 0, 0, 1, 2], ‘seq_lens’: [2, 3], ‘seq_lens_padded’: [3, 3] } ]

Returns:

A dictionary with batched tensors: - ‘input_ids’: tensor of shape [batch_size, seq_len] - stacked token sequences - ‘labels’: tensor of shape [batch_size, seq_len] - stacked labels - ‘position_ids’: tensor of shape [batch_size, seq_len] - stacked position IDs - ‘seq_lens’: tensor of shape [batch_size, max_num_packs] - padded sequence lengths - ‘seq_lens_padded’: tensor of shape [batch_size, max_num_packs] - padded lengths with separators - ‘qkv_format’: str - Always ‘thd’ to indicate THD format

Note: seq_lens and seq_lens_padded are padded with -1000 to handle variable number of packed sequences per example. These sentinel values should be filtered out before use.

Return type:

dict

class nemo_automodel.components.datasets.utils.SFTSingleTurnPreprocessor(tokenizer)

Generic single-turn text-to-text SFT (supervised-fine-tuning) pre-processor.

Parameters:

tokenizer – Pre-trained tokenizer (HF).

Initialization

SFTSingleTurnPreprocessor constructor.

Parameters:

tokenizer – Pretrained tokenizer.

_tokenize_function(examples, dataset)

_compute_dataset_max_len(tokenized_ds)

_pad_function(max_len)

process(raw_dataset, ds)

Main processor entry.

Parameters:

• raw_dataset (datasets.DatasetDict) – the dataset (e.g. returned by load_dataset)

• ds (dataset) – the dataset with get_target method.

Returns:

tokenized + optionally padded datasets (all splits preserved).

Return type:

datasets.DatasetDict