[ENHANCEMENT] S3 data loading for MMapIndexedDataset
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Is your feature request related to a problem? Please describe.**
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
Megatron-LM does not support loading a dataset from S3.
**Describe the solution you'd like**
A clear and concise description of what you want to happen.
I would like to extend MMapIndexedDataset to support S3 data loading.
**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.
A user can download the dataset from S3 to the local file system at the start of training, but that blocks training until the download is complete. Alternatively, a user can store their dataset in a cloud file system. That can work well, but requires managing the cloud file system in addition to S3.
**Proposed implementation**
If you have a proposed implementation for the feature state it here or link to a PR.
I implemented S3 data loading in a private fork of NeMo initially. I have a public, draft PR with that implementation [here](https://github.com/NVIDIA/NeMo/pull/8447). However, NeMo now uses MMapIndexedDataset in Megatron-LM directly, so ~~I would like to port a similar implementation to Megatron-LM~~ I also ported a similar implementation to Megatron-LM [here](https://github.com/NVIDIA/Megatron-LM/pull/729). In particular, the index file is downloaded to the local file system so that it can be memory mapped and the bin file is streamed in chunks from S3. Note that the block shuffling functionality described in the NeMo PR is optional (we can just assume that the user has preshuffled the dataset).
**Additional context**
Add any other context or screenshots about the feature request here.
Contributor guide
Assessment
This issue has not been assessed yet.