GB0029

Graph-Break Type

Short name describing what triggered the graph break

Compilation of intermediate hooks requires compiled autograd

Context

Values or code snippet captured at the break point

var_getattr {self} {name}

Explanation

Explanation of why the graph break was triggered

Dynamo must be in compiled_autograd to register hooks.

Hints

Hints on how to resolve the graph break

• Consider using torch.autograd.Function with a custom backward() method instead of register_hook().

Additional Information

Why This Happens

When you call register_hook() on an intermediate tensor inside compiled code, Dynamo cannot trace through the hook callback during the forward pass. The hook will be invoked during the backward pass, which requires compiled autograd to properly handle.

Workaround: Use torch.autograd.Function

If the backwards hook is actually traceable, instead of registering a hook on an intermediate tensor, wrap the tensor in a custom autograd.Function that applies the same transformation in its backward() method.

Before (causes graph break):

def fn(x):
    y = x * 2
    y.register_hook(lambda grad: grad + 1)  # Graph break!
    return y.sum()

After (compiles with fullgraph=True):

class AddOneToGrad(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        return x  # identity in forward

    @staticmethod
    def backward(ctx, grad):
        return grad + 1  # modify gradient in backward

def fn(x):
    y = x * 2
    y = AddOneToGrad.apply(y)  # No graph break
    return y.sum()

You need to make sure you use the result of the custom autograd function (unlike register_hook which is side-effectful).

Click here to add Additional Info

Back to Registry