Feature/gemma4 support#1899
Draft
cjac wants to merge 4 commits into
Draft
Conversation
…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.
Contributor
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
"Gemma4"model name and"Gemma4ForCausalLM"HuggingFace architecture.build_decoder_configto correctly map Gemma 4's multi-layernorm architecture (post_attention_layernorm,pre_feedforward_layernorm,post_feedforward_layernorm).gemma4routing to leverage the upstream data-drivenhf_config_map.pyfor cleanly propagating model properties likeattn_logit_softcapping,final_logit_softcapping, andquery_pre_attn_scalar.Usage
The Gemma 4 architecture is now supported natively in the existing TensorRT-LLM export flow.
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.Modulemocks (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.).CONTRIBUTING.md: N/AAdditional Information
This unblocks downstream dynamic ONNX graph construction and TensorRT engine compilation for the Gemma 4 MoE model architectures on edge deployments.