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/base.py b/auto_round/algorithms/transforms/base.py index d7093dbc54..00100c188b 100644 --- a/auto_round/algorithms/transforms/base.py +++ b/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 # ------------------------------------------------------------------ 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/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..d63f60be76 100644 --- a/auto_round/algorithms/transforms/spinquant/preprocessor.py +++ b/auto_round/algorithms/transforms/spinquant/preprocessor.py @@ -318,6 +318,271 @@ 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 +944,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 +989,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 +1233,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 +1287,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 +1296,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 +1304,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..775d926dc5 100644 --- a/auto_round/compressors/base.py +++ b/auto_round/compressors/base.py @@ -24,7 +24,6 @@ from auto_round.algorithms.quantization import BaseQuantizer, QuantizationConfig from auto_round.algorithms.transforms import ( BaseRotationConfig, - apply_rotation, ) from auto_round.auto_scheme.gen_auto_scheme import AutoScheme from auto_round.compressors.shard_writer import ShardWriter @@ -308,6 +307,13 @@ 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) + self.static_attention_dtype = kwargs.pop("static_attention_dtype", None) # Attention static dtype if self.static_attention_dtype is not None: @@ -889,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 @@ -904,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() @@ -1097,39 +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. - - 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``. - - ``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. - """ - if not self.rotation_configs: - return - logger.info("Applying Hadamard transform to the model.") - for rotation_cfg in self.rotation_configs: - self.model_context.model = apply_rotation( - self.model_context.model, - rotation_cfg, - data_type=self.quantize_config.data_type, - ) - 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..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"} +_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..e95649a756 100644 --- a/auto_round/compressors/orchestrator.py +++ b/auto_round/compressors/orchestrator.py @@ -372,7 +372,9 @@ 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: for block_name in block_names: pbar.set_description(f"Quantizing {block_name}") @@ -382,12 +384,15 @@ def _quantize_zero_shot(self) -> tuple[torch.nn.Module, dict[str, Any]]: materialize_model_(block) # ── 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") @@ -433,6 +438,9 @@ def _quantize_zero_shot(self) -> tuple[torch.nn.Module, dict[str, Any]]: memory_monitor.log_summary() pbar.update(1) + # ── Pipeline lifecycle: model-level teardown (also finalizes rotation) ─ + self.alg_composer.finalize_run() + cnt = 1 remain_layer_names = [] block_name_set = set(name for block in all_blocks for name in block) @@ -507,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, @@ -588,7 +597,8 @@ 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() pbar.set_description("Quantizing done") pbar.close() @@ -833,6 +843,20 @@ def quantize_block( if not self._post_init_done: self.post_init() + # 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 " + "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