Graph Mode Troubleshooting

This guide helps debug common issues when using Graph execution mode in CoreAI-Opt.

A quantizer.prepare() failure in graph mode happens in one of two stages, and the fix differs sharply between them. The first thing to do is figure out which stage is failing.

Step 1: Diagnose — does torch.export.export succeed?

Quantizer.prepare() first calls torch.export.export to trace the model into an FX graph, then applies quantization annotations on the resulting graph and runs torchao’s prepare_qat_pt2e. To localize the failure, run the export step directly with the same arguments prepare() would use:

import torch

with torch.no_grad():  # matches export_with_no_grad=True (the prepare() default)
    exported_program = torch.export.export(model, example_inputs)

The result of this experiment determines which path to follow:

• Export fails. Go to If torch.export.export fails. The model isn’t torch.export-compatible as written; the workarounds in Steps 2-3 may help.

• Export succeeds but prepare() still fails. Go to If prepare() fails after a successful export.

If torch.export.export fails

Step 2: Try export_with_no_grad=False

The default export_with_no_grad=True wraps the export call in torch.no_grad(). For some models, this context modifies tracing behavior and causes guard failures.

prepared = quantizer.prepare(
    example_inputs=(input_tensor,),
    export_with_no_grad=False,
)

Step 3: Use dynamic shapes for shape-related errors

If the error mentions shape constraints, guards, or symbolic dimensions, the model likely has inputs with variable dimensions (e.g., sequence length, batch size). Specify dynamic_shapes to tell the exporter which dimensions can vary:

from torch.export.dynamic_shapes import Dim

# Example: dynamic batch dimension
prepared = quantizer.prepare(
    example_inputs=(input_tensor,),
    dynamic_shapes={"x": (Dim.AUTO, Dim.STATIC, Dim.STATIC, Dim.STATIC)},
)

# Example: dynamic sequence length with a max constraint
import torch.export

dynamic_shapes = {
    "input_ids": {1: torch.export.Dim("seq_len", max=2048)},
    "attention_mask": {1: torch.export.Dim("seq_len", max=2048)},
}
prepared = quantizer.prepare(
    example_inputs=(input_ids, attention_mask),
    dynamic_shapes=dynamic_shapes,
)

For full details on dynamic shapes, see the PyTorch Export Tutorial – Dynamic Shapes.

If Steps 2-3 don’t resolve the export failure (e.g., the model has data-dependent control flow that torch.export cannot capture), the model definition itself may need to change to become exportable — this is worth fixing at the source, since the same construct can also block conversion via coreai-torch later on. Otherwise, see Fall back to EAGER execution mode below.

If prepare() fails after a successful export

After torch.export.export returns, Quantizer.prepare() applies coreai-opt’s annotation pass and then calls into torch’s prepare_qat_pt2e API. If the error you’re seeing comes from prepare_qat_pt2e itself, it is a torch-side issue — refer to the torchao documentation.

If the error does not come from prepare_qat_pt2e (i.e. it originates inside coreai-opt’s annotation pass), it likely indicates a bug in coreai-opt. Please file an issue on GitHub with the error message and a minimal reproducer. In the meantime, fall back to eager mode below — eager bypasses the entire graph-mode pipeline.

Fall back to EAGER execution mode

EAGER mode bypasses torch.export entirely and uses runtime tracing instead. It is the common fallback for both export failures that can’t be worked around with Steps 2-3 and post-export prepare() failures.

from coreai_opt.quantization import ExecutionMode

config = QuantizerConfig.presets.w8()
config.execution_mode = ExecutionMode.EAGER
quantizer = Quantizer(model, config)
prepared = quantizer.prepare(example_inputs=(input_tensor,))

See Choosing between graph and eager mode for the trade-offs between the two modes.

External Resources

• PyTorch Export Tutorial

• Dynamic Shapes