add missing norm fusion

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

Author: kylesayrs

examples/transform/spinquant_dummy.py

@@ -1,14 +1,13 @@
-from datasets import load_dataset
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 from compressed_tensors.utils import update_parameter_data
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+
 from llmcompressor import oneshot
 from llmcompressor.modifiers.quantization import GPTQModifier, QuantizationModifier
 from llmcompressor.modifiers.transform import SpinQuantModifier
 from llmcompressor.utils import dispatch_for_generation
-from transformers.models.llama.modeling_llama import (
-    LlamaRMSNorm,
-)
 
 hidden_dim = intermediate_dim = 64
 up_dim = 128

src/llmcompressor/modeling/__init__.py

@@ -1,4 +1,4 @@
 # flake8: noqa
 
+from .fuse import *
 from .prepare import *
-from .fuse import *

src/llmcompressor/modeling/fuse.py

@@ -1,8 +1,11 @@
 from typing import Iterable
 
 import torch
-from compressed_tensors import get_execution_device, align_module_device, update_offload_parameter
-
+from compressed_tensors import (
+    align_module_device,
+    get_execution_device,
+    update_offload_parameter,
+)
 from transformers.models.llama.modeling_llama import LlamaRMSNorm
 
 __all__ = ["fuse_norm_linears"]
@@ -22,14 +25,17 @@ def fuse_norm_linears(norm: torch.nn.Module, linears: Iterable[torch.nn.Linear])
         for linear in linears:
             # NOTE: spinquant does this op in float64
             exec_device = get_execution_device(norm)
-            with align_module_device(norm, exec_device), align_module_device(linear, exec_device):
-                
+            with align_module_device(norm, exec_device), align_module_device(
+                linear, exec_device
+            ):
                 weight_dtype = linear.weight.dtype
 
-                new_weight = linear.weight.to(torch.float64) * norm.weight.to(torch.float64)
+                new_weight = linear.weight.to(torch.float64) * norm.weight.to(
+                    torch.float64
+                )
 
                 new_weight = new_weight.to(weight_dtype)
-            
+
             update_offload_parameter(linear, "weight", new_weight)
 
         update_offload_parameter(norm, "weight", torch.ones_like(norm.weight))

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

@@ -1 +1 @@
-from .base import *
+from .base import *

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

@@ -1,14 +1,18 @@
-from typing import Optional, List, Literal, Iterable
+from enum import Enum
+from typing import Iterable, List, Literal, Optional
 
-from compressed_tensors.transform import TransformConfig, TransformScheme, TransformArgs, apply_transform_config
-from pydantic import BaseModel, field_validator, Field
+from compressed_tensors.transform import (
+    TransformArgs,
+    TransformConfig,
+    TransformScheme,
+    apply_transform_config,
+)
+from pydantic import BaseModel, Field, field_validator
+from transformers import PreTrainedModel
 
 from llmcompressor.core import Event, EventType, State
 from llmcompressor.modeling import fuse_norm_linears
 from llmcompressor.modifiers import Modifier
-from enum import Enum
-
-from transformers import PreTrainedModel
 
 
 class SpinQuantMappings(BaseModel):
@@ -29,9 +33,10 @@ class SpinQuantMappings(BaseModel):
     def cast_to_list(cls, value):
         if isinstance(value, str):
             return [value]
-        
+
         return value
-    
+
+
 class NormMapping(BaseModel):
     norm: str
     linears: List[str]
@@ -40,22 +45,18 @@ class NormMapping(BaseModel):
     def cast_to_list(cls, value):
         if isinstance(value, str):
             return [value]
-        
-        return value
 
+        return value
 
 
 llama_spinquant = SpinQuantMappings(
     embedding="re:.*embed_tokens$",
-
     attn_q="re:.*q_proj$",
     attn_k="re:.*k_proj$",
     attn_v="re:.*v_proj$",
     attn_o="re:.*o_proj$",
-
     mlp_in=["re:.*up_proj$", "re:.*gate_proj$"],
     mlp_out="re:.*down_proj$",
-
     lm_head="lm_head",
 )
 
@@ -67,25 +68,31 @@ def cast_to_list(cls, value):
     NormMapping(
         norm="re:.*post_attention_layernorm$",
         linears=["re:.*up_proj$", "re:.*gate_proj$"],
-    )
+    ),
 ]
 
+
 class SpinquantRotation(Enum):
     R1 = "R1"
     R2 = "R2"
     R3 = "R3"
     R4 = "R4"
 
+
 class SpinQuantModifier(Modifier):
     rotations: Iterable[SpinquantRotation] = ("R1", "R2")
-    transform_type: Literal["hadamard", "random-hadamard", "random-matrix"] = Field(default="hadamard")
+    transform_type: Literal["hadamard", "random-hadamard", "random-matrix"] = Field(
+        default="hadamard"
+    )
     randomize: bool = Field(default=False)
     learnable: bool = Field(default=False)
 
     mappings: Optional[SpinQuantMappings] = None
     norm_mappings: Optional[List[NormMapping]] = None
-    
-    transform_config: Optional[TransformConfig] = None  # optional override for more fine-grained control
+
+    transform_config: Optional[TransformConfig] = (
+        None  # optional override for more fine-grained control
+    )
 
     @field_validator("rotations", mode="before")
     def validate_rotations(cls, value):
@@ -101,7 +108,7 @@ def on_initialize(self, state: State, **kwargs) -> bool:
         if self.transform_config is not None:
             if self.mappings is not None:
                 raise ValueError()
-            
+
             return True
 
         config_groups = {}
@@ -129,6 +136,7 @@ def on_start(self, state: State, event: Event, **kwargs):
         # Embedding fusion
         # theoretically, doesn't do anything. Doesn't seem to help model sanity either
         from compressed_tensors import update_offload_parameter
+
         for W in [state.model.model.embed_tokens]:
             W_ = W.weight.data.double()
             W.weight.data = (W_ - W_.mean(dim=-1, keepdim=True)).to(W.weight.data.dtype)
@@ -138,16 +146,24 @@ def on_start(self, state: State, event: Event, **kwargs):
         # TODO: use norm mappings
         # layer norm fusion
         for layer in state.model.model.layers:
-            fuse_norm_linears(layer.input_layernorm, (layer.self_attn.q_proj, layer.self_attn.k_proj, layer.self_attn.v_proj))
-            fuse_norm_linears(layer.post_attention_layernorm, (layer.mlp.gate_proj, layer.mlp.up_proj))
+            fuse_norm_linears(
+                layer.input_layernorm,
+                (
+                    layer.self_attn.q_proj,
+                    layer.self_attn.k_proj,
+                    layer.self_attn.v_proj,
+                ),
+            )
+            fuse_norm_linears(
+                layer.post_attention_layernorm, (layer.mlp.gate_proj, layer.mlp.up_proj)
+            )
+
+        fuse_norm_linears(state.model.model.norm, (state.model.lm_head,))
 
         # needs to happen after the model has been hooked to execute on the GPU
         # otherwise we're applying weight transforms on CPU
         apply_transform_config(state.model, self.transform_config)
 
-
-
-
     def on_event(self, state: State, event: Event, **kwargs):
         if event.type_ == EventType.CALIBRATION_EPOCH_START:
             if not self.started_:
@@ -169,7 +185,6 @@ def on_finalize(self, state: State, **kwargs) -> bool:
 
         return True
 
-
     def _create_r1_scheme(self) -> TransformScheme:
         return TransformScheme(
             type=self.transform_type,
@@ -190,14 +205,14 @@ def _create_r1_scheme(self) -> TransformScheme:
                         self.mappings.attn_k,
                         self.mappings.attn_v,
                         *self.mappings.mlp_in,
-                        self.mappings.lm_head
+                        self.mappings.lm_head,
                     ],
                     location="weight_input",
                     inverse=True,
                 ),
-            ]
+            ],
         )
-    
+
     def _create_r2_scheme(self, model: PreTrainedModel) -> TransformScheme:
         config = model.config
 
@@ -207,7 +222,7 @@ def _create_r2_scheme(self, model: PreTrainedModel) -> TransformScheme:
             head_dim = config.hidden_size // config.num_attention_heads
         else:
             raise NotImplementedError()
-        
+
         return TransformScheme(
             type=self.transform_type,
             randomize=self.randomize,
@@ -223,10 +238,8 @@ def _create_r2_scheme(self, model: PreTrainedModel) -> TransformScheme:
             ],
         )
 
-
     def _create_r3_scheme(self) -> TransformScheme:
         raise NotImplementedError()
 
-
     def _create_r4_scheme(self) -> TransformScheme:
-        raise NotImplementedError()
+        raise NotImplementedError()