diff --git a/.python-version b/.python-version
new file mode 100644
index 0000000..24ee5b1
--- /dev/null
+++ b/.python-version
@@ -0,0 +1 @@
+3.13
diff --git a/.vscode/settings.json b/.vscode/settings.json
index 5036144..7c4aa46 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -22,7 +22,7 @@
"editor.tabSize": 2,
"editor.formatOnSave": true,
"editor.detectIndentation": false,
- "editor.defaultFormatter": "ms-python.black-formatter",
+ "editor.defaultFormatter": "charliermarsh.ruff",
},
"black-formatter.path": ["uvx", "--python=>=3.9,!=3.12.5", "pyink"],
"pylint.enabled": true,
diff --git a/PR_PLAN.md b/PR_PLAN.md
new file mode 100644
index 0000000..2ac33b3
--- /dev/null
+++ b/PR_PLAN.md
@@ -0,0 +1,425 @@
+# Gemma 4 Support for Penzai — Full Implementation Plan
+
+## Context
+
+Gemma 4 introduces four model variants with significant architectural novelty
+compared to Gemma 3. This plan adds support for all four, phased so each phase
+is independently shippable. Checkpoints are loaded from the Flax format
+(google-deepmind/gemma library).
+
+### Variants at a glance
+
+| Variant | Type | Layers | Embed | New features over Gemma 3 |
+|---------|------|--------|-------|---------------------------|
+| 31B | Dense | 60 | 5376 | Per-layer proj dim, partial RoPE, K=V sharing, value norm, skip scale |
+| E4B | Dense+PLE | 42 | 2560 | + Per-Layer Embeddings, KV cache sharing |
+| E2B | Dense+PLE | 35 | 1536 | + Per-Layer Embeddings, KV cache sharing |
+| 26B-A4B | MoE | 30 | 2816 | + 128-expert MoE with top-8 routing, dual-branch FFW |
+
+### Key architectural differences from Gemma 3
+
+1. Different `projection_dim` for local (256) vs global (512) attention layers
+2. Different `num_kv_heads` for local vs global layers (31B: 16 local / 4 global; 26B: 8/2)
+3. Partial RoPE: only 25% of head dims get RoPE in global layers (`ApplyRoPEToSubset`)
+4. K=V sharing: global layers share K and V projections (31B, 26B)
+5. Value normalization: RMSNorm on values (in addition to Q/K norm from Gemma 3)
+6. Learnable skip scale on the attention residual connection
+7. Per-Layer Embeddings (E2B, E4B): separate per-layer token embedding table with gating
+8. KV cache sharing (E2B, E4B): later layers reuse KV projections from earlier layers
+9. MoE (26B-A4B): 128 experts with top-8 routing + dense shared MLP branch
+
+---
+
+## Phase A: Core Infrastructure + gemma4\_31b ✅ DONE
+
+All changes landed and tested (16/16 tests pass, 0 new type errors).
+
+### A1. New config fields on `LlamalikeTransformerConfig`
+
+**File:** `penzai/models/transformer/variants/llamalike_common.py`
+
+Added six backward-compatible optional fields:
+
+```python
+global_projection_dim: int | None = None
+global_num_kv_heads: int | None = None
+global_rope_proportion: float | None = None
+use_value_norm: bool = False
+use_skip_scale: bool = False
+k_eq_v_global: bool = False
+```
+
+### A2. Refactored `_head_info` to be parameterised
+
+Changed `_head_info(config)` → `_head_info(num_kv_heads, query_head_multiplier)`
+so it can be called with per-layer effective values. Updated both call sites.
+
+### A3. Per-layer dimensions in `build_llamalike_attention`
+
+After resolving `attention_type` per layer, the function now also resolves:
+- `projection_dim` → `global_projection_dim` for global layers
+- `effective_num_kv_heads` → `global_num_kv_heads` for global layers
+- `effective_query_head_multiplier` derived from the fixed total query head count
+
+All `Linear.from_config` calls, NamedEinsum shapes, and RoPE layers within the
+function use these effective values.
+
+### A4. Partial RoPE for global layers
+
+When `global_rope_proportion` is set and < 1.0, global layers use
+`pz.nn.ApplyRoPEToSubset` (already in `penzai/nn/embeddings.py`) instead of
+`pz.nn.ApplyRoPE`. Local layers are unaffected.
+
+### A5. Value normalization
+
+When `use_value_norm=True`, an `RMSLayerNorm` is appended to `input_to_value`,
+following the same pattern as QK norm.
+
+### A6. K=V sharing for global layers
+
+When `k_eq_v_global=True` and the layer is global, one `pz.nn.Linear` object is
+created and referenced in both `input_to_key` and `input_to_value` pipelines —
+exactly the embedding-tying pattern used for tied embeddings.
+
+### A7. `ScaledResidual` combinator
+
+**File:** `penzai/nn/combinators.py`
+
+New `@struct.pytree_dataclass` layer:
+
+```python
+class ScaledResidual(layer_base.Layer):
+ delta: layer_base.Layer
+ scale: parameters.ParameterLike
+
+ def __call__(self, value, **side_inputs):
+ return self.delta(value, **side_inputs) + value * self.scale.value
+```
+
+Exported from `penzai/pz/nn.py`.
+
+### A8. `ScaledResidual` in `build_llamalike_block`
+
+When `use_skip_scale=True`, the attention residual uses `ScaledResidual` with a
+learnable scalar parameter instead of `pz.nn.Residual`.
+
+### A9. Per-layer cache axes in `sampling_mode.py`
+
+`KVCachingTransformerLM.from_uncached` now infers `cached_axes` per attention
+layer by inspecting the key `Linear`'s `output_axes`, instead of using a single
+global value from metadata. Backward-compatible for uniform models.
+
+### A10–A12. gemma4\_31b preset, checkpoint loader, docstring
+
+- Added `gemma4_31b` to `_GEMMA_PRESETS` and `_NEEDS_GATING_TRANSPOSE`.
+- Updated `preset_name` Literal type.
+- Auto-detection now checks for `skip_scale` to distinguish Gemma 4 from 3.
+- Checkpoint loader handles value norm, skip scale, K=V sharing, and per-layer
+ attention weight shapes (branches on `attention_type[i]`).
+- Updated module docstring to reference Gemma 4.
+
+### Tests
+
+Two new parameterised test cases (`like_gemma4_31b`) added to
+`tests/models/transformer_llamalike_test.py`:
+- `test_build_and_run_gemma_like_gemma4_31b` — build + forward pass via `jax.eval_shape`
+- `test_build_and_run_sampling_mode_like_gemma4_31b` — sampling mode + decoding loop
+
+All 16 tests pass (14 existing + 2 new).
+
+---
+
+## Phase B: gemma4\_E2B + gemma4\_E4B (Per-Layer Embeddings + KV Cache Sharing) ✅ DONE
+
+All changes landed and tested (18/18 tests pass, 0 new type errors).
+
+### B1. Per-Layer Embedding layer
+
+**File:** `penzai/models/transformer/variants/gemma.py`
+
+New `@struct.pytree_dataclass` layer `PerLayerEmbeddingInject`:
+
+```python
+@struct.pytree_dataclass
+class PerLayerEmbeddingInject(layer_base.Layer):
+ """Injects per-layer embedding information into the residual stream.
+
+ Reads token IDs from side inputs, looks up per-layer embeddings for
+ the current layer, gates them with the current hidden state, projects
+ back to embedding dimension, and adds to the input.
+
+ Attributes:
+ per_layer_table: Shared embedding table for all layers.
+ layer_index: Which layer slice to use from the table.
+ gate: Linear projection from embedding_dim to per_layer_dim (with GELU).
+ projection: Linear projection from per_layer_dim back to embedding_dim.
+ post_norm: RMSLayerNorm applied after projection.
+ per_layer_input_dim: Dimension of per-layer embeddings.
+ token_ids_input_name: Side input key for token IDs.
+ """
+```
+
+**Flow per block:**
+1. Read `token_ids` from side inputs
+2. Index into shared PLE table: `table[{vocabulary: token_ids, layers: layer_index}]`
+3. Scale by `sqrt(per_layer_input_dim)`
+4. Gate: `gelu(gate_linear(hidden_state)) * ple_embedding`
+5. Project back: `projection_linear(gated)`
+6. Normalise and add to residual
+
+### B2. Pass `token_ids` as side input
+
+**File:** `penzai/models/transformer/model_parts.py`
+
+Add `token_ids=tokens` to the side inputs in `TransformerLM.__call__`:
+
+```python
+return self.body(tokens, token_positions=token_positions,
+ token_ids=tokens, **side_inputs)
+```
+
+Backward-compatible: existing layers ignore unknown side inputs.
+
+### B3. KV cache sharing
+
+For E2B/E4B, the last N layers reuse K/V projections from earlier layers.
+
+**Training mode:** Share the same `pz.nn.Linear` Python objects between source
+and target layers (same pattern as embedding tying). Both layers compute the
+projection; the weights are shared by identity.
+
+**Sampling mode:** Shared-KV layers must point to the same `StateVariable` cache.
+Update `sampling_mode.py` to detect when two `Attention` layers share the same
+key `Linear` (by object identity or label) and wire them to the same cache.
+
+### B4. E2B and E4B presets + checkpoint loading
+
+```python
+"gemma4_E2B": dict(
+ num_decoder_blocks=35, vocab_size=262_144,
+ num_kv_heads=1, query_head_multiplier=8,
+ embedding_dim=1536, projection_dim=256, global_projection_dim=512,
+ mlp_hidden_dim=6_144,
+ per_layer_input_dim=256, num_kv_shared_layers=20,
+ use_qk_norm=True, use_value_norm=True,
+ use_post_attn_norm=True, use_post_ffw_norm=True,
+ use_skip_scale=True, global_rope_proportion=0.25,
+ rope_wavelength=1_000_000, local_rope_wavelength=10_000,
+ attention_type=..., # 5×SW(512) + 1×Global
+),
+"gemma4_E4B": dict(
+ num_decoder_blocks=42, vocab_size=262_144,
+ num_kv_heads=2, query_head_multiplier=4,
+ embedding_dim=2560, projection_dim=256, global_projection_dim=512,
+ mlp_hidden_dim=10_240,
+ per_layer_input_dim=256, num_kv_shared_layers=18,
+ use_qk_norm=True, use_value_norm=True,
+ use_post_attn_norm=True, use_post_ffw_norm=True,
+ use_skip_scale=True, global_rope_proportion=0.25,
+ rope_wavelength=1_000_000, local_rope_wavelength=10_000,
+ attention_type=..., # 5×SW(512) + 1×Global
+),
+```
+
+Checkpoint loader extended for PLE weights:
+- `embedder/per_layer_embeddings/w` — `[vocab, layers, ple_dim]`
+- `embedder/per_layer_model_projection/w`
+- `embedder/per_layer_projection_norm/scale`
+- Per block: `layer_N/per_layer_input_gate/w`, `layer_N/per_layer_projection/w`,
+ `layer_N/post_per_layer_input_norm/scale`
+
+### B5. Tests
+
+- Build + forward pass for E2B and E4B configs via `jax.eval_shape`
+- Sampling mode test with KV cache sharing
+- Verify PLE layers receive token IDs and produce correct shapes
+
+---
+
+## Phase C: gemma4\_26b\_a4b (Mixture of Experts) ✅ DONE
+
+All changes landed and tested (20/20 tests pass, 0 new type errors).
+
+### C1. MoE dataflow combinator
+
+**File:** `penzai/nn/mixture_of_experts.py` (new)
+
+Following the `Attention` pattern — a single `@struct.pytree_dataclass` that
+orchestrates dataflow between named sub-components:
+
+```python
+@struct.pytree_dataclass
+class MixtureOfExperts(layer_base.Layer):
+ """A mixture-of-experts dataflow combinator.
+
+ Routes input tokens to a subset of experts via a learned router, runs
+ the selected experts, and combines their outputs with routing weights.
+ Follows the same design philosophy as Attention.
+
+ Attributes:
+ input_to_routing: Maps input to (routing_weights, expert_indices).
+ input_and_routing_to_output: Maps (input, routing_weights,
+ expert_indices) to combined expert output.
+ """
+ input_to_routing: layer_base.Layer
+ input_and_routing_to_output: layer_base.Layer
+
+ def __call__(self, x, **side_inputs):
+ routing_weights, expert_indices = self.input_to_routing(x, **side_inputs)
+ output = self.input_and_routing_to_output(
+ (x, routing_weights, expert_indices), **side_inputs
+ )
+ return output
+```
+
+### C2. Router and expert computation layers
+
+**File:** `penzai/nn/mixture_of_experts.py`
+
+**`MoETopKRouter`** — RMS standardise → learned scale → project to expert
+logits → softmax → top-k → renormalise. Returns `(routing_weights, expert_indices)`.
+
+```python
+@struct.pytree_dataclass
+class MoETopKRouter(layer_base.Layer):
+ router_norm: layer_base.Layer
+ router_scale: parameters.ParameterLike[named_axes.NamedArray]
+ router_logits: layer_base.Layer
+ num_selected_experts: int # metadata, not pytree node
+```
+
+**`MoEGatedExpertComputation`** — Gathers selected expert weights, computes
+GeGLU MLP per expert, scales by per-expert factors, combines with routing
+weights. Expert weights stored as batched NamedArrays with `"experts"` axis.
+
+```python
+@struct.pytree_dataclass
+class MoEGatedExpertComputation(layer_base.Layer):
+ gating_weights: parameters.ParameterLike[named_axes.NamedArray]
+ value_weights: parameters.ParameterLike[named_axes.NamedArray]
+ out_weights: parameters.ParameterLike[named_axes.NamedArray]
+ per_expert_scale: parameters.ParameterLike[named_axes.NamedArray]
+ activation_fn: Any # metadata, not pytree node
+```
+
+Initial implementation uses gather-based dispatch (portable, inspectable).
+Can optimise with `jax.lax.ragged_dot` later.
+
+### C3. `TransformerMoEFeedForward`
+
+**File:** `penzai/models/transformer/model_parts.py`
+
+```python
+@pz.pytree_dataclass(has_implicitly_inherited_fields=True)
+class TransformerMoEFeedForward(pz.nn.Sequential):
+ """Informatively-named Sequential subclass for MoE + dense feedforward."""
+```
+
+### C4. MoE builder + block integration
+
+**File:** `penzai/models/transformer/variants/llamalike_common.py`
+
+New config fields:
+
+```python
+num_experts: int | None = None
+num_selected_experts: int | None = None
+expert_hidden_dim: int | None = None
+dense_moe_shared_hidden_dim: int | None = None
+```
+
+New `build_moe_feedforward()` assembles the dual-branch structure:
+
+```python
+TransformerMoEFeedForward(sublayers=[
+ pz.nn.BranchAndAddTogether(branches=[
+ pz.nn.NamedGroup("dense_branch", [
+ RMSLayerNorm(pre_ffw2_norm),
+ build_llamalike_feedforward(mlp2, hidden=2112),
+ RMSLayerNorm(post_ffw2_norm),
+ ]),
+ pz.nn.NamedGroup("moe_branch", [
+ RMSLayerNorm(pre_ffw_norm),
+ MixtureOfExperts(router, experts),
+ RMSLayerNorm(post_ffw1_norm),
+ ]),
+ ]),
+ RMSLayerNorm(post_ffw_norm),
+])
+```
+
+In `build_llamalike_block`, when `config.num_experts is not None`, use
+`build_moe_feedforward()` in place of `build_llamalike_feedforward()`.
+
+### C5. 26B-A4B preset + checkpoint loading
+
+```python
+"gemma4_26b_a4b": dict(
+ num_decoder_blocks=30, vocab_size=262_144,
+ num_kv_heads=8, global_num_kv_heads=2, query_head_multiplier=2,
+ embedding_dim=2816, projection_dim=256, global_projection_dim=512,
+ mlp_hidden_dim=2_112,
+ num_experts=128, num_selected_experts=8, expert_hidden_dim=704,
+ dense_moe_shared_hidden_dim=2_112,
+ k_eq_v_global=True,
+ use_qk_norm=True, use_value_norm=True,
+ use_post_attn_norm=True, use_post_ffw_norm=True,
+ use_skip_scale=True, global_rope_proportion=0.25,
+ rope_wavelength=1_000_000, local_rope_wavelength=10_000,
+ attention_type=..., # 5×SW(1024) + 1×Global
+),
+```
+
+Checkpoint loader maps MoE weights per layer:
+- `layer_N/mlp/router_logits/w` — `[embed_dim, num_experts]`
+- `layer_N/mlp/router_scale` — `[embed_dim]`
+- `layer_N/mlp/router_norm/scale` — `[embed_dim]`
+- `layer_N/mlp/gating_einsum/w` — `[num_experts, 2, expert_dim, embed_dim]`
+ (split into gating + value along dim 1)
+- `layer_N/mlp/linear/w` — `[num_experts, expert_dim, embed_dim]`
+- `layer_N/mlp/per_expert_scale` — `[num_experts]`
+- `layer_N/mlp2/gating_einsum/w`, `layer_N/mlp2/linear/w` (dense shared MLP)
+
+### C6. Export new layers
+
+**File:** `penzai/pz/nn.py`
+
+```python
+from penzai.nn.mixture_of_experts import (
+ MixtureOfExperts,
+ MoETopKRouter,
+ MoEGatedExpertComputation,
+)
+```
+
+### C7. Tests
+
+- Build + forward pass for 26B-A4B config via `jax.eval_shape`
+- Sampling mode test
+- Unit tests for `MoETopKRouter` and `MoEGatedExpertComputation` in isolation
+
+---
+
+## Files Modified (complete)
+
+| File | Phase | Changes |
+|------|-------|---------|
+| `penzai/models/transformer/variants/llamalike_common.py` | A ✅, C | Config, `_head_info`, attention builder, block builder, MoE builder |
+| `penzai/models/transformer/variants/gemma.py` | A ✅, B | Presets, auto-detect, checkpoint loader, PLE layers |
+| `penzai/nn/combinators.py` | A ✅ | `ScaledResidual` |
+| `penzai/nn/mixture_of_experts.py` | C | **New**: `MixtureOfExperts`, `MoETopKRouter`, `MoEGatedExpertComputation` |
+| `penzai/models/transformer/model_parts.py` | B, C | `token_ids` side input, `TransformerMoEFeedForward` |
+| `penzai/models/transformer/sampling_mode.py` | A ✅ | Per-layer cache axes |
+| `penzai/pz/nn.py` | A ✅, C | Exports |
+| `tests/models/transformer_llamalike_test.py` | A ✅, B, C | New test cases |
+
+## Verification (per phase)
+
+1. **Shape test** — Build each variant with a small test config via `jax.eval_shape`;
+ verify output shape `{"batch": B, "seq": S, "vocabulary": 262144}`
+2. **Forward pass** — Run with random weights and small dims; assert no NaN/shape errors
+3. **Sampling mode** — Convert to `KVCachingTransformerLM`; verify per-layer cache dims
+4. **Backward compat** — Existing Gemma 1/2/3 parameterised tests still pass unchanged
+5. **Type check** — `pyright` on changed files; no new errors beyond pre-existing patterns
+6. **Checkpoint** (if available) — Load real Flax weights, run forward pass, compare
+ with reference implementation
diff --git a/notebooks/gemma_4b.ipynb b/notebooks/gemma_4b.ipynb
new file mode 100644
index 0000000..2b5322c
--- /dev/null
+++ b/notebooks/gemma_4b.ipynb
@@ -0,0 +1,712 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "id": "c8df759c",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import dotenv\n",
+ "import os"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "id": "5247635c",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "True"
+ ]
+ },
+ "execution_count": 2,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "dotenv.load_dotenv(\"/Users/johannesmosig/projects/.env\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "id": "d50fbefc",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import penzai\n",
+ "from __future__ import annotations\n",
+ "from typing import Any\n",
+ "\n",
+ "import os\n",
+ "import dataclasses\n",
+ "import gc\n",
+ "\n",
+ "import jax\n",
+ "import jax.numpy as jnp\n",
+ "import numpy as np\n",
+ "import orbax.checkpoint\n",
+ "from jax.experimental import mesh_utils"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "id": "31f8bdc7",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from pathlib import Path"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "id": "bc195be9",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import treescope\n",
+ "import penzai\n",
+ "from penzai import pz\n",
+ "\n",
+ "from penzai.models import transformer"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 6,
+ "id": "71266e76",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "treescope.basic_interactive_setup(autovisualize_arrays=True)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 7,
+ "id": "9c54c173",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "
"
+ ],
+ "text/plain": [
+ ""
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ ""
+ ],
+ "text/plain": [
+ "'0.1.10'"
+ ]
+ },
+ "execution_count": 7,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "treescope.__version__"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 8,
+ "id": "81f62f44",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Kaggle credentials set.\n"
+ ]
+ }
+ ],
+ "source": [
+ "import kagglehub\n",
+ "try:\n",
+ " kagglehub.config.set_kaggle_credentials(\n",
+ " os.getenv(\"KAGGLE_USERNAME\"), os.getenv(\"KAGGLE_KEY\")\n",
+ " )\n",
+ "except ImportError:\n",
+ " kagglehub.login()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 9,
+ "id": "c2a04b88",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "weights_dir = kagglehub.model_download('google/gemma-4/Flax/gemma-4-e2b')"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 10,
+ "id": "938af6d0",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "\u001b[01;34m/Users/johannesmosig/.cache/kagglehub/models/google/gemma-4/Flax/gemma-4-e2b/1\u001b[0m\n",
+ "├── \u001b[00m_CHECKPOINT_METADATA\u001b[0m\n",
+ "├── \u001b[00m_METADATA\u001b[0m\n",
+ "├── \u001b[01;34md\u001b[0m\n",
+ "│ └── \u001b[00m83b3800e187e664882a289c5499a7728\u001b[0m\n",
+ "├── \u001b[01;34mdescriptor\u001b[0m\n",
+ "│ ├── \u001b[00mdescriptor.pbtxt\u001b[0m\n",
+ "│ └── \u001b[00muuid-4814dfcd-ec60-45b7-9f57-4e270c6d1c85\u001b[0m\n",
+ "├── \u001b[00mmanifest.ocdbt\u001b[0m\n",
+ "└── \u001b[01;34mocdbt.process_0\u001b[0m\n",
+ " ├── \u001b[01;34md\u001b[0m\n",
+ " │ ├── \u001b[00m00721b4cc878d4bcdf1713d373c1a318\u001b[0m\n",
+ " │ ├── \u001b[00m0ab15bbc15c4b9ec2d9486c666103a9a\u001b[0m\n",
+ " │ ├── \u001b[00m0e39f68d27d5c6f613e126f0f42f4c9b\u001b[0m\n",
+ " │ ├── \u001b[00m3a77b1262bbb5c8126656da1bc19942d\u001b[0m\n",
+ " │ ├── \u001b[00m43fdd825eb961d17f67733db4703ec3c\u001b[0m\n",
+ " │ ├── \u001b[00m6f8efca8011fca9525ba70fda5f9b11f\u001b[0m\n",
+ " │ ├── \u001b[00m85ccff2e077eaf6b2b8b6dd12419bf93\u001b[0m\n",
+ " │ ├── \u001b[00m8c2c8f9ec97770675270ded4fcd25ad4\u001b[0m\n",
+ " │ ├── \u001b[00ma3398a043d32e3390cbd768b4692cafe\u001b[0m\n",
+ " │ ├── \u001b[00maca38ba88ebfb0984d6721ba48840540\u001b[0m\n",
+ " │ ├── \u001b[00mb28f3738d864f8b1b70de809e9af7ec7\u001b[0m\n",
+ " │ ├── \u001b[00mbb923f7506f912f4598e7b44dfe95028\u001b[0m\n",
+ " │ ├── \u001b[00me5bb41dc61ba4c04835e3c407d1ba68f\u001b[0m\n",
+ " │ ├── \u001b[00me6d3eb78d0220ec4e3d039c6434cd5be\u001b[0m\n",
+ " │ ├── \u001b[00mf04604465ca526bec1e3efa0fad17125\u001b[0m\n",
+ " │ └── \u001b[00mf255107803f74f22c01cad0a68d52573\u001b[0m\n",
+ " └── \u001b[00mmanifest.ocdbt\u001b[0m\n",
+ "\n",
+ "5 directories, 23 files\n"
+ ]
+ }
+ ],
+ "source": [
+ "! tree {weights_dir}"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 11,
+ "id": "2b0423b5",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# checkpointer = orbax.checkpoint.PyTreeCheckpointer()\n",
+ "# metadata = checkpointer.metadata(weights_dir)\n",
+ "\n",
+ "checkpointer = orbax.checkpoint.PyTreeCheckpointer()\n",
+ "metadata = checkpointer.metadata(weights_dir)\n",
+ "\n",
+ "single_sharding = jax.sharding.SingleDeviceSharding(jax.devices(\"cpu\")[0])\n",
+ "# restore_args = jax.tree.map(\n",
+ "# lambda _: orbax.checkpoint.ArrayRestoreArgs(sharding=single_sharding),\n",
+ "# metadata,\n",
+ "# is_leaf=lambda x: not isinstance(x, dict),\n",
+ "# )\n",
+ "\n",
+ "flat_params = checkpointer.restore(weights_dir)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 12,
+ "id": "d2b6ad0d",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "['Residual', 'Residual', 'PerLayerEmbeddingInject', 'LearnableRescale']\n"
+ ]
+ },
+ {
+ "ename": "AttributeError",
+ "evalue": "'RMSStandardize' object has no attribute 'delta'",
+ "output_type": "error",
+ "traceback": [
+ "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
+ "\u001b[31mAttributeError\u001b[39m Traceback (most recent call last)",
+ "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[12]\u001b[39m\u001b[32m, line 5\u001b[39m\n\u001b[32m 1\u001b[39m model = transformer.variants.gemma.gemma_from_pretrained_checkpoint(flat_params, upcast_activations_to_float32=\u001b[38;5;28;01mFalse\u001b[39;00m)\n\u001b[32m 2\u001b[39m \u001b[38;5;66;03m# Block sublayer types should now be: Residual, Residual, PerLayerEmbeddingInject, LearnableRescale\u001b[39;00m\n\u001b[32m 3\u001b[39m print([type(s).__name__ \u001b[38;5;28;01mfor\u001b[39;00m s \u001b[38;5;28;01min\u001b[39;00m model.body.sublayers[\u001b[32m3\u001b[39m].sublayers])\n\u001b[32m 4\u001b[39m \u001b[38;5;66;03m# A KV-shared layer (e.g., the last block) should show this in its key path:\u001b[39;00m\n\u001b[32m----> \u001b[39m\u001b[32m5\u001b[39m print([type(s).__name__ \u001b[38;5;28;01mfor\u001b[39;00m s \u001b[38;5;28;01min\u001b[39;00m model.body.sublayers[-\u001b[32m3\u001b[39m].sublayers[\u001b[32m0\u001b[39m].delta.sublayers[\u001b[32m1\u001b[39m].input_to_key.__class__.__bases__])\n",
+ "\u001b[31mAttributeError\u001b[39m: 'RMSStandardize' object has no attribute 'delta'"
+ ]
+ }
+ ],
+ "source": [
+ "model = transformer.variants.gemma.gemma_from_pretrained_checkpoint(flat_params, upcast_activations_to_float32=False)\n",
+ "# Block sublayer types should now be: Residual, Residual, PerLayerEmbeddingInject, LearnableRescale\n",
+ "print([type(s).__name__ for s in model.body.sublayers[3].sublayers])\n",
+ "# A KV-shared layer (e.g., the last block) should show this in its key path:\n",
+ "print([type(s).__name__ for s in model.body.sublayers[-3].sublayers[0].delta.sublayers[1].input_to_key.__class__.__bases__])"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "b7ff0362",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# model = transformer.variants.gemma.gemma_from_pretrained_checkpoint(\n",
+ "# flat_params,\n",
+ "# upcast_activations_to_float32=False\n",
+ "# )"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 13,
+ "id": "03dab4a6",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "
"
+ ],
+ "text/plain": [
+ ""
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ ""
+ ],
+ "text/plain": [
+ "0"
+ ]
+ },
+ "execution_count": 13,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "del flat_params\n",
+ "gc.collect()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 20,
+ "id": "e81cc53c",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "
"
+ ],
+ "text/plain": [
+ ""
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "