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
mmap_data_processing_test.py	Reader, datasource, Grain integration, multiprocessing, EOD, and cache behavior.
input_pipeline_utils_test.py	Input/target transforms, segmentation, dataset IDs, and segment merging.
megatron_mmap_compatibility_test.py	Index construction, splitting, blending, CLI, host ordering, and Megatron compatibility.
configs_value_test.py	mmap configuration defaults and validation.
mmap_test_utils.py	Shared non-collected test helpers.

6.2 End-to-end checkpoint restoration smoke test

resume_file_type_test.sh is a single-process TPU-VM mmap_npy checkpoint-restoration smoke test. It validates the smaller but critical end-to-end path in which MaxText saves a Grain iterator state, restores it from the same training directory in a subsequent process, and continues training successfully.

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, verifies from the logs that the checkpoint and iterator state were restored from the same run directory, and confirms that training continues. The script rejects a pre-existing checkpoint directory to avoid accidentally consuming stale state.

This is a functional restoration smoke test, not a determinism test. steps contributes to the mmap_npy sample count and may change the derived epoch bucket, final-epoch handling, and cached indices.

#!/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.

[RissyRan]

Hi reviewers, just an FYI that this PR is from a forked repo. Due to secret key limitations, the Gemini review couldn't be triggered.

[codecov]

Codecov Report

❌ Patch coverage is 93.90739% with 75 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/input_pipeline/_mmap_datasource.py	94.47%	15 Missing and 17 partials ⚠️
src/maxtext/input_pipeline/_mmap_index_utils.py	94.62%	10 Missing and 8 partials ⚠️
src/maxtext/input_pipeline/_megatron_blending.py	91.77%	8 Missing and 5 partials ⚠️
src/maxtext/input_pipeline/input_pipeline_utils.py	93.45%	3 Missing and 4 partials ⚠️
...rc/maxtext/input_pipeline/grain_data_processing.py	90.38%	2 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!

[aireenmei]

Thanks for the detailed PR description, and also update documentations! LGTM in general. Just minor comments:

For the new unit tests, could we consolidate them for easier maintenance?
My suggestion is to use two main new test modules:

• mmap_data_processing_test.py: mmap/mmap_npy datasource behavior, file parsing, sharding, multiprocessing, Grain integration, EOD semantics, and checkpoint behavior.
• megatron_mmap_compatibility_test.py: index construction, splitting, blending, CLI behavior, host-order reconstruction, and comparisons with megatron.core.

The smaller tests could extend existing modules:

• Move MegatronSplitInputsTargets, GenerateDocSegmentIds, and segment-merging tests into input_pipeline_utils_test.py.
• Move mmap configuration default/validation tests into configs_value_test.py.
• Keep mmap_test_utils.py as a shared non-collected helper.

[Yitrus]

Thanks for the detailed PR description, and also update documentations! LGTM in general. Just minor comments:

For the new unit tests, could we consolidate them for easier maintenance? My suggestion is to use two main new test modules:

• mmap_data_processing_test.py: mmap/mmap_npy datasource behavior, file parsing, sharding, multiprocessing, Grain integration, EOD semantics, and checkpoint behavior.
• megatron_mmap_compatibility_test.py: index construction, splitting, blending, CLI behavior, host-order reconstruction, and comparisons with megatron.core.

The smaller tests could extend existing modules:

• Move MegatronSplitInputsTargets, GenerateDocSegmentIds, and segment-merging tests into input_pipeline_utils_test.py.
• Move mmap configuration default/validation tests into configs_value_test.py.
• Keep mmap_test_utils.py as a shared non-collected helper.

Thanks for the suggestion! I consolidated the tests into the two proposed modules, moved the smaller transform/config tests into the existing modules, and kept mmap_test_utils.py as a shared helper.

Validation passed: 312 passed overall and 20 passed for Megatron alignment.

[RissyRan]

I see this error Error: Not adding pull ready because the PR has more than one commit. Please squash your commits..

Could you help squash all 7 commits into one?

[Yitrus]

I see this error Error: Not adding pull ready because the PR has more than one commit. Please squash your commits..

Could you help squash all 7 commits into one?

Done! Squashed all 7 commits into one and force-pushed. The diff is identical to before, just cleaner history now.

[igorts-git]

Hi @Yitrus now the process is failing because there is no license string in the header of the newly added files:

_mmap_datasource.py 
megatron_mmap_compatibility_test.py 
_mmap_index_utils.py 
mmap_data_processing_test.py
_megatron_blending.py 
mmap_test_utils.py
mmap_index_builder.py

[Yitrus]

Hi @Yitrus now the process is failing because there is no license string in the header of the newly added files:

_mmap_datasource.py 
megatron_mmap_compatibility_test.py 
_mmap_index_utils.py 
mmap_data_processing_test.py
_megatron_blending.py 
mmap_test_utils.py
mmap_index_builder.py

Thanks for pointing this out. I’ve added the standard Apache 2.0 license header to all seven newly added files.

[RissyRan]

Hi @igorts-git I tried 2 times of CI/CD rer-run, and pathway unit tests seem down due to connection.

[igorts-git]

The CI problems are now fixed on main, the PR needs to be rebased. Sorry about that.

[RissyRan]

The fix is in #4991

Could you rebase the PR? We will help merge from internal codebase.

[Yitrus]

The fix is in #4991

Could you rebase the PR? We will help merge from internal codebase.

Thanks! I’ve rebased the PR onto the latest main, which includes the fix from #4991. The Pathways CI and all required tests are now passing.
The only remaining failure is Track Test Performance, which ended with Resource not accessible by integration. Could you please rerun this check? Thanks!

[Yitrus]

Hi @aireenmei @igorts-git @RissyRan, just a friendly follow-up on this PR. It now has the required approvals, and the only remaining failure is the "Track Test Performance" check, which failed with Resource not accessible by integration.

Could one of you please help rerun this check and proceed with the internal merge at your convenience? Please let me know if anything else is needed from my side. Thanks!

[aireenmei]

Hi @aireenmei @igorts-git @RissyRan, just a friendly follow-up on this PR. It now has the required approvals, and the only remaining failure is the "Track Test Performance" check, which failed with Resource not accessible by integration.

Could one of you please help rerun this check and proceed with the internal merge at your convenience? Please let me know if anything else is needed from my side. Thanks!

Working on it