fix(cp): preserve gradients when sharding VLM inputs

[HuiyingLi]

What does this PR do ?

Fix context-parallel VLM finetuning when the full-sequence inputs_embeds tensor retains its autograd graph. This affects the shared CP path used by Nemotron Omni and Qwen3.5 after the VLM pre-embedding torch.no_grad() wrapper was removed in #1914.

Related Linear issue: AM-621

Root cause

VLM CP must build text and multimodal embeddings on the complete sequence before sharding. The shared recipe therefore calls the model's pre-embedding path and places the resulting [batch, sequence, hidden] inputs_embeds tensor in the CP batch.

The old torch.no_grad() wrapper made that tensor resizeable, but it also detached the token-embedding and optional multimodal gradient paths. PR #1914 intentionally removed the wrapper so trainable input embeddings and multimodal components receive gradients. Gemma4 uses its own model-owned _cp_make_batch_fn, but Nemotron Omni and Qwen3.5 continue through the shared PyTorch CP buffer path.

PyTorch's public context_parallel(..., buffers=...) implementation is the legacy path: it shards registered buffers in place with resize_/copy_. Resizing an autograd-tracked tensor is forbidden, so training fails before step 0 with:

RuntimeError: cannot resize variables that require grad

The failure is not specific to the latest PyTorch container; the same minimal reproducer fails on the tested PyTorch 2.11 and 2.12 builds.

Fix

• After existing CP padding, detect a gradient-bearing primary sequence tensor.
• Apply the same head-tail load-balancing order out of place with differentiable narrow() and torch.cat() operations.
• Put the local shard back in the batch and exclude only that tensor from PyTorch's mutable buffer list, avoiding both resize_ and double-sharding.
• Continue using the existing CP context for labels, positions, masks, and CP attention dispatch.
• Leave the input_ids and non-gradient buffer paths unchanged.

Backward now propagates from each local shard to the correct full-sequence positions and then to the input embedding or multimodal parameters.

PyTorch CP API direction

PyTorch has a newer internal CP design that avoids this mutation model:

• _context_parallel_shard returns new local shards out of place.
• _ContextParallel applies CP through module wrappers and DTensor dispatch rather than the legacy global SDPA monkey-patch path.
• TorchTitan's CP integration already consumes these internal helpers.

These APIs are still private/underscored and experimental, so migrating Automodel's entire CP attention path to them would be a broader compatibility change. This fix adopts the relevant out-of-place, autograd-preserving sharding behavior while retaining the current CP context for attention and ordinary buffers.

Changelog

• Add an autograd-preserving head-tail shard helper for CP sequence tensors.
• Bypass the legacy in-place buffer sharder for gradient-bearing inputs_embeds.
• Add regression coverage for rank-local ordering and backward gradients.

Test plan

• ruff check nemo_automodel/components/distributed/cp_utils.py tests/unit_tests/distributed/test_cp_utils_inputs_embeds.py
• ruff format --check nemo_automodel/components/distributed/cp_utils.py tests/unit_tests/distributed/test_cp_utils_inputs_embeds.py
• Exact-image CP unit suite: 43 passed
• Two-rank distributed regression: head-tail shards and input gradients match expected positions
• Exact failing environment, unpatched: reproduced on all eight ranks before step 0
• Exact failing environment, patched: Nemotron Omni EP8/CP2 completed 50 training steps and full validation (val_loss=2.0496)
• NeMo-CI: Nemotron Omni EP8/CP2 — 50/50 steps plus validation, passed
• NeMo-CI: Qwen3.6-35B EP8/CP2 — 50/50 steps plus validation, passed
• NeMo-CI: Qwen3.6-27B CP2, two nodes — cleared the CP failure and reached step 33/50; timed out at the 10-minute Slurm limit

Before your PR is "Ready for review"

Pre checks:

• Read and followed the contributor guidelines
• Added regression coverage
• Documentation changes are not required for this internal behavior fix

[copy-pr-bot]

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.

[HuiyingLi]

/claude review

[HuiyingLi]

/ok to test d741578

[akoumpa]

Hi @HuiyingLi, how have you confirmed per-token grad scaling when primary_seq_tensor has pad tokens? I see in line 452 cp_buffers[i] = torch.cat([buf, pad_val], dim=dim) so i'm expecting some ranks to receive pad values, how can we make sure that any gradient scaling takes this into account?

[akoumpa]

two cases I'm worried about (a) maintenance / regression and (b) correctness are:

• gradient scaling with padding + cp
• correctness with HSDP, because i saw it was using local rank -- I think it should be ok, but want to make sure.

[HuiyingLi]

Follow-up: padding and gradient-scaling correctness

The ordering in the training path is what makes the gradient scaling correct for the affected non-PP Qwen/Omni configurations:

1. Before CP padding/sharding, _run_train_optim_step() counts labels != -100 across accumulation microbatches and reduces that count across DP only, not CP. This is intentional because every CP rank initially owns the same full batch.
2. Qwen builds full-sequence inputs_embeds under autograd, including the text embedding and vision paths.
3. make_cp_batch_and_ctx() pads all sequence buffers consistently: zero inputs_embeds, labels=-100, zero position_ids, and padding_mask=True when present.
4. The PR shards the padded gradient-bearing inputs_embeds out of place in PyTorch's head-tail order. PyTorch's legacy CP context still shards labels/position IDs in the same order.
5. Each CP rank computes a summed CE over its non-ignored labels, divided by the global valid-label count N. Backward multiplies by DP*CP, and FSDP averages over that same mesh:

(1 / (DP*CP)) * sum_r[(DP*CP) * sum_{i in shard r}(grad_i) / N]
= sum_i(grad_i) / N

Thus CP-added -100 labels contribute neither to the numerator nor denominator, and ranks with zero valid labels do not bias the result.

Exact 200-padding scaling test

Commit 70526abc strengthens test_inputs_embeds_with_grad_and_cp_padding_preserves_global_token_mean:

• production padding branch: 56 valid rows -> 256 rows, exactly 200 appended rows
• CP=128 simulation, using the production head-tail sharder
• 72 CP ranks receive no valid labels
• local loss is scaled exactly like the recipe and gradients are averaged like FSDP
• resulting float64 loss and embedding gradient match the unpadded reference
• full test file: 24 passed

Real Qwen3.5 distributed check

Configuration: Qwen/Qwen3.5-4B, 8xH100, FSDP2, CP=2, DP=4, global batch=4, MTP disabled for isolation.

For the stress run, the identical 512-row batch was extended with exactly 200 zero embedding rows, 200 -100 labels, and zero position IDs (512 -> 712). After entering the real PyTorch CP context, all 200 labels remained ignored; head-tail sharding assigned 178 forced rows to CP rank 0 and 22 to CP rank 1.

• 512-row control: loss 2.1033506393432617
• 712-row PR path: loss 2.099325656890869
• 712-row detached PyTorch legacy-sharder reference: loss 2.099325656890869

The PR path and PyTorch legacy sharder therefore produce exactly the same forward loss even with 200 forced padding rows, establishing that the out-of-place sharding order/values match PyTorch while retaining autograd.

The 200-row Qwen case is deliberately artificial: CP=2 can naturally append at most 3 rows. Comparing 512 vs. 712 also showed Qwen's BF16 CP kernels have some sequence-length sensitivity (0.19% loss and 4.53% grad-norm difference), but this is identical across the PR and legacy sharding paths and is not introduced by this change. The earlier real 511 -> 512 run exercised Qwen's actual production padding branch.

[HuiyingLi]

/ok to test 70526ab