WIP: Use Dynamo Tracer

Author: Victor-Jung

DeepQuant/Pipeline/Injection.py

@@ -16,7 +16,7 @@
     MHATransformation,
 )
 from DeepQuant.Utils.ConsoleFormatter import ConsoleColor as cc
-from DeepQuant.Utils.CustomTracer import QuantTracer, customBrevitasTrace
+from DeepQuant.Utils.CustomTracer import QuantTracer
 from DeepQuant.Utils.GraphPrinter import GraphModulePrinter
 
 
@@ -40,14 +40,17 @@ def injectCustomForwards(
     executor = TransformationExecutor(transformations, debug=debug, tracer=tracer)
     transformedModel = executor.execute(model, exampleInput)
 
-    fxModel = customBrevitasTrace(
-        root=transformedModel,
-        tracer=tracer,
-    )
-    fxModel.recompile()
+    # fxModel = customBrevitasTrace(
+    #     root=transformedModel,
+    #     tracer=tracer,
+    # )
+    # fxModel.recompile()
+    import IPython; IPython.embed()
+    fxModel = tracer.trace(transformedModel, exampleInput)
+    import IPython; IPython.embed()
 
-    with torch.no_grad():
-        output = fxModel(exampleInput)
+    # with torch.no_grad():
+    output = fxModel(exampleInput)
 
     if torch.allclose(referenceOutput, output, atol=1e-5):
         if debug:

DeepQuant/Pipeline/OriginalTracing.py

@@ -14,14 +14,44 @@
 from DeepQuant.Utils.ConsoleFormatter import ConsoleColor as cc
 from DeepQuant.Utils.GraphPrinter import GraphModulePrinter
 
+from torch._dynamo import allow_in_graph
 
 def traceOriginalModel(
     model: nn.Module, exampleInput: torch.Tensor, debug: bool = False
 ) -> Tuple[nn.Module, torch.Tensor]:
     """Symbolically trace the original model using Brevitas."""
     printer = GraphModulePrinter()
 
-    tracedModel = brevitas_symbolic_trace(model)
+    # tracedModel = brevitas_symbolic_trace(model)
+    graphs = []
+
+    def dynamo_graph_extract_compiler(gm, inputs: torch.Tensor):
+        graphs.append(gm)
+        return gm.forward
+    
+
+    torch._dynamo.reset()
+    torch._dynamo.config.verbose = True
+
+    allow_in_graph(model.inputQuant)
+    allow_in_graph(model.inputQuant.__class__)
+    allow_in_graph(model.inputQuant.forward)
+
+    allow_in_graph(model.inputQuant)
+    allow_in_graph(model.linear1.__class__)
+    allow_in_graph(model.linear1.forward)
+    # JUNGVI: For Philip, dynamo uses the id of the thing passed in allow_in_graph to filter them. But it does not seems to work at least for brevitas layers, IDK if they have smth special...
+
+    import IPython; IPython.embed()
+
+    model_fn = torch.compile(model, backend = dynamo_graph_extract_compiler, dynamic = False)
+    
+    with torch.no_grad():
+        _ = model_fn(exampleInput)
+
+    import IPython; IPython.embed()
+
+    tracedModel = graphs[0]
 
     if debug:
         print(cc.header("1. Original Network"))

DeepQuant/QuantManipulation/QuantNodesDivider.py

@@ -43,8 +43,9 @@ def insertQuantDequantPair(
         Dequant(originalModule, scaleVal, zpVal, bwVal, signed=signedVal),
     )
 
-    with graph.inserting_after(node):
-        quantNode = graph.call_module(quantName, args=(mainArg,))
+    with fxModel.graph.inserting_after(node):
+        import IPython; IPython.embed()
+        quantNode = fxModel.graph.call_module(quantName, args=(mainArg,))
 
     with graph.inserting_after(quantNode):
         dequantNode = graph.call_module(dequantName, args=(quantNode,))

DeepQuant/Utils/CustomTracer.py

@@ -4,18 +4,22 @@
 #
 # Federico Brancasi <[email protected]>
 
-from typing import List, Optional, Type
+from functools import partial
+from typing import List, Optional, Type, Callable
 
+import torch
 import torch.nn as nn
+from torch._dynamo import allow_in_graph
+from torch.fx.graph_module import GraphModule
+
 from brevitas.fx.brevitas_tracer import (
     Tracer,
     _is_brevitas_leaf_module,
     _symbolic_trace,
 )
-from torch.fx.graph_module import GraphModule
 
 
-class QuantTracer(Tracer):
+class QuantTracer():
     """Enhanced tracer with fine-grained control over module tracing."""
 
     def __init__(
@@ -24,34 +28,58 @@ def __init__(
         nonLeafClasses: Optional[List[Type[nn.Module]]] = None,
         debug: bool = False,
     ) -> None:
-        super().__init__()
         self.leafClasses = leafClasses if leafClasses is not None else []
         self.nonLeafClasses = nonLeafClasses if nonLeafClasses is not None else []
         self.debug = debug
 
     def registerLeafModule(self, moduleCls: Type[nn.Module]) -> None:
-        """Register a module class as a leaf module."""
         if moduleCls not in self.leafClasses:
             self.leafClasses.append(moduleCls)
 
     def registerNonLeafModule(self, moduleCls: Type[nn.Module]) -> None:
-        """Register a module class as a non-leaf module."""
         if moduleCls not in self.nonLeafClasses:
             self.nonLeafClasses.append(moduleCls)
 
-    def is_leaf_module(self, m: nn.Module, moduleQualifiedName: str) -> bool:
-        """Determine if a module should be treated as a leaf module."""
-        if any(isinstance(m, lc) for lc in self.leafClasses):
-            return True
-        if any(isinstance(m, nlc) for nlc in self.nonLeafClasses):
-            return False
-        return _is_brevitas_leaf_module(m, moduleQualifiedName)
-
-
-def customBrevitasTrace(
-    root: nn.Module, concreteArgs=None, tracer: Optional[QuantTracer] = None
-) -> GraphModule:
-    """Create an FX GraphModule using the QuantTracer (a custom Brevitas tracer)."""
-    if tracer is None:
-        tracer = QuantTracer()
-    return _symbolic_trace(tracer, root, concreteArgs)
+    def trace(self, model: nn.Module, exampleInput):
+
+        brevitasClasses = [(m, id(m.__class__)) for _, m in model.named_modules() if m.__module__.startswith('brevitas.nn') or m.__module__.startswith('brevitas.core') or m.__module__.startswith('brevitas.proxy')]
+
+        # leafClasses = (set(brevitasClasses) | set(self.leafClasses)) - set(self.nonLeafClasses)
+        leafClasses = brevitasClasses
+        graphs: List[torch.fx.GraphModule] = []
+
+        def dynamo_graph_extract_compiler(gm: GraphModule, inputs: torch.Tensor) -> Callable:
+            graphs.append(gm)
+            return gm.forward
+        
+        torch._dynamo.reset()
+        torch._dynamo.config.verbose = True
+
+        # for op, _ in leafClasses:
+        #     allow_in_graph(op.forward)
+
+        # for _, m in model.named_modules():
+        #     if m.__module__.startswith('brevitas.nn') or m.__module__.startswith('brevitas.core') or m.__module__.startswith('brevitas.proxy'):
+        #        allow_in_graph(m) 
+
+        allow_in_graph(model.inputQuant.__class__)
+        allow_in_graph(model.inputQuant.forward)
+
+        allow_in_graph(model.linear1.__class__)
+        allow_in_graph(model.linear1.forward)
+
+        model_fn = torch.compile(model, backend = dynamo_graph_extract_compiler, dynamic = False)
+        from brevitas.export.inference import quant_inference_mode
+        with torch.no_grad(), quant_inference_mode(model_fn):
+            _ = model_fn(exampleInput)
+
+        import IPython; IPython.embed()
+        return graphs[0]
+
+# def customBrevitasTrace(
+#     root: nn.Module, concreteArgs=None, tracer: Optional[QuantTracer] = None
+# ) -> GraphModule:
+#     """Create an FX GraphModule using the QuantTracer (a custom Brevitas tracer)."""
+#     if tracer is None:
+#         tracer = QuantTracer()
+#     return _symbolic_trace(tracer, root, concreteArgs)