bridge.models.diffusion_gemma.diffusion_gemma_bridge

HF <-> Megatron conversion for DiffusionGemma checkpoints.

Module Contents

Classes

_TiedReplicatedMapping	One Megatron tensor backed by equal decoder and encoder HF tensors.
DiffusionGemmaBridge	Convert the tied DiffusionGemma checkpoint into one Megatron text stack.

Functions

_normalized_diffusion_config

Expose implicit DiffusionGemma text semantics through Gemma4 bridge fields.

API

bridge.models.diffusion_gemma.diffusion_gemma_bridge._normalized_diffusion_config(hf_config)

Expose implicit DiffusionGemma text semantics through Gemma4 bridge fields.

class bridge.models.diffusion_gemma.diffusion_gemma_bridge._TiedReplicatedMapping(

megatron_param: str,

decoder: str,

encoder: str,

)

Bases: megatron.bridge.models.conversion.param_mapping.ReplicatedMapping

One Megatron tensor backed by equal decoder and encoder HF tensors.

Initialization

hf_to_megatron(

hf_weights: Dict[str, torch.Tensor],

megatron_module: torch.nn.Module,

) → torch.Tensor

megatron_to_hf(

megatron_weights: Optional[torch.Tensor],

megatron_module: Optional[torch.nn.Module],

) → Dict[str, torch.Tensor]

resolve(captures)

class bridge.models.diffusion_gemma.diffusion_gemma_bridge.DiffusionGemmaBridge

Bases: megatron.bridge.models.gemma_vl.gemma4_vl_bridge.Gemma4VLBridge

Convert the tied DiffusionGemma checkpoint into one Megatron text stack.

provider_bridge(

hf_pretrained,

) → megatron.bridge.models.diffusion_gemma.diffusion_gemma_provider.DiffusionGemmaModelProvider

_text_config()

_conversion_mode() → str

_hf_layer_prefix() → str

mapping_registry() → megatron.bridge.models.conversion.mapping_registry.MegatronMappingRegistry

classmethod megatron_to_hf_config(

provider: megatron.bridge.models.diffusion_gemma.diffusion_gemma_provider.DiffusionGemmaModelProvider,

) → dict