Use the ColumnMappedTextInstructionDataset

This guide explains how to use ColumnMappedTextInstructionDataset to quickly and flexibly load instruction-answer datasets for LLM fine-tuning, with minimal code changes and support for common tokenization strategies.

The ColumnMappedTextInstructionDataset is a lightweight, plug-and-play helper that lets you train on instruction-answer style corpora without writing custom Python for every new schema. You simply specify which columns map to logical fields like context, question, and answer, and the loader handles the rest automatically. This enables:

• Quick prototyping across diverse instruction datasets
• Schema flexibility without requiring code changes
• Consistent field names for training loops, regardless of dataset source

ColumnMappedTextInstructionDataset is a map-style dataset (torch.utils.data.Dataset): it supports len(ds) and ds[i], and it loads data non-streaming.

It supports two data sources out-of-the-box:

1. Local JSON/JSONL files - pass a single file path or a list of paths on disk. Newline-delimited JSON works great.
2. Hugging Face Hub - point to any dataset repo (org/dataset) that contains the required columns.

For streaming (including Delta Lake / Databricks), use ColumnMappedTextInstructionIterableDataset. The iterable variant always streams by design to avoid accidentally materializing entire datasets to disk/memory.

Quickstart

The fastest way to sanity-check the loader is to point it at an existing Hugging Face dataset and print the first sample. This section provides a minimal, runnable example to help you quickly try out the dataset.

1
from transformers import AutoTokenizer
2
from nemo_automodel.components.datasets.llm.column_mapped_text_instruction_dataset import ColumnMappedTextInstructionDataset
3
4
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")
5
6
ds = ColumnMappedTextInstructionDataset(
7
    path_or_dataset_id="Muennighoff/natural-instructions",
8
    column_mapping={
9
      "context": "definition",
10
      "question": "inputs",
11
      "answer": "targets"
12
    },
13
    tokenizer=tokenizer,
14
    answer_only_loss_mask=True,
15
)
16
17
sample = ds[0]
18
print(sample.keys())
19
20
# Typical keys include: input_ids, labels, attention_mask (and an internal ___PAD_TOKEN_IDS___ helper).
21
# Note: when answer_only_loss_mask=True, prompt tokens are masked in labels with -100
22
# (the standard CrossEntropy "ignore_index").

The code above is intended only for a quick sanity check of the dataset and its tokenization output. For training or production use, configure the dataset using YAML as shown below. YAML offers a reproducible, maintainable, and scalable way to specify dataset and tokenization settings.

Usage Examples

This section provides practical usage examples, including how to load remote datasets, work with local files, and configure pipelines using YAML recipes.

Local JSONL Example

Assume you have a local newline-delimited JSON file at /data/my_corpus.jsonl with the simple schema {instruction, output}. A few sample rows:

1
{"instruction": "Translate 'Hello' to French", "output": "Bonjour"}
2
{"instruction": "Summarize the planet Neptune.", "output": "Neptune is the eighth planet from the Sun."}

You can load it using Python code like:

1
local_ds = ColumnMappedTextInstructionDataset(
2
    path_or_dataset_id=["/data/my_corpus_1.jsonl", "/data/my_corpus_2.jsonl"], # can also be a single path (string)
3
    column_mapping={
4
        "question": "instruction",
5
        "answer": "output",
6
    },
7
    tokenizer=tokenizer,
8
    answer_only_loss_mask=False,  # compute loss over full sequence
9
)
10
11
print(local_ds[0].keys())   # dict_keys(['input_ids', 'labels', 'attention_mask', '___PAD_TOKEN_IDS___'])

You can configure the dataset entirely from your recipe YAML. For example:

1
dataset:
2
  _target_: nemo_automodel.components.datasets.llm.column_mapped_text_instruction_dataset.ColumnMappedTextInstructionDataset
3
  path_or_dataset_id:
4
    - /data/my_corpus_1.jsonl
5
    - /data/my_corpus_2.jsonl
6
  column_mapping:
7
    question: instruction
8
    answer: output
9
  answer_only_loss_mask: false

Remote Dataset Example

In the following section, we demonstrate how to load the instruction-tuning corpus Muennighoff/natural-instructions. The dataset schema is {task_name, id, definition, inputs, targets}.

The following are examples from the training split:

1
{
2
  "task_name": "task001_quoref_question_generation",
3
  "id": "task001-abc123",
4
  "definition": "In this task, you're given passages that...",
5
  "inputs": "Passage: A man is sitting at a piano...",
6
  "targets": "What is the first name of the person who doubted it would be explosive?"
7
}
8
{
9
  "task_name": "task002_math_word_problems",
10
  "id": "task002-def456",
11
  "definition": "Solve the following word problem.",
12
  "inputs": "If there are 3 apples and you take 2...",
13
  "targets": "1"
14
}

For basic QA fine-tuning, we usually map definition → context, inputs → question, and targets → answer as follows:

1
from nemo_automodel.components.datasets.llm.column_mapped_text_instruction_dataset import (
2
    ColumnMappedTextInstructionDataset,
3
)
4
from transformers import AutoTokenizer
5
6
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")
7
8
remote_ds = ColumnMappedTextInstructionDataset(
9
    path_or_dataset_id="Muennighoff/natural-instructions",  # Hugging Face repo ID
10
    column_mapping={
11
        "context": "definition",  # high-level context
12
        "question": "inputs",      # the actual prompt / input
13
        "answer": "targets",       # expected answer string
14
    },
15
    tokenizer=tokenizer,
16
    split="train[:5%]",        # demo slice; omit (i.e., `split="train",`) for full data
17
    answer_only_loss_mask=True,
18
)

You can configure the entire dataset directly from your recipe YAML. For example:

1
# dataset section of your recipe's config.yaml
2
dataset:
3
  _target_: nemo_automodel.components.datasets.llm.column_mapped_text_instruction_dataset.ColumnMappedTextInstructionDataset
4
  path_or_dataset_id: Muennighoff/natural-instructions
5
  split: train
6
  column_mapping:
7
    context: definition
8
    question: inputs
9
    answer: targets
10
  answer_only_loss_mask: true

Streaming / Delta Lake / Databricks

Advanced Options

Tokenization Paths

This section explains how the dataset formats and tokenizes samples.

ColumnMappedTextInstructionDataset produces standard next-token training tensors:

• input_ids
• labels
• attention_mask

When answer_only_loss_mask=True, prompt tokens are masked in labels with -100 (the standard CrossEntropy ignore_index).

The dataset supports two formatting paths:

1. Chat-template path (opt-in): if use_hf_chat_template=True and the tokenizer exposes a chat_template and apply_chat_template, the dataset builds messages like:

   [{"role": "system", "content": <context or "">}, {"role": "user", "content": <question or "">}, {"role": "assistant", "content": <answer>}]

   and tokenizes them via tokenizer.apply_chat_template(..., tokenize=True, return_dict=True).

2. Plain prompt/completion path (default): otherwise the dataset concatenates prompt and answer and tokenizes the result.

In both cases, labels are the next-token targets (shifted by one relative to input_ids). The dataset also includes an internal ___PAD_TOKEN_IDS___ field used downstream for padding.

Parameter Requirements

The following section lists important requirements and caveats for correct usage.

• column_mapping must include answer, and must include at least one of context or question (2- or 3-column mapping only).
• If use_hf_chat_template=True, the tokenizer must support chat templates (chat_template + apply_chat_template).

Slurm Configuration for Distributed Training

For distributed training on Slurm clusters, add a slurm section to your YAML configuration. This section configures the Slurm batch job parameters and automatically generates the appropriate #SBATCH directives.

Slurm Configuration

SLURM jobs are submitted with sbatch directly — no YAML section needed. Copy the reference script, set the CONFIG variable to your YAML, and submit:

1
cp slurm.sub my_cluster.sub
2
# Edit my_cluster.sub — change CONFIG, #SBATCH directives, container, mounts, etc.
3
sbatch my_cluster.sub

All cluster-specific settings (nodes, GPUs, partition, container, mounts, secrets) live in your sbatch script. See the cluster guide for full examples (Pyxis, bare-metal, Apptainer).

Multi-Node Slurm Configuration

When using multiple nodes with Hugging Face datasets:

1. Shared Storage: Ensure all nodes can access the same storage paths
2. HF Cache: Export HF_HOME and HF_DATASETS_CACHE in your sbatch script pointing to shared directories
3. Mounts: Add shared directories as container mounts in your sbatch script

Configure all of this in your sbatch script (my_cluster.sub), not in the YAML.

That’s It!

With the mapping specified, the rest of the NeMo Automodel pipeline (pre-tokenization, packing, collate-fn, etc.) works as usual.