Skip to content

data-preparation.md

Latest commit

 

History

History
159 lines (108 loc) · 7.04 KB

data-preparation.md

File metadata and controls

159 lines (108 loc) · 7.04 KB
Raw

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.