Dataset Overview: LLM, VLM, and Retrieval Datasets in NeMo AutoModel

This page summarizes the datasets supported in NeMo AutoModel for LLM, VLM, and retrieval training and shows how to plug in your own datasets using Python functions or the YAML _target_ mechanism.

• See also: LLM datasets, VLM datasets, and Retrieval dataset for deeper, task-specific guides.

• If a dataset you need is missing, please open a GitHub issue with a short description and example schema so we can prioritize support.

LLM Datasets

NeMo AutoModel supports several common patterns for language modeling and instruction tuning.

HellaSwag (Completion SFT)

• Wrapper: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag
• Use case: single-turn completion-style SFT where a prompt (ctx) is followed by a gold continuation (ending)
• Key args: path_or_dataset, split, num_samples_limit
• Example YAML:

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

SQuAD-Style Question Answering (QA) (Instruction SFT)

• Factory: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
• Use case: instruction/QA tuning with either prompt-and-answer formatting or chat-template formatting

• Example YAML:

1
dataset:
2
  _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
3
  split: train
4
  dataset_name: rajpurkar/squad
5
  start_of_turn_token: "<|assistant|>"

• ColumnMappedTextInstructionDataset (generic instruction SFT)
  • Class: nemo_automodel.components.datasets.llm.column_mapped_text_instruction_dataset.ColumnMappedTextInstructionDataset
  • Use case: quickly adapt instruction datasets by mapping your schema’s columns to context, question, answer
  • Sources: local JSON/JSONL or Hugging Face Hub dataset ID
  • Notes:
    • For tokenizers with chat templates and answer-only loss, you may set answer_only_loss_mask: true and provide start_of_turn_token.
    • Supports streaming mode for large datasets (see Streaming Datasets section below).
    • Map-style, non-streaming dataset (supports len(ds) and ds[i])
    • For streaming (including Delta Lake / Databricks), use ColumnMappedTextInstructionIterableDataset
  • Example YAML:

1
dataset:
2
  _target_: nemo_automodel.components.datasets.llm.column_mapped_text_instruction_dataset.ColumnMappedTextInstructionDataset
3
  path_or_dataset_id: Muennighoff/natural-instructions
4
  split: train
5
  column_mapping:
6
    context: definition
7
    question: inputs
8
    answer: targets
9
  answer_only_loss_mask: true
10
  start_of_turn_token: "<|assistant|>"

See the detailed guide, Column-Mapped Text Instruction Dataset, for more information.

• ChatDataset (multi-turn conversations and tool calling)
  • Class: nemo_automodel.components.datasets.llm.ChatDataset
  • Use case: multi-turn conversations and tool calling in OpenAI chat format
  • Sources: local JSON/JSONL or Hugging Face Hub dataset ID
  • Key args:
    • path_or_dataset_id: path to local file(s) or HuggingFace dataset ID
    • tokenizer: tokenizer instance (required. Must have chat template support)
    • split: dataset split (e.g., “train”, “validation”)
    • name: dataset configuration/subset name
    • seq_length: maximum sequence length for padding/truncation
    • padding: padding strategy (“do_not_pad”, “max_length”, etc.)
    • truncation: truncation strategy (“do_not_truncate”, “longest_first”, etc.)
    • start_of_turn_token: token marking assistant response start (for answer-only loss)
    • chat_template: optional override for tokenizer’s chat template
    • skip_invalid_samples: if true, skip malformed JSONL lines when reading local files (warnings log skip counts); default false fails fast on a bad line
  • Notes:
    • Requires a tokenizer with chat template support
    • Supports both single-turn and multi-turn tool calling
    • Tool definitions are provided in a tools field at the conversation level
    • Tool calls appear in assistant messages via tool_calls field
    • Tool responses use the tool role

ChatDataset (Multi-Turn Conversations and Tool Calling)

• Class: nemo_automodel.components.datasets.llm.ChatDataset
• Use case: multi-turn conversations and tool calling in OpenAI chat format
• Sources: local JSON/JSONL or Hugging Face Hub dataset ID
• Key args:
  • path_or_dataset_id: path to local file(s) or Hugging Face dataset ID
  • tokenizer: tokenizer instance (required; must have chat template support)
  • split: dataset split (e.g., “train”, “validation”)
  • name: dataset configuration/subset name
  • seq_length: maximum sequence length for padding/truncation
  • padding: padding strategy (“do_not_pad”, “max_length”, etc.)
  • truncation: truncation strategy (“do_not_truncate”, “longest_first”, etc.)
  • start_of_turn_token: token marking assistant response start (for answer-only loss)
  • chat_template: optional override for tokenizer’s chat template
  • mask_reasoning_content: optionally exclude rendered reasoning_content tokens from loss
  • skip_invalid_samples: if true, skip malformed JSONL lines when reading local files (warnings log skip counts); default false fails fast on a bad line

• Example YAML:

1
dataset:
2
  _target_: nemo_automodel.components.datasets.llm.ChatDataset
3
  path_or_dataset_id: Salesforce/xlam-function-calling-60k
4
  split: train
5
  tokenizer:
6
    _target_: transformers.AutoTokenizer.from_pretrained
7
    pretrained_model_name_or_path: google/functiongemma-270m-it
8
  seq_length: 2048
9
  start_of_turn_token: "<start_of_turn>"
10
  mask_reasoning_content: false
11
  skip_invalid_samples: false

• Expected data format (OpenAI messages format):

1
{
2
  "messages": [
3
    {
4
      "role": "system",
5
      "content": "You are a helpful assistant."
6
    },
7
    {
8
      "role": "user",
9
      "content": "What's the weather in Seattle and should I bring an umbrella?"
10
    },
11
    {
12
      "role": "assistant",
13
      "reasoning_content": "The user wants weather info and advice. I should call get_weather first, then decide whether an umbrella is needed.",
14
      "content": "",
15
      "tool_calls": [
16
        {
17
          "id": "call_1",
18
          "type": "function",
19
          "function": {
20
            "name": "get_weather",
21
            "arguments": "{\"city\": \"Seattle\"}"
22
          }
23
        }
24
      ]
25
    },
26
    {
27
      "role": "tool",
28
      "tool_call_id": "call_1",
29
      "content": "{\"temperature\": 55, \"condition\": \"rain\", \"precipitation_chance\": 0.85}"
30
    },
31
    {
32
      "role": "assistant",
33
      "reasoning_content": "It is raining with a high precipitation chance, so I should recommend bringing an umbrella.",
34
      "content": "It's currently 55 degrees F and raining in Seattle with an 85% chance of continued precipitation. Yes, definitely bring an umbrella."
35
    }
36
  ],
37
  "tools": [
38
    {
39
      "type": "function",
40
      "function": {
41
        "name": "get_weather",
42
        "description": "Get current weather for a city",
43
        "parameters": {
44
          "type": "object",
45
          "properties": {
46
            "city": {"type": "string"}
47
          },
48
          "required": ["city"]
49
        }
50
      }
51
    }
52
  ]
53
}

• Template requirement example for reasoning_content:

1
{%- if message.reasoning_content %}
2
{% generation %}
3
{{ "<think>\n" + message.reasoning_content + "\n</think>\n" }}
4
{% endgeneration %}
5
{%- endif %}
6
{% generation %}
7
{{ message.content }}
8
{% endgeneration %}

• For single-turn tool calling (one tool call per conversation), omit the tool response and final assistant message:

1
{
2
  "messages": [
3
    {
4
      "role": "user",
5
      "content": "Book a table for two at 7pm in Seattle."
6
    },
7
    {
8
      "role": "assistant",
9
      "content": "",
10
      "tool_calls": [
11
        {
12
          "id": "call_1",
13
          "type": "function",
14
          "function": {
15
            "name": "book_table",
16
            "arguments": "{\"party_size\": 2, \"time\": \"19:00\", \"city\": \"Seattle\"}"
17
          }
18
        }
19
      ]
20
    }
21
  ],
22
  "tools": [
23
    {
24
      "type": "function",
25
      "function": {
26
        "name": "book_table",
27
        "description": "Book a restaurant table",
28
        "parameters": {
29
          "type": "object",
30
          "properties": {
31
            "party_size": {"type": "integer"},
32
            "time": {"type": "string"},
33
            "city": {"type": "string"}
34
          }
35
        }
36
      }
37
    }
38
  ]
39
}

See the Function Calling guide for an end-to-end example with FunctionGemma. For a small reasoning-style chat SFT starting point, see qwen2_5_0p5b_instruct_fineproofs_chat.yaml.

Retrieval (Embedding Fine-Tuning)

• Factory: nemo_automodel.components.datasets.llm.make_retrieval_dataset
• Collator: nemo_automodel.components.datasets.llm.BiEncoderCollator
• Use case: embedding model fine-tuning with (query, positive doc, negative docs) contrastive learning
• Supported schemas:
  • Corpus-ID JSON (Merlin/NeMo-retriever style)
  • Inline-text JSONL (e.g., {"query": "...", "pos_doc": "...", "neg_doc": ["...", "..."]})
• Example YAML:

1
dataset:
2
  _target_: nemo_automodel.components.datasets.llm.make_retrieval_dataset
3
  data_dir_list: /abs/path/to/train.jsonl
4
  data_type: train
5
  n_passages: 5
6
collate_fn:
7
  _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator
8
  q_max_len: 512
9
  p_max_len: 512

See the detailed guide, Retrieval dataset, for more information.

NanoGPT Binary Shards (Pretraining)

• Class: nemo_automodel.components.datasets.llm.nanogpt_dataset.NanogptDataset
• Use case: token-level LM pretraining over .bin shards produced by NanoGPT-style preprocessors (supports legacy and current formats)

Megatron (Pretraining; Interoperable With Pre-Tokenized Megatron Data)

• Class: nemo_automodel.components.datasets.llm.megatron_dataset.MegatronPretraining
• Use case: large-scale LM pretraining over Megatron-LM formatted tokenized corpora
• Interoperability: If your corpus has already been tokenized/indexed for Megatron (i.e., .bin/.idx pairs), you can point AutoModel to those assets directly. No re-tokenization required.
• Key args: paths (single path, glob, weighted list, or per-split dict), seq_length, tokenizer, split, index_mapping_dir, splits_to_build
• Example YAML:

1
dataset:
2
  _target_: nemo_automodel.components.datasets.llm.megatron_dataset.MegatronPretraining
3
  paths: /abs/path/to/processed_data_*_text_document*  # glob or explicit list
4
  index_mapping_dir: /abs/path/to/mapping_dir
5
  tokenizer:
6
    _target_: transformers.AutoTokenizer.from_pretrained
7
    pretrained_model_name_or_path: openai-community/gpt2
8
  seq_length: 1024
9
  split: "0.99, 0.01, 0.00"  # train, validation, test
10
  splits_to_build: "train"

See the detailed pretraining guide, which uses MegatronPretraining data.

Streaming Datasets

Streaming datasets enable processing very large datasets without loading them entirely into memory. This is particularly useful when working with datasets that exceed available RAM or when you want to start training immediately without waiting for the full dataset to download.

What Are Streaming Datasets?

Streaming datasets load and process data incrementally, one batch at a time, rather than loading the entire dataset into memory upfront. This approach:

• Reduces memory footprint: Only the current batch resides in memory
• Enables training on massive datasets: Process terabyte-scale datasets on machines with limited RAM
• Faster startup: Begin training immediately without waiting for full dataset download
• Efficient for remote datasets: Stream directly from Hugging Face Hub without local storage

When to Use Streaming

Use streaming mode when:

• Your dataset is very large (hundreds of GB or TB)
• Available memory is limited compared to dataset size
• You want to start training quickly without downloading the full dataset
• You’re experimenting with a subset of a large dataset

Avoid streaming when:

• Your dataset is small enough to fit comfortably in memory
• You need random access to samples (e.g., for certain sampling strategies)
• You need to know the exact dataset length upfront
• Training requires multiple passes with different orderings

How to Enable Streaming

For ColumnMappedTextInstructionDataset, use the streaming variant by changing the class to ColumnMappedTextInstructionIterableDataset:

1
dataset:
2
  _target_: nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset
3
  path_or_dataset_id: Muennighoff/natural-instructions
4
  split: train
5
  column_mapping:
6
    context: definition
7
    question: inputs
8
    answer: targets
9
  answer_only_loss_mask: true
10
  start_of_turn_token: "<|assistant|>"

For Hugging Face datasets loaded directly, set streaming=True:

1
from datasets import load_dataset
2
3
# Non-streaming (loads entire dataset into memory)
4
dataset = load_dataset("large-dataset/corpus", split="train", streaming=False)
5
6
# Streaming (loads data incrementally)
7
dataset = load_dataset("large-dataset/corpus", split="train", streaming=True)

Streaming Limitations

When using streaming datasets, be aware of these limitations:

1. No random access: You cannot use dataset[index] to access specific samples. Streaming datasets only support iteration.

2. No length information: The len(dataset) operation is not available. You cannot determine the total number of samples upfront.

3. Single-pass iteration: Each iteration consumes the stream. To iterate multiple times, you need to recreate the dataset or use the repeat_on_exhaustion parameter.

4. Limited shuffling: Shuffling is done with a buffer (not the entire dataset), which may not provide perfect randomization.

Distributed Training with Streaming

Streaming datasets support distributed training through sharding:

1
from nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset import (
2
    ColumnMappedTextInstructionIterableDataset
3
)
4
5
dataset = ColumnMappedTextInstructionIterableDataset(
6
    path_or_dataset_id="large-dataset/corpus",
7
    column_mapping={"question": "input", "answer": "output"},
8
    tokenizer=tokenizer,
9
)
10
11
# Shard the dataset across workers
12
dataset = dataset.shard(num_shards=8, index=worker_id)
13
14
# Enable shuffling with a buffer
15
dataset = dataset.shuffle(buffer_size=10000, seed=42)
16
17
# Set epoch for deterministic shuffling across epochs
18
dataset.set_epoch(epoch_num)

Performance Considerations

Memory vs. Speed Trade-offs:

• Streaming reduces memory usage but may be slower than in-memory datasets
• Network latency can impact streaming performance for remote datasets
• Use local caching when repeatedly accessing the same remote dataset

Buffer Size for Shuffling:

• Larger buffers provide better randomization but use more memory
• A buffer size of 10,000-100,000 samples is typically a good balance
• For perfect shuffling, you need a buffer size equal to the dataset size (defeating the purpose of streaming)

Prefetching:

• Most streaming implementations prefetch data in the background
• This helps hide network latency and keeps GPUs busy
• Adjust prefetch settings based on your network speed and batch size

Example: Streaming a Large Dataset

Here’s a complete example of using streaming for a large instruction-tuning dataset:

1
dataset:
2
  _target_: nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset
3
  path_or_dataset_id: HuggingFaceH4/ultrachat_200k
4
  split: train_sft
5
  column_mapping:
6
    question: prompt
7
    answer: completion
8
  answer_only_loss_mask: true
9
  start_of_turn_token: "<|assistant|>"
10
  repeat_on_exhaustion: true  # Automatically restart when stream ends
11
12
dataloader:
13
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
14
  batch_size: 4
15
  num_workers: 4

This configuration:

• Streams the dataset without loading it fully into memory
• Automatically repeats when the stream is exhausted
• Uses multiple workers for efficient data loading
• Applies answer-only loss masking during tokenization

Packed Sequence Support

To reduce padding and improve throughput with variable-length sequences:

1
packed_sequence:
2
  packed_sequence_size: 8192   # > 0 enables packing
3
  split_across_pack: false

Use a collator that pads to an FP8-friendly multiple when training with FP8:

1
dataloader:
2
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
3
  collate_fn:
4
    _target_: nemo_automodel.components.datasets.utils.default_collater
5
    pad_seq_len_divisible: 16

VLM Datasets (Vision/Audio + Language)

VLM datasets are represented as conversations (message lists) that combine text with images or audio and are processed with the model’s AutoProcessor.apply_chat_template and a suitable collate function.

Built-in dataset makers (return lists of conversation dicts):

• RDR items: nemo_automodel.components.datasets.vlm.datasets.make_rdr_dataset (HF: quintend/rdr-items)
• CORD-V2 receipts (Consolidated Receipt Dataset for Post-OCR Parsing): nemo_automodel.components.datasets.vlm.datasets.make_cord_v2_dataset (HF: naver-clova-ix/cord-v2)
• MedPix-VQA (Medical Pixel Question Answering): nemo_automodel.components.datasets.vlm.datasets.make_medpix_dataset
• CommonVoice 17 (CV17) (audio): nemo_automodel.components.datasets.vlm.datasets.make_cv17_dataset

Each example follows the conversation schema expected by apply_chat_template, e.g.:

1
{
2
  "conversation": [
3
    {
4
      "role": "user",
5
      "content": [
6
        {"type": "image", "image": example_image},
7
        {"type": "text",  "text":  "Describe this image."}
8
      ]
9
    },
10
    {
11
      "role": "assistant",
12
      "content": [{"type": "text", "text": ground_truth_text}]
13
    }
14
  ]
15
}

Custom Chat Template

By default, VLM fine-tuning uses the chat template built into the model’s AutoProcessor. To override it, add chat_template under dataset: in your YAML config:

1
dataset:
2
  _target_: nemo_automodel.components.datasets.vlm.datasets.make_medpix_dataset
3
  split: train
4
  chat_template: "{% for msg in messages %}{{ msg.role }}: {{ msg.content }}\n{% endfor %}"

chat_template accepts a Jinja template string, a path to a .jinja file, or a path to a JSON file containing a chat_template key. The override is applied to both the processor and its tokenizer before dataset instantiation.

Collate Functions

• nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
• nemo_automodel.components.datasets.vlm.collate_fns.qwen2_5_collate_fn (Qwen2.5 VL)
• nemo_automodel.components.datasets.vlm.collate_fns.phi4_mm_collate_fn (audio)

Select in your YAML:

1
dataloader:
2
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
3
  batch_size: 1
4
  collate_fn:
5
    _target_: nemo_automodel.components.datasets.vlm.collate_fns.qwen2_5_collate_fn

If you want answer-only loss masking, provide a model-appropriate start_of_response_token to the collate function.

See Gemma-3n and VLM dataset for end-to-end examples.

Diffusion Datasets

Diffusion models don’t train directly on raw images or videos. Instead, the data is first encoded into a compact numerical representation called a latent — this is what the model actually learns from. Text captions are similarly converted into text embeddings that the model uses as conditioning.

This encoding is done once during preprocessing, and the results are saved as cache files (.meta). Training then reads these cache files directly, which is significantly faster than re-encoding on every step.

The built-in preprocessing tool (tools/diffusion/preprocessing_multiprocess.py) handles this conversion. It uses a VAE (Variational Autoencoder) to encode visual data and a text encoder for captions, grouping outputs into resolution-bucketed directories compatible with the multiresolution dataloader.

Dataloader Builders

• Video (T2V): nemo_automodel.components.datasets.diffusion.build_video_multiresolution_dataloader — for Wan 2.1 and HunyuanVideo
• Image (T2I): nemo_automodel.components.datasets.diffusion.build_text_to_image_multiresolution_dataloader — for FLUX.1-dev

Example YAML (Video Dataloader)

1
data:
2
  dataloader:
3
    _target_: nemo_automodel.components.datasets.diffusion.build_video_multiresolution_dataloader
4
    cache_dir: /path/to/processed_meta
5
    model_type: wan
6
    base_resolution: [512, 512]
7
    dynamic_batch_size: false
8
    shuffle: true
9
    drop_last: false
10
    num_workers: 0

See the Diffusion Dataset Preparation guide for full preprocessing instructions and configuration details.

Bring Your Own Dataset

You can integrate custom datasets with zero code changes to NeMo AutoModel by using _target_ in YAML. There are three approaches:

Point to an Existing Class or Function (Dotted Path)

• LLM example (class):

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

• LLM example (factory function):

1
dataset:
2
  _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
3
  split: train
4
  dataset_name: rajpurkar/squad

• VLM example (factory function):

1
dataset:
2
  _target_: nemo_automodel.components.datasets.vlm.datasets.make_cord_v2_dataset
3
  split: train

Point to a Local Python File and Function

1
dataset:
2
  _target_: /abs/path/to/my_custom_dataset.py:build_my_dataset
3
  some_arg: 123
4
  split: train

Where build_my_dataset returns either a datasets.Dataset or a list/iterator of conversation dicts (for VLM).

Use ColumnMappedTextInstructionDataset for Most Instruction Datasets (LLM)

• Ideal when your data has columns like instruction, input, or output but with arbitrary names
• Supports local JSON/JSONL and HF Hub

1
dataset:
2
  _target_: nemo_automodel.components.datasets.llm.column_mapped_text_instruction_dataset.ColumnMappedTextInstructionDataset
3
  path_or_dataset_id: /abs/path/to/*.jsonl  # or org/repo on HF
4
  column_mapping:
5
    context: definition
6
    question: inputs
7
    answer: targets
8
  answer_only_loss_mask: true
9
  start_of_turn_token: "<|assistant|>"

Implement a Minimal Custom Class Pattern (LLM Completion)

If you prefer Python, implement get_context and get_target and reuse the built-in preprocessor:

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_ds = load_dataset(path_or_dataset, split=split)
7
        self.dataset = SFTSingleTurnPreprocessor(tokenizer).process(raw_ds, self)
8
9
    def get_context(self, examples):
10
        return examples["my_context_field"]
11
12
    def get_target(self, examples):
13
        return examples["my_target_field"]

Then reference your class with _target_ in YAML.

Important Considerations

• Chat templates: If your tokenizer has a chat template and you want answer-only loss, provide the correct start_of_turn_token (LLM) or start_of_response_token (VLM collate functions).
• Padding for FP8: If training with FP8, set pad_seq_len_divisible: 16 in your collate function to align sequence lengths.
• Packed sequences: Prefer packed sequences for throughput when fine-tuning LLMs on variable-length corpora.
• Validation: You can define a separate validation_dataset and validation_dataloader block mirroring your training config.

For detailed, end-to-end recipes, browse the example configs under examples/llm_finetune/, examples/llm_pretrain/, and examples/vlm_finetune/.