[bug] IndexedDataset numpy.memmap raises EPERM on WekaFS after ~600 datasets in a large blend; mmap_bin_files=false works around it

Start with megatron/core/datasets/indexed_dataset.py at _IndexReader.init and the mmap_bin_files handling in the dataset configuration; review blended_dataset.py for related cache and mmap options. Reproduce the large WekaFS blend with mmap_bin_files enabled, then compare it with the documented workaround disabled. Done means either a focused mitigation or documentation that accurately explains the parallel-filesystem limitation and validates the existing configuration path.

Problem

When constructing a BlendedMegatronDataset over a large blend (444 prefixes × train+valid = 888 IndexedDataset builds) backed by a WekaFS mount, IndexedDataset.__init__ consistently fails with PermissionError: [Errno 1] Operation not permitted at the numpy.memmap(idx_path, ...) call in _IndexReader.__init__ (Megatron-Core indexed_dataset.py:280).

Failure is reproducible single-rank (no concurrency), after ~590 successful per-dataset builds. The same numpy.memmap calls succeed in isolation on the same files; failure only occurs once enough IndexedDataset instances have been built in one process. Setting BlendedMegatronDatasetConfig.mmap_bin_files=False eliminates the failure: all 888 builds complete and training runs to a clean checkpoint.

The symptom looks like a per-process resource ceiling in the WekaFS client that scales with the number of live mmap regions held by the process. Megatron's IndexedDataset holds two mmaps per dataset (.idx always mmapped via _IndexReader; .bin mmapped via _MMapBinReader when mmap_bin_files=True). Cutting the .bin mmap roughly halves the count and pushes well past the threshold.

We could not find an upstream issue describing this signature. Filing it because the workaround is non-obvious (the existing mmap_bin_files flag is documented for S3, not for parallel-FS mmap exhaustion).

Minimal repro

1. Stage a Megatron `.bin`/`.idx` blend with >600 prefixes on a WekaFS mount.
2. Configure Megatron-Bridge training over that blend with default
   `BlendedMegatronDatasetConfig` (i.e. `mmap_bin_files=True`).
3. Launch training (any rank count, including single rank).
4. During dataset construction, observe a hard failure on one rank around
   the ~590th `IndexedDataset.__init__` call:

   File ".../megatron/core/datasets/indexed_dataset.py", line 280, in __init__
       self.bin_buffer_mmap = numpy.memmap(idx_path, mode="r", order="C")
     File ".../numpy/_core/memmap.py", line 291, in __new__
       mm = mmap.mmap(fid.fileno(), bytes, access=acc, offset=start)
   PermissionError: [Errno 1] Operation not permitted

5. Re-run with `BlendedMegatronDatasetConfig.mmap_bin_files=False`
   (CLI `--no-mmap-bin-files` in Megatron-LM, or set the field directly in
   Megatron-Bridge config). All 888 builds complete; training proceeds.

Expected behavior

IndexedDataset construction over a large blend should complete regardless of dataset count (within memory). At minimum, the workaround should be discoverable: the docstring/help text for mmap_bin_files should mention that disabling it can also be required for parallel-FS-backed data, not only S3.

A direct fix (if root cause is confirmed FS-side) would be a WEKA-client-level fix; from the framework side, prefer not to hold one mmap per dataset for the lifetime of the build phase. Two cheap framework-side mitigations worth considering:

1. Document mmap_bin_files=False as a known mitigation for parallel-filesystem mmap exhaustion (Lustre, WekaFS, etc.), not only S3.
2. Optionally release the .bin mmap between draws (or build it lazily on first sample) so the live mmap count during the build phase scales with concurrent draws, not with blend size.

Affected area

area:data

Regression?

Not sure

Environment

Branch / commit: Megatron-Bridge (NVIDIA-NeMo/Megatron-Bridge, recent main)
Megatron-Core: core/datasets/indexed_dataset.py line 280
(the _IndexReader.__init__ mmap of the .idx file)
Container: nemo:26.04 (NeMo container, NVIDIA published)
Filesystem: WekaFS POSIX mount (/mnt/weka/...)
Cluster: On-prem SLURM, H100/H200-class nodes
GPU / machine: Single-rank repro on 1 GPU; 8-rank repro on 1 node × 8 GPUs
Python / CUDA / PyTorch: container-default (Python 3.12; torch/CUDA per nemo:26.04)
vm.max_map_count: 65530 (far above any observed mmap count)
Blend size: 444 prefixes × {train, valid} = 888 IndexedDataset builds
Data layout: Megatron .bin / .idx pairs on WekaFS, sizes
~MBs (.idx) and ~hundreds of MBs to GBs (.bin) per prefix

Logs

Traceback (most recent call last):
  ...
  File ".../megatron/core/datasets/indexed_dataset.py", line 280, in __init__
    self.bin_buffer_mmap = numpy.memmap(idx_path, mode="r", order="C")
  File ".../numpy/_core/memmap.py", line 291, in __new__
    mm = mmap.mmap(fid.fileno(), bytes, access=acc, offset=start)
PermissionError: [Errno 1] Operation not permitted


Failure point in two distinct runs:

- 8-rank torchrun: rank 4 dies after ~563 successful `_IndexReader` loads.
- 1-rank single-process prepare-cache: dies after 592 successful loads, on `VI_Wikipedia_GGI-head-25.idx`.

After setting `mmap_bin_files=False` (no other changes), the same blend completes all 888 builds, runs the training loop iter 0 → 5, and writes a checkpoint cleanly. No EPERM, no retry needed.

## Already tried (other dataloader optimization flags)

Before settling on `mmap_bin_files=False`, we checked the three upstream
"speed up DataLoader init" flags advertised for large-scale runs. None of
them addresses this failure mode:

| Flag | Tested? | Outcome |
|---|---|---|
| `--data-cache-path` (`config.path_to_cache`) | **Yes — empirically.** Single-rank prepare-cache run with `path_to_cache` first redirected to `/dev/shm`, then to a separate WekaFS cache directory off the data root. | **EPERM still reproduces** at the same `indexed_dataset.py:280` mmap call, after the same ~590-build accumulation. The failing mmap is the per-prefix `.idx` next to `.bin` at the data root — *not* the per-blend `.npy` caches under `path_to_cache`. Redirecting the cache cannot reduce the count of the failing mmaps. |
| `--dataloader-defer-npy-index-mmap` (`config.defer_npy_index_mmap`) | No — reasoned only. | Defers mmap of the per-blend `dataset_index.npy` / `dataset_sample_index.npy` (a handful per blend, see `blended_dataset.py:89` and `:97-104`). Does not touch the 444 per-prefix `.idx`/`.bin` mmaps that drive the accumulation. Wrong granularity for this failure. |
| `--dataloader-fast-cache-load` (`config.fast_cache_load`) | No. | Short-circuits the cache-existence stat check (`blended_dataset.py:140-149`). Cache-hit fast path, orthogonal to mmap count. |

Only `mmap_bin_files=False` reduces the per-process live mmap count
per-dataset, which is what the WekaFS client appears sensitive to.

## Workaround we are using

Set `BlendedMegatronDatasetConfig.mmap_bin_files = False` (equivalent CLI in Megatron-LM: `--no-mmap-bin-files`). This switches `.bin` access from `_MMapBinReader` to `_FileBinReader` (pread-style `os.pread` per draw). `.idx` remains mmapped (the mmap in `_IndexReader.__init__` is unconditional in current Megatron-Core), but halving the per-dataset mmap count is enough to clear the WekaFS threshold for our 444-prefix blend.

Performance caveat we have not yet quantified: `_FileBinReader` does a syscall per sample draw vs. page-fault-on-demand from a mapped region. For dataset construction the cost is negligible; for steady-state throughput we plan to benchmark.