[BUG] Infinite Loop in `_get_num_epochs` Function of `GPTDataset` Class When `num_tokens_per_epoch` is Zero
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Describe the bug**
In `megatron/core/datasets/gpt_dataset.py`, the `GPTDatase` class has a function `_get_num_epochs `which can result in an infinite loop when `num_tokens_per_epoch` is set to 0. This occurs without any explicit notification. This function is invoked in `_build_document_sample_shuffle_indices` when building a `GPTDataset` instance and saving caches.
https://github.com/NVIDIA/Megatron-LM/blob/85bd1f9af2b6a85d5f348509cca5a6251e0cbef1/megatron/core/datasets/gpt_dataset.py#L550-L552
`num_tokens_per_epoch` is set to zero when the split setting allocates a very small percentage of the valid/test dataset split. For example, using a split option like 9998, 0.002, 0.002 can cause `num_tokens_per_epoch` to be zero if the dataset size is insufficient.
https://github.com/NVIDIA/Megatron-LM/blob/85bd1f9af2b6a85d5f348509cca5a6251e0cbef1/examples/mixtral/train_mixtral_8x7b_distributed.sh#L65
Adjusting the configuration to `9990, 8, 2` or `9900, 80, 20` based on the actual data amount resolves this issue.
**To Reproduce**
This issue was encountered while training Mixtral-MoE 8x7B. It can be reproduced as long as the dataset configuration matches the problematic split and is unrelated to the model itself. To replicate, follow the instructions on the [Mixtral-MoE 8x7B training page](https://github.com/NVIDIA/Megatron-LM/blob/main/examples/mixtral/README.md). The following dataset was used:
```
wget https://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/mistral-datasets/wudao_mistralbpe_content_document.bin
wget https://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/mistral-datasets/wudao_mistralbpe_content_document.idx
```
The training was launched on 4 nodes, each containing 8xA100 GPUs.
**Expected behavior**
A clear warning should be provided indicating the possibility of an infinite loop along with guidance to modify the configuration to resolve the issue.
**Proposed fix**
Raise an error at an appropriate place to handle this issue. Possible locations include:
1. In function `_get_num_epochs`, detect if an infinite loop will occur directly.
2. In function `_get_num_tokens_per_epoch` which calculates the variable `num_tokens_per_epoch`.
3. During the initialization of the `GPTDataset` class, if `self.indices` is an empty array, it will cause the `_get_num_tokens_per_epoch` function to return 0
Contributor guide
Assessment
This issue has not been assessed yet.