NVIDIA / NVIDIA/Megatron-LM

[ENHANCEMENT] Preprocessing data that is already partitioned and gzipped

Open
#1,135 1 comment 0 reactions 0 assignees View on GitHub
enhancement
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

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.