.json
+// and data/metadata.json, renders header + token-highlighted examples.
+
+const params = new URLSearchParams(location.search);
+const featureId = parseInt(params.get('feature') ?? '0', 10);
+
+function pad(id, width) {
+ return String(id).padStart(width, '0');
+}
+
+function widthFor(nFeatures) {
+ return Math.max(4, String(nFeatures - 1).length);
+}
+
+async function main() {
+ const md = await fetch('data/metadata.json').then(r => r.json());
+ const w = widthFor(md.n_features);
+ const f = await fetch(`data/features/${pad(featureId, w)}.json`).then(r => r.json());
+
+ document.title = `Feature #${f.feature_id} — MoLT dashboard`;
+ document.getElementById('ckpt').textContent =
+ (md.hf_filename || md.ckpt_path) + ` · layer ${md.layer}`;
+ document.getElementById('title').textContent = `Feature #${f.feature_id}`;
+
+ const meta = document.getElementById('meta');
+ meta.innerHTML = '';
+ addMeta(meta, `tier ${f.tier} · rank ${f.rank}`);
+ addMeta(meta, `activation rate: ${(f.activation_rate * 100).toFixed(3)}%`);
+ addMeta(meta, `max activation: ${f.max_activation.toFixed(3)}`);
+ addMeta(meta, `examples: ${f.examples.length}`);
+
+ // Prev/next nav.
+ const nav = document.createElement('span');
+ if (featureId > 0) {
+ nav.appendChild(link(`?feature=${featureId - 1}`, '← prev'));
+ nav.appendChild(document.createTextNode(' '));
+ }
+ if (featureId < md.n_features - 1) {
+ nav.appendChild(link(`?feature=${featureId + 1}`, 'next →'));
+ }
+ meta.appendChild(nav);
+
+ const root = document.getElementById('examples');
+ if (f.examples.length === 0) {
+ root.innerHTML = 'This transform did not fire on any token in the sampled corpus.
';
+ return;
+ }
+ for (const ex of f.examples) {
+ root.appendChild(renderExample(ex, f.max_activation));
+ }
+}
+
+function addMeta(parent, text) {
+ const s = document.createElement('span');
+ s.textContent = text;
+ parent.appendChild(s);
+}
+
+function link(href, text) {
+ const a = document.createElement('a');
+ a.href = href;
+ a.textContent = text;
+ return a;
+}
+
+function renderExample(ex, globalMax) {
+ const wrap = document.createElement('div');
+ wrap.className = 'example';
+
+ const header = document.createElement('div');
+ header.className = 'example-header';
+ header.textContent = `peak ${ex.peak_activation.toFixed(3)} · position ${ex.peak_token_pos} of ${ex.tokens.length}`;
+ wrap.appendChild(header);
+
+ const tokens = document.createElement('div');
+ tokens.className = 'tokens';
+ // Scale alpha by the per-feature global max so different sequences are
+ // visually comparable within the same feature.
+ const scale = globalMax > 0 ? globalMax : 1.0;
+ for (let i = 0; i < ex.tokens.length; i++) {
+ const span = document.createElement('span');
+ span.className = 'tok';
+ if (i === ex.peak_token_pos) span.classList.add('peak');
+ const a = ex.activations[i];
+ const alpha = Math.max(0, Math.min(1, a / scale));
+ span.style.backgroundColor = `rgba(220, 50, 50, ${alpha.toFixed(3)})`;
+ span.title = `act ${a.toFixed(4)}`;
+ // Preserve raw whitespace inside the span (CSS white-space: pre handles this).
+ span.textContent = ex.tokens[i] === '' ? ' ' : ex.tokens[i];
+ tokens.appendChild(span);
+ }
+ wrap.appendChild(tokens);
+
+ return wrap;
+}
+
+main().catch(err => {
+ document.getElementById('examples').innerHTML =
+ `Failed to load feature data: ${err.message}. ` +
+ `If you opened this with file://, serve the directory first ` +
+ `(e.g. python -m http.server).
`;
+});
diff --git a/crosslayer_transcoder/feature_dash/templates/index.html b/crosslayer_transcoder/feature_dash/templates/index.html
new file mode 100644
index 0000000..c00496f
--- /dev/null
+++ b/crosslayer_transcoder/feature_dash/templates/index.html
@@ -0,0 +1,37 @@
+
+
+
+
+ MoLT Feature Dashboard — index
+
+
+
+ MoLT feature dashboard
+ Loading…
+
+
+
+
+
+
+
+
+
+ | id |
+ tier |
+ rank |
+ rate |
+ max act |
+
+
+
+
+
+
+
diff --git a/crosslayer_transcoder/feature_dash/templates/index.js b/crosslayer_transcoder/feature_dash/templates/index.js
new file mode 100644
index 0000000..c2ee9bb
--- /dev/null
+++ b/crosslayer_transcoder/feature_dash/templates/index.js
@@ -0,0 +1,120 @@
+// Index page. One fetch of metadata.json populates the sortable table.
+
+let allRows = [];
+let sortKey = 'activation_rate';
+let sortDir = 'desc';
+let metadata = null;
+
+function pad(id, width) {
+ return String(id).padStart(width, '0');
+}
+
+async function main() {
+ metadata = await fetch('data/metadata.json').then(r => r.json());
+
+ document.getElementById('meta').textContent =
+ `${metadata.n_features} transforms · ` +
+ `${(metadata.hf_filename || metadata.ckpt_path)} · layer ${metadata.layer} · ` +
+ `${metadata.n_tokens_collected.toLocaleString()} tokens from ${metadata.dashboard_dataset}`;
+
+ // Build the row data.
+ for (let i = 0; i < metadata.n_features; i++) {
+ allRows.push({
+ feature_id: i,
+ tier: metadata.feature_tier[i],
+ rank: metadata.feature_rank[i],
+ activation_rate: metadata.feature_activation_rate[i],
+ max_activation: metadata.feature_max_activation[i],
+ });
+ }
+
+ // Populate tier filter.
+ const uniqueTiers = [...new Set(metadata.feature_tier)].sort((a, b) => a - b);
+ const tierSel = document.getElementById('tier-filter');
+ for (const t of uniqueTiers) {
+ const opt = document.createElement('option');
+ opt.value = String(t);
+ opt.textContent = `tier ${t} (rank ${metadata.ranks[t]})`;
+ tierSel.appendChild(opt);
+ }
+
+ // Wire interactions.
+ document.querySelectorAll('th[data-key]').forEach(th => {
+ th.addEventListener('click', () => {
+ const k = th.dataset.key;
+ if (sortKey === k) {
+ sortDir = sortDir === 'desc' ? 'asc' : 'desc';
+ } else {
+ sortKey = k;
+ sortDir = (k === 'feature_id' || k === 'tier') ? 'asc' : 'desc';
+ }
+ render();
+ });
+ });
+ ['change', 'input'].forEach(evt => {
+ tierSel.addEventListener(evt, render);
+ document.getElementById('min-rate').addEventListener(evt, render);
+ document.getElementById('max-rate').addEventListener(evt, render);
+ });
+
+ render();
+}
+
+function render() {
+ const tier = document.getElementById('tier-filter').value;
+ const minRate = parseFloat(document.getElementById('min-rate').value) || 0;
+ const maxRate = parseFloat(document.getElementById('max-rate').value);
+ const maxRateOk = Number.isFinite(maxRate) ? maxRate : 1;
+
+ const filtered = allRows.filter(r =>
+ (tier === '' || String(r.tier) === tier) &&
+ r.activation_rate >= minRate &&
+ r.activation_rate <= maxRateOk
+ );
+
+ filtered.sort((a, b) => {
+ const av = a[sortKey], bv = b[sortKey];
+ if (av < bv) return sortDir === 'asc' ? -1 : 1;
+ if (av > bv) return sortDir === 'asc' ? 1 : -1;
+ return a.feature_id - b.feature_id;
+ });
+
+ document.querySelectorAll('th[data-key]').forEach(th => {
+ th.classList.remove('sorted-asc', 'sorted-desc');
+ if (th.dataset.key === sortKey) {
+ th.classList.add(sortDir === 'asc' ? 'sorted-asc' : 'sorted-desc');
+ }
+ });
+
+ document.getElementById('count').textContent =
+ `${filtered.length.toLocaleString()} of ${allRows.length.toLocaleString()} shown`;
+
+ const tbody = document.querySelector('#features tbody');
+ // Render up to 2000 rows; the user can filter further if they need more.
+ const cap = 2000;
+ const rows = filtered.slice(0, cap);
+ const html = rows.map(r => {
+ const isDead = r.activation_rate === 0;
+ return `
+ #${r.feature_id} |
+ ${r.tier} |
+ ${r.rank} |
+ ${(r.activation_rate * 100).toFixed(3)}% |
+ ${r.max_activation.toFixed(3)} |
+ `;
+ }).join('');
+ tbody.innerHTML = html;
+ if (filtered.length > cap) {
+ tbody.insertAdjacentHTML(
+ 'beforeend',
+ `| ${filtered.length - cap} more rows hidden — narrow the filter |
`
+ );
+ }
+}
+
+main().catch(err => {
+ document.getElementById('meta').innerHTML =
+ `Failed to load metadata.json: ${err.message}. ` +
+ `If you opened this with file://, serve the directory first ` +
+ `(e.g. python -m http.server).`;
+});
diff --git a/crosslayer_transcoder/model/__init__.py b/crosslayer_transcoder/model/__init__.py
index e4f0aae..6a88618 100644
--- a/crosslayer_transcoder/model/__init__.py
+++ b/crosslayer_transcoder/model/__init__.py
@@ -4,11 +4,13 @@
from .clt import CrossLayerTranscoder
from .clt_lightning import CrossLayerTranscoderModule
+from .molt import Molt
from .topk import BatchTopK, PerLayerBatchTopK, PerLayerTopK
__all__ = [
"CrossLayerTranscoder",
"CrossLayerTranscoderModule",
+ "Molt",
"BatchTopK",
"PerLayerTopK",
"PerLayerBatchTopK",
diff --git a/crosslayer_transcoder/model/clt_lightning.py b/crosslayer_transcoder/model/clt_lightning.py
index 8e4a020..2256028 100644
--- a/crosslayer_transcoder/model/clt_lightning.py
+++ b/crosslayer_transcoder/model/clt_lightning.py
@@ -2,7 +2,7 @@
import os
import subprocess
import time
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union
import lightning as L
import psutil
@@ -23,6 +23,7 @@
Decoder,
)
from crosslayer_transcoder.model.jumprelu import JumpReLU
+from crosslayer_transcoder.model.molt import Molt
from crosslayer_transcoder.model.topk import BatchTopK
@@ -30,7 +31,7 @@ class CrossLayerTranscoderModule(L.LightningModule):
def __init__(
self,
# Pre-constructed modules
- model: CrossLayerTranscoder,
+ model: Union[CrossLayerTranscoder, Molt],
replacement_model: Optional[ReplacementModelAccuracy] = None,
dead_features: Optional[DeadFeatures] = None,
# Training parameters
@@ -85,17 +86,23 @@ def __init__(
self.beta2 = beta2
self.log_metrics_every = log_metrics_every
- assert self.model.encoder.n_layers == self.model.decoder.n_layers, (
- "Encoder and decoder must have the same number of layers"
- )
+ if isinstance(self.model, Molt):
+ self.register_buffer(
+ "last_active",
+ torch.zeros((self.model.n_features,), dtype=torch.long),
+ )
+ else:
+ assert self.model.encoder.n_layers == self.model.decoder.n_layers, (
+ "Encoder and decoder must have the same number of layers"
+ )
- self.register_buffer(
- "last_active",
- torch.zeros(
- (self.model.encoder.n_layers, self.model.encoder.d_features),
- dtype=torch.long,
- ),
- )
+ self.register_buffer(
+ "last_active",
+ torch.zeros(
+ (self.model.encoder.n_layers, self.model.encoder.d_features),
+ dtype=torch.long,
+ ),
+ )
def configure_model(self):
# Apply compilation if requested
@@ -565,3 +572,68 @@ def training_step(self, batch, batch_idx):
torch.cuda.memory._record_memory_history(enabled=None)
exit()
return loss
+
+
+class MoltModule(CrossLayerTranscoderModule):
+ def __init__(
+ self,
+ lambda_sparsity: float = 0.0002,
+ c_sparsity: float = 0.1,
+ use_tanh: bool = True,
+ *args,
+ **kwargs,
+ ):
+ super().__init__(*args, **kwargs)
+ self._lambda = lambda_sparsity
+ self.c = c_sparsity
+ self.use_tanh = use_tanh
+
+ def current_sparsity_penalty(self):
+ n_steps = self.trainer.max_steps
+ current_step = (
+ self.global_step
+ ) # use global step instead of batch idx to work with gradient accumulation
+ cur_lambda = self._lambda * (current_step / n_steps)
+ self.log("training/sparsity_penalty", cur_lambda)
+ return cur_lambda
+
+ def forward(self, batch, layer):
+ return self.model.forward(batch, layer)
+
+ def training_step(self, batch, batch_idx):
+ if batch_idx == 0:
+ self.model.initialize_standardizers(batch)
+ self.log("model/d_latents", self.model.d_latents)
+ self.log("model/n_features", self.model.n_features)
+
+ layer = 8
+
+ # Forward pass
+ resid, mlp_out = batch[:, 0], batch[:, 1]
+ resid = resid[:, layer]
+ mlp_out = mlp_out[:, layer]
+ gate, recons_norm, recons = self.model.forward(resid, layer)
+
+ self.update_dead_features(gate)
+ # Compute MSE loss
+ mse = (recons_norm - self.model.output_standardizer.standardize(mlp_out, layer)) ** 2
+
+ # Compute Sparsity Loss
+ norms = self.model.transform_norm()
+ weighted_norms = norms * gate
+ self.log("model/weighted_norms_mean", weighted_norms.detach().mean().cpu())
+
+ if self.use_tanh:
+ weighted_norms = torch.tanh(weighted_norms * self.c)
+ sparsity = self.current_sparsity_penalty() * weighted_norms.sum(dim=-1).mean()
+ self.log("training/sparsity_loss", sparsity)
+ self.log("L0", (gate > 0.0).float().sum() / gate.shape[0])
+
+ loss = mse.mean() + sparsity
+ self.log("training/mse", mse.mean())
+ self.log("training/loss", loss)
+
+ if batch_idx % self.log_metrics_every == 0:
+ pass
+
+ return loss
diff --git a/crosslayer_transcoder/model/jumprelu.py b/crosslayer_transcoder/model/jumprelu.py
index d21b7ec..c977830 100644
--- a/crosslayer_transcoder/model/jumprelu.py
+++ b/crosslayer_transcoder/model/jumprelu.py
@@ -52,7 +52,8 @@ def backward(ctx, grad_output):
class JumpReLU(SerializableModule):
def __init__(self, theta=0.0, bandwidth=1.0, n_layers=12, d_features=768 * 8):
super().__init__()
- self.theta = nn.Parameter(torch.full((1, n_layers, d_features), theta))
+ shape = (1, n_layers, d_features) if n_layers > 1 else (1, d_features)
+ self.theta = nn.Parameter(torch.full(shape, theta))
self.register_buffer("bandwidth", torch.tensor(bandwidth))
self._init_theta = theta
self.n_layers = n_layers
diff --git a/crosslayer_transcoder/model/molt.py b/crosslayer_transcoder/model/molt.py
new file mode 100644
index 0000000..afb3fd3
--- /dev/null
+++ b/crosslayer_transcoder/model/molt.py
@@ -0,0 +1,93 @@
+import einops
+import torch
+import torch.nn as nn
+from jaxtyping import Float
+
+
+class Molt(nn.Module):
+ def __init__(
+ self,
+ d_acts: int,
+ N: int,
+ nonlinearity: nn.Module,
+ input_standardizer: nn.Module,
+ output_standardizer: nn.Module,
+ ranks: list[int] = [512, 256, 128, 64, 32],
+ ):
+ super().__init__()
+
+ self.d_acts = d_acts
+ self.nonlinearity = nonlinearity
+ self.input_standardizer = input_standardizer
+ self.output_standardizer = output_standardizer
+ Us = []
+ Vs = []
+ rank_multiplier = 1
+ n_features = 0
+ d_latents = 0
+ for rank in ranks:
+ Us.append(nn.Parameter(torch.empty(N * rank_multiplier, rank, d_acts)))
+ Vs.append(nn.Parameter(torch.empty(N * rank_multiplier, d_acts, rank)))
+ n_features += N * rank_multiplier
+ d_latents += N * rank_multiplier * rank
+ rank_multiplier *= 2
+ self.n_features = n_features
+ self.e = nn.Linear(d_acts, n_features)
+ self.Us = nn.ParameterList(Us)
+ self.Vs = nn.ParameterList(Vs)
+
+ print(f"d_latents (transcoder equivalent): {d_latents}")
+ self.d_latents = d_latents
+
+ self.reset_parameters()
+
+ def reset_parameters(self):
+ for U in self.Us:
+ nn.init.xavier_uniform_(U)
+ for V in self.Vs:
+ nn.init.xavier_uniform_(V)
+
+ def transform_norm(self):
+ norms = []
+ for U, V in zip(self.Us, self.Vs):
+ uv = einops.einsum(
+ U,
+ V,
+ "n_transforms d_transform d_acts_out, n_transforms d_acts_in d_transform -> n_transforms d_acts_in d_acts_out",
+ )
+ norms.append(torch.norm(uv, dim=(1, 2)))
+ return torch.cat(norms, dim=0)
+
+ def forward(
+ self, acts: Float[torch.Tensor, "batch_size d_acts"], layer: int
+ ) -> Float[torch.Tensor, "batch_size d_acts"]:
+ acts = self.input_standardizer(acts, layer)
+ pre_actvs = self.e(acts)
+ gate = self.nonlinearity(pre_actvs) # (batch, n_transforms)
+
+ raw_recons = []
+ for U, V in zip(self.Us, self.Vs):
+ latents = einops.einsum(
+ acts,
+ V,
+ "batch d_acts, n_transforms d_acts d_transform -> batch n_transforms d_transform",
+ )
+ raw_recons.append(
+ einops.einsum(
+ latents,
+ U,
+ "batch n_transforms d_transform, n_transforms d_transform d_acts -> batch n_transforms d_acts",
+ )
+ )
+
+ raw_recons = torch.cat(raw_recons, dim=1)
+
+ weighted_recons = gate.unsqueeze(-1) * raw_recons
+ recons_norm = weighted_recons.sum(dim=1)
+
+ recons = self.output_standardizer(recons_norm, layer)
+ return gate, recons_norm, recons
+
+ def initialize_standardizers(self, batch: Float[torch.Tensor, "batch_size io n_layers d_acts"]):
+ self.input_standardizer.initialize_from_batch(batch)
+ self.output_standardizer.initialize_from_batch(batch)
diff --git a/pyproject.toml b/pyproject.toml
index 849e873..52e8de7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -2,6 +2,9 @@
requires = ["hatchling"]
build-backend = "hatchling.build"
+[tool.hatch.build.targets.wheel.force-include]
+"crosslayer_transcoder/feature_dash/templates" = "crosslayer_transcoder/feature_dash/templates"
+
[project]
name = "crosslayer-transcoder"
version = "0.1.0"
diff --git a/tests/test_feature_dash_bundle.py b/tests/test_feature_dash_bundle.py
new file mode 100644
index 0000000..d680bee
--- /dev/null
+++ b/tests/test_feature_dash_bundle.py
@@ -0,0 +1,250 @@
+"""Tests for the single-file `bundle.html` builder.
+
+Two paths to the bundle:
+ - `make_bundle_from_disk(out_dir)`: read an existing dump, write bundle.html.
+ - `make_bundle(collector, meta, tokenizer, out_path, ...)`: build directly
+ from in-memory state, no disk dump needed.
+
+Both should produce a self-contained file (no `fetch(` calls, all data inlined)
+that contains every feature's payload.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+
+import pytest
+import torch
+
+from crosslayer_transcoder.feature_dash.bundle import (
+ default_bundle_filename,
+ make_bundle,
+ make_bundle_from_disk,
+)
+from crosslayer_transcoder.feature_dash.collect import GateCollector
+from crosslayer_transcoder.feature_dash.dump import dump_dashboard
+from crosslayer_transcoder.feature_dash.load import MoltCheckpointMetadata
+
+
+class _StubTokenizer:
+ def decode(self, ids):
+ return f" tok{ids[0]}"
+
+
+def _populate_collector(F: int, K: int, T: int) -> GateCollector:
+ coll = GateCollector(n_features=F, top_k=K, seq_len=T)
+ tok = torch.arange(2 * T, dtype=torch.long).reshape(2, T)
+ g = torch.zeros(2, T, F)
+ g[0, 1, 0] = 1.0
+ g[1, 3, 0] = 2.5
+ g[0, 0, 1] = 0.5
+ coll.update(tok, g)
+ return coll
+
+
+def _meta(F: int, ranks: list[int], N: int) -> MoltCheckpointMetadata:
+ feature_tier: list[int] = []
+ feature_rank: list[int] = []
+ for t, r in enumerate(ranks):
+ n_in_tier = N * (2**t)
+ feature_tier.extend([t] * n_in_tier)
+ feature_rank.extend([r] * n_in_tier)
+ return MoltCheckpointMetadata(
+ ckpt_path="dummy.ckpt",
+ d_acts=8,
+ n_features=F,
+ n_layers=2,
+ ranks=ranks,
+ N=N,
+ feature_tier=feature_tier,
+ feature_rank=feature_rank,
+ base_model_name="stub",
+ training_dataset="stub-ds",
+ global_step=123,
+ epoch=0,
+ )
+
+
+def _extract_inlined_json(html: str, script_id: str) -> dict:
+ """Pull a `` payload."""
+ m = re.search(
+ rf'',
+ html,
+ flags=re.DOTALL,
+ )
+ assert m, f"missing inlined script #{script_id}"
+ raw = m.group(1)
+ # Reverse the `<` hardening before parsing.
+ raw = raw.replace(r"<", "<")
+ return json.loads(raw)
+
+
+def test_bundle_contains_no_fetch_calls(tmp_path: Path):
+ """Bundle must work from file:// — no network/disk fetches allowed."""
+ F, K, T, N = 3, 2, 5, 1
+ coll = _populate_collector(F, K, T)
+ meta = _meta(F, [4, 2], N)
+
+ out = make_bundle(
+ collector=coll,
+ meta=meta,
+ tokenizer=_StubTokenizer(),
+ out_path=tmp_path / "bundle.html",
+ layer=8,
+ window=2,
+ )
+ text = out.read_text()
+ # No fetch(... call anywhere — neither real nor templated.
+ assert "fetch(" not in text
+
+
+def test_bundle_inlines_metadata_and_all_features(tmp_path: Path):
+ F, K, T, N = 3, 2, 5, 1
+ coll = _populate_collector(F, K, T)
+ meta = _meta(F, [4, 2], N)
+
+ bundle_path = make_bundle(
+ collector=coll, meta=meta, tokenizer=_StubTokenizer(),
+ out_path=tmp_path, layer=8, window=2,
+ )
+ # When out_path is a directory, the bundle is named after the checkpoint.
+ # `_meta()` sets ckpt_path="dummy.ckpt" and no hf_filename, so stem="dummy".
+ assert bundle_path == tmp_path / "bundle_dummy.html"
+
+ html = bundle_path.read_text()
+ md = _extract_inlined_json(html, "metadata")
+ feats = _extract_inlined_json(html, "features")
+
+ assert md["n_features"] == F
+ assert md["layer"] == 8
+ assert md["feature_tier"] == [0, 1, 1]
+ assert md["feature_rank"] == [4, 2, 2]
+ assert len(md["feature_activation_rate"]) == F
+ assert len(md["feature_max_activation"]) == F
+
+ assert set(feats.keys()) == {"0", "1", "2"}
+ for k, v in feats.items():
+ assert v["feature_id"] == int(k)
+ assert {"tier", "rank", "activation_rate", "max_activation", "examples"} <= set(v.keys())
+
+ # Feature 0 has two real examples and feature 2 is dead.
+ assert len(feats["0"]["examples"]) == 2
+ assert feats["2"]["examples"] == []
+
+
+def test_bundle_css_is_inlined(tmp_path: Path):
+ F, K, T, N = 3, 2, 5, 1
+ coll = _populate_collector(F, K, T)
+ meta = _meta(F, [4, 2], N)
+ out = make_bundle(
+ collector=coll, meta=meta, tokenizer=_StubTokenizer(),
+ out_path=tmp_path / "bundle.html", layer=8, window=2,
+ )
+ html = out.read_text()
+ # No external stylesheet link, and a non-empty inline " in html
+ assert "table.features" in html # one of our CSS rules
+
+
+def test_bundle_from_disk_reads_dumped_dashboard(tmp_path: Path):
+ F, K, T, N = 3, 2, 5, 1
+ coll = _populate_collector(F, K, T)
+ meta = _meta(F, [4, 2], N)
+
+ dump_dashboard(
+ collector=coll, meta=meta, tokenizer=_StubTokenizer(),
+ out_dir=tmp_path, layer=8, window=2, copy_assets=False,
+ )
+ bundle_path = make_bundle_from_disk(tmp_path)
+
+ # ckpt_path="dummy.ckpt" -> bundle_dummy.html
+ assert bundle_path == tmp_path / "bundle_dummy.html"
+ html = bundle_path.read_text()
+ md = _extract_inlined_json(html, "metadata")
+ feats = _extract_inlined_json(html, "features")
+ assert md["n_features"] == F
+ assert len(feats) == F
+
+
+def test_bundle_from_disk_errors_when_data_missing(tmp_path: Path):
+ with pytest.raises(FileNotFoundError):
+ make_bundle_from_disk(tmp_path)
+
+
+def test_default_bundle_filename_prefers_hf_filename():
+ F, N = 3, 1
+ meta = _meta(F, [4, 2], N)
+ meta.hf_filename = "gpt2-molt-lam-0_00015-50M.ckpt"
+ meta.ckpt_path = "/some/random/local/path.ckpt"
+ assert default_bundle_filename(meta) == "bundle_gpt2-molt-lam-0_00015-50M.html"
+
+
+def test_default_bundle_filename_falls_back_to_ckpt_path():
+ F, N = 3, 1
+ meta = _meta(F, [4, 2], N)
+ meta.hf_filename = None
+ meta.ckpt_path = "/runs/exp42/checkpoint-final.ckpt"
+ assert default_bundle_filename(meta) == "bundle_checkpoint-final.html"
+
+
+def test_default_bundle_filename_handles_unsafe_chars():
+ F, N = 3, 1
+ meta = _meta(F, [4, 2], N)
+ meta.hf_filename = "weird/name with spaces.ckpt"
+ # The stem is "name with spaces" (Path.stem already drops the directory).
+ # Spaces are sanitised to '-'.
+ assert default_bundle_filename(meta) == "bundle_name-with-spaces.html"
+
+
+def test_make_bundle_explicit_filepath_is_respected(tmp_path: Path):
+ F, K, T, N = 3, 2, 5, 1
+ coll = _populate_collector(F, K, T)
+ meta = _meta(F, [4, 2], N)
+
+ explicit = tmp_path / "my-custom-name.html"
+ out = make_bundle(
+ collector=coll, meta=meta, tokenizer=_StubTokenizer(),
+ out_path=explicit, layer=8, window=2,
+ )
+ # When the user passes a file path, we don't override it.
+ assert out == explicit
+
+
+def test_bundle_handles_script_tag_in_token_text(tmp_path: Path):
+ """A token string containing `` must not break the bundle.
+
+ The sanitiser replaces `<` with `\\u003c`; the JS in the bundle parses the
+ JSON via `JSON.parse`, which decodes the unicode escape back. We just
+ verify here that the raw HTML doesn't contain a literal `` inside
+ the data payload.
+ """
+
+ class _NastyTokenizer:
+ def decode(self, ids):
+ return "BAD"
+
+ F, K, T, N = 3, 2, 5, 1
+ coll = _populate_collector(F, K, T)
+ meta = _meta(F, [4, 2], N)
+
+ out = make_bundle(
+ collector=coll, meta=meta, tokenizer=_NastyTokenizer(),
+ out_path=tmp_path / "bundle.html", layer=8, window=2,
+ )
+ html = out.read_text()
+ # Find the data scripts and check their content is sanitised.
+ m = re.search(
+ r'',
+ html, flags=re.DOTALL,
+ )
+ assert m
+ payload = m.group(1)
+ # The sanitised payload must not contain a literal "') == 3
diff --git a/tests/test_feature_dash_collect.py b/tests/test_feature_dash_collect.py
new file mode 100644
index 0000000..dcbdb85
--- /dev/null
+++ b/tests/test_feature_dash_collect.py
@@ -0,0 +1,158 @@
+"""Tests for the GateCollector accumulator (no LM, no dataset).
+
+Synthetic gate batches go in, we check that firing counts, activation rates,
+top-K ordering, and per-token activation traces come out right.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from crosslayer_transcoder.feature_dash.collect import GateCollector
+
+
+F = 3 # features
+K = 4 # top-K
+T = 6 # seq_len
+
+
+def _zero_gates(B: int) -> torch.Tensor:
+ return torch.zeros(B, T, F)
+
+
+def test_firing_count_and_total_tokens():
+ coll = GateCollector(n_features=F, top_k=K, seq_len=T)
+ tok = torch.arange(2 * T, dtype=torch.long).reshape(2, T)
+ g = _zero_gates(2)
+ g[0, 1, 0] = 0.5
+ g[0, 4, 0] = 1.5
+ g[1, 2, 0] = 0.0 # 0 doesn't count as active (gate > 0)
+ g[1, 3, 1] = 0.7
+
+ coll.update(tok, g)
+
+ assert coll.total_tokens == 2 * T
+ assert coll.firing_count.tolist() == [2, 1, 0]
+ rates = coll.activation_rate()
+ assert torch.allclose(rates, torch.tensor([2 / 12, 1 / 12, 0.0]))
+
+
+def test_max_activation_tracks_global_max_across_batches():
+ coll = GateCollector(n_features=F, top_k=K, seq_len=T)
+ g1 = _zero_gates(1)
+ g1[0, 0, 0] = 1.0
+ coll.update(torch.zeros(1, T, dtype=torch.long), g1)
+
+ g2 = _zero_gates(1)
+ g2[0, 0, 0] = 0.5
+ g2[0, 0, 1] = 4.2
+ coll.update(torch.zeros(1, T, dtype=torch.long), g2)
+
+ # Once a batch lands, max_activation rises from -inf to >=0 even for
+ # non-firing features (max of -inf and 0 is 0). That's fine for the
+ # dashboard — feature_summary surfaces 0.0 for dead features.
+ assert torch.allclose(
+ coll.max_activation, torch.tensor([1.0, 4.2, 0.0]), atol=1e-6
+ )
+
+
+def test_topk_orders_by_per_sequence_peak():
+ coll = GateCollector(n_features=F, top_k=K, seq_len=T)
+ # 5 sequences, all activate feature 0 with different peaks at different positions.
+ B = 5
+ tok = torch.arange(B * T, dtype=torch.long).reshape(B, T)
+ g = _zero_gates(B)
+ peaks_in = [0.1, 3.0, 2.0, 0.5, 5.0]
+ pos_in = [0, 2, 5, 1, 4]
+ for b, (p, pos) in enumerate(zip(peaks_in, pos_in)):
+ g[b, pos, 0] = p
+ coll.update(tok, g)
+
+ summary = coll.feature_summary(0)
+
+ # Top-K = 4: should drop the 0.1 peak.
+ assert summary.top_peaks == pytest.approx([5.0, 3.0, 2.0, 0.5])
+ assert summary.activation_rate == pytest.approx(5 / (B * T))
+ assert summary.max_activation == 5.0
+ # Token ids of top sequence should match seq with peak 5.0 (b=4).
+ assert summary.top_token_ids[0] == tok[4].tolist()
+ # The activation trace's argmax should match where we placed the peak.
+ trace0 = summary.top_activations[0]
+ assert max(range(T), key=lambda i: trace0[i]) == pos_in[4]
+
+
+def test_topk_merges_across_multiple_batches():
+ coll = GateCollector(n_features=F, top_k=K, seq_len=T)
+ # Batch 1: peaks 1, 2, 3 for feature 0.
+ tok1 = torch.full((3, T), 11, dtype=torch.long)
+ g1 = _zero_gates(3)
+ g1[0, 0, 0] = 1.0
+ g1[1, 0, 0] = 2.0
+ g1[2, 0, 0] = 3.0
+ coll.update(tok1, g1)
+
+ # Batch 2: peaks 0.5, 4.0 for feature 0. Top-4 across the two batches
+ # should be {4.0, 3.0, 2.0, 1.0}, dropping 0.5.
+ tok2 = torch.full((2, T), 22, dtype=torch.long)
+ g2 = _zero_gates(2)
+ g2[0, 0, 0] = 0.5
+ g2[1, 0, 0] = 4.0
+ coll.update(tok2, g2)
+
+ summary = coll.feature_summary(0)
+ assert summary.top_peaks == [4.0, 3.0, 2.0, 1.0]
+ # Top sequence is from batch 2 (its tokens are all 22).
+ assert summary.top_token_ids[0] == [22] * T
+ # The 1.0-peak entry is from batch 1 (tokens all 11).
+ assert summary.top_token_ids[3] == [11] * T
+
+
+def test_dead_feature_summary_is_empty():
+ coll = GateCollector(n_features=F, top_k=K, seq_len=T)
+ g = _zero_gates(2)
+ g[0, 0, 0] = 1.0 # only feature 0 fires
+ coll.update(torch.zeros(2, T, dtype=torch.long), g)
+
+ dead = coll.feature_summary(2)
+ assert dead.activation_rate == 0.0
+ assert dead.max_activation == 0.0
+ assert dead.top_peaks == []
+ assert dead.top_token_ids == []
+ assert dead.top_activations == []
+
+
+def test_per_token_trace_is_feature_specific():
+ """The trace stored for feature i must be gates[..., i] for that sequence,
+ not any other feature."""
+ coll = GateCollector(n_features=F, top_k=K, seq_len=T)
+ tok = torch.zeros(1, T, dtype=torch.long)
+ g = _zero_gates(1)
+ g[0, :, 0] = torch.arange(T, dtype=torch.float32) # feat 0: 0..T-1
+ g[0, :, 1] = torch.arange(T, dtype=torch.float32) * 10 # feat 1: 0,10,20,...
+ coll.update(tok, g)
+
+ s0 = coll.feature_summary(0)
+ s1 = coll.feature_summary(1)
+ assert s0.top_activations[0] == list(range(T))
+ assert s1.top_activations[0] == [i * 10 for i in range(T)]
+
+
+def test_update_rejects_wrong_shape():
+ coll = GateCollector(n_features=F, top_k=K, seq_len=T)
+ bad_g = torch.zeros(1, T + 1, F)
+ try:
+ coll.update(torch.zeros(1, T + 1, dtype=torch.long), bad_g)
+ except ValueError:
+ return
+ raise AssertionError("expected ValueError for seq_len mismatch")
+
+
+def test_update_rejects_wrong_feature_count():
+ coll = GateCollector(n_features=F, top_k=K, seq_len=T)
+ bad_g = torch.zeros(1, T, F + 1)
+ try:
+ coll.update(torch.zeros(1, T, dtype=torch.long), bad_g)
+ except ValueError:
+ return
+ raise AssertionError("expected ValueError for feature count mismatch")
diff --git a/tests/test_feature_dash_dump.py b/tests/test_feature_dash_dump.py
new file mode 100644
index 0000000..baf3e63
--- /dev/null
+++ b/tests/test_feature_dash_dump.py
@@ -0,0 +1,292 @@
+"""Tests for windowing (Step 3) and JSON dump (Step 4).
+
+The windowing helper is checked with a stub tokenizer so the test stays CPU
+and offline. The dump test builds a tiny GateCollector + MoltCheckpointMetadata
+by hand, runs `dump_dashboard`, and checks the file layout and JSON schema.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+import torch
+
+from crosslayer_transcoder.feature_dash.collect import (
+ FeatureSummary,
+ GateCollector,
+ window_example,
+ window_feature_summary,
+)
+from crosslayer_transcoder.feature_dash.dump import dump_dashboard
+from crosslayer_transcoder.feature_dash.load import MoltCheckpointMetadata
+
+
+class _StubTokenizer:
+ """Minimal HF-tokenizer-like object: id -> string with a leading space."""
+
+ def decode(self, ids):
+ # Match HF behavior of single-id decode for byte-level BPE: returns
+ # the raw token text including leading whitespace where applicable.
+ return f" tok{ids[0]}"
+
+
+# ---- Step 3 -----------------------------------------------------------------
+
+
+def test_window_example_centers_on_peak():
+ ids = list(range(10)) # 0..9
+ acts = [0.0] * 10
+ acts[6] = 5.0 # peak at index 6
+
+ out = window_example(ids, acts, _StubTokenizer(), window=2)
+
+ # window=2 -> indices [4, 5, 6, 7, 8] = 5 tokens
+ assert out["tokens"] == [" tok4", " tok5", " tok6", " tok7", " tok8"]
+ assert out["activations"] == [0.0, 0.0, 5.0, 0.0, 0.0]
+ assert out["peak_token_pos"] == 2 # index of 6 inside the window
+ assert out["peak_activation"] == 5.0
+
+
+def test_window_example_clips_at_left_edge():
+ ids = list(range(10))
+ acts = [0.0] * 10
+ acts[1] = 3.0 # peak near the start
+
+ out = window_example(ids, acts, _StubTokenizer(), window=4)
+
+ # window=4 around pos 1 -> [0..5], length 6, peak at pos 1
+ assert out["tokens"] == [f" tok{i}" for i in range(6)]
+ assert out["peak_token_pos"] == 1
+ assert out["peak_activation"] == 3.0
+
+
+def test_window_example_clips_at_right_edge():
+ ids = list(range(10))
+ acts = [0.0] * 10
+ acts[9] = 7.0 # peak at the end
+
+ out = window_example(ids, acts, _StubTokenizer(), window=3)
+
+ # window=3 around pos 9 -> [6..9], length 4, peak at last pos within window
+ assert out["tokens"] == [" tok6", " tok7", " tok8", " tok9"]
+ assert out["peak_token_pos"] == 3
+ assert out["peak_activation"] == 7.0
+
+
+def test_window_example_rejects_mismatched_lengths():
+ with pytest.raises(ValueError):
+ window_example([1, 2, 3], [0.1, 0.2], _StubTokenizer(), window=1)
+
+
+def test_window_feature_summary_passes_through_metadata():
+ summary = FeatureSummary(
+ feature_id=42,
+ activation_rate=0.05,
+ max_activation=2.5,
+ top_peaks=[2.5, 1.0],
+ top_token_ids=[[10, 11, 12], [20, 21, 22]],
+ top_activations=[[0.0, 2.5, 0.0], [1.0, 0.0, 0.0]],
+ )
+ out = window_feature_summary(summary, _StubTokenizer(), window=10)
+
+ assert out["feature_id"] == 42
+ assert out["activation_rate"] == 0.05
+ assert out["max_activation"] == 2.5
+ assert len(out["examples"]) == 2
+ assert out["examples"][0]["peak_activation"] == 2.5
+ assert out["examples"][1]["peak_activation"] == 1.0
+
+
+# ---- Step 4 -----------------------------------------------------------------
+
+
+def _populate_collector(F: int, K: int, T: int) -> GateCollector:
+ coll = GateCollector(n_features=F, top_k=K, seq_len=T)
+ # feature 0 fires on two sequences with different peaks
+ tok = torch.arange(2 * T, dtype=torch.long).reshape(2, T)
+ g = torch.zeros(2, T, F)
+ g[0, 1, 0] = 1.0
+ g[1, 3, 0] = 2.5
+ g[0, 0, 1] = 0.5 # feature 1 fires once
+ coll.update(tok, g)
+ return coll
+
+
+def _meta(F: int, ranks: list[int], N: int) -> MoltCheckpointMetadata:
+ feature_tier: list[int] = []
+ feature_rank: list[int] = []
+ for t, r in enumerate(ranks):
+ n_in_tier = N * (2**t)
+ feature_tier.extend([t] * n_in_tier)
+ feature_rank.extend([r] * n_in_tier)
+ return MoltCheckpointMetadata(
+ ckpt_path="dummy.ckpt",
+ d_acts=8,
+ n_features=F,
+ n_layers=2,
+ ranks=ranks,
+ N=N,
+ feature_tier=feature_tier,
+ feature_rank=feature_rank,
+ base_model_name="stub",
+ training_dataset="stub-ds",
+ global_step=123,
+ epoch=0,
+ )
+
+
+def test_dump_dashboard_writes_metadata_and_features(tmp_path: Path):
+ F, K, T, N = 3, 2, 5, 1
+ ranks = [4, 2] # tier 0: 1 feat, tier 1: 2 feats -> 3
+ coll = _populate_collector(F, K, T)
+ meta = _meta(F, ranks, N)
+
+ data_dir = dump_dashboard(
+ collector=coll,
+ meta=meta,
+ tokenizer=_StubTokenizer(),
+ out_dir=tmp_path,
+ layer=8,
+ dataset_name="Skylion007/openwebtext",
+ window=2,
+ )
+
+ assert data_dir == tmp_path / "data"
+ assert (data_dir / "metadata.json").is_file()
+
+ feat_dir = data_dir / "features"
+ files = sorted(feat_dir.glob("*.json"))
+ assert len(files) == F
+ # Filenames are zero-padded to width=4 (max(4, len("2"))).
+ assert [p.name for p in files] == ["0000.json", "0001.json", "0002.json"]
+
+ md = json.loads((data_dir / "metadata.json").read_text())
+ assert md["schema_version"] == 1
+ assert md["n_features"] == F
+ assert md["ranks"] == ranks
+ assert md["layer"] == 8
+ assert md["dashboard_dataset"] == "Skylion007/openwebtext"
+ assert md["n_tokens_collected"] == 2 * T
+ assert md["seq_len"] == T
+ assert md["top_k"] == K
+ assert md["window"] == 2
+ assert md["feature_tier"] == [0, 1, 1]
+ assert md["feature_rank"] == [4, 2, 2]
+ # Per-feature aggregates so the index page renders from one fetch.
+ assert len(md["feature_activation_rate"]) == F
+ assert len(md["feature_max_activation"]) == F
+ assert md["feature_activation_rate"][0] == pytest.approx(2 / (2 * T))
+ assert md["feature_max_activation"][0] == pytest.approx(2.5)
+ assert md["feature_activation_rate"][2] == 0.0 # dead feature
+ assert md["feature_max_activation"][2] == 0.0 # max clamped from -inf to 0
+
+
+def test_dump_dashboard_copies_render_assets(tmp_path: Path):
+ F, K, T, N = 3, 2, 5, 1
+ ranks = [4, 2]
+ coll = _populate_collector(F, K, T)
+ meta = _meta(F, ranks, N)
+
+ dump_dashboard(
+ collector=coll,
+ meta=meta,
+ tokenizer=_StubTokenizer(),
+ out_dir=tmp_path,
+ layer=8,
+ window=2,
+ )
+ # Static assets land alongside data/.
+ for rel in [
+ "index.html",
+ "dashboard.html",
+ "assets/dashboard.css",
+ "assets/dashboard.js",
+ "assets/index.js",
+ ]:
+ assert (tmp_path / rel).is_file(), f"missing {rel}"
+
+
+def test_dump_dashboard_skips_assets_when_disabled(tmp_path: Path):
+ F, K, T, N = 3, 2, 5, 1
+ coll = _populate_collector(F, K, T)
+ meta = _meta(F, [4, 2], N)
+
+ dump_dashboard(
+ collector=coll,
+ meta=meta,
+ tokenizer=_StubTokenizer(),
+ out_dir=tmp_path,
+ layer=8,
+ window=2,
+ copy_assets=False,
+ )
+ assert not (tmp_path / "index.html").exists()
+ assert not (tmp_path / "assets").exists()
+ # data/ still written.
+ assert (tmp_path / "data" / "metadata.json").is_file()
+
+
+def test_dump_dashboard_feature_payload_schema(tmp_path: Path):
+ F, K, T, N = 3, 2, 5, 1
+ ranks = [4, 2]
+ coll = _populate_collector(F, K, T)
+ meta = _meta(F, ranks, N)
+
+ dump_dashboard(
+ collector=coll,
+ meta=meta,
+ tokenizer=_StubTokenizer(),
+ out_dir=tmp_path,
+ layer=8,
+ window=2,
+ )
+ feat0 = json.loads((tmp_path / "data" / "features" / "0000.json").read_text())
+
+ # Required schema keys.
+ assert set(feat0.keys()) >= {
+ "feature_id",
+ "tier",
+ "rank",
+ "activation_rate",
+ "max_activation",
+ "examples",
+ }
+ assert feat0["feature_id"] == 0
+ assert feat0["tier"] == 0
+ assert feat0["rank"] == 4
+ assert feat0["activation_rate"] == pytest.approx(2 / (2 * T))
+ assert feat0["max_activation"] == pytest.approx(2.5)
+
+ # Two examples, sorted desc by peak.
+ assert len(feat0["examples"]) == 2
+ assert feat0["examples"][0]["peak_activation"] == pytest.approx(2.5)
+ assert feat0["examples"][1]["peak_activation"] == pytest.approx(1.0)
+
+ ex0 = feat0["examples"][0]
+ assert len(ex0["tokens"]) == len(ex0["activations"])
+ # peak_token_pos must index the argmax of the windowed activations.
+ assert (
+ max(range(len(ex0["activations"])), key=lambda i: ex0["activations"][i])
+ == ex0["peak_token_pos"]
+ )
+
+
+def test_dump_dashboard_dead_feature_yields_no_examples(tmp_path: Path):
+ F, K, T, N = 3, 2, 5, 1
+ ranks = [4, 2]
+ coll = _populate_collector(F, K, T) # feature 2 never fires
+ meta = _meta(F, ranks, N)
+
+ dump_dashboard(
+ collector=coll,
+ meta=meta,
+ tokenizer=_StubTokenizer(),
+ out_dir=tmp_path,
+ layer=8,
+ window=2,
+ )
+ feat2 = json.loads((tmp_path / "data" / "features" / "0002.json").read_text())
+ assert feat2["activation_rate"] == 0.0
+ assert feat2["examples"] == []
diff --git a/tests/test_feature_dash_integration.py b/tests/test_feature_dash_integration.py
new file mode 100644
index 0000000..38cb2ba
--- /dev/null
+++ b/tests/test_feature_dash_integration.py
@@ -0,0 +1,189 @@
+"""Integration smoke test for the feature-dash pipeline.
+
+Wires together: build tiny MoLT -> save Lightning ckpt -> load_molt ->
+GateCollector with synthetic gates -> dump_dashboard -> assert the full file
+contract holds.
+
+We bypass `collect_features` (which would need a real GPT-2 + HF dataset) by
+feeding the collector hand-crafted batches. The integration we care about
+here is the file pipeline: that the loaded MoLT's metadata flows correctly
+into the JSON, and that all the per-feature contracts the renderer relies on
+are satisfied.
+
+CLI smoke (Step 6) is also exercised here via argparse `--help`.
+"""
+
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+import lightning as L
+import pytest
+import torch
+from transformers import GPT2TokenizerFast
+
+from crosslayer_transcoder.feature_dash import (
+ GateCollector,
+ dump_dashboard,
+ load_molt,
+)
+from crosslayer_transcoder.model.clt_lightning import MoltModule
+from crosslayer_transcoder.model.jumprelu import JumpReLU
+from crosslayer_transcoder.model.molt import Molt
+from crosslayer_transcoder.model.standardize import (
+ DimensionwiseInputStandardizer,
+ DimensionwiseOutputStandardizer,
+)
+
+
+D_ACTS = 16
+N_LAYERS_LM = 4
+N = 2
+RANKS = [4, 2] # tier 0: 2 feats, tier 1: 4 feats -> 6 features total
+LAYER = 1
+SEQ_LEN = 8
+TOP_K = 3
+
+
+def _build_molt() -> Molt:
+ n_features = N * 1 + N * 2 # 6
+ nonlin = JumpReLU(theta=0.05, bandwidth=1.0, n_layers=1, d_features=n_features)
+ in_std = DimensionwiseInputStandardizer(n_layers=N_LAYERS_LM, activation_dim=D_ACTS)
+ out_std = DimensionwiseOutputStandardizer(
+ n_layers=N_LAYERS_LM, activation_dim=D_ACTS
+ )
+ fake_batch = torch.randn(8, 2, N_LAYERS_LM, D_ACTS)
+ in_std.initialize_from_batch(fake_batch)
+ out_std.initialize_from_batch(fake_batch)
+
+ return Molt(
+ d_acts=D_ACTS,
+ N=N,
+ ranks=RANKS,
+ nonlinearity=nonlin,
+ input_standardizer=in_std,
+ output_standardizer=out_std,
+ )
+
+
+def _save_lightning_checkpoint(tmp_path: Path, molt: Molt) -> str:
+ module = MoltModule(model=molt)
+ trainer = L.Trainer(
+ accelerator="cpu",
+ devices=1,
+ logger=False,
+ enable_checkpointing=False,
+ enable_progress_bar=False,
+ enable_model_summary=False,
+ max_steps=0,
+ )
+ trainer.strategy.connect(module)
+ ckpt = tmp_path / "molt.ckpt"
+ trainer.save_checkpoint(str(ckpt))
+ return str(ckpt)
+
+
+def _populate_collector_with_molt_gates(
+ molt: Molt, collector: GateCollector, n_batches: int = 4, batch_size: int = 4
+) -> None:
+ """Run synthetic resid through the loaded MoLT to get realistic gates."""
+ torch.manual_seed(0)
+ for _ in range(n_batches):
+ # Random residuals; the standardizer will re-center them.
+ resid = torch.randn(batch_size, SEQ_LEN, D_ACTS)
+ resid_std = molt.input_standardizer(resid, LAYER)
+ gates = molt.nonlinearity(molt.e(resid_std)) # (B, T, F)
+ # Synthetic token ids: just whatever, the renderer doesn't care.
+ token_ids = torch.randint(
+ 0, 1000, (batch_size, SEQ_LEN), dtype=torch.long
+ )
+ collector.update(token_ids, gates.float())
+
+
+def test_full_pipeline_contract(tmp_path: Path):
+ """Load -> collect -> dump produces a directory the renderer can consume."""
+ src_molt = _build_molt()
+ ckpt_path = _save_lightning_checkpoint(tmp_path, src_molt)
+
+ molt, meta = load_molt(ckpt_path, device="cpu")
+ assert meta.n_features == src_molt.n_features
+
+ collector = GateCollector(
+ n_features=meta.n_features, top_k=TOP_K, seq_len=SEQ_LEN
+ )
+ _populate_collector_with_molt_gates(molt, collector)
+
+ tokenizer = GPT2TokenizerFast.from_pretrained("openai-community/gpt2")
+
+ out_dir = tmp_path / "dash"
+ dump_dashboard(
+ collector=collector,
+ meta=meta,
+ tokenizer=tokenizer,
+ out_dir=out_dir,
+ layer=LAYER,
+ dataset_name="synthetic",
+ window=2,
+ )
+
+ # --- File layout ---
+ for rel in [
+ "data/metadata.json",
+ "index.html",
+ "dashboard.html",
+ "assets/dashboard.css",
+ "assets/dashboard.js",
+ "assets/index.js",
+ ]:
+ assert (out_dir / rel).is_file(), f"missing {rel}"
+ feat_files = sorted((out_dir / "data" / "features").glob("*.json"))
+ assert len(feat_files) == meta.n_features
+
+ # --- metadata.json invariants the renderer relies on ---
+ md = json.loads((out_dir / "data" / "metadata.json").read_text())
+ assert md["schema_version"] == 1
+ assert md["n_features"] == meta.n_features
+ assert len(md["feature_activation_rate"]) == meta.n_features
+ assert len(md["feature_max_activation"]) == meta.n_features
+ assert len(md["feature_tier"]) == meta.n_features
+ assert len(md["feature_rank"]) == meta.n_features
+ assert all(0.0 <= r <= 1.0 for r in md["feature_activation_rate"])
+ assert all(m >= 0.0 for m in md["feature_max_activation"])
+
+ # --- per-feature invariants ---
+ for fp in feat_files:
+ d = json.loads(fp.read_text())
+ assert {"feature_id", "tier", "rank", "activation_rate",
+ "max_activation", "examples"} <= set(d.keys())
+ assert 0.0 <= d["activation_rate"] <= 1.0
+ for ex in d["examples"]:
+ assert {"peak_activation", "peak_token_pos", "tokens", "activations"} <= set(ex.keys())
+ assert len(ex["tokens"]) == len(ex["activations"])
+ assert 0 <= ex["peak_token_pos"] < len(ex["tokens"])
+ # peak_token_pos must index the argmax of the windowed activations.
+ argmax = max(
+ range(len(ex["activations"])),
+ key=lambda i: ex["activations"][i],
+ )
+ assert argmax == ex["peak_token_pos"]
+ # Peak activation must be > 0 (dead-feature filter in feature_summary).
+ assert ex["peak_activation"] > 0
+
+
+def test_cli_help_runs():
+ """Make sure argparse construction (and therefore imports) don't blow up."""
+ result = subprocess.run(
+ [sys.executable, "-m", "crosslayer_transcoder.feature_dash", "--help"],
+ capture_output=True,
+ text=True,
+ timeout=60,
+ )
+ assert result.returncode == 0, result.stderr
+ assert "feature dashboard" in result.stdout.lower()
+ # Both source flags should appear.
+ assert "--hf-filename" in result.stdout
+ assert "--local-ckpt" in result.stdout
+ assert "--out" in result.stdout
diff --git a/tests/test_feature_dash_load.py b/tests/test_feature_dash_load.py
new file mode 100644
index 0000000..9ceb59a
--- /dev/null
+++ b/tests/test_feature_dash_load.py
@@ -0,0 +1,124 @@
+"""Smoke test for the MoLT feature-dash checkpoint loader.
+
+Builds a tiny MoLT, wraps it in a MoltModule, saves a Lightning checkpoint via
+trainer.save_checkpoint, then loads it back with `load_molt` and checks:
+
+ - architecture is recovered (d_acts, n_features, n_layers, ranks, N),
+ - standardizer buffers round-trip,
+ - JumpReLU theta round-trips,
+ - forward output matches the source model bit-for-bit on a fixed input,
+ - the per-feature tier/rank index lines up with the rank tier sizes.
+"""
+
+from __future__ import annotations
+
+import lightning as L
+import pytest
+import torch
+
+from crosslayer_transcoder.feature_dash.load import infer_molt_arch, load_molt
+from crosslayer_transcoder.model.clt_lightning import MoltModule
+from crosslayer_transcoder.model.jumprelu import JumpReLU
+from crosslayer_transcoder.model.molt import Molt
+from crosslayer_transcoder.model.standardize import (
+ DimensionwiseInputStandardizer,
+ DimensionwiseOutputStandardizer,
+)
+
+
+D_ACTS = 16
+N_LAYERS = 4
+N = 3
+RANKS = [8, 4] # tier 0: 3 transforms, tier 1: 6 transforms -> 9 features
+LAYER = 2
+
+
+def _build_molt() -> Molt:
+ n_features = N * 1 + N * 2 # = 9
+ nonlin = JumpReLU(theta=0.05, bandwidth=1.0, n_layers=1, d_features=n_features)
+ in_std = DimensionwiseInputStandardizer(n_layers=N_LAYERS, activation_dim=D_ACTS)
+ out_std = DimensionwiseOutputStandardizer(n_layers=N_LAYERS, activation_dim=D_ACTS)
+
+ fake_batch = torch.randn(8, 2, N_LAYERS, D_ACTS)
+ in_std.initialize_from_batch(fake_batch)
+ out_std.initialize_from_batch(fake_batch)
+
+ return Molt(
+ d_acts=D_ACTS,
+ N=N,
+ ranks=RANKS,
+ nonlinearity=nonlin,
+ input_standardizer=in_std,
+ output_standardizer=out_std,
+ )
+
+
+def _save_lightning_checkpoint(tmp_path, molt: Molt) -> str:
+ module = MoltModule(model=molt)
+ trainer = L.Trainer(
+ accelerator="cpu",
+ devices=1,
+ logger=False,
+ enable_checkpointing=False,
+ enable_progress_bar=False,
+ enable_model_summary=False,
+ max_steps=0,
+ )
+ # `strategy.connect` wires the module so save_checkpoint works without
+ # actually running `fit`.
+ trainer.strategy.connect(module)
+ ckpt_path = tmp_path / "molt.ckpt"
+ trainer.save_checkpoint(str(ckpt_path))
+ return str(ckpt_path)
+
+
+def test_infer_molt_arch_from_state_dict():
+ molt = _build_molt()
+ sd = {f"model.{k}": v for k, v in molt.state_dict().items()}
+ # mimic what Lightning would add
+ sd["last_active"] = torch.zeros((molt.n_features,), dtype=torch.long)
+
+ arch = infer_molt_arch(sd)
+ assert arch == {
+ "d_acts": D_ACTS,
+ "n_features": N * 1 + N * 2,
+ "n_layers": N_LAYERS,
+ "ranks": RANKS,
+ "N": N,
+ }
+
+
+def test_load_molt_roundtrip(tmp_path):
+ torch.manual_seed(0)
+ src = _build_molt()
+ ckpt_path = _save_lightning_checkpoint(tmp_path, src)
+
+ loaded, meta = load_molt(ckpt_path, device="cpu")
+
+ assert meta.d_acts == D_ACTS
+ assert meta.n_features == src.n_features
+ assert meta.ranks == RANKS
+ assert meta.N == N
+ assert meta.n_layers == N_LAYERS
+
+ # Tier index has the right tier counts.
+ assert meta.feature_tier.count(0) == N
+ assert meta.feature_tier.count(1) == 2 * N
+ assert meta.feature_rank[0] == RANKS[0]
+ assert meta.feature_rank[-1] == RANKS[1]
+ assert len(meta.feature_tier) == src.n_features
+
+ # Forward output matches bit-for-bit on a fixed input.
+ src.eval()
+ loaded.eval()
+ x = torch.randn(5, D_ACTS)
+ with torch.no_grad():
+ g_src, rn_src, r_src = src(x, layer=LAYER)
+ g_loaded, rn_loaded, r_loaded = loaded(x, layer=LAYER)
+ assert torch.allclose(g_src, g_loaded)
+ assert torch.allclose(rn_src, rn_loaded)
+ assert torch.allclose(r_src, r_loaded)
+
+ # Loaded model must be in eval and have grads disabled.
+ assert not loaded.training
+ assert all(not p.requires_grad for p in loaded.parameters())
diff --git a/tests/test_molt_smoke.py b/tests/test_molt_smoke.py
new file mode 100644
index 0000000..ebb28f6
--- /dev/null
+++ b/tests/test_molt_smoke.py
@@ -0,0 +1,110 @@
+"""Smoke tests for the MoLT (Mixture of Low-rank Transcoders) port.
+
+CPU test exercises the model wiring (low-rank transforms + JumpReLU + standardizers).
+GPU test runs a fp32 and a 16-mixed forward+backward+optimizer step on synthetic
+activations and checks that all parameters and the loss remain finite — mirrors
+what the full-config CLI smoke covers, without the data generator or wandb.
+"""
+
+import pytest
+import torch
+
+from crosslayer_transcoder.model.jumprelu import JumpReLU
+from crosslayer_transcoder.model.molt import Molt
+from crosslayer_transcoder.model.standardize import (
+ DimensionwiseInputStandardizer,
+ DimensionwiseOutputStandardizer,
+)
+
+
+D_ACTS = 64
+N_LAYERS = 12
+N = 4
+RANKS = [8, 4]
+B = 4
+
+
+def _build_molt(device: torch.device) -> Molt:
+ n_features = N + 2 * N
+ nonlin = JumpReLU(theta=0.03, bandwidth=1.0, n_layers=1, d_features=n_features)
+ in_std = DimensionwiseInputStandardizer(n_layers=N_LAYERS, activation_dim=D_ACTS)
+ out_std = DimensionwiseOutputStandardizer(n_layers=N_LAYERS, activation_dim=D_ACTS)
+
+ fake_batch = torch.randn(B, 2, N_LAYERS, D_ACTS)
+ in_std.initialize_from_batch(fake_batch)
+ out_std.initialize_from_batch(fake_batch)
+
+ return Molt(
+ d_acts=D_ACTS,
+ N=N,
+ ranks=RANKS,
+ nonlinearity=nonlin,
+ input_standardizer=in_std,
+ output_standardizer=out_std,
+ ).to(device)
+
+
+def test_molt_cpu_forward():
+ torch.manual_seed(0)
+ m = _build_molt(torch.device("cpu"))
+ acts = torch.randn(B, D_ACTS)
+
+ gate, recons_norm, recons = m(acts, layer=8)
+
+ assert gate.shape == (B, m.n_features)
+ assert recons_norm.shape == (B, D_ACTS)
+ assert recons.shape == (B, D_ACTS)
+ assert torch.isfinite(gate).all()
+ assert torch.isfinite(recons).all()
+
+
+def _run_train_step(device, autocast_dtype):
+ """Forward + backward + optimizer step. Mirrors MoltModule.training_step
+ without depending on the Lightning Trainer."""
+ torch.manual_seed(0)
+ m = _build_molt(device)
+ optim = torch.optim.Adam(m.parameters(), lr=2e-4)
+ scaler = torch.amp.GradScaler("cuda", enabled=autocast_dtype is torch.float16)
+
+ resid = torch.randn(B, D_ACTS, device=device)
+ mlp_out = torch.randn(B, D_ACTS, device=device)
+
+ optim.zero_grad(set_to_none=True)
+ if autocast_dtype is None:
+ gate, recons_norm, _ = m(resid, layer=8)
+ target = m.output_standardizer.standardize(mlp_out, 8)
+ mse = ((recons_norm - target) ** 2).mean()
+ norms = m.transform_norm()
+ sparsity = torch.tanh(norms * gate * 100.0).sum(dim=-1).mean() * 1.5e-4
+ loss = mse + sparsity
+ loss.backward()
+ optim.step()
+ else:
+ with torch.amp.autocast("cuda", dtype=autocast_dtype):
+ gate, recons_norm, _ = m(resid, layer=8)
+ target = m.output_standardizer.standardize(mlp_out, 8)
+ mse = ((recons_norm - target) ** 2).mean()
+ norms = m.transform_norm()
+ sparsity = torch.tanh(norms * gate * 100.0).sum(dim=-1).mean() * 1.5e-4
+ loss = mse + sparsity
+ scaler.scale(loss).backward()
+ scaler.step(optim)
+ scaler.update()
+
+ return loss, m
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_molt_gpu_fp32_train_step():
+ loss, m = _run_train_step(torch.device("cuda"), autocast_dtype=None)
+ assert torch.isfinite(loss), f"non-finite loss: {loss.item()}"
+ for name, p in m.named_parameters():
+ assert torch.isfinite(p).all(), f"non-finite param after step: {name}"
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_molt_gpu_amp_train_step():
+ loss, m = _run_train_step(torch.device("cuda"), autocast_dtype=torch.float16)
+ assert torch.isfinite(loss), f"non-finite loss: {loss.item()}"
+ for name, p in m.named_parameters():
+ assert torch.isfinite(p).all(), f"non-finite param after AMP step: {name}"