add blockwise rotation.

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)
+

auto_round/algorithms/transforms/base.py

@@ -113,6 +113,89 @@ 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
     # ------------------------------------------------------------------

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:

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
     # ------------------------------------------------------------------