Add graph mode debugging hints and troubleshooting doc

docs/src/debugging/graph_mode_troubleshooting.md

@@ -0,0 +1,90 @@
+# Graph Mode Troubleshooting
+
+This guide helps debug common issues when using [Graph execution mode](../quantization/overview.md#two-execution-modes-graph-and-eager) 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:
+
+```python
+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](#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-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.
+
+```python
+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:
+
+```python
+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](https://docs.pytorch.org/tutorials/intermediate/torch_export_tutorial.html#constraints-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](https://github.com/apple/coreai-torch) later on. Otherwise, see [Fall back to EAGER execution mode](#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](https://docs.pytorch.org/ao/stable/).
+
+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](#fall-back-to-eager-execution-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.
+
+```python
+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](../quantization/overview.md#choosing-between-graph-and-eager-mode) for the trade-offs between the two modes.
+
+## External Resources
+
+- [PyTorch Export Tutorial](https://docs.pytorch.org/tutorials/intermediate/torch_export_tutorial.html)
+- [Dynamic Shapes](https://docs.pytorch.org/tutorials/intermediate/torch_export_tutorial.html#constraints-dynamic-shapes)

docs/src/index.md

@@ -38,12 +38,20 @@ palettization/index
 
 utils/joint_compression
 utils/mixed_precision
-utils/model_inspection
 utils/activation_comparison
 utils/casting
 utils/coreai_compression
 ```
 
+```{toctree}
+:maxdepth: 1
+:caption: Debugging
+:hidden:
+
+debugging/model_inspection
+debugging/graph_mode_troubleshooting
+```
+
 ```{toctree}
 :maxdepth: 1
 :caption: Reference

docs/src/palettization/config.md

@@ -116,7 +116,7 @@ op_config = OpKMeansPalettizerConfig(
 
 ## Examples
 
-Several examples below configure specific module types or module names. To determine these for your model, use {class}`~coreai_opt.inspection.ModelInspector` with `execution_mode="eager"` — see [Inspecting Model Structure](../utils/model_inspection.md). Palettization supports eager mode only.
+Several examples below configure specific module types or module names. To determine these for your model, use {class}`~coreai_opt.inspection.ModelInspector` with `execution_mode="eager"` — see [Inspecting Model Structure](../debugging/model_inspection.md). Palettization supports eager mode only.
 
 ### Apply 4-bit palettization globally, 8-bit to linear layers
 

docs/src/quantization/config.md

@@ -184,7 +184,7 @@ The defaults are:
 
 In [Quantization Overview](overview.md) we saw how to use the default `W_INT8_A_INT8` config. [Config classes and their defaults](#config-classes-and-their-defaults) described the default settings in `QuantizerConfig()`, `ModuleQuantizerConfig()`, and `OpQuantizerConfig()`. Let us now see how to configure quantization when non-default settings are desired.
 
-Several examples below configure specific module names, module types, op names, or op types. To determine these for your model, see [Inspecting Model Structure](../utils/model_inspection.md).
+Several examples below configure specific module names, module types, op names, or op types. To determine these for your model, see [Inspecting Model Structure](../debugging/model_inspection.md).
 
 ### Example: `W_MXFP4_A_FP8` applied to all supported ops
 
@@ -874,4 +874,4 @@ inspector = ModelInspector(
 print(inspector.format_summary())
 ```
 
-See [Inspecting Model Structure](../utils/model_inspection.md) for full usage, examples, and a comparison of graph and eager mode op naming.
+See [Inspecting Model Structure](../debugging/model_inspection.md) for full usage, examples, and a comparison of graph and eager mode op naming.

docs/src/quantization/overview.md

@@ -203,7 +203,7 @@ The two modes are expected to produce very similar models for weight-only quanti
 
 A few scenarios where `eager` mode may need to be used instead of `graph`:
 
-- If you run into any errors during the `prepare` call which, under the hood, invokes the `torch.export.export` and `torchao`'s `prepare_qat_pt2e`/`convert_pt2e` APIs.
+- If you run into any errors during the `prepare` call which, under the hood, invokes the `torch.export.export` and `torchao`'s `prepare_qat_pt2e`/`convert_pt2e` APIs. See [Graph Mode Troubleshooting](../debugging/graph_mode_troubleshooting.md) for common export errors and workarounds before falling back to eager mode.
 - When `torch.nn.Module` needs to be provided as an input, instead of `ExportedProgram` to the conversion API of [coreai-torch](https://github.com/apple/coreai-torch). This happens when the `coreai-torch` conversion needs to "externalize" certain sub-modules to map them to _composite ops_ for better runtime performance.
 
 #### Weights and activations quantization

src/coreai_opt/_utils/torch_utils.py

@@ -403,12 +403,13 @@ def export_model(
     Args:
         model (torch.nn.Module): The model to export.
         example_inputs (tuple[Any, ...]): Example inputs for tracing.
-        dynamic_shapes: Dynamic shapes specification for torch.export.
+        dynamic_shapes (dict[str, Any] | tuple[Any] | list[Any] | None):
+                        Dynamic shapes specification for torch.export.
                         Can be a dict mapping input names to dynamic dimensions,
                         a tuple/list of dynamic shapes per input, or None for
                         static shapes. Used to specify which dimensions can vary
                         at runtime during model export.
-        export_with_no_grad: Whether to call torch.export.export within a
+        export_with_no_grad (bool): Whether to call torch.export.export within a
                              torch.no_grad() context.
 
     Returns:
@@ -429,8 +430,27 @@ def export_model(
                 exported_model = exported_program.module()
         return exported_model
     except Exception as e:
+        no_grad_hint = (
+            "  - Try export_with_no_grad=False — the torch.no_grad() context can "
+            "modify tracing behavior for some models.\n"
+            if export_with_no_grad
+            else "  - Try export_with_no_grad=True — exporting with torch.no_grad() "
+            "simplifies the traced graph.\n"
+        )
+        _dynamic_shapes_url = (
+            "https://docs.pytorch.org"
+            "/tutorials/intermediate/torch_export_tutorial.html"
+            "#constraints-dynamic-shapes"
+        )
         raise RuntimeError(
-            f"Failed to trace the model with torch.export.export(), received error: {e}"
+            f"Failed to trace the model with torch.export.export(), received error: "
+            f"{e}\n\n"
+            f"Debugging hints:\n"
+            f"{no_grad_hint}"
+            f"  - If the error mentions shape constraints, try specifying "
+            f"dynamic_shapes. See: {_dynamic_shapes_url}\n"
+            f"  - As a last resort, consider using EAGER execution mode: "
+            f"config.execution_mode = ExecutionMode.EAGER"
         ) from e
 
 

src/coreai_opt/quantization/_graph/_annotation_utils.py

@@ -40,6 +40,7 @@
 from coreai_opt._utils.version_utils import version_ge as _version_ge
 from coreai_opt.config.compression_config import ModuleConfigDict
 from coreai_opt.config.spec import CompressionTargetTensor
+from coreai_opt.quantization._graph._utils import get_source_module_name
 from coreai_opt.quantization.config import ModuleQuantizerConfig
 from coreai_opt.quantization.config.quantization_config import (
     _ACTIVATION_SPEC_DICT,
@@ -948,8 +949,40 @@ def _propagate_output_qspec(
 
 
 def _get_call_function_node_from_partition(partition: SourcePartition) -> torch.fx.Node:
-    """Return the first call_function node in the partition."""
-    return [node for node in partition.nodes if node.op == "call_function"][0]
+    """
+    Given a partition, return the call function node associated with the partition.
+
+    We expect there to be only one call function node in the partition.
+    """
+    call_function_nodes = [node for node in partition.nodes if node.op == "call_function"]
+    if len(call_function_nodes) != 1:
+        # torch.export's insert_deferred_runtime_asserts synthesizes one SymInt mul per
+        # shape-runtime assertion, all sharing one torch_fn tag, so several can collapse
+        # into a single partition. They carry no tensor value to annotate, so picking any
+        # one of them is safe here; downstream floating-point filtering no-ops on SymInt.
+        if call_function_nodes and all(
+            isinstance(node.meta.get("val"), torch.SymInt) for node in call_function_nodes
+        ):
+            return call_function_nodes[0]
+
+        module_names = {
+            name
+            for node in call_function_nodes
+            if (name := get_source_module_name(node)) is not None
+        }
+        module_hint = ""
+        if module_names:
+            module_hint = (
+                f"\nSource module(s): {', '.join(sorted(module_names))}. "
+                f"Consider excluding this module from quantization via "
+                f"module_name_configs."
+            )
+        error_msg = (
+            f"Expected exactly 1 call function node in source partition but got "
+            f"{call_function_nodes}.{module_hint}"
+        )
+        raise RuntimeError(error_msg)
+    return call_function_nodes[0]
 
 
 def match_pattern_with_sequential_partitions(

tests/quantization/test_graph_mode_quantizer.py

@@ -184,6 +184,20 @@ def test_prepare_with_dynamic_shapes(
         output = prepared_model(simple_model_input)
         assert output.shape == (1, 10)
 
+    def test_export_error_message_contains_hints(self, basic_config):
+        """Test that export failure error messages contain debugging hints."""
+
+        class DataDependentModel(nn.Module):
+            def forward(self, x):
+                if x.sum() > 0:
+                    return x * 2
+                return x * 3
+
+        model = DataDependentModel()
+        quantizer = Quantizer(model, basic_config)
+        with pytest.raises(RuntimeError, match="Debugging hints"):
+            quantizer.prepare(example_inputs=(torch.randn(2, 2),))
+
     def test_finalize_with_none_model_arg(
         self, simple_conv_linear_model, basic_config, simple_model_input
     ):

tests/quantization/test_graph_mode_utils.py

@@ -4,11 +4,16 @@
 # be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
 
 from functools import wraps
+from unittest.mock import MagicMock
 
+import pytest
 import torch
 import torch.nn as nn
 
 from coreai_opt.quantization import Quantizer, QuantizerConfig
+from coreai_opt.quantization._graph._annotation_utils import (
+    _get_call_function_node_from_partition,
+)
 from coreai_opt.quantization._graph._utils import (
     _has_no_disallowed_kwargs,
     restore_kwargs,
@@ -178,3 +183,58 @@ def test_strip_and_restore_round_trip(self):
 
         restore_kwargs(graph, saved)
         assert custom_node.kwargs == metadata
+
+
+class TestGetCallFunctionNodeFromPartition:
+    def test_single_node_returns_node(self):
+        """Single call_function node in partition is returned successfully."""
+        node = MagicMock()
+        node.op = "call_function"
+        partition = MagicMock()
+        partition.nodes = [node]
+
+        result = _get_call_function_node_from_partition(partition)
+        assert result is node
+
+    def test_multi_node_error_includes_module_name(self):
+        """Multi-node partition error includes module name from nn_module_stack."""
+        module_fqn = "encoder.layer.0.attention.self"
+        nodes = []
+        for _ in range(3):
+            node = MagicMock()
+            node.op = "call_function"
+            node.meta = {"nn_module_stack": {"key": (module_fqn, type)}}
+            nodes.append(node)
+
+        partition = MagicMock()
+        partition.nodes = nodes
+
+        with pytest.raises(RuntimeError, match=module_fqn):
+            _get_call_function_node_from_partition(partition)
+
+    def test_multi_node_error_suggests_module_name_configs(self):
+        """Multi-node partition error suggests using module_name_configs."""
+        node = MagicMock()
+        node.op = "call_function"
+        node.meta = {"nn_module_stack": {"key": ("model.layer1", type)}}
+
+        partition = MagicMock()
+        partition.nodes = [node, node]
+
+        with pytest.raises(RuntimeError, match="module_name_configs"):
+            _get_call_function_node_from_partition(partition)
+
+    def test_multi_node_error_without_module_stack(self):
+        """Multi-node partition error still works without nn_module_stack."""
+        nodes = []
+        for _ in range(2):
+            node = MagicMock()
+            node.op = "call_function"
+            node.meta = {}
+            nodes.append(node)
+
+        partition = MagicMock()
+        partition.nodes = nodes
+
+        with pytest.raises(RuntimeError, match="Expected exactly 1 call function node"):
+            _get_call_function_node_from_partition(partition)