Upload gemma4.py with huggingface_hub

gemma4.py

@@ -0,0 +1,652 @@
+"""
+Gemma 4 E2B — clean PyTorch forward pass (text model only).
+
+Architecture:
+  - 35 decoder layers, hidden_size=1536, vocab=262144
+  - 8 Q heads, 1 KV head (MQA)
+  - Sliding attention layers (0-3, 5-8, 10-13, 15-18, 20-23, 25-28, 30-33):
+      head_dim=256, sliding_window=512, rope_theta=10000
+  - Full attention layers (every 5th: 4,9,14,19,24,29,34):
+      head_dim=512, partial_rotary_factor=0.25 (only first 128 of 512 dims rotated),
+      rope_theta=1000000
+  - MLP (all layers): GeGLU, intermediate_size=6144
+  - Per-layer auxiliary stream (full details below)
+  - layer_scalar: per-layer learned scalar multiplied onto residual contributions
+  - QK RMSNorm before RoPE, attn_scale=1.0
+  - Final: RMSNorm + tied lm_head + logit softcapping at 30.0
+
+Per-layer auxiliary stream:
+  Model-level (computed once, before all layers):
+    1. embed_tokens_per_layer(input_ids)          → [B, T, 35*256]  (vocab lookup)
+    2. per_layer_model_projection(x_embed)         → [B, T, 35*256]  (project hidden→aux)
+       scaled by hidden_size**-0.5
+    3. per_layer_projection_norm (RMSNorm(256)) on the projection slice per layer
+    4. Combine: per_layer_inputs = (embed_aux + proj_aux) * (1/sqrt(2))
+       reshaped to [B, T, 35, 256]
+
+  Per-layer (at layer i):
+    per_layer_input_i = per_layer_inputs[:, :, i, :]      # [B, T, 256]
+    x_normed = input_layernorm(x)
+    gate  = sigmoid(per_layer_input_gate(x_normed))       # [B, T, 256]
+    gated = gate * per_layer_input_i                      # [B, T, 256]
+    out   = per_layer_projection(gated)                   # [B, T, 1536]  (256→1536)
+    x     = x + post_per_layer_input_norm(out)
+
+  Weight shapes in checkpoint:
+    per_layer_model_projection.weight : [8960, 1536]   (Linear 1536→8960)
+    per_layer_projection_norm.weight  : [256]           (RMSNorm on 256-dim slices)
+    layers.i.per_layer_input_gate.weight  : [256, 1536] (Linear 1536→256)
+    layers.i.per_layer_projection.weight  : [1536, 256] (Linear 256→1536)
+    layers.i.post_per_layer_input_norm.weight : [1536]  (RMSNorm on 1536-dim output)
+"""
+
+import math
+import os
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from safetensors import safe_open
+from transformers import AutoTokenizer
+
+# ── device / dtype ────────────────────────────────────────────────────────────
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE  = torch.bfloat16
+
+# ── model path ────────────────────────────────────────────────────────────────
+# Try known HF repo caches in order; first one that exists wins. Override with
+# $GEMMA4_HF_REPO to point at an arbitrary repo cache (e.g., "google/gemma-4-e2b-it").
+_HUB_ROOT = Path(os.path.expanduser("~/.cache/huggingface/hub"))
+_REPO_CANDIDATES = (
+    os.environ.get("GEMMA4_HF_REPO", ""),
+    "gg-hf-gg/gemma-4-E2B",
+    "google/gemma-4-e2b-it",
+)
+
+
+def _resolve_model_paths():
+    """Return (snapshot_dir, safetensors_path). Picks first available repo+snapshot
+    that actually contains a .safetensors file. Iterates ALL snapshots per repo
+    before moving to the next repo — iterdir() order is not deterministic and HF
+    may keep multiple snapshots where only one has weights blob-resolved.
+    """
+    for repo in _REPO_CANDIDATES:
+        if not repo:
+            continue
+        repo_cache = _HUB_ROOT / ("models--" + repo.replace("/", "--"))
+        snap_root = repo_cache / "snapshots"
+        if not snap_root.is_dir():
+            continue
+        for snap in sorted(p for p in snap_root.iterdir() if p.is_dir()):
+            # Prefer model.safetensors (single-file) else any .safetensors
+            sft = snap / "model.safetensors"
+            if not sft.exists():
+                candidates = sorted(snap.glob("*.safetensors"))
+                if not candidates:
+                    continue
+                sft = candidates[0]
+            return snap, sft
+    raise FileNotFoundError(
+        "No Gemma-4 E2B HF cache found. Tried: " + ", ".join(r for r in _REPO_CANDIDATES if r)
+        + ". Run `hf download google/gemma-4-e2b-it` or set GEMMA4_HF_REPO."
+    )
+
+
+MODEL_DIR, SAFETENSORS_BLOB = _resolve_model_paths()
+
+# ── architecture constants ────────────────────────────────────────────────────
+N_LAYERS       = 35
+HIDDEN_SIZE    = 1536
+VOCAB_SIZE     = 262144
+N_Q_HEADS      = 8
+N_KV_HEADS     = 1
+HEAD_DIM_SLIDE = 256          # sliding attention head dim
+HEAD_DIM_FULL  = 512          # full attention head dim
+PER_LAYER_DIM  = 256          # per-layer auxiliary stream width per layer
+INTERMEDIATE        = 6144    # MLP intermediate size (layers 0-14)
+INTERMEDIATE_WIDE   = 12288   # double-wide MLP intermediate size (layers 15-34)
+# Layers 15-34 use double-wide MLP (use_double_wide_mlp=True in config)
+DOUBLE_WIDE_START   = 15
+SLIDING_WINDOW = 512
+ROPE_THETA_SLIDE  = 10_000.0
+ROPE_THETA_FULL   = 1_000_000.0
+PARTIAL_ROT_FULL  = 0.25      # only first floor(512*0.25)=128 dims get RoPE
+RMS_EPS           = 1e-6
+LOGIT_CAP         = 30.0
+ATTN_SCALE        = 1.0       # QK are RMSNorm'd, so no sqrt(d) scaling needed
+
+# Per-layer projection scale: hidden_size**-0.5 (applied to per_layer_model_projection output)
+PER_LAYER_PROJ_SCALE = HIDDEN_SIZE ** -0.5
+# Input combination scale: 1/sqrt(2) (mix embed aux + model projection)
+PER_LAYER_INPUT_SCALE = math.sqrt(0.5)  # = 1/sqrt(2)
+
+# Full-attention layers: every 5th layer (0-indexed: 4,9,14,19,24,29,34)
+FULL_ATTN_LAYERS = frozenset(range(4, N_LAYERS, 5))
+
+
+def is_full_attention(layer_idx: int) -> bool:
+    """Return True if layer_idx uses full (global) attention."""
+    return layer_idx in FULL_ATTN_LAYERS
+
+
+# ── RMSNorm ───────────────────────────────────────────────────────────────────
+
+class RMSNorm(nn.Module):
+    """RMSNorm with weight * normed, weight initialized to ones."""
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_f32 = x.float()
+        normed = x_f32 * torch.rsqrt(x_f32.pow(2).mean(-1, keepdim=True) + RMS_EPS)
+        return (normed * self.weight.float()).to(x.dtype)
+
+
+# ── RoPE ─────────────────────────────────────────────────────────────────────
+
+def build_rope_freqs(
+    head_dim: int,
+    max_seq: int,
+    theta: float,
+    device: torch.device,
+    n_rot_pairs: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Build cos/sin tables of shape [max_seq, head_dim].
+
+    For full-attention layers with partial rotation, only the first
+    n_rot_pairs*2 positions carry actual frequencies; the rest are zeros
+    (NoPE — no positional encoding for those dims).
+
+    Args:
+        head_dim:    total head dimension
+        max_seq:     maximum sequence length to precompute
+        theta:       RoPE base frequency
+        device:      target device
+        n_rot_pairs: if set, only compute real freqs for this many pairs;
+                     remaining dims get freq=0 (cos=1, sin=0 → identity).
+    """
+    half = head_dim // 2
+    if n_rot_pairs is None:
+        n_rot_pairs = half
+
+    # Build frequencies only for the pairs that actually rotate
+    inv_freq = 1.0 / (theta ** (
+        torch.arange(0, n_rot_pairs, device=device).float() / half
+    ))  # shape [n_rot_pairs]
+
+    # Pad with zeros for the remaining pairs (NoPE: cos=1, sin=0)
+    if n_rot_pairs < half:
+        inv_freq = torch.cat([
+            inv_freq,
+            torch.zeros(half - n_rot_pairs, device=device),
+        ])  # [half]
+
+    t = torch.arange(max_seq, device=device).float()
+    freqs = torch.outer(t, inv_freq)          # [T, half]
+    freqs = torch.cat([freqs, freqs], dim=-1) # [T, head_dim]
+    return freqs.cos(), freqs.sin()
+
+
+def apply_rope(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Apply rotary embeddings.
+
+    Args:
+        x:   [B, H, T, head_dim]
+        cos: [T, head_dim]  (broadcastable)
+        sin: [T, head_dim]
+    """
+    half = x.shape[-1] // 2
+    x1, x2 = x[..., :half], x[..., half:]
+    rotated = torch.cat([-x2, x1], dim=-1)
+    T = x.shape[2]
+    cos_ = cos[:T].unsqueeze(0).unsqueeze(0).to(x.dtype)  # [1,1,T,D]
+    sin_ = sin[:T].unsqueeze(0).unsqueeze(0).to(x.dtype)
+    return x * cos_ + rotated * sin_
+
+
+# ── Attention ─────────────────────────────────────────────────────────────────
+
+class Attention(nn.Module):
+    """
+    Multi-query attention (8 Q heads, 1 KV head).
+
+    Sliding layers: head_dim=256, local window=512.
+    Full layers:    head_dim=512, causal (no window restriction).
+    """
+
+    def __init__(self, layer_idx: int):
+        super().__init__()
+        self.layer_idx   = layer_idx
+        self.full_attn   = is_full_attention(layer_idx)
+        self.head_dim    = HEAD_DIM_FULL if self.full_attn else HEAD_DIM_SLIDE
+        hd               = self.head_dim
+
+        self.q_proj = nn.Linear(HIDDEN_SIZE, N_Q_HEADS  * hd, bias=False)
+        self.k_proj = nn.Linear(HIDDEN_SIZE, N_KV_HEADS * hd, bias=False)
+        self.v_proj = nn.Linear(HIDDEN_SIZE, N_KV_HEADS * hd, bias=False)
+        self.o_proj = nn.Linear(N_Q_HEADS   * hd, HIDDEN_SIZE, bias=False)
+
+        self.q_norm = RMSNorm(hd)
+        self.k_norm = RMSNorm(hd)
+
+    def forward(
+        self,
+        x: torch.Tensor,    # [B, T, D]
+        cos: torch.Tensor,  # [T, head_dim]
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        B, T, _ = x.shape
+        hd = self.head_dim
+
+        q = self.q_proj(x).view(B, T, N_Q_HEADS,  hd).transpose(1, 2)  # [B,Hq,T,hd]
+        k = self.k_proj(x).view(B, T, N_KV_HEADS, hd).transpose(1, 2)  # [B,1,T,hd]
+        v = self.v_proj(x).view(B, T, N_KV_HEADS, hd).transpose(1, 2)
+
+        # Per-head QK normalisation (before RoPE)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # Rotary position embeddings
+        q = apply_rope(q, cos, sin)
+        k = apply_rope(k, cos, sin)
+
+        # Expand KV to match Q heads (MQA)
+        k = k.expand(B, N_Q_HEADS, T, hd)
+        v = v.expand(B, N_Q_HEADS, T, hd)
+
+        if self.full_attn:
+            # Standard causal attention, no window restriction
+            out = F.scaled_dot_product_attention(
+                q, k, v,
+                is_causal=True,
+                scale=ATTN_SCALE,
+            )
+        else:
+            # Sliding window causal attention.
+            # attn_mask[i, j] = True means query-position i CAN attend to key-position j.
+            # Causal: j <= i  (can only attend to past/current positions)
+            # Window: i - j < SLIDING_WINDOW
+            idx = torch.arange(T, device=x.device)
+            # idx.unsqueeze(0)  = [1, T] broadcast as j (key) axis
+            # idx.unsqueeze(1)  = [T, 1] broadcast as i (query) axis
+            # mask[i, j] = True iff j <= i AND i - j < SLIDING_WINDOW
+            attn_mask = (
+                (idx.unsqueeze(0) <= idx.unsqueeze(1)) &          # j <= i (causal)
+                (idx.unsqueeze(1) - idx.unsqueeze(0) < SLIDING_WINDOW)  # i - j < W
+            )  # [T_q, T_k]
+            out = F.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask=attn_mask,
+                scale=ATTN_SCALE,
+            )
+
+        out = out.transpose(1, 2).contiguous().view(B, T, N_Q_HEADS * hd)
+        return self.o_proj(out)
+
+
+# ── MLP (GeGLU) ───────────────────────────────────────────────────────────────
+
+class MLP(nn.Module):
+    """
+    GeGLU feed-forward network.
+
+    Layers 0-14:  intermediate_size=6144
+    Layers 15-34: intermediate_size=12288 (double-wide)
+    """
+
+    def __init__(self, layer_idx: int):
+        super().__init__()
+        inter = INTERMEDIATE_WIDE if layer_idx >= DOUBLE_WIDE_START else INTERMEDIATE
+        self.gate_proj = nn.Linear(HIDDEN_SIZE, inter, bias=False)
+        self.up_proj   = nn.Linear(HIDDEN_SIZE, inter, bias=False)
+        self.down_proj = nn.Linear(inter, HIDDEN_SIZE, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate = F.gelu(self.gate_proj(x), approximate="tanh")
+        return self.down_proj(gate * self.up_proj(x))
+
+
+# ── Decoder layer ─────────────────────────────────────────────────────────────
+
+class Gemma4TextLayer(nn.Module):
+    """
+    Single Gemma 4 decoder layer.
+
+    Execution order (per forward call):
+      1. Per-layer auxiliary stream injection
+      2. Self-attention block (pre/post norm, residual scaled by layer_scalar)
+      3. MLP block (pre/post norm, residual scaled by layer_scalar)
+
+    Per-layer auxiliary stream injection:
+      Receives per_layer_input [B,T,256] = combined embed+projection for this layer.
+        x_normed = input_layernorm(x)
+        gate     = sigmoid(per_layer_input_gate(x_normed))   # [B,T,256]
+        gated    = gate * per_layer_input                     # [B,T,256]
+        out_1536 = per_layer_projection(gated)               # [B,T,1536]
+        x        = x + post_per_layer_input_norm(out_1536)
+    """
+
+    def __init__(self, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+
+        # Attention
+        self.self_attn = Attention(layer_idx)
+
+        # MLP (double-wide for layers 15+)
+        self.mlp = MLP(layer_idx)
+
+        # Layer norms
+        self.input_layernorm            = RMSNorm(HIDDEN_SIZE)
+        self.post_attention_layernorm   = RMSNorm(HIDDEN_SIZE)
+        self.pre_feedforward_layernorm  = RMSNorm(HIDDEN_SIZE)
+        self.post_feedforward_layernorm = RMSNorm(HIDDEN_SIZE)
+        self.post_per_layer_input_norm  = RMSNorm(HIDDEN_SIZE)
+
+        # Per-layer auxiliary stream weights:
+        #   per_layer_input_gate:  Linear(1536→256),  weight=[256, 1536]
+        #   per_layer_projection:  Linear(256→1536),  weight=[1536, 256]
+        self.per_layer_input_gate  = nn.Linear(HIDDEN_SIZE, PER_LAYER_DIM, bias=False)
+        self.per_layer_projection  = nn.Linear(PER_LAYER_DIM, HIDDEN_SIZE, bias=False)
+
+        # Scalar multiplier for attention and MLP residual contributions
+        self.layer_scalar = nn.Parameter(torch.ones(1))
+
+    def forward(
+        self,
+        x: torch.Tensor,               # [B, T, D]
+        cos: torch.Tensor,             # RoPE tables for this layer type
+        sin: torch.Tensor,
+        per_layer_input: torch.Tensor, # [B, T, 256] combined embed+projection for this layer
+    ) -> torch.Tensor:
+
+        scalar = self.layer_scalar.to(x.dtype)
+
+        # ── 1. Per-layer auxiliary stream injection ──────────────────────────
+        # Gate uses the model's hidden activation (gelu_pytorch_tanh), matching
+        # the Gemma3n reference implementation.
+        # The layer_scalar multiplies all residual contributions (per-layer, attn, MLP).
+        x_normed = self.input_layernorm(x)
+        gate  = F.gelu(self.per_layer_input_gate(x_normed), approximate="tanh")  # [B,T,256]
+        gated = gate * per_layer_input                                            # [B,T,256]
+        out   = self.per_layer_projection(gated)                                  # [B,T,1536]
+        x     = x + scalar * self.post_per_layer_input_norm(out)
+
+        # ── 2. Self-attention ────────────────────────────────────────────────
+        # Apply input_layernorm again after the per-layer injection
+        h = self.self_attn(self.input_layernorm(x), cos, sin)
+        x = x + scalar * self.post_attention_layernorm(h)
+
+        # ── 3. MLP ───────────────────────────────────────────────────────────
+        h = self.mlp(self.pre_feedforward_layernorm(x))
+        x = x + scalar * self.post_feedforward_layernorm(h)
+
+        return x
+
+
+# ── Full model ─────────────────────────────────────────────────────────────────
+
+class Gemma4ForCausalLM(nn.Module):
+    """
+    Gemma 4 E2B text model (causal LM head, no vision/audio).
+
+    Tied embeddings: embed_tokens.weight is shared with lm_head.
+    Output logits are softcapped: 30 * tanh(logits / 30).
+
+    Per-layer auxiliary stream is computed model-level before layer iteration:
+      - embed_tokens_per_layer lookup:    [B,T,35*256]
+      - per_layer_model_projection:       Linear(1536→35*256)
+      - per_layer_projection_norm:        RMSNorm(256) per layer-slice
+      - combine:  per_layer_inputs = (embed_aux + proj_scaled) * (1/sqrt(2))
+    """
+
+    def __init__(self):
+        super().__init__()
+
+        # Token embeddings
+        self.embed_tokens           = nn.Embedding(VOCAB_SIZE, HIDDEN_SIZE)
+        self.embed_tokens_per_layer = nn.Embedding(VOCAB_SIZE, N_LAYERS * PER_LAYER_DIM)
+
+        # Final norm
+        self.norm = RMSNorm(HIDDEN_SIZE)
+
+        # Transformer layers
+        self.layers = nn.ModuleList([Gemma4TextLayer(i) for i in range(N_LAYERS)])
+
+        # Model-level per-layer projection (hidden → all layer aux dims at once)
+        # weight shape: [35*256, 1536] = [8960, 1536]
+        self.per_layer_model_projection = nn.Linear(
+            HIDDEN_SIZE, N_LAYERS * PER_LAYER_DIM, bias=False
+        )
+        # Norm applied to per-layer projection slices [256]
+        self.per_layer_projection_norm = RMSNorm(PER_LAYER_DIM)
+
+        # RoPE tables (computed lazily)
+        self._rope_slide_cos: torch.Tensor | None = None
+        self._rope_slide_sin: torch.Tensor | None = None
+        self._rope_full_cos:  torch.Tensor | None = None
+        self._rope_full_sin:  torch.Tensor | None = None
+        self._rope_seq:       int = 0
+
+    @staticmethod
+    def is_full_attention(layer_idx: int) -> bool:
+        return is_full_attention(layer_idx)
+
+    def _ensure_rope(self, seq_len: int, device: torch.device) -> None:
+        """Precompute (or extend) RoPE tables on demand."""
+        if self._rope_slide_cos is not None and self._rope_seq >= seq_len:
+            return
+        max_seq = max(seq_len, 2048)
+
+        # Sliding layers: head_dim=256, full rotation
+        cs, sn = build_rope_freqs(HEAD_DIM_SLIDE, max_seq, ROPE_THETA_SLIDE, device)
+        self._rope_slide_cos = cs
+        self._rope_slide_sin = sn
+
+        # Full-attention layers: head_dim=512, partial_rotary_factor=0.25.
+        # 512 * 0.25 = 128 dims rotated = 64 rotation pairs (half=256, 64 of 256 pairs).
+        n_rot = int(HEAD_DIM_FULL * PARTIAL_ROT_FULL) // 2  # = 64
+        cf, sf = build_rope_freqs(
+            HEAD_DIM_FULL, max_seq, ROPE_THETA_FULL, device, n_rot_pairs=n_rot
+        )
+        self._rope_full_cos = cf
+        self._rope_full_sin = sf
+        self._rope_seq = max_seq
+
+    def _compute_per_layer_inputs(
+        self, input_ids: torch.Tensor, x_embed: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Precompute per-layer auxiliary inputs for all 35 layers.
+
+        Returns:
+            per_layer_inputs: [B, T, N_LAYERS, PER_LAYER_DIM]
+        """
+        B, T = input_ids.shape
+
+        # 1. Token-based per-layer embeddings (vocabulary lookup)
+        # Scaled by sqrt(PER_LAYER_DIM)=16, matching Gemma3n's ScaledWordEmbedding convention
+        embed_aux = self.embed_tokens_per_layer(input_ids).to(x_embed.dtype)
+        embed_aux = embed_aux * math.sqrt(PER_LAYER_DIM)           # scale by sqrt(256)=16
+        # embed_aux: [B, T, 35*256]  reshape → [B, T, 35, 256]
+        embed_aux = embed_aux.view(B, T, N_LAYERS, PER_LAYER_DIM)
+
+        # 2. Hidden-state projection: project x_embed to [B, T, 35*256]
+        proj_all  = self.per_layer_model_projection(x_embed)  # [B, T, 35*256]
+        proj_all  = proj_all * PER_LAYER_PROJ_SCALE            # scale by 1/sqrt(hidden)
+        proj_all  = proj_all.view(B, T, N_LAYERS, PER_LAYER_DIM)
+        # Apply RMSNorm(256) to each layer slice
+        proj_all  = self.per_layer_projection_norm(proj_all)   # broadcast over [B,T,N]
+
+        # 3. Combine: (embed_aux + proj_normed) * (1/sqrt(2))
+        per_layer_inputs = (embed_aux + proj_all) * PER_LAYER_INPUT_SCALE
+
+        return per_layer_inputs  # [B, T, 35, 256]
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            input_ids: [B, T] long tensor
+
+        Returns:
+            logits: [B, T, vocab_size] with softcapping applied
+        """
+        B, T = input_ids.shape
+        self._ensure_rope(T, input_ids.device)
+
+        # Token embeddings scaled by sqrt(hidden_size)
+        x = self.embed_tokens(input_ids) * math.sqrt(HIDDEN_SIZE)  # [B,T,D]
+
+        # Compute per-layer auxiliary inputs (uses unmodified x_embed)
+        per_layer_inputs = self._compute_per_layer_inputs(input_ids, x)
+
+        for i, layer in enumerate(self.layers):
+            per_layer_i = per_layer_inputs[:, :, i, :]  # [B, T, 256]
+
+            if is_full_attention(i):
+                cos, sin = self._rope_full_cos, self._rope_full_sin
+            else:
+                cos, sin = self._rope_slide_cos, self._rope_slide_sin
+
+            x = layer(x, cos, sin, per_layer_i)
+
+        x = self.norm(x)
+
+        # Tied lm_head: F.linear(x, embed_tokens.weight)
+        logits = F.linear(x, self.embed_tokens.weight.to(x.dtype))  # [B,T,V]
+
+        # Logit softcapping
+        logits = LOGIT_CAP * torch.tanh(logits / LOGIT_CAP)
+        return logits
+
+    @classmethod
+    def load_weights(
+        cls,
+        safetensors_path: str | Path,
+        device: str = "cpu",
+    ) -> "Gemma4ForCausalLM":
+        """
+        Load from the safetensors checkpoint.
+
+        Weight names in the file follow the pattern:
+            model.language_model.X  →  self.X
+        """
+        model  = cls()
+        path   = str(safetensors_path)
+        prefix = "model.language_model."
+        state  = {}
+
+        with safe_open(path, framework="pt", device=device) as f:
+            for key in f.keys():
+                if not key.startswith(prefix):
+                    continue
+                local_key = key[len(prefix):]  # strip "model.language_model."
+                state[local_key] = f.get_tensor(key)
+
+        missing, unexpected = model.load_state_dict(state, strict=False)
+        if missing:
+            print(f"[load_weights] {len(missing)} missing keys (first 5): {missing[:5]}")
+        if unexpected:
+            print(f"[load_weights] {len(unexpected)} unexpected keys (first 5): {unexpected[:5]}")
+
+        model = model.to(dtype=DTYPE)
+        return model
+
+
+# ── Convenience loader ─────────────────────────────────────────────────────────
+
+def load_gemma4(
+    device: str | None = None,
+) -> tuple[Gemma4ForCausalLM, AutoTokenizer]:
+    """
+    Load the Gemma 4 E2B model and tokenizer.
+
+    Returns:
+        (model, tokenizer)  — model is in eval mode on `device`.
+    """
+    if device is None:
+        device = DEVICE
+
+    print(f"Loading Gemma 4 E2B from {SAFETENSORS_BLOB} ...")
+    model = Gemma4ForCausalLM.load_weights(SAFETENSORS_BLOB, device=device)
+    model = model.to(device).eval()
+
+    print(f"Loading tokenizer from {MODEL_DIR} ...")
+    tokenizer = AutoTokenizer.from_pretrained(str(MODEL_DIR), local_files_only=True)
+
+    return model, tokenizer
+
+
+# ── PPL evaluation ─────────────────────────────────────────────────────────────
+
+def ppl_on_text(
+    model: Gemma4ForCausalLM,
+    tokenizer: AutoTokenizer,
+    text: str,
+    device: str | None = None,
+    max_length: int = 1024,
+) -> float:
+    """
+    Compute token-level perplexity on `text`.
+
+    Args:
+        model:      Gemma4ForCausalLM in eval mode
+        tokenizer:  matching AutoTokenizer
+        text:       input string
+        device:     device for inference (defaults to DEVICE)
+        max_length: truncate to this many tokens
+
+    Returns:
+        perplexity (float)
+    """
+    if device is None:
+        device = DEVICE
+
+    enc = tokenizer(text, return_tensors="pt", truncation=True, max_length=max_length)
+    input_ids = enc["input_ids"].to(device)
+
+    with torch.no_grad():
+        logits = model(input_ids)           # [1, T, V]
+
+    # Shift: predict token t+1 from position t
+    shift_logits = logits[0, :-1, :]        # [T-1, V]
+    shift_labels = input_ids[0, 1:]         # [T-1]
+
+    log_probs = F.log_softmax(shift_logits.float(), dim=-1)
+    nll = -log_probs.gather(1, shift_labels.unsqueeze(1)).squeeze(1).mean()
+    return nll.exp().item()
+
+
+# ── main ──────────────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    _WIKI_TEXT = (
+        "The transformer architecture was introduced in the paper "
+        "'Attention Is All You Need' by Vaswani et al. in 2017. "
+        "It relies entirely on self-attention mechanisms, dispensing with "
+        "recurrence and convolutions entirely. Transformers have since become "
+        "the dominant architecture for natural language processing, powering "
+        "models such as BERT, GPT, T5, and the Gemma family. "
+        "The key innovation is the multi-head attention mechanism, which allows "
+        "the model to jointly attend to information from different representation "
+        "subspaces at different positions. This is complemented by position-wise "
+        "feed-forward networks and residual connections with layer normalisation. "
+        "Large language models built on this architecture are trained on massive "
+        "corpora using next-token prediction (autoregressive language modelling) "
+        "or masked language modelling. They exhibit emergent capabilities such as "
+        "few-shot and zero-shot generalisation across a wide variety of tasks."
+    )
+
+    model, tokenizer = load_gemma4()
+
+    ppl = ppl_on_text(model, tokenizer, _WIKI_TEXT)
+    print(f"\nPerplexity on sample text: {ppl:.2f}  (target: ~17–18 for bfloat16)")