bridge.models.diffusion_gemma.modeling_diffusion_gemma

Megatron DiffusionGemma encoder/decoder model.

Module Contents

Classes

DiffusionGemmaLanguageModelEmbedding	Gemma token embedding with HF DiffusionGemma’s scale and padding semantics.
DiffusionGemmaRMSNorm	RMSNorm with DiffusionGemma’s FP32 normalization and optional scale.
DiffusionGemmaSelfConditioning	Decoder self-conditioning block from the pinned HF implementation.
DiffusionGemmaEncoderOutput	Encoder outputs and optional per-layer attention K/V.
DiffusionGemmaBlockDiffusionOutput	Decoder logits plus the encoder outputs used to produce them.
DiffusionGemmaModel	Gemma 4 VL multimodal encoder plus DiffusionGemma decoder-only modules.

Functions

_attention_qkv
_replace_attention_args

API

class bridge.models.diffusion_gemma.modeling_diffusion_gemma.DiffusionGemmaLanguageModelEmbedding

Bases: megatron.bridge.models.gemma.gemma3_provider.Gemma3LanguageModelEmbedding

Gemma token embedding with HF DiffusionGemma’s scale and padding semantics.

padding_idx: Optional[int]

None

forward(

input_ids: torch.Tensor,

position_ids: torch.Tensor,

tokentype_ids: int = None,

) → torch.Tensor

class bridge.models.diffusion_gemma.modeling_diffusion_gemma.DiffusionGemmaRMSNorm(

dim: int,

eps: float = 1e-06,

with_scale: bool = True,

)

Bases: torch.nn.Module

RMSNorm with DiffusionGemma’s FP32 normalization and optional scale.

Initialization

forward(hidden_states: torch.Tensor) → torch.Tensor

class bridge.models.diffusion_gemma.modeling_diffusion_gemma.DiffusionGemmaSelfConditioning(

hidden_size: int,

intermediate_size: int,

eps: float,

)

Bases: torch.nn.Module

Decoder self-conditioning block from the pinned HF implementation.

Initialization

forward(

inputs_embeds: torch.Tensor,

self_conditioning_signal: torch.Tensor,

) → torch.Tensor

class bridge.models.diffusion_gemma.modeling_diffusion_gemma.DiffusionGemmaEncoderOutput

Encoder outputs and optional per-layer attention K/V.

last_hidden_state: torch.Tensor

None

key_values: Optional[list[tuple[torch.Tensor, torch.Tensor]]]

None

class bridge.models.diffusion_gemma.modeling_diffusion_gemma.DiffusionGemmaBlockDiffusionOutput

Decoder logits plus the encoder outputs used to produce them.

logits: torch.Tensor

None

encoder_last_hidden_state: torch.Tensor

None

encoder_loss: torch.Tensor | None

None

class bridge.models.diffusion_gemma.modeling_diffusion_gemma.DiffusionGemmaModel(

config,

pre_process: bool = True,

post_process: bool = True,

vp_stage: Optional[int] = None,

)

Bases: megatron.bridge.models.gemma_vl.modeling_gemma4_vl.Gemma4VLModel

Gemma 4 VL multimodal encoder plus DiffusionGemma decoder-only modules.

The Megatron language model is the single tied text stack. The HF checkpoint stores its tied tensors under model.decoder.*; the encoder reuses them.

Initialization

_without_output_head()

_tp_group()

_compute_attention_mask(

input_ids: torch.Tensor,

mm_token_type_ids: Optional[torch.Tensor] = None,

) → Optional[torch.Tensor]

encode(

*,

return_key_values: bool = False,

attention_mask: Optional[torch.Tensor] = None,

**kwargs,

) → bridge.models.diffusion_gemma.modeling_diffusion_gemma.DiffusionGemmaEncoderOutput

Run the multimodal encoder without materializing vocabulary logits.

_decoder_embeddings(

decoder_input_ids: torch.Tensor,

self_conditioning_logits: Optional[torch.Tensor],

self_conditioning_mask: Optional[torch.Tensor],

) → torch.Tensor

_decoder_attention(

key_values,

encoder_attention_mask: Optional[torch.Tensor],

canvas_length: int,

)

decode(

decoder_input_ids: torch.Tensor,

key_values: list[tuple[torch.Tensor, torch.Tensor]],

encoder_length: int,

encoder_attention_mask: Optional[torch.Tensor] = None,

self_conditioning_logits: Optional[torch.Tensor] = None,

self_conditioning_mask: Optional[torch.Tensor] = None,

) → torch.Tensor

Denoise one canvas from captured encoder K/V and return logits [B, S, V].

forward(

input_ids: torch.LongTensor,

decoder_input_ids: torch.LongTensor,

position_ids: Optional[torch.LongTensor] = None,

pixel_values: Optional[torch.Tensor] = None,

image_position_ids: Optional[torch.LongTensor] = None,

mm_token_type_ids: Optional[torch.LongTensor] = None,

encoder_attention_mask: Optional[torch.Tensor] = None,

self_conditioning_logits: Optional[torch.Tensor] = None,

self_conditioning_mask: Optional[torch.Tensor] = None,

encoder_target_ids: Optional[torch.Tensor] = None,

encoder_target_mask: Optional[torch.Tensor] = None,

) → bridge.models.diffusion_gemma.modeling_diffusion_gemma.DiffusionGemmaBlockDiffusionOutput

freeze(

freeze_language_model: bool,

freeze_vision_model: bool,

freeze_vision_projection: bool,

freeze_audio_model: bool = False,

freeze_audio_projection: bool = False,

freeze_self_conditioning: bool = False,

)

bridge.models.diffusion_gemma.modeling_diffusion_gemma._attention_qkv(args, kwargs)

bridge.models.diffusion_gemma.modeling_diffusion_gemma._replace_attention_args(

args,

kwargs,

query,

key,

value,

attention_mask,

)