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

1. 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.
2. Batched tokenization: Documents are tokenized in batches (controlled by tokenization_batch_size) to avoid out-of-memory issues on large datasets.
3. 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).
4. Index output: A .idx file stores metadata including sequence lengths, byte offsets, and document boundaries for efficient random access during training.

Quick Start

1
from nemo_curator.core.client import RayClient
2
from nemo_curator.pipeline import Pipeline
3
from nemo_curator.stages.text.io.reader import JsonlReader
4
from nemo_curator.stages.text.io.writer.megatron_tokenizer import MegatronTokenizerWriter
5
6
# Initialize Ray client
7
ray_client = RayClient()
8
ray_client.start()
9
10
# Define pipeline stages
11
stages = [
12
    JsonlReader(
13
        file_paths="/path/to/data",
14
        fields=["text"],
15
    ),
16
    MegatronTokenizerWriter(
17
        path="/path/to/output",
18
        model_identifier="nvidia/NVIDIA-Nemotron-Nano-12B-v2",
19
        append_eod=True,
20
    ),
21
]
22
23
# Create and run the pipeline
24
pipeline = Pipeline(
25
    name="megatron-tokenize",
26
    description="Tokenize dataset for Megatron-LM.",
27
    stages=stages,
28
)
29
30
results = pipeline.run()
31
32
ray_client.stop()

Configuration

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

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:

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

Complete Pipeline Example

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

1
from nemo_curator.core.client import RayClient
2
from nemo_curator.pipeline import Pipeline
3
from nemo_curator.stages.text.io.reader import ParquetReader
4
from nemo_curator.stages.text.io.writer.megatron_tokenizer import MegatronTokenizerWriter
5
6
ray_client = RayClient()
7
ray_client.start()
8
9
stages = [
10
    ParquetReader(
11
        file_paths="datasets/tinystories/",
12
    ),
13
    MegatronTokenizerWriter(
14
        path="datasets/tinystories-tokens/",
15
        model_identifier="nvidia/NVIDIA-Nemotron-Nano-12B-v2",
16
        append_eod=True,
17
        tokenization_batch_size=2000,
18
    ),
19
]
20
21
pipeline = Pipeline(
22
    name="megatron-tokenize",
23
    description="Tokenize TinyStories for Megatron-LM.",
24
    stages=stages,
25
)
26
27
results = pipeline.run()
28
29
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.