Integrate sequence-packing scheduler into training loop and wire varlen dataset into GPT pretraining - #6742
Open
ilml wants to merge 4 commits into
Open
Integrate sequence-packing scheduler into training loop and wire varlen dataset into GPT pretraining#6742ilml wants to merge 4 commits into
ilml wants to merge 4 commits into
Conversation
Contributor
Author
|
/ok to test b9a71e3 |
2 tasks
Contributor
Author
|
/ok to test c467af1 |
cuichenx
approved these changes
Aug 21, 2026
| data_parallel_size=mpu.get_data_parallel_world_size()) | ||
| elif args.dataloader_type == 'single': | ||
| if args.hybrid_context_parallel: | ||
| if args.hybrid_context_parallel and args.sequence_packing_scheduler is None: |
Contributor
There was a problem hiding this comment.
not in this PR but we should change everything to "dynamic" now?
zhongbozhu
approved these changes
Aug 21, 2026
Contributor
Author
|
/ok to test 095c361 |
Contributor
Author
|
/ok to test 0a0f9e3 |
Contributor
Author
|
/ok to test 4ccffde |
Contributor
Author
|
/ok to test c7208e1 |
Contributor
Author
|
/ok to test ac3d429 |
Contributor
Author
|
/ok to test 9268336 |
…aining Split 7/10 from NVIDIA#3386 (sequence packing / THD E2E support). Wires wrap_data_iterator into train_step/evaluate with dynamic num_microbatches, adds global seqlen-stats plumbing for FLOPs accounting, gates the HybridCP sampler, and restructures pretrain_gpt.forward_step for packed (THD) batches with padding mask. Original changes by @xiaoyao0115 in NVIDIA#3386. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: ilml <tolong@nvidia.com>
…vider Split 10/10 from NVIDIA#3386 (sequence packing / THD E2E support). Adds --use-varlen-dataset / --varlen-sbhd-validation / --varlen-mock-dataset-config-json, their validation (auto-selecting the dp_balanced scheduler), the dataset-provider wiring in pretrain_gpt, and the remaining dataloader/unpack 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>
…in_hybrid Per review: mirror the pretrain_gpt integration for HybridModel, whose forward already accepts packed_seq_params and padding_mask (packed THD batches are exercised today by hybrid SFT). - get_batch takes the get_batch_on_this_rank_for_sequence_packing path when --sequence-packing-scheduler is set, same as pretrain_gpt. - forward_step handles the scheduler's 6/7-tuple batches and threads padding_mask through to the model; the legacy SFT/inter-document path is unchanged apart from the enclosing tuple dispatch. - The dataset provider builds VarlenDataset/MockVarlenDataset under --use-varlen-dataset and passes the varlen config fields through. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: ilml <tolong@nvidia.com>
ilml
force-pushed
the
split/3386-training-r2
branch
from
August 27, 2026 16:14
9268336 to
97e73af
Compare
Contributor
Author
|
/ok to test 97e73af09b03256de883166a86590fa77e73ec99 |
@ilml, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/ |
jaredcasper
approved these changes
Aug 27, 2026
ilml
enabled auto-merge
August 27, 2026 18:51
Contributor
Author
|
/ok to test 97e73af |
Contributor
Author
|
/ok to test 992c2b7 |
The repo black config skips magic trailing commas; apply exactly the hunk the lint job suggests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: ilml <tolong@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Integrates the sequence-packing scheduler into
train_step/evaluateand wires the varlen dataset into the pretraining CLI, dataloader, and dataset provider (--use-varlen-dataset) — for bothpretrain_gptandpretrain_hybrid(per review:HybridModel.forwardalready acceptspacked_seq_params/padding_mask, so the hybrid entry point gets the same batch-fetch, forward_step, and provider wiring; note hybrid THD wiring goes beyond the original #3386, which only wired GPT). Final PR of the #3386 split series tracked in the Dynamic Context Parallelism project — after this, the full THD/sequence-packing E2E from #3386 is landed. Original changes by @xiaoyao0115 in #3386.Series status
Remaining PRs of the #3386 split, tracked in the Dynamic Context Parallelism project. All PRs target
main— no stacked bases. Reviews run in parallel; merges go in order:pretrain_gpt+pretrain_hybrid)#6734 carries a temporary copy of #5902's commit, and #6742 carries #6734's commits (review its last three commits) — the duplicates drop out automatically as parents merge and the branches are rebased.
Merged so far: #5901, #5903, #6625, #5908, #6684, #5902, #6734. This is the final PR of the series.
Contribution process