NVIDIA-NeMo / NVIDIA-NeMo/Curator

Add additional write stage for merging multiple Megatron `IndexedDataset` file pairs into one

Open
#1,399 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

community-request enhancement waiting-on-maintainers
Dominant language
Python
Stars
1.8k
Forks
328
Avg merge
4d 5h
Merged PRs (30d)
30

Description

Is your feature request related to a problem? Please describe.

Currently, when I use the MegatronTokenizerWriter, it creates individual .bin/.idx file pairs, for each worker. This is obviously needed for efficiency, but when running on a machine with a lot of vCPUs, I got >100 different file pairs - which make it difficult to work with.

Describe the solution you'd like

It would be very helpful to be able to add some post-processing stage which merges all of them into a single output file. It can either be part of the same stage, or an additional optional stage in the pipeline.

Describe alternatives you've considered

A workaround solution is to run the merge_datasets.py script from Megatron-LM. Ideally if this could be part of the pipeline, it would simplify the process - especially since installing megatron-core requires a lot of dependencies, and it would be simpler to work within a single env with NeMo-Curator.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the MegatronTokenizerWriter and review the referenced Megatron-LM merge_datasets.py script to understand the expected merge behavior. Add an optional post-processing stage that combines the worker-generated .bin/.idx pairs into one output pair, then verify that it works within a NeMo-Curator pipeline without requiring megatron-core.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data-engineering
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.