Skip to content

Fixed dp-balanced THD packing scheduler - #5993

Draft
FDecaYed wants to merge 1 commit into
NVIDIA:mainfrom
FDecaYed:codex/fixed-thd-packing-scheduler
Draft

Fixed dp-balanced THD packing scheduler#5993
FDecaYed wants to merge 1 commit into
NVIDIA:mainfrom
FDecaYed:codex/fixed-thd-packing-scheduler

Conversation

@FDecaYed

Copy link
Copy Markdown
Contributor

Fixed dp-balanced THD packing scheduler

Summary

Add the fixed, deterministic data-parallel-balanced scheduler for packed THD
batches, including its training-loop contract and process-group forwarding.

Scope and non-goals

  • Balance compatible variable-length samples across data-parallel ranks.
  • Preserve THD metadata and use the model-provided process-group collection.
  • Do not include dynamic context parallelism, SFT-specific packing, dataset
    generation, or DSv4 model code.

Provenance

This local recut is reconstructed from the frozen main baseline
bb5647a9bdd0, not a replay. It references master #5795 and retains only the
fixed scheduler portion of #5679. Credit goes to the original contributors of
#5679: @ilml. The DSv4 integration in #5795 is by @hxbai.

Dependencies

None beyond the frozen baseline used for this local recut. This is a reusable
prerequisite outside the nine-PR DSv4-specific series.

Tests

  • Deterministic packing and data-parallel load balance.
  • THD metadata, padding, and process-group forwarding.
  • Training-step integration and rejection of unsupported scheduler modes.

Publication

Publish in the first prerequisite round on then-current main. This is not a
stacked-review PR.

@copy-pr-bot

copy-pr-bot Bot commented Jul 23, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

Copy link
Copy Markdown
Contributor Author

/ok to test 0be6b6b

Copy link
Copy Markdown
Contributor Author

/ok to test 49bd7a3

@FDecaYed

Copy link
Copy Markdown
Contributor Author

/ok to test 37a68b3

Signed-off-by: Deyu Fu <deyuf@nvidia.com>
@FDecaYed
FDecaYed force-pushed the codex/fixed-thd-packing-scheduler branch from 37a68b3 to 8040550 Compare July 24, 2026 08:47
@copy-pr-bot

copy-pr-bot Bot commented Jul 24, 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.

@FDecaYed

Copy link
Copy Markdown
Contributor Author

/ok to test 8040550

@asolergi-nv

Copy link
Copy Markdown
Contributor

/claude review

data_iterator,
num_microbatches: int,
dp_group,
tp_group,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Typo: "emtpy" → "empty".

Suggested change
tp_group,
# Get a batch from data_iterator or create an empty batch.

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.

2 participants