NVIDIA / NVIDIA/Megatron-LM

Support time-based training without requiring large `train_iters`

Open
#3,488 4 comments 0 reactions 0 assignees View on GitHub
bug community-request waiting-on-customer
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

Hello all, thank you for the nice library.

**Describe the bug**

When train_iters is large relative to global_batch_size (e.g. time-based training on a single node with exit_duration_in_mins), _build_document_index in gpt_dataset.py:577 allocates a numpy.mgrid array of shape (num_epochs, num_documents) that OOMs:

numpy.core._exceptions._ArrayMemoryError: Unable to allocate 373. GiB
for an array with shape (2, 255884, 97879) and data type int64

**Steps/Code to reproduce bug**

Single node, GBS=8, large train_iters:

```
python pretrain_gpt.py \
--train-iters 10000000 \
--global-batch-size 8 \
--seq-length 4096 \
--exit-duration-in-mins 60 \
... # any dataset with >100k documents
```

The high train_iters / GBS ratio requires many epochs, causing _build_document_index to pre-allocate the full num_epochs × num_documents grid.

**Expected behavior**

Training should start and run until exit_duration_in_mins without OOM on the document index. Either:
- Cap the document index allocation to a reasonable number of epochs and extend lazily
- Support train_iters=None when exit_duration_in_mins is set, building epochs on demand

**Additional context**

Common on HPC clusters where jobs use walltime limits and exit_duration_in_mins is the natural stopping mechanism. Users are forced to carefully tune `train_iters` between "large enough to not stop early" and "small enough to not OOM the document index".

Using `train_samples` doesn't help since Bridge converts it to `train_iters` via `train_samples // global_batch_size`.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.