[None][feat] Port 13 AutoDeploy custom models to sharding IR + opt them in via registry - #14778
Merged
greg-kwasniewski1 merged 26 commits intoJun 9, 2026
Merged
Conversation
Add explicit sharding-hint ops to the canonical Gemma modeling file:
- torch.ops.auto_deploy.torch_linear_simple for q/k/v/o and MLP projections,
with tp_mode (colwise / rowwise), tp_min_local_shape=head_dim on Q/K/V
(GQA-safe), and layer_type ("mha" / "mlp").
- torch.ops.auto_deploy.view (tp_scaled_dim=2, layer_type="mha") on
Q/K/V/attn_output reshapes whose head-count dim scales with TP.
- torch.ops.auto_deploy.all_reduce after the rowwise o_proj and MLP down_proj.
Model logic is untouched: Gemma's (1+weight) RMSNorm load hook, the
embedding sqrt(hidden_size) normalization, the explicit attention_mask
path with is_causal=False, and the parameter / module hierarchy all
remain identical to the legacy file.
Validated via the offline equivalence test from NVIDIA#13963 on a tiny 4-layer
Gemma config across tp-only / ep-only / tep / attn-dp parallelism configs.
Refs NVIDIA#14642
Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Add explicit sharding-hint ops to the canonical Mistral modeling file:
- torch.ops.auto_deploy.torch_linear_simple for q/k/v/o and MLP projections,
with tp_mode (colwise / rowwise), tp_min_local_shape=head_dim on Q/K/V
(GQA-safe), and layer_type ("mha" / "mlp").
- torch.ops.auto_deploy.view (tp_scaled_dim=2, layer_type="mha") on
Q/K/V/attn_output reshapes whose head-count dim scales with TP.
- torch.ops.auto_deploy.all_reduce after the rowwise o_proj and MLP down_proj.
Model logic is untouched: sliding-window attention, is_causal=True
torch_attention call, parameter / module hierarchy, RMSNorm semantics
all remain identical to the legacy file.
Validated via the offline equivalence test from NVIDIA#13963 on a tiny 4-layer
Mistral config across tp-only / ep-only / tep / attn-dp parallelism configs.
Refs NVIDIA#14642
Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Add explicit sharding-hint ops to the canonical Granite modeling file:
- torch.ops.auto_deploy.torch_linear_simple for q/k/v/o and MLP projections,
with tp_mode (colwise / rowwise), tp_min_local_shape=head_dim on Q/K/V
(GQA-safe), and layer_type ("mha" / "mlp").
- torch.ops.auto_deploy.view (tp_scaled_dim=2, layer_type="mha") on
Q/K/V/attn_output reshapes whose head-count dim scales with TP.
- torch.ops.auto_deploy.all_reduce after the rowwise o_proj and MLP down_proj.
Model logic is untouched: Granite's embedding_multiplier, residual_multiplier,
attention_multiplier (used as the attention scaling kwarg), and logits_scaling
are all preserved.
Validated via the offline equivalence test from NVIDIA#13963 on a tiny 4-layer
Granite config across tp-only / ep-only / tep / attn-dp parallelism configs.
Refs NVIDIA#14642
Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Add explicit sharding-hint ops to the canonical HunYuan Dense V1 modeling file:
- torch.ops.auto_deploy.torch_linear_simple for q/k/v/o and MLP projections,
with tp_mode (colwise / rowwise), tp_min_local_shape=head_dim on Q/K/V
(GQA-safe), and layer_type ('mha' / 'mlp').
- torch.ops.auto_deploy.view (tp_scaled_dim=2, layer_type='mha') on
Q/K/V/attn_output reshapes whose head-count dim scales with TP.
- torch.ops.auto_deploy.all_reduce after the rowwise o_proj and MLP down_proj.
Model logic is untouched: HunYuan's QK-normalization (RMSNorm on Q/K applied
after RoPE), Dynamic NTK-Alpha RoPE scaling, and SiLU-gated MLP are preserved.
Validated via the offline equivalence test from NVIDIA#13963 on a tiny 4-layer
HunYuan Dense config across tp-only / ep-only / tep / attn-dp parallelism
configs.
Refs NVIDIA#14642
Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Add explicit sharding-hint ops to the canonical Seed-OSS modeling file:
- torch.ops.auto_deploy.torch_linear_simple for q/k/v/o and MLP projections,
with tp_mode (colwise / rowwise), tp_min_local_shape=head_dim on Q/K/V
(GQA-safe), and layer_type ('mha' / 'mlp').
- torch.ops.auto_deploy.view (tp_scaled_dim=2, layer_type='mha') on
Q/K/V/attn_output reshapes whose head-count dim scales with TP.
- torch.ops.auto_deploy.all_reduce after the rowwise o_proj and MLP down_proj.
Model logic (GQA semantics, attention_bias / attention_out_bias asymmetry,
standard RoPE, RMSNorm, residuals) is untouched.
Validated via the offline equivalence test from NVIDIA#13963 on a tiny 4-layer
Seed-OSS config across tp-only / ep-only / tep / attn-dp parallelism configs.
Refs NVIDIA#14642
Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Add explicit sharding-hint ops to the canonical Gemma 2 modeling file:
- torch.ops.auto_deploy.torch_linear_simple for q/k/v/o and MLP projections,
with tp_mode (colwise / rowwise), tp_min_local_shape=head_dim on Q/K/V
(GQA-safe), and layer_type ('mha' / 'mlp').
- torch.ops.auto_deploy.view (tp_scaled_dim=2, layer_type='mha') on
Q/K/V/attn_output reshapes whose head-count dim scales with TP.
- torch.ops.auto_deploy.all_reduce after the rowwise o_proj and MLP down_proj.
Model logic is untouched: Gemma 2's (1+weight) RMSNorm scaling, 4 layer norms
per decoder layer, attention logit softcapping, alternating sliding window /
full attention, custom query_pre_attn_scalar attention scaling, embedding
sqrt(hidden_size) normalization, and final logit softcapping are all preserved.
Validated via the offline equivalence test from NVIDIA#13963 on a tiny 4-layer
Gemma 2 config across tp-only / ep-only / tep / attn-dp parallelism configs.
Refs NVIDIA#14642
Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Add explicit sharding-hint ops to the canonical Cohere/Cohere2 modeling file:
- torch.ops.auto_deploy.torch_linear_simple for q/k/v/o and MLP projections,
with tp_mode (colwise / rowwise), tp_min_local_shape=head_dim on Q/K/V
(GQA-safe), and layer_type ('mha' / 'mlp').
- torch.ops.auto_deploy.view (tp_scaled_dim=2, layer_type='mha') on
Q/K/V/attn_output reshapes whose head-count dim scales with TP.
- torch.ops.auto_deploy.all_reduce after the rowwise o_proj and MLP down_proj.
Model logic (parallel attention + MLP pattern, CohereLayerNorm, interleaved
RoPE via torch_rope_with_qk_interleaving, Cohere2 sliding window / conditional
RoPE, optional CohereLayerNorm-based Q/K-norm applied before RoPE on the
per-head head_dim axis) is untouched. Q/K-norm weights live on head_dim, which
is not a TP-sharded dim, so they remain replicated and require no hints.
Validated via the offline equivalence test from NVIDIA#13963 on a tiny 4-layer
Cohere config across tp-only / ep-only / tep / attn-dp parallelism configs.
Refs NVIDIA#14642
Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Add explicit sharding-hint ops to the canonical EXAONE modeling file:
- torch.ops.auto_deploy.torch_linear_simple for q/k/v/out and MLP projections
(c_fc_0 / c_fc_1 / c_proj), with tp_mode (colwise / rowwise),
tp_min_local_shape=head_dim on Q/K/V (GQA-safe), and layer_type
('mha' / 'mlp').
- torch.ops.auto_deploy.view (tp_scaled_dim=2, layer_type='mha') on
Q/K/V/attn_output reshapes whose head-count dim scales with TP.
- torch.ops.auto_deploy.all_reduce after the rowwise out_proj and MLP c_proj.
Model logic (EXAONE-specific naming, llama3-style RoPE scaling, bundled
ExaoneConfig, ExaoneAttentionBlock checkpoint-key wrapper) is untouched.
Validated via the offline equivalence test from NVIDIA#13963 on a tiny 4-layer
EXAONE config across tp-only / ep-only / tep / attn-dp parallelism configs.
Refs NVIDIA#14642
Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Add explicit sharding-hint ops to the canonical Llama 3 modeling file:
- torch.ops.auto_deploy.torch_linear_simple for q/k/v/o and MLP projections,
with tp_mode (colwise / rowwise), tp_min_local_shape=head_dim on Q/K/V
(GQA-safe), and layer_type ('mha' / 'mlp').
- torch.ops.auto_deploy.view (tp_scaled_dim=2, layer_type='mha') on
Q/K/V/attn_output reshapes whose head-count dim scales with TP.
- torch.ops.auto_deploy.all_reduce after the rowwise o_proj and MLP down_proj.
Model logic (Llama 3 / 3.1 / 3.2 / 3.3 GQA, SwiGLU, RMSNorm, RoPE with
llama3-style scaling via init_rope_inv_freq) is untouched.
Validated via the offline equivalence test from NVIDIA#13963 on a tiny 4-layer
Llama 3 config across tp-only / ep-only / tep / attn-dp parallelism configs.
Refs NVIDIA#14642
Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
modeling_llama3_ir.py was an earlier sharding-IR variant that diverged from the canonical modeling_llama3.py before a rope-handling fix landed (see commit that introduced _rope_utils.init_rope_inv_freq). The _ir variant indexes ROPE_INIT_FUNCTIONS["default"] directly, but recent transformers versions no longer have the "default" key (valid keys: dynamic / linear / llama3 / longrope / proportional / yarn), so loading the model raises `KeyError: 'default'`. Since the canonical modeling_llama3.py is now itself sharding-IR-aware (previous commit) and routes RoPE init through the version-safe init_rope_inv_freq helper, the _ir variant is redundant and broken. Confirmed no in-repo importers (grep against tensorrt_llm/, tests/, examples/). Refs NVIDIA#14642 Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
…ant custom_ops import Now that sharding IR is the AutoDeploy default, the skill no longer needs to advocate for it in every ported file. Update the allowlist (Step 0 / A4) and self-audit table (Step 12) to require a single-line module-level header noting the file uses sharding IR, per-class 'Sharding strategy:' blocks on MLP/Attention/MoE etc., and NO 'Shardable custom ops used:' trailer. Also drop the prior A4 'add custom_ops side-effect import' rule. The line 'from ... import custom_ops # noqa: F401' is redundant: torch.ops.auto_deploy.* is already registered by importing the tensorrt_llm._torch.auto_deploy package, which any test or runtime path does before loading a modeling file. New ports must not add it (F8). Older reference models still carry the redundant import and the verbose docstring; do not replicate that pattern. Step 2 is rewritten to explicitly say 'do NOT add the side-effect import'. Refs NVIDIA#14642 Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
…instructions Skills should tell the agent what TO do, not what NOT to do. The previous commit (2e7d34b) used negative instructions ('do not add ...') and introduced an F8 row in the self-audit table for the redundant custom_ops import. If the allowlist only names allowed changes, the agent never adds anything outside it — no need for a 'do not add' counterweight. - Strips negative phrasing from A4 (docstring rule). A4 now lists only what to add: a single-line module header + per-class Sharding strategy blocks. - Removes the old Step 2 ('Do NOT add a custom_ops side-effect import') entirely; remaining steps renumbered down by one (Step 11 was Step 12, etc.). Cross-references in body and validation checklist updated. - Drops F8 from the self-audit table. A4 description trimmed to the affirmative form. Refs NVIDIA#14642 Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
…ted files Match the convention now codified in the ad-sharding-ir-port skill (commit 2baa14b): one-line module-level docstring header '"""<Model> model (sharding IR).' followed by the original source-of-truth / HF link block; per-class 'Sharding strategy:' blocks already in place. Removes from the 9 ported modeling files: - The 6-line 'This is a sharding-aware rewrite ... deterministic, node-local sharding' paragraph that advocated for sharding IR. - The 4-line 'Shardable custom ops used:' trailer. - The 'from ... import custom_ops # noqa: F401' side-effect import line, which is redundant: torch.ops.auto_deploy.* is already registered by importing the tensorrt_llm._torch.auto_deploy package, which any test or runtime path does before loading a modeling file. Net effect: -117 lines, no semantic change. Equivalence test re-confirmed on modeling_gemma.py (tp-only) after the slim. Refs NVIDIA#14642 Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Add explicit sharding-hint ops to the canonical DeepSeek-V2 modeling file:
- torch.ops.auto_deploy.torch_linear_simple for q_a/q_b/q, kv_a, o_proj,
MLP gate/up/down, and shared-expert MLP, with tp_mode (none / colwise /
rowwise) and layer_type ('mla' / 'mlp' / 'moe').
- torch.ops.auto_deploy.view (tp_scaled_dim=2, layer_type='mla') on the
Q reshape and post-attention output reshape whose head-count dim scales
with TP.
- torch_mla called with enable_sharding=True, layer_type='mla' so
_apply_hint_mla shards kv_b_proj.weight column-wise per head.
- torch.ops.auto_deploy.all_reduce after rowwise o_proj (layer_type='mla')
and a single all_reduce after the (routed + shared) MoE merge point
(layer_type='moe'). Shared expert MLP constructed with
add_all_reduce=False, layer_type='moe' so its closing all_reduce is
deferred.
- torch_moe carries layer_type='moe'.
Model logic (MLA semantics, softmax-based router gate, RMSNorm, residuals)
is untouched.
Validated via the offline equivalence test from NVIDIA#13963 on a tiny 4-layer
DeepSeek-V2 config across tp-only / ep-only / tep / attn-dp parallelism
configs.
Refs NVIDIA#14642
Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Add explicit sharding-hint ops to the canonical GLM4 MoE modeling file:
- torch.ops.auto_deploy.torch_linear_simple for q/k/v/o, MLP gate/up/down
and shared-expert MLP, with tp_mode (colwise / rowwise),
tp_min_local_shape=head_dim on K/V (GQA-safe), and layer_type
('mha' / 'mlp' / 'moe').
- torch.ops.auto_deploy.view (tp_scaled_dim=2, layer_type='mha') on
Q/K/V/attn_output reshapes whose head-count dim scales with TP.
- torch.ops.auto_deploy.all_reduce after rowwise attn o_proj (layer_type=
'mha') and a single all_reduce after the (routed + shared) MoE merge
point (layer_type='moe'). Shared expert MLP constructed with
add_all_reduce=False, layer_type='moe' so its closing all_reduce is
deferred.
- torch_moe carries layer_type='moe'. The fused noaux_tc_op router gate
is TP-replicated and kept verbatim.
Model logic (GQA + partial RoPE, QK norm, RMSNorm, residuals) is
untouched.
Validated via the offline equivalence test from NVIDIA#13963 on a tiny 4-layer
GLM4 MoE config across tp-only / ep-only / tep / attn-dp parallelism
configs.
Refs NVIDIA#14642
Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Add explicit sharding-hint ops to the canonical Qwen3 MoE modeling file:
- torch.ops.auto_deploy.torch_linear_simple for q/k/v/o and dense MLP
gate/up/down, with tp_mode (colwise / rowwise),
tp_min_local_shape=head_dim on K/V (GQA-safe), and layer_type
('mha' / 'mlp').
- torch.ops.auto_deploy.view (tp_scaled_dim=2, layer_type='mha') on
Q/K/V/attn_output reshapes whose head-count dim scales with TP.
- torch.ops.auto_deploy.all_reduce after rowwise attn o_proj
(layer_type='mha'), after rowwise MLP down (layer_type='mlp'), and
after torch_moe (layer_type='moe'). No shared expert in Qwen3 MoE, so
a single MoE all_reduce suffices.
- torch_moe carries layer_type='moe'. The plain nn.Linear router gate is
TP-replicated and kept verbatim.
Model logic (GQA, per-head Q/K RMSNorm, RoPE, residuals) is untouched.
Validated via the offline equivalence test from NVIDIA#13963 on a tiny 4-layer
Qwen3 MoE config across tp-only / ep-only / tep / attn-dp parallelism
configs.
Refs NVIDIA#14642
Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Add explicit sharding-hint ops to the canonical Llama 4 modeling file:
- torch.ops.auto_deploy.torch_linear_simple for q/k/v/o and MLP
gate/up/down (including the shared expert), with tp_mode (colwise /
rowwise), tp_min_local_shape=head_dim on K/V (GQA-safe), and layer_type
('mha' / 'mlp' / 'moe').
- torch.ops.auto_deploy.view (tp_scaled_dim=2, layer_type='mha') on
Q/K/V/attn_output reshapes whose head-count dim scales with TP.
- torch.ops.auto_deploy.all_reduce after rowwise attn o_proj
(layer_type='mha'), after rowwise MLP down on dense layers
(layer_type='mlp'), and a single all_reduce after the (routed + shared)
MoE merge point (layer_type='moe'). Shared expert MLP constructed with
add_all_reduce=False, layer_type='moe' so its closing all_reduce is
deferred.
- The routed MoE remains expressed as stacked-weight torch.bmm; the
pattern_matcher stage rewrites it into torch_moe (which carries
layer_type='moe' by default) before sharding runs. The sigmoid router
is a plain nn.Linear and stays TP-replicated.
Model logic (GQA + complex RoPE, NoPE layers, L2 QK norm, attn
temperature tuning, residuals) is untouched.
Validated via the offline equivalence test from NVIDIA#13963 on a tiny 4-layer
Llama 4 config across tp-only / ep-only / tep / attn-dp parallelism
configs.
Refs NVIDIA#14642
Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
…orted model families Add 'enable_sharder_ir.yaml' to the yaml_extra of every model_registry/models.yaml entry whose config class maps to one of the 13 modeling files ported in this PR (llama3, mistral, granite, hunyuan_dense, seed_oss, gemma2, cohere, exaone, gemma, deepseek_v2, glm4_moe, qwen3_moe, llama4). Without this, those entries continue to run through the legacy detect_sharding pipeline. detect_sharding does recognize torch_linear_simple via is_any_lin_op, so it still shards correctly, but the explicit sharding-hint kwargs on the IR-ported ops are dormant metadata under that path. This commit makes the registry actually exercise the apply_sharding_hints transform on those models, matching the opt-in pattern PR NVIDIA#13478 used for the first four families (qwen3, qwen3_5_moe, deepseek, nemotron_h). 70 entries updated. Excluded by design (still on legacy via no override): - 'multimodal.yaml' entries: vision/audio variants take a separate path. - 'simple_shard_only.yaml' entries: Llama-3_1-Nemotron-51B, -Ultra-253B (and -FP8), Qwen3-30B-A3B, Qwen3-235B-A22B - explicit opt-outs of the regular sharding pass. Refs NVIDIA#14642 Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
|
PR_Github #51184 [ run ] triggered by Bot. Commit: |
Contributor
📝 WalkthroughWalkthroughThis PR transitions AutoDeploy from legacy PyTorch execution to sharding-IR across 13 text-only models by replacing direct linear/view operations with ChangesAutoDeploy Sharding IR Transition
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/auto_deploy/model_registry/models.yaml`:
- Line 6: The models.yaml's yaml_extra lists include enable_sharder_ir.yaml but
several dense-family model entries (the multimodal/simple_shard_only groups and
the default_ws_2, default_ws_4, default_ws_8 blocks) are missing that include;
verify whether the omission is intentional and if not add
"enable_sharder_ir.yaml" to the yaml_extra array for the specific model entries
named (default_ws_2 entries: nvidia/NVIDIA-Nemotron-Nano-12B-v2,
nvidia/Nemotron-H-8B-Base-8K, nvidia/Nemotron-H-8B-Reasoning-128K; default_ws_4
entries: neuralmagic/Mixtral-8x7B-Instruct-v0.1-FP8,
nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3, nvidia/Nemotron-H-47B-Reasoning-128K,
nvidia/OpenReasoning-Nemotron-32B; default_ws_8 entry:
mistralai/Mixtral-8x22B-Instruct-v0.1) by adding the filename to each entry's
yaml_extra list to match the other configs (or document why these remain
excluded).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9618c3e5-3d2d-4a4c-913e-3752b3ef379c
📒 Files selected for processing (16)
.claude/skills/ad-sharding-ir-port/SKILL.mdexamples/auto_deploy/model_registry/models.yamltensorrt_llm/_torch/auto_deploy/models/custom/modeling_cohere.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek_v2.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_exaone.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_gemma.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_gemma2.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_granite.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_hunyuan_dense.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_llama3.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_llama3_ir.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_llama4.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_mistral.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_moe.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_seed_oss.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/auto_deploy/models/custom/modeling_llama3_ir.py
Collaborator
|
PR_Github #51184 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #51637 [ run ] triggered by Bot. Commit: |
Signed-off-by: Grzegorz Kwasniewski <213329731+greg-kwasniewski1@users.noreply.github.com>
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
|
PR_Github #51670 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51637 [ run ] completed with state |
Collaborator
|
PR_Github #51670 [ run ] completed with state
|
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast --post-merge |
Collaborator
|
PR_Github #51846 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51846 [ run ] completed with state
|
…ng files CI runs `pytest .../test_sharding_ir_equivalence.py` with no flags via `test_unittests_v2[unittest/auto_deploy/multigpu/transformations]`. The existing `sharding_ir_modeling_file` fixture skipped the test when `--sharding-ir-modeling-file` was absent, so CI silently exercised zero modeling files -- the test has been dead-code on CI since it landed. Switch the conftest to materialize parametrize off CLI options: - `--sharding-ir-modeling-file` supplied: single-file mode (legacy debug invocation, unchanged behavior). Dist config defaults to `tep` (4 ranks). - both flags absent: auto-discovery mode. Parametrize over every `modeling_*.py` in `tensorrt_llm/_torch/auto_deploy/models/custom/` (37 files today) at the cheapest dist config (`tp-only`, 2 ranks). - `--sharding-ir-dist-config` may pin a single config in either mode. Test ids stay compact: `test_sharding_ir_equivalence[qwen3_next-tp-only]` (modeling_id strips the `modeling_` prefix and `.py` suffix). Failure localization is per-file with one grep. Non-IR-marked modeling files pass as a no-op identity: `apply_sharding_hints` iterates the graph, finds no `torch.ops.auto_deploy.all_reduce` markers and no hint kwargs, leaves the graph untouched -- sharded == unsharded by construction. So the 18 legacy files trivially pass; the 19 IR-marked files exercise the full equivalence path. A hypothetical future regression on any IR port fails immediately, with the modeling file path in the test id. Cost: ~10-15 min on the H100-4_GPUs-AutoDeploy-1 stage at `tp-only` (2 ranks, tiny configs). The full matrix (4 dist configs × 37 files) remains available for explicit per-file debugging. Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast --post-merge |
Resolves conflict in examples/auto_deploy/model_registry/models.yaml.
Main commented out 10 of 11 large-model registry entries (DeepSeek-R1,
DeepSeek-Coder-V2, Qwen3-VL-8B, Qwen2-VL-72B, CodeLlama-70b,
Llama-3.2-90B-Vision, gpt-oss-120b, Llama-4-{Scout,Maverick},
Nemotron-Super-120B-{BF16-BF16KV,A12B}) with documented failure reasons
(OOM / timeout / IndexError NVIDIA#14681 / AttributeError / DistStoreError /
CC "Function Not Found"). Only DeepSeek-R1-0528 stays active.
Take main's version verbatim. Rationale:
- Main's gating reflects current runtime reality. Re-enabling entries
that demonstrably fail in CI would regress everyone.
- Two entries (DeepSeek-Coder-V2, CodeLlama-70b) had `enable_sharder_ir.yaml`
added by this PR; main's commented-out blocks drop that. The opt-in is
moot while the entries are disabled, and NVIDIA#14786 deletes ALL 93
`enable_sharder_ir.yaml` references anyway (the auto-detect dispatcher
makes them redundant), so preserving them here would be churn that gets
reverted.
- The single active entry (DeepSeek-R1-0528) is identical on both sides.
No other files conflict; rest of the merge auto-resolved.
Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #52573 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52573 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #52589 [ run ] triggered by Bot. Commit: |
…r gate Auto-discovery added in this PR parametrizes test_sharding_ir_equivalence over every modeling_*.py. Of 17 failures, 5 fail in spec_from_modeling_file (config class not discoverable / no *ForCausalLM registered), ~7 fail in build_ir_model with tiny config (Padding_idx vs num_embeddings, int vs list intermediate_size, etc.), and ~3 are legacy files with non-IR auto_deploy ops where strip_sharding_hints schema-strips their hint kwargs and spuriously mutates gm_sharded. Two minimal edits handle all three: 1. Parent-process pre-flight: spec_from_modeling_file + build_ir_model on CPU before spawn. Convert config-resolution RuntimeError into pytest.skip (Cat A); convert TypeError/AttributeError/KeyError/ValueError/AssertionError from tiny-config instantiation into pytest.skip (Cat C). Per-rank GPU instantiation in worker stays authoritative. 2. Worker IR-marker gate after torch_export_to_gm: if no torch.ops.auto_deploy.all_reduce nodes, return early. gm_sharded stays an identity copy of gm_unsharded, comparison passes trivially. Matches conftest docstring: "Legacy (non-IR-marked) modeling files pass as a no-op identity (apply_sharding_hints finds no markers and leaves the graph unchanged, so sharded == unsharded by construction)". _has_ir_markers is the in-test equivalent of the follow-up PR's has_sharding_ir_markers; when that PR lands the auto-detect default sharding pipeline, this gate becomes the same predicate that routes legacy files to detect_sharding. Local verification on 5xH100, world_size=2 tp-only: - 13 ported files: 13 PASSED, 0 FAILED. - 17 previously-failing: 5 PASSED, 12 SKIPPED, 0 FAILED. Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #52593 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52589 [ run ] completed with state |
greg-kwasniewski1
added a commit
to nv-auto-deploy/TensorRT-LLM
that referenced
this pull request
Jun 7, 2026
…etect Pre-resolves the NVIDIA#14778 dependency. Conflict resolution: - examples/auto_deploy/model_registry/models.yaml: keep this PR's HEAD side. NVIDIA#14786 makes sharding IR the default + drops enable_sharder_ir.yaml opt-in references; NVIDIA#14778's per-entry opt-ins are obsoleted by this PR. - tests/unittest/auto_deploy/multigpu/transformations/library/conftest.py: combine both PRs: * NVIDIA#14778's _discover_modeling_files() helper + pytest_generate_tests (auto-discovery + matrix parametrize) -- needed for the multigpu CI coverage on every modeling file. * NVIDIA#14786's sharding_ir_eagle_draft fixture -- on-demand Eagle draft test target, distinct from the auto-discovered file matrix. Drop NVIDIA#14786's stale sharding_ir_dist_config @pytest.fixture (now superseded by pytest_generate_tests parametrize on the same name). Restore `import pytest` lost by the import-block auto-merge. Verified: pytest --collect-only collects 39 tests cleanly on the merged conftest. Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Collaborator
|
PR_Github #52593 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #52752 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52752 [ run ] completed with state |
bmarimuthu-nv
approved these changes
Jun 9, 2026
bmarimuthu-nv
left a comment
bmarimuthu-nv
left a comment
Collaborator
There was a problem hiding this comment.
LGTM, thanks @greg-kwasniewski1 !
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.
Summary
Ports 13 AutoDeploy custom modeling files (
tensorrt_llm/_torch/auto_deploy/models/custom/modeling_*.py) to use sharding IR, validated by the offline equivalence test introduced in #13963. Addsenable_sharder_ir.yamlto 70 model_registry entries so they actually exerciseapply_sharding_hintsat runtime — same opt-in pattern PR #13478 used for its 4 families.This is scope-restricted: the default sharding pipeline in
default.yamlstays on legacy (detect_sharding); the flip toapply_sharding_hintsas default is deferred to a follow-up PR (see Future work).Ported models (13)
modeling_gemma.pymodeling_mistral.pymodeling_granite.pymodeling_hunyuan_dense.pymodeling_seed_oss.pymodeling_gemma2.pymodeling_cohere.pymodeling_exaone.pymodeling_llama3.pymodeling_deepseek_v2.pymodeling_glm4_moe.pymodeling_qwen3_moe.pymodeling_llama4.pyAlso deletes
modeling_llama3_ir.py— a stale sharding-IR variant that diverged from canonicalmodeling_llama3.pybefore the rope-handling fix landed; it raisesKeyError: 'default'on current transformers (ROPE_INIT_FUNCTIONS["default"]no longer exists).Changes
all_reduceinsertion → docstring update)..claude/skills/ad-sharding-ir-port/SKILL.md):Sharding strategy:blocks (no advocacy paragraph, no redundantfrom ... import custom_ops).modeling_llama3_ir.py.examples/auto_deploy/model_registry/models.yaml): addsenable_sharder_ir.yamlto 70 entries whose config class maps to one of the 13 ported families. Skipped 13 entries flagged asmultimodal.yamlorsimple_shard_only.yaml.Tests / validation
pytest tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py --sharding-ir-modeling-file <path>) passes 4/4 dist configs (tp-only / ep-only / tep / attn-dp) on every ported file. Re-validation sweep will run alongside CI.SHARDING_IR_SABOTAGE=1) confirmed working on a representative model — it correctly flips PASS to FAIL when collectives are removed.Deferred (out of scope; follow-up PRs)
trust_remote_codeconfigs / padding-idx tiny-config quirks / OTHER_FAIL; 3 structural blockers — starcoder2's attention TP shape, olmo3's flat-projection RMSNorm under TP, gpt_oss's per-headsinksparameter; 5 MoE+multimodal explicitly out of scope per scope decision; 2 baselineOTHER_FAILto investigate separately).default.yamland the newenable_legacy_sharding.yamloverride: realized that flipping the default would regress ~30+ registry entries that map to non-IR-ported modeling files (Phi-4, Qwen2.5, etc.). Deferred until either (a) more models are IR-ported, or (b) the legacy override is wired in everywhere needed.Risk
Library-visible, low risk. Each ported file is exercised by an offline equivalence test that compares sharded vs unsharded outputs on tiny configs across four parallelism layouts. Multi-GPU AutoDeploy CI stages mandatory (see Test plan).
Test plan
DGX_H100-4_GPUs-AutoDeploy-1andDGX_B200-4_GPUs-AutoDeploy-1CI stages green (triggered below)Fixes #14642
Summary by CodeRabbit
Release Notes
New Features
Documentation
Chores