From 143d6e8e56ce1b097222ce5c140c56ba72d10f7b Mon Sep 17 00:00:00 2001 From: Utkarsh Simha Date: Wed, 1 Jul 2026 21:46:20 -0700 Subject: [PATCH 1/4] Add graph mode debugging hints and troubleshooting doc Surfaces actionable hints (export_with_no_grad toggle, dynamic_shapes, eager fallback) in the torch.export.export failure message, and includes the source module name when a graph partition annotation fails so it can be excluded via module_name_configs. Adds a Graph Mode Troubleshooting doc under a new Debugging doc section. --- .../debugging/graph_mode_troubleshooting.md | 90 +++++++++++++++++++ .../{utils => debugging}/model_inspection.md | 0 docs/src/index.md | 10 ++- docs/src/quantization/config.md | 4 +- docs/src/quantization/overview.md | 2 +- src/coreai_opt/_utils/torch_utils.py | 26 +++++- .../quantization/_graph/_annotation_utils.py | 28 +++++- .../quantization/test_graph_mode_quantizer.py | 14 +++ tests/quantization/test_graph_mode_utils.py | 60 +++++++++++++ 9 files changed, 225 insertions(+), 9 deletions(-) create mode 100644 docs/src/debugging/graph_mode_troubleshooting.md rename docs/src/{utils => debugging}/model_inspection.md (100%) diff --git a/docs/src/debugging/graph_mode_troubleshooting.md b/docs/src/debugging/graph_mode_troubleshooting.md new file mode 100644 index 0000000..be4afef --- /dev/null +++ b/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 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), 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/) and report against torch. + +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 on our end. **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) diff --git a/docs/src/utils/model_inspection.md b/docs/src/debugging/model_inspection.md similarity index 100% rename from docs/src/utils/model_inspection.md rename to docs/src/debugging/model_inspection.md diff --git a/docs/src/index.md b/docs/src/index.md index 92b8f7c..07d6c4d 100644 --- a/docs/src/index.md +++ b/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 diff --git a/docs/src/quantization/config.md b/docs/src/quantization/config.md index 2adf49a..4f0f2d1 100644 --- a/docs/src/quantization/config.md +++ b/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. diff --git a/docs/src/quantization/overview.md b/docs/src/quantization/overview.md index 49eceaa..44f46de 100644 --- a/docs/src/quantization/overview.md +++ b/docs/src/quantization/overview.md @@ -202,7 +202,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 diff --git a/src/coreai_opt/_utils/torch_utils.py b/src/coreai_opt/_utils/torch_utils.py index 915891c..6966c1c 100644 --- a/src/coreai_opt/_utils/torch_utils.py +++ b/src/coreai_opt/_utils/torch_utils.py @@ -390,12 +390,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: @@ -416,8 +417,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 diff --git a/src/coreai_opt/quantization/_graph/_annotation_utils.py b/src/coreai_opt/quantization/_graph/_annotation_utils.py index fe304ff..315be17 100644 --- a/src/coreai_opt/quantization/_graph/_annotation_utils.py +++ b/src/coreai_opt/quantization/_graph/_annotation_utils.py @@ -41,6 +41,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, @@ -914,8 +915,31 @@ 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: + 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( diff --git a/tests/quantization/test_graph_mode_quantizer.py b/tests/quantization/test_graph_mode_quantizer.py index 81d3d3a..3f5d739 100644 --- a/tests/quantization/test_graph_mode_quantizer.py +++ b/tests/quantization/test_graph_mode_quantizer.py @@ -183,6 +183,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 ): diff --git a/tests/quantization/test_graph_mode_utils.py b/tests/quantization/test_graph_mode_utils.py index 13f838f..e59b95d 100644 --- a/tests/quantization/test_graph_mode_utils.py +++ b/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) From ed8263dc004fadd987380792419bcd46e06da836 Mon Sep 17 00:00:00 2001 From: usimha Date: Tue, 14 Jul 2026 16:56:38 -0700 Subject: [PATCH 2/4] docs: fix stale model_inspection.md link in palettization config The link still pointed at the old utils/ location after this branch moved the doc to debugging/. --- docs/src/palettization/config.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/palettization/config.md b/docs/src/palettization/config.md index e83f777..9718299 100644 --- a/docs/src/palettization/config.md +++ b/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 From f397d039869cbd3f15a59b6dcce3bef96ff8b957 Mon Sep 17 00:00:00 2001 From: usimha <135899523+u-simha@users.noreply.github.com> Date: Tue, 14 Jul 2026 17:55:38 -0700 Subject: [PATCH 3/4] docs: address review feedback on graph mode troubleshooting guide Link "Graph execution mode" to the execution-modes section, note that export failures may warrant fixing the model definition (since the same construct can block coreai-torch conversion later), and drop the "our"/"report against torch" phrasing per review. --- docs/src/debugging/graph_mode_troubleshooting.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/src/debugging/graph_mode_troubleshooting.md b/docs/src/debugging/graph_mode_troubleshooting.md index be4afef..3da2086 100644 --- a/docs/src/debugging/graph_mode_troubleshooting.md +++ b/docs/src/debugging/graph_mode_troubleshooting.md @@ -1,6 +1,6 @@ # Graph Mode Troubleshooting -This guide helps debug common issues when using Graph execution mode in CoreAI-Opt. +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. @@ -61,13 +61,13 @@ prepared = quantizer.prepare( 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), see [Fall back to EAGER execution mode](#fall-back-to-eager-execution-mode) below. +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/) and report against torch. +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 on our end. **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. +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 From f8ce1f4ce8e2c4b8ba0d18a7a5c4b248ee284aac Mon Sep 17 00:00:00 2001 From: usimha <135899523+u-simha@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:16:30 -0700 Subject: [PATCH 4/4] fix: don't error on SymInt-only multi-node partitions in annotation utils torch.export's insert_deferred_runtime_asserts synthesizes one SymInt mul per shape-runtime assertion under a shared torch_fn tag, collapsing several into a single SourcePartition. The new strict "exactly 1 call function node" check treated this as an error, regressing test_prepare_with_symint_mul_partition_collision (and CI). These SymInt nodes carry no tensor value to annotate, so picking any one of them is safe, matching the pre-existing lenient behavior for this case. --- src/coreai_opt/quantization/_graph/_annotation_utils.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/coreai_opt/quantization/_graph/_annotation_utils.py b/src/coreai_opt/quantization/_graph/_annotation_utils.py index 315be17..1a526fd 100644 --- a/src/coreai_opt/quantization/_graph/_annotation_utils.py +++ b/src/coreai_opt/quantization/_graph/_annotation_utils.py @@ -922,6 +922,15 @@ def _get_call_function_node_from_partition(partition: SourcePartition) -> torch. """ 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