From b010ac06228cfdf6dcc3b393a6ef903fa3256bb9 Mon Sep 17 00:00:00 2001
From: lkk12014402 <kakao.lv@intel.com>
Date: Fri, 24 Jul 2026 11:44:58 +0000
Subject: [PATCH 1/4] add blockwise rotation.

Signed-off-by: lkk12014402 <kakao.lv@intel.com>
---
 auto_round/algorithms/transforms/base.py      |  86 +++
 .../algorithms/transforms/spinquant/apply.py  |  61 ++
 .../transforms/spinquant/preprocessor.py      | 571 ++++++++++++++----
 auto_round/autoround.py                       |   1 +
 auto_round/compressors/base.py                | 104 +++-
 auto_round/compressors/entry.py               |   2 +-
 auto_round/compressors/orchestrator.py        |  26 +
 7 files changed, 710 insertions(+), 141 deletions(-)

diff --git a/auto_round/algorithms/transforms/base.py b/auto_round/algorithms/transforms/base.py
index d7093dbc54..0f915b5329 100644
--- a/auto_round/algorithms/transforms/base.py
+++ b/auto_round/algorithms/transforms/base.py
@@ -113,6 +113,92 @@ def apply_to_model(
             The transformed model.
         """
 
+    # ------------------------------------------------------------------
+    # Layer-wise rotation interface (optional)
+    # ------------------------------------------------------------------
+
+    @property
+    def supports_layerwise(self) -> bool:
+        """Whether this rotation algorithm supports layer-wise execution.
+
+        When ``True``, the compressor can call :meth:`prepare_layerwise`
+        during ``post_init`` and then :meth:`rotate_layer` per-block inside
+        the ``_quantize_blocks`` loop — avoiding the need to load the
+        entire model onto GPU at once.
+        """
+        return False
+
+    def prepare_layerwise(
+        self,
+        model: torch.nn.Module,
+        data_type: str = "mx_fp",
+        **kwargs: Any,
+    ) -> "BaseRotation":
+        """Prepare for layer-wise rotation without modifying model weights.
+
+        Called once during ``post_init`` when ``layerwise_rotation=True``.
+        Implementations should initialise rotation matrices (as model
+        buffers) and any other lightweight state, but must **not** modify
+        model weights or register hooks yet.
+
+        Args:
+            model: The model to prepare (not modified).
+            data_type: Quantization data type.
+            **kwargs: Algorithm-specific arguments (e.g. ``dataloader``).
+
+        Returns:
+            ``self``, so the compressor can later call :meth:`rotate_layer`.
+
+        Raises:
+            NotImplementedError: If the algorithm does not support
+                layer-wise rotation.
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support layer-wise rotation. "
+            f"Use full-model rotation instead."
+        )
+
+    def rotate_layer(
+        self,
+        layer: torch.nn.Module,
+        layer_idx: int,
+        **kwargs: Any,
+    ) -> None:
+        """Apply rotation to a single decoder layer.
+
+        Called per-block in the ``_quantize_blocks`` loop, after the block
+        is materialised and placed on the target device, **before**
+        reference-output collection.
+
+        Implementations should:
+        - Rotate weights of target modules inside *layer*.
+        - Register any online hooks (R1/R3/R4) needed at inference time.
+        - Be idempotent — calling twice on the same layer should be safe.
+
+        Args:
+            layer: A single decoder layer, already on the target device.
+            layer_idx: Zero-based index of this layer in the model.
+            **kwargs: Algorithm-specific arguments.
+
+        Raises:
+            NotImplementedError: If the algorithm does not support
+                layer-wise rotation.
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support layer-wise rotation."
+        )
+
+    def finalize_layerwise(self, model: torch.nn.Module) -> None:
+        """Post-loop cleanup after all layers have been rotated.
+
+        Called once after the ``_quantize_blocks`` loop completes.
+        Default implementation is a no-op.
+
+        Args:
+            model: The fully-rotated model.
+        """
+        pass
+
     # ------------------------------------------------------------------
     # Factory
     # ------------------------------------------------------------------
diff --git a/auto_round/algorithms/transforms/spinquant/apply.py b/auto_round/algorithms/transforms/spinquant/apply.py
index e8e4c8f686..c112f4e3df 100644
--- a/auto_round/algorithms/transforms/spinquant/apply.py
+++ b/auto_round/algorithms/transforms/spinquant/apply.py
@@ -106,6 +106,67 @@ def apply_to_model(
         preprocessor = SpinQuantPreprocessor(model, self.config)
         return preprocessor.preprocess(dataloader)
 
+    # ------------------------------------------------------------------
+    # Layer-wise rotation interface (block-wise quantization)
+    # ------------------------------------------------------------------
+
+    @property
+    def supports_layerwise(self) -> bool:
+        """SpinQuant supports layer-wise rotation (Online R1 + R2/R3/R4)."""
+        return True
+
+    def prepare_layerwise(
+        self,
+        model: torch.nn.Module,
+        data_type: str = "mx_fp",
+        **kwargs: Any,
+    ) -> "SpinQuantRotation":
+        """Prepare for layer-wise rotation: init R matrices only.
+
+        Creates a :class:`SpinQuantPreprocessor`, calls its
+        :meth:`~SpinQuantPreprocessor.prepare` method (which initialises
+        rotation matrices without modifying weights), and stores the
+        preprocessor for later :meth:`rotate_layer` calls.
+
+        Args:
+            model: The model to prepare (weights not modified).
+            data_type: Quantization data type (informational).
+            **kwargs: ``dataloader`` forwarded when ``trainable_rotation=True``.
+
+        Returns:
+            ``self``, for chaining.
+        """
+        from auto_round.algorithms.transforms.spinquant.preprocessor import (
+            SpinQuantPreprocessor,
+        )
+
+        dataloader = kwargs.get("dataloader")
+        self._preprocessor = SpinQuantPreprocessor(model, self.config)
+        self._preprocessor.prepare(dataloader)
+        return self
+
+    def rotate_layer(
+        self,
+        layer: torch.nn.Module,
+        layer_idx: int,
+        **kwargs: Any,
+    ) -> None:
+        """Apply rotation to a single decoder layer.
+
+        Delegates to :meth:`SpinQuantPreprocessor.rotate_layer`.
+        """
+        if not hasattr(self, "_preprocessor"):
+            raise RuntimeError(
+                "prepare_layerwise() must be called before rotate_layer(). "
+                "The preprocessor has not been initialised."
+            )
+        self._preprocessor.rotate_layer(layer, layer_idx)
+
+    def finalize_layerwise(self, model: torch.nn.Module) -> None:
+        """Cleanup after all layers have been rotated."""
+        if hasattr(self, "_preprocessor"):
+            self._preprocessor.finalize()
+
     # ------------------------------------------------------------------
     # SerializerMixin — Save side
     # ------------------------------------------------------------------
diff --git a/auto_round/algorithms/transforms/spinquant/preprocessor.py b/auto_round/algorithms/transforms/spinquant/preprocessor.py
index c816c3d98b..6762fbd586 100644
--- a/auto_round/algorithms/transforms/spinquant/preprocessor.py
+++ b/auto_round/algorithms/transforms/spinquant/preprocessor.py
@@ -318,6 +318,273 @@ def preprocess(self, dataloader: Optional[Any] = None) -> nn.Module:
         logger.info("[SpinQuant] Preprocessing complete!")
         return self.model
 
+    # ------------------------------------------------------------------
+    # Layer-wise API (for block-lifecycle / block-wise quantization)
+    # ------------------------------------------------------------------
+
+    def prepare(self, dataloader: Optional[Any] = None) -> None:
+        """Global preparation for layer-wise (block-wise) rotation.
+
+        Performs all lightweight, non-destructive steps:
+        - Validates dimensions
+        - Initialises rotation matrices (stored as model buffers)
+        - Optionally trains rotation matrices (requires full model + dataloader)
+        - Pre-computes rotation state shared across layers
+
+        Does **NOT** modify model weights or register hooks — that is done
+        per-layer by :meth:`rotate_layer`.
+
+        This method is the layer-wise counterpart of :meth:`preprocess`.
+
+        Args:
+            dataloader: Required only when ``trainable_rotation=True``.
+
+        Raises:
+            ValueError: If ``online_r1_rotation=False`` while R1 is enabled
+                (offline R1 is incompatible with layer-wise mode because it
+                rotates the inter-layer hidden-state space via embed_tokens /
+                lm_head, which breaks pre-cached block inputs).
+        """
+        # Layer-wise rotation only works with Online R1 (offline R1 rotates the
+        # shared residual stream, which is incompatible with pre-cached block
+        # inputs).
+        if self.config.r1 and not self.config.online_r1_rotation:
+            raise ValueError(
+                "[SpinQuant] Layer-wise rotation requires online_r1_rotation=True. "
+                "Offline R1 changes inter-layer hidden state space, which is "
+                "incompatible with pre-cached block inputs. Use full-model "
+                "rotation (layerwise_rotation=False) for offline R1."
+            )
+
+        logger.info("[SpinQuant] Preparing for layer-wise rotation...")
+        logger.info(
+            f"[SpinQuant] Model architecture: hidden_size={self.hidden_size}, "
+            f"head_dim={self.head_dim}, intermediate_size={self.intermediate_size}"
+        )
+        logger.info(
+            f"[SpinQuant] Rotation config: R1={self.config.r1}, R2={self.config.r2}, "
+            f"R3={self.config.r3}, R4={self.config.r4}, "
+            f"online_r1={self.config.online_r1_rotation}"
+        )
+
+        # Step 1: Validate dimensions
+        self._validate_dimensions()
+
+        # Step 2: Initialise rotation matrices (lightweight — only R matrix buffers)
+        logger.info("[SpinQuant] Initialising rotation matrices...")
+        self._init_rotation_matrices()
+
+        # Step 3: Train if requested (needs full model forward/backward)
+        if self.config.trainable_rotation or self.config.trainable_smooth:
+            if dataloader is None:
+                raise ValueError("dataloader required when trainable=True")
+            logger.warning(
+                "[SpinQuant] Layer-wise mode with trainable_rotation: training "
+                "requires a full model forward/backward pass (higher memory). "
+                "Consider pre-training R matrices separately for large models."
+            )
+            # Trainable smooth needs norm replacement first
+            if self.config.trainable_smooth:
+                self._replace_norms_with_trainable()
+            self._train_rotations(dataloader)
+
+        # Pre-compute shared state consumed by rotate_layer()
+        if self.config.r1 and self.config.online_r1_rotation:
+            self._r1_shared = self._compute_online_r1_shared()
+        else:
+            self._r1_shared = None
+
+        if self.config.r2 and self.head_dim > 0:
+            R2_head = self._get_rotation_tensor("spinquant_R2_head")
+            if R2_head is not None:
+                R2 = R2_head.data.to(torch.float64)
+                self._r2_shared = (R2, R2.t())
+            else:
+                self._r2_shared = None
+        else:
+            self._r2_shared = None
+
+        if self.config.r4 and self.r4_rotation_size > 0:
+            use_random_r4 = self.config.random_r4
+            R4 = None
+            if use_random_r4:
+                R4_matrix = getattr(self.model, "spinquant_R4_matrix", None)
+                if R4_matrix is None:
+                    raise RuntimeError("[SpinQuant] random_r4=True but spinquant_R4_matrix buffer not found.")
+                R4 = R4_matrix.to(torch.float64)
+            self._r4_shared = (self.r4_rotation_size, use_random_r4, R4)
+        else:
+            self._r4_shared = None
+
+        # Store config on model for downstream serialization
+        self.model._rotation_config = self.config
+        self.model._spinquant_config = self.config  # legacy alias
+
+        logger.info("[SpinQuant] Layer-wise preparation complete (no weights modified yet).")
+
+    def rotate_layer(self, layer: nn.Module, layer_idx: int) -> None:
+        """Apply all configured rotations to a single decoder layer.
+
+        This is the per-block entry point called by the compressor inside the
+        block-wise quantization loop. The layer must already be materialised
+        and on the target device.
+
+        Operations performed (in order):
+        1. **R1 (Online)**: rotate target-module weights + register hooks
+        2. **R2 (Offline fuse)**: fuse into v_proj output + o_proj input
+        3. **R4 (Offline fuse)**: fuse into down_proj input
+        4. **R3 (Hook)**: monkeypatch apply_rotary_pos_emb
+        5. **R4 (Hook)**: register forward_pre_hook on down_proj
+
+        Args:
+            layer: A single decoder layer, already on the target device.
+            layer_idx: Zero-based index for logging purposes.
+        """
+        if not hasattr(self, "_r1_shared"):
+            raise RuntimeError(
+                "[SpinQuant] prepare() must be called before rotate_layer(). "
+                "The layer-wise state has not been initialised."
+            )
+        if not (hasattr(layer, "self_attn") and hasattr(layer, "mlp")):
+            logger.warning(f"[SpinQuant] Layer {layer_idx}: no self_attn/mlp, skipping rotation.")
+            return
+
+        layer_device = next(layer.parameters()).device
+
+        # ── R1: Online weight rotation + hooks ──
+        if self._r1_shared is not None:
+            self._apply_online_r1_to_layer(layer, *self._r1_shared)
+
+        # ── R2: Offline fuse into v_proj/o_proj ──
+        if self._r2_shared is not None:
+            self._fuse_r2_to_layer(layer, *self._r2_shared)
+
+        # ── R4: Offline fuse into down_proj ──
+        if self._r4_shared is not None:
+            self._fuse_r4_to_layer(layer, *self._r4_shared)
+
+        # ── R3: Monkeypatch attention (after RoPE) ──
+        if self.config.r3 and self.head_dim > 0:
+            self._rotate_layer_r3_hook(layer, layer_idx)
+
+        # ── R4: Hook on down_proj ──
+        if self.config.r4 and self.r4_rotation_size > 0:
+            self._rotate_layer_r4_hook(layer, layer_idx, layer_device)
+
+        logger.debug(f"[SpinQuant] Layer {layer_idx}: rotation applied.")
+
+    def _rotate_layer_r3_hook(self, layer: nn.Module, layer_idx: int) -> None:
+        """Apply the R3 monkeypatch to a single layer's attention module.
+
+        Consistent with the model-level ``register_spinquant_hooks`` R3 path.
+        """
+        if not is_pow2(self.head_dim):
+            return
+
+        from auto_round.algorithms.transforms.spinquant.monkeypatch import (
+            add_qk_rotation_after_rope,
+        )
+
+        attn = layer.self_attn
+        if not (hasattr(attn, "q_proj") and hasattr(attn, "k_proj")):
+            return
+        if getattr(attn, "_spinquant_r3_patched", False):
+            return  # Already patched (idempotent)
+
+        random_r3 = self.config.random_r3
+        r3_matrix = getattr(self.model, "spinquant_R3_head", None)
+
+        try:
+            wrapper = add_qk_rotation_after_rope(attn, rope_function_name="apply_rotary_pos_emb")
+            if random_r3 and r3_matrix is not None:
+                wrapper.set_matrix(r3_matrix)
+            else:
+                wrapper.set_hadamard(None, self.head_dim)
+            attn._spinquant_r3_patched = True
+            self._hook_handles.append(("r3_monkeypatch", f"layer_{layer_idx}", attn, wrapper))
+        except ValueError as e:
+            logger.warning(f"[SpinQuant] R3 monkeypatch failed for layer {layer_idx}: {e}")
+
+    def _rotate_layer_r4_hook(self, layer: nn.Module, layer_idx: int, layer_device: torch.device) -> None:
+        """Register the R4 ``forward_pre_hook`` on a single layer's down_proj.
+
+        Consistent with the model-level ``register_spinquant_hooks`` R4 path.
+        """
+        mlp = layer.mlp
+        if not hasattr(mlp, "down_proj"):
+            return
+
+        r4_size = self.r4_rotation_size
+        use_random = self.config.random_r4
+        need_block = r4_size < self.intermediate_size
+        module = mlp.down_proj
+
+        if use_random:
+            R4_matrix = getattr(self.model, "spinquant_R4_matrix", None)
+            if R4_matrix is None:
+                return
+            R4 = R4_matrix.to(device=layer_device, dtype=torch.float32)
+
+            def _make_hook(R, rot_size, block_mode):
+                def hook(mod, args):
+                    x = args[0]
+                    R_local = R.to(x.device, dtype=x.dtype)
+                    if block_mode:
+                        shape = x.shape
+                        x = x.reshape(*shape[:-1], -1, rot_size)
+                        x = (x @ R_local).reshape(shape)
+                    else:
+                        x = x @ R_local
+                    return (x,) + args[1:]
+
+                return hook
+
+            hook = _make_hook(R4, r4_size, need_block)
+        else:
+            try:
+                had_K_mat, had_K_val = get_hadamard_K(r4_size)
+            except ValueError:
+                return
+            had_K_mat = had_K_mat.to(device=layer_device, dtype=torch.float32)
+
+            def _make_hook_butterfly(had_mat, k_val, rot_size, block_mode):
+                def hook(mod, args):
+                    x = args[0]
+                    if block_mode:
+                        shape = x.shape
+                        x = x.reshape(*shape[:-1], -1, rot_size)
+                        x = matmul_hadU(x, hadamard_K=had_mat.to(x.device), K=k_val)
+                        x = x.reshape(shape)
+                    else:
+                        x = matmul_hadU(x, hadamard_K=had_mat.to(x.device), K=k_val)
+                    return (x,) + args[1:]
+
+                return hook
+
+            hook = _make_hook_butterfly(had_K_mat, had_K_val, r4_size, need_block)
+
+        hook._spinquant_hook = True
+        handle = module.register_forward_pre_hook(hook)
+        self._hook_handles.append(handle)
+
+    def finalize(self) -> None:
+        """Post-loop cleanup after all layers have been rotated layer-wise.
+
+        Mirrors :meth:`_cleanup` but for the layer-wise path.
+        """
+        self.model.eval()
+        for p in self.model.parameters():
+            p.requires_grad = False
+        self.rotation_params.clear()
+        self.smooth_params.clear()
+        self._rotated_modules.clear()
+
+        n_r1 = len(self._r1_hook_handles)
+        n_other = len(self._hook_handles)
+        logger.info(
+            f"[SpinQuant] Layer-wise rotation finalized. " f"Active hooks: {n_r1} R1 + {n_other} R3/R4."
+        )
+
     def _validate_dimensions(self) -> None:
         """Validate dimension requirements and disable rotations that can't work."""
         # R1: check r1_rotation_size divides hidden_size and is power of 2
@@ -679,6 +946,34 @@ def _apply_online_r1(self) -> None:
             "need to save/reload the model."
         )
 
+        r1_size, use_random, R1_full, hadamard_K, K = self._compute_online_r1_shared()
+
+        n_rotated = 0
+        n_hooked = 0
+
+        for layer in self._get_layers():
+            nr, nh = self._apply_online_r1_to_layer(layer, r1_size, use_random, R1_full, hadamard_K, K)
+            n_rotated += nr
+            n_hooked += nh
+
+        mode_str = "random matrix (x @ R)" if use_random else "deterministic butterfly"
+        logger.info(
+            f"[SpinQuant] Online R1: rotated {n_rotated} target modules, "
+            f"registered {n_hooked} activation hooks "
+            f"(rotation_size={r1_size}, mode={mode_str}, "
+            f"lm_head/embed_tokens/o_proj/down_proj unchanged)"
+        )
+
+    def _compute_online_r1_shared(self):
+        """Compute R1 rotation state shared across all layers (online R1).
+
+        Returns a tuple ``(r1_size, use_random, R1_full, hadamard_K, K)`` that
+        can be passed to :meth:`_apply_online_r1_to_layer`.  Extracted so that
+        both full-model (:meth:`_apply_online_r1`) and layer-wise
+        (:meth:`rotate_layer`) paths share identical R1 conventions.
+        """
+        r1_size = self.r1_rotation_size
+        use_random = self.config.random_r1
         model_device = next(self.model.parameters()).device
 
         # For random R1: use the stored full matrix
@@ -696,88 +991,86 @@ def _apply_online_r1(self) -> None:
             R1_full = None
             hadamard_K, K = get_hadamard_K(r1_size)
             hadamard_K = hadamard_K.to(model_device)
+        return r1_size, use_random, R1_full, hadamard_K, K
 
-        n_rotated = 0
-        n_hooked = 0
+    def _apply_online_r1_to_layer(self, layer, r1_size, use_random, R1_full, hadamard_K, K):
+        """Apply online R1 to a single decoder layer: rotate target-module
+        weights and register the matching activation ``forward_pre_hook`` s.
 
-        for layer in self._get_layers():
-            if not (hasattr(layer, "self_attn") and hasattr(layer, "mlp")):
-                continue
+        Returns ``(n_rotated, n_hooked)``.
+        """
+        if not (hasattr(layer, "self_attn") and hasattr(layer, "mlp")):
+            return 0, 0
 
-            layer_device = next(layer.parameters()).device
+        layer_device = next(layer.parameters()).device
 
-            attn = layer.self_attn
-            mlp = layer.mlp
+        attn = layer.self_attn
+        mlp = layer.mlp
 
-            # Target modules: (parent_module, attr_name)
-            target_specs = []
-            for proj_name in ("q_proj", "k_proj", "v_proj"):
-                if hasattr(attn, proj_name):
-                    target_specs.append((attn, proj_name))
-            for proj_name in ("gate_proj", "up_proj"):
-                if hasattr(mlp, proj_name):
-                    target_specs.append((mlp, proj_name))
-
-            for parent, attr_name in target_specs:
-                module = getattr(parent, attr_name)
-                dtype = module.weight.data.dtype
-                in_features = module.weight.shape[-1]
-
-                if use_random:
-                    # Random R1: explicit matrix multiply
-                    R = R1_full.to(layer_device)
-                    if r1_size == in_features:
-                        W = module.weight.data.to(torch.float64)
-                        module.weight.data = (W @ R).to(dtype)
-                    elif in_features % r1_size == 0:
-                        # Block rotation must use W @ R (not W @ R.T) to stay
-                        # consistent with the online hook (x @ R). The random
-                        # Hadamard matrix is orthonormal but NOT symmetric, so
-                        # rotate_in_channels_ (which applies R.T) would break
-                        # equivalence. Pass R.T so it computes W @ (R.T).T = W @ R.
-                        rotate_in_channels_(module, R_in=R.T)
-                    else:
-                        raise ValueError(
-                            f"Online R1: in_features={in_features} not compatible " f"with r1_rotation_size={r1_size}"
-                        )
+        # Target modules: (parent_module, attr_name)
+        target_specs = []
+        for proj_name in ("q_proj", "k_proj", "v_proj"):
+            if hasattr(attn, proj_name):
+                target_specs.append((attn, proj_name))
+        for proj_name in ("gate_proj", "up_proj"):
+            if hasattr(mlp, proj_name):
+                target_specs.append((mlp, proj_name))
+
+        n_rotated = 0
+        n_hooked = 0
+        for parent, attr_name in target_specs:
+            module = getattr(parent, attr_name)
+            dtype = module.weight.data.dtype
+            in_features = module.weight.shape[-1]
+
+            if use_random:
+                # Random R1: explicit matrix multiply
+                R = R1_full.to(layer_device)
+                if r1_size == in_features:
+                    W = module.weight.data.to(torch.float64)
+                    module.weight.data = (W @ R).to(dtype)
+                elif in_features % r1_size == 0:
+                    # Block rotation must use W @ R (not W @ R.T) to stay
+                    # consistent with the online hook (x @ R). The random
+                    # Hadamard matrix is orthonormal but NOT symmetric, so
+                    # rotate_in_channels_ (which applies R.T) would break
+                    # equivalence. Pass R.T so it computes W @ (R.T).T = W @ R.
+                    rotate_in_channels_(module, R_in=R.T)
                 else:
-                    # Deterministic Hadamard: butterfly algorithm
-                    had_K_local = hadamard_K.to(layer_device)
-                    if r1_size == in_features:
-                        module.weight.data = matmul_hadU(module.weight.data, hadamard_K=had_K_local, K=K).to(dtype)
-                    elif in_features % r1_size == 0:
-                        R_block = had_K_local.to(torch.float64)
-                        if R_block.shape[0] != r1_size:
-                            had_1, _ = get_hadamard_K(r1_size // K)
-                            R_block = torch.kron(
-                                had_K_local.to(device="cpu", dtype=torch.float64),
-                                had_1.to(device="cpu", dtype=torch.float64),
-                            )
-                        R_block = R_block / math.sqrt(r1_size)
-                        rotate_in_channels_(module, R_in=R_block)
-                    else:
-                        raise ValueError(
-                            f"Online R1: in_features={in_features} not compatible " f"with r1_rotation_size={r1_size}"
+                    raise ValueError(
+                        f"Online R1: in_features={in_features} not compatible " f"with r1_rotation_size={r1_size}"
+                    )
+            else:
+                # Deterministic Hadamard: butterfly algorithm
+                had_K_local = hadamard_K.to(layer_device)
+                if r1_size == in_features:
+                    module.weight.data = matmul_hadU(module.weight.data, hadamard_K=had_K_local, K=K).to(dtype)
+                elif in_features % r1_size == 0:
+                    R_block = had_K_local.to(torch.float64)
+                    if R_block.shape[0] != r1_size:
+                        had_1, _ = get_hadamard_K(r1_size // K)
+                        R_block = torch.kron(
+                            had_K_local.to(device="cpu", dtype=torch.float64),
+                            had_1.to(device="cpu", dtype=torch.float64),
                         )
-                n_rotated += 1
-
-                # Register forward_pre_hook for online activation rotation
-                if use_random:
-                    hook = self._make_online_r1_hook_matrix(R1_full.to(layer_device), r1_size, in_features)
+                    R_block = R_block / math.sqrt(r1_size)
+                    rotate_in_channels_(module, R_in=R_block)
                 else:
-                    hook = self._make_online_r1_hook_butterfly(r1_size, in_features, hadamard_K.to(layer_device), K)
-                hook._spinquant_hook = True  # tag for selective removal
-                handle = module.register_forward_pre_hook(hook)
-                self._r1_hook_handles.append(handle)
-                n_hooked += 1
+                    raise ValueError(
+                        f"Online R1: in_features={in_features} not compatible " f"with r1_rotation_size={r1_size}"
+                    )
+            n_rotated += 1
 
-        mode_str = "random matrix (x @ R)" if use_random else "deterministic butterfly"
-        logger.info(
-            f"[SpinQuant] Online R1: rotated {n_rotated} target modules, "
-            f"registered {n_hooked} activation hooks "
-            f"(rotation_size={r1_size}, mode={mode_str}, "
-            f"lm_head/embed_tokens/o_proj/down_proj unchanged)"
-        )
+            # Register forward_pre_hook for online activation rotation
+            if use_random:
+                hook = self._make_online_r1_hook_matrix(R1_full.to(layer_device), r1_size, in_features)
+            else:
+                hook = self._make_online_r1_hook_butterfly(r1_size, in_features, hadamard_K.to(layer_device), K)
+            hook._spinquant_hook = True  # tag for selective removal
+            handle = module.register_forward_pre_hook(hook)
+            self._r1_hook_handles.append(handle)
+            n_hooked += 1
+        return n_rotated, n_hooked
 
     @staticmethod
     def _make_online_r1_hook_butterfly(r1_size, in_features, hadamard_K, K):
@@ -942,34 +1235,41 @@ def _fuse_r2_rotation(self) -> None:
 
         n_fused = 0
         for layer in self._get_layers():
-            if not hasattr(layer, "self_attn"):
-                continue
-            attn = layer.self_attn
+            n_fused += self._fuse_r2_to_layer(layer, R2, R2_T)
 
-            # v_proj: W_new = R2^T @ W per head on output dimension
-            if hasattr(attn, "v_proj"):
-                W = attn.v_proj.weight.data
-                dtype = W.dtype
-                W = W.to(torch.float64)
-                n_heads = W.shape[0] // self.head_dim
-                W_reshaped = W.reshape(n_heads, self.head_dim, W.shape[1])
-                W_reshaped = torch.einsum("ij,kjl->kil", R2_T, W_reshaped)
-                attn.v_proj.weight.data = W_reshaped.reshape(W.shape).to(dtype)
-
-            # o_proj: W_new = W @ R2 per head on input dimension
-            # (R2^{-1} = R2^T on the activation side ↔ W @ R2 on the weight side)
-            if hasattr(attn, "o_proj"):
-                W = attn.o_proj.weight.data
-                dtype = W.dtype
-                W = W.to(torch.float64)
-                n_heads = W.shape[1] // self.head_dim
-                W_reshaped = W.reshape(W.shape[0], n_heads, self.head_dim)
-                W_reshaped = torch.einsum("ijk,kl->ijl", W_reshaped, R2)
-                attn.o_proj.weight.data = W_reshaped.reshape(W.shape).to(dtype)
+        logger.info(f"[SpinQuant] R2 fused into {n_fused} layers (v_proj out + o_proj in, head_dim={self.head_dim})")
 
-            n_fused += 1
+    def _fuse_r2_to_layer(self, layer, R2, R2_T) -> int:
+        """Fuse R2 per-head rotation into a single layer's v_proj/o_proj.
 
-        logger.info(f"[SpinQuant] R2 fused into {n_fused} layers (v_proj out + o_proj in, head_dim={self.head_dim})")
+        Returns 1 if the layer was fused, else 0.
+        """
+        if not hasattr(layer, "self_attn"):
+            return 0
+        attn = layer.self_attn
+
+        # v_proj: W_new = R2^T @ W per head on output dimension
+        if hasattr(attn, "v_proj"):
+            W = attn.v_proj.weight.data
+            dtype = W.dtype
+            W = W.to(torch.float64)
+            n_heads = W.shape[0] // self.head_dim
+            W_reshaped = W.reshape(n_heads, self.head_dim, W.shape[1])
+            W_reshaped = torch.einsum("ij,kjl->kil", R2_T.to(W.device), W_reshaped)
+            attn.v_proj.weight.data = W_reshaped.reshape(W.shape).to(dtype)
+
+        # o_proj: W_new = W @ R2 per head on input dimension
+        # (R2^{-1} = R2^T on the activation side ↔ W @ R2 on the weight side)
+        if hasattr(attn, "o_proj"):
+            W = attn.o_proj.weight.data
+            dtype = W.dtype
+            W = W.to(torch.float64)
+            n_heads = W.shape[1] // self.head_dim
+            W_reshaped = W.reshape(W.shape[0], n_heads, self.head_dim)
+            W_reshaped = torch.einsum("ijk,kl->ijl", W_reshaped, R2.to(W.device))
+            attn.o_proj.weight.data = W_reshaped.reshape(W.shape).to(dtype)
+
+        return 1
 
     def _fuse_r4_rotation(self) -> None:
         """Fuse R4 rotation into down_proj's input side.
@@ -989,6 +1289,7 @@ def _fuse_r4_rotation(self) -> None:
         use_random = self.config.random_r4
 
         # Get the rotation matrix
+        R4 = None
         if use_random:
             R4_matrix = getattr(self.model, "spinquant_R4_matrix", None)
             if R4_matrix is None:
@@ -997,38 +1298,7 @@ def _fuse_r4_rotation(self) -> None:
 
         n_fused = 0
         for layer in self._get_layers():
-            if not hasattr(layer, "mlp"):
-                continue
-            mlp = layer.mlp
-            if hasattr(mlp, "down_proj"):
-                W = mlp.down_proj.weight.data
-                dtype = W.dtype
-
-                if use_random:
-                    # Random: explicit W @ R per block
-                    W = W.to(torch.float64)
-                    if r4_size == W.shape[1]:
-                        mlp.down_proj.weight.data = (W @ R4).to(dtype)
-                    else:
-                        out_feat, in_feat = W.shape
-                        n_blocks = in_feat // r4_size
-                        W_reshaped = W.reshape(out_feat, n_blocks, r4_size)
-                        W_reshaped = torch.einsum("ijk,kl->ijl", W_reshaped, R4)
-                        mlp.down_proj.weight.data = W_reshaped.reshape(out_feat, in_feat).to(dtype)
-                else:
-                    # Deterministic: matmul_hadU (butterfly algorithm)
-                    # matmul_hadU operates on the last dimension — for weight
-                    # shape [out, in], last dim = in_features which is what
-                    # we want to rotate (input channels of down_proj).
-                    if r4_size == W.shape[1]:
-                        mlp.down_proj.weight.data = matmul_hadU(W).to(dtype)
-                    else:
-                        out_feat, in_feat = W.shape
-                        n_blocks = in_feat // r4_size
-                        W_reshaped = W.reshape(out_feat, n_blocks, r4_size)
-                        W_rotated = matmul_hadU(W_reshaped)
-                        mlp.down_proj.weight.data = W_rotated.reshape(out_feat, in_feat).to(dtype)
-                n_fused += 1
+            n_fused += self._fuse_r4_to_layer(layer, r4_size, use_random, R4)
 
         mode_str = "random x @ R" if use_random else "deterministic butterfly"
         logger.info(
@@ -1036,6 +1306,47 @@ def _fuse_r4_rotation(self) -> None:
             f"(r4_rotation_size={r4_size}, mode={mode_str})"
         )
 
+    def _fuse_r4_to_layer(self, layer, r4_size, use_random, R4) -> int:
+        """Fuse R4 rotation into a single layer's down_proj input side.
+
+        Returns 1 if the layer was fused, else 0.
+        """
+        if not hasattr(layer, "mlp"):
+            return 0
+        mlp = layer.mlp
+        if not hasattr(mlp, "down_proj"):
+            return 0
+
+        W = mlp.down_proj.weight.data
+        dtype = W.dtype
+
+        if use_random:
+            # Random: explicit W @ R per block
+            W = W.to(torch.float64)
+            R4 = R4.to(W.device)
+            if r4_size == W.shape[1]:
+                mlp.down_proj.weight.data = (W @ R4).to(dtype)
+            else:
+                out_feat, in_feat = W.shape
+                n_blocks = in_feat // r4_size
+                W_reshaped = W.reshape(out_feat, n_blocks, r4_size)
+                W_reshaped = torch.einsum("ijk,kl->ijl", W_reshaped, R4)
+                mlp.down_proj.weight.data = W_reshaped.reshape(out_feat, in_feat).to(dtype)
+        else:
+            # Deterministic: matmul_hadU (butterfly algorithm)
+            # matmul_hadU operates on the last dimension — for weight
+            # shape [out, in], last dim = in_features which is what
+            # we want to rotate (input channels of down_proj).
+            if r4_size == W.shape[1]:
+                mlp.down_proj.weight.data = matmul_hadU(W).to(dtype)
+            else:
+                out_feat, in_feat = W.shape
+                n_blocks = in_feat // r4_size
+                W_reshaped = W.reshape(out_feat, n_blocks, r4_size)
+                W_rotated = matmul_hadU(W_reshaped)
+                mlp.down_proj.weight.data = W_rotated.reshape(out_feat, in_feat).to(dtype)
+        return 1
+
     # ------------------------------------------------------------------
     # Step 8: Cleanup
     # ------------------------------------------------------------------
diff --git a/auto_round/autoround.py b/auto_round/autoround.py
index 2b67300dec..124f3b1da3 100644
--- a/auto_round/autoround.py
+++ b/auto_round/autoround.py
@@ -43,6 +43,7 @@
     "ignore_layers",
     "quant_lm_head",
     "to_quant_block_names",
+    "layerwise_rotation",
     "model_free",
     "disable_model_free",
     "model_dtype",
diff --git a/auto_round/compressors/base.py b/auto_round/compressors/base.py
index f0060b54c8..f00ef55ba8 100644
--- a/auto_round/compressors/base.py
+++ b/auto_round/compressors/base.py
@@ -25,6 +25,7 @@
 from auto_round.algorithms.transforms import (
     BaseRotationConfig,
     apply_rotation,
+    normalize_rotation_config,
 )
 from auto_round.auto_scheme.gen_auto_scheme import AutoScheme
 from auto_round.compressors.shard_writer import ShardWriter
@@ -308,6 +309,16 @@ def __init__(
         if device is not None:
             logger.warning("`device` is deprecated, please use `device_map` instead")
 
+        # Layer-wise (block-wise) rotation: when True, rotation matrices are
+        # only prepared up-front and applied per decoder block inside the
+        # block quantization loop instead of full-model up-front. Saves memory
+        # and pairs with block-wise quantization. Only SpinQuant/QuaRot with
+        # ``online_r1_rotation=True`` supports it.
+        self.layerwise_rotation = kwargs.pop("layerwise_rotation", False)
+        # Prepared per-layer rotation instances (populated by ``_apply_rotations``
+        # when ``layerwise_rotation=True``). Empty in full-model mode.
+        self._rotation_transforms: "list[BaseRotation]" = []
+
         self.static_attention_dtype = kwargs.pop("static_attention_dtype", None)
         # Attention static dtype
         if self.static_attention_dtype is not None:
@@ -1100,36 +1111,109 @@ def _resolve_formats(self) -> None:
     def _apply_rotations(self) -> None:
         """Phase 4.5 – Apply Hadamard / rotation transforms to the model.
 
+        Two modes are supported:
+
+        - **Full-model** (default): each rotation config is applied to the
+          entire model immediately via
+          :func:`~auto_round.algorithms.transforms.apply_rotation`, so the
+          model leaves this method with rotated weights and any online hooks
+          already installed.
+        - **Layer-wise** (``layerwise_rotation=True``): for configs whose
+          rotation algorithm reports ``supports_layerwise``, only the rotation
+          matrices are initialised here (lightweight, weights untouched). The
+          prepared rotation instances are stored in ``self._rotation_transforms``
+          and the actual per-block rotation is deferred to :meth:`_on_block_ready`
+          inside the block-quantization loop. Configs that do not support
+          layer-wise mode transparently fall back to full-model rotation.
+
         Preconditions:
           - Phase 3 complete: model topology is final (``apply_patches`` has
             replaced / merged layers, e.g. MoE experts), so rotation operates
             on the same modules that quantization will later see.
           - Phase 4 complete: ``self.layer_config`` is built; rotation only
-            transforms weights and does not change layer names, so this
-            ordering matches the old arch where rotation ran after
-            ``configure_layer_config``.
+            transforms weights and does not change layer names.
           - ``self.quantize_config.data_type`` is final (rotation backend
             dispatch depends on it).
 
-        Work performed:
-          - Iterates ``self.rotation_configs`` and calls
-            :func:`~auto_round.algorithms.transforms.apply_rotation` on the
-            model for each config.
-
         Postconditions:
-          - ``self.model_context.model`` carries the rotated weights and any
-            inserted online-Hadamard hooks.
+          - Full-model configs: ``self.model_context.model`` carries the
+            rotated weights and any inserted online-Hadamard hooks.
+          - Layer-wise configs: ``self._rotation_transforms`` holds the prepared
+            rotation instances; model weights are unchanged until the block loop.
         """
+        # Reset any state from a previous run (idempotent post_init).
+        self._rotation_transforms = []
+
         if not self.rotation_configs:
             return
+
+        from auto_round.algorithms.transforms.base import BaseRotation
+
         logger.info("Applying Hadamard transform to the model.")
         for rotation_cfg in self.rotation_configs:
+            if self.layerwise_rotation:
+                normalised = normalize_rotation_config(rotation_cfg)
+                if normalised is None:
+                    continue
+                rotation = BaseRotation.from_config(normalised)
+                if rotation.supports_layerwise:
+                    logger.info(
+                        "[Rotation] Layer-wise mode: preparing R matrices only "
+                        "(rotation deferred to per-block hook)."
+                    )
+                    rotation.prepare_layerwise(
+                        self.model_context.model,
+                        data_type=self.quantize_config.data_type,
+                    )
+                    self._rotation_transforms.append(rotation)
+                    continue
+                logger.warning(
+                    f"[Rotation] {rotation.__class__.__name__} does not support "
+                    f"layer-wise mode. Falling back to full-model rotation."
+                )
             self.model_context.model = apply_rotation(
                 self.model_context.model,
                 rotation_cfg,
                 data_type=self.quantize_config.data_type,
             )
 
+    # ------------------------------------------------------------------
+    # Block Lifecycle Hooks (layer-wise / block-wise rotation)
+    # ------------------------------------------------------------------
+
+    def _on_block_ready(self, block: torch.nn.Module, block_name, block_idx: int) -> None:
+        """Block lifecycle hook: apply layer-wise rotation to a block.
+
+        Called after a block is materialised and moved on-device, before
+        reference-output collection. No-op unless ``layerwise_rotation=True``
+        and at least one rotation config supports layer-wise mode.
+
+        Args:
+            block: The decoder block, already on the target device.
+            block_name: Block name(s) in the model tree. ``str`` (nblocks=1)
+                or ``list[str]`` (nblocks>1, WrapperMultiblock).
+            block_idx: Zero-based index in block_names (step of nblocks).
+        """
+        if not self._rotation_transforms:
+            return
+
+        if isinstance(block_name, (list, tuple)):
+            sub_modules = list(block.layers) if hasattr(block, "layers") else [block]
+            for j, sub_mod in enumerate(sub_modules):
+                for t in self._rotation_transforms:
+                    t.rotate_layer(sub_mod, layer_idx=block_idx + j)
+        else:
+            for t in self._rotation_transforms:
+                t.rotate_layer(block, layer_idx=block_idx)
+
+    def _finalize_block_processing(self, model) -> None:
+        """Finalize layer-wise rotation after all blocks are processed.
+
+        Safe to call even when no rotation transforms are active (no-op).
+        """
+        for t in self._rotation_transforms:
+            t.finalize_layerwise(model)
+
     def _patch_model(self) -> None:
         """Phase 3 – Model structure patching.
 
diff --git a/auto_round/compressors/entry.py b/auto_round/compressors/entry.py
index ce6b826e8c..892c99837c 100644
--- a/auto_round/compressors/entry.py
+++ b/auto_round/compressors/entry.py
@@ -23,7 +23,7 @@
 from auto_round.utils.device_manager import normalize_default_device_map
 
 _ENTRY_ROUTE_KWARGS = {"model_free", "disable_model_free", "disable_opt_rtn"}
-_ENTRY_COMPRESSOR_KWARGS = {"scale_dtype", "ignore_layers", "quant_lm_head", "to_quant_block_names"}
+_ENTRY_COMPRESSOR_KWARGS = {"scale_dtype", "ignore_layers", "quant_lm_head", "to_quant_block_names", "layerwise_rotation"}
 _ENTRY_BASE_KWARGS = {
     "format",
     "dataset",
diff --git a/auto_round/compressors/orchestrator.py b/auto_round/compressors/orchestrator.py
index d050e3861f..949ea3eb5b 100644
--- a/auto_round/compressors/orchestrator.py
+++ b/auto_round/compressors/orchestrator.py
@@ -231,6 +231,9 @@ def _quantize_blocks(
 
             m, _, _ = self.alg_composer.dispatch_block(m, input_ids, input_others)
 
+            # ── Layer-wise rotation: rotate this block before reference collection ──
+            self._on_block_ready(m, block_name_or_names, i)
+
             # ── Pipeline lifecycle: per-block setup ───────────────────────────
             from auto_round.algorithms.composer import BlockContext
 
@@ -373,6 +376,7 @@ def _quantize_zero_shot(self) -> tuple[torch.nn.Module, dict[str, Any]]:
 
         all_blocks = self.quant_block_list or get_block_names(self.model)
         pbar = tqdm(range(sum(len(block) for block in all_blocks)))
+        _zs_block_idx = 0
         for block_names in all_blocks:
             for block_name in block_names:
                 pbar.set_description(f"Quantizing {block_name}")
@@ -381,6 +385,10 @@ def _quantize_zero_shot(self) -> tuple[torch.nn.Module, dict[str, Any]]:
                 # ── Infrastructure: materialize ───────────────────────────
                 materialize_model_(block)
 
+                # ── Layer-wise rotation: rotate this block before quantization ──
+                self._on_block_ready(block, block_name, _zs_block_idx)
+                _zs_block_idx += 1
+
                 # ── Pure algorithm ────────────────────────────────────────
                 ctx = BlockContext(
                     model=self.model,
@@ -433,6 +441,8 @@ def _quantize_zero_shot(self) -> tuple[torch.nn.Module, dict[str, Any]]:
                 memory_monitor.log_summary()
                 pbar.update(1)
 
+        self._finalize_block_processing(self.model)
+
         cnt = 1
         remain_layer_names = []
         block_name_set = set(name for block in all_blocks for name in block)
@@ -590,6 +600,8 @@ def _quantize_data_driven(self) -> tuple[torch.nn.Module, dict[str, Any]]:
 
         # ── Pipeline lifecycle: finalize_quantization (model-level teardown)
         self.alg_composer.finalize_run()
+        # ── Layer-wise rotation: finalize after all blocks processed ─
+        self._finalize_block_processing(self.model_context.model)
         pbar.set_description("Quantizing done")
         pbar.close()
         if self.compress_context.low_cpu_mem_usage:
@@ -833,6 +845,20 @@ def quantize_block(
         if not self._post_init_done:
             self.post_init()
 
+        # Layer-wise rotation is driven by the internal block loop in
+        # ``quantize()`` (rotate on ``_on_block_ready`` → cleanup on
+        # ``_finalize_block_processing``). This externally-driven single-block API
+        # cannot guarantee that lifecycle, and rotating here would desync the
+        # caller's own reference/teacher outputs (collected on the un-rotated
+        # block). Fail loudly instead of producing silently wrong results.
+        if self.layerwise_rotation and self._rotation_transforms:
+            raise NotImplementedError(
+                "layerwise_rotation=True is not supported through the single-block "
+                "quantize_block() API (e.g. LLM-Compressor). Use the full AutoRound "
+                "quantize() entry point, or disable layerwise_rotation to apply "
+                "full-model rotation up-front."
+            )
+
         # ── Zero-shot (RTN) path: no calibration data needed ──────────────────
         if not self.need_calib:
             from auto_round.algorithms.composer import BlockContext

From a3cf1eb1cd1316d415624c331d12008564273b89 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 24 Jul 2026 11:46:53 +0000
Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 auto_round/algorithms/transforms/base.py                  | 7 ++-----
 .../algorithms/transforms/spinquant/preprocessor.py       | 4 +---
 auto_round/compressors/entry.py                           | 8 +++++++-
 3 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/auto_round/algorithms/transforms/base.py b/auto_round/algorithms/transforms/base.py
index 0f915b5329..00100c188b 100644
--- a/auto_round/algorithms/transforms/base.py
+++ b/auto_round/algorithms/transforms/base.py
@@ -154,8 +154,7 @@ def prepare_layerwise(
                 layer-wise rotation.
         """
         raise NotImplementedError(
-            f"{self.__class__.__name__} does not support layer-wise rotation. "
-            f"Use full-model rotation instead."
+            f"{self.__class__.__name__} does not support layer-wise rotation. " f"Use full-model rotation instead."
         )
 
     def rotate_layer(
@@ -184,9 +183,7 @@ def rotate_layer(
             NotImplementedError: If the algorithm does not support
                 layer-wise rotation.
         """
-        raise NotImplementedError(
-            f"{self.__class__.__name__} does not support layer-wise rotation."
-        )
+        raise NotImplementedError(f"{self.__class__.__name__} does not support layer-wise rotation.")
 
     def finalize_layerwise(self, model: torch.nn.Module) -> None:
         """Post-loop cleanup after all layers have been rotated.
diff --git a/auto_round/algorithms/transforms/spinquant/preprocessor.py b/auto_round/algorithms/transforms/spinquant/preprocessor.py
index 6762fbd586..d63f60be76 100644
--- a/auto_round/algorithms/transforms/spinquant/preprocessor.py
+++ b/auto_round/algorithms/transforms/spinquant/preprocessor.py
@@ -581,9 +581,7 @@ def finalize(self) -> None:
 
         n_r1 = len(self._r1_hook_handles)
         n_other = len(self._hook_handles)
-        logger.info(
-            f"[SpinQuant] Layer-wise rotation finalized. " f"Active hooks: {n_r1} R1 + {n_other} R3/R4."
-        )
+        logger.info(f"[SpinQuant] Layer-wise rotation finalized. " f"Active hooks: {n_r1} R1 + {n_other} R3/R4.")
 
     def _validate_dimensions(self) -> None:
         """Validate dimension requirements and disable rotations that can't work."""
diff --git a/auto_round/compressors/entry.py b/auto_round/compressors/entry.py
index 892c99837c..f0736f8c4f 100644
--- a/auto_round/compressors/entry.py
+++ b/auto_round/compressors/entry.py
@@ -23,7 +23,13 @@
 from auto_round.utils.device_manager import normalize_default_device_map
 
 _ENTRY_ROUTE_KWARGS = {"model_free", "disable_model_free", "disable_opt_rtn"}
-_ENTRY_COMPRESSOR_KWARGS = {"scale_dtype", "ignore_layers", "quant_lm_head", "to_quant_block_names", "layerwise_rotation"}
+_ENTRY_COMPRESSOR_KWARGS = {
+    "scale_dtype",
+    "ignore_layers",
+    "quant_lm_head",
+    "to_quant_block_names",
+    "layerwise_rotation",
+}
 _ENTRY_BASE_KWARGS = {
     "format",
     "dataset",

From b161861fe1d9a5c03e5d8b1faab3352e958a40b2 Mon Sep 17 00:00:00 2001
From: lkk <33276950+lkk12014402@users.noreply.github.com>
Date: Mon, 27 Jul 2026 14:07:58 +0800
Subject: [PATCH 3/4] fix lint issue.

---
 auto_round/compressors/base.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/auto_round/compressors/base.py b/auto_round/compressors/base.py
index f00ef55ba8..4c1b226773 100644
--- a/auto_round/compressors/base.py
+++ b/auto_round/compressors/base.py
@@ -16,7 +16,7 @@
 import os
 import sys
 from dataclasses import asdict, dataclass, fields
-from typing import Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 import torch
 from transformers import AutoConfig, set_seed
@@ -69,6 +69,9 @@
 from auto_round.utils.device_manager import device_manager
 from auto_round.utils.offload import OffloadManager
 
+if TYPE_CHECKING:
+    from auto_round.algorithms.transforms import BaseRotation
+
 
 @dataclass
 class SerializedCompressorConfig:

From 8e6c4fd602c4e6fcd3da4b82fe56b5dd2418ea56 Mon Sep 17 00:00:00 2001
From: lkk12014402 <kaokao.lv@intel.com>
Date: Tue, 4 Aug 2026 08:31:33 +0000
Subject: [PATCH 4/4] fix the coupling problem of rotation.

Signed-off-by: lkk12014402 <kaokao.lv@intel.com>
---
 auto_round/algorithms/composer.py             | 127 +++++++++++++++++-
 .../algorithms/transforms/hadamard/config.py  |  21 ---
 auto_round/compressors/base.py                | 125 ++---------------
 auto_round/compressors/orchestrator.py        |  36 +++--
 4 files changed, 152 insertions(+), 157 deletions(-)

diff --git a/auto_round/algorithms/composer.py b/auto_round/algorithms/composer.py
index bf399f8142..6df0adc27a 100644
--- a/auto_round/algorithms/composer.py
+++ b/auto_round/algorithms/composer.py
@@ -118,10 +118,18 @@ def __init__(self, configs: list, orchestrator: "BaseOrchestrator" = None) -> No
         """
         from auto_round.algorithms.quantization.base import BaseQuantizer
         from auto_round.algorithms.quantization.config import QuantizationConfig
-        from auto_round.algorithms.transforms.base import BasePreprocessor
+        from auto_round.algorithms.transforms.base import BasePreprocessor, BaseRotationConfig
 
         configs = list(configs)
 
+        # Rotation configs travel in the same config list but are not pipeline
+        # members (they are ``BaseRotationConfig``, not ``QuantizationConfig``).
+        # Capture them here so the composer owns the full rotation lifecycle
+        # (see ``apply_model_transforms`` / ``finalize_run``) and the orchestrator
+        # stays rotation-agnostic.
+        self._rotation_configs = [c for c in configs if isinstance(c, BaseRotationConfig)]
+        self._layerwise_rotation = bool(getattr(orchestrator, "layerwise_rotation", False))
+
         _, block_quantizer_configs = split_quantization_configs(configs)
         if not block_quantizer_configs:
             from auto_round.algorithms.quantization.rtn.config import RTNConfig
@@ -197,6 +205,10 @@ def __init__(self, configs: list, orchestrator: "BaseOrchestrator" = None) -> No
                 self.block_quantizer.bind_block_forward_runner(self.block_forward)
         self.scheme = getattr(orchestrator, "scheme_context", None)
 
+        # Rotation lifecycle state (populated by apply_model_transforms)
+        self._rotation_transforms: list = []
+        self._rotation_prepared: bool = False
+
     # ── Internal hook helpers (act_max calibration) ───────────────────────────
 
     def _register_act_max_hooks(self, block: "torch.nn.Module") -> list:
@@ -370,6 +382,12 @@ def compress_block(
         """
         block_forward_fn = self.block_forward
 
+        # ── Step 0: Layer-wise rotation (before any reference/calibration) ────
+        # Rotates this block's weights and installs online hooks so all downstream
+        # calibration and reference collection operate on the rotated block. No-op
+        # unless layer-wise rotation is active.
+        self._run_block_ready_transforms(block, block_ctx)
+
         # ── Step 1: Preprocessor calibration (e.g. AWQ activation stats) ──────
         with torch.no_grad():
             pre_hooks = []
@@ -545,3 +563,110 @@ def prepare_run(self, composer: "AlgorithmComposer" = None):
     def finalize_run(self):
         for alg in self.members():
             alg.finalize_run()
+        # Rotation teardown is part of the model-level finalize stage.
+        self._finalize_rotation(self._owning_model())
+
+    # ------------------------------------------------------------------
+    # Rotation lifecycle (owned entirely by the composer)
+    # ------------------------------------------------------------------
+    #
+    # Rotation is a model-level pre-quantisation transform. Full-model rotation
+    # must run *before* calibration data is cached, which is earlier than the
+    # per-member ``prepare_run`` stage; layer-wise rotation instead prepares its
+    # matrices here and rotates each block from within ``compress_block``. Both
+    # are driven internally so the orchestrator only calls the single generic
+    # entry point :meth:`apply_model_transforms`.
+
+    def _resolve_rotation_data_type(self) -> str:
+        """Best-effort resolution of the quantization data_type for rotation dispatch."""
+        if self.scheme is not None and getattr(self.scheme, "data_type", None):
+            return self.scheme.data_type
+        if self.block_quantizer is not None:
+            return getattr(self.block_quantizer.config, "data_type", "mx_fp")
+        return "mx_fp"
+
+    def _owning_model(self) -> "torch.nn.Module | None":
+        """Return the live model driven by this pipeline (via the block quantizer binding)."""
+        if self.block_quantizer is not None:
+            return getattr(self.block_quantizer, "model", None)
+        return None
+
+    def apply_model_transforms(self, model: "torch.nn.Module") -> "torch.nn.Module":
+        """Apply model-level pre-quantisation transforms (rotation) to *model*.
+
+        Generic entry point invoked once by the orchestrator before calibration
+        caching / the block loop. For full-model rotation the model is rotated
+        immediately and returned; for layer-wise rotation only the rotation
+        matrices are initialised and the per-block work is deferred to
+        :meth:`compress_block`. Idempotent — repeated calls are a no-op.
+
+        Returns:
+            The (possibly mutated) model.
+        """
+        if self._rotation_prepared:
+            return model
+
+        self._rotation_transforms = []
+        if not self._rotation_configs:
+            self._rotation_prepared = True
+            return model
+
+        from auto_round.algorithms.transforms import apply_rotation, normalize_rotation_config
+        from auto_round.algorithms.transforms.base import BaseRotation
+
+        data_type = self._resolve_rotation_data_type()
+        logger.info("Applying Hadamard transform to the model.")
+        for rotation_cfg in self._rotation_configs:
+            if self._layerwise_rotation:
+                normalised = normalize_rotation_config(rotation_cfg)
+                if normalised is None:
+                    continue
+                rotation = BaseRotation.from_config(normalised)
+                if rotation.supports_layerwise:
+                    logger.info(
+                        "[Rotation] Layer-wise mode: preparing R matrices only "
+                        "(rotation deferred to per-block hook)."
+                    )
+                    rotation.prepare_layerwise(model, data_type=data_type)
+                    self._rotation_transforms.append(rotation)
+                    continue
+                logger.warning(
+                    f"[Rotation] {rotation.__class__.__name__} does not support "
+                    f"layer-wise mode. Falling back to full-model rotation."
+                )
+            model = apply_rotation(model, rotation_cfg, data_type=data_type)
+
+        self._rotation_prepared = True
+        return model
+
+    def _run_block_ready_transforms(self, block: "torch.nn.Module", block_ctx: "BlockContext") -> None:
+        """Apply layer-wise rotation to a block before reference collection.
+
+        Called as the first step of :meth:`compress_block`. No-op when no
+        layer-wise rotation transforms are active. Uses the block's global
+        index (``block_ctx.block_index``) as the rotation layer index.
+        """
+        if not self._rotation_transforms:
+            return
+
+        block_idx = block_ctx.block_index
+        block_names = block_ctx.block_names
+        if isinstance(block_names, (list, tuple)) and len(block_names) > 1:
+            sub_modules = list(block.layers) if hasattr(block, "layers") else [block]
+            for j, sub_mod in enumerate(sub_modules):
+                for t in self._rotation_transforms:
+                    t.rotate_layer(sub_mod, layer_idx=block_idx + j)
+        else:
+            for t in self._rotation_transforms:
+                t.rotate_layer(block, layer_idx=block_idx)
+
+    def _finalize_rotation(self, model: "torch.nn.Module") -> None:
+        """Finalize layer-wise rotation after all blocks are processed (no-op when inactive)."""
+        for t in self._rotation_transforms:
+            t.finalize_layerwise(model)
+
+    @property
+    def has_layerwise_rotation(self) -> bool:
+        """Whether layer-wise rotation transforms are active."""
+        return bool(self._rotation_transforms)
+
diff --git a/auto_round/algorithms/transforms/hadamard/config.py b/auto_round/algorithms/transforms/hadamard/config.py
index 9e46ce0ffe..bf8639ef32 100644
--- a/auto_round/algorithms/transforms/hadamard/config.py
+++ b/auto_round/algorithms/transforms/hadamard/config.py
@@ -85,27 +85,6 @@ class RotationConfig(BaseModel, BaseRotationConfig):
 
     model_config = {"arbitrary_types_allowed": True}
 
-    def __init__(self, **data: Any) -> None:
-        """Initialize a Hadamard rotation configuration.
-
-        Args:
-            algorithm: Canonical algorithm name used for registry lookup.
-            backend: Rotation backend to use. ``auto`` lets AutoRound pick
-                an implementation, ``inplace`` uses QuaRot-style online
-                rotation, and ``transform`` uses the transform backend.
-            block_size: Grouped Hadamard block size. None keeps the backend
-                default behavior.
-            hadamard_type: Hadamard transform variant, such as
-                ``hadamard``, ``random_hadamard``, or ``quarot_hadamard``.
-            fuse_online_to_weight: Whether online Hadamard rotation should
-                be fused into the weights when supported.
-            allow_online_rotation: Whether online activation rotation is
-                allowed.
-            random_seed: Internal flag used by random Hadamard paths.
-            **data: Additional Pydantic field values forwarded to BaseModel.
-        """
-        super().__init__(**data)
-
     @field_validator("backend")
     @classmethod
     def _validate_backend(cls, v: str) -> str:
diff --git a/auto_round/compressors/base.py b/auto_round/compressors/base.py
index 4c1b226773..775d926dc5 100644
--- a/auto_round/compressors/base.py
+++ b/auto_round/compressors/base.py
@@ -16,7 +16,7 @@
 import os
 import sys
 from dataclasses import asdict, dataclass, fields
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import Any, Optional, Union
 
 import torch
 from transformers import AutoConfig, set_seed
@@ -24,8 +24,6 @@
 from auto_round.algorithms.quantization import BaseQuantizer, QuantizationConfig
 from auto_round.algorithms.transforms import (
     BaseRotationConfig,
-    apply_rotation,
-    normalize_rotation_config,
 )
 from auto_round.auto_scheme.gen_auto_scheme import AutoScheme
 from auto_round.compressors.shard_writer import ShardWriter
@@ -69,9 +67,6 @@
 from auto_round.utils.device_manager import device_manager
 from auto_round.utils.offload import OffloadManager
 
-if TYPE_CHECKING:
-    from auto_round.algorithms.transforms import BaseRotation
-
 
 @dataclass
 class SerializedCompressorConfig:
@@ -318,9 +313,6 @@ def __init__(
         # and pairs with block-wise quantization. Only SpinQuant/QuaRot with
         # ``online_r1_rotation=True`` supports it.
         self.layerwise_rotation = kwargs.pop("layerwise_rotation", False)
-        # Prepared per-layer rotation instances (populated by ``_apply_rotations``
-        # when ``layerwise_rotation=True``). Empty in full-model mode.
-        self._rotation_transforms: "list[BaseRotation]" = []
 
         self.static_attention_dtype = kwargs.pop("static_attention_dtype", None)
         # Attention static dtype
@@ -903,7 +895,6 @@ def post_init(self) -> None:
         self._resolve_formats()
         self._patch_model()
         self._build_layer_config()
-        self._apply_rotations()
 
         # Reclaim temporaries from Phases 1-4 (scheme resolution, format
         # parsing, model patching, layer-config walk) before Phase 5
@@ -918,6 +909,14 @@ def post_init(self) -> None:
         # so _build_composer must run first.
         self._build_composer()
 
+        # Phase 4.5 – Model-level pre-quantisation transforms (rotation).
+        # applies full-model rotation up-front (or prepares
+        # layer-wise rotation matrices). Runs here so every entry point — the full
+        # quantize() loop, the zero-shot loop, and the external single-block
+        # quantize_block() API — sees a consistently transformed model before any
+        # calibration data is collected.
+        self.model_context.model = self.alg_composer.apply_model_transforms(self.model_context.model)
+
         # Set block_forward torch compile for block forward
         # Final trim after all init phases.
         gc.collect()
@@ -1111,112 +1110,6 @@ def _resolve_formats(self) -> None:
             for _lname, _lval in _gguf_layer_cfg.items():
                 self.layer_config.setdefault(_lname, _lval)
 
-    def _apply_rotations(self) -> None:
-        """Phase 4.5 – Apply Hadamard / rotation transforms to the model.
-
-        Two modes are supported:
-
-        - **Full-model** (default): each rotation config is applied to the
-          entire model immediately via
-          :func:`~auto_round.algorithms.transforms.apply_rotation`, so the
-          model leaves this method with rotated weights and any online hooks
-          already installed.
-        - **Layer-wise** (``layerwise_rotation=True``): for configs whose
-          rotation algorithm reports ``supports_layerwise``, only the rotation
-          matrices are initialised here (lightweight, weights untouched). The
-          prepared rotation instances are stored in ``self._rotation_transforms``
-          and the actual per-block rotation is deferred to :meth:`_on_block_ready`
-          inside the block-quantization loop. Configs that do not support
-          layer-wise mode transparently fall back to full-model rotation.
-
-        Preconditions:
-          - Phase 3 complete: model topology is final (``apply_patches`` has
-            replaced / merged layers, e.g. MoE experts), so rotation operates
-            on the same modules that quantization will later see.
-          - Phase 4 complete: ``self.layer_config`` is built; rotation only
-            transforms weights and does not change layer names.
-          - ``self.quantize_config.data_type`` is final (rotation backend
-            dispatch depends on it).
-
-        Postconditions:
-          - Full-model configs: ``self.model_context.model`` carries the
-            rotated weights and any inserted online-Hadamard hooks.
-          - Layer-wise configs: ``self._rotation_transforms`` holds the prepared
-            rotation instances; model weights are unchanged until the block loop.
-        """
-        # Reset any state from a previous run (idempotent post_init).
-        self._rotation_transforms = []
-
-        if not self.rotation_configs:
-            return
-
-        from auto_round.algorithms.transforms.base import BaseRotation
-
-        logger.info("Applying Hadamard transform to the model.")
-        for rotation_cfg in self.rotation_configs:
-            if self.layerwise_rotation:
-                normalised = normalize_rotation_config(rotation_cfg)
-                if normalised is None:
-                    continue
-                rotation = BaseRotation.from_config(normalised)
-                if rotation.supports_layerwise:
-                    logger.info(
-                        "[Rotation] Layer-wise mode: preparing R matrices only "
-                        "(rotation deferred to per-block hook)."
-                    )
-                    rotation.prepare_layerwise(
-                        self.model_context.model,
-                        data_type=self.quantize_config.data_type,
-                    )
-                    self._rotation_transforms.append(rotation)
-                    continue
-                logger.warning(
-                    f"[Rotation] {rotation.__class__.__name__} does not support "
-                    f"layer-wise mode. Falling back to full-model rotation."
-                )
-            self.model_context.model = apply_rotation(
-                self.model_context.model,
-                rotation_cfg,
-                data_type=self.quantize_config.data_type,
-            )
-
-    # ------------------------------------------------------------------
-    # Block Lifecycle Hooks (layer-wise / block-wise rotation)
-    # ------------------------------------------------------------------
-
-    def _on_block_ready(self, block: torch.nn.Module, block_name, block_idx: int) -> None:
-        """Block lifecycle hook: apply layer-wise rotation to a block.
-
-        Called after a block is materialised and moved on-device, before
-        reference-output collection. No-op unless ``layerwise_rotation=True``
-        and at least one rotation config supports layer-wise mode.
-
-        Args:
-            block: The decoder block, already on the target device.
-            block_name: Block name(s) in the model tree. ``str`` (nblocks=1)
-                or ``list[str]`` (nblocks>1, WrapperMultiblock).
-            block_idx: Zero-based index in block_names (step of nblocks).
-        """
-        if not self._rotation_transforms:
-            return
-
-        if isinstance(block_name, (list, tuple)):
-            sub_modules = list(block.layers) if hasattr(block, "layers") else [block]
-            for j, sub_mod in enumerate(sub_modules):
-                for t in self._rotation_transforms:
-                    t.rotate_layer(sub_mod, layer_idx=block_idx + j)
-        else:
-            for t in self._rotation_transforms:
-                t.rotate_layer(block, layer_idx=block_idx)
-
-    def _finalize_block_processing(self, model) -> None:
-        """Finalize layer-wise rotation after all blocks are processed.
-
-        Safe to call even when no rotation transforms are active (no-op).
-        """
-        for t in self._rotation_transforms:
-            t.finalize_layerwise(model)
-
     def _patch_model(self) -> None:
         """Phase 3 – Model structure patching.
 
diff --git a/auto_round/compressors/orchestrator.py b/auto_round/compressors/orchestrator.py
index 949ea3eb5b..e95649a756 100644
--- a/auto_round/compressors/orchestrator.py
+++ b/auto_round/compressors/orchestrator.py
@@ -231,9 +231,6 @@ def _quantize_blocks(
 
             m, _, _ = self.alg_composer.dispatch_block(m, input_ids, input_others)
 
-            # ── Layer-wise rotation: rotate this block before reference collection ──
-            self._on_block_ready(m, block_name_or_names, i)
-
             # ── Pipeline lifecycle: per-block setup ───────────────────────────
             from auto_round.algorithms.composer import BlockContext
 
@@ -375,6 +372,7 @@ def _quantize_zero_shot(self) -> tuple[torch.nn.Module, dict[str, Any]]:
                 tied_weights_layers.append(lm_head_name)
 
         all_blocks = self.quant_block_list or get_block_names(self.model)
+
         pbar = tqdm(range(sum(len(block) for block in all_blocks)))
         _zs_block_idx = 0
         for block_names in all_blocks:
@@ -385,17 +383,16 @@ def _quantize_zero_shot(self) -> tuple[torch.nn.Module, dict[str, Any]]:
                 # ── Infrastructure: materialize ───────────────────────────
                 materialize_model_(block)
 
-                # ── Layer-wise rotation: rotate this block before quantization ──
-                self._on_block_ready(block, block_name, _zs_block_idx)
-                _zs_block_idx += 1
-
                 # ── Pure algorithm ────────────────────────────────────────
+                # ``block_index`` carries the global block index so compress_block
+                # can drive layer-wise rotation with the correct layer_idx.
                 ctx = BlockContext(
                     model=self.model,
                     block_names=[block_name],
                     block_name=block_name,
-                    block_index=0,
+                    block_index=_zs_block_idx,
                 )
+                _zs_block_idx += 1
                 # ── MoE scale alignment for FP8 dispatch efficiency ────────────────
                 if is_nv_fp(self.act_data_type) or not self.act_dynamic:
                     set_amax_for_all_moe_layers(block, attr_name="act_max")
@@ -441,7 +438,8 @@ def _quantize_zero_shot(self) -> tuple[torch.nn.Module, dict[str, Any]]:
                 memory_monitor.log_summary()
                 pbar.update(1)
 
-        self._finalize_block_processing(self.model)
+        # ── Pipeline lifecycle: model-level teardown (also finalizes rotation) ─
+        self.alg_composer.finalize_run()
 
         cnt = 1
         remain_layer_names = []
@@ -517,6 +515,7 @@ def _quantize_data_driven(self) -> tuple[torch.nn.Module, dict[str, Any]]:
             )
         else:
             logger.info("start to cache block inputs")
+
         all_inputs = self.cache_data(
             to_cache_block_names,
             self.calibration_context.nsamples,
@@ -598,10 +597,9 @@ def _quantize_data_driven(self) -> tuple[torch.nn.Module, dict[str, Any]]:
                     f"but got {len(self.formats)} formats."
                 )
 
-        # ── Pipeline lifecycle: finalize_quantization (model-level teardown)
+        # ── Pipeline lifecycle: finalize_quantization (model-level teardown,
+        #    which also finalizes any layer-wise rotation) ──────────────────
         self.alg_composer.finalize_run()
-        # ── Layer-wise rotation: finalize after all blocks processed ─
-        self._finalize_block_processing(self.model_context.model)
         pbar.set_description("Quantizing done")
         pbar.close()
         if self.compress_context.low_cpu_mem_usage:
@@ -845,13 +843,13 @@ def quantize_block(
         if not self._post_init_done:
             self.post_init()
 
-        # Layer-wise rotation is driven by the internal block loop in
-        # ``quantize()`` (rotate on ``_on_block_ready`` → cleanup on
-        # ``_finalize_block_processing``). This externally-driven single-block API
-        # cannot guarantee that lifecycle, and rotating here would desync the
-        # caller's own reference/teacher outputs (collected on the un-rotated
-        # block). Fail loudly instead of producing silently wrong results.
-        if self.layerwise_rotation and self._rotation_transforms:
+        # Layer-wise rotation is driven by the internal block loop inside
+        # ``AlgorithmComposer.compress_block`` (rotate as step 0, cleanup in
+        # ``finalize_run``). This externally-driven single-block API cannot
+        # guarantee that lifecycle, and rotating here would desync the caller's
+        # own reference/teacher outputs (collected on the un-rotated block).
+        # Fail loudly instead of producing silently wrong results.
+        if self.layerwise_rotation and self.alg_composer.has_layerwise_rotation:
             raise NotImplementedError(
                 "layerwise_rotation=True is not supported through the single-block "
                 "quantize_block() API (e.g. LLM-Compressor). Use the full AutoRound "