From ed819b7597cb8a74aa2f664110888e35ce617b10 Mon Sep 17 00:00:00 2001 From: Santosh Bhavani Date: Wed, 8 Apr 2026 17:30:29 -0700 Subject: [PATCH 1/5] docs: add data loading at scale guide Add documentation for configuring Megatron data pipeline at 256-512+ node scale, covering index building architecture, cache pre-warming, and the key flags (--dataloader-fast-cache-load, --dataloader-defer-npy-index-mmap, --per-dataset-sequences-path). Includes troubleshooting section and guidance on consolidating small dataset files. --- docs/index.md | 1 + docs/user-guide/data-loading.md | 132 ++++++++++++++++++++++++++++++++ 2 files changed, 133 insertions(+) create mode 100644 docs/user-guide/data-loading.md diff --git a/docs/index.md b/docs/index.md index 4b75ed2c0c8..079fdfbecae 100644 --- a/docs/index.md +++ b/docs/index.md @@ -48,6 +48,7 @@ get-started/quickstart :caption: Basic Usage user-guide/data-preparation +user-guide/data-loading user-guide/training-examples user-guide/parallelism-guide ``` diff --git a/docs/user-guide/data-loading.md b/docs/user-guide/data-loading.md new file mode 100644 index 00000000000..b197e8e48cb --- /dev/null +++ b/docs/user-guide/data-loading.md @@ -0,0 +1,132 @@ + + +# Data Loading at Scale + +This guide covers how Megatron's data pipeline works and how to configure it for efficient training at 256 nodes and beyond. At this scale, the primary bottlenecks are **index building**, **filesystem metadata operations**, and **barrier synchronization** -- not raw data bandwidth. + +## How Data Loading Works + +Understanding the architecture helps explain why specific flags matter. + +Megatron builds three index arrays for each dataset: a **document index** (shuffled document order), a **sample index** (mapping samples to document offsets), and a **shuffle index** (final sample permutation). This happens once during initialization: + +1. **Rank 0** builds all three indices and writes them to a cache directory as `.npy` files. +2. All ranks synchronize at a `torch.distributed.barrier()`. +3. **All other ranks** load the cached indices via memory-mapped reads (`numpy.load(mmap_mode='r')`). + +After initialization, data access is **read-only and lock-free**. Every data-parallel rank reads a disjoint, contiguous chunk of the shuffle index. No cross-rank coordination is needed during training because all ranks derive the same deterministic permutation from a shared random seed. + +## The Problem at 256+ Nodes + +Three things break down at large node counts: + +1. **Barrier synchronization**: All ranks block while rank 0 builds indices. On a 512-node job, this means 4,095 GPUs sit idle. +2. **Filesystem metadata storms**: When blending many datasets, thousands of simultaneous `open()` and `stat()` calls from all ranks can overwhelm NFS/Lustre metadata servers. +3. **Simultaneous memory-mapping**: All ranks `mmap` three large `.npy` files at once after the barrier, causing a burst of page faults and I/O. + +## Baseline: Establish Maximum Achievable Performance + +Before tuning data loading, establish a performance ceiling by running with `--mock-data`. This bypasses the data pipeline entirely and shows the maximum throughput your configuration can achieve without any dataloader overhead. The gap between `--mock-data` performance and real-data performance tells you exactly how much time the dataloader is costing you. + +## Recommended Configuration + +### Step 1: Consolidate dataset files + +A common issue at scale is having datasets split across many small file prefixes. Thousands of 100 MB files perform significantly worse than tens of 10 GB+ files, both for building dataset caches and for runtime file access. + +Use the merge tool to consolidate: + +```bash +python tools/merge_datasets.py --input --output +``` + +**Target at least 10 GB per file.** This reduces the number of file descriptors, metadata lookups, and index-building work at initialization. + +### Step 2: Pre-warm the dataset cache + +Build the index cache as a separate step before training. This can be done either within the first training job (the traditional approach -- rank 0 builds, others wait) or as a dedicated pre-build step: + +```bash +# Pre-build per-dataset sequence counts +python tools/build_sequences_per_dataset.py \ + --data-path \ + --output sequences.json +``` + +### Step 3: Launch training with optimized data loading + +Once the cache is ready, enable the fast-path flags: + +```bash +torchrun --nproc_per_node=8 --nnodes=512 ... pretrain_gpt.py \ + --dataloader-fast-cache-load \ + --dataloader-defer-npy-index-mmap \ + --per-dataset-sequences-path sequences.json \ + --data-cache-path /path/to/cache \ + --num-workers 2 \ + ... +``` + +### Flag reference + +| Flag | Default | Recommendation | What it does | +|------|---------|----------------|-------------| +| `--dataloader-fast-cache-load` | off | **On** | Skips the rank-0 barrier by assuming the cache already exists. All ranks build their dataset views in parallel. This is the single biggest win at scale. | +| `--dataloader-defer-npy-index-mmap` | off | **On** | Defers memory-mapping of `.npy` index files until first access. When combined with `--num-workers > 0`, index loading is overlapped with the training iteration rather than blocking startup. | +| `--per-dataset-sequences-path` | None | **Set** | Points to a JSON file mapping each dataset path to its `(sequence_count, document_count)`. Replaces per-file metadata reads with a single JSON lookup. Critical when blending hundreds of datasets. Generate with `tools/build_sequences_per_dataset.py`. See the [PR description](https://github.com/NVIDIA/Megatron-LM/pull/2445) for the expected file format. | +| `--data-cache-path` | None | **Set** | Directory where index `.npy` files are cached. Must be on shared storage for multi-node jobs so all ranks can read it. | +| `--num-workers` | 2 | **Keep as small as necessary** | Number of DataLoader worker processes. The goal is to satisfy: *time to process a batch > time to prepare a batch*. This hides dataloader work behind the training step. Increasing beyond what's needed wastes CPU and memory. | +| `--no-mmap-bin-files` | mmap on | **Test both** | Memory-mapping `.bin` files leverages the OS page cache, but the optimal setting is filesystem-dependent. Some large-scale production configurations disable mmap. Test with and without to determine what works best for your storage. | + +### Object storage (S3 / Multi-Storage Client) + +When data lives on S3 or MSC rather than a POSIX filesystem: + +- **Index files** (`.idx`) are small and downloaded once by rank 0, then shared via the filesystem cache. +- **Binary data files** (`.bin`) are streamed on-demand in 256 MB chunks, avoiding the need to download entire files. +- Set `--no-mmap-bin-files` since memory-mapping doesn't apply to object storage. +- The rank-0-builds-then-barrier pattern still applies for index construction. + +## Scaling Characteristics + +| Aspect | Behavior | Why it works | +|--------|----------|-------------| +| **Cross-rank contention** | None after init | All index files are read-only; `numpy.memmap` uses OS page cache with no locking | +| **Sampling determinism** | All ranks produce the same permutation | Shared `numpy.random.RandomState(seed)` with epoch-based seed variation | +| **Data-parallel sharding** | Each DP rank gets a disjoint chunk | `indices[dp_rank :: dp_size]` -- no overlap, no coordination | +| **Index broadcast** | Via shared filesystem, not collectives | Rank 0 writes `.npy` files; other ranks read them. No explicit `torch.distributed.broadcast` | + +## Troubleshooting + +**Symptom: Training hangs at startup for minutes** +- Likely cause: Rank 0 is building indices while all other ranks wait at the barrier. +- Fix: Pre-warm the cache and enable `--dataloader-fast-cache-load`. + +**Symptom: Metadata server errors or slow `open()` calls** +- Likely cause: Thousands of ranks simultaneously opening per-dataset files for metadata. +- Fix: Use `--per-dataset-sequences-path` to consolidate metadata into a single JSON file. Also consolidate small dataset files using `tools/merge_datasets.py`. + +**Symptom: Spike in I/O at training start, then normal** +- Likely cause: All ranks simultaneously memory-mapping index files after the barrier. +- Fix: Enable `--dataloader-defer-npy-index-mmap` to overlap index loading with training. + +**Symptom: Slow data loading during training (not just startup)** +- Run with `--mock-data` to confirm the dataloader is the bottleneck. +- Check `--num-workers` is > 0 to hide prefetching behind the training step. +- Verify data is on fast storage (local SSD or high-bandwidth parallel filesystem). +- Test with `--no-mmap-bin-files` -- the optimal setting depends on your filesystem. +- Check if datasets are split across many small files. Merge to 10 GB+ per file with `tools/merge_datasets.py`. + +## Related Resources + +- [PR #2445](https://github.com/NVIDIA/Megatron-LM/pull/2445): Original implementation of fast cache load, deferred mmap, and per-dataset sequences optimizations. +- [PR #4080](https://github.com/NVIDIA/Megatron-LM/pull/4080): Dedicated script for pre-building dataset caches (pending review). +- [`tools/merge_datasets.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/tools/merge_datasets.py): Merge multiple small dataset files into larger ones. +- [`tools/build_sequences_per_dataset.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/tools/build_sequences_per_dataset.py): Generate the `--per-dataset-sequences-path` JSON file. From e71a5b7d7783d853c2001db2162ee421bd128289 Mon Sep 17 00:00:00 2001 From: Santosh Bhavani Date: Thu, 9 Apr 2026 07:45:05 -0700 Subject: [PATCH 2/5] docs: align data loading guide with offline cache prep --- docs/user-guide/data-loading.md | 52 ++++++++++++++++++++++++--------- 1 file changed, 38 insertions(+), 14 deletions(-) diff --git a/docs/user-guide/data-loading.md b/docs/user-guide/data-loading.md index b197e8e48cb..6d9c746a05a 100644 --- a/docs/user-guide/data-loading.md +++ b/docs/user-guide/data-loading.md @@ -21,7 +21,7 @@ Megatron builds three index arrays for each dataset: a **document index** (shuff 2. All ranks synchronize at a `torch.distributed.barrier()`. 3. **All other ranks** load the cached indices via memory-mapped reads (`numpy.load(mmap_mode='r')`). -After initialization, data access is **read-only and lock-free**. Every data-parallel rank reads a disjoint, contiguous chunk of the shuffle index. No cross-rank coordination is needed during training because all ranks derive the same deterministic permutation from a shared random seed. +After initialization, data access is **read-only and lock-free**. Each data-parallel rank consumes a disjoint subset of samples, and no cross-rank coordination is needed during training because all ranks derive the same deterministic permutation from a shared random seed. ## The Problem at 256+ Nodes @@ -41,26 +41,49 @@ Before tuning data loading, establish a performance ceiling by running with `--m A common issue at scale is having datasets split across many small file prefixes. Thousands of 100 MB files perform significantly worse than tens of 10 GB+ files, both for building dataset caches and for runtime file access. -Use the merge tool to consolidate: +Use the merge tool to consolidate datasets stored as many small prefixes in one directory: ```bash -python tools/merge_datasets.py --input --output +python tools/merge_datasets.py \ + --input /path/to/input-directory \ + --output-prefix /path/to/output/merged ``` **Target at least 10 GB per file.** This reduces the number of file descriptors, metadata lookups, and index-building work at initialization. -### Step 2: Pre-warm the dataset cache +### Step 2: Pre-build the dataset cache -Build the index cache as a separate step before training. This can be done either within the first training job (the traditional approach -- rank 0 builds, others wait) or as a dedicated pre-build step: +Build the GPT dataset cache as a separate step before training. This avoids the usual "rank 0 builds, everyone else waits" startup path and is the recommended workflow for large jobs: + +```bash +python tools/prepare_cache.py \ + --data-path \ + --split 99,1,0 \ + --data-cache-path /path/to/cache \ + --global-batch-size \ + --seq-length \ + ... +``` + +If your later training job does not set `--global-batch-size`, or you are preparing the cache on a machine that does not match the future training topology, also pass: + +```bash +--prepare-cache-world-size +``` + +This keeps the prepared cache aligned with the sample counts expected by training. + +### Step 3: Optionally pre-build per-dataset metadata + +When blending many datasets, generate the `--per-dataset-sequences-path` JSON ahead of time to avoid one metadata read per file prefix at startup: ```bash -# Pre-build per-dataset sequence counts python tools/build_sequences_per_dataset.py \ --data-path \ - --output sequences.json + --per-dataset-sequences-path sequences.json ``` -### Step 3: Launch training with optimized data loading +### Step 4: Launch training with optimized data loading Once the cache is ready, enable the fast-path flags: @@ -80,7 +103,7 @@ torchrun --nproc_per_node=8 --nnodes=512 ... pretrain_gpt.py \ |------|---------|----------------|-------------| | `--dataloader-fast-cache-load` | off | **On** | Skips the rank-0 barrier by assuming the cache already exists. All ranks build their dataset views in parallel. This is the single biggest win at scale. | | `--dataloader-defer-npy-index-mmap` | off | **On** | Defers memory-mapping of `.npy` index files until first access. When combined with `--num-workers > 0`, index loading is overlapped with the training iteration rather than blocking startup. | -| `--per-dataset-sequences-path` | None | **Set** | Points to a JSON file mapping each dataset path to its `(sequence_count, document_count)`. Replaces per-file metadata reads with a single JSON lookup. Critical when blending hundreds of datasets. Generate with `tools/build_sequences_per_dataset.py`. See the [PR description](https://github.com/NVIDIA/Megatron-LM/pull/2445) for the expected file format. | +| `--per-dataset-sequences-path` | None | **Set when blending many datasets** | Points to a JSON file mapping each dataset path to its `(sequence_count, document_count)`. Replaces per-file metadata reads with a single JSON lookup. Generate with `tools/build_sequences_per_dataset.py`. | | `--data-cache-path` | None | **Set** | Directory where index `.npy` files are cached. Must be on shared storage for multi-node jobs so all ranks can read it. | | `--num-workers` | 2 | **Keep as small as necessary** | Number of DataLoader worker processes. The goal is to satisfy: *time to process a batch > time to prepare a batch*. This hides dataloader work behind the training step. Increasing beyond what's needed wastes CPU and memory. | | `--no-mmap-bin-files` | mmap on | **Test both** | Memory-mapping `.bin` files leverages the OS page cache, but the optimal setting is filesystem-dependent. Some large-scale production configurations disable mmap. Test with and without to determine what works best for your storage. | @@ -89,10 +112,10 @@ torchrun --nproc_per_node=8 --nnodes=512 ... pretrain_gpt.py \ When data lives on S3 or MSC rather than a POSIX filesystem: -- **Index files** (`.idx`) are small and downloaded once by rank 0, then shared via the filesystem cache. +- **Index files** (`.idx`) are cached locally under `object_storage_cache_path`. - **Binary data files** (`.bin`) are streamed on-demand in 256 MB chunks, avoiding the need to download entire files. - Set `--no-mmap-bin-files` since memory-mapping doesn't apply to object storage. -- The rank-0-builds-then-barrier pattern still applies for index construction. +- Ensure the index-cache path is visible wherever the later dataset construction will run. ## Scaling Characteristics @@ -100,14 +123,14 @@ When data lives on S3 or MSC rather than a POSIX filesystem: |--------|----------|-------------| | **Cross-rank contention** | None after init | All index files are read-only; `numpy.memmap` uses OS page cache with no locking | | **Sampling determinism** | All ranks produce the same permutation | Shared `numpy.random.RandomState(seed)` with epoch-based seed variation | -| **Data-parallel sharding** | Each DP rank gets a disjoint chunk | `indices[dp_rank :: dp_size]` -- no overlap, no coordination | +| **Data-parallel sharding** | Each DP rank gets a disjoint subset of samples | No overlap during training; assignment happens in the sampler rather than via extra dataset coordination | | **Index broadcast** | Via shared filesystem, not collectives | Rank 0 writes `.npy` files; other ranks read them. No explicit `torch.distributed.broadcast` | ## Troubleshooting **Symptom: Training hangs at startup for minutes** - Likely cause: Rank 0 is building indices while all other ranks wait at the barrier. -- Fix: Pre-warm the cache and enable `--dataloader-fast-cache-load`. +- Fix: Pre-build the cache with `tools/prepare_cache.py` and enable `--dataloader-fast-cache-load`. **Symptom: Metadata server errors or slow `open()` calls** - Likely cause: Thousands of ranks simultaneously opening per-dataset files for metadata. @@ -127,6 +150,7 @@ When data lives on S3 or MSC rather than a POSIX filesystem: ## Related Resources - [PR #2445](https://github.com/NVIDIA/Megatron-LM/pull/2445): Original implementation of fast cache load, deferred mmap, and per-dataset sequences optimizations. -- [PR #4080](https://github.com/NVIDIA/Megatron-LM/pull/4080): Dedicated script for pre-building dataset caches (pending review). +- [PR #4080](https://github.com/NVIDIA/Megatron-LM/pull/4080): Adds `tools/prepare_cache.py` for offline GPT dataset cache preparation. +- [`tools/prepare_cache.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/tools/prepare_cache.py): Pre-build GPT dataset caches ahead of training. - [`tools/merge_datasets.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/tools/merge_datasets.py): Merge multiple small dataset files into larger ones. - [`tools/build_sequences_per_dataset.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/tools/build_sequences_per_dataset.py): Generate the `--per-dataset-sequences-path` JSON file. From 33c195e68987e2e8e5b3964b7866558e4fd72b15 Mon Sep 17 00:00:00 2001 From: Santosh Bhavani Date: Thu, 9 Apr 2026 09:56:41 -0700 Subject: [PATCH 3/5] docs: tighten data loading troubleshooting guidance --- docs/user-guide/data-loading.md | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/docs/user-guide/data-loading.md b/docs/user-guide/data-loading.md index 6d9c746a05a..08c27f989d1 100644 --- a/docs/user-guide/data-loading.md +++ b/docs/user-guide/data-loading.md @@ -134,7 +134,7 @@ When data lives on S3 or MSC rather than a POSIX filesystem: **Symptom: Metadata server errors or slow `open()` calls** - Likely cause: Thousands of ranks simultaneously opening per-dataset files for metadata. -- Fix: Use `--per-dataset-sequences-path` to consolidate metadata into a single JSON file. Also consolidate small dataset files using `tools/merge_datasets.py`. +- Fix: Use `--per-dataset-sequences-path` to consolidate metadata into a single JSON file. **Symptom: Spike in I/O at training start, then normal** - Likely cause: All ranks simultaneously memory-mapping index files after the barrier. @@ -142,10 +142,9 @@ When data lives on S3 or MSC rather than a POSIX filesystem: **Symptom: Slow data loading during training (not just startup)** - Run with `--mock-data` to confirm the dataloader is the bottleneck. -- Check `--num-workers` is > 0 to hide prefetching behind the training step. -- Verify data is on fast storage (local SSD or high-bandwidth parallel filesystem). +- If startup, not steady-state throughput, is the main issue, try `--dataloader-defer-npy-index-mmap`. +- If you are blending many dataset prefixes, try `--per-dataset-sequences-path`. - Test with `--no-mmap-bin-files` -- the optimal setting depends on your filesystem. -- Check if datasets are split across many small files. Merge to 10 GB+ per file with `tools/merge_datasets.py`. ## Related Resources From 4459dfc62e2f8508d8ffc1acbfc47ec8bde7b311 Mon Sep 17 00:00:00 2001 From: Santosh Bhavani Date: Thu, 9 Apr 2026 10:04:02 -0700 Subject: [PATCH 4/5] docs: move data loading guide to advanced features --- docs/index.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/index.md b/docs/index.md index 079fdfbecae..0dbf7d2e3b7 100644 --- a/docs/index.md +++ b/docs/index.md @@ -48,7 +48,6 @@ get-started/quickstart :caption: Basic Usage user-guide/data-preparation -user-guide/data-loading user-guide/training-examples user-guide/parallelism-guide ``` @@ -73,6 +72,7 @@ user-guide/features/dist_optimizer user-guide/features/optimizer_cpu_offload user-guide/features/pipeline_parallel_layout user-guide/features/fine_grained_activation_offloading +user-guide/data-loading user-guide/features/megatron_energon user-guide/features/megatron_rl user-guide/features/tokenizers @@ -104,4 +104,4 @@ apidocs/index.rst :caption: Resources advanced/index -``` \ No newline at end of file +``` From 7323f9cb131652b43c3d172c1e94f90cb82c9084 Mon Sep 17 00:00:00 2001 From: Santosh Bhavani Date: Tue, 14 Apr 2026 09:35:53 -0700 Subject: [PATCH 5/5] docs: remove customer-specific metadata storm guidance --- docs/user-guide/data-loading.md | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/docs/user-guide/data-loading.md b/docs/user-guide/data-loading.md index 08c27f989d1..1f0d544317c 100644 --- a/docs/user-guide/data-loading.md +++ b/docs/user-guide/data-loading.md @@ -9,7 +9,7 @@ # Data Loading at Scale -This guide covers how Megatron's data pipeline works and how to configure it for efficient training at 256 nodes and beyond. At this scale, the primary bottlenecks are **index building**, **filesystem metadata operations**, and **barrier synchronization** -- not raw data bandwidth. +This guide covers how Megatron's data pipeline works and how to configure it for efficient training at 256 nodes and beyond. At this scale, the primary bottlenecks are **index building** and **barrier synchronization** -- not raw data bandwidth. ## How Data Loading Works @@ -28,8 +28,7 @@ After initialization, data access is **read-only and lock-free**. Each data-para Three things break down at large node counts: 1. **Barrier synchronization**: All ranks block while rank 0 builds indices. On a 512-node job, this means 4,095 GPUs sit idle. -2. **Filesystem metadata storms**: When blending many datasets, thousands of simultaneous `open()` and `stat()` calls from all ranks can overwhelm NFS/Lustre metadata servers. -3. **Simultaneous memory-mapping**: All ranks `mmap` three large `.npy` files at once after the barrier, causing a burst of page faults and I/O. +2. **Simultaneous memory-mapping**: All ranks `mmap` three large `.npy` files at once after the barrier, causing a burst of page faults and I/O. ## Baseline: Establish Maximum Achievable Performance @@ -132,10 +131,6 @@ When data lives on S3 or MSC rather than a POSIX filesystem: - Likely cause: Rank 0 is building indices while all other ranks wait at the barrier. - Fix: Pre-build the cache with `tools/prepare_cache.py` and enable `--dataloader-fast-cache-load`. -**Symptom: Metadata server errors or slow `open()` calls** -- Likely cause: Thousands of ranks simultaneously opening per-dataset files for metadata. -- Fix: Use `--per-dataset-sequences-path` to consolidate metadata into a single JSON file. - **Symptom: Spike in I/O at training start, then normal** - Likely cause: All ranks simultaneously memory-mapping index files after the barrier. - Fix: Enable `--dataloader-defer-npy-index-mmap` to overlap index loading with training.