Data preparation guide

Data preparation guide#

Primus routes training through Megatron-LM, TorchTitan, and MaxText. Each backend expects its own data format and preprocessing pipeline. This guide summarizes how to prepare data, how to use mock data for smoke tests, and which environment variables commonly apply.

Scripts referenced below live under the Primus repository root, for example:

  • examples/megatron/preprocess_data.py

  • examples/megatron/prepare.py

  • examples/megatron/prepare_bookcorpus_megatron_dataset.py

  • examples/torchtitan/prepare.py

1. Overview#

Backend

Format

Typical entry

Megatron

Indexed .bin + .idx datasets

data_path, train_data_path, tokenizer args

TorchTitan

Hugging Face datasets + local assets

training.dataset, training.dataset_path, model.hf_assets_path

MaxText

TFDS / Hugging Face / Grain / synthetic

dataset_type, paths per pipeline

All backends support synthetic or mock data for configuration and scaling tests without large downloads.

2. Mock data (testing)#

Megatron#

Set in the trainer module:

mock_data: true

Default in primus/configs/modules/megatron/trainer_base.yaml is false. When true, training uses generated data matching configured dimensions so you can validate YAML, parallelism, and throughput without real corpora.

TorchTitan#

training:
  mock_data: true

Default in primus/configs/modules/torchtitan/pre_trainer.yaml is true (useful for quick runs; set false and supply real datasets for production).

MaxText#

Use dataset_type: synthetic (or other synthetic paths in MaxText configs). See third_party/maxtext/src/MaxText/configs/base.yml and model YAMLs under third_party/maxtext/src/MaxText/configs/.

3. Megatron data pipeline#

Inputs#

  • Raw JSON or JSONL text (one JSON object per line for JSONL).

  • Optional sentence splitting via NLTK when --split-sentences is used (requires NLTK data; see Environment variables).

Preprocessing: examples/megatron/preprocess_data.py#

The script tokenizes input and writes Megatron indexed datasets (.bin + .idx). It uses build_tokenizer from Primus’s Megatron tokenizer integration and accepts tokenizer flags from _add_tokenizer_args.

Important arguments (from the script’s argparse):

Argument

Description

--input

Path to input JSON (required).

--json-keys

Keys to read (default text).

--output-prefix

Output path without suffix; produces `{prefix}{key}{document

--workers

Number of worker processes (required).

--partitions

Split input for parallel preprocessing (default 1).

--split-sentences

Run NLTK sentence splitting before encode.

--append-eod

Append end-of-document token.

Example (mirrors examples/megatron/prepare.py for BookCorpus-style flows):

python3 examples/megatron/preprocess_data.py \
  --input /path/to/train.json \
  --tokenizer-type HuggingFaceTokenizer \
  --tokenizer-model /path/to/tokenizer \
  --output-prefix /path/to/out/bookcorpus_train \
  --workers "$(nproc)" \
  --split-sentences \
  --partitions 2

Configuring training runs#

Parameter

Notes

data_path

Single path or weighted blend: 0.5 /path/a 0.5 /path/b

train_data_path, valid_data_path, test_data_path

Separate splits when used

split

Train/valid/test ratio string, e.g. "99,1,0" (default in trainer_base.yaml) or "98,2,0" for train/valid/test

dataloader_type

Megatron dataloader type; default in trainer_base.yaml is null (set explicitly in experiments as needed)

BookCorpus example scripts#

  • examples/megatron/prepare_bookcorpus_megatron_dataset.py—downloads BookCorpus to JSON via Hugging Face datasets, optional --out-dir.

  • examples/megatron/prepare.py—orchestrates download, train/valid split, and calls preprocess_data.py with tokenizer settings from Primus config; respects TOKENIZED_TRAIN_DATA_PATH / TOKENIZED_EVAL_DATA_PATH for output locations.

Tokenizers#

Tokenizer type and model path are set on the model preset (for example tokenizer_type, tokenizer_model in primus/configs/models/megatron/language_model.yaml comments list Llama2Tokenizer, HuggingFaceTokenizer, etc.).

4. TorchTitan data pipeline#

TorchTitan uses Hugging Face datasets style identifiers and local paths.

Key

Default (pre_trainer.yaml)

Description

training.dataset

c4

Dataset identifier for TorchTitan loaders.

training.dataset_path

null

Local directory for dataset assets when needed.

Tokenizer and model assets are resolved from model.hf_assets_path (or equivalent in your model preset). The preparation script examples/torchtitan/prepare.py:

  • Resolves the TorchTitan checkout path.

  • Runs scripts/download_hf_assets.py inside TorchTitan to fetch tokenizer assets for a given repo_id.

  • Uses HF_TOKEN when the model or dataset is gated.

5. MaxText data pipeline#

MaxText configuration is defined in upstream YAML (for example third_party/maxtext/src/MaxText/configs/base.yml).

Parameter

Meaning

dataset_type

One of synthetic, hf, grain, tfds (per base.yml comments).

hf_path, hf_data_dir, hf_train_files

Hugging Face pipeline inputs when dataset_type: hf.

per_device_batch_size

Batch sizing on each device.

packing

Sequence packing for efficiency (default True in base.yml).

See MaxText’s data input documentation for Grain and TFDS specifics.

6. Environment variables for data#

Variable

Usage

TOKENIZED_DATA_PATH / PRIMUS_TOKENIZED_DATA_PATH

Tokenized dataset locations for Megatron hooks and examples (see docs/03-configuration-reference/environment-variables.md).

TOKENIZED_TRAIN_DATA_PATH, TOKENIZED_EVAL_DATA_PATH

Override output paths in examples/megatron/prepare.py.

DATA_PATH

General data root used in scripts and CI-style launches.

HF_TOKEN

Required for gated Hugging Face models and some datasets (TorchTitan prepare.py, Kubernetes examples in examples/README.md).

HF_HOME

Hugging Face cache directory (used in examples/megatron/prepare.py).

NLTK_DATA

NLTK tokenizer data directory for sentence splitting in preprocess_data.py when NLTK_DATA is set.

Summary#

  1. Use mock or synthetic data to validate configs and performance before investing in large preprocessing jobs.

  2. For Megatron, convert JSON/JSONL to .bin/.idx with preprocess_data.py and point data_path or split paths at the outputs.

  3. For TorchTitan, set training.dataset / dataset_path and run examples/torchtitan/prepare.py to fetch tokenizer assets.

  4. For MaxText, configure dataset_type and per_device_batch_size per upstream base.yml and model YAMLs.

Contents