[None][feat] Add Gemma 4 12B Unified (encoder-free multimodal) support - #15768
Conversation
Adds TensorRT-LLM PyTorch-backend support for google/gemma-4-12B(-it), the encoder-free "unified" member of the Gemma 4 family (HF architecture Gemma4UnifiedForConditionalGeneration, model_type gemma4_unified). - modeling_gemma4_unified.py (new): subclasses Gemma4ForConditionalGeneration; reuses Gemma4ForCausalLM for the text core (built from text_config) and adds the encoder-free multimodal front-end -- Gemma4UnifiedVisionEmbedder (LayerNorm -> Dense -> LayerNorm + factorized 2D positional embedding -> RMSNorm -> Linear) for vision and the existing Gemma4MultimodalEmbedder for audio. No vision/audio encoder towers. - Register the architecture (register_auto_model + _GEMMA4_ARCHITECTURES), the HF weight mapper (register_mapper), and the input processor (register_input_processor, model_type "gemma4_unified"). - Unit test (registration + weight-key routing) + l0_b200 test-list entry + supported-models.md rows (L + I + A). Validated on B200: coherent text, accurate image description, and exact audio transcription. Requires transformers>=5.10 at runtime for the gemma4_unified config; no transformers/requirements pin is changed in the repo. Signed-off-by: tianruih <tianruih@nvidia.com>
…new__ attrs)
Use MODEL_CLASS_MAPPER_MAPPING[f'{arch}_HF'] (there is no get_model_mapper); set embed_vision/embed_audio=None on the __new__-built wrapper since MM load_weights consults them. 4/4 pass standalone (unittest).
Signed-off-by: tianruih <tianruih@nvidia.com>
12B Unified needs the gemma4_unified config (transformers>=5.10). TRT-LLM keeps its transformers pin unchanged; the l0-registered test now skips cleanly under the pinned env and runs only when transformers>=5.10 is installed. Signed-off-by: tianruih <tianruih@nvidia.com>
…at/gemma4-12b-unified
📝 WalkthroughWalkthroughAdds ChangesGemma4 Unified Encoder-Free Multimodal Model
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/models/modeling_utils.py (1)
845-864: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winError message understates the version requirement for the unified arch.
Gemma4UnifiedForConditionalGenerationis now part of_GEMMA4_ARCHITECTURES, but the shared error message tells users to installtransformers>=5.5.0. The unified model actually requirestransformers>=5.10(matching thepytest.mark.skipifgate intest_modeling_gemma4_unified.py). A user on e.g. 5.6 would be told to upgrade to a version they already have, yet resolution still fails.Consider differentiating the minimum version for the unified arch.
🤖 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 `@tensorrt_llm/_torch/models/modeling_utils.py` around lines 845 - 864, The Gemma4 version check in get_model_architecture currently uses one shared error for all _GEMMA4_ARCHITECTURES, but Gemma4UnifiedForConditionalGeneration needs a higher transformers minimum than the other Gemma4 variants. Update the fallback branch in get_model_architecture to distinguish the unified architecture from the rest and emit a version requirement that matches the test gating in test_modeling_gemma4_unified.py, while keeping the existing behavior for the other Gemma4 arch names.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/models/modeling_gemma4_unified.py (1)
52-52: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd missing annotations and use Python 3.10 generics.
Several new functions omit return annotations, and the file imports legacy
typing.Dict/List/Optional. Please add explicit returns (-> Nonewhere applicable,-> torch.Tensorfor feature helpers) and usedict[...],list[...], andT | Nonefor the new annotations. As per coding guidelines, “Always annotate functions” and “Prefer built-in typeslist,dict,tupleovertyping.List,typing.Dict,typing.Tuple; use|syntax instead oftyping.Union.”Also applies to: 93-93, 135-135, 188-188, 261-261, 272-272, 283-283, 383-383
🤖 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 `@tensorrt_llm/_torch/models/modeling_gemma4_unified.py` at line 52, The new annotations in modeling_gemma4_unified.py still use legacy typing imports and miss explicit return types in several helpers. Update the affected functions such as the feature helpers and initializer-style methods to include returns like -> None or -> torch.Tensor, and replace typing.Dict/List/Optional with built-in generics and union syntax (dict[...], list[...], T | None). Make sure the signatures for the affected symbols are fully annotated and remove any remaining legacy typing usage in the file.Source: Coding guidelines
🤖 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.
Inline comments:
In `@tensorrt_llm/_torch/models/modeling_gemma4_unified.py`:
- Around line 305-348: The multimodal assembly in modeling_gemma4_unified is
losing per-request pairing between features and their masks/position ids because
it collects metadata into separate lists before concatenation. Update the
request loop to store each modality as a tuple of its tensor plus its
mask/position metadata, then build the batched tensors from those paired records
so audio masks and vision/video position ids stay aligned with the correct
features. In the feature-construction path around _get_image_features and
_get_audio_features, validate that required metadata is present per request
before concatenating, instead of relying on parallel list indexes.
- Around line 261-263: The unified vision path in `_get_image_features` and
`Gemma4UnifiedVisionEmbedder.forward` can receive `image_position_ids=None`, but
the embedder still clamps it unconditionally. Add an explicit validation in
`_get_image_features` before calling `self.embed_vision(...)` so unified vision
either raises a clear `ValueError` when positions are required or routes through
a supported no-position path. Make sure the fix covers the call site around
`image_position_ids` and the `forward` logic that currently assumes positions
are always present.
- Around line 325-363: The issue is that modeling_gemma4_unified.py is passing
modality-separated mm_embeds into find_input_mm_embeds, which expects either one
already-concatenated embeddings tensor or one tensor per request. Update the
multimodal merge path in the block that builds mm_embeds and calls
find_input_mm_embeds so the embeddings are combined in placeholder order before
the helper is invoked, or construct one tensor per MultimodalParams instead.
Keep the existing token-ID gathering logic around all_mm_token_ids,
fuse_token_ids, and mm_token_type_ids, but ensure the final mm_embeds shape
matches what find_input_mm_embeds expects.
---
Outside diff comments:
In `@tensorrt_llm/_torch/models/modeling_utils.py`:
- Around line 845-864: The Gemma4 version check in get_model_architecture
currently uses one shared error for all _GEMMA4_ARCHITECTURES, but
Gemma4UnifiedForConditionalGeneration needs a higher transformers minimum than
the other Gemma4 variants. Update the fallback branch in get_model_architecture
to distinguish the unified architecture from the rest and emit a version
requirement that matches the test gating in test_modeling_gemma4_unified.py,
while keeping the existing behavior for the other Gemma4 arch names.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_gemma4_unified.py`:
- Line 52: The new annotations in modeling_gemma4_unified.py still use legacy
typing imports and miss explicit return types in several helpers. Update the
affected functions such as the feature helpers and initializer-style methods to
include returns like -> None or -> torch.Tensor, and replace
typing.Dict/List/Optional with built-in generics and union syntax (dict[...],
list[...], T | None). Make sure the signatures for the affected symbols are
fully annotated and remove any remaining legacy typing usage in the file.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4eab05f7-731e-4643-9297-4c40b617099c
📒 Files selected for processing (7)
docs/source/models/supported-models.mdtensorrt_llm/_torch/models/__init__.pytensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.pytensorrt_llm/_torch/models/modeling_gemma4_unified.pytensorrt_llm/_torch/models/modeling_utils.pytests/integration/test_lists/test-db/l0_b200.ymltests/unittest/_torch/modeling/test_modeling_gemma4_unified.py
Signed-off-by: tianruih <tianruih@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #56572 [ run ] triggered by Bot. Commit: |
… HF) _get_token_type_mask treated every non-zero mm_token_type_id as a bidirectional blob, making AUDIO (type 3) bidirectional. HF Gemma4 makes only vision (image=1/video=2) bidirectional and leaves audio causal. Restrict the bidirectional blob to vision types so audio attends causally, matching HF. Image/video behavior unchanged; no-op for variants without audio (e.g. 31B). Signed-off-by: tianruih <tianruih@nvidia.com>
|
/bot kill |
|
PR_Github #56575 [ kill ] triggered by Bot. Commit: |
|
PR_Github #56572 [ run ] completed with state |
|
PR_Github #56575 [ kill ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #56583 [ run ] triggered by Bot. Commit: |
|
PR_Github #56583 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #56608 [ run ] triggered by Bot. Commit: |
|
PR_Github #56608 [ run ] completed with state |
81fa69f to
83d752c
Compare
|
PR_Github #57181 [ run ] triggered by Bot. Commit: |
|
PR_Github #57181 [ run ] completed with state
|
…enizer The engine-side input processor is constructed without a tokenizer and does no text preprocessing; building the vendored processor there would dereference a None tokenizer. Install the fallback only on instances that carry one. Signed-off-by: tianruih <tianruih@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #57198 [ run ] triggered by Bot. Commit: |
|
PR_Github #57198 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57321 [ run ] triggered by Bot. Commit: |
The executor precomputes multimodal token indices from model.mm_token_ids; the parent class sets the backing _mm_token_ids in its __init__, which this class skips. Without it the executor falls back to the >= vocab_size heuristic, which finds no positions for Gemma4's in-vocab soft-token ids. Signed-off-by: tianruih <tianruih@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #57325 [ run ] triggered by Bot. Commit: |
|
PR_Github #57321 [ run ] completed with state |
|
PR_Github #57325 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57353 [ run ] triggered by Bot. Commit: |
|
PR_Github #57353 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57450 [ run ] triggered by Bot. Commit: |
|
PR_Github #57450 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #57486 [ run ] triggered by Bot. Commit: |
|
PR_Github #57486 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57525 [ run ] triggered by Bot. Commit: |
|
PR_Github #57525 [ run ] completed with state |
NVIDIA#15768) Signed-off-by: tianruih <tianruih@nvidia.com>
5 participants
Description
Adds TensorRT-LLM (PyTorch backend) support for
google/gemma-4-12B/-it— the dense11.95B "Unified" member of the Gemma 4 family
(HF
architectures=["Gemma4UnifiedForConditionalGeneration"],model_type="gemma4_unified").This builds on the Gemma 4 multimodal support added in #12932.
Unlike the standard Gemma 4 VLMs (26B/31B), the 12B Unified model is encoder-free: it has no
ViT vision tower or Conformer audio tower. Raw pixel patches (48×48×3) and audio frames are
projected directly into the LM embedding space through lightweight linear pipelines (the entire
multimodal front-end is 11 tensors; 666 of 677 checkpoint tensors are the text backbone).
The port is intentionally minimal and config-driven — latest
main's Gemma 4 text core alreadysupports every text feature the 12B needs:
Gemma4ForCausalLMfromtext_config(per-layer head_dim256/512, interleaved VSWA,
k_eq_vMQA on global layers, p-RoPE, per-layerlayer_scalar,final-logit softcap, tied embeddings — all already on
main).Gemma4UnifiedVisionEmbedder(LayerNorm → Dense 6912→3840 →LayerNorm + factorized-2D positional embedding → RMSNorm → Linear) and a reused
Gemma4MultimodalEmbedderfor audio — composed from TRT-LLMLinear/LayerNormmodules(TP/quant-aware), overriding
_get_{image,audio}_featuresto skip the towers.Gemma4ForConditionalGeneration, reusingpost_config/get_sub_model_config/infer_max_seq_len/get_model_defaultsandGemma4InputProcessor(HFAutoProcessor→Gemma4UnifiedProcessor; output keys match).(
use_bidirectional_attention="vision") is engaged by passingmm_token_type_ids, identical to26B/31B.
Transformers requirement: the
gemma4_unifiedconfig class ships intransformers>=5.10, whichthe user installs at runtime. This PR does not change the repo's transformers pin; the new l0
unit test is
skipif(transformers < 5.10)so CI on the pinned version skips it cleanly.New files
tensorrt_llm/_torch/models/modeling_gemma4_unified.pyGemma4UnifiedForConditionalGenerationwrapper: encoder-free vision/audio embedders + text-core reuse +Gemma4UnifiedInputProcessor;@register_auto_model+@register_input_processor.tests/unittest/_torch/modeling/test_modeling_gemma4_unified.pyPlus registry/doc one-liners:
_GEMMA4_ARCHITECTURES(modeling_utils.py), models__init__.py, theGemma 4 HF weight mapper (
@register_mapper),l0_b200.yml, andsupported-models.md(L+I+A rows).Test Coverage
tests/unittest/_torch/modeling/test_modeling_gemma4_unified.py,registered in
tests/integration/test_lists/test-db/l0_b200.yml; skips whentransformers < 5.10.transcription) generation all validated.
Summary by CodeRabbit
New Features
Documentation
Tests