Skip to content

Integrate sequence-packing scheduler into training loop and GPT pretraining - #5907

Closed
ilml wants to merge 2 commits into
NVIDIA:pull-request/5906from
ilml:split/3386-07-scheduler-training-integration
Closed

Integrate sequence-packing scheduler into training loop and GPT pretraining#5907
ilml wants to merge 2 commits into
NVIDIA:pull-request/5906from
ilml:split/3386-07-scheduler-training-integration

Conversation

@ilml

@ilml ilml commented Jul 20, 2026

Copy link
Copy Markdown
Contributor
  • I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Wires wrap_data_iterator into train_step/evaluate with dynamic num_microbatches, adds global seqlen-stats plumbing for variable-length FLOPs accounting, gates the HybridCP sampler on the new scheduler flag, and restructures pretrain_gpt.forward_step for packed (THD) batches with padding mask.

Part 07/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): #5902, #6589, #5905.

Split series (#3386)

Part PR Hard deps (merge first)
01 #5901
02 #5902
03 #5903
04 #6589 #5903
05 #5905 #6589
06 #5906 #5902, #5903
07 #5907 #5902, #6589, #5905
08 #5908
09 #5909 #5906, #5908
10 #5910 #5907, #5909

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 #5906.

Review notes

  • data_samplers.py identity-collate or-chain includes use_vanilla_collate_fn, which is defined nowhere in the series — candidate for deletion in review. The varlen flags in the same chain are getattr-guarded and land later in the series.

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

@copy-pr-bot

copy-pr-bot Bot commented Jul 20, 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 marked this pull request as ready for review July 22, 2026 01:02
@ilml
ilml requested review from a team as code owners July 22, 2026 01:02
@ilml
ilml force-pushed the split/3386-07-scheduler-training-integration branch 3 times, most recently from 9216d74 to 7ed21ee Compare July 30, 2026 21:19
@ilml
ilml force-pushed the split/3386-07-scheduler-training-integration branch from 7ed21ee to b25804f Compare August 3, 2026 20:39
@ilml
ilml force-pushed the split/3386-07-scheduler-training-integration branch 3 times, most recently from 9c5b861 to 2a59632 Compare August 10, 2026 22:09
ilml and others added 2 commits August 10, 2026 22:09
…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>
Follow-up to dropping the unused mirror field in NVIDIA#5903: the dataset
config no longer carries this knob, so core_gpt_dataset_config_from_args
must not set it. The scheduler is still selected from
args.sequence_packing_scheduler via ModelParallelConfig.

Original changes by @xiaoyao0115 in NVIDIA#3386.

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

ilml commented Aug 17, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 2a59632

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