opentau.policies.pi07.gemma3_with_expert

Gemma 3 backbone with Gemma-v1 action expert, for the PI06 policy.

Mirrors paligemma_with_expert.py but:

• the vision-language backbone is Gemma3ForConditionalGeneration (Gemma 3 4B, SigLIP-400m/14 + Gemma 3 text, 34 interleaved sliding-window/global layers);

• the action expert is a Gemma-v1 GemmaForCausalLM with the AdaRMS and gated-residual patches applied by opentau.utils.transformers_patch.

The per-layer attention loop below concatenates backbone and expert queries/ keys/values along the sequence dimension at every layer (the MoE-like pattern introduced in π0), so the expert can cross-attend to the backbone’s activations at every depth. Gemma 3 specifics (q_norm/k_norm, pre/post feedforward RMSNorms, per-layer local vs global RoPE, sliding-window attention) are all honored.

transformers_patch is imported at module load so the expert path picks up adaptive RMSNorm and _gated_residual. The Gemma 3 backbone remains stock — its layer-norms return a plain tensor and are used without a cond= argument.

Functions

apply_rope(x, positions[, max_wavelength])

Applies RoPE to x with the given positions and base wavelength.

Classes

Gemma3WithExpertConfig([gemma3_config, ...])	Configuration wrapper bundling a Gemma 3 VLM config and a Gemma-v1 expert config.
Gemma3WithExpertModel(config)	Gemma 3 VLM interleaved layer-wise with a Gemma-v1 action expert.
InterleavedDecoderLayer(backbone_layer, ...)	One step of pi07's interleaved Gemma 3 backbone + Gemma-v1 expert decoder.

class opentau.policies.pi07.gemma3_with_expert.Gemma3WithExpertConfig(gemma3_config: dict | None = None, gemma_expert_config: dict | None = None, freeze_vision_encoder: bool = True, train_expert_only: bool = True, attention_implementation: str = 'eager', load_pretrained_gemma3: bool = False, discrete_action_vocab_size: int | None = None, dropout: float = 0.1, gradient_checkpointing: bool = False, disable_action_expert: bool = False, disable_internal_bf16_cast: bool = False, **kwargs)

Bases: PretrainedConfig

Configuration wrapper bundling a Gemma 3 VLM config and a Gemma-v1 expert config.

__init__(gemma3_config: dict | None = None, gemma_expert_config: dict | None = None, freeze_vision_encoder: bool = True, train_expert_only: bool = True, attention_implementation: str = 'eager', load_pretrained_gemma3: bool = False, discrete_action_vocab_size: int | None = None, dropout: float = 0.1, gradient_checkpointing: bool = False, disable_action_expert: bool = False, disable_internal_bf16_cast: bool = False, **kwargs)

Initializes the configuration.

Parameters:

• gemma3_config – Optional Gemma 3 config dict. Defaults to the google/gemma-3-4b-pt topology.

• gemma_expert_config – Optional Gemma-v1 action-expert config dict. Defaults to a ~860M-parameter Gemma with AdaRMS enabled.

• freeze_vision_encoder – Freeze the SigLIP tower during training.

• train_expert_only – Only update the expert and its heads.

• attention_implementation – “eager”, “sdpa”, or “fa2”. “fa2” is not implemented and falls back to eager with a warning. “sdpa” dispatches to torch.nn.functional.scaled_dot_product_attention; see the per-layer note about Gemma 3’s interleaved local/global pattern in forward() — π0.7 deliberately keeps the same block-causal mask at every layer, so the SDPA call sees a regular bool mask and takes the standard fused path.

• load_pretrained_gemma3 – Whether to pull pretrained Gemma 3 weights from the Hub (only recommended when He-initializing the expert).

• discrete_action_vocab_size – FAST tokenizer vocab size.

• dropout – Dropout probability applied in the per-layer loop.

• gradient_checkpointing – Wrap each interleaved decoder-layer body in torch.utils.checkpoint.checkpoint during training. Trades roughly one extra forward pass per step (~25-33% compute) for a large slice of activation memory per rank, enabling larger per-rank batch sizes. Only safe under plain DDP (MULTI_GPU), single-process (NO), or DeepSpeed ZeRO-1/2 — see the train.py guard. Defaults to False.

• disable_action_expert – When True, skip instantiating the Gemma-v1 action expert entirely (self.gemma_expert is None). Use this for callers that never feed the expert stream (e.g. the high-level planner, which always passes inputs_embeds=[prefix_embs, None]) — saves ~860M parameters of dead weight on disk and in memory. forward will raise if a non-None expert input is provided when the expert is disabled. Incompatible with train_expert_only=True. Defaults to False.

• disable_internal_bf16_cast – When True, skip the in-__init__ call to to_bfloat16_like_physical_intelligence so the model’s params remain in their constructed dtype (fp32 on a fresh instance). The intended consumer is the FSDP path, which needs fp32 outer params so MixedPrecision( param_dtype=bf16, reduce_dtype=bf16) can downcast on the fly while keeping the fp32 outer copy for AdamW to step on (mirrors DeepSpeed’s bf16-live / fp32-master regime). Other backends (DDP, single, DeepSpeed ZeRO-1/2) still want the cast applied here, then layer fp32 master back on top via MasterWeightOptimizer / BF16_Optimizer. Defaults to False.

• **kwargs – Passed to PretrainedConfig.

model_type: str = 'Gemma3WithExpertModel'

sub_configs: dict[str, type['PretrainedConfig']] = {'gemma3_config': <class 'transformers.models.auto.configuration_auto.AutoConfig'>, 'gemma_expert_config': <class 'transformers.models.auto.configuration_auto.AutoConfig'>}

class opentau.policies.pi07.gemma3_with_expert.Gemma3WithExpertModel(config: Gemma3WithExpertConfig)

Bases: PreTrainedModel

Gemma 3 VLM interleaved layer-wise with a Gemma-v1 action expert.

__init__(config: Gemma3WithExpertConfig)

Args: config ([PretrainedConfig]):

Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [~PreTrainedModel.from_pretrained] method to load the model weights.

config_class

alias of Gemma3WithExpertConfig

eager_attention_forward(attention_mask: Tensor, batch_size: int, head_dim: int, query_states: Tensor, key_states: Tensor, value_states: Tensor, scaling: float | None = None) → Tensor

Standard eager scaled-dot-product attention. attention_mask is a boolean 2D mask of shape (B, Q, K) (True = attend).

embed_discrete_actions(actions: Tensor) → Tensor

embed_image(image: Tensor) → Tensor

Runs the SigLIP tower + multimodal projector to obtain image tokens.

Mirrors Gemma3ForConditionalGeneration.get_image_features: feed pixel_values through the vision tower and run multi_modal_projector on the resulting last_hidden_state, returning a plain (B, mm_tokens_per_image, text_hidden_size) tensor.

When the vision tower has been wrapped with space-time attention by SpaceTimeSiglipVideoEncoder (low-level component), suppress the temporal sublayer here — single-image inputs have no time axis to attend over. The context manager is a no-op when no wrappers are present.

embed_language_tokens(tokens: Tensor) → Tensor

Embed token ids through Gemma 3’s shared text embedding table.

forward(attention_mask: Tensor | None = None, position_ids: LongTensor | None = None, past_key_values: list[FloatTensor] | Cache | None = None, inputs_embeds: list[FloatTensor] | None = None, n_cross_att_tokens: int | None = None, use_cache: bool | None = None, fill_kv_cache: bool | None = None, adarms_cond: list[Tensor] | None = None) → tuple[list[FloatTensor | None], list[FloatTensor] | Cache | None]

Interleaved per-layer forward for the Gemma 3 backbone and Gemma-v1 expert.

The two streams (index 0 = backbone, index 1 = expert) share each layer’s attention — queries and KVs are concatenated along the sequence axis. When one stream’s embeddings are None the other runs alone, pulling KVs for the missing stream from past_key_values when use_cache=True.

Parameters:

• attention_mask – 2D boolean mask of shape (B, Q, K_total). See opentau.policies.pi05.modeling_pi05.make_att_2d_masks.

• position_ids – (B, L_total) token positions, used for RoPE.

• past_key_values – Per-layer KV cache populated on a previous call.

• inputs_embeds – [backbone_embeds, expert_embeds]. Either may be None.

• n_cross_att_tokens – Number of prefix tokens to retain in the cache (must be provided when fill_kv_cache=True).

• use_cache – Read KVs from past_key_values (prefix cross-attention).

• fill_kv_cache – Write this call’s KVs into past_key_values.

• adarms_cond – Per-stream AdaRMS conditioning tensors [None, cond].

Returns:

A pair (outputs_embeds, past_key_values) where outputs_embeds is a two-element list mirroring the inputs_embeds layout.

get_attention_interface()

Returns the attention implementation function based on config.

Dispatches on self.config.attention_implementation:

• "eager": per-layer matmul-softmax-matmul in fp32 (historical default; see eager_attention_forward).

• "sdpa": torch.nn.functional.scaled_dot_product_attention which on A100 + bf16 dispatches to FlashAttention-2 / mem-efficient backends. Note π0.7 keeps the same block-causal mask at every layer (sliding window deliberately not enforced — see the comment near apply_rope), so SDPA sees a regular bool mask and does not need a per-layer mask shape branch.

• "fa2": accepted for backward compatibility; falls back to eager with a warning emitted at config validation time.

sdpa_attention_forward(attention_mask: Tensor, batch_size: int, head_dim: int, query_states: Tensor, key_states: Tensor, value_states: Tensor, scaling: float | None = None) → Tensor

SDPA attention forward pass using F.scaled_dot_product_attention.

Same output shape and semantics as eager_attention_forward but delegates the scores-softmax-matmul chain to PyTorch’s fused SDPA kernel. On A100 + bf16 PyTorch typically dispatches to FlashAttention-2. Q/K are not upcast to float32 (modern attention kernels accumulate the softmax in fp32 internally); training dynamics match eager within bf16 reassociation noise.

Parameters:

• attention_mask – Boolean mask of shape (B, Q, K_total); True = attend. π0.7 keeps the same block-causal mask at every layer.

• batch_size – Batch size.

• head_dim – Per-head dimension.

• query_states – (B, Q, num_attention_heads, head_dim).

• key_states – (B, K_total, num_key_value_heads, head_dim).

• value_states – (B, K_total, num_key_value_heads, head_dim).

• scaling – Override for the QK scaling factor. Gemma 3 uses query_pre_attn_scalar ** -0.5 rather than the default head_dim ** -0.5; the caller passes this through.

Returns:

Attention output of shape (B, Q, num_attention_heads * head_dim).

Return type:

torch.Tensor

set_requires_grad() → None

to_bfloat16_like_physical_intelligence() → None

train(mode: bool = True)

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

class opentau.policies.pi07.gemma3_with_expert.InterleavedDecoderLayer(backbone_layer: Module, expert_layer: Module | None, layer_type: str, rope_local: float, rope_global: float, query_pre_attn_scaling: float, head_dim: int, dropout: Module, attention_interface_provider)

Bases: Module

One step of pi07’s interleaved Gemma 3 backbone + Gemma-v1 expert decoder.

Bundles the backbone layer and expert layer for one decoder step into a single nn.Module so that FSDP / ZeRO-3 see a coherent forward unit. Their all-gather hooks fire on this module’s forward and prefetch every sub-component (input_layernorm, self_attn.{q,k,v}_proj, mlp, …) at once. The pre-refactor design called per-sub-component on the wrapped Gemma3DecoderLayer directly (layer.input_layernorm(...)), which bypassed FSDP’s hooks and hung at NCCL all-gather with mismatched sizes across ranks (different ranks would request different params).

This module owns its backbone and expert layers as submodules. The enclosing Gemma3WithExpertModel empties the original gemma3.language_model.model.layers and gemma_expert.model.layers after handing them off so each parameter is registered exactly once and FSDP / ZeRO-3 don’t double-shard.

Parameters:

• backbone_layer – One Gemma3DecoderLayer from the Gemma 3 text decoder.

• expert_layer – The corresponding GemmaDecoderLayer from the Gemma-v1 action expert, or None when Gemma3WithExpertConfig.disable_action_expert is set. Single-stream callers (e.g. the high-level planner) always pass inputs_embeds[1] is None and forward short-circuits before indexing the expert side, so a None placeholder is safe.

• layer_type – Per-layer attention type from gemma3_config.text_config.layer_types — "sliding_attention" selects rope_local; anything else uses rope_global.

• rope_local – RoPE base for sliding-window layers.

• rope_global – RoPE base for global layers.

• query_pre_attn_scaling – Pre-softmax scaling factor (typically head_dim ** -0.5).

• head_dim – Per-head dimension shared by both streams.

• dropout – Dropout module applied to attention output and MLP output.

• attention_interface_provider – Zero-arg callable that returns the current attention dispatch function (parent.get_attention_interface()). Resolved at every forward rather than captured at init so tests / runtime code that monkey-patches the parent’s eager_attention_forward / sdpa_attention_forward see their patches honoured. Stored as a plain callable attribute — not a submodule — so it doesn’t appear in state_dict.

__init__(backbone_layer: Module, expert_layer: Module | None, layer_type: str, rope_local: float, rope_global: float, query_pre_attn_scaling: float, head_dim: int, dropout: Module, attention_interface_provider)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs_embeds: list[FloatTensor | None], attention_mask: Tensor | None, position_ids: LongTensor | None, past_key_values: dict | None, n_cross_att_tokens: int | None, use_cache: bool | None, fill_kv_cache: bool | None, adarms_cond: list[Tensor | None], batch_size: int, layer_idx: int) → list[FloatTensor | None]

Run one interleaved decoder layer.

Body extracted from Gemma3WithExpertModel._run_layer. The KV-cache write is idempotent under gradient-checkpointing recompute because each call writes the same layer_idx key with the same K/V.

opentau.policies.pi07.gemma3_with_expert.apply_rope(x: Tensor, positions: Tensor, max_wavelength: float = 10000.0) → Tensor

Applies RoPE to x with the given positions and base wavelength.

Parameters:

• x – Tensor of shape (B, L, H, D).

• positions – Tensor of shape (B, L).

• max_wavelength – RoPE base frequency. Gemma 3 uses 10_000 for sliding (local) layers and 1_000_000 for full (global) layers; Gemma-v1 expert uses 10_000.

Returns:

RoPE-transformed tensor, same shape / dtype as the input.