[None][chore] Model update 260308 - #12011
Conversation
Signed-off-by: Tal Cherckez <127761168+tcherckez-nvidia@users.noreply.github.com>
331ade4 to
4fd6537
Compare
Signed-off-by: Tal Cherckez <127761168+tcherckez-nvidia@users.noreply.github.com>
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
/bot skip --comment "model list update" |
📝 WalkthroughWalkthroughThe PR refactors MoE weight handling from per-expert weight lists to fused weight tensors across the TensorRT-LLM auto-deploy pipeline, including model initialization, graph transformations, quantization, and sharding. This enables more efficient weight materialization and parameter management during model export and deployment. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ❌ 3❌ Failed checks (2 warnings, 1 inconclusive)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs). 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: 11
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py (1)
161-166:⚠️ Potential issue | 🟡 MinorRemove the unused
Expertclass.The
Expertclass (lines 161–166) is no longer used byMoEOpModelor anywhere else in the codebase after the refactor to fused parameters. Delete it to eliminate dead code.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py` around lines 161 - 166, The Expert class is dead code after the fused-parameter refactor; remove the unused class definition named Expert from the tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py file so only active models like MoEOpModel remain; simply delete the entire class Expert (constructor and its parameters w1/w2/w3) to eliminate the unused symbol and update any imports or references if they exist elsewhere.tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py (1)
269-282:⚠️ Potential issue | 🔴 CriticalThis tensor-only
torch_moesignature breaks the BMM fusion path.
MatchBmmMoePatternintensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.pystill emitstorch_moe(..., w1_nodes, w2_nodes, w3_nodes)with Python lists. After this change, that path no longer satisfies the op contract and will fail the first time a Llama4-style MoE graph executes.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py` around lines 269 - 282, The new strict tensor-only signature for torch_moe breaks the existing BMM-fusion emitter (MatchBmmMoePattern) which still passes Python lists for w1_nodes/w2_nodes/w3_nodes; change the torch_moe signature to accept sequences of tensors instead of single tensors so it matches emitted calls: update the parameters w1_weight, w2_weight, w3_weight to types like Sequence[torch.Tensor] (or List[torch.Tensor]) and ensure the implementation and the `@torch.library.custom_op` declaration accept and handle Python lists/tuples of tensors so the existing fused_moe emission path (MatchBmmMoePattern) continues to work without changing its outputs.tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py (1)
719-739:⚠️ Potential issue | 🔴 Critical
fuse_nvfp4_moeis still incompatible with these fused inputs.This path now emits
torch_quant_nvfp4_moewith single stackedget_attrtensors, but_stack_nvfp4_moe_weights()later in the same file still does_stack(w1_list, dim=0)/_stack(w2_list, dim=0)on the extracted args. As soon asfuse_nvfp4_moeruns after this transform, it will try to treat a single fused node like an expert list and fail.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py` around lines 719 - 739, The fuse_nvfp4_moe path now emits a single stacked get_attr tensor (torch_quant_nvfp4_moe) but _stack_nvfp4_moe_weights() still assumes it will receive a list and calls torch.stack(list, dim=0), which will break when given an already stacked tensor; update _stack_nvfp4_moe_weights to accept either a pre-stacked tensor or a list by checking the argument type (e.g., isinstance(arg, torch.Tensor) or torch.fx.Node type for get_attr) and returning the tensor as-is when stacked, otherwise perform the existing torch.stack(w_list, dim=0) behavior; ensure callers like fused_moe / fuse_nvfp4_moe and places that extract get_attr continue to work with both forms.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/auto_deploy/models/patches/mixtral.py (1)
115-125: Multiple patches chainAutoModelForCausalLM.from_config— verify import order.This file, along with
qwen3.pyandminimax_m2.py, all patchAutoModelForCausalLM.from_configat module load time. Each stores a reference to the previous implementation before overwriting, which creates a chain.The behavior depends on import order: if
mixtral.pyis imported beforeqwen3.py, the chain isqwen3 → mixtral → original. Eachprepare_fused_moe_blockscorrectly checkstype(module).__name__so only relevant blocks are processed.This pattern works correctly as long as all patch files are imported consistently. Consider documenting this import-order dependency or consolidating the patching logic in a single location if this becomes fragile.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/models/patches/mixtral.py` around lines 115 - 125, Multiple modules (mixtral.py, qwen3.py, minimax_m2.py) patch AutoModelForCausalLM.from_config by chaining wrappers (_from_config_prev/_from_config_patched) which makes behavior import-order dependent; either add a comment near AutoModelForCausalLM.from_config (and in mixtral.py's _from_config_patched) documenting the import-order dependency and that prepare_fused_moe_blocks only affects matching module types, or consolidate all patches into a single central initializer that registers/combines all preparers (calling prepare_fused_moe_blocks and the equivalents from qwen3/minimax_m2) to remove fragile ordering. Ensure references to AutoModelForCausalLM.from_config, _from_config_prev, _from_config_patched, and prepare_fused_moe_blocks are updated accordingly.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py`:
- Around line 703-717: The pre-load hook _load_hook captures loop variables
(w1_name,w2_name,w3_name and w1_keys/w2_keys/w3_keys) causing all hooks to use
the last pattern; change the registration to bind the current pattern values
when registering the hook (e.g., create the hook via a small factory or use
default args) so each gm._register_load_state_dict_pre_hook receives a closure
that references that layer's specific fused_key and expert_keys rather than the
loop-final ones; update the registration site where _load_hook is defined to
capture the current w*_name and w*_keys per-iteration.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/quantize_moe.py`:
- Around line 65-71: The fused-path attribute naming in _quantize_moe_node uses
a default fused_key_counter=0 which causes name collisions; change the default
so callers cannot accidentally reuse 0 (either remove the default or set it to
None) and implement a deterministic unique counter when None (e.g., attach an
incrementing counter to the GraphModule like gm._quantize_fused_counter and
increment it each time you need a new fused_key) and apply the same change to
the other quantize-MoE helper that uses quant_moe_nvfp4_{fused_key_counter}_...
(the block around lines 124-130) so every fused MoE node gets a unique
fused_key_counter instead of always 0.
- Around line 219-235: The fused-quantized path flattens expert scale groups by
using args_list.extend(lst), breaking the expected grouped layout used
downstream; replace that flattening with preserving each expert's scale list
(e.g., use args_list.append for each lst) so scale_node_lists (and names like
scale_buf_names) remain as grouped list arguments passed into
gm.graph.call_function (the new_node creation) matching the per-expert path
layout and keeping downstream sharding logic intact.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py`:
- Around line 1712-1738: The get_partition helper is only defined inside the
per-expert else: branch so when use_fused_weights is true and scale_names is
non-empty the code calls get_partition before it exists; move or hoist the
get_partition definition (the function that computes expert_start/expert_end and
returns (partition, remainder)) out of the per-expert else block so it is
defined alongside _weight_list_from_arg and available to both fused-weight and
per-expert paths; ensure references to get_partition in the scale sharding logic
use this single hoisted function and that its signature matches existing calls
(lst, world_size, rank).
- Around line 1597-1610: The rebuilt call to torch.ops.auto_deploy.torch_moe
currently hardcodes the first six positional args and drops any additional
positional arguments from the original node (so flags like is_gated_mlp, act_fn,
apply_routing_on_input are lost); fix by appending the original trailing
positional arguments when constructing new_node — i.e., build the new args tuple
using args[0], args[1], args[2], gm.graph.get_attr(key_w1),
gm.graph.get_attr(key_w2), gm.graph.get_attr(key_w3) and then extend it with the
remaining original positional args (slice node.args starting at index 6) before
calling gm.graph.call_function(torch.ops.auto_deploy.torch_moe), while still
passing through kwargs.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_modeling_nemotron_h.py`:
- Around line 112-114: The test currently loads only the custom model
(NemotronHForCausalLM) so it compares the custom implementation against itself;
restore the original intent by loading both the HF AutoModelForCausalLM path and
the custom NemotronH path via _load_nemotron_moe_layer: call
_load_nemotron_moe_layer(model_name) to get the HF/reference module and call
_load_nemotron_moe_layer(model_name, custom_model_cls=NemotronHForCausalLM) to
get the custom module, copy the state_dict from the HF/reference module into the
custom module (or vice versa) using load_state_dict, then run the forward passes
and assert outputs between the HF/reference module and the custom module so the
test validates semantic parity rather than comparing the same implementation to
itself.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_moe_weight_loading.py`:
- Around line 27-48: The tests (test_moe_load_pre_stacked_format,
test_moe_load_per_expert_single_format and the similar ones at 141-162)
currently build incoming state_dict values from the same tensors that
setup_fused_moe_weights already populated, so a no-op load_state_dict still
passes; change the tests to provide state_dict values that differ from the
current expert/fused tensors (e.g., create cloned tensors and modify them—add a
small constant or random noise—or use zeros/ones tensors of the same shape)
before calling moe_module.load_state_dict(..., strict=False), then assert that
getattr(moe_module, "_fused_weight") (and similarly "_w_weight"/"_w_bias" where
applicable) equals those provided state_dict tensors; this ensures the MoE load
hook actually updates the fused params rather than relying on pre-population
from setup_fused_moe_weights.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_quant_moe_per_expert_vs_fused.py`:
- Around line 47-52: The reference quantize_per_expert currently casts
per-expert scales to w.dtype which diverges from the fused path; update
quantize_per_expert (and its use of fp8_scale) to keep scales in float32 (e.g.,
ensure scales = torch.stack([...]).clamp(min=eps).to(dtype=torch.float32,
device=device) or simply avoid casting to w.dtype), then use those float32
scales when computing w_q while still producing w_q in dtype_fp8; adjust
references to fp8_scale, E, device, dtype_fp8, and eps accordingly so the scales
remain float32.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_quant_moe.py`:
- Around line 146-147: The test decorator for
test_quantized_moe_fused_path_emits_trtllm_op currently only gates on
fp8_compatible(); update the pytest.mark.skipif condition to also check that the
TRTLLM custom op is available (e.g. change skipif(not fp8_compatible()) to
skipif(not fp8_compatible() or not trtllm_ops_available()) or equivalent), and
apply the same change to the other two test decorators that assert/load
trtllm_quant_fp8_moe_fused so tests are skipped when TRTLLM ops are missing;
reference the test name test_quantized_moe_fused_path_emits_trtllm_op and the
fp8_compatible() guard when locating the spots to update.
---
Outside diff comments:
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py`:
- Around line 269-282: The new strict tensor-only signature for torch_moe breaks
the existing BMM-fusion emitter (MatchBmmMoePattern) which still passes Python
lists for w1_nodes/w2_nodes/w3_nodes; change the torch_moe signature to accept
sequences of tensors instead of single tensors so it matches emitted calls:
update the parameters w1_weight, w2_weight, w3_weight to types like
Sequence[torch.Tensor] (or List[torch.Tensor]) and ensure the implementation and
the `@torch.library.custom_op` declaration accept and handle Python lists/tuples
of tensors so the existing fused_moe emission path (MatchBmmMoePattern)
continues to work without changing its outputs.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py`:
- Around line 719-739: The fuse_nvfp4_moe path now emits a single stacked
get_attr tensor (torch_quant_nvfp4_moe) but _stack_nvfp4_moe_weights() still
assumes it will receive a list and calls torch.stack(list, dim=0), which will
break when given an already stacked tensor; update _stack_nvfp4_moe_weights to
accept either a pre-stacked tensor or a list by checking the argument type
(e.g., isinstance(arg, torch.Tensor) or torch.fx.Node type for get_attr) and
returning the tensor as-is when stacked, otherwise perform the existing
torch.stack(w_list, dim=0) behavior; ensure callers like fused_moe /
fuse_nvfp4_moe and places that extract get_attr continue to work with both
forms.
In `@tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py`:
- Around line 161-166: The Expert class is dead code after the fused-parameter
refactor; remove the unused class definition named Expert from the
tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py file so only
active models like MoEOpModel remain; simply delete the entire class Expert
(constructor and its parameters w1/w2/w3) to eliminate the unused symbol and
update any imports or references if they exist elsewhere.
---
Nitpick comments:
In `@tensorrt_llm/_torch/auto_deploy/models/patches/mixtral.py`:
- Around line 115-125: Multiple modules (mixtral.py, qwen3.py, minimax_m2.py)
patch AutoModelForCausalLM.from_config by chaining wrappers
(_from_config_prev/_from_config_patched) which makes behavior import-order
dependent; either add a comment near AutoModelForCausalLM.from_config (and in
mixtral.py's _from_config_patched) documenting the import-order dependency and
that prepare_fused_moe_blocks only affects matching module types, or consolidate
all patches into a single central initializer that registers/combines all
preparers (calling prepare_fused_moe_blocks and the equivalents from
qwen3/minimax_m2) to remove fragile ordering. Ensure references to
AutoModelForCausalLM.from_config, _from_config_prev, _from_config_patched, and
prepare_fused_moe_blocks are updated accordingly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: a886edb4-78c0-429d-8c1f-5e0ebe1be7a8
📒 Files selected for processing (23)
examples/auto_deploy/model_registry/models.yamltensorrt_llm/_torch/auto_deploy/config/default.yamltensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_nemotron_h.pytensorrt_llm/_torch/auto_deploy/models/patches/common.pytensorrt_llm/_torch/auto_deploy/models/patches/granitemoe.pytensorrt_llm/_torch/auto_deploy/models/patches/minimax_m2.pytensorrt_llm/_torch/auto_deploy/models/patches/mixtral.pytensorrt_llm/_torch/auto_deploy/models/patches/qwen3.pytensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.pytensorrt_llm/_torch/auto_deploy/transform/library/quantize_moe.pytensorrt_llm/_torch/auto_deploy/transform/library/sharding.pytensorrt_llm/_torch/auto_deploy/utils/_graph.pytensorrt_llm/_torch/auto_deploy/utils/node_utils.pytests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.pytests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/moe/test_ad_moe_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_minimax_m2_patches.pytests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_modeling_nemotron_h.pytests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_moe_weight_loading.pytests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_moe_fusion.pytests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_quant_moe.pytests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_quant_moe_per_expert_vs_fused.py
💤 Files with no reviewable changes (3)
- tensorrt_llm/_torch/auto_deploy/config/default.yaml
- tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py
- examples/auto_deploy/model_registry/models.yaml
tcherckez-nvidia
changed the title
|
PR_Github #38155 [ skip ] triggered by Bot. Commit: |
|
PR_Github #38155 [ skip ] completed with state |
Signed-off-by: Tal Cherckez <127761168+tcherckez-nvidia@users.noreply.github.com>
Signed-off-by: Tal Cherckez <127761168+tcherckez-nvidia@users.noreply.github.com>
Signed-off-by: Tal Cherckez <127761168+tcherckez-nvidia@users.noreply.github.com>
Signed-off-by: Tal Cherckez <127761168+tcherckez-nvidia@users.noreply.github.com>
3 participants
Summary by CodeRabbit
New Features
Refactor
Chores
Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.