feat(backward): explain the freed autograd graph in terms of invokes, and cover batched-invoke gradients

[Hotragn]

Rebased onto 0.8 as requested in #696 — this supersedes #694, and folds in the salvageable half of #693. Verified the problem still reproduces on 0.8 before porting.

Summary

Calling .backward() in more than one invoke of a trace fails inside torch.autograd:

with model.trace() as tracer:
    with tracer.invoke(prompt_a):
        a1 = model.transformer.h[5].output
        with model.output.logits.sum().backward():
            grad_a = a1.grad.save()

    with tracer.invoke(prompt_b):
        a1 = model.transformer.h[5].output
        with model.output.logits.sum().backward():   # RuntimeError
            grad_b = a1.grad.save()

RuntimeError: Trying to backward through the graph a second time (or directly
access saved tensors after they have already been freed). ...

Accurate, but it points away from the cause. Each invoke calls .backward() exactly once, so "a second time" reads as describing someone else's code. And the real reason is a fact about nnsight rather than torch: invokes are not separate runs. Every invoke's input is combined into one batch and the model is called once, so the whole trace has a single autograd graph, and the first .backward() frees it for everything after it.

docs/gotchas/backward.md does cover retain_graph, but as "if you call .backward() more than once on overlapping graphs". One backward per invoke does not read as that.

Changes

_explain_freed_graph translates that one error — matching autograd's message (a bare RuntimeError, no type or code to key off, so string matching is the only option and it is kept narrow), stating the cause in nnsight's terms, and showing the fix while appending torch's original text:

RuntimeError: The autograd graph for this trace has already been freed.

Every invoke in a trace contributes to a single batched forward pass, so all
invokes share one autograd graph. An earlier `.backward()` in the same trace
freed the graph that this one needs.

Pass `retain_graph=True` to every `.backward()` except the last one:

    with model.trace() as tracer:
        with tracer.invoke(prompt_a):
            ...
            with loss.backward(retain_graph=True):
                grad_a = a1.grad.save()
        with tracer.invoke(prompt_b):
            ...
            with loss.backward():
                grad_b = a1.grad.save()

The same applies to two `.backward()` calls inside one invoke. If you do not
need the gradients together, run each backward pass in its own trace instead.

Original error from torch.autograd: Trying to backward through the graph a
second time ...

When the failing call already passed retain_graph=True, it says the earlier one did not, instead of suggesting a flag that is already set. Every other RuntimeError is re-raised untouched.

Batched-invoke gradient coverage, which the suite was missing entirely. test_backward.py only exercised single-invoke gradients, and test_batching.py is entirely @torch.no_grad() — so nothing asserted that an invoke's gradient is its own rows. TestBackwardAcrossInvokes pins that against plain-autograd references:

• two invokes, each reading its own gradient
• three invokes, parametrised over which one takes the gradient (first / middle / last), so each batch offset is exercised rather than only the final one
• a negative control: the result must not match another invoke's gradient
• edit isolation: a.grad = a.grad * 3 in one invoke leaves the other untouched

Invokes get different row counts (2 / 3 / 1), so a slice read from the wrong offset shows up as a shape mismatch even where the values would be close. Built on the existing MLP + _BatchEnvoy pattern, which keeps the assertions exact — allclose(atol=1e-6) — and adds no model download; the whole file runs in ~16s.

Behaviour is unchanged

retain_graph=True across invokes already worked, and each invoke's gradient was already correct — verified before writing the diagnostic. This makes the failure legible and locks the correctness in.

Not done here

nnsight could retain the graph implicitly for all but the last invoke and make this disappear. That trades peak memory for convenience on every batched trace and changes a default, so it seemed like your call rather than something to fold into an error message. Happy to do it if you'd prefer.

Verification

Without the diagnostic, 2 of the 15 tests in test_backward.py fail; with it, 15 pass. Full CPU suite on this branch: 857 passed, 7 skipped, 2 failed — the two failures are test_serialization.py::TestScopeFiltering on Python 3.14, which predate this branch and are fixed in #698.

_{0.8 @ 1f974f0, Python 3.14.3, transformers 5.15.1, torch 2.9.1, CPU.}

[Hotragn]

Rebased onto current 0.8 now that #698 is in, so the caveat in the description is resolved — the full CPU suite is green on this branch:

854 passed, 7 skipped, 1 xfailed, 27 warnings

No other changes; same single commit, just replayed onto 8f7546c with no conflicts.