NVIDIA-NeMo / NVIDIA-NeMo/Curator
Support batchsize/process_batch in the minHash stage
Nobody has claimed this yet.
- 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.
GPUs typically do better with larger batches (1-3GB) for processing but CPU stages do better with smaller (256MB or so) batches or run into OOMs.
To avoid this we'd like to implelement/support proccess_batch for the MinHashStage to ensure that we can use a smaller blocksize during read but saturate the GPU better for minhash.
Describe the solution you'd like
Implement process_batch for the minHash stage that reads the files represented by all the incoming task.data one by one and then concat it. There's still some penalty of having to read it one by one (read not saturated), but the minhash and write can be on larger blocks.
Describe alternatives you've considered
Similar approach to #1332 but arguably simpler since it's the typically process/process_batch type stage.
Additional context
Add any other context or screenshots about the feature request here.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the MinHashStage entry point and inspect how its existing process path handles task.data. Compare the process/process_batch pattern described in issue #1332, then define done as process_batch reading each represented file, concatenating the inputs, and allowing larger MinHash and write blocks while retaining smaller read blocks.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- data-engineering, performance
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 65/100