wehub-resource-sync
470 lines
19 KiB
Python
470 lines
19 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Gemma 4 Unified multimodal model (encoder-free image + audio + video).
|
|
|
|
The Unified Gemma4 variant has no SigLIP vision tower and no audio tower.
|
|
Raw pixel patches are projected directly to LM space via a Dense+LayerNorm
|
|
pipeline with factorized 2D positional embeddings (Gemma4UnifiedVisionEmbedder),
|
|
then routed through the same Gemma4MultimodalEmbedder used by the tower-based
|
|
variant. Audio inputs are raw waveform frames projected directly through the
|
|
multimodal embedder.
|
|
|
|
This module subclasses Gemma4ForConditionalGeneration from gemma4_mm rather
|
|
than reimplementing it from scratch. Only the multimodal pipeline differs;
|
|
the language model, MTP integration, bidirectional attention helpers,
|
|
embedding/forward path, and LoRA support are all inherited unchanged.
|
|
"""
|
|
|
|
import math
|
|
from collections.abc import Iterable, Mapping
|
|
|
|
import torch
|
|
from torch import nn
|
|
from transformers.models.gemma4_unified.configuration_gemma4_unified import (
|
|
Gemma4UnifiedConfig,
|
|
)
|
|
from transformers.models.gemma4_unified.processing_gemma4_unified import (
|
|
Gemma4UnifiedProcessor,
|
|
)
|
|
|
|
from vllm.config import VllmConfig
|
|
from vllm.config.multimodal import VideoDummyOptions
|
|
from vllm.model_executor.layers.linear import ColumnParallelLinear
|
|
from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
|
|
from vllm.model_executor.models.gemma4_mm import (
|
|
_SUPPORTED_SOFT_TOKENS,
|
|
_VIDEO_MAX_FRAMES,
|
|
_VIDEO_MAX_SOFT_TOKENS,
|
|
Gemma4AudioInputs,
|
|
Gemma4DummyInputsBuilder,
|
|
Gemma4ForConditionalGeneration,
|
|
Gemma4ImageInputs,
|
|
Gemma4ImagePixelInputs,
|
|
Gemma4MultimodalEmbedder,
|
|
Gemma4MultiModalProcessor,
|
|
Gemma4ProcessingInfo,
|
|
_get_max_soft_tokens,
|
|
)
|
|
from vllm.model_executor.models.module_mapping import MultiModelKeys
|
|
from vllm.multimodal import MULTIMODAL_REGISTRY
|
|
|
|
from .utils import (
|
|
AutoWeightsLoader,
|
|
WeightsMapper,
|
|
init_vllm_registered_model,
|
|
maybe_prefix,
|
|
)
|
|
|
|
# Re-export so tests/code targeting the unified variant can import from here
|
|
# rather than reaching into gemma4_mm.
|
|
__all__ = [
|
|
"Gemma4ImagePixelInputs",
|
|
"Gemma4UnifiedVisionEmbedder",
|
|
"Gemma4UnifiedProcessingInfo",
|
|
"Gemma4UnifiedForConditionalGeneration",
|
|
]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Encoder-free vision embedder
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class Gemma4UnifiedVisionEmbedder(nn.Module):
|
|
"""Encoder-free vision embedder for Gemma4 Unified variants.
|
|
|
|
Projects raw pixel patches to LM space via dense projection and
|
|
factorized 2D positional embeddings. Replaces the SigLIP vision
|
|
tower used by the tower-based Gemma4 variant.
|
|
|
|
Pipeline: raw patches → LN₁ → Dense → LN₂ → +factorized_posemb → LN₃.
|
|
"""
|
|
|
|
def __init__(self, config, quant_config=None, prefix=""):
|
|
super().__init__()
|
|
patch_dim = config.model_patch_size**2 * 3
|
|
mm_embed_dim = config.mm_embed_dim
|
|
|
|
self.patch_ln1 = nn.LayerNorm(patch_dim)
|
|
self.patch_dense = ColumnParallelLinear(
|
|
patch_dim,
|
|
mm_embed_dim,
|
|
bias=True,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.patch_dense",
|
|
gather_output=True,
|
|
)
|
|
self.patch_ln2 = nn.LayerNorm(mm_embed_dim)
|
|
|
|
self.pos_embedding = nn.Parameter(
|
|
torch.zeros(config.mm_posemb_size, 2, mm_embed_dim)
|
|
)
|
|
self.pos_norm = nn.LayerNorm(mm_embed_dim)
|
|
|
|
def _factorized_posemb(self, positions_xy: torch.Tensor) -> torch.Tensor:
|
|
clamped_pos = positions_xy.clamp(min=0).long()
|
|
valid_mask = positions_xy != -1
|
|
|
|
pos_embs = torch.zeros(
|
|
*positions_xy.shape[:-1],
|
|
self.pos_embedding.shape[-1],
|
|
device=positions_xy.device,
|
|
dtype=self.pos_embedding.dtype,
|
|
)
|
|
for i in range(2):
|
|
axis_pe = self.pos_embedding[:, i, :][clamped_pos[..., i]]
|
|
mask = valid_mask[..., i].unsqueeze(-1).to(axis_pe.dtype)
|
|
pos_embs = pos_embs + (axis_pe * mask)
|
|
return pos_embs
|
|
|
|
def forward(
|
|
self,
|
|
pixel_values: torch.Tensor,
|
|
pixel_position_ids: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
hidden_states = self.patch_ln1(pixel_values.to(self.pos_embedding.dtype))
|
|
hidden_states, _ = self.patch_dense(hidden_states)
|
|
hidden_states = self.patch_ln2(hidden_states)
|
|
|
|
pos_embs = self._factorized_posemb(pixel_position_ids)
|
|
hidden_states = hidden_states + pos_embs
|
|
hidden_states = self.pos_norm(hidden_states)
|
|
return hidden_states
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Processing info
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class Gemma4UnifiedProcessingInfo(Gemma4ProcessingInfo):
|
|
"""ProcessingInfo for the Gemma4 Unified variant.
|
|
|
|
Two field-name differences from the tower-based parent:
|
|
* config → ``Gemma4UnifiedConfig`` (not ``Gemma4Config``)
|
|
* vision_config.``num_soft_tokens`` (not ``default_output_length``)
|
|
|
|
Everything else (token sequencing, audio limits, video frame budget,
|
|
parser construction) is inherited unchanged.
|
|
"""
|
|
|
|
def get_hf_config(self):
|
|
return self.ctx.get_hf_config(Gemma4UnifiedConfig)
|
|
|
|
def get_hf_processor(self, **kwargs: object) -> Gemma4UnifiedProcessor:
|
|
return self.ctx.get_hf_processor(
|
|
Gemma4UnifiedProcessor,
|
|
**kwargs,
|
|
)
|
|
|
|
def get_mm_max_tokens_per_item(
|
|
self, seq_len: int, mm_counts: Mapping[str, int]
|
|
) -> Mapping[str, int] | None:
|
|
config = self.get_hf_config()
|
|
# Unified field is `num_soft_tokens`. Tower-based parent uses
|
|
# `default_output_length`, hence the override.
|
|
tokens_per_image = config.vision_config.num_soft_tokens
|
|
merged_kwargs = self.ctx.get_merged_mm_kwargs({})
|
|
val, _ = _get_max_soft_tokens(merged_kwargs)
|
|
if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS:
|
|
tokens_per_image = val
|
|
tokens: dict[str, int] = {"image": tokens_per_image}
|
|
if config.audio_config is not None:
|
|
processor = self.get_hf_processor()
|
|
tokens["audio"] = processor.audio_seq_length
|
|
num_frames = _VIDEO_MAX_FRAMES
|
|
mm_config = self.ctx.model_config.get_multimodal_config()
|
|
video_opts = mm_config.limit_per_prompt.get("video")
|
|
if (
|
|
isinstance(video_opts, VideoDummyOptions)
|
|
and video_opts.num_frames is not None
|
|
):
|
|
num_frames = min(num_frames, video_opts.num_frames)
|
|
tokens["video"] = num_frames * (_VIDEO_MAX_SOFT_TOKENS + 2 + 6)
|
|
return tokens
|
|
|
|
def _compute_num_soft_tokens(
|
|
self,
|
|
image_width: int,
|
|
image_height: int,
|
|
max_soft_tokens: int | None = None,
|
|
) -> int:
|
|
vision_cfg = self.get_hf_config().vision_config
|
|
patch_size = vision_cfg.patch_size
|
|
pooling_kernel_size = vision_cfg.pooling_kernel_size
|
|
|
|
if max_soft_tokens is None:
|
|
max_soft_tokens = vision_cfg.num_soft_tokens
|
|
|
|
unit = patch_size * pooling_kernel_size
|
|
max_patches = max_soft_tokens * pooling_kernel_size**2
|
|
num_patches_orig = (image_height / patch_size) * (image_width / patch_size)
|
|
scale = math.sqrt(max_patches / num_patches_orig)
|
|
target_h = max(unit, int(math.floor(image_height * scale / unit)) * unit)
|
|
target_w = max(unit, int(math.floor(image_width * scale / unit)) * unit)
|
|
num_patches = (target_h // patch_size) * (target_w // patch_size)
|
|
num_soft_tokens = num_patches // (pooling_kernel_size**2)
|
|
return min(num_soft_tokens, max_soft_tokens)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main model
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@MULTIMODAL_REGISTRY.register_processor(
|
|
Gemma4MultiModalProcessor,
|
|
info=Gemma4UnifiedProcessingInfo,
|
|
dummy_inputs=Gemma4DummyInputsBuilder,
|
|
)
|
|
class Gemma4UnifiedForConditionalGeneration(Gemma4ForConditionalGeneration):
|
|
"""Encoder-free Gemma4 (Unified) for conditional generation.
|
|
|
|
Inherits multimodal embedding routing, PLE handling, bidirectional
|
|
attention helpers, language-model forward, LoRA, and pipeline-parallel
|
|
support from :class:`Gemma4ForConditionalGeneration`. Overrides only:
|
|
|
|
* ``__init__`` — builds the encoder-free vision embedder instead of
|
|
SigLIP/audio towers (LightOnOCR-style: ``nn.Module.__init__`` +
|
|
full rebuild, no ``super().__init__()``).
|
|
* ``hf_to_vllm_mapper`` — adds the ``model.vision_embedder.`` prefix.
|
|
* ``_process_image_input`` / ``_process_video_input`` /
|
|
``_process_audio_input`` — encoder-free projection paths.
|
|
* ``load_weights`` — ignore-prefix list excludes the absent towers.
|
|
* ``get_mm_mapping`` — no tower entries.
|
|
"""
|
|
|
|
hf_to_vllm_mapper = WeightsMapper(
|
|
orig_to_new_prefix={
|
|
"model.embed_audio.": "embed_audio.",
|
|
"model.embed_vision.": "embed_vision.",
|
|
"model.language_model.": "language_model.model.",
|
|
"model.vision_embedder.": "vision_embedder.",
|
|
"lm_head.": "language_model.lm_head.",
|
|
"model": "language_model.model",
|
|
}
|
|
)
|
|
|
|
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
|
# LightOnOCR-style rebuild: do NOT call super().__init__ — that
|
|
# would build a SigLIP vision tower and an audio tower we don't
|
|
# need. Initialize nn.Module directly and assemble the
|
|
# encoder-free pipeline below.
|
|
nn.Module.__init__(self)
|
|
config = vllm_config.model_config.hf_config
|
|
quant_config = vllm_config.quant_config
|
|
multimodal_config = vllm_config.model_config.multimodal_config
|
|
self.config = config
|
|
self.quant_config = quant_config
|
|
self.multimodal_config = multimodal_config
|
|
|
|
# No towers — set to None so inherited load_weights / get_mm_mapping
|
|
# and any tower-aware logic short-circuits.
|
|
self.vision_tower = None
|
|
self.audio_tower = None
|
|
|
|
# ---- Encoder-free vision embedder ----
|
|
self.vision_embedder = (
|
|
Gemma4UnifiedVisionEmbedder(
|
|
config.vision_config,
|
|
quant_config=quant_config,
|
|
prefix=maybe_prefix(prefix, "vision_embedder"),
|
|
)
|
|
if config.vision_config is not None
|
|
else None
|
|
)
|
|
self.embed_vision = (
|
|
Gemma4MultimodalEmbedder(
|
|
config.vision_config,
|
|
config.text_config,
|
|
)
|
|
if config.vision_config is not None
|
|
else None
|
|
)
|
|
|
|
# ---- Encoder-free audio embedder ----
|
|
self.embed_audio = (
|
|
Gemma4MultimodalEmbedder(
|
|
config.audio_config,
|
|
config.text_config,
|
|
)
|
|
if config.audio_config is not None
|
|
else None
|
|
)
|
|
|
|
# ---- Language model (vLLM optimised) ----
|
|
with self._mark_language_model(vllm_config):
|
|
self.language_model: Gemma4ForCausalLM = init_vllm_registered_model(
|
|
vllm_config=vllm_config,
|
|
hf_config=config.text_config,
|
|
prefix=maybe_prefix(prefix, "language_model"),
|
|
architectures=["Gemma4ForCausalLM"],
|
|
)
|
|
|
|
# PLE is disabled for the unified variant (text config defaults
|
|
# hidden_size_per_layer_input to 0). Skip the buffer.
|
|
ple_dim = getattr(
|
|
config.text_config,
|
|
"hidden_size_per_layer_input",
|
|
None,
|
|
)
|
|
if ple_dim is not None and ple_dim > 0:
|
|
embed = self.language_model.model.embed_tokens
|
|
self.per_layer_embeddings = torch.zeros(
|
|
vllm_config.scheduler_config.max_num_batched_tokens,
|
|
config.text_config.num_hidden_layers,
|
|
ple_dim,
|
|
device=next(embed.parameters()).device,
|
|
dtype=vllm_config.model_config.dtype,
|
|
)
|
|
else:
|
|
self.per_layer_embeddings = None
|
|
|
|
self.make_empty_intermediate_tensors = (
|
|
self.language_model.make_empty_intermediate_tensors
|
|
)
|
|
|
|
# --- Precompute full-attention layer indices for bidi clearing ---
|
|
self._full_attn_layer_idxs: frozenset[int] = frozenset()
|
|
text_config = config.text_config
|
|
if getattr(text_config, "use_bidirectional_attention", None) == "vision":
|
|
layer_types = getattr(text_config, "layer_types", None)
|
|
if layer_types:
|
|
self._full_attn_layer_idxs = frozenset(
|
|
i for i, lt in enumerate(layer_types) if lt != "sliding_attention"
|
|
)
|
|
|
|
# --- MixtureOfExperts delegation to language_model ---
|
|
self.moe_layers = self.language_model.moe_layers
|
|
self.num_moe_layers = self.language_model.num_moe_layers
|
|
self.num_logical_experts = self.language_model.num_logical_experts
|
|
self.num_physical_experts = self.language_model.num_physical_experts
|
|
self.num_local_physical_experts = self.language_model.num_local_physical_experts
|
|
self.num_routed_experts = self.language_model.num_routed_experts
|
|
self.num_expert_groups = self.language_model.num_expert_groups
|
|
self.num_shared_experts = self.language_model.num_shared_experts
|
|
self.num_redundant_experts = self.language_model.num_redundant_experts
|
|
self.set_eplb_state = self.language_model.set_eplb_state
|
|
|
|
gen_cfg = vllm_config.model_config.try_get_generation_config()
|
|
self._suppress_token_ids = gen_cfg.get("suppress_tokens") if gen_cfg else None
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Multimodal processing (encoder-free overrides)
|
|
# ------------------------------------------------------------------ #
|
|
|
|
def _process_image_input(
|
|
self,
|
|
image_input: Gemma4ImageInputs,
|
|
) -> list[torch.Tensor]:
|
|
"""Project raw image patches directly to LM space.
|
|
|
|
No vision tower: each image's pre-patchified pixel values are
|
|
embedded via Gemma4UnifiedVisionEmbedder, projected through
|
|
Gemma4MultimodalEmbedder, and padding patches (pp == -1) are
|
|
stripped per image.
|
|
"""
|
|
pixel_values = image_input["pixel_values"]
|
|
pixel_position_ids = image_input["pixel_position_ids"]
|
|
target_dtype = self.embed_vision.embedding_projection.weight.dtype
|
|
|
|
per_image_features: list[torch.Tensor] = []
|
|
for pv, pp in zip(pixel_values, pixel_position_ids, strict=True):
|
|
pv = pv.unsqueeze(0)
|
|
pp = pp.unsqueeze(0)
|
|
embedded = self.vision_embedder(pv, pp)
|
|
projected = self.embed_vision(embedded.to(target_dtype))
|
|
padding_mask = (pp.squeeze(0) == -1).all(dim=-1)
|
|
valid_features = projected.squeeze(0)[~padding_mask]
|
|
per_image_features.append(valid_features)
|
|
return per_image_features
|
|
|
|
def _process_video_input(
|
|
self,
|
|
video_input: dict[str, torch.Tensor],
|
|
) -> list[torch.Tensor]:
|
|
"""Project video frames to LM space, one frame at a time.
|
|
|
|
Frames are split per video, each frame is embedded + projected,
|
|
and per-frame valid embeddings are concatenated per video.
|
|
"""
|
|
pixel_values = video_input["pixel_values_videos"]
|
|
pixel_position_ids = video_input["pixel_position_ids_videos"]
|
|
frame_counts = video_input["video_frame_counts"]
|
|
target_dtype = self.embed_vision.embedding_projection.weight.dtype
|
|
|
|
if isinstance(frame_counts, torch.Tensor):
|
|
fc_list = frame_counts.tolist()
|
|
else:
|
|
fc_list = list(frame_counts)
|
|
|
|
pv_per_video = torch.split(pixel_values, fc_list, dim=0)
|
|
pp_per_video = torch.split(pixel_position_ids, fc_list, dim=0)
|
|
|
|
per_video_embeddings: list[torch.Tensor] = []
|
|
for pv_chunk, pp_chunk in zip(pv_per_video, pp_per_video):
|
|
frame_embs: list[torch.Tensor] = []
|
|
for i in range(pv_chunk.shape[0]):
|
|
pv = pv_chunk[i].unsqueeze(0)
|
|
pp = pp_chunk[i].unsqueeze(0)
|
|
embedded = self.vision_embedder(pv, pp)
|
|
projected = self.embed_vision(embedded.to(target_dtype))
|
|
padding_mask = (pp.squeeze(0) == -1).all(dim=-1)
|
|
frame_embs.append(projected.squeeze(0)[~padding_mask])
|
|
per_video_embeddings.append(torch.cat(frame_embs, dim=0))
|
|
return per_video_embeddings
|
|
|
|
def _process_audio_input(
|
|
self,
|
|
audio_input: Gemma4AudioInputs,
|
|
) -> list[torch.Tensor]:
|
|
"""Project raw waveform-frame features directly to LM space.
|
|
|
|
No audio tower: the per-frame raw features are passed straight
|
|
through the multimodal embedder, then padding is stripped.
|
|
"""
|
|
input_features = audio_input["input_features_padded"].squeeze(1)
|
|
input_features_mask = audio_input["input_features_mask"].squeeze(1)
|
|
|
|
target_dtype = self.embed_audio.embedding_projection.weight.dtype
|
|
audio_features = self.embed_audio(input_features.to(target_dtype))
|
|
per_audio: list[torch.Tensor] = []
|
|
for enc, mask in zip(audio_features, input_features_mask, strict=True):
|
|
per_audio.append(enc[mask])
|
|
return per_audio
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Weight loading
|
|
# ------------------------------------------------------------------ #
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
ignore_prefixes = [
|
|
# Vestigial Gemma3n-style embedding tables not used by
|
|
# Gemma4MultimodalEmbedder (which has only projection + norm).
|
|
"embed_vision.embedding.",
|
|
"embed_audio.embedding.",
|
|
]
|
|
if self.embed_audio is None:
|
|
ignore_prefixes.append("embed_audio.")
|
|
|
|
loader = AutoWeightsLoader(
|
|
self,
|
|
ignore_unexpected_prefixes=ignore_prefixes,
|
|
)
|
|
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# LoRA / multimodal mapping
|
|
# ------------------------------------------------------------------ #
|
|
|
|
def get_mm_mapping(self) -> MultiModelKeys:
|
|
"""Module prefix mapping for the encoder-free model (no towers)."""
|
|
connectors = ["embed_vision"]
|
|
if self.embed_audio is not None:
|
|
connectors.append("embed_audio")
|
|
return MultiModelKeys.from_string_field(
|
|
language_model="language_model",
|
|
connector=connectors,
|
|
tower_model=[],
|
|
)
|