better registery logic

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

Author: kylesayrs

src/llmcompressor/modifiers/transform/spinquant/base.py

@@ -15,8 +15,8 @@
 from llmcompressor.modeling import fuse_norm_linears, normalize_embedding
 from llmcompressor.modifiers import Modifier
 
-from .mappings import SPINQUANT_MAPPING_REGISTRY, SpinQuantMappings
-from .norm_mappings import NORM_MAPPING_REGISTRY, NormMapping
+from .mappings import SpinQuantMapping, infer_mapping_from_model
+from .norm_mappings import NormMapping, infer_norm_mapping_from_model
 
 
 class SpinquantRotation(str, Enum):
@@ -36,9 +36,7 @@ class SpinQuantModifier(Modifier, use_enum_values=True):
 
     # norm mappings separate from spinquant mappings to allow users to
     # override spinquant mappings with transform_config without overriding norms
-    # we can combine these mappings, but it requires some more validation logic
-    # maybe there's a reason to keep if other modifiers want norm fusing, idk
-    mappings: Optional[SpinQuantMappings] = Field(default=None, exclude=True)
+    mappings: Optional[SpinQuantMapping] = Field(default=None, exclude=True)
     norm_mappings: Optional[List[NormMapping]] = Field(default=None, exclude=True)
 
     # optional override for more fine-grained control
@@ -53,8 +51,8 @@ def validate_rotations(cls, value):
 
     def on_initialize(self, state: State, **kwargs) -> bool:
         # TODO: more validation
-        self.mappings = SPINQUANT_MAPPING_REGISTRY[state.model.__class__.__name__]
-        self.norm_mappings = NORM_MAPPING_REGISTRY[state.model.__class__.__name__]
+        self.mappings = infer_mapping_from_model(state.model)
+        self.norm_mappings = infer_norm_mapping_from_model(state.model)
 
         if self.transform_config is not None:
             if self.mappings is not None:

src/llmcompressor/modifiers/transform/spinquant/mappings.py

@@ -1,9 +1,13 @@
 from typing import Dict, List, Optional
 
+from loguru import logger
 from pydantic import BaseModel, Field, field_validator
+from transformers import PreTrainedModel
 
+__all__ = ["SpinQuantMapping", "infer_mapping_from_model"]
 
-class SpinQuantMappings(BaseModel):
+
+class SpinQuantMapping(BaseModel):
     embedding: str
 
     attn_q: str
@@ -25,7 +29,7 @@ def cast_to_list(cls, value):
         return value
 
 
-_default_mappings = SpinQuantMappings(
+_default_mappings = SpinQuantMapping(
     embedding="re:.*embed_tokens$",
     attn_q="re:.*q_proj$",
     attn_k="re:.*k_proj$",
@@ -37,6 +41,17 @@ def cast_to_list(cls, value):
 )
 
 
-SPINQUANT_MAPPING_REGISTRY: Dict[str, SpinQuantMappings] = {
+SPINQUANT_MAPPING_REGISTRY: Dict[str, SpinQuantMapping] = {
     "LlamaForCausalLM": _default_mappings,
 }
+
+
+def infer_mapping_from_model(model: PreTrainedModel) -> SpinQuantMapping:
+    architecture = model.__class__.__name__
+    if architecture not in SPINQUANT_MAPPING_REGISTRY:
+        logger.info(
+            f"Unrecognized model architecture {architecture}. "
+            "Falling back to default mappings"
+        )
+
+    return SPINQUANT_MAPPING_REGISTRY.get(architecture, _default_mappings)

src/llmcompressor/modifiers/transform/spinquant/norm_mappings.py

@@ -1,6 +1,10 @@
 from typing import Dict, List
 
+from loguru import logger
 from pydantic import BaseModel, field_validator
+from transformers import PreTrainedModel
+
+__all__ = ["infer_norm_mapping_from_model"]
 
 
 class NormMapping(BaseModel):
@@ -15,7 +19,7 @@ def cast_to_list(cls, value):
         return value
 
 
-_default_norm_mappings = [
+_default_mappings = [
     NormMapping(
         norm="re:.*input_layernorm$",
         linears=["re:.*q_proj$", "re:.*k_proj$", "re:.*v_proj$"],
@@ -31,5 +35,16 @@ def cast_to_list(cls, value):
 ]
 
 NORM_MAPPING_REGISTRY: Dict[str, NormMapping] = {
-    "LlamaForCausalLM": _default_norm_mappings,
+    "LlamaForCausalLM": _default_mappings,
 }
+
+
+def infer_norm_mapping_from_model(model: PreTrainedModel) -> List[NormMapping]:
+    architecture = model.__class__.__name__
+    if architecture not in NORM_MAPPING_REGISTRY:
+        logger.info(
+            f"Unrecognized model architecture {architecture}. "
+            "Falling back to default mappings"
+        )
+
+    return NORM_MAPPING_REGISTRY.get(architecture, _default_mappings)

src/llmcompressor/transformers/sparsification/compressed_tensors_utils.py

@@ -14,7 +14,6 @@
     register_offload_parameter,
 )
 from loguru import logger
-from safetensors.torch import storage_ptr
 from transformers import PreTrainedModel
 
 from llmcompressor.core import active_session