Capture speed, by kind of use¶

What each way of using TorchLens costs against a plain forward pass of the same model, measured on 2026-11-03. TorchLens records every operation the model runs, so most kinds of capture cost far more than a forward hook; this page says how much, for each kind, including where TorchLens is slow. Every number is a median wall time from the harness in benchmarks/capture/, and every one comes from a result file in data/ (one file per run; see Hosts, versions and method).

The machines were shared and loaded, so absolute times are slower than on an idle machine. Compare rows within a workload, not across workloads or hosts.

Summary¶

Each cell gives the ratio to a plain forward of the same workload on the same host, then the absolute median in brackets. For the LSTM tagger and ResNet-18 the reference is the harness's plain-forward arm (nothing attached). The GPT-2 runs had no plain-forward arm, so their reference is the forward-hook arm of the same run, which on the other models timed within a quarter of a plain forward. "Not measured" means no run covered that cell; Not measured yet lists them.

The workloads behind the cells. LSTM tagger and ResNet-18: one input; a sweep is eight steering strengths; tl.extract_dataset reads one module over 16 stimuli in batches of 4; the gradient is taken at one module output. GPT-2 small: the recorder and trace rows cache every block's output for 256 sequences of 32 tokens in batches of 16 (16 forwards); a sweep is 12 steered forwards (3 layers by 4 strengths) on a batch of 16 sequences of 32 tokens. The ResNet-18 GPU column ran on a workstation with an RTX 2080 Ti; its tl.extract_dataset cell ran on a second RTX 2080 Ti machine after a fix (see tl.extract_dataset).

Two patterns run through the table:

• TorchLens adds a roughly fixed amount of time per recorded operation, so the ratio depends on how fast the plain forward is. A full trace of ResNet-18 adds about 696 ms on the CPU worker (710.60 ms against 15.01 ms) and about 425 ms on the GPU workstation (427.04 ms against 2.30 ms); the GPU forward is about 6.5 times faster, so the GPU ratio is about four times larger. GPT-2 small on one loaded CPU thread takes about 1.8 s per forward, so the trace's added cost is a small multiple there and about 300 times the forward on a GPU.
• Editing without recording is cheap; recording is not. The bound live edit stays within about two times a plain forward on every model measured. Every kind that records lands between 11 and 241 times a plain forward on the LSTM tagger and ResNet-18, and the replay check costs more.

A plain forward after a trace¶

What it measures: a forward pass with nothing recorded, run after a trace in the same process. import torchlens leaves torch untouched; the first capture installs TorchLens's torch wrappers, and they stay installed until torchlens.backends.torch.wrappers.unwrap_torch() (or a capture with CaptureOptions(unwrap_when_done=True)) removes them. Every forward your process runs afterwards, including code that never calls TorchLens, goes through them. This row says what that costs.

Numbers: 1.02x on the LSTM tagger on CPU (1.50 ms against 1.47 ms), 1.04x on ResNet-18 on GPU (2.39 ms against 2.30 ms) and 1.03x on the LSTM tagger on GPU (0.36 ms against 0.35 ms). On ResNet-18 on the loaded CPU worker the median was 1.47x (22.11 ms against 15.01 ms), but the lowest process medians were 15.41 ms against 14.32 ms, and two identical plain-forward arms on that worker differed by 54%, so this cell is mostly the worker's noise.

Known costs: the one-time cost of the first capture in a process (installing the wrappers) is not in any row here; every arm ran 3 warmup calls first.

Live edits at hook speed¶

What it does: tl.when(site, edit).bind(model) applies an edit inside the model's own forward and records nothing. Each timed call builds the bound edit and runs it, the same way the hook arm registers its hook, runs, and removes it.

When to use it: steering or patching over many forwards, or through generation, when you do not need a record of each forward.

Numbers: within about two times a plain forward on every model measured (see the summary). The first run's rows (LSTM tagger on CPU, and GPT-2 small on CPU):

On GPT-2 small the bound edit and the forward hook land 0.46 s apart on a 27 s sweep, inside the spread of three repeats (hooks 25.36 to 28.00 s, bound edit 26.80 to 27.74 s), and the steered logits are identical (largest absolute difference 0.0). On the LSTM tagger, whose plain forward takes about 1.5 ms, the bound edit's fixed setup cost per call is visible. A plain forward after a trace, with TorchLens's wrappers still installed and nothing recorded, runs at plain-forward speed.

Known costs: a bound edit records nothing, so it cannot give you activations or gradients; pair it with a hook or a capture for those.

The fast recorder, tl.record¶

What it does: tl.record(model, x, save=...) runs the forward through TorchLens's wrapped operations and keeps events and payloads for the selected sites, with no graph postprocessing. It returns a Recording; call Recording.to_trace() later if you need the graph.

When to use it: pulling activations from many forwards in a loop, when you know the sites.

Numbers: 11x a plain forward on ResNet-18 on CPU (150.06 ms), 43x on GPU (98.55 ms), 50x on the LSTM tagger on GPU (18.16 ms against 0.36 ms), and 50x on GPT-2 small on GPU (8.35 s against 168 ms for forward hooks). It costs a fraction of a trace of the same sites: 3.4 times less on ResNet-18 on CPU (150.06 ms against 505.21 ms), 4.0 times less on GPU (98.55 ms against 389.71 ms) and about 6 times less on GPT-2 small on GPU (8.35 s against 49.67 s).

Known costs: it still runs every operation through the wrappers, so it stays well above a forward hook. Backward capture needs a trace; these runs did not time tl.record with gradients.

Traces: selected sites or every activation¶

What it does: tl.trace(model, x) records every operation's metadata and graph edges, checks them, and saves every activation. tl.trace(model, x, save=...) records the same graph and keeps only the selected payloads. tl.extract(model, x, {...}) reads named module outputs from one such call.

When to use it: you need operations no module returns, you do not yet know the site, you want the graph, or you want the recording checked against the forward.

Numbers: 37x to 58x a plain forward on the LSTM tagger and ResNet-18 on CPU, 164x to 186x on ResNet-18 on GPU, and 296x on GPT-2 small on GPU; on GPT-2 small on one loaded CPU thread, where each forward takes about 1.8 s, 3.4x (see the summary).

Saving only selected sites keeps fewer tensors but takes about the same time, because the structure of every operation is still recorded: on ResNet-18 on GPU a full trace took 427.04 ms and a trace keeping one module 381.29 ms; on the LSTM tagger on CPU 81.38 ms and 83.49 ms.

Known costs: the trace's fixed cost dominates on small models (the LSTM tagger's 1.5 ms forward becomes about 80 ms).

Reruns of a finished trace¶

Measured on 2026-11-05, on a shared 32-core Intel Xeon 2.80 GHz Linux CPU worker, one torch thread per process, load average 14 to 19. Python 3.10, torch 2.13.0 (CPU build), transformers 5.18.0, and a development build of TorchLens (2.38.1) from that day. Raw data: data/capture_2026-11-05-rerun.json. Each LSTM tagger and ResNet-18 row is three processes (GPT-2 two); each process times up to 15 samples after 3 warmup calls (GPT-2: up to 5 after 1) and reports its median, and the table gives the median of the process medians with the lowest and highest in brackets. The GPT-2 small rows use the GPT-2 small architecture with random weights (no download). The timing script is not yet part of benchmarks/. The forward-hook rows ran in a second job on the same worker; each one names that job's own plain forward.

What it does: trace.run(model, new_x) reruns the model on new input and returns a new trace with the same saved sites. On a trace that saved a selection of sites (save=tl.module(...) or tl.func(...)) it takes the guarded fast rerun, the model's native forward with the staged hooks, if any, writing the saved sites into a new trace, with or without a staged edit; a trace that saved every activation recaptures, and when a guard refuses (a changed call structure, a value replacement, a non-module target) it falls back to recapturing the whole forward, with the reason in last_run["fast_refused"]. The tables below were measured on the earlier build, in which an unedited trace.run recaptured and then copied each operation's state into the new trace, costing 1.69 times a fresh tl.trace on ResNet-18, 1.57 times on GPT-2 small and 1.53 times on the LSTM tagger; the fast engine's timings are in the next subsection. trace.run_transaction(inputs=new_x, fast=True) takes the same engine whether or not an edit is staged, as long as the capture saved a selection, and raises instead of falling back.

When to use it: to read the same sites on many new inputs, or to rerun a staged edit, run_transaction(inputs=new_x, fast=True) costs 1.3 to 3 times a plain forward on these models.

Correctness: the fast rerun with a staged edit matched a plain forward with the same edit applied by a forward hook (largest absolute difference 0 on GPT-2 small and ResNet-18, 1.9e-08 on the LSTM tagger), and the unedited fast rerun's saved site matched a fresh tl.trace of the new input exactly on all three models.

Saving more sites on the rerun, trace.run_transaction(inputs=new_x, save=[...module addresses]), took about as long as the same saves, measured on a second worker (two processes each; plain forward 10.61 and 10.73 ms on ResNet-18): ResNet-18 637.45 and 731.43 ms against 692.01 and 719.65 ms; LSTM tagger 58.84 and 58.65 ms against 60.71 and 60.09 ms. trace.run takes no save=.

Known costs: on GPT-2 small the plain-forward arm timed slower than the forward-hook and fast arms in both jobs, and the cause was not found; compare GPT-2 rows with the fresh trace rather than the plain forward. A save= selector such as tl.module(...) on run_transaction is refused in this build (module addresses work). The first capture in a process took 1.7 to 2.4 s on ResNet-18, 1.3 to 2.7 s on the LSTM tagger and 4.9 to 5.5 s on GPT-2 small (one call each, installing the torch wrappers).

trace.run on new inputs with the fast engine (development build)¶

A development build of TorchLens leading to 3.0 sends trace.run(model, new_x) on an unedited trace that saved a selection through the guarded fast engine instead of a full recapture. Both builds were timed in two jobs that ran at the same time on the same shared 32-core Intel Xeon 2.80 GHz Linux CPU worker (one torch thread per process, load average 6 to 15), with the same models, inputs, repeats and saved site as the tables above. The plain forward timed differently in the two jobs, so compare the "times plain forward" columns, not the milliseconds across columns.

Correctness: the run reported engine rerun before and guarded_fast after on every call, with no fallback. Before and after, the saved site of the new trace equalled a fresh tl.trace of the new input exactly (largest absolute difference 0) on all three models, and the new trace's output matched a plain forward (0 on ResNet-18 and GPT-2 small, 3.0e-08 on the LSTM tagger).

After the change, trace.run costs less than a fresh tl.trace on every model (0.11 times on ResNet-18, 0.51 on GPT-2 small, 0.16 on the LSTM tagger) but still more than run_transaction(inputs=new_x, fast=True) on the two small models (about 28 ms more on ResNet-18 and 7 ms more on the LSTM tagger per call).

Stopping a rerun early¶

trace.run_transaction(inputs=new_x, fast=True, until=<saved module>) stops the fast rerun once that module has run. The stopped module's output equalled a fresh tl.trace of the new input exactly on all three models. Without fast=True the same until= still costs about a full trace. Same worker and build as above; each row is the median of up to 15 samples (GPT-2: 5) in one process.

Replaying a saved trace without the model¶

A trace captured with CaptureOptions(replayable=True) and saved with trace.save(path, level="runnable", include_weights=True) can be loaded with tl.load(path) and run on new input with loaded.run_transaction(inputs=new_x), with no model object. Without include_weights=True the file holds no weights and the replay runs on a fresh random initialization (TorchLens warns runnable_random_init_run). With the weights, the replay's output equalled model(new_x) (largest absolute difference 0 on ResNet-18 and GPT-2 small, 3e-08 on the LSTM tagger), and every run reported path faithfulness verified with the state taken from the file. Measured on 2026-11-05 on a shared 32-core Intel Xeon 2.80 GHz Linux CPU worker, one torch thread per process, load average 14 to 21, same build as above, in one process per model (median of up to 15 samples, GPT-2 up to 5). The worker was busier than for the rows above, so compare these rows with each other only.

One-time costs per model: ResNet-18 save with weights 5.08 s (a 46.3 MiB file) and load 3.31 s; GPT-2 small 26.85 s (627.8 MiB) and 14.42 s; LSTM tagger 1.15 s (1.7 MiB) and 498 ms. The replay costs about what a fresh trace does on ResNet-18 and the LSTM tagger (0.84 and 0.88 times) and about twice a fresh trace on GPT-2 small; the replayable capture itself is 4 to 5 times a default trace. GPT-2 small saves only with use_cache=False: with the default key-value cache the save is refused (output_container_contract_missing).

Edit sweeps with a recording per edit¶

What it does: one tl.record(..., intervene=tl.when(...)) or one tl.trace(..., intervene=tl.when(...)) per edit, each returning the edited output.

When to use it: you need a record of each edited forward (activations, a graph, a check). If you only need the edited outputs, the bound live edit is 9 to 107 times faster on the LSTM tagger and ResNet-18. On GPT-2 small on CPU, where each forward takes seconds, the gap is under two times.

Numbers: on ResNet-18 on CPU, eight strengths took 195.36 ms bound, 1.79 s with a recording per strength and 5.26 s with a trace per strength; on GPU 33.63 ms, 1.06 s and 3.61 s.

Features over a stimulus set, tl.extract_dataset¶

What it does: runs a trace per batch and gathers the requested module's output across batches, in stimulus order.

Numbers: 37x a plain loop over the batches on the LSTM tagger on CPU (356.00 ms against 9.69 ms), 21x on ResNet-18 on CPU (2.69 s against 129.86 ms) and 241x on ResNet-18 on GPU (1.99 s against 8.25 ms). On the same GPU machine one tl.record per batch took 428.00 ms, about a fifth.

Known costs: the capture cost is paid per batch, so large batches spread it. On GPU the first run failed with inputs on two devices; the cause was fixed the same day and the GPU cell above is the fixed run, on a second RTX 2080 Ti machine.

Gradients¶

What it does: tl.trace(..., capture=CaptureOptions(save_grads=...)) then trace.backward(...) records the forward graph and the gradient at the selected module.

Numbers: 34x a plain forward and backward on the LSTM tagger on CPU (117.50 ms against 3.49 ms), 12x on ResNet-18 on CPU (980.27 ms against 80.91 ms) and 62x on GPU (551.37 ms against 8.88 ms). A forward hook plus a tensor hook took 70.31 ms on ResNet-18 on CPU, and Captum's layer gradient 39.27 ms.

Known costs: the LSTM tagger's gradient cells could not run on GPU for any tool (cuDNN's RNN backward needs training mode, and the harness runs in eval mode).

Operations on a finished trace¶

The same runs timed what you do with a trace after capture. Medians of three processes, from the same result files:

The one-call tl.summary(model, x) is measured separately, with its own tool, on Summary performance.

Recording a generation¶

Recording a generation keeps a record of every decode step. Measured on 2026-11-03 on the same 32-core Intel Xeon worker, four torch threads, one process per token count, on a development build of TorchLens from that day, with a 70-million-parameter GPT-NeoX-style model (EleutherAI/pythia-70m, float32), greedy decoding from a 7-token prompt, and the output of every block's MLP and of the final layer norm saved at every step. Each row is one call of tl.record(model, ids, save=sites, generate_kwargs={"max_new_tokens": n, "min_new_tokens": n, "do_sample": False, "use_cache": True}), where sites joins tl.module(...) for each saved module with |. One process per row, no repeats. Peak process memory is the operating system's maximum resident set size; saved activations are the summed saved bytes of every step's record. No run timed a plain model.generate(...) for comparison. Raw data: data/capture_2026-11-03-generation.json and data/capture_2026-11-03-gpu2.json.

Current behaviour: process memory grows linearly, about 2.7 MiB per token, and the time per token stays flat (the 500- and 1,000-token rows from the fix's own check, the 2,000-token row from the third run):

The same 2,000-token run on the RTX 2080 Ti node (four torch threads) took 269 ms per token, peaked at 6.8 GiB of process memory and 458 MiB of GPU memory, and saved the same 27.4 MiB of activations.

Before a fix on 2026-11-03 the recorder kept every decode pass's whole output, including the model's key-value cache, so retained memory grew with the square of the generated tokens and the time per token grew with the run:

The saved activations were a small part of that growth: 6.9 MiB of saved tensors at 500 tokens against 5.3 GiB of process memory; 13.8 MiB of saved tensors at 1,000 tokens against 15.7 GiB.

For long generations, record a shorter run or save fewer sites.

Where TorchLens is slower¶

A capture records every operation's metadata and checks it, so it costs much more than a hook that reads one tensor. On a small model the capture's fixed cost dominates. The first run's rows (LSTM tagger on CPU):

What this means in practice:

• To read one module's output, a forward hook is the faster tool. A capture is worth its cost when you need operations no module returns, do not yet know the site, or want the recording checked against the forward.
• To sweep many edits, use the bound edit (tl.when(...).bind(model)), not a trace per edit.
• To pull selected activations in a loop, use tl.record, not tl.trace.
• Batch feature extraction pays the capture cost per batch. Large batches spread it.
• Saving only selected sites (save=) keeps fewer tensors, not less time: the structure of every operation is still recorded, which is why the one-module trace costs about what the full trace does.
• On GPU the ratios are larger, because the plain forward is faster and TorchLens's added time is not.
• torchvision's fx feature extraction is faster than a plain forward of ResNet-18 for one module, because it builds a graph that stops at that module.
• To time a model itself, locate the slow part with TorchLens, then measure it clean with torch.utils.benchmark.Timer.

Not measured yet¶

• Stopping a fresh capture early with stop_after= (stopping a fast rerun early is in Stopping a rerun early).
• A fast rerun of a trace on new inputs on the 3 Nov workloads (the 5 Nov reruns are in Reruns of a finished trace).
• A trace of metadata only, saving no activations.
• tl.record on the LSTM tagger on CPU, and every kind except the recorder, the trace and the sweeps on GPT-2 small (the GPT-2 runs used their own workloads and had no plain-forward arm).
• The one-time cost of the first capture in a process.
• A plain model.generate(...) as the reference for recording a generation, and generation past 2,000 tokens.
• fx on the LSTM tagger: its data-dependent control flow cannot be traced by fx.
• Gradient arms on the LSTM tagger on GPU: cuDNN's RNN backward needs training mode.

Hosts, versions and method¶

Each run's result file names its host and versions; the page's tables quote the files' medians.

• First run (LSTM tagger on CPU; GPT-2 small live edits on CPU): a shared 32-core Intel Xeon 2.80 GHz Linux CPU worker, one torch thread per process, load average about 40. Python 3.10, torch 2.13.0 (CPU build), transformers 5.18.0, TorchLens 2.36.3, nnsight 0.8.1. Raw data: data/capture_2026-11-03.json.
• Second run, CPU rows (ResNet-18 and GPT-2 small): a shared 32-core Intel Xeon 2.80 GHz Linux CPU worker, one torch thread per process, load average 7 to 28. Versions as in the first run, with torchvision 0.28.0 and Captum 0.9.0. Raw data: data/capture_2026-11-03-run2.json.
• Second run, GPU rows (ResNet-18): a 20-core Intel Core i9-9900X 3.50 GHz Linux workstation with an NVIDIA GeForce RTX 2080 Ti, one torch thread per process, load average about 10. Versions as above except torch 2.13.0 and torchvision 0.28.0 built for CUDA 12.6. Raw data: data/capture_2026-11-03-gpu.json.
• Third run, GPT-2 small and LSTM tagger on GPU: a 256-core AMD EPYC 7662 Linux GPU node with an NVIDIA L40 (46 GiB), one torch thread per process, load average 56 to 74 (a shared node). Raw data: data/capture_2026-11-03-gpu3.json.
• Third run, tl.extract_dataset on GPU and the GPU generation run: a 48-core Intel Xeon Gold 6126 2.60 GHz Linux GPU node with an NVIDIA GeForce RTX 2080 Ti (11 GiB), one torch thread per process, load average 1 to 3. Raw data: data/capture_2026-11-03-gpu2.json.
• Third run versions on both GPU nodes: Python 3.10, torch 2.13.0 built for CUDA 12.6, transformers 5.18.0, nnsight 0.8.1, TransformerLens 4.3.0, and a development build of TorchLens from that day.
• Reruns of a finished trace (2026-11-05): see their section.
• Recording a generation: see its section.

Method. Each LSTM tagger and ResNet-18 row is three separate processes (two for some first-run rows, as each row says); each process times up to 20 samples after 3 warmup calls (as few as 3 for the slowest TorchLens cells, which reach the 20 s budget first) and reports its median. The tables give the median of the process medians, and in brackets the lowest and highest. Each GPT-2 row is one process that runs the whole workload three times; the tables give the median of the three wall times and, in brackets, the fastest and slowest. The CPU GPT-2 rows come from two runs at different times on the shared worker, and its load changed between them, so each run carries its own forward-hook row: compare a tool with the hook row of the same run. Every GPU arm synchronises the GPU before it stops its clock, and each GPU job had one GPU to itself.

Each tool's output was compared with the forward-hook arm's output and matched: bit for bit for every TorchLens and nnsight arm on ResNet-18, and with a largest absolute difference of 0.0 on GPT-2. The CPU worker was loaded while the second run's rows ran: two identical plain-forward arms timed on it differed by 54%, against 5% on the GPU workstation. TransformerLens runs its own GPT-2 implementation loaded with the same weights, so its rows time a different model code path with hooks on hook_resid_post; every other GPT-2 row runs the transformers model.

All measurements¶

Every row the summary draws on, by run, with the peer tools measured beside TorchLens.

GPT-2 small on CPU (second run)¶

ResNet-18 on CPU (second run)¶

ResNet-18 on GPU (second run, RTX 2080 Ti workstation)¶

The fx rows are faster than the plain forward because create_feature_extractor builds a graph that stops at the requested module, so it runs only part of ResNet-18. The other arms run the whole network. On GPU the tl.extract_dataset arm failed in this run with "Expected all tensors to be on the same device" (cuda:0 and cpu): in-memory results were copied to the host while the model stayed on the GPU. It was fixed the same day; the last table in this section measured the fixed arm.

GPT-2 small on GPU (third run, L40 node)¶

LSTM tagger on GPU (third run, L40 node)¶

The LSTM tagger's gradient arms could not run on GPU for any tool (cuDNN's RNN backward is only available in training mode, and the harness runs the model in eval mode), and fx refuses the model's data-dependent control flow. Of the peer tools, nnsight ran; Captum is a gradient arm; TransformerLens has no LSTM.

ResNet-18 on GPU after the tl.extract_dataset fix (third run, RTX 2080 Ti node)¶

Reproduce¶

pip install -e ".[dev,test]" nnsight==0.8.1 captum==0.9.0
OMP_NUM_THREADS=1 python benchmarks/capture/driver.py --smoke --warmups 3 --samples 20 \
    --budget-s 20 --processes 3 --no-sweep --subjects lstm-tagger --out capture-lstm.json
OMP_NUM_THREADS=1 python -m benchmarks.capture.mi_workloads.matrix --model gpt2 \
    --tools hooks,tl_bind --workloads steer_sweep --out mi-gpt2.jsonl --resdir mi-res
# second run: ResNet-18 (add --device cuda for the GPU rows) and GPT-2 small full forwards
OMP_NUM_THREADS=1 python benchmarks/capture/driver.py --smoke --warmups 3 --samples 20 \
    --budget-s 20 --threads 1 --processes 3 --no-sweep --subjects resnet18 --out capture-resnet18.json
OMP_NUM_THREADS=1 python -m benchmarks.capture.mi_workloads.matrix --model gpt2 --timeout 1500 \
    --tools hooks,nnsight,tl_record,tl_trace --workloads cache,steer_sweep,patch_sweep \
    --out mi-gpt2.jsonl --resdir mi-res -- --slow-rep-s 0
# third run: the same on a GPU (TL_BENCH_DEVICE moves the GPT-2 workloads; the driver takes --device)
pip install transformer-lens==4.3.0
TL_BENCH_DEVICE=cuda OMP_NUM_THREADS=1 python -m benchmarks.capture.mi_workloads.matrix --model gpt2 \
    --timeout 1500 --tools hooks,nnsight,tlens,tl_record,tl_trace --workloads cache,steer_sweep,patch_sweep \
    --out mi-gpu-gpt2.jsonl --resdir mi-res -- --slow-rep-s 0
OMP_NUM_THREADS=1 python benchmarks/capture/driver.py --device cuda --smoke --warmups 3 --samples 20 \
    --budget-s 20 --threads 1 --processes 3 --no-sweep --subjects lstm-tagger,resnet18 --out capture-gpu.json