[model] support gemma4_unified

[Jintao-Huang]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces support for the gemma4_unified model, updating documentation, configuration parsing, and model constants, alongside implementing the Gemma4UnifiedVit and Gemma4UnifiedBridge classes. The review feedback identifies two critical issues in the new code: first, self.vision_tower and self.audio_tower are not initialized in Gemma4UnifiedVit.prepare_model, which will cause an AttributeError during the forward pass; second, the state dict conversion logic in Gemma4UnifiedBridge contains overlapping prefix checks that prevent more specific keys from being mapped correctly due to incorrect evaluation order.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

In Gemma4UnifiedVit.prepare_model, self.vision_tower and self.audio_tower are not initialized. Since Gemma4UnifiedVit inherits from Gemma4Vit and uses its get_inputs_embeds method (which accesses self.vision_tower and self.audio_tower via get_image_features/get_audio_features), omitting their initialization will lead to AttributeError during the forward pass. Additionally, the weight mapping relies on these attributes being present.

Please initialize self.vision_tower and self.audio_tower using AutoModel.from_config as done in Gemma4Vit.

Suggested change

	def prepare_model(self, hf_config: PretrainedConfig):
	from transformers.models.gemma4_unified.modeling_gemma4_unified import (Gemma4UnifiedModel,
	Gemma4UnifiedMultimodalEmbedder,
	Gemma4UnifiedVisionEmbedder)
	dtype = hf_config.torch_dtype
	self.embed_vision = (
	Gemma4UnifiedVisionEmbedder(hf_config.vision_config, hf_config.text_config).to(dtype)
	if hf_config.vision_config is not None else None)
	self.embed_audio = (
	Gemma4UnifiedMultimodalEmbedder(hf_config.audio_config, hf_config.text_config).to(dtype)
	if hf_config.audio_config is not None else None)
	self.register_buffer('embed_scale', torch.tensor(hf_config.hidden_size**0.5).to(dtype), persistent=False)
	self.model_cls = Gemma4UnifiedModel
	def prepare_model(self, hf_config: PretrainedConfig):
	from transformers.models.gemma4_unified.modeling_gemma4_unified import (Gemma4UnifiedModel,
	Gemma4UnifiedMultimodalEmbedder,
	Gemma4UnifiedVisionEmbedder)
	self.vision_tower = AutoModel.from_config(hf_config.vision_config) if hf_config.vision_config is not None else None
	self.audio_tower = AutoModel.from_config(hf_config.audio_config) if hf_config.audio_config is not None else None
	dtype = self.vision_tower.dtype if self.vision_tower is not None else hf_config.torch_dtype
	self.embed_vision = (
	Gemma4UnifiedVisionEmbedder(hf_config.vision_config, hf_config.text_config).to(dtype)
	if hf_config.vision_config is not None else None)
	self.embed_audio = (
	Gemma4UnifiedMultimodalEmbedder(hf_config.audio_config, hf_config.text_config).to(dtype)
	if hf_config.audio_config is not None else None)
	self.register_buffer('embed_scale', torch.tensor(hf_config.hidden_size**0.5).to(dtype), persistent=False)
	self.model_cls = Gemma4UnifiedModel

[Jintao-Huang]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for the gemma4_unified model type, updating configuration parsers, constants, and model definitions, and refactoring attention mask creation in Gemma4. Additionally, it simplifies state dictionary conversion logic in DeepSeek-V4 and Qwen3-Emb. Feedback on these changes highlights a potential AttributeError in Gemma4UnifiedVit.prepare_model due to missing attribute initializations (vision_embedder, vision_tower, and audio_tower), and suggests a performance optimization to avoid redundant computations of the vision group mask when creating sliding and full attention masks.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

When use_bidirectional_attention == 'vision', _create_attention_mask is called twice: once for sliding_attention and once for full_attention. Inside _create_attention_mask, the same_vision_group tensor is computed using multiple heavy operations (torch.roll, torch.cumsum, torch.where, etc.). Since mm_token_type_ids is identical for both calls, this results in redundant calculations in every forward pass.

Consider precomputing same_vision_group once in forward and passing it to _create_attention_mask to improve efficiency:

        full_attention = sliding_attention = attention_mask
        if self.text_config.use_bidirectional_attention == 'vision':
            same_vision_group = None
            if mm_token_type_ids is not None:
                is_vision = mm_token_type_ids > 0
                is_prev_vision = torch.roll(is_vision, shifts=1, dims=-1)
                is_prev_vision[:, 0] = False
                vision_group_ids = torch.cumsum((is_vision & ~is_prev_vision).int(), dim=1) - 1
                vision_group_ids = torch.where(is_vision, vision_group_ids, torch.full_like(vision_group_ids, -1))
                q_group = vision_group_ids.unsqueeze(1).unsqueeze(-1)
                k_group = vision_group_ids.unsqueeze(1).unsqueeze(-2)
                same_vision_group = (q_group == k_group) & (q_group >= 0) & (k_group >= 0)
            sliding_attention = self._create_attention_mask(sliding_attention, same_vision_group, is_sliding=True)
            full_attention = self._create_attention_mask(full_attention, same_vision_group, is_sliding=False)
        kwargs['attention_mask'] = {'sliding_attention': sliding_attention, 'full_attention': full_attention}

And update _create_attention_mask accordingly:

    def _create_attention_mask(self, attention_mask, same_vision_group, is_sliding: bool):
        if is_sliding:
            window_size = self.text_config.sliding_window - 1
            seq_len = attention_mask.shape[-1]
            window_mask = torch.ones(seq_len, seq_len, dtype=torch.bool, device=attention_mask.device)
            window_mask = ~torch.triu(window_mask, diagonal=-window_size)
            attention_mask = attention_mask | window_mask
        if same_vision_group is not None:
            attention_mask = attention_mask & ~same_vision_group
        return attention_mask