feat(input): Megatron MMap Data Formats for MaxText Grain

[Yitrus]

1. Summary

This proposal adds mmap and mmap_npy data formats to MaxText's Grain input pipeline so training can consume Megatron-LM text-preprocessing .bin/.idx datasets directly. mmap is a simpler Grain-native reader; mmap_npy uses Megatron-style indices to reproduce document shuffle, fixed-length sample construction, sample shuffle, and multi-dataset blending order.

The goal of mmap_npy is element-by-element equivalence to Megatron GPTDataset for raw L + 1 token samples. It does not introduce a separate training runtime: data still reaches the JAX model through Grain transforms, batching, and MultiHostDataLoadIterator.

Indices can be built offline with a CLI or automatically on the first training run. The runtime cache is keyed by an epoch bucket; on a cache miss every host deterministically computes indices in memory while only host 0 atomically persists the shared cache, so no barrier is required.

2. Motivation and scope

Megatron-LM is widely used for large-scale pretraining, and users often already own .bin/.idx data assets. Without a direct reader, moving to MaxText requires data conversion or accepting a different sample order, which confounds convergence, throughput, and numerical comparisons.

This design covers Megatron text indexed datasets only. Input data must contain EOD tokens written during preprocessing with --append-eod; runtime code does not insert EOD again. ArrayRecord, TFRecord, Parquet, HF, TFDS, and other existing training paths remain unchanged.

3. User-facing formats and configuration

3.1 Select a format

The implementation continues to use dataset_type=grain and adds grain_file_type=mmap and grain_file_type=mmap_npy. mmap reads .bin/.idx and uses EOD-aware Grain transforms; it is suitable for simple training or exploration but does not promise complete Megatron GPTDataset sample-order parity. mmap_npy reads the same data format and loads or builds three .npy indices to provide Megatron sample-order compatibility.

/cache/wiki_indices/
  <hash>-document_index.npy
  <hash>-sample_index.npy
  <hash>-shuffle_index.npy

3.2 Configure mmap/mmap_npy

The table below lists the configuration needed to use mmap/mmap_npy.

Parameter	Configuration location	Values and units	Notes
grain_file_type	src/maxtext/configs/base.yml src/maxtext/configs/types.py::GrainDataset	mmap: direct .bin/.idx reader. mmap_npy: Megatron-style indexed reader.	mmap_npy is the compatibility path for Megatron GPT sample order; mmap is the simpler path.
grain_train_files grain_eval_files	src/maxtext/configs/base.yml src/maxtext/configs/types.py::GrainDataset	mmap: a prefix/directory, or weighted entries separated by ;. mmap_npy: npy_dir|bin_prefix_or_directory; weighted entries use ,weight and ;. Path strings and dimensionless weights.	npy_dir contains only index files; the .bin/.idx prefix or directory contains tokens. Mixture weights must be non-negative and have a positive total.
max_target_length	src/maxtext/configs/base.yml src/maxtext/configs/types.py::MaxTextConfig	Positive integer L, in tokens.	mmap_npy reads L+1 raw tokens and emits length-L model fields. It must match the intended Megatron sequence length.
mmap_eod_id	src/maxtext/configs/base.yml src/maxtext/configs/types.py::MMapDataset	Integer token ID; default 0.	Must equal the EOD token written during Megatron preprocessing with --append-eod. It replaced the earlier pad_id name because EOD and padding differ semantically; bos_id is not used by this path.
mmap_split_sentences	src/maxtext/configs/base.yml src/maxtext/configs/types.py::MMapDataset	Boolean; default false.	Set true only when the indexed dataset was produced with Megatron --split-sentences; otherwise the reader reports a layout mismatch.
mmap_npy_split	src/maxtext/configs/base.yml src/maxtext/configs/types.py::MMapDataset	Empty string disables splitting. Comma-separated ratios such as 99,1 or 98,1,1 enable splitting.	Applies only to mmap_npy. Training uses split 0; evaluation uses split 1 when a split is configured. Boundaries use Megatron-compatible rounding and document indices retain global document IDs.
blend_cache_dir	src/maxtext/configs/base.yml src/maxtext/configs/types.py::MMapDataset	Empty string disables persistence. Otherwise a writable directory path.	Used only for multi-dataset mmap_npy blending. It stores runtime-generated global blend dispatch indices; cache-write failure falls back to in-memory indices.
blend_index_dir	src/maxtext/configs/base.yml src/maxtext/configs/types.py::MMapDataset	Empty string disables offline loading. Otherwise a directory containing dataset_index.npy and dataset_sample_index.npy, or one uniquely identifiable equivalent pair.	Used only for multi-dataset mmap_npy blending and takes precedence over blend_cache_dir. The pair must match dataset order and requested size; invalid files fall back to a rebuilt dispatch.
reset_attention_mask	src/maxtext/configs/base.yml src/maxtext/configs/types.py::MMapDataset	Boolean; default true.	For mmap/mmap_npy transforms, true starts a new attention segment after EOD and restarts the next position at 0; false keeps one continuous segment and positions. The standard Grain pretraining path does not read this field.
eod_mask_loss	src/maxtext/configs/base.yml src/maxtext/configs/types.py::MMapDataset	Boolean; default false.	For mmap/mmap_npy transforms, true masks positions whose input token is EOD; false keeps them in loss. It is independent of attention reset. The simple mmap path has a ShiftData limitation that still masks EOD-related loss when this field is false; use mmap_npy for strict Megatron loss semantics.
packing_max_segments_per_sample	src/maxtext/configs/base.yml src/maxtext/configs/types.py::MMapDataset	Integer divisor; default 25. <=0 disables short-segment merging.	When reset is enabled, the threshold is max_target_length // packing_max_segments_per_sample tokens. EOD-derived boundaries at or below the threshold are merged; this is independent of Grain packing itself. Set <=0 to retain every EOD boundary.

3.3 Prepare and run a dataset

The following is a single-dataset example that strictly preserves Megatron GPT next-token, EOD, and index-order semantics.

# mmap_npy_megatron_gpt.yml
# If this file is outside src/maxtext/configs/, use an absolute base_config path.
base_config: "base.yml"

dataset_type: "grain"
grain_file_type: "mmap_npy"
grain_train_files: "/cache/wiki_indices|/data/wiki_text_document"

# Must match the reference Megatron data job.
max_target_length: 2048
mmap_eod_id: 2
data_shuffle_seed: 1234
reset_attention_mask: true
eod_mask_loss: false
packing_max_segments_per_sample: 0

# Match the actual indexed-data layout. No train/eval document split here.
mmap_split_sentences: false
mmap_npy_split: ""

# These values determine the mmap_npy training index size.
global_batch_size_to_load: 1024
steps: 1000

# Set for the target environment.
base_output_directory: "gs://YOUR_BUCKET/maxtext-output"

python3 -m maxtext.trainers.pre_train.train \
  src/maxtext/configs/mmap_npy_megatron_gpt.yml

The offline tool tools/data_processing/mmap_index_builder.py is a CLI wrapper around the core library. It is not a required training prerequisite: a runtime cache miss builds the same indices automatically. Offline construction is useful when users want to inspect artifacts before training, publish indices to a read-only directory, or avoid index computation during first-job startup.

The following single-dataset command corresponds to the YAML above.

python3 tools/data_processing/mmap_index_builder.py convert \
  --input /data/wiki_text_document \
  --output-dir /cache/wiki_indices \
  --seq-length 2048 \
  --num-samples 1024000 \
  --seed 1234 \
  --add-extra-token 1

Multi-dataset offline construction uses the blend subcommand.

python3 tools/data_processing/mmap_index_builder.py blend \
  --datasets '/data/wiki_text_document,0.7;/data/code_text_document,0.3' \
  --output-dir /cache/wiki_code_blend \
  --seq-length 2048 \
  --total-samples 1024000 \
  --seed 1234 \
  --margin 0.5 \
  --add-extra-token 1

Offline blend artifacts require both child-index directories and the root dispatch directory; add the following snippet to the same experiment YAML.

grain_train_files: "/cache/wiki_code_blend/dataset_0|/data/wiki_text_document,0.7;/cache/wiki_code_blend/dataset_1|/data/code_text_document,0.3"
blend_index_dir: "/cache/wiki_code_blend"

4. Compatibility contract

4.1 model batches

For ordinary pretraining, optional features can add dataset_id, image, or SFT/DPO fields, but both standard Grain and mmap_npy deliver the following six core fields to the model.

The following example uses L=4, PAD=0, EOD=99, reset_attention_mask=True, and eod_mask_loss=False. The raw L+1 tokens are [10, 11, 99, 20, 21], where 99 is the EOD boundary between two documents.

Standard Grain without packing does not treat EOD as a document boundary: valid tokens are in segment 1, padding is in segment 0, and positions increase continuously. ShiftData left-shifts targets and pads the tail and a zero loss mask at the last position.

Standard Grain, without packing
  inputs:                [10, 11, 99, 20]
  targets:               [11, 99, 20,  0]
  inputs_segmentation:   [ 1,  1,  1,  1]
  targets_segmentation:  [ 1,  1,  1,  0]
  inputs_position:       [ 0,  1,  2,  3]
  targets_position:      [ 0,  1,  2,  3]

Packing is a separate feature: it places multiple short, independent samples into one fixed-length sequence to reduce padding. Each original sample restarts positions at 0, and segment boundaries are defined by the Grain packer's sample-concatenation points rather than automatically by EOD.

mmap_npy forms next-token pairs directly from real L+1 tokens: inputs=tokens[:-1] and targets=tokens[1:]，whereas mmap_npy retains a real final target and computes its loss. With reset enabled, EOD still belongs to the preceding document; the following 20 starts a new attention segment and position resets to 0. With eod_mask_loss=False, every target contributes to supervision.

Megatron mmap_npy
  inputs:                [10, 11, 99, 20]
  targets:               [11, 99, 20, 21]
  inputs_segmentation:   [ 1,  1,  1,  2]
  targets_segmentation:  [ 1,  1,  1,  1]
  inputs_position:       [ 0,  1,  2,  0]
  targets_position:      [ 0,  1,  2,  0]

Consequently, standard Grain has a padded final target, The ordinary standard-Grain pretraining path does not read reset_attention_mask or eod_mask_loss, and it neither resets positions nor increments input segments at EOD.

Both simple mmap and mmap_npy use EOD-aware attention/position transforms; their main difference is next-token sample construction. mmap still uses ShiftData afterwards, so with eod_mask_loss=False it masks EOD-related loss because mmap_eod_id also acts as an ignored ID; use mmap_npy when strict Megatron loss semantics are required.

4.2 Megatron parity boundary

mmap_npy promises Megatron GPTDataset data-consumption semantics: with the same data, configuration, and random seed, each position in the global sample stream reads the same raw tokens and constructs semantically corresponding model fields.

This promise requires grain_file_type=mmap_npy and matching .bin/.idx, EOD ID, sequence length, shuffle seed, split/blend configuration, and requested sample count. Conventional Megatron GPT EOD semantics additionally require reset_attention_mask=true, eod_mask_loss=false, and packing_max_segments_per_sample=0.

5. Architecture and runtime behavior

5.1 Dataloader placement

The diagram shows where the new logic belongs: offline preprocessing ends at .bin/.idx; mmap_npy generates or loads Megatron indices before the Grain runtime, then joins the existing Grain transforms, batching, and multi-host loader. The optional offline CLI only produces indices; training can still rebuild them when the cache is absent.

5.2 Megatron-parity implementation map

Stage	Megatron-aligned behavior	MaxText implementation
Indexed input	Reads the same Megatron .bin/.idx format and relies on the EOD token written during preprocessing.	MMapIndexedDataset in _mmap_datasource.py parses the index and performs token reads. Runtime does not insert EOD.
Child index construction	Uses one seed for Megatron-compatible document shuffle, L+1 sample construction, and sample shuffle.	build_indices() in _mmap_index_utils.py implements these three steps and produces the document, sample, and shuffle indices. MegatronNpyDataSource consumes the resulting arrays for random access.
Dataset split	Applies Megatron-style ratios and preserves global document IDs, so sample boundaries still address the original dataset correctly.	The index builder applies mmap_npy_split to select the documents for a split and writes their global IDs into document_index.npy. MegatronNpyDataSource uses those IDs to read the corresponding documents from the original .bin/.idx dataset.
Weighted blending	Chooses a child dataset with Megatron greedy error-minimization and pins the dispatch to the requested total sample count.	build_blending_indices() and MegatronBlendedDataSource in _megatron_blending.py; offline/cache dispatch files are an optional MaxText delivery mechanism.
Model-field construction	Retains the real final next-token label, and applies EOD-aware attention segmentation, loss masking, and position reset according to the selected options.	GenerateDocSegmentIds and MegatronSplitInputsTargets in input_pipeline_utils.py construct the six model fields.
Grain and distributed delivery	The deterministic global sample stream can be partitioned across hosts without rebuilding independent child streams.	_mmap_pretrain_pipeline() and make_grain_train_iterator() attach the source to Grain; MultiHostDataLoadIterator transfers local batches to the JAX global mesh. This runtime is MaxText-specific rather than a Megatron dataloader replica.

6. Validation, robustness, and tests

6.1 Unit tests

Unit tests cover reader format and corrupted-input checks, offline and runtime indices, data sources, EOD transforms, and alignment with real Megatron helpers/GPTDataset. The table below describes the primary purpose of each test file.

Test file	Primary coverage
tests/unit/mmap_index_builder_test.py	Builds and persists document, sample, and shuffle indices; validates shard discovery, split document IDs, the convert CLI, and offline blend-index output.
tests/unit/mmap_data_processing_test.py	Validates .idx/.bin parsing, reader/source boundaries and failures, sentence-split layout, simple mmap Grain pipeline integration, EOD semantics, and the no-runtime-EOD-insertion rule.
tests/unit/megatron_npy_datasource_test.py	Covers mmap_npy path resolution, index discovery, random access and fixed-length output, pickle safety, runtime index-cache lifecycle, and pipeline integration.
tests/unit/megatron_split_dataset_id_test.py	Verifies that split index files retain global document IDs and that each configured split reads the intended source documents.
tests/unit/packing_max_segments_test.py	Verifies the short-segment threshold derived from packing_max_segments_per_sample, including disabled and invalid settings.
tests/unit/megatron_dataloader_alignment_test.py	Compares document/sample/shuffle indices, token samples, blend dispatch, split rounding, loss-mask/position semantics, and host-stride reassembly against Megatron helpers and GPTDataset.

6.2 End-to-end checkpoint-resume smoke test

resume_file_type_test.sh is a single-process TPU-VM mmap_npy checkpoint-resume smoke test. It validates the smaller but critical end-to-end path that a saved Grain iterator state can be restored by a subsequent MaxText run in the same training directory.

The script requires a TPU-VM checkout on the target branch and caller-provided GRAIN_TRAIN_FILES, MMAP_EOD_ID, TOKENIZER_PATH, and VOCAB_SIZE. It fixes grain_file_type=mmap_npy, disables the elastic iterator, and runs two trainings with the same run_name and output directory: Run A trains for STEPS_A and saves a final checkpoint, after which the script confirms that a single-process Grain iterator state exists. Run B raises the total step count to STEPS_B, checks the log for restoration from that run directory, and continues training. The script rejects a pre-existing checkpoint directory so that it cannot accidentally consume stale state.

#!/usr/bin/env bash
set -euo pipefail

MAXTEXT_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
EXPECTED_BRANCH="${EXPECTED_BRANCH:-feature/megatron-mmap-grain-file-type}"
CURRENT_BRANCH="$(git -C "${MAXTEXT_ROOT}" branch --show-current)"
if [[ "${CURRENT_BRANCH}" != "${EXPECTED_BRANCH}" ]]; then
  echo "ERROR: expected branch ${EXPECTED_BRANCH}, found ${CURRENT_BRANCH}." >&2
  echo "Switch the TPU VM checkout before running this local script." >&2
  exit 2
fi

# Set VENV_PATH when the TPU VM uses a dedicated virtual environment. Leaving
# it empty uses the active Python environment as-is.
VENV_PATH="${VENV_PATH:-}"
if [[ -n "${VENV_PATH}" ]]; then
  # shellcheck disable=SC1090
  source "${VENV_PATH}/bin/activate"
fi

# MEGATRON_TRAIN_FILES is retained as an alias for prior local invocations.
GRAIN_TRAIN_FILES="${GRAIN_TRAIN_FILES:-${MEGATRON_TRAIN_FILES:-}}"
GRAIN_TRAIN_FILES="${GRAIN_TRAIN_FILES:?Set GRAIN_TRAIN_FILES to 'npy_dir|bin_prefixes'.}"
MMAP_EOD_ID="${MMAP_EOD_ID:?Set MMAP_EOD_ID to the EOD token stored in the .bin files.}"
TOKENIZER_PATH="${TOKENIZER_PATH:?Set TOKENIZER_PATH to a local or Hugging Face tokenizer path.}"
VOCAB_SIZE="${VOCAB_SIZE:?Set VOCAB_SIZE to cover every token ID in the dataset.}"

OUTPUT_DIR="${OUTPUT_DIR:-/tmp/megatron_mmap_resume_test}"
RUN_NAME="${RUN_NAME:-megatron-mmap-resume-$(date +%Y%m%d-%H%M%S)}"
LOG_A="${LOG_A:-${OUTPUT_DIR}/${RUN_NAME}.runA.log}"
LOG_B="${LOG_B:-${OUTPUT_DIR}/${RUN_NAME}.runB.log}"

GRAIN_FILE_TYPE="${GRAIN_FILE_TYPE:-mmap_npy}"
TOKENIZER_TYPE="${TOKENIZER_TYPE:-huggingface}"
DATA_SEED="${DATA_SEED:-42}"
PER_DEVICE_BATCH="${PER_DEVICE_BATCH:-1}"
SEQ_LEN="${SEQ_LEN:-128}"
WEIGHT_DTYPE="${WEIGHT_DTYPE:-bfloat16}"
STEPS_A="${STEPS_A:-10}"
STEPS_B="${STEPS_B:-20}"
CHECKPOINT_PERIOD="${CHECKPOINT_PERIOD:-${STEPS_A}}"
MODEL_NAME="${MODEL_NAME:-gpt3-52k}"
MMAP_SPLIT_SENTENCES="${MMAP_SPLIT_SENTENCES:-False}"
BLEND_CACHE_DIR="${BLEND_CACHE_DIR:-}"
BLEND_INDEX_DIR="${BLEND_INDEX_DIR:-}"
MMAP_NPY_SPLIT="${MMAP_NPY_SPLIT:-}"

if [[ "${GRAIN_FILE_TYPE}" != "mmap_npy" ]]; then
  echo "ERROR: this parity/checkpoint smoke requires GRAIN_FILE_TYPE=mmap_npy." >&2
  exit 2
fi

if (( STEPS_A <= 0 || STEPS_B <= STEPS_A || CHECKPOINT_PERIOD <= 0 )); then
  echo "ERROR: require STEPS_A > 0, STEPS_B > STEPS_A, and CHECKPOINT_PERIOD > 0." >&2
  exit 2
fi

RUN_DIR="${OUTPUT_DIR}/${RUN_NAME}"
if [[ -e "${RUN_DIR}/checkpoints" ]]; then
  echo "ERROR: ${RUN_DIR}/checkpoints already exists; choose a fresh RUN_NAME." >&2
  exit 2
fi

mkdir -p "${OUTPUT_DIR}"
export PYTHONPATH="${MAXTEXT_ROOT}/src:${PYTHONPATH:-}"
export PYTHONUNBUFFERED=1

extra_data_args=()
if [[ -n "${BLEND_CACHE_DIR}" ]]; then
  extra_data_args+=("blend_cache_dir=${BLEND_CACHE_DIR}")
fi
if [[ -n "${BLEND_INDEX_DIR}" ]]; then
  extra_data_args+=("blend_index_dir=${BLEND_INDEX_DIR}")
fi
if [[ -n "${MMAP_NPY_SPLIT}" ]]; then
  extra_data_args+=("mmap_npy_split=${MMAP_NPY_SPLIT}")
fi

run_train() {
  local steps="$1"
  local logfile="$2"
  echo "----- launching: steps=${steps}; log=${logfile} -----"
  python3 -m maxtext.trainers.pre_train.train \
    "${MAXTEXT_ROOT}/src/maxtext/configs/base.yml" \
    model_name="${MODEL_NAME}" \
    override_model_config=True \
    run_name="${RUN_NAME}" \
    base_output_directory="${OUTPUT_DIR}" \
    dataset_type=grain \
    grain_file_type="${GRAIN_FILE_TYPE}" \
    grain_train_files="${GRAIN_TRAIN_FILES}" \
    grain_use_elastic_iterator=False \
    mmap_eod_id="${MMAP_EOD_ID}" \
    mmap_split_sentences="${MMAP_SPLIT_SENTENCES}" \
    data_shuffle_seed="${DATA_SEED}" \
    per_device_batch_size="${PER_DEVICE_BATCH}" \
    max_target_length="${SEQ_LEN}" \
    steps="${steps}" \
    eval_interval=0 \
    enable_checkpointing=True \
    async_checkpointing=False \
    checkpoint_period="${CHECKPOINT_PERIOD}" \
    save_checkpoint_on_completion=True \
    tokenizer_type="${TOKENIZER_TYPE}" \
    tokenizer_path="${TOKENIZER_PATH}" \
    tokenize_train_data=False \
    vocab_size="${VOCAB_SIZE}" \
    weight_dtype="${WEIGHT_DTYPE}" \
    "${extra_data_args[@]}" \
    2>&1 | tee "${logfile}"
}

echo "=== Megatron mmap_npy / Grain checkpoint-resume smoke test ==="
echo "  branch        : ${CURRENT_BRANCH}"
echo "  run_name      : ${RUN_NAME}"
echo "  output_dir    : ${RUN_DIR}"
echo "  grain type    : ${GRAIN_FILE_TYPE}"
echo "  train files   : ${GRAIN_TRAIN_FILES}"
echo "  sequence len  : ${SEQ_LEN}"
echo "  Run A / Run B : ${STEPS_A} / ${STEPS_B} steps"
echo

run_train "${STEPS_A}" "${LOG_A}"

ITER_STATE="$(find "${RUN_DIR}/checkpoints" -path '*/iter/process_0-of-1.json' -print -quit 2>/dev/null || true)"
if [[ -z "${ITER_STATE}" ]]; then
  echo "ERROR: Run A completed without a single-process Grain iterator checkpoint under iter/." >&2
  exit 1
fi
echo "Run A saved Grain iterator state: ${ITER_STATE}"

run_train "${STEPS_B}" "${LOG_B}"

if ! grep -q "restoring from this run's directory step" "${LOG_B}"; then
  echo "ERROR: Run B did not log restoration from the existing run directory." >&2
  exit 1
fi

echo
echo "=== PASS: Megatron mmap_npy / Grain checkpoint-resume smoke test ==="
echo "Run B restored an existing checkpoint and continued training."
echo "Run A log: ${LOG_A}"
echo "Run B log: ${LOG_B}"

Checklist

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests and provided reproduction details above where applicable.
• I have made corresponding documentation changes, including adding the new page to the relevant toctree.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.