Skip to content

Feature/gemma4 support#1899

Draft
cjac wants to merge 4 commits into
NVIDIA:mainfrom
LLC-Technologies-Collier:feature/gemma4-support
Draft

Feature/gemma4 support#1899
cjac wants to merge 4 commits into
NVIDIA:mainfrom
LLC-Technologies-Collier:feature/gemma4-support

Conversation

@cjac

@cjac cjac commented Jul 3, 2026

Copy link
Copy Markdown

What does this PR do?

Type of change: New feature, new tests

This PR introduces native support for the Gemma 4 model family (heterogeneous MoE/dense architecture) into the PyTorch export and TensorRT-LLM conversion pipelines.

Specifically, it:

  • Registers the "Gemma4" model name and "Gemma4ForCausalLM" HuggingFace architecture.
  • Updates build_decoder_config to correctly map Gemma 4's multi-layernorm architecture (post_attention_layernorm, pre_feedforward_layernorm, post_feedforward_layernorm).
  • Aligns gemma4 routing to leverage the upstream data-driven hf_config_map.py for cleanly propagating model properties like attn_logit_softcapping, final_logit_softcapping, and query_pre_attn_scalar.
  • Adds Google LLC and contributors copyright attributions to the modified files.

Usage

The Gemma 4 architecture is now supported natively in the existing TensorRT-LLM export flow.

from transformers import AutoModelForCausalLM
from modelopt.torch.export import export_tensorrt_llm_checkpoint

# Load the Gemma 4 model
model = AutoModelForCausalLM.from_pretrained("google/gemma-4-26b-it", torch_dtype="auto")

# Export to TensorRT-LLM checkpoint format
export_tensorrt_llm_checkpoint(
    model, 
    export_dir="/path/to/output_trt_ckpt", 
    decoder_type="gemma4"
)

Testing

Added a dedicated unit test suite: tests/unit/torch/export/test_export_gemma4.py.

To avoid requiring multi-gigabyte HuggingFace checkpoints during CI runs, the test suite utilizes structurally accurate nn.Module mocks (DummyGemma4Layer, DummyRMSNorm, DummyAttention) to validate the exporter's type-checking, layernorm mapping, and config translation pathways in milliseconds.

Tested locally via pytest tests/unit/torch/export/test_export_gemma4.py, passing 100% green.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

  • Is this change backward compatible?: ✅
  • If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
  • Did you write any new necessary tests?: ✅
  • Did you update Changelog?: ❌ (Will update if requested during review)
  • Did you get Claude approval on this PR?: N/A

Additional Information

This unblocks downstream dynamic ONNX graph construction and TensorRT engine compilation for the Gemma 4 MoE model architectures on edge deployments.

cjac and others added 4 commits July 3, 2026 14:15
…nd conversion pipelines

Extend the PyTorch export and TensorRT-LLM conversion utilities inside `nvidia-modelopt` to natively support the Gemma 4 (heterogeneous MoE/dense) model family.

These changes align the ModelOpt checkpoint translation logic with Gemma 4's multi-layernorm architecture and logit softcapping properties, preventing export-time schema crashes.

1. **`modelopt/torch/export/layer_utils.py`**:
   - Integrated `"gemma4"` into layernorm node routing.
   - Configured appropriate mappings for `"post_attention_layernorm"`, `"pre_feedforward_layernorm"`, and `"post_feedforward_layernorm"` blocks.

2. **`modelopt/torch/export/model_utils.py`**:
   - Registered `"Gemma4": "gemma4"` mapping in the global model-to-type map `MODEL_NAME_TO_TYPE`.

3. **`modelopt/torch/export/tensorrt_llm_utils.py`**:
   - Mapped `"gemma4"` to `"Gemma4ForCausalLM"` inside `MODEL_NAME_TO_HF_ARCH_MAP`.
   - Propagated model-specific properties (including attention logit soft-capping bounds) during conversion to `tensorrt_llm` layouts, aligning Gemma 4 with Gemma 2 and Gemma 3 execution profiles.
This commit introduces a dedicated, standalone unit test suite for the Gemma 4 model family within the PyTorch export pipeline. By utilizing mock modules, these tests validate the recent Gemma 4 architectural integration without requiring large external HuggingFace checkpoints.

Key test coverage includes:
- **Model Registration:** Verifies correct mapping of `Gemma4` to the `gemma4` type and `Gemma4ForCausalLM` HuggingFace architecture.
- **Layernorm Configuration:** Uses dynamically structured mock modules (`DummyRMSNorm`, `DummyAttention`, `DummyLinear`) to validate that `build_decoder_config` properly identifies and extracts configurations for `post_attention_layernorm`, `pre_feedforward_layernorm`, and `post_feedforward_layernorm` blocks.
- **Logit Softcapping:** Ensures `convert_to_tensorrt_llm_config` accurately propagates `attn_logit_softcapping`, `final_logit_softcapping`, and `query_pre_attn_scalar` properties into the resulting TensorRT-LLM configurations.
@copy-pr-bot

copy-pr-bot Bot commented Jul 3, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@coderabbitai

coderabbitai Bot commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 30f03093-02d5-4341-9de3-2c903da5f9d3

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant