TorchInstruments

Record per-layer activation and output-gradient statistics without changing your training loop.

from datetime import timedelta
from torchinstruments import inject_observer, remove_observer

inject_observer(model, interval=timedelta(minutes=1), output_dir="stats")
try:
    train(model)
finally:
    remove_observer(model)

The first forward is sampled immediately, then once per interval. The default interval is one minute. Evaluation works too, including torch.no_grad() and torch.inference_mode().

Calling a child directly, such as model.encoder(x) inside a Lightning step, also works. When the observed root is bypassed, selected layers sample independently at the given interval. Each selected module's forward() is intercepted once, so both module(x) and module.forward(x) are captured. Removal restores the original methods.

Output

Everything is created automatically:

history.parquet, result.json, index.md, and TensorBoard update after each sampled forward or backward. They are readable without calling remove_observer(). Removal detaches hooks, closes resources, and cleans up intermediate history.parts/ chunks. Interrupted runs retain the last published files and completed chunks. Each refresh rewrites the Parquet snapshot and aggregates collected history, so its cost grows with run length.

Statistics

Each sampled tensor reports mean, population standard deviation, min, max, p25, p50, p75, zero fraction, and nonfinite fraction. Gradients are recorded separately. Train/eval context is captured automatically, and missing measurements carry reasons.

JSON aggregates each metric's history, including previous/recent window means. These summarize sampled statistics, not pooled tensor entries. Exact observations stay in Parquet. There are no scores or diagnoses. Give index.md to an LLM as the starting point.

Layer selection

All modules, including composite blocks and the root, contribute scalar history when called. The root has the empty layer name "". Histograms cover all selected modules by default. Set max_histogram_modules to a positive integer to explicitly limit histogram work and event size. To focus the dashboard, use histogram_selector:

inject_observer(
    model,
    histogram_selector=lambda name, module: name.startswith("encoder.blocks.7."),
)

Use selector similarly to restrict all collection. Histogram focus must match an observed module. Use selector=leaf_modules() (imported from torchinstruments) for leaf-only coverage. Other existing extension arguments remain available; ordinary usage needs none of them.

Examples and limits

• Examples: ordinary training, Lightning, and controlled model problems.
• LLM guide: how to interpret measurements and query history.
• Benchmarks: interval-based usage and separate stress tests.
• Changelog: version changes and schema migration.

CPU, native MPS, in-place activations, and non-reentrant checkpointing are tested. Reentrant checkpoint internals lack gradient coverage. Only the first backward is recorded, with the caller's loss scaling. CUDA performance and torch.compile compatibility are unverified. The interval amortizes collection cost; sampled forwards still scan their tensors.

MIT licensed. Author: Vadym Stupakov vadim.stupakov@gmail.com.