06 — Experiment tracking and training diagnostics

Three flows: (a) a training run with a per-channel gradient audit, (b) a hyperparameter study with a composite objective and pruning, and (c) following or resuming a study from disk. All of it lands in .soma/, for a front end to read.

import json, pathlib, time
import torch

torch.manual_seed(0)  # the numbers in this notebook are reproducible
import torch.nn as nn

import soma
from soma import ChannelConfig, DifferentiableFilter, Graph, search

# A benign PyTorch warning about backward hooks (the audit uses them
# on purpose); the test suite silences it anyway.
import warnings
warnings.filterwarnings("ignore", message="Full backward hook is firing")

(a) A training run, audited

track_run creates .soma/runs/<run_id>/ — a manifest with git and host, a status with a heartbeat, the graph’s topology, events/metrics.jsonl. gradient_audit with channels= adds per-channel diagnostics: dead channels, dormant ones (Sokar 2023), ignored ones (gradient starvation), and leakage between groups (CKA).

class Encoder(DifferentiableFilter):
    _cache_version = "nb06-encoder-v1"
    lr: float = search(1e-3, 1e-1, scale="log")

    def build_module(self, input_shape):
        return nn.Sequential(nn.Linear(input_shape[-1], 16), nn.ReLU(), nn.Linear(16, 8))

    def output_shape(self, input_shape):
        return (*input_shape[:-1], 8)

g = Graph()
g.node("encoder", Encoder())
x = torch.randn(64, 12)
y = torch.randn(64, 8)
g.materialize(x)
g.train()
g.make_optimizer(lr=0.01)

Adam (
Parameter Group 0
    amsgrad: False
    betas: (0.9, 0.999)
    capturable: False
    decoupled_weight_decay: False
    differentiable: False
    eps: 1e-08
    foreach: None
    fused: None
    lr: 0.01
    maximize: False
    weight_decay: 0
)

with g.track_run("baseline", tags=["demo"]) as run:
    cfg = ChannelConfig(snapshot_every=10,
                        groups={"encoder": {"a": range(0, 4), "b": range(4, 8)}})
    with g.gradient_audit(channels=cfg) as audit:
        module = dict(g.filters())["encoder"]._module
        for epoch in range(5):
            run.log_epoch(epoch, total=5)
            with g.context() as ctx:
                g.zero_grad()
                loss = ((module(x) - y) ** 2).mean()
                g.backward(ctx, loss)   # the audit's snapshot + StepCompleted
            g.step(ctx)
            run.log("loss", loss.detach().item(), step=epoch)
    print(audit.report().pretty())

run_dir = pathlib.Path(run.dir)
print(sorted(p.name for p in run_dir.iterdir()))

                  filter   steps      act|μ|       act σ      |out∂|        |θ∂|         |θ|         ∂/θ                     flags
----------------------------------------------------------------------------------------------------------------------------------
                 encoder       5   1.949e-01   2.314e-01   9.288e-02   3.302e-01   2.843e+00   1.158e-01                   HEALTHY
['diagnostics', 'events.jsonl', 'fingerprint.json', 'graph.json', 'graph.mmd', 'manifest.json', 'metrics.jsonl', 'status.json']

# Everything a front end needs: events, metrics and diagnostics
print((run_dir / "graph.mmd").read_text())
print(json.loads((run_dir / "status.json").read_text()))
print((run_dir / "diagnostics" / "report.json").read_text()[:400])

graph LR
    encoder[encoder]

{'state': 'completed', 'updated_at': '2026-08-04T00:34:15.551753905Z', 'heartbeat_at': '2026-08-04T00:34:15.551753905Z', 'finished_at': '2026-08-04T00:34:15.551753905Z'}
{
  "n_steps": 5,
  "filters": [
    {
      "filter": "encoder",
      "n_steps": 5,
      "metrics": {
        "act_mean_abs": 0.19487954676151276,
        "act_std": 0.23137009441852568,
        "act_zero_frac": 0.0,
        "act_zero_frac_max": 0.0,
        "act_sat_frac_max": 0.0,
        "out_grad_norm": 0.09288433492183686,
        "out_grad_max": 0.01320467609912157,
        "param_grad_no

(b) A study: grid, a composite objective, and pruning

The space comes from the filters’ own search() descriptors (graph.search_space(), named node.param). The objective can be a callable over the metrics; pruning uses the median stopping rule through trial.report().

study = g.study("demo-grid", strategy="grid", n_trials=3,
                objective=lambda m: m["fit"] - 0.1 * m["cost"],
                direction="maximize", pruning=("median", 2))

def train(trial):
    g.apply_params(trial.params)
    lr = trial["encoder.lr"]
    for step in range(8):
        fit = 1.0 - abs(lr - 0.01) * 10 + step * 0.01
        if trial.report("fit", fit, step):
            return None                      # podado
    return {"fit": fit, "cost": lr * 100}

# Live events: aggregate progress only (TrialMetric arrives one per
# report(); imprimirlos todos es ruido).
def progress(e):
    if e["event_type"] == "StudyProgress":
        print(f" → {e['completed']}/{e['total']} trials, best={e['best_value']:.3f}")

study.run(train, on_event=progress)
time.sleep(0.3)   # the event callback is async: let the last lines drain
for t in study.trials:
    mark = "✂" if t["state"] == "pruned" else "✓"
    print(f" {mark} {t['id']}  {t['state']}")
best = study.best_trial
print(f"best: {best['id']}  params={best['params']}  score={best['metrics']['score']:.3f}")
print("run dir:", study.run_dir)

 → 1/3 trials, best=0.970
 → 2/3 trials, best=0.970
 → 3/3 trials, best=0.970
 ✓ trial_0000  completed
 ✓ trial_0001  completed
 ✓ trial_0002  completed
best: trial_0000  params={'encoder.lr': 0.0010000000000000002}  score=0.970
run dir: .soma/runs/study_20260804T003415_ea87

(c) Following and resuming from disk

study.json is rewritten atomically after every trial, so any machine with access to the directory can load the state, and resume=True picks up exactly where it stopped — without repeating grid points.

reloaded = soma.Study.load(study.run_dir)
print(reloaded.progress, len(reloaded.trials))
# reloaded.run(train, resume=True)   # would carry on if trials remained

for exp in soma.experiments():
    print(exp["name"], exp["metrics"], exp["tags"])

1.0 3
baseline {'loss': 1.0386008024215698} ['demo', 'run:run_20260804T003415_ba6e']
demo-grid {'cost': 0.10000000000000002, 'fit': 0.98, 'score': 0.97} ['run:study_20260804T003415_ea87']