fix(cp): preserve gradients when sharding VLM inputs - #2931
Merged
Conversation
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
HuiyingLi
marked this pull request as ready for review
July 6, 2026 06:01
Contributor
Author
|
/claude review |
Contributor
Author
|
/ok to test d741578 |
HuiyingLi
added a commit
that referenced
this pull request
Jul 7, 2026
…tach Depends on #2931 (out-of-place sharding of grad-bearing inputs_embeds in the generic CP path); assumes it merges. With the resize_() constraint handled there, the remaining grad-blocking workarounds around CP pre-embedding are obsolete and harmful: - prepare_cp_forward loses pre_embed_no_grad. Its two users were the VLM eval site (redundant: _run_validation_epoch is already @torch.no_grad()) and the VLM KD student prep, where blocking gradients to trainable input embeddings and the vision tower is the same defect class #1914 removed from the train path. - minimax_m3_vl stops detaching its pre-embedded inputs_embeds — the detach existed only to survive context_parallel's in-place resize and silently froze the embeddings/vision tower under CP; this aligns it with qwen3_5/qwen3_5_moe/nemotron_omni. Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
akoumpa
reviewed
Jul 8, 2026
| # mutate only the integer/mask buffers. | ||
| primary_seq_tensor = cp_buffers[0] | ||
| if primary_seq_tensor.requires_grad: | ||
| batch[primary_key] = _shard_grad_buffer_for_cp(primary_seq_tensor, cp_seq_dims[0], cp_mesh) |
Contributor
There was a problem hiding this comment.
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?
Contributor
There was a problem hiding this comment.
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.
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Contributor
Author
Follow-up: padding and gradient-scaling correctnessThe ordering in the training path is what makes the gradient scaling correct for the affected non-PP Qwen/Omni configurations:
Thus CP-added Exact 200-padding scaling testCommit
Real Qwen3.5 distributed checkConfiguration: For the stress run, the identical 512-row batch was extended with exactly 200 zero embedding rows, 200
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 |
Contributor
Author
|
/ok to test 70526ab |
3 tasks
akoumpa
previously approved these changes
Jul 13, 2026
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Contributor
Author
|
/ok to test d22f77a |
hemildesai
approved these changes
Jul 14, 2026
akoumpa
approved these changes
Jul 14, 2026
akoumpa
added a commit
that referenced
this pull request
Jul 23, 2026
…#2937) * refactor(distributed): introduce CPSharder, retire private CP batch keys Models that own their CP batch sharding now return a CPSharder dataclass (under the 'cp_sharder' batch key) from prepare_model_inputs_for_cp, replacing the private batch-key side channel (_cp_make_batch_fn, _cp_metadata_seq_dims, _cp_metadata_pad_values, _cp_full_logits_grad_touch). - components/distributed/cp_sharder.py: CPSharder contract (shard_batch + local_token_global_indices required; token-tensor shard/gather synthesized from the indices; finalize_loss hook), the shared contiguous-shard batch implementation (merges the gemma4/dsv4 pad-table copies, parameterized by pad_multiple / extra_seq_keys / synthesize_packed_seq_ids), and full_logits_grad_touch (moved from the recipe loss site). - gemma4_moe: cp_batch.py delegates to the shared sharder; vision-group-id metadata moves from private batch keys to explicit sharder args. - deepseek_v4: deletes its duplicated pad/shard body, delegates to the shared implementation (THD guard + _dsv4_cp_group injection kept). - glm_moe_dsa: hook returns a CPSharder (packed_thd layout). - cp_utils.make_cp_batch_and_ctx dispatches on 'cp_sharder'; the legacy _cp_make_batch_fn batch key still works behind a DeprecationWarning. - llm/train_ft.py consumes the sharder's finalize_loss instead of the _cp_full_logits_grad_touch flag. Part of #2879. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(models): unify prepare_model_inputs_for_cp signature and invocation All eight models now share one hook signature: prepare_model_inputs_for_cp(batch: dict, *, num_chunks: int = 1) -> dict. Legacy per-key kwarg calls (input_ids=..., pixel_values=...) are repacked by normalize_prepare_cp_args behind a DeprecationWarning for one release. - Every model's forward(_pre_embed_only=True) interception builds the batch dict internally (no deprecation from internal calls); glm_moe_dsa gains the interception it was missing, so all model-owned CP models are reachable through __call__ (FSDP2 unshard hooks fire during pre-embed). - num_chunks is a real keyword parameter everywhere instead of being smuggled through **kwargs (previously read by only dsv4/glm). - ModelCapabilities gains cp_style ('none'|'pre_embed'|'model_owned') and cp_layout (diagnostic) so downstream libraries get a reliable capability signal instead of hasattr(model, 'prepare_model_inputs_for_cp'), which cannot distinguish pre-embed VLMs from models that own CP sharding. qwen3_5_moe declares cp_style='pre_embed' while keeping its conservative supports_cp=False gate (hook exists and is exercised by the ep8/cp2 recipe). - llm/train_ft.py passes the batch dict to the hook. Part of #2879. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(recipes): collapse CP dispatch into prepare_cp_forward (#2879) Adds cp_utils.prepare_cp_forward — a single CP dispatch (magi / model-owned CPSharder / TE-THD / generic torch context_parallel) returning (ctx, batch, cp_sharder) — and collapses the per-recipe branching into one call at every CP site: - llm/train_ft.py: the magi branch + model-owned hook + make_cp_batch_and_ctx block (the exact range pinned by #2879) becomes one prepare_cp_forward call; the returned sharder feeds finalize_loss. - vlm/finetune.py train + eval: the duplicated _cp_active/VLM_INPUT_KEYS pre-embed blocks move into the dispatcher (invoke_pre_embed / drop_mm_inputs / pre_embed_no_grad express the PP-stage and eval variants). - vlm/kd.py: same, with the teacher-compat check as an on_pre_embedded callback. - llm/kd.py (both sites): dispatch through prepare_cp_forward with invoke_pre_embed=False (KD has not wired model-owned CP). The pre-embed hook is now invoked uniformly through model.__call__(_pre_embed_only=True, ...) for LLM models too, so FSDP2 pre-forward hooks fire during pre-embed. Raw multimodal inputs are dropped only when the hook returns inputs_embeds; sharder-only hooks (DSV4/GLM) keep input_ids intact. Tests: recipe wiring tests retarget their make_cp_batch_and_ctx patches to cp_utils; cp test files gain an autouse no-dist fixture so a TP test's process group can no longer leak into fake-mesh rank resolution. Closes #2879. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): keep model-specific CP keys out of the shared sharder The shared contiguous sharder had absorbed Gemma4 specifics when its implementation moved out of gemma4_moe/cp_batch.py: the pad-sentinel table listed mm_token_type_ids / per_layer_inputs / _packed_seq_ids, and the _packed_seq_ids synthesis (a Gemma4 manual-CP-attention need) ran for every model behind a synthesize_packed_seq_ids flag. Model-specific logic belongs to models: the shared table now covers only the universal keys (input_ids, inputs_embeds, padding_mask; labels/position_ids/ loss_mask handled separately), and everything else arrives through the extra_seq_keys/extra_pad_values arguments. Gemma4's cp_batch.py owns its key bundle and the _packed_seq_ids synthesis again, running them before delegating to the shared sharder (the attention_mask->padding_mask conversion is exposed as a shared idempotent helper so the synthesis still sees padding_mask). DSV4's wrapper drops the now-removed flag; its only remaining quirk is the generic pad_multiple argument, whose compress-ratio derivation already lives in deepseek_v4/cp.py. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): move magi llm/vlm prep dispatch into MagiState prepare_cp_forward knew three magi internals: the two per-domain method names, their differing signatures, and the domain switch that existed only to pick between them. Backend specifics belong to the backend: MagiState gains a uniform prepare_batch(model, batch, *, device_mesh, domain, is_thd, pad_id, num_chunks) that owns the llm/vlm split, and the dispatcher's magi branch shrinks to a duck-typed call knowing only the (ctx, batch) contract. No behavior change; the branch remains selected by magi.enabled and the model-hook interaction rule (llm+magi skips pre-embed) is unchanged. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): dispatch magi at the TE rung of make_cp_batch_and_ctx magi was still special-cased in prepare_cp_forward: an early-return branch bypassing the prep chain, a domain parameter existing only to pick between two magi methods, and per-call threading of recipe-static arguments. Make magi behave like TE: everything recipe-static (domain, cp group, device mesh, HF-vs-custom) binds once at setup_magi, and MagiState exposes make_cp_batch(cp_mesh, batch, *, padding_token_id, num_chunks, is_thd, model=None) — the same shape and dispatch rung as make_cp_batch_for_te, returning (implicitly nullcontext,) the dispatched batch, active at cp<=1 like the TE prep (packing conversion / mask-spec activation). model is passed opaquely for magi's per-step key/spec stamping on attention modules (the HF attention interface cannot receive the key through kwargs; module attributes are the multi-model-safe channel per the direction of #2622). prepare_cp_forward loses the magi branch and the domain parameter; the llm-magi hook-skip rule now reads the bound magi.domain. MagiState. prepare_batch (added one commit ago, never released) is replaced by make_cp_batch. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): drop the CP full-logits grad touch Removes the zero-valued full-logits loss term (finalize_loss / full_logits_grad_touch, formerly the _cp_full_logits_grad_touch batch flag) from the DSV4/GLM model-owned CP paths, the CPSharder contract, and the llm recipe's loss site, matching its end-to-end removal in #2731. The CPSharder finalize_loss slot goes with it: with no remaining user it would be a speculative hook; it can return with a concrete consumer. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * test: fix CP patch points missed by the dispatch refactor Two test files sat outside the refactor's per-commit test sweeps and kept asserting pre-refactor internals: - test_finetune_vlm_helpers.py monkeypatched vlm.finetune.make_cp_batch_and_ctx, which the recipe no longer imports since the prepare_cp_forward collapse; retarget the 19 patch sites to cp_utils and widen the fakes for the dispatcher's extra arguments. - test_glm_moe_dsa_tilelang.py still asserted the retired _cp_make_batch_fn/_cp_full_logits_grad_touch batch keys; assert the CPSharder contract instead. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(models): drop the prepare_model_inputs_for_cp legacy-kwargs shim The normalize_prepare_cp_args deprecation shim protected callers of the old per-key form (input_ids=..., pixel_values=...), but no such callers exist: all recipes and forward interceptions already pass the batch dict, NeMo-RL main never invokes the hook, and its unmerged gemma4-cp draft is slated to move to the planned public CP interface (#2861). Keeping the shim only kept the signature loose. Remove the shim, tighten all eight hooks to prepare_model_inputs_for_cp(batch: dict, *, num_chunks: int = 1), and convert the remaining legacy-style callers (tests) to the dict form. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): drop pre-embed no_grad wrappers and minimax detach Depends on #2931 (out-of-place sharding of grad-bearing inputs_embeds in the generic CP path); assumes it merges. With the resize_() constraint handled there, the remaining grad-blocking workarounds around CP pre-embedding are obsolete and harmful: - prepare_cp_forward loses pre_embed_no_grad. Its two users were the VLM eval site (redundant: _run_validation_epoch is already @torch.no_grad()) and the VLM KD student prep, where blocking gradients to trainable input embeddings and the vision tower is the same defect class #1914 removed from the train path. - minimax_m3_vl stops detaching its pre-embedded inputs_embeds — the detach existed only to survive context_parallel's in-place resize and silently froze the embeddings/vision tower under CP; this aligns it with qwen3_5/qwen3_5_moe/nemotron_omni. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): pass the whole batch to the CP pre-embed hook The dispatcher filtered hook inputs through VLM_INPUT_KEYS — a central union-of-all-models registry of multimodal input keys that every new model had to extend, and the last piece of model knowledge living in cp_utils. Now the batch dict rides through model.__call__ as an opaque _cp_batch kwarg and the model reads the keys it needs; VLM_INPUT_KEYS is gone from the CP dispatch entirely (its one legitimate remaining use — dropping raw multimodal inputs on PP stages without embeddings — returns to the VLM recipe, a PP concern). Consumed-key removal uses a return channel, not in-place mutation: a hook returns None for every raw input it consumed (e.g. into inputs_embeds) and the dispatcher removes those keys from the batch. In-place pops looked simpler but break silently under FSDP2, whose forward-kwargs cast can hand the hook a rebuilt copy of the batch dict — caught by the gemma4-26B end-to-end run, not by unit tests, since test fakes are not FSDP-wrapped. Keys both consumed and re-emitted (gemma4's mm_token_type_ids) work naturally: the returned real value wins over the None marker. The eight forward interceptions collapse to one uniform line; per-model tests assert each hook's consumed-key markers. Verified: affected unit selection green (1314 passed; remaining failures are pre-existing fused-CE/tilelang environment issues, identical on the clean tree), and the gemma4-26B ep8+cp2 50-step run is bit-exact against the pre-refactor baseline. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): remove the deprecated _cp_make_batch_fn fallback The legacy private batch key was kept behind a DeprecationWarning for out-of-tree callers, but none exist: NeMo-RL main never attaches it, and the unmerged gemma4-cp draft — the only known user — will be rebuilt on the planned public CP interface (#2861). The cp_sharder dispatch is now the single model-owned CP entry point. Tests that exercised dispatch through the legacy key now construct a CPSharder directly; stale docstring references updated. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): pass CPSharder as an explicit parameter, not a batch key The sharder still transited between prepare_cp_forward and make_cp_batch_and_ctx inside the batch dict — a leftover of the retired function-pointer-in-a-dict plumbing that hid it from type checkers and debuggers. make_cp_batch_and_ctx now takes cp_sharder: CPSharder | None explicitly; the training batch stays pure tensors. The one remaining dict hop — the hook's return through model.__call__ — is inherent to the FSDP interception path and documented on the hook contract. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(models): drop the unused cp_style/cp_layout capability flags Nothing reads them: the CP dispatcher gates on hasattr(model, 'prepare_model_inputs_for_cp') and the runtime capability gate reads supports_cp. cp_layout also duplicated CPSharder.layout — two declarations of the same fact. They were groundwork for the public downstream CP interface (#2861); reintroduce them there together with their consumer (the plan's backend/layout fields) instead of carrying dead declarations. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): prune dead surface from the CP dispatch and sharder Final simplification sweep over the PR, applying the zero-consumer rule: - prepare_cp_forward returns (ctx, batch): the third element (cp_sharder) lost its only consumer when the grad-touch finalize_loss was removed. - on_pre_embedded callback removed: its single user — the VLM KD teacher-compat check — reads only the hidden dim, which sequence sharding never changes, so the recipe checks batch['inputs_embeds'] after the call instead of through a callback. - prepare_inputs_embeds_for_cp thin wrappers (gemma4_moe, nemotron_omni) deleted: no production callers. - CPSharder's shard/gather_token_tensor_fn override slots removed: no model fills them; the first real override (magi's undispatch) arrives with the public plan (#2861) and the slots return with it. The verb surface itself (local_token_global_indices + default shard/gather) stays: it is the model-provided layout fact the #2861 plan verbs (gather_token_tensor, zpqiu's shard_token_tensor, target alignment) are built from, and cannot be synthesized framework-side later without another cross-model contract change. - cp_utils: the two duplicated local mesh-size helpers merge into one module-level _mesh_dim_size. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): make every CP backend a CPSharder make_cp_batch_and_ctx now resolves a single CPSharder (model-owned > magi > TE > generic) and calls shard_batch — the per-backend branching collapses into _resolve_cp_sharder: - the generic torch context_parallel path becomes the framework's default sharder: shard_batch_load_balanced (body moved verbatim) with layout="round_robin" and closed-form round_robin_local_indices matching torch's 2*cp head-tail chunk pairing, so the token-tensor shard/gather verbs now work for the load-balanced layout too - magi and TE/THD become framework-built sharders wrapping their existing batch prep; their token layouts are data-dependent (cu_seqlens partitioning / dispatch solver), so local_token_global_indices is None and the token verbs fail loudly instead of sharding the wrong slice - resolution preserves the prior rung semantics exactly: model-owned and generic shard only at cp>1, magi/TE also run at cp<=1 (THD packing conversion / mask-spec activation) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): move shard_batch_load_balanced next to its layout math cp_sharder.py now hosts every pure-torch shard_batch implementation (contiguous + round-robin load-balanced) alongside their index maps, so a layout's pieces live in one file; cp_utils keeps dispatch, the TE prep, and the torch-CP transport machinery (create_context_parallel_ctx / get_train_context stay put — NeMo-RL imports them from cp_utils and tests patch them there; the moved function binds them at call time to avoid a module-level import cycle). Function body unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): drop unused mesh params from the contiguous CP prep _prepare_manual_cp_batch never read cp_mesh/tp_mesh; rename it _prepare_contiguous_cp_batch to match the model-owned-contiguous terminology ("manual" predates the CPSharder contract). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * feat(distributed): capture data-dependent CP index maps at shard time TE/THD and magi token layouts depend on batch content, so their sharders cannot provide local_token_global_indices as a pure function of (cp_mesh, seq_len) — but the partition is computed during the shard itself (tex.thd_get_partitioned_indices; magi's get_position_ids). Keep it: - make_cp_batch_for_te(return_local_indices=True) returns the partition it applied (identity arange when CP is inactive; None in chunked mode, where each chunk is its own token space) - MagiState.make_cp_batch(return_local_indices=True) returns the dispatch positions on the paths that dispatch (HF single-sequence, custom packed) - the framework THD/magi sharders install the captured map via the new captured_token_indices wrapper, which validates the requested stream length so a mismatched tensor cannot be silently mis-sharded; the token verbs work after the first shard_batch and raise before it Framework sharders are built per resolution, so a capture never leaks across steps. The THD partition itself is unchanged (computed once from input_ids instead of per key — all token keys share the same stream). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * feat(distributed): return the resolved CPSharder from the CP dispatch prepare_cp_forward and make_cp_batch_and_ctx now return (ctx, batch, sharder) so callers — in particular downstream libraries (#2861) — hold the resolved sharder and can keep per-token tensors aligned with the sharded inputs via its token verbs. When no CP prep applies the dispatch returns a layout="none" identity sharder (passthrough shard_batch, arange index map) instead of nothing, so consumer code is branch-free across cp_size 1 and N. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(distributed): preserve CP preparation contracts Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * Revert "fix(distributed): preserve CP preparation contracts" This reverts commit 3985a7c23e5794d041680a581c8e3121af15d657. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * feat(distributed): capture pad facts at shard time; token verbs take caller coordinates Padding is an internal detail of the CP layout, but its facts (how long the padded stream is, where each input token went) were produced inside shard_batch and thrown away - so a consumer co-sharding advantages/masks or gathering token logprobs had to reproduce them, and on the contiguous layouts a custom pad_multiple opened a silent-misalignment window (a plausible length passed the divisibility check but did not match the sharded stream). Capture them on the sharder instead, same mechanism as the captured THD indices: - original_seq_len / padded_seq_len: measured by the resolver closures (none, round-robin) and by shard_batch_contiguous via a new record_on parameter that the model hooks pass (gemma4/dsv4/glm construct their sharder first, then bind shard_batch with record_on=sharder) - flat-stream (THD) layouts: the BSHD->THD flatten is a pure reshape, so the TE/GLM shards capture the pre-flatten input_row_shape and the verbs translate between row and stream coordinates - DSV4 packed repad genuinely repositions tokens, so it now also returns the input->rebuilt-row position map (-1 = dropped input pad slot) and the sharder captures it as input_token_stream_positions Verb behavior (field presence only - never branched on layout): - shard_token_tensor(t, fill=...): accepts original-length tensors (right-padded with the explicit fill), input-row tensors (flattened), input-coordinate tensors on repositioned layouts (scattered via the map), or already-padded tensors; any other length raises - gather_token_tensor(t, trim=True, fill=...): validates the gathered length against the captured facts, then restores the caller's original coordinates (slice / un-flatten / map back with fill for dropped slots) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * feat(distributed): capture magi dispatch facts for the token verbs Completes the shard-facts capture for the magi backend, in the resolver closure (magi internals untouched): the HF single-sequence path pads at the tail of the global order, so original/padded lengths are captured and trim restores the caller's [1, S]; the packed path over a pure THD flatten captures the pre-flatten row shape when the dispatch added no pad (padded == rows x cols). With this, every backend answers the token verbs in the caller's coordinates, which also absorbs the previously planned extra_token_keys batch channel — the captured facts let shard_token_tensor reproduce the packed transform post-hoc. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): shrink the CP dispatch surface Four prunes, each with the audit that motivated it: - one public entry point: make_cp_batch_and_ctx becomes private (_make_cp_batch_and_ctx, kept as the internal dispatch seam); its four remaining production callers (dllm x3, minimax functional driver) move to prepare_cp_forward(None, ...) - identical behavior, the hook gate short-circuits without a model. The public CP story is now one function in, one hook per model, one sharder out. - drop the cp_size override on prepare_cp_forward: introduced by #2590 as a recipe-side convenience read of the config; the mesh is the runtime truth and is always available at the call site. The hook gate now reads the mesh only (the hook-gate unit test fakes a mesh instead of a config). - drop CPSharder.layout: its only consumers were error messages and test assertions - by this PR's own zero-consumer rule it goes. Tests assert behavior (which index map / which shard_batch) instead of a tag; the no-branching rule strengthens to "the sharder carries no backend tag". - merge the duplicated mesh probes: _get_submesh becomes the module-level helper and _mesh_dim_size delegates to it. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * test(distributed): multi-rank functional test for the CPSharder token verbs The unit suite only reaches gather_token_tensor's identity early-returns; this torchrun driver (2+ GPUs) runs the real collectives: fill-shard -> differentiable gather -> trim round-trips a caller-coordinate tensor through the round-robin layout, backward routes each global position's gradient to the owning rank in local head-tail order, and the sum-over-consumers semantics of the differentiable all-gather (local grad = world_size x the single-loss grad for a replicated full-sequence loss) is pinned down explicitly - the factor replicated-loss CP trainers must normalize for. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): rename CPSharder to ContextParallelismSharder Review ask (akoumpa): spell the class name out. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): report shard facts as a return value, install centrally The shard-time facts (captured indices, lengths, row shape, position map) were installed by mutating the sharder from three different places: resolver closures assigning fields on outer variables, a record_on side-channel parameter threaded through the contiguous impl, and two-step hook constructions (build the sharder, then partial-bind it into its own shard_batch). Replace all three with one seam: - shard_batch's contract becomes (cp_mesh, tp_mesh, batch, ...) -> (ctx, batch, ShardFacts | None); every implementation reports what it learned as a small frozen dataclass - the dispatch is the single installation point (sharder.install_shard_facts) — the temporal rule "token verbs raise before the first shard" now lives in one place - record_on is gone, hooks construct their sharder in one expression, and the round-robin/identity resolver closures collapse into plain constructions over module-level impls (shard_batch_identity returns) No behavior change: batch tensor math is untouched; unit suites and the multi-rank token-verb functional test pass unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): store the shard layout whole; rename ShardFacts to ShardLayout Follow-up to the facts-return contract: the sharder mirrored the report's four fields and installed them one by one — the same data twice plus a transcription method. Store the report as a single field instead (sharder.shard_layout, plain assignment in the dispatch), read it uniformly in the verbs via an all-None placeholder, and move the reported-partition length validation to the read side (_indices) — which deletes install_shard_facts and the captured_token_indices wrapper outright. Net -24 lines. ShardLayout (nee ShardFacts) now matches the ecosystem's naming for layout metadata riding alongside a batch. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(distributed): inline the contiguous shard helpers; drop dead code - _prepare_contiguous_cp_batch and _make_contiguous_shard_cp_batch each had exactly one production consumer (shard_batch_contiguous itself); the split existed to ferry six locals through a tuple and eleven keywords through a signature. Inline both: one linear function (normalize -> pad -> slice), both ferry signatures gone. The diffcov tests now exercise the same branches through the public function. - the inline also retires the legacy _cp_metadata_seq_dims / _cp_metadata_pad_values batch keys - their only remaining reader was a test; production passes extra_seq_keys/extra_pad_values explicitly. Private CP batch keys are now fully gone. - rename the ShardLayout-holding locals from 'facts' to 'layout' to match the class rename. - delete _build_position_ids (cp_utils): zero production consumers, only its own two tests kept it alive. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * chore(distributed): fix stale wording in the no-layout placeholder comment Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(cp): allow kwargs-only pre-embed call on DSV4/GLM forwards prepare_cp_forward invokes hook-aware models through __call__ with only keyword arguments (_pre_embed_only=True, _cp_batch=batch, num_chunks=n) so FSDP2 pre-forward hooks run. DeepseekV4ForCausalLM.forward and GlmMoeDsaForCausalLM.forward declared input_ids as a required positional, so the call died at signature binding (TypeError: missing 'input_ids') before the _pre_embed_only branch could run - caught by the 16-node deepseek_v4_flash_cp_tulu3 nemo-ci parity run under pp4/cp8. Default input_ids to None, matching the other six pre-embed models (step3p7 has the identical MTP star-args shape with the None default). Adds a signature-binding contract test over all eight pre-embed models. Also registers ci nodes/time for glm_5.2_tulu3_32k_tilelang_cp8: cp8 x pp4 x ep64 needs 256 GPUs, but with no ci block the recipe generates a 1-node CI job that always fails the world-size check (sibling glm_5.2_hellaswag_pp already registers nodes: 32). Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(cp): restore minimax pre-embed detach - PP microbatches share the graph 545dc204 dropped the detach as part of unifying the grad-carrying pre-embed contract (the original resize_()-rejects-grad reason was obsoleted by the out-of-place grad shard). But minimax_m3_vl_sft_cp2_medpix_2k runs cp2 x pp4: prepare_cp_forward pre-embeds the WHOLE batch once, the PP schedule then splits it into microbatches, and every microbatch backward traverses the same vision/embedding graph - the second one raises 'Trying to backward through the graph a second time' (nemo-ci job 367479592; the pre-refactor baseline passes 20/20 steps on the same stack). The other pre-embed VLMs (qwen3_5/qwen3_5_moe/nemotron_omni) train at pp1 where a shared graph is backwarded exactly once, which is why the parity matrix only caught this on minimax. The embeddings/vision tower are frozen for this recipe family, so the detach loses no gradient; vision training under CP+PP would need per-microbatch pre-embedding and is out of scope here. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(cp): fail loud when PP receives grad-carrying pre-embed inputs_embeds Grad-carrying inputs_embeds from the CP pre-embed hook cannot work under pipeline parallelism: the batch is pre-embedded once, the PP schedule splits it into microbatches, and each microbatch backward traverses the same shared pre-embed graph - the second one raises the cryptic 'Trying to backward through the graph a second time'. minimax detaches again since 7d9e4bc7; the other pre-embed VLMs (qwen3_5, qwen3_5_moe, nemotron_omni, gemma4, step3p7) do not detach but currently have no pp>1 CP recipes. This guard turns the future foot-gun into an actionable error at the dispatch boundary. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * Revert "fix(cp): fail loud when PP receives grad-carrying pre-embed inputs_embeds" This reverts commit 37a10bb355a658a30f659d1e4904d9b7e492a589. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * feat(cp): aux-only shard_batch + in-forward sequence shard helper Add shard_batch_aux_only and shard_sequence_for_cp to the CP sharder, the framework-owned pieces for Megatron-style per-microbatch CP: a model embeds and sequence-shards its own primary stream inside forward while the no-grad aux streams (labels/position_ids/loss_mask/padding_mask) ride the same round-robin context_parallel context. shard_batch_aux_only mirrors shard_batch_load_balanced but excludes the primary stream from the CP buffer list, leaving input_ids/inputs_embeds full-length in the batch. shard_sequence_for_cp pads and round-robin index_selects a full-length tensor inside a forward (differentiable, so gradients reach the embeddings/vision tower). Shared padding/normalization helpers are extracted so shard_batch_load_balanced stays behavior-identical. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(cp): sink minimax_m3_vl pre-embed into forward (per-microbatch CP shard) MiniMax M3 VL's prepare_model_inputs_for_cp becomes sharder-only: it returns a ContextParallelismSharder backed by shard_batch_aux_only (round-robin shards only labels/position_ids/loss_mask/padding_mask and installs the ring-SDPA context) and consumes nothing, leaving input_ids and the multimodal inputs full-length in the batch. The forward now embeds + splices vision on the full sequence and shards the result with shard_sequence_for_cp per microbatch, so the embeddings and vision tower are trainable under CP and the removed detach (the PP×CP shared pre-embed double-backward workaround) is no longer needed: each microbatch owns its own graph. The MoE parallelizer's apply_cp hands the CP submesh to the model (model.cp_mesh) so the forward can build this rank's shard. get_pipeline_stage_metas reports the asymmetric stage-0 layout (full-length token-id input, local sharded-length outputs) required by text CP×PP. Image/video microbatch chunking under CP×PP is out of scope and raises NotImplementedError; text-only CP×PP is the supported combined topology. pp=1 numerics are equivalent by construction: shard_sequence_for_cp selects the same round-robin token positions with the same zero padding as the old dispatch-level load-balanced shard. Verified: tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_vlm.py (36 passed with test_cp_sharder.py in the auto2606rc9 container). CP×PP numerics require multi-GPU validation (cp2/pp4 nemo-ci). Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(cp): sink qwen3_5 pre-embed into forward (per-microbatch CP shard) Qwen3.5 dense VLM's prepare_model_inputs_for_cp becomes sharder-only: it returns a ContextParallelismSharder backed by shard_batch_aux_only plus the full-sequence mRoPE position_ids (computed by get_rope_index), which the aux shard slices on the mRoPE axis. Embedding and the image/video multimodal splice move into forward via the new _embed_and_splice_for_cp helper; when CP is active the first stage embeds the full sequence, keeps this rank's round-robin chunk pair via shard_sequence_for_cp, and feeds self.model exactly what the old dispatch-level pre-embed did (inputs_embeds sharded, input_ids None) -- so the decoder, MTP, and loss are bit-identical by construction at pp=1. Embeddings and the vision tower are now trainable under CP. Adds cp_mesh (installed by the parallelizer) and a get_pipeline_stage_metas that reports the asymmetric stage-0 layout under CP (full token-id input, local sharded outputs) and reduces to the framework-default symmetric shapes at cp_size==1, so PP-without-CP (e.g. tp4pp4) is unchanged. Image/video microbatch chunking under CP*PP raises NotImplementedError; text-only CP*PP is supported. Verified: tests/unit_tests/models/qwen3_5/ (73 passed) in the auto2606rc9 container. CP and CP*PP numerics need multi-GPU validation. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(cp): sink qwen3_5_moe pre-embed into forward (per-microbatch CP shard) Mirrors the qwen3_5 dense migration for the MoE VLM: prepare_model_inputs_for_cp becomes sharder-only (ContextParallelismSharder + full mRoPE position_ids for the aux shard), embedding and the image/video splice move into forward via _embed_and_splice_for_cp, and the first stage under CP embeds the full sequence and round-robin shards it (shard_sequence_for_cp) before feeding self.model the same inputs_embeds/input_ids=None the old dispatch-level pre-embed produced -- bit-identical downstream (decoder, MTP, lm_head) at pp=1. The CP-aware GatedDeltaNet linear attention and ring SDPA consume the same round-robin layout as before. Flips the stale supports_cp capability False->True (CP is a real config: qwen3_6_35b_medpix_ep8cp2_4k.yaml), adds cp_mesh (installed by Qwen3_5ParallelizationStrategy) and a CP-aware get_pipeline_stage_metas (default-symmetric at cp_size==1). Image/video microbatch chunking under CP*PP raises NotImplementedError. Verified: tests/unit_tests/models/qwen3_5_moe/ (161 passed) in the auto2606rc9 container. CP and CP*PP numerics need multi-GPU validation. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(cp): sink step3p7 pre-embed into forward (per-microbatch CP shard) Step3.7 VLM's prepare_model_inputs_for_cp becomes sharder-only (returns a ContextParallelismSharder backed by shard_batch_aux_only, consumes nothing; Step3.7 uses plain 1-D positions so the aux shard injects/slices them). The forward embeds via get_multimodal_embeddings + prepare_inputs_embeds on the full sequence and round-robin shards the result (shard_sequence_for_cp) before the backbone, MTP-source-embed build, and lm_head -- feeding the same inputs_embeds/input_ids=None the old dispatch-level pre-embed produced, so downstream (incl. MTP embeds rolled from inputs_embeds) is bit-identical at pp=1. get_pipeline_stage_metas is extended (not replaced) to report local (padded/cp) sequence lengths for the sharded stage outputs and the propagated MTP hidden states, keeping the full-length token-id input on the first stage; cp_size==1 preserves the existing symmetric MTP metas. Image chunking under CP*PP raises NotImplementedError. cp_mesh comes from the MoE parallelizer's apply_cp. Note: step3p7 has no CP example config today (all cp_size=1), so its CP path is unexercised; this preserves the prior round-robin pre-embed mechanism by construction. Verified: tests/unit_tests/models/step3p7/ (46 passed) in the auto2606rc9 container; CP/CP*PP numerics need multi-GPU validation. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(cp): sink nemotron_omni pre-embed into forward (per-microbatch CP shard) NemotronOmni's prepare_model_inputs_for_cp becomes sharder-only (returns a ContextParallelismSharder backed by shard_batch_aux_only, consumes nothing; plain 1-D positions are injected/sharded by the aux shard). The forward already contained the image/video/audio splice (the inputs_embeds is None block); it now round-robin shards the spliced embeddings (shard_sequence_for_cp) before the LM, feeding the same sharded inputs_embeds the old dispatch-level pre-embed produced -- bit-identical downstream at pp=1 (NemotronOmni is CP-only, supports_pp=False). Embeddings and the vision/audio encoders are now trainable under CP. cp_mesh comes from the MoE parallelizer's apply_cp. The multimodal-splice unit tests move from asserting the hook's return to asserting the inputs_embeds the forward hands the LM (captured via a stub LM), preserving behavioral coverage of every modality; adds a CP-shard shape check. Verified: tests/unit_tests/models/nemotron_omni/ (45 passed) in the auto2606rc9 container. CP numerics need multi-GPU validation. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * test(cp): GPU forward-equivalence for the minimax VL in-forward CP shard Multi-rank (2 GPU) L1 driver that exercises the sunk CP path end to end: prepare_cp_forward runs the sharder-only prepare_model_inputs_for_cp hook (shard_batch_aux_only + ring-SDPA context) and the forward embeds + round-robin shards inputs_embeds via shard_sequence_for_cp. Asserts the unsharded cp2 logits match the cp1 eager forward. Verified on 2x H100 (auto2606rc9 container): top1=1.0000, loss eager 5.555040 vs cp 5.555128 (|dloss|=8.7e-5) -- numerically confirms the pilot migration (commit 7fa36586) that CPU unit tests could not. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * test(cp): cp2xpp2 layer-2 verification for the in-forward pre-embed sink Config-swappable (NEMO_CP_PP_MODEL={minimax,step3p7}) driver over the real AutoPipeline split + schedule.step under cp2xpp2 (4 GPU) and cp2xpp1 (2 GPU) with tiny random-init text-only configs. Exercises the full sunk layer-2 contract: sharder-only hook, in-forward embed + shard_sequence_for_cp, asymmetric get_pipeline_stage_metas, per-microbatch backward, and (step3p7) the MTP stage metas + per-depth MTP loss threaded through the schedule. Asserts 20 clean steps (no double-backward), finite loss, embed_tokens gradients, and MTP-used for step3p7. Verified on H100 (auto2606rc9), bf16, cp2xpp2 vs cp2xpp1: minimax 5.588573 vs 5.588573 (mtp_used=False) step3p7 7.875539 vs 7.875540 (mtp_used=True) Exact-in-bf16 layout invariance across the PP split, embeddings trainable, and the MTP-under-PPxCP path (step3p7's one structural delta) confirmed clean. This closes the layer-2 gap the cp-only mechanism proof could not: per-microbatch in-forward shard eliminates the shared-pre-embed double-backward. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(cp): support minimax images under CP×PP via the media side channel The sunk minimax CP path rejected images under cp>1 && pp>1 with a NotImplementedError -- a regression: the pre-sink branch ran minimax_m3_vl_sft_cp2_medpix_2k cp2×pp4 (nemo-ci 367493325). The repo already has the per-microbatch pixel side channel (prepare_vlm_media_for_pp -> stage_vlm_media_for_pp -> stage-0 _vlm_*_chunks), and minimax's forward already pulls those chunks; the guard just blocked it. Remove the guard so the existing chunk pull -> embed+splice -> shard_sequence_for_cp path runs per microbatch. Media rides the side channel, not the stage tensor stream, so get_pipeline_stage_metas is unchanged (first-stage input stays input_ids [mb, S]). The deeper blocker: with embed+splice sunk into forward, the vision tower's bidirectional patch attention now runs inside the ring-SDPA context_parallel context, where torch's load-balanced ring all-gathers Q/K/V and rejects the non-causal attention ("Load balancing requires is_causal=True"). Add cp_dispatcher_suspended(cp_mesh) in cp_utils, which suspends the legacy context_parallel SDPA monkeypatch around the vision forward and restores it for the sharded text decoder; wrap minimax's vision_tower call with it. Verified on 4x H100 (auto2606rc9), bf16, 2 images of different sizes across 2 samples: cp2×pp2 loss 5.609686 == cp2×pp1 5.609686 (exact), 20 clean steps (no double-backward), embed AND vision-tower gradients finite -- vision is trainable under CP now. Committed run_cp_pp_image_sink.py. Text-only cp2×pp2 (5.588573) and the minimax CPU unit suite (42 passed) are unregressed. Note: the vision-in-ring issue is shared by the other sunk VLMs (qwen3_5, qwen3_5_moe ep8cp2 medpix, nemotron_omni); they need the same one-line cp_dispatcher_suspended wrap on their vision/audio splice -- follow-up, helper is in place. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(cp): suspend the CP ring around vision/audio in qwen3_5, qwen3_5_moe, nemotron_omni Completes the sink-introduced image×CP fix across the migrated VLMs. Like minimax (558afcde), these models splice vision/audio in-forward, so their encoders' bidirectional attention runs inside the ring-SDPA context_parallel context when CP is active and torch's load-balanced ring rejects/mis-gathers the non-causal attention. Wrap each encoder call in cp_dispatcher_suspended(self.cp_mesh): - qwen3_5 / qwen3_5_moe: the get_image_features / get_video_features block in _embed_and_splice_for_cp (qwen3_5_moe ep8cp2 medpix is a shipped image config); - nemotron_omni: extract_feature_dynamic / extract_feature / extract_video_feature / extract_sound_feature in the forward splice. Verified: run_cp_dispatcher_suspend.py (2 GPU) directly confirms the helper contract -- a non-causal SDPA inside a real ring context fails/gathers without the suspend, runs as a plain local SDPA with it, and the ring is restored for the sharded causal decoder afterward (without_suspend_failed=True, suspend_ok=True, ring_restored_after=True). The helper is also proven end-to-end by the minimax image cp2×pp2 run (558afcde, vision grads finite). Per-model wrap placement is covered by the qwen3_5/qwen3_5_moe/nemotron_omni CP unit suites (42 passed). Note: a tiny-config end-to-end GPU image run for qwen3_5_moe was not built (its HF visual-tower tiny config is heavy); the medpix ep8cp2 nemo-ci recipe is the end-to-end confirmation, and the helper + minimax proofs cover the mechanism. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(cp): stage VLM media for PP when CP pre-embed is sunk into forward The sunk VLMs (minimax_m3_vl, qwen3_5, qwen3_5_moe, step3p7) now embed and sequence-shard per microbatch inside forward and pull their media from the PP side channel (_vlm_pixel_values_chunks), populated by stage_vlm_media_for_pp. But the VLM recipe gated media staging off (pp_n_microbatches=None) whenever a model exposed prepare_model_inputs_for_cp under CP -- an assumption that only holds for recipe-level pre-embedders (gemma4), which emit inputs_embeds before the schedule so raw media never rides schedule.step. For the sunk models the sharder-only hook leaves raw pixel_values/image_grid_thw in the batch, so with staging disabled torch pipelining row-chunks pixel_values and image_grid_thw independently. The per-microbatch patch slice then disagrees with the grid metadata the vision RoPE uses, crashing in _apply_vision_rope (e.g. "size of tensor a (156) must match tensor b (160)" on the minimax_m3_vl_sft_tulu3_text_cp8_16k cp8xpp4 run, where fake-image injection gives every packed text sample a placeholder image). Fix: mark the sunk models with cp_preembed_in_forward=True and stage VLM media for PP whenever a model is not a recipe-level CP pre-embedder, so torch never sees raw media to row-chunk and each microbatch pulls a grid-aligned chunk via chunk_vlm_media. gemma4 (no flag) keeps its recipe-level pre-embed path. Adds a regression case to test_setup_skips_pp_media_prechunk_when_cp_preembeds_vlm_inputs covering a sunk model under CP (expects staging on). Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(cp): invoke the sharder-only hook on every PP stage for sunk models The sunk VLMs expose a sharder-only CP hook (embed + sequence-shard happen in their forward). The VLM recipe only invoked that hook on the first PP stage (invoke_pre_embed = has_first_stage). On non-first stages the hook was skipped, so prepare_cp_forward fell through to the generic round-robin sharder, which shards the primary stream (input_ids) to the local length. The recipe then fed that already-local length to pp.update_seq_len, and the CP-aware get_pipeline_stage_metas divided it by cp a SECOND time -> non-first-stage recv buffers became S/cp², so the inter-stage P2P truncated the incoming hidden to S/cp². A later sparse-DSA attention layer then hit a RoPE size mismatch (freqs built from the correctly S/cp-sharded position_ids vs the S/cp² hidden): "The size of tensor a (512) must match the size of tensor b (1024)" on the after-side cp8×pp4 / cp2×pp4 CI runs (jobs 367636608/367636609). Fix: invoke the sharder-only hook on ALL PP stages when the model is sunk (cp_preembed_in_forward=True). The hook does no compute and consumes nothing, so it installs the aux-only sharder on every stage; input_ids stays full-length, update_seq_len sees the FULL seq_len on every rank, and the stage metas divide by cp exactly once. Recipe-level pre-embedders (gemma4) keep their first-stage-only gate, and the non-first-stage media drop is decoupled onto the unchanged is_first_or_no_pp condition so their behavior is identical. Reproduced and fixed on GPU (tiny minimax with sparse DSA, cp2×pp2, replicating the recipe's per-stage invoke_pre_embed gate): before, non-first stages crash with h=(1,64) vs freqs=(1,128); after, all stages shard once (256->128) and the step passes. Adds regression tests pinning that a sunk model invokes its hook on non-first PP stages (keeping input_ids full) while a recipe-level pre-embedder does not. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * debug(cp): TEMP NEMO_CP_DEBUG shape logging in minimax forward (REVERT before merge) Env-gated (NEMO_CP_DEBUG=1) shape logging that localized the text-decoder RoPE double-shard (jobs 367636608/367636609): CondGen.forward enter/post-embed/ post-shard (rank0), CausalLM.forward layer-loop h/freqs/position (rank0 or on h≠freqs), and cp_sparse_attn._cp_forward x/freqs (cp_rank0 or on x≠freqs), with self-flagging "*** MISMATCH ***" lines so it captures the crashing (non-first) rank. Root cause fixed in the parent commit; keep this separate and revert it before merge. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * feat(cp): contiguous aux-only shard + in-forward contiguous slice Add the contiguous-CP counterparts of the round-robin per-microbatch sink (shard_batch_aux_only / shard_sequence_for_cp): shard_batch_aux_only_contiguous shards only the no-grad aux streams (labels/position_ids/loss_mask/padding_mask plus model-provided extra_seq_keys) and leaves the primary stream full-length, and slice_sequence_for_cp_contiguous pads and keeps this rank's contiguous seq_start:seq_end slice of a full-length tensor inside a model forward (differentiable, so gradients reach the embeddings/vision tower). shard_batch_contiguous and shard_batch_aux_only_contiguous now share one core (_shard_batch_contiguous_impl) gated on a single shard_primary switch, so the aux-only slice is bit-for-bit equivalent to the dispatch-level contiguous shard of the primary. Unit tests pin that equivalence (aux-only + in-forward slice == shard_batch_contiguous), the full-length primary, differentiability, the cp_size*max(pad_multiple,2) divisor, and the cp<=1 identity. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(cp): sink gemma4 pre-embed into forward (per-microbatch contiguous CP shard) Migrate Gemma4 (dense E-series/31B + MoE 26B) from the recipe-level CP pre-embedder to the Megatron-style sunk pattern: prepare_model_inputs_for_cp is now sharder-only (cp_preembed_in_forward=True) and consumes nothing; the model's own forward embeds, splices vision, builds per_layer_inputs and the flex-ring mask metadata on the full microbatch sequence, then keeps this rank's contiguous slice (_cp_sunk_prepare_inputs -> slice_sequence_for_cp_contiguous). This makes the embeddings and vision tower trainable under CP and removes the PP×CP shared pre-embed graph. The sharder-only hook returns a ContextParallelismSharder whose shard_batch (_cp_shard_batch_aux_only -> make_contiguous_aux_only_shard_cp_batch_and_ctx) shards only the no-grad aux streams (labels/position_ids/loss_mask/padding_mask) plus the synthesized _packed_seq_ids document map, and leaves input_ids / pixel_values / mm_token_type_ids full-length for the forward. It also records cp_mesh and installs the p2p flex ring (the one model-owned seam that reliably hands Gemma4 the CP submesh, since dense variants are not guaranteed to run the MoE apply_cp). _packed_seq_ids is sharded here rather than in the forward because its pad-region zeros depend on the global pad tail; every other ring metadata stream is a per-token or cumsum-over-full-then-slice quantity the forward slices to the identical contiguous layout. The flex-ring mask inputs reach the ring exactly as before -- only the shard call-site moved from the dispatch into the forward. Both dense (metadata stashed on the ring-hooked self_attn as _cp_dense_metadata) and MoE (metadata threaded through the backend kwargs) paths embed+splice+slice in-forward; the old "CP + pixel_values requires pre-computed inputs_embeds" NotImplementedError is removed (pixels are now spliced in-forward). Recipes exercising per_layer_inputs are the E-series (E2B/E4B, hidden_size_per_layer_input); 26B MoE and 31B dense have none. Gemma4 uses its own flex ring (not a torch-CP SDPA patch), so no cp_dispatcher_suspended is needed. gemma4 has no PP wiring in-forward, so this is CP-only; per_layer_inputs never coexists with PP (E-series is supports_pp=False). Correctness rests on a structural bit-exactness argument, unit-verified: the in-forward contiguous slice (slice_sequence_for_cp_contiguous) is proven bit- identical to the previous dispatch-level shard_batch_contiguous slice (same positions, pad sentinels, and cp_size*2 divisor), and the forward reuses the identical embed / vision-splice / per_layer / gemma4_vision_group_ids functions, so the per-rank tensors reaching the ring are unchanged. A 2-GPU cp2 harness (run_gemma4_vl_cp_sink.py) drives the sunk path end to end (E-series per_layer + vision-bidirectional mask) and confirms it runs without desync with the cp2 loss matching the cp1 eager forward to 4.2e-2 (bf16 flex ring vs HF eager SDPA). Real-recipe bit-exact validation runs on nemo-ci. Unit tests updated to the sharder-only contract (hook returns only cp_sharder, consumes nothing; forward owns the embed/splice/per_layer slice); gemma4 CP unit suites green (54 passed) and the broad gemma4 model suite green (343 passed). Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(cp): keep the contiguous shard divisibility guard on the padded primary The shared _shard_batch_contiguous_impl derived padded_seq_len from seq_len + pad_len, which is always divisible by cp_size, so the post-pad divisibility ValueError became unreachable. Read the primary tensor's actual length for the primary-inclusive shard (restoring the guard and the previous shard_batch_contiguous behavior), and keep the intended padded length only for the aux-only path where the primary is deliberately left full-length. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(cp): drop the _pre_embed_only __call__ pre-embed protocol Now that every CP model is sharder-only (the vision pre-embed / sequence shard runs in the model's own forward per microbatch), the recipe-level pre-embed protocol has no remaining job. prepare_cp_forward previously routed the hook through model.__call__(_pre_embed_only=True, _cp_batch=batch, num_chunks=n) so FSDP2 forward pre-hooks would unshard vision weights during an in-hook embed; with the embed gone the hook touches no weights, so this indirection is dead. - prepare_cp_forward calls model.prepare_model_inputs_for_cp(batch, num_chunks=n) directly as a plain method and reads cp_sharder from the result. The consumed-key merge loop is deleted: every hook returns only {"cp_sharder": ...} and consumes nothing (input_ids and multimodal inputs stay in the batch for the forward). - The `if _pre_embed_only: return self.prepare_model_inputs_for_cp(...)` branch is removed from all 8 CP models (deepseek_v4, glm_moe_dsa, gemma4_moe, step3p7, qwen3_5, qwen3_5_moe, minimax_m3_vl, nemotron_omni), and the `_pre_embed_only` forward parameter is removed from the two that declared it (gemma4_moe, nemotron_omni). Recipes already invoke the hook only through prepare_cp_forward, so nothing else calls the old protocol. cp_preembed_in_forward is KEPT: it is a separate concern from the deleted protocol -- the VLM recipe reads it to stage multimodal media for PP and to invoke the (no-compute) hook on every PP stage for sunk models, which is load-bearing for CP×PP VLM runs. It is not vacuous: it distinguishes a sunk VLM (media stays in the batch for the forward) from a non-hook model. test_cp_pre_embed_protocol.py is narrowed from the kwargs-only forward-binding contract to the new contract (each model exposes prepare_model_inputs_for_cp binding (instance, batch, num_chunks)). The per-model tests that drove the removed forward branch are converted to direct prepare_model_inputs_for_cp calls or deleted; the recipe-wiring spies are rewritten to sharder-only hooks invoked directly. Full CP unit regression green in the container (per-model + protocol + dispatch: 229; recipe-wiring: 203; the only failures are a pre-existing cut_cross_entropy-not-installed environment issue). Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(cp): remove production-dead make_contiguous_shard_cp_batch_and_ctx Census after the gemma4 sunk-CP migration. make_contiguous_shard_cp_batch_and_ctx (gemma4 cp_batch.py) has zero production consumers: gemma4 now shards only the aux streams via make_contiguous_aux_only_shard_cp_batch_and_ctx and embeds / slices its primary + per_layer_inputs + ring metadata in-forward. Its only remaining references were two test files, so by the zero-consumer rule (tests are not consumers) it is deleted, along with the _GEMMA4_SEQ_KEYS / _GEMMA4_PAD_VALUES tables and the now-unused shard_batch_contiguous import it wrapped. Its test coverage is repointed to the living public entry so it survives on a production path: test_cp_utils' _contiguous_sharder fixture and test_gemma4_2b4b_cp's per_layer_inputs shard tests now call cp_sharder.shard_batch_contiguous directly with the model-provided per-token keys (per_layer_inputs / _packed_seq_ids / mm_token_type_ids) -- the same public contiguous shard DSV4/Gemma4 wrap. _synthesize_single_document_seq_ids stays (it is used by the surviving aux-only sharder and its own test) and keeps the module's _cm import alive. test_cp_utils / test_gemma4_2b4b_cp green in the container (44). Census record (the other suspects re-grepped: all LIVE, kept with consumers): - shard_batch_load_balanced + its inputs_embeds/requires_grad branch (via _shard_grad_buffer_for_cp): live -- the generic round-robin CP path in _resolve_cp_sharder; the grad-primary branch is the documented NeMo-RL inputs_embeds consumer. - _shard_grad_buffer_for_cp: live (shard_batch_load_balanced, above). - invoke_pre_embed (prepare_cp_forward param): live -- recipes pass it (llm/kd, vlm/kd, vlm/finetune) to gate the hook on PP-no-embed / KD paths. - VLM pixel-key popping (VLM_INPUT_KEYS): live -- vlm/finetune drops media on non-first PP stages; train_dspark derives its mm-key set from it. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(cp): remove the KD teacher-compat guard dead after the pre-embed sink Census follow-up to the sunk-CP migration. _validate_cp_pre_embed_teacher_compatibility (vlm/kd.py) guarded the old shared-mesh KD path where the student PRE-embedded inputs_embeds and the teacher consumed that same batch, so their input-embedding hidden sizes had to match. Now every student is sharder-only: the CP hook consumes nothing, input_ids (not inputs_embeds) rides CP, and the teacher embeds those input_ids with its own table (test_vlm_kd_cp_prepare_shards_input_ids_and_teacher_embeds_them pins exactly this). So the recipe guard at _forward_backward_step is unreachable -- `batch.get("inputs_embeds")` is always None after prepare_cp_forward -- and the failure mode it caught (teacher fed student embeds of a mismatched width) can no longer occur. Audit: the guarded branch is the only production caller of _validate_cp_pre_embed_teacher_compatibility, which is the only caller of _get_model_input_embedding_dim. Both are otherwise referenced solely by their own tests (dead code wearing a seatbelt). Deleted: the two helpers, the dead recipe branch, and their orphaned unit tests + the 6 fake-embedding model classes that only fed them. Kept coverage lands on real behavior: the sunk-contract KD test (input_ids shards, teacher embeds it) stays green. Full KD test file: 2 passed. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * debug(cp): TEMP env-gated per-param grad dump for the 26b regression (REVERT before merge) Dump every trainable param's post-backward grad once, on rank 0, at the first non-PP backward, to localize the 26b bit-exactness regression to a specific param. Two modes: - NEMO_GRAD_DUMP=stdout -> one grep-able line per param in the job log ("[GRAD_DUMP] <name> sum=<fp64 sum> norm=<fp32 norm>"), so two CI job traces (pre-sink vs sunk) can be diffed with logs/cmp_grads.py without shared storage (eos: only the job log is retrievable). - NEMO_GRAD_DUMP=<path> -> torch.save a {name: grad} .pt (local runs). full_tensor() (the FSDP DTensor all-gather) runs on ALL ranks before the rank-0-only print/save, so it is not rank-gated. Rationale: a cp2 single-fwd+bwd grad probe (dp2xcp2 via the real parallelize_model, frozen embed/vision to mirror 26b) proves the gemma4 sunk migration is bit-exact old-vs-new across text-MoE (0/82), MoE+vision-splice (0/83), FSDP (0/41) and MoE grad-accumulation (0/41); the forward is bit-exact at step 0. The 26b divergence therefore lives in the FSDP/EP gradient path -- specifically EP (deepep/ep8), which is unrunnable locally (no deepep/grouped_gemm) -- so this dump localizes it directly on the real ep8cp2 run. The recipe backward line it hooks is identical in the pre-sink snapshot, so cherry-pick this commit onto 479e8a97 too. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(cp): revert TEMP debug + collapse the pre-embed PP gate Remove the two TEMP debug commits kept for the 26b bit-exactness investigation (now closed: real gemma4_26b_a4b_moe_medpix_ep8cp2 recipe is bit-identical pre-sink vs sunk on the same image, 0/60 expert grads differ): - revert NEMO_CP_DEBUG shape logging in minimax_m3_vl (model.py, cp_sparse_attn.py) - revert the env-gated per-param NEMO_GRAD_DUMP block in the VLM recipe Also collapse the now-vacuous _pre_embed_here local in _forward_backward_step to an unconditional invoke_pre_embed=True. Every PP-capable VLM is sunk (cp_preembed_in_forward) so _model_sunk is always True, and the only recipe-level pre-embedder (gemma4) runs without PP so _is_first_or_no_pp is always True; the False branch was dead in production, exercised only by a synthetic non-sunk spy. Drop that control test and its _SpyVLM helper (a test wearing a seatbelt for removed behavior). The invoke_pre_embed parameter on prepare_cp_forward is kept — KD recipes remain live consumers. CP regression + vlm wiring/helpers/KD-mesh helpers: 266 passed. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(cp): collapse the vacuous recipe_level_cp_preembed PP gate Commit B sank every PP-capable VLM into its own forward (cp_preembed_in_forward=True), so the recipe_level_cp_preembed category — a PP+CP model with a CP hook but NOT sunk — has zero members. Its gate was therefore always False, and `if self.pp_enabled and not recipe_level_cp_preembed` was just `if self.pp_enabled`. Collapse it and keep the 156-vs-160 lesson in a one-line comment (media must ride the per-microbatch PP side channel, never schedule.step). Drop the dead parametrize case + its non-sunk _StageWithCPPrepare spy; the sunk-under-CP staging case stays as the 156-vs-160 regression guard. Census of cp_preembed_in_forward after this commit: the only remaining reader in the recipe (_model_sunk) was already removed in 6effec2b, and this removes the last one (line 576). The flag is now zero-read inside nemo_automodel — it survives only as the class-attr capability marker on the 5 sunk VLMs (qwen3_5, qwen3_5_moe, minimax_m3_vl, step3p7, gemma4_moe). Left in place: a model capability declaration with possible out-of-repo consumers (NeMo-RL / molt introspection) is a separate, cross-cutting prune, not this recipe change. vlm wiring + cp regression: 265 passed. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(cp): remove the now-zero-consumer cp_preembed_in_forward flag Its last two readers were removed on this branch: the recipe's every-stage hook gate (_model_sunk, dropped in 6effec2b, hook now invoked unconditionally) and the media-staging gate (recipe_level_cp_preembed, dropped in 122b80f1, media now always staged under PP). With zero readers, the flag is dead. Census / never-published evidence: cp_preembed_in_forward was introduced this week entirely on the unpushed huiyingl/refactor/cp-unify branch — it has never appeared in any published commit, so no out-of-repo consumer (NeMo-RL, molt, KD) can possibly read it. Zero in-repo readers + zero possible external readers = textbook zero-consumer deletion. This also resolves the prior inconsistency where the flag was set on 5 sunk VLMs (qwen3_5, qwen3_5_moe, minimax_m3_vl, step3p7, gemma4_moe) but NOT on nemotron_omni. Removed: the 5 class-attr definitions + their dead explanatory comments, the two dangling prose references (the finetune.py hook comment and the test docstring), and the two vestigial mock attrs on the wiring-test doubles. Repo grep for cp_preembed_in_forward is now empty. The models still pull media from the PP side channel in forward (documented on _pp_keep_self_forward) — only the recipe-facing marker is gone. cp regression + wiring + the 4 touched-model tests: 319 passed. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * docs(cp): fix stale + shorten overlong branch-added CP comments Comments/docstrings only — no code changes. Audits comment lines this branch ADDED (git diff origin/main...HEAD) for two problems. STALE (false after the sink landed): - finetune.py _forward_backward_step: dropped "pre-embed hook is routed through __call__ so FSDP2 forward pre-hooks fire and unshard the vision tower" — since commit C the hook is a plain method call, sharder-only, touches no weights. Rewrote to current truth. OVERLONG (compressed to the constraint, keeping the non-obvious WHY — bit-exactness with the old dispatch, aux-stream alignment, differentiability): - finetune.py every-stage-hook block: 11 -> 6 lines (kept the S/cp² RoPE-mismatch reason; dropped the moot "recipe-level pre-embedder (gemma4)" line — gemma4 is sunk now). - qwen3_5 / qwen3_5_moe / step3p7 / nemotron_omni forward CP comment: 6 -> 4 each. - minimax forward CP comment 5 -> 3; prepare_model_inputs_for_cp docstring 9 -> 6 (dropped "Megatron-style" / "PP×CP shared pre-embed graph no longer exists" history narration). - cp_utils sharder-only contract comment: 6 -> 4. Deferred: the cp_sharder.py contiguous-shard function docstrings are long but get merged/rewritten by the follow-up shard_batch_contiguous collapse, so they are compressed there rather than churned…
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
3 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.
What does this PR do ?
Fix context-parallel VLM finetuning when the full-sequence
inputs_embedstensor retains its autograd graph. This affects the shared CP path used by Nemotron Omni and Qwen3.5 after the VLM pre-embeddingtorch.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_embedstensor 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 withresize_/copy_. Resizing an autograd-tracked tensor is forbidden, so training fails before step 0 with: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
narrow()andtorch.cat()operations.resize_and double-sharding.input_idsand 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_shardreturns new local shards out of place._ContextParallelapplies CP through module wrappers and DTensor dispatch rather than the legacy global SDPA monkey-patch path.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
inputs_embeds.Test plan
ruff check nemo_automodel/components/distributed/cp_utils.py tests/unit_tests/distributed/test_cp_utils_inputs_embeds.pyruff format --check nemo_automodel/components/distributed/cp_utils.py tests/unit_tests/distributed/test_cp_utils_inputs_embeds.pyval_loss=2.0496)Before your PR is "Ready for review"
Pre checks: