refactor(distributed): unify CP input prep and dispatch across models by HuiyingLi · Pull Request #2937 · NVIDIA-NeMo/Automodel

HuiyingLi · 2026-07-06T17:26:01Z

What does this PR do ?

Unifies Automodel's context-parallelism plumbing behind one contract and one dispatch: every CP backend is a CPSharder, every recipe forward goes through a single prepare_cp_forward call, 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 torch context_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 CPSharder contract (components/distributed/cp_sharder.py) — a dataclass of two required callables, a diagnostic label, and 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 (THD cu_seqlens partitioning, magi's dispatch solver) install the partition their shard_batch just computed.
layout — diagnostic only; no framework code may branch on it.

Single dispatch, sharder returned — prepare_cp_forward / make_cp_batch_and_ctx return (ctx, batch, sharder). Resolution order: model-owned > magi > TE > generic round-robin > a layout="none" 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 uniform prepare_model_inputs_for_cp(batch: dict, *, num_chunks=1) hook, invoked through model(_pre_embed_only=True, ...) so FSDP2 pre-forward hooks fire; the hook returns {key: None} for raw inputs it consumed (FSDP2's kwargs cast can hand the hook a copy, so in-place pops don't survive). 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_batch captures its facts on the sharder (original_seq_len / padded_seq_len; the pre-flatten input_row_shape for 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:

down: 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);
up: 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_batch sits 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_fn fallback, per-key hook kwargs, cp_style/cp_layout capability flags, the CP full-logits grad touch (superseded by #2731), and the pre-embed no_grad wrappers + minimax detach (superseded by #2931) are all removed outright.

Validation

50-step branch-vs-base loss-parity runs, same node/container/seed/data (wandb: Nemo-automodel/huiyingl_workspace, runs cpunify-*):

CP flavor	config	result
Model-owned (flex ring, MoE, ep8+cp2)	gemma4-26B-A4B medpix 4k	bit-exact — 50/50 steps, mean \|Δloss\| = 0.0 (re-verified after every numerics-adjacent commit)
Generic torch SDPA CP (cp2)	qwen3 tulu3 chat	pass — mean \|Δloss\| ~0.01% at every checkpoint: original refactor, backend-sharder unification, and the shard-facts capture (post-main-merge)
MagiAttention (packed, ep8+cp2)	qwen3-MoE-30B	pass — mean \|Δloss\| 0.2–0.3% across three re-verifications, within the magi/MoE bf16 noise band PR #2384 established

1500+ affected unit tests pass; new L0 tests cover the layout math (closed-form indices vs torch's chunk pairing, shard→gather round trips), sharder resolution order, shard-time captures per backend, caller-coordinate verb round trips (incl. the DSV4 repad position map), and the mismatched-length guards.

Not reproducible in this environment: dsv4/minimax multi-node-only configs; gemma4-31B packed config (preprocessed dataset artifact). Known pre-existing issues found while validating (identical on both trees, not addressed here): torch-2.12 context_parallel vs grad-carrying inputs_embeds on resize-padding.

Before your PR is "Ready for review"

Pre checks:

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

Additional Information

Closes Make prepare_model_inputs_for_cp consistent across all models #2879; delivers the downstream contract for Feature Request: Generic Automodel CP interface for downstream libraries e.g. NeMo RL #2861 (no separate Plan object: every field it would bind is one call away from the returned sharder — backend = sharder.layout, lengths are captured facts, verbs are sharder methods)
Merged with main (includes fix(deepseek_v4): support packed THD with context parallel #2731 DSV4 packed THD CP, translated onto the shared contiguous sharder + extra_seq_keys; its packed repad now also emits the input-position map for the token verbs)
NeMo-RL usage sketch: ctx, batch, sharder = prepare_cp_forward(model, mesh, batch) → co-shard per-token tensors with sharder.shard_token_tensor(t, fill=0) → local per-token loss → sharder.gather_token_tensor(lp, trim=True) only where full-sequence logprobs are needed
Assumes the semantics of fix(cp): preserve gradients when sharding VLM inputs #2931 (no-grad wrapper removal); rebase coordination with feat(magi): honor AttnMaskSpec on the HF attention backend #2622 (magi_attn_utils)

🤖 Generated with Claude Code

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>

copy-pr-bot · 2026-07-06T17:26:04Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

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>

HuiyingLi · 2026-07-07T07:59:38Z

Pushed two updates:

97b533df drops the CP full-logits grad touch end-to-end (the finalize_loss slot on CPSharder, the full_logits_grad_touch helper, the DSV4/GLM wirings, and the recipe loss-site consumption), matching its removal in fix(deepseek_v4): support packed THD with context parallel #2731 — this also shrinks the eventual rebase between the two PRs.
01f02d5d fixes two test files that were still asserting pre-refactor internals (test_finetune_vlm_helpers.py patched the recipe-module make_cp_batch_and_ctx binding that moved into cp_utils; one test_glm_moe_dsa_tilelang.py test asserted the retired private batch keys).

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>

akoumpa · 2026-07-10T06:46:26Z

+
+
+@dataclass
+class CPSharder:


Suggested change

class CPSharder:

class ContextParallelismSharder:

…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>

This reverts commit 3985a7c. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

akoumpa · 2026-07-10T09:01:55Z

/claude review

…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>

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>

chtruong814 · 2026-07-10T14:30:17Z

/claude review

chtruong814 · 2026-07-10T16:02:47Z

/claude review

github-actions · 2026-07-10T16:13:46Z

+            key: None
+            for key in ("input_ids", "pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw")
+        }
+        return {**consumed, "inputs_embeds": inputs_embeds}


The prior implementation returned inputs_embeds.detach() here, with a comment documenting a concrete reason: torch's context_parallel shards buffers via in-place resize_() and rejects tensors that require grad, so the un-detached embeds could raise at CP-shard time. This PR drops the .detach() (and that comment) and only documents the new None-consumed-keys mechanism, which doesn't address the original concern.

Sibling VLM models (step3p7, nemotron_omni, qwen3_5) return un-detached inputs_embeds on the same round-robin context_parallel path, so this may be a safe harmonization — but only if MiniMax's embedding/vision path is likewise frozen (or grad-requiring buffers are accepted). Please confirm a cp_size>1 MiniMax run still shards without error; if the embeddings can require grad here, restore the .detach().

HuiyingLi added 6 commits July 5, 2026 10:48

HuiyingLi added 2 commits July 7, 2026 00:42

HuiyingLi added 5 commits July 7, 2026 02:17

HuiyingLi mentioned this pull request Jul 8, 2026

fix(deepseek_v4): support packed THD with context parallel #2731

Merged

HuiyingLi and others added 6 commits July 8, 2026 06:24

HuiyingLi force-pushed the huiyingl/refactor/cp-unify branch from ae293cf to 749391a Compare July 8, 2026 16:45

HuiyingLi linked an issue Jul 9, 2026 that may be closed by this pull request

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

Open

akoumpa reviewed Jul 10, 2026

View reviewed changes

akoumpa and others added 4 commits July 10, 2026 00:03

Merge branch 'huiyingl/refactor/cp-unify' of https://github.com/NVIDI…

c3014ad

…A-NeMo/Automodel into huiyingl/refactor/cp-unify Signed-off-by: HuiyingLi <willwin.lee@gmail.com> # Conflicts: # nemo_automodel/components/distributed/cp_utils.py

fix(distributed): preserve CP preparation contracts

3985a7c

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

akoumpa force-pushed the huiyingl/refactor/cp-unify branch from 32d609e to 3985a7c Compare July 10, 2026 08:38

Revert "fix(distributed): preserve CP preparation contracts"

78a987c

This reverts commit 3985a7c. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

HuiyingLi and others added 2 commits July 10, 2026 04:48

github-actions Bot reviewed Jul 10, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

refactor(distributed): unify CP input prep and dispatch across models#2937

refactor(distributed): unify CP input prep and dispatch across models#2937
HuiyingLi wants to merge 26 commits into
mainfrom
huiyingl/refactor/cp-unify

HuiyingLi commented Jul 6, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jul 6, 2026

Uh oh!

HuiyingLi commented Jul 7, 2026

Uh oh!

akoumpa Jul 10, 2026

Uh oh!

akoumpa commented Jul 10, 2026

Uh oh!

chtruong814 commented Jul 10, 2026

Uh oh!

chtruong814 commented Jul 10, 2026

Uh oh!

github-actions Bot Jul 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

HuiyingLi commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Changelog

Validation

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Jul 6, 2026

Uh oh!

HuiyingLi commented Jul 7, 2026

Uh oh!

akoumpa Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

akoumpa commented Jul 10, 2026

Uh oh!

chtruong814 commented Jul 10, 2026

Uh oh!

chtruong814 commented Jul 10, 2026

Uh oh!

github-actions Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

HuiyingLi commented Jul 6, 2026 •

edited

Loading