[training, perf] feat: Sync sequence packing (THD) E2E support from MLM PR 3386 - #3634
Closed
cuichenx wants to merge 1 commit into
Closed
[training, perf] feat: Sync sequence packing (THD) E2E support from MLM PR 3386#3634cuichenx wants to merge 1 commit into
cuichenx wants to merge 1 commit into
Conversation
…LM PR 3386 Mirror the training-side changes from NVIDIA/Megatron-LM PR #3386 ("Add E2E support for THD format") into the Bridge training loop. The mcore-side data scheduler that this code depends on (``wrap_data_iterator``, ``get_batch_on_this_rank_for_sequence_packing``) ships as part of that PR; this change is intentionally ahead of the submodule bump and the new code paths are only entered when ``model_config.sequence_packing`` / ``model_config.sequence_packing_scheduler`` are set. Changes: - ``flop_utils.num_floating_point_operations``: refactor signature to accept ``seqlen_sum_this_global_batch`` and ``seqlen_squared_sum_this_global_batch`` for THD packed accounting. When unset, both default to the BSHD geometry (``batch_size * seq_length`` and ``batch_size * seq_length**2``), preserving exact numerical equivalence for existing fixed-length callers and tests. Helpers ``mlp_layer_flops``, ``moe_layer_flops``, ``attn_layer_flops``, ``mamba_layer_flops``, ``gdn_layer_flops``, and ``hybrid_flops`` are updated in lockstep. The standard ``transformer_flops`` body decomposes self-attention into a token-linear part (* seqlen_sum) and a sequence-quadratic part (* seqlen_squared_sum); MLA, MHA/GQA, SWA, and GDN-override paths are all migrated. - ``train.train_step``: when ``model_config.sequence_packing`` is enabled, delegate to ``wrap_data_iterator`` to obtain the packed iterator, the scheduler-determined microbatch count, and the seqlen sums needed for accurate FLOPs. The function now returns those two metrics as additional fields; ``train`` consumes them and uses them for per-step FLOPs accounting in the packed case while keeping the cached ``num_floating_point_operations_model * batch_size`` fast path for BSHD. - ``gpt_step.get_batch``: when ``cfg.model.sequence_packing_scheduler`` is set, call ``get_batch_on_this_rank_for_sequence_packing`` (which handles TP/PP broadcast and CP partitioning internally) and project the resulting ``PackedSeqParams`` into Bridge's flat 10-tuple return shape. Dispatch happens before the standard PP-stage gating so middle PP stages still receive the cu_seqlens metadata the helper broadcasts. Out of scope for this PR (Phase 1 only): mock SFT dataset support and the ``--sft-mock-dataset-config-json`` validation logic from MLM PR 3386 are not ported. Those will come in a follow-up if needed. Depends on: NVIDIA/Megatron-LM#3386 landing and the Megatron-Core submodule being bumped to a commit that includes ``megatron.core.datasets.data_schedule`` and the new ``ModelParallelConfig`` / ``TransformerConfig`` fields. Signed-off-by: Chen Cui <chcui@nvidia.com>
|
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. |
yaoyu-33
added
the
area:training
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
2 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.
Summary
Phase 1 port of the Bridge-training-side counterpart of NVIDIA/Megatron-LM#3386 ("Add E2E support for THD format"). Wires sequence-packing (THD) end-to-end through the Bridge training loop and FLOPs accounting. The new code paths are only entered when
model_config.sequence_packing/model_config.sequence_packing_schedulerare set, so existing BSHD recipes are unaffected.This PR depends on NVIDIA/Megatron-LM#3386 and the Megatron-Core submodule bump that ships:
megatron.core.datasets.data_schedule.wrap_data_iteratormegatron.core.datasets.data_schedule.get_batch_on_this_rank_for_sequence_packingModelParallelConfig.sequence_packingandTransformerConfig.sequence_packing_scheduler/max_seqlen_per_dp_cp_rank/hybrid_context_parallelUntil that lands and is bumped here, the new branches will fail at runtime when enabled. Imports are guarded (lazy) so the test suite and existing recipes still work without it.
Marked as draft to make this dependency explicit. Will undraft after the mcore bump.
Changes
training/utils/flop_utils.pynum_floating_point_operationsaccepts optionalseqlen_sum_this_global_batch/seqlen_squared_sum_this_global_batch— both default to the BSHD geometry so existing callers and tests are exactly equivalent. Internal helpers (mlp/moe/attn/mamba/gdn/hybrid_layer_flops) refactored in lockstep.transformer_flopsdecomposes self-attention into a token-linear (× seqlen_sum) part and a sequence-quadratic (× seqlen_squared_sum) part across MLA, MHA/GQA, SWA, and GDN-override paths.training/train.pytrain_stepcallswrap_data_iteratorwhen packing is enabled, returns 2 extra metrics (seqlen_sum_this_global_batch,seqlen_squared_sum_this_global_batch);trainunpacks them and uses them for per-step FLOPs in the packed case while keeping the cachednum_floating_point_operations_model * batch_sizefast path for BSHD.training/gpt_step.pyget_batchdispatches toget_batch_on_this_rank_for_sequence_packingwhencfg.model.sequence_packing_scheduleris set; the helper handles TP/PP broadcast and CP partitioning internally, so dispatch happens before the standard PP-stage gating.PackedSeqParamsis projected into Bridge's flat 10-tuple.Out of scope (deferred)
The MLM PR also adds:
MockSFTLowLevelDataset/MockSFTDataset(depends onpandas, only used for variable-length benchmarking)--sft-mock-dataset-config-jsonvalidation logic and default-config plumbing invalidate_argsThese are skipped intentionally for Phase 1 and can be ported in a follow-up if the SFT mock-data flow is needed downstream.
Test plan
sequence_packing_scheduleron a small Llama recipe, confirm packed iterator integration works end-to-end.seqlen_sum/seqlen_squared_sumare derived frombatch_size).tests/unit_tests/training/utils/test_flop_utils.py— should pass without modification (BSHD path).uv run pre-commit run --all-files(auto-formatter has already run via local pre-commit hook on this commit).