commit

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/llmcompressor/modifiers/awq/mappings.py

@@ -133,18 +133,9 @@ class AWQMapping:
 ]
 
 _bloom_mappings = [
-    AWQMapping(
-        "re:.*input_layernorm$",
-        ["re:.*query_key_value$"]
-    ),
-    AWQMapping(
-        "re:.*post_attention_layernorm$",
-        ["re:.*dense_h_to_4h$"]
-    ),
-    AWQMapping(
-        "re:.*gelu_impl$",
-        ["re:.*dense_4h_to_h$"]
-    ),
+    AWQMapping("re:.*input_layernorm$", ["re:.*query_key_value$"]),
+    AWQMapping("re:.*post_attention_layernorm$", ["re:.*dense_h_to_4h$"]),
+    AWQMapping("re:.*gelu_impl$", ["re:.*dense_4h_to_h$"]),
     # Note: AutoAWQ excludes this mapping, based on researcher's post in
     # https://github.com/mit-han-lab/llm-awq/issues/2#issuecomment-1606297469
     # AWQMapping(

src/llmcompressor/pipelines/sequential/helpers.py

@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple
 
 import torch
+from accelerate.hooks import remove_hook_from_module
 from compressed_tensors.quantization import find_name_or_class_matches
 from compressed_tensors.utils import (
     has_offloaded_params,
@@ -177,13 +178,12 @@ def __init__(self, ancestors: Set[Module], offloaded: Set[Module]):
 
         # check unlikely case that ancestors have direct params which are offloaded
         offloaded_ancestors = offloaded & ancestors
-        if offloaded_ancestors:
-            names = set(module.__class__.__name__ for module in offloaded_ancestors)
+        for ancestor in offloaded_ancestors:
+            remove_hook_from_module(ancestor, recurse=False)
+            self.offloaded.remove(ancestor)
             logger.warning(
-                "The following modules are call graph ancestors of sequential targets,"
-                f"but also contain offloaded modules: {names}.\n"
-                "These modules will not be traced, and any sequential target children "
-                "will be executed jointly, which may lead to OOM errors"
+                f"Direct parameters attached to {ancestor.__class__.__name__} have "
+                "been onloaded in order to ensure safe graph capture and execution"
             )
 
     def create_arg(self, a: Any) -> Argument:

tests/llmcompressor/transformers/tracing/test_models.py

@@ -42,7 +42,13 @@
             "text",
             [],
         ),
-        ("RedHatAI/DeepSeek-V3-BF16", AutoModelForCausalLM, None, "text", []),
+        (
+            "ibm-granite/granite-20b-code-instruct-8k",
+            AutoModelForCausalLM,
+            None,
+            "text",
+            [],
+        )("RedHatAI/DeepSeek-V3-BF16", AutoModelForCausalLM, None, "text", []),
         # --- vision ---
         (
             "HuggingFaceM4/Idefics3-8B-Llama3",