Skip to content

Add DP-balanced sequence-packing scheduler and wrap_data_iterator entry point - #6626

Merged
ilml merged 2 commits into
NVIDIA:pull-request/6625from
ilml:split/3386-05-packing-scheduler
Aug 18, 2026
Merged

Add DP-balanced sequence-packing scheduler and wrap_data_iterator entry point#6626
ilml merged 2 commits into
NVIDIA:pull-request/6625from
ilml:split/3386-05-packing-scheduler

Conversation

@ilml

@ilml ilml commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Note

Replaces #5905, which was auto-closed when its review-time base (a copy-pr-bot mirror ref) was deleted after #5903 merged. Same branch, rebased onto main past #5903; prior review history lives in #5905.

  • I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Adds BasePackingScheduler, DpBalancedScheduler, PackingSchedulerEnum, and the wrap_data_iterator entry point (schedule, pack, reroute, broadcast metadata, return dynamic num_microbatches + FLOPs stats), plus the readme section and test_wrap_dataloader unit tests.

Part 05/10 of splitting #3386 (Add E2E support for THD format; dev-branch PR #2924). Original changes by @xiaoyao0115 in #3386 — split into functionally self-contained PRs to ease review. Hard dependencies (must merge first): #6625.

Split series (#3386)

Part PR Hard deps (merge first)
01 #5901
02 #5902
03 #5903
04 #6625 #5903
05 #6626 #6625
06 #6627 #5902, #5903
07 #6628 #5902, #6625, #6626
08 #5908
09 #6629 #6627, #5908
10 #6630 #6628, #6629

Branches are stacked linearly (each on the previous) so every PR shows a clean own-diff once its base is retargeted to the copy-pr-bot pull-request/<parent> ref; until then the Files-changed view of a stacked PR includes its ancestors — its own change is the last commit.
Stacked on #6625.

Review notes

  • Doc nit carried from the original PR: the readme mentions the raw TE name thd_get_partitioned_indices instead of the wrapper get_thd_partitioned_indices.

Issue tracking

Linked issue: Related to #3386

Contribution process

Pre-checks

  • I have added relevant unit tests
  • I have added relevant functional tests
  • I have added proper typing to my code Typing guidelines
  • I have added relevant documentation
  • I have run the autoformatter.sh on my PR

🤖 Generated with Claude Code

ilml and others added 2 commits August 18, 2026 02:59
…ry point

Split 5/10 from NVIDIA#3386 (sequence packing / THD E2E support). Adds
BasePackingScheduler, DpBalancedScheduler, PackingSchedulerEnum, the
wrap_data_iterator entry point, docs, and the wrap_dataloader unit
tests.

Original changes by @xiaoyao0115 in NVIDIA#3386.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: ilml <tolong@nvidia.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: ilml <tolong@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Aug 18, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@ilml

ilml commented Aug 18, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 050f033

@ilml
ilml changed the base branch from main to pull-request/6625 August 18, 2026 16:32
@ilml
ilml marked this pull request as ready for review August 18, 2026 17:17
@ilml
ilml requested review from a team as code owners August 18, 2026 17:17
@ilml
ilml merged commit 7027fa8 into NVIDIA:pull-request/6625 Aug 18, 2026
95 of 97 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants