PT2E quantization silently does nothing when the model's forward is under torch.no_grad()

[john-rocky]

Summary

PT2E quantizes nothing when the exported forward runs under torch.no_grad(). No error, no warning. The .pte comes out the size of the fp32 build with correlation 1.000000 against fp32, which reads as "this model does not benefit from int8" rather than "int8 never ran".

Repro

import torch
import torch.nn as nn
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (
    XNNPACKQuantizer, get_symmetric_quantization_config)
from executorch.exir import to_edge_transform_and_lower
from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e


class MLP(nn.Module):
    def __init__(self, d=512, n=4):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(d, d) for _ in range(n)])

    def body(self, x):
        for layer in self.layers:
            x = torch.relu(layer(x))
        return x

    def forward(self, x):
        return self.body(x)


class MLPNoGrad(MLP):
    @torch.no_grad()
    def forward(self, x):
        return self.body(x)


def build(model, inputs, capture_under_no_grad=False):
    if capture_under_no_grad:
        with torch.no_grad():
            gm = torch.export.export(model, inputs).module()
    else:
        gm = torch.export.export(model, inputs).module()

    linear = sum(1 for n in gm.graph.nodes if n.op == "call_function"
                 and str(n.target).startswith("aten.linear"))
    hop = sorted({str(n.target) for n in gm.graph.nodes
                  if n.op == "call_function" and "set_grad" in str(n.target)})

    quantizer = XNNPACKQuantizer()
    quantizer.set_operator_type(
        torch.ops.aten.linear.default,
        get_symmetric_quantization_config(is_per_channel=True, is_dynamic=True))
    prepared = prepare_pt2e(gm, quantizer)
    with torch.no_grad():
        prepared(*inputs)
    converted = convert_pt2e(prepared)

    int8 = sum(1 for _, t in list(converted.named_buffers())
               + list(converted.named_parameters()) if t.dtype == torch.int8)
    ep = torch.export.export(converted, inputs)
    pte = to_edge_transform_and_lower(ep, partitioner=[XnnpackPartitioner()]).to_executorch()
    return linear, hop, int8, len(pte.buffer)


inputs = (torch.randn(8, 512),)
torch.manual_seed(0)
weights = MLP().eval().state_dict()
cases = (("plain forward", MLP, False),
         ("@torch.no_grad() forward", MLPNoGrad, False),
         ("@torch.no_grad(), captured under no_grad", MLPNoGrad, True))
for name, cls, under in cases:
    model = cls().eval()
    model.load_state_dict(weights)
    linear, hop, int8, size = build(model, inputs, under)
    print(f"{name:42s} aten.linear at top level={linear}  HOP={hop}  "
          f"int8 tensors after convert_pt2e={int8}  .pte={size / 1e6:.2f}MB")

Output — executorch 1.4.0, torch 2.13.0, macOS arm64:

plain forward                              aten.linear at top level=4  HOP=[]  int8 tensors after convert_pt2e=4  .pte=1.07MB
@torch.no_grad() forward                   aten.linear at top level=0  HOP=['wrap_with_set_grad_enabled']  int8 tensors after convert_pt2e=0  .pte=4.21MB
@torch.no_grad(), captured under no_grad   aten.linear at top level=4  HOP=[]  int8 tensors after convert_pt2e=4  .pte=1.07MB

All three arms load the same weights and take the same input. Only the decorator and the grad mode at capture change.

What happens

torch.export.export(model, inputs).module() records the no_grad region as a wrap_with_set_grad_enabled higher-order op, and the whole body moves into its subgraph. The top-level graph has no aten.linear left to annotate. Inside the HOP submodule the four aten.linear nodes are still there and carry zero quantize/dequantize nodes after convert_pt2e — the annotator never reached them. Same result with a global static config, so it is not specific to the dynamic recipe.

Row 3 is the workaround. The HOP is only built when the requested grad mode differs from the ambient one — torch/_export/passes/replace_set_grad_with_hop_pass.py:46, _is_set_grad_enabled_sub_mod(..., omit_if_same_with_ambient=True) at line 110 — so capturing under torch.no_grad() inlines the body and quantization proceeds normally.

Why it seems worth a guard

The documented flow is the one that hits it. docs/source/quantization-overview.md and docs/source/tutorial-xnnpack-delegate-lowering.md both call torch.export.export(model, sample_inputs).module() and say nothing about grad mode.

The decorator is common in the models people convert. In transformers 5.15.0, @torch.no_grad() sits directly on forward in 38 model directories, and on get_image_embeddings in 8 more (sam, sam2, edgetam, ...). I measured two of them: the Sam2Model and EdgeTamModel image encoders, both reached through get_image_embeddings.

That is how I found it. SAM2.1-hiera-tiny's vision encoder is 98.5% nn.Linear by parameter count, and its int8 build kept coming out at the same 109.2 MB as fp32 with correlation 1.000000. I spent a session reading the XNNPACK serializer for a bug that was not there. With the capture under no_grad, the same export gives 51 quantized linears and 29.1 MB.

Suggestions

1. Warn from prepare_pt2e when the graph handed to it contains a subgraph the annotator will not descend into. A quantizer that matches zero nodes is worth a warning on its own.
2. One line in quantization-overview.md: capture inside with torch.no_grad(): when the model's forward sets its own grad mode.

Happy to send either as a PR.