[Tracing] Update ignored functions list (#1599)

## Purpose ##
* Fix tracing for model definitions introduced as part of
`transformers==4.53`
* Resolves #1603

## Background ##
In the latest transformers release, this change landed which changed the
name of the function which generates the causal mask.
huggingface/transformers#37866

## Changes ##
* Extend the list of function names to ignore during tracing,
specifically targeting functions which create causal masks
* Update debugger tool to use ignore list from `DatasetArguments`
* Update Tracer to skip masking function as part of autowrapping any
functions which were not caught by the autowrapper

## Testing ##
* `tests/llmcompressor/transformers/tracing/test_models.py` now passes
with the latest `transformers==4.53`

---------

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

Author: kylesayrs

src/llmcompressor/args/dataset_arguments.py

@@ -181,7 +181,17 @@ class DatasetArguments(CustomDatasetArguments):
         },
     )
     tracing_ignore: List[str] = field(
-        default_factory=lambda: ["_update_causal_mask"],
+        default_factory=lambda: [
+            "_update_causal_mask",
+            "create_causal_mask",
+            "make_causal_mask",
+            "get_causal_mask",
+            "mask_interface",
+            "mask_function",
+            "_prepare_4d_causal_attention_mask",
+            "_prepare_fsmt_decoder_inputs",
+            "_prepare_4d_causal_attention_mask_with_cache_position",
+        ],
         metadata={
             "help": "List of functions to ignore during tracing, either "
             "{module}.{method_name} or {function_name}"

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

@@ -2,7 +2,7 @@
 import inspect
 from collections import deque
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple
 
 import torch
 from compressed_tensors.quantization import find_name_or_class_matches
@@ -169,10 +169,12 @@ 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
+        super().__init__(autowrap_functions=_get_autowrap_functions())
+
         # check unlikely case that ancestors have direct params which are offloaded
         offloaded_ancestors = offloaded & ancestors
         if offloaded_ancestors:
@@ -531,3 +533,12 @@ def dispatch_for_sequential(model: PreTrainedModel) -> PreTrainedModel:
         logger.warning("CUDA is not available! Compressing model on CPU instead")
 
     return model
+
+
+def _get_autowrap_functions() -> Tuple[Callable[[Any], Any], ...]:
+    try:
+        from transformers.masking_utils import LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING
+
+        return tuple(LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING.values())
+    except ImportError:
+        return tuple()

src/llmcompressor/transformers/tracing/debug.py

@@ -22,7 +22,7 @@ def parse_args():
     parser.add_argument("--model_id", type=str, required=True, help="The stub of the model to load")  # noqa: E501
     parser.add_argument("--model_class", type=str, required=True, help="The class name of the model")  # noqa: E501
     parser.add_argument("--sequential_targets", type=str, nargs="*", default=None, metavar="TARGET", help="List of targets for sequential tracing")  # noqa: E501
-    parser.add_argument("--ignore", type=str, nargs="*", default=["_update_causal_mask"], metavar="PATTERN", help="List of patterns to ignore during tracing")  # noqa: E501
+    parser.add_argument("--ignore", type=str, nargs="*", default=DatasetArguments().tracing_ignore, metavar="PATTERN", help="List of patterns to ignore during tracing")  # noqa: E501
     parser.add_argument("--modality", type=str, default="text", help="Modality of calibration dataset, defaults to text")  # noqa: E501
     parser.add_argument("--trust_remote_code", type=bool, default=False, help="Whether to trust model remote code")  # noqa: E501
     parser.add_argument("--skip_weights", type=bool, default=True, help="Whether to load the model with dummy weights")  # noqa: E501
@@ -34,7 +34,7 @@ def trace(
     model_id: str,
     model_class: Type[PreTrainedModel],
     sequential_targets: Optional[Union[List[str], str]] = None,
-    ignore: Union[List[str], str] = ["_update_causal_mask"],
+    ignore: Union[List[str], str] = DatasetArguments().tracing_ignore,
     modality: str = "text",
     trust_remote_code: bool = True,
     skip_weights: bool = True,