[https://nvbugs/6379636][fix] Fix Gemma4 MoE weight loading by 2ez4bz · Pull Request #16108 · NVIDIA/TensorRT-LLM

2ez4bz · 2026-07-08T06:40:34Z

Summary by CodeRabbit

New Features
- Added support for Gemma 4 NVFP4 checkpoint formats in expert weight mapping, improving compatibility with newer model exports.
- Expanded multimodal accuracy coverage for Gemma 4 on supported hardware.
Bug Fixes
- Improved handling of unsupported expert tensor layouts by raising an error instead of silently accepting them.
- Updated accuracy baselines for the Gemma 4 NVFP4 multimodal configuration.

Description

This commit fixes some issues in MoE weight loading for Gemma4, and adds an E2E accuracy test for it.

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)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
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.

coderabbitai · 2026-07-08T06:46:29Z

📝 Walkthrough

Walkthrough

Adds NVFP4 per-expert MoE checkpoint key remapping to Gemma4HfWeightMapper with validation for unsupported expert keys, corresponding unit tests, a new B200-only NVFP4 accuracy integration test for Gemma4-26B-A4B, an updated MMMU accuracy reference, and a B200 test list entry.

Changes

Gemma4 NVFP4 MoE remapping

Layer / File(s)	Summary
NVFP4 expert key remap and validation `tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py`	Adds `expert_projection_map` and regex-based matching to translate NVFP4 per-expert `gate/up/down_proj` tensor keys into TRT-LLM `moe.experts.{id}.{w1,w3,w2}` naming, and raises `ValueError` for unrecognized `experts.*` keys.
Weight mapper unit tests `tests/unittest/_torch/modeling/test_modeling_gemma4.py`	Removes transformers version-gating skip logic and adds `TestGemma4HfWeightMapper` covering NVFP4, VLM prefix, fused BF16 layouts, and unsupported-key error handling.

Gemma4 NVFP4 B200 accuracy test

Layer / File(s)	Summary
Updated accuracy reference `tests/integration/defs/accuracy/references/mmmu.yaml`	Replaces the prior accuracy value with a quant-specific NVFP4/FP8 KV cache baseline for google/gemma-4-26B-A4B-it.
New accuracy test and test list entry `tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py`, `tests/integration/test_lists/test-db/l0_b200.yml`	Adds `TestGemma4_26B_A4B::test_nvfp4`, a B200-only test asserting NVFP4 quant config and evaluating MMMU accuracy, and registers it in the B200 post_merge test list.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Suggested reviewers: jieli-matrix, yechank-nvidia, kaiyux

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description covers the main goal, but the Test Coverage section is empty and lacks specific tests for the new Gemma4 paths.	Add the relevant tests, such as the Gemma4 unit and integration tests, and briefly state what each one validates.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title matches the main change and includes the NVBugs link and fix type.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

🧹 Nitpick comments (1)

tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py (1)
404-406: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Mutable dict class attribute (RUF012).

EXTRA_EVALUATOR_KWARGS is a mutable dict literal assigned directly as a class attribute; static analysis flags this as RUF012. Low risk here since it's read-only, but consider annotating as ClassVar[dict] to make the immutability intent explicit and silence the lint warning.
🤖 Prompt for 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.

In `@tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py` around
lines 404 - 406, `EXTRA_EVALUATOR_KWARGS` is a mutable dict class attribute that
triggers RUF012; update the test class in `test_llm_api_pytorch_multimodal` to
make the intent explicit by annotating this attribute as a `ClassVar[dict]` (or
equivalent typed class-level constant) so the linter recognizes it as
intentionally shared and read-only.
Source: Linters/SAST tools

🤖 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.

Nitpick comments:
In `@tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py`:
- Around line 404-406: `EXTRA_EVALUATOR_KWARGS` is a mutable dict class
attribute that triggers RUF012; update the test class in
`test_llm_api_pytorch_multimodal` to make the intent explicit by annotating this
attribute as a `ClassVar[dict]` (or equivalent typed class-level constant) so
the linter recognizes it as intentionally shared and read-only.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 1ac85bc0-ef63-4fbf-8084-1f07799d0709

📥 Commits

Reviewing files that changed from the base of the PR and between fd271e0 and 0c6435b.

📒 Files selected for processing (5)

tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py
tests/integration/defs/accuracy/references/mmmu.yaml
tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py
tests/integration/test_lists/test-db/l0_b200.yml
tests/unittest/_torch/modeling/test_modeling_gemma4.py

Hudayday

LGTM

lfr-0531

Added a minor comment, but LGTM overall.

2ez4bz · 2026-07-08T17:16:17Z

/bot run

tensorrt-cicd · 2026-07-08T17:22:36Z

PR_Github #58269 [ run ] triggered by Bot. Commit: 0c6435b Link to invocation

tburt-nv · 2026-07-08T20:48:44Z

/bot run

tensorrt-cicd · 2026-07-08T20:54:33Z

PR_Github #58298 [ run ] triggered by Bot. Commit: 0c6435b Link to invocation

tensorrt-cicd · 2026-07-08T20:54:46Z

PR_Github #58269 [ run ] completed with state ABORTED. Commit: 0c6435b

Link to invocation

tensorrt-cicd · 2026-07-09T02:12:57Z

PR_Github #58298 [ run ] completed with state SUCCESS. Commit: 0c6435b
/LLM/main/L0_MergeRequest_PR pipeline #46933 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

2ez4bz · 2026-07-09T04:49:03Z

/bot run

tensorrt-cicd · 2026-07-09T04:54:44Z

PR_Github #58372 [ run ] triggered by Bot. Commit: 0c6435b Link to invocation

tensorrt-cicd · 2026-07-09T05:51:38Z

PR_Github #58372 [ run ] completed with state FAILURE. Commit: 0c6435b
/LLM/main/L0_MergeRequest_PR pipeline #46998 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

2ez4bz · 2026-07-09T15:59:18Z

/bot run

tensorrt-cicd · 2026-07-09T16:05:51Z

PR_Github #58487 [ run ] triggered by Bot. Commit: 908f825 Link to invocation

tensorrt-cicd · 2026-07-09T16:22:27Z

PR_Github #58487 [ run ] completed with state FAILURE. Commit: 908f825
/LLM/main/L0_MergeRequest_PR pipeline #47096 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

2ez4bz · 2026-07-09T20:00:25Z

/bot run

tensorrt-cicd · 2026-07-09T20:08:57Z

PR_Github #58508 [ run ] triggered by Bot. Commit: 908f825 Link to invocation

2ez4bz · 2026-07-09T23:19:15Z

/bot run

tensorrt-cicd · 2026-07-09T23:25:29Z

PR_Github #58529 [ run ] triggered by Bot. Commit: c284593 Link to invocation

tensorrt-cicd · 2026-07-09T23:25:46Z

PR_Github #58508 [ run ] completed with state ABORTED. Commit: 908f825
/LLM/main/L0_MergeRequest_PR pipeline #47115 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

This commit fixes some issues in MoE weight loading for Gemma4, and adds an E2E accuracy test for it. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>

2ez4bz · 2026-07-10T03:47:29Z

/bot run

tensorrt-cicd · 2026-07-10T03:53:02Z

PR_Github #58583 [ run ] triggered by Bot. Commit: bab1853 Link to invocation

tensorrt-cicd · 2026-07-10T03:53:06Z

PR_Github #58529 [ run ] completed with state ABORTED. Commit: c284593
/LLM/main/L0_MergeRequest_PR pipeline #47131 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

2ez4bz · 2026-07-10T06:16:49Z

/bot run

tensorrt-cicd · 2026-07-10T06:22:44Z

PR_Github #58615 [ run ] triggered by Bot. Commit: bab1853 Link to invocation

tensorrt-cicd · 2026-07-10T06:22:49Z

PR_Github #58583 [ run ] completed with state ABORTED. Commit: bab1853
/LLM/main/L0_MergeRequest_PR pipeline #47175 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

tensorrt-cicd · 2026-07-10T12:00:07Z

PR_Github #58615 [ run ] completed with state SUCCESS. Commit: bab1853
/LLM/main/L0_MergeRequest_PR pipeline #47205 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

dongfengy · 2026-07-11T06:26:22Z

@2ez4bz Heads-up: TestGemma4MoE::test_bf16 added here appears to fail deterministically on the DGX_H100-4_GPUs-AutoDeploy-1 shard of every recent L0_Test-x86_64-Multi-GPU run:

FAILED accuracy/test_llm_api_autodeploy.py::TestGemma4MoE::test_bf16 - ValueError: Not registered specs: {'dtype': 'auto', 'quant_algo': None, 'kv_cache_quant_algo': None, 'spec_dec_algo': None, 'extra_acc_spec': None}

Cause: references/mmmu.yaml registers only a quant_algo: NVFP4 / kv_cache_quant_algo: FP8 entry for google/gemma-4-26B-A4B-it, but test_bf16 evaluates MMMU with the plain bf16 spec, which has no registered entry, so the accuracy harness raises before scoring.

Observed identically in at least four unrelated PR runs (builds #2327/#2328/#2329/#2330 of LLM/main/L0_Test-x86_64-Multi-GPU), so it currently blocks pre-merge CI for any PR that triggers the multi-GPU stage. Could you add the bf16 MMMU reference or waive the test? Thanks!

🤖 Generated with Claude Code

…nce) accuracy/test_llm_api_autodeploy.py::TestGemma4MoE::test_bf16 evaluates MMMU with the plain bf16 spec, but references/mmmu.yaml only registers a NVFP4/FP8-KV entry for google/gemma-4-26B-A4B-it (added in NVIDIA#16108), so the accuracy harness raises ValueError (Not registered specs) before scoring. The failure is deterministic and PR-independent: it failed identically in every recent LLM/main/L0_Test-x86_64-Multi-GPU run (builds 2327-2336), blocking pre-merge CI for any PR that triggers the multi-GPU stage. Waive until the bf16 MMMU reference is added. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>

…26B-A4B-it PR NVIDIA#16108 replaced the plain unquantized entry for google/gemma-4-26B-A4B-it with only an NVFP4 entry, which broke TestGemma4MoE::test_bf16: the autodeploy bf16 test looks up the model with no quant (defaults: dtype='auto', quant_algo=None, kv_cache_quant_algo=None, ...) and no reference row matched, raising ValueError: Not registered specs. Add the pre-NVIDIA#16108 bf16 baseline (56.667) back alongside the NVFP4 entry so both TestGemma4MoE::test_bf16 (H100 autodeploy, bf16) and the new PyTorch NVFP4 test resolve their thresholds. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>

…racy: 56.667` alongside the existing… (#16310) Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>

2ez4bz requested review from a team as code owners July 8, 2026 06:40

2ez4bz requested review from aswinvisva and yechank-nvidia July 8, 2026 06:40

github-actions Bot assigned 2ez4bz Jul 8, 2026

coderabbitai Bot reviewed Jul 8, 2026

View reviewed changes

Hudayday approved these changes Jul 8, 2026

View reviewed changes

lfr-0531 approved these changes Jul 8, 2026

View reviewed changes

Comment thread tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py

litaotju approved these changes Jul 8, 2026

View reviewed changes

2ez4bz enabled auto-merge (squash) July 8, 2026 15:51

2ez4bz changed the title ~~[None][fix] Fix Gemma4 MoE weight loading~~ [https://6379636][fix] Fix Gemma4 MoE weight loading Jul 8, 2026

2ez4bz changed the title ~~[https://6379636][fix] Fix Gemma4 MoE weight loading~~ [https://nvbugs/6379636][fix] Fix Gemma4 MoE weight loading Jul 8, 2026

2ez4bz force-pushed the dev-gemma4-moe-fix branch from 0c6435b to 908f825 Compare July 9, 2026 15:55

2ez4bz force-pushed the dev-gemma4-moe-fix branch from 908f825 to c284593 Compare July 9, 2026 23:19

[None][fix] Fix Gemma4 MoE weight loading

bab1853

This commit fixes some issues in MoE weight loading for Gemma4, and adds an E2E accuracy test for it. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>

2ez4bz force-pushed the dev-gemma4-moe-fix branch from c284593 to bab1853 Compare July 10, 2026 03:47

2ez4bz merged commit fd9166c into NVIDIA:main Jul 10, 2026
8 checks passed

dongfengy mentioned this pull request Jul 12, 2026

[None][test] Waive TestGemma4MoE::test_bf16 (missing bf16 MMMU reference) #16289

Closed

1 task

trtllm-agent mentioned this pull request Jul 13, 2026

[https://nvbugs/6445322][fix] Restore the pre-#16108 baseline - accuracy: 56.667 alongside the existing… #16310

Merged

2 tasks

2ez4bz pushed a commit that referenced this pull request Jul 14, 2026

[https://nvbugs/6445322][fix] Restore the pre-#16108 baseline `- accu…

e9f101b

…racy: 56.667` alongside the existing… (#16310) Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>

Uh oh!

Conversation

2ez4bz commented Jul 8, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

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

Uh oh!

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Hudayday left a comment

Choose a reason for hiding this comment

Uh oh!

lfr-0531 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

2ez4bz commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

tburt-nv commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

2ez4bz commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

2ez4bz commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

2ez4bz commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

2ez4bz commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

2ez4bz commented Jul 10, 2026

Uh oh!

tensorrt-cicd commented Jul 10, 2026

Uh oh!

tensorrt-cicd commented Jul 10, 2026

Uh oh!

2ez4bz commented Jul 10, 2026

Uh oh!

tensorrt-cicd commented Jul 10, 2026

Uh oh!

tensorrt-cicd commented Jul 10, 2026

Uh oh!

tensorrt-cicd commented Jul 10, 2026

Uh oh!

Uh oh!

dongfengy commented Jul 11, 2026

Uh oh!

Reviewers

Assignees

Labels

2ez4bz commented Jul 8, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 8, 2026 •

edited

Loading