TorchInstruments
Turn “accuracy stalled” into evidence about what changed inside the model.
A loss curve says that a run is underperforming. It does not say whether gradients started collapsing at one layer, activations developed a heavy tail, scale began drifting, or a module entered a new unstable regime. TorchInstruments follows those internal signals while training continues normally.
inject once → train normally → stats.json updates live → ask a narrower question
Quick start
from torchinstruments import inject_observer
inject_observer(model, output_dir="stats")
train(model)
There is no telemetry call inside the training loop. Open stats/index.md, or tell a
filesystem-capable LLM:
Read stats/index.md and stats/stats.json. Find the strongest evidence for activation drift,
gradient collapse, heavy-tail growth, oscillation, or a regime change. For each finding, give
the exact layer and indicators, plausible mechanisms, missing evidence, and the smallest
controlled experiment.
TorchInstruments supports Python 3.11 and newer. Run the complete local example with:
git clone https://github.com/Red-Eyed/torchinstruments.git
cd torchinstruments
uv sync --dev
uv run examples/basic_training.py
What the result looks like
Suppose a modified model stalls below its baseline. TorchInstruments can support a result such as:
Observed:
encoder.blocks.7.projoutput RMS rose from0.82to1.71. Its fast/slow EMA gap reached0.083, linear slope is positive withR²=0.94,p999_abs_to_rmsdoubled, and output-gradient RMS fell from0.0081to0.0002with a large drawdown.What it suggests: scale and tail growth coincide with weakening gradient flow at block 7. Normalization, residual scaling, or saturation there is more plausible than a model-wide optimizer failure.
Next experiment: restore the previous normalization or residual scale at block 7 only, while holding the seed and data order fixed.
The evidence narrows the hypothesis space. It does not pretend that correlation proves cause.
More than mean and standard deviation
Two tensor distributions can have the same mean and standard deviation while having completely different tails. The default sampled profile therefore includes:
- location and scale: mean, standard deviation, RMS, extrema, mean absolute value, L1/L2 norms;
- shape: median, quartiles,
p01–p999, skewness, excess kurtosis, and central ranges; - tails:
p99_abs,p999_abs, max/RMS ratios, and mass beyond three standard deviations; - prevalence: finite, zero, positive, and negative fractions;
- concentration: normalized magnitude entropy and effective support;
- optional fixed-bin histograms with explicit underflow, overflow, and non-finite counts.
For important metrics such as RMS, finite fraction, skewness, kurtosis, and tail ratios, TorchInstruments maintains technical-analysis-like indicators over sampled forwards:
- fast and slow EMAs plus their absolute and relative gap;
- momentum over several horizons;
- linear slope and
R²; - exponentially weighted change and volatility;
- z-score against prior behavior;
- drawdown, runup, and historical-range position;
- directional up/down balance;
- CUSUM regime-change scores;
- lag-one autocorrelation, oscillation fraction, and consecutive directional runs.
Every series records its observation count and warm-up state so an LLM can distinguish mature evidence from a three-sample coincidence.
One live file, not thousands of samples
stats/
index.md
stats.json
stats.json is the single canonical telemetry record. It contains run metadata, the module
catalog, forward and backward layer summaries, current distributions, temporal indicators,
mergeable histograms, observer overhead, and bounded error summaries. It is atomically replaced
after sampled forward and backward observations, so readers never see a partial file.
No per-sample files or raw tensors are persisted. First/latest/extreme values retain their sample IDs and timestamps. Memory used by temporal windows, metric series, and error identities is explicitly bounded.
What is monitored by default
| Boundary | Default behavior |
|---|---|
| Sampling | First root forward after each 60-second monotonic interval |
| Modules | Leaf modules, without redundant container outputs |
| Forward | Every tensor in selected-module outputs, including nested structures |
| Backward | Gradient with respect to every differentiable selected-module output |
| Distribution | Rich finite-value shape, tail, prevalence, and concentration statistics |
| Temporal analysis | Bounded trend, momentum, volatility, regime, and stability indicators |
| Histograms | Disabled; fixed ranges are recommended for exact live aggregation |
| Persistence | One strict, atomically updated stats.json |
| Failures | Warn by default and aggregate instrumentation errors in telemetry |
The current release does not monitor module inputs, grad_input, parameters, parameter gradients,
losses, optimizer state, or optimizer updates. It never calls a sampled parameter change an
optimizer update.
Direct forward() calls
Normal PyTorch module(...) dispatch uses native hooks. If model code literally invokes
module.forward(...), enable reversible direct-forward capture:
inject_observer(model, output_dir="stats", capture_direct_forwards=True)
The root and recursively selected modules are wrapped once, so mixed module(...) and
module.forward(...) execution is observed exactly once. remove_observer(model) restores the
previous instance attributes.
Histograms
Histograms are opt-in because they cost more than scalar reductions. Fixed bins are exactly mergeable across the run:
from torchinstruments import histogram, inject_observer
inject_observer(
model,
histograms=[
histogram(
bins=64,
value_range=(-8.0, 8.0),
every_n_samples=10,
),
],
)
Dynamic-bin histograms retain their latest distribution, but cannot be merged after their edges change. The live JSON records that limitation explicitly.
Lightning and TensorBoard
Use the same externally owned Lightning logger for task metrics and model telemetry:
from lightning.pytorch.loggers import TensorBoardLogger
from torchinstruments import CompositeSink, DirectorySink, TensorBoardSink, inject_observer
logger = TensorBoardLogger(save_dir="logs", name="experiment")
inject_observer(
model.network,
sink=CompositeSink(
DirectorySink("stats"),
TensorBoardSink(logger),
),
)
TorchInstruments never closes the caller's logger. TensorBoard receives live per-sample events;
stats.json retains bounded online indicators rather than the complete dashboard history. The
tested MNIST example
demonstrates the complete integration.
Configure indicator windows
The default path needs no configuration. Research-specific horizons remain explicit at the sink boundary:
from torchinstruments import DirectorySink, IndicatorConfig, LiveAggregator, inject_observer
config = IndicatorConfig(
momentum_horizons=(1, 10, 100),
recent_window=100,
warmup_observations=100,
)
sink = DirectorySink("stats", aggregator_factory=lambda: LiveAggregator(config))
inject_observer(model, sink=sink)
max_series, max_tensor_paths, max_module_calls, max_histograms, and
max_error_summaries provide hard bounds for dynamic model structure and error messages. The
live record counts observations omitted by each structural limit.
Research and LLM use cases
TorchInstruments helps investigate:
- where gradient signal first weakens or amplifies;
- whether activation scale is drifting or merely oscillating;
- whether skew, kurtosis, or extreme-to-RMS ratios are growing;
- which module first emits non-finite values;
- whether an architecture change creates a new internal regime;
- whether rare outliers make a layer quantization-hostile;
- whether internal behavior is stable enough to redirect investigation toward data, loss, or evaluation.
See the LLM analysis guide and research workflows for evidence-constrained prompts and baseline-versus-candidate experiments.
Safety and compatibility
- Injection adds no parameters, buffers, or modules;
state_dict()remains unchanged. - Outputs and gradients remain bit-identical in the test suite.
- Unsampled callbacks perform only a cheap context lookup.
- Sampled reductions happen on the tensor device; only compact results move to CPU.
- Raw activations and gradients are never written to disk.
- Duplicate injection raises
ObserverAlreadyAttachedError. - Python 3.11–3.14 and PyTorch 2.0+ are declared; CUDA, Accelerate, distributed output, and
torch.compileremain unclaimed until dedicated compatibility tests exist.
The core wheel depends only on PyTorch and the Python standard library. Lightning, TensorBoard, and torchvision are development/example dependencies.
License
TorchInstruments is released under the MIT License.
Citation
If TorchInstruments supports your research or engineering work, cite it as:
@software{stupakov_2026_torchinstruments,
author = {Vadym Stupakov},
title = {TorchInstruments: Passive PyTorch Model Telemetry},
year = {2026},
version = {0.5.0},
url = {https://github.com/Red-Eyed/torchinstruments}
}
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distributions
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file torchinstruments-0.5.0-py3-none-any.whl.
File metadata
- Download URL: torchinstruments-0.5.0-py3-none-any.whl
- Upload date:
- Size: 52.9 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via:
uv/0.12.1 {"installer":{"name":"uv","version":"0.12.1","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
448d9a484c778862e807d4587d9627bc6ce2674ce3c63a52a82db649cc73637c
|
|
| MD5 |
8493ed085f95e96407246e72292b2493
|
|
| BLAKE2b-256 |
aefe7dc4a771278b588a3d6cdfbb83408c60a48da7d54ba1bebc503bb0b3a32d
|