[Core] Add dynamic context-parallel THD scheduling - #5679
Closed
ilml wants to merge 1 commit into
Closed
Conversation
Contributor
Author
|
/ok to test |
@ilml, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/ |
Contributor
Author
|
/ok to test 554f423 |
guihong-nv
reviewed
Jul 8, 2026
guihong-nv
reviewed
Jul 8, 2026
parthmannan
requested changes
Jul 10, 2026
ilml
marked this pull request as draft
July 14, 2026 20:05
ilml
force-pushed
the
codex/dynamic-cp-thd-core
branch
from
July 14, 2026 20:22
554f423 to
b38b574
Compare
Contributor
Author
|
/ok to test b38b574 |
Add dynamic DPxCP scheduling, per-microbatch process-group selection, and THD real-versus-padded boundary handling. Preserve the existing fixed-CP scheduler and tokenizer API while rejecting unsupported CUDA graph, FSDP, MoE, MLA runtime, and hybrid-model combinations. Co-authored-by: xiaoyao0115 <1804647152@qq.com> Signed-off-by: ilml <tolong@nvidia.com>
ilml
force-pushed
the
codex/dynamic-cp-thd-core
branch
from
July 14, 2026 21:33
938e6f6 to
533deb7
Compare
WalkthroughAdds dynamic context-parallel process groups and configurable sequence-packing schedulers, integrates packed batches through training and attention, extends SFT dataset support, and forwards SFT chat-template arguments. ChangesSequence packing and dynamic context parallelism
SFT tokenizer template arguments
Padding-mask alignment
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Training
participant Scheduler
participant DCP
participant Attention
participant TransformerEngine
Training->>Scheduler: wrap data iterator
Scheduler->>DCP: reroute scheduled samples
DCP-->>Scheduler: return packed microbatches
Scheduler-->>Training: provide iterator and token statistics
Training->>Attention: forward packed sequence parameters
Attention->>TransformerEngine: use dynamic CP group
TransformerEngine-->>Attention: restore original CP group
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Contributor
Author
|
/ok to test 533deb7 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
megatron/core/extensions/transformer_engine.py (1)
1794-1905: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRestore-on-success only: exception during attention leaves
self.cp_grouppermanently corrupted.The dynamic-CP group swap (lines 1794-1816) is only reverted at lines 1897-1904, after
super().forward(...). If that call raises for any reason, the restore never runs and this persistent module's CP group/stream stays pointed at the dynamic-CP value for every subsequent forward call on this layer (including non-dynamic-CP iterations).🔧 Proposed fix: guarantee restoration with try/finally
original_cp_group = self.cp_group original_cp_global_ranks = self.cp_global_ranks - if packed_seq_params is not None: - # If Dynamic CP group is provided, update TE DPA CP group - if packed_seq_params.local_cp_size is not None: - if packed_seq_params.local_cp_size == 1: - super().set_context_parallel_group(None, None, None, self.cp_comm_type) - else: - assert ( - packed_seq_params.cp_group is not None - ), "cp_group is not set in packed_seq_params for dynamic CP" - self.cp_group = packed_seq_params.cp_group - if TEDotProductAttention.cp_stream is None: - TEDotProductAttention.cp_stream = torch.cuda.Stream() - super().set_context_parallel_group( - self.cp_group, - torch.distributed.get_process_group_ranks(self.cp_group), - TEDotProductAttention.cp_stream, - self.cp_comm_type, - ) - self.kept_packed_seq_params.discard("cp_group") - self.kept_packed_seq_params.discard("local_cp_size") + dynamic_cp_active = packed_seq_params is not None and packed_seq_params.local_cp_size is not None + if packed_seq_params is not None: + # If Dynamic CP group is provided, update TE DPA CP group + if dynamic_cp_active: + if packed_seq_params.local_cp_size == 1: + super().set_context_parallel_group(None, None, None, self.cp_comm_type) + else: + assert ( + packed_seq_params.cp_group is not None + ), "cp_group is not set in packed_seq_params for dynamic CP" + self.cp_group = packed_seq_params.cp_group + if TEDotProductAttention.cp_stream is None: + TEDotProductAttention.cp_stream = torch.cuda.Stream() + super().set_context_parallel_group( + self.cp_group, + torch.distributed.get_process_group_ranks(self.cp_group), + TEDotProductAttention.cp_stream, + self.cp_comm_type, + ) + self.kept_packed_seq_params.discard("cp_group") + self.kept_packed_seq_params.discard("local_cp_size") + + try:...then wrap through the end of both
super().forward(...)branches, and change the restore block to:- if packed_seq_params is not None and packed_seq_params.local_cp_size is not None: - super().set_context_parallel_group( - original_cp_group, - original_cp_global_ranks, - TEDotProductAttention.cp_stream, - self.cp_comm_type, - ) + finally: + if dynamic_cp_active: + super().set_context_parallel_group( + original_cp_group, + original_cp_global_ranks, + TEDotProductAttention.cp_stream, + self.cp_comm_type, + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@megatron/core/extensions/transformer_engine.py` around lines 1794 - 1905, Guarantee restoration of the original context-parallel state when attention execution fails. In the forward method, wrap both `super().forward(...)` branches and all intervening processing through the return in a `try/finally`, and move the existing restoration using `original_cp_group`, `original_cp_global_ranks`, and `TEDotProductAttention.cp_stream` into the `finally` block so dynamic-CP changes are reverted on success and exceptions.
🧹 Nitpick comments (1)
megatron/core/datasets/data_schedule_utils.py (1)
312-328: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winReuse
gid2local_idfor O(1) membership checks instead of tensorin.
gid2local_id(line 312) already provides O(1) lookup, butsend_ids_sorted(and the similar filter insend_seq_lens) test membership viagid in global_ids_this_rank, a tensor scan. For large batches this becomes O(N·M) instead of O(N).♻️ Proposed refactor
send_ids_sorted = [ - gid for d in dp_ranks for gid in combined_sample_id_groups[d] if gid in global_ids_this_rank + gid for d in dp_ranks for gid in combined_sample_id_groups[d] if gid in gid2local_id ]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@megatron/core/datasets/data_schedule_utils.py` around lines 312 - 328, Update the `send_ids_sorted` filter and the corresponding filter in `send_seq_lens` to use `gid2local_id` for membership checks instead of scanning `global_ids_this_rank` with `gid in`. Preserve the existing ordering and output values while ensuring each membership test is O(1).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@megatron/core/datasets/data_schedule_utils.py`:
- Around line 857-862: Replace the `assert False` guard in
`align_sample_id_groups` with an explicit exception that remains active under
Python optimization, while preserving the existing failure message and boundary
condition when no tail microbatch can provide enough IDs.
- Around line 522-559: Define or import _DYNAMIC_CP_WORKLOAD_CAP_DELTA at module
scope before next_hdp_group_packing_aware computes cap. Ensure the symbol is
available for every non-empty scheduler call and retains the intended
workload-cap delta value, or source it from the existing configuration mechanism
if one is already established.
In `@megatron/core/parallel_state.py`:
- Around line 428-434: Update the group-size generation in the parallel-group
construction logic around group_sizes and create_group so every generated size
evenly divides len(ranks), preventing partial trailing groups from being stored
under a larger key. Either reject non-power-of-two domains unless min_cp_size
selects the full domain, or filter group_sizes to divisors of len(ranks),
preserving valid full-domain behavior.
In `@megatron/training/arguments.py`:
- Around line 1389-1393: Update the deprecation warning in the
hybrid_context_parallel argument handling to pass stacklevel=2, matching the
existing warning in model_parallel_config.py and ensuring callers receive the
correct warning location.
- Around line 1400-1403: Update the dynamic_context_parallel validation to check
the normalized args.cuda_graph_impl value instead of accessing
args.enable_cuda_graph, which may have been deleted earlier in the
argument-processing function. Preserve the existing assertion message and reject
the combination when CUDA Graph is enabled.
In `@pretrain_gpt.py`:
- Line 129: Initialize is_hybrid_cp from args.dynamic_context_parallel instead
of hardcoding it to False, so get_batch() uses the dynamic CP behavior
consistently with get_batch_on_this_tp_rank(), get_batch_on_this_cp_rank(), and
dummy_train_step().
In `@pretrain_hybrid.py`:
- Around line 455-460: Update the unsupported-feature guard in
pretrain_hybrid.py to reject either args.dynamic_context_parallel or the
deprecated args.hybrid_context_parallel alias before dataset configuration is
built. Preserve the existing ValueError message and behavior for both flags.
---
Outside diff comments:
In `@megatron/core/extensions/transformer_engine.py`:
- Around line 1794-1905: Guarantee restoration of the original context-parallel
state when attention execution fails. In the forward method, wrap both
`super().forward(...)` branches and all intervening processing through the
return in a `try/finally`, and move the existing restoration using
`original_cp_group`, `original_cp_global_ranks`, and
`TEDotProductAttention.cp_stream` into the `finally` block so dynamic-CP changes
are reverted on success and exceptions.
---
Nitpick comments:
In `@megatron/core/datasets/data_schedule_utils.py`:
- Around line 312-328: Update the `send_ids_sorted` filter and the corresponding
filter in `send_seq_lens` to use `gid2local_id` for membership checks instead of
scanning `global_ids_this_rank` with `gid in`. Preserve the existing ordering
and output values while ensuring each membership test is O(1).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Enterprise
Run ID: 0da1dc8d-ba7e-40d4-840a-198bc594de12
📒 Files selected for processing (30)
megatron/core/datasets/data_schedule.pymegatron/core/datasets/data_schedule_utils.pymegatron/core/datasets/gpt_dataset.pymegatron/core/datasets/readme.mdmegatron/core/extensions/transformer_engine.pymegatron/core/model_parallel_config.pymegatron/core/models/gpt/gpt_model.pymegatron/core/packed_seq_params.pymegatron/core/parallel_state.pymegatron/core/pipeline_parallel/schedules.pymegatron/core/tokenizers/text/libraries/sft_tokenizer.pymegatron/core/tokenizers/text/text_tokenizer.pymegatron/core/transformer/attention.pymegatron/core/transformer/moe/moe_layer.pymegatron/core/transformer/transformer_config.pymegatron/elastification/pretrain_hybrid_flex.pymegatron/training/arguments.pymegatron/training/datasets/data_samplers.pymegatron/training/datasets/sft_dataset.pymegatron/training/datasets/utils.pymegatron/training/initialize.pymegatron/training/training.pypretrain_gpt.pypretrain_hybrid.pytests/unit_tests/models/test_hybrid_moe_model.pytests/unit_tests/test_model_parallel_config.pytests/unit_tests/test_parallel_state.pytests/unit_tests/test_sequence_packing.pytests/unit_tests/tokenizers/test_tokenizer.pytests/unit_tests/transformer/moe/test_moe_layer.py
💤 Files with no reviewable changes (1)
- megatron/core/pipeline_parallel/schedules.py
Comment on lines
+522
to
+559
| def next_hdp_group_packing_aware( | ||
| sample_seqlens: List[Tuple[int, int]], | ||
| total_gpus: int, | ||
| max_seq_len_per_rank: int, | ||
| min_cp_size: int = 1, | ||
| ) -> Tuple[List[List[int]], List[Tuple[int, int]], List[float], List[List[int]]]: | ||
| """Form one DCP microbatch with packing-aware CP group selection. | ||
|
|
||
| This differs from the legacy DCP scheduler in two ways: | ||
| 1. Short sequences may use a larger CP group than their minimum required | ||
| CP size when that lowers the critical-path rank workload. | ||
| 2. Candidate placements are bounded by ``tall * max_seq_len_per_rank``, | ||
| the per-rank workload upper bound for packing sequences no longer than | ||
| the local tallest sequence in the microbatch. | ||
|
|
||
| The scheduler keeps the legacy invariant that each returned microbatch has | ||
| no empty DPxCP rank after the fill step. For non-power-of-two DPxCP layouts, | ||
| it falls back to the full DPxCP group if power-of-two expansion cannot fill | ||
| every rank. | ||
| """ | ||
| if not sample_seqlens: | ||
| return ( | ||
| [[] for _ in range(total_gpus)], | ||
| [], | ||
| [0.0 for _ in range(total_gpus)], | ||
| [[] for _ in range(total_gpus)], | ||
| ) | ||
|
|
||
| def cp_min_fn(seq_len: int) -> int: | ||
| return dcp_gpus_needed(seq_len, max_seq_len_per_rank, min_cp_size) | ||
|
|
||
| def workload(seq_len: int, cp_size: int) -> float: | ||
| return (seq_len * seq_len) / cp_size | ||
|
|
||
| sample_seqlens = sorted(sample_seqlens, key=lambda x: x[1], reverse=True) | ||
| local_tall = sample_seqlens[0][1] | ||
| cap = float(local_tall) * float(max_seq_len_per_rank) * (1.0 + _DYNAMIC_CP_WORKLOAD_CAP_DELTA) | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Searching for _DYNAMIC_CP_WORKLOAD_CAP_DELTA references..."
rg -n "_DYNAMIC_CP_WORKLOAD_CAP_DELTA" .
echo
echo "Inspecting the relevant file region..."
file="megatron/core/datasets/data_schedule_utils.py"
wc -l "$file"
sed -n '520,565p' "$file"Repository: NVIDIA/Megatron-LM
Length of output: 2395
Define _DYNAMIC_CP_WORKLOAD_CAP_DELTA before using it
next_hdp_group_packing_aware references _DYNAMIC_CP_WORKLOAD_CAP_DELTA, but this module never defines or imports it. Any non-empty call will raise NameError when computing cap, breaking the dynamic-CP scheduler path. Add the missing module-level constant or pass it in from config.
🧰 Tools
🪛 Ruff (0.15.21)
[error] 558-558: Undefined name _DYNAMIC_CP_WORKLOAD_CAP_DELTA
(F821)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@megatron/core/datasets/data_schedule_utils.py` around lines 522 - 559, Define
or import _DYNAMIC_CP_WORKLOAD_CAP_DELTA at module scope before
next_hdp_group_packing_aware computes cap. Ensure the symbol is available for
every non-empty scheduler call and retains the intended workload-cap delta
value, or source it from the existing configuration mechanism if one is already
established.
Source: Linters/SAST tools
Comment on lines
+857
to
+862
| attempts_since_split = 0 | ||
| while remainder > 0: | ||
| if i < 0: | ||
| if attempts_since_split >= len(sample_id_groups): | ||
| assert False, 'align_sample_id_groups: no tail microbatch has enough ids to split' | ||
| i = len(sample_id_groups) - 1 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
assert False is stripped under python -O.
Ruff (B011) flags this. If assertions are disabled, the attempts_since_split >= len(sample_id_groups) guard silently no-ops instead of failing, leaving i negative and changing loop behavior unexpectedly rather than failing loudly.
🛡️ Proposed fix
- if attempts_since_split >= len(sample_id_groups):
- assert False, 'align_sample_id_groups: no tail microbatch has enough ids to split'
+ if attempts_since_split >= len(sample_id_groups):
+ raise AssertionError(
+ 'align_sample_id_groups: no tail microbatch has enough ids to split'
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| attempts_since_split = 0 | |
| while remainder > 0: | |
| if i < 0: | |
| if attempts_since_split >= len(sample_id_groups): | |
| assert False, 'align_sample_id_groups: no tail microbatch has enough ids to split' | |
| i = len(sample_id_groups) - 1 | |
| attempts_since_split = 0 | |
| while remainder > 0: | |
| if i < 0: | |
| if attempts_since_split >= len(sample_id_groups): | |
| raise AssertionError( | |
| 'align_sample_id_groups: no tail microbatch has enough ids to split' | |
| ) | |
| i = len(sample_id_groups) - 1 |
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 861-861: Do not assert False (python -O removes these calls), raise AssertionError()
Replace assert False
(B011)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@megatron/core/datasets/data_schedule_utils.py` around lines 857 - 862,
Replace the `assert False` guard in `align_sample_id_groups` with an explicit
exception that remains active under Python optimization, while preserving the
existing failure message and boundary condition when no tail microbatch can
provide enough IDs.
Source: Linters/SAST tools
Comment on lines
+428
to
434
| group_sizes = [ | ||
| 2**i for i in range(int(log2(len(ranks))) + 1) if min_cp_size <= 2**i < len(ranks) | ||
| ] | ||
| for group_size in group_sizes: | ||
| for i in range(0, len(ranks), group_size): | ||
| group = create_group( | ||
| ranks[i : i + group_size], |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Reject domains that cannot be evenly partitioned into every generated group size.
The current evenness check permits a six-rank DP×CP domain. Line 432 then creates a final two-rank slice for group_size=4 and stores it under key 4, so callers requesting CP size 4 receive a size-2 group. Require a power-of-two domain unless the minimum selects the full domain, or only generate sizes that divide len(ranks).
Proposed fix
- assert len(ranks_with_cp) % 2 == 0 or min_dynamic_context_parallel_size == len(
- ranks_with_cp
- ), (
- "Dynamic context parallel requires an even DP x CP domain unless the "
- "minimum group size selects the full domain"
- )
+ is_power_of_two = len(ranks_with_cp) > 0 and (
+ len(ranks_with_cp) & (len(ranks_with_cp) - 1)
+ ) == 0
+ assert (
+ is_power_of_two
+ or min_dynamic_context_parallel_size == len(ranks_with_cp)
+ ), (
+ "Dynamic context parallel requires a power-of-two DP x CP domain "
+ "unless the minimum group size selects the full domain"
+ )Also applies to: 934-945
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@megatron/core/parallel_state.py` around lines 428 - 434, Update the
group-size generation in the parallel-group construction logic around
group_sizes and create_group so every generated size evenly divides len(ranks),
preventing partial trailing groups from being stored under a larger key. Either
reject non-power-of-two domains unless min_cp_size selects the full domain, or
filter group_sizes to divisors of len(ranks), preserving valid full-domain
behavior.
Comment on lines
1389
to
+1393
| if args.hybrid_context_parallel: | ||
| assert not args.pipeline_model_parallel_size > 1, 'Hybrid context parallelism not supported with pipeline parallelism' | ||
| assert not args.enable_cuda_graph, 'Hybrid context parallelism not supported with CUDA Graph' | ||
| assert not args.use_megatron_fsdp, 'Hybrid context parallelism not supported with Megatron FSDP' | ||
| assert args.dataloader_type == 'single', 'Hybrid context parallelism only supported with single dataloader type' | ||
| assert args.calculate_per_token_loss, 'Hybrid context parallelism must be used with --calculate-per-token-loss' | ||
| warnings.warn( | ||
| '--hybrid-context-parallel is deprecated; use --dynamic-context-parallel.', | ||
| DeprecationWarning, | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Missing stacklevel on deprecation warning.
Flagged by Ruff (B028). The equivalent deprecation warning in megatron/core/model_parallel_config.py (hybrid_context_parallel → dynamic_context_parallel) already uses stacklevel=2; this one should match for consistent, useful warning locations.
🔧 Proposed fix
if args.hybrid_context_parallel:
warnings.warn(
'--hybrid-context-parallel is deprecated; use --dynamic-context-parallel.',
DeprecationWarning,
+ stacklevel=2,
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if args.hybrid_context_parallel: | |
| assert not args.pipeline_model_parallel_size > 1, 'Hybrid context parallelism not supported with pipeline parallelism' | |
| assert not args.enable_cuda_graph, 'Hybrid context parallelism not supported with CUDA Graph' | |
| assert not args.use_megatron_fsdp, 'Hybrid context parallelism not supported with Megatron FSDP' | |
| assert args.dataloader_type == 'single', 'Hybrid context parallelism only supported with single dataloader type' | |
| assert args.calculate_per_token_loss, 'Hybrid context parallelism must be used with --calculate-per-token-loss' | |
| warnings.warn( | |
| '--hybrid-context-parallel is deprecated; use --dynamic-context-parallel.', | |
| DeprecationWarning, | |
| ) | |
| if args.hybrid_context_parallel: | |
| warnings.warn( | |
| '--hybrid-context-parallel is deprecated; use --dynamic-context-parallel.', | |
| DeprecationWarning, | |
| stacklevel=2, | |
| ) |
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 1390-1390: No explicit stacklevel keyword argument found
Set stacklevel=2
(B028)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@megatron/training/arguments.py` around lines 1389 - 1393, Update the
deprecation warning in the hybrid_context_parallel argument handling to pass
stacklevel=2, matching the existing warning in model_parallel_config.py and
ensuring callers receive the correct warning location.
Source: Linters/SAST tools
Comment on lines
+1400
to
+1403
| if args.dynamic_context_parallel: | ||
| assert not args.enable_cuda_graph, ( | ||
| 'Dynamic context parallelism is not supported with CUDA Graph' | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
args.enable_cuda_graph may already be deleted here, causing AttributeError instead of a clean assertion.
Earlier in this same function (lines 605-610), when --enable-cuda-graph is set, the code does args.cuda_graph_impl = "local"; del args.enable_cuda_graph. If a user then also passes --dynamic-context-parallel, this line's assert not args.enable_cuda_graph reads a deleted Namespace attribute and raises AttributeError rather than the intended assertion message. args.cuda_graph_impl is already normalized at this point and is safe to check instead.
🔧 Proposed fix
- assert not args.enable_cuda_graph, (
- 'Dynamic context parallelism is not supported with CUDA Graph'
- )
+ assert args.cuda_graph_impl == "none", (
+ 'Dynamic context parallelism is not supported with CUDA Graph'
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if args.dynamic_context_parallel: | |
| assert not args.enable_cuda_graph, ( | |
| 'Dynamic context parallelism is not supported with CUDA Graph' | |
| ) | |
| if args.dynamic_context_parallel: | |
| assert args.cuda_graph_impl == "none", ( | |
| 'Dynamic context parallelism is not supported with CUDA Graph' | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@megatron/training/arguments.py` around lines 1400 - 1403, Update the
dynamic_context_parallel validation to check the normalized args.cuda_graph_impl
value instead of accessing args.enable_cuda_graph, which may have been deleted
earlier in the argument-processing function. Preserve the existing assertion
message and reject the combination when CUDA Graph is enabled.
| vp_stage=vp_stage, | ||
| ) | ||
| is_hybrid_cp = args.hybrid_context_parallel | ||
| is_hybrid_cp = False |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate relevant files and inspect the surrounding code with line numbers.
git ls-files | rg '(^|/)(pretrain_gpt\.py|training\.py)$'
echo '--- pretrain_gpt.py (around line 129 and get_batch helpers) ---'
nl -ba pretrain_gpt.py | sed -n '90,220p'
echo '--- training.py (dummy_train_step and skip-iteration handling) ---'
nl -ba megatron/training/training.py | sed -n '2220,2315p'Repository: NVIDIA/Megatron-LM
Length of output: 306
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Show only the relevant sections.
echo '--- pretrain_gpt.py: is_hybrid_cp and CP group selection ---'
rg -n -A6 -B6 'is_hybrid_cp|get_batch_on_this_cp_rank|get_batch_on_this_tp_rank|dynamic_context_parallel|hybrid_context_parallel' pretrain_gpt.py
echo '--- training.py: dummy_train_step and skip-iteration handling ---'
rg -n -A8 -B8 'dummy_train_step|dynamic_context_parallel|sequence_packing_scheduler|skip-iteration|skip iteration|sequence packing' megatron/training/training.pyRepository: NVIDIA/Megatron-LM
Length of output: 10788
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the helper signatures and branches that consume is_hybrid_cp.
echo '--- helper definitions in megatron/core/utils.py ---'
rg -n -A20 -B8 'def get_batch_on_this_tp_rank|def get_batch_on_this_cp_rank|is_hybrid_cp|dynamic_cp_group_func|hybrid_cp_group_func' megatron/core/utils.py
echo '--- any other call sites of get_batch_on_this_cp_rank ---'
rg -n -A3 -B3 'get_batch_on_this_cp_rank\(' .
echo '--- any other uses of dynamic_context_parallel in pretrain_gpt.py ---'
rg -n -A3 -B3 'dynamic_context_parallel' pretrain_gpt.py megatron/training/training.pyRepository: NVIDIA/Megatron-LM
Length of output: 23276
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Print the exact helper definitions with line numbers.
for f in megatron/core/utils.py; do
echo "--- $f ---"
grep -n "def get_batch_on_this_tp_rank\|def get_batch_on_this_cp_rank" -n "$f"
sed -n '1,260p' "$f" | nl -ba | sed -n '1,260p'
doneRepository: NVIDIA/Megatron-LM
Length of output: 299
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Target just the helper bodies.
python3 - <<'PY'
from pathlib import Path
import re
p = Path("megatron/core/utils.py")
text = p.read_text()
for fn in ["get_batch_on_this_tp_rank", "get_batch_on_this_cp_rank"]:
m = re.search(rf"^def {fn}\b.*?(?=^def |\Z)", text, re.S | re.M)
if not m:
print(f"missing {fn}")
continue
start = text[:m.start()].count("\n") + 1
end = text[:m.end()].count("\n") + 1
print(f"--- {fn} lines {start}-{end} ---")
for i, line in enumerate(text[m.start():m.end()].splitlines(), start=start):
print(f"{i:4d}: {line}")
PYRepository: NVIDIA/Megatron-LM
Length of output: 17496
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
import re
p = Path("megatron/core/utils.py")
text = p.read_text()
for fn in ["get_batch_on_this_tp_rank", "get_batch_on_this_cp_rank"]:
m = re.search(rf"^def {fn}\b.*?(?=^def |\Z)", text, re.S | re.M)
if not m:
print(f"missing {fn}")
continue
start = text[:m.start()].count("\n") + 1
print(f"--- {fn} ---")
for i, line in enumerate(text[m.start():m.end()].splitlines(), start=start):
if "is_hybrid_cp" in line or "dynamic_cp_group_func" in line or "hybrid_cp_group_func" in line or "dynamic_cp" in line or "cp_group" in line:
print(f"{i:4d}: {line}")
PYRepository: NVIDIA/Megatron-LM
Length of output: 1958
Set is_hybrid_cp from args.dynamic_context_parallel here. get_batch_on_this_tp_rank() and get_batch_on_this_cp_rank() branch on this flag to broadcast local_cp_size and select the hybrid CP group, so hardcoding False disables the dynamic CP path in get_batch() and diverges from dummy_train_step().
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pretrain_gpt.py` at line 129, Initialize is_hybrid_cp from
args.dynamic_context_parallel instead of hardcoding it to False, so get_batch()
uses the dynamic CP behavior consistently with get_batch_on_this_tp_rank(),
get_batch_on_this_cp_rank(), and dummy_train_step().
Comment on lines
+455
to
+460
| if args.dynamic_context_parallel: | ||
| raise ValueError( | ||
| "Dynamic context parallelism is not supported by pretrain_hybrid.py. " | ||
| "Hybrid/Mamba layers do not yet implement per-microbatch context groups " | ||
| "and padding semantics; use pretrain_gpt.py for DCP training." | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Guard doesn't cover the deprecated hybrid_context_parallel alias.
This check only inspects args.dynamic_context_parallel. But later in this same file the dataset config is built with hybrid_context_parallel=args.hybrid_context_parallel, and GPTDatasetConfig.__post_init__ silently promotes hybrid_context_parallel=True into dynamic_context_parallel=True. A user who still passes the legacy --hybrid-context-parallel flag bypasses this guard entirely and ends up in the exact unsupported state ("Hybrid/Mamba layers do not yet implement per-microbatch context groups and padding semantics") the guard is meant to block.
🐛 Proposed fix
- if args.dynamic_context_parallel:
+ if args.dynamic_context_parallel or getattr(args, 'hybrid_context_parallel', False):
raise ValueError(
"Dynamic context parallelism is not supported by pretrain_hybrid.py. "
"Hybrid/Mamba layers do not yet implement per-microbatch context groups "
"and padding semantics; use pretrain_gpt.py for DCP training."
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if args.dynamic_context_parallel: | |
| raise ValueError( | |
| "Dynamic context parallelism is not supported by pretrain_hybrid.py. " | |
| "Hybrid/Mamba layers do not yet implement per-microbatch context groups " | |
| "and padding semantics; use pretrain_gpt.py for DCP training." | |
| ) | |
| if args.dynamic_context_parallel or getattr(args, 'hybrid_context_parallel', False): | |
| raise ValueError( | |
| "Dynamic context parallelism is not supported by pretrain_hybrid.py. " | |
| "Hybrid/Mamba layers do not yet implement per-microbatch context groups " | |
| "and padding semantics; use pretrain_gpt.py for DCP training." | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pretrain_hybrid.py` around lines 455 - 460, Update the unsupported-feature
guard in pretrain_hybrid.py to reject either args.dynamic_context_parallel or
the deprecated args.hybrid_context_parallel alias before dataset configuration
is built. Preserve the existing ValueError message and behavior for both flags.
yaoyu-33
approved these changes
Jul 15, 2026
Phlip79
reviewed
Jul 15, 2026
Member
There was a problem hiding this comment.
Can you make the equivalent change to hybrid_model.py?
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
6 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.
Training loss:
Throughput:
Summary
dp_balanced) and dynamic context parallelism (default_dynamic_cp).User-facing flags
--dynamic-context-parallel--calculate-per-token-loss,--max-seqlen-per-dp-cp-rank,--dataloader-type single, and--cuda-graph-impl none. Megatron FSDP, MoE, MLA, and hybrid/Mamba entrypoints are not supported.--min-dynamic-context-parallel-size NNselects the full domain.--sequence-packing-scheduler default_dynamic_cpDP x CP x --max-seqlen-per-dp-cp-rank. It is selected automatically by--dynamic-context-parallel.--sequence-packing-scheduler dp_balanced--hybrid-context-parallel--dynamic-context-parallel.--max-seqlen-per-dp-cp-rank N--sft-mock-dataset-config-json VALUE--sft --mock-data.Dynamic CP requires Transformer Engine 2.9 or newer.
Stack and scope
Validation
uv run isort --check-onlyon changed Python filesuv run ruff checkon changed Python filesgit diff --check533deb7b7e68a7ed5daa988fcfdbe0af73172cc1The available environment does not have a usable PyTorch test installation and has only two GPUs, so the distributed numerical tests were not run locally. CI was requested on the signed commit.
Summary by CodeRabbit
New Features
Bug Fixes
Documentation