Hook ABI v5 goldens for units, anatomy, mid and the S lens; population artifacts, bf16 norm fix, --dtype

AGENTS.md

@@ -95,7 +95,8 @@ python scripts/migrate_lens_v7.py --home DIR [--dry-run]  # pre-5.4 lens homes
 python scripts/regenerate_bundled.py [TARGETS...]         # bundled corpora + neutral baseline
 ```
 
-Model-loading verbs also take `-d/--device` and `-q/--quantize 4bit|8bit`;
+Model-loading verbs also take `-d/--device`, `-q/--quantize 4bit|8bit` and
+`--dtype float32|bfloat16|float16`;
 most read verbs take `-j/--json`. The root parser has exactly these nine
 verbs; a bare verb or group prints its help and exits 0, and `drowse MODEL` is
 an argparse error. Precedence is CLI > YAML > runner default; every verb with
@@ -238,11 +239,17 @@ interface here:
   (`loom/recipe.py`), the tree filter (`loom/filter.py`), the sampler, the
   measurement envelope, the affine lowering, a term's layers and shares
   (`SteeringComposer.term_layers`, the validate response), DLS, the Fisher fit
-  and topology;
+  and topology, and (hook ABI v5, pinned by the `units-tiny-models-v1` goldens
+  of `tests/support/browser_units.py`) the in-block sites, unit mean-ablation,
+  the unit tap, unit activity, token anatomy, the R/S lens reads, the S
+  readout, the static unit reads and the neuron push;
 - names it mirrors: `RESERVED_NAMESPACES`, `store/manifold_schema.py`
   `METHOD_FIT_MODES`, the manifold sidecar and `.drowse` formats, lens format
   v7 (`sites.out`, `out.layer-L` shards);
-- Python it calls: `lens.store.save_lens_at`/`LensProvenance`, and the
+- Python it calls: `lens.store.save_lens_at`/`LensProvenance`,
+  `units.reads.UnitReader.population_artifacts`/`null_basis`/`head_kept_ids`
+  (its offline unit packs), the fixture generator in
+  `tests/support/browser_units.py`, and the
   lens estimator (today the removed `lens.fit.fit_jacobian_lens`; its hosted
   fit moving to `lens.fit.fit_lens(kind=, sites=("out",))` is a pending
   follow-up, `drowse/lens/AGENTS.md`).

README.md

@@ -344,6 +344,7 @@ Common options:
 |---|---:|---|
 | `-d`, `--device` | `auto` | `cuda`, `mps`, `cpu`, or automatic selection |
 | `-q`, `--quantize` | none | `4bit` or `8bit` bitsandbytes quantization on CUDA |
+| `--dtype` | device default | `float32`, `bfloat16` or `float16` weights (default bf16 on CUDA/MPS, fp32 on CPU); a fit records the dtype it ran at |
 | `-p`, `--probes` | `all` | Bundled probe categories, `all`, or `none` |
 | `-H`, `--host` | `127.0.0.1` | Bind address; non-loopback requires an API key |
 | `-P`, `--port` | `8000` | Bind port |

drowse/cli/AGENTS.md

@@ -95,18 +95,21 @@ validate FILE` adds `ensure_vectors_installed(strict=True, dry_run=True)`
 ## Sessions
 
 `open_session(args, *, roster=False)` is the one construction path: narrated on
-stderr, built from `device`, `quantize`, `system_prompt`, `max_tokens`, `dls =
+stderr, built from `device`, `dtype`, `quantize`, `system_prompt`, `max_tokens`, `dls =
 not --no-dls`, `compile`, `cuda_graphs`, `return_top_k`, closed on exit.
 
 - `roster=True` only for verbs that read probe readings: `serve`, `experiment
   fan`, `experiment transcript run`. `-p`: unset or `all` → `None` (tagged
   concept axes plus already-fitted bundled multi-node manifolds), `none` or
   empty → `[]`, categories verbatim. Everything else constructs with
   `probes=[]`, which also skips the construction-time neutral capture and fits.
-- `-d/-q` (`MODEL_LOAD`) exist exactly on the verbs that load weights: `serve`,
+- `-d/-q/--dtype` (`MODEL_LOAD`) exist exactly on the verbs that load weights: `serve`,
   `manifold extract|generate|fit`, `template score`, `lens fit|top|decompose`,
-  `sae train`, `unit read|top|find|stats`, every `experiment` verb. `manifold transfer` has neither but loads
-  a model (default device) when it must capture a missing neutral cache.
+  `sae train`, `unit read|top|find|stats`, every `experiment` verb. `manifold transfer` has none but loads
+  a model (default device and dtype) when it must capture a missing neutral cache.
+  `--dtype` is long-only and reaches the loader and the model-source fingerprint,
+  so a fit's identity records the dtype it ran at and the no-op proofs
+  (`lens fit`, `manifold extract|fit`) key on it.
 
 ## Fail before the load
 

drowse/cli/context.py

@@ -192,11 +192,12 @@ def open_session(
     """
     from drowse.session.session import DrowseSession
 
-    tail = f" ({args.quantize})" if args.quantize else ""
-    note(f"Loading model: {args.model}{tail}")
+    tail = ", ".join(filter(None, (args.quantize, args.dtype)))
+    note(f"Loading model: {args.model}{f' ({tail})' if tail else ''}")
     session = DrowseSession.from_pretrained(
         args.model,
         device=args.device,
+        dtype=args.dtype,
         quantize=args.quantize,
         probes=resolve_probes(args.probes) if roster else [],
         system_prompt=getattr(args, "system_prompt", None),

drowse/cli/spec.py

@@ -112,8 +112,10 @@ def smoothing(value: str) -> "str | float":
              help="Device: auto (detect), cuda, mps, or cpu (default: auto)")
 QUANTIZE = arg("-q", "--quantize", choices=["4bit", "8bit"], default=None,
                help="Quantization mode (default: bf16/fp16)")
-#: Every verb that loads a model takes these two.
-MODEL_LOAD = (DEVICE, QUANTIZE)
+DTYPE = arg("--dtype", choices=["float32", "bfloat16", "float16"], default=None,
+            help="Weight dtype (default: bfloat16 on CUDA/MPS, float32 on CPU)")
+#: Every verb that loads a model takes these three.
+MODEL_LOAD = (DEVICE, QUANTIZE, DTYPE)
 
 PROBES = arg(
     "-p", "--probes", nargs="*", default=None,

drowse/cli/verbs/lens.py

@@ -104,6 +104,7 @@ def lens_fit_is_noop(
         seq_len=args.seq_len,
         source_fingerprint=lambda: model_source_fingerprint(
             args.model, quantize=args.quantize, device=args.device,
+            dtype=args.dtype,
         ),
         docs=docs,
         corpus_spec=None if hf_corpus is None else hf_corpus.spec,

drowse/cli/verbs/manifold.py

@@ -69,7 +69,7 @@ def _fit_noop(
     return preflight_manifold_fit_noop(
         folder, model_id=args.model, layer_indices=layers,
         sae=sae, fit_mode=fit_mode, hyperparams=hyperparams,
-        quantize=args.quantize, device=args.device,
+        quantize=args.quantize, device=args.device, dtype=args.dtype,
     )
 
 

drowse/fit/preflight.py

@@ -69,6 +69,7 @@ def preflight_manifold_fit_noop(
     hyperparams: Mapping[str, Any] | None = None,
     quantize: str | None = None,
     device: str = "auto",
+    dtype: str | None = None,
 ) -> ManifoldFitProof | None:
     """Prove an existing fit is exact for ``model_id`` without loading weights.
 
@@ -100,7 +101,9 @@ def preflight_manifold_fit_noop(
         return None
     if mf.fit_mode == "baked":
         return None
-    source_fingerprint = model_source_fingerprint(model_id, quantize=quantize, device=device)
+    source_fingerprint = model_source_fingerprint(
+        model_id, quantize=quantize, device=device, dtype=dtype,
+    )
     if source_fingerprint is None:
         return None
     loaded_fingerprint = _proven_neutral_loaded_fingerprint(model_id, source_fingerprint)

drowse/model/blocks.py

@@ -71,10 +71,9 @@
 _MOE_NAMES = frozenset({"experts", "router", "shared_expert", "shared_experts", "gate"})
 
 _GELU_C = math.sqrt(2.0 / math.pi)
-#: Probe-match tolerances.  Wide enough for bf16 modules, far below the O(1)
-#: gap between the shapes being told apart.
+#: Probe-match tolerance, relative to the expected value.  Wide enough for
+#: bf16 modules, far below the O(1) gap between the shapes being told apart.
 _MATCH_RTOL = 3e-2
-_OFFSET_GAP = 0.25
 
 
 # -- activations -------------------------------------------------------------
@@ -435,24 +434,61 @@ def _norm_spec(
     r = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + float(eps))
     implied = y / (x * r)  # the per-channel multiplier the module applied
     w = None if weight is None else weight.detach().float().cpu()
-    candidates = (0.0, 1.0) if w is not None else (0.0,)
-    for offset in candidates:
+    # Each convention's error is judged relative to its own gain: a bf16
+    # module rounds a gain of ~1,000 by a few units, so an absolute gap would
+    # refuse the right convention on large gains (gemma-3-270m's MLP output
+    # norms).
+    fits: dict[float, float] = {}
+    for offset in (0.0, 1.0) if w is not None else (0.0,):
         gain = torch.ones(dim) if w is None else w + offset
-        error = (implied - gain).abs()
-        if (
-            float(error.median()) < _OFFSET_GAP
-            and bool((error <= _MATCH_RTOL * gain.abs().clamp_min(1.0)).all())
-        ):
-            return NormSpec(
-                module=module, dim=dim, eps=float(eps), offset=offset,
-                weighted=w is not None,
-            )
-    raise UnsupportedBlockError(
-        f"{where} ({kind}) is not an RMSNorm with a plain or 1+w weight",
-        layer=layer,
+        error = (implied - gain).abs() / gain.abs().clamp_min(1.0)
+        if bool((error <= _MATCH_RTOL).all()):
+            fits[offset] = float(error.median())
+    if not fits:
+        raise UnsupportedBlockError(
+            f"{where} ({kind}) is not an RMSNorm with a plain or 1+w weight",
+            layer=layer,
+        )
+    offset = min(fits, key=fits.__getitem__)
+    if len(fits) > 1:
+        # Every gain is large enough that both conventions fit within the
+        # module's rounding: the module itself, run with a zero weight,
+        # answers its offset exactly (a plain norm 0, a 1+w norm the bare
+        # normalization).  The closer median is the fallback when it can't.
+        probed = _zero_weight_offset(module, x.to(device=device, dtype=dtype), x * r)
+        if probed is not None:
+            offset = probed
+    return NormSpec(
+        module=module, dim=dim, eps=float(eps), offset=offset,
+        weighted=w is not None,
     )
 
 
+def _zero_weight_offset(
+    module: nn.Module, x: torch.Tensor, normalized: torch.Tensor,
+) -> float | None:
+    """The offset an RMSNorm adds to its weight, read by calling it with a
+    zero weight (without touching its own), or ``None`` when that call
+    doesn't answer one of the two conventions."""
+    weight = getattr(module, "weight", None)
+    if not isinstance(weight, torch.Tensor):
+        return None
+    try:
+        with torch.no_grad():
+            y = torch.func.functional_call(
+                module, {"weight": torch.zeros_like(weight)}, (x,),
+            ).float().cpu()
+    except Exception:  # a hook-wrapped or exotic module: fall back
+        return None
+    if tuple(y.shape) != tuple(normalized.shape):
+        return None
+    implied = y / normalized
+    for offset in (0.0, 1.0):
+        if bool(((implied - offset).abs() <= _MATCH_RTOL).all()):
+            return offset
+    return None
+
+
 def _activation_kind(act: Any, layer: int) -> ActivationKind:
     if not callable(act):
         raise UnsupportedBlockError("MLP has no activation ('act_fn')", layer=layer)

drowse/units/AGENTS.md

@@ -295,6 +295,17 @@ dictionary="neurons")` (`NeuronWrites`: each layer's writes, signed, in the
 whitened metric) with each neuron's pole and label. `stats(unit)` is the
 firing block alone.
 
+`population_artifacts(layer)` is the browser seam's offline half: everything
+a reader without a vocabulary pass needs at one layer, host fp32 tensors keyed
+by name, a key absent where its lens half is. Per neuron population the R
+`rho` (the final norm's scale of each transported row), `lse` (`lse_neg` for
+writes) and contrast `mean`/`std`; the S `norm` (`‖u‖`) and `footprint_std`
+(`sqrt(uᵀ Σ_E u)`) and contrast (S fires pooled); the heads' `ov` contrast
+per family over `head_kept_ids()`; the head norms' mean derivatives; the
+whitened write norms `find` divides by. drowse.ai's pack builder calls it
+(and the public `null_basis()`) instead of the private builders, and
+`tests/fixtures/browser/unit-artifacts-*` pins its output.
+
 **Firing.** `Firing.scale` is what the unit's readings divide by:
 `{"source": "catalog" | "neutral", "detail", "top", "value"}` (`top` the
 source's `act_top` of the pole or head, `value` the soft scale over it).

drowse/units/reads.py

@@ -967,7 +967,7 @@ def _hubless(
 
     # --------------------------------------------------------------- neurons
 
-    def _null_basis(self) -> torch.Tensor:
+    def null_basis(self) -> torch.Tensor:
         """The bottom :data:`NULL_SHARE` right-singular directions of
         ``W_U diag(γ_final)``, ``[d, m]`` on the device (built once)."""
         null = self._null
@@ -1013,7 +1013,7 @@ def neuron(self, ref: NeuronRef, *, k: int = 8, hide_hubs: bool = False) -> Neur
                 )
             logits = self._masked_logit_lens(w)
             skew, kurt = logit_moments(logits[None])
-            null = self._null_basis()
+            null = self.null_basis()
             scalars = torch.stack([
                 w.norm(), cos_ab, (null.T @ w).pow(2).sum() / w.pow(2).sum().clamp_min(1e-30),
                 skew[0], kurt[0],
@@ -1895,6 +1895,87 @@ def find(
 
     # ------------------------------------------------------------- population
 
+    def population_artifacts(self, layer: int) -> dict[str, torch.Tensor]:
+        """Everything a reader without a pass over the vocabulary needs at
+        ``layer``, fp32 on the host and keyed by name; a key is absent where
+        its lens half is (the R and S matrices the reader's lenses hold at the
+        site).  What drowse.ai's pack builder ships per layer:
+
+        - population scalars, per neuron population (``writes`` at ``out``,
+          ``fires+``/``fires-`` at ``mid``): R ``.rho`` (the final norm's
+          scale of each transported row, ``rsqrt(mean((J x)²) + ε_f)``),
+          ``.lse`` (and ``.lse_neg`` for ``writes``); S ``.norm`` (``‖u‖``,
+          ``u = x K``) and ``.footprint_std`` (``sqrt(uᵀ Σ_E u)``), both
+          floored as ``top`` floors them;
+        - the contrast statistics ``hide_hubs`` and ``top`` z-score by:
+          ``<family>.<population>.mean``/``.std`` ``[V]`` (S ``fires`` is one
+          pooled population of both poles), and per family the heads'
+          ``<family>.heads.ov.mean``/``.std`` over the head-read vocabulary
+          (``head_kept_ids``) where both of a head read's factors exist;
+        - ``heads.norm_derivative.<role>`` ``[heads, hd, hd]``: the head norms'
+          mean derivatives a head read applies (query heads for ``query``, KV
+          heads otherwise), for a block with head norms;
+        - ``neurons.write_norms`` ``[d_ff]``: ``‖W w_n‖``, the whitened write
+          norms ``find`` divides its scores by (with a neutral baseline).
+        """
+        from drowse.lens.decompose import NeuronWrites
+
+        layer = int(layer)
+        if not 0 <= layer < self.n_layers:
+            raise RequestError(f"layer {layer} is outside the model's {self.n_layers} layers")
+        out: dict[str, torch.Tensor] = {}
+        with self._section("unit artifacts"), torch.inference_mode():
+            lenses = self._lenses()
+            block = self.runtime.block(layer)
+            emb = self._embeddings()
+            eps = float(self.runtime.spec().final_norm.eps)
+            writes, plus, minus = self._neuron_populations(block)
+            places: tuple[tuple[str, str, dict[str, torch.Tensor]], ...] = (
+                ("writes", "out", {"writes": writes}),
+                ("fires", "mid", {"fires+": plus, "fires-": minus}),
+            )
+            for place, site, populations in places:
+                matrix, label = self._r_matrix(lenses, layer, site)
+                stats = None if label is None else self._neuron_contrast(lenses, layer, "R", place)
+                if stats is not None:
+                    for name, rows in populations.items():
+                        x = rows if matrix is None else rows @ matrix.T
+                        out[f"R.{name}.rho"] = torch.rsqrt(x.pow(2).mean(-1) + eps)
+                        parts = ("mean", "std", "lse", "lse_neg") if place == "writes" else ("mean", "std", "lse")
+                        for part in parts:
+                            out[f"R.{name}.{part}"] = stats.get(f"{name}.{part}")
+                s_matrix = self._matrix(lenses.s, layer, site)
+                stats = None if s_matrix is None else self._neuron_contrast(lenses, layer, "S", place)
+                if stats is not None and s_matrix is not None:
+                    for name, rows in populations.items():
+                        u = self._s_covectors(rows, s_matrix)
+                        out[f"S.{name}.norm"] = u.norm(dim=-1).clamp_min(1e-12)
+                        out[f"S.{name}.footprint_std"] = emb.variance(u).clamp_min(1e-30).sqrt()
+                    for part in ("mean", "std"):
+                        out[f"S.{place}.{part}"] = stats.get(f"{place}.{part}")
+            norms = self._head_norms(block)
+            for role, value in norms.items():
+                if value is not None:
+                    out[f"heads.norm_derivative.{role}"] = value
+            heads = range(block.n_heads)
+            if self._b_factor(lenses, block, heads, norms) is not None:
+                for family in ("R", "S"):
+                    if self._a_factor(lenses, block, heads, family) is not None:
+                        mean, std = self._head_contrast(lenses, block, family, norms).mean_std("ov")
+                        out[f"{family}.heads.ov.mean"] = mean
+                        out[f"{family}.heads.ov.std"] = std
+            whitener = None if self._whitener is None else self._whitener()
+            if whitener is not None and layer in whitener:
+                out["neurons.write_norms"] = NeuronWrites(neuron_writes(block), whitener, layer).norms()
+            _drain(self.device)
+            return {name: value.detach().to("cpu", torch.float32).contiguous() for name, value in out.items()}
+
+    def head_kept_ids(self) -> torch.Tensor:
+        """The vocabulary ids a head read streams (its ``ov`` statistics'
+        index), on the host."""
+        with torch.inference_mode():
+            return self._head_kept().cpu()
+
     def layer_contrast(self, layer: int) -> dict[str, Any]:
         """Build (or load) one layer's neuron contrast statistics for both
         families and places: the whole-layer read ``hide_hubs`` and ``top``