bridge.models.gemma.gemma4_bridge

Megatron Bridge for Gemma 4 text-only (CausalLM).

Supports all Gemma 4 text variants:

• MoE (enable_moe_block=True): Gemma4ForCausalLM (26B-A4B and similar)

• Dense (enable_moe_block=False): same HF class, dispatched via Gemma4DenseProvider

Usage::

AutoBridge.from_hf_pretrained(“google/gemma-4-26B-A4B”) └─ Gemma4Bridge (registered for Gemma4ForCausalLM) ├─ provider_bridge() MoE → Gemma4ModelProvider │ Dense → Gemma4DenseProvider └─ mapping_registry() MoE path → _moe_mapping_registry() Dense path → _dense_mapping_registry()

Module Contents

Classes

_Gemma4QKVMapping	QKV mapping tolerating missing v_proj on global attention layers (K=V).
_Gemma4DenseQKVMapping	QKV mapping tolerating missing k_proj AND v_proj on shared-KV layers.
Gemma4Bridge	Megatron Bridge for Gemma 4 text-only (CausalLM).

Functions

_infer_attn_pattern	Infer (sliding, global) interleaved attention pattern from layer_types list.
_layer_types_from_provider	Reconstruct the Hugging Face per-layer attention pattern.
_rope_parameters_from_provider	Reconstruct Gemma 4’s dual local/global RoPE configuration.

API

class bridge.models.gemma.gemma4_bridge._Gemma4QKVMapping(*args, **kwargs)

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

QKV mapping tolerating missing v_proj on global attention layers (K=V).

Initialization

class bridge.models.gemma.gemma4_bridge._Gemma4DenseQKVMapping(*args, **kwargs)

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

QKV mapping tolerating missing k_proj AND v_proj on shared-KV layers.

Initialization

bridge.models.gemma.gemma4_bridge._infer_attn_pattern(layer_types: list[str]) → tuple[int, int]

Infer (sliding, global) interleaved attention pattern from layer_types list.

bridge.models.gemma.gemma4_bridge._layer_types_from_provider(

provider: megatron.bridge.models.gemma.gemma4_provider.Gemma4ModelProvider | megatron.bridge.models.gemma.gemma4_provider.Gemma4DenseProvider,

) → list[str]

Reconstruct the Hugging Face per-layer attention pattern.

bridge.models.gemma.gemma4_bridge._rope_parameters_from_provider(

provider: megatron.bridge.models.gemma.gemma4_provider.Gemma4ModelProvider | megatron.bridge.models.gemma.gemma4_provider.Gemma4DenseProvider,

) → dict[str, dict]

Reconstruct Gemma 4’s dual local/global RoPE configuration.

class bridge.models.gemma.gemma4_bridge.Gemma4Bridge

Bases: megatron.bridge.models.conversion.model_bridge.MegatronModelBridge

Megatron Bridge for Gemma 4 text-only (CausalLM).

Dispatches to Dense or MoE path based on enable_moe_block in HF config.

_CONDITIONAL_MOE_FIELDS

‘frozenset(…)’

_should_map_hf_config_field(

hf_config: Any,

hf_name: str,

megatron_name: str,

value: Any,

) → bool

provider_bridge(

hf_pretrained: megatron.bridge.models.hf_pretrained.causal_lm.PreTrainedCausalLM,

) → Gemma4ModelProvider | Gemma4DenseProvider

_text_config() → Any | None

Return the text config used to dispatch dense vs MoE behavior.

_is_dense_config() → bool

Return whether the current HF config describes a dense Gemma 4 model.

_build_dense_provider(

hf_config,

) → megatron.bridge.models.gemma.gemma4_provider.Gemma4DenseProvider

Build a Gemma4DenseProvider from HF config.

_build_moe_provider(

hf_config,

) → megatron.bridge.models.gemma.gemma4_provider.Gemma4ModelProvider

Build a Gemma4ModelProvider from HF config (MoE path).

classmethod megatron_to_hf_config(

provider: megatron.bridge.models.gemma.gemma4_provider.Gemma4ModelProvider | megatron.bridge.models.gemma.gemma4_provider.Gemma4DenseProvider,

) → dict

Preserve Gemma 4 architecture fields affected by provider conversion.

maybe_modify_converted_hf_weight(

task,

converted_weights_dict,

hf_state_dict,

)

Un-fuse fused weights and drop synthesized keys on export.

maybe_modify_loaded_hf_weight(

hf_param: str | dict[str, str],

hf_state_dict: Mapping[str, torch.Tensor],

) → torch.Tensor

Handle special weight loading for Gemma 4.

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

_dense_mapping_registry(

megatron_prefix: str = '',

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

Parameter mappings for the Dense variant.

_hf_layer_prefix() → str

Text-only CausalLM: weights at model.*; override in VL subclass.

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

Parameter mappings for the MoE variant.

_split_qkv_linear_out_weight(megatron_model, linear_out_weight)

Detect global vs sliding layers by tensor size for LoRA export.