Fix runtime CP group handling in attention for hybrid CP by ilml · Pull Request #6821 · NVIDIA/Megatron-LM

ilml · 2026-08-24T22:14:36Z

What does this PR do?

Two fixes for runtime (hybrid/dynamic) CP in the attention stack, where a
model built with static context_parallel_size == 1 binds a CP group per
microbatch via PackedSeqParams.cp_group:

Lazily create TE's auxiliary CP stream. The constructor only allocates
TEDotProductAttention.cp_stream when static CP > 1, so the runtime branch
passed cp_stream=None into TE and AttnFuncWithCPAndKVP2P crashed with
AttributeError: 'NoneType' object has no attribute 'wait_event'.
Route the runtime CP group to RoPE. apply_rotary_pos_emb derives
(cp_size, cp_rank) for the fused THD RoPE kernel from its cp_group
argument, but attention always passed the static pg_collection.cp. Under
runtime CP each rank holds only its group's zigzag slice while cu_seqlens
describe the full sequence, so the kernel indexed positions past the
tensor: illegal memory accesses on short sub-samples (caught with
CUDA_LAUNCH_BLOCKING=1 in fused_rope_forward), and wrong positions,
NaN gradients, then a softmax-lse device assert on long ones. This is the
dangerous one: with different luck it trains-but-wrong.

Test

New unit tests assert (a) RoPE receives the runtime group when
packed_seq_params.cp_group is bound and the static group otherwise, and
(b) the CP stream is lazily created and handed to TE's
set_context_parallel_group.

Validation

End to end on 4x GB200: 8B GPT, 100k-token THD workload, mixed CP1/2/4
groups, finite losses/grad norms. One of three independent fixes required to
make hybrid CP runnable (siblings: schema-bridge and rerun-iterator PRs);
supersedes the combined draft #6816.

Two fixes for models built with context_parallel_size == 1 that bind a CP group per microbatch at runtime (hybrid/dynamic context parallelism): 1. TEDotProductAttention only allocates its auxiliary CP CUDA stream in the constructor when static CP > 1, so the runtime branch passed cp_stream=None into TE and AttnFuncWithCPAndKVP2P crashed with AttributeError: 'NoneType' object has no attribute 'wait_event'. Create the stream lazily when a runtime CP group is bound. 2. apply_rotary_pos_emb derives (cp_size, cp_rank) for the fused THD RoPE kernel from its cp_group argument, but attention always passed the static pg_collection.cp group. Under hybrid CP each rank holds only its runtime group's zigzag slice while cu_seqlens describe the full sequence, so the kernel indexed positions past the tensor: illegal memory accesses on short sub-samples, and wrong positions, NaN gradients and a softmax-lse device assert on long ones. Prefer packed_seq_params.cp_group when the runtime path binds one. Unit tests cover both: RoPE must receive the runtime group (and the static group when none is bound), and the CP stream must be created lazily and passed to TE. Validated end to end on 4x GB200 (8B GPT, 100k-token THD profile, mixed CP1/2/4 groups, finite losses and grad norms). Signed-off-by: Tom Long <tolong@nvidia.com>

copy-pr-bot · 2026-08-24T22:14:41Z

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.

Review follow-up: instead of a RoPE-local variable, re-assign pg_collection.cp to the runtime CP group so every CP consumer in the forward benefits. Because hybrid CP binds a different group per microbatch and local_cp_size == 1 sub-samples carry no group, the build-time group (captured at construction) is restored whenever no runtime group is bound - otherwise a CP1 microbatch following a CP2 one would inherit the stale larger group and recreate the wrong-positions bug. The unit test now also asserts the bind and the restore on the collection itself. Signed-off-by: Tom Long <tolong@nvidia.com>

ilml · 2026-08-24T22:35:57Z

Done — took @parthmannan's suggestion and went one step further: the runtime group is now bound on pg_collection.cp itself so any CP consumer in the forward benefits. One subtlety this surfaced: hybrid CP binds a different group per microbatch, and local_cp_size == 1 sub-samples carry no group at all — a bare re-assign would leave a CP1 microbatch that follows a CP2 one seeing the stale larger group (wrong RoPE positions again). So the module captures the build-time group at construction and restores it whenever no runtime group is bound. The unit test now asserts both the bind and the restore on the collection; re-validated on GB200 (2 passed).

Keeps the tests consistent with the runtime-CP contract (local_cp_size == 1 sub-samples carry no cp_group) that NVIDIA#6825 asserts. Signed-off-by: Tom Long <tolong@nvidia.com>

ilml · 2026-08-24T23:10:37Z

Verified and audited, per the question above:

local_cp_size == 1 carries no CP group by construction: get_batch_on_this_cp_rank only binds hybrid_cp_group when local_cp_size > 1.
Swept every consumer of packed_seq_params.cp_group / local_cp_size in core for the "group always present" assumption: the TE wrapper handles and asserts the None case; gpt_model's three rotary call sites are safe (the packed_seq=True path never CP-slices — per-rank slicing happens in the fused THD kernel, which is what this PR fixes); MLA/absorbed-MLA hard-assert hybrid CP off. No other latent group-assumption bugs found.
The contract was only asserted in one direction ("no group ⇒ size 1"). Follow-up Assert the runtime-CP group contract both ways #6825 adds the converse ("size 1 ⇒ no group"), documents the invariant on PackedSeqParams, and adds a unit test — so a future producer that starts binding size-1 groups fails loudly instead of silently flipping the convention. Also updated this PR's tests to bind runtime groups with local_cp_size = 2 for consistency with that contract.

zhongbozhu

LGTM

ilml · 2026-09-04T17:56:10Z

/ok to test 9493661

ilml · 2026-09-04T21:43:48Z

/ok to test 14a19cc

The repo black config skips magic trailing commas; apply exactly the hunks the lint job suggests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: ilml <tolong@nvidia.com>

nemo-automation-bot · 2026-09-05T00:57:57Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/33934670269

ilml mentioned this pull request Aug 24, 2026

Fix hybrid context parallel end-to-end integration #6816

Closed

ilml marked this pull request as ready for review August 24, 2026 22:22

ilml requested review from a team as code owners August 24, 2026 22:22

github-actions Bot added the complexity: low label Aug 24, 2026

parthmannan reviewed Aug 24, 2026

View reviewed changes

Comment thread megatron/core/transformer/attention.py Resolved

parthmannan approved these changes Aug 24, 2026

View reviewed changes

Phlip79 approved these changes Aug 24, 2026

View reviewed changes

ilml mentioned this pull request Aug 24, 2026

Assert the runtime-CP group contract both ways #6825

Merged

Use local_cp_size > 1 when binding a runtime CP group in tests

9493661

Keeps the tests consistent with the runtime-CP contract (local_cp_size == 1 sub-samples carry no cp_group) that NVIDIA#6825 asserts. Signed-off-by: Tom Long <tolong@nvidia.com>

zhongbozhu approved these changes Aug 25, 2026

View reviewed changes

svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Aug 25, 2026

lhb8125 mentioned this pull request Sep 1, 2026

Megatron-LM Monthly Activity Report — August 2026 #7006

Open

lhb8125 mentioned this pull request Sep 4, 2026

Megatron-LM Monthly Activity Report — August 2026 (Final, through August 31, 2026) #6921

Open

jaredcasper approved these changes Sep 4, 2026

View reviewed changes

svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Sep 4, 2026

ilml enabled auto-merge September 4, 2026 17:54

copy-pr-bot Bot temporarily deployed to public September 4, 2026 18:00 Inactive

Collapse parallel_attention test calls for black

14a19cc

The repo black config skips magic trailing commas; apply exactly the hunks the lint job suggests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: ilml <tolong@nvidia.com>

copy-pr-bot Bot temporarily deployed to test September 4, 2026 21:44 Inactive

copy-pr-bot Bot temporarily deployed to public September 4, 2026 21:47 Inactive

ilml added this pull request to the merge queue Sep 5, 2026

Merged via the queue into NVIDIA:main with commit df4996f Sep 5, 2026
101 of 102 checks passed

ilml deleted the fix/hybrid-cp-runtime-attention branch September 5, 2026 03:59

lhb8125 mentioned this pull request Sep 7, 2026

Megatron-LM Monthly Activity Report — September 2026 (Final, through September 30, 2026) #7100

Open

lhb8125 mentioned this pull request Oct 1, 2026

Megatron-LM Monthly Activity Report — September 2026 #7773

Open

Navigation Menu

Repository navigation

Fix runtime CP group handling in attention for hybrid CP - #6821

Fix runtime CP group handling in attention for hybrid CP#6821
ilml merged 4 commits into
NVIDIA:mainfrom
ilml:fix/hybrid-cp-runtime-attention

ilml commented Aug 24, 2026

Uh oh!

copy-pr-bot Bot commented Aug 24, 2026

Uh oh!

ilml commented Aug 24, 2026

Uh oh!

Uh oh!

ilml commented Aug 24, 2026

Uh oh!

zhongbozhu left a comment

Uh oh!

ilml commented Sep 4, 2026

Uh oh!

ilml commented Sep 4, 2026

Uh oh!

nemo-automation-bot Bot commented Sep 5, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Conversation

ilml commented Aug 24, 2026

What does this PR do?

Test

Validation

Uh oh!

copy-pr-bot Bot commented Aug 24, 2026

Uh oh!

ilml commented Aug 24, 2026

Uh oh!

Uh oh!

ilml commented Aug 24, 2026

Uh oh!

zhongbozhu left a comment

Choose a reason for hiding this comment

Uh oh!

ilml commented Sep 4, 2026

Uh oh!

ilml commented Sep 4, 2026

Uh oh!

nemo-automation-bot Bot commented Sep 5, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants