[None][feat] Gemma4 MM: native vision + audio towers - #14300
Conversation
- modeling_gemma4_vision.py: Gemma4VisionModel inherits Attention base class, participates in TRT-LLM backend dispatch (FLASHINFER for LM, TRTLLM for vision). - modeling_gemma4_audio.py: stub for Phase 2 audio tower migration. - modeling_gemma4mm.py: construct native TRT-LLM vision tower instead of HF AutoModel.from_config(). Pre-nvidia/main merge snapshot. Signed-off-by: Hudayday <tianruih@nvidia.com>
modeling_gemma4_vision.py: drop ``[:]`` slicing on clamp scalar buffers in
``load_weights`` and ``_remap_clamp_buffers`` (use ``.clone()`` instead).
HF registers clamp buffers as 0-dim ``torch.tensor(-float("inf"))``; ``[:]``
on a 0-dim tensor raises and breaks ``use_clipped_linears=True`` loads.
Also switch ``seq_lens`` pinning to ``maybe_pin_memory`` per repo policy.
modeling_gemma4mm.py: add no-op ``_check_and_adjust_experts_implementation``
on ``Gemma4ForConditionalGeneration``. HF transformers 5.5.3
``PreTrainedModel.__init__`` dispatches this on VL wrappers and raises when
the wrapper has no direct MoE layers (mirrors ``Qwen3VLModelBase`` override).
TRT-LLM manages expert implementations independently.
test_gemma4_multimodal.py: full Tier 0/1/2 rewrite against the post-Option B
vision tower.
- Tier 0 ``TestGemma4VisionTower``: build / forward random-weights (fp32 —
bf16 + head_dim=32 + trtllm-gen unfused MHA fallback at sm_100 is NaN-
prone) / load-weights remap / clamp buffer parity.
- Tier 1 ``TestGemma4VisionPipeline``: HF vs TRT vision tower parity.
- Tier 1 ``TestGemma4ForConditionalGeneration``: vision_tower attach/skip.
- Tier 2 ``TestGemma4MultimodalE2E`` via shared ``TestModelingMultimodal``
skeleton; gated by ``_gemma4_e2e_model_available()`` which also requires
``{LLM_MODELS_ROOT}/multimodals/test_data/seashore.png`` (skeleton image
fixture used by ``modality="image"``).
B200 (sm_100, c163c72 wheel + transformers 5.5.3) result: Tier 0/1 12/12
PASS, Tier 2 13 PASS + 1 skip (E2E gracefully gated on missing fixture).
Signed-off-by: Hudayday <tianruih@nvidia.com>
HF Gemma4RMSNorm computes `w * x / rms(x)` (plain w), NOT the `(1 + w) * x / rms(x)` convention used by Gemma3 LLM. The previous audio+vision ports re-implemented RMSNorm by analogy with Gemma3, which silently corrupted activations because the Gemma4 checkpoint weights are not zero-centered (mean=+2.7, absmax=9.2). Compounded across ~36 norm applications in the conformer, this produced grammatical-but-wrong German translations: E4B CoVoST BLEU=1.90 vs HF baseline 24.54. Fix: replace inline norm classes with thin adapters inheriting `transformers.models.gemma4.modeling_gemma4.Gemma4RMSNorm`, preserving the existing kwarg + dtype + has_weights call signature. Vision uses `with_scale=has_weights` for the weightless v_norm slot. Drop the 2D reshape helper in vision _head_norm (HF norm is pure PyTorch, no flashinfer dispatcher constraint on shape). Also wire mm wrapper to use the native `Gemma4AudioModel` (with `load_weights`) instead of `AutoModel.from_config`. Verification (B200): E4B CoVoST BLEU = 24.69 (job 2206264), slightly above HF baseline 24.54. 31B vision (job 2206265) past the original `flashinfer_rmsnorm` dispatcher crash on bf16 1152-wide SigLIP tensors; remaining 31B issue is a separate downstream attention resize bug unrelated to RMSNorm. Tests: add `TestGemma4MMTowerRMSNormConvention` (numerical regression locking in plain-w convention, inheritance + dtype + weightless guards) and `TestGemma4AudioTowerStructure` (module exports + mm wiring) in tests/unittest/_torch/modeling/test_modeling_gemma4.py. Signed-off-by: Hudayday <tianruih@nvidia.com>
Signed-off-by: Hudayday <tianruih@nvidia.com>
Vision tower's tower-local TRTLLM attn metadata was pulling max_num_tokens from the LLM model_config (defaults to LLM's 16384) and keeping max_num_requests=8192. The C++ attention op sizes its workspace as ~max_num_requests * heads * max_num_tokens^2 * bytes, which scaled to TB-range and crashed with "Tried to allocate 1397.40 GiB". Vision is called per-image (modeling_gemma4mm._get_image_features loops and feeds one image at a time), so batch and seq are vision-scale, not LLM-scale. Bind to (max_num_requests=1, max_num_tokens=8192) — matches the actual call shape and keeps the workspace at ~tens of MB. Signed-off-by: Hudayday <tianruih@nvidia.com>
Removes the eight DEBUG_PROBE_BEGIN/END blocks (and the _DEBUG_AUDIO flag, _audio_stats helper, unused os import) from Gemma4AudioModel.forward / load_weights. The probes were gated by an env var and never fired under normal runs; carrying them in source is noise. Signed-off-by: Hudayday <tianruih@nvidia.com>
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughThis PR replaces HuggingFace ChangesGemma4 Audio/Vision Towers and Integration
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/models/modeling_gemma4_audio.py (1)
247-256: 💤 Low valueUnused tuple unpacking in
_extract_block_context.The unpacked variables
batch_size,seq_len,num_heads,head_dimare never used. Replace with underscores to silence the static analysis warning.♻️ Suggested fix
def _extract_block_context(self, hidden_states: torch.Tensor) -> torch.Tensor: """Extract overlapping context windows of ``context_size`` per block.""" - batch_size, seq_len, num_heads, head_dim = hidden_states.shape + _, _, _, _ = hidden_states.shape # Shape validated by caller hidden_states = F.pad(🤖 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_audio.py` around lines 247 - 256, In _extract_block_context replace the unused tuple unpacking "batch_size, seq_len, num_heads, head_dim = hidden_states.shape" with underscore placeholders (e.g. _, _, _, _ = hidden_states.shape) or simply omit the assignment to silence the static-analysis warning; update the line in the method _extract_block_context in modeling_gemma4_audio.py so only necessary values are captured or ignored to avoid unused variable warnings while preserving the rest of the tensor padding, unfold, movedim, and contiguous operations.
🤖 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 `@tensorrt_llm/_torch/models/modeling_gemma4_audio.py`:
- Around line 247-256: In _extract_block_context replace the unused tuple
unpacking "batch_size, seq_len, num_heads, head_dim = hidden_states.shape" with
underscore placeholders (e.g. _, _, _, _ = hidden_states.shape) or simply omit
the assignment to silence the static-analysis warning; update the line in the
method _extract_block_context in modeling_gemma4_audio.py so only necessary
values are captured or ignored to avoid unused variable warnings while
preserving the rest of the tensor padding, unfold, movedim, and contiguous
operations.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2eaf5b81-46cf-4cec-a63b-30ab01df1413
📒 Files selected for processing (6)
tensorrt_llm/_torch/models/modeling_gemma4_audio.pytensorrt_llm/_torch/models/modeling_gemma4_vision.pytensorrt_llm/_torch/models/modeling_gemma4mm.pytests/integration/test_lists/test-db/l0_b200.ymltests/unittest/_torch/modeling/test_gemma4_multimodal.pytests/unittest/_torch/modeling/test_modeling_gemma4.py
|
PR_Github #49169 [ run ] triggered by Bot. Commit: |
Two test fixes for failures in L0_Test build 2316:
1. TestGemma4MMTowerRMSNormConvention::test_vision_attention_uses_native_rmsnorm_adapter
Failure: AssertionError on '_Gemma4VisionRMSNorm(hidden_size=' not
found in source. The guard read the vision module source as a string
and did a literal substring match. yapf wraps long ctor calls so the
actual source has '_Gemma4VisionRMSNorm(\n hidden_size=...'
— the assertion passed in CI before the call sites were reformatted
and broke once the formatter ran. Switch the search to a regex that
tolerates whitespace + newline between '(' and 'hidden_size=', and
keep a negative lookbehind so a future direct TRT-LLM RMSNorm
construction still trips the guard.
2. TestGemma4VisionTower::test_forward_random_weights
Failure: AssertionError 'Vision tower output contains NaN'. The test
built the TRT-LLM vision tower with PyTorch-default init on fused
qkv_proj and pushed activations through the TRTLLM attention backend.
head_dim=32 + uninitialised fused weights saturate the attention
kernel and produce NaN even in fp32 (the dtype switch from bf16 to
fp32 in an earlier commit did not help). Production code always
calls load_weights(hf_state_dict) before forward, so the appropriate
sanity-check init is HF defaults. Seed the TRT-LLM tower via an HF
Gemma4VisionModel state_dict (same path test_forward_parity_no_clamp
uses) and keep shape + no-NaN assertions.
Local sbatch verification on B200 (job 2233456):
TestGemma4VisionTower::test_forward_random_weights PASSED
TestGemma4MMTowerRMSNormConvention
::test_vision_attention_uses_native_rmsnorm_adapter PASSED
test_modeling_gemma4.py (full file) 74/74 PASSED
Pre-existing main-side flake / CI infra failures from build 2316 that
also showed up in build 2317 without this PR's diff and are NOT
addressed here:
- tests/unittest/auto_deploy/singlegpu/models/test_qwen3_5_moe.py
::test_vision_attention_matches_reference
::test_vision_block_matches_reference
::test_vlm_wrapper_delta_is_request_scoped_no_cross_call_leakage
- tests/perf/test_perf_sanity.py
::test_e2e[aggr_upload-llama3_1_8b_fp8_ad_hopper-llama3_1_8b_ad_ws1_1k1k]
Signed-off-by: Hudayday <tianruih@nvidia.com>
|
PR_Github #49169 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49302 [ run ] triggered by Bot. Commit: |
|
PR_Github #49302 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49343 [ run ] triggered by Bot. Commit: |
|
PR_Github #49343 [ run ] completed with state |
Addresses 5 reviewer comments on PR 14300:
C1+C2: Drop leading underscore from private classes per reviewer style
preference. Renames `_Gemma4{Audio,Vision}RMSNorm` and
`_{Audio,Vision}Output` to their public counterparts. Test imports
updated.
C4+C5: Lift the per-image Python loop in `_get_image_features` to a
single batched `vision_tower(pixel_values, image_position_ids)` call —
matches the pattern in Gemma3vl / Qwen2VL / Qwen3VL / LlavaNext /
Nemotron-Nano-dyn-res. The inner `Gemma4VisionModel.forward` already
supported B>1 via varlen `cu_seqlens` (per-image `seq_lens_cpu` from
`valid.sum(dim=-1)`); only the caller-side loop had to go.
`max_num_requests` on the tower-local `attn_metadata` is bumped from a
hardcoded 1 to the LLM-side value (typically 8192) — see below for the
workspace story.
FMHA enablement: vision attention previously fell into the unfused MHA
path (`Fall back to unfused MHA because of unsupported data type`) on
all 4 Gemma4 variants. qkv reaching `thop.attention` was not bf16/half
due to HF `Gemma4RMSNorm` / RoPE intermediate fp32 promotion, and
26B/31B vision head_dim=72 is not in the trtllm-gen sm100a cubin set
(H64/H80/H128/H256/H512). Two changes restore the FMHA path:
- Explicit dtype cast on q/k/v to the qkv_proj weight dtype right
before `forward_impl` keeps the FMHA dispatcher's mType check happy.
- For variants whose HF head_dim is not FMHA-supported (currently
26B/31B vision: 72), `Gemma4VisionAttention.__init__` overrides
self.head_dim to the next supported size (72 to 80), and
`Gemma4VisionModel._pad_attention_head_dim` zero-pads q/k/v/o_proj
weights along the head dim at load time. RMSNorm + RoPE operate on
the unpadded HF head_dim — zeros in K/V don't contribute to QK^T,
zeros in V produce zero output channels, zeros in o_proj's last
columns ignore those channels. Net: byte-equivalent to unpadded HF
math, plus ~11% extra compute (80/72) on 26B/31B.
Empirical (B200, samples=30, gemma4_main_c163c72867 wheel + this
overlay):
variant head_dim FMHA fallback workspace MMMU acc CoVoST BLEU
E2B 64 0 ~32 MB 42.22 21.87
E4B 64 0 ~32 MB 48.00 24.69
26B-A4B 72 -> 80 0 ~32 MB 56.67 -
31B 72 -> 80 0 ~32 MB 65.22 -
Workspace at max_num_requests=1 pre-fix (unfused) was 273 MB on 31B at
3135 patches/image. FMHA on path is linear in tokens (no qk_buf, no
attention_mask buffer) so the bump to LLM-side max_num_requests keeps
the workspace ~tens of MB regardless of B. E4B CoVoST BLEU 24.69
matches the HF tower baseline 24.54 — audio path quality preserved.
C3: `valid.sum(dim=-1).to(torch.int).cpu()` sync is API-mandated.
AttentionMetadata.prompt_lens is typed Optional[List[int]] (host);
same pattern in modeling_siglip.py, modeling_qwen2vl.py,
modeling_nemotron_nano.extract_feature_dynamic. No code change; reply
documents the contract.
Tests added:
- `TestGemma4VisionCrossImageBatching` in test_modeling_gemma4.py:
structural guards that `_get_image_features` contains exactly one
vision_tower call (no per-image for loop) and that the vision
attn_metadata does not hardcode max_num_requests=1.
- `TestGemma4VisionHeadDimPadding` in test_gemma4_multimodal.py:
verifies attention overrides head_dim to 80 when HF head_dim is 72,
skips override when head_dim is already FMHA-supported, zero-pads
qkv_proj/o_proj weights at load, and matches HF forward output at
atol/rtol=5e-2.
- `test_forward_batched_matches_per_image_loop` in
test_gemma4_multimodal.py: B=3 batched call output matches B=1
looped x 3 concat to atol=1e-4 (guards against future cross-image
leakage in shared attn_metadata state).
Signed-off-by: Hudayday <tianruih@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #49441 [ run ] triggered by Bot. Commit: |
|
PR_Github #49441 [ run ] completed with state |
…CPU-side
Follow-up to PR 14300 review (yechank-nvidia):
> But it seems Gemma4 derives those values from GPU tensor to CPU which
> occurs CPU-GPU sync. and hurt performance very badly. Is there anyway
> to derive those numbers from host tensors?
Earlier reply conflated two different sync patterns:
- SigLip / Qwen2VL / Nemotron-Nano derive `prompt_lens` from CPU-side
processor metadata (e.g. Qwen2VL's `grid_thw.tolist()` on a CPU tensor)
so the subsequent pinned-memory H2D in `attn_metadata.prepare()` is
non-blocking.
- Gemma4 vision was reducing `pixel_position_ids` (GPU-resident for RoPE)
with `valid.sum(-1).cpu()` — a real GPU→CPU blocking sync that stalled
the CUDA stream per forward.
Fix mirrors the Qwen2VL pattern:
Gemma4InputProcessor (CPU)
pixel_position_ids ───── [executor moves to GPU later] ────► RoPE
image_seq_lens ─── List[int], never tensorised to GPU ──► attn_metadata
`Gemma4InputProcessor.__call__` now computes the per-image valid-patch
count on CPU at the moment it inserts the `-1` padding sentinels into
`image_position_ids` (both tensors are still host-resident at that point;
the executor pipeline ships them to GPU later). The count is stored as a
Python `List[int]` under `multimodal_data["image"]["image_seq_lens"]`
(and the analogous `image_data["video"]` field for video frames) — Python
lists aren't touched by the device-migration step, so the value stays CPU
through the entire pipeline.
`Gemma4ForConditionalGeneration.forward` collects per-request lists and
flattens with `list.extend(...)` (no torch ops), then passes through to
`_get_image_features` as a new optional kwarg.
`Gemma4VisionModel.forward` now picks `seq_lens_cpu` from one of two CPU
branches (both `device="cpu"`, neither touches GPU):
- caller-supplied list → `torch.tensor(list, dtype=torch.int, device="cpu")`
- no list (synthetic / unit-test input with no -1 padding) →
`torch.full((B,), N, dtype=torch.int, device="cpu")`
The previous `valid.sum(dim=-1).to(torch.int).cpu()` GPU reduction +
blocking D2H is gone entirely (no fallback path — the assert on
`len(image_seq_lens) == B` catches plumbing bugs, and the all-valid
default covers synthetic inputs). RoPE on GPU is untouched —
`pixel_position_ids` still flows GPU-side into `rotary_emb`.
`attn_metadata.prepare()` does the subsequent pinned-memory H2D async
(non-blocking) — same as SigLip / Qwen2VL / Nemotron-Nano.
Empirical (B200, 30-sample MMMU + 20-sample CoVoST, gemma4_main_c163c72867
wheel + this overlay):
variant task FMHA fall-back workspace accuracy
E2B MMMU 0 ~32 MB 42.22
E4B MMMU 0 ~32 MB 48.00
26B-A4B MMMU 0 ~32 MB 58.89
31B MMMU 0 ~32 MB 65.22
E2B CoVoST en_de 0 - BLEU 21.87
E4B CoVoST en_de 0 - BLEU 24.69
No regression vs the previous commit (c157778). Sync count per
MMMU 30-sample run drops by ~1000 D2H syncs (one per `_get_image_features`
call); per-sync cost is ~50-100 us on B200, so ~50-100 ms total wallclock
saved per 30 samples — modest in throughput-bound benchmarks, more visible
in latency-sensitive paths.
New unit test added: `TestGemma4VisionTower.test_forward_image_seq_lens_kwarg_matches_default`
verifies byte-equivalent output between the explicit-list path and the
default all-valid path (atol=1e-4).
Signed-off-by: Hudayday <tianruih@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #49565 [ run ] triggered by Bot. Commit: |
…kill gaps Two fixes flagged by the multimodal-model-onboarding skill (NVIDIA#13842) audit against this PR: 1. Drop `transformers.models.gemma4` direct imports. `Gemma4VisionRMSNorm` and `Gemma4AudioRMSNorm` previously inherited from `transformers.models.gemma4.modeling_gemma4.Gemma4RMSNorm`. The skill's "Multimodal encoder — port to TRT-LLM modules" rule: > No imports of HF modules or computations from > transformers.models.<family> — they couple the file to a specific > transformers release and break silently on upstream refactors. Both classes are now ~12-line native `nn.Module` subclasses implementing the plain `w * x / rms` Gemma4 convention (not the Gemma3 LLM `(1 + w) * x / rms` variant) with fp32 accumulator and cast-back to input dtype — mathematically identical to the HF class. 2. Declare `multimodal_data_device_paths`. The skill's Contract 2 / Phase 2 checklist requires every multimodal wrapper to list the dotted paths in `multimodal_data` that the engine should ship to GPU. Anything not listed stays CPU-resident. Added on `Gemma4ForConditionalGeneration`: lists `pixel_values`, `image_position_ids`, `pixel_values_videos`, `audio_features`, and `audio_features_mask` as GPU paths. Critically, `image_seq_lens` and `video.image_seq_lens` (the CPU `List[int]` from the previous commit's D2H-sync elimination) are NOT listed — the prior commit's invariant ("Python list isn't device-migrated") is now an explicit architectural declaration instead of an implicit assumption. Empirical (B200, gemma4_main_c163c72867 wheel + this overlay): variant task FMHA fall-back accuracy E2B MMMU 30s 0 42.22 E4B MMMU 30s 0 48.00 26B-A4B MMMU 30s 0 57.44 (= prior 58.89 within seed noise) 31B MMMU 30s 0 65.22 E2B CoVoST en_de 20s 0 BLEU 21.87 E4B CoVoST en_de 20s 0 BLEU 24.69 (= HF baseline 24.54) All within seed noise of the prior commit (`001a6a01bd`); no perf or accuracy regression introduced by porting the RMSNorm. Tests updated: - `test_audio_rmsnorm_inherits_from_hf_class` / `test_vision_rmsnorm_inherits_from_hf_class` rewritten as `test_audio_rmsnorm_does_not_import_hf_class` / `test_vision_rmsnorm_does_not_import_hf_class`: guard the new invariant (no `transformers.models.gemma4` imports in either modeling file; class MRO has no HF module). - `test_audio_rmsnorm_matches_hf_class_with_same_weights` / `test_vision_rmsnorm_*` continue to verify byte-identical output vs HF reference at fp32 with shared weights — confirms math equivalence. Signed-off-by: Hudayday <tianruih@nvidia.com>
|
/bot run --disable-fail-fast |
yechank-nvidia
left a comment
There was a problem hiding this comment.
Thanks all for the hard work!
|
PR_Github #49594 [ run ] triggered by Bot. Commit: |
|
PR_Github #49565 [ run ] completed with state |
|
PR_Github #49594 [ run ] completed with state |
Signed-off-by: Hudayday <tianruih@nvidia.com>
Signed-off-by: Hudayday <tianruih@nvidia.com>
Signed-off-by: Hudayday <tianruih@nvidia.com>
Summary
Per multimedia team's PR feedback on the initial Gemma4 MM port, this PR replaces the temporary
AutoModel.from_configHF fallback with native TensorRT-LLM implementations of the Gemma4 vision and audio(Conformer) towers, so that the multimodal path goes through TRT-LLM's own attention/MLP modules rather than re-entering HuggingFace at runtime.
modeling_gemma4_vision.py) — replacesAutoModel.from_config; vision attention workspace now bound to the vision scale rather than inheriting the LM-side bound (avoids the5589 GiB resize previously observed on 31B MMMU).
modeling_gemma4_audio.py) — replacesAutoModel.from_config; chunked local attention + rel-pos bias ported directly from the HF reference.w * x / rms, not the Gemma3 LLM-side(1 + w) * x / rms. Wiring the correct norm convention is what bringsBLEU back into parity (E4B CoVoST2 1.90 → ~41.9 in earlier validation runs).
modeling_gemma4mm.pyupdated to import the native towers instead of the HF fallback.tests/unittest/_torch/multimodal/test_gemma4_multimodal.pyregistered inl0_b200.yml, covering Tier 0/1/2 vision-tower parity plus the new RMSNorm regression tests.DEBUG_AUDIO_TOWERenv-gated stats) have been scrubbed before this commit; only Signed-off-by trailers in the history.Alignment with HF tower baseline
The fix-mode (native towers) results have been aligned with the HF-fallback baseline. All completed fix-mode cells so far are within ±2pp of the HF baseline (well inside MMMU/CoVoST seed noise):
Summary by CodeRabbit
Release Notes
New Features
Tests