Add sequence_packing_scheduler config field, CLI arg, and validation - #5902
Add sequence_packing_scheduler config field, CLI arg, and validation#5902ilml wants to merge 2 commits into
Conversation
Split 2/10 from NVIDIA#3386 (sequence packing / THD E2E support). Adds the sequence_packing_scheduler knob to ModelParallelConfig, its TransformerConfig validation (TE>=2.9 pin, variable_seq_lengths, alltoall dispatcher), the explicit CLI args, and validate_args checks. Feature-flagged no-op until the scheduler lands. Original changes by @xiaoyao0115 in NVIDIA#3386. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: ilml <tolong@nvidia.com>
|
/ok to test 30a81d8 |
…-config Signed-off-by: Tom Long <tolong@nvidia.com> # Conflicts: # megatron/training/arguments.py # tests/unit_tests/models/test_hybrid_moe_model.py
|
/ok to test d76a634 |
|
|
||
| # Packed-sequence buffer-size check. Placed after varlen scheduler | ||
| # auto-select so it validates the final resolved scheduler. | ||
| if args.sequence_packing_scheduler is not None: |
There was a problem hiding this comment.
from my understanding, --calculate-per-token-loss should be required when sequence_packing_scheduler is set. for example, if two packed microbatches contain 1 and 9 valid tokens, the default path in schedules.py gives the two microbatches equal weight (i.e., weight 1/2 for the first token, and 1/18 for each of the other tokens) instead of token weights of 1:9, so gradients depend on how samples happen to be packed. the --hybrid-context-parallel already does this validation near line 1385.
| self.variable_seq_lengths = True | ||
|
|
||
| # TODO(tailaim): add support for other dispatcher types | ||
| assert self.moe_token_dispatcher_type == "alltoall", ( |
There was a problem hiding this comment.
this check also runs for dense models, where num_moe_experts is None and moe_token_dispatcher_type is unused; its default is allgather. for example, TransformerConfig(..., sequence_packing_scheduler="dp_balanced", max_seqlen_per_dp_cp_rank=4096) fails even though it has no MoE dispatcher; please require alltoall only when num_moe_experts is set and add dense and MoE coverage.
|
/ok to test d76a634 |
6 participants
What does this PR do?
Adds the
sequence_packing_schedulerknob toModelParallelConfig, itsTransformerConfig.__post_init__validation (TE>=2.9 pin,variable_seq_lengths, alltoall MoE dispatcher), the explicit--sequence-packing-scheduler/--max-seqlen-per-dp-cp-rank/--hybrid-context-parallelargs with factory exclude-list entries, and the packed-sequence buffer-size check invalidate_args. Fully feature-flagged: setting the flag validates but is a no-op until the scheduler lands.Part 02/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): none — independently mergeable.
Split series (#3386)
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.Issue tracking
Linked issue: Related to #3386
Contribution process
Pre-checks
🤖 Generated with Claude Code
Series status (updated)
Remaining PRs of the #3386 split, tracked in the Dynamic Context Parallelism project. Reviews run in parallel; merge order:
Merged so far: #5901, #5903, #5908, #6625.