TPS-free 2D bucket estimation and filtering by pzelasko · Pull Request #11738 · NVIDIA-NeMo/NeMo

pzelasko · 2025-01-02T17:17:28Z

What does this PR do ?

Improves the UX of estimating 2D buckets and preparing the training configuration. Key changes:

2D bucket estimation auto-finds the right max-token-per-second setting for each bucket separately instead of assuming a global constant max-TPS. The per-bucket max-TPS is determined as 4*stddev(bucket_tps).
Training with 2D bucketing works without setting the max_tps filter.
The data that doesn't fit the 2D buckets (either 1st or 2nd dim) is filtered out automatically during sampling.
- strict mode (default): first allocate the duration bucket, and then try to allocate to max_tokens sub-bucket. If an example doesn't fit, discard it.
- lenient mode (model.train_ds.bucketing_2d_strict_mode=False): find any bucket that will fit a given example. That means token-per-second outliers are pushed to buckets with higher durations increasing the padding but reducing the amount of discarded data. Use at your own risk - for the setups I tested it so far with, it may cause training instability (likely due to inclusion of lower-quality data).
Fixes some issues with estimate_duration_bins_2d.py script.
For data with complex TPS distributions: max-TPS can be defined as a list with the same length as the list of buckets for fine-grained control. A best-guess setting is suggested by estimate_duration_bins_2d.py.
Renamed tarred_random_access to skip_missing_manifest_entries and adjusted the logic to reduce CPU memory usage and accelerate loading time.

Collection: ASR, TTS, SpeechLLM

Changelog

Add specific line by line info of high level changes in this PR.

Usage

You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

New Feature
Bugfix
Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

Related to # (issue)

Signed-off-by: Piotr Żelasko <petezor@gmail.com>

nemo/collections/common/data/lhotse/sampling.py

 import bisect
 import logging
 import math
+from bisect import bisect_left, bisect_right


Signed-off-by: Piotr Żelasko <petezor@gmail.com>

tests/collections/common/test_2d_bucketing_constraint.py

-from lhotse.testing.dummies import DummyManifest, dummy_cut
-from nemo.collections.common.data.lhotse.sampling import FixedBucketBatchSizeConstraint2D
+from lhotse.testing.dummies import dummy_cut
+from lhotse.testing.random import deterministic_rng


Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

Signed-off-by: pzelasko <pzelasko@users.noreply.github.com>

Signed-off-by: Piotr Żelasko <petezor@gmail.com>

…com/nvidia/nemo into 2d-bucketing-and-tps-improvements-2

tests/collections/common/test_2d_bucketing_constraint.py

+    filter_2d = BucketingFilter(constraint)
+    cut = make_cut(duration=2.0, num_tokens=20)
+    assert filter_2d(cut) == False
+    assert constraint.select_bucket(constraint.max_seq_len_buckets, cut) == None


Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

Signed-off-by: Piotr Żelasko <petezor@gmail.com>

nithinraok

LGTM

nithinraok · 2025-01-10T14:33:50Z

nemo/collections/common/data/lhotse/dataloader.py


+    if config.get("tarred_random_access", False):
+        logging.warning(
+            "Option 'tarred_random_access' is deprecated and replaced with 'skip_missing_manifest_entries'.",


may be also add version from which this would be removed

Signed-off-by: Piotr Zelasko <pzelasko@cs-oci-ord-login-01.cm.cluster>

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

Signed-off-by: Piotr Żelasko <petezor@gmail.com>

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

Signed-off-by: Piotr Żelasko <petezor@gmail.com>

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

github-actions · 2025-01-22T15:01:42Z

beep boop 🤖: 🚨 The following files must be fixed before merge!

Your code was analyzed with PyLint. The following annotations have been identified:

************* Module nemo.collections.common.data.lhotse.sampling
nemo/collections/common/data/lhotse/sampling.py:15:0: W0611: Unused import bisect (unused-import)
nemo/collections/common/data/lhotse/sampling.py:16:0: W0611: Unused import logging (unused-import)
nemo/collections/common/data/lhotse/sampling.py:18:0: W0611: Unused bisect_right imported from bisect (unused-import)

-----------------------------------
Your code has been rated at 9.87/10

Mitigation guide:

Add sensible and useful docstrings to functions and methods
For trivial methods like getter/setters, consider adding # pylint: disable=C0116 inside the function itself
To disable multiple functions/methods at once, put a # pylint: disable=C0116 before the first and a # pylint: enable=C0116 after the last.

By applying these rules, we reduce the occurance of this message in future.

Thank you for improving NeMo's documentation!

github-actions · 2025-01-23T05:07:39Z

[🤖]: Hi @pzelasko 👋,

We wanted to let you know that a CICD pipeline for this PR just finished successfully

So it might be time to merge this PR or get some approvals

I'm just a bot so I'll leave it you what to do next.

//cc @pablo-garay @ko3n1g

nithinraok

LGTM

* TPS-free 2D bucket estimation and filtering Signed-off-by: Piotr Żelasko <petezor@gmail.com> * fixes Signed-off-by: Piotr Żelasko <petezor@gmail.com> * extra unit test Signed-off-by: Piotr Żelasko <petezor@gmail.com> * Fixes in 2D bin estimation script and OOMptimizer Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com> * Apply isort and black reformatting Signed-off-by: pzelasko <pzelasko@users.noreply.github.com> * Support list values for max_tps/max_tpt and bucketing Signed-off-by: Piotr Żelasko <petezor@gmail.com> * fix max_tps/max_tpt config parsing for list values Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com> * Transition: tarred_random_access -> skip_missing_manifest_entries Signed-off-by: Piotr Żelasko <petezor@gmail.com> * Improve warning messages Signed-off-by: Piotr Żelasko <petezor@gmail.com> * Support single unified BPE tokenizer for Canary2 Signed-off-by: Piotr Zelasko <pzelasko@cs-oci-ord-login-01.cm.cluster> * bugfix +es to unified tokenizer support for Canary Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com> * CanaryBPETokenizer Signed-off-by: Piotr Żelasko <petezor@gmail.com> * Revert changes to fetching pad id in AED models Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com> * Update docs Signed-off-by: Piotr Żelasko <petezor@gmail.com> * fix tests Signed-off-by: Piotr Żelasko <petezor@gmail.com> * Remove max_tps and max_duration from OOMptimizer output Signed-off-by: Piotr Żelasko <petezor@gmail.com> --------- Signed-off-by: Piotr Żelasko <petezor@gmail.com> Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com> Signed-off-by: pzelasko <pzelasko@users.noreply.github.com> Signed-off-by: Piotr Zelasko <pzelasko@cs-oci-ord-login-01.cm.cluster> Co-authored-by: pzelasko <pzelasko@users.noreply.github.com> Signed-off-by: Parth Mannan <pmannan@nvidia.com>

* TPS-free 2D bucket estimation and filtering Signed-off-by: Piotr Żelasko <petezor@gmail.com> * fixes Signed-off-by: Piotr Żelasko <petezor@gmail.com> * extra unit test Signed-off-by: Piotr Żelasko <petezor@gmail.com> * Fixes in 2D bin estimation script and OOMptimizer Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com> * Apply isort and black reformatting Signed-off-by: pzelasko <pzelasko@users.noreply.github.com> * Support list values for max_tps/max_tpt and bucketing Signed-off-by: Piotr Żelasko <petezor@gmail.com> * fix max_tps/max_tpt config parsing for list values Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com> * Transition: tarred_random_access -> skip_missing_manifest_entries Signed-off-by: Piotr Żelasko <petezor@gmail.com> * Improve warning messages Signed-off-by: Piotr Żelasko <petezor@gmail.com> * Support single unified BPE tokenizer for Canary2 Signed-off-by: Piotr Zelasko <pzelasko@cs-oci-ord-login-01.cm.cluster> * bugfix +es to unified tokenizer support for Canary Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com> * CanaryBPETokenizer Signed-off-by: Piotr Żelasko <petezor@gmail.com> * Revert changes to fetching pad id in AED models Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com> * Update docs Signed-off-by: Piotr Żelasko <petezor@gmail.com> * fix tests Signed-off-by: Piotr Żelasko <petezor@gmail.com> * Remove max_tps and max_duration from OOMptimizer output Signed-off-by: Piotr Żelasko <petezor@gmail.com> --------- Signed-off-by: Piotr Żelasko <petezor@gmail.com> Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com> Signed-off-by: pzelasko <pzelasko@users.noreply.github.com> Signed-off-by: Piotr Zelasko <pzelasko@cs-oci-ord-login-01.cm.cluster> Co-authored-by: pzelasko <pzelasko@users.noreply.github.com> Signed-off-by: Abhinav Garg <abhgarg@nvidia.com>

* TPS-free 2D bucket estimation and filtering Signed-off-by: Piotr Żelasko <petezor@gmail.com> * fixes Signed-off-by: Piotr Żelasko <petezor@gmail.com> * extra unit test Signed-off-by: Piotr Żelasko <petezor@gmail.com> * Fixes in 2D bin estimation script and OOMptimizer Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com> * Apply isort and black reformatting Signed-off-by: pzelasko <pzelasko@users.noreply.github.com> * Support list values for max_tps/max_tpt and bucketing Signed-off-by: Piotr Żelasko <petezor@gmail.com> * fix max_tps/max_tpt config parsing for list values Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com> * Transition: tarred_random_access -> skip_missing_manifest_entries Signed-off-by: Piotr Żelasko <petezor@gmail.com> * Improve warning messages Signed-off-by: Piotr Żelasko <petezor@gmail.com> * Support single unified BPE tokenizer for Canary2 Signed-off-by: Piotr Zelasko <pzelasko@cs-oci-ord-login-01.cm.cluster> * bugfix +es to unified tokenizer support for Canary Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com> * CanaryBPETokenizer Signed-off-by: Piotr Żelasko <petezor@gmail.com> * Revert changes to fetching pad id in AED models Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com> * Update docs Signed-off-by: Piotr Żelasko <petezor@gmail.com> * fix tests Signed-off-by: Piotr Żelasko <petezor@gmail.com> * Remove max_tps and max_duration from OOMptimizer output Signed-off-by: Piotr Żelasko <petezor@gmail.com> --------- Signed-off-by: Piotr Żelasko <petezor@gmail.com> Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com> Signed-off-by: pzelasko <pzelasko@users.noreply.github.com> Signed-off-by: Piotr Zelasko <pzelasko@cs-oci-ord-login-01.cm.cluster> Co-authored-by: pzelasko <pzelasko@users.noreply.github.com> Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

TPS-free 2D bucket estimation and filtering

42892d8

Signed-off-by: Piotr Żelasko <petezor@gmail.com>

github-actions bot added the common label Jan 2, 2025

fixes

8291b09

Signed-off-by: Piotr Żelasko <petezor@gmail.com>

github-advanced-security bot found potential problems Jan 2, 2025

View reviewed changes

nemo/collections/common/data/lhotse/sampling.py

import bisect

import logging

import math

from bisect import bisect_left, bisect_right

Check notice

Code scanning / CodeQL

Unused import Note

Import of 'bisect_right' is not used.

extra unit test

48230c5

Signed-off-by: Piotr Żelasko <petezor@gmail.com>

github-advanced-security bot found potential problems Jan 3, 2025

View reviewed changes

pzelasko and others added 4 commits January 6, 2025 10:18

Fixes in 2D bin estimation script and OOMptimizer

c48b813

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

Apply isort and black reformatting

83cd8df

Signed-off-by: pzelasko <pzelasko@users.noreply.github.com>

Support list values for max_tps/max_tpt and bucketing

ab2f4d6

Signed-off-by: Piotr Żelasko <petezor@gmail.com>

Merge branch '2d-bucketing-and-tps-improvements-2' of https://github.…

565a04a

…com/nvidia/nemo into 2d-bucketing-and-tps-improvements-2

github-advanced-security bot found potential problems Jan 6, 2025

View reviewed changes

pzelasko and others added 2 commits January 6, 2025 10:50

fix max_tps/max_tpt config parsing for list values

c9f0b71

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

Transition: tarred_random_access -> skip_missing_manifest_entries

d8b2d75

Signed-off-by: Piotr Żelasko <petezor@gmail.com>

pzelasko requested a review from nithinraok January 9, 2025 18:04

pzelasko marked this pull request as ready for review January 9, 2025 18:04

Improve warning messages

d8f4dc7

Signed-off-by: Piotr Żelasko <petezor@gmail.com>

nithinraok reviewed Jan 10, 2025

View reviewed changes

nithinraok previously approved these changes Jan 10, 2025

View reviewed changes

pzelasko added the Run CICD label Jan 10, 2025

Support single unified BPE tokenizer for Canary2

bff9690

Signed-off-by: Piotr Zelasko <pzelasko@cs-oci-ord-login-01.cm.cluster>

pzelasko dismissed nithinraok’s stale review via bff9690 January 10, 2025 16:04

bugfix +es to unified tokenizer support for Canary

cab0f8b

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

github-actions bot added the ASR label Jan 10, 2025

pzelasko added Run CICD and removed Run CICD labels Jan 10, 2025

pzelasko and others added 2 commits January 12, 2025 17:22

CanaryBPETokenizer

acf5ea4

Signed-off-by: Piotr Żelasko <petezor@gmail.com>

Revert changes to fetching pad id in AED models

9fced0e

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

pzelasko added Run CICD and removed Run CICD labels Jan 13, 2025

Update docs

4c13728

Signed-off-by: Piotr Żelasko <petezor@gmail.com>

pzelasko requested a review from nithinraok January 17, 2025 01:24

artbataev requested review from ko3n1g and pablo-garay as code owners January 17, 2025 01:25

github-actions bot added core Changes to NeMo Core NLP CI Multi Modal labels Jan 17, 2025

pzelasko force-pushed the 2d-bucketing-and-tps-improvements-2 branch from 49a1536 to a238da2 Compare January 17, 2025 01:27

github-actions bot removed core Changes to NeMo Core NLP CI Multi Modal labels Jan 17, 2025

Merge branch 'main' into 2d-bucketing-and-tps-improvements-2

2b79d2d

pzelasko added Run CICD and removed Run CICD labels Jan 17, 2025

pzelasko enabled auto-merge (squash) January 17, 2025 01:28

Fix integration test

0ab2c09

Signed-off-by: Piotr Żelasko <petezor@gmail.com>

pzelasko added Run CICD and removed Run CICD labels Jan 21, 2025

fix estimate_duration_bins_2d.py for BPE tokenizer without prompt

e9e17e2

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

pzelasko added Run CICD and removed Run CICD labels Jan 22, 2025

nithinraok approved these changes Jan 23, 2025

View reviewed changes

pzelasko merged commit 7e24313 into main Jan 23, 2025
205 of 213 checks passed

pzelasko deleted the 2d-bucketing-and-tps-improvements-2 branch January 23, 2025 19:24

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

TPS-free 2D bucket estimation and filtering#11738

TPS-free 2D bucket estimation and filtering#11738
pzelasko merged 20 commits intomainfrom
2d-bucketing-and-tps-improvements-2

pzelasko commented Jan 2, 2025 •

edited

Loading

Uh oh!

Check notice

Check notice

Check notice

nithinraok left a comment

Uh oh!

nithinraok Jan 10, 2025

Uh oh!

github-actions bot commented Jan 22, 2025

Uh oh!

github-actions bot commented Jan 23, 2025

Uh oh!

nithinraok left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

pzelasko commented Jan 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Changelog

Usage

GitHub Actions CI

Before your PR is "Ready for review"

Who can review?

Additional Information

Uh oh!

Check notice

Check notice

Check notice

nithinraok left a comment

Choose a reason for hiding this comment

Uh oh!

nithinraok Jan 10, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jan 22, 2025

Uh oh!

github-actions bot commented Jan 23, 2025

Uh oh!

nithinraok left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pzelasko commented Jan 2, 2025 •

edited

Loading