initial tracing of model (#42)

Add torch.fx tracing that can be used when replacing layers with compressed variants, or by calling the function directly

Author: nroope

src/pquant/__init__.py

@@ -19,7 +19,14 @@
         pdp_config,
         wanda_config,
     )
-    from .core.torch import activations, layers, optimizers, pruning_methods, quantizer
+    from .core.torch import (
+        activations,
+        layers,
+        optimizers,
+        pruning_methods,
+        quantizer,
+        tracing,
+    )
     from .core.torch.layers import (
         add_compression_layers,
         apply_final_compression,
@@ -29,9 +36,10 @@
         load_torch_hgq_model,
         post_training_prune,
     )
+    from .core.torch.tracing import check_quantization, print_quantization_check
     from .core.torch.train import train_model
 
-    _forwards = ["activations", "layers", "quantizer", "optimizers"]
+    _forwards = ["activations", "layers", "quantizer", "optimizers", "tracing"]
 
     for name in _forwards:
         mod = importlib.import_module(f".core.torch.{name}", package="pquant")
@@ -57,6 +65,8 @@
     _forwards.append("load_from_dictionary")
     _forwards.append("get_ebops")
     _forwards.append("load_torch_hgq_model")
+    _forwards.append("check_quantization")
+    _forwards.append("print_quantization_check")
     _forwards.append("PQConfig")
     __all__ = _forwards
 

src/pquant/core/torch/convert_to_onnx.py

@@ -53,6 +53,7 @@
     PQLayerNorm,
     PQMultiheadAttention,
 )
+from pquant.core.torch.quantizer import Quantizer  # noqa: E402
 
 # ---------------------------------------------------------------------------
 # QONNX Quant node
@@ -1200,6 +1201,15 @@ def _emit_module(
             module, prefix, current, current, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights
         )
         return out
+    if isinstance(module, Quantizer):
+        # Standalone quantizer (e.g. an auto-inserted missing quantizer or a
+        # constant-matrix quantizer): emit a single QDQ node from its k/i/f.
+        k, i, f = module.get_quantization_bits()
+        new_nodes, out = quant_fn(
+            prefix, current, module.round_mode, k, i, f, initializers, overflow_mode=getattr(module, "overflow", "SAT")
+        )
+        nodes.extend(new_nodes)
+        return out
     raise TypeError(f"Unsupported module type for ONNX export: {type(module).__name__}")
 
 
@@ -1481,6 +1491,8 @@ class _PQTracer(_fx.Tracer):
         PQAvgPool1d,
         PQAvgPool2d,
         PQMultiheadAttention,
+        PQActivation,
+        Quantizer,
     )
 
     def is_leaf_module(self, m: nn.Module, qualname: str) -> bool:
@@ -1516,13 +1528,16 @@ def convert_to_onnx_fx(
     # need to expand torch's two-arg .transpose(d0, d1) into a full ONNX perm.
     from torch.fx.passes.shape_prop import ShapeProp
 
+    # Build the probe tensor on the model's own device so ShapeProp doesn't hit a
+    # device mismatch when a default device (e.g. CUDA) is set via torch.set_default_device.
+    device = next((p.device for p in model.parameters()), None)
     with torch.no_grad():
-        ShapeProp(gm).propagate(torch.zeros(1, *input_shape))
+        ShapeProp(gm).propagate(torch.zeros(1, *input_shape, device=device))
 
     onnx_nodes: list[onnx.NodeProto] = []
     initializers: list[onnx.TensorProto] = []
     node_to_name: dict[_fx.Node, str] = {}
-    output_name: str = ""
+    output_names: list[str] = []
 
     def _res(arg) -> str:
         if isinstance(arg, _fx.Node):
@@ -1680,6 +1695,11 @@ def _resolve_perm_dims(args, rank: int) -> list[int]:
                 onnx_nodes.append(oh.make_node("Relu", inputs=[_res(node.args[0])], outputs=[out]))
                 node_to_name[node] = out
 
+            elif fn in (_F.sigmoid, torch.sigmoid):
+                out = f"{node.name}_sigmoid"
+                onnx_nodes.append(oh.make_node("Sigmoid", inputs=[_res(node.args[0])], outputs=[out]))
+                node_to_name[node] = out
+
             elif fn is torch.flatten:
                 start_dim = node.args[1] if len(node.args) > 1 else node.kwargs.get("start_dim", 0)
                 out = f"{node.name}_flatten"
@@ -1739,28 +1759,29 @@ def _resolve_perm_dims(args, rank: int) -> list[int]:
 
         elif node.op == "output":
             ret = node.args[0]
-            if isinstance(ret, _fx.Node):
-                val = node_to_name[ret]
+            rets = list(ret) if isinstance(ret, (tuple, list)) else [ret]
+            for r in rets:
+                if not isinstance(r, _fx.Node):
+                    raise TypeError("FX ONNX export: unsupported (non-tensor) model output")
+                val = node_to_name[r]
                 # MHA nodes store a tuple (out, avg_attn); expose the attention output.
-                output_name = val[0] if isinstance(val, tuple) else val
-            elif isinstance(ret, (tuple, list)) and len(ret) == 1:
-                val = node_to_name[ret[0]]
-                output_name = val[0] if isinstance(val, tuple) else val
-            else:
-                raise TypeError("Only single-output models are supported for FX ONNX export")
+                output_names.append(val[0] if isinstance(val, tuple) else val)
 
     with torch.no_grad():
-        dummy_out = model(torch.zeros(1, *input_shape))
-    output_shape = [None] + list(dummy_out.shape[1:])
+        dummy_out = model(torch.zeros(1, *input_shape, device=device))
+    dummy_outs = list(dummy_out) if isinstance(dummy_out, (tuple, list)) else [dummy_out]
 
     batch_dim = oh.make_tensor_value_info("input", TensorProto.FLOAT, [None, *input_shape])
-    output_vi = oh.make_tensor_value_info(output_name, TensorProto.FLOAT, output_shape)
+    output_vis = [
+        oh.make_tensor_value_info(name, TensorProto.FLOAT, [None] + list(t.shape[1:]))
+        for name, t in zip(output_names, dummy_outs)
+    ]
 
     onnx_graph = oh.make_graph(
         nodes=onnx_nodes,
         name="pquant_onnx_fx",
         inputs=[batch_dim],
-        outputs=[output_vi],
+        outputs=output_vis,
         initializer=initializers,
     )
 

src/pquant/core/torch/layers.py

@@ -644,9 +644,15 @@ def extra_repr(self):
         return s.format(**self.__dict__)
 
 
-def add_compression_layers(model, config, input_shape=None):
+def add_compression_layers(model, config, input_shape=None, add_missing_quantizers=False):
     model = add_quantized_activations_to_model_layer(model, config)
     model = add_pruning_to_model(model, config)
+    if add_missing_quantizers:
+        # Imported here (not at module top) to avoid a circular import: tracing.py
+        # imports the layer classes defined in this module.
+        from pquant.core.torch.tracing import check_quantization
+
+        model = check_quantization(model, add_missing_quantizers=True, config=config)
     model.to("cuda")
     if input_shape is not None:
         model(torch.rand(input_shape).to("cuda"))