bridge.models.gemma_vl.modeling_gemma4_vl#

Gemma 4 Vision-Language model.

Vision-Language model (Gemma4VLModel):

  • HuggingFace Gemma4 vision tower + multimodal embedder

  • Megatron-Core GPT language model (Dense or MoE)

Text-only (Dense/MoE) layer specs and providers live in:

  • megatron.bridge.models.gemma.modeling_gemma4

  • megatron.bridge.models.gemma.gemma4_provider

Module Contents#

Classes#

_SimpleVisionEmbedder

Fallback Gemma4 vision projector for transformers versions without the HF class.

_SimpleAudioEmbedder

Fallback Gemma4 audio projector for transformers versions without the HF class.

Gemma4VLModel

Gemma 4 Vision-Language-Audio model.

Functions#

_keep_hf_precision_buffers_in_fp32

Keep HF non-persistent precision-sensitive buffers in fp32 after casts.

API#

bridge.models.gemma_vl.modeling_gemma4_vl._keep_hf_precision_buffers_in_fp32(module: torch.nn.Module) → None#

Keep HF non-persistent precision-sensitive buffers in fp32 after casts.

HF Gemma4 registers buffers such as vision RoPE inv_freq and audio inv_timescales as non-persistent fp32 buffers. A plain module.to(dtype=bf16) casts them to bf16, but from_pretrained(torch_dtype=bf16) keeps them in fp32.

class bridge.models.gemma_vl.modeling_gemma4_vl._SimpleVisionEmbedder(
vision_hidden: int,
text_hidden: int,
eps: float,
)#

Bases: torch.nn.Module

Fallback Gemma4 vision projector for transformers versions without the HF class.

Initialization

forward(x)#
class bridge.models.gemma_vl.modeling_gemma4_vl._SimpleAudioEmbedder(
audio_proj_dim: int,
text_hidden: int,
eps: float,
)#

Bases: torch.nn.Module

Fallback Gemma4 audio projector for transformers versions without the HF class.

Initialization

forward(x)#
class bridge.models.gemma_vl.modeling_gemma4_vl.Gemma4VLModel(
config: megatron.bridge.models.gpt_provider.GPTModelProvider,
pre_process: bool = True,
post_process: bool = True,
vp_stage: Optional[int] = None,
)#

Bases: megatron.core.transformer.module.MegatronModule

Gemma 4 Vision-Language-Audio model.

Wraps HF vision/audio towers + multimodal projectors with a Megatron-Core GPT language model (Dense or MoE).

Forward flow: 1. Embed text tokens via language model embedding 2. If pixel_values: vision_tower → embed_vision → scatter at image_token_id positions 3. If input_features: audio_tower → embed_audio → scatter at audio_token_id positions 4. Forward through language model decoder

Initialization

_init_embed_vision(config)#

Initialize the multimodal embedder (vision → language projection).

_init_embed_audio(config)#

Initialize the audio projector (audio encoder output → language space).

Gemma4’s embed_audio mirrors embed_vision: parameter-free RMSNorm followed by a linear projection from audio_config.output_proj_dims to text hidden_size.

set_input_tensor(input_tensor) → None#
get_image_features(pixel_values, image_position_ids=None, **kwargs)#

Extract and project image features using HF vision tower + embedder.

get_audio_features(input_features, **kwargs)#

Extract and project audio features using HF audio tower + embedder.

_scatter_modality_features(
inputs_embeds: torch.Tensor,
input_ids: torch.LongTensor,
features: torch.Tensor,
token_id: int,
modality_name: str,
) → torch.Tensor#

Scatter projected modality features into the embedding at special token positions.

forward(
input_ids: torch.LongTensor = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
pixel_values: Optional[torch.Tensor] = None,
image_position_ids: Optional[torch.LongTensor] = None,
input_features: Optional[torch.Tensor] = None,
labels: Optional[torch.Tensor] = None,
runtime_gather_output: Optional[bool] = None,
packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None,
*,
loss_mask: Optional[torch.Tensor] = None,
) → torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]#

Forward pass combining HF vision/audio encoders with Megatron language model.

freeze(
freeze_language_model: bool,
freeze_vision_model: bool,
freeze_vision_projection: bool,
freeze_audio_model: bool = False,
freeze_audio_projection: bool = False,
)#

Freeze model modules for fine-tuning.

_compute_attention_mask(
input_ids: torch.Tensor,
) → Optional[torch.Tensor]#

Compute HF-style attention masks for full and sliding Gemma4 layers.