Sync Dynamic-CP feature from dev to main - #5279
Closed
ilml wants to merge 1 commit into
Closed
Conversation
Port the Dynamic CP (dynamic context parallelism) feature from the dev branch, superseding the Hybrid Context Parallel implementation (NVIDIA#2282) with the evolved design from dev PRs NVIDIA#2000, NVIDIA#3405, NVIDIA#4226 and follow-up fixes. - Copy the packing schedulers (BasePackingScheduler, DpBalancedScheduler, DefaultDynamicCPScheduler), wrap_data_iterator and get_batch_on_this_rank_for_sequence_packing into megatron/core/datasets/data_schedule.py, and add megatron/core/datasets/data_schedule_utils.py (both byte-identical to dev), keeping HybridCPDataLoaderWrapper for compatibility. - Rename --hybrid-context-parallel to --dynamic-context-parallel; keep hybrid_context_parallel as a deprecated alias on ModelParallelConfig. - Add min_dynamic_context_parallel_size and sequence_packing_scheduler knobs with validation. - Add get_thd_partitioned_indices to transformer_engine.py and thread packed-seq/cu_seqlens plumbing through attention, MLA, MTP, mamba/GDN and the GPT model, matching dev. - Replace the HybridCPDataLoaderWrapper wiring in training.py and pretrain_gpt.py with dev's wrap_data_iterator pattern, and update parallel_state dynamic DPxCP group helpers (byte-identical to dev). - Update main-only callers (pretrain_hybrid.py, tools/prepare_cache.py, megatron/elastification/pretrain_hybrid_flex.py) and existing unit tests for the rename. Co-Authored-By: Claude Fable 5 <noreply@anthropic.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
1 participant
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
Ports the Dynamic CP (dynamic context parallelism) feature from the
devbranch tomain, superseding the Hybrid Context Parallel implementation (#2282) with the evolved design from dev PRs #2000 (Dynamic CP part 2), #3405 (THD+rope fix and Dynamic-CP naming), #4226 (minor improvements), and follow-up fixes.What's included
Core feature files — byte-identical copies from
dev:megatron/core/datasets/data_schedule.py—BasePackingScheduler,DpBalancedScheduler,DefaultDynamicCPScheduler,wrap_data_iterator,get_batch_on_this_rank_for_sequence_packing(keepsHybridCPDataLoaderWrapperfor compatibility)megatron/core/datasets/data_schedule_utils.py— new module (scheduling/rerouting/packing helpers)megatron/core/parallel_state.py— dynamic DPxCP group helpers (get_dynamic_data_context_parallel_groups)megatron/core/packed_seq_params.py,megatron/core/datasets/readme.mdSurgical ports (only dynamic-CP hunks; unrelated dev-only features excluded):
megatron/training/arguments.py—--dynamic-context-parallel,--min-dynamic-context-parallel-size,--sequence-packing-scheduler+ validationmegatron/core/model_parallel_config.py— new config fields;hybrid_context_parallelkept as a deprecated aliasmegatron/core/extensions/transformer_engine.py— addsget_thd_partitioned_indicesand dynamic-CP group selection inTEDotProductAttentionattention.py,multi_latent_attention.py,absorbed_mla.py,multi_token_prediction.py,gpt_model.py,gated_delta_net.py,mamba_mixer.py,mamba_context_parallel.py,schedules.py,gpt_dataset.py,sft_dataset.py,data_samplers.py,core/utils.py,transformer_config.py,initialize.pymegatron/training/training.py/pretrain_gpt.py— replaceHybridCPDataLoaderWrapperwiring with dev'swrap_data_iteratorpatternRename fallout (main-only callers updated):
pretrain_hybrid.py,megatron/elastification/pretrain_hybrid_flex.py—args.dynamic_context_parallel+ renamed group helper importtools/prepare_cache.py— drops the removedhybrid_context_paralleldataset-config kwarg (mirrors dev)test_parallel_state.py,test_get_batch.py,test_hybrid_moe_model.pygolden config). No new tests added.Breaking change
The CLI flag
--hybrid-context-parallelis renamed to--dynamic-context-parallel(matchingdev). TheModelParallelConfig.hybrid_context_parallelfield remains as a deprecated alias that maps todynamic_context_parallelwith aDeprecationWarning.Verification
python3 -m py_compile.origin/dev.black/isortclean on all changedmegatron/core/andtests/files (CI formatting scope).get_hybrid_data_context_parallel_groups,args.hybrid_context_parallel) outside the intentional deprecated-alias sites.🤖 Generated with Claude Code