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 abbreviated rows are illustrative examples of that schema. Their IDs and text are not copied from the live 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.

• A 2-column column_mapping must have exactly the logical keys {answer, context} or {answer, question}. A 3-column mapping must have exactly {context, question, answer}. Each value names the corresponding source-dataset column.
• If use_hf_chat_template=True, the dataset uses the tokenizer’s chat template when chat_template is present and apply_chat_template is callable; otherwise it falls back to plain prompt/completion formatting.

Slurm Configuration for Distributed Training

For distributed training on Slurm clusters, use the checked-in root-level slurm.sub script and submit it directly with sbatch. Recipe YAML does not select a Slurm launcher or generate #SBATCH directives.

Slurm Configuration

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. Dataset access: Give every node network access to the dataset or a populated node-local cache
2. Optional shared cache: To avoid duplicate downloads, point HF_HOME or HF_DATASETS_CACHE at a shared directory
3. Mounts: If you use shared storage inside a container, mount that directory on every node

Set cache environment variables and container mounts in your sbatch script (my_cluster.sub), not in the recipe YAML. Network access is a cluster property.

That’s It!

With the mapping specified, the dataset tokenizes each sample lazily when ds[i] is accessed; downstream packing and collation then work as usual.