Fix (equalize): Fix LayerwiseActivationRotation (#1413)

Author: pablomlago

src/brevitas/graph/equalize.py

@@ -1774,6 +1774,17 @@ def find_module_by_name(self, model: nn.Module, regions: List[Region], prefix: s
                 full_name = prefix + '.' + name if prefix != '' else name
                 self.find_module_by_name(module, regions, full_name)
 
+    def transform_model(
+            self, model: nn.Module, rewriters: List[Transform], delay_rewriters: bool) -> nn.Module:
+        # In some circumstances, it might be useful to apply model transformations at a later moment
+        # The user should not be resposible for this in any case
+        if delay_rewriters:
+            return model
+        if is_model_offloaded_accelerate(model):
+            return apply_rewriters_accelerate(model, rewriters)
+        else:
+            return apply_rewriters(model, rewriters)
+
 
 class GraphRotationEqualization(RotationEqualization):
 
@@ -2005,16 +2016,6 @@ def apply(self,
         else:
             return graph_model
 
-    def transform_model(self, model, rewriters, delay_rewriters):
-        # In some circumstances, it might be useful to apply model transformations at a later moment
-        # The user should not be resposible for this in any case
-        if delay_rewriters:
-            return model
-        if is_model_offloaded_accelerate(model):
-            return apply_rewriters_accelerate(model, rewriters)
-        else:
-            return apply_rewriters(model, rewriters)
-
 
 @torch.no_grad()
 def apply_rewriters(
@@ -2114,15 +2115,17 @@ def __init__(
         self.supported_sinks = (nn.Linear)
 
     def apply(self, model: nn.Module) -> nn.Module:
+        regions: List[Region] = []
+        rewriters: List[Transform] = []
 
         blacklist_orphan_layers = self.blacklist_layers + self.layers_to_expand
-        regions: List[Region] = []
         self.find_module(model, regions, blacklist_layers=blacklist_orphan_layers)
         expanded_regions = []
         self.find_module_by_name(model, expanded_regions)
 
         if len(expanded_regions) > 0:
             regions.extend(expanded_regions)
         if len(regions) > 0:
-            _compute_rotations(model, regions, expansion_step=self.expansion_step)
+            rewriters.extend(_compute_rotations(model, regions, expansion_step=self.expansion_step))
+        model = self.transform_model(model, rewriters, delay_rewriters=False)
         return model

src/brevitas_examples/llm/main.py

@@ -335,9 +335,11 @@ def quantize_llm(args, extra_args=None):
         model = eq.apply(model)
         remove_hooks(model)
     elif args.rotation == 'layerwise':
+        model = offload_model(model)
         eq = LayerwiseActivationRotation(
             layers_to_expand=layers_to_expand, expansion_step=args.expansion_step)
         model = eq.apply(model)
+        remove_hooks(model)
     elif args.rotation == 'fused_no_fx':
         fused_rotation_no_fx(model, calibration_loader, args)
 

tests/brevitas_examples/test_llm_cases.py

@@ -332,6 +332,14 @@ class LLMQuantLayerTypeCases:
                     "<class 'brevitas.nn.equalized_layer.EqualizedModule'>",
                 "model.layers.0.self_attn.q_proj.layer":
                     "<class 'brevitas.nn.quant_linear.QuantLinear'>",},},
+        {
+            "model": "hf-internal-testing/tiny-random-LlamaForCausalLM",
+            "rotation": "layerwise",
+            "exp_layer_types": {
+                "model.layers.0.self_attn.q_proj":
+                    "<class 'brevitas.nn.equalized_layer.RotatedModule'>",
+                "model.layers.0.self_attn.q_proj.layer":
+                    "<class 'brevitas.nn.quant_linear.QuantLinear'>",},},
         {
             "model": "hf-internal-testing/tiny-random-MistralForCausalLM",
             "quantize_last_layer": True,
@@ -360,6 +368,7 @@ class LLMQuantLayerTypeCases:
             "mistral-fp8_fnuz",
             "llama-mxfp8",
             "llama-int8-act_equalization=layerwise",
+            "llama-int8-rotation=layerwise",
             "mistral-int8-quant-last-layer",
             "llama-int8-svd_quant",
             "opt-quant-sdpa",],)