Support time-based training without requiring large `train_iters`
- 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
Assessment
This issue has not been assessed yet.