PyTorch inference_mode memory savings need autograd review

Direct answer

torch.inference_mode() is not guaranteed valid for every custom torch.autograd.Function path. It is safe for standard inference where no custom backward logic depends on saved tensors or version counters. When a workload saves tensors via save_for_backward or mutates tensors in place after forward, torch.no_grad() is the safer default. Between the two, steady-state request overhead is broadly similar; inference_mode() can reduce memory by skipping version tracking, but it may add minor sanity-check overhead and carries compatibility risk with custom autograd code.

What is confirmed vs. what needs review

Confirmed across PyTorch 2.x: inference_mode() produces inference tensors with requires_grad=False and disables view/version tracking. Operations that try to re-enable gradients on those tensors can raise errors or silently detach, depending on the surrounding Function implementation.

Version-dependent and not guaranteed: whether save_for_backward accepts inference tensors and whether later mutations on those tensors are allowed. Exact semantics must be confirmed against the deployed PyTorch release.

Decision path for this service

1. Audit custom autograd.Function classes for save_for_backward calls and in-place mutations after the forward pass.
2. If any exist, wrap only the standard model forward in inference_mode() and keep custom paths under no_grad(), or use no_grad() globally until compatibility is proven.
3. Separate cold-start from steady-state: import-time and first-request costs include context initialization; later requests reflect per-call overhead.

Compatibility check

Run a minimal forward pass in the deployed release with a custom Function that saves an input tensor and attempts a backward-relevant mutation. Capture any RuntimeError or unexpected detachment.

import torch, time

class CustomFn(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x * 2

    @staticmethod
    def backward(ctx, grad):
        x, = ctx.saved_tensors
        return grad * 2

def probe(mode_name, ctx_mgr):
    torch.cuda.reset_peak_memory_stats()
    t0 = time.time()
    with ctx_mgr():
        x = torch.randn(1024, device="cuda", requires_grad=True)
        y = CustomFn.apply(x)
        z = y.sum()
    dt = time.time() - t0
    peak = torch.cuda.max_memory_allocated()
    print(f"{mode_name}: time={dt:.6f}s peak={peak}")

probe("inference_mode", torch.inference_mode)
probe("no_grad", torch.no_grad)

If the custom path raises under inference_mode(), that confirms incompatibility for this workload. If it succeeds, re-test after any PyTorch upgrade because behavior is version-dependent.

Cold-start vs. steady-state measurement

Measure the first request after import separately from subsequent requests. Use torch.cuda.max_memory_allocated() for peak memory and wall-clock timing for latency. Expect one-time initialization cost in the first call; steady-state differences are usually small.

# cold-start
t0 = time.time()
with torch.inference_mode():
    out = model(inp)
print("cold:", time.time() - t0)

# steady-state
for _ in range(10):
    t0 = time.time()
    with torch.inference_mode():
        out = model(inp)
    print("steady:", time.time() - t0)

Missing diagnostic that changes the recommendation

Do any custom autograd.Function paths call save_for_backward on tensors that will be created inside inference_mode(), or mutate tensors in place after forward? If yes, inference_mode() is not valid for those paths without further review; if no, it is likely safe and may offer modest memory savings.