skip mask functions in tracing

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/llmcompressor/pipelines/sequential/ast_utils/auto_wrapper.py

@@ -180,15 +180,9 @@ def _wrap_if_possible(self, node: ast.AST) -> Union[ast.AST, ast.Assign, ast.Cal
             return node
 
         if isinstance(node, ast.stmt):
-            logger.debug("---- Autowrapper ----")
-            logger.debug(ast.unparse(node))
-            logger.debug("---------------------")
             return self._wrap_stmt(node)
 
         elif isinstance(node, ast.expr):
-            logger.debug("---- Autowrapper ----")
-            logger.debug(ast.unparse(node))
-            logger.debug("---------------------")
             return self._wrap_expr(node)
 
         else:
@@ -254,6 +248,11 @@ def _wrap_stmt(self, node: ast.stmt) -> ast.Assign:
         # update local names with newly returned values
         self._local_names |= returns
 
+        # log newly created function definition
+        logger.debug("---- Autowrapper ----")
+        logger.debug(ast.unparse(ast.fix_missing_locations(fn_def)))
+        logger.debug("---------------------")
+
         return assign_call
 
     def _wrap_expr(self, node: ast.expr) -> ast.Call:

src/llmcompressor/pipelines/sequential/helpers.py

@@ -18,6 +18,7 @@
 from torch.nn import Module
 from transformers import PreTrainedModel
 from transformers.configuration_utils import PretrainedConfig
+from transformers.masking_utils import LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING
 from transformers.utils.fx import HFTracer
 
 from llmcompressor.modifiers import Modifier
@@ -169,10 +170,13 @@ class SequentialTracer(HFTracer):
     """
 
     def __init__(self, ancestors: Set[Module], offloaded: Set[Module]):
-        super().__init__()
         self.ancestors = ancestors
         self.offloaded = offloaded
 
+        # skip any mask creation functions not already caught by the autowrapper
+        autowrap_functions = tuple(LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING.values())
+        super().__init__(autowrap_functions=autowrap_functions)
+
         # check unlikely case that ancestors have direct params which are offloaded
         offloaded_ancestors = offloaded & ancestors
         if offloaded_ancestors: