[https://nvbugs/5496960][fix] Fix Gemma model forward.

[hyukn]

Fix the missing condition for the Gemma V1 model path.

Summary by CodeRabbit

• Bug Fixes
  • Corrects application of normalization only when inter-layer norms are enabled, preventing misconfigured behavior.
  • Aligns epsilon values with the selected configuration for more stable and predictable outputs.
  • Improves numerical consistency across model configurations, reducing variance between runs.
  • Enhances reliability of residual and post-residual normalization, leading to smoother training/inference in supported setups.

[coderabbitai]

📝 Walkthrough

Walkthrough

Gated normalization weights and epsilon for AllReduceParams in GemmaDecoderLayer.forward based on inter_layernorms, added norm_pre_residual_weight from post_layernorm.weight, and defaulted eps to 1e-06 when inter_layernorms is False, altering fusion parameter composition and control flow.

Changes

Cohort / File(s)	Summary
Decoder layer norm fusion gating tensorrt_llm/models/gemma/model.py	In GemmaDecoderLayer.forward: gated norm_weight by inter_layernorms; added norm_pre_residual_weight using post_layernorm.weight.value; set eps to pre_feedforward_layernorm.eps when inter_layernorms is True else 1e-06; updated AllReduceParams construction accordingly.

Sequence Diagram(s)

sequenceDiagram
    participant DL as GemmaDecoderLayer.forward
    participant LN as LayerNorms
    participant AR as AllReduceParams (fusion)

    DL->>LN: Read pre/post layernorm weights
    alt inter_layernorms == True
        Note over DL,AR: Use inter-layer norms
        DL->>AR: Construct with norm_weight, norm_pre_residual_weight
        DL->>AR: eps = pre_feedforward_layernorm.eps
    else inter_layernorms == False
        Note over DL,AR: No inter-layer norms
        DL->>AR: Construct with norm_weight = None
        DL->>AR: norm_pre_residual_weight = post_layernorm.weight.value
        DL->>AR: eps = 1e-06
    end
    DL->>AR: Invoke fused residual RMS pre/post norm
    AR-->>DL: Normalized/Reduced outputs

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

✨ Finishing Touches

• 📝 Generate Docstrings

🧪 Generate unit tests

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

CodeRabbit Commands (Invoked using PR/Issue comments)

Type @coderabbitai help to get the list of available commands.

Other keywords and placeholders

• Add @coderabbitai ignore or @coderabbit ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai or @coderabbitai title anywhere in the PR title to generate the title automatically.

Status, Documentation and Community

• Visit our Status Page to check the current availability of CodeRabbit.
• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[hyukn]

/bot run --disable-fail-fast

[coderabbitai]

Actionable comments posted: 0

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/models/gemma/model.py (1)

193-195: Guard post_feedforward_layernorm and fall back to post_layernorm in the fused‐MLP branch.

Replace in tensorrt_llm/models/gemma/model.py (around line 216):

-    norm_pre_residual_weight=self.post_feedforward_layernorm.weight.value,
+    norm_pre_residual_weight=(
+        self.post_feedforward_layernorm.weight.value
+        if self.config.inter_layernorms
+        else self.post_layernorm.weight.value
+    ),

🧹 Nitpick comments (1)

tensorrt_llm/models/gemma/model.py (1)

163-165: Prefer config-derived epsilon over hard-coded 1e-06.

Use the existing LN’s eps (or config.norm_epsilon) to keep behavior consistent with model config and avoid silent drift if defaults change.

Apply:

-                eps=self.pre_feedforward_layernorm.eps
-                if self.config.inter_layernorms else 1e-06,
+                eps=(self.pre_feedforward_layernorm.eps
+                     if self.config.inter_layernorms
+                     else self.post_layernorm.eps),

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

📥 Commits

Reviewing files that changed from the base of the PR and between 9eecdf2 and ad22d15.

📒 Files selected for processing (1)

• tensorrt_llm/models/gemma/model.py (1 hunks)

🧰 Additional context used

📓 Path-based instructions (4)

**/*

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Filenames compiled into a target must be case-insensitively unique

Files:

• tensorrt_llm/models/gemma/model.py

**/*.{h,hpp,hh,hxx,cc,cpp,cxx,cu,cuh,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Use spaces, not tabs; indent 4 spaces

Files:

• tensorrt_llm/models/gemma/model.py

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: Code must target Python 3.8+
Indent with 4 spaces; do not use tabs (Python)
Maintain module namespace on import: prefer from package.subpackage import foo; use foo.Symbol()
Python filenames use snake_case
Python class names use PascalCase
Python functions and methods use snake_case
Python local variables use snake_case; if starting with a number concept, prefix with k (e.g., k_99th_percentile)
Python global variables use G_ prefix with UPPER_SNAKE_CASE
Python constants use UPPER_SNAKE_CASE
Avoid shadowing variables from outer scopes
Initialize all externally visible class members in init
For public interfaces, prefer docstrings over comments; comments should be for in-function or file-local interfaces
Use Google-style docstrings for classes and functions (Sphinx-parsable)
Document attributes and variables inline with docstrings immediately after assignment
Avoid reflection when a non-reflective approach suffices
Limit except clauses to specific exceptions where possible
When using try/except for duck-typing, keep try body minimal and move logic to else

Files:

• tensorrt_llm/models/gemma/model.py

**/*.{cpp,cc,cxx,h,hpp,hh,hxx,cu,cuh,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Prepend NVIDIA copyright header (current year) to all source files

Files:

• tensorrt_llm/models/gemma/model.py

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

• GitHub Check: Pre-commit Check

🔇 Additional comments (2)

tensorrt_llm/models/gemma/model.py (2)

160-162: Good catch: gate norm_weight by inter_layernorms.

Passing None for norm_weight on Gemma V1 avoids feeding a non-existent pre-FFN LN into the fusion path. This aligns with the PR goal.

---

140-146: Resolve FIXME around residual tuple unpacking in reduce_fusion path.

Please verify that hidden_states carries (hidden, residual) for local_layer_idx > 0 and that residual semantics match the fusion op’s expectations.

I can help trace call sites that set next_layer_input_layernorm_args/residual tuples if you share the relevant files.

[hyukn]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #17583 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #17583 [ run ] completed with state SUCCESS
/LLM/release-1.0/L0_MergeRequest_PR pipeline #347 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.