[Tracing] Support tracing of Gemma3 [#1248]

src/llmcompressor/pipelines/sequential/helpers.py

@@ -1,11 +1,12 @@
 import inspect
 from collections import deque
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Set, Union
+from typing import Any, Callable, Dict, List, Optional, Set, Union
 
 from compressed_tensors import has_offloaded_params
 from compressed_tensors.quantization import find_name_or_class_matches
 from torch.fx import Graph, GraphModule, Node
+from torch.fx.graph import PythonCode
 from torch.fx.proxy import Argument
 from torch.nn import Module
 from transformers import PreTrainedModel
@@ -32,16 +33,33 @@ class Subgraph:
     graph: Graph
     input_names: Set[str]
     consumed_names: Set[str]
+    _code: Optional[PythonCode] = None
 
-    def compile_forward(self) -> Callable[[Any], Any]:
+    def forward(self, *args, **kwargs) -> Dict[str, Any]:
         """
-        Generate and compile code for executing this subgraph
+        Execute the operations within the subgraph
 
-        :return: function which, when called, executes this subgraph
+        :param \\*args: argument inputs to subgraph forward function
+        :param \\**kwargs: keyword inputs to subgraph forward function
+        :return keyword outputs of subgraph forward function (non-consumed variables):
         """
-        code = self.graph.python_code("self")
-        exec(code.src, code.globals)
-        return code.globals.get("forward")
+        if self._code is None:
+            self._code = self.graph.python_code("self")
+            exec(self._code.src, self._code.globals)
+
+        forward_fn = self._code.globals.get("forward")
+
+        try:
+            outputs = forward_fn(*args, **kwargs)
+        except Exception as exception:
+            raise RuntimeError(
+                "Raised an exception during execution of the following code:\n"
+                f"```\n{add_line_numbers(self._code.src)}\n```\n"
+                "This is likely due to a violation of shape assumptions made when "
+                "tracing"
+            ) from exception
+
+        return outputs
 
 
 def trace_subgraphs(
@@ -376,3 +394,9 @@ def match_modules(model: Module, target_names: List[str]) -> Set[Module]:
         for name, module in model.named_modules()
         if find_name_or_class_matches(name, module, target_names)
     )
+
+
+def add_line_numbers(text: str) -> str:
+    lines = text.splitlines()
+    numbered_lines = [f"{i + 1} {line}" for i, line in enumerate(lines)]
+    return "\n".join(numbered_lines)

src/llmcompressor/pipelines/sequential/pipeline.py

@@ -61,13 +61,10 @@ def run_pipeline(
             calib_desc = f"({subgraph_index + 1}/{num_subgraphs}): Calibrating"
             prop_desc = f"({subgraph_index + 1}/{num_subgraphs}): Propagating"
 
-            # compile subgraph forward function
-            forward_function = subgraph.compile_forward()
-
             # do an preliminary pass to trigger modifier hooks
             for batch_index in tqdm.tqdm(range(len(dataloader)), desc=calib_desc):
                 inputs = intermediates.fetch(batch_index, subgraph.input_names)
-                forward_function(model, **inputs)
+                subgraph.forward(model, **inputs)
 
             # TODO: replace with a lifecycle event
             if callback_modifier:
@@ -78,7 +75,7 @@ def run_pipeline(
             with HooksMixin.disable_hooks():
                 for batch_index in tqdm.tqdm(range(len(dataloader)), desc=prop_desc):
                     inputs = intermediates.fetch(batch_index, subgraph.input_names)
-                    output = forward_function(model, **inputs)
+                    output = subgraph.forward(model, **inputs)
 
                     if subgraph_index < num_subgraphs - 1:
                         intermediates.update(batch_index, output)

src/llmcompressor/transformers/tracing/__init__.py

@@ -1,3 +1,6 @@
+from .gemma3 import (
+    Gemma3ForConditionalGeneration as TraceableGemma3ForConditionalGeneration,
+)
 from .llava import (
     LlavaForConditionalGeneration as TraceableLlavaForConditionalGeneration,
 )
@@ -11,12 +14,13 @@
     Idefics3ForConditionalGeneration as TraceableIdefics3ForConditionalGeneration,
 )
 from .qwen2_5_vl import (
-    Qwen2_5_VLForConditionalGeneration as TraceableQwen2_5_VLForConditionalGeneration
+    Qwen2_5_VLForConditionalGeneration as TraceableQwen2_5_VLForConditionalGeneration,
 )
 from .debug import get_model_class
 
 __all__ = [
     "get_model_class",
+    "TraceableGemma3ForConditionalGeneration",
     "TraceableLlavaForConditionalGeneration",
     "TraceableMllamaForConditionalGeneration",
     "TraceableQwen2VLForConditionalGeneration",