NVIDIA / NVIDIA/Megatron-LM

[BUG] Infinite Loop in `_get_num_epochs` Function of `GPTDataset` Class When `num_tokens_per_epoch` is Zero

Open
#957 3 comments 2 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

**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

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.