[TRTLLM-13429][feat] Switch DeepSeek/NemotronH/Qwen3/Qwen3.5-MoE to sharding-IR canonical models#13478
[TRTLLM-13429][feat] Switch DeepSeek/NemotronH/Qwen3/Qwen3.5-MoE to sharding-IR canonical models#13478greg-kwasniewski1 wants to merge 2 commits intoNVIDIA:mainfrom
Conversation
|
/bot run --stage-list "A10-Build_Docs, A10-PackageSanityCheck-PY310-UB2204, A100X-PackageSanityCheck-PY312-UB2404, A30-AutoDeploy-1, H100_PCIe-AutoDeploy-1, DGX_B200-AutoDeploy-1, A100X-PyTorch-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_B200-4_GPUs-AutoDeploy-1" |
📝 WalkthroughWalkthroughThis change consolidates four model families from dual-file implementations to IR-fied single files. The non-IR implementations gain sharding-aware custom ops carrying explicit sharding metadata, the parallel Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 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}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
.claude/skills/ad-model-onboard/SKILL.md (1)
556-559:⚠️ Potential issue | 🟠 MajorResolve conflicting guidance on
torch.ops.trtllm.*usage.Line 556 requires preserving specific
torch.ops.trtllm.*router ops, but Line 558 says to never usetrtllm_*. This contradiction can lead to incorrect ports. Please carve out an explicit exception in the “Key Gotchas” rule for the noaux_tc/dsv3 router case.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In @.claude/skills/ad-model-onboard/SKILL.md around lines 556 - 559, The "Key Gotchas" rule contradicts earlier guidance about preserving specific router ops; add an explicit exception stating that for noaux_tc / DeepSeek-V3 style routers you must keep torch.ops.trtllm.noaux_tc_op and torch.ops.trtllm.dsv3_router_gemm_op exactly as in the non-IR base (router gate is TP-REPLICATED, no sharding hints), while retaining the general prohibition on trtllm_* in AutoDeploy (i.e., only allow these two trtllm ops when the source model uses them verbatim and do not introduce vanilla replacements), and update the SKILL.md section text to reflect this exception so both rules are consistent.tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py (1)
9-37:⚠️ Potential issue | 🔴 CriticalAdd
modeling_llama3_irto_MODEL_MODULESor restore theAD_USE_IR_MODELSgate.The file
modeling_llama3_ir.pyexists but is not discoverable through the custom models__init__.py— it's neither listed in_MODEL_MODULESnor guarded by theAD_USE_IR_MODELSenvironment variable that tests still set. The module self-registers viaAutoModelForCausalLMFactory, but removal of the environment-gated registration breaks the intended conditional loading mechanism referenced in the tests.Either add
"modeling_llama3_ir": ["Llama3ForCausalLM"]to_MODEL_MODULES(consistent with other converted IR variants like deepseek, nemotron_h, qwen3, and qwen3_5_moe), or restore the conditional gate if IR-only loading is required.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py` around lines 9 - 37, The custom models registry is missing the modeling_llama3_ir entry so modeling_llama3_ir (which defines Llama3ForCausalLM and self-registers via AutoModelForCausalLMFactory) isn't discoverable; fix by either adding "modeling_llama3_ir": ["Llama3ForCausalLM"] to the _MODEL_MODULES dict in __init__.py (so the importlib loop loads it like modeling_deepseek/modeling_qwen3 entries) or restore the AD_USE_IR_MODELS environment-gate around the registration/import of modeling_llama3_ir (reintroduce the AD_USE_IR_MODELS check used by tests) so the module is conditionally registered as intended.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In @.claude/skills/ad-model-onboard/SKILL.md:
- Around line 503-509: The fenced code block that begins with the table header
"| Hunk lines (in IR file) | Summary of change | Category | Verdict |" is
missing a language tag which triggers MD040; update that fenced block by adding
a language identifier (e.g., "markdown") right after the opening backticks so
the block becomes ```markdown, ensuring the table renders and linters stop
flagging it; locate the block by searching for the exact table header text
within SKILL.md.
In `@tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py`:
- Around line 1-4: This file is missing the required NVIDIA copyright header at
the top; add the standard NVIDIA source-file header (with updated year if this
is a modification) as the very first lines of
tensorrt_llm._torch.auto_deploy.models.custom.__init__ before any imports, then
keep the existing imports and _logger definition unchanged so symbols like
importlib, logging and _logger remain intact.
In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py`:
- Line 1: Update the copyright header year from 2025 to 2026 in the file
modeling_deepseek.py: locate the top-of-file copyright comment and change the
year range to include 2026 so the header reflects the file modification in 2026.
---
Outside diff comments:
In @.claude/skills/ad-model-onboard/SKILL.md:
- Around line 556-559: The "Key Gotchas" rule contradicts earlier guidance about
preserving specific router ops; add an explicit exception stating that for
noaux_tc / DeepSeek-V3 style routers you must keep torch.ops.trtllm.noaux_tc_op
and torch.ops.trtllm.dsv3_router_gemm_op exactly as in the non-IR base (router
gate is TP-REPLICATED, no sharding hints), while retaining the general
prohibition on trtllm_* in AutoDeploy (i.e., only allow these two trtllm ops
when the source model uses them verbatim and do not introduce vanilla
replacements), and update the SKILL.md section text to reflect this exception so
both rules are consistent.
In `@tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py`:
- Around line 9-37: The custom models registry is missing the modeling_llama3_ir
entry so modeling_llama3_ir (which defines Llama3ForCausalLM and self-registers
via AutoModelForCausalLMFactory) isn't discoverable; fix by either adding
"modeling_llama3_ir": ["Llama3ForCausalLM"] to the _MODEL_MODULES dict in
__init__.py (so the importlib loop loads it like
modeling_deepseek/modeling_qwen3 entries) or restore the AD_USE_IR_MODELS
environment-gate around the registration/import of modeling_llama3_ir
(reintroduce the AD_USE_IR_MODELS check used by tests) so the module is
conditionally registered as intended.
🪄 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: c69a2cab-991f-4250-b5a0-f23027751349
📒 Files selected for processing (9)
.claude/skills/ad-model-onboard/SKILL.mdtensorrt_llm/_torch/auto_deploy/models/custom/__init__.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek_ir.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_nemotron_h.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_nemotron_h_ir.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe_ir.py
💤 Files with no reviewable changes (2)
- tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek_ir.py
- tensorrt_llm/_torch/auto_deploy/models/custom/modeling_nemotron_h_ir.py
| ``` | ||
| | Hunk lines (in IR file) | Summary of change | Category | Verdict | | ||
| |---|---|---|---| | ||
| | 234-240 | F.linear → torch_linear_simple, tp_mode="colwise" | A1 + A2 | OK | | ||
| | 264-340 | noaux_tc_op replaced with vanilla PyTorch | F1 | REVERTED to base | | ||
| | ... | ... | ... | ... | | ||
| ``` |
There was a problem hiding this comment.
Add a language tag to the fenced code block.
Line 503 starts a fenced block without a language, which triggers markdownlint MD040 and weakens rendering/tooling support.
💡 Proposed fix
-```
+```markdown
| Hunk lines (in IR file) | Summary of change | Category | Verdict |
|---|---|---|---|
| 234-240 | F.linear → torch_linear_simple, tp_mode="colwise" | A1 + A2 | OK |
| 264-340 | noaux_tc_op replaced with vanilla PyTorch | F1 | REVERTED to base |
| ... | ... | ... | ... |</details>
<details>
<summary>🧰 Tools</summary>
<details>
<summary>🪛 markdownlint-cli2 (0.22.1)</summary>
[warning] 503-503: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
</details>
</details>
<details>
<summary>🤖 Prompt for AI Agents</summary>
Verify each finding against the current code and only fix it if needed.
In @.claude/skills/ad-model-onboard/SKILL.md around lines 503 - 509, The fenced
code block that begins with the table header "| Hunk lines (in IR file) |
Summary of change | Category | Verdict |" is missing a language tag which
triggers MD040; update that fenced block by adding a language identifier (e.g.,
"markdown") right after the opening backticks so the block becomes ```markdown,
ensuring the table renders and linters stop flagging it; locate the block by
searching for the exact table header text within SKILL.md.
</details>
<!-- fingerprinting:phantom:triton:hawk:c9d0c262-c58a-444d-85ee-34a1dfb042e6 -->
<!-- d98c2f50 -->
<!-- This is an auto-generated comment by CodeRabbit -->
| import importlib | ||
| import logging | ||
| import os | ||
|
|
||
| _logger = logging.getLogger(__name__) |
There was a problem hiding this comment.
Add the required NVIDIA copyright header.
This modified Python source still starts directly with imports, so it no longer satisfies the repo’s source-file header requirement.
♻️ Proposed fix
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
import importlib
import logging📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| import importlib | |
| import logging | |
| import os | |
| _logger = logging.getLogger(__name__) | |
| # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. | |
| import importlib | |
| import logging | |
| _logger = logging.getLogger(__name__) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py` around lines 1 -
4, This file is missing the required NVIDIA copyright header at the top; add the
standard NVIDIA source-file header (with updated year if this is a modification)
as the very first lines of
tensorrt_llm._torch.auto_deploy.models.custom.__init__ before any imports, then
keep the existing imports and _logger definition unchanged so symbols like
importlib, logging and _logger remain intact.
| @@ -1,18 +1,52 @@ | |||
| # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | |||
There was a problem hiding this comment.
Update the copyright year.
This file is modified in 2026, but the header still stops at 2025.
♻️ Proposed fix
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | |
| # Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py` at line
1, Update the copyright header year from 2025 to 2026 in the file
modeling_deepseek.py: locate the top-of-file copyright comment and change the
year range to include 2026 so the header reflects the file modification in 2026.
|
PR_Github #45593 [ run ] triggered by Bot. Commit: |
|
/bot run --stage-list "A10-Build_Docs, A10-PackageSanityCheck-PY310-UB2204, A100X-PackageSanityCheck-PY312-UB2404, A30-AutoDeploy-1, H100_PCIe-AutoDeploy-1, DGX_B200-AutoDeploy-1, A100X-PyTorch-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_B200-4_GPUs-AutoDeploy-1" |
|
PR_Github #45594 [ run ] triggered by Bot. Commit: |
|
PR_Github #45594 [ run ] completed with state
|
|
/bot run --stage-list "A10-Build_Docs, A10-PackageSanityCheck-PY310-UB2204, A100X-PackageSanityCheck-PY312-UB2404, A30-AutoDeploy-1, H100_PCIe-AutoDeploy-1, DGX_B200-AutoDeploy-1, A100X-PyTorch-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_B200-4_GPUs-AutoDeploy-1" |
|
PR_Github #45730 [ run ] triggered by Bot. Commit: |
|
PR_Github #45730 [ run ] completed with state
|
|
/bot run --stage-list "A10-Build_Docs, A10-PackageSanityCheck-PY310-UB2204, A100X-PackageSanityCheck-PY312-UB2404, A30-AutoDeploy-1, H100_PCIe-AutoDeploy-1, DGX_B200-AutoDeploy-1, A100X-PyTorch-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_B200-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #45765 [ run ] triggered by Bot. Commit: |
… to sharding-IR canonical Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> Made-with: Cursor
…py staging Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> Made-with: Cursor
greg-kwasniewski1
force-pushed
the
gk/switch_to_ir_models
branch
from
46f9725 to
72cae82
Compare
|
/bot run --stage-list "A10-Build_Docs, A10-PackageSanityCheck-PY310-UB2204, A100X-PackageSanityCheck-PY312-UB2404, A30-AutoDeploy-1, H100_PCIe-AutoDeploy-1, DGX_B200-AutoDeploy-1, A100X-PyTorch-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_B200-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #45771 [ run ] triggered by Bot. Commit: |
2 participants
Summary
Fixes #13429.
Replace the legacy non-IR
modeling_*.pyfor four model architectures with their sharding-IR variants. The IR-aware implementation (usingapply_sharding_hintsfor TP/EP/BMM via canonical-op kwargs liketp_mode,layer_type,output_sizes,tp_min_local_shape,tp_scaled_dim,shardable) is now the canonical implementation. TheAD_USE_IR_MODELSenv-var gate is removed.File swap
git rm modeling_X.py && git mv modeling_X_ir.py modeling_X.pyfor:modeling_deepseek.py— DeepSeek V3, MLA + MoEmodeling_nemotron_h.py— NemotronH, hybrid Mamba + MoEmodeling_qwen3.py— Qwen3 dense (was previously only available as_ir; now registered in__init__.py)modeling_qwen3_5_moe.py— Qwen3.5 MoE, GatedDeltaNet + Gated MHA + MoEContinuation of #12419 (sharding-IR introduction) and #13272 (post-#12419 cleanup), per the staged plan to migrate model implementations to the hint-driven sharder one cohort at a time.
Skill update for future regenerations
.claude/skills/ad-model-onboard/SKILL.md:noaux_tc"Key Gotcha" bullet — was telling agents to substitute fused trtllm calls with vanilla PyTorch; corrected to KEEPtorch.ops.trtllm.{noaux_tc_op, dsv3_router_gemm_op}verbatim (no AD transform recovers vanilla → fused; vanilla replacement is ~17x more kernel launches and loses the FP8-friendly router GEMM).diff -uand classify every hunk before reporting done.position_idsrule now carries an "IR-port exception" pointing at the contract.Audit summary
Each of the four IR files was regenerated from its non-IR base by parallel
generalPurposesubagents using the updated skill, then audited at the parent level. Every hunk classifies as A1-A5 (allowed). Cross-checks: sameclasscount, samenn.Parameter/register_buffer/register_load_state_dict_pre_hookcount, samedef forward(count, bothtorch.ops.trtllm.*call sites inDeepSeekV3MoEGateandNemotronHTopkRouterpreserved verbatim (thenoaux_tcpaths are the original fused kernels), and theposition_idscontract preserved per architecture.Test plan
small/fp8/deepseek/r1.yamlwith the IR sharder +world_size=2on 2x H100: exit=0, 51 sharding nodes, 10 prompts generated cleanly. (AD_USE_IR_MODELS=1no longer needed since IR is canonical.)ci-guidelines.md) — will trigger via/bot run --stage-list ...after PR opens./bot runonce the targeted run is green.Refs: #12419 (sharding-IR introduction), #13271 (post-#12419 cleanup feature), #13272 (post-#12419 cleanup PR), #13429 (this feature's tracking issue).
Made with Cursor
Summary by CodeRabbit
New Features
Refactor
Documentation