Save and Export Text Data

After processing your text datasets with NeMo Curator, use writer stages to export curated data for downstream use. Curator provides writers for common formats (JSONL, Parquet) as well as specialized writers for training frameworks.

Megatron Tokenization

MegatronTokenizerWriter tokenizes text documents and writes the .bin and .idx files required by Megatron-LM for data loading during pretraining. This replaces the need to run Megatron’s preprocess_data.py script separately and integrates tokenization directly into your curation pipeline.

How It Works

Tokenizer loading: Downloads and loads a Hugging Face tokenizer specified by model_identifier. The tokenizer is downloaded once per node and loaded once per worker.
Batched tokenization: Documents are tokenized in batches (controlled by tokenization_batch_size) to avoid out-of-memory issues on large datasets.
Binary output: Tokenized data is written to a .bin file containing packed token IDs. Vocabulary sizes above 65,536 use 4 bytes per token (int32); smaller vocabularies use 2 bytes (uint16).
Index output: A .idx file stores metadata including sequence lengths, byte offsets, and document boundaries for efficient random access during training.

Quick Start

from nemo_curator.core.client import RayClient
from nemo_curator.pipeline import Pipeline
from nemo_curator.stages.text.io.reader import JsonlReader
from nemo_curator.stages.text.io.writer.megatron_tokenizer import MegatronTokenizerWriter

# Initialize Ray client
ray_client = RayClient()
ray_client.start()

# Define pipeline stages
stages = [
    JsonlReader(
        file_paths="/path/to/data",
        fields=["text"],
    ),
    MegatronTokenizerWriter(
        path="/path/to/output",
        model_identifier="nvidia/NVIDIA-Nemotron-Nano-12B-v2",
        append_eod=True,
    ),
]

# Create and run the pipeline
pipeline = Pipeline(
    name="megatron-tokenize",
    description="Tokenize dataset for Megatron-LM.",
    stages=stages,
)

results = pipeline.run()

ray_client.stop()

Configuration

Parameter	Type	Default	Description
`path`	str	Required	Output directory for `.bin` and `.idx` files
`model_identifier`	str	Required	Hugging Face model identifier or local path for the tokenizer
`text_field`	str	`"text"`	Name of the column containing text to tokenize
`append_eod`	bool	`False`	Append the tokenizer’s EOS token at the end of each document
`tokenization_batch_size`	int	`1000`	Number of documents to tokenize per batch before writing to disk
`cache_dir`	str \| None	`None`	Local cache directory for the downloaded tokenizer
`hf_token`	str \| None	`None`	Hugging Face API token for accessing gated models

Output Format

The writer produces paired files for each input partition:

output/
├── {hash_1}.bin    # Packed token IDs (binary)
├── {hash_1}.idx    # Sequence metadata (lengths, offsets, document boundaries)
├── {hash_2}.bin
├── {hash_2}.idx

File format details

.bin file: Contains concatenated token IDs for all documents in the partition. Token IDs are stored as int32 (4 bytes) when the tokenizer vocabulary exceeds 65,536 tokens, or uint16 (2 bytes) for smaller vocabularies such as GPT-2.

.idx file: Contains a fixed header followed by per-sequence metadata:

9-byte magic header (MMIDIDX\x00\x00)
8-byte version number
1-byte dtype code
8-byte sequence count
8-byte document count
Per-sequence lengths: 4-byte int32 array (one entry per sequence)
Per-sequence byte offsets: 8-byte int64 array (one entry per sequence)
Document boundary indices: 8-byte int64 array (sequence count + 1 entries)

These files are directly compatible with Megatron-LM’s MMapIndexedDataset data loader.

End-of-Document Tokens

When append_eod=True, the tokenizer’s EOS token is appended to the end of each document’s token sequence. This is consistent with the behavior of Megatron’s preprocess_data.py and is required for some training configurations that use document boundaries for attention masking.

If the tokenizer does not define an EOS token, append_eod is automatically disabled with a warning.

Using Different Tokenizers

MegatronTokenizerWriter supports any tokenizer available through Hugging Face’s AutoTokenizer:

Standard Model

MegatronTokenizerWriter(
    path="output/",
    model_identifier="nvidia/NVIDIA-Nemotron-Nano-12B-v2",
    append_eod=True,
)

Gated Model

MegatronTokenizerWriter(
    path="output/",
    model_identifier="meta-llama/Llama-3.1-8B",
    hf_token="hf_...",
    append_eod=True,
)

Local Tokenizer

MegatronTokenizerWriter(
    path="output/",
    model_identifier="/local/path/to/tokenizer",
    append_eod=True,
)

Complete Pipeline Example

This example reads the TinyStories dataset from Parquet files and tokenizes it for Megatron-LM:

from nemo_curator.core.client import RayClient
from nemo_curator.pipeline import Pipeline
from nemo_curator.stages.text.io.reader import ParquetReader
from nemo_curator.stages.text.io.writer.megatron_tokenizer import MegatronTokenizerWriter

ray_client = RayClient()
ray_client.start()

stages = [
    ParquetReader(
        file_paths="datasets/tinystories/",
    ),
    MegatronTokenizerWriter(
        path="datasets/tinystories-tokens/",
        model_identifier="nvidia/NVIDIA-Nemotron-Nano-12B-v2",
        append_eod=True,
        tokenization_batch_size=2000,
    ),
]

pipeline = Pipeline(
    name="megatron-tokenize",
    description="Tokenize TinyStories for Megatron-LM.",
    stages=stages,
)

results = pipeline.run()

ray_client.stop()

A runnable version of this example is available in the tutorials directory.

For more information on using tokenized data with Megatron-LM, see the Related Tools page.