[ENHANCEMENT] Preprocessing data that is already partitioned and gzipped
- 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.**
Using [tools/preprocess_data.py](https://github.com/NVIDIA/Megatron-LM/blob/76f9f48939ba5ecff0fed7bfbd4204df05d3e4da/tools/preprocess_data.py), I would like to preprocess the [Dolma](https://allenai.github.io/dolma/) dataset, which is already split into ~2000 `.json.gz` files, without rewriting the entire dataset to intermediate partition `.json` files. To achieve this, I would like to set `args.workers` << `args.partitions` because it is infeasible to have 2000+ workers.
**Describe the solution you'd like**
Removing the following constraint to support preprocessing data that is already partitioned: https://github.com/NVIDIA/Megatron-LM/blob/76f9f48939ba5ecff0fed7bfbd4204df05d3e4da/tools/preprocess_data.py#L343
**Describe alternatives you've considered**
Possible alternative: preprocessing the files in separate batches where for each batch `args.workers % args.partitions == 0`. This can be achieved by creating symlinks to each file that follow the appropriate partition file naming convention (`_{id}.json`) and running preprocess_data.py independently for each batch.
**Proposed implementation**
A possible solution is to modify the preprocessing logic to remove the constraint `args.workers % args.partitions == 0` and directly process each of the input files.
**Additional context**
Something to note is that when the input is of the form`*.json.gz` then extension is currently calculated as ".gz" and therefore the partition file pattern is `*.json_{id}.gz` https://github.com/NVIDIA/Megatron-LM/blob/76f9f48939ba5ecff0fed7bfbd4204df05d3e4da/tools/preprocess_data.py#L252
Contributor guide
Assessment
This issue has not been assessed yet.