refactor(distributed): unify CP input prep and dispatch across models - #2937
Merged
Conversation
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>
…ocation
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>
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>
…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>
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>
…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>
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>
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>
Contributor
Author
|
Pushed two updates:
Verified in the auto2604 container: affected unit-test selection green (1107 passed / 42 skipped), ruff format+check clean. |
…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>
…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>
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>
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>
… 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>
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>
…arder 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>
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>
…out 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>
… 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>
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>
HuiyingLi
force-pushed
the
huiyingl/refactor/cp-unify
branch
from
July 8, 2026 16:45
ae293cf to
749391a
Compare
akoumpa
reviewed
Jul 10, 2026
…unify Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> # Conflicts: # nemo_automodel/components/distributed/cp_utils.py # nemo_automodel/components/models/deepseek_v4/cp.py # nemo_automodel/components/models/deepseek_v4/model.py # nemo_automodel/components/models/gemma4_moe/model.py # nemo_automodel/recipes/llm/train_ft.py # tests/unit_tests/models/deepseek_v4/test_dsv4_cp_batch.py # tests/unit_tests/models/glm_moe_dsa/test_glm_moe_dsa_tilelang.py # tests/unit_tests/recipes/test_train_ft.py
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>
…A-NeMo/Automodel into huiyingl/refactor/cp-unify Signed-off-by: HuiyingLi <willwin.lee@gmail.com> # Conflicts: # nemo_automodel/components/distributed/cp_utils.py
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
akoumpa
force-pushed
the
huiyingl/refactor/cp-unify
branch
from
July 10, 2026 08:38
32d609e to
3985a7c
Compare
cp_dispatcher_suspended is called by VLM forward paths while media encoders run outside the text CP attention. In fake/no-outer-context unit tests the torch context_parallel dispatcher is not active on entry, so the helper must not re-enable it with the fake mesh on exit. Guard the re-enable path by checking whether SDPA was the torch context_parallel wrapper on entry, and add a Nemotron Omni CP regression that asserts a fake CP forward leaves global SDPA unchanged. Validation: - pytest tests/unit_tests/models/nemotron_omni/test_nemotron_omni_cp.py tests/unit_tests/models/nemotron_parse/test_nemotron_parse_model.py::test_nemotron_parse_forward_with_stub_encoder - pytest tests/unit_tests/models/nemotron_omni/test_nemotron_omni_cp.py tests/unit_tests/models/nemotron_parse/test_nemotron_parse_model.py::test_nemotron_parse_forward_with_stub_encoder tests/unit_tests/models/nemotron_v3/test_nemotron_v3_model.py::TestNemotronV3Model::test_model_forward_shape tests/unit_tests/models/qwen3_5/test_qwen3_5_mtp.py::TestQwen3_5MTPModel::test_mtp_disabled_matches_hf_forward tests/unit_tests/speculative/test_dflash_core.py::test_forward_returns_finite_loss_and_grads_flow_to_draft - pytest tests/unit_tests/models/glm_moe_dsa/test_glm_moe_dsa_cp.py tests/unit_tests/models/glm_moe_dsa/test_glm_moe_dsa_tilelang.py -k "glm_dsa_prepare_model_inputs_for_cp_binds_batch_sharder or make_glm_dsa_packed_cp_batch or shard_glm_dsa_packed_cp_batch" - pytest tests/unit_tests/models/nemotron_omni/test_nemotron_omni_cp.py tests/unit_tests/distributed/test_cp_utils.py::test_prepare_cp_forward_merges_model_hook_batch_updates - ruff format --check on touched files - ruff check on touched files - git diff --check Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Contributor
Author
|
/ok to test 49f897a |
Contributor
|
/claude review |
9 tasks
Contributor
|
/ok to test 65dcfb7 |
akoumpa
previously approved these changes
Jul 22, 2026
Contributor
Author
|
/ok to test c7bb59b |
* refactor(distributed): simplify CP sharder API Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * refactor(distributed): group context parallel modules Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * refactor(cp): simplify THD recipe policy Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * refactor(cp): construct strategy sharders directly Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * refactor(cp): remove legacy sharder construction Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * test(cp): update VLM sharder mocks Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> --------- Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> Signed-off-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Contributor
Author
|
/ok to test faec4aa |
…unify Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Contributor
Author
|
/ok to test 4d8970d |
Contributor
3 tasks
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
4 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 ?
Unifies Automodel's context-parallelism plumbing behind one contract and one dispatch: every CP backend is a
ContextParallelismSharder, every recipe forward goes through a singleprepare_cp_forwardcall, and the resolved sharder is returned to the caller — making it the public downstream CP interface proposed in #2861. Model-specific CP logic lives in model directories; the framework builds sharders for the paths it owns (generic torchcontext_parallel, TE/THD, MagiAttention). Closes #2879; delivers the contract for #2861.Changelog
Reviewable sequentially; later commits also delete the transitional compatibility earlier ones introduced, so the final tree carries no shims. By theme:
The
ContextParallelismShardercontract (components/distributed/cp_sharder.py) — a dataclass of two required callables plus shard-time-captured facts, replacing the private batch keys (_cp_make_batch_fn,_cp_metadata_*,_cp_full_logits_grad_touch):shard_batch(cp_mesh, tp_mesh, batch, *, loss_mask=None, padding_token_id=0) -> (ctx_factory, batch)local_token_global_indices(cp_mesh, padded_seq_len, device) -> LongTensor— the global position of every local token, the universal layout coordinate system.shard_token_tensor/gather_token_tensor(differentiable all-gather + reorder) are synthesized from it. Closed-form for contiguous/round-robin; data-dependent layouts (THDcu_seqlenspartitioning, magi's dispatch solver) install the partition theirshard_batchjust computed.The sharder carries no backend tag — nothing may branch on which backend produced it (a
layoutdiagnostic string existed transiently and was removed as consumer-less).Single dispatch, sharder returned — the public entry point is exactly one function:
prepare_cp_forward(model, device_mesh, batch, ...) -> (ctx, batch, sharder)(the oldmake_cp_batch_and_ctxis now the private internal seam). Resolution order: model-owned > magi > TE > generic round-robin > an identity sharder, so consumer code is branch-free across cp_size 1 and N. Model-owned CP models (gemma4_moe, deepseek_v4, glm_moe_dsa) return their sharder from the uniformprepare_model_inputs_for_cp(batch: dict, *, num_chunks=1)hook, invoked by the dispatch as a plain method call (the hook is sharder-only — it touches no weights, so no FSDP2 unshard routing is needed). The generic torch path's round-robin indices are closed-form (torch's 2·cp head-tail chunk pairing).Caller-coordinate token verbs (the #2861 asks) — padding is an internal detail of the CP layout, so
shard_batchcaptures its facts on the sharder (original_seq_len/padded_seq_len; the pre-flatteninput_row_shapefor THD, whose BSHD→THD flatten is a pure reshape; the input→rebuilt-row position map for DSV4's packed repad, which genuinely repositions tokens). The verbs then accept and return tensors in the caller's coordinates:sharder.shard_token_tensor(advantages, fill=0)— original-length / row-shaped / input-coordinate tensors are padded, flattened, or scattered exactly like the model inputs before sharding; alignment is by construction (zpqiu's co-sharding ask);sharder.gather_token_tensor(local_logprobs, trim=True)— token-level gather without materializing[B,S,V]logits, restored to the original length/shape (the issue's gather ask).Any length that cannot be transformed unambiguously raises — this closes a silent-misalignment window on contiguous layouts with a custom
pad_multiple(a plausible length passed the cp-divisibility check but did not match the sharded stream).MagiAttention as a first-class backend — recipe-static facts bind once at
setup_magi;MagiState.make_cp_batchsits at the TE dispatch rung (both also run at cp≤1 for packing conversion / mask-spec activation) and returns its dispatch positions (get_position_ids) as the sharder's index map.Deletions — no deprecation shims survive: the legacy
_cp_make_batch_fnfallback, per-key hook kwargs,cp_style/cp_layoutcapability flags, the CP full-logits grad touch (superseded by #2731), and the pre-embedno_gradwrappers + minimax detach (superseded by #2931) are all removed outright.Validation
Full per-run ledger (every link): validation comment. Summary:
Re-validation on the final tip
49f897a1(2026-07-21; one node/container, before = merge-basecdaaff91, 20 steps, same seed/data) — every single-node CP path and three non-CP controls are bit-exact at logged precision (0/21 steps differ):Multi-node / at-scale (nemo-ci, 20-step before/after — the paths a single node cannot exercise):
(gemma4 26b/31b CI pairs are 20/20 bit-exact on both torch-2.12 and 2.13 stacks; the single cross-image 0.672% delta was traced to per-build CI image drift via a same-image control — 60/60 expert grads bit-identical. Full matrix in the ledger comment.)
Related fixes validated: transformers-5.12 minimax config regression (#3134) and torch-2.13 SAC×FSDP2 (#3133), evidence in the ledger.
1500+ affected unit tests pass; committed GPU mechanism tests (torchrun, 2–4 GPUs) cover the CP token verbs over real NCCL, minimax cp2 forward-equivalence, the cp2×pp2 layer-2 sink harness (incl. MTP and image variants), and the ring-suspend contract.
Before your PR is "Ready for review"
Pre checks:
Additional Information
extra_seq_keys; its packed repad now also emits the input-position map for the token verbs)ctx, batch, sharder = prepare_cp_forward(model, mesh, batch)→ co-shard per-token tensors withsharder.shard_token_tensor(t, fill=0)→ local per-token loss →sharder.gather_token_tensor(lp, trim=True)only where full-sequence logprobs are neededmake_cp_batch_and_ctxdirectly (2-tuple return) or invoked the pre-embed hook viamodel(_pre_embed_only=True, **mm_kwargs)(e.g. molt's FSDP worker) must move toprepare_cp_forward— one call replaces both, plus the hand-rolled pad/index/gather mathtests/functional_tests/context_parallel/run_cp_sharder_token_verbs.py) pins the token verbs' collective semantics, including the sum-over-consumers gradient of the differentiable gather🤖 Generated with Claude Code
Update: per-microbatch CP pre-embed sinking (Megatron-style)
Following the Megatron-LM/Megatron-Bridge pattern (embed + vision splice + CP shard inside the model forward, per microbatch), all VLM CP pre-embedding moved from the recipe dispatch into each model's forward:
shard_batch_aux_only(no-grad aux streams shard at dispatch; input streams stay full-length) + differentiableshard_sequence_for_cpinside forward.shard_batch_contiguous(shard_primary=False)aux sharding + in-forward contiguous slice; per_layer_inputs built per microbatch.cp_dispatcher_suspended(bidirectional attention must not be ring-patched)._pre_embed_only__call__protocol (all 8 models), the recipe-level pre-embed gates, and the transitionalcp_preembed_in_forwardmarker — the dispatch now calls the hook as a plain method.get_pipeline_stage_metashook; the sharder-only hook is invoked on every PP stage soupdate_seq_lensees a consistent sequence length on all ranks.Validation: committed GPU tests (CP verbs over NCCL; cp2 forward-equivalence; cp2×pp2 layer-2 harness incl. MTP and image variants — no double-backward, cp2×pp2 == cp2×pp1 bf16-exact, embed/vision grads flow) plus a full real-recipe before/after parity ledger with nemo-ci/wandb links — see the validation comment: #2937 (comment)