test(lm): fix the batched-gradient assertion that has kept CI red for a month

[Hotragn]

Summary

tests/test_lm.py::TestGradients::test_backward_with_multiple_invokers has failed on every push to main and dev since it landed in #671 (2026-07-25). The test job has been red for a month:

FAILED tests/test_lm.py::TestGradients::test_backward_with_multiple_invokers[cpu] - assert False
======= 1 failed, 243 passed, 2 skipped, 1 xpassed, 6 warnings in 44.13s =======

The gradient is correct. The assertion is not satisfiable.

Why the assertion cannot pass

assert torch.allclose(ref_grad, batched_grad, atol=1e-5)

d(sum lm_head.output)/d(h[5].attn.c_proj.output) reaches ~4e4 in float32, where a single ULP is already ~4e-3 — 400× the tolerance. Running the same prompt batched vs. unbatched reorders the reductions through six layers of backward, so the two differ by ~1e-5 relative (0.7 absolute) regardless of what nnsight does.

Two independent controls confirm nnsight is not the source:

1. Pure HuggingFace, no nnsight. retain_grad() on the same submodule, logits[i:i+1].sum().backward(), same batch shape — reproduces the discrepancy bit for bit:

   nnsight  single-vs-batched: abs diff max = 1.07958984375   rel diff max = 0.06253758
   pure HF  single-vs-batched: abs diff max = 1.07958984375   rel diff max = 0.06253758
   

2. float64. The same comparison, same code, dtype=torch.float64:

   float32:  nnsight batched vs single-prompt reference -> absmax 6.6e-01, relmax 4.1e-02
   float64:  nnsight batched vs single-prompt reference -> absmax 2.7e-10, relmax 5.1e-11
   

   The error collapses by nine orders of magnitude when the mantissa grows. It is float32 reduction order, not a slicing bug.

I also checked the forward pass separately, in case left-padding was involved: nnsight's left-padded batch matches pure HF's left-padded batch exactly (1.37e-04 on logits of magnitude 155, i.e. ~1e-6 relative), and dropping the derived position_ids moves that to 9.08e+01 — so #673's derivation is doing its job.

The test also could not catch its own bug class

Both invokes used the same prompt:

prompt = "The quick brown fox jumps"
with tracer.invoke(prompt): ...
with tracer.invoke(prompt): ...   # grad taken here

With identical inputs, returning invoke 0's gradient — or the entire batch — is indistinguishable from success. The test guards MissedProviderError (#664) but not "is this my gradient", which is the harder half.

It also only ever exercised the last invoke of two, so the batch-slice offset under test was always the same one.

Changes

• Assert on a normalised norm ratio (_grad_rel_err) instead of allclose. It is scale-free and insensitive to individual near-zero entries, where cancellation makes the elementwise relative error large (0.06) even for a correct tensor. Measured spread is ~1e-5; threshold is 1e-4.
• Distinct prompts per invoke, so a wrong row is observable.
• Negative control: the result must not match any other invoke's gradient. Measured ~1.2 relative — four orders above the threshold.
• Parametrised over which invoke takes the gradient: first/last of two, and first/middle/last of three, so each batch-slice offset is exercised rather than only the final one.

Does the new test have teeth?

Mutation test — force the gradient to always be read from batch row 0:

-shape, stride, offset = tensor.shape, tensor.stride(), tensor.storage_offset()
+shape, stride, offset = tensor.shape, tensor.stride(), 0  # MUTANT

new test vs mutant:
  FAILED ...test_backward_with_multiple_invokers[cpu-2-1]
  FAILED ...test_backward_with_multiple_invokers[cpu-3-1]
  FAILED ...test_backward_with_multiple_invokers[cpu-3-2]
  3 failed, 2 passed          <- exactly the cases whose target is not row 0

old test vs mutant:
  1 failed                    <- but it also failed without the mutant, so it
                                 cannot distinguish the two

Result

Full CI test list on CPU, this branch:

248 passed, 2 skipped, 1 xpassed, 6 warnings in 146.87s

Green, with strictly more coverage than before (5 parametrised cases where there was 1).

Test-only change — no source files touched.

_{Run with transformers 5.15.1 / torch 2.9.1 on CPU.}

[Hotragn]

Closing this — it doesn't apply to 0.8, which is where I've moved per #696.

test_backward_with_multiple_invokers doesn't exist on 0.8; tests/test_backward.py is a rewrite against a small MLP with a plain-autograd reference_grad, and the unsatisfiable atol=1e-5-on-4e4-magnitude assertion is gone with it. So the red CI this PR was fixing is a dev/main-only problem. If dev is still going to see releases it's a one-line change to make there, but it seemed wrong to leave a PR open against a branch you're not developing on.

Worth noting what is missing on 0.8, though: there is now no batched-invoke gradient coverage at all. test_backward.py only exercises single-invoke gradients, and test_batching.py is entirely @torch.no_grad() — so nothing asserts that an invoke's gradient is its own rows. I've carried that half over to #699, parametrised over first/middle/last of two and three invokes with a negative control, using the existing MLP + _BatchEnvoy pattern so the assertions are exact (atol=1e-6) rather than fighting float32 reduction order.