`bridge.models.gemma.modeling_gemma4`#

Gemma 4 Dense and MoE layer specs, attention, positional embeddings, and helpers.

Dense (E2B, E4B, and 31B) layer specification:

4-norm transformer structure (input, post-attn, pre-MLP, post-MLP)
Dual RoPE (sliding θ=10000, global θ=1000000 with partial rotation)
Per-Layer Embeddings (PLE)
Shared KV cache (last N layers)

MoE layer specification:

TE-based transformer layer with per-layer output scaling
Dual RoPE with separate local/global embeddings
Heterogeneous sliding/global attention with independent head dims

Module Contents#

Classes#

`Gemma4RMSNorm`	HF Gemma4-compatible RMSNorm.
`Gemma4DenseMLP`	Keep both MCore projections on Gemma 4’s layer-specific FFN width.
`Gemma4MoERouter`	Token router for Gemma-4 Dense MoE block.
`Gemma4MoEExperts`	Sparse expert collection for Gemma-4 Dense MoE block.
`Gemma4DenseTransformerLayerSubmodules`	TransformerLayerSubmodules extended with Gemma-4 Dense post-sublayer norms.
`Gemma4DenseSelfAttention`	SelfAttention subclass for Gemma-4 Dense.
`Gemma4DenseTransformerLayer`	Transformer layer implementing Gemma-4 Dense 4-norm residual structure.
`_Gemma4ProportionalRotaryEmbedding`	Gemma-4 full-attention RoPE with proportional partial rotation.
`Gemma4DenseRotaryEmbedding`	Dual-theta RoPE for Gemma-4 Dense (sliding θ=10000, global θ=1000000 partial).
`Gemma4TransformerLayer`	Gemma 4 MoE transformer layer with per-layer output scaling and extra post-norms.
`Gemma4TopKRouter`	Gemma 4 MoE router with per-expert scaling.
`Gemma4MoELayer`	Gemma 4 MoE layer with post-routed-expert and post-shared-expert normalization.
`Gemma4OutputLayer`	Mixin that applies final_logit_softcapping after the output linear layer.
`Gemma4SelfAttention`	Gemma 4 MoE self attention with heterogeneous sliding/global layers.
`Gemma4TEDotProductAttention`	Gemma 4 MoE core attention — switches between sliding and global window.
`Gemma4RotaryEmbedding`	Gemma 4 MoE position RoPE — dual local/global embeddings.

Functions#

`_gemma4_rms_norm`	Apply the exact RMSNorm expression used by Hugging Face Gemma 4.
`_mark_sequence_parallel_parameter`	Mark replicated parameters whose gradients span sequence-parallel shards.
`_gemma4_dense_ffn_hidden_size`	Return the HF FFN width for a one-indexed dense Gemma 4 layer.
`_is_gemma4_sliding_layer`	Return whether a Gemma4 layer uses sliding attention.
`wire_gemma4_kv_sharing`	Wire shared-KV source references between Gemma4DenseSelfAttention layers.
`get_gemma4_layer_spec`	Return a ModuleSpec for a Gemma-4 Dense transformer layer (local/non-TE).
`_attach_ple_modules`	Add PLE embedding / projection / norm modules to a GPTModel instance.
`_compute_per_layer_inputs`	Compute per_layer_inputs of shape [b, s_local, num_layers, ple_dim], or None.
`_gemma4_layer_input`
`_gemma4_layer_accepts_input_ids`
`_gemma4_checkpointed_forward`	MCore recompute helper variant that carries Gemma4 PLE through checkpoint args.
`_patch_ple_block_threading`	Patch one Gemma4 decoder instance to thread PLE inputs through clean MCore.
`_install_ple_forward`	Patch model.forward() to compute PLE and inject as per_layer_inputs.
`_logit_softcapping`
`_install_tied_kv`	Mark global attention layers that require K=V weight tying.
`_gemma4_block_spec`	Build Gemma 4 MoE block spec with patched attention, layer, and MoE modules.

Data#

`HAVE_TE`
`RMSNorm`
`gemma4_layer_spec`

API#

bridge.models.gemma.modeling_gemma4.HAVE_TE#: None

bridge.models.gemma.modeling_gemma4._gemma4_rms_norm( hidden_states: torch.Tensor, weight: torch.Tensor | None, eps: float, ) → torch.Tensor#: Apply the exact RMSNorm expression used by Hugging Face Gemma 4.

bridge.models.gemma.modeling_gemma4._mark_sequence_parallel_parameter( parameter: torch.nn.Parameter, config: megatron.core.transformer.transformer_config.TransformerConfig, ) → None#: Mark replicated parameters whose gradients span sequence-parallel shards.

class bridge.models.gemma.modeling_gemma4.Gemma4RMSNorm( config: megatron.core.transformer.transformer_config.TransformerConfig, hidden_size: int, eps: float = 1e-06, with_scale: bool = True, )#

Bases: torch.nn.Module

HF Gemma4-compatible RMSNorm.

Gemma4 uses torch.pow(mean_squared, -0.5) rather than rsqrt. The forward values are very close, but using the same expression keeps parity tests stable for block/model gradients.

Parameters:: with_scale – If False, no learnable weight is created (matches HF’s with_scale=False used e.g. in the MoE router norm).

Initialization

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

bridge.models.gemma.modeling_gemma4.RMSNorm#: None

class bridge.models.gemma.modeling_gemma4.Gemma4DenseMLP(

config: megatron.core.transformer.transformer_config.TransformerConfig,

submodules: megatron.core.transformer.mlp.MLPSubmodules,

ffn_hidden_size: int | None = None,

**kwargs: object,

)#

Bases: megatron.core.transformer.mlp.MLP

Keep both MCore projections on Gemma 4’s layer-specific FFN width.

MCore’s FC1 accepts an explicit width, while FC2 reads it from the config. Gemma 4 E2B doubles the final shared layers, so a shallow config copy keeps both projections consistent without replacing MCore’s optimized kernels.

Initialization

bridge.models.gemma.modeling_gemma4._gemma4_dense_ffn_hidden_size( config: megatron.core.transformer.transformer_config.TransformerConfig, layer_number: int, ) → int#: Return the HF FFN width for a one-indexed dense Gemma 4 layer.

class bridge.models.gemma.modeling_gemma4.Gemma4MoERouter( config: megatron.core.transformer.transformer_config.TransformerConfig, )#

Bases: torch.nn.Module

Token router for Gemma-4 Dense MoE block.

Mirrors HF Gemma4TextRouter:

Scaleless RMSNorm → multiply by learnable per-dim scale × 1/√hidden_size
Linear projection → softmax → top-k selection
Normalize top-k weights; apply per-expert learned scale

Initialization

forward( hidden_states: torch.Tensor, ) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]#

class bridge.models.gemma.modeling_gemma4.Gemma4MoEExperts( config: megatron.core.transformer.transformer_config.TransformerConfig, )#

Bases: torch.nn.Module

Sparse expert collection for Gemma-4 Dense MoE block.

Mirrors HF Gemma4TextExperts.

Initialization

forward( hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor, ) → torch.Tensor#

class bridge.models.gemma.modeling_gemma4.Gemma4DenseTransformerLayerSubmodules#

Bases: megatron.core.transformer.transformer_layer.TransformerLayerSubmodules

TransformerLayerSubmodules extended with Gemma-4 Dense post-sublayer norms.

post_self_attn_layernorm: megatron.core.transformer.transformer_layer.LayerNormBuilder#: None

post_mlp_layernorm: megatron.core.transformer.transformer_layer.LayerNormBuilder#: None

post_per_layer_input_norm: megatron.core.transformer.transformer_layer.LayerNormBuilder#: None

bridge.models.gemma.modeling_gemma4._is_gemma4_sliding_layer( config: megatron.core.transformer.transformer_config.TransformerConfig, layer_number: int, ) → bool#: Return whether a Gemma4 layer uses sliding attention.

class bridge.models.gemma.modeling_gemma4.Gemma4DenseSelfAttention(

config: megatron.core.transformer.transformer_config.TransformerConfig,

submodules,

layer_number: int,

*args,

**kwargs,

)#

Bases: megatron.core.transformer.attention.SelfAttention

SelfAttention subclass for Gemma-4 Dense.

Extends SelfAttention with:

v_norm: scaleless RMSNorm on value states
attention_k_eq_v: full-attention layers reuse K projection for V
Shared KV cache: last N layers reuse K/V from an earlier layer

Initialization

sharded_state_dict( prefix: str = '', sharded_offsets: tuple = (), metadata=None, )#: Separate sliding and global layers in the checkpoint.

_v_norm(value: torch.Tensor) → torch.Tensor#

_get_k_eq_v_query_key_value_tensors( hidden_states: torch.Tensor, key_value_states=None, ) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]#

get_query_key_value_tensors( hidden_states: torch.Tensor, key_value_states=None, output_gate: bool = False, split_qkv: bool = True, )#

forward(

hidden_states: torch.Tensor,

attention_mask: torch.Tensor,

*args,

**kwargs,

)#

class bridge.models.gemma.modeling_gemma4.Gemma4DenseTransformerLayer(

config: megatron.core.transformer.transformer_config.TransformerConfig,

submodules: bridge.models.gemma.modeling_gemma4.Gemma4DenseTransformerLayerSubmodules,

layer_number: int = 1,

**kwargs: object,

)#

Bases: megatron.core.transformer.transformer_layer.TransformerLayer

Transformer layer implementing Gemma-4 Dense 4-norm residual structure.

Differences from the standard TransformerLayer:

post_self_attn_layernorm: applied to attention output before residual add.
post_mlp_layernorm: applied to MLP output before residual add.
Dual RoPE: selects sliding or full-attention embedding per layer.
PLE: per-layer embedding residual block after attention + MLP.
Optional local MoE block (Step 5, enabled by enable_moe_block=True).

Initialization

forward(*args, **kwargs)#

_forward_attention(

hidden_states: torch.Tensor,

attention_mask: Optional[torch.Tensor] = None,

inference_context: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None,

rotary_pos_emb=None,

rotary_pos_cos: Optional[torch.Tensor] = None,

rotary_pos_sin: Optional[torch.Tensor] = None,

rotary_pos_cos_sin=None,

attention_bias: Optional[torch.Tensor] = None,

packed_seq_params=None,

sequence_len_offset: Optional[torch.Tensor] = None,

inference_params=None,

**kwargs,

)#

_forward_mlp( hidden_states: torch.Tensor, inference_context: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None, padding_mask: Optional[torch.Tensor] = None, ) → torch.Tensor#

bridge.models.gemma.modeling_gemma4.wire_gemma4_kv_sharing(model: torch.nn.Module) → None#

Wire shared-KV source references between Gemma4DenseSelfAttention layers.

Must be called once after the model is fully constructed.

bridge.models.gemma.modeling_gemma4.get_gemma4_layer_spec( config: Optional[megatron.core.transformer.transformer_config.TransformerConfig] = None, ) → megatron.core.transformer.spec_utils.ModuleSpec#: Return a ModuleSpec for a Gemma-4 Dense transformer layer (local/non-TE).

bridge.models.gemma.modeling_gemma4.gemma4_layer_spec#: ‘get_gemma4_layer_spec(…)’

class bridge.models.gemma.modeling_gemma4._Gemma4ProportionalRotaryEmbedding( kv_channels: int, partial_rotary_factor: float, rotary_interleaved: bool = False, seq_len_interpolation_factor: Optional[float] = None, rotary_base: float = 1000000.0, use_cpu_initialization: bool = False, cp_group: Optional[torch.distributed.ProcessGroup] = None, )#

Bases: megatron.core.models.common.embeddings.rotary_pos_embedding.RotaryEmbedding

Gemma-4 full-attention RoPE with proportional partial rotation.

Initialization

class bridge.models.gemma.modeling_gemma4.Gemma4DenseRotaryEmbedding( config: megatron.core.transformer.transformer_config.TransformerConfig, rotary_percent: float = 1.0, seq_len_interpolation_factor: Optional[float] = None, use_cpu_initialization: bool = False, cp_group: Optional[torch.distributed.ProcessGroup] = None, )#

Bases: torch.nn.Module

Dual-theta RoPE for Gemma-4 Dense (sliding θ=10000, global θ=1000000 partial).

Initialization

forward( max_seq_len: int, offset: int = 0, packed_seq: bool = False, cp_group: Optional[torch.distributed.ProcessGroup] = None, )#: Return (emb_sliding, emb_full).

get_rotary_seq_len(*args, **kwargs) → int#

get_cos_sin(max_seq_len: int, offset: int = 0)#

bridge.models.gemma.modeling_gemma4._attach_ple_modules( model: torch.nn.Module, config: megatron.core.transformer.transformer_config.TransformerConfig, provider: megatron.bridge.models.gemma.gemma4_provider.Gemma4DenseProvider, ) → None#: Add PLE embedding / projection / norm modules to a GPTModel instance.

bridge.models.gemma.modeling_gemma4._compute_per_layer_inputs( model: torch.nn.Module, input_ids: torch.Tensor, decoder_input: torch.Tensor, ) → torch.Tensor | None#: Compute per_layer_inputs of shape [b, s_local, num_layers, ple_dim], or None.

bridge.models.gemma.modeling_gemma4._gemma4_layer_input( per_layer_inputs: torch.Tensor | None, layer: torch.nn.Module, ) → torch.Tensor | None#

bridge.models.gemma.modeling_gemma4._gemma4_layer_accepts_input_ids(layer: torch.nn.Module) → bool#

bridge.models.gemma.modeling_gemma4._gemma4_checkpointed_forward( self: torch.nn.Module, hidden_states: torch.Tensor, attention_mask: torch.Tensor, context: Tensor | None, context_mask: Tensor | None, rotary_pos_emb: torch.Tensor, attention_bias: Tensor | None, packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams, use_inner_quantization_context: bool, padding_mask: Tensor | None = None, extract_layer_indices: set[int] | None = None, layer_offset: int = 0, input_ids: Tensor | None = None, per_layer_inputs: Tensor | None = None, )#: MCore recompute helper variant that carries Gemma4 PLE through checkpoint args.

bridge.models.gemma.modeling_gemma4._patch_ple_block_threading(decoder: torch.nn.Module) → None#

Patch one Gemma4 decoder instance to thread PLE inputs through clean MCore.

Clean Megatron-Core’s GPTModel already forwards extra_block_kwargs to its decoder, but TransformerBlock does not know Gemma4’s per_layer_inputs. This patch is deliberately instance-scoped: it only affects the Gemma4 decoder created by this provider and leaves the TransformerBlock class unchanged.

bridge.models.gemma.modeling_gemma4._install_ple_forward(model: torch.nn.Module) → None#: Patch model.forward() to compute PLE and inject as per_layer_inputs.

class bridge.models.gemma.modeling_gemma4.Gemma4TransformerLayer(

config: megatron.core.transformer.transformer_config.TransformerConfig,

submodules: megatron.core.transformer.transformer_layer.TransformerLayerSubmodules,

layer_number: int = 1,

**kwargs: object,

)#

Bases: megatron.core.transformer.transformer_layer.TransformerLayer

Gemma 4 MoE transformer layer with per-layer output scaling and extra post-norms.

Initialization

_forward_mlp( hidden_states: torch.Tensor, inference_context: megatron.core.inference.contexts.BaseInferenceContext | None = None, padding_mask: torch.Tensor | None = None, packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None = None, ) → torch.Tensor#: Run HF’s separate shared-expert, routed-expert, and router inputs.

_forward_post_mlp( mlp_output_with_bias: tuple[torch.Tensor, torch.Tensor | None], residual: torch.Tensor, ) → torch.Tensor#

class bridge.models.gemma.modeling_gemma4.Gemma4TopKRouter(

config: megatron.core.transformer.transformer_config.TransformerConfig,

**kwargs: object,

)#

Bases: megatron.core.transformer.moe.router.TopKRouter

Gemma 4 MoE router with per-expert scaling.

Initialization

gating(input: torch.Tensor) → torch.Tensor#: Apply HF’s scaleless RMSNorm and learned router scale before projection.

routing( logits: torch.Tensor, padding_mask: torch.Tensor | None = None, input_ids: torch.Tensor | None = None, ) → tuple[torch.Tensor, torch.Tensor | None]#

class bridge.models.gemma.modeling_gemma4.Gemma4MoELayer(

config: megatron.core.transformer.transformer_config.TransformerConfig,

submodules: object,

**kwargs: object,

)#

Bases: megatron.core.transformer.moe.moe_layer.MoELayer

Gemma 4 MoE layer with post-routed-expert and post-shared-expert normalization.

Initialization

forward_with_separate_inputs( expert_input: torch.Tensor, shared_expert_input: torch.Tensor, router_input: torch.Tensor, padding_mask: torch.Tensor | None = None, ) → tuple[torch.Tensor, torch.Tensor | None]#: Run the MoE using the three independently normalized HF inputs.

postprocess( output: torch.Tensor, shared_expert_output: torch.Tensor | None, ) → torch.Tensor#

bridge.models.gemma.modeling_gemma4._logit_softcapping( logits: torch.Tensor, scale: float | None, ) → torch.Tensor#

class bridge.models.gemma.modeling_gemma4.Gemma4OutputLayer#

Bases: torch.nn.Module

Mixin that applies final_logit_softcapping after the output linear layer.

forward(*args, **kwargs)#

bridge.models.gemma.modeling_gemma4._install_tied_kv( model: torch.nn.Module, provider: megatron.bridge.models.gemma.gemma4_provider.Gemma4ModelProvider, ) → None#: Mark global attention layers that require K=V weight tying.

bridge.models.gemma.modeling_gemma4._gemma4_block_spec(config, use_transformer_engine=True, **kwargs)#: Build Gemma 4 MoE block spec with patched attention, layer, and MoE modules.

class bridge.models.gemma.modeling_gemma4.Gemma4SelfAttention(

config: megatron.core.transformer.transformer_config.TransformerConfig,

layer_number: int,

**kwargs,

)#

Bases: megatron.core.transformer.attention.SelfAttention

Gemma 4 MoE self attention with heterogeneous sliding/global layers.

Initialization

sharded_state_dict(prefix='', sharded_offsets=(), metadata=None)#: Override to separate sliding and global layers in the checkpoint.

_get_tied_query_key_value_tensors( hidden_states: torch.Tensor, key_value_states: torch.Tensor | None = None, ) → tuple[torch.Tensor, torch.Tensor, torch.Tensor]#: Return independently normalized K and V from Gemma 4’s shared raw projection.

get_query_key_value_tensors(

hidden_states,

key_value_states=None,

**kwargs,

)#: Override to apply v_norm and enforce K=V tying for global attention.

forward( hidden_states: torch.Tensor, attention_mask: torch.Tensor, key_value_states: Optional[torch.Tensor] = None, inference_context: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None, rotary_pos_emb: Optional[torch.Tensor] = None, rotary_pos_cos: Optional[torch.Tensor] = None, rotary_pos_sin: Optional[torch.Tensor] = None, rotary_pos_cos_sin: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, attention_bias: Optional[torch.Tensor] = None, packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None, sequence_len_offset: Optional[int] = None, *, inference_params: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None, ) → Tuple[torch.Tensor, torch.Tensor]#

class bridge.models.gemma.modeling_gemma4.Gemma4TEDotProductAttention(

config: megatron.core.transformer.transformer_config.TransformerConfig,

layer_number: int,

attn_mask_type: megatron.core.transformer.enums.AttnMaskType,

attention_type: str,

attention_dropout: Optional[float] = None,

**kwargs,

)#

Bases: bridge.models.gemma.modeling_gemma4.TEDotProductAttention

Gemma 4 MoE core attention — switches between sliding and global window.

Initialization

class bridge.models.gemma.modeling_gemma4.Gemma4RotaryEmbedding(

rotary_base: int = 1000000,

rotary_base_local: int = 10000,

global_kv_channels: int = 512,

global_rotary_percent: float = 0.25,

**kwargs,

)#

Bases: megatron.core.models.common.embeddings.rotary_pos_embedding.RotaryEmbedding

Gemma 4 MoE position RoPE — dual local/global embeddings.

Initialization

forward( max_seq_len: int, offset: int = 0, packed_seq: bool = False, cp_group: torch.distributed.ProcessGroup | None = None, ) → tuple[torch.Tensor, torch.Tensor]#

_forward_cached( max_seq_len: int, offset: int = 0, packed_seq: bool = False, ) → tuple[torch.Tensor, torch.Tensor]#

bridge.models.gemma.modeling_gemma4#

Module Contents#

Classes#

Functions#

Data#

API#

`bridge.models.gemma.modeling_gemma4`#