Caching Activations Over a Dataset¶
Most interpretability work that goes past a single prompt starts the same way: run a corpus through a model, keep some activations, write them to disk, and then do something with them — train a probe, fit a dictionary, compute feature statistics. The modelling is the interesting part. The collection is where runs die: at hour three, at 40,000 of 50,000 sequences, out of memory, or — worse — quietly, with tensors on disk that no longer match the model that produced them.
This tutorial builds that collection step end to end on GPT-2 and a small slice of WikiText. It is sized to run in Colab, and structured so that the only thing that changes at 50,000 sequences is a constant. We will:
- capture many modules per forward pass with
tracer.cache(modules=[...]), and measure what an unfiltered cache costs; - wrap collection in
torch.no_grad()and measure what forgetting costs; - shard to disk with a manifest, kill the run halfway, and resume it;
- anchor the cache — re-run a random sample of rows and check the stored tensors match a fresh forward pass bit-for-bit — and then show the re-check that will never be bit-exact, so you don't chase it;
- watch GPU memory across the whole run, because a leak that is invisible at 100 sequences is fatal at 50,000;
- handle padding, which silently poisons every reduction over a variable-length batch;
- and then actually use the data: a layer sweep that localises duplicate-token detection to attention layer 0, against a control that sits at exactly chance.
A cache that has drifted from the model is worse than no cache: everything downstream is wrong and nothing raises. Roughly half of this tutorial is the checks that keep that from happening.
Setup¶
try:
import google.colab
is_colab = True
except ImportError:
is_colab = False
if is_colab:
!pip install -U nnsight datasets scikit-learn
import gc, json, os, shutil, tempfile, time
from IPython.display import clear_output
import numpy as np
import pandas as pd
import torch
import nnsight
from nnsight import TransformersModel
import plotly.express as px
import plotly.io as pio
pio.renderers.default = "colab" if is_colab else "plotly_mimetype+notebook_connected+colab+notebook"
# The corpus builder below tokenizes paragraphs longer than GPT-2's context; that is fine
# (we repack them ourselves) but it makes transformers warn once per batch.
import transformers
transformers.logging.set_verbosity_error()
import warnings
warnings.filterwarnings("ignore", message=".*IProgress not found.*")
/home/localjadenfk/miniconda3/envs/ndif2/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html from .autonotebook import tqdm as notebook_tqdm
We use GPT-2 small: 12 blocks, 768 hidden dimensions, and small enough that the pipeline — not the model — is what we are measuring. Everything below works unchanged on a larger model; the constants (batch size, shard size, dtype) are what you re-tune.
model = TransformersModel(
"openai-community/gpt2",
device_map="auto",
dispatch=True,
task="text-generation", # explicit task; required when HF is in offline mode
)
clear_output()
DEVICE = next(model.parameters()).device
print(model.config.n_layer, "blocks,", model.config.n_embd, "hidden dims, on", DEVICE)
12 blocks, 768 hidden dims, on cuda:0
CUDA = DEVICE.type == "cuda"
def gpu_mib(kind="alloc"):
# Current or peak allocated GPU memory in MiB (0 on CPU).
if not CUDA:
return 0.0
fn = {"alloc": torch.cuda.memory_allocated,
"peak": torch.cuda.max_memory_allocated,
"reserved": torch.cuda.memory_reserved}[kind]
return fn() / 2**20
def reset_peak():
if CUDA:
torch.cuda.reset_peak_memory_stats()
def sync():
if CUDA:
torch.cuda.synchronize()
print(f"model resident: {gpu_mib():.1f} MiB")
model resident: 475.5 MiB
The corpus¶
Two things make a corpus job easy that are easy to skip.
Pack, don't pad. We concatenate the tokenized text into one long stream and cut it into
fixed-length blocks. Every row then has exactly SEQ_LEN real tokens, so there are no pad
positions to mask, no attention-mask bookkeeping, and no ragged batches. This is what SAE
and probing pipelines normally do. When you can't pack — because rows are semantic units,
like chat turns — you get padding, and the section Padding
below shows exactly how badly that bites.
Tokenize once, up front. Tokenization is CPU work and re-doing it inside the capture
loop puts it on the critical path. GPT-2's vocabulary is 50,257, so the whole corpus fits
in uint16 — 128 KB per 1,000 sequences of 64 tokens.
We take 2,048 blocks of 64 tokens = 131,072 tokens. That is small on purpose; it takes seconds and the numbers extrapolate linearly, which we check at the end.
from datasets import load_dataset
SEQ_LEN = 64
N_SEQ = 2048
ds = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="train")
need = N_SEQ * SEQ_LEN
buf = np.empty(need + 4096, dtype=np.uint16)
n, batch = 0, []
for row in ds:
if len(row["text"]) < 40: # skip blank lines and section headers
continue
batch.append(row["text"])
if len(batch) == 256:
for seq in model.tokenizer(batch)["input_ids"]:
k = min(len(seq), len(buf) - n)
buf[n:n + k] = np.asarray(seq[:k], dtype=np.uint16)
n += k
batch = []
if n >= need:
break
corpus = buf[:need].reshape(N_SEQ, SEQ_LEN)
clear_output()
print(f"corpus: {corpus.shape} {corpus.dtype} = {corpus.nbytes/1e3:.0f} kB on disk, "
f"{corpus.size:,} tokens")
print(repr(model.tokenizer.decode(corpus[0])[:180]) + " ...")
corpus: (2048, 64) uint16 = 262 kB on disk, 131,072 tokens ' Senjō no Valkyria 3 : Unrecorded Chronicles ( Japanese : 戦場のヴァルキュリア3 , lit . Valkyria of the Battlefield 3 ) , commonly referred to as Valkyria Chronicles III outside Japan , is a' ...
Row 0 starts mid-sentence, which is the point of packing: blocks are windows into a token stream, not documents.
One batch, one cache¶
tracer.cache(modules=[...]) records the output of every listed module for the whole run,
in one forward pass. It sits on the same path the interleaver already uses to deliver
values, so there are no per-module hooks to register and unregister, and it observes values
after any intervention you apply.
Two things to know before the first call:
- Cache keys are module paths prefixed with
model.— the Envoy you pass in asmodel.transformer.h[6]comes back as"model.transformer.h.6". Printcache.keys()once rather than guessing. - A transformers-5 decoder block's
.outputis a bare tensor, butattn.outputis a tuple. Both appear in the same cache. Unwrap defensively.
# The modules we will collect for the rest of this notebook.
# transformer.drop -- dropout after wte + wpe, i.e. the embeddings (identity in eval)
# h[0].attn, h[0].mlp -- the two sublayers of block 0, to see where a feature appears
# h[0..11] -- the residual stream after each block
MODULES = ([model.transformer.drop, model.transformer.h[0].attn, model.transformer.h[0].mlp]
+ [model.transformer.h[i] for i in range(model.config.n_layer)])
NAMES = (["model.transformer.drop", "model.transformer.h.0.attn", "model.transformer.h.0.mlp"]
+ [f"model.transformer.h.{i}" for i in range(model.config.n_layer)])
LABELS = ["embeddings", "attn.0", "mlp.0"] + [f"block {i}" for i in range(model.config.n_layer)]
BATCH = 32
ids = torch.from_numpy(corpus[:BATCH].astype(np.int64))
with torch.no_grad():
with model.trace({"input_ids": ids}) as tracer:
cache = tracer.cache(modules=MODULES, dtype=torch.float16)
print(list(cache.keys())[:5], "...")
print()
for name in ["model.transformer.drop", "model.transformer.h.0.attn", "model.transformer.h.6"]:
out = cache[name].output
kind = type(out).__name__
shape = out[0].shape if isinstance(out, tuple) else out.shape
print(f"{name:30s} {kind:7s} {tuple(shape)} {(out[0] if isinstance(out, tuple) else out).dtype}"
f" on {(out[0] if isinstance(out, tuple) else out).device}")
['model.transformer.drop', 'model.transformer.h.0.attn', 'model.transformer.h.0.mlp', 'model.transformer.h.0', 'model.transformer.h.1'] ... model.transformer.drop Tensor (32, 64, 768) torch.float16 on cpu model.transformer.h.0.attn tuple (32, 64, 768) torch.float16 on cpu model.transformer.h.6 Tensor (32, 64, 768) torch.float16 on cpu
h.0.attn came back as a tuple and the blocks as bare tensors — in the same cache, from
the same trace. Indexing a block's output with [0] (the old nnsight idiom) would silently
hand you row 0 of the batch instead of the hidden states. One helper removes the hazard:
def unwrap(value):
# Blocks return a tensor; attention returns a tuple. Take the hidden states either way.
return value[0] if isinstance(value, tuple) else value
print(unwrap(cache["model.transformer.h.0.attn"].output).shape)
torch.Size([32, 64, 768])
Note also that dtype=torch.float16 and the default device=torch.device("cpu") are
applied on the way out of the trace. Halving on the way out halves both the PCIe traffic
and the bytes on disk, and — we check this later — it is exactly equivalent to capturing in
fp32 and calling .half() afterwards.
Two decisions that set the cost of the whole run¶
Before writing a pipeline, two flags decide whether it is viable. Both are easy to get wrong because both are silent when wrong: the run still finishes, it just costs several times more than it should.
Decision 1: always pass modules=.
tracer.cache() with no arguments caches every module in the model. For GPT-2 that is
151 modules — every Conv1D, every LayerNorm, every dropout, and lm_head, whose output
is [batch, seq, 50257]. All of it is copied device-to-host on every call. Let's measure
both, at the batch size the pipeline will actually use.
def cache_bytes(cache):
total = 0
for key in cache.keys():
value = cache[key].output
for tensor in (value if isinstance(value, (tuple, list)) else [value]):
if isinstance(tensor, torch.Tensor):
total += tensor.numel() * tensor.element_size()
return total
def time_cache(modules):
times = []
for _ in range(3):
sync(); t0 = time.time()
with torch.no_grad():
with model.trace({"input_ids": ids}) as tracer:
c = (tracer.cache(dtype=torch.float16) if modules is None
else tracer.cache(modules=modules, dtype=torch.float16))
sync(); times.append(time.time() - t0)
return min(times), len(list(c.keys())), cache_bytes(c)
for label, mods in [("cache() ", None), ("cache(modules=[...])", MODULES)]:
dt, n_mod, nbytes = time_cache(mods)
print(f"{label} {dt*1000:7.1f} ms/batch {n_mod:3d} modules "
f"{nbytes/2**20:7.1f} MiB copied to host")
_ = gc.collect()
cache() 775.7 ms/batch 151 modules 925.4 MiB copied to host cache(modules=[...]) 57.2 ms/batch 15 modules 45.0 MiB copied to host
On the A6000 this notebook ran on: 776 ms and 925 MiB per batch unfiltered, 57 ms and
45 MiB filtered — about 14x the time and 20x the bytes, for 136 modules we will never read.
Over the 64 batches of the pipeline below that is 50 seconds of pure copying instead of 4. The research run behind this tutorial measured ~3 GB of blocking device-to-host copy per
call at batch 128; adding modules= took one collection loop from 5 minutes to 3 seconds.
Cache the modules you will read, and nothing else.
Decision 2: wrap collection in torch.no_grad().
A trace is a normal forward pass, so it runs with autograd on unless you say otherwise.
Every intermediate the backward pass would need is kept alive, and any activation you hold
comes back carrying a live grad_fn that pins the graph behind it. .detach() after the
fact does not help: the graph was already built during the forward.
with model.trace({"input_ids": ids}) as tracer:
h = nnsight.save(model.transformer.h[6].output)
print(f"no torch.no_grad(): grad_fn={type(h.grad_fn).__name__} requires_grad={h.requires_grad} device={h.device}")
with torch.no_grad():
with model.trace({"input_ids": ids}) as tracer:
h = nnsight.save(model.transformer.h[6].output)
print(f"with torch.no_grad(): grad_fn={h.grad_fn} requires_grad={h.requires_grad}")
del h
_ = gc.collect()
no torch.no_grad(): grad_fn=AddBackward0 requires_grad=True device=cuda:0 with torch.no_grad(): grad_fn=None requires_grad=False
That is the mechanism. Here is the price, measured on the same work — 256 sequences, 15 modules, everything moved to CPU and freed each batch:
import contextlib
def capture_pass(n_seq, grad):
if CUDA:
torch.cuda.empty_cache()
reset_peak()
ctx = contextlib.nullcontext() if grad else torch.no_grad()
sync(); t0 = time.time()
with ctx:
for i in range(0, n_seq, BATCH):
batch_ids = torch.from_numpy(corpus[i:i + BATCH].astype(np.int64))
with model.trace({"input_ids": batch_ids}) as tracer:
c = tracer.cache(modules=MODULES, dtype=torch.float16)
stacked = torch.stack([unwrap(c[name].output) for name in NAMES])
del c, stacked
sync()
return gpu_mib("peak"), n_seq * SEQ_LEN / (time.time() - t0)
peak_ng, tps_ng = capture_pass(256, grad=False)
peak_g, tps_g = capture_pass(256, grad=True)
print(f"torch.no_grad(): peak {peak_ng:7.1f} MiB {tps_ng:8.0f} tok/s")
print(f"autograd on : peak {peak_g:7.1f} MiB {tps_g:8.0f} tok/s")
print(f" -> {peak_g/max(peak_ng,1e-9):.2f}x peak memory, "
f"{tps_ng/tps_g:.2f}x throughput")
torch.no_grad(): peak 1026.3 MiB 26464 tok/s
autograd on : peak 3050.6 MiB 35450 tok/s
-> 2.97x peak memory, 0.75x throughput
~3x peak GPU memory — 1.0 GiB with no_grad, 3.0 GiB without.
Trust the memory ratio and not the throughput ratio. Peak allocation is deterministic and reproduced to the tenth of a MiB across runs; the timing is noise on a shared GPU, and across two runs of this notebook it came out 1.4x slower with autograd once and marginally faster the next time. The reason the time penalty is small at all is that we move each batch to the host and free it, so the graph is discarded every batch instead of accumulating — a pipeline that holds saved GPU tensors pays far more (earlier measurements on that shape put it near 18x memory and 11x time).
Either way, 3.0 GiB versus 1.0 GiB is the difference between fitting your batch size on the
GPU and not, and nothing in the run hints that you are paying for gradients you will never
use. Note that .detach() on the captured tensor does not help — tracer.cache(detach=True)
is already the default, and the memory above was measured with it. The graph is built
during the forward pass; only torch.no_grad() stops that.
tracer.cache() versus per-module .save()¶
You can write the same capture by hand: read each module's .output inside the trace, move
it, and save the list. Is one faster? Let's interleave A/B/A/B so drift on a shared GPU
doesn't decide the answer.
BLOCKS = [model.transformer.h[i] for i in range(model.config.n_layer)]
BLOCK_NAMES = [f"model.transformer.h.{i}" for i in range(model.config.n_layer)]
N_BENCH = 512
def run_cache():
reset_peak(); sync(); t0 = time.time()
with torch.no_grad():
for i in range(0, N_BENCH, BATCH):
b = torch.from_numpy(corpus[i:i + BATCH].astype(np.int64))
with model.trace({"input_ids": b}) as tracer:
c = tracer.cache(modules=BLOCKS, dtype=torch.float16)
stacked = torch.stack([c[name].output for name in BLOCK_NAMES]); del c, stacked
sync(); return N_BENCH * SEQ_LEN / (time.time() - t0), gpu_mib("peak")
def run_save():
reset_peak(); sync(); t0 = time.time()
with torch.no_grad():
for i in range(0, N_BENCH, BATCH):
b = torch.from_numpy(corpus[i:i + BATCH].astype(np.int64))
with model.trace({"input_ids": b}) as tracer:
hs = nnsight.save([blk.output.to("cpu", torch.float16) for blk in BLOCKS])
stacked = torch.stack(hs); del hs, stacked
sync(); return N_BENCH * SEQ_LEN / (time.time() - t0), gpu_mib("peak")
run_cache(); run_save() # warm up, then interleave
for _ in range(3):
(tc, pc), (ts, ps) = run_cache(), run_save()
print(f"cache {tc:8.0f} tok/s peak {pc:7.1f} MiB | "
f"save {ts:8.0f} tok/s peak {ps:7.1f} MiB")
cache 37676 tok/s peak 1026.3 MiB | save 37564 tok/s peak 1026.3 MiB
cache 37194 tok/s peak 1026.3 MiB | save 37166 tok/s peak 1026.3 MiB
cache 37394 tok/s peak 1026.3 MiB | save 37448 tok/s peak 1026.3 MiB
# ...and they produce the same bytes.
with torch.no_grad():
with model.trace({"input_ids": ids}) as tracer:
c = tracer.cache(modules=[model.transformer.h[6]], dtype=torch.float16)
with torch.no_grad():
with model.trace({"input_ids": ids}) as tracer:
saved = nnsight.save([model.transformer.h[6].output.to("cpu", torch.float16)])
print("cache == save, bit-for-bit:", torch.equal(c["model.transformer.h.6"].output, saved[0]))
del c, saved
_ = gc.collect()
cache == save, bit-for-bit: True
Throughput and peak memory are the same to within run-to-run noise, and the tensors are identical. The choice is about ergonomics, not cost:
.save()keeps the move explicit — you write.to("cpu", torch.float16)yourself, so you can see exactly what crosses the bus and when. It also requires you to know which modules return tuples.tracer.cache()wins when you don't want to hard-code module names, when you want inputs as well as outputs (include_inputs=True), and above all in a generation loop, where it appends one entry per decoding step automatically instead of needing a per-step save.
One flag is worth knowing about even though we do not pass it. tracer.cache(...) copies
synchronously by default (non_blocking=False). non_blocking=True enqueues the
device-to-host copy and returns immediately; nothing synchronises before you read the
result, and reading a CPU tensor does not synchronise CUDA, so you can read a tensor whose
copy is still in flight. It fails silently, it fails worse under GPU contention, and it is
how a whole shard set becomes garbage. Leave the default alone unless you synchronise
yourself.
The pipeline¶
Now the real thing. The design is four decisions, and none of them are nnsight's job — which is the point: nothing global survives a trace, so a capture loop is just a loop.
Step 1: shard. Write every SHARD_SEQ sequences to their own file. Shards are the unit
of resumption, the unit of parallel reading later, and they keep peak host memory to one
shard rather than the whole corpus.
Step 2: write atomically. torch.save to shard_0007.pt.tmp, then os.replace it onto
shard_0007.pt. os.replace is atomic within a filesystem, so a file with the final name
is always complete. A kill -9 mid-write leaves a .tmp that the next run ignores.
Step 3: keep a manifest. One JSON line per completed shard — index, filename, byte
count, and which corpus rows it holds — flushed and fsynced. On restart we trust a shard
only if the manifest lists it and the file on disk is exactly the recorded size.
Step 4: log memory per shard. Allocated, peak and reserved GPU memory at every shard boundary. This costs nothing and it is the only way to distinguish "we are fine" from "we are 4% into the leak".
Store the token ids alongside the activations. Labels for probing come from the token stream, and a shard that cannot say which tokens it came from is much less useful later.
The version people write first:
# Works at 100 sequences. Do not run this at 50,000.
all_acts = []
for i in range(0, N_SEQ, BATCH):
batch_ids = torch.from_numpy(corpus[i:i + BATCH].astype(np.int64))
with model.trace({"input_ids": batch_ids}) as tracer:
cache = tracer.cache() # every module
all_acts.append({k: cache[k].output for k in cache.keys()})
torch.save(all_acts, "acts.pt")
Four problems, in increasing order of how much they hurt: it caches every module (Decision 1); it runs with autograd on (Decision 2); it holds the entire corpus in host RAM before a single byte reaches disk; and if anything fails at sequence 49,000 — an OOM, a preempted node, a full disk — there is nothing on disk to resume from.
SHARD_SEQ = 128 # sequences per shard -> 16 shards
OUT_DIR = os.path.join(tempfile.mkdtemp(prefix="nnsight_acts_"), "shards")
print("writing shards to", OUT_DIR)
bytes_per_seq = len(NAMES) * SEQ_LEN * model.config.n_embd * 2 # fp16
print(f"expected: {N_SEQ // SHARD_SEQ} shards, "
f"{bytes_per_seq * N_SEQ / 2**30:.2f} GiB total "
f"({bytes_per_seq / 2**10:.0f} KiB per sequence)")
writing shards to /tmp/nnsight_acts_ec41yaip/shards expected: 16 shards, 2.81 GiB total (1440 KiB per sequence)
def capture_batch(batch_ids):
# One forward pass -> {module name: [batch, seq, hidden] fp16 on the host}.
with torch.no_grad():
with model.trace({"input_ids": batch_ids}) as tracer:
cache = tracer.cache(modules=MODULES, dtype=torch.float16)
return {name: unwrap(cache[name].output) for name in NAMES}
def write_shard(out_dir, index, payload, rows):
# Atomic write + one fsync'd manifest line. Never leaves a half-written shard.
fname = f"shard_{index:04d}.pt"
tmp = os.path.join(out_dir, fname + ".tmp")
torch.save(payload, tmp)
os.replace(tmp, os.path.join(out_dir, fname)) # atomic within a filesystem
nbytes = os.path.getsize(os.path.join(out_dir, fname))
with open(os.path.join(out_dir, "manifest.jsonl"), "a") as fh:
fh.write(json.dumps({"shard": index, "file": fname, "bytes": nbytes,
"rows": [int(rows[0]), int(rows[-1]) + 1]}) + "\n")
fh.flush()
os.fsync(fh.fileno())
def completed_shards(out_dir):
# Shards the manifest claims AND whose file is present at the recorded size.
done = {}
manifest = os.path.join(out_dir, "manifest.jsonl")
if os.path.exists(manifest):
for line in open(manifest):
record = json.loads(line)
path = os.path.join(out_dir, record["file"])
if os.path.exists(path) and os.path.getsize(path) == record["bytes"]:
done[record["shard"]] = record
return done
def collect(out_dir, crash_at=None):
os.makedirs(out_dir, exist_ok=True)
done = completed_shards(out_dir)
print(f"[resume] {len(done)} of {N_SEQ // SHARD_SEQ} shards already complete")
log, tokens, t0 = [], 0, time.time()
for s in range(N_SEQ // SHARD_SEQ):
if s in done:
continue
if crash_at is not None and s == crash_at:
raise RuntimeError(f"simulated crash at shard {s}")
rows = np.arange(s * SHARD_SEQ, (s + 1) * SHARD_SEQ)
parts = {name: [] for name in NAMES}
for b in range(0, SHARD_SEQ, BATCH):
batch_ids = torch.from_numpy(corpus[rows[b:b + BATCH]].astype(np.int64))
captured = capture_batch(batch_ids)
for name in NAMES:
parts[name].append(captured[name])
del captured # free the batch before the next one
payload = {"acts": {name: torch.cat(parts[name]) for name in NAMES},
"tokens": torch.from_numpy(corpus[rows].astype(np.int32)),
"rows": torch.from_numpy(rows)}
del parts
write_shard(out_dir, s, payload, rows)
del payload
tokens += SHARD_SEQ * SEQ_LEN
log.append({"shard": s, "allocated": gpu_mib("alloc"), "peak": gpu_mib("peak"),
"reserved": gpu_mib("reserved"), "elapsed": time.time() - t0,
"tokens": tokens})
return log, tokens / max(time.time() - t0, 1e-9)
Now we break it on purpose. A pipeline that has never been interrupted has never had
its resumption tested. We run with crash_at=8, which raises in the middle of the run, and
then look at what is on disk.
reset_peak()
try:
collect(OUT_DIR, crash_at=8)
except RuntimeError as e:
print("run died:", e)
print()
print("on disk:", sorted(os.listdir(OUT_DIR))[:4], "...", len(os.listdir(OUT_DIR)), "entries")
print()
for line in list(open(os.path.join(OUT_DIR, "manifest.jsonl")))[:3]:
print(line.strip())
[resume] 0 of 16 shards already complete
run died: simulated crash at shard 8
on disk: ['manifest.jsonl', 'shard_0000.pt', 'shard_0001.pt', 'shard_0002.pt'] ... 9 entries
{"shard": 0, "file": "shard_0000.pt", "bytes": 188783738, "rows": [0, 128]}
{"shard": 1, "file": "shard_0001.pt", "bytes": 188783738, "rows": [128, 256]}
{"shard": 2, "file": "shard_0002.pt", "bytes": 188783738, "rows": [256, 384]}
Eight shards and a manifest. To make the restart test honest, we also plant the wreckage a
real kill -9 would leave: a half-written temporary file, and a manifest line for a shard
whose file was truncated.
# a torn temp file, as if the process died inside torch.save
with open(os.path.join(OUT_DIR, "shard_0008.pt.tmp"), "wb") as fh:
fh.write(b"\0" * 1_000_000)
# a shard that made it into the manifest but is short on disk
victim = os.path.join(OUT_DIR, "shard_0007.pt")
full = os.path.getsize(victim)
with open(victim, "r+b") as fh:
fh.truncate(full - 4096)
done = completed_shards(OUT_DIR)
print("manifest lines:", sum(1 for _ in open(os.path.join(OUT_DIR, "manifest.jsonl"))))
print("shards accepted on restart:", sorted(done))
print("shard 7 rejected (size mismatch):", 7 not in done)
manifest lines: 8 shards accepted on restart: [0, 1, 2, 3, 4, 5, 6] shard 7 rejected (size mismatch): True
Shard 7 is in the manifest but is 4 kB short, so the size check rejects it and it will be
rewritten. The .tmp file is never a candidate, because only files with the final name are
ever considered. This is the whole of resumability: about twenty lines, none of them
nnsight-specific.
Re-running the identical call now picks up where it stopped.
os.remove(os.path.join(OUT_DIR, "shard_0008.pt.tmp")) # a real run would sweep stale .tmp too
memlog, tok_per_s = collect(OUT_DIR)
shards = sorted(f for f in os.listdir(OUT_DIR) if f.endswith(".pt"))
disk_bytes = sum(os.path.getsize(os.path.join(OUT_DIR, f)) for f in os.listdir(OUT_DIR))
print(f"\n{len(shards)}/{N_SEQ // SHARD_SEQ} shards, no stale .tmp: "
f"{not any(f.endswith('.tmp') for f in os.listdir(OUT_DIR))}")
print(f"{disk_bytes:,} bytes = {disk_bytes / 2**30:.2f} GiB on disk")
print(f"{tok_per_s:,.0f} tok/s end to end, including torch.save")
[resume] 7 of 16 shards already complete
16/16 shards, no stale .tmp: True 3,020,541,156 bytes = 2.81 GiB on disk 18,209 tok/s end to end, including torch.save
Seven shards were skipped (0–6), and shards 7–15 were written — shard 7 because we truncated it. Note what did not happen: no global state carried over from the crashed run, no cache had to be invalidated, and the resumed segment needed no special code path.
The throughput number is end-to-end and includes writing ~2.8 GiB to disk, which for this
module list is the dominant cost — 15 modules at fp16 is 1.4 MiB of activations per
sequence, or about 22 kB per token. Capture without the disk writes ran at roughly
26,000 tok/s in the no_grad cell above; the writes are what you are paying for here, which
is the usual shape of a corpus job once the module list is more than a layer or two.
scale = 50_000 / N_SEQ
print(f"extrapolated to 50,000 sequences ({50_000 * SEQ_LEN / 1e6:.1f}M tokens):")
print(f" wall clock ~{50_000 * SEQ_LEN / tok_per_s / 60:.1f} min (at the rate measured above)")
print(f" disk ~{disk_bytes * scale / 2**30:.0f} GiB")
print(f" ...for one layer only, fp16: ~{disk_bytes * scale / len(NAMES) / 2**30:.1f} GiB")
extrapolated to 50,000 sequences (3.2M tokens): wall clock ~2.9 min (at the rate measured above) disk ~69 GiB ...for one layer only, fp16: ~4.6 GiB
This is why storage tiering matters at real scale. Sixty-nine gigabytes for 3.2M tokens is a lot of disk to fill before you know whether the experiment works. The pipeline this tutorial is derived from stored one layer token-level for all 50,000 sequences (4.6 GiB), per-layer pooled means for all 50,000 (0.9 GiB), and all layers token-level for only the first 4,096 (4.5 GiB) — 10.7 GiB total instead of 69. Decide which slice you need at full scale and which you only need a sample of, and store them separately.
Did anything leak?¶
The failure this section exists to prevent looks like this: the run is fine for an hour, then dies at 80%. Anything that accumulates per trace — a retained graph, an uncollected reference cycle, a growing list — is invisible over 100 prompts and fatal over 50,000. We logged memory at every shard boundary, so we can simply look.
df = pd.DataFrame(memlog)
fig = px.line(df, x="shard", y=["allocated", "peak", "reserved"], markers=True,
labels={"value": "GPU memory (MiB)", "shard": "shard index",
"variable": ""},
title="GPU memory at every shard boundary")
fig.update_yaxes(rangemode="tozero")
fig.show()
print("allocated:", df["allocated"].min(), "->", df["allocated"].max(), "MiB",
f"(spread {df['allocated'].max() - df['allocated'].min():.3f} MiB)")
print("peak :", df["peak"].min(), "->", df["peak"].max(), "MiB")