Info

Last Execution: 2026-07-25

Package	Version
nnsight	0.8
Python	3.12.13
torch	2.13.0+cu126
transformers	5.15.0

Gradients¶

nnsight lets you read and edit gradients through the with tensor.backward(): context. Opening it runs the real backward pass interleaved with the body of the block, so the block can read and replace the .grad of any tensor as the gradient reaches it. This is the foundation for attribution methods like integrated gradients and attribution patching.

Two rules shape everything below:

• Capture forward tensors before the backward block. By the time autograd runs, the forward pass is over — grab any .output/.input you want gradients for first, then open with loss.backward(): and read .grad on them.
• .grad lives on tensors, not modules. There is no module.grad: capture the tensor from .output, then read its .grad.

Setup¶

In [1]:

from nnsight.modeling.transformers import TransformersModel

model = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True)

from nnsight.modeling.transformers import TransformersModel model = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True)

Accessing Gradients¶

Capture a tensor during the forward pass, then enter a with loss.backward(): context and read .grad on the tensor you captured. A tensor read from .output is already in the autograd graph, so no requires_grad_(True) is needed. Reading hs.grad parks the block until autograd produces that gradient.

In [2]:

with model.trace("The Eiffel Tower is in the city of"):

    hs = model.transformer.h[-1].output       # capture during the forward

    loss = model.output.logits.sum()

    with loss.backward():                      # real backward, interleaved
        grad = hs.grad.clone().save()          # read the gradient flowing into hs

print(f"Gradient shape: {grad.shape}")
print(f"Gradient mean: {grad.abs().mean():.4f}")

with model.trace("The Eiffel Tower is in the city of"): hs = model.transformer.h[-1].output # capture during the forward loss = model.output.logits.sum() with loss.backward(): # real backward, interleaved grad = hs.grad.clone().save() # read the gradient flowing into hs print(f"Gradient shape: {grad.shape}") print(f"Gradient mean: {grad.abs().mean():.4f}")

Gradient shape: torch.Size([1, 10, 768])
Gradient mean: 472.7237

Request .grad on the tensor you captured directly

Read .grad on the exact tensor you captured (hs.grad), not on a slice or index of it (hs[0].grad). An indexing view is a new tensor whose gradient isn't the one autograd delivers, and requesting it raises OutOfOrderError.

Gradients at Multiple Layers¶

You can capture gradients at multiple points. Because backprop reaches the deepest layer first, request .grad in the reverse of the forward order — later layers first. Requesting an earlier-forward tensor's gradient before a later one raises OutOfOrderError.

In [3]:

with model.trace("The Eiffel Tower is in the city of"):

    hs_early = model.transformer.h[0].output
    hs_late = model.transformer.h[-1].output

    with model.output.logits.sum().backward():
        # Reverse-forward order: late layer first, then early
        grad_late = hs_late.grad.clone().save()
        grad_early = hs_early.grad.clone().save()

print(f"Layer 0 gradient mean:  {grad_early.abs().mean():.4f}")
print(f"Layer 11 gradient mean: {grad_late.abs().mean():.4f}")

with model.trace("The Eiffel Tower is in the city of"): hs_early = model.transformer.h[0].output hs_late = model.transformer.h[-1].output with model.output.logits.sum().backward(): # Reverse-forward order: late layer first, then early grad_late = hs_late.grad.clone().save() grad_early = hs_early.grad.clone().save() print(f"Layer 0 gradient mean: {grad_early.abs().mean():.4f}") print(f"Layer 11 gradient mean: {grad_late.abs().mean():.4f}")

Layer 0 gradient mean:  7570.3960
Layer 11 gradient mean: 472.7237

Modifying Gradients¶

You can intervene on gradients just like forward-pass activations. Assigning hs.grad = ... inside the block replaces the gradient that flows onward to earlier layers and weights.

In [4]:

with model.trace("The Eiffel Tower is in the city of"):

    hs = model.transformer.h[-1].output

    with model.output.logits.sum().backward():
        grad_before = hs.grad.clone().save()
        hs.grad = hs.grad * 0                   # ablate the gradient downstream
        grad_after = hs.grad.clone().save()

print(f"Before ablation - mean: {grad_before.abs().mean():.4f}")
print(f"After ablation - mean:  {grad_after.abs().mean():.4f}")

with model.trace("The Eiffel Tower is in the city of"): hs = model.transformer.h[-1].output with model.output.logits.sum().backward(): grad_before = hs.grad.clone().save() hs.grad = hs.grad * 0 # ablate the gradient downstream grad_after = hs.grad.clone().save() print(f"Before ablation - mean: {grad_before.abs().mean():.4f}") print(f"After ablation - mean: {grad_after.abs().mean():.4f}")

Before ablation - mean: 472.7237
After ablation - mean:  0.0000

Multiple Backward Passes¶

Pass retain_graph=True on all but the last backward call to keep the computation graph alive for additional passes.

In [5]:

with model.trace("The Eiffel Tower is in the city of"):

    hs = model.transformer.h[-1].output
    logits = model.output.logits

    with logits.sum().backward(retain_graph=True):
        grad1 = hs.grad.clone().save()

    with (logits.sum() * 2).backward():
        grad2 = hs.grad.clone().save()

print(f"First backward grad mean:  {grad1.abs().mean():.4f}")
print(f"Second backward grad mean: {grad2.abs().mean():.4f}")  # 2x the first

with model.trace("The Eiffel Tower is in the city of"): hs = model.transformer.h[-1].output logits = model.output.logits with logits.sum().backward(retain_graph=True): grad1 = hs.grad.clone().save() with (logits.sum() * 2).backward(): grad2 = hs.grad.clone().save() print(f"First backward grad mean: {grad1.abs().mean():.4f}") print(f"Second backward grad mean: {grad2.abs().mean():.4f}") # 2x the first

First backward grad mean:  472.7237
Second backward grad mean: 945.4473

Standalone Backward¶

with tensor.backward(): is independent of the forward trace — it works anywhere, as long as the tensors' autograd graph is still alive. Save the forward tensors, then open the backward block afterward.

In [6]:

with model.trace("The Eiffel Tower is in the city of"):

    hs = model.transformer.h[-1].output.save()
    logits = model.output.logits.save()

# Backward pass outside the trace
with logits.sum().backward():
    grad = hs.grad.clone().save()

print(f"Standalone backward grad shape: {grad.shape}")

with model.trace("The Eiffel Tower is in the city of"): hs = model.transformer.h[-1].output.save() logits = model.output.logits.save() # Backward pass outside the trace with logits.sum().backward(): grad = hs.grad.clone().save() print(f"Standalone backward grad shape: {grad.shape}")

Standalone backward grad shape: torch.Size([1, 10, 768])

How the backward context works

When you import nnsight, it patches torch.Tensor.backward at import time so it can be used as a with context manager. For the duration of one backward run, torch.Tensor.grad is replaced by a property: reading t.grad registers a self-removing autograd hook and parks the block until the gradient arrives; writing t.grad = v swaps a replacement into that same channel. A bare tensor.backward() with no with block falls through to vanilla PyTorch unchanged.