feat: add gemma4 moe ring CP - #1914

khazic · 2026-04-20T12:27:26Z

What does this PR do?

Adds a Gemma4-specific Context Parallelism (CP) path for dense and MoE VLM
fine-tuning, without changing the generic CP behavior of any other model.

Gemma4 cannot use the existing causal ring-attention CP path for this workload:
its decoder mask combines image-token bidirectional attention, sliding-window
attention, and packed-sequence document boundaries, and its large head
dimension (256) rules out FlashAttention-2 and the SDPA flash/mem-efficient
kernels. This PR routes only Gemma4 decoder self-attention through a model-owned
p2p ring FlexAttention path, plugged into a small model-agnostic seam in the CP
infrastructure.

At a high level, each CP rank keeps its local query shard. K/V (and the token
metadata needed to build masks) are rotated around the CP ranks via point-to-point
ring exchange — they are never fully all-gathered. For each ring step the active
Gemma4 layer builds a torch.compile(flex_attention) block mask (full / sliding,
packed-document boundaries, image bidirectional groups, padding rows), computes
attention against the current K/V chunk, and the per-chunk outputs are combined with
an online-softmax (logsumexp) merge. Only the local query output shard is returned.

Design

CP attention routing is model-agnostic (apply_cp / attach_cp_attention_hooks):

TE DotProductAttention → TE's own context-parallel group (unchanged).
Non-TE attention exposing run_cp_manual_attention → the model owns its CP
transport + compute through that seam (Gemma4's flex ring).
Any other non-TE attention → the classic DTensor SDPA CP path.

cp_utils contains no model names — only the TE-vs-not split and the seam; the
Gemma4 ring lives entirely in components/models/gemma4_moe/cp_attention.py and is
attached via attach_gemma4_cp_ring_attention. This keeps the path easy to extend to
other models that need their own CP attention.

Changelog

components/models/gemma4_moe/cp_attention.py (new)
- Gemma4 p2p ring FlexAttention CP: ring K/V collection + p2p exchange, per-chunk
  FlexAttention with layer-aware masks (causal, sliding-window, packed-document
  boundaries, image-token bidirectional groups, padding-only query rows), online
  softmax merge, a custom autograd Function with ring grad exchange, the
  run_cp_manual_attention entry, and attach_gemma4_cp_ring_attention.
components/models/gemma4_moe/model.py
- CP-aware forward for dense and MoE paths.
- prepare_model_inputs_for_cp / prepare_inputs_embeds_for_cp to build text+image
  inputs_embeds on the unsharded sequence before CP sharding.
- Packed / vision-bidirectional FlexAttention block-mask mapping; pad-token resolution.
components/distributed/cp_utils.py
- Model-agnostic CP attention seam + SDPA-swap hooks; manual contiguous sequence
  sharding for the model-owned path; single-document _packed_seq_ids synthesis;
  padding / position-id helpers. Generic models keep the DTensor context_parallel
  path.
components/moe/parallelizer.py
- Capability-based apply_cp routing (TE vs hooks; model-owned vs generic DTensor).
_transformers/infrastructure.py
- Attach the dense (non-TE) CP hooks when cp_size > 1.
recipes/vlm/finetune.py
- Pre-compute inputs_embeds before CP sharding for Gemma4 + image inputs.
- Sum validation label-token counts over CP (the VLM recipe shards labels before
  counting), so val_loss divides like-for-like.
examples/vlm_finetune/gemma4/gemma4_26b_a4b_moe_medpix_ep8cp2_4k.yaml
- Runnable Gemma4-26B-A4B MoE MedPix EP8×CP2 config at 4k sequence length.
Tests
- Comprehensive unit tests for the ring attention (mask builders, ring/grad p2p
  exchange, online-softmax merge, autograd, manual-attention entry), the CP wiring
  seam and hooks, the packed/vision mask mapping, parallelizer routing, the
  infrastructure CP-hook attach, and the validation token-counting fix —
  bringing the PR's changed lines to ~100% unit coverage (CPU-only, no GPU
  required; FlexAttention exercised in CPU-eager mode).

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?

Test Results

Targeted behavior

Gemma4 CP2 routes through the compiled FlexAttention p2p ring. Observed log line
(EP8×CP2, 2k seq, 26B-A4B MoE):
- Gemma4 CP using compiled flex_attention p2p ring. Q=(1, 16, 1024, 256) K_local=(1, 8, 1024, 256) head_dim=256->256 cp_rank=0 cp_size=2
Packed-sequence mask correctness target:
- CP2 forward(A alone) == CP2 forward(packed [A, B])[:, A slice]
- CP2 forward(B alone) == CP2 forward(packed [A, B])[:, B slice]
Text-only, image, packed, and no-pack inputs all route through the same Gemma4 CP
ring implementation, with layer-specific mask construction.
Unit tests verify the ring forward and gradients match a full-attention SDPA
reference at cp_size=1, and the online-softmax merge matches a single full
softmax.
ep8cp1 ep8cp2 parity runs:
https://wandb.ai/Nemo-automodel/huiyingl_workspace/runs/plvv3r6b
https://wandb.ai/Nemo-automodel/huiyingl_workspace/runs/itjikprm

single node 32k seqlen on mock https://wandb.ai/Nemo-automodel/huiyingl_workspace/runs/bm3lch7d

copy-pr-bot · 2026-04-20T12:27:29Z

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.

khazic · 2026-04-20T12:27:38Z

Heads-up on validation status: I'm currently running end-to-end tests on 4 nodes (32x A100-80GB) for both the MoE path (gemma4_26b_a4b_moe / gemma-4-E4B) and the 31B dense path with cp_size=2. Will post the loss-parity numbers (CP=1 vs CP=2) and step-time figures once the runs finish. The mock-dataset configs in this PR are already validated locally; please hold off on merging until the multi-node runs come back green.

khazic · 2026-04-20T15:02:02Z

4-node smoke-test results (32x H100-80GB)

End-to-end CP verification on 4 nodes (32 x H100-80GB) with Gemma4-E4B MoE (~7B total, 4B activated).

Loss parity: CP=1 vs CP=2 (matched per-GPU shard size)

To verify the ring-attention + flex_attention path is mathematically equivalent to the non-CP baseline, I ran two matched configurations (same per-GPU S_local, different cp_size):

step	CP=1, seq=2048	CP=2, seq=4096	Δ
0	16.9419	17.1810	0.24
1	10.1633	10.4407	0.28
2	8.4085	8.7233	0.31
3	7.1558	7.4028	0.24

Trajectories overlap within noise. Residual Δ is from different global batch sizes (32 vs 16) sampling different mock examples per step plus reduction-order FP differences. grad_norm tracks similarly: 1256 -> 68 (CP=1) vs 1136 -> 77 (CP=2).

Topology / infra notes

cp_size=2, dp_size=16 (CP within node over NVLink, DP across nodes over TCP)
activation_checkpointing: true, PYTORCH_ALLOC_CONF=expandable_segments:True
NCCL_IB_HCA=^=all + NCCL_NET_PLUGIN=none were needed to work around a NCCL 2.27.5 bug in the K8s container where the proxy thread kept attempting IB-type handshakes even with NCCL_IB_DISABLE=1. Not related to this PR.

31B dense: not yet validated

I have not yet run the 31B dense config end-to-end. Earlier attempts OOM'd during backward on a 32x 80GB cluster with pure FSDP+CP (weights + Adam state + flex_attention backward workspace together exceed per-GPU budget). 31B long-context training almost certainly needs TP alongside CP to shard the hidden dimension as well; that combination is not covered by this PR and will be a follow-up. The 31B yamls in this PR are provided as configuration templates; they are not expected to run as-is without either TP or a larger cluster.

Why all-gather + flex_attention instead of ring-attention

The all-gather variant here is deliberately simple: each rank gathers full K/V along the sequence dimension then runs a compiled flex_attention with a per-rank causal block mask. We chose this over a ring-attention implementation for two reasons:

Correctness under mixed DTensor/plain-tensor paths. PyTorch's native context_parallel DTensor dispatch into SDPA kept hitting a "mixed torch.Tensor and DTensor" crash in the attention inner_fn (the compiled decoder propagates a local-tensor Q while K/V return as DTensors after all-gather). All-gather-then-local-SDPA avoids the mixed path entirely and works on every PyTorch version we tested.
Reused infra. flex_attention already gives us Flash-style O(N) memory via its Triton tiled kernel and accepts arbitrary per-rank block masks; we just had to wrap it in torch.compile and pass kernel_options to shrink block sizes when head_dim >= 256 (the Gemma4 31B global-attention-layer variant has head_dim=512).

A proper ring-attention implementation would save the S_full K/V memory per rank (O(N) in seq instead of O(N/CP)), and could let us use Flash Attention directly (because each ring chunk is a standard causal attention). That is the right long-term direction but is a larger change; keeping it out of this PR so the diff stays reviewable.

Known limitations (follow-up PR candidates)

create_block_mask materializes a dense [S_local, S_full] bool mask, which allocates ~1 GiB at seq=16K and OOMs before flex_attention even starts. Needs a sparse or streaming construction.
CE loss materializes a fp32 [B*S_local, 262144] logit tensor (~4 GiB at seq=8K). Chunked CE would unblock longer contexts.
Ring-attention path (see above) to eliminate the S_full K/V memory footprint.

khazic · 2026-04-20T15:14:09Z

A few more notes to round out the review context.

Ring attention is the right long-term direction

The all-gather variant in this PR is a pragmatic starting point, not the ideal design. A proper ring attention implementation (each rank keeps only its local K/V shard and rotates them in a ring) would be strictly better along every axis:

Memory: avoids the per-rank S_full K/V materialization, so attention memory goes from O(S_full/CP) (current) back to O(S_local) per rank — the actual reason people reach for CP in the first place.
Compute path compatibility: each ring chunk becomes a standard causal (or fully-open) attention on contiguous Q, K, V of the same length, so Flash Attention accepts it directly. We then do not need flex_attention at all and recover its best-case throughput.
Communication overlap: ring P2P sends can be overlapped with per-chunk Flash attention, which the all-gather pattern cannot.

I chose the all-gather + flex_attention path in this PR because it (a) avoids the "mixed torch.Tensor / DTensor" crash in PyTorch's native context_parallel dispatch on current PyTorch versions, and (b) reuses flex_attention's already-tested block-sparse kernel so the CP-specific code stays small and reviewable. The follow-up should be a real ring attention wrapper; I'm happy to take that as a separate PR.

VLM path: image and no-image both validated

The recipe-level changes in finetune.py handle both paths:

Text-only (no pixel_values in the batch): make_cp_batch_and_ctx shards input_ids, labels, and position_ids directly. The gemma4_4b_cp2_mock.yaml smoke test runs this path (mock dataset with num_images_per_sample: 0); loss-parity numbers are in the previous comment.
With images (pixel_values present): we cannot CP-shard input_ids before image features are merged, because image tokens are concentrated at the start of the sequence and the masked_scatter inside HF Gemma4 cannot redistribute image features across CP ranks correctly (each rank would see a different number of image-token slots). Solution: at the recipe layer (finetune.py:920-954), when CP is active and pixel_values is present, we pre-compute inputs_embeds on the full unsharded sequence via a _pre_embed_only=True pass through model.__call__, then hand inputs_embeds (not input_ids) to make_cp_batch_and_ctx. Also validated on a mock VLM dataset.
We route the pre-embed through model.__call__ (not direct method call) specifically so the FSDP2 pre-forward hook fires and all-gathers the sharded embed_tokens weight into a Replicate placement — calling the method directly bypasses the hook and we hit a mixed DTensor / plain-tensor error. This is why prepare_inputs_embeds_for_cp is exposed via _pre_embed_only instead of called directly.

Debug decisions worth flagging for the review

These are the non-obvious ones we landed on during a multi-day debugging session; a reviewer may want to push back or suggest a better shape.

flex_attention must be called through torch.compile. Without torch.compile the eager implementation materializes the full [B, nH, S_local, S_full] scores matrix and OOMs immediately on any real sequence length — the exact opposite of flex_attention's purpose. We lazily compile once per process and cache (cp_utils.py:160-171). A fresh warning appeared in PyTorch 2.10 that surfaces this; without the compile wrapper I spent an embarrassing amount of time misdiagnosing the OOM as a memory-pressure problem.
enable_gqa=True must be passed explicitly to flex_attention. Gemma 4 uses GQA (E4B has 16 Q / 2 KV heads; 31B has 32 Q / 16 KV). Our fallback SDPA path already auto-enabled this flag, but the initial flex_attention call did not, which caused it to always raise ValueError: Expect query and key/value to have the same number of heads, silently falling back to SDPA MATH and OOMing. One-line fix in cp_utils.py:260-266.
Triton kernel_options must be shrunk for large head_dim. Gemma 4 31B's "global attention" layer variant uses head_dim=512. With default flex_attention BLOCK_M/BLOCK_N=128 and num_stages=3, the required Triton shared memory is ~200 KB, which exceeds A100/H100's ~163 KB limit; the kernel errors out with No valid triton configs. OutOfMemoryError: out of resource: triton_tem_fused_flex_attention_0 Required: 200704 Hardware limit: 166912. We now pass kernel_options={"BLOCK_M": 32, "BLOCK_N": 32, "num_stages": 1, ...} whenever head_dim >= 256 (cp_utils.py:265-283). Throughput trade-off: 32x32 blocks are ~2-3x slower than the default for the same kernel, but they fit. Could be made adaptive in a future PR.
Silent except-catch on the flex_attention call is intentional — with a logger. The primary path is wrapped in try: ... except Exception as err: so that any Triton compile failure (e.g. on an unsupported shape) falls back to SDPA rather than crashing the run, but we log the exception type, shapes, and traceback once per process (cp_utils.py:266-284) so the silent fallback does not hide an actual kernel bug. This was added after a multi-hour chase where flex_attention was failing silently and the only visible symptom was an unrelated-looking 1 GiB OOM deep inside SDPA.
Manual sequence slicing, not context_parallel DTensor dispatch. PyTorch's torch.distributed.tensor.experimental.context_parallel wraps Q/K/V as DTensors and relies on the DTensor SDPA dispatch to trigger ring communication. On the PyTorch versions we tested this breaks because the compiled decoder propagates a plain-tensor Q through the graph while K/V all-gather comes back as DTensors, and the inner_fn trips a "mixed torch.Tensor and DTensor" assert. We moved to manual sequence slicing in make_cp_batch_and_ctx and put the CP communication inside an SDPA hook (attach_cp_sdpa_hooks), which keeps activations outside of attention as plain tensors end-to-end.
Gemma4 per-layer-input 640 GiB broadcast. Orthogonal to CP but hit during this work. Older transformers versions of Gemma4TextModel.get_per_layer_inputs call the fallback path [B, S, 1, H] == [1, 1, vocab, H] → [B, S, vocab, H] bool when input_ids is None, which is ~640 GiB on Gemma4-E4B. _patch_gemma4_per_layer_inputs in infrastructure.py:109-144 replaces it with a (inputs_embeds @ embed_tokens.T).argmax(-1) reverse lookup (~1 GiB, and correct as long as Gemma's inputs_embeds = embed_tokens(ids) * sqrt(H) invariant holds, which it does on all Gemma 4 variants).

Happy to split any of the above into their own PRs if the reviewer prefers smaller diffs.

HuiyingLi · 2026-04-20T16:09:26Z

/ok to test 1699fd2

HuiyingLi · 2026-04-20T23:39:43Z

/claude review

claude · 2026-04-20T23:42:37Z

+    # from the propagated hidden-states vs plain K/V after allgather).
+    # Ring attention is still performed correctly because attach_cp_sdpa_hooks
+    # (registered in infrastructure.py for all CP>1 cases) wraps Q/K/V as
+    # sequence-sharded DTensors at the SDPA call site; the DTensor SDPA
+    # dispatch then handles the ring communication.  All activations outside
+    # the SDPA call remain plain tensors.


Stale comment: this describes the old DTensor-based approach, but _cp_sdpa now uses manual all-gather of K/V + per-rank causal masking — it no longer wraps Q/K/V as DTensors or relies on DTensor SDPA dispatch.

Suggested change

# from the propagated hidden-states vs plain K/V after allgather).

# Ring attention is still performed correctly because attach_cp_sdpa_hooks

# (registered in infrastructure.py for all CP>1 cases) wraps Q/K/V as

# sequence-sharded DTensors at the SDPA call site; the DTensor SDPA

# dispatch then handles the ring communication. All activations outside

# the SDPA call remain plain tensors.

# Ring attention is still performed correctly because attach_cp_sdpa_hooks

# (registered in infrastructure.py for all CP>1 cases) replaces SDPA with

# a manual all-gather of K/V followed by per-rank causal masking. All

# activations outside the SDPA call remain plain tensors.

claude · 2026-04-20T23:42:41Z

+                batch.pop("image_grid_thw", None)
+                batch.pop("video_grid_thw", None)
+                batch.pop("n_images_per_sample", None)
+                with torch.no_grad():
+                    batch["inputs_embeds"] = _model(
+                        input_ids=_input_ids,
+                        pixel_values=_pixel_values,


Nit/footgun: torch.no_grad() here means the returned inputs_embeds won't carry gradients. This is fine when freeze_embeddings: true (as in all the CP yaml configs), but if a future user runs CP with unfrozen embeddings, gradients won't flow back through the embedding layer and the embeddings will silently stop training.

Consider adding a comment or an assertion to make this constraint explicit, e.g.:

# no_grad is safe here because embeddings and vision tower are frozen under CP; # if unfreezing, this must be revisited to preserve gradient flow.

claude · 2026-04-20T23:42:51Z

Review

Good work overall — the manual all-gather approach is well-reasoned and the defensive NotImplementedError guards for unsupported CP+pixel_values paths are helpful.

Two inline comments:

Stale comment in cp_utils.py (lines 497–502): Still describes the old DTensor-wrapping approach, but _cp_sdpa now does manual all-gather + per-rank masking. Suggested replacement text inline.
torch.no_grad() footgun in finetune.py: The pre-embed path runs under no_grad, which silently breaks gradient flow to embeddings if someone later runs CP with freeze_embeddings: false. Worth a guarding comment.

Missing test coverage: The PR description mentions unit tests for prepare_inputs_embeds_for_cp, but no such test exists in the diff. The two new tests cover the mm_token_type_ids nulling logic, but prepare_inputs_embeds_for_cp (the image-feature merge helper) is untested — consider adding a basic test that verifies image features are correctly scattered into the embeddings tensor.

HuiyingLi · 2026-04-21T02:04:46Z

4-node smoke-test results (32x H100-80GB)

End-to-end CP verification on 4 nodes (32 x H100-80GB) with Gemma4-E4B MoE (~7B total, 4B activated).

Loss parity: CP=1 vs CP=2 (matched per-GPU shard size)

To verify the ring-attention + flex_attention path is mathematically equivalent to the non-CP baseline, I ran two matched configurations (same per-GPU S_local, different cp_size):

step CP=1, seq=2048 CP=2, seq=4096 Δ
0 16.9419 17.1810 0.24
1 10.1633 10.4407 0.28
2 8.4085 8.7233 0.31
3 7.1558 7.4028 0.24
Trajectories overlap within noise. Residual Δ is from different global batch sizes (32 vs 16) sampling different mock examples per step plus reduction-order FP differences. grad_norm tracks similarly: 1256 -> 68 (CP=1) vs 1136 -> 77 (CP=2).

Topology / infra notes

cp_size=2, dp_size=16 (CP within node over NVLink, DP across nodes over TCP)

activation_checkpointing: true, PYTORCH_ALLOC_CONF=expandable_segments:True

NCCL_IB_HCA=^=all + NCCL_NET_PLUGIN=none were needed to work around a NCCL 2.27.5 bug in the K8s container where the proxy thread kept attempting IB-type handshakes even with NCCL_IB_DISABLE=1. Not related to this PR.

31B dense: not yet validated

I have not yet run the 31B dense config end-to-end. Earlier attempts OOM'd during backward on a 32x 80GB cluster with pure FSDP+CP (weights + Adam state + flex_attention backward workspace together exceed per-GPU budget). 31B long-context training almost certainly needs TP alongside CP to shard the hidden dimension as well; that combination is not covered by this PR and will be a follow-up. The 31B yamls in this PR are provided as configuration templates; they are not expected to run as-is without either TP or a larger cluster.

Why all-gather + flex_attention instead of ring-attention

The all-gather variant here is deliberately simple: each rank gathers full K/V along the sequence dimension then runs a compiled flex_attention with a per-rank causal block mask. We chose this over a ring-attention implementation for two reasons:

Correctness under mixed DTensor/plain-tensor paths. PyTorch's native context_parallel DTensor dispatch into SDPA kept hitting a "mixed torch.Tensor and DTensor" crash in the attention inner_fn (the compiled decoder propagates a local-tensor Q while K/V return as DTensors after all-gather). All-gather-then-local-SDPA avoids the mixed path entirely and works on every PyTorch version we tested.

Reused infra. flex_attention already gives us Flash-style O(N) memory via its Triton tiled kernel and accepts arbitrary per-rank block masks; we just had to wrap it in torch.compile and pass kernel_options to shrink block sizes when head_dim >= 256 (the Gemma4 31B global-attention-layer variant has head_dim=512).

A proper ring-attention implementation would save the S_full K/V memory per rank (O(N) in seq instead of O(N/CP)), and could let us use Flash Attention directly (because each ring chunk is a standard causal attention). That is the right long-term direction but is a larger change; keeping it out of this PR so the diff stays reviewable.

Known limitations (follow-up PR candidates)

create_block_mask materializes a dense [S_local, S_full] bool mask, which allocates ~1 GiB at seq=16K and OOMs before flex_attention even starts. Needs a sparse or streaming construction.

CE loss materializes a fp32 [B*S_local, 262144] logit tensor (~4 GiB at seq=8K). Chunked CE would unblock longer contexts.

Ring-attention path (see above) to eliminate the S_full K/V memory footprint.

We can do 31B in a separate PR. Do you mind removing the 31B template config

…ep in cp_contiguous_shard Stop splitting make_cp_batch_and_ctx into shared helpers. Its load-balanced (non-manual) path is now byte-identical to upstream; the only addition is a single early `if batch.pop("_cp_manual", False)` hop into the model-owned contiguous-shard path. The gemma4 pre-shard prep (attention_mask -> padding_mask conversion, position_ids injection, labels resolution) moves into cp_contiguous_shard.make_contiguous_shard_cp_batch_and_ctx, so the classic CP path reverts to main exactly and all gemma4-specific logic is contained in cp_contiguous_shard. Removed: _prepare_cp_batch_common, _make_context_parallel_cp_batch_and_ctx, and the module-level _get_submesh/_get_mesh_size (nested back inside the function as in main). No behavior change: 26B gemma4 ep8 cp1-vs-cp2 5-step losses are bit-identical to the pre-refactor cp2 run; parity Δ < 0.02/step. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

_run_validation_epoch counted num_label_tokens per-CP-shard and all-reduced both loss and tokens with include_cp=True. That is equivalent but diverges from recipes/llm/train_ft.py (and main), which count the full label tokens *before* CP sharding, reduce the (CP-sharded) loss with include_cp=True, and reduce the token count over DP only -- every CP rank holds the identical full count, so summing it over CP would scale val_loss down by cp_size. Revert to the train_ft.py convention: count (batch["labels"] != -100).sum() pre-shard, reduce total_tokens/total_num_label_tokens without include_cp. Test reverts to test_run_validation_epoch_does_not_sum_tokens_over_cp accordingly. Shrinks the vlm/finetune.py diff vs main to a single line (the pre-embed no_grad removal, which is a real CP fix and is kept). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

…ke TE/DSV4) Move the CP mask handling for Gemma4 into its model-owned ring: the ring pre-hook now nulls the (CP-invalid) full-sequence attention_mask and forces is_causal=True itself, instead of depending on cp_utils' generic attach_context_parallel_hooks mask-strip. attach_context_parallel_hooks now self-skips modules exposing run_cp_manual_attention (mirroring attach_cp_attention_hooks), and the MoE parallelizer only attaches the generic hooks for non-model-owned attention. This makes Gemma4 fully self-contained for CP attention -- transport AND masking -- the same way TE's DotProductAttention is in DSV4. Semantically a no-op: the ring receives identical inputs (attention_mask=None, is_causal=True) either way -- verified by instrumenting the pre-hook. The 26B ep8 cp1-vs-cp2 5-step losses stay within the CP parity band (<0.02/step). Note: cp2 grad_norm/loss shift at the 4th-5th decimal vs the pre-edit cached run (e.g. step0 grad_norm 58.9111 -> 58.8312). This is a torch.compile/inductor recompilation artifact, NOT a masking change: any edit to cp_attention.py -- even a bare `_ = 0` in the pre-hook -- reproduces the identical shift, because editing the source busts the kernel cache and re-autotunes. Both values are valid within fp tolerance. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

…d-aid apply_model_infrastructure ran its generic CP attention block (line ~632) for every non-TE model with cp_size>1 -- including MoE models, which already get their full CP setup from the MoE parallelizer's apply_cp (via _shard_ep_fsdp). For Gemma4 this meant the generic hooks were applied twice (the source of the earlier "clobber"), which only worked because cp_utils carried skip-guards. Gate the infra CP block on ep_size <= 1 so MoE models are handled solely by apply_cp -- mirroring how TE/DSV4 is already excluded via _uses_te_attention. With that, Gemma4 (and any model-owned/MoE CP model) no longer reaches attach_context_parallel_hooks at all, so its run_cp_manual_attention skip-guard there was pure dead weight; revert it. cp_utils' attach_context_parallel_hooks is now byte-identical to upstream again. No behavior change: 26B ep8 cp2 5-step losses are bit-identical to the prior commit (the :632 pass was already a no-op for Gemma4 via the skip-guards). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

…ssic path attach_cp_sdpa_hooks was renamed to attach_cp_attention_hooks and grown a run_cp_manual_attention skip-guard + _cp_uses_attention_hook marker + vision exclusion. But gemma4 is model-owned (its ring is installed by setup_cp_attention) and never reaches this function: apply_cp only calls it when needs_generic_cp_attention (gemma4 takes the model-owned branch), and infra's CP block now skips MoE. DSV4 is TE. So nothing on this branch routes through the classic DTensor-SDPA path's renamed form. Restore the function to upstream verbatim (name + body) and point the dense (infra) and generic-MoE (parallelizer) callers + their tests back at attach_cp_sdpa_hooks. The classic-CP path is now byte-identical to main; gemma4 is unaffected (it does not use it). The vision-exclusion helper / mask-strip marker in attach_context_parallel_hooks are a separate follow-up. Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

…ion frozenset gemma4 is model-owned and never reaches attach_context_parallel_hooks (apply_cp calls it only for needs_generic_cp_attention; infra's CP block now skips MoE). So its two divergences from upstream were unneeded here: - the _cp_uses_attention_hook fast-path branch was dead (only gemma4's ring sets the flag, and gemma4 never hits this hook; the flag stays for gemma4's decoder). - the _CP_NON_TEXT_MODULE_PATH_PARTS vision exclusion was protecting gemma4's vision tower at the infra full-model pass, now moot since that pass skips MoE. Restore the function verbatim to main, delete the frozenset + _is_cp_non_text / _is_cp_attention_module_name helpers, and drop the 3 tests that covered the removed behavior. cp_utils.py's only remaining diff vs main is the _cp_manual dispatch in make_cp_batch_and_ctx (gemma4's contiguous-shard batch seam). Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

… + stale comment The CP-attention block carried a stale comment claiming routing happens "by capability at attention time ... in cp_utils" -- that capability skip-guard was removed; routing is decided here in the loop (TE / model-owned setup_cp_attention / generic). Rewrite the comment to match. Also drop the redundant needs_cp_attention_hooks flag: it was just (model_owned or generic) and produced a misleading "Attached CP attention hooks (model-owned)" log even though model-owned attention attaches nothing in this block. Gate the generic-hook attach on needs_generic_cp_attention directly and log model-owned vs generic accurately. No behavior change: gemma4 (model-owned), DSV4 (TE), and generic CP all route exactly as before; 62 parallelizer unit tests pass. Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

…in apply_cp Neither flag was needed by the models that exist: gemma4 routes via the model-owned setup_cp_attention branch, DSV4 via the TE DotProductAttention branch. needs_generic_cp_attention only gated a generic plain-SDPA MoE CP fallback that nothing on this branch triggers (main's apply_cp didn't support non-TE attention at all), and model_owned_cp_attention was log-only. Remove both flags and the post-loop generic-hook attach. A full/sliding block that is neither TE nor model-owned now logs a warning (fails loudly, like main) instead of silently relying on a never-exercised path. apply_cp's CP-attention routing is now just: TE -> context_parallel_group; setup_cp_attention -> model-owned; else -> warn. Tests updated to assert no generic hooks are attached and that model-owned routing calls setup_cp_attention. No behavior change for gemma4 (model-owned) or DSV4 (TE); 136 unit tests pass. Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

… (infra) The CP PR dropped main's is_compile_enabled gate so attach_cp_sdpa_hooks (the DTensor SDPA re-wrap) ran for all CP>1 -- that was for the original manual-slice gemma4 design (commit abd9d69). gemma4 has since moved to its own ring (setup_cp_attention) and no longer reaches this dense-CP block at all (MoE, skipped by the ep_size<=1 guard). The dense batch path is back to upstream context_parallel, where the SDPA re-wrap is only needed under torch.compile (Dynamo drops DTensor metadata through the compiled graph); without compile, context_parallel dispatches SDPA natively. Restore main's gate: only attach attach_cp_sdpa_hooks when compile is enabled. infra's dense-CP block now matches main except the ep_size<=1 guard (skip MoE, handled by apply_cp) and the _cp_enabled flag. gemma4/DSV4 unaffected (neither uses this block); 131 unit tests pass + 26B cp2 5-step regression matches the validated reference exactly. Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

apply_cp no longer routes generic attention through the dense hooks (that path was removed); it warns instead. Fix the ep_size-guard comment accordingly and note it is load-bearing: non-TE MoE (e.g. Gemma4) is not excluded by _uses_te_attention, so without the ep_size gate this pass would mask-strip Gemma4's vision tower and clobber its model-owned ring. Also trim the verbose attach_cp_sdpa_hooks compile-gate comment. Comment-only; no behavior change. Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

_make_contiguous_shard_cp_batch enumerated the batch-resident sequence keys three times -- the known_sequence_keys set, the per-key padding block, and the per-key slicing block -- so adding a field meant editing three places. Collapse the batch-resident keys into a single seq_pad_values table (key -> pad sentinel, all seq-dim 1) that drives padding, slicing, and the known-key set. position_ids (special pad), labels and loss_mask (carried as locals) stay handled explicitly. Behavior-preserving: same keys, seq dims, and pad values. 26B ep8 cp2 5-step is bit-identical to the reference (losses and grad_norms); 36 cp_utils unit tests pass. Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

… via callable cp_contiguous_shard was used only by Gemma4 (sole _cp_manual setter; DSV4 is TE and never touches it). Move it into the model: components/distributed/ cp_contiguous_shard.py -> components/models/gemma4_moe/cp_batch.py. To keep cp_utils model-agnostic (no distributed->model import) and avoid touching the recipe, dispatch is now a callback: Gemma4's prepare_model_inputs_for_cp attaches make_contiguous_shard_cp_batch_and_ctx to the batch as _cp_make_batch_fn, and make_cp_batch_and_ctx invokes whatever sharding callable the batch carries (replacing the _cp_manual flag + lazy import). vlm/finetune.py untouched. Gemma4 now fully owns its CP transport (ring), masking, and batch sharding; cp_utils carries only the generic callable seam. Behavior-preserving: 26B ep8 cp2 5-step bit-identical to reference; cp_utils unit tests pass. Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

Order the nemo_automodel imports alphabetically (distributed before models) in the cp_utils test files touched by the cp_batch move. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

HuiyingLi · 2026-06-13T17:52:52Z

/claude review

claude · 2026-06-13T17:56:34Z

@@ -0,0 +1,248 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.


Nit: copyright year is 2020 — looks like a copy-paste from an older file. Every other new file in this PR uses 2025.

Suggested change

# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.

# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

Addresses Claude review nit: the moved file inherited a 2020 header from the original cp_contiguous_shard.py; bump to 2025 to match the other new files in this PR. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

HuiyingLi · 2026-06-13T17:58:21Z

/claude review

HuiyingLi · 2026-06-13T17:58:22Z

/ok to test f736e09

claude · 2026-06-13T18:01:22Z

+def _force_repeat_kv_for_sdpa():
+    """Disable native SDPA GQA so the CP hook can own grouped-query handling."""
+    try:
+        from transformers.integrations import sdpa_attention
+    except ImportError:
+        yield
+        return
+
+    original_use_gqa = sdpa_attention.use_gqa_in_sdpa
+    sdpa_attention.use_gqa_in_sdpa = lambda _attention_mask, _key: False
+    try:


nit: _force_repeat_kv_for_sdpa is defined and tested but never called in production code — grep shows usage only in test_gemma4_model_cp.py. Is this intended for a follow-up, or was it meant to be wired into the CP forward path? If it's not needed yet, consider removing it to avoid confusion.

Addresses Claude review: this context manager (disable HF sdpa_attention GQA so a hook owns grouped-query handling) was never wired into production. Gemma4's ring CP path owns GQA itself via enable_gqa in its flex/SDPA call (cp_attention.py), so the manager is dead -- only its own swap/restore unit tests referenced it. Remove it, its two tests, and the now-unused contextlib / sys imports. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

HuiyingLi · 2026-06-13T18:13:35Z

Good catch — _force_repeat_kv_for_sdpa was dead (the ring CP path owns GQA via enable_gqa in its flex/SDPA call, never HF's sdpa_attention). Removed it, its two tests, and the now-unused imports in dadd6df.

/claude review

HuiyingLi · 2026-06-13T18:13:36Z

/ok to test dadd6df

claude · 2026-06-13T18:18:36Z

@@ -0,0 +1,619 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.


LGTM — thorough implementation with strong test coverage. The ring attention design (model-owned CP seam, p2p ring FlexAttention, online-softmax merge, custom autograd backward with gradient exchange) is clean and well-isolated from the generic CP infrastructure.

A few observations (all non-blocking):

ModelCapabilities: MoE variant already has supports_cp=True — no staleness issue. The new ep8cp2 recipe YAML validates it.

SDPA monkey-patching (lines 586/591): The _pre_hook / _post_hook pair that swaps F.scaled_dot_product_attention is safe because layers forward sequentially and the post-hook has always_call=True, but it's worth noting that this pattern wouldn't survive concurrent layer execution (e.g. if someone tried to parallelize layers within a pipeline stage). Not a concern for current usage.

torch.no_grad() removal in the recipe: Intentional and correct — enables gradient flow through trainable multimodal projectors during CP pre-embed. The test correctly verifies requires_grad propagation.

Double masked_fill for empty query rows: Applied both per-chunk (line 297) and after merge (line 368). The post-merge application is the authoritative one; the per-chunk one is defensive. Harmless redundancy.

hemildesai · 2026-06-11T23:58:56Z


        model_parts = model.parts if hasattr(model, "parts") else [model]
        for mp in model_parts:
+            mp._cp_enabled = True


This will set it for all models right?

@hemildesai deleted, thank you!

hemildesai · 2026-06-15T02:48:01Z

+    # full-sequence embeddings and null out local attention masks.
    # With CP the mask is not sharded along the sequence dim and TE asserts
    # "Padding mask not supported with context parallelism!".
+    model._cp_enabled = True


Where is this used?

this one is used by gemma4. the one below by qwen35. I have another PR on qwen35, will unify the two flags here. They are checks before vlm CP, if the placeholder img tokens are already pre-embd.

The dense/non-EP CP block in apply_model_infrastructure set mp._cp_enabled=True on every model part, but nothing reads it on that path: the only _cp_enabled readers are gemma4_moe and qwen3_5_moe, both MoE models that run ep>1 and get the flag from the MoE parallelizer's apply_cp instead (this block is gated ep_size<=1). No dense model reads it. Remove the dead set; keep the mask-strip / DTensor-SDPA hooks the dense CP path actually needs. Addresses review feedback. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

Remove the comment explaining the is_compile_enabled gate; it no longer adds clarity. No behavior change. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

khazic requested review from HuiyingLi, ZhiyuLi-Nvidia, adil-a, akoumpa, hemildesai and pthombre as code owners April 20, 2026 12:27

github-actions Bot added the community-request label Apr 20, 2026

khazic mentioned this pull request Apr 20, 2026

Does Gemma 4 support Context Parallelism (CP)? #1912

Closed

copy-pr-bot Bot temporarily deployed to test April 20, 2026 16:09 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci April 20, 2026 16:09 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci April 20, 2026 16:31 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci April 20, 2026 16:50 Failure

claude Bot reviewed Apr 20, 2026

View reviewed changes

HuiyingLi and others added 13 commits June 13, 2026 02:51

claude Bot reviewed Jun 13, 2026

View reviewed changes

hemildesai reviewed Jun 15, 2026

View reviewed changes

HuiyingLi and others added 2 commits June 14, 2026 22:19

hemildesai approved these changes Jun 15, 2026

View reviewed changes

sharonyu-115 mentioned this pull request Jun 30, 2026

Feature Request: Generic Automodel CP interface for downstream libraries e.g. NeMo RL #2861

Closed

HuiyingLi mentioned this pull request Jul 6, 2026

fix(cp): preserve gradients when sharding VLM inputs #2931

Merged

12 tasks

HuiyingLi mentioned this pull request Jul 14, 2026

fix(attention): constrain large-head FFPA Flex fallback #3064

Open

3 tasks

		@@ -0,0 +1,248 @@
		# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.

	# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
	# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

		@@ -0,0 +1,619 @@
		# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

Uh oh!

Conversation

khazic commented Apr 20, 2026 • edited by HuiyingLi Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Design

Changelog

Before your PR is "Ready for review"

Test Results

Targeted behavior

Uh oh!

copy-pr-bot Bot commented Apr 20, 2026

Uh oh!

khazic commented Apr 20, 2026

Uh oh!

khazic commented Apr 20, 2026

4-node smoke-test results (32x H100-80GB)

Loss parity: CP=1 vs CP=2 (matched per-GPU shard size)

Topology / infra notes

31B dense: not yet validated

Why all-gather + flex_attention instead of ring-attention

Known limitations (follow-up PR candidates)

Uh oh!

khazic commented Apr 20, 2026

Ring attention is the right long-term direction

VLM path: image and no-image both validated

Debug decisions worth flagging for the review

Uh oh!

HuiyingLi commented Apr 20, 2026

Uh oh!

HuiyingLi commented Apr 20, 2026

Uh oh!

claude Bot Apr 20, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Apr 20, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot commented Apr 20, 2026

Review

Uh oh!

HuiyingLi commented Apr 21, 2026

4-node smoke-test results (32x H100-80GB)

Loss parity: CP=1 vs CP=2 (matched per-GPU shard size)

Topology / infra notes

31B dense: not yet validated

Why all-gather + flex_attention instead of ring-attention

Known limitations (follow-up PR candidates)

Uh oh!

HuiyingLi commented Jun 13, 2026

Uh oh!

claude Bot Jun 13, 2026

Choose a reason for hiding this comment

Uh oh!

HuiyingLi commented Jun 13, 2026

Uh oh!

HuiyingLi commented Jun 13, 2026

Uh oh!

claude Bot Jun 13, 2026

Choose a reason for hiding this comment

Uh oh!

HuiyingLi commented Jun 13, 2026

Uh oh!

HuiyingLi commented Jun 13, 2026

Uh oh!

claude Bot Jun 13, 2026

Choose a reason for hiding this comment

Uh oh!

hemildesai Jun 11, 2026

Choose a reason for hiding this comment

Uh oh!

HuiyingLi Jun 15, 2026

Choose a reason for hiding this comment

Uh oh!

hemildesai Jun 15, 2026

Choose a reason for hiding this comment

Uh oh!

HuiyingLi Jun 15, 2026

Choose a reason for hiding this comment

Uh oh!

khazic commented Apr 20, 2026 •

edited by HuiyingLi

Loading