Integrate Your Own Text Dataset

This guide shows you how to integrate your own dataset into NeMo AutoModel for training. You will learn about two main dataset types: completion datasets for language modeling (such as HellaSwag) and instruction datasets for question-answering tasks (such as SQuAD). This guide covers how to create custom datasets by implementing the required methods and preprocessing functions, and finally shows you how to specify your own data logic using YAML configuration with file paths, allowing you to define custom dataset processing without modifying the main codebase.

Quick Start Summary

Types of Supported Datasets

NeMo AutoModel supports a variety of datasets, depending on the task.

Completion Datasets

Completion datasets are single text sequences designed for language modeling where the model learns to predict the next token given a context. These datasets typically contain a context (prompt) and a target (completion) that the model should learn to generate.

Example: HellaSwag

The HellaSwag dataset is a popular completion dataset used for commonsense reasoning. It contains situations with multiple-choice endings where the model must choose the most plausible continuation.

HellaSwag dataset structure:

• Context (ctx): A situation or scenario description
• Endings: Multiple possible completions (four options)
• Label: Index of the correct ending

Example:

Context: "A man is sitting at a piano in a large room."
Endings: [
  "He starts playing a beautiful melody.",
  "He eats a sandwich while sitting there.",
  "He suddenly becomes invisible.",
  "He transforms into a robot."
]
Label: 0  # First ending is correct

Preprocess with SFTSingleTurnPreprocessor

NeMo AutoModel provides the SFTSingleTurnPreprocessor class to handle completion datasets. This processor:

1. Extracts context and target using get_context() and get_target().
2. Tokenizes and cleans the context and target separately.
3. Concatenates the context and target into one sequence.
4. Creates a loss mask with -100 for the context and target IDs for the target.
5. Pads the sequences to an equal length.

Create Your Own Completion Dataset

To adapt your dataset into this format, define a class such as the following:

1
from datasets import load_dataset
2
from nemo_automodel.components.datasets.utils import SFTSingleTurnPreprocessor
3
4
class MyCompletionDataset:
5
    def __init__(self, path_or_dataset, tokenizer, split="train"):
6
        raw_datasets = load_dataset(path_or_dataset, split=split)
7
        processor = SFTSingleTurnPreprocessor(tokenizer)
8
        self.dataset = processor.process(raw_datasets, self)
9
10
    def get_context(self, examples):
11
        """Extract context/prompt from your dataset"""
12
        return examples["context_field"]  # Replace with your context field
13
14
    def get_target(self, examples):
15
        """Extract target/completion from your dataset"""
16
        return examples["target_field"]   # Replace with your target field
17
18
    def __getitem__(self, index):
19
        return self.dataset[index]
20
21
    def __len__(self):
22
        return len(self.dataset)

Instruction Datasets

Instruction datasets are question-answer pairs where the model learns to respond to specific instructions or questions. These datasets are structured as context-question pairs with corresponding answers, making them ideal for teaching models to follow instructions and provide accurate responses.

Example: SQuAD

The SQuAD (Stanford Question Answering Dataset) is a popular instruction dataset for reading comprehension. It contains questions based on Wikipedia articles along with their answers.

SQuAD dataset structure:

• Context: A paragraph of text from Wikipedia
• Question: A question about the context
• Answers: The correct answer with its position in the context

Create Your Own Instruction Dataset

The squad.py file contains the implementation for processing the SQuAD dataset into a format suitable for instruction tuning. It defines a dataset class and preprocessing functions that extract the context, question, and answer fields, concatenate them into a prompt-completion format, and apply tokenization, padding, and loss masking. This serves as a template for building custom instruction datasets by following a similar structure and adapting the extraction logic to your dataset’s schema.

Based on the SQuAD implementation in squad.py, you can create your own instruction dataset using the make_squad_dataset pattern:

1
from datasets import load_dataset
2
3
def make_my_instruction_dataset(
4
    tokenizer,
5
    seq_length=None,
6
    limit_dataset_samples=None,
7
    split="train",
8
    dataset_name="your-dataset-name",
9
):
10
    if limit_dataset_samples:
11
        split = f"{split}[:{limit_dataset_samples}]"
12
13
    dataset = load_dataset(dataset_name, split=split)
14
15
    return dataset.map(
16
        your_own_fmt_fn,  # Your formatting function
17
        batched=False,
18
        remove_columns=dataset.column_names,
19
    )

YAML-Based Custom Dataset Configuration

NeMo AutoModel supports YAML-based dataset specification using the _target_ key. This allows you to reference dataset-building classes or functions using either of the following methods:

1. Python Dotted Path

1
dataset:
2
  _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag
3
  path_or_dataset: rowan/hellaswag
4
  split: train

2. File Path and Function Name

<file-path>:<function-name>

where:

• <file-path>: The absolute path to a Python file containing your dataset function
• <function-name>: The name of the function to call from that file

1
dataset:
2
  _target_: /path/to/your/custom_dataset.py:build_my_dataset
3
  num_blocks: 111

This calls build_my_dataset() from the specified file with the other keys (such as num_blocks) as arguments. This approach allows you to integrate custom datasets using configuration alone, with no need to alter the codebase or package structure.

Packed Sequence Support in NeMo AutoModel

NeMo AutoModel supports packed sequences, which is a technique to optimize training with variable-length sequences (such as text) by minimizing padding.

What Is a Packed Sequence?

Instead of padding each sequence to a fixed length, which wastes computation on [PAD] tokens, packed sequences perform the following actions:

• Concatenate short sequences into a single continuous sequence.
• Separate sequences with special tokens (such as [EOS]).
• Track lengths using an “attention mask” to prevent cross-sequence information leakage.

Benefits

Packed sequences offer the following benefits:

• They reduce redundant computation on padding tokens, which leads to faster training.
• They enable larger effective batch sizes, which leads to better GPU utilization.
• They are especially useful for language modeling and text datasets.

Enable Packed Sequences in NeMo AutoModel

To enable packed sequences, add these keys to your recipe’s YAML config:

packed_sequence:
   # Set packed_sequence_size > 0 to run with packed sequences
   packed_sequence_size: 1024
   packing_strategy: thd
   # Optional: tokenize the dataset across this many processes before packing
   num_proc: 8

The packed_sequence block accepts the following settings:

• packed_sequence_size: Defines the total token length of each packed sequence. Higher values require more GPU memory.
• packing_strategy: Selects thd (the default sequence-length-aware format) or neat. With thd, AutoModel disables packing when the model does not support seq_lens.
• prepacked: If true, treats the dataset as already packed and skips recipe-side packing.
• max_packs: Optionally caps the number of packed samples created.
• drop_long_samples: For neat, drops samples longer than packed_sequence_size when true. Otherwise, the packer raises an error.
• num_proc: Specifies the number of worker processes used to tokenize the dataset before packing. Packing first runs one full pass over the dataset, which tokenizes every sample serially. Setting num_proc > 1 front-loads that tokenization across processes, producing identical, order-preserving output. This defaults to 1 (serial). The speedup grows with the per-sample tokenization cost (such as long multi-turn chats and heavy chat templates) and available CPU cores. This runs independently on each data-parallel rank, so configure it so that local_ranks * num_proc does not oversubscribe the node’s CPUs.

Troubleshooting Tips

• Tokenization Mismatch? Ensure your tokenizer aligns with the model’s expected inputs.
• Dataset too large? Use limit_dataset_samples in your YAML config to load a subset, which is useful for quick debugging.
• Loss not decreasing? Verify that your loss mask correctly ignores prompt tokens.