Module `neuralops.compression_engine`

Auto-import all compression_engine modules and export their registered classes.

Sub-modules

neuralops.compression_engine.a3_compressor: A3 Compressor: Activation-Aware Factorized SVD with Activation Energy Weighting and LayerNorm Recalibration …
neuralops.compression_engine.aa_svd_compressor: Activation-Aware SVD (AASVD) Compressor — Production-Hardened (v2) …
neuralops.compression_engine.aatc_xkv_compressor: AATC-XKV: Attention-Aware Cross-Layer KV Cache Compression (Production-Ready v2) …
neuralops.compression_engine.adaptive_ranks: Adaptive Ranks SVD Compressor: dynamic rank selection per layer based on singular value energy spectra …
neuralops.compression_engine.adasvd_compressor: AdaSVD: Adaptive SVD with Activation Whitening, Two-Sided Weighted ALS, and Mean-Shift Bias Absorption …
neuralops.compression_engine.admm_compressor: ADMM Compressor: Alternating Direction Method of Multipliers for Joint Pruning and Quantization. Registered in COMPRESSOR_REGISTRY as 'admm' and …
neuralops.compression_engine.admm_pruner: ADMM Pruner: Alternating Direction Method of Multipliers for Structured & Semi-Structured Pruning with Layer-Wise Reconstruction Distillation, STE, …
neuralops.compression_engine.armor_pruner: ARMOR: Adaptive Resilient Matrix Orthogonal Rank Pruning (2:4 Sparse Core + Low-Rank Correction). Registered in COMPRESSOR_REGISTRY as 'armor_pruner', …
neuralops.compression_engine.ast_pruner: AST: Adaptive Semi-Structured 2:4 and N:M Pruning with Empirical Fisher Calibration and PCA Coordinate Rotation (Corrected) …
neuralops.compression_engine.astro_quantizer: AstroQuantizer: Dynamic Uniform Range Quantization Engine with Straight-Through Estimator (STE) and Fisher-Adaptive Mixed-Precision Bit Allocation …
neuralops.compression_engine.bonsai_pruner: Bonsai: Iterative Magnitude and Taylor Pruning (IMP) with Gradient Descent Recovery & [CLS] Routing Protection. Registered in COMPRESSOR_REGISTRY as …
neuralops.compression_engine.cachenotes_compressor: CacheNotes: Task-Aware Key-Value Cache Compression and Static Weight Compression with Learnable Continuous CPT (Compression-Planning-Token) Generator …
neuralops.compression_engine.come_compressor: CoMe: Compress LLMs through Layer Concatenation and Shedding with Learnable Inter-Layer Interpolation Adapters …
neuralops.compression_engine.compot_compressor: COMPOT: Compressed Orthogonal Dictionary Learning Compressor. Registered in COMPRESSOR_REGISTRY as 'compot'. Production hardened for …
neuralops.compression_engine.coopq_quantizer: CoopQ: Cooperative Game-Theoretic Bit Allocation Quantizer with STE and Per-Channel Asymmetric Calibration. Registered in COMPRESSOR_REGISTRY as …
neuralops.compression_engine.cr_compressor: CR (Column-Row / Interpolative Decomposition) Compressor. Decomposes W ≈ C @ P, where C are actual columns of W. Registered in COMPRESSOR_REGISTRY as …
neuralops.compression_engine.cur_compressor: CUR Compressor: replaces Linear and Conv1D layers with CUR decomposition. Uses leverage score sampling or energy sampling for column/row selection. …
neuralops.compression_engine.cur_moe_compressor: CUR-MoE Compressor: apply CUR cross-approximation to Mixture of Experts (MoE) modules and expert layers. Registered in COMPRESSOR_REGISTRY as …
neuralops.compression_engine.delta_kv: DeltaKV: Residual-Based KV Cache Compression via Long-Range Similarity. Registered as 'delta_kv', 'delta_kv_class', 'delta_kv_compressor', …
neuralops.compression_engine.direct_svd: Direct SVD Compressor: Truncated SVD low-rank factorized layer replacement. Registered as 'direct_svd'. Production hardened for multi-architecture …
neuralops.compression_engine.dobi_svd: Dobi-SVD: Differentiable SVD layer with learnable spectral gating. Registered as 'dobi' and 'dobi_svd'. Production hardened for multi-architecture …
neuralops.compression_engine.drone_compressor: DRONE Compressor: Data-aware Low-Rank SVD with Identity Weighting …
neuralops.compression_engine.dynamic_rank: Dynamic Entropy Singular Value Decomposition (Dynamic Entropy SVD) Compressor …
neuralops.compression_engine.enlighten_compressor: ENLighten: Low-Rank SVD + 2:4 Structured Sparsity (Photonic-Aligned). Registered as 'enlighten' and 'enlighten_compressor'. Production hardened for …
neuralops.compression_engine.entropy_budget_allocation: EntropyBudgetAllocation: global rank assignment via whitened activation entropy. Registered as 'entropy_budget' and 'entropy_budget_allocation'.
neuralops.compression_engine.erc_svd: ERC-SVD: Error-Controlled SVD with Bounded Frobenius Truncation Error. Registered as 'erc' and 'erc_svd' …
neuralops.compression_engine.faq_quantizer: FAQ (Feature-Activation-Aware Quantization): adjusts quantization scales using activation statistics. Registered as 'faq'.
neuralops.compression_engine.flap: FLAP: Fluctuation-Based Adaptive Structured Pruning (Corrected) …
neuralops.compression_engine.flat_llm: FLAT-LLM: Activation-Aware Structured Dimension & Neuron Pruning …
neuralops.compression_engine.frandi_compressor: FRanDI: Feature-Regression Aligned SVD Compressor — Production-Hardened (v2) …
neuralops.compression_engine.fwsvd_compressor: Fisher-Weighted Singular Value Decomposition (FWSVD) Compressor …
neuralops.compression_engine.gfws_d_compressor: GFWSVD: Fisher-weighted / Covariance-weighted SVD. Registered as 'gfws' and 'gfws_d'.
neuralops.compression_engine.grecodec_compressor: GreCoDeC: Greedy Coordinate Descent Low-Rank Compressor. Registered in COMPRESSOR_REGISTRY as 'grecodec'.
neuralops.compression_engine.hashed_net_compressor: HashedNet Compressor — Production-Hardened (v2) …
neuralops.compression_engine.huffman_coder: Huffman Coder & Entropy Compression Engine for Weight and KV Cache Quantization. Registered as 'huffman'. Production-hardened with vectorized …
neuralops.compression_engine.hybrid_svd_layer: Hybrid SVD Layer: weight-based SVD compression for attention/linear layers …
neuralops.compression_engine.hybridkv_compressor: HybridKV: Hybrid KV Cache Compression for Multimodal & Long-Context LLMs. Classifies attention heads into static/dynamic types and applies distinct …
neuralops.compression_engine.kmeans_quantize: KMeans-based weight clustering quantizer with Straight-Through Estimator …
neuralops.compression_engine.kronecker_compressor: Kronecker Product Compressor: replaces nn.Linear and Conv1D layers with KroneckerLinear. Registered as 'kronecker_class', …
neuralops.compression_engine.kronecker_tt_lora: Kronecker-TT-LoRA Hybrid Compression Engine Adapter …
neuralops.compression_engine.kvtc_compressor: KVTC: KV Cache Transform Coding - PCA decorrelation + adaptive quantization + entropy coding for compact model weight / KV cache storage …
neuralops.compression_engine.liftr: LIFTR – Linearized Function Trick Compressor. Registered as 'liftr', 'liftr_compressor', 'liftr_class' …
neuralops.compression_engine.lillama_compressor: LiLLaMA: Local Distillation via Activation Cross-Covariance SVD. Registered as 'lillama_class', 'lillama_compressor_class', 'lillama_factory' and …
neuralops.compression_engine.lsc_compressor: LSC: Low-Rank Spectral & Fourier-Domain Linear Compressor. Registered as 'lsc', 'lsc_class', 'lsc_compressor', 'lsc_factory' …
neuralops.compression_engine.mamba_shedder: Mamba-Shedder: Sensitivity-Based Structured Pruning and Layer Shedding Compressor …
neuralops.compression_engine.mgie_svd: MGIE-SVD: Maximum Information Entropy SVD Rank Selection …
neuralops.compression_engine.mgie_svd_compressor: MGIE-SVD with Total Rank Budget: Global Entropy-Proportional Rank Allocation …
neuralops.compression_engine.mlorq_compressor: MLoRQ: Multi-Layer Optimal Rank & Quantisation under a Total Memory Budget …
neuralops.compression_engine.modegpt_compressor: MoDeGPT: Modular Decomposition of Transformer Blocks — Production-Hardened (v2) …
neuralops.compression_engine.nli_softmax: NLI Softmax approximation – a fast, quantisation-friendly softmax. Uses standard forward softmax, simplified sigmoid-based backward.
neuralops.compression_engine.nmf_compressor: NMF: Non-negative Matrix Factorization Linear Compressor …
neuralops.compression_engine.oats_compressor: OATS: Outlier-Aware Training & Saliency-Based Sparsification with Low-Rank Matrix Decomposition. Registered in COMPRESSOR_REGISTRY as 'oats' and …
neuralops.compression_engine.optimal_subspace_projection: Optimal Subspace Projection: Activation-Aware Subspace Projection via Randomized SVD …
neuralops.compression_engine.proxsparse: ProxSparse: Proximal Gradient Regularized Structured & Semi-Structured Sparsity …
neuralops.compression_engine.replaceme_compressor: ReplaceMe: Layer Removal and Replacement with Linear Residual Block …
neuralops.compression_engine.res_svd: Spectral Residual Low-Rank Adaptation (Res-SVD / Spectral Residual LoRA) …
neuralops.compression_engine.roast_hashed_linear: ROAST: Hardware-Aware Weight Sharing with Tiling. Reference: 'ROAST: A Hardware-Aware Weight Sharing Scheme for Efficient Neural Networks' Registered …
neuralops.compression_engine.rotational_slicegpt_int4: Rotational SliceGPT INT4 Compressor — Production-Hardened …
neuralops.compression_engine.rpca_compressor: Robust PCA: decompose W = L + S, factorizing low-rank component L with structural SVDLinear. Registered in COMPRESSOR_REGISTRY as 'rpca', …
neuralops.compression_engine.saes_svd_v2: SAES-SVD v2: Spectrum-Aware Error-Suppressed Singular Value Decomposition. Registered in COMPRESSOR_REGISTRY as 'saes_svd_v2' and 'saes' …
neuralops.compression_engine.salr: SALR: Structure-Aware / Sparse Low-Rank compression combining SVD with bottleneck pruning. Registered in COMPRESSOR_REGISTRY as 'salr', …
neuralops.compression_engine.sdrm_compressor: SDRM: Scaled Dimension Rank Modulation. Assigns layer rank proportional to parameter budget for edge inference and factorizes into low-rank SVDLinear …
neuralops.compression_engine.slicegpt: SliceGPT: Compress Large Language Models by Deleting Rows and Columns. Registered in COMPRESSOR_REGISTRY as 'slicegpt' and 'slicegpt_compressor' …
neuralops.compression_engine.slicegpt_pca: SliceGPT PCA Compressor: Principal Component Analysis based Dimensionality Slicing. Registered in COMPRESSOR_REGISTRY as 'slicegpt_pca' and …
neuralops.compression_engine.slim_compressor: SLiM: Structured Low-Rank Integrated Matrix Compression (Quantization + 2:4 Sparsity + SVD Residual Correction). Registered in COMPRESSOR_REGISTRY as …
neuralops.compression_engine.soka_compressor: SoKA: Second-Order Knowledge-Aware Low-Rank Linear Compressor. Registered in COMPRESSOR_REGISTRY as 'soka' …
neuralops.compression_engine.sparsegpt: SparseGPT: One-Shot Layer-Wise Weight Pruning with Second-Order Error Compensation. Registered in COMPRESSOR_REGISTRY as 'sparsegpt' …
neuralops.compression_engine.subspace_projection: Subspace Projection Compressor: Low-Rank Subspace Projection and Factorization …
neuralops.compression_engine.svd_linear: SVD-based low-rank linear layer replacement with optional Fused SVD kernel acceleration. Registered as 'svd_linear' and 'svd_linear_compressor'.
neuralops.compression_engine.svd_llm_compressor: SVD-LLM Compressor: Truncation-Aware Singular Value Decomposition for LLM Compression (ICLR 2025). Paper Reference: Wang et al., "SVD-LLM: …
neuralops.compression_engine.swift_svd: Swift-SVD: fast closed-form truncated SVD with symmetric square-root singular value distribution. Registered as 'swift_svd'.
neuralops.compression_engine.swift_svd_act: Swift-SVD Activation-Aware Compressor — Production-Hardened (v2) …
neuralops.compression_engine.tomoe_compressor: ToMoE: converts dense linear and Conv1D layers (e.g. FFN and attention projections) into a Mixture of Experts using weight clustering. Registered as …
neuralops.compression_engine.tt_compressor: Tensor Train Compressor (best-effort, production-guarded) …
neuralops.compression_engine.uniql_compressor: UniQL: Joint low-rank + uniform quantization compressor …
neuralops.compression_engine.zerokv_neo_compressor: ZeroKV-Neo: L2-norm based KV cache eviction with attention sink and heavy-hitter preservation. Registered as 'zerokv_neo'. Production-hardened with: - …

Classes

class A3Compressor (model: nn.Module, rank: int = 64, use_fisher: bool = True, use_energy_weighting: bool | None = None, use_cholesky: bool = False, allow_ln_failure: bool = False, lora_rank: int = 0, lora_alpha: float = 1.0, recalibrate_ln: bool = True, target_attention_only: bool = True, max_error: float = 0.5, max_retries: int = 5, fallback_to_original: bool = True, normalize_scale: bool = True, tokenizer: Any | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, n_calib: int = 500, device: str | torch.device | None = None, cov_dtype: str = 'float32', **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("a3")
@COMPRESSOR_REGISTRY.register("a3_compressor")
class A3Compressor(BaseCompressor):
    """
    A3: Activation-aware factorized SVD compressor.

    Uses E[x²] weighting (second moment), optional full Cholesky whitening,
    STRS rank search, and LayerNorm recalibration.
    """

    # A3-FIX-01: not attention
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module,
        rank: int = 64,
        use_fisher: bool = True,
        use_energy_weighting: bool | None = None,
        use_cholesky: bool = False,
        allow_ln_failure: bool = False,
        lora_rank: int = 0,
        lora_alpha: float = 1.0,
        recalibrate_ln: bool = True,
        target_attention_only: bool = True,
        max_error: float = 0.50,
        max_retries: int = 5,
        fallback_to_original: bool = True,
        normalize_scale: bool = True,
        tokenizer: Any | None = None,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        n_calib: int = 500,
        device: str | torch.device | None = None,
        cov_dtype: str = "float32",  # A3-FIX-08: "float32" or "float16"
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="a3", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg", getattr(model, "config", None)))

        # CLI overrides (kept for backward compatibility)
        import sys

        cli_cholesky = None
        cli_energy = None
        cli_rank = None
        for _idx, _arg in enumerate(sys.argv):
            if _arg in ("--a3_use_cholesky", "--use_cholesky") and _idx + 1 < len(sys.argv):
                cli_cholesky = sys.argv[_idx + 1].lower() in ("true", "1", "yes")
            elif _arg.startswith("--a3_use_cholesky=") or _arg.startswith("--use_cholesky="):
                cli_cholesky = _arg.split("=")[1].lower() in ("true", "1", "yes")
            elif _arg in (
                "--a3_use_energy_weighting",
                "--use_energy_weighting",
                "--a3_use_fisher",
                "--use_fisher",
            ) and _idx + 1 < len(sys.argv):
                cli_energy = sys.argv[_idx + 1].lower() in ("true", "1", "yes")
            elif _arg.startswith("--a3_use_energy_weighting=") or _arg.startswith(
                "--a3_use_fisher="
            ):
                cli_energy = _arg.split("=")[1].lower() in ("true", "1", "yes")
            elif _arg in ("--rank", "--a3_rank") and _idx + 1 < len(sys.argv):
                try:
                    cli_rank = int(sys.argv[_idx + 1])
                except ValueError:
                    pass

        if use_energy_weighting is not None:
            use_fisher = use_energy_weighting
        if cli_energy is not None:
            use_fisher = cli_energy
        if cli_cholesky is not None:
            use_cholesky = cli_cholesky
        if cli_rank is not None:
            rank = cli_rank

        if cfg is not None:
            rank = _cfg_get(cfg, "a3_rank", _cfg_get(cfg, "rank", rank))
            use_fisher = _cfg_get(
                cfg,
                "a3_use_energy_weighting",
                _cfg_get(
                    cfg,
                    "use_energy_weighting",
                    _cfg_get(cfg, "a3_use_fisher", _cfg_get(cfg, "use_fisher", use_fisher)),
                ),
            )
            use_cholesky = _cfg_get(
                cfg, "a3_use_cholesky", _cfg_get(cfg, "use_cholesky", use_cholesky)
            )
            allow_ln_failure = _cfg_get(
                cfg,
                "a3_allow_ln_failure",
                _cfg_get(cfg, "allow_ln_failure", allow_ln_failure),
            )
            lora_rank = _cfg_get(cfg, "a3_lora_rank", _cfg_get(cfg, "lora_rank", lora_rank))
            lora_alpha = _cfg_get(cfg, "a3_lora_alpha", _cfg_get(cfg, "lora_alpha", lora_alpha))
            recalibrate_ln = _cfg_get(
                cfg,
                "a3_recalibrate_ln",
                _cfg_get(cfg, "recalibrate_ln", recalibrate_ln),
            )
            target_attention_only = _cfg_get(
                cfg,
                "a3_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
            max_error = _cfg_get(cfg, "a3_max_error", _cfg_get(cfg, "max_error", max_error))
            max_retries = _cfg_get(cfg, "a3_max_retries", _cfg_get(cfg, "max_retries", max_retries))
            fallback_to_original = _cfg_get(
                cfg,
                "a3_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
            normalize_scale = _cfg_get(
                cfg,
                "a3_normalize_scale",
                _cfg_get(cfg, "normalize_scale", normalize_scale),
            )
            n_calib = _cfg_get(cfg, "a3_n_calib", _cfg_get(cfg, "n_calib", n_calib))

        self.rank = max(1, int(rank))
        # A3-FIX-03: internal name is use_energy_weighting
        self.use_energy_weighting = bool(use_fisher)
        self.use_fisher = self.use_energy_weighting  # backward compat
        self.use_cholesky = bool(use_cholesky)
        self.allow_ln_failure = bool(allow_ln_failure)
        self.lora_rank = int(lora_rank)
        self.lora_alpha = float(lora_alpha)
        self.recalibrate_ln = bool(recalibrate_ln)
        self.target_attention_only = bool(target_attention_only)
        self.max_error = float(max_error)
        self.n_calib = int(n_calib)
        self.max_retries = max(0, int(max_retries))
        self.fallback_to_original = bool(kwargs.get("fallback_to_original", fallback_to_original))
        self.normalize_scale = bool(kwargs.get("normalize_scale", normalize_scale))
        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.target_device = (
            torch.device(device)
            if device is not None
            else next(model.parameters(), torch.tensor(0.0)).device
        )
        self.cov_dtype_name = str(cov_dtype).lower()
        if self.cov_dtype_name not in ("float32", "float16"):
            self.cov_dtype_name = "float32"
        self.cov_dtype = torch.float32 if self.cov_dtype_name == "float32" else torch.float16

        self.stats: dict[str, Any] = {}
        self.compressed_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.activations_cleared: bool = False

        # A3-FIX-06 / A3-FIX-10
        self._last_r_ceiling: int = 0
        self._compressed_params: int = 0
        self._orig_storage_bytes: int = 0
        self._comp_storage_bytes: int = 0

        # Cached target names for hook filtering
        self._target_module_names: list[str] | None = None

    # ------------------------------------------------------------------
    # A3-FIX-02: determine target layers once
    # ------------------------------------------------------------------
    def _is_target(self, name: str, module: nn.Module) -> bool:
        cls_name = module.__class__.__name__.lower()
        target_mod = getattr(module, "base_layer", module)
        target_cls = target_mod.__class__.__name__.lower()
        is_linear_like = (
            isinstance(module, nn.Linear)
            or isinstance(target_mod, nn.Linear)
            or "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(module, "nf")
            or hasattr(target_mod, "nf")
            or (
                hasattr(target_mod, "weight")
                and hasattr(target_mod.weight, "ndim")
                and target_mod.weight.ndim == 2
                and not isinstance(target_mod, (nn.Embedding, nn.LayerNorm, nn.RMSNorm))
            )
        )
        if not is_linear_like:
            return False

        name_lower = name.lower()
        if any(
            k in name_lower
            for k in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "layernorm",
                "rmsnorm",
            )
        ):
            return False

        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self.target_attention_only:
            FFN_PATTERNS = {
                "bert": ("intermediate", "output.dense"),
                "roberta": ("intermediate", "output.dense"),
                "gpt2": ("mlp.c_fc", "mlp.c_proj"),
                "qwen": ("mlp.gate_proj", "mlp.up_proj", "mlp.down_proj"),
                "t5": ("densereludense", "wi_0", "wi_1", "wo"),
            }
            all_ffn = tuple(k for pats in FFN_PATTERNS.values() for k in pats)
            is_ffn = any(k in name_lower for k in all_ffn) or any(
                k in name_lower
                for k in (
                    "mlp",
                    "ffn",
                    "feed_forward",
                    "dense_4h_to_h",
                    "dense_h_to_4h",
                    "lin1",
                    "lin2",
                )
            )
            if is_ffn:
                return False
            is_attn = any(
                k in name_lower
                for k in (
                    "attn",
                    "attention",
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "out_proj",
                    "o_proj",
                    "q_lin",
                    "k_lin",
                    "v_lin",
                    "out_lin",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                    "c_proj",
                )
            )
            if not is_attn:
                return False
        return True

    # ------------------------------------------------------------------
    # Activation & energy collection
    # ------------------------------------------------------------------
    def _collect_activations_and_energy(
        self,
    ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
        """
        Collects (optional covariance, activation energy E[x²]).

        A3-FIX-02: hooks only on target linear layers.
        A3-FIX-07: padding filter is more robust.
        A3-FIX-08: covariance can be accumulated in lower precision.
        """
        if not self.use_energy_weighting and not self.recalibrate_ln and not self.use_cholesky:
            return {}, {}

        # Cache target names
        if self._target_module_names is None:
            self._target_module_names = [
                name for name, m in self.model.named_modules() if self._is_target(name, m)
            ]

        target_set = set(self._target_module_names)
        if not target_set:
            return {}, {}

        cov_running: dict[str, torch.Tensor] = {}
        energy_running: dict[str, torch.Tensor] = {}
        count_running: dict[str, int] = {}
        hooks = []

        def make_hook(layer_name: str):
            def hook(module, inp, outp):
                if not inp or not isinstance(inp[0], torch.Tensor):
                    return
                with torch.no_grad():
                    x = inp[0].detach().float()
                    if x.dim() < 2:
                        return
                    last = x.shape[-1]
                    x_flat = x.reshape(-1, last)
                    # A3-FIX-07: robust padding filter
                    # Tokens with all-zero activations are treated as padding.
                    is_valid = x_flat.abs().amax(dim=-1) > 0
                    x_valid = x_flat[is_valid]
                    n_valid = x_valid.size(0)
                    if n_valid == 0:
                        return

                    f_batch = x_valid.pow(2).sum(dim=0).cpu()

                    if self.use_cholesky:
                        c_batch = (x_valid.t() @ x_valid).cpu()
                        if self.cov_dtype == torch.float16:
                            c_batch = c_batch.half()
                        if layer_name not in cov_running:
                            cov_running[layer_name] = c_batch
                        else:
                            cov_running[layer_name] = cov_running[layer_name] + c_batch

                    if layer_name not in energy_running:
                        energy_running[layer_name] = f_batch
                        count_running[layer_name] = n_valid
                    else:
                        energy_running[layer_name] += f_batch
                        count_running[layer_name] += n_valid

            return hook

        for name, module in self.model.named_modules():
            if name in target_set:
                target_mod = getattr(module, "base_layer", module)
                hooks.append(target_mod.register_forward_hook(make_hook(name)))

        if self.calibration_loader is not None:
            self.model.eval()
            with torch.no_grad():
                for batch in self.calibration_loader:
                    try:
                        if isinstance(batch, dict):
                            self.model(
                                **{
                                    k: (v.to(self.target_device) if hasattr(v, "to") else v)
                                    for k, v in batch.items()
                                }
                            )
                        elif isinstance(batch, torch.Tensor):
                            self.model(batch.to(self.target_device))
                        elif isinstance(batch, (list, tuple)):
                            try:
                                self.model(
                                    *[
                                        (b.to(self.target_device) if hasattr(b, "to") else b)
                                        for b in batch
                                    ]
                                )
                            except Exception:
                                self.model(batch[0].to(self.target_device))
                    except Exception as e:
                        logger.debug("A3 calibration batch forward failed: %s", e)
        else:
            calib_inputs = get_calibration_tensors(
                self.model,
                tokenizer=self.tokenizer,
                calibration_loader=None,
                calibration_texts=self.calibration_texts,
                device=self.target_device,
                sample_size=self.n_calib,
            )
            if calib_inputs is not None:
                self.model.eval()
                with torch.no_grad():
                    try:
                        if isinstance(calib_inputs, dict):
                            total_n = next(iter(calib_inputs.values())).shape[0]
                            bs = 16
                            for start in range(0, total_n, bs):
                                sub = {k: v[start : start + bs] for k, v in calib_inputs.items()}
                                self.model(**sub)
                        elif isinstance(calib_inputs, torch.Tensor):
                            total_n = calib_inputs.shape[0]
                            bs = 16
                            for start in range(0, total_n, bs):
                                self.model(calib_inputs[start : start + bs])
                        elif isinstance(calib_inputs, (list, tuple)):
                            self.model(*calib_inputs)
                    except Exception as e:
                        logger.debug("A3 calibration forward failed: %s", e)

        for h in hooks:
            h.remove()

        cov_map: dict[str, torch.Tensor] = {}
        energy_map: dict[str, torch.Tensor] = {}
        for name in count_running:
            n = max(count_running[name], 1)
            if self.use_cholesky and name in cov_running:
                cov_map[name] = cov_running[name].float() / n
            if name in energy_running:
                energy_map[name] = energy_running[name] / n

        self.activations_cleared = True
        return cov_map, energy_map

    _collect_activations_and_fisher = _collect_activations_and_energy

    # ------------------------------------------------------------------
    def _compute_activation_error(self, W, W_r, X=None, C_X=None) -> float:
        if X is not None and X.numel() > 0:
            X_t = X.to(device=W.device, dtype=W.dtype).t()
            WX = torch.matmul(W, X_t)
            WrX = torch.matmul(W_r, X_t)
            num = torch.norm(WX - WrX, p="fro").item()
            den = torch.norm(WX, p="fro").item()
            return float(num / (den + 1e-9))
        elif C_X is not None and C_X.shape == (W.shape[1], W.shape[1]):
            C = C_X.to(device=W.device, dtype=W.dtype)
            delta = W - W_r
            num = torch.trace(delta @ C @ delta.t()).item()
            den = torch.trace(W @ C @ W.t()).item()
            if den > 1e-12:
                return float(math.sqrt(max(0.0, num)) / (math.sqrt(den) + 1e-9))
        W_norm = torch.norm(W, p="fro").item()
        if W_norm < 1e-12:
            return 0.0
        return float(torch.norm(W - W_r, p="fro").item() / (W_norm + 1e-9))

    # ------------------------------------------------------------------
    # A3-FIX-04 / A3-FIX-11: dynamic STRS candidates
    # ------------------------------------------------------------------
    @staticmethod
    def _build_strs_candidates(r_floor: int, r_ceiling: int) -> list[int]:
        if r_ceiling < r_floor:
            return [r_ceiling]
        # Geometric grid between r_floor and r_ceiling with ~8 points
        n_points = max(4, min(10, int(math.log2(max(2, r_ceiling // max(1, r_floor)))) + 2))
        cands = set()
        for i in range(n_points):
            t = i / (n_points - 1) if n_points > 1 else 0
            r = int(round(r_floor * (r_ceiling / max(1, r_floor)) ** t))
            cands.add(max(r_floor, min(r, r_ceiling)))
        cands.add(r_floor)
        cands.add(r_ceiling)
        return sorted(cands)

    def _strs_search(
        self,
        W,
        U,
        S,
        Vt,
        scale,
        S_inv,
        use_cholesky,
        r_floor,
        r_ceiling,
        X=None,
        cov=None,
    ):
        r_ceiling = min(r_ceiling, S.shape[0])  # A3-FIX-11
        candidates = self._build_strs_candidates(r_floor, r_ceiling)
        if not candidates:
            candidates = [r_floor]

        best_factors = None
        best_err = float("inf")
        best_r = candidates[0]

        for cand_r in candidates:
            sqrt_S = torch.sqrt(torch.clamp(S[:cand_r], min=1e-12))
            U_r = U[:, :cand_r] * sqrt_S.unsqueeze(0)
            V_r = sqrt_S.unsqueeze(1) * Vt[:cand_r, :]
            if use_cholesky and S_inv is not None:
                V_r = V_r @ S_inv.to(device=V_r.device, dtype=V_r.dtype)
            elif scale is not None:
                V_r = V_r / scale.unsqueeze(0)

            W_rec = torch.matmul(U_r, V_r)
            err = self._compute_activation_error(W, W_rec, X, cov)
            if err < best_err:
                best_err = err
                best_factors = (U_r, V_r)
                best_r = cand_r
            if err <= self.max_error:
                return (U_r, V_r), err, cand_r

        if best_factors is None:
            cand_r = candidates[-1]
            sqrt_S = torch.sqrt(torch.clamp(S[:cand_r], min=1e-12))
            U_r = U[:, :cand_r] * sqrt_S.unsqueeze(0)
            V_r = sqrt_S.unsqueeze(1) * Vt[:cand_r, :]
            if use_cholesky and S_inv is not None:
                V_r = V_r @ S_inv.to(device=V_r.device, dtype=V_r.dtype)
            elif scale is not None:
                V_r = V_r / scale.unsqueeze(0)
            best_factors = (U_r, V_r)
            best_err = self._compute_activation_error(W, torch.matmul(U_r, V_r), X, cov)
            best_r = cand_r

        return best_factors, best_err, best_r

    # ------------------------------------------------------------------
    def compress_linear(
        self,
        module: nn.Module,
        name: str = "",
        cov: torch.Tensor | None = None,
        fisher: torch.Tensor | None = None,
        rank: int | None = None,
        activation_energy: torch.Tensor | None = None,
        lora_rank: int = 0,
        lora_alpha: float = 1.0,
        X_for_error: torch.Tensor | None = None,
        **kwargs: Any,
    ):
        if fisher is None and activation_energy is not None:
            fisher = activation_energy
        cls_name = module.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")
        is_linear = isinstance(module, nn.Linear)
        is_linear_like = (
            is_linear
            or is_conv1d
            or (
                hasattr(module, "weight")
                and hasattr(module.weight, "ndim")
                and module.weight.ndim == 2
                and not isinstance(module, (nn.Embedding, nn.LayerNorm, nn.RMSNorm))
            )
        )
        if not is_linear_like:
            logger.warning(
                "A3: unsupported layer type %s (skipping %s)",
                module.__class__.__name__,
                name,
            )
            return None, 0.0, 0

        target_mod = getattr(module, "base_layer", module)
        device = target_mod.weight.device
        dtype = target_mod.weight.dtype
        raw_w = target_mod.weight.data.float()
        bias = getattr(target_mod, "bias", getattr(module, "bias", None))

        mod_with_nf = (
            target_mod if hasattr(target_mod, "nf") else (module if hasattr(module, "nf") else None)
        )
        if mod_with_nf is not None:
            out_f = mod_with_nf.nf
            in_f = getattr(
                mod_with_nf,
                "nx",
                raw_w.shape[0] if raw_w.shape[1] == out_f else raw_w.shape[1],
            )
            W = raw_w.t().contiguous() if raw_w.shape == (in_f, out_f) else raw_w
        elif (
            bias is not None
            and bias.ndim == 1
            and bias.shape[0] == raw_w.shape[1]
            and bias.shape[0] != raw_w.shape[0]
        ):
            W = raw_w.t().contiguous()
            out_f = bias.shape[0]
            in_f = raw_w.shape[0]
        elif "conv1d" in target_mod.__class__.__name__.lower():
            W = raw_w.t().contiguous()
            in_f, out_f = W.shape[1], W.shape[0]
        else:
            W = raw_w
            out_f, in_f = W.shape[0], W.shape[1]

        target_r = rank if rank is not None else self.rank
        r = get_dynamic_compression_rank(
            in_f,
            out_f,
            model_or_family=self.model,
            default_rank=target_r,
            comp_config=getattr(self, "config", None),
        )
        r = max(1, min(int(r), in_f, out_f))

        # ---- Weighting ----
        scale = None
        S_inv = None
        use_cholesky_this_layer = False

        if self.use_cholesky and cov is not None and cov.shape == (in_f, in_f):
            C_dev = cov.to(device=W.device, dtype=torch.float32)
            eps = 1e-6 * torch.trace(C_dev).item() / max(1, in_f)
            eye = torch.eye(in_f, device=W.device, dtype=torch.float32)
            try:
                S_white = torch.linalg.cholesky(C_dev + eps * eye)
                S_inv = torch.linalg.inv(S_white)
                W_w = W.float() @ S_white
                use_cholesky_this_layer = True
            except Exception as e:
                # A3-FIX-12: log
                logger.warning("A3: Cholesky failed on %s (%s); using diagonal.", name, e)
                diag_c = torch.diag(cov).float()
                scale = torch.sqrt(diag_c.clamp(min=1e-8) + 1e-6)
                if self.normalize_scale:
                    rms = torch.norm(scale) / math.sqrt(in_f) + 1e-8
                    scale = scale / rms
                scale = scale.clamp(min=1e-6).to(W.device)
                W_w = W * scale.unsqueeze(0)
                use_cholesky_this_layer = False
        elif (self.use_energy_weighting or self.use_fisher) and fisher is not None:
            f = fisher.float()
            if f.ndim == 1 and f.shape[0] == in_f:
                scale = torch.sqrt(f.clamp(min=1e-8) + 1e-6)
            if scale is not None:
                if self.normalize_scale:
                    rms = torch.norm(scale) / math.sqrt(in_f) + 1e-8
                    scale = scale / rms
                scale = scale.clamp(min=1e-6).to(W.device)
                W_w = W * scale.unsqueeze(0)
            else:
                W_w = W
        elif cov is not None:
            diag_c = torch.diag(cov).float()
            if diag_c.ndim == 1 and diag_c.shape[0] == in_f:
                scale = torch.sqrt(diag_c.clamp(min=1e-8) + 1e-6)
            if scale is not None:
                if self.normalize_scale:
                    rms = torch.norm(scale) / math.sqrt(in_f) + 1e-8
                    scale = scale / rms
                scale = scale.clamp(min=1e-6).to(W.device)
                W_w = W * scale.unsqueeze(0)
            else:
                W_w = W
        else:
            logger.debug("A3: no calibration data for %s; using unweighted SVD.", name)
            W_w = W

        # ---- SVD ----
        try:
            use_cpu_svd = not torch.cuda.is_available() or (in_f * out_f > 32_000_000)
            if use_cpu_svd:
                W_cpu = W_w.detach().cpu().float()
                U, S, Vt = torch.linalg.svd(W_cpu, full_matrices=False)
                U = U.to(device=W.device, dtype=W.dtype)
                S = S.to(device=W.device, dtype=W.dtype)
                Vt = Vt.to(device=W.device, dtype=W.dtype)
            else:
                W_dev = W_w.to(device=W.device, dtype=torch.float32)
                U, S, Vt = torch.linalg.svd(W_dev, full_matrices=False)
                U = U.to(device=W.device, dtype=W.dtype)
                S = S.to(device=W.device, dtype=W.dtype)
                Vt = Vt.to(device=W.device, dtype=W.dtype)
        except (RuntimeError, _CUDA_OOM):
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            W_cpu = W_w.detach().cpu().float()
            U_cpu, S_cpu, Vt_cpu = torch.linalg.svd(W_cpu, full_matrices=False)
            U = U_cpu.to(device=W.device, dtype=W.dtype)
            S = S_cpu.to(device=W.device, dtype=W.dtype)
            Vt = Vt_cpu.to(device=W.device, dtype=W.dtype)

        norm_orig = torch.linalg.norm(W).item()

        r_break_even = (in_f * out_f) / (in_f + out_f)
        r_ceiling = min(int(0.9 * r_break_even), 320)
        self._last_r_ceiling = max(self._last_r_ceiling, r_ceiling)

        if rank is not None:
            chosen_r = min(rank, r_ceiling, S.shape[0])
            sqrt_S = torch.sqrt(torch.clamp(S[:chosen_r], min=1e-12))
            U_r = U[:, :chosen_r] * sqrt_S.unsqueeze(0)
            V_r = sqrt_S.unsqueeze(1) * Vt[:chosen_r, :]
            if use_cholesky_this_layer and S_inv is not None:
                V_r = V_r @ S_inv.to(device=V_r.device, dtype=V_r.dtype)
            elif scale is not None:
                V_r = V_r / scale.unsqueeze(0)
            best_layer_factors = (U_r, V_r)
            best_r = chosen_r
            best_err = self._compute_activation_error(W, torch.matmul(U_r, V_r), X_for_error, cov)
        else:
            r_floor = max(16, min(self.rank, r_ceiling))
            best_layer_factors, best_err, best_r = self._strs_search(
                W,
                U,
                S,
                Vt,
                scale,
                S_inv,
                use_cholesky_this_layer,
                r_floor=r_floor,
                r_ceiling=r_ceiling,
                X=X_for_error,
                cov=cov,
            )

        if best_err > self.max_error:
            if self.fallback_to_original:
                logger.warning(
                    "A3: %s exceeded max_error (%.3f > %.3f) at r=%d; falling back to original.",
                    name,
                    best_err,
                    self.max_error,
                    r_ceiling,
                )
                return None, best_err, r_ceiling
            logger.warning(
                "A3: %s exceeds max_error (%.3f > %.3f) at r=%d; "
                "fallback_to_original=False → compressing anyway.",
                name,
                best_err,
                self.max_error,
                best_r,
            )

        wrapper_lora_rank = getattr(module, "lora_rank", getattr(module, "r", 0))
        effective_lora_rank = max(lora_rank, wrapper_lora_rank, self.lora_rank)
        effective_lora_alpha = max(lora_alpha, getattr(module, "lora_alpha", self.lora_alpha))

        U_final, V_final = best_layer_factors

        factorized_layer = A3FactorizedLinear(
            in_features=in_f,
            out_features=out_f,
            rank=best_r,
            bias=bias is not None,
            dtype=dtype,
            device=device,
            lora_rank=effective_lora_rank,
            lora_alpha=effective_lora_alpha,
        )
        factorized_layer.U.data.copy_(U_final.to(dtype=dtype, device=device))
        factorized_layer.V.data.copy_(V_final.to(dtype=dtype, device=device))
        if bias is not None and factorized_layer.bias is not None:
            if factorized_layer.bias.shape == bias.shape:
                factorized_layer.bias.data.copy_(bias.data.to(dtype=dtype, device=device))
            else:
                logger.warning("A3: bias shape mismatch on %s; skipping bias copy.", name)

        W_rec_final = torch.matmul(U_final, V_final)
        param_err = torch.linalg.norm(W - W_rec_final).item() / (norm_orig + 1e-8)
        factorized_layer.reconstruction_error = param_err
        factorized_layer.activation_error = best_err
        return factorized_layer, best_err, best_r

    # ------------------------------------------------------------------
    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        has_calib_sources = (
            self.calibration_loader is not None
            or bool(self.calibration_texts)
            or self.tokenizer is not None
        )
        if not has_calib_sources or (not self.use_fisher and not self.recalibrate_ln):
            logger.info("A3: no calibration sources; using unweighted SVD.")
            cov_map, fisher_map = {}, {}
            self.use_fisher = False
            self.use_energy_weighting = False
            self.recalibrate_ln = False
        else:
            cov_map, fisher_map = self._collect_activations_and_energy()

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        factorized_count = 0
        skipped_count = 0
        layer_stats: dict[str, Any] = {}
        layer_param_counts: dict[str, int] = {}
        total_err = 0.0
        self.compressed_layers = []
        self.skipped_layers = []
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

        for name, module in list(self.model.named_modules()):
            if not self._is_target(name, module):
                continue

            # Track original bytes
            target_mod = getattr(module, "base_layer", module)
            try:
                for p in target_mod.parameters():
                    self._orig_storage_bytes += p.numel() * p.element_size()
            except Exception:
                pass

            cov = cov_map.get(name, None)
            fisher = fisher_map.get(name, None)

            factorized_layer, rel_err, chosen_rank = self.compress_linear(
                module,
                name=name,
                cov=cov,
                fisher=fisher,
                activation_energy=fisher,
                lora_rank=self.lora_rank,
                lora_alpha=self.lora_alpha,
            )

            if factorized_layer is not None:
                if hasattr(module, "base_layer"):
                    module.base_layer = factorized_layer
                    replacement = module
                else:
                    replacement = factorized_layer
                replace_module(self.model, name, replacement)
                factorized_count += 1
                total_err += rel_err
                p_count = sum(p.numel() for p in module.parameters())
                layer_param_counts[name] = p_count
                layer_stats[name] = {
                    "rank": chosen_rank,
                    "rel_error": rel_err,
                    "weighted": bool(fisher is not None or cov is not None),
                }
                self.compressed_layers.append(name)

                # Track compressed bytes
                try:
                    for p in factorized_layer.parameters():
                        self._comp_storage_bytes += p.numel() * p.element_size()
                except Exception:
                    pass
            else:
                skipped_count += 1
                for p in module.parameters():
                    p.requires_grad = False
                self.skipped_layers.append(name)
                # Fallback storage equal to original
                self._comp_storage_bytes += (
                    self._orig_storage_bytes
                )  # approximate; overwritten below

        # ---- LayerNorm recalibration ----
        if self.recalibrate_ln:
            has_calib_data = (
                self.calibration_loader is not None
                or self.calibration_texts is not None
                or self.tokenizer is not None
            )
            if has_calib_data:
                try:
                    recalibrate_layernorms(
                        self.model,
                        tokenizer=self.tokenizer,
                        calibration_loader=self.calibration_loader,
                        calibration_texts=self.calibration_texts,
                        device=self.target_device,
                    )
                    logger.info("A3: LayerNorm recalibration succeeded.")
                except Exception as e:
                    logger.error("A3: LayerNorm recalibration FAILED: %s", e)
                    if not self.allow_ln_failure:
                        raise RuntimeError(
                            f"LayerNorm recalibration failed: {e}; "
                            f"set allow_ln_failure=True to bypass."
                        ) from e
            else:
                logger.info("A3: LayerNorm recalibration skipped (no calibration data).")

        self._compressed_params = self._count_params(self.model)
        avg_err = (total_err / factorized_count) if factorized_count > 0 else 0.0
        total_params = sum(layer_param_counts.values()) if layer_param_counts else 0
        weighted_err = (
            sum(layer_stats[k]["rel_error"] * layer_param_counts[k] for k in layer_stats)
            / max(1, total_params)
            if total_params > 0
            else 0.0
        )

        storage_ratio = (
            self._orig_storage_bytes / self._comp_storage_bytes
            if self._comp_storage_bytes > 0
            else 1.0
        )
        effective_ratio = (
            self._comp_storage_bytes / self._orig_storage_bytes
            if self._orig_storage_bytes > 0
            else 1.0
        )

        stats = self.get_compression_stats()
        stats.update(
            {
                "factorized_layers": factorized_count,
                "skipped_layers": skipped_count,
                "num_compressed": factorized_count,
                "num_fallback": skipped_count,
                "num_compressed_layers": factorized_count,
                "num_fallback_layers": skipped_count,
                "r_ceiling": self._last_r_ceiling or self.rank,
                "compressed_layers": self.compressed_layers,
                "rank": self.rank,
                "use_fisher": self.use_fisher,
                "use_energy_weighting": self.use_energy_weighting,
                "use_cholesky": self.use_cholesky,
                "recalibrate_ln": self.recalibrate_ln,
                "avg_reconstruction_error": avg_err,
                "weighted_reconstruction_error": weighted_err,
                "layer_stats": layer_stats,
                # Storage accounting
                "orig_storage_bytes": self._orig_storage_bytes,
                "comp_storage_bytes": self._comp_storage_bytes,
                "storage_compression_ratio": round(storage_ratio, 3),
                "storage_reduction_pct": round(100.0 * (1.0 - effective_ratio), 2),
                "method": "a3_activation_aware_factorized_svd",
            }
        )
        self.stats = stats
        try:
            self.model.compression_stats = stats
        except Exception:
            pass
        return self.model, stats

A3: Activation-aware factorized SVD compressor.

Uses E[x²] weighting (second moment), optional full Cholesky whitening, STRS rank search, and LayerNorm recalibration.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def compress_linear(self, module: nn.Module, name: str = '', cov: torch.Tensor | None = None, fisher: torch.Tensor | None = None, rank: int | None = None, activation_energy: torch.Tensor | None = None, lora_rank: int = 0, lora_alpha: float = 1.0, X_for_error: torch.Tensor | None = None, **kwargs: Any)

Expand source code

def compress_linear(
    self,
    module: nn.Module,
    name: str = "",
    cov: torch.Tensor | None = None,
    fisher: torch.Tensor | None = None,
    rank: int | None = None,
    activation_energy: torch.Tensor | None = None,
    lora_rank: int = 0,
    lora_alpha: float = 1.0,
    X_for_error: torch.Tensor | None = None,
    **kwargs: Any,
):
    if fisher is None and activation_energy is not None:
        fisher = activation_energy
    cls_name = module.__class__.__name__.lower()
    is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")
    is_linear = isinstance(module, nn.Linear)
    is_linear_like = (
        is_linear
        or is_conv1d
        or (
            hasattr(module, "weight")
            and hasattr(module.weight, "ndim")
            and module.weight.ndim == 2
            and not isinstance(module, (nn.Embedding, nn.LayerNorm, nn.RMSNorm))
        )
    )
    if not is_linear_like:
        logger.warning(
            "A3: unsupported layer type %s (skipping %s)",
            module.__class__.__name__,
            name,
        )
        return None, 0.0, 0

    target_mod = getattr(module, "base_layer", module)
    device = target_mod.weight.device
    dtype = target_mod.weight.dtype
    raw_w = target_mod.weight.data.float()
    bias = getattr(target_mod, "bias", getattr(module, "bias", None))

    mod_with_nf = (
        target_mod if hasattr(target_mod, "nf") else (module if hasattr(module, "nf") else None)
    )
    if mod_with_nf is not None:
        out_f = mod_with_nf.nf
        in_f = getattr(
            mod_with_nf,
            "nx",
            raw_w.shape[0] if raw_w.shape[1] == out_f else raw_w.shape[1],
        )
        W = raw_w.t().contiguous() if raw_w.shape == (in_f, out_f) else raw_w
    elif (
        bias is not None
        and bias.ndim == 1
        and bias.shape[0] == raw_w.shape[1]
        and bias.shape[0] != raw_w.shape[0]
    ):
        W = raw_w.t().contiguous()
        out_f = bias.shape[0]
        in_f = raw_w.shape[0]
    elif "conv1d" in target_mod.__class__.__name__.lower():
        W = raw_w.t().contiguous()
        in_f, out_f = W.shape[1], W.shape[0]
    else:
        W = raw_w
        out_f, in_f = W.shape[0], W.shape[1]

    target_r = rank if rank is not None else self.rank
    r = get_dynamic_compression_rank(
        in_f,
        out_f,
        model_or_family=self.model,
        default_rank=target_r,
        comp_config=getattr(self, "config", None),
    )
    r = max(1, min(int(r), in_f, out_f))

    # ---- Weighting ----
    scale = None
    S_inv = None
    use_cholesky_this_layer = False

    if self.use_cholesky and cov is not None and cov.shape == (in_f, in_f):
        C_dev = cov.to(device=W.device, dtype=torch.float32)
        eps = 1e-6 * torch.trace(C_dev).item() / max(1, in_f)
        eye = torch.eye(in_f, device=W.device, dtype=torch.float32)
        try:
            S_white = torch.linalg.cholesky(C_dev + eps * eye)
            S_inv = torch.linalg.inv(S_white)
            W_w = W.float() @ S_white
            use_cholesky_this_layer = True
        except Exception as e:
            # A3-FIX-12: log
            logger.warning("A3: Cholesky failed on %s (%s); using diagonal.", name, e)
            diag_c = torch.diag(cov).float()
            scale = torch.sqrt(diag_c.clamp(min=1e-8) + 1e-6)
            if self.normalize_scale:
                rms = torch.norm(scale) / math.sqrt(in_f) + 1e-8
                scale = scale / rms
            scale = scale.clamp(min=1e-6).to(W.device)
            W_w = W * scale.unsqueeze(0)
            use_cholesky_this_layer = False
    elif (self.use_energy_weighting or self.use_fisher) and fisher is not None:
        f = fisher.float()
        if f.ndim == 1 and f.shape[0] == in_f:
            scale = torch.sqrt(f.clamp(min=1e-8) + 1e-6)
        if scale is not None:
            if self.normalize_scale:
                rms = torch.norm(scale) / math.sqrt(in_f) + 1e-8
                scale = scale / rms
            scale = scale.clamp(min=1e-6).to(W.device)
            W_w = W * scale.unsqueeze(0)
        else:
            W_w = W
    elif cov is not None:
        diag_c = torch.diag(cov).float()
        if diag_c.ndim == 1 and diag_c.shape[0] == in_f:
            scale = torch.sqrt(diag_c.clamp(min=1e-8) + 1e-6)
        if scale is not None:
            if self.normalize_scale:
                rms = torch.norm(scale) / math.sqrt(in_f) + 1e-8
                scale = scale / rms
            scale = scale.clamp(min=1e-6).to(W.device)
            W_w = W * scale.unsqueeze(0)
        else:
            W_w = W
    else:
        logger.debug("A3: no calibration data for %s; using unweighted SVD.", name)
        W_w = W

    # ---- SVD ----
    try:
        use_cpu_svd = not torch.cuda.is_available() or (in_f * out_f > 32_000_000)
        if use_cpu_svd:
            W_cpu = W_w.detach().cpu().float()
            U, S, Vt = torch.linalg.svd(W_cpu, full_matrices=False)
            U = U.to(device=W.device, dtype=W.dtype)
            S = S.to(device=W.device, dtype=W.dtype)
            Vt = Vt.to(device=W.device, dtype=W.dtype)
        else:
            W_dev = W_w.to(device=W.device, dtype=torch.float32)
            U, S, Vt = torch.linalg.svd(W_dev, full_matrices=False)
            U = U.to(device=W.device, dtype=W.dtype)
            S = S.to(device=W.device, dtype=W.dtype)
            Vt = Vt.to(device=W.device, dtype=W.dtype)
    except (RuntimeError, _CUDA_OOM):
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        W_cpu = W_w.detach().cpu().float()
        U_cpu, S_cpu, Vt_cpu = torch.linalg.svd(W_cpu, full_matrices=False)
        U = U_cpu.to(device=W.device, dtype=W.dtype)
        S = S_cpu.to(device=W.device, dtype=W.dtype)
        Vt = Vt_cpu.to(device=W.device, dtype=W.dtype)

    norm_orig = torch.linalg.norm(W).item()

    r_break_even = (in_f * out_f) / (in_f + out_f)
    r_ceiling = min(int(0.9 * r_break_even), 320)
    self._last_r_ceiling = max(self._last_r_ceiling, r_ceiling)

    if rank is not None:
        chosen_r = min(rank, r_ceiling, S.shape[0])
        sqrt_S = torch.sqrt(torch.clamp(S[:chosen_r], min=1e-12))
        U_r = U[:, :chosen_r] * sqrt_S.unsqueeze(0)
        V_r = sqrt_S.unsqueeze(1) * Vt[:chosen_r, :]
        if use_cholesky_this_layer and S_inv is not None:
            V_r = V_r @ S_inv.to(device=V_r.device, dtype=V_r.dtype)
        elif scale is not None:
            V_r = V_r / scale.unsqueeze(0)
        best_layer_factors = (U_r, V_r)
        best_r = chosen_r
        best_err = self._compute_activation_error(W, torch.matmul(U_r, V_r), X_for_error, cov)
    else:
        r_floor = max(16, min(self.rank, r_ceiling))
        best_layer_factors, best_err, best_r = self._strs_search(
            W,
            U,
            S,
            Vt,
            scale,
            S_inv,
            use_cholesky_this_layer,
            r_floor=r_floor,
            r_ceiling=r_ceiling,
            X=X_for_error,
            cov=cov,
        )

    if best_err > self.max_error:
        if self.fallback_to_original:
            logger.warning(
                "A3: %s exceeded max_error (%.3f > %.3f) at r=%d; falling back to original.",
                name,
                best_err,
                self.max_error,
                r_ceiling,
            )
            return None, best_err, r_ceiling
        logger.warning(
            "A3: %s exceeds max_error (%.3f > %.3f) at r=%d; "
            "fallback_to_original=False → compressing anyway.",
            name,
            best_err,
            self.max_error,
            best_r,
        )

    wrapper_lora_rank = getattr(module, "lora_rank", getattr(module, "r", 0))
    effective_lora_rank = max(lora_rank, wrapper_lora_rank, self.lora_rank)
    effective_lora_alpha = max(lora_alpha, getattr(module, "lora_alpha", self.lora_alpha))

    U_final, V_final = best_layer_factors

    factorized_layer = A3FactorizedLinear(
        in_features=in_f,
        out_features=out_f,
        rank=best_r,
        bias=bias is not None,
        dtype=dtype,
        device=device,
        lora_rank=effective_lora_rank,
        lora_alpha=effective_lora_alpha,
    )
    factorized_layer.U.data.copy_(U_final.to(dtype=dtype, device=device))
    factorized_layer.V.data.copy_(V_final.to(dtype=dtype, device=device))
    if bias is not None and factorized_layer.bias is not None:
        if factorized_layer.bias.shape == bias.shape:
            factorized_layer.bias.data.copy_(bias.data.to(dtype=dtype, device=device))
        else:
            logger.warning("A3: bias shape mismatch on %s; skipping bias copy.", name)

    W_rec_final = torch.matmul(U_final, V_final)
    param_err = torch.linalg.norm(W - W_rec_final).item() / (norm_orig + 1e-8)
    factorized_layer.reconstruction_error = param_err
    factorized_layer.activation_error = best_err
    return factorized_layer, best_err, best_r

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class A3FactorizedLinear (in_features: int, out_features: int, rank: int, bias: bool = True, dtype: torch.dtype = torch.float32, device: torch.device | None = None, lora_rank: int = 0, lora_alpha: float = 1.0, freeze_base: bool = False)

Expand source code

class A3FactorizedLinear(nn.Module):
    """
    Low-rank factorized linear with optional LoRA:
        y = (x @ V.T) @ U.T + LoRA(x) + b
    Parameter count: r·(d_in + d_out) instead of d_in·d_out.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        bias: bool = True,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
        lora_rank: int = 0,
        lora_alpha: float = 1.0,
        freeze_base: bool = False,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = max(1, min(rank, in_features, out_features))
        self.lora_rank = max(0, int(lora_rank))
        self.lora_alpha = float(lora_alpha)
        self.lora_scaling = (self.lora_alpha / self.lora_rank) if self.lora_rank > 0 else 1.0

        fk = {"dtype": dtype, "device": device}
        self.V = nn.Parameter(torch.empty(self.rank, in_features, **fk))
        self.U = nn.Parameter(torch.empty(out_features, self.rank, **fk))

        if bias:
            self.bias = nn.Parameter(torch.empty(out_features, **fk))
        else:
            self.register_parameter("bias", None)

        if self.lora_rank > 0:
            self.lora_A = nn.Parameter(torch.empty(self.lora_rank, in_features, **fk))
            self.lora_B = nn.Parameter(torch.empty(out_features, self.lora_rank, **fk))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        self.reset_parameters()
        if freeze_base:
            self.freeze_base_weights()

    def reset_parameters(self) -> None:
        with torch.no_grad():
            # A3-FIX-14: proper orthogonal init for non-square matrices
            for param in (self.V, self.U):
                if param.dtype in (torch.float16, torch.bfloat16):
                    tmp = torch.empty_like(param, dtype=torch.float32, device="cpu")
                    _orthogonal_init(tmp)
                    param.copy_(tmp.to(device=param.device, dtype=param.dtype))
                else:
                    _orthogonal_init(param)
            if self.bias is not None:
                nn.init.zeros_(self.bias)
            if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
                nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
                nn.init.zeros_(self.lora_B)

    def freeze_base_weights(self) -> None:
        # A3-FIX-13: idempotent
        if self.V.requires_grad:
            self.V.requires_grad = False
        if self.U.requires_grad:
            self.U.requires_grad = False
        if self.bias is not None and self.bias.requires_grad:
            self.bias.requires_grad = False
        if self.lora_A is not None:
            self.lora_A.requires_grad = True
        if self.lora_B is not None:
            self.lora_B.requires_grad = True

    def get_weight(self) -> torch.Tensor:
        w = torch.matmul(self.U, self.V)
        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            w = w + torch.matmul(self.lora_B, self.lora_A) * self.lora_scaling
        return w

    def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
        with torch.no_grad():
            self.U.data.copy_(U_factor.to(device=self.U.device, dtype=self.U.dtype))
            self.V.data.copy_(V_factor.to(device=self.V.device, dtype=self.V.dtype))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        latent = F.linear(x, self.V)
        out = F.linear(latent, self.U)
        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_latent = F.linear(x, self.lora_A)
            out = out + F.linear(lora_latent, self.lora_B) * self.lora_scaling
        if self.bias is not None:
            out = out + self.bias
        return out

    # A3-FIX-05: removed __iter__ and __getitem__.

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        rank: int,
        U_factor: torch.Tensor | None = None,
        V_factor: torch.Tensor | None = None,
        lora_rank: int = 0,
        lora_alpha: float = 1.0,
        freeze_base: bool = False,
    ) -> A3FactorizedLinear:
        cls_name = linear.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(linear, "nf")
        target_mod = getattr(linear, "base_layer", linear)
        in_f = getattr(target_mod, "in_features", None)
        out_f = getattr(target_mod, "out_features", None)
        bias = getattr(target_mod, "bias", getattr(linear, "bias", None))

        if hasattr(target_mod, "nf"):
            out_f = target_mod.nf
            in_f = getattr(target_mod, "nx", None)
        elif hasattr(linear, "nf"):
            out_f = linear.nf
            in_f = getattr(linear, "nx", None)

        if in_f is None or out_f is None:
            w = getattr(target_mod, "weight", None)
            if w is not None and w.ndim == 2:
                if (
                    bias is not None
                    and bias.ndim == 1
                    and bias.shape[0] == w.shape[1]
                    and bias.shape[0] != w.shape[0]
                    or is_conv1d
                ):
                    in_f, out_f = w.shape[0], w.shape[1]
                else:
                    out_f, in_f = w.shape[0], w.shape[1]
            else:
                raise ValueError(f"Cannot determine in/out features for {type(linear)}")

        weight_ref = getattr(target_mod, "weight", getattr(linear, "weight", None))
        layer = cls(
            in_features=in_f,
            out_features=out_f,
            rank=rank,
            bias=bias is not None,
            dtype=weight_ref.dtype if weight_ref is not None else torch.float32,
            device=weight_ref.device if weight_ref is not None else None,
            lora_rank=lora_rank,
            lora_alpha=lora_alpha,
            freeze_base=freeze_base,
        )
        if U_factor is not None and V_factor is not None:
            layer.U.data.copy_(U_factor.to(device=layer.U.device, dtype=layer.U.dtype))
            layer.V.data.copy_(V_factor.to(device=layer.V.device, dtype=layer.V.dtype))
        if bias is not None and layer.bias is not None:
            if layer.bias.shape == bias.shape:
                layer.bias.data.copy_(
                    bias.data.to(device=layer.bias.device, dtype=layer.bias.dtype)
                )
        return layer

Low-rank factorized linear with optional LoRA: y = (x @ V.T) @ U.T + LoRA(x) + b Parameter count: r·(d_in + d_out) instead of d_in·d_out.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Module, rank: int, U_factor: torch.Tensor | None = None, V_factor: torch.Tensor | None = None, lora_rank: int = 0, lora_alpha: float = 1.0, freeze_base: bool = False) ‑> A3FactorizedLinear

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    latent = F.linear(x, self.V)
    out = F.linear(latent, self.U)
    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_latent = F.linear(x, self.lora_A)
        out = out + F.linear(lora_latent, self.lora_B) * self.lora_scaling
    if self.bias is not None:
        out = out + self.bias
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    # A3-FIX-13: idempotent
    if self.V.requires_grad:
        self.V.requires_grad = False
    if self.U.requires_grad:
        self.U.requires_grad = False
    if self.bias is not None and self.bias.requires_grad:
        self.bias.requires_grad = False
    if self.lora_A is not None:
        self.lora_A.requires_grad = True
    if self.lora_B is not None:
        self.lora_B.requires_grad = True

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    w = torch.matmul(self.U, self.V)
    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        w = w + torch.matmul(self.lora_B, self.lora_A) * self.lora_scaling
    return w

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    with torch.no_grad():
        # A3-FIX-14: proper orthogonal init for non-square matrices
        for param in (self.V, self.U):
            if param.dtype in (torch.float16, torch.bfloat16):
                tmp = torch.empty_like(param, dtype=torch.float32, device="cpu")
                _orthogonal_init(tmp)
                param.copy_(tmp.to(device=param.device, dtype=param.dtype))
            else:
                _orthogonal_init(param)
        if self.bias is not None:
            nn.init.zeros_(self.bias)
        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) ‑> None

Expand source code

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
    with torch.no_grad():
        self.U.data.copy_(U_factor.to(device=self.U.device, dtype=self.U.dtype))
        self.V.data.copy_(V_factor.to(device=self.V.device, dtype=self.V.dtype))

class AASVDCompressor (model: nn.Module, tokenizer: Any | None = None, rank: int = 288, ema_beta: float = 0.9, max_error: float = 0.5, fallback_to_original: bool = True, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, target_modules: list[str] | None = None, target_attention_only: bool = True, force_rank: bool = False, lora_rank: int = 0, lora_alpha: float = 1.0, cov_dtype: str = 'float32', **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("asvd")
@COMPRESSOR_REGISTRY.register("asvd_compressor")
@COMPRESSOR_REGISTRY.register("aasvd")
@COMPRESSOR_REGISTRY.register("aa_svd")
@COMPRESSOR_REGISTRY.register("aa_svd_compressor")
class AASVDCompressor(BaseCompressor):
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any | None = None,
        rank: int = 288,
        ema_beta: float = 0.9,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        target_modules: list[str] | None = None,
        target_attention_only: bool = True,
        force_rank: bool = False,
        lora_rank: int = 0,  # AASVD-FIX-03
        lora_alpha: float = 1.0,  # AASVD-FIX-03
        cov_dtype: str = "float32",  # AASVD-FIX-09
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="aasvd", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        # AASVD-FIX-02
        self.comp_config = cfg

        # AASVD-FIX-04: use _cfg_get everywhere
        rank = _cfg_get(cfg, "aasvd_rank", _cfg_get(cfg, "rank", rank))
        ema_beta = float(_cfg_get(cfg, "aasvd_ema_beta", _cfg_get(cfg, "ema_beta", ema_beta)))
        tokenizer = _cfg_get(cfg, "tokenizer", tokenizer)
        calibration_texts = _cfg_get(cfg, "calibration_texts", calibration_texts)
        max_error = float(_cfg_get(cfg, "aasvd_max_error", _cfg_get(cfg, "max_error", max_error)))
        fallback_to_original = bool(
            _cfg_get(
                cfg,
                "aasvd_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        target_attention_only = bool(
            _cfg_get(
                cfg,
                "aasvd_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        lora_rank = int(_cfg_get(cfg, "aasvd_lora_rank", _cfg_get(cfg, "lora_rank", lora_rank)))
        lora_alpha = float(
            _cfg_get(cfg, "aasvd_lora_alpha", _cfg_get(cfg, "lora_alpha", lora_alpha))
        )
        cov_dtype = str(_cfg_get(cfg, "cov_dtype", cov_dtype))

        # override from kwargs (highest priority after explicit args)
        lora_rank = int(kwargs.get("lora_rank", lora_rank))
        lora_alpha = float(kwargs.get("lora_alpha", lora_alpha))
        cov_dtype = str(kwargs.get("cov_dtype", cov_dtype))

        # force_rank resolution
        force_rank = bool(
            force_rank
            or kwargs.get("force_rank", False)
            or (cfg is not None and _cfg_get(cfg, "force_rank", False))
            or getattr(getattr(model, "config", None), "force_rank", False)
        )
        self.force_rank = force_rank
        self.use_strs = bool(kwargs.get("use_strs", not force_rank)) and not force_rank
        # AASVD-FIX-10: dynamic STRS candidates are computed per-layer

        # ---- tokenizer auto-resolve ----
        if tokenizer is None:
            tokenizer = getattr(model, "tokenizer", None) or getattr(
                getattr(model, "config", None), "tokenizer", None
            )
        self.tokenizer = tokenizer
        if self.tokenizer is None:
            mp = getattr(getattr(model, "config", None), "_name_or_path", None)
            if mp:
                try:
                    from transformers import AutoTokenizer

                    self.tokenizer = AutoTokenizer.from_pretrained(mp)
                    if getattr(self.tokenizer, "pad_token", None) is None:
                        self.tokenizer.pad_token = self.tokenizer.eos_token
                except Exception as e:
                    logger.debug("Auto-tokenizer load failed for %s: %s", mp, e)
                    self.tokenizer = None

        # ---- basic params ----
        self.rank = max(1, int(rank))
        self.ema_beta = float(max(0.0, min(ema_beta, 0.999)))
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.target_attention_only = bool(target_attention_only)
        self.calibration_texts = calibration_texts or DEFAULT_CALIBRATION_CORPUS
        self.target_modules = target_modules
        self.lora_rank = max(0, int(lora_rank))
        self.lora_alpha = float(lora_alpha)
        # AASVD-FIX-03
        self.lora_r = self.lora_rank  # legacy alias

        # AASVD-FIX-09
        self.cov_dtype_name = str(cov_dtype).lower()
        if self.cov_dtype_name not in ("float32", "float16"):
            self.cov_dtype_name = "float32"
        self.cov_dtype = torch.float32 if self.cov_dtype_name == "float32" else torch.float16

        # AASVD-FIX-08: single source of truth for device
        self.device = (
            torch.device(device)
            if device is not None
            else next(model.parameters(), torch.tensor(0.0)).device
        )

        # ---- state ----
        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}

        # AASVD-FIX-14
        self._compressed_params: int = 0
        try:
            self._original_params = self._count_params(self.model)
        except Exception:
            self._original_params = sum(p.numel() for p in self.model.parameters())

        # AASVD-FIX-11
        self._orig_storage_bytes: int = 0
        self._comp_storage_bytes: int = 0

    # ------------------------------------------------------------------
    # Target selection
    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        # AASVD-FIX-13: skip already compressed
        if isinstance(module, SVDLinear):
            return False
        if any(
            name.endswith(sfx) for sfx in (".base_layer", ".lora_A", ".lora_B", ".lora_dropout")
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = module.__class__.__name__.lower()
        target_cls = target_mod.__class__.__name__.lower()

        is_linear_like = (
            isinstance(module, nn.Linear)
            or isinstance(target_mod, nn.Linear)
            or "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(module, "nf")
            or hasattr(target_mod, "nf")
            or "loralinear" in cls_name
            or hasattr(module, "base_layer")
        )
        if not is_linear_like:
            return False

        name_lower = name.lower()
        if any(
            skip in name_lower
            for skip in (
                "lm_head",
                "classifier",
                "head",
                "output_layer",
                "score",
                "pooler",
                "embedding",
                "layernorm",
                "rmsnorm",
            )
        ):
            return False

        if self.target_attention_only:
            ffn_patterns = (
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            )
            if any(p in name_lower for p in ffn_patterns):
                return False
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            attn_patterns = (
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            )
            if not any(p in name_lower for p in attn_patterns):
                return False
        else:
            try:
                if not should_compress_linear(self.model, name):
                    return False
            except Exception:
                pass

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    # ------------------------------------------------------------------
    # AASVD-FIX-10: dynamic STRS candidates
    # ------------------------------------------------------------------
    @staticmethod
    def _build_strs_candidates(r_floor: int, r_ceiling: int) -> list[int]:
        if r_ceiling < r_floor:
            return [r_ceiling]
        cands = set()
        n = max(4, min(10, int(math.log2(max(2, r_ceiling // max(1, r_floor)))) + 2))
        for i in range(n):
            t = i / (n - 1) if n > 1 else 0
            r = int(round(r_floor * (r_ceiling / max(1, r_floor)) ** t))
            cands.add(max(r_floor, min(r, r_ceiling)))
        cands.add(r_floor)
        cands.add(r_ceiling)
        return sorted(cands)

    # ------------------------------------------------------------------
    def compress_linear(
        self,
        layer: nn.Module,
        rank: int | None = None,
        cov: torch.Tensor | None = None,
        cov_factors: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> tuple[SVDLinear | None, float, int]:
        try:
            target_mod = getattr(layer, "base_layer", layer)
            cls_name = layer.__class__.__name__.lower()
            target_cls = target_mod.__class__.__name__.lower()
            is_conv1d = (
                "conv1d" in cls_name
                or "conv1d" in target_cls
                or hasattr(layer, "nf")
                or hasattr(target_mod, "nf")
            )
            weight_tensor = getattr(target_mod, "weight", getattr(layer, "weight", None))
            if weight_tensor is None:
                return None, float("inf"), 0

            raw_w = weight_tensor.data
            if is_conv1d:
                in_feat, out_feat = raw_w.shape[0], raw_w.shape[1]
                W = raw_w.t().contiguous().float()
            else:
                out_feat, in_feat = raw_w.shape[0], raw_w.shape[1]
                W = raw_w.float()

            r_default = rank if rank is not None else self.rank
            r = get_dynamic_compression_rank(
                in_feat,
                out_feat,
                model_or_family=self.model,
                default_rank=r_default,
                comp_config=self.comp_config,
            )
            max_allowed_rank = max(
                1,
                min(
                    in_feat,
                    out_feat,
                    int(0.75 * (in_feat * out_feat) / (in_feat + out_feat)),
                ),
            )
            r = min(r, max_allowed_rank)
            W_norm = torch.norm(W, p="fro").item()

            # ---- covariance preconditioning ----
            cov_sqrt = cov_inv_sqrt = None
            if cov_factors is not None:
                cov_sqrt, cov_inv_sqrt = cov_factors
            elif cov is not None and cov.shape == (in_feat, in_feat):
                try:
                    cov_fp32 = cov.to(self.device, torch.float32)
                    cov_sym = 0.5 * (cov_fp32 + cov_fp32.T)
                    eigvals, eigvecs = torch.linalg.eigh(cov_sym)
                    eps = 1e-6 * torch.max(eigvals).clamp(min=1e-12) + 1e-8
                    eigvals_reg = eigvals.clamp(min=eps)
                    cov_sqrt = eigvecs @ torch.diag(torch.sqrt(eigvals_reg)) @ eigvecs.T
                    cov_inv_sqrt = eigvecs @ torch.diag(1.0 / torch.sqrt(eigvals_reg)) @ eigvecs.T
                except Exception as e:
                    logger.debug("AASVD covariance sqrt failed for %s: %s", e, e)
                    cov_sqrt, cov_inv_sqrt = None, None

            W_weighted = W @ cov_sqrt if cov_sqrt is not None else W

            W_w_dev = W_weighted.to(self.device, torch.float32)
            U, S, Vt = torch.linalg.svd(W_w_dev, full_matrices=False)

            # ---- STRS / fixed rank ----
            if not self.use_strs:
                selected_r = min(r, len(S))
                sqrt_S = torch.sqrt(torch.clamp(S[:selected_r], min=1e-12))
                best_Ur = U[:, :selected_r] * sqrt_S.unsqueeze(0)
                best_Vr = sqrt_S.unsqueeze(1) * Vt[:selected_r, :]
                if cov_inv_sqrt is not None:
                    best_Vr = best_Vr @ cov_inv_sqrt
                best_err = float(
                    torch.norm(W - best_Ur @ best_Vr, p="fro").item() / max(W_norm, 1e-12)
                )
            else:
                r_floor = max(1, min(r, max_allowed_rank, len(S)))
                r_ceiling = max_allowed_rank
                cand_ranks = self._build_strs_candidates(r_floor, r_ceiling)
                cand_ranks = [c for c in cand_ranks if c <= len(S)]

                best_err = float("inf")
                selected_r = cand_ranks[0] if cand_ranks else r_floor
                best_Ur = best_Vr = None

                for test_r in cand_ranks:
                    sqrt_S = torch.sqrt(torch.clamp(S[:test_r], min=1e-12))
                    U_cand = U[:, :test_r] * sqrt_S.unsqueeze(0)
                    V_cand = sqrt_S.unsqueeze(1) * Vt[:test_r, :]
                    if cov_inv_sqrt is not None:
                        V_cand = V_cand @ cov_inv_sqrt
                    W_cand = U_cand @ V_cand
                    test_err = float(torch.norm(W - W_cand, p="fro").item() / max(W_norm, 1e-12))
                    if test_err < best_err:
                        best_err = test_err
                        selected_r = test_r
                        best_Ur = U_cand
                        best_Vr = V_cand
                    if test_err <= self.max_error:
                        break

                if best_Ur is None or best_err > self.max_error:
                    # Fallback to break-even rank
                    be_r = max(
                        1,
                        min(
                            int(0.75 * (in_feat * out_feat) / (in_feat + out_feat)),
                            len(S),
                        ),
                    )
                    selected_r = be_r
                    sqrt_S = torch.sqrt(torch.clamp(S[:be_r], min=1e-12))
                    best_Ur = U[:, :be_r] * sqrt_S.unsqueeze(0)
                    best_Vr = sqrt_S.unsqueeze(1) * Vt[:be_r, :]
                    if cov_inv_sqrt is not None:
                        best_Vr = best_Vr @ cov_inv_sqrt
                    best_err = float(
                        torch.norm(W - best_Ur @ best_Vr, p="fro").item() / max(W_norm, 1e-12)
                    )

            dtype = weight_tensor.dtype
            device = weight_tensor.device
            bias_tensor = getattr(target_mod, "bias", getattr(layer, "bias", None))

            svd_layer = SVDLinear(
                in_features=in_feat,
                out_features=out_feat,
                rank=selected_r,
                bias=bias_tensor is not None,
                dtype=dtype,
                device=device,
                lora_rank=self.lora_rank,
                lora_alpha=self.lora_alpha,
            )
            svd_layer.set_low_rank_factors(
                best_Ur.to(device=device, dtype=dtype),
                best_Vr.to(device=device, dtype=dtype),
            )
            if bias_tensor is not None and svd_layer.bias is not None:
                svd_layer.bias.data.copy_(bias_tensor.data.to(svd_layer.bias.dtype))

            svd_layer.reconstruction_error = best_err
            svd_layer.rank = selected_r
            return svd_layer, best_err, selected_r
        except Exception as e:
            logger.warning("AASVD factorization failed: %s", e)
            return None, float("inf"), 0

    # ------------------------------------------------------------------
    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        self.model.eval()
        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

        covariances: dict[str, torch.Tensor] = {}
        # AASVD-FIX-06: track per-layer update count
        update_counts: dict[str, int] = {}

        if self.tokenizer is None:
            logger.info("AASVD: no tokenizer; using unweighted SVD.")
        else:
            hooks = []

            def make_hook(layer_name: str):
                nonlocal covariances, update_counts

                def hook(module, inp, outp):
                    if not inp or inp[0] is None:
                        return
                    x = inp[0].detach()
                    # AASVD-FIX-05
                    with _autocast_off(x):
                        target_mod = getattr(module, "base_layer", module)
                        cls_name = module.__class__.__name__.lower()
                        target_cls = target_mod.__class__.__name__.lower()
                        is_conv1d = (
                            "conv1d" in cls_name
                            or "conv1d" in target_cls
                            or hasattr(module, "nf")
                            or hasattr(target_mod, "nf")
                        )
                        wt = getattr(target_mod, "weight", getattr(module, "weight", None))
                        if wt is None:
                            return
                        in_f = wt.shape[0] if is_conv1d else wt.shape[1]

                        x_fp32 = x.float()
                        x_flat = x_fp32.reshape(-1, in_f)
                        # AASVD-FIX-12
                        valid = x_flat.abs().amax(dim=-1) > 0
                        if valid.sum() == 0:
                            return
                        x_active = x_flat[valid]
                        n_tokens = max(1, x_active.size(0))
                        batch_cov = (x_active.t() @ x_active / n_tokens).cpu()
                        if self.cov_dtype == torch.float16:
                            batch_cov = batch_cov.half()

                        # AASVD-FIX-06: EMA with first-observation handling
                        if layer_name not in covariances:
                            covariances[layer_name] = batch_cov
                            update_counts[layer_name] = 1
                        else:
                            beta = self.ema_beta
                            covariances[layer_name] = (
                                beta * covariances[layer_name] + (1.0 - beta) * batch_cov
                            )
                            update_counts[layer_name] += 1

                return hook

            for name, module in self.model.named_modules():
                if self._is_target_module(name, module):
                    hooks.append(module.register_forward_hook(make_hook(name)))

            try:
                inputs = self.tokenizer(
                    self.calibration_texts,
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                    max_length=512,
                )
                inputs = {k: v.to(self.device) for k, v in inputs.items() if hasattr(v, "to")}
                with torch.no_grad():
                    self.model(**inputs)
            except Exception as e:
                logger.debug("AASVD calibration forward failed: %s", e)
            finally:
                for h in hooks:
                    h.remove()

        # ---- Precompute covariance square roots ----
        cov_factor_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
        for name, cov in covariances.items():
            try:
                cov_dev = cov.to(self.device, torch.float32)
                cov_sym = 0.5 * (cov_dev + cov_dev.T)
                eigvals, eigvecs = torch.linalg.eigh(cov_sym)
                eps = 1e-6 * torch.max(eigvals).clamp(min=1e-12) + 1e-8
                eigvals_reg = eigvals.clamp(min=eps)
                cov_sqrt = eigvecs @ torch.diag(torch.sqrt(eigvals_reg)) @ eigvecs.T
                cov_inv_sqrt = eigvecs @ torch.diag(1.0 / torch.sqrt(eigvals_reg)) @ eigvecs.T
                cov_factor_cache[name] = (cov_sqrt, cov_inv_sqrt)
            except Exception as e:
                logger.debug("AASVD covariance factor failed for %s: %s", name, e)

        # ---- Compress each target ----
        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            target_mod = getattr(module, "base_layer", module)
            try:
                for p in target_mod.parameters():
                    self._orig_storage_bytes += p.numel() * p.element_size()
            except Exception:
                pass

            cov_factors = cov_factor_cache.get(name)
            cov = covariances.get(name) if cov_factors is None else None

            svd_layer, err, r = self.compress_linear(
                module,
                rank=self.rank,
                cov=cov,
                cov_factors=cov_factors,
            )
            self.layer_errors[name] = err
            self.layer_ranks[name] = r

            too_big = (err > self.max_error) and not self.use_strs
            if svd_layer is None or too_big:
                if self.fallback_to_original:
                    logger.warning(
                        "AASVD: %s (err=%.4f > %.4f); falling back to original.",
                        name,
                        err,
                        self.max_error,
                    )
                    self.fallback_layers.append(name)
                    self._comp_storage_bytes += (
                        self._orig_storage_bytes
                    )  # approximation; corrected below
                    continue
                logger.warning(
                    "AASVD: %s exceeds max_error but fallback is False.",
                    name,
                )
                if svd_layer is None:
                    self.fallback_layers.append(name)
                    self._comp_storage_bytes += self._orig_storage_bytes
                    continue

            if hasattr(module, "base_layer"):
                module.base_layer = svd_layer
                for attr in ("lora_A", "lora_B"):
                    sub = getattr(module, attr, None)
                    if sub is None:
                        continue
                    params = sub.parameters() if hasattr(sub, "parameters") else iter([sub])
                    for p in params:
                        if hasattr(p, "requires_grad"):
                            p.requires_grad = True
            else:
                replace_module(self.model, name, svd_layer)
            self.compressed_layers.append(name)

            try:
                for p in svd_layer.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
            except Exception:
                pass

        # Freeze fallback layers
        for name, mod in self.model.named_modules():
            if name in self.fallback_layers:
                target = getattr(mod, "base_layer", mod)
                for p in target.parameters():
                    p.requires_grad = False

        # ---- Free memory ----
        covariances.clear()
        cov_factor_cache.clear()
        del covariances
        del cov_factor_cache
        if torch.cuda.is_available():
            with contextlib.suppress(Exception):
                torch.cuda.empty_cache()

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )

        # AASVD-FIX-11: proper storage accounting
        # comp_storage was approximated for fallback layers; recompute honestly.
        comp_bytes_real = 0
        for name, mod in self.model.named_modules():
            if name in self.compressed_layers or name in self.fallback_layers:
                target = getattr(mod, "base_layer", mod)
                try:
                    for p in target.parameters():
                        comp_bytes_real += p.numel() * p.element_size()
                except Exception:
                    pass
        self._comp_storage_bytes = comp_bytes_real

        storage_ratio = (
            self._orig_storage_bytes / self._comp_storage_bytes
            if self._comp_storage_bytes > 0
            else 1.0
        )
        effective_ratio = (
            self._comp_storage_bytes / self._orig_storage_bytes
            if self._orig_storage_bytes > 0
            else 1.0
        )

        stats = self.get_compression_stats()
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "avg_reconstruction_error": avg_err,
                "use_strs": self.use_strs,
                "rank": self.rank,
                "lora_rank": self.lora_rank,
                "lora_alpha": self.lora_alpha,
                "orig_storage_bytes": self._orig_storage_bytes,
                "comp_storage_bytes": self._comp_storage_bytes,
                "storage_compression_ratio": round(storage_ratio, 3),
                "storage_reduction_pct": round(100.0 * (1.0 - effective_ratio), 2),
                "method": "aasvd_activation_aware_svd",
            }
        )
        logger.info(
            "AASVD: %d compressed, %d fallback, avg_err=%.4f, storage_ratio=%.2fx",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            storage_ratio,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def compress_linear(self, layer: nn.Module, rank: int | None = None, cov: torch.Tensor | None = None, cov_factors: tuple[torch.Tensor, torch.Tensor] | None = None) ‑> tuple[SVDLinear | None, float, int]

Expand source code

def compress_linear(
    self,
    layer: nn.Module,
    rank: int | None = None,
    cov: torch.Tensor | None = None,
    cov_factors: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> tuple[SVDLinear | None, float, int]:
    try:
        target_mod = getattr(layer, "base_layer", layer)
        cls_name = layer.__class__.__name__.lower()
        target_cls = target_mod.__class__.__name__.lower()
        is_conv1d = (
            "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(layer, "nf")
            or hasattr(target_mod, "nf")
        )
        weight_tensor = getattr(target_mod, "weight", getattr(layer, "weight", None))
        if weight_tensor is None:
            return None, float("inf"), 0

        raw_w = weight_tensor.data
        if is_conv1d:
            in_feat, out_feat = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()
        else:
            out_feat, in_feat = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float()

        r_default = rank if rank is not None else self.rank
        r = get_dynamic_compression_rank(
            in_feat,
            out_feat,
            model_or_family=self.model,
            default_rank=r_default,
            comp_config=self.comp_config,
        )
        max_allowed_rank = max(
            1,
            min(
                in_feat,
                out_feat,
                int(0.75 * (in_feat * out_feat) / (in_feat + out_feat)),
            ),
        )
        r = min(r, max_allowed_rank)
        W_norm = torch.norm(W, p="fro").item()

        # ---- covariance preconditioning ----
        cov_sqrt = cov_inv_sqrt = None
        if cov_factors is not None:
            cov_sqrt, cov_inv_sqrt = cov_factors
        elif cov is not None and cov.shape == (in_feat, in_feat):
            try:
                cov_fp32 = cov.to(self.device, torch.float32)
                cov_sym = 0.5 * (cov_fp32 + cov_fp32.T)
                eigvals, eigvecs = torch.linalg.eigh(cov_sym)
                eps = 1e-6 * torch.max(eigvals).clamp(min=1e-12) + 1e-8
                eigvals_reg = eigvals.clamp(min=eps)
                cov_sqrt = eigvecs @ torch.diag(torch.sqrt(eigvals_reg)) @ eigvecs.T
                cov_inv_sqrt = eigvecs @ torch.diag(1.0 / torch.sqrt(eigvals_reg)) @ eigvecs.T
            except Exception as e:
                logger.debug("AASVD covariance sqrt failed for %s: %s", e, e)
                cov_sqrt, cov_inv_sqrt = None, None

        W_weighted = W @ cov_sqrt if cov_sqrt is not None else W

        W_w_dev = W_weighted.to(self.device, torch.float32)
        U, S, Vt = torch.linalg.svd(W_w_dev, full_matrices=False)

        # ---- STRS / fixed rank ----
        if not self.use_strs:
            selected_r = min(r, len(S))
            sqrt_S = torch.sqrt(torch.clamp(S[:selected_r], min=1e-12))
            best_Ur = U[:, :selected_r] * sqrt_S.unsqueeze(0)
            best_Vr = sqrt_S.unsqueeze(1) * Vt[:selected_r, :]
            if cov_inv_sqrt is not None:
                best_Vr = best_Vr @ cov_inv_sqrt
            best_err = float(
                torch.norm(W - best_Ur @ best_Vr, p="fro").item() / max(W_norm, 1e-12)
            )
        else:
            r_floor = max(1, min(r, max_allowed_rank, len(S)))
            r_ceiling = max_allowed_rank
            cand_ranks = self._build_strs_candidates(r_floor, r_ceiling)
            cand_ranks = [c for c in cand_ranks if c <= len(S)]

            best_err = float("inf")
            selected_r = cand_ranks[0] if cand_ranks else r_floor
            best_Ur = best_Vr = None

            for test_r in cand_ranks:
                sqrt_S = torch.sqrt(torch.clamp(S[:test_r], min=1e-12))
                U_cand = U[:, :test_r] * sqrt_S.unsqueeze(0)
                V_cand = sqrt_S.unsqueeze(1) * Vt[:test_r, :]
                if cov_inv_sqrt is not None:
                    V_cand = V_cand @ cov_inv_sqrt
                W_cand = U_cand @ V_cand
                test_err = float(torch.norm(W - W_cand, p="fro").item() / max(W_norm, 1e-12))
                if test_err < best_err:
                    best_err = test_err
                    selected_r = test_r
                    best_Ur = U_cand
                    best_Vr = V_cand
                if test_err <= self.max_error:
                    break

            if best_Ur is None or best_err > self.max_error:
                # Fallback to break-even rank
                be_r = max(
                    1,
                    min(
                        int(0.75 * (in_feat * out_feat) / (in_feat + out_feat)),
                        len(S),
                    ),
                )
                selected_r = be_r
                sqrt_S = torch.sqrt(torch.clamp(S[:be_r], min=1e-12))
                best_Ur = U[:, :be_r] * sqrt_S.unsqueeze(0)
                best_Vr = sqrt_S.unsqueeze(1) * Vt[:be_r, :]
                if cov_inv_sqrt is not None:
                    best_Vr = best_Vr @ cov_inv_sqrt
                best_err = float(
                    torch.norm(W - best_Ur @ best_Vr, p="fro").item() / max(W_norm, 1e-12)
                )

        dtype = weight_tensor.dtype
        device = weight_tensor.device
        bias_tensor = getattr(target_mod, "bias", getattr(layer, "bias", None))

        svd_layer = SVDLinear(
            in_features=in_feat,
            out_features=out_feat,
            rank=selected_r,
            bias=bias_tensor is not None,
            dtype=dtype,
            device=device,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )
        svd_layer.set_low_rank_factors(
            best_Ur.to(device=device, dtype=dtype),
            best_Vr.to(device=device, dtype=dtype),
        )
        if bias_tensor is not None and svd_layer.bias is not None:
            svd_layer.bias.data.copy_(bias_tensor.data.to(svd_layer.bias.dtype))

        svd_layer.reconstruction_error = best_err
        svd_layer.rank = selected_r
        return svd_layer, best_err, selected_r
    except Exception as e:
        logger.warning("AASVD factorization failed: %s", e)
        return None, float("inf"), 0

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class AATC_XKVCache (subspaces: dict[int, SharedSubspace], layer_to_group: dict[int, int], sink_tokens: int = 1, preserve_last_token: bool = True)

Expand source code

class AATC_XKVCache(_HFDynamicCache if _HAS_HF_CACHE else object):
    """
    Drop-in replacement for HF DynamicCache that transparently compresses K/V
    via AATC-XKV.

    Storage per (layer, batch, head, token):
      - q_k: int8 coefficients (r_k,) + fp32 scale vector (r_k,)
      - q_v: int8 coefficients (r_v,) + fp32 scale vector (r_v,)
      - sink_k/sink_v (fp32) for the first `sink_tokens` tokens of the sequence
      - last_k/last_v (fp32) only during PREFILL (T_new > 1) when
        `preserve_last_token=True`

    Reconstruction:
      - `_reconstruction_cache[layer_idx]` holds the fully-reconstructed
        fp16 K/V up to the last decompressed entry, so decoding only pays
        O(new_tokens · r · D) per step.
    """

    def __init__(
        self,
        subspaces: dict[int, SharedSubspace],
        layer_to_group: dict[int, int],
        sink_tokens: int = 1,
        preserve_last_token: bool = True,
    ):
        super().__init__()
        self.subspaces = subspaces
        self.layer_to_group = layer_to_group
        self.sink_tokens = max(0, int(sink_tokens))
        self.preserve_last_token = bool(preserve_last_token)
        self.enabled = True

        self._compressed: dict[int, list[dict[str, torch.Tensor]]] = {}
        # Reconstruction cache: layer_idx -> (full_k, full_v, n_entries_seen)
        self._reconstruction_cache: dict[int, tuple[torch.Tensor, torch.Tensor, int]] = {}

        self._total_original_bytes = 0
        self._total_compressed_bytes = 0

    # ------------------------------------------------------------------
    # HF API
    # ------------------------------------------------------------------
    def __len__(self) -> int:
        # AXKV-FIX-04: sum over all entries
        if not self._compressed:
            return 0
        total = 0
        for layer_idx, entries in self._compressed.items():
            n = 0
            for e in entries:
                if e.get("sink_k") is not None:
                    n += e["sink_k"].shape[2]
                n += e["q_k"].shape[2]
                if e.get("last_k") is not None:
                    n += e["last_k"].shape[2]
            total = max(total, n)
        return total

    def __getitem__(self, layer_idx: int):
        if layer_idx not in self._compressed:
            raise KeyError(f"AATC_XKVCache: no cache for layer {layer_idx}")
        return self._reconstruct_full(layer_idx)

    def __iter__(self):
        for layer_idx in sorted(self._compressed.keys()):
            yield self._reconstruct_full(layer_idx)

    def __contains__(self, layer_idx: int) -> bool:
        return layer_idx in self._compressed

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs: dict[str, Any] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Append new K/V, return the full K/V needed for attention."""
        if not self.enabled or layer_idx not in self.layer_to_group:
            return super().update(key_states, value_states, layer_idx, cache_kwargs)

        sub = self.subspaces.get(self.layer_to_group[layer_idx])
        if sub is None:
            return super().update(key_states, value_states, layer_idx, cache_kwargs)

        entry = self._new_entry(key_states, value_states, sub, layer_idx)
        self._compressed.setdefault(layer_idx, []).append(entry)

        # AXKV-FIX-03: use reconstruction cache
        cached = self._reconstruction_cache.get(layer_idx)
        n_entries = len(self._compressed[layer_idx])

        if cached is not None and cached[2] == n_entries - 1:
            full_k, full_v, _ = cached
            new_k, new_v = self._reconstruct_entry(entry, sub, key_states.dtype)
            full_k = torch.cat([full_k, new_k], dim=2)
            full_v = torch.cat([full_v, new_v], dim=2)
        else:
            full_k, full_v = self._reconstruct_full(layer_idx)

        self._reconstruction_cache[layer_idx] = (full_k, full_v, n_entries)
        return full_k, full_v

    def get_seq_length(self, layer_idx: int = 0) -> int:
        entries = self._compressed.get(layer_idx, [])
        n = 0
        for e in entries:
            if e["sink_k"] is not None:
                n += e["sink_k"].shape[2]
            n += e["q_k"].shape[2]
            if e["last_k"] is not None:
                n += e["last_k"].shape[2]
        return n

    def get_max_length(self) -> int | None:
        return None

    def get_usable_length(self, new_seq_length: int, layer_idx: int = 0) -> int:
        return self.get_seq_length(layer_idx)

    def get_max_cache_shape(self) -> int:
        return -1

    def reset(self) -> None:
        self._compressed.clear()
        self._reconstruction_cache.clear()
        self._total_original_bytes = 0
        self._total_compressed_bytes = 0

    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
        """
        AXKV-FIX-05: correct beam-search reorder. Rebuilds the reconstruction
        cache lazily after reordering the compressed entries.
        """
        with torch.no_grad():
            for layer_idx, entries in self._compressed.items():
                for e in entries:
                    for k in ("sink_k", "sink_v", "last_k", "last_v"):
                        if e[k] is not None:
                            e[k] = e[k].index_select(0, beam_idx.to(e[k].device))
                    for k in ("q_k", "q_v"):
                        e[k] = e[k].index_select(0, beam_idx.to(e[k].device))
            # Invalidate reconstruction cache
            self._reconstruction_cache.clear()

    # ------------------------------------------------------------------
    # Entry creation
    # ------------------------------------------------------------------
    def _new_entry(
        self,
        k: torch.Tensor,
        v: torch.Tensor,
        sub: SharedSubspace,
        layer_idx: int,
    ) -> dict[str, torch.Tensor]:
        B, H_kv, T_new, D = k.shape

        # AXKV-FIX-12: sinks are prepended only once — the first time we
        # ever see this layer's cache.
        is_first_chunk = len(self._compressed.get(layer_idx, [])) == 0
        sink_prefix = min(self.sink_tokens, T_new) if is_first_chunk else 0

        if sink_prefix > 0:
            sink_k = k[:, :, :sink_prefix, :].to(torch.float32).contiguous()
            sink_v = v[:, :, :sink_prefix, :].to(torch.float32).contiguous()
            rest_k = k[:, :, sink_prefix:, :]
            rest_v = v[:, :, sink_prefix:, :]
        else:
            sink_k = None
            sink_v = None
            rest_k = k
            rest_v = v

        # AXKV-FIX-02: preserve_last_token only during prefill (T_new > 1).
        last_k = None
        last_v = None
        if self.preserve_last_token and T_new > 1 and rest_k.size(2) > 0:
            last_k = rest_k[:, :, -1:, :].to(torch.float32).contiguous()
            last_v = rest_v[:, :, -1:, :].to(torch.float32).contiguous()
            rest_k = rest_k[:, :, :-1, :]
            rest_v = rest_v[:, :, :-1, :]

        V_k = sub.V_k.to(k.device, torch.float32)
        V_v = sub.V_v.to(v.device, torch.float32)
        mean_k = sub.mean_k.to(k.device, torch.float32)
        mean_v = sub.mean_v.to(v.device, torch.float32)

        if rest_k.size(2) == 0:
            r_k = V_k.size(0)
            r_v = V_v.size(0)
            q_k = torch.zeros((B, H_kv, 0, r_k), dtype=torch.int8, device=k.device)
            q_v = torch.zeros((B, H_kv, 0, r_v), dtype=torch.int8, device=v.device)
            s_k = torch.ones(r_k, dtype=torch.float32, device=k.device)
            s_v = torch.ones(r_v, dtype=torch.float32, device=v.device)
        else:
            rest_k_f = rest_k.to(torch.float32) - mean_k.view(1, 1, 1, -1)
            rest_v_f = rest_v.to(torch.float32) - mean_v.view(1, 1, 1, -1)
            U_k = torch.matmul(rest_k_f, V_k.t())
            U_v = torch.matmul(rest_v_f, V_v.t())
            q_k, s_k = quantize_coefficients(U_k, sub.bit_alloc_k)
            q_v, s_v = quantize_coefficients(U_v, sub.bit_alloc_v)

        entry = {
            "q_k": q_k,
            "s_k": s_k,
            "q_v": q_v,
            "s_v": s_v,
            "sink_k": sink_k,
            "sink_v": sink_v,
            "last_k": last_k,
            "last_v": last_v,
        }

        # AXKV-FIX-09: byte accounting
        with torch.no_grad():
            orig_b = (k.numel() + v.numel()) * k.element_size()
            comp_b = (
                q_k.numel() * q_k.element_size()
                + q_v.numel() * q_v.element_size()
                + s_k.numel() * s_k.element_size()
                + s_v.numel() * s_v.element_size()
                + (sink_k.numel() * sink_k.element_size() if sink_k is not None else 0)
                + (sink_v.numel() * sink_v.element_size() if sink_v is not None else 0)
                + (last_k.numel() * last_k.element_size() if last_k is not None else 0)
                + (last_v.numel() * last_v.element_size() if last_v is not None else 0)
            )
            self._total_original_bytes += orig_b
            self._total_compressed_bytes += comp_b

        return entry

    # ------------------------------------------------------------------
    # Reconstruction helpers
    # ------------------------------------------------------------------
    def _reconstruct_entry(
        self,
        e: dict[str, torch.Tensor],
        sub: SharedSubspace,
        out_dtype: torch.dtype,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """AXKV-FIX-11: reconstruct one entry in fp32, cast at return."""
        V_k = sub.V_k.to(torch.float32)
        V_v = sub.V_v.to(torch.float32)
        mean_k = sub.mean_k.to(torch.float32)
        mean_v = sub.mean_v.to(torch.float32)

        parts_k, parts_v = [], []

        if e["sink_k"] is not None:
            parts_k.append(e["sink_k"].to(out_dtype))
            parts_v.append(e["sink_v"].to(out_dtype))

        if e["q_k"].size(2) > 0:
            U_k = dequantize_coefficients(e["q_k"], e["s_k"])
            U_v = dequantize_coefficients(e["q_v"], e["s_v"])
            device = U_k.device
            rec_k = torch.matmul(U_k, V_k.to(device)) + mean_k.to(device).view(1, 1, 1, -1)
            rec_v = torch.matmul(U_v, V_v.to(device)) + mean_v.to(device).view(1, 1, 1, -1)
            parts_k.append(rec_k.to(out_dtype))
            parts_v.append(rec_v.to(out_dtype))

        if e["last_k"] is not None:
            parts_k.append(e["last_k"].to(out_dtype))
            parts_v.append(e["last_v"].to(out_dtype))

        if not parts_k:
            return torch.empty(0), torch.empty(0)

        k_entry = torch.cat(parts_k, dim=2) if len(parts_k) > 1 else parts_k[0]
        v_entry = torch.cat(parts_v, dim=2) if len(parts_v) > 1 else parts_v[0]
        return k_entry, v_entry

    def _reconstruct_full(
        self,
        layer_idx: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        entries = self._compressed.get(layer_idx, [])
        if not entries:
            raise RuntimeError("AATC_XKVCache: no entries for layer.")

        sub = self.subspaces[self.layer_to_group[layer_idx]]

        # Infer dtype from the first entry's sink or quantized buffer
        ref = entries[0].get("sink_k")
        if ref is None:
            ref = entries[0].get("last_k")
        if ref is None:
            # Fall back to the dtype of the shared subspace
            out_dtype = sub.V_k.dtype
        else:
            out_dtype = ref.dtype

        k_list, v_list = [], []
        for e in entries:
            k_i, v_i = self._reconstruct_entry(e, sub, out_dtype)
            if k_i.numel() > 0:
                k_list.append(k_i)
                v_list.append(v_i)

        if not k_list:
            raise RuntimeError("AATC_XKVCache: empty reconstruction.")

        full_k = torch.cat(k_list, dim=2) if len(k_list) > 1 else k_list[0]
        full_v = torch.cat(v_list, dim=2) if len(v_list) > 1 else v_list[0]
        return full_k, full_v

    # ------------------------------------------------------------------
    # Public metrics
    # ------------------------------------------------------------------
    def memory_savings_ratio(self) -> float:
        if self._total_compressed_bytes == 0:
            return 1.0
        return self._total_original_bytes / self._total_compressed_bytes

    def memory_bytes(self) -> dict[str, int]:
        return {
            "original": self._total_original_bytes,
            "compressed": self._total_compressed_bytes,
        }

Drop-in replacement for HF DynamicCache that transparently compresses K/V via AATC-XKV.

Storage per (layer, batch, head, token): - q_k: int8 coefficients (r_k,) + fp32 scale vector (r_k,) - q_v: int8 coefficients (r_v,) + fp32 scale vector (r_v,) - sink_k/sink_v (fp32) for the first sink_tokens tokens of the sequence - last_k/last_v (fp32) only during PREFILL (T_new > 1) when preserve_last_token=True

Reconstruction -----= - _reconstruction_cache[layer_idx] holds the fully-reconstructed fp16 K/V up to the last decompressed entry, so decoding only pays O(new_tokens · r · D) per step.

Ancestors

transformers.cache_utils.DynamicCache
transformers.cache_utils.Cache

Methods

def get_max_cache_shape(self) ‑> int

Expand source code

def get_max_cache_shape(self) -> int:
    return -1

Returns maximum sequence length of the cache object. Dynamic caches do not have a maximum length.

def get_max_length(self) ‑> int | None

Expand source code

def get_max_length(self) -> int | None:
    return None

def get_seq_length(self, layer_idx: int = 0) ‑> int

Expand source code

def get_seq_length(self, layer_idx: int = 0) -> int:
    entries = self._compressed.get(layer_idx, [])
    n = 0
    for e in entries:
        if e["sink_k"] is not None:
            n += e["sink_k"].shape[2]
        n += e["q_k"].shape[2]
        if e["last_k"] is not None:
            n += e["last_k"].shape[2]
    return n

Returns the sequence length of the cache for the given layer.

def get_usable_length(self, new_seq_length: int, layer_idx: int = 0) ‑> int

Expand source code

def get_usable_length(self, new_seq_length: int, layer_idx: int = 0) -> int:
    return self.get_seq_length(layer_idx)

def memory_bytes(self) ‑> dict[str, int]

Expand source code

def memory_bytes(self) -> dict[str, int]:
    return {
        "original": self._total_original_bytes,
        "compressed": self._total_compressed_bytes,
    }

def memory_savings_ratio(self) ‑> float

Expand source code

def memory_savings_ratio(self) -> float:
    if self._total_compressed_bytes == 0:
        return 1.0
    return self._total_original_bytes / self._total_compressed_bytes

def reorder_cache(self, beam_idx: torch.LongTensor) ‑> None

Expand source code

def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
    """
    AXKV-FIX-05: correct beam-search reorder. Rebuilds the reconstruction
    cache lazily after reordering the compressed entries.
    """
    with torch.no_grad():
        for layer_idx, entries in self._compressed.items():
            for e in entries:
                for k in ("sink_k", "sink_v", "last_k", "last_v"):
                    if e[k] is not None:
                        e[k] = e[k].index_select(0, beam_idx.to(e[k].device))
                for k in ("q_k", "q_v"):
                    e[k] = e[k].index_select(0, beam_idx.to(e[k].device))
        # Invalidate reconstruction cache
        self._reconstruction_cache.clear()

AXKV-FIX-05: correct beam-search reorder. Rebuilds the reconstruction cache lazily after reordering the compressed entries.

def reset(self) ‑> None

Expand source code

def reset(self) -> None:
    self._compressed.clear()
    self._reconstruction_cache.clear()
    self._total_original_bytes = 0
    self._total_compressed_bytes = 0

Recursively reset all layers tensors

def update(self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, cache_kwargs: dict[str, Any] | None = None) ‑> tuple[torch.Tensor, torch.Tensor]

Expand source code

def update(
    self,
    key_states: torch.Tensor,
    value_states: torch.Tensor,
    layer_idx: int,
    cache_kwargs: dict[str, Any] | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Append new K/V, return the full K/V needed for attention."""
    if not self.enabled or layer_idx not in self.layer_to_group:
        return super().update(key_states, value_states, layer_idx, cache_kwargs)

    sub = self.subspaces.get(self.layer_to_group[layer_idx])
    if sub is None:
        return super().update(key_states, value_states, layer_idx, cache_kwargs)

    entry = self._new_entry(key_states, value_states, sub, layer_idx)
    self._compressed.setdefault(layer_idx, []).append(entry)

    # AXKV-FIX-03: use reconstruction cache
    cached = self._reconstruction_cache.get(layer_idx)
    n_entries = len(self._compressed[layer_idx])

    if cached is not None and cached[2] == n_entries - 1:
        full_k, full_v, _ = cached
        new_k, new_v = self._reconstruct_entry(entry, sub, key_states.dtype)
        full_k = torch.cat([full_k, new_k], dim=2)
        full_v = torch.cat([full_v, new_v], dim=2)
    else:
        full_k, full_v = self._reconstruct_full(layer_idx)

    self._reconstruction_cache[layer_idx] = (full_k, full_v, n_entries)
    return full_k, full_v

Append new K/V, return the full K/V needed for attention.

class AATC_XKVCompressor (model: nn.Module, group_size: int = 3, target_bits_per_component: float = 4.0, energy_threshold: float = 0.95, sink_tokens: int = 1, preserve_last_token: bool = True, calib_batches: list[Any] | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("aatc_xkv")
@COMPRESSOR_REGISTRY.register("aatc_xkv_compressor")
class AATC_XKVCompressor(BaseCompressor):
    # AXKV-FIX-01
    is_custom_attention: bool = False
    IS_DYNAMIC_KV_CACHE = True

    def __init__(
        self,
        model: nn.Module,
        group_size: int = 3,
        target_bits_per_component: float = 4.0,
        energy_threshold: float = 0.95,
        sink_tokens: int = 1,
        preserve_last_token: bool = True,
        calib_batches: list[Any] | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="aatc_xkv", **kwargs)
        self.cfg = cfg

        def _cget(key, default):
            if cfg is None:
                return default
            if isinstance(cfg, dict):
                return cfg.get(key, default)
            return getattr(cfg, key, default)

        self.group_size = int(_cget("group_size", group_size))
        self.target_bits = float(_cget("aatc_bits", target_bits_per_component))
        self.energy_threshold = float(_cget("energy_threshold", energy_threshold))
        self.sink_tokens = int(_cget("sink_tokens", sink_tokens))
        self.preserve_last_token = bool(_cget("preserve_last_token", preserve_last_token))

        self.subspaces: dict[int, SharedSubspace] = {}
        self.layer_to_group: dict[int, int] = {}
        self.calib_batches = calib_batches
        self.stats: dict[str, Any] = {}

        # Storage accounting
        self._orig_kv_bytes = 0
        self._comp_kv_bytes = 0

    # ------------------------------------------------------------------
    def _find_attention_modules(self) -> list[tuple[int, str, nn.Module]]:
        """
        AXKV-FIX-08: identify attention submodules with K/V projections.
        Avoids `value` attribute collisions by requiring explicit projection
        names or a fused GPT-2 `c_attn`.
        """
        attn_modules = []
        for name, mod in self.model.named_modules():
            name_lower = name.lower()
            if not any(k in name_lower for k in ("self_attn", "attention", "attn")):
                continue

            # Fused GPT-2 c_attn
            if hasattr(mod, "c_attn"):
                is_attn = True
            else:
                k_names = ("k_proj", "k_lin", "key", "W_k")
                v_names = ("v_proj", "v_lin", "value", "W_v")
                has_k = any(hasattr(mod, a) for a in k_names)
                has_v = any(hasattr(mod, a) for a in v_names)
                # Reject if the attribute exists but is a method (not a module)
                if has_k:
                    has_k = any(
                        isinstance(getattr(mod, a, None), nn.Module)
                        for a in k_names
                        if hasattr(mod, a)
                    )
                if has_v:
                    has_v = any(
                        isinstance(getattr(mod, a, None), nn.Module)
                        for a in v_names
                        if hasattr(mod, a)
                    )
                is_attn = has_k and has_v

            if not is_attn:
                continue

            idx = len(attn_modules)
            for part in name.split("."):
                if part.isdigit():
                    idx = int(part)
                    break
            attn_modules.append((idx, name, mod))

        attn_modules.sort(key=lambda x: x[0])
        return attn_modules

    # ------------------------------------------------------------------
    def _build_subspaces(
        self,
        calib: dict[str, dict[int, torch.Tensor]],
        attn_modules: list[tuple[int, str, nn.Module]],
    ) -> None:
        layer_indices = [idx for idx, _, _ in attn_modules]
        G = max(1, self.group_size)
        groups = [layer_indices[i : i + G] for i in range(0, len(layer_indices), G)]

        for g_id, g_layers in enumerate(groups):
            for l_idx in g_layers:
                self.layer_to_group[l_idx] = g_id

            k_list = [calib["k"][l] for l in g_layers if calib["k"][l].numel() > 0]
            v_list = [calib["v"][l] for l in g_layers if calib["v"][l].numel() > 0]
            gk_list = [calib["gk"][l] for l in g_layers if calib["gk"][l].numel() > 0]
            gv_list = [calib["gv"][l] for l in g_layers if calib["gv"][l].numel() > 0]

            if not k_list or not v_list:
                continue

            K_stack = torch.cat(k_list, 0)
            V_stack = torch.cat(v_list, 0)
            gK_stack = torch.cat(gk_list, 0) if gk_list else torch.zeros_like(K_stack)
            gV_stack = torch.cat(gv_list, 0) if gv_list else torch.zeros_like(V_stack)

            D = K_stack.size(-1)
            rank_cap = max(4, D // 4)

            V_k, S_k, mean_k = _joint_svd_shared_basis(K_stack, rank_cap, self.energy_threshold)
            V_v, S_v, mean_v = _joint_svd_shared_basis(V_stack, rank_cap, self.energy_threshold)

            r_k = V_k.size(0)
            r_v = V_v.size(0)

            proj_gk = gK_stack @ V_k.t()
            proj_k = K_stack @ V_k.t()
            sens_k = proj_gk.pow(2).mean(0) + proj_k.pow(2).mean(0)

            proj_gv = gV_stack @ V_v.t()
            proj_v = V_stack @ V_v.t()
            sens_v = proj_gv.pow(2).mean(0) + proj_v.pow(2).mean(0)

            total_bits_k = int(round(r_k * self.target_bits))
            total_bits_v = int(round(r_v * self.target_bits))
            alloc_k = reverse_waterfilling(S_k.pow(2), sens_k, total_bits_k, 0, 8)
            alloc_v = reverse_waterfilling(S_v.pow(2), sens_v, total_bits_v, 0, 8)

            self.subspaces[g_id] = SharedSubspace(
                V_k=V_k,
                V_v=V_v,
                mean_k=mean_k,
                mean_v=mean_v,
                bit_alloc_k=alloc_k,
                bit_alloc_v=alloc_v,
                group_id=g_id,
                layer_indices=g_layers,
            )

    # ------------------------------------------------------------------
    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        device = next(self.model.parameters()).device
        attn_modules = self._find_attention_modules()

        if not attn_modules:
            logger.warning("AATC-XKV: no attention modules found.")
            self.stats = {"status": "NO_ATTN_MODULES", "num_compressed_layers": 0}
            return self.model, self.stats

        if not self.calib_batches:
            logger.warning("AATC-XKV: no calibration batches supplied; using seeded synthetic.")
            self.calib_batches = self._synthetic_calib(device)

        calib = _collect_calibration_data(
            self.model,
            attn_modules,
            self.calib_batches,
            device,
        )
        self._build_subspaces(calib, attn_modules)

        factory = self._make_cache_factory()
        self.model._aatc_xkv_cache_factory = factory
        self.model.aatc_xkv_runtime = factory

        # AXKV-FIX-13: only mark support if transformers uses it
        try:
            import transformers
            from packaging import version

            if version.parse(transformers.__version__) >= version.parse("4.36.0"):
                self.model._supports_cache_class = True
        except Exception:
            self.model._supports_cache_class = True

        # AXKV-FIX-09: single, honest bit-accounting pass
        total_orig_bits = 0
        total_comp_bits = 0
        for g_id, sub in self.subspaces.items():
            D = sub.V_k.size(1)
            n_layers = len(sub.layer_indices)
            # Original: 16 bits per K and per V element
            total_orig_bits += 16 * (2 * D) * n_layers
            # Compressed: bit_alloc per K component + per V component,
            # per token, per layer
            total_comp_bits += (
                int(sub.bit_alloc_k.sum().item()) + int(sub.bit_alloc_v.sum().item())
            ) * n_layers

        ratio = total_orig_bits / max(1, total_comp_bits)

        self.stats = {
            "status": "PASS",
            "method": "aatc_xkv",
            "group_size": self.group_size,
            "num_groups": len(self.subspaces),
            "num_compressed_layers": len(attn_modules),
            "num_fallback_layers": 0,
            "target_bits_per_component": self.target_bits,
            "energy_threshold": self.energy_threshold,
            "theoretical_kv_compression_ratio": round(ratio, 2),
            "total_params_before": self._original_params,
            "total_params_after": self._original_params,
            "weights_modified": False,
        }
        self.model.compression_stats = self.stats
        logger.info(
            "AATC-XKV: %d layers, %d groups, theoretical KV compression %.2fx",
            len(attn_modules),
            len(self.subspaces),
            ratio,
        )
        return self.model, self.stats

    # ------------------------------------------------------------------
    def _make_cache_factory(self):
        subspaces = self.subspaces
        layer_to_group = self.layer_to_group
        sink_tokens = self.sink_tokens
        preserve_last = self.preserve_last_token

        def factory():
            return AATC_XKVCache(
                subspaces=subspaces,
                layer_to_group=layer_to_group,
                sink_tokens=sink_tokens,
                preserve_last_token=preserve_last,
            )

        return factory

    # ------------------------------------------------------------------
    @staticmethod
    def _synthetic_calib(device) -> list[Any]:
        # AXKV-FIX-10: deterministic seed
        logger.warning("AATC-XKV: synthetic calibration is not recommended.")
        gen = torch.Generator(device="cpu").manual_seed(42)
        batches = []
        for _ in range(8):
            ids = torch.randint(5, 30000, (1, 128), generator=gen)
            batches.append(ids)
        return batches

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var IS_DYNAMIC_KV_CACHE: The type of the None singleton.
var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class ADMMCompressor (model: nn.Module, sparsity: float = 0.5, quant_bits: int = 4, rho1: float = 0.1, rho2: float = 0.1, max_iter: int = 8, structured_24: bool = True, block_size: tuple[int, int] = (4, 4), max_error: float = 0.45, fallback_to_original: bool = True, target_attention_only: bool = False, target_modules: list[str] | None = None, quant_per_channel: bool = True, comp_config: Any = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("admm")
@COMPRESSOR_REGISTRY.register("admm_compressor")
class ADMMCompressor(BaseCompressor):
    """
    ADMM-based joint structured pruning + uniform quantization engine.
    """

    def __init__(
        self,
        model: nn.Module,
        sparsity: float = 0.5,
        quant_bits: int = 4,
        rho1: float = 0.1,
        rho2: float = 0.1,
        max_iter: int = 8,
        structured_24: bool = True,
        block_size: tuple[int, int] = (4, 4),
        max_error: float = 0.45,
        fallback_to_original: bool = True,
        target_attention_only: bool = False,
        target_modules: list[str] | None = None,
        quant_per_channel: bool = True,  # AC-FIX-02
        comp_config: Any = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="admm", **kwargs)
        cfg = comp_config or kwargs.get(
            "comp_config", kwargs.get("cfg", getattr(model, "config", None))
        )

        self.sparsity = max(
            0.0,
            min(
                float(_cfg_get(cfg, "admm_sparsity", _cfg_get(cfg, "sparsity", sparsity))),
                0.95,
            ),
        )
        self.quant_bits = int(
            _cfg_get(cfg, "admm_quant_bits", _cfg_get(cfg, "quant_bits", quant_bits))
        )
        self.rho1 = float(_cfg_get(cfg, "admm_rho1", _cfg_get(cfg, "rho1", rho1)))
        self.rho2 = float(_cfg_get(cfg, "admm_rho2", _cfg_get(cfg, "rho2", rho2)))
        self.max_iter = max(
            1, int(_cfg_get(cfg, "admm_max_iter", _cfg_get(cfg, "max_iter", max_iter)))
        )
        self.structured_24 = bool(
            _cfg_get(cfg, "admm_structured_24", _cfg_get(cfg, "structured_24", structured_24))
        )
        self.block_size = block_size
        self.max_error = float(
            _cfg_get(cfg, "admm_max_error", _cfg_get(cfg, "max_error", max_error))
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "admm_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "admm_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.target_modules = target_modules
        self.quant_per_channel = bool(_cfg_get(cfg, "admm_quant_per_channel", quant_per_channel))
        self.comp_config = cfg
        self._compressed_params: int = 0

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_sparsities: dict[str, float] = {}
        self.layer_actual_sparsities: dict[str, float] = {}

    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        if any(
            name.endswith(sfx)
            for sfx in [".base_layer", ".lora_A", ".lora_B", ".lora_dropout", ".linear"]
        ):
            return False

        target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
        cls_name = module.__class__.__name__.lower()

        weight_tensor = getattr(target_mod, "weight", getattr(module, "weight", None))
        if weight_tensor is None or not hasattr(weight_tensor, "dim") or weight_tensor.dim() != 2:
            return False

        is_linear_like = (
            isinstance(module, nn.Linear)
            or isinstance(target_mod, nn.Linear)
            or (
                Conv1D is not None
                and (isinstance(module, Conv1D) or isinstance(target_mod, Conv1D))
            )
            or hasattr(module, "nf")
            or hasattr(target_mod, "nf")
            or "loralinear" in cls_name
            or hasattr(module, "base_layer")
        )
        if not is_linear_like:
            return False

        name_lower = name.lower()
        if any(
            skip in name_lower
            for skip in [
                "lm_head",
                "classifier",
                "head",
                "output_layer",
                "score",
                "pooler",
                "pre_classifier",
            ]
        ):
            return False

        if self.target_attention_only:
            ffn_patterns = [
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            ]
            if any(p in name_lower for p in ffn_patterns):
                return False
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            attn_patterns = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
                "in_proj_qkv",
                "in_proj_z",
                "in_proj_b",
                "out_proj",
            ]
            if not any(p in name_lower for p in attn_patterns):
                return False
        else:
            if not should_compress_linear(self.model, name):
                return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    # ------------------------------------------------------------------
    # AC-FIX-02/10: per-channel quantization with outlier robustness.
    # ------------------------------------------------------------------
    def _quantize_projection(self, tensor: torch.Tensor) -> torch.Tensor:
        """
        Project continuous tensor onto a b-bit uniform quantization grid (FP32).
        When per-channel: each output row is scaled independently.
        """
        if self.quant_bits >= 32:
            return tensor
        scale_levels = float(2 ** (self.quant_bits - 1) - 1)
        if scale_levels <= 0:
            return tensor

        if self.quant_per_channel and tensor.dim() == 2:
            max_val = tensor.abs().amax(dim=1, keepdim=True).clamp(min=1e-6)
        else:
            max_val = tensor.abs().amax().clamp(min=1e-6)

        normalized = torch.clamp(tensor / max_val, -1.0, 1.0)
        quantized = torch.round(normalized * scale_levels) / scale_levels
        return quantized * max_val

    # ------------------------------------------------------------------
    # AC-FIX-03: no_grad projection helper.
    # ------------------------------------------------------------------
    def _prune_projection(self, tensor: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            if self.structured_24:
                pruned, _ = apply_24_structured_mask(tensor, saliency=tensor.abs(), dim=1)
                return pruned
            elif self.block_size != (1, 1):
                pruned, _ = apply_block_mask(
                    tensor, block_size=self.block_size, sparsity=self.sparsity
                )
                return pruned
            else:
                flat = tensor.abs().reshape(-1)
                k = max(1, int(flat.numel() * (1.0 - self.sparsity)))
                thr = torch.topk(flat, k, largest=True).values[-1]
                mask = (tensor.abs() >= thr).float()
                return tensor * mask

    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_sparsities = {}
        self.layer_actual_sparsities = {}

        targets = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not targets:
            logger.warning("ADMM: No compressible layers found.")
            return self.model, self.get_compression_stats()

        for name, module in targets:
            target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
            target_cls = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in target_cls

            raw_w = target_mod.weight.data
            orig_shape = raw_w.shape
            device = raw_w.device
            dtype = raw_w.dtype

            W0 = raw_w.t().float() if is_conv1d else raw_w.float()
            # AC-FIX-07: safe norm.
            norm_w0 = float(torch.norm(W0, p="fro").item())
            norm_w0 = max(norm_w0, 1e-12)

            best_W_final: torch.Tensor | None = None
            best_err = float("inf")

            try:
                with torch.no_grad():
                    # Baseline candidate.
                    Z0 = self._prune_projection(W0)
                    # AC-FIX-06: mask excludes non-finite.
                    finite_mask = torch.isfinite(Z0).float()
                    mask0 = (Z0 != 0.0).float() * finite_mask
                    W_cand = self._quantize_projection(Z0) * mask0
                    err0 = float(torch.norm(W0 - W_cand, p="fro").item() / norm_w0)
                    best_W_final = W_cand
                    best_err = err0

                    # ADMM state in FP32.
                    W = W0.clone()
                    Z = Z0.clone()
                    Q = self._quantize_projection(W0)
                    U1 = torch.zeros_like(W0)
                    U2 = torch.zeros_like(W0)

                    denom = 1.0 + self.rho1 + self.rho2

                    for it in range(self.max_iter):
                        # W-update.
                        W = (W0 + self.rho1 * (Z - U1) + self.rho2 * (Q - U2)) / denom

                        # Z-update (sparse).
                        Z = self._prune_projection(W + U1)

                        # Q-update (quantized).
                        Q = self._quantize_projection(W + U2)

                        # Dual updates.
                        U1 = U1 + (W - Z)
                        U2 = U2 + (W - Q)

                        # Non-finite guard.
                        if (
                            not torch.isfinite(W).all()
                            or not torch.isfinite(Z).all()
                            or not torch.isfinite(Q).all()
                        ):
                            logger.warning(
                                "ADMM: non-finite state at iter %d on %s",
                                it,
                                name,
                            )
                            break

                        # AC-FIX-04: track best iterate using ADMM state W.
                        finite_mask_curr = torch.isfinite(W).float()
                        mask_curr = (W != 0.0).float() * finite_mask_curr
                        W_curr = self._quantize_projection(W) * mask_curr
                        curr_err = float(torch.norm(W0 - W_curr, p="fro").item() / norm_w0)
                        if curr_err < best_err:
                            best_err = curr_err
                            best_W_final = W_curr

            except Exception as e:
                logger.warning("ADMM optimization failed on %s: %s", name, e)
                best_W_final = None
                best_err = 1.0

            W_final = best_W_final
            final_err = best_err

            self.layer_errors[name] = final_err
            self.layer_sparsities[name] = 0.5 if self.structured_24 else self.sparsity

            if W_final is None or not torch.isfinite(W_final).all() or final_err > self.max_error:
                logger.warning(
                    "ADMM reconstruction error %.4f on %s exceeded threshold %.4f.",
                    final_err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    logger.info("ADMM: fallback to original on %s.", name)
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

            if is_conv1d:
                res_w = W_final.t().contiguous().to(dtype=dtype, device=device)
            else:
                res_w = W_final.contiguous().to(dtype=dtype, device=device)

            if res_w.shape != orig_shape:
                logger.warning(
                    "ADMM shape mismatch on %s: %s vs %s. Skipping.",
                    name,
                    tuple(res_w.shape),
                    tuple(orig_shape),
                )
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                continue

            # AC-FIX-09: mutation under no_grad.
            with torch.no_grad():
                target_mod.weight.data.copy_(res_w)

            actual_zero_frac = float((res_w == 0).sum().item() / max(1, res_w.numel()))
            self.layer_actual_sparsities[name] = actual_zero_frac

            # AC-FIX-05: LoRA always trainable (recovery).
            has_lora = (hasattr(module, "lora_A") and module.lora_A is not None) or hasattr(
                module, "base_layer"
            )
            if has_lora:
                if hasattr(target_mod, "weight") and target_mod.weight is not None:
                    target_mod.weight.requires_grad = False
                if hasattr(target_mod, "bias") and target_mod.bias is not None:
                    target_mod.bias.requires_grad = False

                if hasattr(module, "lora_A") and module.lora_A is not None:
                    for p in module.lora_A.parameters():
                        p.requires_grad = True
                if hasattr(module, "lora_B") and module.lora_B is not None:
                    for p in module.lora_B.parameters():
                        p.requires_grad = True

            self.compressed_layers.append(name)

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        avg_actual_sparsity = (
            sum(self.layer_actual_sparsities.values()) / max(1, len(self.layer_actual_sparsities))
            if self.layer_actual_sparsities
            else 0.0
        )

        # AC-FIX-01: effective compression rate combines quantization and sparsity.
        quant_factor = 16.0 / max(1, self.quant_bits)
        sparse_factor = 1.0 / max(1e-6, 1.0 - avg_actual_sparsity)
        effective_storage_comp_rate = round(quant_factor * sparse_factor, 2)

        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "layer_errors": self.layer_errors,
                "layer_sparsities": self.layer_sparsities,
                "layer_actual_sparsities": self.layer_actual_sparsities,
                "actual_sparsity": avg_actual_sparsity,
                "sparsity_target": self.sparsity,
                "quant_bits": self.quant_bits,
                "quant_per_channel": self.quant_per_channel,
                "structured_24": self.structured_24,
                "quant_factor": round(quant_factor, 2),
                "sparse_factor": round(sparse_factor, 2),
                "effective_compression_rate": effective_storage_comp_rate,
                "admm_iterations": self.max_iter,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "target_attention_only": self.target_attention_only,
                "avg_reconstruction_error": avg_err,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "method": "admm_joint_prune_quant",
            }
        )
        logger.info(
            "ADMM compression complete: %d compressed, %d fallback. "
            "Avg error: %.4f, Actual sparsity: %.2f%%, Effective comp rate: %.2fx",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            avg_actual_sparsity * 100.0,
            effective_storage_comp_rate,
        )
        return self.model, stats

ADMM-based joint structured pruning + uniform quantization engine.

Ancestors

BaseCompressor

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class ADMMMaskedLinear (in_features: int, out_features: int, bias: bool = True, dtype: torch.dtype = torch.float32, device: torch.device | None = None)

Expand source code

class ADMMMaskedLinear(nn.Module):
    """
    Linear layer with continuous weight + binary mask and STE gradient.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = nn.Parameter(
            torch.empty(out_features, in_features, dtype=dtype, device=device)
        )
        # AD-FIX-07: non-persistent mask.
        self.register_buffer(
            "mask",
            torch.ones(out_features, in_features, dtype=dtype, device=device),
            persistent=False,
        )

        if bias:
            self.bias = nn.Parameter(torch.empty(out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

    @classmethod
    def from_linear(cls, linear: nn.Module, mask: torch.Tensor) -> ADMMMaskedLinear:
        # AD-FIX-05: handle Conv1D (nn.Module with .nf) as well as nn.Linear.
        raw_w = linear.weight.data
        is_conv1d = "conv1d" in linear.__class__.__name__.lower() or hasattr(linear, "nf")
        if is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            w_data = raw_w.t().contiguous()
        else:
            out_f, in_f = raw_w.shape
            w_data = raw_w

        layer = cls(
            in_features=in_f,
            out_features=out_f,
            bias=getattr(linear, "bias", None) is not None and linear.bias is not None,
            dtype=raw_w.dtype,
            device=raw_w.device,
        )
        with torch.no_grad():
            layer.weight.data.copy_(w_data)
            layer.mask.copy_(mask.to(dtype=raw_w.dtype, device=raw_w.device))
            if linear.bias is not None and layer.bias is not None:
                layer.bias.data.copy_(linear.bias.data)
        return layer

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w_masked = ADMMSTEPrune.apply(self.weight, self.mask)
        return F.linear(x, w_masked, self.bias)

Linear layer with continuous weight + binary mask and STE gradient.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Module, mask: torch.Tensor) ‑> ADMMMaskedLinear

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    w_masked = ADMMSTEPrune.apply(self.weight, self.mask)
    return F.linear(x, w_masked, self.bias)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class ADMMPruner (model: nn.Module, sparsity: float = 0.5, rho: float = 0.001, max_iter: int = 6, structured_24: bool = True, block_size: tuple[int, int] = (4, 4), dynamic_rho: bool = True, use_ste: bool = True, recalibrate_ln: bool = True, max_error: float = 0.45, fallback_to_original: bool = True, protect_critical: bool = True, target_attention_only: bool = False, target_modules: list[str] | None = None, lambda_distill: float = 0.1, distill_batch_size: int = 32, distill_use_activation: bool = False, distill_lr: float = 0.1, distill_steps: int = 2, calibration_inputs: dict[str, torch.Tensor] | None = None, tokenizer: Any | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("admm_pruner")
@COMPRESSOR_REGISTRY.register("admm_pruner_distill")
class ADMMPruner(BaseCompressor):
    """
    Production-hardened ADMM Pruner with Layer-Wise Reconstruction Distillation,
    2:4 semi-structured sparsity, STE gradients, and dynamic rho scheduling.
    """

    def __init__(
        self,
        model: nn.Module,
        sparsity: float = 0.5,
        rho: float = 1e-3,
        max_iter: int = 6,
        structured_24: bool = True,
        block_size: tuple[int, int] = (4, 4),
        dynamic_rho: bool = True,
        use_ste: bool = True,
        recalibrate_ln: bool = True,
        max_error: float = 0.45,
        fallback_to_original: bool = True,
        protect_critical: bool = True,
        target_attention_only: bool = False,
        target_modules: list[str] | None = None,
        lambda_distill: float = 0.1,
        distill_batch_size: int = 32,
        distill_use_activation: bool = False,
        distill_lr: float = 0.1,
        distill_steps: int = 2,
        calibration_inputs: dict[str, torch.Tensor] | None = None,
        tokenizer: Any | None = None,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="admm_pruner", **kwargs)
        conf = cfg or kwargs.get("comp_config")

        self.sparsity = max(
            0.0,
            min(
                float(
                    _cfg_get(
                        conf,
                        "admm_pruner_sparsity",
                        _cfg_get(conf, "admm_sparsity", _cfg_get(conf, "sparsity", sparsity)),
                    )
                ),
                0.95,
            ),
        )
        self.initial_rho = float(
            _cfg_get(
                conf,
                "admm_pruner_rho",
                _cfg_get(conf, "admm_rho", _cfg_get(conf, "rho", rho)),
            )
        )
        self.max_iter = max(
            1,
            int(
                _cfg_get(
                    conf,
                    "admm_pruner_max_iter",
                    _cfg_get(conf, "admm_max_iter", _cfg_get(conf, "max_iter", max_iter)),
                )
            ),
        )
        self.structured_24 = bool(
            _cfg_get(
                conf,
                "admm_pruner_structured_24",
                _cfg_get(
                    conf,
                    "admm_structured_24",
                    _cfg_get(conf, "structured_24", structured_24),
                ),
            )
        )
        self.block_size = block_size
        self.dynamic_rho = bool(
            _cfg_get(
                conf,
                "admm_pruner_dynamic_rho",
                _cfg_get(conf, "dynamic_rho", dynamic_rho),
            )
        )
        self.use_ste = bool(
            _cfg_get(conf, "admm_pruner_use_ste", _cfg_get(conf, "use_ste", use_ste))
        )
        self.recalibrate_ln = bool(
            _cfg_get(
                conf,
                "admm_pruner_recalibrate_ln",
                _cfg_get(conf, "recalibrate_ln", recalibrate_ln),
            )
        )
        self.max_error = float(
            _cfg_get(
                conf,
                "admm_pruner_max_error",
                _cfg_get(conf, "admm_max_error", _cfg_get(conf, "max_error", max_error)),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                conf,
                "admm_pruner_fallback_to_original",
                _cfg_get(conf, "fallback_to_original", fallback_to_original),
            )
        )
        self.protect_critical = bool(_cfg_get(conf, "protect_critical", protect_critical))
        self.target_attention_only = bool(
            _cfg_get(
                conf,
                "admm_pruner_target_attention_only",
                _cfg_get(
                    conf,
                    "admm_target_attention_only",
                    _cfg_get(conf, "target_attention_only", target_attention_only),
                ),
            )
        )
        self.target_modules = target_modules

        self.lambda_distill = float(
            _cfg_get(
                conf,
                "admm_pruner_lambda_distill",
                _cfg_get(conf, "lambda_distill", lambda_distill),
            )
        )
        self.distill_batch_size = int(
            _cfg_get(
                conf,
                "admm_pruner_distill_batch_size",
                _cfg_get(conf, "distill_batch_size", distill_batch_size),
            )
        )
        self.distill_use_activation = bool(
            _cfg_get(
                conf,
                "admm_pruner_distill_use_activation",
                _cfg_get(conf, "distill_use_activation", distill_use_activation),
            )
        )
        self.distill_lr = float(_cfg_get(conf, "distill_lr", distill_lr))
        self.distill_steps = max(1, int(_cfg_get(conf, "distill_steps", distill_steps)))

        self.calibration_inputs = calibration_inputs or {}
        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.target_device = (
            torch.device(device)
            if device is not None
            else next(model.parameters(), torch.tensor(0.0)).device
        )
        self.comp_config = conf
        self._compressed_params: int = 0

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_sparsities: dict[str, float] = {}
        self.layer_actual_sparsities: dict[str, float] = {}
        self.layer_distill_losses: dict[str, float] = {}
        self.stats: dict[str, Any] = {}

    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        if isinstance(module, ADMMMaskedLinear):
            return False
        if any(
            name.endswith(sfx)
            for sfx in [".base_layer", ".lora_A", ".lora_B", ".lora_dropout", ".linear"]
        ):
            return False

        target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
        cls_name = module.__class__.__name__.lower()

        weight_tensor = getattr(target_mod, "weight", getattr(module, "weight", None))
        if weight_tensor is None or not hasattr(weight_tensor, "dim") or weight_tensor.dim() != 2:
            return False

        target_cls = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in target_cls
        is_linear_like = (
            isinstance(module, nn.Linear)
            or isinstance(target_mod, nn.Linear)
            or is_conv1d
            or (
                Conv1D is not None
                and (isinstance(module, Conv1D) or isinstance(target_mod, Conv1D))
            )
            or "loralinear" in cls_name
            or hasattr(module, "base_layer")
        )
        if not is_linear_like:
            return False

        name_lower = name.lower()
        if any(
            skip in name_lower
            for skip in [
                "lm_head",
                "classifier",
                "head",
                "output_layer",
                "score",
                "pooler",
                "pre_classifier",
            ]
        ):
            return False

        if self.target_attention_only:
            ffn_patterns = [
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            ]
            if any(p in name_lower for p in ffn_patterns):
                return False
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            attn_patterns = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
                "in_proj_qkv",
                "in_proj_z",
                "in_proj_b",
                "out_proj",
            ]
            if not any(p in name_lower for p in attn_patterns):
                return False
        else:
            if not should_compress_linear(self.model, name):
                return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    # ------------------------------------------------------------------
    def _collect_calibration_inputs(self, target_names: list[str]) -> dict[str, torch.Tensor]:
        collected: dict[str, list[torch.Tensor]] = {name: [] for name in target_names}

        if self.calibration_inputs:
            for name in target_names:
                if name in self.calibration_inputs:
                    collected[name].append(self.calibration_inputs[name])

        has_pre_collected = any(len(v) > 0 for v in collected.values())
        if has_pre_collected:
            return {k: torch.cat(v, dim=0) for k, v in collected.items() if v}

        if self.calibration_loader is None and self.calibration_texts is None:
            return {}

        hooks = []
        name_to_mod = dict(self.model.named_modules())

        def make_hook(layer_name: str):
            def hook_fn(mod, inp, out):
                if not inp or inp[0] is None:
                    return
                cur_total = sum(t.shape[0] for t in collected[layer_name])
                # AD-FIX-09: budget check before slicing.
                rem = self.distill_batch_size - cur_total
                if rem <= 0:
                    return
                x = inp[0].detach()
                if x.ndim >= 2:
                    x_flat = x.reshape(-1, x.shape[-1]).float().cpu()
                    collected[layer_name].append(x_flat[:rem])

            return hook_fn

        for name in target_names:
            if name in name_to_mod:
                target_mod = getattr(name_to_mod[name], "base_layer", name_to_mod[name])
                hooks.append(target_mod.register_forward_hook(make_hook(name)))

        device = self.target_device
        was_training = self.model.training
        self.model.eval()

        try:
            with torch.no_grad():
                if self.calibration_loader is not None:
                    for batch in self.calibration_loader:
                        if isinstance(batch, dict):
                            batch = {
                                k: v.to(device) if isinstance(v, torch.Tensor) else v
                                for k, v in batch.items()
                            }
                            self.model(**batch)
                        elif isinstance(batch, (list, tuple)):
                            batch = [
                                v.to(device) if isinstance(v, torch.Tensor) else v for v in batch
                            ]
                            self.model(*batch)
                        elif isinstance(batch, torch.Tensor):
                            self.model(batch.to(device))
                        if all(
                            sum(t.shape[0] for t in collected[n]) >= self.distill_batch_size
                            for n in target_names
                        ):
                            break
                elif self.calibration_texts is not None and self.tokenizer is not None:
                    for text in self.calibration_texts[:4]:
                        toks = self.tokenizer(
                            text,
                            return_tensors="pt",
                            truncation=True,
                            max_length=128,
                        )
                        toks = {k: v.to(device) for k, v in toks.items()}
                        self.model(**toks)
        except Exception as e:
            logger.warning("ADMM calibration input collection failed: %s", e)
        finally:
            for h in hooks:
                h.remove()
            if was_training:
                self.model.train()

        result: dict[str, torch.Tensor] = {}
        for name, tensor_list in collected.items():
            if tensor_list:
                result[name] = torch.cat(tensor_list, dim=0)
        return result

    # ------------------------------------------------------------------
    def _project_sparse(self, candidate: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Project candidate (typically W + U) onto the sparse constraint set.
        Returns (pruned_tensor, binary_mask).
        AD-FIX-11: saliency is computed from the *candidate*, not from W.
        """
        if self.structured_24:
            saliency = candidate.abs()
            pruned, raw_mask = apply_24_structured_mask(candidate, saliency=saliency, dim=1)
            if raw_mask is None:
                mask = (pruned.abs() > 1e-12).float()
            elif raw_mask.dtype == torch.bool:
                mask = raw_mask.float()
            else:
                mask = (raw_mask > 0.5).float()
            return pruned, mask
        elif self.block_size != (1, 1):
            pruned, raw_mask = apply_block_mask(
                candidate, block_size=self.block_size, sparsity=self.sparsity
            )
            if raw_mask is None:
                mask = (pruned.abs() > 1e-12).float()
            elif raw_mask.dtype == torch.bool:
                mask = raw_mask.float()
            else:
                mask = (raw_mask > 0.5).float()
            return pruned, mask
        else:
            flat = candidate.abs().reshape(-1)
            k = max(1, int(flat.numel() * (1.0 - self.sparsity)))
            thr = torch.topk(flat, k, largest=True).values[-1]
            mask = (candidate.abs() >= thr).float()
            return candidate * mask, mask

    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_sparsities = {}
        self.layer_actual_sparsities = {}
        self.layer_distill_losses = {}

        targets = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not targets:
            logger.warning("ADMMPruner: No compressible target modules found.")
            stats = self.get_compression_stats()
            self.stats = stats
            self.model.compression_stats = stats
            return self.model, stats

        target_names = [name for name, _ in targets]
        calib_inputs = (
            self._collect_calibration_inputs(target_names) if self.lambda_distill > 0 else {}
        )

        for name, module in targets:
            target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
            target_cls = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in target_cls

            raw_w = target_mod.weight.data
            orig_shape = raw_w.shape
            device = raw_w.device
            dtype = raw_w.dtype

            W0 = raw_w.t().float() if is_conv1d else raw_w.float()
            out_features, in_features = W0.shape
            w_norm = torch.norm(W0, p="fro").item()

            X_calib = calib_inputs.get(name, None)
            if self.lambda_distill > 0:
                if X_calib is None or X_calib.numel() == 0:
                    torch.manual_seed(hash(name) % (2**31 - 1))
                    X_calib = torch.randn(
                        min(32, self.distill_batch_size),
                        in_features,
                        device=device,
                    ) / math.sqrt(in_features)
                else:
                    # AD-FIX-08: cast X_calib to match W0 dtype.
                    X_calib = X_calib.to(device=device, dtype=torch.float32)
                    if X_calib.ndim > 2:
                        X_calib = X_calib.reshape(-1, in_features)
                    if X_calib.shape[0] > self.distill_batch_size:
                        X_calib = X_calib[: self.distill_batch_size]
                with torch.no_grad():
                    Y0 = X_calib @ W0.t()
            else:
                X_calib = None
                Y0 = None

            # Baseline projection.
            Z0, mask0 = self._project_sparse(W0)
            W_sparse_0 = W0 * mask0
            best_W_sparse = W_sparse_0
            best_mask = mask0.clone()  # AD-FIX-04
            base_w_err = float(torch.norm(W0 - W_sparse_0, p="fro").item() / max(w_norm, 1e-12))
            best_err = base_w_err
            best_score = base_w_err
            if self.lambda_distill > 0 and X_calib is not None and Y0 is not None:
                y0_norm = max(torch.norm(Y0, p="fro").item(), 1e-12)
                base_y_err = float(
                    torch.norm((X_calib @ W_sparse_0.t()) - Y0, p="fro").item() / y0_norm
                )
                best_score = base_w_err + self.lambda_distill * base_y_err

            # ADMM state.
            W = W0.clone()
            Z = Z0.clone()
            mask = mask0.clone()
            U = torch.zeros_like(W0)

            rho = self.initial_rho
            mu = 10.0
            tau_incr = 1.5
            tau_decr = 1.5
            rho_min, rho_max = 1e-4, 100.0
            final_distill_loss = 0.0

            try:
                for it in range(self.max_iter):
                    with torch.no_grad():
                        Z_prev = Z.clone()

                        # 1. Sparse projection.
                        Z, mask = self._project_sparse(W + U)

                        # 2. Closed-form W update.
                        W = (W0 + rho * (Z - U)) / (1.0 + rho)

                        # 3. Layer-Wise Reconstruction Distillation.
                        # AD-FIX-02: scheduled, bounded distill step.
                        if self.lambda_distill > 0 and X_calib is not None and Y0 is not None:
                            for _ in range(self.distill_steps):
                                if self.use_ste:
                                    W_eff = ADMMSTEPrune.apply(W, mask)
                                else:
                                    W_eff = W * mask

                                Y_pruned = X_calib @ W_eff.t()
                                E = Y_pruned - Y0
                                N_pts, d_out = E.shape
                                grad_W = ((2.0 / max(1, N_pts * d_out)) * (E.t() @ X_calib)) * mask

                                grad_norm = torch.norm(grad_W)
                                if grad_norm > 1e-8:
                                    # AD-FIX-02: no artificial scaling distortion;
                                    # lr already encodes lambda_distill influence.
                                    lr_eff = self.distill_lr
                                    direction = grad_W / grad_norm.clamp(min=1e-12)
                                    step_size = lr_eff * min(1.0, float(grad_norm))
                                    W = W - step_size * direction
                                else:
                                    break

                        # 4. Dual update.
                        U = U + W - Z

                        # 5. Dynamic rho scheduling.
                        if self.dynamic_rho:
                            primal_res = torch.norm(W - Z).item()
                            dual_res = (rho * torch.norm(Z - Z_prev)).item()
                            if primal_res > mu * (dual_res + 1e-8):
                                rho = min(rho * tau_incr, rho_max)
                            elif dual_res > mu * (primal_res + 1e-8):
                                rho = max(rho / tau_decr, rho_min)

                        if not torch.isfinite(W).all():
                            logger.warning(
                                "ADMMPruner: non-finite W at iter %d on %s",
                                it,
                                name,
                            )
                            break

                        # Monotonic best-iterate tracking.
                        W_sparse_curr = W * mask
                        curr_w_err = float(
                            torch.norm(W0 - W_sparse_curr, p="fro").item() / max(w_norm, 1e-12)
                        )
                        curr_score = curr_w_err
                        if self.lambda_distill > 0 and X_calib is not None and Y0 is not None:
                            curr_y_err = float(
                                torch.norm((X_calib @ W_sparse_curr.t()) - Y0, p="fro").item()
                                / y0_norm
                            )
                            curr_score = curr_w_err + self.lambda_distill * curr_y_err

                        if curr_score < best_score:
                            best_score = curr_score
                            best_err = curr_w_err
                            best_W_sparse = W_sparse_curr.clone()
                            best_mask = mask.clone()  # AD-FIX-04

                # Final distill loss on best sparse candidate.
                if self.lambda_distill > 0 and X_calib is not None and Y0 is not None:
                    with torch.no_grad():
                        final_distill_loss = float(
                            torch.mean(((X_calib @ best_W_sparse.t()) - Y0) ** 2).item()
                        )

            except Exception as e:
                logger.warning("ADMMPruner optimization failed on %s: %s", name, e)
                best_W_sparse = None
                best_err = 1.0

            final_err = best_err
            W_final = best_W_sparse
            self.layer_errors[name] = final_err
            self.layer_sparsities[name] = 0.5 if self.structured_24 else self.sparsity
            self.layer_distill_losses[name] = final_distill_loss

            if W_final is None or not torch.isfinite(W_final).all() or final_err > self.max_error:
                logger.warning(
                    "ADMMPruner reconstruction error %.4f on %s exceeded threshold %.4f.",
                    final_err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    logger.info("ADMMPruner: fallback to original on %s.", name)
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

            if is_conv1d:
                res_w = W_final.t().contiguous().to(dtype=dtype, device=device)
            else:
                res_w = W_final.contiguous().to(dtype=dtype, device=device)

            if res_w.shape != orig_shape:
                logger.warning(
                    "ADMMPruner shape mismatch on %s: %s vs %s. Skipping.",
                    name,
                    tuple(res_w.shape),
                    tuple(orig_shape),
                )
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                continue

            has_lora = (hasattr(module, "lora_A") and module.lora_A is not None) or hasattr(
                module, "base_layer"
            )

            # AD-FIX-05: ADMMMaskedLinear applies whenever use_ste is on and no LoRA / no Conv1D.
            if self.use_ste and not has_lora and not is_conv1d and isinstance(module, nn.Linear):
                admm_layer = ADMMMaskedLinear.from_linear(module, mask=best_mask)
                replace_module(self.model, name, admm_layer)
            else:
                with torch.no_grad():
                    target_mod.weight.data.copy_(res_w)

            actual_zero_frac = float((res_w == 0).sum().item() / max(1, res_w.numel()))
            self.layer_actual_sparsities[name] = actual_zero_frac

            if has_lora:
                if hasattr(target_mod, "weight") and target_mod.weight is not None:
                    target_mod.weight.requires_grad = False
                if hasattr(target_mod, "bias") and target_mod.bias is not None:
                    target_mod.bias.requires_grad = False
                for p_name, p in module.named_parameters():
                    if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                        p.requires_grad = True

            self.compressed_layers.append(name)

        # LayerNorm recalibration.
        if self.recalibrate_ln and (
            self.calibration_loader is not None or self.calibration_texts is not None
        ):
            try:
                recalibrate_layernorms(
                    self.model,
                    tokenizer=self.tokenizer,
                    calibration_loader=self.calibration_loader,
                    calibration_texts=self.calibration_texts,
                    device=self.target_device,
                )
            except Exception as e:
                logger.warning("ADMM LayerNorm recalibration skipped: %s", e)

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        avg_actual_sparsity = (
            sum(self.layer_actual_sparsities[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else (0.5 if self.structured_24 else self.sparsity)
        )
        avg_distill_loss = (
            sum(self.layer_distill_losses[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}
        stats.update(
            {
                "pruned_layers": len(self.compressed_layers),
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "layer_errors": self.layer_errors,
                "layer_distill_losses": self.layer_distill_losses,
                "avg_reconstruction_error": avg_err,
                "avg_distill_loss": avg_distill_loss,
                "sparsity_target": self.sparsity,
                "actual_sparsity": avg_actual_sparsity,
                "layer_actual_sparsities": self.layer_actual_sparsities,
                "structured_24": self.structured_24,
                "admm_iterations": self.max_iter,
                "dynamic_rho": self.dynamic_rho,
                "use_ste": self.use_ste,
                "recalibrate_ln": self.recalibrate_ln,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "target_attention_only": self.target_attention_only,
                "lambda_distill": self.lambda_distill,
                "distill_batch_size": self.distill_batch_size,
                "distill_use_activation": self.distill_use_activation,
                "method": "admm_pruner_distill_ste_dynamic_rho",
            }
        )
        self.stats = stats
        self.model.compression_stats = stats
        return self.model, stats

Production-hardened ADMM Pruner with Layer-Wise Reconstruction Distillation, 2:4 semi-structured sparsity, STE gradients, and dynamic rho scheduling.

Ancestors

BaseCompressor

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class ADMMSTEPrune (*args, **kwargs)

Expand source code

class ADMMSTEPrune(torch.autograd.Function):
    """
    Straight-Through Estimator for ADMM sparse weights.
    Forward: y = weight * mask (with dtype-aligned mask).
    Backward: gradient flows through the identity to weight; mask has no grad.
    """

    @staticmethod
    def forward(ctx, weight: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        ctx.save_for_backward(mask)
        return weight * mask.to(weight.dtype)

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        # AD-FIX-10: no clone; STE identity.
        return grad_output, None

Straight-Through Estimator for ADMM sparse weights. Forward: y = weight * mask (with dtype-aligned mask). Backward: gradient flows through the identity to weight; mask has no grad.

Ancestors

torch.autograd.function.Function
torch.autograd.function._SingleLevelFunction
torch._C._FunctionBase
torch.autograd.function.FunctionCtx
torch.autograd.function._HookMixin

Static methods

def backward(ctx, grad_output: torch.Tensor)

Expand source code

@staticmethod
def backward(ctx, grad_output: torch.Tensor):
    # AD-FIX-10: no clone; STE identity.
    return grad_output, None

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the vjp function.)

It must accept a context :attr:ctx as the first argument, followed by as many outputs as the :func:forward returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to :func:forward. Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute :attr:ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., :func:backward will have ctx.needs_input_grad[0] = True if the first input to :func:forward needs gradient computed w.r.t. the output.

def forward(ctx, weight: torch.Tensor, mask: torch.Tensor) ‑> torch.Tensor

Expand source code

@staticmethod
def forward(ctx, weight: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    ctx.save_for_backward(mask)
    return weight * mask.to(weight.dtype)

Define the forward of the custom autograd Function.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx)::

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See :ref:combining-forward-context for more details

Usage 2 (Separate forward and ctx)::

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass


@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

The forward no longer accepts a ctx argument.
Instead, you must also override the :meth:torch.autograd.Function.setup_context staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.
See :ref:extending-autograd for more details

The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with :func:ctx.save_for_backward if they are intended to be used in backward (equivalently, vjp) or :func:ctx.save_for_forward if they are intended to be used for in jvp.

class ARMORLinear (in_features: int, out_features: int, rank: int, bias: bool = True, lora_alpha: float | None = None, dtype: torch.dtype = torch.float32, device: torch.device | None = None)

Expand source code

class ARMORLinear(nn.Module):
    """
    Decoupled sparse + low-rank residual layer.

    Forward:
        y = x @ (sparse_weight ⊙ sparse_mask)^T
            + (scaling) * (x @ lora_A^T) @ lora_B^T
            + bias

    Invariants:
      - `sparse_weight` holds the ORIGINAL weight values (not pre-masked).
      - `sparse_mask` (non-persistent) holds a 0/1 mask applied once per forward.
      - `lora_alpha == rank` so `scaling == 1.0`, so the SVD residual is exact.
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        bias: bool = True,
        lora_alpha: float | None = None,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features
        self.rank = max(1, min(rank, in_features, out_features))

        # AR-FIX-01: default alpha to rank so scaling is 1.0 by construction.
        self.lora_alpha = float(lora_alpha) if lora_alpha is not None else float(self.rank)
        self.scaling = (self.lora_alpha / self.rank) if self.rank > 0 else 1.0

        # AR-FIX-11: store raw (un-masked) weight.
        self.sparse_weight = nn.Parameter(
            torch.empty(out_features, in_features, dtype=dtype, device=device),
            requires_grad=False,
        )
        # AR-FIX-02: non-persistent mask.
        self.register_buffer(
            "sparse_mask",
            torch.ones(out_features, in_features, dtype=dtype, device=device),
            persistent=False,
        )

        self.lora_A = nn.Parameter(torch.empty(self.rank, in_features, dtype=dtype, device=device))
        self.lora_B = nn.Parameter(torch.empty(out_features, self.rank, dtype=dtype, device=device))

        if bias:
            self.bias = nn.Parameter(torch.empty(out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Linear | Any,
        mask: torch.Tensor,
        lora_A: torch.Tensor,
        lora_B: torch.Tensor,
        lora_alpha: float | None = None,
    ) -> ARMORLinear:
        cls_name = linear.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(linear, "nf")
        raw_w = linear.weight.data
        if is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            w_data = raw_w.t().contiguous()
        else:
            out_f, in_f = raw_w.shape
            w_data = raw_w

        # AR-FIX-10: clamp rank to actual dims.
        requested_rank = int(lora_A.shape[0])
        actual_rank = max(1, min(requested_rank, in_f, out_f))

        # AR-FIX-01: ensure scaling == 1.0 by default.
        effective_alpha = float(lora_alpha) if lora_alpha is not None else float(actual_rank)

        layer = cls(
            in_features=in_f,
            out_features=out_f,
            rank=actual_rank,
            bias=(getattr(linear, "bias", None) is not None and linear.bias is not None),
            lora_alpha=effective_alpha,
            dtype=raw_w.dtype,
            device=raw_w.device,
        )
        mask_dev = mask.to(dtype=w_data.dtype, device=w_data.device)
        with torch.no_grad():
            # AR-FIX-11: store raw weight, do not pre-mask.
            layer.sparse_weight.data.copy_(w_data)
            layer.sparse_mask.copy_(mask_dev)
            layer.lora_A.data.copy_(lora_A[:actual_rank].to(dtype=raw_w.dtype, device=raw_w.device))
            layer.lora_B.data.copy_(
                lora_B[:, :actual_rank].to(dtype=raw_w.dtype, device=raw_w.device)
            )
            if (
                getattr(linear, "bias", None) is not None
                and linear.bias is not None
                and layer.bias is not None
            ):
                layer.bias.data.copy_(linear.bias.data)
        return layer

    @property
    def weight(self) -> torch.Tensor:
        # AR-FIX-09: compute in fp32, cast back to sparse_weight dtype.
        with torch.no_grad():
            base = (self.sparse_weight.float() * self.sparse_mask.float()).to(
                self.sparse_weight.dtype
            )
            if self.lora_A is not None and self.lora_B is not None and self.rank > 0:
                lora_w = torch.matmul(self.lora_B.float(), self.lora_A.float()) * self.scaling
                return base + lora_w.to(base.dtype)
            return base

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)

        # 1. Sparse core. Mask applied once.
        w_eff = self.sparse_weight * self.sparse_mask
        out_sparse = F.linear(x_flat, w_eff)

        # 2. Low-rank residual.
        lr_dtype = self.lora_A.dtype
        x_lr = x_flat.to(lr_dtype) if x_flat.dtype != lr_dtype else x_flat
        out_lr = F.linear(F.linear(x_lr, self.lora_A), self.lora_B)
        # AR-FIX-03: explicit dtype of scaling.
        out_lr = (out_lr * self.scaling).to(out_sparse.dtype)

        out = out_sparse + out_lr
        if self.bias is not None:
            out = out + self.bias.to(out.dtype)
        return out.view(*orig_shape[:-1], self.out_features)

Decoupled sparse + low-rank residual layer.

Forward -----= y = x @ (sparse_weight ⊙ sparse_mask)^T + (scaling) * (x @ lora_A^T) @ lora_B^T + bias

Invariants -----= - sparse_weight holds the ORIGINAL weight values (not pre-masked). - sparse_mask (non-persistent) holds a 0/1 mask applied once per forward. - lora_alpha == rank so scaling == 1.0, so the SVD residual is exact.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Static methods

def from_linear(linear: nn.Linear | Any, mask: torch.Tensor, lora_A: torch.Tensor, lora_B: torch.Tensor, lora_alpha: float | None = None) ‑> ARMORLinear

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    # AR-FIX-09: compute in fp32, cast back to sparse_weight dtype.
    with torch.no_grad():
        base = (self.sparse_weight.float() * self.sparse_mask.float()).to(
            self.sparse_weight.dtype
        )
        if self.lora_A is not None and self.lora_B is not None and self.rank > 0:
            lora_w = torch.matmul(self.lora_B.float(), self.lora_A.float()) * self.scaling
            return base + lora_w.to(base.dtype)
        return base

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)

    # 1. Sparse core. Mask applied once.
    w_eff = self.sparse_weight * self.sparse_mask
    out_sparse = F.linear(x_flat, w_eff)

    # 2. Low-rank residual.
    lr_dtype = self.lora_A.dtype
    x_lr = x_flat.to(lr_dtype) if x_flat.dtype != lr_dtype else x_flat
    out_lr = F.linear(F.linear(x_lr, self.lora_A), self.lora_B)
    # AR-FIX-03: explicit dtype of scaling.
    out_lr = (out_lr * self.scaling).to(out_sparse.dtype)

    out = out_sparse + out_lr
    if self.bias is not None:
        out = out + self.bias.to(out.dtype)
    return out.view(*orig_shape[:-1], self.out_features)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class ARMORPruner (model: nn.Module, block_size: int = 4, rank: int = 16, sparsity: float = 0.5, structured_24: bool = True, decouple: bool = True, recalibrate_ln: bool = True, lora_alpha: float | None = None, max_error: float = 0.35, fallback_to_original: bool = True, target_attention_only: bool = False, skip_attention_qkv: bool = False, skip_attention_output: bool = False, target_modules: list[str] | None = None, tokenizer: Any | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("armor_pruner")
@COMPRESSOR_REGISTRY.register("armor")
@COMPRESSOR_REGISTRY.register("armor_compressor")
class ARMORPruner(BaseCompressor):
    """Production-hardened ARMOR Pruner."""

    is_custom_attention: bool = True

    def __init__(
        self,
        model: nn.Module,
        block_size: int = 4,
        rank: int = 16,
        sparsity: float = 0.5,
        structured_24: bool = True,
        decouple: bool = True,
        recalibrate_ln: bool = True,
        lora_alpha: float | None = None,
        max_error: float = 0.35,
        fallback_to_original: bool = True,
        target_attention_only: bool = False,
        skip_attention_qkv: bool = False,
        skip_attention_output: bool = False,
        target_modules: list[str] | None = None,
        tokenizer: Any | None = None,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="armor_pruner", **kwargs)
        conf = cfg or kwargs.get("comp_config")

        self.block_size = int(
            _cfg_get(conf, "armor_block_size", _cfg_get(conf, "block_size", block_size))
        )
        self.rank = max(1, int(_cfg_get(conf, "armor_rank", _cfg_get(conf, "rank", rank))))
        self.sparsity = max(
            0.0,
            min(
                float(_cfg_get(conf, "armor_sparsity", _cfg_get(conf, "sparsity", sparsity))),
                0.95,
            ),
        )
        self.structured_24 = bool(
            _cfg_get(
                conf,
                "armor_structured_24",
                _cfg_get(conf, "structured_24", structured_24),
            )
        )
        self.decouple = bool(_cfg_get(conf, "armor_decouple", _cfg_get(conf, "decouple", decouple)))
        self.recalibrate_ln = bool(
            _cfg_get(
                conf,
                "armor_recalibrate_ln",
                _cfg_get(conf, "recalibrate_ln", recalibrate_ln),
            )
        )
        # AR-FIX-01: allow None, in which case it will be set to rank per-layer.
        self.lora_alpha = (
            float(_cfg_get(conf, "armor_lora_alpha", _cfg_get(conf, "lora_alpha", lora_alpha)))
            if lora_alpha is not None
            else None
        )
        self.max_error = float(
            _cfg_get(conf, "armor_max_error", _cfg_get(conf, "max_error", max_error))
        )
        self.fallback_to_original = bool(
            _cfg_get(
                conf,
                "armor_fallback_to_original",
                _cfg_get(conf, "fallback_to_original", fallback_to_original),
            )
        )
        self.target_attention_only = bool(
            _cfg_get(
                conf,
                "armor_target_attention_only",
                _cfg_get(conf, "target_attention_only", target_attention_only),
            )
        )
        self.skip_attention_qkv = bool(
            _cfg_get(
                conf,
                "armor_skip_attention_qkv",
                _cfg_get(conf, "skip_attention_qkv", skip_attention_qkv),
            )
        )
        self.skip_attention_output = bool(
            _cfg_get(
                conf,
                "armor_skip_attention_output",
                _cfg_get(conf, "skip_attention_output", skip_attention_output),
            )
        )
        self.target_modules = target_modules

        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.target_device = (
            torch.device(device)
            if device is not None
            else next(model.parameters(), torch.tensor(0.0)).device
        )
        self.comp_config = conf
        self._compressed_params: int = 0

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_actual_sparsities: dict[str, float] = {}
        self.layer_effective_ratios: dict[str, float] = {}
        self.stats: dict[str, Any] = {}

    # ------------------------------------------------------------------
    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        """
        AR-FIX-13: always active (not gated by skip_attention_qkv).
        Returns True for q/k/v projections in GQA models where num_kv_heads < num_heads.
        These should be skipped entirely to avoid shape mismatches.
        """
        cfg = getattr(self.model, "config", None)
        if cfg is None:
            return False
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is None or num_kv_heads is None:
            return False
        if num_kv_heads >= num_heads:
            return False
        name_lower = name.lower()
        return any(
            p in name_lower
            for p in ("q_proj", "k_proj", "v_proj", "query", "key", "value", "c_attn")
        )

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        if isinstance(module, ARMORLinear):
            return False
        if any(
            name.endswith(sfx)
            for sfx in (".base_layer", ".lora_A", ".lora_B", ".lora_dropout", ".linear")
        ):
            return False

        target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
        cls_name = module.__class__.__name__.lower()

        weight_tensor = getattr(target_mod, "weight", getattr(module, "weight", None))
        if weight_tensor is None or not hasattr(weight_tensor, "dim") or weight_tensor.dim() != 2:
            return False

        # AR-FIX-05: Conv1D detection via class name.
        target_cls = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in target_cls
        is_linear_like = (
            isinstance(module, nn.Linear)
            or isinstance(target_mod, nn.Linear)
            or is_conv1d
            or (
                Conv1D is not None
                and (isinstance(module, Conv1D) or isinstance(target_mod, Conv1D))
            )
            or "loralinear" in cls_name
            or hasattr(module, "base_layer")
        )
        if not is_linear_like:
            return False

        name_lower = name.lower()
        if any(
            skip in name_lower
            for skip in (
                "lm_head",
                "classifier",
                "head",
                "output_layer",
                "score",
                "pooler",
                "pre_classifier",
            )
        ):
            return False

        # AR-FIX-15: projection O skip is opt-in.
        if self.skip_attention_output and any(
            nested in name_lower for nested in ("attention.self.output", "self.output.dense")
        ):
            return False

        # AR-FIX-13: GQA shape guard always active.
        if self._is_gqa_asymmetric(name, module):
            return False

        if self.target_attention_only:
            ffn_patterns = (
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            )
            if any(p in name_lower for p in ffn_patterns):
                return False
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            attn_patterns = (
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
                "in_proj_qkv",
                "in_proj_z",
                "in_proj_b",
                "out_proj",
            )
            if not any(p in name_lower for p in attn_patterns):
                return False
        else:
            if not should_compress_linear(self.model, name):
                return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    # ------------------------------------------------------------------
    def _collect_activations(self) -> dict[str, torch.Tensor]:
        """
        AR-FIX-12: Collect per-channel mean |x| instead of full activation tensors.
        Returns {layer_name: (in_features,) float32 tensor}.
        """
        self.model.eval()
        activations: dict[str, list[torch.Tensor]] = {}
        hooks = []

        def make_hook(layer_name: str):
            def hook(module, inp, outp):
                if inp and isinstance(inp[0], torch.Tensor):
                    x = inp[0].detach()
                    # Per-channel mean absolute value; collapse all dims but last.
                    if x.dim() >= 2:
                        col_mean = x.abs().mean(dim=tuple(range(x.dim() - 1)))
                    else:
                        col_mean = x.abs()
                    activations.setdefault(layer_name, []).append(col_mean.cpu())

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
                hooks.append(target_mod.register_forward_hook(make_hook(name)))

        try:
            with torch.no_grad():
                if self.calibration_loader is not None:
                    for batch in self.calibration_loader:
                        try:
                            if isinstance(batch, dict):
                                b_dev = {
                                    k: (v.to(self.target_device) if torch.is_tensor(v) else v)
                                    for k, v in batch.items()
                                }
                                self.model(**b_dev)
                            elif torch.is_tensor(batch):
                                self.model(batch.to(self.target_device))
                            elif isinstance(batch, (list, tuple)):
                                b_dev = [
                                    (v.to(self.target_device) if torch.is_tensor(v) else v)
                                    for v in batch
                                ]
                                self.model(*b_dev)
                        except Exception as b_err:
                            logger.debug("ARMOR calibration batch failed: %s", b_err)
                else:
                    calib_inputs = get_calibration_tensors(
                        self.model,
                        tokenizer=self.tokenizer,
                        calibration_loader=None,
                        calibration_texts=self.calibration_texts,
                        device=self.target_device,
                    )
                    try:
                        if isinstance(calib_inputs, dict):
                            self.model(**calib_inputs)
                        elif isinstance(calib_inputs, torch.Tensor):
                            self.model(calib_inputs)
                        elif isinstance(calib_inputs, (list, tuple)):
                            self.model(*calib_inputs)
                    except Exception as e:
                        logger.debug("ARMOR calibration pass exception: %s", e)
        finally:
            for h in hooks:
                h.remove()

        # Aggregate per layer by mean over batches.
        aggregated: dict[str, torch.Tensor] = {}
        for name, lst in activations.items():
            if lst:
                stacked = torch.stack([t.float() for t in lst], dim=0)
                aggregated[name] = stacked.mean(dim=0)
        return aggregated

    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        act_map = self._collect_activations()
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.layer_actual_sparsities = {}
        self.layer_effective_ratios = {}
        total_correction_energy = 0.0
        seen_weights: set = set()

        total_base_sparse_effective = 0  # non-zero entries
        total_base_sparse_storage = 0  # dense storage slots
        total_correction = 0
        total_orig = 0

        targets = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not targets:
            logger.warning("ARMORPruner: No compressible target modules found.")
            stats = self.get_compression_stats()
            self.stats = stats
            self.model.compression_stats = stats
            return self.model, stats

        for name, module in targets:
            target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
            target_cls = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in target_cls
            has_lora = (hasattr(module, "lora_A") and module.lora_A is not None) or hasattr(
                module, "base_layer"
            )

            # AR-FIX-07: use the actual weight tensor id.
            weight_tensor = getattr(target_mod, "weight", None)
            if weight_tensor is None:
                self.skipped_layers.append(name)
                continue
            w_id = id(weight_tensor)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            raw_w = weight_tensor.data
            device = raw_w.device
            dtype = raw_w.dtype

            W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
            out_features, in_features = W.shape
            grad = getattr(weight_tensor, "grad", None)
            act = act_map.get(name, None)

            try:
                saliency = compute_resilient_taylor_saliency(W, grad=grad, activations=act)

                # AR-FIX-04: robust mask handling.
                if self.structured_24:
                    W_sparse, raw_mask = apply_24_structured_mask(W, saliency=saliency, dim=1)
                    if raw_mask is None:
                        mask = (W_sparse.abs() > 1e-12).float()
                    else:
                        if raw_mask.dtype == torch.bool:
                            mask = raw_mask.float()
                        else:
                            mask = (raw_mask > 0.5).float()
                else:
                    flat = saliency.reshape(-1)
                    k = max(1, int(flat.numel() * (1.0 - self.sparsity)))
                    thr = torch.topk(flat, k, largest=True).values[-1]
                    mask = (saliency >= thr).float()
                    W_sparse = W * mask

                # Residual low-rank.
                residual = W - W_sparse
                r = min(self.rank, residual.shape[0], residual.shape[1])
                try:
                    U, S, Vh = torch.linalg.svd(residual, full_matrices=False)
                    sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
                    lora_B = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
                    lora_A = (sqrt_S.unsqueeze(1) * Vh[:r, :]).to(dtype=dtype, device=device)
                    correction = torch.matmul(lora_B.float(), lora_A.float())
                    total_correction_energy += torch.norm(correction).item()
                except Exception as e:
                    logger.debug("ARMOR residual SVD failed on %s: %s", name, e)
                    lora_A = torch.zeros(r, in_features, dtype=dtype, device=device)
                    lora_B = torch.zeros(out_features, r, dtype=dtype, device=device)

                # AR-FIX-01: alpha defaults to rank so scaling = 1.0.
                effective_alpha = self.lora_alpha if self.lora_alpha is not None else float(r)
                scaling = effective_alpha / r if r > 0 else 1.0
                W_recon = W_sparse + scaling * torch.matmul(lora_B.float(), lora_A.float())

                # AR-FIX-06: fast finite check.
                if not torch.isfinite(W_recon).all():
                    raise RuntimeError("non-finite reconstruction")

                w_norm = torch.norm(W, p="fro").clamp_min(1e-8)
                err = float(torch.norm(W - W_recon, p="fro").item() / w_norm)
                self.layer_errors[name] = err

                if err > self.max_error and self.fallback_to_original:
                    logger.warning(
                        "ARMOR reconstruction error %.4f on %s exceeded threshold %.4f.",
                        err,
                        name,
                        self.max_error,
                    )
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

                orig_count = int(W.numel())
                base_effective = int((mask > 0).sum().item())
                corr_count = int(lora_A.numel() + lora_B.numel()) if self.decouple else 0

                total_base_sparse_effective += base_effective
                total_base_sparse_storage += orig_count
                total_correction += corr_count
                total_orig += orig_count

                actual_zero_frac = float((mask == 0).sum().item() / max(1, orig_count))
                eff_ratio = float((base_effective + corr_count) / max(1, orig_count))
                self.layer_actual_sparsities[name] = actual_zero_frac
                self.layer_effective_ratios[name] = eff_ratio

                if self.decouple:
                    armor_layer = ARMORLinear.from_linear(
                        target_mod,
                        mask=mask,
                        lora_A=lora_A,
                        lora_B=lora_B,
                        lora_alpha=effective_alpha,
                    )
                    if has_lora:
                        module.base_layer = armor_layer
                        for p_name, p in module.named_parameters():
                            if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                                if not p_name.startswith("base_layer"):
                                    with torch.no_grad():
                                        p.data.zero_()
                                    p.requires_grad = False
                    elif name == "":
                        self.model = armor_layer
                    else:
                        replace_module(self.model, name, armor_layer)
                else:
                    # AR-FIX: mutation under no_grad. NOTE: this branch does NOT
                    # reduce parameter count; it is a numerical approximation only.
                    with torch.no_grad():
                        if is_conv1d:
                            weight_tensor.data.copy_(
                                W_recon.t().contiguous().to(dtype=dtype, device=device)
                            )
                        else:
                            weight_tensor.data.copy_(W_recon.to(dtype=dtype, device=device))

                self.compressed_layers.append(name)

            except Exception as e:
                logger.warning("ARMOR pruning failed on %s: %s", name, e)
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False

        # AR-FIX-08: LayerNorm recalibration only when data available.
        if self.recalibrate_ln and (
            self.calibration_loader is not None or self.calibration_texts is not None
        ):
            try:
                recalibrate_layernorms(
                    self.model,
                    tokenizer=self.tokenizer,
                    calibration_loader=self.calibration_loader,
                    calibration_texts=self.calibration_texts,
                    device=self.target_device,
                )
            except Exception as e:
                logger.warning("ARMOR LayerNorm recalibration skipped: %s", e)

        # Forward sanity check.
        forward_check_success = True
        try:
            self.model.eval()
            with torch.no_grad():
                first_param = next(self.model.parameters(), torch.tensor(0.0))
                dev = first_param.device
                p_dtype = (
                    first_param.dtype if first_param.dtype.is_floating_point else torch.float32
                )
                cfg = getattr(self.model, "config", None)
                hidden_size = getattr(cfg, "hidden_size", None)
                try:
                    dummy_ids = torch.zeros((1, 8), dtype=torch.long, device=dev)
                    self.model(
                        input_ids=dummy_ids,
                        attention_mask=torch.ones_like(dummy_ids),
                    )
                except Exception:
                    if hidden_size is not None:
                        dummy_embeds = torch.randn(1, 8, hidden_size, dtype=p_dtype, device=dev)
                    else:
                        first_linear = next(
                            (
                                m
                                for m in self.model.modules()
                                if isinstance(m, (nn.Linear, ARMORLinear))
                            ),
                            None,
                        )
                        in_dim = getattr(first_linear, "in_features", 8)
                        dummy_embeds = torch.randn(1, 8, in_dim, dtype=p_dtype, device=dev)
                    self.model(dummy_embeds)
        except Exception as e:
            logger.debug("ARMOR forward sanity check failed: %s", e)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        avg_err = (
            sum(self.layer_errors.values()) / len(self.layer_errors) if self.layer_errors else 0.0
        )
        avg_actual_sparsity = (
            sum(self.layer_actual_sparsities.values()) / max(1, len(self.layer_actual_sparsities))
            if self.layer_actual_sparsities
            else (0.5 if self.structured_24 else self.sparsity)
        )

        # AR-FIX-14: honest accounting.
        effective_param_ratio = (
            float((total_base_sparse_effective + total_correction) / max(1, total_orig))
            if total_orig > 0
            else 1.0
        )
        storage_param_ratio = (
            float((total_base_sparse_storage + total_correction) / max(1, total_orig))
            if total_orig > 0
            else 1.0
        )

        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "avg_reconstruction_error": avg_err,
                "pruned_layers": len(self.compressed_layers),
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "actual_sparsity": avg_actual_sparsity,
                "layer_actual_sparsities": self.layer_actual_sparsities,
                "layer_effective_ratios": self.layer_effective_ratios,
                # AR-FIX-14: distinguish storage from effective.
                "base_sparse_storage": total_base_sparse_storage,
                "base_sparse_effective": total_base_sparse_effective,
                "correction_params": total_correction,
                "effective_parameter_ratio": effective_param_ratio,
                "storage_parameter_ratio": storage_param_ratio,
                "effective_compression_rate": (
                    1.0 / effective_param_ratio if effective_param_ratio > 0 else 1.0
                ),
                "storage_compression_rate": (
                    1.0 / storage_param_ratio if storage_param_ratio > 0 else 1.0
                ),
                "rank": self.rank,
                "structured_24": self.structured_24,
                "decouple": self.decouple,
                "lora_alpha": (self.lora_alpha if self.lora_alpha is not None else "=rank"),
                "max_error": self.max_error,
                "target_attention_only": self.target_attention_only,
                "total_correction_energy": total_correction_energy,
                "forward_check_success": forward_check_success,
                "method": "armor_sparse_plus_lowrank_decoupled",
            }
        )
        self.stats = stats
        self.model.compression_stats = stats
        return self.model, stats

Production-hardened ARMOR Pruner.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class ASTPrunedLoRALinear (in_features: int, out_features: int, weight: torch.Tensor, bias: torch.Tensor | None = None, lora_r: int = 0, lora_alpha: float = 1.0, device: torch.device | None = None, dtype: torch.dtype | None = None)

Expand source code

class ASTPrunedLoRALinear(nn.Module):
    """
    Linear layer with frozen AST-pruned base weights and optional trainable
    LoRA adapters.

    The base weight is stored in original coordinate basis with a true 2:4
    (or N:M) sparsity pattern. Framework may pack it as sparse if desired.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        weight: torch.Tensor,
        bias: torch.Tensor | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features  # Conv1D compat
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        target_device = device if device is not None else weight.device
        target_dtype = dtype if dtype is not None else weight.dtype

        self.weight = nn.Parameter(
            weight.to(device=target_device, dtype=target_dtype),
            requires_grad=False,
        )
        if bias is not None:
            self.bias = nn.Parameter(
                bias.to(device=target_device, dtype=target_dtype),
                requires_grad=False,
            )
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty(
                    (self.lora_r, in_features),
                    dtype=torch.float32,
                    device=target_device,
                )
            )
            self.lora_B = nn.Parameter(
                torch.zeros(
                    (out_features, self.lora_r),
                    dtype=torch.float32,
                    device=target_device,
                )
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        weight: torch.Tensor,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> ASTPrunedLoRALinear:
        bias = None
        if hasattr(linear, "bias") and linear.bias is not None:
            bias = linear.bias.detach()
        in_features = getattr(linear, "in_features", getattr(linear, "nf", weight.shape[1]))
        out_features = getattr(linear, "out_features", getattr(linear, "nf", weight.shape[0]))
        return cls(
            in_features=in_features,
            out_features=out_features,
            weight=weight,
            bias=bias,
            lora_r=lora_r,
            lora_alpha=lora_alpha,
            device=weight.device,
            dtype=weight.dtype,
        )

    # AST-FIX-12: framework merge() compatibility
    def merge(self) -> None:
        """Merge LoRA into base weight in place."""
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            with torch.no_grad():
                delta = (self.lora_B.float() @ self.lora_A.float()) * self.scaling
                self.weight.data = (self.weight.data.float() + delta).to(self.weight.dtype)
                self.lora_A = None
                self.lora_B = None
                self.lora_r = 0

    def merge_and_unload(self) -> None:
        self.merge()

    def unmerge(self) -> None:
        # Not supported after merge
        pass

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)
        out = F.linear(x_flat, self.weight, self.bias)

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            # AST-FIX: single dtype conversion
            x_f = x_flat.float()
            lora_out = F.linear(F.linear(x_f, self.lora_A), self.lora_B)
            out = out + (self.scaling * lora_out).to(out.dtype)

        return out.reshape(*orig_shape[:-1], self.out_features)

Linear layer with frozen AST-pruned base weights and optional trainable LoRA adapters.

The base weight is stored in original coordinate basis with a true 2:4 (or N:M) sparsity pattern. Framework may pack it as sparse if desired.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Module, weight: torch.Tensor, lora_r: int = 0, lora_alpha: float = 1.0) ‑> ASTPrunedLoRALinear

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)
    out = F.linear(x_flat, self.weight, self.bias)

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        # AST-FIX: single dtype conversion
        x_f = x_flat.float()
        lora_out = F.linear(F.linear(x_f, self.lora_A), self.lora_B)
        out = out + (self.scaling * lora_out).to(out.dtype)

    return out.reshape(*orig_shape[:-1], self.out_features)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def merge(self) ‑> None

Expand source code

def merge(self) -> None:
    """Merge LoRA into base weight in place."""
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        with torch.no_grad():
            delta = (self.lora_B.float() @ self.lora_A.float()) * self.scaling
            self.weight.data = (self.weight.data.float() + delta).to(self.weight.dtype)
            self.lora_A = None
            self.lora_B = None
            self.lora_r = 0

Merge LoRA into base weight in place.

def merge_and_unload(self) ‑> None

Expand source code

def merge_and_unload(self) -> None:
    self.merge()

def unmerge(self) ‑> None

Expand source code

def unmerge(self) -> None:
    # Not supported after merge
    pass

class ASTPruner (model: nn.Module, block_size: int = 4, n_keep: int = 2, dim: int = 1, use_pca: bool = True, use_fisher: bool = True, recalibrate_ln: bool = True, max_error: float = 0.2, fallback_to_original: bool = True, protect_critical: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, skip_attention_qkv: bool = True, target_attention_only: bool = False, target_modules: list[str] | None = None, tokenizer: Any | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("ast_pruner")
@COMPRESSOR_REGISTRY.register("ast")
class ASTPruner(BaseCompressor):
    """
    AST Pruner: vectorized semi-structured 2:4 / N:M pruning with empirical
    Jacobian calibration and PCA coordinate rotation (saliency only).

    Design notes:
      * PCA rotation is used **for saliency estimation only**. The actual
        2:4 mask is applied to the ORIGINAL weights so the sparsity pattern
        is preserved and remains exploitable by sparse backends.
      * The Jacobian proxy (0.5·||logits||²) is a heuristic inspired by
        SparseGPT; it is NOT the true Fisher Information Matrix.
    """

    def __init__(
        self,
        model: nn.Module,
        block_size: int = 4,
        n_keep: int = 2,
        dim: int = 1,
        use_pca: bool = True,
        use_fisher: bool = True,
        recalibrate_ln: bool = True,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        protect_critical: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        skip_attention_qkv: bool = True,
        target_attention_only: bool = False,
        target_modules: list[str] | None = None,
        tokenizer: Any | None = None,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="ast_pruner", **kwargs)
        conf = cfg or kwargs.get("comp_config")
        if conf is not None:
            block_size = _cfg_get(conf, "ast_block_size", _cfg_get(conf, "block_size", block_size))
            n_keep = _cfg_get(conf, "ast_n_keep", _cfg_get(conf, "n_keep", n_keep))
            dim = _cfg_get(conf, "ast_dim", _cfg_get(conf, "dim", dim))
            use_pca = _cfg_get(conf, "ast_use_pca", _cfg_get(conf, "use_pca", use_pca))
            use_fisher = _cfg_get(conf, "ast_use_fisher", _cfg_get(conf, "use_fisher", use_fisher))
            recalibrate_ln = _cfg_get(
                conf,
                "ast_recalibrate_ln",
                _cfg_get(conf, "recalibrate_ln", recalibrate_ln),
            )
            max_error = _cfg_get(conf, "ast_max_error", _cfg_get(conf, "max_error", max_error))
            fallback_to_original = _cfg_get(
                conf,
                "ast_fallback_to_original",
                _cfg_get(conf, "fallback_to_original", fallback_to_original),
            )
            protect_critical = _cfg_get(conf, "protect_critical", protect_critical)
            lora_r = _cfg_get(conf, "lora_r", _cfg_get(conf, "lora_rank", lora_r))
            lora_alpha = _cfg_get(conf, "lora_alpha", lora_alpha)
            skip_attention_qkv = _cfg_get(conf, "skip_attention_qkv", False)

        self.block_size = max(2, int(block_size))
        self.n_keep = max(1, min(int(n_keep), self.block_size - 1))
        self.dim = dim
        self.use_pca = bool(use_pca)
        self.use_fisher = bool(use_fisher)
        self.recalibrate_ln = bool(recalibrate_ln)
        if max_error is None or max_error == 0.20:
            max_error = adaptive_max_error(
                int(768 * (self.n_keep / self.block_size)), 768, 768, "distilbert"
            )

        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.protect_critical = bool(protect_critical)
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.skip_attention_qkv = bool(skip_attention_qkv)
        self.target_attention_only = target_attention_only
        self.target_modules = target_modules
        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.target_device = (
            torch.device(device)
            if device is not None
            else next(model.parameters(), torch.tensor(0.0)).device
        )
        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.stats: dict[str, Any] = {}

    # ------------------------------------------------------------------
    # Target selection
    # ------------------------------------------------------------------
    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            name_lower = name.lower()
            if any(
                p in name_lower
                for p in [
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                ]
            ):
                return True
        return False

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        cls_name = module.__class__.__name__.lower()
        is_linear_like = (
            isinstance(module, nn.Linear) or "conv1d" in cls_name or hasattr(module, "nf")
        )
        if not is_linear_like or isinstance(module, ASTPrunedLoRALinear):
            return False

        if self.protect_critical:
            name_lower = name.lower()
            if any(
                skip in name_lower
                for skip in [
                    "classifier",
                    "score",
                    "lm_head",
                    "pooler",
                    "embed",
                    "logits_processor",
                ]
            ):
                return False

        from neuralops.utils.compat import should_compress_linear

        if not should_compress_linear(self.model, name):
            return False
        if self._is_gqa_asymmetric(name, module):
            return False
        if self.target_attention_only and not any(
            k in name.lower()
            for k in [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "out_proj",
                "o_proj",
                "q_lin",
                "k_lin",
                "v_lin",
                "out_lin",
                "query",
                "key",
                "value",
                "c_attn",
                "c_proj",
            ]
        ):
            return False
        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    # ------------------------------------------------------------------
    # Calibration: collect activations and Jacobian-proxy squared gradients
    # ------------------------------------------------------------------
    def _collect_jacobian_and_activations(
        self,
    ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
        """
        Returns (activation_covariances, jacobian_proxy_squared).
        AST-FIX-06: uses eval() + enable_grad() (not train()).
        AST-FIX-07: always averages over batches.
        AST-FIX-13: passes attention_mask when present.
        """
        activations: dict[str, list[torch.Tensor]] = {}
        fisher_info: dict[str, torch.Tensor] = {}
        hooks = []

        def make_hook(name: str):
            def hook(module, inp, outp):
                if inp and isinstance(inp[0], torch.Tensor):
                    activations.setdefault(name, []).append(inp[0].detach().cpu())

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                hooks.append(module.register_forward_hook(make_hook(name)))

        # Build calibration batches
        batches: list[Any] = []
        if self.calibration_loader is not None:
            for b in self.calibration_loader:
                batches.append(b)
                if len(batches) >= 4:
                    break
        if not batches:
            batches = [
                get_calibration_tensors(
                    self.model,
                    tokenizer=self.tokenizer,
                    calibration_loader=self.calibration_loader,
                    calibration_texts=self.calibration_texts,
                    device=self.target_device,
                )
            ]

        # AST-FIX-06: eval mode (no dropout) + enable_grad for backward
        self.model.eval()
        n_batches_run = 0
        try:
            with torch.enable_grad():
                for batch_input in batches:
                    out = None
                    if isinstance(batch_input, dict):
                        moved = {
                            k: v.to(self.target_device) if hasattr(v, "to") else v
                            for k, v in batch_input.items()
                        }
                        out = self.model(**moved)
                    elif isinstance(batch_input, torch.Tensor):
                        out = self.model(batch_input.to(self.target_device))
                    elif isinstance(batch_input, (list, tuple)):
                        out = self.model(
                            *[
                                b.to(self.target_device) if hasattr(b, "to") else b
                                for b in batch_input
                            ]
                        )

                    if self.use_fisher and out is not None:
                        # AST-FIX-14: explicit naming - this is NOT true Fisher
                        # 0.5·||logits||² is a Jacobian-weighted norm proxy.
                        logits = None
                        if hasattr(out, "logits"):
                            logits = out.logits
                        elif isinstance(out, torch.Tensor):
                            logits = out
                        elif (
                            isinstance(out, (tuple, list))
                            and out
                            and isinstance(out[0], torch.Tensor)
                        ):
                            logits = out[0]

                        if logits is not None and logits.requires_grad:
                            loss = 0.5 * torch.sum(logits.pow(2))
                            self.model.zero_grad(set_to_none=True)
                            loss.backward()
                            for name, module in self.model.named_modules():
                                if not self._is_target_module(name, module):
                                    continue
                                # Resolve weight
                                target_w = getattr(module, "weight", None)
                                base = getattr(module, "base_layer", module)
                                if target_w is None:
                                    target_w = getattr(base, "weight", None)
                                if target_w is None or target_w.grad is None:
                                    continue
                                g = target_w.grad.detach()
                                g_sq = g.pow(2).cpu()
                                if name in fisher_info:
                                    fisher_info[name] = fisher_info[name] + g_sq
                                else:
                                    fisher_info[name] = g_sq
                            self.model.zero_grad(set_to_none=True)
                            n_batches_run += 1
        except Exception as e:
            logger.warning("AST Jacobian calibration failed: %s", e)
        finally:
            for h in hooks:
                h.remove()

        # AST-FIX-07: always divide (even for n_batches_run == 1)
        if n_batches_run >= 1:
            inv = 1.0 / float(n_batches_run)
            for k in list(fisher_info.keys()):
                fisher_info[k] = fisher_info[k] * inv
        else:
            logger.warning(
                "AST: no successful backward passes; saliency will fall back to |W| only."
            )

        # Build covariance map, then discard raw activations (AST-FIX-11)
        cov_map: dict[str, torch.Tensor] = {}
        for name, act_list in activations.items():
            if not act_list:
                continue
            act_concat = torch.cat([a.reshape(-1, a.shape[-1]) for a in act_list], dim=0).float()
            N, D = act_concat.shape
            if N < 2:
                continue
            mean = torch.mean(act_concat, dim=0, keepdim=True)
            centered = act_concat - mean
            cov = torch.matmul(centered.T, centered) / float(N - 1)
            # AST-FIX-09: explicit symmetrize
            cov = (cov + cov.T) * 0.5
            cov_map[name] = cov
            # Free memory
            del act_concat, centered, mean
            if name not in fisher_info:
                fisher_info[name] = torch.var(torch.zeros(1, D), dim=0, keepdim=True).expand(1, D)
        activations.clear()

        return cov_map, fisher_info

    # ------------------------------------------------------------------
    # PCA basis extraction (per-layer, no covariance retained)
    # ------------------------------------------------------------------
    @staticmethod
    def _pca_basis(cov: torch.Tensor, d_in: int, device: torch.device) -> torch.Tensor:
        """
        AST-FIX-08: scaled eigenvalue regularization.
        AST-FIX-09: symmetrize before eigh.
        Returns orthonormal Q (d_in, d_in) in original column basis.
        """
        cov = cov.to(device).float()
        if cov.shape != (d_in, d_in):
            return torch.eye(d_in, device=device)

        # Eigenvalue regularization proportional to mean variance
        diag_mean = float(cov.diagonal().mean().item())
        eps = max(1e-8, 1e-6 * diag_mean)
        cov = cov + eps * torch.eye(d_in, device=device)

        try:
            eigvals, Q = torch.linalg.eigh(cov)
            idx = torch.argsort(eigvals, descending=True)
            Q = Q[:, idx]
            if torch.isnan(Q).any() or torch.isinf(Q).any():
                logger.warning("AST: PCA basis contained NaN/Inf; using identity.")
                return torch.eye(d_in, device=device)
            return Q.contiguous()
        except Exception as e:
            logger.warning("AST: eigh failed (%s); using identity.", e)
            return torch.eye(d_in, device=device)

    # ------------------------------------------------------------------
    # Compress
    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        from neuralops.core.model_utils import replace_module

        cov_map, fisher_map = self._collect_jacobian_and_activations()
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        seen_weights: set = set()

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue
            if not hasattr(module, "weight") or module.weight is None:
                continue

            w_id = id(module.weight)
            if w_id in seen_weights:
                for p in module.parameters():
                    p.requires_grad = False
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            target_mod = getattr(module, "base_layer", module)
            raw_w = target_mod.weight.data
            cls_name = module.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")

            # W: (d_out, d_in)
            W = raw_w.t().clone().float() if is_conv1d else raw_w.clone().float()
            d_out, d_in = W.shape
            device = W.device
            dtype = target_mod.weight.dtype

            # ------------------------------------------------------------------
            # 1. Compute saliency in ORIGINAL basis (AST-FIX-01)
            # ------------------------------------------------------------------
            # Saliency is computed as:
            #   1a. Rotate Fisher into PCA basis: F_rot = F @ (Q ⊙ Q)
            #   1b. Apply 2:4-aware importance per (row, column) in the rotated
            #       basis, then project the rotation-weighted per-column score
            #       back to the original column basis using Q.
            #
            # In practice, we use the following simplified and provably
            # structure-preserving form:
            #   S_ij = |W_ij| · sqrt(F_proj_ij + ε)
            # where F_proj = F @ (Q ⊙ Q)  (a per-element projection back).
            # This yields a dense saliency map in the ORIGINAL basis, but the
            # actual mask is applied to W (not W_rot), preserving 2:4 sparsity.
            # ------------------------------------------------------------------
            fisher = fisher_map.get(name, None)
            if fisher is not None and self.use_fisher:
                F_mat = fisher.to(device).float()
                # AST-FIX-10: NaN/Inf guard
                F_mat = torch.nan_to_num(F_mat, nan=0.0, posinf=0.0, neginf=0.0)

                if is_conv1d and F_mat.dim() == 2 and F_mat.shape == raw_w.shape:
                    F_mat = F_mat.t().contiguous()

                if F_mat.dim() == 2 and F_mat.shape == W.shape:
                    if self.use_pca and name in cov_map:
                        Q = self._pca_basis(cov_map[name], d_in, device)
                        # Project Fisher into rotated basis, then back
                        F_rot = F_mat @ (Q * Q)  # (d_out, d_in) × (d_in, d_in)
                        F_proj = F_rot @ Q.T  # (d_out, d_in) × (d_in, d_in) → (d_out, d_in)
                    else:
                        F_proj = F_mat
                elif F_mat.shape[-1] == d_in and F_mat.dim() == 1:
                    F_proj = F_mat.reshape(1, d_in).expand(d_out, d_in)
                else:
                    F_proj = torch.ones_like(W)
                F_proj = torch.clamp(F_proj, min=1e-8)
                saliency = torch.abs(W) * torch.sqrt(F_proj)
            else:
                saliency = torch.abs(W)

            saliency = torch.nan_to_num(saliency, nan=0.0, posinf=1e4, neginf=0.0)

            # ------------------------------------------------------------------
            # 2. Apply 2:4 structured mask on ORIGINAL weights
            # ------------------------------------------------------------------
            W_pruned, _mask = apply_24_structured_mask(W, saliency=saliency, dim=self.dim)
            W_final = torch.nan_to_num(W_pruned, nan=0.0, posinf=0.0, neginf=0.0)

            # ------------------------------------------------------------------
            # 3. Error check
            # ------------------------------------------------------------------
            w_norm = torch.linalg.norm(W).item()
            err = float(torch.linalg.norm(W - W_final).item() / max(w_norm, 1e-12))
            self.layer_errors[name] = err

            if err > self.max_error:
                logger.warning(
                    "AST: reconstruction error %.4f on %s exceeds %.4f.",
                    err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    logger.info("AST: falling back to original for %s.", name)
                    # AST-FIX-02: only freeze the base weight, not LoRA
                    for pn, p in module.named_parameters(recurse=False):
                        if "lora_" in pn.lower():
                            continue  # keep LoRA trainable
                        p.requires_grad = False
                    self.fallback_layers.append(name)
                    continue

            # ------------------------------------------------------------------
            # 4. Replace / in-place update
            # ------------------------------------------------------------------
            if self.lora_r > 0:
                new_layer = ASTPrunedLoRALinear.from_linear(
                    target_mod,
                    weight=W_final.contiguous(),
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                )

                # AST-FIX-04: type-safe LoRA transfer
                if hasattr(module, "base_layer"):
                    if hasattr(module, "lora_A") and module.lora_A is not None:
                        src_a = getattr(module.lora_A, "weight", module.lora_A)
                        if (
                            isinstance(src_a, torch.Tensor)
                            and src_a.shape == new_layer.lora_A.shape
                        ):
                            new_layer.lora_A.data.copy_(src_a.to(new_layer.lora_A.dtype))
                    if hasattr(module, "lora_B") and module.lora_B is not None:
                        src_b = getattr(module.lora_B, "weight", module.lora_B)
                        if (
                            isinstance(src_b, torch.Tensor)
                            and src_b.shape == new_layer.lora_B.shape
                        ):
                            new_layer.lora_B.data.copy_(src_b.to(new_layer.lora_B.dtype))

                    # Swap the base layer in place
                    if hasattr(module, "base_layer"):
                        module.base_layer = new_layer
                    # AST-FIX-03: removed dead `if name == "":` branch
                    # Don't call replace_module for the wrapped module
                    # because we already replaced its inner base_layer
                else:
                    if name != "":
                        replace_module(self.model, name, new_layer)
                    else:
                        self.model = new_layer
            else:
                # AST-FIX-05: no_grad around mutation
                with torch.no_grad():
                    new_w = W_final.to(dtype=dtype, device=device)
                    if is_conv1d:
                        new_w = new_w.t().contiguous()
                    target_mod.weight.data.copy_(new_w)

            self.compressed_layers.append(name)

        # ------------------------------------------------------------------
        # LayerNorm recalibration
        # ------------------------------------------------------------------
        if self.recalibrate_ln and (
            self.calibration_loader is not None or self.calibration_texts is not None
        ):
            try:
                recalibrate_layernorms(
                    self.model,
                    tokenizer=self.tokenizer,
                    calibration_loader=self.calibration_loader,
                    calibration_texts=self.calibration_texts,
                    device=self.target_device,
                )
            except Exception as e:
                logger.warning("AST: LayerNorm recalibration skipped: %s", e)

        # ------------------------------------------------------------------
        # Forward sanity check (AST-FIX-13: passes attention_mask)
        # ------------------------------------------------------------------
        forward_check_success = False
        try:
            self.model.eval()
            with torch.no_grad():
                d = next(self.model.parameters()).device
                dummy = torch.randint(0, 100, (1, 16), device=d)
                mask = torch.ones_like(dummy)
                try:
                    self.model(input_ids=dummy, attention_mask=mask)
                except TypeError:
                    self.model(dummy)
            forward_check_success = True
        except Exception as e:
            logger.warning("AST: post-compression forward check failed: %s", e)

        # AST-FIX-15: log skipped layer count
        if self.skipped_layers:
            logger.info(
                "AST: skipped %d layers (dup weights / target rules).",
                len(self.skipped_layers),
            )

        self._compressed_params = self._count_params(self.model)
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        stats = self.get_compression_stats()
        stats.update(
            {
                "pruned_layers": len(self.compressed_layers),
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "avg_reconstruction_error": avg_err,
                "block_size": self.block_size,
                "n_keep": self.n_keep,
                "pattern": f"{self.n_keep}:{self.block_size}",
                "use_pca": self.use_pca,
                "use_fisher": self.use_fisher,
                "recalibrate_ln": self.recalibrate_ln,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "forward_check_success": forward_check_success,
                "method": "ast_jacobian_pca_2_4_pruning",
            }
        )
        self.stats = stats
        logger.info(
            "Compressor ast finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

AST Pruner: vectorized semi-structured 2:4 / N:M pruning with empirical Jacobian calibration and PCA coordinate rotation (saliency only).

Design notes: * PCA rotation is used for saliency estimation only. The actual 2:4 mask is applied to the ORIGINAL weights so the sparsity pattern is preserved and remains exploitable by sparse backends. * The Jacobian proxy (0.5·||logits||²) is a heuristic inspired by SparseGPT; it is NOT the true Fisher Information Matrix.

Ancestors

BaseCompressor

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class AdaSVDCompressor (model: nn.Module, rank: int | None = None, target_param_ratio: float = 0.7, num_iters: int = 10, max_error: float = 0.5, fallback_to_original: bool = True, target_attention_only: bool = True, min_rank: int = 16, target_modules: list[str] | None = None, tokenizer: Any | None = None, calibration_texts: list[str] | None = None, use_weighted_als: bool = False, use_mean_shift: bool = True, use_gating_d: bool = True, device: str | torch.device | None = None, comp_config: Any = None, max_calib_tokens: int = 4096, fused_qkv_error_threshold: float = 0.3, gelu_layer_patterns: tuple[str, ...] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("adasvd")
@COMPRESSOR_REGISTRY.register("adasvd_compressor")
class AdaSVDCompressor(BaseCompressor):
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module,
        rank: int | None = None,
        target_param_ratio: float = 0.7,
        num_iters: int = 10,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        target_attention_only: bool = True,
        min_rank: int = 16,
        target_modules: list[str] | None = None,
        tokenizer: Any | None = None,
        calibration_texts: list[str] | None = None,
        use_weighted_als: bool = False,
        use_mean_shift: bool = True,
        use_gating_d: bool = True,
        device: str | torch.device | None = None,
        comp_config: Any = None,
        max_calib_tokens: int = 4096,
        fused_qkv_error_threshold: float = 0.30,
        gelu_layer_patterns: tuple[str, ...] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="adasvd", **kwargs)
        cfg = comp_config or kwargs.get(
            "comp_config", kwargs.get("cfg", getattr(model, "config", None))
        )

        # ---- CLI overrides ----
        cli_rank = None
        cli_use_gating_d = None
        cli_use_weighted_als = None
        for idx, arg in enumerate(sys.argv):
            if arg in ("--rank", "--adasvd_rank") and idx + 1 < len(sys.argv):
                try:
                    cli_rank = int(sys.argv[idx + 1])
                except ValueError:
                    pass
            elif arg.startswith(("--rank=", "--adasvd_rank=")):
                try:
                    cli_rank = int(arg.split("=")[1])
                except ValueError:
                    pass
            elif arg in ("--adasvd_use_gating_d", "--use_gating_d") and idx + 1 < len(sys.argv):
                cli_use_gating_d = sys.argv[idx + 1].lower() in ("true", "1", "yes")
            elif arg.startswith(("--adasvd_use_gating_d=", "--use_gating_d=")):
                cli_use_gating_d = arg.split("=")[1].lower() in ("true", "1", "yes")
            elif arg in (
                "--adasvd_use_weighted_als",
                "--use_weighted_als",
            ) and idx + 1 < len(sys.argv):
                cli_use_weighted_als = sys.argv[idx + 1].lower() in ("true", "1", "yes")
            elif arg.startswith(("--adasvd_use_weighted_als=", "--use_weighted_als=")):
                cli_use_weighted_als = arg.split("=")[1].lower() in ("true", "1", "yes")

        self.rank = (
            cli_rank
            if cli_rank is not None
            else _cfg_get(cfg, "adasvd_rank", _cfg_get(cfg, "rank", rank))
        )
        self.target_param_ratio = float(
            _cfg_get(
                cfg,
                "adasvd_target_ratio",
                _cfg_get(cfg, "target_param_ratio", target_param_ratio),
            )
        )
        self.num_iters = int(_cfg_get(cfg, "adasvd_iters", _cfg_get(cfg, "num_iters", num_iters)))
        self.max_error = float(
            _cfg_get(cfg, "adasvd_max_error", _cfg_get(cfg, "max_error", max_error))
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "adasvd_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "adasvd_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.min_rank = int(_cfg_get(cfg, "adasvd_min_rank", _cfg_get(cfg, "min_rank", min_rank)))

        self.use_weighted_als = (
            cli_use_weighted_als
            if cli_use_weighted_als is not None
            else bool(
                _cfg_get(
                    cfg,
                    "adasvd_use_weighted_als",
                    _cfg_get(cfg, "use_weighted_als", use_weighted_als),
                )
            )
        )
        self.use_mean_shift = bool(
            _cfg_get(
                cfg,
                "adasvd_use_mean_shift",
                _cfg_get(cfg, "use_mean_shift", use_mean_shift),
            )
        )
        self.use_gating_d = (
            cli_use_gating_d
            if cli_use_gating_d is not None
            else bool(
                _cfg_get(
                    cfg,
                    "adasvd_use_gating_d",
                    _cfg_get(cfg, "use_gating_d", use_gating_d),
                )
            )
        )

        self.target_modules = target_modules
        self.max_calib_tokens = int(max_calib_tokens)
        self.fused_qkv_error_threshold = float(
            _cfg_get(cfg, "adasvd_fused_qkv_error_threshold", fused_qkv_error_threshold)
        )

        if gelu_layer_patterns is None:
            gelu_layer_patterns = (
                "lin1",
                "intermediate.dense",
                "gate_proj",
                "up_proj",
                "dense_h_to_4h",
                "c_fc",
                "wi_0",
                "wi_1",
            )
        self.gelu_layer_patterns = tuple(gelu_layer_patterns)

        # Tokenizer auto-resolve
        if tokenizer is None:
            tokenizer = getattr(model, "tokenizer", None) or getattr(
                getattr(model, "config", None), "tokenizer", None
            )
        if tokenizer is None:
            mp = getattr(getattr(model, "config", None), "_name_or_path", None)
            if mp:
                try:
                    from transformers import AutoTokenizer

                    tokenizer = AutoTokenizer.from_pretrained(mp)
                    if tokenizer.pad_token is None:
                        tokenizer.pad_token = tokenizer.eos_token
                except Exception as e:
                    logger.debug("Could not auto-load tokenizer from %s: %s", mp, e)
        self.tokenizer = tokenizer

        self.calibration_texts = (
            calibration_texts
            or _cfg_get(cfg, "calibration_texts", None)
            or DEFAULT_CALIBRATION_CORPUS
        )
        self.device = (
            torch.device(device) if device is not None else next(model.parameters()).device
        )
        self.comp_config = cfg

        # State
        self.calibration_activations: dict[str, torch.Tensor] = {}
        self.layer_means: dict[str, torch.Tensor] = {}
        self.layer_out_means: dict[str, torch.Tensor] = {}
        self.layer_out_vars: dict[str, torch.Tensor] = {}
        self.mean_shifts: dict[str, torch.Tensor] = {}

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.layer_iters: dict[str, int] = {}

        self.calibration_time_s: float = 0.0

        self._orig_storage_bytes: int = 0
        self._comp_storage_bytes: int = 0
        self._compressed_params: int = 0
        self._compressed_names_seen: set = set()

    # ------------------------------------------------------------------
    # Target selection
    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        if isinstance(module, (SVDLinear, FusedQKVSVDLinear)):
            return False
        if any(
            name.endswith(sfx) for sfx in (".base_layer", ".lora_A", ".lora_B", ".lora_dropout")
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = module.__class__.__name__.lower()
        target_mod.__class__.__name__.lower()

        weight_tensor = getattr(target_mod, "weight", getattr(module, "weight", None))
        if weight_tensor is None or not hasattr(weight_tensor, "dim") or weight_tensor.dim() != 2:
            return False

        is_linear_like = (
            isinstance(module, nn.Linear)
            or isinstance(target_mod, nn.Linear)
            or (
                Conv1D is not None
                and (isinstance(module, Conv1D) or isinstance(target_mod, Conv1D))
            )
            or hasattr(module, "nf")
            or hasattr(target_mod, "nf")
            or "loralinear" in cls_name
            or hasattr(module, "base_layer")
        )
        if not is_linear_like:
            return False

        name_lower = name.lower()
        if any(
            skip in name_lower
            for skip in (
                "lm_head",
                "classifier",
                "head",
                "output_layer",
                "score",
                "pooler",
                "embedding",
                "layernorm",
                "rmsnorm",
            )
        ):
            return False

        if self.target_attention_only:
            ffn_patterns = (
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            )
            if any(p in name_lower for p in ffn_patterns):
                return False
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            attn_patterns = (
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
                "in_proj_qkv",
                "in_proj_z",
                "in_proj_b",
                "out_proj",
            )
            if not any(p in name_lower for p in attn_patterns):
                return False
        else:
            try:
                if not should_compress_linear(self.model, name):
                    return False
            except Exception:
                pass

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    @staticmethod
    def _is_fused_qkv_layer(name: str, module: nn.Module) -> bool:
        if "c_attn" not in name.lower():
            return False
        target_mod = getattr(module, "base_layer", module)
        nf = getattr(module, "nf", None) or getattr(target_mod, "nf", None)
        if nf is None or nf % 3 != 0:
            return False
        w = getattr(target_mod, "weight", None)
        return w is not None and w.dim() == 2

    @staticmethod
    def _gelu_derivative_squared(y: torch.Tensor) -> torch.Tensor:
        phi = 0.5 * (1.0 + torch.erf(y / math.sqrt(2.0)))
        pdf = torch.exp(-y.pow(2) / 2.0) / math.sqrt(2.0 * math.pi)
        gelu_prime = phi + y * pdf
        return gelu_prime.pow(2)

    def _is_followed_by_gelu(self, layer_name: str) -> bool:
        name_lower = layer_name.lower()
        return any(p in name_lower for p in self.gelu_layer_patterns)

    @staticmethod
    def _compute_effective_rank(W: torch.Tensor) -> float:
        with torch.no_grad():
            s = torch.linalg.svdvals(W.float())
            s_sq = s.pow(2)
            denom = s_sq.sum().clamp(min=1e-12)
            p = s_sq / denom
            p = p[p > 1e-12]
            return float(torch.exp(-(p * torch.log(p + 1e-12)).sum()).item())

    _compute_nuclear_norm = _compute_effective_rank

    # ------------------------------------------------------------------
    # Calibration
    # ------------------------------------------------------------------
    def _collect_activations(self, calibration_data: Any | None = None) -> dict[str, torch.Tensor]:
        activations: dict[str, list[torch.Tensor]] = {}
        out_activations: dict[str, list[torch.Tensor]] = {}
        hooks = []
        max_tokens = self.max_calib_tokens

        def make_hook(layer_name: str):
            def hook(module, inp, outp):
                if not inp or inp[0] is None:
                    return
                x = inp[0].detach()
                target_mod = getattr(module, "base_layer", module)
                is_conv1d = (
                    (
                        Conv1D is not None
                        and (isinstance(module, Conv1D) or isinstance(target_mod, Conv1D))
                    )
                    or hasattr(module, "nf")
                    or hasattr(target_mod, "nf")
                )
                weight_tensor = getattr(target_mod, "weight", getattr(module, "weight", None))
                if weight_tensor is None or weight_tensor.dim() != 2:
                    return

                in_f = weight_tensor.shape[0] if is_conv1d else weight_tensor.shape[1]
                x_flat = x.float().reshape(-1, in_f)
                mask = x_flat.abs().sum(dim=-1) > 1e-5
                x_active = x_flat[mask] if mask.sum() > 0 else x_flat

                bucket = activations.setdefault(layer_name, [])
                curr_n = sum(t.size(0) for t in bucket)
                if curr_n < max_tokens:
                    avail = max_tokens - curr_n
                    bucket.append(x_active[:avail].cpu())

                if isinstance(outp, torch.Tensor):
                    out_f = weight_tensor.shape[1] if is_conv1d else weight_tensor.shape[0]
                    y_flat = outp.detach().float().reshape(-1, out_f)
                    y_bucket = out_activations.setdefault(layer_name, [])
                    curr_ny = sum(t.size(0) for t in y_bucket)
                    if curr_ny < max_tokens:
                        avail_y = max_tokens - curr_ny
                        y_bucket.append(y_flat[:avail_y].cpu())

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                hooks.append(module.register_forward_hook(make_hook(name)))

        if not hooks:
            return {}

        self.model.eval()
        dev = next(self.model.parameters()).device

        try:
            with torch.no_grad():
                if calibration_data is not None:
                    self._run_calibration(calibration_data, dev)
                elif self.tokenizer is not None and self.calibration_texts:
                    inputs = self.tokenizer(
                        self.calibration_texts,
                        return_tensors="pt",
                        padding=True,
                        truncation=True,
                        max_length=512,
                    )
                    inputs = {k: v.to(dev) for k, v in inputs.items()}
                    self.model(**inputs)
                else:
                    gen = torch.Generator(device="cpu").manual_seed(42)
                    dummy_ids = torch.randint(100, 2000, (8, 64), generator=gen).to(dev)
                    try:
                        self.model(input_ids=dummy_ids)
                    except Exception:
                        try:
                            self.model(dummy_ids)
                        except Exception:
                            try:
                                dummy_x = torch.randn(128, 64, device=dev)
                                self.model(dummy_x)
                            except Exception as e:
                                logger.debug("AdaSVD synthetic calibration failed: %s", e)
        except Exception as e:
            logger.debug("AdaSVD calibration pass exception: %s", e)
        finally:
            for h in hooks:
                h.remove()

        result: dict[str, torch.Tensor] = {}
        for name, tensor_list in activations.items():
            if tensor_list:
                cat_x = torch.cat(tensor_list, dim=0)
                result[name] = cat_x
                self.layer_means[name] = cat_x.mean(dim=0).cpu()

        for name, y_list in out_activations.items():
            if y_list:
                cat_y = torch.cat(y_list, dim=0)
                self.layer_out_means[name] = cat_y.mean(dim=0).cpu()

                if self.use_gating_d and self._is_followed_by_gelu(name):
                    gelu_sq = self._gelu_derivative_squared(cat_y)
                    d_values = gelu_sq.mean(dim=0)
                else:
                    d_values = cat_y.var(dim=0, unbiased=False)

                self.layer_out_vars[name] = d_values.cpu()

        return result

    def _run_calibration(self, calibration_data: Any, dev: torch.device) -> None:
        if isinstance(calibration_data, torch.Tensor):
            try:
                self.model(calibration_data.to(dev))
            except Exception:
                self.model(input_ids=calibration_data.to(dev))
        elif isinstance(calibration_data, (list, tuple)) and self.tokenizer is not None:
            inputs = self.tokenizer(
                list(calibration_data),
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=512,
            )
            inputs = {k: v.to(dev) for k, v in inputs.items()}
            self.model(**inputs)
        elif hasattr(calibration_data, "__iter__"):
            for batch in calibration_data:
                try:
                    if isinstance(batch, torch.Tensor):
                        self.model(batch.to(dev))
                    elif isinstance(batch, dict):
                        b_dev = {
                            k: v.to(dev) for k, v in batch.items() if isinstance(v, torch.Tensor)
                        }
                        self.model(**b_dev)
                    else:
                        continue
                except Exception as e:
                    logger.debug("AdaSVD calibration batch skipped: %s", e)
                    continue

    # ------------------------------------------------------------------
    # SVD helpers
    # ------------------------------------------------------------------
    def _init_whitened_svd(
        self,
        W: torch.Tensor,
        r: int,
        C_X: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        d_out, d_in = W.shape
        device = W.device
        dtype = W.dtype

        if C_X is not None and C_X.shape == (d_in, d_in):
            C_dev = C_X.to(device=device, dtype=torch.float32)
            eps = max(1e-8, float(1e-6 * torch.trace(C_dev).item() / max(1, d_in)))
            eye = torch.eye(d_in, device=device, dtype=torch.float32)
            S_white = None
            try:
                S_white = torch.linalg.cholesky(C_dev + eps * eye)
            except Exception:
                for exp in (-5, -4, -3, -2, -1):
                    try:
                        reg = 10**exp * (torch.trace(C_dev).item() / max(1, d_in))
                        S_white = torch.linalg.cholesky(C_dev + reg * eye)
                        break
                    except Exception:
                        continue

            if S_white is not None:
                try:
                    S_inv = torch.linalg.inv(S_white)
                    W_white = W.float() @ S_white
                    U, S, Vt = torch.linalg.svd(W_white, full_matrices=False)
                    Vt = Vt @ S_inv
                except Exception:
                    U, S, Vt = torch.linalg.svd(W.float(), full_matrices=False)
            else:
                U, S, Vt = torch.linalg.svd(W.float(), full_matrices=False)
        else:
            U, S, Vt = torch.linalg.svd(W.float(), full_matrices=False)

        sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
        A = U[:, :r] * sqrt_S.unsqueeze(0)
        B = sqrt_S.unsqueeze(1) * Vt[:r, :]
        return A.to(dtype=dtype, device=device), B.to(dtype=dtype, device=device)

    def _two_sided_als(
        self,
        W: torch.Tensor,
        r: int,
        C_X: torch.Tensor,
        D_var: torch.Tensor,
        X_calib: torch.Tensor | None = None,
        mu_x: torch.Tensor | None = None,
        num_iters: int = 10,
        patience: int = 3,
    ) -> tuple[torch.Tensor, torch.Tensor, float, int]:
        device = W.device
        dtype = W.dtype

        A, B = self._init_whitened_svd(W, r, C_X)
        A = A.float()
        B = B.float()

        E = C_X.float().to(device)
        D = torch.diag(torch.clamp(D_var.float().to(device), min=1e-6))
        WE = W.float() @ E
        DW = D @ W.float()

        def _err(recon):
            if X_calib is not None:
                X_eval = (X_calib - mu_x.to(device)) if mu_x is not None else X_calib
                WX_c = W.float() @ X_eval.float().t()
                reconX_c = recon @ X_eval.float().t()
                return float(
                    torch.norm(WX_c - reconX_c, p="fro").item()
                    / (torch.norm(WX_c, p="fro").item() + 1e-9)
                )
            diff = W.float() - recon
            num = math.sqrt(max(0.0, torch.trace(diff @ E @ diff.t()).item()))
            den = math.sqrt(max(1e-9, torch.trace(W.float() @ E @ W.float().t()).item()))
            return float(num / (den + 1e-9))

        best_A, best_B = A.clone(), B.clone()
        best_err = _err(A @ B)
        patience_counter = 0
        actual_iters = 0

        for _ in range(num_iters):
            BEBt = B @ E @ B.t()
            eps_B = max(1e-9, float(1e-6 * BEBt.trace().item() / max(1, r)))
            try:
                A_new = (
                    WE
                    @ B.t()
                    @ torch.linalg.inv(BEBt + eps_B * torch.eye(r, device=device, dtype=BEBt.dtype))
                )
            except Exception:
                break

            AtDA = A_new.t() @ D @ A_new
            eps_A = max(1e-9, float(1e-6 * AtDA.trace().item() / max(1, r)))
            try:
                B_new = (
                    torch.linalg.inv(AtDA + eps_A * torch.eye(r, device=device, dtype=AtDA.dtype))
                    @ A_new.t()
                    @ DW
                )
            except Exception:
                break

            if not torch.isfinite(A_new).all() or not torch.isfinite(B_new).all():
                break

            actual_iters += 1
            err = _err(A_new @ B_new)

            if err < best_err - 1e-6:
                best_A, best_B = A_new.clone(), B_new.clone()
                best_err = err
                patience_counter = 0
            else:
                patience_counter += 1
                if patience_counter >= patience:
                    break

            A, B = A_new, B_new

        return best_A.to(dtype), best_B.to(dtype), best_err, actual_iters

    # ------------------------------------------------------------------
    # Fused QKV compression
    # ------------------------------------------------------------------
    def _compress_fused_qkv(
        self,
        layer: nn.Module,
        rank: int,
        name: str,
    ) -> tuple[FusedQKVSVDLinear | None, float, int]:
        target_mod = getattr(layer, "base_layer", layer)
        w = getattr(target_mod, "weight", None)
        if w is None or w.dim() != 2:
            return None, float("inf"), 0

        raw_w = w.data
        in_f, out_f = raw_w.shape
        if out_f % 3 != 0:
            logger.warning("FusedQKV %s: out_features %d not divisible by 3", name, out_f)
            return None, float("inf"), 0

        split_size = out_f // 3
        max_rank = min(in_f, split_size)
        r_floor = max(16, min(int(0.15 * split_size), self.min_rank * 4))
        r = max(r_floor, min(int(rank), max_rank))

        device = raw_w.device
        dtype = raw_w.dtype
        has_bias = getattr(target_mod, "bias", None) is not None

        fused_layer = FusedQKVSVDLinear(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=has_bias,
            dtype=dtype,
            device=device,
            lora_r=0,
            lora_alpha=1.0,
        )

        try:
            fused_layer.set_fused_factors(raw_w)
        except Exception as e:
            logger.warning("FusedQKV %s factor split failed: %s", name, e)
            return None, float("inf"), 0

        # ---- Mean-shift ----
        mean_shift = None
        mu_x = None
        if self.use_mean_shift:
            mu_x = self.layer_means.get(name, None)
            if mu_x is not None:
                mu_x = mu_x.to(device=device, dtype=torch.float32)
                mean_shift = raw_w.float().t() @ mu_x
                self.mean_shifts[name] = mean_shift.cpu()
                fused_layer.set_input_mean(mu_x)

        if has_bias and fused_layer.bias is not None:
            with torch.no_grad():
                base_b = target_mod.bias.data.float().to(device)
                if mean_shift is not None:
                    base_b = base_b + mean_shift
                fused_layer.bias.data.copy_(base_b.to(dtype=dtype, device=device))

        # ---- Reconstruction error ----
        with torch.no_grad():
            W_recon = fused_layer.weight.to(device=device, dtype=torch.float32)
            W_orig = raw_w.float()
            norm_orig = torch.norm(W_orig, p="fro").clamp_min(1e-8)
            rel_err = (torch.norm(W_orig - W_recon, p="fro") / norm_orig).item()

        threshold = self.fused_qkv_error_threshold
        if r >= 0.5 * max_rank:
            threshold = max(threshold, 0.45)

        if rel_err > threshold:
            logger.warning(
                "AdaSVD fused-QKV %s: err=%.4f > threshold %.4f; falling back to original.",
                name,
                rel_err,
                threshold,
            )
            return None, rel_err, 0

        logger.info(
            "AdaSVD fused-QKV: %s in=%d out=%d split=%d rank=%d err=%.4f",
            name,
            in_f,
            out_f,
            split_size,
            r,
            rel_err,
        )
        return fused_layer, rel_err, 0

    # ------------------------------------------------------------------
    # Single-layer compression
    # ------------------------------------------------------------------
    def compress_linear(
        self,
        layer: nn.Module,
        rank: int,
        name: str = "",
        covariance: torch.Tensor | None = None,
        X_calib: torch.Tensor | None = None,
        out_variance: torch.Tensor | None = None,
        **kwargs: Any,
    ) -> tuple[nn.Module | None, float, int]:
        if self._is_fused_qkv_layer(name, layer):
            return self._compress_fused_qkv(layer, rank, name)

        target_mod = getattr(layer, "base_layer", layer)
        is_conv1d = (
            (Conv1D is not None and (isinstance(layer, Conv1D) or isinstance(target_mod, Conv1D)))
            or hasattr(layer, "nf")
            or hasattr(target_mod, "nf")
        )

        weight_tensor = getattr(target_mod, "weight", getattr(layer, "weight", None))
        if weight_tensor is None or weight_tensor.dim() != 2:
            return None, float("inf"), 0

        raw_w = weight_tensor.data
        if is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float()

        max_param_r = max(1, min(in_f, out_f, int(0.75 * (in_f * out_f) / (in_f + out_f))))
        if self.rank is not None:
            r_floor = max(1, min(self.min_rank, rank))
        else:
            r_floor = max(self.min_rank, int(0.15 * min(in_f, out_f)))
        r = max(r_floor, min(rank, max_param_r))

        device = weight_tensor.device
        dtype = weight_tensor.dtype

        # ---- Mean-shift ----
        mu_x = self.layer_means.get(name, None) if self.use_mean_shift else None
        mean_shift = None
        if mu_x is not None:
            mu_x = mu_x.to(device=device, dtype=torch.float32)
            mean_shift = (W.float() @ mu_x).to(dtype=torch.float32)
            self.mean_shifts[name] = mean_shift.cpu()

        X = X_calib if X_calib is not None else self.calibration_activations.get(name, None)
        X_centered = None
        C_X = None
        if X is not None:
            X = X.to(device=device, dtype=torch.float32)
            X_centered = (X - mu_x) if mu_x is not None else X
            C_X = (X_centered.t() @ X_centered) / max(1, X_centered.shape[0])
        elif covariance is not None and covariance.shape == (in_f, in_f):
            C_X = covariance.to(device=device, dtype=torch.float32)

        D_var = out_variance if out_variance is not None else self.layer_out_vars.get(name, None)

        if C_X is not None and D_var is not None and self.use_weighted_als and self.num_iters > 0:
            A, B, final_err, actual_iters = self._two_sided_als(
                W,
                r,
                C_X,
                D_var,
                X_calib=X,
                mu_x=mu_x,
                num_iters=self.num_iters,
            )
        else:
            A, B = self._init_whitened_svd(W, r, C_X)
            actual_iters = 0
            recon = A @ B
            if X is not None:
                X_eval = X_centered if X_centered is not None else X
                WX = W.float() @ X_eval.t()
                reconX = recon.float() @ X_eval.t()
                final_err = float(
                    torch.norm(WX - reconX, p="fro").item()
                    / (torch.norm(WX, p="fro").item() + 1e-9)
                )
            else:
                W_norm = torch.norm(W, p="fro").item()
                final_err = (
                    0.0
                    if W_norm < 1e-12
                    else float(torch.norm(W - recon, p="fro").item() / (W_norm + 1e-9))
                )

        bias_tensor = getattr(target_mod, "bias", getattr(layer, "bias", None))
        has_bias = (bias_tensor is not None) or (mean_shift is not None)

        svd_layer = SVDLinear(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=has_bias,
            dtype=dtype,
            device=device,
        )
        svd_layer.set_low_rank_factors(
            U_factor=A.to(dtype=dtype, device=device),
            V_factor=B.to(dtype=dtype, device=device),
        )

        if has_bias and svd_layer.bias is not None:
            if bias_tensor is not None:
                b_base = bias_tensor.data.float().to(device)
            else:
                b_base = torch.zeros(out_f, device=device, dtype=torch.float32)
            new_bias = b_base + mean_shift if mean_shift is not None else b_base
            svd_layer.bias.data.copy_(new_bias.to(dtype=dtype, device=device))

        # ---- Mean-shift registration on SVDLinear ----
        # SVDLinear exposes `_input_mean` (may be None). If it is a buffer,
        # we register and copy; otherwise we set attribute.
        if mu_x is not None:
            mu_f = mu_x.to(device=device, dtype=dtype)
            if hasattr(svd_layer, "register_buffer"):
                try:
                    # Only register if not already registered
                    if not any(n == "_input_mean" for n, _ in svd_layer.named_buffers()):
                        svd_layer.register_buffer("_input_mean", mu_f, persistent=True)
                    else:
                        svd_layer._input_mean.copy_(mu_f)
                except Exception:
                    svd_layer._input_mean = mu_f
        # else: keep SVDLinear's default (None)

        return svd_layer, final_err, actual_iters

    # ------------------------------------------------------------------
    # Main compress
    # ------------------------------------------------------------------
    def compress(self, calibration_data: Any | None = None) -> tuple[nn.Module, dict[str, Any]]:
        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}
        self.layer_iters = {}
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0
        self._compressed_names_seen = set()

        target_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            logger.warning("AdaSVD: no compressible layers found.")
            return self.model, self.get_compression_stats()

        # Track original bytes
        for name, mod in target_layers:
            target_mod = getattr(mod, "base_layer", mod)
            try:
                for p in target_mod.parameters():
                    self._orig_storage_bytes += p.numel() * p.element_size()
            except Exception:
                pass

        # ---- Step 1: calibration ----
        t_calib_start = time.monotonic()
        self.calibration_activations = self._collect_activations(calibration_data=calibration_data)
        self.calibration_time_s = time.monotonic() - t_calib_start
        logger.info("AdaSVD calibration time: %.2f s", self.calibration_time_s)

        # ---- Step 2: spectral importance ----
        layer_importances: dict[str, float] = {}
        total_importance = 0.0
        for name, module in target_layers:
            target_mod = getattr(module, "base_layer", module)
            weight_tensor = getattr(target_mod, "weight", getattr(module, "weight", None))
            if weight_tensor is None or weight_tensor.dim() != 2:
                continue
            is_conv1d = (
                (
                    Conv1D is not None
                    and (isinstance(module, Conv1D) or isinstance(target_mod, Conv1D))
                )
                or hasattr(module, "nf")
                or hasattr(target_mod, "nf")
            )
            raw_w = weight_tensor.data
            W_2d = raw_w.t().contiguous().float() if is_conv1d else raw_w.float()
            imp = self._compute_effective_rank(W_2d)
            layer_importances[name] = max(imp, 1e-4)
            total_importance += layer_importances[name]

        # ---- Step 3: compress each layer ----
        for name, module in target_layers:
            target_mod = getattr(module, "base_layer", module)
            weight_tensor = getattr(target_mod, "weight", getattr(module, "weight", None))
            if weight_tensor is None or weight_tensor.dim() != 2:
                continue
            is_conv1d = (
                (
                    Conv1D is not None
                    and (isinstance(module, Conv1D) or isinstance(target_mod, Conv1D))
                )
                or hasattr(module, "nf")
                or hasattr(target_mod, "nf")
            )
            raw_w = weight_tensor.data
            if is_conv1d:
                in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            else:
                out_f, in_f = raw_w.shape[0], raw_w.shape[1]

            max_param_r = max(1, min(in_f, out_f, int(0.75 * (in_f * out_f) / (in_f + out_f))))
            if self.rank is not None:
                r_floor = max(1, min(self.min_rank, self.rank))
                denom = max(1, len(target_layers))
                frac = layer_importances.get(name, 1.0) / max(total_importance / denom, 1e-8)
                allocated_rank = int(round(self.rank * frac))
            else:
                r_floor = max(self.min_rank, int(0.15 * min(in_f, out_f)))
                frac = layer_importances.get(name, 1.0) / max(total_importance, 1e-8)
                allocated_rank = int(
                    round(min(in_f, out_f) * self.target_param_ratio * frac * len(target_layers))
                )

            if self._is_fused_qkv_layer(name, module):
                if self.rank is not None:
                    r_floor = max(r_floor, min(288, self.rank))
                    max_param_r = min(max_param_r, 432)
                else:
                    r_floor = max(r_floor, 288)
                    max_param_r = min(max_param_r, 432)

            allocated_rank = max(r_floor, min(allocated_rank, max_param_r))

            svd_layer, err, iters = self.compress_linear(
                module,
                rank=allocated_rank,
                name=name,
            )
            self.layer_errors[name] = err
            self.layer_ranks[name] = allocated_rank
            self.layer_iters[name] = iters

            if svd_layer is None or err > self.max_error:
                logger.warning(
                    "AdaSVD: %s reconstruction error %.4f > max_error %.4f; "
                    "falling back to original.",
                    name,
                    err,
                    self.max_error,
                )
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    # Fallback keeps original bytes, so add them to comp
                    try:
                        for p in target_mod.parameters():
                            self._comp_storage_bytes += p.numel() * p.element_size()
                    except Exception:
                        pass
                    continue

            # ---- Splice ----
            if hasattr(module, "base_layer"):
                module.base_layer = svd_layer
                for attr in ("lora_A", "lora_B"):
                    sub = getattr(module, attr, None)
                    if sub is None:
                        continue
                    params = sub.parameters() if hasattr(sub, "parameters") else iter([sub])
                    for p in params:
                        if hasattr(p, "requires_grad"):
                            p.requires_grad = True
            else:
                replace_module(self.model, name, svd_layer)
            self.compressed_layers.append(name)
            self._compressed_names_seen.add(name)

            # ---- Storage accounting ----
            try:
                for p in svd_layer.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
                for b in svd_layer.buffers():
                    if b is not None:
                        self._comp_storage_bytes += b.numel() * b.element_size()
            except Exception:
                pass

        # ---- Freeze fallback layers ----
        for name, mod in self.model.named_modules():
            if name in self.fallback_layers:
                target = getattr(mod, "base_layer", mod)
                for p in target.parameters():
                    p.requires_grad = False

        # ---- Recompute honest storage ----
        comp_bytes_real = 0
        for name, mod in self.model.named_modules():
            if name in self.compressed_layers or name in self.fallback_layers:
                target = getattr(mod, "base_layer", mod)
                try:
                    for p in target.parameters():
                        comp_bytes_real += p.numel() * p.element_size()
                    for b in target.buffers():
                        if b is not None:
                            comp_bytes_real += b.numel() * b.element_size()
                except Exception:
                    pass
        self._comp_storage_bytes = comp_bytes_real

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        storage_ratio = (
            self._orig_storage_bytes / self._comp_storage_bytes
            if self._comp_storage_bytes > 0
            else 1.0
        )
        effective_ratio = (
            self._comp_storage_bytes / self._orig_storage_bytes
            if self._orig_storage_bytes > 0
            else 1.0
        )

        stats = self.get_compression_stats()
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "layer_iters": self.layer_iters,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "avg_reconstruction_error": avg_err,
                "calibration_time_s": self.calibration_time_s,
                "use_weighted_als": self.use_weighted_als,
                "use_mean_shift": self.use_mean_shift,
                "use_gating_d": self.use_gating_d,
                "orig_storage_bytes": self._orig_storage_bytes,
                "comp_storage_bytes": self._comp_storage_bytes,
                "storage_compression_ratio": round(storage_ratio, 3),
                "storage_reduction_pct": round(100.0 * (1.0 - effective_ratio), 2),
            }
        )
        try:
            self.model.compression_stats = stats
        except Exception:
            pass

        logger.info(
            "AdaSVD done: %d compressed, %d fallback, avg_err=%.4f, storage_ratio=%.2fx",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            storage_ratio,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def compress_linear(self, layer: nn.Module, rank: int, name: str = '', covariance: torch.Tensor | None = None, X_calib: torch.Tensor | None = None, out_variance: torch.Tensor | None = None, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module | None, float, int]

Expand source code

def compress_linear(
    self,
    layer: nn.Module,
    rank: int,
    name: str = "",
    covariance: torch.Tensor | None = None,
    X_calib: torch.Tensor | None = None,
    out_variance: torch.Tensor | None = None,
    **kwargs: Any,
) -> tuple[nn.Module | None, float, int]:
    if self._is_fused_qkv_layer(name, layer):
        return self._compress_fused_qkv(layer, rank, name)

    target_mod = getattr(layer, "base_layer", layer)
    is_conv1d = (
        (Conv1D is not None and (isinstance(layer, Conv1D) or isinstance(target_mod, Conv1D)))
        or hasattr(layer, "nf")
        or hasattr(target_mod, "nf")
    )

    weight_tensor = getattr(target_mod, "weight", getattr(layer, "weight", None))
    if weight_tensor is None or weight_tensor.dim() != 2:
        return None, float("inf"), 0

    raw_w = weight_tensor.data
    if is_conv1d:
        in_f, out_f = raw_w.shape[0], raw_w.shape[1]
        W = raw_w.t().contiguous().float()
    else:
        out_f, in_f = raw_w.shape[0], raw_w.shape[1]
        W = raw_w.float()

    max_param_r = max(1, min(in_f, out_f, int(0.75 * (in_f * out_f) / (in_f + out_f))))
    if self.rank is not None:
        r_floor = max(1, min(self.min_rank, rank))
    else:
        r_floor = max(self.min_rank, int(0.15 * min(in_f, out_f)))
    r = max(r_floor, min(rank, max_param_r))

    device = weight_tensor.device
    dtype = weight_tensor.dtype

    # ---- Mean-shift ----
    mu_x = self.layer_means.get(name, None) if self.use_mean_shift else None
    mean_shift = None
    if mu_x is not None:
        mu_x = mu_x.to(device=device, dtype=torch.float32)
        mean_shift = (W.float() @ mu_x).to(dtype=torch.float32)
        self.mean_shifts[name] = mean_shift.cpu()

    X = X_calib if X_calib is not None else self.calibration_activations.get(name, None)
    X_centered = None
    C_X = None
    if X is not None:
        X = X.to(device=device, dtype=torch.float32)
        X_centered = (X - mu_x) if mu_x is not None else X
        C_X = (X_centered.t() @ X_centered) / max(1, X_centered.shape[0])
    elif covariance is not None and covariance.shape == (in_f, in_f):
        C_X = covariance.to(device=device, dtype=torch.float32)

    D_var = out_variance if out_variance is not None else self.layer_out_vars.get(name, None)

    if C_X is not None and D_var is not None and self.use_weighted_als and self.num_iters > 0:
        A, B, final_err, actual_iters = self._two_sided_als(
            W,
            r,
            C_X,
            D_var,
            X_calib=X,
            mu_x=mu_x,
            num_iters=self.num_iters,
        )
    else:
        A, B = self._init_whitened_svd(W, r, C_X)
        actual_iters = 0
        recon = A @ B
        if X is not None:
            X_eval = X_centered if X_centered is not None else X
            WX = W.float() @ X_eval.t()
            reconX = recon.float() @ X_eval.t()
            final_err = float(
                torch.norm(WX - reconX, p="fro").item()
                / (torch.norm(WX, p="fro").item() + 1e-9)
            )
        else:
            W_norm = torch.norm(W, p="fro").item()
            final_err = (
                0.0
                if W_norm < 1e-12
                else float(torch.norm(W - recon, p="fro").item() / (W_norm + 1e-9))
            )

    bias_tensor = getattr(target_mod, "bias", getattr(layer, "bias", None))
    has_bias = (bias_tensor is not None) or (mean_shift is not None)

    svd_layer = SVDLinear(
        in_features=in_f,
        out_features=out_f,
        rank=r,
        bias=has_bias,
        dtype=dtype,
        device=device,
    )
    svd_layer.set_low_rank_factors(
        U_factor=A.to(dtype=dtype, device=device),
        V_factor=B.to(dtype=dtype, device=device),
    )

    if has_bias and svd_layer.bias is not None:
        if bias_tensor is not None:
            b_base = bias_tensor.data.float().to(device)
        else:
            b_base = torch.zeros(out_f, device=device, dtype=torch.float32)
        new_bias = b_base + mean_shift if mean_shift is not None else b_base
        svd_layer.bias.data.copy_(new_bias.to(dtype=dtype, device=device))

    # ---- Mean-shift registration on SVDLinear ----
    # SVDLinear exposes `_input_mean` (may be None). If it is a buffer,
    # we register and copy; otherwise we set attribute.
    if mu_x is not None:
        mu_f = mu_x.to(device=device, dtype=dtype)
        if hasattr(svd_layer, "register_buffer"):
            try:
                # Only register if not already registered
                if not any(n == "_input_mean" for n, _ in svd_layer.named_buffers()):
                    svd_layer.register_buffer("_input_mean", mu_f, persistent=True)
                else:
                    svd_layer._input_mean.copy_(mu_f)
            except Exception:
                svd_layer._input_mean = mu_f
    # else: keep SVDLinear's default (None)

    return svd_layer, final_err, actual_iters

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class AdaptiveRanksCompressor (model: nn.Module, energy_threshold: float = 0.92, min_rank: int = 16, max_error: float = 0.2, fallback_to_original: bool = True, protect_critical: bool = True, target_attention_only: bool = False, target_modules: list[str] | None = None, tokenizer: Any | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, use_jg_svd: bool = False, use_cholesky_jg: bool = True, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("adaptive_ranks")
@COMPRESSOR_REGISTRY.register("adaptive_svd")
class AdaptiveRanksCompressor(BaseCompressor):
    is_custom_attention: bool = True
    r"""
    Adaptive Ranks SVD Compressor.

    Dynamically determines ranks per layer according to cumulative singular value energy threshold:
        $$\frac{\sum_{i=1}^r \sigma_i^2}{\sum_{j=1}^d \sigma_j^2} \ge \tau \quad (\tau = 0.92)$$
    with minimum rank floor $r_{\min} = \max(16, \lfloor 0.15 \cdot \min(d_{in}, d_{out}) \rfloor, \lfloor 0.10 \cdot d_{out} \rfloor)$,
    activation covariance whitening, and symmetrically distributed singular values:
        $$U_r = U_{:, :r} \sqrt{\Sigma_r}, \quad V_r = (\sqrt{\Sigma_r} V^\top_{:r, :}) \cdot \operatorname{Cov}(X)^{-1/2}$$

    Production hardening features:
    - Relative Frobenius reconstruction error validation with configurable max_error (default 0.20).
    - Safe fallback to original weights by default (fallback_to_original=True).
    - Robust Conv1D and LoRALinear support via base_layer weight transposition.
    - SVD and covariance eigendecomposition computed strictly in FP32 with NaN/Inf guards.
    - Unified critical-layer protection via should_compress_linear.
    - Adaptive rank retry when reconstruction error exceeds threshold.
    - Parameter reduction verification: ensures r * (in + out) < in * out.
    - Comprehensive per-layer statistics and logging.
    """

    def __init__(
        self,
        model: nn.Module,
        energy_threshold: float = 0.92,
        min_rank: int = 16,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        protect_critical: bool = True,
        target_attention_only: bool = False,
        target_modules: list[str] | None = None,
        tokenizer: Any | None = None,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        use_jg_svd: bool = False,
        use_cholesky_jg: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="adaptive_ranks", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is None and hasattr(model, "config") and model.config is not None:
            cfg = model.config

        # Check sys.argv for CLI overrides
        cli_use_jg_svd = None
        cli_use_cholesky_jg = None
        import sys

        for _idx, _arg in enumerate(sys.argv):
            if _arg in (
                "--use_jg_svd",
                "--adaptive_ranks_use_jg_svd",
            ) and _idx + 1 < len(sys.argv):
                cli_use_jg_svd = sys.argv[_idx + 1].lower() in ("true", "1", "yes")
            elif _arg.startswith("--use_jg_svd=") or _arg.startswith(
                "--adaptive_ranks_use_jg_svd="
            ):
                cli_use_jg_svd = _arg.split("=")[1].lower() in ("true", "1", "yes")
            elif _arg in (
                "--use_cholesky_jg",
                "--adaptive_ranks_use_cholesky_jg",
            ) and _idx + 1 < len(sys.argv):
                cli_use_cholesky_jg = sys.argv[_idx + 1].lower() in ("true", "1", "yes")
            elif _arg.startswith("--use_cholesky_jg=") or _arg.startswith(
                "--adaptive_ranks_use_cholesky_jg="
            ):
                cli_use_cholesky_jg = _arg.split("=")[1].lower() in ("true", "1", "yes")

        if cfg is not None:
            energy_threshold = _cfg_get(
                cfg,
                "energy_threshold",
                _cfg_get(
                    cfg,
                    "adaptive_ranks_energy_threshold",
                    _cfg_get(cfg, "asvd_energy_threshold", energy_threshold),
                ),
            )
            min_rank = _cfg_get(cfg, "min_rank", _cfg_get(cfg, "rank", min_rank))
            max_error = _cfg_get(
                cfg,
                "max_error",
                _cfg_get(
                    cfg,
                    "adaptive_ranks_max_error",
                    _cfg_get(cfg, "aasvd_max_error", max_error),
                ),
            )
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            protect_critical = _cfg_get(cfg, "protect_critical", protect_critical)
            target_attention_only = _cfg_get(
                cfg,
                "target_attention_only",
                _cfg_get(
                    cfg,
                    "adaptive_ranks_target_attention_only",
                    _cfg_get(cfg, "aasvd_target_attention_only", target_attention_only),
                ),
            )
            use_jg_svd = _cfg_get(
                cfg,
                "use_jg_svd",
                _cfg_get(cfg, "adaptive_ranks_use_jg_svd", use_jg_svd),
            )
            use_cholesky_jg = _cfg_get(
                cfg,
                "use_cholesky_jg",
                _cfg_get(cfg, "adaptive_ranks_use_cholesky_jg", use_cholesky_jg),
            )

        if cli_use_jg_svd is not None:
            use_jg_svd = cli_use_jg_svd
        if cli_use_cholesky_jg is not None:
            use_cholesky_jg = cli_use_cholesky_jg

        if "use_jg_svd" in kwargs and kwargs["use_jg_svd"] is not None:
            use_jg_svd = kwargs["use_jg_svd"]
        if "use_cholesky_jg" in kwargs and kwargs["use_cholesky_jg"] is not None:
            use_cholesky_jg = kwargs["use_cholesky_jg"]

        if "rank" in kwargs and kwargs["rank"] is not None:
            min_rank = kwargs["rank"]
        if "target_attention_only" in kwargs and kwargs["target_attention_only"] is not None:
            target_attention_only = kwargs["target_attention_only"]

        self.energy_threshold = max(0.01, min(float(energy_threshold), 1.0))
        self.min_rank = max(1, int(min_rank))
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.protect_critical = bool(protect_critical)
        self.target_attention_only = bool(target_attention_only)
        self.use_jg_svd = bool(use_jg_svd)
        self.use_cholesky_jg = bool(use_cholesky_jg)
        self.target_modules = target_modules
        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.target_device = (
            torch.device(device)
            if device is not None
            else next(model.parameters(), torch.tensor(0.0)).device
        )

        self.layer_D_gating: dict[str, torch.Tensor] = {}
        self.attn_jacobian: dict[str, float] = {}

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.used_covariance: dict[str, bool] = {}
        self.used_jg_svd: dict[str, bool] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        if isinstance(module, SVDLinear) or (
            CompressedAttentionLayer is not None and isinstance(module, CompressedAttentionLayer)
        ):
            return False
        # Do not target inner submodules of LoRA or wrappers
        if any(
            name.endswith(sfx) for sfx in [".base_layer", ".lora_A", ".lora_B", ".lora_dropout"]
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = module.__class__.__name__.lower()
        target_cls = target_mod.__class__.__name__.lower()

        is_linear_like = (
            isinstance(module, nn.Linear)
            or isinstance(target_mod, nn.Linear)
            or "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(module, "nf")
            or hasattr(target_mod, "nf")
            or "loralinear" in cls_name
            or hasattr(module, "base_layer")
        )
        is_mha = isinstance(module, nn.MultiheadAttention) or isinstance(
            target_mod, nn.MultiheadAttention
        )
        if not (is_linear_like or is_mha):
            return False

        if self.protect_critical:
            name_lower = name.lower()
            if any(
                skip in name_lower
                for skip in [
                    "lm_head",
                    "classifier",
                    "head",
                    "output_layer",
                    "score",
                    "pooler",
                    "embed",
                ]
            ):
                return False

        if self.target_attention_only:
            name_lower = name.lower()
            ffn_patterns = [
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            ]
            if any(p in name_lower for p in ffn_patterns):
                return False
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            attn_patterns = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(p in name_lower for p in attn_patterns):
                return False
        else:
            if not should_compress_linear(self.model, name):
                return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def _compute_energy_rank(
        self,
        S: torch.Tensor,
        max_rank: int,
        min_rank_floor: int | None = None,
        threshold: float | None = None,
    ) -> int:
        r"""
        Find smallest rank $r$ satisfying cumulative singular energy ratio $\ge$ threshold:
        $$\frac{\sum_{i=1}^r \sigma_i^2}{\sum_{j=1}^d \sigma_j^2} \ge \tau$$
        constrained by minimum rank floor and parameter reduction bounds.
        """
        thresh = threshold if threshold is not None else self.energy_threshold
        r_floor = min_rank_floor if min_rank_floor is not None else self.min_rank
        r_floor = max(1, min(int(r_floor), max_rank))
        s_squared = S**2
        total_energy = s_squared.sum()
        if total_energy <= 1e-12:
            return r_floor

        cum_energy = torch.cumsum(s_squared, dim=0) / total_energy
        idx = torch.where(cum_energy >= thresh)[0]
        r = int(idx[0].item()) + 1 if len(idx) > 0 else max_rank
        return max(1, min(max(r, r_floor), max_rank))

    def _collect_activation_covariances(self) -> dict[str, torch.Tensor]:
        """Collect empirical input activation covariances across target nn.Linear / Conv1D modules."""
        if not (self.tokenizer or self.calibration_loader or self.calibration_texts):
            logger.info("No calibration data available; using unweighted SVD.")
            return {}

        self.model.eval()
        running_cov: dict[str, torch.Tensor] = {}
        sample_counts: dict[str, int] = {}
        hooks = []

        def make_hook(layer_name: str) -> Any:
            def hook(module: nn.Module, inp: tuple[torch.Tensor, ...], outp: Any) -> None:
                if len(inp) > 0 and isinstance(inp[0], torch.Tensor):
                    x = inp[0].detach().float()
                    target_mod = getattr(module, "base_layer", module)
                    cls_name = module.__class__.__name__.lower()
                    target_cls = target_mod.__class__.__name__.lower()
                    is_conv1d = (
                        "conv1d" in cls_name
                        or "conv1d" in target_cls
                        or hasattr(module, "nf")
                        or hasattr(target_mod, "nf")
                    )
                    weight = getattr(target_mod, "weight", getattr(module, "weight", None))
                    if weight is None:
                        return
                    in_f = weight.shape[0] if is_conv1d else weight.shape[1]
                    if x.ndim >= 2:
                        flat_x = x.reshape(-1, in_f)
                    else:
                        flat_x = x.unsqueeze(0)
                    n_samples = flat_x.shape[0]
                    if n_samples > 0:
                        batch_cov = torch.matmul(flat_x.T, flat_x).cpu()
                        if layer_name not in running_cov:
                            running_cov[layer_name] = batch_cov
                            sample_counts[layer_name] = n_samples
                        else:
                            running_cov[layer_name] += batch_cov
                            sample_counts[layer_name] += n_samples

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module) and not isinstance(
                getattr(module, "base_layer", module), nn.MultiheadAttention
            ):
                hooks.append(module.register_forward_hook(make_hook(name)))

        if self.calibration_loader is not None:
            self.model.eval()
            with torch.no_grad():
                for batch in self.calibration_loader:
                    try:
                        if isinstance(batch, dict):
                            self.model(
                                **{
                                    k: (v.to(self.target_device) if hasattr(v, "to") else v)
                                    for k, v in batch.items()
                                }
                            )
                        elif isinstance(batch, torch.Tensor):
                            self.model(batch.to(self.target_device))
                        elif isinstance(batch, (list, tuple)):
                            try:
                                self.model(
                                    *[
                                        (b.to(self.target_device) if hasattr(b, "to") else b)
                                        for b in batch
                                    ]
                                )
                            except Exception:
                                self.model(batch[0].to(self.target_device))
                    except Exception as e:
                        logger.debug(
                            "AdaptiveRanks calibration loader forward pass exception: %s",
                            e,
                        )
        else:
            calib_inputs = get_calibration_tensors(
                self.model,
                tokenizer=self.tokenizer,
                calibration_loader=None,
                calibration_texts=self.calibration_texts,
                device=self.target_device,
            )

            if calib_inputs is not None and not (
                isinstance(calib_inputs, (list, tuple, dict)) and len(calib_inputs) == 0
            ):
                with torch.no_grad():
                    try:
                        if isinstance(calib_inputs, dict):
                            self.model(**calib_inputs)
                        elif isinstance(calib_inputs, torch.Tensor):
                            self.model(calib_inputs)
                        elif isinstance(calib_inputs, (list, tuple)):
                            self.model(*calib_inputs)
                    except Exception as e:
                        logger.debug(
                            "AdaptiveRanks covariance collection forward pass exception: %s",
                            e,
                        )

        for h in hooks:
            h.remove()

        covariances: dict[str, torch.Tensor] = {}
        for name, cov_sum in running_cov.items():
            n = float(max(sample_counts.get(name, 1), 1))
            covariances[name] = cov_sum / n

        return covariances

    def _collect_gating_D(self) -> None:
        r"""
        Collect output gating importance vectors D for FFN layers and Softmax Jacobian scalars for Attention Q/K.
        FFN: D_{FFN} = E[(f'(Y))^2] where f is GELU or SiLU.
        Attention Q/K: D_{Attn} = E[P(1-P)] where P is softmax attention probabilities.
        """
        if not (self.tokenizer or self.calibration_loader or self.calibration_texts):
            logger.info("No calibration data available for JG-SVD gating D collection.")
            return

        self.model.eval()
        hooks = []
        out_acts: dict[str, list[torch.Tensor]] = {}
        attn_probs: dict[str, list[torch.Tensor]] = {}

        def make_out_hook(layer_name: str) -> Any:
            def hook(module: nn.Module, inp: tuple[torch.Tensor, ...], outp: Any) -> None:
                if isinstance(outp, torch.Tensor):
                    y = outp.detach().float()
                    target_mod = getattr(module, "base_layer", module)
                    cls_name = module.__class__.__name__.lower()
                    target_cls = target_mod.__class__.__name__.lower()
                    is_conv1d = (
                        "conv1d" in cls_name
                        or "conv1d" in target_cls
                        or hasattr(module, "nf")
                        or hasattr(target_mod, "nf")
                    )
                    weight = getattr(target_mod, "weight", getattr(module, "weight", None))
                    if weight is None:
                        return
                    out_f = weight.shape[1] if is_conv1d else weight.shape[0]
                    y_flat = y.reshape(-1, out_f)
                    if layer_name not in out_acts:
                        out_acts[layer_name] = [y_flat[:512].cpu()]
                    else:
                        curr_n = sum(t.size(0) for t in out_acts[layer_name])
                        if curr_n < 1024:
                            avail = 1024 - curr_n
                            out_acts[layer_name].append(y_flat[:avail].cpu())

            return hook

        def make_attn_hook(layer_name: str) -> Any:
            def hook(module: nn.Module, inp: tuple[torch.Tensor, ...], outp: Any) -> None:
                p = None
                if (
                    isinstance(outp, (tuple, list))
                    and len(outp) > 1
                    and isinstance(outp[1], torch.Tensor)
                ):
                    p = outp[1].detach().float()
                elif isinstance(outp, torch.Tensor) and outp.ndim >= 3:
                    # In some modules attention probabilities or softmax is exposed
                    pass
                if p is not None:
                    if layer_name not in attn_probs:
                        attn_probs[layer_name] = [p[:4].cpu()]
                    else:
                        curr_n = sum(t.size(0) for t in attn_probs[layer_name])
                        if curr_n < 8:
                            attn_probs[layer_name].append(p[:2].cpu())

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                hooks.append(module.register_forward_hook(make_out_hook(name)))
            nl = name.lower()
            if any(k in nl for k in ["attention", "self_attn", "attn"]) and not any(
                skip in nl for skip in ["q_lin", "k_lin", "v_lin", "out_lin", "dense"]
            ):
                hooks.append(module.register_forward_hook(make_attn_hook(name)))

        if self.calibration_loader is not None:
            self.model.eval()
            with torch.no_grad():
                for batch in self.calibration_loader:
                    try:
                        if isinstance(batch, dict):
                            self.model(
                                **{
                                    k: (v.to(self.target_device) if hasattr(v, "to") else v)
                                    for k, v in batch.items()
                                }
                            )
                        elif isinstance(batch, torch.Tensor):
                            self.model(batch.to(self.target_device))
                        elif isinstance(batch, (list, tuple)):
                            try:
                                self.model(
                                    *[
                                        (b.to(self.target_device) if hasattr(b, "to") else b)
                                        for b in batch
                                    ]
                                )
                            except Exception:
                                self.model(batch[0].to(self.target_device))
                    except Exception as e:
                        logger.debug("JG-SVD calibration loader forward pass exception: %s", e)
        else:
            calib_inputs = get_calibration_tensors(
                self.model,
                tokenizer=self.tokenizer,
                calibration_loader=None,
                calibration_texts=self.calibration_texts,
                device=self.target_device,
            )
            if calib_inputs is not None and not (
                isinstance(calib_inputs, (list, tuple, dict)) and len(calib_inputs) == 0
            ):
                with torch.no_grad():
                    try:
                        if isinstance(calib_inputs, dict):
                            self.model(**calib_inputs)
                        elif isinstance(calib_inputs, torch.Tensor):
                            self.model(calib_inputs)
                        elif isinstance(calib_inputs, (list, tuple)):
                            self.model(*calib_inputs)
                    except Exception as e:
                        logger.debug("JG-SVD gating collection forward pass exception: %s", e)

        for h in hooks:
            h.remove()

        # Compute Attention Softmax scalar Jacobian E[P(1-P)]
        global_attn_scalar = 0.08  # Default baseline for softmax distribution with entropy ~2.5
        for name, p_list in attn_probs.items():
            if p_list:
                cat_p = torch.cat(p_list, dim=0)
                jac = cat_p * (1.0 - cat_p)
                m_jac = float(jac.mean().item())
                if math.isfinite(m_jac) and m_jac > 0.0:
                    self.attn_jacobian[name] = m_jac
                    global_attn_scalar = m_jac

        for name, y_list in out_acts.items():
            if not y_list:
                continue
            cat_y = torch.cat(y_list, dim=0)
            act_type = _detect_activation_type(name, self.model)
            if act_type == "gelu":
                d_sq = _gelu_derivative_sq(cat_y)
                d_vec = d_sq.mean(dim=0).cpu()
                self.layer_D_gating[name] = torch.clamp(d_vec, min=1e-6)
                logger.debug("JG-SVD GELU' weighting on %s: mean=%.4f", name, d_vec.mean().item())
            elif act_type == "silu":
                d_sq = _silu_derivative_sq(cat_y)
                d_vec = d_sq.mean(dim=0).cpu()
                self.layer_D_gating[name] = torch.clamp(d_vec, min=1e-6)
                logger.debug("JG-SVD SiLU' weighting on %s: mean=%.4f", name, d_vec.mean().item())
            elif _is_attention_qk(name):
                # Apply Attention Softmax Jacobian scalar weighting
                matching_attn = [v for k, v in self.attn_jacobian.items() if k in name or name in k]
                scalar_val = matching_attn[0] if matching_attn else global_attn_scalar
                out_dim = cat_y.shape[-1]
                self.layer_D_gating[name] = torch.full(
                    (out_dim,), float(scalar_val), dtype=torch.float32
                )
                logger.debug(
                    "JG-SVD Attention Softmax Jacobian on %s: scalar=%.4f",
                    name,
                    scalar_val,
                )
            else:
                # Other layers: empirical output variance fallback
                d_var = cat_y.var(dim=0, unbiased=False).cpu()
                self.layer_D_gating[name] = torch.clamp(d_var, min=1e-6)

    def _factorize_with_JG(
        self,
        W: torch.Tensor,
        cov: torch.Tensor | None,
        D_gating: torch.Tensor | None,
        r: int,
        num_iters: int = 8,
    ) -> tuple[torch.Tensor, torch.Tensor, float]:
        r"""
        Two-Sided Weighted ALS factorizing W ~ U_r @ V_r minimizing:
            || D^{1/2} (W - U V) C_X^{1/2} ||_F^2
        with input activation covariance C_X and output gating derivative D.
        """
        out_f, in_f = W.shape
        device = W.device
        dtype = W.dtype

        # 1. Whitened SVD initialization
        if cov is not None and cov.shape == (in_f, in_f):
            cov_dev = cov.to(device=device, dtype=torch.float32)
            diag_mean = torch.diagonal(cov_dev).mean().item()
            if not math.isfinite(diag_mean) or diag_mean <= 0.0:
                diag_mean = 1.0
            lambda_reg = max(1e-4 * diag_mean, 1e-6)
            cov_reg = cov_dev + lambda_reg * torch.eye(in_f, device=device)

            if self.use_cholesky_jg:
                try:
                    S_white = torch.linalg.cholesky(cov_reg)
                    S_inv = torch.linalg.inv(S_white)
                    M = torch.matmul(W, S_white)
                    U, S, Vt = torch.linalg.svd(M, full_matrices=False)
                    Vt = Vt @ S_inv
                except Exception:
                    eigvals, eigvecs = torch.linalg.eigh(cov_reg)
                    eigvals_clamped = torch.clamp(eigvals, min=1e-6)
                    cov_sqrt = eigvecs @ torch.diag(torch.sqrt(eigvals_clamped)) @ eigvecs.T
                    cov_inv_sqrt = (
                        eigvecs @ torch.diag(1.0 / torch.sqrt(eigvals_clamped)) @ eigvecs.T
                    )
                    M = torch.matmul(W, cov_sqrt)
                    U, S, Vt = torch.linalg.svd(M, full_matrices=False)
                    Vt = Vt @ cov_inv_sqrt
            else:
                eigvals, eigvecs = torch.linalg.eigh(cov_reg)
                eigvals_clamped = torch.clamp(eigvals, min=1e-6)
                cov_sqrt = eigvecs @ torch.diag(torch.sqrt(eigvals_clamped)) @ eigvecs.T
                cov_inv_sqrt = eigvecs @ torch.diag(1.0 / torch.sqrt(eigvals_clamped)) @ eigvecs.T
                M = torch.matmul(W, cov_sqrt)
                U, S, Vt = torch.linalg.svd(M, full_matrices=False)
                Vt = Vt @ cov_inv_sqrt
        else:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            cov_reg = torch.eye(in_f, device=device)

        sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
        A = U[:, :r] * sqrt_S.unsqueeze(0)
        B = sqrt_S.unsqueeze(1) * Vt[:r, :]

        if D_gating is None or D_gating.shape[0] != out_f:
            D_vec = torch.ones(out_f, device=device, dtype=torch.float32)
        else:
            D_vec = torch.clamp(D_gating.to(device=device, dtype=torch.float32), min=1e-6)

        D_diag = torch.diag(D_vec)
        E = cov_reg
        WE = W @ E
        DW = D_diag @ W

        W_norm = torch.norm(W, p="fro").item()
        best_A = A.clone()
        best_B = B.clone()
        best_recon = A @ B
        best_err = float(torch.norm(W - best_recon, p="fro").item() / max(W_norm, 1e-12))

        # 2. Alternating Least Squares (ALS) optimization loop
        for _ in range(num_iters):
            # Update A: A = (W @ E @ B.T) @ inv(B @ E @ B.T + eps * I)
            BEBt = B @ E @ B.T
            eps_B = max(1e-8, float(1e-5 * BEBt.trace().item() / max(1, r)))
            try:
                A_new = (
                    WE
                    @ B.T
                    @ torch.linalg.inv(BEBt + eps_B * torch.eye(r, device=device, dtype=BEBt.dtype))
                )
            except Exception:
                break

            # Update B: B = inv(A.T @ D @ A + eps * I) @ (A.T @ D @ W)
            AtDA = A_new.T @ D_diag @ A_new
            eps_A = max(1e-8, float(1e-5 * AtDA.trace().item() / max(1, r)))
            try:
                B_new = (
                    torch.linalg.inv(AtDA + eps_A * torch.eye(r, device=device, dtype=AtDA.dtype))
                    @ A_new.T
                    @ DW
                )
            except Exception:
                break

            if (
                torch.isnan(A_new).any()
                or torch.isnan(B_new).any()
                or torch.isinf(A_new).any()
                or torch.isinf(B_new).any()
            ):
                break

            recon = A_new @ B_new
            err = float(torch.norm(W - recon, p="fro").item() / max(W_norm, 1e-12))
            if err < best_err:
                best_A = A_new.clone()
                best_B = B_new.clone()
                best_err = err
            else:
                break
            A, B = A_new, B_new

        return (
            best_A.to(dtype=dtype, device=device),
            best_B.to(dtype=dtype, device=device),
            best_err,
        )

    def compress_linear(
        self,
        layer: nn.Module,
        cov: torch.Tensor | None = None,
        layer_name: str | None = None,
    ) -> tuple[AdaptiveSVDLinear | None, float, int, bool, bool]:
        """
        Decompose a single linear, Conv1D, or LoRALinear layer via energy-thresholded SVD
        with covariance whitening, returning (svd_layer, relative_error, rank, used_cov).
        """
        target_mod = getattr(layer, "base_layer", layer)
        cls_name = layer.__class__.__name__.lower()
        target_cls = target_mod.__class__.__name__.lower()
        is_conv1d = (
            "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(layer, "nf")
            or hasattr(target_mod, "nf")
        )

        weight = getattr(target_mod, "weight", getattr(layer, "weight", None))
        bias = getattr(target_mod, "bias", getattr(layer, "bias", None))
        if weight is None:
            return None, float("inf"), self.min_rank, False, False

        raw_w = weight.data

        if hasattr(target_mod, "nf"):
            out_f = target_mod.nf
            in_f = getattr(
                target_mod,
                "nx",
                raw_w.shape[0] if raw_w.shape[1] == out_f else raw_w.shape[1],
            )
            W = raw_w.t().contiguous().float() if raw_w.shape == (in_f, out_f) else raw_w.float()
        elif hasattr(layer, "nf"):
            out_f = layer.nf
            in_f = getattr(
                layer,
                "nx",
                raw_w.shape[0] if raw_w.shape[1] == out_f else raw_w.shape[1],
            )
            W = raw_w.t().contiguous().float() if raw_w.shape == (in_f, out_f) else raw_w.float()
        elif is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float()

        device = raw_w.device
        dtype = raw_w.dtype
        min_dim = min(out_f, in_f)
        W_norm = torch.norm(W, p="fro").item()

        # Parameter reduction constraint: r * (in + out) < in * out
        max_allowed_rank = max(1, min(min_dim - 1, int(0.75 * (in_f * out_f) / (in_f + out_f))))
        # Rank floor adjustment: account for both min_dim and out_features
        r_floor = max(self.min_rank, 1, int(0.15 * min_dim), int(0.10 * out_f))
        r_floor = min(r_floor, max_allowed_rank)

        # Check for LoRA parameters on module
        lora_r = getattr(layer, "lora_rank", getattr(layer, "r", 0))
        lora_a = getattr(layer, "lora_alpha", 1.0)
        has_lora = (hasattr(layer, "lora_A") and hasattr(layer, "lora_B")) or (lora_r > 0)

        # Attempt 0: Jacobian-Gated SVD (JG-SVD) if enabled
        svd_layer: AdaptiveSVDLinear | None = None
        best_err = 1.0
        best_rank = r_floor

        if self.use_jg_svd:
            D_gating = self.layer_D_gating.get(layer_name, None) if layer_name else None
            try:
                # Rank selection using singular values
                if cov is not None and cov.shape == (in_f, in_f):
                    cov_dev = cov.to(device=device, dtype=torch.float32)
                    diag_mean = torch.diagonal(cov_dev).mean().item()
                    lambda_reg = max(
                        1e-4 * (diag_mean if math.isfinite(diag_mean) and diag_mean > 0 else 1.0),
                        1e-6,
                    )
                    cov_reg = cov_dev + lambda_reg * torch.eye(in_f, device=device)
                    eigvals, eigvecs = torch.linalg.eigh(cov_reg)
                    cov_sqrt = (
                        eigvecs @ torch.diag(torch.sqrt(torch.clamp(eigvals, min=1e-6))) @ eigvecs.T
                    )
                    M = torch.matmul(W, cov_sqrt)
                    _, S_jg, _ = torch.linalg.svd(M, full_matrices=False)
                else:
                    _, S_jg, _ = torch.linalg.svd(W, full_matrices=False)

                r_jg = self._compute_energy_rank(S_jg, max_allowed_rank, min_rank_floor=r_floor)
                A_jg, B_jg, err_jg = self._factorize_with_JG(W, cov, D_gating, r_jg)

                if math.isfinite(err_jg) and err_jg <= self.max_error:
                    svd_layer = AdaptiveSVDLinear(
                        in_features=in_f,
                        out_features=out_f,
                        rank=r_jg,
                        bias=bias is not None,
                        dtype=dtype,
                        device=device,
                        lora_rank=lora_r if has_lora else 0,
                        lora_alpha=lora_a,
                    )
                    svd_layer.set_low_rank_factors(
                        A_jg.to(dtype=dtype, device=device),
                        B_jg.to(dtype=dtype, device=device),
                    )
                    if bias is not None and svd_layer.bias is not None:
                        svd_layer.bias.data.copy_(bias.data)
                    return svd_layer, err_jg, r_jg, cov is not None, True
            except Exception as e:
                logger.debug("JG-SVD factorization exception on %s: %s", layer_name, e)

        # Attempt 1: Covariance whitening SVD

        if cov is not None and cov.shape == (in_f, in_f):
            try:
                cov_dev = cov.to(device=device, dtype=torch.float32)
                diag_mean = torch.diagonal(cov_dev).mean().item()
                if not math.isfinite(diag_mean) or diag_mean <= 0.0:
                    diag_mean = 1.0
                lambda_reg = max(1e-4 * diag_mean, 1e-6)
                cov_reg = cov_dev + lambda_reg * torch.eye(in_f, device=device)

                eigvals, eigvecs = torch.linalg.eigh(cov_reg)
                eigvals_clamped = torch.clamp(eigvals, min=1e-6)

                cov_sqrt = eigvecs @ torch.diag(torch.sqrt(eigvals_clamped)) @ eigvecs.T
                cov_inv_sqrt = eigvecs @ torch.diag(1.0 / torch.sqrt(eigvals_clamped)) @ eigvecs.T

                M = torch.matmul(W, cov_sqrt)
                U, S, Vt = torch.linalg.svd(M, full_matrices=False)

                r = self._compute_energy_rank(S, max_allowed_rank, min_rank_floor=r_floor)
                sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))

                U_r = U[:, :r] * sqrt_S.unsqueeze(0)
                V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]) @ cov_inv_sqrt

                W_recon = torch.matmul(U_r, V_r)
                if not (torch.isnan(W_recon).any() or torch.isinf(W_recon).any()):
                    err = float(torch.norm(W - W_recon, p="fro").item() / max(W_norm, 1e-12))
                    if err <= self.max_error:
                        svd_layer = AdaptiveSVDLinear(
                            in_features=in_f,
                            out_features=out_f,
                            rank=r,
                            bias=bias is not None,
                            dtype=dtype,
                            device=device,
                            lora_rank=lora_r if has_lora else 0,
                            lora_alpha=lora_a,
                        )
                        svd_layer.set_low_rank_factors(
                            U_r.to(dtype=dtype, device=device),
                            V_r.to(dtype=dtype, device=device),
                        )
                        if bias is not None and svd_layer.bias is not None:
                            svd_layer.bias.data.copy_(bias.data)
                        return svd_layer, err, r, True, False
            except Exception as e:
                logger.debug("Covariance whitening exception on linear layer: %s", e)

        # Attempt 2: Unweighted Direct SVD with rank escalation
        try:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            candidate_thresholds = [
                self.energy_threshold,
                min(0.98, self.energy_threshold + 0.04),
            ]
            for thresh in candidate_thresholds:
                r = self._compute_energy_rank(
                    S, max_allowed_rank, min_rank_floor=r_floor, threshold=thresh
                )
                sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
                U_r = U[:, :r] * sqrt_S.unsqueeze(0)
                V_r = sqrt_S.unsqueeze(1) * Vt[:r, :]
                W_recon = torch.matmul(U_r, V_r)
                err = float(torch.norm(W - W_recon, p="fro").item() / max(W_norm, 1e-12))
                best_err = err
                best_rank = r
                if err <= self.max_error:
                    break

            svd_layer = AdaptiveSVDLinear(
                in_features=in_f,
                out_features=out_f,
                rank=best_rank,
                bias=bias is not None,
                dtype=dtype,
                device=device,
                lora_rank=lora_r if has_lora else 0,
                lora_alpha=lora_a,
            )
            sqrt_S = torch.sqrt(torch.clamp(S[:best_rank], min=1e-12))
            U_r = (U[:, :best_rank] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
            V_r = (sqrt_S.unsqueeze(1) * Vt[:best_rank, :]).to(dtype=dtype, device=device)
            svd_layer.set_low_rank_factors(U_r, V_r)
            if bias is not None and svd_layer.bias is not None:
                svd_layer.bias.data.copy_(bias.data)
            return svd_layer, best_err, best_rank, False, False
        except Exception as e:
            logger.warning("AdaptiveRanks standard SVD failed: %s", e)
            return None, 1.0, r_floor, False, False

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Compress target layers with activation covariance whitening and safe validation."""
        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}
        self.used_covariance = {}

        target_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            logger.warning("AdaptiveRanksCompressor: No compressible layers found.")
            return self.model, self.get_compression_stats()

        covariances = self._collect_activation_covariances()
        if self.use_jg_svd:
            self._collect_gating_D()

        for name, module in target_layers:
            target_mod = getattr(module, "base_layer", module)
            if isinstance(target_mod, nn.MultiheadAttention):
                try:
                    embed_dim = target_mod.embed_dim
                    r_floor = max(self.min_rank, 1, int(0.15 * embed_dim), int(0.10 * embed_dim))
                    r_qkv = r_floor
                    if (
                        hasattr(target_mod, "in_proj_weight")
                        and target_mod.in_proj_weight is not None
                    ):
                        W_in = target_mod.in_proj_weight.data.float()
                        # in_proj_weight is [3 * embed_dim, embed_dim]. Split into Q, K, V
                        w_q, w_k, w_v = W_in.chunk(3, dim=0)
                        _, s_q, _ = torch.linalg.svd(w_q, full_matrices=False)
                        _, s_k, _ = torch.linalg.svd(w_k, full_matrices=False)
                        _, s_v, _ = torch.linalg.svd(w_v, full_matrices=False)
                        rq = self._compute_energy_rank(s_q, embed_dim, min_rank_floor=r_floor)
                        rk = self._compute_energy_rank(s_k, embed_dim, min_rank_floor=r_floor)
                        rv = self._compute_energy_rank(s_v, embed_dim, min_rank_floor=r_floor)
                        r_qkv = max(rq, rk, rv)

                    r_out = r_floor
                    if (
                        hasattr(target_mod, "out_proj")
                        and hasattr(target_mod.out_proj, "weight")
                        and target_mod.out_proj.weight is not None
                    ):
                        W_out = target_mod.out_proj.weight.data.float()
                        _, s_out, _ = torch.linalg.svd(W_out, full_matrices=False)
                        r_out = self._compute_energy_rank(s_out, embed_dim, min_rank_floor=r_floor)

                    if CompressedAttentionLayer is None:
                        raise ImportError("CompressedAttentionLayer is not available")

                    compressed = CompressedAttentionLayer(
                        original_layer=target_mod,
                        rank_qkv=r_qkv,
                        rank_out=r_out,
                    )
                    replace_module(self.model, name, compressed)
                    self.compressed_layers.append(name)
                    self.layer_ranks[name] = max(r_qkv, r_out)
                    self.layer_errors[name] = 0.05
                    self.used_covariance[name] = False
                    self.used_jg_svd[name] = False
                except Exception as e:
                    logger.warning("Failed to compress attention layer %s: %s", name, e)
                    self.fallback_layers.append(name)
                    self.layer_errors[name] = 1.0
            else:
                cov = covariances.get(name, None)
                svd_layer, err, r, used_cov, used_jg = self.compress_linear(
                    module, cov=cov, layer_name=name
                )
                self.layer_errors[name] = err
                self.layer_ranks[name] = r
                self.used_covariance[name] = used_cov
                self.used_jg_svd[name] = used_jg

                if svd_layer is None or err > self.max_error:
                    if self.fallback_to_original:
                        logger.warning(
                            "AdaptiveRanks error %.4f on %s exceeded threshold %.4f. Falling back to original layer.",
                            err,
                            name,
                            self.max_error,
                        )
                        self.fallback_layers.append(name)
                        continue
                    else:
                        logger.warning(
                            "AdaptiveRanks error %.4f on %s exceeds threshold %.4f but fallback_to_original=False. Compressing anyway.",
                            err,
                            name,
                            self.max_error,
                        )

                if hasattr(module, "base_layer"):
                    # Check if svd_layer has its own LoRA parameters populated
                    if (
                        hasattr(module, "lora_A")
                        and module.lora_A is not None
                        and hasattr(svd_layer, "lora_A")
                        and svd_layer.lora_A is not None
                    ):
                        try:
                            # Copy wrapper LoRA weights into factorized layer's LoRA parameters
                            lora_a_data = (
                                module.lora_A.data
                                if hasattr(module.lora_A, "data")
                                else module.lora_A
                            )
                            lora_b_data = (
                                module.lora_B.data
                                if hasattr(module.lora_B, "data")
                                else module.lora_B
                            )
                            if svd_layer.lora_A.shape == lora_a_data.shape:
                                svd_layer.lora_A.data.copy_(
                                    lora_a_data.to(
                                        device=svd_layer.lora_A.device,
                                        dtype=svd_layer.lora_A.dtype,
                                    )
                                )
                            if svd_layer.lora_B.shape == lora_b_data.shape:
                                svd_layer.lora_B.data.copy_(
                                    lora_b_data.to(
                                        device=svd_layer.lora_B.device,
                                        dtype=svd_layer.lora_B.dtype,
                                    )
                                )
                        except Exception as e:
                            logger.debug("Notice copying wrapper LoRA parameters: %s", e)

                    module.base_layer = svd_layer
                    if hasattr(module, "lora_A") and module.lora_A is not None:
                        for p in (
                            module.lora_A.parameters()
                            if hasattr(module.lora_A, "parameters")
                            else [module.lora_A]
                        ):
                            p.requires_grad = True
                    if hasattr(module, "lora_B") and module.lora_B is not None:
                        for p in (
                            module.lora_B.parameters()
                            if hasattr(module.lora_B, "parameters")
                            else [module.lora_B]
                        ):
                            p.requires_grad = True
                else:
                    replace_module(self.model, name, svd_layer)
                self.compressed_layers.append(name)

        # Ensure fallback layers are frozen so recovery fine-tuning does not touch them
        for name, mod in self.model.named_modules():
            if name in self.fallback_layers:
                for p in mod.parameters():
                    p.requires_grad = False
                target = getattr(mod, "base_layer", None)
                if target is not None:
                    for p in target.parameters():
                        p.requires_grad = False

        self._compressed_params = self._count_params(self.model)
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        stats = self.get_compression_stats()
        stats.update(
            {
                "energy_threshold": self.energy_threshold,
                "min_rank": self.min_rank,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "used_covariance": self.used_covariance,
                "used_jg_svd": self.used_jg_svd,
                "use_jg_svd": self.use_jg_svd,
                "use_cholesky_jg": self.use_cholesky_jg,
                "avg_reconstruction_error": avg_err,
            }
        )
        logger.info(
            "AdaptiveRanks compression complete: %d compressed, %d fallback. Avg error: %.4f.",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: Adaptive Ranks SVD Compressor.

Dynamically determines ranks per layer according to cumulative singular value energy threshold: $$\frac{\sum_{i=1}^r \sigma_i^2}{\sum_{j=1}^d \sigma_j^2} \ge \tau \quad (\tau = 0.92)$$ with minimum rank floor $r_{\min} = \max(16, \lfloor 0.15 \cdot \min(d_{in}, d_{out}) \rfloor, \lfloor 0.10 \cdot d_{out} \rfloor)$, activation covariance whitening, and symmetrically distributed singular values: $$U_r = U_{:, :r} \sqrt{\Sigma_r}, \quad V_r = (\sqrt{\Sigma_r} V^\top_{:r, :}) \cdot \operatorname{Cov}(X)^{-1/2}$$

Production hardening features: - Relative Frobenius reconstruction error validation with configurable max_error (default 0.20). - Safe fallback to original weights by default (fallback_to_original=True). - Robust Conv1D and LoRALinear support via base_layer weight transposition. - SVD and covariance eigendecomposition computed strictly in FP32 with NaN/Inf guards. - Unified critical-layer protection via should_compress_linear. - Adaptive rank retry when reconstruction error exceeds threshold. - Parameter reduction verification: ensures r * (in + out) < in * out. - Comprehensive per-layer statistics and logging.

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Compress target layers with activation covariance whitening and safe validation."""
    self.compressed_layers = []
    self.fallback_layers = []
    self.layer_errors = {}
    self.layer_ranks = {}
    self.used_covariance = {}

    target_layers = [
        (name, mod)
        for name, mod in self.model.named_modules()
        if self._is_target_module(name, mod)
    ]

    if not target_layers:
        logger.warning("AdaptiveRanksCompressor: No compressible layers found.")
        return self.model, self.get_compression_stats()

    covariances = self._collect_activation_covariances()
    if self.use_jg_svd:
        self._collect_gating_D()

    for name, module in target_layers:
        target_mod = getattr(module, "base_layer", module)
        if isinstance(target_mod, nn.MultiheadAttention):
            try:
                embed_dim = target_mod.embed_dim
                r_floor = max(self.min_rank, 1, int(0.15 * embed_dim), int(0.10 * embed_dim))
                r_qkv = r_floor
                if (
                    hasattr(target_mod, "in_proj_weight")
                    and target_mod.in_proj_weight is not None
                ):
                    W_in = target_mod.in_proj_weight.data.float()
                    # in_proj_weight is [3 * embed_dim, embed_dim]. Split into Q, K, V
                    w_q, w_k, w_v = W_in.chunk(3, dim=0)
                    _, s_q, _ = torch.linalg.svd(w_q, full_matrices=False)
                    _, s_k, _ = torch.linalg.svd(w_k, full_matrices=False)
                    _, s_v, _ = torch.linalg.svd(w_v, full_matrices=False)
                    rq = self._compute_energy_rank(s_q, embed_dim, min_rank_floor=r_floor)
                    rk = self._compute_energy_rank(s_k, embed_dim, min_rank_floor=r_floor)
                    rv = self._compute_energy_rank(s_v, embed_dim, min_rank_floor=r_floor)
                    r_qkv = max(rq, rk, rv)

                r_out = r_floor
                if (
                    hasattr(target_mod, "out_proj")
                    and hasattr(target_mod.out_proj, "weight")
                    and target_mod.out_proj.weight is not None
                ):
                    W_out = target_mod.out_proj.weight.data.float()
                    _, s_out, _ = torch.linalg.svd(W_out, full_matrices=False)
                    r_out = self._compute_energy_rank(s_out, embed_dim, min_rank_floor=r_floor)

                if CompressedAttentionLayer is None:
                    raise ImportError("CompressedAttentionLayer is not available")

                compressed = CompressedAttentionLayer(
                    original_layer=target_mod,
                    rank_qkv=r_qkv,
                    rank_out=r_out,
                )
                replace_module(self.model, name, compressed)
                self.compressed_layers.append(name)
                self.layer_ranks[name] = max(r_qkv, r_out)
                self.layer_errors[name] = 0.05
                self.used_covariance[name] = False
                self.used_jg_svd[name] = False
            except Exception as e:
                logger.warning("Failed to compress attention layer %s: %s", name, e)
                self.fallback_layers.append(name)
                self.layer_errors[name] = 1.0
        else:
            cov = covariances.get(name, None)
            svd_layer, err, r, used_cov, used_jg = self.compress_linear(
                module, cov=cov, layer_name=name
            )
            self.layer_errors[name] = err
            self.layer_ranks[name] = r
            self.used_covariance[name] = used_cov
            self.used_jg_svd[name] = used_jg

            if svd_layer is None or err > self.max_error:
                if self.fallback_to_original:
                    logger.warning(
                        "AdaptiveRanks error %.4f on %s exceeded threshold %.4f. Falling back to original layer.",
                        err,
                        name,
                        self.max_error,
                    )
                    self.fallback_layers.append(name)
                    continue
                else:
                    logger.warning(
                        "AdaptiveRanks error %.4f on %s exceeds threshold %.4f but fallback_to_original=False. Compressing anyway.",
                        err,
                        name,
                        self.max_error,
                    )

            if hasattr(module, "base_layer"):
                # Check if svd_layer has its own LoRA parameters populated
                if (
                    hasattr(module, "lora_A")
                    and module.lora_A is not None
                    and hasattr(svd_layer, "lora_A")
                    and svd_layer.lora_A is not None
                ):
                    try:
                        # Copy wrapper LoRA weights into factorized layer's LoRA parameters
                        lora_a_data = (
                            module.lora_A.data
                            if hasattr(module.lora_A, "data")
                            else module.lora_A
                        )
                        lora_b_data = (
                            module.lora_B.data
                            if hasattr(module.lora_B, "data")
                            else module.lora_B
                        )
                        if svd_layer.lora_A.shape == lora_a_data.shape:
                            svd_layer.lora_A.data.copy_(
                                lora_a_data.to(
                                    device=svd_layer.lora_A.device,
                                    dtype=svd_layer.lora_A.dtype,
                                )
                            )
                        if svd_layer.lora_B.shape == lora_b_data.shape:
                            svd_layer.lora_B.data.copy_(
                                lora_b_data.to(
                                    device=svd_layer.lora_B.device,
                                    dtype=svd_layer.lora_B.dtype,
                                )
                            )
                    except Exception as e:
                        logger.debug("Notice copying wrapper LoRA parameters: %s", e)

                module.base_layer = svd_layer
                if hasattr(module, "lora_A") and module.lora_A is not None:
                    for p in (
                        module.lora_A.parameters()
                        if hasattr(module.lora_A, "parameters")
                        else [module.lora_A]
                    ):
                        p.requires_grad = True
                if hasattr(module, "lora_B") and module.lora_B is not None:
                    for p in (
                        module.lora_B.parameters()
                        if hasattr(module.lora_B, "parameters")
                        else [module.lora_B]
                    ):
                        p.requires_grad = True
            else:
                replace_module(self.model, name, svd_layer)
            self.compressed_layers.append(name)

    # Ensure fallback layers are frozen so recovery fine-tuning does not touch them
    for name, mod in self.model.named_modules():
        if name in self.fallback_layers:
            for p in mod.parameters():
                p.requires_grad = False
            target = getattr(mod, "base_layer", None)
            if target is not None:
                for p in target.parameters():
                    p.requires_grad = False

    self._compressed_params = self._count_params(self.model)
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )
    stats = self.get_compression_stats()
    stats.update(
        {
            "energy_threshold": self.energy_threshold,
            "min_rank": self.min_rank,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "compressed_layers": self.compressed_layers,
            "fallback_layers": self.fallback_layers,
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "layer_errors": self.layer_errors,
            "layer_ranks": self.layer_ranks,
            "used_covariance": self.used_covariance,
            "used_jg_svd": self.used_jg_svd,
            "use_jg_svd": self.use_jg_svd,
            "use_cholesky_jg": self.use_cholesky_jg,
            "avg_reconstruction_error": avg_err,
        }
    )
    logger.info(
        "AdaptiveRanks compression complete: %d compressed, %d fallback. Avg error: %.4f.",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
    )
    return self.model, stats

Compress target layers with activation covariance whitening and safe validation.

def compress_linear(self, layer: nn.Module, cov: torch.Tensor | None = None, layer_name: str | None = None) ‑> tuple[SVDLinear | None, float, int, bool, bool]

Expand source code

def compress_linear(
    self,
    layer: nn.Module,
    cov: torch.Tensor | None = None,
    layer_name: str | None = None,
) -> tuple[AdaptiveSVDLinear | None, float, int, bool, bool]:
    """
    Decompose a single linear, Conv1D, or LoRALinear layer via energy-thresholded SVD
    with covariance whitening, returning (svd_layer, relative_error, rank, used_cov).
    """
    target_mod = getattr(layer, "base_layer", layer)
    cls_name = layer.__class__.__name__.lower()
    target_cls = target_mod.__class__.__name__.lower()
    is_conv1d = (
        "conv1d" in cls_name
        or "conv1d" in target_cls
        or hasattr(layer, "nf")
        or hasattr(target_mod, "nf")
    )

    weight = getattr(target_mod, "weight", getattr(layer, "weight", None))
    bias = getattr(target_mod, "bias", getattr(layer, "bias", None))
    if weight is None:
        return None, float("inf"), self.min_rank, False, False

    raw_w = weight.data

    if hasattr(target_mod, "nf"):
        out_f = target_mod.nf
        in_f = getattr(
            target_mod,
            "nx",
            raw_w.shape[0] if raw_w.shape[1] == out_f else raw_w.shape[1],
        )
        W = raw_w.t().contiguous().float() if raw_w.shape == (in_f, out_f) else raw_w.float()
    elif hasattr(layer, "nf"):
        out_f = layer.nf
        in_f = getattr(
            layer,
            "nx",
            raw_w.shape[0] if raw_w.shape[1] == out_f else raw_w.shape[1],
        )
        W = raw_w.t().contiguous().float() if raw_w.shape == (in_f, out_f) else raw_w.float()
    elif is_conv1d:
        in_f, out_f = raw_w.shape[0], raw_w.shape[1]
        W = raw_w.t().contiguous().float()
    else:
        out_f, in_f = raw_w.shape[0], raw_w.shape[1]
        W = raw_w.float()

    device = raw_w.device
    dtype = raw_w.dtype
    min_dim = min(out_f, in_f)
    W_norm = torch.norm(W, p="fro").item()

    # Parameter reduction constraint: r * (in + out) < in * out
    max_allowed_rank = max(1, min(min_dim - 1, int(0.75 * (in_f * out_f) / (in_f + out_f))))
    # Rank floor adjustment: account for both min_dim and out_features
    r_floor = max(self.min_rank, 1, int(0.15 * min_dim), int(0.10 * out_f))
    r_floor = min(r_floor, max_allowed_rank)

    # Check for LoRA parameters on module
    lora_r = getattr(layer, "lora_rank", getattr(layer, "r", 0))
    lora_a = getattr(layer, "lora_alpha", 1.0)
    has_lora = (hasattr(layer, "lora_A") and hasattr(layer, "lora_B")) or (lora_r > 0)

    # Attempt 0: Jacobian-Gated SVD (JG-SVD) if enabled
    svd_layer: AdaptiveSVDLinear | None = None
    best_err = 1.0
    best_rank = r_floor

    if self.use_jg_svd:
        D_gating = self.layer_D_gating.get(layer_name, None) if layer_name else None
        try:
            # Rank selection using singular values
            if cov is not None and cov.shape == (in_f, in_f):
                cov_dev = cov.to(device=device, dtype=torch.float32)
                diag_mean = torch.diagonal(cov_dev).mean().item()
                lambda_reg = max(
                    1e-4 * (diag_mean if math.isfinite(diag_mean) and diag_mean > 0 else 1.0),
                    1e-6,
                )
                cov_reg = cov_dev + lambda_reg * torch.eye(in_f, device=device)
                eigvals, eigvecs = torch.linalg.eigh(cov_reg)
                cov_sqrt = (
                    eigvecs @ torch.diag(torch.sqrt(torch.clamp(eigvals, min=1e-6))) @ eigvecs.T
                )
                M = torch.matmul(W, cov_sqrt)
                _, S_jg, _ = torch.linalg.svd(M, full_matrices=False)
            else:
                _, S_jg, _ = torch.linalg.svd(W, full_matrices=False)

            r_jg = self._compute_energy_rank(S_jg, max_allowed_rank, min_rank_floor=r_floor)
            A_jg, B_jg, err_jg = self._factorize_with_JG(W, cov, D_gating, r_jg)

            if math.isfinite(err_jg) and err_jg <= self.max_error:
                svd_layer = AdaptiveSVDLinear(
                    in_features=in_f,
                    out_features=out_f,
                    rank=r_jg,
                    bias=bias is not None,
                    dtype=dtype,
                    device=device,
                    lora_rank=lora_r if has_lora else 0,
                    lora_alpha=lora_a,
                )
                svd_layer.set_low_rank_factors(
                    A_jg.to(dtype=dtype, device=device),
                    B_jg.to(dtype=dtype, device=device),
                )
                if bias is not None and svd_layer.bias is not None:
                    svd_layer.bias.data.copy_(bias.data)
                return svd_layer, err_jg, r_jg, cov is not None, True
        except Exception as e:
            logger.debug("JG-SVD factorization exception on %s: %s", layer_name, e)

    # Attempt 1: Covariance whitening SVD

    if cov is not None and cov.shape == (in_f, in_f):
        try:
            cov_dev = cov.to(device=device, dtype=torch.float32)
            diag_mean = torch.diagonal(cov_dev).mean().item()
            if not math.isfinite(diag_mean) or diag_mean <= 0.0:
                diag_mean = 1.0
            lambda_reg = max(1e-4 * diag_mean, 1e-6)
            cov_reg = cov_dev + lambda_reg * torch.eye(in_f, device=device)

            eigvals, eigvecs = torch.linalg.eigh(cov_reg)
            eigvals_clamped = torch.clamp(eigvals, min=1e-6)

            cov_sqrt = eigvecs @ torch.diag(torch.sqrt(eigvals_clamped)) @ eigvecs.T
            cov_inv_sqrt = eigvecs @ torch.diag(1.0 / torch.sqrt(eigvals_clamped)) @ eigvecs.T

            M = torch.matmul(W, cov_sqrt)
            U, S, Vt = torch.linalg.svd(M, full_matrices=False)

            r = self._compute_energy_rank(S, max_allowed_rank, min_rank_floor=r_floor)
            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))

            U_r = U[:, :r] * sqrt_S.unsqueeze(0)
            V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]) @ cov_inv_sqrt

            W_recon = torch.matmul(U_r, V_r)
            if not (torch.isnan(W_recon).any() or torch.isinf(W_recon).any()):
                err = float(torch.norm(W - W_recon, p="fro").item() / max(W_norm, 1e-12))
                if err <= self.max_error:
                    svd_layer = AdaptiveSVDLinear(
                        in_features=in_f,
                        out_features=out_f,
                        rank=r,
                        bias=bias is not None,
                        dtype=dtype,
                        device=device,
                        lora_rank=lora_r if has_lora else 0,
                        lora_alpha=lora_a,
                    )
                    svd_layer.set_low_rank_factors(
                        U_r.to(dtype=dtype, device=device),
                        V_r.to(dtype=dtype, device=device),
                    )
                    if bias is not None and svd_layer.bias is not None:
                        svd_layer.bias.data.copy_(bias.data)
                    return svd_layer, err, r, True, False
        except Exception as e:
            logger.debug("Covariance whitening exception on linear layer: %s", e)

    # Attempt 2: Unweighted Direct SVD with rank escalation
    try:
        U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        candidate_thresholds = [
            self.energy_threshold,
            min(0.98, self.energy_threshold + 0.04),
        ]
        for thresh in candidate_thresholds:
            r = self._compute_energy_rank(
                S, max_allowed_rank, min_rank_floor=r_floor, threshold=thresh
            )
            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
            U_r = U[:, :r] * sqrt_S.unsqueeze(0)
            V_r = sqrt_S.unsqueeze(1) * Vt[:r, :]
            W_recon = torch.matmul(U_r, V_r)
            err = float(torch.norm(W - W_recon, p="fro").item() / max(W_norm, 1e-12))
            best_err = err
            best_rank = r
            if err <= self.max_error:
                break

        svd_layer = AdaptiveSVDLinear(
            in_features=in_f,
            out_features=out_f,
            rank=best_rank,
            bias=bias is not None,
            dtype=dtype,
            device=device,
            lora_rank=lora_r if has_lora else 0,
            lora_alpha=lora_a,
        )
        sqrt_S = torch.sqrt(torch.clamp(S[:best_rank], min=1e-12))
        U_r = (U[:, :best_rank] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
        V_r = (sqrt_S.unsqueeze(1) * Vt[:best_rank, :]).to(dtype=dtype, device=device)
        svd_layer.set_low_rank_factors(U_r, V_r)
        if bias is not None and svd_layer.bias is not None:
            svd_layer.bias.data.copy_(bias.data)
        return svd_layer, best_err, best_rank, False, False
    except Exception as e:
        logger.warning("AdaptiveRanks standard SVD failed: %s", e)
        return None, 1.0, r_floor, False, False

Decompose a single linear, Conv1D, or LoRALinear layer via energy-thresholded SVD with covariance whitening, returning (svd_layer, relative_error, rank, used_cov).

class AdaptiveTransformerCompressor (model: nn.Module, energy_threshold: float = 0.92, min_rank: int = 16, max_error: float = 0.2, fallback_to_original: bool = True, protect_critical: bool = True, target_attention_only: bool = False, target_modules: list[str] | None = None, tokenizer: Any | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, use_jg_svd: bool = False, use_cholesky_jg: bool = True, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("adaptive_ranks")
@COMPRESSOR_REGISTRY.register("adaptive_svd")
class AdaptiveRanksCompressor(BaseCompressor):
    is_custom_attention: bool = True
    r"""
    Adaptive Ranks SVD Compressor.

    Dynamically determines ranks per layer according to cumulative singular value energy threshold:
        $$\frac{\sum_{i=1}^r \sigma_i^2}{\sum_{j=1}^d \sigma_j^2} \ge \tau \quad (\tau = 0.92)$$
    with minimum rank floor $r_{\min} = \max(16, \lfloor 0.15 \cdot \min(d_{in}, d_{out}) \rfloor, \lfloor 0.10 \cdot d_{out} \rfloor)$,
    activation covariance whitening, and symmetrically distributed singular values:
        $$U_r = U_{:, :r} \sqrt{\Sigma_r}, \quad V_r = (\sqrt{\Sigma_r} V^\top_{:r, :}) \cdot \operatorname{Cov}(X)^{-1/2}$$

    Production hardening features:
    - Relative Frobenius reconstruction error validation with configurable max_error (default 0.20).
    - Safe fallback to original weights by default (fallback_to_original=True).
    - Robust Conv1D and LoRALinear support via base_layer weight transposition.
    - SVD and covariance eigendecomposition computed strictly in FP32 with NaN/Inf guards.
    - Unified critical-layer protection via should_compress_linear.
    - Adaptive rank retry when reconstruction error exceeds threshold.
    - Parameter reduction verification: ensures r * (in + out) < in * out.
    - Comprehensive per-layer statistics and logging.
    """

    def __init__(
        self,
        model: nn.Module,
        energy_threshold: float = 0.92,
        min_rank: int = 16,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        protect_critical: bool = True,
        target_attention_only: bool = False,
        target_modules: list[str] | None = None,
        tokenizer: Any | None = None,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        use_jg_svd: bool = False,
        use_cholesky_jg: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="adaptive_ranks", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is None and hasattr(model, "config") and model.config is not None:
            cfg = model.config

        # Check sys.argv for CLI overrides
        cli_use_jg_svd = None
        cli_use_cholesky_jg = None
        import sys

        for _idx, _arg in enumerate(sys.argv):
            if _arg in (
                "--use_jg_svd",
                "--adaptive_ranks_use_jg_svd",
            ) and _idx + 1 < len(sys.argv):
                cli_use_jg_svd = sys.argv[_idx + 1].lower() in ("true", "1", "yes")
            elif _arg.startswith("--use_jg_svd=") or _arg.startswith(
                "--adaptive_ranks_use_jg_svd="
            ):
                cli_use_jg_svd = _arg.split("=")[1].lower() in ("true", "1", "yes")
            elif _arg in (
                "--use_cholesky_jg",
                "--adaptive_ranks_use_cholesky_jg",
            ) and _idx + 1 < len(sys.argv):
                cli_use_cholesky_jg = sys.argv[_idx + 1].lower() in ("true", "1", "yes")
            elif _arg.startswith("--use_cholesky_jg=") or _arg.startswith(
                "--adaptive_ranks_use_cholesky_jg="
            ):
                cli_use_cholesky_jg = _arg.split("=")[1].lower() in ("true", "1", "yes")

        if cfg is not None:
            energy_threshold = _cfg_get(
                cfg,
                "energy_threshold",
                _cfg_get(
                    cfg,
                    "adaptive_ranks_energy_threshold",
                    _cfg_get(cfg, "asvd_energy_threshold", energy_threshold),
                ),
            )
            min_rank = _cfg_get(cfg, "min_rank", _cfg_get(cfg, "rank", min_rank))
            max_error = _cfg_get(
                cfg,
                "max_error",
                _cfg_get(
                    cfg,
                    "adaptive_ranks_max_error",
                    _cfg_get(cfg, "aasvd_max_error", max_error),
                ),
            )
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            protect_critical = _cfg_get(cfg, "protect_critical", protect_critical)
            target_attention_only = _cfg_get(
                cfg,
                "target_attention_only",
                _cfg_get(
                    cfg,
                    "adaptive_ranks_target_attention_only",
                    _cfg_get(cfg, "aasvd_target_attention_only", target_attention_only),
                ),
            )
            use_jg_svd = _cfg_get(
                cfg,
                "use_jg_svd",
                _cfg_get(cfg, "adaptive_ranks_use_jg_svd", use_jg_svd),
            )
            use_cholesky_jg = _cfg_get(
                cfg,
                "use_cholesky_jg",
                _cfg_get(cfg, "adaptive_ranks_use_cholesky_jg", use_cholesky_jg),
            )

        if cli_use_jg_svd is not None:
            use_jg_svd = cli_use_jg_svd
        if cli_use_cholesky_jg is not None:
            use_cholesky_jg = cli_use_cholesky_jg

        if "use_jg_svd" in kwargs and kwargs["use_jg_svd"] is not None:
            use_jg_svd = kwargs["use_jg_svd"]
        if "use_cholesky_jg" in kwargs and kwargs["use_cholesky_jg"] is not None:
            use_cholesky_jg = kwargs["use_cholesky_jg"]

        if "rank" in kwargs and kwargs["rank"] is not None:
            min_rank = kwargs["rank"]
        if "target_attention_only" in kwargs and kwargs["target_attention_only"] is not None:
            target_attention_only = kwargs["target_attention_only"]

        self.energy_threshold = max(0.01, min(float(energy_threshold), 1.0))
        self.min_rank = max(1, int(min_rank))
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.protect_critical = bool(protect_critical)
        self.target_attention_only = bool(target_attention_only)
        self.use_jg_svd = bool(use_jg_svd)
        self.use_cholesky_jg = bool(use_cholesky_jg)
        self.target_modules = target_modules
        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.target_device = (
            torch.device(device)
            if device is not None
            else next(model.parameters(), torch.tensor(0.0)).device
        )

        self.layer_D_gating: dict[str, torch.Tensor] = {}
        self.attn_jacobian: dict[str, float] = {}

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.used_covariance: dict[str, bool] = {}
        self.used_jg_svd: dict[str, bool] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        if isinstance(module, SVDLinear) or (
            CompressedAttentionLayer is not None and isinstance(module, CompressedAttentionLayer)
        ):
            return False
        # Do not target inner submodules of LoRA or wrappers
        if any(
            name.endswith(sfx) for sfx in [".base_layer", ".lora_A", ".lora_B", ".lora_dropout"]
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = module.__class__.__name__.lower()
        target_cls = target_mod.__class__.__name__.lower()

        is_linear_like = (
            isinstance(module, nn.Linear)
            or isinstance(target_mod, nn.Linear)
            or "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(module, "nf")
            or hasattr(target_mod, "nf")
            or "loralinear" in cls_name
            or hasattr(module, "base_layer")
        )
        is_mha = isinstance(module, nn.MultiheadAttention) or isinstance(
            target_mod, nn.MultiheadAttention
        )
        if not (is_linear_like or is_mha):
            return False

        if self.protect_critical:
            name_lower = name.lower()
            if any(
                skip in name_lower
                for skip in [
                    "lm_head",
                    "classifier",
                    "head",
                    "output_layer",
                    "score",
                    "pooler",
                    "embed",
                ]
            ):
                return False

        if self.target_attention_only:
            name_lower = name.lower()
            ffn_patterns = [
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            ]
            if any(p in name_lower for p in ffn_patterns):
                return False
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            attn_patterns = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(p in name_lower for p in attn_patterns):
                return False
        else:
            if not should_compress_linear(self.model, name):
                return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def _compute_energy_rank(
        self,
        S: torch.Tensor,
        max_rank: int,
        min_rank_floor: int | None = None,
        threshold: float | None = None,
    ) -> int:
        r"""
        Find smallest rank $r$ satisfying cumulative singular energy ratio $\ge$ threshold:
        $$\frac{\sum_{i=1}^r \sigma_i^2}{\sum_{j=1}^d \sigma_j^2} \ge \tau$$
        constrained by minimum rank floor and parameter reduction bounds.
        """
        thresh = threshold if threshold is not None else self.energy_threshold
        r_floor = min_rank_floor if min_rank_floor is not None else self.min_rank
        r_floor = max(1, min(int(r_floor), max_rank))
        s_squared = S**2
        total_energy = s_squared.sum()
        if total_energy <= 1e-12:
            return r_floor

        cum_energy = torch.cumsum(s_squared, dim=0) / total_energy
        idx = torch.where(cum_energy >= thresh)[0]
        r = int(idx[0].item()) + 1 if len(idx) > 0 else max_rank
        return max(1, min(max(r, r_floor), max_rank))

    def _collect_activation_covariances(self) -> dict[str, torch.Tensor]:
        """Collect empirical input activation covariances across target nn.Linear / Conv1D modules."""
        if not (self.tokenizer or self.calibration_loader or self.calibration_texts):
            logger.info("No calibration data available; using unweighted SVD.")
            return {}

        self.model.eval()
        running_cov: dict[str, torch.Tensor] = {}
        sample_counts: dict[str, int] = {}
        hooks = []

        def make_hook(layer_name: str) -> Any:
            def hook(module: nn.Module, inp: tuple[torch.Tensor, ...], outp: Any) -> None:
                if len(inp) > 0 and isinstance(inp[0], torch.Tensor):
                    x = inp[0].detach().float()
                    target_mod = getattr(module, "base_layer", module)
                    cls_name = module.__class__.__name__.lower()
                    target_cls = target_mod.__class__.__name__.lower()
                    is_conv1d = (
                        "conv1d" in cls_name
                        or "conv1d" in target_cls
                        or hasattr(module, "nf")
                        or hasattr(target_mod, "nf")
                    )
                    weight = getattr(target_mod, "weight", getattr(module, "weight", None))
                    if weight is None:
                        return
                    in_f = weight.shape[0] if is_conv1d else weight.shape[1]
                    if x.ndim >= 2:
                        flat_x = x.reshape(-1, in_f)
                    else:
                        flat_x = x.unsqueeze(0)
                    n_samples = flat_x.shape[0]
                    if n_samples > 0:
                        batch_cov = torch.matmul(flat_x.T, flat_x).cpu()
                        if layer_name not in running_cov:
                            running_cov[layer_name] = batch_cov
                            sample_counts[layer_name] = n_samples
                        else:
                            running_cov[layer_name] += batch_cov
                            sample_counts[layer_name] += n_samples

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module) and not isinstance(
                getattr(module, "base_layer", module), nn.MultiheadAttention
            ):
                hooks.append(module.register_forward_hook(make_hook(name)))

        if self.calibration_loader is not None:
            self.model.eval()
            with torch.no_grad():
                for batch in self.calibration_loader:
                    try:
                        if isinstance(batch, dict):
                            self.model(
                                **{
                                    k: (v.to(self.target_device) if hasattr(v, "to") else v)
                                    for k, v in batch.items()
                                }
                            )
                        elif isinstance(batch, torch.Tensor):
                            self.model(batch.to(self.target_device))
                        elif isinstance(batch, (list, tuple)):
                            try:
                                self.model(
                                    *[
                                        (b.to(self.target_device) if hasattr(b, "to") else b)
                                        for b in batch
                                    ]
                                )
                            except Exception:
                                self.model(batch[0].to(self.target_device))
                    except Exception as e:
                        logger.debug(
                            "AdaptiveRanks calibration loader forward pass exception: %s",
                            e,
                        )
        else:
            calib_inputs = get_calibration_tensors(
                self.model,
                tokenizer=self.tokenizer,
                calibration_loader=None,
                calibration_texts=self.calibration_texts,
                device=self.target_device,
            )

            if calib_inputs is not None and not (
                isinstance(calib_inputs, (list, tuple, dict)) and len(calib_inputs) == 0
            ):
                with torch.no_grad():
                    try:
                        if isinstance(calib_inputs, dict):
                            self.model(**calib_inputs)
                        elif isinstance(calib_inputs, torch.Tensor):
                            self.model(calib_inputs)
                        elif isinstance(calib_inputs, (list, tuple)):
                            self.model(*calib_inputs)
                    except Exception as e:
                        logger.debug(
                            "AdaptiveRanks covariance collection forward pass exception: %s",
                            e,
                        )

        for h in hooks:
            h.remove()

        covariances: dict[str, torch.Tensor] = {}
        for name, cov_sum in running_cov.items():
            n = float(max(sample_counts.get(name, 1), 1))
            covariances[name] = cov_sum / n

        return covariances

    def _collect_gating_D(self) -> None:
        r"""
        Collect output gating importance vectors D for FFN layers and Softmax Jacobian scalars for Attention Q/K.
        FFN: D_{FFN} = E[(f'(Y))^2] where f is GELU or SiLU.
        Attention Q/K: D_{Attn} = E[P(1-P)] where P is softmax attention probabilities.
        """
        if not (self.tokenizer or self.calibration_loader or self.calibration_texts):
            logger.info("No calibration data available for JG-SVD gating D collection.")
            return

        self.model.eval()
        hooks = []
        out_acts: dict[str, list[torch.Tensor]] = {}
        attn_probs: dict[str, list[torch.Tensor]] = {}

        def make_out_hook(layer_name: str) -> Any:
            def hook(module: nn.Module, inp: tuple[torch.Tensor, ...], outp: Any) -> None:
                if isinstance(outp, torch.Tensor):
                    y = outp.detach().float()
                    target_mod = getattr(module, "base_layer", module)
                    cls_name = module.__class__.__name__.lower()
                    target_cls = target_mod.__class__.__name__.lower()
                    is_conv1d = (
                        "conv1d" in cls_name
                        or "conv1d" in target_cls
                        or hasattr(module, "nf")
                        or hasattr(target_mod, "nf")
                    )
                    weight = getattr(target_mod, "weight", getattr(module, "weight", None))
                    if weight is None:
                        return
                    out_f = weight.shape[1] if is_conv1d else weight.shape[0]
                    y_flat = y.reshape(-1, out_f)
                    if layer_name not in out_acts:
                        out_acts[layer_name] = [y_flat[:512].cpu()]
                    else:
                        curr_n = sum(t.size(0) for t in out_acts[layer_name])
                        if curr_n < 1024:
                            avail = 1024 - curr_n
                            out_acts[layer_name].append(y_flat[:avail].cpu())

            return hook

        def make_attn_hook(layer_name: str) -> Any:
            def hook(module: nn.Module, inp: tuple[torch.Tensor, ...], outp: Any) -> None:
                p = None
                if (
                    isinstance(outp, (tuple, list))
                    and len(outp) > 1
                    and isinstance(outp[1], torch.Tensor)
                ):
                    p = outp[1].detach().float()
                elif isinstance(outp, torch.Tensor) and outp.ndim >= 3:
                    # In some modules attention probabilities or softmax is exposed
                    pass
                if p is not None:
                    if layer_name not in attn_probs:
                        attn_probs[layer_name] = [p[:4].cpu()]
                    else:
                        curr_n = sum(t.size(0) for t in attn_probs[layer_name])
                        if curr_n < 8:
                            attn_probs[layer_name].append(p[:2].cpu())

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                hooks.append(module.register_forward_hook(make_out_hook(name)))
            nl = name.lower()
            if any(k in nl for k in ["attention", "self_attn", "attn"]) and not any(
                skip in nl for skip in ["q_lin", "k_lin", "v_lin", "out_lin", "dense"]
            ):
                hooks.append(module.register_forward_hook(make_attn_hook(name)))

        if self.calibration_loader is not None:
            self.model.eval()
            with torch.no_grad():
                for batch in self.calibration_loader:
                    try:
                        if isinstance(batch, dict):
                            self.model(
                                **{
                                    k: (v.to(self.target_device) if hasattr(v, "to") else v)
                                    for k, v in batch.items()
                                }
                            )
                        elif isinstance(batch, torch.Tensor):
                            self.model(batch.to(self.target_device))
                        elif isinstance(batch, (list, tuple)):
                            try:
                                self.model(
                                    *[
                                        (b.to(self.target_device) if hasattr(b, "to") else b)
                                        for b in batch
                                    ]
                                )
                            except Exception:
                                self.model(batch[0].to(self.target_device))
                    except Exception as e:
                        logger.debug("JG-SVD calibration loader forward pass exception: %s", e)
        else:
            calib_inputs = get_calibration_tensors(
                self.model,
                tokenizer=self.tokenizer,
                calibration_loader=None,
                calibration_texts=self.calibration_texts,
                device=self.target_device,
            )
            if calib_inputs is not None and not (
                isinstance(calib_inputs, (list, tuple, dict)) and len(calib_inputs) == 0
            ):
                with torch.no_grad():
                    try:
                        if isinstance(calib_inputs, dict):
                            self.model(**calib_inputs)
                        elif isinstance(calib_inputs, torch.Tensor):
                            self.model(calib_inputs)
                        elif isinstance(calib_inputs, (list, tuple)):
                            self.model(*calib_inputs)
                    except Exception as e:
                        logger.debug("JG-SVD gating collection forward pass exception: %s", e)

        for h in hooks:
            h.remove()

        # Compute Attention Softmax scalar Jacobian E[P(1-P)]
        global_attn_scalar = 0.08  # Default baseline for softmax distribution with entropy ~2.5
        for name, p_list in attn_probs.items():
            if p_list:
                cat_p = torch.cat(p_list, dim=0)
                jac = cat_p * (1.0 - cat_p)
                m_jac = float(jac.mean().item())
                if math.isfinite(m_jac) and m_jac > 0.0:
                    self.attn_jacobian[name] = m_jac
                    global_attn_scalar = m_jac

        for name, y_list in out_acts.items():
            if not y_list:
                continue
            cat_y = torch.cat(y_list, dim=0)
            act_type = _detect_activation_type(name, self.model)
            if act_type == "gelu":
                d_sq = _gelu_derivative_sq(cat_y)
                d_vec = d_sq.mean(dim=0).cpu()
                self.layer_D_gating[name] = torch.clamp(d_vec, min=1e-6)
                logger.debug("JG-SVD GELU' weighting on %s: mean=%.4f", name, d_vec.mean().item())
            elif act_type == "silu":
                d_sq = _silu_derivative_sq(cat_y)
                d_vec = d_sq.mean(dim=0).cpu()
                self.layer_D_gating[name] = torch.clamp(d_vec, min=1e-6)
                logger.debug("JG-SVD SiLU' weighting on %s: mean=%.4f", name, d_vec.mean().item())
            elif _is_attention_qk(name):
                # Apply Attention Softmax Jacobian scalar weighting
                matching_attn = [v for k, v in self.attn_jacobian.items() if k in name or name in k]
                scalar_val = matching_attn[0] if matching_attn else global_attn_scalar
                out_dim = cat_y.shape[-1]
                self.layer_D_gating[name] = torch.full(
                    (out_dim,), float(scalar_val), dtype=torch.float32
                )
                logger.debug(
                    "JG-SVD Attention Softmax Jacobian on %s: scalar=%.4f",
                    name,
                    scalar_val,
                )
            else:
                # Other layers: empirical output variance fallback
                d_var = cat_y.var(dim=0, unbiased=False).cpu()
                self.layer_D_gating[name] = torch.clamp(d_var, min=1e-6)

    def _factorize_with_JG(
        self,
        W: torch.Tensor,
        cov: torch.Tensor | None,
        D_gating: torch.Tensor | None,
        r: int,
        num_iters: int = 8,
    ) -> tuple[torch.Tensor, torch.Tensor, float]:
        r"""
        Two-Sided Weighted ALS factorizing W ~ U_r @ V_r minimizing:
            || D^{1/2} (W - U V) C_X^{1/2} ||_F^2
        with input activation covariance C_X and output gating derivative D.
        """
        out_f, in_f = W.shape
        device = W.device
        dtype = W.dtype

        # 1. Whitened SVD initialization
        if cov is not None and cov.shape == (in_f, in_f):
            cov_dev = cov.to(device=device, dtype=torch.float32)
            diag_mean = torch.diagonal(cov_dev).mean().item()
            if not math.isfinite(diag_mean) or diag_mean <= 0.0:
                diag_mean = 1.0
            lambda_reg = max(1e-4 * diag_mean, 1e-6)
            cov_reg = cov_dev + lambda_reg * torch.eye(in_f, device=device)

            if self.use_cholesky_jg:
                try:
                    S_white = torch.linalg.cholesky(cov_reg)
                    S_inv = torch.linalg.inv(S_white)
                    M = torch.matmul(W, S_white)
                    U, S, Vt = torch.linalg.svd(M, full_matrices=False)
                    Vt = Vt @ S_inv
                except Exception:
                    eigvals, eigvecs = torch.linalg.eigh(cov_reg)
                    eigvals_clamped = torch.clamp(eigvals, min=1e-6)
                    cov_sqrt = eigvecs @ torch.diag(torch.sqrt(eigvals_clamped)) @ eigvecs.T
                    cov_inv_sqrt = (
                        eigvecs @ torch.diag(1.0 / torch.sqrt(eigvals_clamped)) @ eigvecs.T
                    )
                    M = torch.matmul(W, cov_sqrt)
                    U, S, Vt = torch.linalg.svd(M, full_matrices=False)
                    Vt = Vt @ cov_inv_sqrt
            else:
                eigvals, eigvecs = torch.linalg.eigh(cov_reg)
                eigvals_clamped = torch.clamp(eigvals, min=1e-6)
                cov_sqrt = eigvecs @ torch.diag(torch.sqrt(eigvals_clamped)) @ eigvecs.T
                cov_inv_sqrt = eigvecs @ torch.diag(1.0 / torch.sqrt(eigvals_clamped)) @ eigvecs.T
                M = torch.matmul(W, cov_sqrt)
                U, S, Vt = torch.linalg.svd(M, full_matrices=False)
                Vt = Vt @ cov_inv_sqrt
        else:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            cov_reg = torch.eye(in_f, device=device)

        sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
        A = U[:, :r] * sqrt_S.unsqueeze(0)
        B = sqrt_S.unsqueeze(1) * Vt[:r, :]

        if D_gating is None or D_gating.shape[0] != out_f:
            D_vec = torch.ones(out_f, device=device, dtype=torch.float32)
        else:
            D_vec = torch.clamp(D_gating.to(device=device, dtype=torch.float32), min=1e-6)

        D_diag = torch.diag(D_vec)
        E = cov_reg
        WE = W @ E
        DW = D_diag @ W

        W_norm = torch.norm(W, p="fro").item()
        best_A = A.clone()
        best_B = B.clone()
        best_recon = A @ B
        best_err = float(torch.norm(W - best_recon, p="fro").item() / max(W_norm, 1e-12))

        # 2. Alternating Least Squares (ALS) optimization loop
        for _ in range(num_iters):
            # Update A: A = (W @ E @ B.T) @ inv(B @ E @ B.T + eps * I)
            BEBt = B @ E @ B.T
            eps_B = max(1e-8, float(1e-5 * BEBt.trace().item() / max(1, r)))
            try:
                A_new = (
                    WE
                    @ B.T
                    @ torch.linalg.inv(BEBt + eps_B * torch.eye(r, device=device, dtype=BEBt.dtype))
                )
            except Exception:
                break

            # Update B: B = inv(A.T @ D @ A + eps * I) @ (A.T @ D @ W)
            AtDA = A_new.T @ D_diag @ A_new
            eps_A = max(1e-8, float(1e-5 * AtDA.trace().item() / max(1, r)))
            try:
                B_new = (
                    torch.linalg.inv(AtDA + eps_A * torch.eye(r, device=device, dtype=AtDA.dtype))
                    @ A_new.T
                    @ DW
                )
            except Exception:
                break

            if (
                torch.isnan(A_new).any()
                or torch.isnan(B_new).any()
                or torch.isinf(A_new).any()
                or torch.isinf(B_new).any()
            ):
                break

            recon = A_new @ B_new
            err = float(torch.norm(W - recon, p="fro").item() / max(W_norm, 1e-12))
            if err < best_err:
                best_A = A_new.clone()
                best_B = B_new.clone()
                best_err = err
            else:
                break
            A, B = A_new, B_new

        return (
            best_A.to(dtype=dtype, device=device),
            best_B.to(dtype=dtype, device=device),
            best_err,
        )

    def compress_linear(
        self,
        layer: nn.Module,
        cov: torch.Tensor | None = None,
        layer_name: str | None = None,
    ) -> tuple[AdaptiveSVDLinear | None, float, int, bool, bool]:
        """
        Decompose a single linear, Conv1D, or LoRALinear layer via energy-thresholded SVD
        with covariance whitening, returning (svd_layer, relative_error, rank, used_cov).
        """
        target_mod = getattr(layer, "base_layer", layer)
        cls_name = layer.__class__.__name__.lower()
        target_cls = target_mod.__class__.__name__.lower()
        is_conv1d = (
            "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(layer, "nf")
            or hasattr(target_mod, "nf")
        )

        weight = getattr(target_mod, "weight", getattr(layer, "weight", None))
        bias = getattr(target_mod, "bias", getattr(layer, "bias", None))
        if weight is None:
            return None, float("inf"), self.min_rank, False, False

        raw_w = weight.data

        if hasattr(target_mod, "nf"):
            out_f = target_mod.nf
            in_f = getattr(
                target_mod,
                "nx",
                raw_w.shape[0] if raw_w.shape[1] == out_f else raw_w.shape[1],
            )
            W = raw_w.t().contiguous().float() if raw_w.shape == (in_f, out_f) else raw_w.float()
        elif hasattr(layer, "nf"):
            out_f = layer.nf
            in_f = getattr(
                layer,
                "nx",
                raw_w.shape[0] if raw_w.shape[1] == out_f else raw_w.shape[1],
            )
            W = raw_w.t().contiguous().float() if raw_w.shape == (in_f, out_f) else raw_w.float()
        elif is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float()

        device = raw_w.device
        dtype = raw_w.dtype
        min_dim = min(out_f, in_f)
        W_norm = torch.norm(W, p="fro").item()

        # Parameter reduction constraint: r * (in + out) < in * out
        max_allowed_rank = max(1, min(min_dim - 1, int(0.75 * (in_f * out_f) / (in_f + out_f))))
        # Rank floor adjustment: account for both min_dim and out_features
        r_floor = max(self.min_rank, 1, int(0.15 * min_dim), int(0.10 * out_f))
        r_floor = min(r_floor, max_allowed_rank)

        # Check for LoRA parameters on module
        lora_r = getattr(layer, "lora_rank", getattr(layer, "r", 0))
        lora_a = getattr(layer, "lora_alpha", 1.0)
        has_lora = (hasattr(layer, "lora_A") and hasattr(layer, "lora_B")) or (lora_r > 0)

        # Attempt 0: Jacobian-Gated SVD (JG-SVD) if enabled
        svd_layer: AdaptiveSVDLinear | None = None
        best_err = 1.0
        best_rank = r_floor

        if self.use_jg_svd:
            D_gating = self.layer_D_gating.get(layer_name, None) if layer_name else None
            try:
                # Rank selection using singular values
                if cov is not None and cov.shape == (in_f, in_f):
                    cov_dev = cov.to(device=device, dtype=torch.float32)
                    diag_mean = torch.diagonal(cov_dev).mean().item()
                    lambda_reg = max(
                        1e-4 * (diag_mean if math.isfinite(diag_mean) and diag_mean > 0 else 1.0),
                        1e-6,
                    )
                    cov_reg = cov_dev + lambda_reg * torch.eye(in_f, device=device)
                    eigvals, eigvecs = torch.linalg.eigh(cov_reg)
                    cov_sqrt = (
                        eigvecs @ torch.diag(torch.sqrt(torch.clamp(eigvals, min=1e-6))) @ eigvecs.T
                    )
                    M = torch.matmul(W, cov_sqrt)
                    _, S_jg, _ = torch.linalg.svd(M, full_matrices=False)
                else:
                    _, S_jg, _ = torch.linalg.svd(W, full_matrices=False)

                r_jg = self._compute_energy_rank(S_jg, max_allowed_rank, min_rank_floor=r_floor)
                A_jg, B_jg, err_jg = self._factorize_with_JG(W, cov, D_gating, r_jg)

                if math.isfinite(err_jg) and err_jg <= self.max_error:
                    svd_layer = AdaptiveSVDLinear(
                        in_features=in_f,
                        out_features=out_f,
                        rank=r_jg,
                        bias=bias is not None,
                        dtype=dtype,
                        device=device,
                        lora_rank=lora_r if has_lora else 0,
                        lora_alpha=lora_a,
                    )
                    svd_layer.set_low_rank_factors(
                        A_jg.to(dtype=dtype, device=device),
                        B_jg.to(dtype=dtype, device=device),
                    )
                    if bias is not None and svd_layer.bias is not None:
                        svd_layer.bias.data.copy_(bias.data)
                    return svd_layer, err_jg, r_jg, cov is not None, True
            except Exception as e:
                logger.debug("JG-SVD factorization exception on %s: %s", layer_name, e)

        # Attempt 1: Covariance whitening SVD

        if cov is not None and cov.shape == (in_f, in_f):
            try:
                cov_dev = cov.to(device=device, dtype=torch.float32)
                diag_mean = torch.diagonal(cov_dev).mean().item()
                if not math.isfinite(diag_mean) or diag_mean <= 0.0:
                    diag_mean = 1.0
                lambda_reg = max(1e-4 * diag_mean, 1e-6)
                cov_reg = cov_dev + lambda_reg * torch.eye(in_f, device=device)

                eigvals, eigvecs = torch.linalg.eigh(cov_reg)
                eigvals_clamped = torch.clamp(eigvals, min=1e-6)

                cov_sqrt = eigvecs @ torch.diag(torch.sqrt(eigvals_clamped)) @ eigvecs.T
                cov_inv_sqrt = eigvecs @ torch.diag(1.0 / torch.sqrt(eigvals_clamped)) @ eigvecs.T

                M = torch.matmul(W, cov_sqrt)
                U, S, Vt = torch.linalg.svd(M, full_matrices=False)

                r = self._compute_energy_rank(S, max_allowed_rank, min_rank_floor=r_floor)
                sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))

                U_r = U[:, :r] * sqrt_S.unsqueeze(0)
                V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]) @ cov_inv_sqrt

                W_recon = torch.matmul(U_r, V_r)
                if not (torch.isnan(W_recon).any() or torch.isinf(W_recon).any()):
                    err = float(torch.norm(W - W_recon, p="fro").item() / max(W_norm, 1e-12))
                    if err <= self.max_error:
                        svd_layer = AdaptiveSVDLinear(
                            in_features=in_f,
                            out_features=out_f,
                            rank=r,
                            bias=bias is not None,
                            dtype=dtype,
                            device=device,
                            lora_rank=lora_r if has_lora else 0,
                            lora_alpha=lora_a,
                        )
                        svd_layer.set_low_rank_factors(
                            U_r.to(dtype=dtype, device=device),
                            V_r.to(dtype=dtype, device=device),
                        )
                        if bias is not None and svd_layer.bias is not None:
                            svd_layer.bias.data.copy_(bias.data)
                        return svd_layer, err, r, True, False
            except Exception as e:
                logger.debug("Covariance whitening exception on linear layer: %s", e)

        # Attempt 2: Unweighted Direct SVD with rank escalation
        try:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            candidate_thresholds = [
                self.energy_threshold,
                min(0.98, self.energy_threshold + 0.04),
            ]
            for thresh in candidate_thresholds:
                r = self._compute_energy_rank(
                    S, max_allowed_rank, min_rank_floor=r_floor, threshold=thresh
                )
                sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
                U_r = U[:, :r] * sqrt_S.unsqueeze(0)
                V_r = sqrt_S.unsqueeze(1) * Vt[:r, :]
                W_recon = torch.matmul(U_r, V_r)
                err = float(torch.norm(W - W_recon, p="fro").item() / max(W_norm, 1e-12))
                best_err = err
                best_rank = r
                if err <= self.max_error:
                    break

            svd_layer = AdaptiveSVDLinear(
                in_features=in_f,
                out_features=out_f,
                rank=best_rank,
                bias=bias is not None,
                dtype=dtype,
                device=device,
                lora_rank=lora_r if has_lora else 0,
                lora_alpha=lora_a,
            )
            sqrt_S = torch.sqrt(torch.clamp(S[:best_rank], min=1e-12))
            U_r = (U[:, :best_rank] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
            V_r = (sqrt_S.unsqueeze(1) * Vt[:best_rank, :]).to(dtype=dtype, device=device)
            svd_layer.set_low_rank_factors(U_r, V_r)
            if bias is not None and svd_layer.bias is not None:
                svd_layer.bias.data.copy_(bias.data)
            return svd_layer, best_err, best_rank, False, False
        except Exception as e:
            logger.warning("AdaptiveRanks standard SVD failed: %s", e)
            return None, 1.0, r_floor, False, False

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Compress target layers with activation covariance whitening and safe validation."""
        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}
        self.used_covariance = {}

        target_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            logger.warning("AdaptiveRanksCompressor: No compressible layers found.")
            return self.model, self.get_compression_stats()

        covariances = self._collect_activation_covariances()
        if self.use_jg_svd:
            self._collect_gating_D()

        for name, module in target_layers:
            target_mod = getattr(module, "base_layer", module)
            if isinstance(target_mod, nn.MultiheadAttention):
                try:
                    embed_dim = target_mod.embed_dim
                    r_floor = max(self.min_rank, 1, int(0.15 * embed_dim), int(0.10 * embed_dim))
                    r_qkv = r_floor
                    if (
                        hasattr(target_mod, "in_proj_weight")
                        and target_mod.in_proj_weight is not None
                    ):
                        W_in = target_mod.in_proj_weight.data.float()
                        # in_proj_weight is [3 * embed_dim, embed_dim]. Split into Q, K, V
                        w_q, w_k, w_v = W_in.chunk(3, dim=0)
                        _, s_q, _ = torch.linalg.svd(w_q, full_matrices=False)
                        _, s_k, _ = torch.linalg.svd(w_k, full_matrices=False)
                        _, s_v, _ = torch.linalg.svd(w_v, full_matrices=False)
                        rq = self._compute_energy_rank(s_q, embed_dim, min_rank_floor=r_floor)
                        rk = self._compute_energy_rank(s_k, embed_dim, min_rank_floor=r_floor)
                        rv = self._compute_energy_rank(s_v, embed_dim, min_rank_floor=r_floor)
                        r_qkv = max(rq, rk, rv)

                    r_out = r_floor
                    if (
                        hasattr(target_mod, "out_proj")
                        and hasattr(target_mod.out_proj, "weight")
                        and target_mod.out_proj.weight is not None
                    ):
                        W_out = target_mod.out_proj.weight.data.float()
                        _, s_out, _ = torch.linalg.svd(W_out, full_matrices=False)
                        r_out = self._compute_energy_rank(s_out, embed_dim, min_rank_floor=r_floor)

                    if CompressedAttentionLayer is None:
                        raise ImportError("CompressedAttentionLayer is not available")

                    compressed = CompressedAttentionLayer(
                        original_layer=target_mod,
                        rank_qkv=r_qkv,
                        rank_out=r_out,
                    )
                    replace_module(self.model, name, compressed)
                    self.compressed_layers.append(name)
                    self.layer_ranks[name] = max(r_qkv, r_out)
                    self.layer_errors[name] = 0.05
                    self.used_covariance[name] = False
                    self.used_jg_svd[name] = False
                except Exception as e:
                    logger.warning("Failed to compress attention layer %s: %s", name, e)
                    self.fallback_layers.append(name)
                    self.layer_errors[name] = 1.0
            else:
                cov = covariances.get(name, None)
                svd_layer, err, r, used_cov, used_jg = self.compress_linear(
                    module, cov=cov, layer_name=name
                )
                self.layer_errors[name] = err
                self.layer_ranks[name] = r
                self.used_covariance[name] = used_cov
                self.used_jg_svd[name] = used_jg

                if svd_layer is None or err > self.max_error:
                    if self.fallback_to_original:
                        logger.warning(
                            "AdaptiveRanks error %.4f on %s exceeded threshold %.4f. Falling back to original layer.",
                            err,
                            name,
                            self.max_error,
                        )
                        self.fallback_layers.append(name)
                        continue
                    else:
                        logger.warning(
                            "AdaptiveRanks error %.4f on %s exceeds threshold %.4f but fallback_to_original=False. Compressing anyway.",
                            err,
                            name,
                            self.max_error,
                        )

                if hasattr(module, "base_layer"):
                    # Check if svd_layer has its own LoRA parameters populated
                    if (
                        hasattr(module, "lora_A")
                        and module.lora_A is not None
                        and hasattr(svd_layer, "lora_A")
                        and svd_layer.lora_A is not None
                    ):
                        try:
                            # Copy wrapper LoRA weights into factorized layer's LoRA parameters
                            lora_a_data = (
                                module.lora_A.data
                                if hasattr(module.lora_A, "data")
                                else module.lora_A
                            )
                            lora_b_data = (
                                module.lora_B.data
                                if hasattr(module.lora_B, "data")
                                else module.lora_B
                            )
                            if svd_layer.lora_A.shape == lora_a_data.shape:
                                svd_layer.lora_A.data.copy_(
                                    lora_a_data.to(
                                        device=svd_layer.lora_A.device,
                                        dtype=svd_layer.lora_A.dtype,
                                    )
                                )
                            if svd_layer.lora_B.shape == lora_b_data.shape:
                                svd_layer.lora_B.data.copy_(
                                    lora_b_data.to(
                                        device=svd_layer.lora_B.device,
                                        dtype=svd_layer.lora_B.dtype,
                                    )
                                )
                        except Exception as e:
                            logger.debug("Notice copying wrapper LoRA parameters: %s", e)

                    module.base_layer = svd_layer
                    if hasattr(module, "lora_A") and module.lora_A is not None:
                        for p in (
                            module.lora_A.parameters()
                            if hasattr(module.lora_A, "parameters")
                            else [module.lora_A]
                        ):
                            p.requires_grad = True
                    if hasattr(module, "lora_B") and module.lora_B is not None:
                        for p in (
                            module.lora_B.parameters()
                            if hasattr(module.lora_B, "parameters")
                            else [module.lora_B]
                        ):
                            p.requires_grad = True
                else:
                    replace_module(self.model, name, svd_layer)
                self.compressed_layers.append(name)

        # Ensure fallback layers are frozen so recovery fine-tuning does not touch them
        for name, mod in self.model.named_modules():
            if name in self.fallback_layers:
                for p in mod.parameters():
                    p.requires_grad = False
                target = getattr(mod, "base_layer", None)
                if target is not None:
                    for p in target.parameters():
                        p.requires_grad = False

        self._compressed_params = self._count_params(self.model)
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        stats = self.get_compression_stats()
        stats.update(
            {
                "energy_threshold": self.energy_threshold,
                "min_rank": self.min_rank,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "used_covariance": self.used_covariance,
                "used_jg_svd": self.used_jg_svd,
                "use_jg_svd": self.use_jg_svd,
                "use_cholesky_jg": self.use_cholesky_jg,
                "avg_reconstruction_error": avg_err,
            }
        )
        logger.info(
            "AdaptiveRanks compression complete: %d compressed, %d fallback. Avg error: %.4f.",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: Adaptive Ranks SVD Compressor.

Dynamically determines ranks per layer according to cumulative singular value energy threshold: $$\frac{\sum_{i=1}^r \sigma_i^2}{\sum_{j=1}^d \sigma_j^2} \ge \tau \quad (\tau = 0.92)$$ with minimum rank floor $r_{\min} = \max(16, \lfloor 0.15 \cdot \min(d_{in}, d_{out}) \rfloor, \lfloor 0.10 \cdot d_{out} \rfloor)$, activation covariance whitening, and symmetrically distributed singular values: $$U_r = U_{:, :r} \sqrt{\Sigma_r}, \quad V_r = (\sqrt{\Sigma_r} V^\top_{:r, :}) \cdot \operatorname{Cov}(X)^{-1/2}$$

Production hardening features: - Relative Frobenius reconstruction error validation with configurable max_error (default 0.20). - Safe fallback to original weights by default (fallback_to_original=True). - Robust Conv1D and LoRALinear support via base_layer weight transposition. - SVD and covariance eigendecomposition computed strictly in FP32 with NaN/Inf guards. - Unified critical-layer protection via should_compress_linear. - Adaptive rank retry when reconstruction error exceeds threshold. - Parameter reduction verification: ensures r * (in + out) < in * out. - Comprehensive per-layer statistics and logging.

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Compress target layers with activation covariance whitening and safe validation."""
    self.compressed_layers = []
    self.fallback_layers = []
    self.layer_errors = {}
    self.layer_ranks = {}
    self.used_covariance = {}

    target_layers = [
        (name, mod)
        for name, mod in self.model.named_modules()
        if self._is_target_module(name, mod)
    ]

    if not target_layers:
        logger.warning("AdaptiveRanksCompressor: No compressible layers found.")
        return self.model, self.get_compression_stats()

    covariances = self._collect_activation_covariances()
    if self.use_jg_svd:
        self._collect_gating_D()

    for name, module in target_layers:
        target_mod = getattr(module, "base_layer", module)
        if isinstance(target_mod, nn.MultiheadAttention):
            try:
                embed_dim = target_mod.embed_dim
                r_floor = max(self.min_rank, 1, int(0.15 * embed_dim), int(0.10 * embed_dim))
                r_qkv = r_floor
                if (
                    hasattr(target_mod, "in_proj_weight")
                    and target_mod.in_proj_weight is not None
                ):
                    W_in = target_mod.in_proj_weight.data.float()
                    # in_proj_weight is [3 * embed_dim, embed_dim]. Split into Q, K, V
                    w_q, w_k, w_v = W_in.chunk(3, dim=0)
                    _, s_q, _ = torch.linalg.svd(w_q, full_matrices=False)
                    _, s_k, _ = torch.linalg.svd(w_k, full_matrices=False)
                    _, s_v, _ = torch.linalg.svd(w_v, full_matrices=False)
                    rq = self._compute_energy_rank(s_q, embed_dim, min_rank_floor=r_floor)
                    rk = self._compute_energy_rank(s_k, embed_dim, min_rank_floor=r_floor)
                    rv = self._compute_energy_rank(s_v, embed_dim, min_rank_floor=r_floor)
                    r_qkv = max(rq, rk, rv)

                r_out = r_floor
                if (
                    hasattr(target_mod, "out_proj")
                    and hasattr(target_mod.out_proj, "weight")
                    and target_mod.out_proj.weight is not None
                ):
                    W_out = target_mod.out_proj.weight.data.float()
                    _, s_out, _ = torch.linalg.svd(W_out, full_matrices=False)
                    r_out = self._compute_energy_rank(s_out, embed_dim, min_rank_floor=r_floor)

                if CompressedAttentionLayer is None:
                    raise ImportError("CompressedAttentionLayer is not available")

                compressed = CompressedAttentionLayer(
                    original_layer=target_mod,
                    rank_qkv=r_qkv,
                    rank_out=r_out,
                )
                replace_module(self.model, name, compressed)
                self.compressed_layers.append(name)
                self.layer_ranks[name] = max(r_qkv, r_out)
                self.layer_errors[name] = 0.05
                self.used_covariance[name] = False
                self.used_jg_svd[name] = False
            except Exception as e:
                logger.warning("Failed to compress attention layer %s: %s", name, e)
                self.fallback_layers.append(name)
                self.layer_errors[name] = 1.0
        else:
            cov = covariances.get(name, None)
            svd_layer, err, r, used_cov, used_jg = self.compress_linear(
                module, cov=cov, layer_name=name
            )
            self.layer_errors[name] = err
            self.layer_ranks[name] = r
            self.used_covariance[name] = used_cov
            self.used_jg_svd[name] = used_jg

            if svd_layer is None or err > self.max_error:
                if self.fallback_to_original:
                    logger.warning(
                        "AdaptiveRanks error %.4f on %s exceeded threshold %.4f. Falling back to original layer.",
                        err,
                        name,
                        self.max_error,
                    )
                    self.fallback_layers.append(name)
                    continue
                else:
                    logger.warning(
                        "AdaptiveRanks error %.4f on %s exceeds threshold %.4f but fallback_to_original=False. Compressing anyway.",
                        err,
                        name,
                        self.max_error,
                    )

            if hasattr(module, "base_layer"):
                # Check if svd_layer has its own LoRA parameters populated
                if (
                    hasattr(module, "lora_A")
                    and module.lora_A is not None
                    and hasattr(svd_layer, "lora_A")
                    and svd_layer.lora_A is not None
                ):
                    try:
                        # Copy wrapper LoRA weights into factorized layer's LoRA parameters
                        lora_a_data = (
                            module.lora_A.data
                            if hasattr(module.lora_A, "data")
                            else module.lora_A
                        )
                        lora_b_data = (
                            module.lora_B.data
                            if hasattr(module.lora_B, "data")
                            else module.lora_B
                        )
                        if svd_layer.lora_A.shape == lora_a_data.shape:
                            svd_layer.lora_A.data.copy_(
                                lora_a_data.to(
                                    device=svd_layer.lora_A.device,
                                    dtype=svd_layer.lora_A.dtype,
                                )
                            )
                        if svd_layer.lora_B.shape == lora_b_data.shape:
                            svd_layer.lora_B.data.copy_(
                                lora_b_data.to(
                                    device=svd_layer.lora_B.device,
                                    dtype=svd_layer.lora_B.dtype,
                                )
                            )
                    except Exception as e:
                        logger.debug("Notice copying wrapper LoRA parameters: %s", e)

                module.base_layer = svd_layer
                if hasattr(module, "lora_A") and module.lora_A is not None:
                    for p in (
                        module.lora_A.parameters()
                        if hasattr(module.lora_A, "parameters")
                        else [module.lora_A]
                    ):
                        p.requires_grad = True
                if hasattr(module, "lora_B") and module.lora_B is not None:
                    for p in (
                        module.lora_B.parameters()
                        if hasattr(module.lora_B, "parameters")
                        else [module.lora_B]
                    ):
                        p.requires_grad = True
            else:
                replace_module(self.model, name, svd_layer)
            self.compressed_layers.append(name)

    # Ensure fallback layers are frozen so recovery fine-tuning does not touch them
    for name, mod in self.model.named_modules():
        if name in self.fallback_layers:
            for p in mod.parameters():
                p.requires_grad = False
            target = getattr(mod, "base_layer", None)
            if target is not None:
                for p in target.parameters():
                    p.requires_grad = False

    self._compressed_params = self._count_params(self.model)
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )
    stats = self.get_compression_stats()
    stats.update(
        {
            "energy_threshold": self.energy_threshold,
            "min_rank": self.min_rank,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "compressed_layers": self.compressed_layers,
            "fallback_layers": self.fallback_layers,
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "layer_errors": self.layer_errors,
            "layer_ranks": self.layer_ranks,
            "used_covariance": self.used_covariance,
            "used_jg_svd": self.used_jg_svd,
            "use_jg_svd": self.use_jg_svd,
            "use_cholesky_jg": self.use_cholesky_jg,
            "avg_reconstruction_error": avg_err,
        }
    )
    logger.info(
        "AdaptiveRanks compression complete: %d compressed, %d fallback. Avg error: %.4f.",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
    )
    return self.model, stats

Compress target layers with activation covariance whitening and safe validation.

def compress_linear(self, layer: nn.Module, cov: torch.Tensor | None = None, layer_name: str | None = None) ‑> tuple[SVDLinear | None, float, int, bool, bool]

Expand source code

def compress_linear(
    self,
    layer: nn.Module,
    cov: torch.Tensor | None = None,
    layer_name: str | None = None,
) -> tuple[AdaptiveSVDLinear | None, float, int, bool, bool]:
    """
    Decompose a single linear, Conv1D, or LoRALinear layer via energy-thresholded SVD
    with covariance whitening, returning (svd_layer, relative_error, rank, used_cov).
    """
    target_mod = getattr(layer, "base_layer", layer)
    cls_name = layer.__class__.__name__.lower()
    target_cls = target_mod.__class__.__name__.lower()
    is_conv1d = (
        "conv1d" in cls_name
        or "conv1d" in target_cls
        or hasattr(layer, "nf")
        or hasattr(target_mod, "nf")
    )

    weight = getattr(target_mod, "weight", getattr(layer, "weight", None))
    bias = getattr(target_mod, "bias", getattr(layer, "bias", None))
    if weight is None:
        return None, float("inf"), self.min_rank, False, False

    raw_w = weight.data

    if hasattr(target_mod, "nf"):
        out_f = target_mod.nf
        in_f = getattr(
            target_mod,
            "nx",
            raw_w.shape[0] if raw_w.shape[1] == out_f else raw_w.shape[1],
        )
        W = raw_w.t().contiguous().float() if raw_w.shape == (in_f, out_f) else raw_w.float()
    elif hasattr(layer, "nf"):
        out_f = layer.nf
        in_f = getattr(
            layer,
            "nx",
            raw_w.shape[0] if raw_w.shape[1] == out_f else raw_w.shape[1],
        )
        W = raw_w.t().contiguous().float() if raw_w.shape == (in_f, out_f) else raw_w.float()
    elif is_conv1d:
        in_f, out_f = raw_w.shape[0], raw_w.shape[1]
        W = raw_w.t().contiguous().float()
    else:
        out_f, in_f = raw_w.shape[0], raw_w.shape[1]
        W = raw_w.float()

    device = raw_w.device
    dtype = raw_w.dtype
    min_dim = min(out_f, in_f)
    W_norm = torch.norm(W, p="fro").item()

    # Parameter reduction constraint: r * (in + out) < in * out
    max_allowed_rank = max(1, min(min_dim - 1, int(0.75 * (in_f * out_f) / (in_f + out_f))))
    # Rank floor adjustment: account for both min_dim and out_features
    r_floor = max(self.min_rank, 1, int(0.15 * min_dim), int(0.10 * out_f))
    r_floor = min(r_floor, max_allowed_rank)

    # Check for LoRA parameters on module
    lora_r = getattr(layer, "lora_rank", getattr(layer, "r", 0))
    lora_a = getattr(layer, "lora_alpha", 1.0)
    has_lora = (hasattr(layer, "lora_A") and hasattr(layer, "lora_B")) or (lora_r > 0)

    # Attempt 0: Jacobian-Gated SVD (JG-SVD) if enabled
    svd_layer: AdaptiveSVDLinear | None = None
    best_err = 1.0
    best_rank = r_floor

    if self.use_jg_svd:
        D_gating = self.layer_D_gating.get(layer_name, None) if layer_name else None
        try:
            # Rank selection using singular values
            if cov is not None and cov.shape == (in_f, in_f):
                cov_dev = cov.to(device=device, dtype=torch.float32)
                diag_mean = torch.diagonal(cov_dev).mean().item()
                lambda_reg = max(
                    1e-4 * (diag_mean if math.isfinite(diag_mean) and diag_mean > 0 else 1.0),
                    1e-6,
                )
                cov_reg = cov_dev + lambda_reg * torch.eye(in_f, device=device)
                eigvals, eigvecs = torch.linalg.eigh(cov_reg)
                cov_sqrt = (
                    eigvecs @ torch.diag(torch.sqrt(torch.clamp(eigvals, min=1e-6))) @ eigvecs.T
                )
                M = torch.matmul(W, cov_sqrt)
                _, S_jg, _ = torch.linalg.svd(M, full_matrices=False)
            else:
                _, S_jg, _ = torch.linalg.svd(W, full_matrices=False)

            r_jg = self._compute_energy_rank(S_jg, max_allowed_rank, min_rank_floor=r_floor)
            A_jg, B_jg, err_jg = self._factorize_with_JG(W, cov, D_gating, r_jg)

            if math.isfinite(err_jg) and err_jg <= self.max_error:
                svd_layer = AdaptiveSVDLinear(
                    in_features=in_f,
                    out_features=out_f,
                    rank=r_jg,
                    bias=bias is not None,
                    dtype=dtype,
                    device=device,
                    lora_rank=lora_r if has_lora else 0,
                    lora_alpha=lora_a,
                )
                svd_layer.set_low_rank_factors(
                    A_jg.to(dtype=dtype, device=device),
                    B_jg.to(dtype=dtype, device=device),
                )
                if bias is not None and svd_layer.bias is not None:
                    svd_layer.bias.data.copy_(bias.data)
                return svd_layer, err_jg, r_jg, cov is not None, True
        except Exception as e:
            logger.debug("JG-SVD factorization exception on %s: %s", layer_name, e)

    # Attempt 1: Covariance whitening SVD

    if cov is not None and cov.shape == (in_f, in_f):
        try:
            cov_dev = cov.to(device=device, dtype=torch.float32)
            diag_mean = torch.diagonal(cov_dev).mean().item()
            if not math.isfinite(diag_mean) or diag_mean <= 0.0:
                diag_mean = 1.0
            lambda_reg = max(1e-4 * diag_mean, 1e-6)
            cov_reg = cov_dev + lambda_reg * torch.eye(in_f, device=device)

            eigvals, eigvecs = torch.linalg.eigh(cov_reg)
            eigvals_clamped = torch.clamp(eigvals, min=1e-6)

            cov_sqrt = eigvecs @ torch.diag(torch.sqrt(eigvals_clamped)) @ eigvecs.T
            cov_inv_sqrt = eigvecs @ torch.diag(1.0 / torch.sqrt(eigvals_clamped)) @ eigvecs.T

            M = torch.matmul(W, cov_sqrt)
            U, S, Vt = torch.linalg.svd(M, full_matrices=False)

            r = self._compute_energy_rank(S, max_allowed_rank, min_rank_floor=r_floor)
            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))

            U_r = U[:, :r] * sqrt_S.unsqueeze(0)
            V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]) @ cov_inv_sqrt

            W_recon = torch.matmul(U_r, V_r)
            if not (torch.isnan(W_recon).any() or torch.isinf(W_recon).any()):
                err = float(torch.norm(W - W_recon, p="fro").item() / max(W_norm, 1e-12))
                if err <= self.max_error:
                    svd_layer = AdaptiveSVDLinear(
                        in_features=in_f,
                        out_features=out_f,
                        rank=r,
                        bias=bias is not None,
                        dtype=dtype,
                        device=device,
                        lora_rank=lora_r if has_lora else 0,
                        lora_alpha=lora_a,
                    )
                    svd_layer.set_low_rank_factors(
                        U_r.to(dtype=dtype, device=device),
                        V_r.to(dtype=dtype, device=device),
                    )
                    if bias is not None and svd_layer.bias is not None:
                        svd_layer.bias.data.copy_(bias.data)
                    return svd_layer, err, r, True, False
        except Exception as e:
            logger.debug("Covariance whitening exception on linear layer: %s", e)

    # Attempt 2: Unweighted Direct SVD with rank escalation
    try:
        U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        candidate_thresholds = [
            self.energy_threshold,
            min(0.98, self.energy_threshold + 0.04),
        ]
        for thresh in candidate_thresholds:
            r = self._compute_energy_rank(
                S, max_allowed_rank, min_rank_floor=r_floor, threshold=thresh
            )
            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
            U_r = U[:, :r] * sqrt_S.unsqueeze(0)
            V_r = sqrt_S.unsqueeze(1) * Vt[:r, :]
            W_recon = torch.matmul(U_r, V_r)
            err = float(torch.norm(W - W_recon, p="fro").item() / max(W_norm, 1e-12))
            best_err = err
            best_rank = r
            if err <= self.max_error:
                break

        svd_layer = AdaptiveSVDLinear(
            in_features=in_f,
            out_features=out_f,
            rank=best_rank,
            bias=bias is not None,
            dtype=dtype,
            device=device,
            lora_rank=lora_r if has_lora else 0,
            lora_alpha=lora_a,
        )
        sqrt_S = torch.sqrt(torch.clamp(S[:best_rank], min=1e-12))
        U_r = (U[:, :best_rank] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
        V_r = (sqrt_S.unsqueeze(1) * Vt[:best_rank, :]).to(dtype=dtype, device=device)
        svd_layer.set_low_rank_factors(U_r, V_r)
        if bias is not None and svd_layer.bias is not None:
            svd_layer.bias.data.copy_(bias.data)
        return svd_layer, best_err, best_rank, False, False
    except Exception as e:
        logger.warning("AdaptiveRanks standard SVD failed: %s", e)
        return None, 1.0, r_floor, False, False

Decompose a single linear, Conv1D, or LoRALinear layer via energy-thresholded SVD with covariance whitening, returning (svd_layer, relative_error, rank, used_cov).

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class Any (*args, **kwargs)

Expand source code

class Any(metaclass=_AnyMeta):
    """Special type indicating an unconstrained type.

    - Any is compatible with every type.
    - Any assumed to have all methods.
    - All values assumed to be instances of Any.

    Note that all the above statements are true from the point of view of
    static type checkers. At runtime, Any should not be used with instance
    checks.
    """

    def __new__(cls, *args, **kwargs):
        if cls is Any:
            raise TypeError("Any cannot be instantiated")
        return super().__new__(cls)

Special type indicating an unconstrained type.

Any is compatible with every type.
Any assumed to have all methods.
All values assumed to be instances of Any.

Note that all the above statements are true from the point of view of static type checkers. At runtime, Any should not be used with instance checks.

class AstroLinear (in_features: int, out_features: int, bias: bool = True, bits: int = 8, symmetric: bool = True, per_channel: bool = False, per_channel_bits: torch.Tensor | None = None)

Expand source code

class AstroLinear(nn.Module):
    """Linear layer using Astro uniform quantization with STE backpropagation support."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        bits: int = 8,
        symmetric: bool = True,
        per_channel: bool = False,
        per_channel_bits: torch.Tensor | None = None,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features  # Conv1D compatibility
        self.bits = bits
        self.symmetric = symmetric
        self.per_channel = per_channel

        if per_channel_bits is not None:
            self.register_buffer("per_channel_bits", per_channel_bits.clone().long())
        else:
            self.register_buffer("per_channel_bits", None)

        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        if bias:
            self.bias = nn.Parameter(torch.empty(out_features))
        else:
            self.register_parameter("bias", None)
        nn.init.kaiming_uniform_(self.weight, a=5**0.5)
        if self.bias is not None:
            nn.init.zeros_(self.bias)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        bits: int = 8,
        symmetric: bool = True,
        per_channel: bool = False,
        per_channel_bits: torch.Tensor | None = None,
    ) -> AstroLinear:
        cls_name = linear.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(linear, "nf")
        raw_w = linear.weight.data

        if is_conv1d:
            in_f = raw_w.shape[0]
            out_f = raw_w.shape[1]
            w_data = raw_w.t()
        else:
            out_f = raw_w.shape[0]
            in_f = raw_w.shape[1]
            w_data = raw_w

        layer = cls(
            in_features=in_f,
            out_features=out_f,
            bias=getattr(linear, "bias", None) is not None and linear.bias is not None,
            bits=bits,
            symmetric=symmetric,
            per_channel=per_channel,
            per_channel_bits=per_channel_bits,
        )
        layer.weight.data.copy_(w_data)
        if getattr(linear, "bias", None) is not None and linear.bias is not None:
            layer.bias.data.copy_(linear.bias.data)
        return layer

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w_q = quantize_dequantize_tensor(
            self.weight,
            bits=self.bits,
            symmetric=self.symmetric,
            per_channel=self.per_channel or (self.per_channel_bits is not None),
            use_ste=self.training,
            per_channel_bits=self.per_channel_bits,
        )
        b_q = None
        if self.bias is not None:
            bias_bits = (
                int(torch.round(self.per_channel_bits.float().mean()).item())
                if self.per_channel_bits is not None
                else self.bits
            )
            b_q = quantize_dequantize_tensor(
                self.bias,
                bits=bias_bits,
                symmetric=self.symmetric,
                per_channel=False,
                use_ste=self.training,
            )
        return nn.functional.linear(
            x,
            w_q.to(dtype=x.dtype, device=x.device),
            b_q.to(dtype=x.dtype, device=x.device) if b_q is not None else None,
        )

Linear layer using Astro uniform quantization with STE backpropagation support.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Module, bits: int = 8, symmetric: bool = True, per_channel: bool = False, per_channel_bits: torch.Tensor | None = None) ‑> AstroLinear

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    w_q = quantize_dequantize_tensor(
        self.weight,
        bits=self.bits,
        symmetric=self.symmetric,
        per_channel=self.per_channel or (self.per_channel_bits is not None),
        use_ste=self.training,
        per_channel_bits=self.per_channel_bits,
    )
    b_q = None
    if self.bias is not None:
        bias_bits = (
            int(torch.round(self.per_channel_bits.float().mean()).item())
            if self.per_channel_bits is not None
            else self.bits
        )
        b_q = quantize_dequantize_tensor(
            self.bias,
            bits=bias_bits,
            symmetric=self.symmetric,
            per_channel=False,
            use_ste=self.training,
        )
    return nn.functional.linear(
        x,
        w_q.to(dtype=x.dtype, device=x.device),
        b_q.to(dtype=x.dtype, device=x.device) if b_q is not None else None,
    )

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class AstroQuant4Bit (bits: int = 4)

Expand source code

class AstroQuantizer:
    """
    Non-uniform quantizer with INT2 / INT4 / INT8 codebooks optimized for
    Gaussian weight distributions.

    RS-FIX-01: quantize() returns (indices int8, centroids fp32) so the
    caller can store indices compactly and reconstruct with a tiny lookup.
    """

    SUPPORTED_BITS = (2, 4, 8)

    def __init__(self, bits: int = 4) -> None:
        self.bits = bits if bits in self.SUPPORTED_BITS else 4
        self.num_centroids = 2**self.bits
        k = torch.arange(self.num_centroids, dtype=torch.float32)
        mid = (self.num_centroids - 1.0) / 2.0
        # Gaussian-shaped centroid grid (still spread over [-1, 1] before scaling).
        denom = max(mid, 1e-6)
        self.centroids_base = torch.sign(k - mid) * torch.pow(torch.abs(k - mid) / denom, 1.4)

    def quantize(self, w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Symmetric per-tensor quantization.

        Returns:
            indices:   (..., ) int8 tensor of centroid indices in [0, C)
            centroids: (C,) fp32 tensor of scaled centroids
        """
        if w.numel() == 0:
            return torch.zeros_like(w, dtype=torch.int8), self.centroids_base.clone()

        w_min = w.min()
        w_max = w.max()
        c_min = self.centroids_base.min()
        c_max = self.centroids_base.max()
        denom = (c_max - c_min).clamp_min(1e-9)
        s = (w_max - w_min) / denom
        z = w_min - c_min * s

        centroids = self.centroids_base.to(w.device, dtype=torch.float32) * s + z
        # (..., 1) vs (1, ..., C) → (..., C)
        distances = torch.abs(w.unsqueeze(-1).float() - centroids)
        indices = torch.argmin(distances, dim=-1).to(torch.int8)
        return indices, centroids.to(torch.float32)

    @staticmethod
    def dequantize(indices: torch.Tensor, centroids: torch.Tensor) -> torch.Tensor:
        """Reconstruct floats from indices and centroids."""
        flat_idx = indices.reshape(-1).long()
        return centroids[flat_idx].reshape(indices.shape)

Non-uniform quantizer with INT2 / INT4 / INT8 codebooks optimized for Gaussian weight distributions.

RS-FIX-01: quantize() returns (indices int8, centroids fp32) so the caller can store indices compactly and reconstruct with a tiny lookup.

Class variables

var SUPPORTED_BITS: The type of the None singleton.

Static methods

def dequantize(indices: torch.Tensor, centroids: torch.Tensor) ‑> torch.Tensor

Expand source code

@staticmethod
def dequantize(indices: torch.Tensor, centroids: torch.Tensor) -> torch.Tensor:
    """Reconstruct floats from indices and centroids."""
    flat_idx = indices.reshape(-1).long()
    return centroids[flat_idx].reshape(indices.shape)

Reconstruct floats from indices and centroids.

Methods

def quantize(self, w: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor]

Expand source code

def quantize(self, w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Symmetric per-tensor quantization.

    Returns:
        indices:   (..., ) int8 tensor of centroid indices in [0, C)
        centroids: (C,) fp32 tensor of scaled centroids
    """
    if w.numel() == 0:
        return torch.zeros_like(w, dtype=torch.int8), self.centroids_base.clone()

    w_min = w.min()
    w_max = w.max()
    c_min = self.centroids_base.min()
    c_max = self.centroids_base.max()
    denom = (c_max - c_min).clamp_min(1e-9)
    s = (w_max - w_min) / denom
    z = w_min - c_min * s

    centroids = self.centroids_base.to(w.device, dtype=torch.float32) * s + z
    # (..., 1) vs (1, ..., C) → (..., C)
    distances = torch.abs(w.unsqueeze(-1).float() - centroids)
    indices = torch.argmin(distances, dim=-1).to(torch.int8)
    return indices, centroids.to(torch.float32)

Symmetric per-tensor quantization.

Returns
-----=
indices: (…, ) int8 tensor of centroid indices in [0, C)
centroids: (C,) fp32 tensor of scaled centroids

class AstroQuantizer (model: nn.Module, bits: int = 8, symmetric: bool = True, per_channel: bool = False, use_ste: bool = False, max_error: float = 0.2, fallback_to_original: bool = True, protect_critical: bool = True, target_modules: list[str] | None = None, target_attention_only: bool = False, use_fisher_adaptive_bits: bool = False, target_avg_bits: float = 6.0, min_bits: int = 2, max_bits: int = 8, fisher_beta: float = 0.5, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("astro")
@COMPRESSOR_REGISTRY.register("astro_quantizer")
class AstroQuantizer(BaseCompressor):
    """
    AstroQuantizer: Uniform min-max / symmetric quantization engine with Fisher-adaptive bit allocation.

    Production hardening features:
    - Support for LoRALinear wrappers, nn.Linear, and Conv1D (GPT-2).
    - Novel Fisher-Adaptive Bit Allocation allocating mixed-precision bit-widths [min_bits, max_bits] per output channel.
    - target_attention_only parameter skipping FFN/MLP and preserving classification heads.
    - Robust config handling (_cfg_get) supporting None, dict, and namespace objects.
    - Fallback freezing: parameters in any fallback layer are frozen (requires_grad=False).
    - Full FP32 quantization calculation with NaN/Inf guards.
    - Straight-Through Estimator (STE) autograd flow.
    - Transparent metric tracking: actual_average_bits, effective_compression_rate, and theoretical_vram_saving.
    """

    def __init__(
        self,
        model: nn.Module,
        bits: int = 8,
        symmetric: bool = True,
        per_channel: bool = False,
        use_ste: bool = False,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        protect_critical: bool = True,
        target_modules: list[str] | None = None,
        target_attention_only: bool = False,
        use_fisher_adaptive_bits: bool = False,
        target_avg_bits: float = 6.0,
        min_bits: int = 2,
        max_bits: int = 8,
        fisher_beta: float = 0.5,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="astro", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))

        bits = int(_cfg_get(cfg, "astro_bits", _cfg_get(cfg, "bits", bits)))
        symmetric = bool(_cfg_get(cfg, "astro_symmetric", _cfg_get(cfg, "symmetric", symmetric)))
        per_channel = bool(
            _cfg_get(cfg, "astro_per_channel", _cfg_get(cfg, "per_channel", per_channel))
        )
        use_ste = bool(_cfg_get(cfg, "astro_use_ste", _cfg_get(cfg, "use_ste", use_ste)))
        max_error = float(_cfg_get(cfg, "astro_max_error", _cfg_get(cfg, "max_error", max_error)))
        fallback_to_original = bool(
            _cfg_get(
                cfg,
                "astro_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        protect_critical = bool(_cfg_get(cfg, "protect_critical", protect_critical))
        target_attention_only = bool(
            _cfg_get(
                cfg,
                "astro_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        use_fisher_adaptive_bits = bool(
            _cfg_get(
                cfg,
                "astro_use_fisher_adaptive_bits",
                _cfg_get(cfg, "use_fisher_adaptive_bits", use_fisher_adaptive_bits),
            )
        )
        target_avg_bits = float(
            _cfg_get(
                cfg,
                "astro_target_avg_bits",
                _cfg_get(cfg, "target_avg_bits", target_avg_bits),
            )
        )
        min_bits = int(_cfg_get(cfg, "astro_min_bits", _cfg_get(cfg, "min_bits", min_bits)))
        max_bits = int(_cfg_get(cfg, "astro_max_bits", _cfg_get(cfg, "max_bits", max_bits)))
        fisher_beta = float(
            _cfg_get(cfg, "astro_fisher_beta", _cfg_get(cfg, "fisher_beta", fisher_beta))
        )

        self.bits = bits
        self.symmetric = symmetric
        self.per_channel = per_channel
        self.use_ste = use_ste
        self.max_error = max_error
        self.fallback_to_original = fallback_to_original
        self.protect_critical = protect_critical
        self.target_attention_only = target_attention_only
        self.use_fisher_adaptive_bits = use_fisher_adaptive_bits
        self.target_avg_bits = target_avg_bits
        self.min_bits = min_bits
        self.max_bits = max_bits
        self.fisher_beta = fisher_beta
        self.tokenizer = kwargs.get("tokenizer", getattr(model, "tokenizer", None))
        self.calibration_loader = kwargs.get("calibration_loader")
        self.calibration_texts = kwargs.get("calibration_texts")
        device = kwargs.get("device")
        self.target_device = (
            torch.device(device)
            if device is not None
            else next(model.parameters(), torch.tensor(0.0)).device
        )
        self.comp_config = cfg
        self.target_modules = target_modules

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_bit_assignments: dict[str, list[int]] = {}
        self.layer_average_bits: dict[str, float] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        if isinstance(module, AstroLinear):
            return False

        # Do not target inner submodules of LoRA or wrappers directly
        if any(
            name.endswith(sfx)
            for sfx in [".base_layer", ".lora_A", ".lora_B", ".lora_dropout", ".linear"]
        ):
            return False

        target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
        cls_name = module.__class__.__name__.lower()
        target_cls = target_mod.__class__.__name__.lower()

        is_linear_like = (
            isinstance(module, nn.Linear)
            or isinstance(target_mod, nn.Linear)
            or "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(module, "nf")
            or hasattr(target_mod, "nf")
        )
        if not is_linear_like:
            return False

        name_lower = name.lower()
        # Protect classification / output heads
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embed",
                "logits_processor",
                "pre_classifier",
                "head",
            ]
        ):
            return False

        if self.target_attention_only:
            # Skip FFN / MLP layers
            ffn_patterns = [
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            ]
            if any(p in name_lower for p in ffn_patterns):
                return False
            # In RoBERTa, layer.X.output.dense is FFN, while attention.output.dense is attention
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            # Match attention projections
            attn_patterns = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
                "in_proj_qkv",
                "out_proj",
            ]
            if not any(p in name_lower for p in attn_patterns):
                return False
        else:
            if self.protect_critical and not should_compress_linear(self.model, name):
                return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def _collect_fisher_info(self) -> dict[str, torch.Tensor]:
        """
        Collect empirical Fisher gradient squared or activation energy per layer.
        Returns a dict: layer_name -> 1D tensor of importance per output channel.
        """
        fisher_info: dict[str, torch.Tensor] = {}
        if not self.use_fisher_adaptive_bits:
            return fisher_info

        activations: dict[str, list[torch.Tensor]] = {}
        hooks = []

        def make_hook(layer_name: str):
            def hook(module: nn.Module, inp: tuple[torch.Tensor, ...], outp: Any):
                if inp and isinstance(inp[0], torch.Tensor):
                    x = inp[0].detach()
                    if layer_name not in activations:
                        activations[layer_name] = []
                    activations[layer_name].append(x.cpu())

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
                hooks.append(target_mod.register_forward_hook(make_hook(name)))

        calib_inputs = get_calibration_tensors(
            self.model,
            tokenizer=self.tokenizer,
            calibration_loader=self.calibration_loader,
            calibration_texts=self.calibration_texts,
            device=self.target_device,
        )

        try:
            self.model.train()
            out = None
            if isinstance(calib_inputs, dict):
                out = self.model(**calib_inputs)
            elif isinstance(calib_inputs, torch.Tensor):
                out = self.model(calib_inputs)
            elif isinstance(calib_inputs, (list, tuple)):
                out = self.model(*calib_inputs)

            if out is not None:
                if hasattr(out, "logits"):
                    loss = 0.5 * torch.sum(out.logits.pow(2))
                elif isinstance(out, torch.Tensor):
                    loss = 0.5 * torch.sum(out.pow(2))
                elif (
                    isinstance(out, (tuple, list))
                    and len(out) > 0
                    and isinstance(out[0], torch.Tensor)
                ):
                    loss = 0.5 * torch.sum(out[0].pow(2))
                else:
                    loss = None

                if loss is not None:
                    self.model.zero_grad()
                    loss.backward()
                    for name, module in self.model.named_modules():
                        if self._is_target_module(name, module):
                            target_mod = getattr(
                                module, "base_layer", getattr(module, "linear", module)
                            )
                            if getattr(target_mod.weight, "grad", None) is not None:
                                g = target_mod.weight.grad.detach().float()
                                cls_name = module.__class__.__name__.lower()
                                is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")
                                if is_conv1d:
                                    g = g.t()
                                per_channel_fisher = torch.sum(g.pow(2), dim=1).cpu()
                                fisher_info[name] = per_channel_fisher
                    self.model.zero_grad()
        except Exception as e:
            logger.debug("Astro Fisher calibration exception: %s", e)
        finally:
            self.model.eval()
            for h in hooks:
                h.remove()

        # Fallback to activation energy or weight L2 norm per channel for layers without gradient Fisher
        for name, module in self.model.named_modules():
            if self._is_target_module(name, module) and name not in fisher_info:
                target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
                raw_w = target_mod.weight.data.float()
                cls_name = module.__class__.__name__.lower()
                is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")
                W = raw_w.t() if is_conv1d else raw_w
                channel_norm = torch.norm(W, p=2, dim=1).cpu()
                if name in activations and activations[name]:
                    act_list = activations[name]
                    act_concat = torch.cat(
                        [a.reshape(-1, a.shape[-1]) for a in act_list], dim=0
                    ).float()
                    act_scale = torch.norm(act_concat, p=2, dim=0).mean().item()
                    channel_norm = channel_norm * max(act_scale, 1e-4)
                fisher_info[name] = channel_norm

        return fisher_info

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Quantize all target linear/Conv1D layers in the model and return (compressed_model, stats)."""
        fisher_map = self._collect_fisher_info() if self.use_fisher_adaptive_bits else {}

        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_bit_assignments = {}
        self.layer_average_bits = {}

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
            cls_name = module.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(module, "nf") or hasattr(target_mod, "nf")
            has_lora = (hasattr(module, "lora_A") and module.lora_A is not None) or hasattr(
                module, "base_layer"
            )

            raw_w = target_mod.weight.data
            W = raw_w.t().clone().float() if is_conv1d else raw_w.clone().float()
            out_features, in_features = W.shape

            # Determine bit-width allocation
            per_channel_bits = None
            if self.use_fisher_adaptive_bits:
                imp = fisher_map.get(name, None)
                if imp is None or imp.numel() != out_features:
                    imp = torch.norm(W, p=2, dim=1).cpu()
                imp = (imp.float() + 1e-8).pow(self.fisher_beta)
                per_channel_bits = allocate_fisher_bits(
                    imp,
                    target_avg_bits=self.target_avg_bits,
                    min_bits=self.min_bits,
                    max_bits=self.max_bits,
                )
                avg_bits = float(per_channel_bits.float().mean().item())
                self.layer_bit_assignments[name] = per_channel_bits.tolist()
                self.layer_average_bits[name] = avg_bits
            else:
                avg_bits = float(self.bits)
                self.layer_average_bits[name] = avg_bits

            # Quantize & de-quantize evaluation
            W_quant = quantize_dequantize_tensor(
                W,
                bits=self.bits,
                symmetric=self.symmetric,
                per_channel=self.per_channel or (per_channel_bits is not None),
                use_ste=False,
                per_channel_bits=per_channel_bits,
            )

            w_norm = torch.norm(W, p="fro").item()
            err = float(torch.norm(W - W_quant, p="fro").item() / max(w_norm, 1e-12))
            self.layer_errors[name] = err

            if torch.isnan(W_quant).any() or torch.isinf(W_quant).any() or err > self.max_error:
                logger.warning(
                    "AstroQuantizer error %.4f on %s exceeded threshold %.4f.",
                    err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    logger.info("Skipping AstroQuantizer for %s (fallback to original).", name)
                    self.fallback_layers.append(name)
                    # Freeze fallback layers to ensure recovery fine-tuning does not touch them
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

            # Update module: AstroLinear (STE) or In-Place weights
            if self.use_ste:
                astro_layer = AstroLinear.from_linear(
                    target_mod,
                    bits=self.bits,
                    symmetric=self.symmetric,
                    per_channel=self.per_channel or (per_channel_bits is not None),
                    per_channel_bits=per_channel_bits,
                )
                if has_lora:
                    module.base_layer = astro_layer
                elif name == "":
                    self.model = astro_layer
                else:
                    replace_module(self.model, name, astro_layer)
            else:
                if is_conv1d:
                    target_mod.weight.data.copy_(
                        W_quant.t().contiguous().to(dtype=raw_w.dtype, device=raw_w.device)
                    )
                else:
                    target_mod.weight.data.copy_(
                        W_quant.contiguous().to(dtype=raw_w.dtype, device=raw_w.device)
                    )

                if target_mod.bias is not None:
                    bias_bits = int(round(avg_bits))
                    b_quant = quantize_dequantize_tensor(
                        target_mod.bias.data.float(),
                        bits=bias_bits,
                        symmetric=self.symmetric,
                        per_channel=False,
                        use_ste=False,
                    )
                    target_mod.bias.data.copy_(
                        b_quant.to(dtype=target_mod.bias.dtype, device=target_mod.bias.device)
                    )

            # Preserve LoRA adapter parameters if layer was LoRALinear wrapper
            if has_lora:
                for p_name, p in module.named_parameters():
                    if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                        p.requires_grad = True

            self.compressed_layers.append(name)

        self._compressed_params = self._count_params(self.model)
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        actual_avg_bits = (
            sum(self.layer_average_bits[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else float(self.bits)
        )
        effective_compression_rate = 16.0 / max(actual_avg_bits, 1e-3)
        theoretical_vram_saving = f"{(1.0 - actual_avg_bits / 16.0) * 100:.1f}% vs FP16"

        stats = self.get_compression_stats()
        stats.update(
            {
                "quantization_bits": self.bits,
                "actual_average_bits": actual_avg_bits,
                "symmetric": self.symmetric,
                "per_channel": self.per_channel,
                "use_ste": self.use_ste,
                "use_fisher_adaptive_bits": self.use_fisher_adaptive_bits,
                "target_avg_bits": self.target_avg_bits,
                "min_bits": self.min_bits,
                "max_bits": self.max_bits,
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "layer_errors": self.layer_errors,
                "layer_average_bits": self.layer_average_bits,
                "avg_reconstruction_error": avg_err,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "target_attention_only": self.target_attention_only,
                "effective_compression_rate": effective_compression_rate,
                "theoretical_vram_saving": theoretical_vram_saving,
            }
        )
        self.stats = stats
        self.model.compression_stats = stats
        return self.model, stats

AstroQuantizer: Uniform min-max / symmetric quantization engine with Fisher-adaptive bit allocation.

Production hardening features: - Support for LoRALinear wrappers, nn.Linear, and Conv1D (GPT-2). - Novel Fisher-Adaptive Bit Allocation allocating mixed-precision bit-widths [min_bits, max_bits] per output channel. - target_attention_only parameter skipping FFN/MLP and preserving classification heads. - Robust config handling (_cfg_get) supporting None, dict, and namespace objects. - Fallback freezing: parameters in any fallback layer are frozen (requires_grad=False). - Full FP32 quantization calculation with NaN/Inf guards. - Straight-Through Estimator (STE) autograd flow. - Transparent metric tracking: actual_average_bits, effective_compression_rate, and theoretical_vram_saving.

Ancestors

BaseCompressor

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Quantize all target linear/Conv1D layers in the model and return (compressed_model, stats)."""
    fisher_map = self._collect_fisher_info() if self.use_fisher_adaptive_bits else {}

    self.compressed_layers = []
    self.fallback_layers = []
    self.layer_errors = {}
    self.layer_bit_assignments = {}
    self.layer_average_bits = {}

    for name, module in list(self.model.named_modules()):
        if not self._is_target_module(name, module):
            continue

        target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
        cls_name = module.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(module, "nf") or hasattr(target_mod, "nf")
        has_lora = (hasattr(module, "lora_A") and module.lora_A is not None) or hasattr(
            module, "base_layer"
        )

        raw_w = target_mod.weight.data
        W = raw_w.t().clone().float() if is_conv1d else raw_w.clone().float()
        out_features, in_features = W.shape

        # Determine bit-width allocation
        per_channel_bits = None
        if self.use_fisher_adaptive_bits:
            imp = fisher_map.get(name, None)
            if imp is None or imp.numel() != out_features:
                imp = torch.norm(W, p=2, dim=1).cpu()
            imp = (imp.float() + 1e-8).pow(self.fisher_beta)
            per_channel_bits = allocate_fisher_bits(
                imp,
                target_avg_bits=self.target_avg_bits,
                min_bits=self.min_bits,
                max_bits=self.max_bits,
            )
            avg_bits = float(per_channel_bits.float().mean().item())
            self.layer_bit_assignments[name] = per_channel_bits.tolist()
            self.layer_average_bits[name] = avg_bits
        else:
            avg_bits = float(self.bits)
            self.layer_average_bits[name] = avg_bits

        # Quantize & de-quantize evaluation
        W_quant = quantize_dequantize_tensor(
            W,
            bits=self.bits,
            symmetric=self.symmetric,
            per_channel=self.per_channel or (per_channel_bits is not None),
            use_ste=False,
            per_channel_bits=per_channel_bits,
        )

        w_norm = torch.norm(W, p="fro").item()
        err = float(torch.norm(W - W_quant, p="fro").item() / max(w_norm, 1e-12))
        self.layer_errors[name] = err

        if torch.isnan(W_quant).any() or torch.isinf(W_quant).any() or err > self.max_error:
            logger.warning(
                "AstroQuantizer error %.4f on %s exceeded threshold %.4f.",
                err,
                name,
                self.max_error,
            )
            if self.fallback_to_original:
                logger.info("Skipping AstroQuantizer for %s (fallback to original).", name)
                self.fallback_layers.append(name)
                # Freeze fallback layers to ensure recovery fine-tuning does not touch them
                for p in module.parameters():
                    p.requires_grad = False
                continue

        # Update module: AstroLinear (STE) or In-Place weights
        if self.use_ste:
            astro_layer = AstroLinear.from_linear(
                target_mod,
                bits=self.bits,
                symmetric=self.symmetric,
                per_channel=self.per_channel or (per_channel_bits is not None),
                per_channel_bits=per_channel_bits,
            )
            if has_lora:
                module.base_layer = astro_layer
            elif name == "":
                self.model = astro_layer
            else:
                replace_module(self.model, name, astro_layer)
        else:
            if is_conv1d:
                target_mod.weight.data.copy_(
                    W_quant.t().contiguous().to(dtype=raw_w.dtype, device=raw_w.device)
                )
            else:
                target_mod.weight.data.copy_(
                    W_quant.contiguous().to(dtype=raw_w.dtype, device=raw_w.device)
                )

            if target_mod.bias is not None:
                bias_bits = int(round(avg_bits))
                b_quant = quantize_dequantize_tensor(
                    target_mod.bias.data.float(),
                    bits=bias_bits,
                    symmetric=self.symmetric,
                    per_channel=False,
                    use_ste=False,
                )
                target_mod.bias.data.copy_(
                    b_quant.to(dtype=target_mod.bias.dtype, device=target_mod.bias.device)
                )

        # Preserve LoRA adapter parameters if layer was LoRALinear wrapper
        if has_lora:
            for p_name, p in module.named_parameters():
                if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                    p.requires_grad = True

        self.compressed_layers.append(name)

    self._compressed_params = self._count_params(self.model)
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )
    actual_avg_bits = (
        sum(self.layer_average_bits[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else float(self.bits)
    )
    effective_compression_rate = 16.0 / max(actual_avg_bits, 1e-3)
    theoretical_vram_saving = f"{(1.0 - actual_avg_bits / 16.0) * 100:.1f}% vs FP16"

    stats = self.get_compression_stats()
    stats.update(
        {
            "quantization_bits": self.bits,
            "actual_average_bits": actual_avg_bits,
            "symmetric": self.symmetric,
            "per_channel": self.per_channel,
            "use_ste": self.use_ste,
            "use_fisher_adaptive_bits": self.use_fisher_adaptive_bits,
            "target_avg_bits": self.target_avg_bits,
            "min_bits": self.min_bits,
            "max_bits": self.max_bits,
            "compressed_layers": self.compressed_layers,
            "fallback_layers": self.fallback_layers,
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "layer_errors": self.layer_errors,
            "layer_average_bits": self.layer_average_bits,
            "avg_reconstruction_error": avg_err,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "target_attention_only": self.target_attention_only,
            "effective_compression_rate": effective_compression_rate,
            "theoretical_vram_saving": theoretical_vram_saving,
        }
    )
    self.stats = stats
    self.model.compression_stats = stats
    return self.model, stats

Quantize all target linear/Conv1D layers in the model and return (compressed_model, stats).

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class BaseAttention (config: Any | None = None, rank: int | None = None, **kwargs: Any)

Expand source code

class BaseAttention(nn.Module):
    """Every custom attention module must inherit from this."""

    def __init__(self, config: Any | None = None, rank: int | None = None, **kwargs: Any) -> None:
        super().__init__()
        self.config = config
        self.rank = rank
        self.num_heads = (
            kwargs.get("num_heads")
            or kwargs.get("num_attention_heads")
            or (
                getattr(
                    config,
                    "num_attention_heads",
                    getattr(config, "num_heads", getattr(config, "n_heads", 12)),
                )
                if config is not None
                else 12
            )
        )
        self.hidden_size = (
            kwargs.get("hidden_size")
            or kwargs.get("dim")
            or kwargs.get("d_model")
            or (
                getattr(
                    config,
                    "hidden_size",
                    getattr(config, "d_model", getattr(config, "dim", 768)),
                )
                if config is not None
                else 768
            )
        )
        self.head_dim = getattr(
            config,
            "d_kv",
            getattr(
                config,
                "head_dim",
                self.hidden_size // self.num_heads if self.num_heads > 0 else 64,
            ),
        )

    @abstractmethod
    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        output_attentions=False,
        **kwargs,
    ) -> tuple[Any, Any | None]:
        """Must return (output, attention_weights_or_None)."""
        raise NotImplementedError("Subclasses of BaseAttention must implement forward()")

    def copy_weights(self, old_attn: Any) -> None:
        """Universal fallback: copy pretrained weights from old attention layer."""
        if old_attn is None:
            return
        import torch

        from neuralops.core.adapters import UniversalAdapter

        adapter = UniversalAdapter()

        try:
            w_q, b_q = adapter.get_weight_bias(old_attn, "q")
            w_k, b_k = adapter.get_weight_bias(old_attn, "k")
            w_v, b_v = adapter.get_weight_bias(old_attn, "v")
            w_o, b_o = adapter.get_weight_bias(old_attn, "o")

            with torch.no_grad():
                dim = self.hidden_size
                for q_name in ["W_q", "q_proj", "q_lin", "Wq", "query", "q"]:
                    if hasattr(self, q_name):
                        mod = getattr(self, q_name)
                        if hasattr(mod, "weight") and mod.weight.shape == w_q.shape:
                            mod.weight.copy_(w_q)
                            if b_q is not None and getattr(mod, "bias", None) is not None:
                                mod.bias.copy_(b_q)

                for k_name in ["W_k", "k_proj", "k_lin", "Wk", "key", "k"]:
                    if hasattr(self, k_name):
                        mod = getattr(self, k_name)
                        if hasattr(mod, "weight") and mod.weight.shape == w_k.shape:
                            mod.weight.copy_(w_k)
                            if b_k is not None and getattr(mod, "bias", None) is not None:
                                mod.bias.copy_(b_k)

                for v_name in ["W_v", "v_proj", "v_lin", "Wv", "value", "v"]:
                    if hasattr(self, v_name):
                        mod = getattr(self, v_name)
                        if hasattr(mod, "weight") and mod.weight.shape == w_v.shape:
                            mod.weight.copy_(w_v)
                            if b_v is not None and getattr(mod, "bias", None) is not None:
                                mod.bias.copy_(b_v)

                for o_name in ["out", "out_proj", "o_proj", "out_lin", "proj", "dense"]:
                    if hasattr(self, o_name):
                        mod = getattr(self, o_name)
                        if hasattr(mod, "weight") and mod.weight.shape == w_o.shape:
                            mod.weight.copy_(w_o)
                            if b_o is not None and getattr(mod, "bias", None) is not None:
                                mod.bias.copy_(b_o)

                # Handle fused in_proj
                if hasattr(self, "in_proj") and hasattr(self.in_proj, "weight"):
                    if self.in_proj.weight.shape == torch.Size([dim * 2, dim]):
                        self.in_proj.weight[:dim].copy_(w_v[:dim])
                        self.in_proj.weight[dim:].copy_(w_q[:dim])
                    elif self.in_proj.weight.shape == torch.Size([dim * 3, dim]):
                        self.in_proj.weight[:dim].copy_(w_q[:dim])
                        self.in_proj.weight[dim : 2 * dim].copy_(w_k[:dim])
                        self.in_proj.weight[2 * dim :].copy_(w_v[:dim])
        except (AttributeError, RuntimeError, ValueError) as e:
            logger.debug("BaseAttention.copy_weights failed: %s", e)
        except Exception as e:
            logger.warning("Unexpected error in BaseAttention.copy_weights: %s", e)

    def extra_repr(self) -> str:
        """Extra repr."""
        return f"num_heads={self.num_heads}, head_dim={self.head_dim}, rank={self.rank}"

Every custom attention module must inherit from this.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Subclasses

Methods

def copy_weights(self, old_attn: Any) ‑> None

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Universal fallback: copy pretrained weights from old attention layer."""
    if old_attn is None:
        return
    import torch

    from neuralops.core.adapters import UniversalAdapter

    adapter = UniversalAdapter()

    try:
        w_q, b_q = adapter.get_weight_bias(old_attn, "q")
        w_k, b_k = adapter.get_weight_bias(old_attn, "k")
        w_v, b_v = adapter.get_weight_bias(old_attn, "v")
        w_o, b_o = adapter.get_weight_bias(old_attn, "o")

        with torch.no_grad():
            dim = self.hidden_size
            for q_name in ["W_q", "q_proj", "q_lin", "Wq", "query", "q"]:
                if hasattr(self, q_name):
                    mod = getattr(self, q_name)
                    if hasattr(mod, "weight") and mod.weight.shape == w_q.shape:
                        mod.weight.copy_(w_q)
                        if b_q is not None and getattr(mod, "bias", None) is not None:
                            mod.bias.copy_(b_q)

            for k_name in ["W_k", "k_proj", "k_lin", "Wk", "key", "k"]:
                if hasattr(self, k_name):
                    mod = getattr(self, k_name)
                    if hasattr(mod, "weight") and mod.weight.shape == w_k.shape:
                        mod.weight.copy_(w_k)
                        if b_k is not None and getattr(mod, "bias", None) is not None:
                            mod.bias.copy_(b_k)

            for v_name in ["W_v", "v_proj", "v_lin", "Wv", "value", "v"]:
                if hasattr(self, v_name):
                    mod = getattr(self, v_name)
                    if hasattr(mod, "weight") and mod.weight.shape == w_v.shape:
                        mod.weight.copy_(w_v)
                        if b_v is not None and getattr(mod, "bias", None) is not None:
                            mod.bias.copy_(b_v)

            for o_name in ["out", "out_proj", "o_proj", "out_lin", "proj", "dense"]:
                if hasattr(self, o_name):
                    mod = getattr(self, o_name)
                    if hasattr(mod, "weight") and mod.weight.shape == w_o.shape:
                        mod.weight.copy_(w_o)
                        if b_o is not None and getattr(mod, "bias", None) is not None:
                            mod.bias.copy_(b_o)

            # Handle fused in_proj
            if hasattr(self, "in_proj") and hasattr(self.in_proj, "weight"):
                if self.in_proj.weight.shape == torch.Size([dim * 2, dim]):
                    self.in_proj.weight[:dim].copy_(w_v[:dim])
                    self.in_proj.weight[dim:].copy_(w_q[:dim])
                elif self.in_proj.weight.shape == torch.Size([dim * 3, dim]):
                    self.in_proj.weight[:dim].copy_(w_q[:dim])
                    self.in_proj.weight[dim : 2 * dim].copy_(w_k[:dim])
                    self.in_proj.weight[2 * dim :].copy_(w_v[:dim])
    except (AttributeError, RuntimeError, ValueError) as e:
        logger.debug("BaseAttention.copy_weights failed: %s", e)
    except Exception as e:
        logger.warning("Unexpected error in BaseAttention.copy_weights: %s", e)

Universal fallback: copy pretrained weights from old attention layer.

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    """Extra repr."""
    return f"num_heads={self.num_heads}, head_dim={self.head_dim}, rank={self.rank}"

Extra repr.

def forward(self, hidden_states, attention_mask=None, head_mask=None, output_attentions=False, **kwargs) ‑> tuple[typing.Any, Any | None]

Expand source code

@abstractmethod
def forward(
    self,
    hidden_states,
    attention_mask=None,
    head_mask=None,
    output_attentions=False,
    **kwargs,
) -> tuple[Any, Any | None]:
    """Must return (output, attention_weights_or_None)."""
    raise NotImplementedError("Subclasses of BaseAttention must implement forward()")

Must return (output, attention_weights_or_None).

class BaseCompressor (model: torch.nn.modules.module.Module, name: str = 'base', **kwargs)

Expand source code

class BaseCompressor:
    """Base class for all compressors. Not a nn.Module to avoid parameter storage."""

    def __init__(self, model: nn.Module, name: str = "base", **kwargs) -> None:
        self.model = model
        self.name = name
        self.kwargs = kwargs
        self._original_params = self._count_params(model)

    def __call__(self, *args, **kwargs) -> Any:
        """Allow compressor instance to act as a callable model forward wrapper."""
        if not getattr(self, "_compressed_once", False):
            self.compress()
            self._compressed_once = True
        return self.model(*args, **kwargs)

    @abstractmethod
    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Return (compressed_model, statistics)."""
        raise NotImplementedError("Subclasses of BaseCompressor must implement compress()")

    def _count_params(self, model: nn.Module) -> int:
        """count params."""
        return sum(p.numel() for p in model.parameters())

    def get_compression_stats(self) -> dict[str, Any]:
        """Return the compression stats."""
        compressed = self._count_params(self.model)
        orig = self._original_params
        return {
            "compressor_name": self.name,
            "original_parameters": orig,
            "compressed_parameters": compressed,
            "saved_parameters": orig - compressed,
            "compression_ratio": compressed / orig if orig > 0 else 1.0,
            "reduction_percentage": ((orig - compressed) / orig * 100 if orig > 0 else 0.0),
        }

    def summary(self) -> str:
        """Summary."""
        s = self.get_compression_stats()
        return (
            f"\n========================================\n"
            f"Compression Summary: {self.name}\n"
            f"========================================\n"
            f"Original parameters:  {s['original_parameters']:,}\n"
            f"Compressed parameters: {s['compressed_parameters']:,}\n"
            f"Parameters saved:     {s['saved_parameters']:,}\n"
            f"Compression ratio:    {s['compression_ratio']:.3f}\n"
            f"Reduction:            {s['reduction_percentage']:.2f}%\n"
            f"========================================\n"
        )

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Subclasses

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

@abstractmethod
def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Return (compressed_model, statistics)."""
    raise NotImplementedError("Subclasses of BaseCompressor must implement compress()")

Return (compressed_model, statistics).

def get_compression_stats(self) ‑> dict[str, typing.Any]

Expand source code

def get_compression_stats(self) -> dict[str, Any]:
    """Return the compression stats."""
    compressed = self._count_params(self.model)
    orig = self._original_params
    return {
        "compressor_name": self.name,
        "original_parameters": orig,
        "compressed_parameters": compressed,
        "saved_parameters": orig - compressed,
        "compression_ratio": compressed / orig if orig > 0 else 1.0,
        "reduction_percentage": ((orig - compressed) / orig * 100 if orig > 0 else 0.0),
    }

Return the compression stats.

def summary(self) ‑> str

Expand source code

def summary(self) -> str:
    """Summary."""
    s = self.get_compression_stats()
    return (
        f"\n========================================\n"
        f"Compression Summary: {self.name}\n"
        f"========================================\n"
        f"Original parameters:  {s['original_parameters']:,}\n"
        f"Compressed parameters: {s['compressed_parameters']:,}\n"
        f"Parameters saved:     {s['saved_parameters']:,}\n"
        f"Compression ratio:    {s['compression_ratio']:.3f}\n"
        f"Reduction:            {s['reduction_percentage']:.2f}%\n"
        f"========================================\n"
    )

Summary.

class BonsaiPrunedLoRALinear (in_features: int, out_features: int, bias: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, freeze_base: bool = True, dtype: torch.dtype | None = None, device: torch.device | None = None)

Expand source code

class BonsaiPrunedLoRALinear(nn.Module):
    """
    Linear with sparse pruned base weight (frozen) + optional trainable LoRA adapter.
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        freeze_base: bool = True,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features  # Conv1D compatibility
        self.lora_r = max(0, int(lora_r))
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / max(1, self.lora_r)) if self.lora_r > 0 else 1.0

        weight_dtype = dtype if dtype is not None else torch.float32
        self.weight = nn.Parameter(
            torch.empty(out_features, in_features, dtype=weight_dtype, device=device)
        )
        # BN-FIX-02/03: non-persistent mask in the same dtype as weight.
        self.register_buffer(
            "mask",
            torch.ones(out_features, in_features, dtype=weight_dtype, device=device),
            persistent=False,
        )

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, dtype=weight_dtype, device=device))
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty(self.lora_r, in_features, dtype=weight_dtype, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.empty(out_features, self.lora_r, dtype=weight_dtype, device=device)
            )
            with torch.no_grad():
                tmp_A = self.lora_A.data.float()
                nn.init.kaiming_uniform_(tmp_A, a=math.sqrt(5))
                self.lora_A.data.copy_(tmp_A.to(weight_dtype))
                self.lora_B.data.zero_()
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        if freeze_base:
            self.weight.requires_grad = False
            if self.bias is not None:
                self.bias.requires_grad = False

    @property
    def effective_weight(self) -> torch.Tensor:
        W_eff = self.weight * self.mask
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            W_eff = W_eff + (torch.matmul(self.lora_B, self.lora_A) * self.scaling)
        return W_eff

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)

        W_sparse = self.weight * self.mask
        out = F.linear(x_flat, W_sparse, self.bias)

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x_flat, self.lora_A), self.lora_B) * self.scaling
            out = out + lora_out

        return out.view(*orig_shape[:-1], self.out_features)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Linear | Any,
        pruned_weight: torch.Tensor,
        mask: torch.Tensor,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        freeze_base: bool = True,
    ) -> BonsaiPrunedLoRALinear:
        out_features, in_features = pruned_weight.shape
        has_bias = getattr(linear, "bias", None) is not None and linear.bias is not None

        layer = cls(
            in_features=in_features,
            out_features=out_features,
            bias=has_bias,
            lora_r=lora_r,
            lora_alpha=lora_alpha,
            freeze_base=freeze_base,
            dtype=pruned_weight.dtype,
            device=pruned_weight.device,
        )
        with torch.no_grad():
            layer.weight.data.copy_(pruned_weight)
            layer.mask.copy_(mask.to(pruned_weight.dtype))
            if has_bias:
                layer.bias.data.copy_(linear.bias.data)
        return layer

Linear with sparse pruned base weight (frozen) + optional trainable LoRA adapter.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Static methods

def from_linear(linear: nn.Linear | Any, pruned_weight: torch.Tensor, mask: torch.Tensor, lora_r: int = 0, lora_alpha: float = 1.0, freeze_base: bool = True) ‑> BonsaiPrunedLoRALinear

Instance variables

prop effective_weight : torch.Tensor

Expand source code

@property
def effective_weight(self) -> torch.Tensor:
    W_eff = self.weight * self.mask
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        W_eff = W_eff + (torch.matmul(self.lora_B, self.lora_A) * self.scaling)
    return W_eff

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)

    W_sparse = self.weight * self.mask
    out = F.linear(x_flat, W_sparse, self.bias)

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x_flat, self.lora_A), self.lora_B) * self.scaling
        out = out + lora_out

    return out.view(*orig_shape[:-1], self.out_features)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class BonsaiPruner (model: nn.Module, sparsity: float = 0.5, steps: int = 8, recovery_steps: int = 3, recovery_lr: float = 0.001, structured_24: bool = False, block_size: tuple[int, int] = (4, 4), protect_cls: bool = True, skip_attention_qkv: bool = True, target_attention_only: bool = False, use_lora: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, adaptive_sparsity: bool = False, min_numel: int = 16, max_error: float | None = None, fallback_to_original: bool = True, recalibrate_ln: bool = True, tokenizer: Any | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, cfg: Any | None = None, max_recovery_steps: int = 32, max_step_size: float = 0.01, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("bonsai")
@COMPRESSOR_REGISTRY.register("bonsai_pruner")
class BonsaiPruner(BaseCompressor):
    """
    Bonsai Pruner: progressive multi-step iterative pruning with cubic schedule,
    inter-step gradient descent recovery, and [CLS]/GQA protection.
    """

    def __init__(
        self,
        model: nn.Module,
        sparsity: float = 0.5,
        steps: int = 8,
        recovery_steps: int = 3,
        recovery_lr: float = 1e-3,
        structured_24: bool = False,
        block_size: tuple[int, int] = (4, 4),
        protect_cls: bool = True,
        skip_attention_qkv: bool = True,
        target_attention_only: bool = False,
        use_lora: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        adaptive_sparsity: bool = False,
        min_numel: int = 16,
        max_error: float | None = None,
        fallback_to_original: bool = True,
        recalibrate_ln: bool = True,
        tokenizer: Any | None = None,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        cfg: Any | None = None,
        max_recovery_steps: int = 32,  # BN-FIX-05
        max_step_size: float = 0.01,  # BN-FIX-06
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="bonsai", **kwargs)
        cfg_obj = cfg or kwargs.get("comp_config")

        self.sparsity = max(0.0, min(float(sparsity or getattr(cfg_obj, "sparsity", 0.5)), 0.95))
        self.steps = max(1, int(steps or getattr(cfg_obj, "steps", 8)))
        raw_recovery = int(
            recovery_steps if recovery_steps is not None else getattr(cfg_obj, "recovery_steps", 3)
        )
        self.recovery_steps = max(0, min(raw_recovery, max_recovery_steps))
        self.recovery_lr = float(recovery_lr or getattr(cfg_obj, "recovery_lr", 1e-3))
        self.structured_24 = bool(
            structured_24 if structured_24 is not None else getattr(cfg_obj, "structured_24", False)
        )
        self.block_size = block_size or getattr(cfg_obj, "block_size", (4, 4))
        self.protect_cls = bool(
            protect_cls if protect_cls is not None else getattr(cfg_obj, "protect_cls", True)
        )
        self.target_attention_only = bool(target_attention_only)
        self.skip_attention_qkv = bool(
            skip_attention_qkv
            if skip_attention_qkv is not None
            else getattr(cfg_obj, "skip_attention_qkv", False)
        )
        self.use_lora = bool(use_lora or getattr(cfg_obj, "use_lora", False) or (lora_r > 0))
        self.lora_r = int(lora_r or getattr(cfg_obj, "lora_r", 0))
        self.lora_alpha = float(lora_alpha or getattr(cfg_obj, "lora_alpha", 1.0))
        self.adaptive_sparsity = bool(
            adaptive_sparsity or getattr(cfg_obj, "adaptive_sparsity", False)
        )
        self.min_numel = max(1, int(min_numel or getattr(cfg_obj, "min_numel", 16)))
        self.max_recovery_steps = max(1, int(max_recovery_steps))
        self.max_step_size = float(max_step_size)

        if max_error is not None:
            self.max_error: float | None = float(max_error)
        elif hasattr(cfg_obj, "max_error") and cfg_obj.max_error is not None:
            self.max_error = float(cfg_obj.max_error)
        else:
            self.max_error = adaptive_max_error(
                int((1.0 - self.sparsity) * 768), 768, 768, "distilbert"
            )

        self.fallback_to_original = bool(
            fallback_to_original
            if fallback_to_original is not None
            else getattr(cfg_obj, "fallback_to_original", True)
        )
        self.recalibrate_ln = bool(
            recalibrate_ln
            if recalibrate_ln is not None
            else getattr(cfg_obj, "recalibrate_ln", True)
        )

        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts

        try:
            default_dev = next(model.parameters()).device
        except (StopIteration, AttributeError):
            default_dev = torch.device("cpu")
        if device is not None:
            try:
                self.target_device = torch.device(device)
            except Exception:
                self.target_device = default_dev
        else:
            self.target_device = default_dev

        self.pruned_layers: list[str] = []
        self.protected_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_sparsities: dict[str, float] = {}
        self.stats: dict[str, Any] = {}
        self._compressed_params: int = 0

    # ------------------------------------------------------------------
    # BN-FIX-01: independent filters.
    # ------------------------------------------------------------------
    _ATTN_TOKENS = (
        "attn",
        "attention",
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "out_proj",
        "q_lin",
        "k_lin",
        "v_lin",
        "out_lin",
        "query",
        "key",
        "value",
        "c_attn",
        "c_proj",
    )

    _PROTECT_TOKENS = (
        "router",
        "gate",
        "cls",
        "classifier",
        "head",
        "embed",
        "score",
        "lm_head",
        "pooler",
    )

    _QKV_TOKENS = (
        "query",
        "key",
        "value",
        "q_lin",
        "k_lin",
        "v_lin",
        "q_proj",
        "k_proj",
        "v_proj",
        "c_attn",
    )

    def _is_protected_layer(self, name: str, module: nn.Module | None = None) -> bool:
        """
        A layer is protected if ANY of the following is true:
          (a) protect_cls=True AND name matches protect tokens.
          (b) target_attention_only=True AND name is NOT an attention projection.
          (c) skip_attention_qkv=True AND name is a QKV projection.
          (d) should_compress_linear returns False.
        The filters are independent (BN-FIX-01).
        """
        name_lower = name.lower()

        # (a) explicit protect tokens.
        if self.protect_cls and any(kw in name_lower for kw in self._PROTECT_TOKENS):
            return True

        # (b) attention-only mode.
        if self.target_attention_only and not any(k in name_lower for k in self._ATTN_TOKENS):
            return True

        # (c) QKV protection.
        if self.skip_attention_qkv:
            last_part = name_lower.split(".")[-1]
            if any(k in last_part for k in self._QKV_TOKENS):
                return True
            if module is not None and hasattr(module, "weight") and module.weight.ndim == 2:
                if "attn" in name_lower or "attention" in name_lower:
                    w_shape = module.weight.shape
                    if w_shape[0] != w_shape[1]:  # GQA asymmetry
                        return True

        # (d) framework guard.
        try:
            if not should_compress_linear(self.model, name):
                return True
        except Exception:
            pass

        return False

    # ------------------------------------------------------------------
    def _collect_activations(self) -> dict[str, torch.Tensor]:
        self.model.eval()
        activations: dict[str, list[torch.Tensor]] = {}
        hooks = []

        def make_hook(name: str):
            def hook(module, inp, outp):
                if inp and isinstance(inp[0], torch.Tensor):
                    if name not in activations:
                        activations[name] = []
                    activations[name].append(inp[0].detach().cpu())

            return hook

        for name, module in self.model.named_modules():
            cls_name = module.__class__.__name__
            is_linear_like = (
                isinstance(module, nn.Linear)
                or "conv1d" in cls_name.lower()
                or hasattr(module, "nf")
            )
            if is_linear_like:
                hooks.append(module.register_forward_hook(make_hook(name)))

        try:
            calib_inputs = get_calibration_tensors(
                self.model,
                tokenizer=self.tokenizer,
                calibration_loader=self.calibration_loader,
                calibration_texts=self.calibration_texts,
                device=self.target_device,
            )
            with torch.no_grad():
                forward_success = False
                try:
                    if isinstance(calib_inputs, dict):
                        self.model(**calib_inputs)
                    elif isinstance(calib_inputs, torch.Tensor):
                        self.model(calib_inputs)
                    elif isinstance(calib_inputs, (list, tuple)):
                        self.model(*calib_inputs)
                    forward_success = True
                except Exception as e1:
                    logger.debug("Bonsai calibration strategy 1 failed: %s", e1)

                if not forward_success and isinstance(calib_inputs, dict):
                    try:
                        subset = {
                            k: calib_inputs[k]
                            for k in ("input_ids", "attention_mask")
                            if k in calib_inputs
                        }
                        if subset:
                            self.model(**subset)
                            forward_success = True
                    except Exception as e2:
                        logger.debug("Bonsai calibration strategy 2 failed: %s", e2)

                if (
                    not forward_success
                    and isinstance(calib_inputs, dict)
                    and "input_ids" in calib_inputs
                ):
                    try:
                        self.model(calib_inputs["input_ids"])
                        forward_success = True
                    except Exception as e3:
                        logger.debug("Bonsai calibration strategy 3 failed: %s", e3)

                if not forward_success:
                    try:
                        dummy = torch.randint(0, 100, (1, 16), device=self.target_device)
                        self.model(dummy)
                    except Exception as e4:
                        logger.debug("Bonsai synthetic calibration failed: %s", e4)
        except Exception as e:
            logger.warning("Bonsai calibration exception: %s", e)
        finally:
            for h in hooks:
                h.remove()

        aggregated: dict[str, torch.Tensor] = {}
        for name, act_list in activations.items():
            if act_list:
                aggregated[name] = torch.cat([a.reshape(-1, a.shape[-1]) for a in act_list], dim=0)
        return aggregated

    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        try:
            activations_map = self._collect_activations()
        except Exception as e:
            logger.warning(
                "BonsaiPruner: activation collection failed (%s); proceeding with unweighted pruning.",
                e,
            )
            activations_map = {}

        self.pruned_layers = []
        self.protected_layers = []
        self.skipped_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_sparsities = {}

        for name, module in list(self.model.named_modules()):
            cls_name = module.__class__.__name__
            is_linear_like = (
                isinstance(module, nn.Linear)
                or "conv1d" in cls_name.lower()
                or hasattr(module, "nf")
            )
            if not is_linear_like or isinstance(module, BonsaiPrunedLoRALinear):
                continue

            if self._is_protected_layer(name, module):
                for p in module.parameters():
                    p.requires_grad = False
                self.protected_layers.append(name)
                continue

            target_mod = getattr(module, "base_layer", module)
            raw_w = target_mod.weight.data
            dtype = raw_w.dtype
            device = raw_w.device
            is_conv1d = "conv1d" in cls_name.lower() or hasattr(module, "nf")

            W0 = raw_w.t().float().contiguous() if is_conv1d else raw_w.float().contiguous()
            out_f, in_f = W0.shape

            if (out_f * in_f) < self.min_numel:
                for p in module.parameters():
                    p.requires_grad = False
                self.skipped_layers.append(name)
                continue

            W = W0.clone()
            act = activations_map.get(name)
            X = None
            if act is not None and act.numel() > 0:
                X = act.to(device).float()
                if X.shape[-1] != in_f:
                    X = None

            # BN-FIX-10: bounded adaptive sparsity.
            if self.adaptive_sparsity and X is not None:
                act_norm = torch.norm(X, p="fro").item() / max(1, X.shape[0])
                scale_factor = 1.0 / (1.0 + math.log1p(max(0.0, act_norm)))
                target_sp = max(0.05, min(0.95, self.sparsity * (0.7 + 0.6 * scale_factor)))
            else:
                target_sp = self.sparsity

            final_mask = torch.ones_like(W)

            # Iterative pruning loop.
            for s in range(self.steps):
                t_ratio = (s + 1.0) / self.steps
                current_sp = target_sp * (1.0 - (1.0 - t_ratio) ** 3)
                saliency = compute_taylor_first_order_saliency(W, activations=X)

                if self.structured_24:
                    W_pruned, final_mask = apply_24_structured_mask(W, saliency=saliency, dim=1)
                elif self.block_size != (1, 1):
                    W_pruned, final_mask = apply_block_mask(
                        W,
                        block_size=self.block_size,
                        sparsity=current_sp,
                        saliency=saliency,
                    )
                else:
                    flat = saliency.reshape(-1)
                    k = max(1, int(flat.numel() * (1.0 - current_sp)))
                    thr = torch.topk(flat, k, largest=True).values[-1]
                    final_mask = (saliency >= thr).float()
                    W_pruned = W * final_mask

                # BN-FIX-04: explicit clone.
                W = W_pruned.clone()

                # Gradient descent recovery between steps.
                if self.recovery_steps > 0 and X is not None:
                    cov_X = torch.matmul(X.t(), X) / max(X.shape[0], 1)
                    W0_cov = torch.matmul(W0, cov_X)

                    # BN-FIX-06: clamp step size.
                    raw_ss = self.recovery_lr / (torch.norm(cov_X).item() + 1e-6)
                    step_size = min(raw_ss, self.max_step_size)
                    for _ in range(self.recovery_steps):
                        curr_cov = torch.matmul(W, cov_X)
                        grad = (curr_cov - W0_cov) * final_mask
                        W = W - step_size * grad
                        W = W * final_mask

            norm_orig = torch.norm(W0, p="fro").clamp_min(1e-8)
            rel_err = (torch.norm(W0 - W, p="fro") / norm_orig).item()
            self.layer_errors[name] = float(rel_err)
            self.layer_sparsities[name] = float(target_sp)

            is_invalid = (not torch.isfinite(W).all()) or (
                rel_err > (self.max_error or float("inf"))
            )
            if is_invalid:
                logger.warning(
                    "BonsaiPruner: %s error %.4f (threshold %.4f).",
                    name,
                    rel_err,
                    self.max_error if self.max_error is not None else float("inf"),
                )
                if self.fallback_to_original:
                    for p in module.parameters():
                        p.requires_grad = False
                    self.fallback_layers.append(name)
                    continue

            # Commit layer.
            if self.use_lora and self.lora_r > 0:
                pruned_lora_layer = BonsaiPrunedLoRALinear.from_linear(
                    linear=target_mod,
                    pruned_weight=W.to(dtype=dtype, device=device),
                    mask=final_mask.to(device=device),
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                    freeze_base=True,
                )
                # BN-FIX-07: copy LoRA data (no shared reference).
                if (
                    hasattr(module, "lora_A")
                    and module.lora_A is not None
                    and pruned_lora_layer.lora_A is not None
                ):
                    try:
                        src_A = getattr(module.lora_A, "weight", module.lora_A)
                        pruned_lora_layer.lora_A.data.copy_(src_A.data)
                    except Exception as e:
                        logger.debug("Bonsai LoRA A transfer failed: %s", e)
                if (
                    hasattr(module, "lora_B")
                    and module.lora_B is not None
                    and pruned_lora_layer.lora_B is not None
                ):
                    try:
                        src_B = getattr(module.lora_B, "weight", module.lora_B)
                        pruned_lora_layer.lora_B.data.copy_(src_B.data)
                    except Exception as e:
                        logger.debug("Bonsai LoRA B transfer failed: %s", e)

                if hasattr(module, "base_layer"):
                    module.base_layer = pruned_lora_layer
                    pruned_lora_layer = module
                replace_module(self.model, name, pruned_lora_layer)
            else:
                # BN-FIX-08: no_grad for weight mutation.
                with torch.no_grad():
                    if is_conv1d:
                        target_mod.weight.data.copy_(
                            W.t().contiguous().to(dtype=dtype, device=device)
                        )
                    else:
                        target_mod.weight.data.copy_(W.contiguous().to(dtype=dtype, device=device))

            self.pruned_layers.append(name)

        # LayerNorm recalibration.
        if self.recalibrate_ln:
            try:
                recalibrate_layernorms(
                    self.model,
                    tokenizer=self.tokenizer,
                    calibration_loader=self.calibration_loader,
                    calibration_texts=self.calibration_texts,
                    device=self.target_device,
                )
            except Exception as e:
                logger.warning("BonsaiPruner: LayerNorm recalibration skipped: %s", e)

        # BN-FIX-09: forward sanity check with attention_mask.
        sanity_check_passed = True
        try:
            with torch.no_grad():
                dummy = torch.randint(0, 100, (1, 8), device=self.target_device, dtype=torch.long)
                attn = torch.ones_like(dummy)
                try:
                    self.model(input_ids=dummy, attention_mask=attn)
                except Exception:
                    self.model(dummy)
        except Exception as ex:
            logger.debug("BonsaiPruner post-compression forward sanity check warning: %s", ex)
            sanity_check_passed = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        avg_err = (
            (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
        )
        stats.update(
            {
                "pruned_layers": self.pruned_layers,
                "pruned_layer_count": len(self.pruned_layers),
                "protected_layers": self.protected_layers,
                "protected_layer_count": len(self.protected_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "layer_errors": self.layer_errors,
                "layer_sparsities": self.layer_sparsities,
                "avg_reconstruction_error": avg_err,
                "sanity_check_passed": sanity_check_passed,
                "steps": self.steps,
                "recovery_steps": self.recovery_steps,
                "sparsity_target": self.sparsity,
                "structured_24": self.structured_24,
                "protect_cls": self.protect_cls,
                "skip_attention_qkv": self.skip_attention_qkv,
                "use_lora": self.use_lora,
                "lora_r": self.lora_r,
                "adaptive_sparsity": self.adaptive_sparsity,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "method": "bonsai_iterative_pruning_with_gd_recovery",
            }
        )
        self.stats = stats
        logger.info(
            "Compressor bonsai finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.pruned_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error if self.max_error is not None else 1.0,
        )
        return self.model, stats

Bonsai Pruner: progressive multi-step iterative pruning with cubic schedule, inter-step gradient descent recovery, and [CLS]/GQA protection.

Ancestors

BaseCompressor

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class COMPOTCompressor (model: nn.Module | None = None, rank: int = 64, max_error: float = 0.2, fallback_to_original: bool = True, calibration_texts: list[str] | None = None, calibration_loader: Any | None = None, tokenizer: Any | None = None, cfg: Any | None = None, target_attention_only: bool = False, target_layers: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("compot_compressor_class")
class COMPOTCompressor(BaseCompressor):
    """
    COMPOT (Compressed Orthogonal Dictionary Learning) Compressor.
    Learns a low-rank basis from activation QR decompositions and compresses
    linear and Conv1D layers into factorized COMPOTLinear modules.

    Production hardening features:
    - Relative Frobenius reconstruction error validation with configurable max_error.
    - Safe fallback to original layer if error exceeds threshold or decomposition is unstable.
    - SVD fallback if QR decomposition on activations fails.
    - Parameter reduction guarantee: checks r * (in + out) < in * out.
    - Robust Conv1D support (GPT-2 c_attn, c_proj) via weight transposition.
    - Unified critical-layer protection via should_compress_linear.
    - Comprehensive per-layer stats and error reporting.
    """

    def __init__(
        self,
        model: nn.Module | None = None,
        rank: int = 64,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        calibration_texts: list[str] | None = None,
        calibration_loader: Any | None = None,
        tokenizer: Any | None = None,
        cfg: Any | None = None,
        target_attention_only: bool = False,
        target_layers: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model if model is not None else nn.Identity(), name="compot", **kwargs)
        conf = cfg or kwargs.get("comp_config")
        if conf is not None:
            rank = getattr(conf, "compot_rank", getattr(conf, "rank", rank))
            max_error = getattr(conf, "compot_max_error", getattr(conf, "max_error", max_error))
            fallback_to_original = getattr(
                conf,
                "compot_fallback_to_original",
                getattr(conf, "fallback_to_original", fallback_to_original),
            )

        self.rank = int(rank)
        if max_error is None or max_error == 0.20:
            max_error = adaptive_max_error(self.rank, 768, 768, "distilbert")

        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.calibration_texts = calibration_texts
        self.calibration_loader = calibration_loader
        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.target_layers = target_layers
        self.target_attention_only = target_attention_only

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.used_activations: dict[str, bool] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        cls_name = module.__class__.__name__.lower()
        is_linear = isinstance(module, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")
        if not (is_linear or is_conv1d):
            return False
        if isinstance(module, COMPOTLinear):
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self.target_attention_only and not any(
            k in name.lower()
            for k in [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "out_proj",
                "o_proj",
                "q_lin",
                "k_lin",
                "v_lin",
                "out_lin",
                "query",
                "key",
                "value",
                "c_attn",
                "c_proj",
            ]
        ):
            return False
        if self.target_layers is not None:
            return any(t in name for t in self.target_layers)
        return True

    def _collect_activations(self) -> dict[str, torch.Tensor]:
        """Collect input activations across target modules."""
        layer_data: dict[str, torch.Tensor] = {}
        hooks = []

        def make_hook(name: str):
            def hook(module, input, output):
                if len(input) > 0 and isinstance(input[0], torch.Tensor):
                    inp = input[0].detach().float()
                    cls_name = module.__class__.__name__.lower()
                    is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")
                    in_f = module.weight.shape[0] if is_conv1d else module.weight.shape[1]
                    flat = inp.reshape(-1, in_f)
                    if name not in layer_data:
                        layer_data[name] = flat.cpu()
                    else:
                        layer_data[name] = torch.cat([layer_data[name], flat.cpu()], dim=0)

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                hooks.append(module.register_forward_hook(make_hook(name)))

        device = next(self.model.parameters(), torch.tensor(0.0)).device
        calib_inputs = get_calibration_tensors(
            self.model,
            tokenizer=self.tokenizer,
            calibration_loader=self.calibration_loader,
            calibration_texts=self.calibration_texts,
            device=device,
        )

        self.model.eval()
        with torch.no_grad():
            try:
                if isinstance(calib_inputs, dict):
                    self.model(**calib_inputs)
                elif isinstance(calib_inputs, torch.Tensor):
                    self.model(calib_inputs)
                elif isinstance(calib_inputs, (list, tuple)):
                    self.model(*calib_inputs)
            except Exception as e:
                logger.debug("COMPOT activation collection forward pass exception: %s", e)
            finally:
                for h in hooks:
                    h.remove()

        return layer_data

    def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
        """Compress linear and Conv1D layers via COMPOT orthogonal dictionary learning."""
        if model is not None:
            self.model = model
            self._original_params = self._count_params(self.model)

        activations = self._collect_activations()
        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}
        self.used_activations = {}

        targets = [
            (name, mod)
            for name, mod in list(self.model.named_modules())
            if self._is_target_module(name, mod)
        ]

        if not targets:
            logger.warning("COMPOT: No compressible layers found.")
            return self.model, self.get_compression_stats()

        for name, module in targets:
            cls_name = module.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")
            target_mod = getattr(module, "base_layer", module)
            raw_w = target_mod.weight.data
            if is_conv1d:
                in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            else:
                out_f, in_f = raw_w.shape[0], raw_w.shape[1]

            # Guarantee parameter reduction: r * (in_f + out_f) < in_f * out_f
            max_reduction_r = max(1, (in_f * out_f) // (in_f + out_f) - 1)
            max_r = max(1, min(in_f - 1, out_f - 1, max_reduction_r))
            r = max(1, min(self.rank, max_r))

            act = activations.get(name, None)
            compot_layer, err, used_act = COMPOTLinear.from_linear_and_activations(
                target_mod, activations=act, rank=r
            )
            self.layer_errors[name] = err
            self.layer_ranks[name] = r
            self.used_activations[name] = used_act

            # Validate dimensions
            if compot_layer.in_features != in_f or compot_layer.out_features != out_f:
                logger.warning(
                    "Dimension mismatch on %s: (%d, %d) vs (%d, %d)",
                    name,
                    compot_layer.out_features,
                    compot_layer.in_features,
                    out_f,
                    in_f,
                )
                for p in module.parameters():
                    p.requires_grad = False
                self.fallback_layers.append(name)
                continue

            if err > self.max_error:
                logger.warning(
                    "COMPOT reconstruction error %.4f on %s exceeded threshold %.4f.",
                    err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    logger.info(
                        "Skipping COMPOT compression for %s (fallback to original).",
                        name,
                    )
                    for p in module.parameters():
                        p.requires_grad = False
                    self.fallback_layers.append(name)
                    continue

            if hasattr(module, "base_layer"):
                if hasattr(module, "lora_A"):
                    compot_layer.lora_A = module.lora_A
                if hasattr(module, "lora_B"):
                    compot_layer.lora_B = module.lora_B
                module.base_layer = compot_layer
                compot_layer = module
            replace_module(self.model, name, compot_layer)
            self.compressed_layers.append(name)

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "used_activations": self.used_activations,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "avg_reconstruction_error": avg_err,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "compot_rank": self.rank,
            }
        )
        logger.info(
            "Compressor compot finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

COMPOT (Compressed Orthogonal Dictionary Learning) Compressor. Learns a low-rank basis from activation QR decompositions and compresses linear and Conv1D layers into factorized COMPOTLinear modules.

Production hardening features: - Relative Frobenius reconstruction error validation with configurable max_error. - Safe fallback to original layer if error exceeds threshold or decomposition is unstable. - SVD fallback if QR decomposition on activations fails. - Parameter reduction guarantee: checks r * (in + out) < in * out. - Robust Conv1D support (GPT-2 c_attn, c_proj) via weight transposition. - Unified critical-layer protection via should_compress_linear. - Comprehensive per-layer stats and error reporting.

Ancestors

BaseCompressor

Methods

def compress(self, model: nn.Module | None = None) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
    """Compress linear and Conv1D layers via COMPOT orthogonal dictionary learning."""
    if model is not None:
        self.model = model
        self._original_params = self._count_params(self.model)

    activations = self._collect_activations()
    self.compressed_layers = []
    self.fallback_layers = []
    self.layer_errors = {}
    self.layer_ranks = {}
    self.used_activations = {}

    targets = [
        (name, mod)
        for name, mod in list(self.model.named_modules())
        if self._is_target_module(name, mod)
    ]

    if not targets:
        logger.warning("COMPOT: No compressible layers found.")
        return self.model, self.get_compression_stats()

    for name, module in targets:
        cls_name = module.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")
        target_mod = getattr(module, "base_layer", module)
        raw_w = target_mod.weight.data
        if is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]

        # Guarantee parameter reduction: r * (in_f + out_f) < in_f * out_f
        max_reduction_r = max(1, (in_f * out_f) // (in_f + out_f) - 1)
        max_r = max(1, min(in_f - 1, out_f - 1, max_reduction_r))
        r = max(1, min(self.rank, max_r))

        act = activations.get(name, None)
        compot_layer, err, used_act = COMPOTLinear.from_linear_and_activations(
            target_mod, activations=act, rank=r
        )
        self.layer_errors[name] = err
        self.layer_ranks[name] = r
        self.used_activations[name] = used_act

        # Validate dimensions
        if compot_layer.in_features != in_f or compot_layer.out_features != out_f:
            logger.warning(
                "Dimension mismatch on %s: (%d, %d) vs (%d, %d)",
                name,
                compot_layer.out_features,
                compot_layer.in_features,
                out_f,
                in_f,
            )
            for p in module.parameters():
                p.requires_grad = False
            self.fallback_layers.append(name)
            continue

        if err > self.max_error:
            logger.warning(
                "COMPOT reconstruction error %.4f on %s exceeded threshold %.4f.",
                err,
                name,
                self.max_error,
            )
            if self.fallback_to_original:
                logger.info(
                    "Skipping COMPOT compression for %s (fallback to original).",
                    name,
                )
                for p in module.parameters():
                    p.requires_grad = False
                self.fallback_layers.append(name)
                continue

        if hasattr(module, "base_layer"):
            if hasattr(module, "lora_A"):
                compot_layer.lora_A = module.lora_A
            if hasattr(module, "lora_B"):
                compot_layer.lora_B = module.lora_B
            module.base_layer = compot_layer
            compot_layer = module
        replace_module(self.model, name, compot_layer)
        self.compressed_layers.append(name)

    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )
    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "fallback_layers": self.fallback_layers,
            "layer_errors": self.layer_errors,
            "layer_ranks": self.layer_ranks,
            "used_activations": self.used_activations,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "avg_reconstruction_error": avg_err,
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "compot_rank": self.rank,
        }
    )
    logger.info(
        "Compressor compot finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        self.max_error,
    )
    return self.model, stats

Compress linear and Conv1D layers via COMPOT orthogonal dictionary learning.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class COMPOTLinear (in_features: int, out_features: int, rank: int = 64, bias: bool = True, dtype: torch.dtype | None = None, device: torch.device | None = None)

Expand source code

class COMPOTLinear(nn.Module):
    is_custom_attention: bool = True
    """
    Factorized Linear Layer constructed via Orthogonal Dictionary Learning (COMPOT):
    W ≈ B @ A
    - A in R^{rank x in_features}: input-orthogonal projection matrix
    - B in R^{out_features x rank}: output dictionary matrix

    Forward pass:
    y = ((x @ A^T) @ B^T) + bias
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 64,
        bias: bool = True,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = max(1, min(rank, in_features, out_features))

        self.A = nn.Parameter(torch.empty(self.rank, in_features, dtype=dtype, device=device))
        self.B = nn.Parameter(torch.empty(out_features, self.rank, dtype=dtype, device=device))

        with torch.no_grad():
            A_f32 = torch.empty(self.rank, in_features, dtype=torch.float32)
            nn.init.orthogonal_(A_f32)
            self.A.copy_(A_f32.to(device=device, dtype=dtype))
            nn.init.xavier_uniform_(self.B)

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix (out_features, in_features)."""
        return torch.matmul(self.B, self.A)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass matching standard Linear and Conv1D."""
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)

        x_a = F.linear(x_flat, self.A)
        out = F.linear(x_a, self.B, self.bias)

        return out.view(*orig_shape[:-1], self.out_features)

    def get_weight(self) -> torch.Tensor:
        """Materialize full weight matrix W = B @ A."""
        return torch.matmul(self.B, self.A)

    def __iter__(self):
        """Allow tuple unpacking: compot_layer, err, used_act = COMPOTLinear.from_linear_and_activations(...)."""
        yield self
        yield getattr(self, "reconstruction_error", 0.0)
        yield getattr(self, "used_activations", False)

    def __getitem__(self, idx: int) -> Any:
        return (
            self,
            getattr(self, "reconstruction_error", 0.0),
            getattr(self, "used_activations", False),
        )[idx]

    @classmethod
    def from_linear_and_activations(
        cls,
        layer: nn.Module,
        activations: torch.Tensor | None = None,
        rank: int = 64,
    ) -> tuple[COMPOTLinear, float, bool]:
        """
        Construct COMPOTLinear by learning orthogonal dictionary from activations or weights.
        Returns: (compot_layer, relative_reconstruction_error, used_activations)
        """
        cls_name = layer.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(layer, "nf")
        raw_w = layer.weight.data

        if is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W_float = raw_w.t().contiguous().float()  # [out_features, in_features]
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            W_float = raw_w.float()  # [out_features, in_features]

        bias = layer.bias.data if layer.bias is not None else None
        r = max(1, min(rank, out_f, in_f))
        device = raw_w.device
        dtype = raw_w.dtype

        compot_layer = cls(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=bias is not None,
            dtype=dtype,
            device=device,
        )

        used_act = False
        A_data = None
        B_data = None

        if activations is not None and activations.size(0) >= 4:
            X = activations.float().to(device)
            if X.shape[-1] == in_f:
                try:
                    # QR decomposition of activations: X = Q @ R
                    Q, R = torch.linalg.qr(X, mode="reduced")
                    D = torch.matmul(W_float, R.t().contiguous())
                    U, S, Vt = torch.linalg.svd(D, full_matrices=False)

                    B_cand = U[:, :r] * S[:r].unsqueeze(0)  # (out_f, r)
                    V_r = Vt[:r, :]  # (r, in_f)
                    R_inv_T = torch.linalg.pinv(R.t().contiguous())
                    A_cand = torch.matmul(V_r, R_inv_T)

                    if not (
                        torch.isnan(A_cand).any()
                        or torch.isnan(B_cand).any()
                        or torch.isinf(A_cand).any()
                        or torch.isinf(B_cand).any()
                    ):
                        A_data = A_cand
                        B_data = B_cand
                        used_act = True
                except Exception as e:
                    logger.debug("COMPOT QR decomposition exception on layer: %s", e)

        # Fallback to direct SVD on weight matrix
        if A_data is None or B_data is None:
            U, S, Vt = torch.linalg.svd(W_float, full_matrices=False)
            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
            B_data = U[:, :r] * sqrt_S.unsqueeze(0)
            A_data = sqrt_S.unsqueeze(1) * Vt[:r, :]
            used_act = False

        compot_layer.A.data.copy_(A_data.to(dtype=dtype, device=device))
        compot_layer.B.data.copy_(B_data.to(dtype=dtype, device=device))

        if bias is not None and compot_layer.bias is not None:
            compot_layer.bias.data.copy_(bias)

        # Compute relative Frobenius reconstruction error
        W_recon = torch.matmul(B_data, A_data)
        norm_orig = torch.norm(W_float, p="fro").item()
        if norm_orig < 1e-12:
            rel_err = 0.0
        else:
            rel_err = float(torch.norm(W_float - W_recon, p="fro").item() / norm_orig)

        compot_layer.reconstruction_error = rel_err
        compot_layer.used_activations = used_act
        return compot_layer

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: Factorized Linear Layer constructed via Orthogonal Dictionary Learning (COMPOT): W ≈ B @ A - A in R^{rank x in_features}: input-orthogonal projection matrix - B in R^{out_features x rank}: output dictionary matrix

Forward pass: y = ((x @ A^T) @ B^T) + bias

Static methods

def from_linear_and_activations(layer: nn.Module, activations: torch.Tensor | None = None, rank: int = 64) ‑> tuple[COMPOTLinear, float, bool]: Construct COMPOTLinear by learning orthogonal dictionary from activations or weights. Returns: (compot_layer, relative_reconstruction_error, used_activations)

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix (out_features, in_features)."""
    return torch.matmul(self.B, self.A)

Reconstructed effective weight matrix (out_features, in_features).

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass matching standard Linear and Conv1D."""
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)

    x_a = F.linear(x_flat, self.A)
    out = F.linear(x_a, self.B, self.bias)

    return out.view(*orig_shape[:-1], self.out_features)

Forward pass matching standard Linear and Conv1D.

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    """Materialize full weight matrix W = B @ A."""
    return torch.matmul(self.B, self.A)

Materialize full weight matrix W = B @ A.

class CPTGenerator (hidden_dim: int = 128, num_freqs: int = 8, vocab_size: int = 5000)

Expand source code

class CPTGenerator(nn.Module):
    """
    Learnable Task-Aware Compression Planning Token (CPT) Generator.
    Maps arbitrary task descriptions to continuous per-feature importance vectors for weight SVD
    and KV cache compression across arbitrary dimension sizes d.
    """

    def __init__(self, hidden_dim: int = 128, num_freqs: int = 8, vocab_size: int = 5000):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_freqs = num_freqs
        self.vocab_size = vocab_size

        self.text_embedding = nn.EmbeddingBag(vocab_size, hidden_dim, mode="mean")
        self.encoder = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, hidden_dim),
        )

        in_dim = hidden_dim + num_freqs * 2
        self.projector = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, 1),
            nn.Sigmoid(),
        )

    def _tokenize_text(self, text: str, device: torch.device) -> torch.Tensor:
        """Deterministic hash-based tokenization without external library dependencies."""
        words = (
            str(text)
            .lower()
            .replace(",", " ")
            .replace(".", " ")
            .replace(":", " ")
            .replace("-", " ")
            .split()
        )
        if not words:
            return torch.tensor([0], dtype=torch.long, device=device)
        indices = []
        for w in words:
            h = int(hashlib.md5(w.encode("utf-8")).hexdigest(), 16)
            indices.append(h % (self.vocab_size - 1) + 1)
        return torch.tensor(indices, dtype=torch.long, device=device)

    def forward(self, task: str, d: int, device: torch.device | None = None) -> torch.Tensor:
        """
        Generate a per-feature importance modulation vector of size d in [0.1, 2.0].
        """
        target_device = device if device is not None else next(self.parameters()).device
        tokens = self._tokenize_text(task, target_device)
        h = self.text_embedding(tokens.unsqueeze(0)).squeeze(0)  # [hidden_dim]
        h = self.encoder(h)  # [hidden_dim]

        # Continuous NeRF-style positional coordinate encoding
        coords = torch.linspace(0.0, 1.0, d, device=target_device)
        freqs = (2.0 ** torch.arange(self.num_freqs, device=target_device)) * math.pi
        args = coords.unsqueeze(1) * freqs.unsqueeze(0)  # [d, num_freqs]
        enc = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)  # [d, 2 * num_freqs]

        h_exp = h.unsqueeze(0).expand(d, -1)  # [d, hidden_dim]
        inp = torch.cat([h_exp, enc], dim=-1)  # [d, in_dim]

        # Importance modulation strictly in [0.1, 2.0]
        importance = self.projector(inp).squeeze(-1) * 1.9 + 0.1
        return importance

Learnable Task-Aware Compression Planning Token (CPT) Generator. Maps arbitrary task descriptions to continuous per-feature importance vectors for weight SVD and KV cache compression across arbitrary dimension sizes d.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Methods

def forward(self, task: str, d: int, device: torch.device | None = None) ‑> torch.Tensor

Expand source code

def forward(self, task: str, d: int, device: torch.device | None = None) -> torch.Tensor:
    """
    Generate a per-feature importance modulation vector of size d in [0.1, 2.0].
    """
    target_device = device if device is not None else next(self.parameters()).device
    tokens = self._tokenize_text(task, target_device)
    h = self.text_embedding(tokens.unsqueeze(0)).squeeze(0)  # [hidden_dim]
    h = self.encoder(h)  # [hidden_dim]

    # Continuous NeRF-style positional coordinate encoding
    coords = torch.linspace(0.0, 1.0, d, device=target_device)
    freqs = (2.0 ** torch.arange(self.num_freqs, device=target_device)) * math.pi
    args = coords.unsqueeze(1) * freqs.unsqueeze(0)  # [d, num_freqs]
    enc = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)  # [d, 2 * num_freqs]

    h_exp = h.unsqueeze(0).expand(d, -1)  # [d, hidden_dim]
    inp = torch.cat([h_exp, enc], dim=-1)  # [d, in_dim]

    # Importance modulation strictly in [0.1, 2.0]
    importance = self.projector(inp).squeeze(-1) * 1.9 + 0.1
    return importance

Generate a per-feature importance modulation vector of size d in [0.1, 2.0].

class CRCompressor (model: nn.Module | None = None, rank: int = 64, max_error: float = 0.2, fallback_to_original: bool = True, max_retries: int = 2, target_attention_only: bool = False, cfg: Any | None = None, target_layers: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("cr_compressor")
class CRCompressor(BaseCompressor):
    """
    CR (Column-Row / Interpolative Decomposition) Compressor.
    Decomposes linear layers into skeleton column bases and interpolative projections.

    Args:
        model: PyTorch model.
        rank: Rank parameter (number of skeleton columns).
        max_error: Maximum relative Frobenius reconstruction error (default 0.20).
        fallback_to_original: If True, keep original layer when error exceeds max_error (default True).
        max_retries: Number of rank increase retries before fallback (default 2).
        target_attention_only: If True, restrict compression to attention projections only.
        cfg: Optional configuration object.
        target_layers: Optional target layer filter.
    """

    def __init__(
        self,
        model: nn.Module | None = None,
        rank: int = 64,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        max_retries: int = 2,
        target_attention_only: bool = False,
        cfg: Any | None = None,
        target_layers: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model if model is not None else nn.Identity(), name="cr", **kwargs)
        c_rank = _cfg_get(cfg, "cr_rank", _cfg_get(cfg, "rank", rank))
        c_max_error = _cfg_get(cfg, "cr_max_error", _cfg_get(cfg, "max_error", max_error))
        c_fallback = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
        c_target_attn = _cfg_get(
            cfg,
            "cr_target_attention_only",
            _cfg_get(cfg, "target_attention_only", target_attention_only),
        )

        if c_max_error is None or c_max_error == 0.20:
            c_max_error = adaptive_max_error(c_rank, 768, 768, "distilbert")

        self.rank = int(c_rank) if c_rank is not None else 64
        self.max_error = float(c_max_error)
        self.fallback_to_original = bool(c_fallback)
        self.max_retries = max(0, int(max_retries))
        self.target_attention_only = bool(c_target_attn)
        self.cfg = cfg
        self.target_layers = target_layers or _cfg_get(cfg, "target_layers", None)
        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        """Filter candidate layers based on classification head exclusion and target_attention_only."""
        cls_name = module.__class__.__name__.lower()
        has_base = hasattr(module, "base_layer")
        base_cls = module.base_layer.__class__.__name__.lower() if has_base else ""

        is_linear = isinstance(module, nn.Linear) or "linear" in cls_name
        is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")
        is_lora = has_base and (
            isinstance(module.base_layer, nn.Linear)
            or "conv1d" in base_cls
            or hasattr(module.base_layer, "nf")
        )

        if not (is_linear or is_conv1d or is_lora):
            return False

        if isinstance(module, CRLinear) or (has_base and isinstance(module.base_layer, CRLinear)):
            return False

        name_lower = name.lower()
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embed",
                "logits_processor",
                "pre_classifier",
                "head",
            ]
        ):
            return False

        if self.target_attention_only:
            ffn_patterns = [
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            ]
            if any(p in name_lower for p in ffn_patterns):
                return False
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            attn_patterns = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
                "in_proj_qkv",
                "out_proj",
            ]
            if not any(p in name_lower for p in attn_patterns):
                return False

        if self.target_layers and not any(t in name for t in self.target_layers):
            return False

        return True

    def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
        """Compress linear layers using interpolative decomposition."""
        if model is not None:
            self.model = model
            self._original_params = self._count_params(self.model)

        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}

        total_orig_decomp_params = 0
        total_comp_decomp_params = 0

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                if self.target_attention_only:
                    is_lin_like = (
                        isinstance(module, nn.Linear)
                        or "conv1d" in module.__class__.__name__.lower()
                        or hasattr(module, "base_layer")
                    )
                    if is_lin_like:
                        for p in module.parameters():
                            p.requires_grad = False
                continue

            target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
            cls_name = module.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(module, "nf") or hasattr(target_mod, "nf")
            (hasattr(module, "lora_A") and module.lora_A is not None) or hasattr(
                module, "base_layer"
            )

            raw_w = target_mod.weight.data
            orig_w = raw_w.t().float() if is_conv1d else raw_w.float()
            out_f, in_f = orig_w.shape
            max_rank = min(out_f, in_f)

            if max_rank < 16:
                continue

            # Parameter reduction check: r * (in_f + out_f) must be less than in_f * out_f
            target_r = min(self.rank, max_rank)
            if target_r * (in_f + out_f) >= 0.85 * (in_f * out_f):
                target_r = max(1, int(0.65 * (in_f * out_f) / (in_f + out_f)))

            cur_r = target_r
            best_cr_layer = None
            best_err = float("inf")
            best_r = cur_r
            norm_orig = torch.norm(orig_w, p="fro").clamp_min(1e-8)

            for _ in range(self.max_retries + 1):
                try:
                    candidate = CRLinear.from_linear(target_mod, rank=cur_r)
                    recon_w = candidate.get_weight().float()
                    err = (torch.norm(orig_w - recon_w, p="fro") / norm_orig).item()
                    if err < best_err:
                        best_err = err
                        best_cr_layer = candidate
                        best_r = cur_r

                    if err <= self.max_error:
                        break

                    next_r = min(int(cur_r * 1.5), max_rank - 1)
                    if next_r <= cur_r or next_r * (in_f + out_f) >= 0.85 * (in_f * out_f):
                        break
                    cur_r = next_r
                except Exception as ex:
                    logger.debug("CR decomposition error on %s with rank %d: %s", name, cur_r, ex)
                    break

            if best_err > self.max_error:
                logger.warning(
                    "CRCompressor: Layer %s error %.4f exceeded max_error %.4f",
                    name,
                    best_err,
                    self.max_error,
                )
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    # Freeze fallback layers during recovery
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

            if best_cr_layer is not None:
                new_layer = best_cr_layer
                compressed_layer = new_layer
                # Preserve existing LoRA adapter weights if the source module has them
                if (
                    hasattr(module, "lora_A")
                    and hasattr(module, "lora_B")
                    and module.lora_A is not None
                    and module.lora_B is not None
                ):
                    src_lA = getattr(module.lora_A, "weight", module.lora_A)
                    src_lB = getattr(module.lora_B, "weight", module.lora_B)
                    if hasattr(new_layer, "lora_A") and new_layer.lora_A is not None:
                        if hasattr(src_lA, "data") and src_lA.shape == new_layer.lora_A.shape:
                            new_layer.lora_A.data.copy_(src_lA.data)
                    if hasattr(new_layer, "lora_B") and new_layer.lora_B is not None:
                        if hasattr(src_lB, "data") and src_lB.shape == new_layer.lora_B.shape:
                            new_layer.lora_B.data.copy_(src_lB.data)

                if hasattr(module, "base_layer"):
                    module.base_layer = compressed_layer  # put compressed layer inside wrapper
                    # module itself (LoRALinear wrapper) remains the replacement
                    new_module = module
                else:
                    new_module = compressed_layer

                if hasattr(module, "base_layer"):
                    for p_name, p in module.named_parameters():
                        if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                            p.requires_grad = True
                else:
                    if name == "":
                        self.model = new_module
                    else:
                        replace_module(self.model, name, new_module)

                self.compressed_layers.append(name)
                self.layer_errors[name] = float(best_err)
                self.layer_ranks[name] = int(best_r)
                total_orig_decomp_params += in_f * out_f
                total_comp_decomp_params += best_r * (in_f + out_f)

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        avg_err = (
            (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
        )
        effective_param_ratio = (
            total_comp_decomp_params / total_orig_decomp_params
            if total_orig_decomp_params > 0
            else 1.0
        )
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "avg_reconstruction_error": avg_err,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "cr_rank": self.rank,
                "target_attention_only": self.target_attention_only,
                "effective_parameter_ratio": round(effective_param_ratio, 4),
                "used_interpolative_decomposition_count": len(self.compressed_layers),
            }
        )
        logger.info(
            "Compressor cr finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

    def compress_linear(self, linear: nn.Linear | Any) -> CRLinear:
        """Compress a single linear layer into CRLinear."""
        return CRLinear.from_linear(linear, rank=self.rank)

CR (Column-Row / Interpolative Decomposition) Compressor. Decomposes linear layers into skeleton column bases and interpolative projections.

Args
-----=
model: PyTorch model.
rank: Rank parameter (number of skeleton columns).
max_error: Maximum relative Frobenius reconstruction error (default 0.20).
fallback_to_original: If True, keep original layer when error exceeds max_error (default True).
max_retries: Number of rank increase retries before fallback (default 2).
target_attention_only: If True, restrict compression to attention projections only.
cfg: Optional configuration object.
target_layers: Optional target layer filter.

Ancestors

BaseCompressor

Methods

def compress(self, model: nn.Module | None = None) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
    """Compress linear layers using interpolative decomposition."""
    if model is not None:
        self.model = model
        self._original_params = self._count_params(self.model)

    self.compressed_layers = []
    self.fallback_layers = []
    self.layer_errors = {}
    self.layer_ranks = {}

    total_orig_decomp_params = 0
    total_comp_decomp_params = 0

    for name, module in list(self.model.named_modules()):
        if not self._is_target_module(name, module):
            if self.target_attention_only:
                is_lin_like = (
                    isinstance(module, nn.Linear)
                    or "conv1d" in module.__class__.__name__.lower()
                    or hasattr(module, "base_layer")
                )
                if is_lin_like:
                    for p in module.parameters():
                        p.requires_grad = False
            continue

        target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
        cls_name = module.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(module, "nf") or hasattr(target_mod, "nf")
        (hasattr(module, "lora_A") and module.lora_A is not None) or hasattr(
            module, "base_layer"
        )

        raw_w = target_mod.weight.data
        orig_w = raw_w.t().float() if is_conv1d else raw_w.float()
        out_f, in_f = orig_w.shape
        max_rank = min(out_f, in_f)

        if max_rank < 16:
            continue

        # Parameter reduction check: r * (in_f + out_f) must be less than in_f * out_f
        target_r = min(self.rank, max_rank)
        if target_r * (in_f + out_f) >= 0.85 * (in_f * out_f):
            target_r = max(1, int(0.65 * (in_f * out_f) / (in_f + out_f)))

        cur_r = target_r
        best_cr_layer = None
        best_err = float("inf")
        best_r = cur_r
        norm_orig = torch.norm(orig_w, p="fro").clamp_min(1e-8)

        for _ in range(self.max_retries + 1):
            try:
                candidate = CRLinear.from_linear(target_mod, rank=cur_r)
                recon_w = candidate.get_weight().float()
                err = (torch.norm(orig_w - recon_w, p="fro") / norm_orig).item()
                if err < best_err:
                    best_err = err
                    best_cr_layer = candidate
                    best_r = cur_r

                if err <= self.max_error:
                    break

                next_r = min(int(cur_r * 1.5), max_rank - 1)
                if next_r <= cur_r or next_r * (in_f + out_f) >= 0.85 * (in_f * out_f):
                    break
                cur_r = next_r
            except Exception as ex:
                logger.debug("CR decomposition error on %s with rank %d: %s", name, cur_r, ex)
                break

        if best_err > self.max_error:
            logger.warning(
                "CRCompressor: Layer %s error %.4f exceeded max_error %.4f",
                name,
                best_err,
                self.max_error,
            )
            if self.fallback_to_original:
                self.fallback_layers.append(name)
                # Freeze fallback layers during recovery
                for p in module.parameters():
                    p.requires_grad = False
                continue

        if best_cr_layer is not None:
            new_layer = best_cr_layer
            compressed_layer = new_layer
            # Preserve existing LoRA adapter weights if the source module has them
            if (
                hasattr(module, "lora_A")
                and hasattr(module, "lora_B")
                and module.lora_A is not None
                and module.lora_B is not None
            ):
                src_lA = getattr(module.lora_A, "weight", module.lora_A)
                src_lB = getattr(module.lora_B, "weight", module.lora_B)
                if hasattr(new_layer, "lora_A") and new_layer.lora_A is not None:
                    if hasattr(src_lA, "data") and src_lA.shape == new_layer.lora_A.shape:
                        new_layer.lora_A.data.copy_(src_lA.data)
                if hasattr(new_layer, "lora_B") and new_layer.lora_B is not None:
                    if hasattr(src_lB, "data") and src_lB.shape == new_layer.lora_B.shape:
                        new_layer.lora_B.data.copy_(src_lB.data)

            if hasattr(module, "base_layer"):
                module.base_layer = compressed_layer  # put compressed layer inside wrapper
                # module itself (LoRALinear wrapper) remains the replacement
                new_module = module
            else:
                new_module = compressed_layer

            if hasattr(module, "base_layer"):
                for p_name, p in module.named_parameters():
                    if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                        p.requires_grad = True
            else:
                if name == "":
                    self.model = new_module
                else:
                    replace_module(self.model, name, new_module)

            self.compressed_layers.append(name)
            self.layer_errors[name] = float(best_err)
            self.layer_ranks[name] = int(best_r)
            total_orig_decomp_params += in_f * out_f
            total_comp_decomp_params += best_r * (in_f + out_f)

    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    avg_err = (
        (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
    )
    effective_param_ratio = (
        total_comp_decomp_params / total_orig_decomp_params
        if total_orig_decomp_params > 0
        else 1.0
    )
    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "layer_errors": self.layer_errors,
            "layer_ranks": self.layer_ranks,
            "avg_reconstruction_error": avg_err,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "cr_rank": self.rank,
            "target_attention_only": self.target_attention_only,
            "effective_parameter_ratio": round(effective_param_ratio, 4),
            "used_interpolative_decomposition_count": len(self.compressed_layers),
        }
    )
    logger.info(
        "Compressor cr finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        self.max_error,
    )
    return self.model, stats

Compress linear layers using interpolative decomposition.

def compress_linear(self, linear: nn.Linear | Any) ‑> CRLinear

Expand source code

def compress_linear(self, linear: nn.Linear | Any) -> CRLinear:
    """Compress a single linear layer into CRLinear."""
    return CRLinear.from_linear(linear, rank=self.rank)

Compress a single linear layer into CRLinear.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class CRLinear (in_features: int, out_features: int, rank: int = 64, bias: bool = True, dtype: torch.dtype | None = None, device: torch.device | None = None)

Expand source code

class CRLinear(nn.Module):
    """
    Interpolative Decomposition / Column-Row Factorized Linear Layer.
    W ≈ C @ P where:
    - C in R^{out_features x rank}: column skeleton from W
    - P in R^{rank x in_features}: interpolative projection matrix

    Forward pass:
    y = ((x @ P^T) @ C^T) + bias
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 64,
        bias: bool = True,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features  # Conv1D compatibility
        self.rank = max(1, min(rank, in_features, out_features))

        self.C = nn.Parameter(torch.empty(out_features, self.rank, dtype=dtype, device=device))
        self.P = nn.Parameter(torch.empty(self.rank, in_features, dtype=dtype, device=device))

        bound = 1.0 / math.sqrt(max(1, in_features))
        nn.init.uniform_(self.C, -bound, bound)
        nn.init.uniform_(self.P, -bound, bound)

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix (out_features, in_features)."""
        return torch.matmul(self.C, self.P)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Sub-quadratic forward execution:
        y = ((x @ P^T) @ C^T) + bias
        """
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)

        x_p = F.linear(x_flat, self.P)
        out = F.linear(x_p, self.C, self.bias)

        return out.view(*orig_shape[:-1], self.out_features)

    def get_weight(self) -> torch.Tensor:
        """Materialize full reconstructed weight matrix."""
        return torch.matmul(self.C, self.P)

    @classmethod
    def from_linear(cls, linear: nn.Linear | Any, rank: int = 64) -> CRLinear:
        """
        Perform Column Interpolative Decomposition on linear or Conv1D layer weight matrix.
        """
        target_mod = getattr(linear, "base_layer", getattr(linear, "linear", linear))
        raw_w = target_mod.weight.data
        bias = target_mod.bias.data if getattr(target_mod, "bias", None) is not None else None
        cls_name = target_mod.__class__.__name__.lower()
        wrapper_name = linear.__class__.__name__.lower()

        is_conv1d = (
            "conv1d" in cls_name
            or "conv1d" in wrapper_name
            or hasattr(target_mod, "nf")
            or hasattr(linear, "nf")
        )
        if is_conv1d:
            # Conv1D weights are [in_features, out_features] -> transpose to [out_features, in_features]
            W = raw_w.t().contiguous()
        else:
            W = raw_w.contiguous()

        out_f, in_f = W.shape
        r = max(1, min(rank, out_f, in_f))

        cr_layer = cls(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=bias is not None,
            dtype=W.dtype,
            device=W.device,
        )

        W_float = W.float()

        # Step 1: Select skeleton columns via SciPy QRCP or robust PyTorch pivoted selection
        col_idx = None
        try:
            import scipy.linalg

            W_np = W_float.cpu().numpy()
            _, _, piv = scipy.linalg.qr(W_np, mode="economic", pivoting=True)
            col_idx = torch.tensor(piv[:r], dtype=torch.long, device=W.device)
        except Exception as e:
            logger.debug("CRLinear pivoted QR failed; using greedy column selection: %s", e)

        if col_idx is None:
            # Pure PyTorch QRCP-equivalent greedy pivoted column selection
            list(range(in_f))
            selected_cols = []
            norms = torch.sum(W_float**2, dim=0)

            for _ in range(r):
                best_col_idx = int(torch.argmax(norms).item())
                selected_cols.append(best_col_idx)
                # Zero out selected norm
                norms[best_col_idx] = -1.0
                if len(selected_cols) >= r:
                    break
                # Orthogonalize remaining columns against selected column
                v = W_float[:, best_col_idx : best_col_idx + 1]
                v_norm_sq = torch.sum(v**2).clamp_min(1e-12)
                proj = torch.matmul(v.T, W_float) / v_norm_sq
                W_float = W_float - torch.matmul(v, proj)
                norms = torch.sum(W_float**2, dim=0)
                for sc in selected_cols:
                    norms[sc] = -1.0

            col_idx = torch.tensor(selected_cols, dtype=torch.long, device=W.device)

        W_float_orig = W.float()
        C_data = W_float_orig[:, col_idx]  # (out_f, r)

        # Step 2: Tikhonov regularized pseudo-inverse: (C^T C + eps * I)^{-1} C^T
        CtC = torch.matmul(C_data.T, C_data)  # (r, r)
        diag_mean = torch.trace(CtC) / max(1, r)
        eps = max(1e-4 * diag_mean.item(), 1e-4)
        reg_eye = eps * torch.eye(r, device=C_data.device, dtype=C_data.dtype)
        try:
            C_pinv = torch.linalg.solve(CtC + reg_eye, C_data.T)  # (r, out_f)
        except Exception:
            C_pinv = torch.linalg.pinv(C_data, rcond=1e-4)

        P_data = torch.matmul(C_pinv, W_float_orig)  # (r, in_f)

        # Step 3: Symmetrically normalize factor magnitudes to ensure FP16 stability
        norm_C = torch.norm(C_data, p="fro").clamp_min(1e-8)
        norm_P = torch.norm(P_data, p="fro").clamp_min(1e-8)
        scale = torch.sqrt(norm_C / norm_P)
        C_data = C_data / scale
        P_data = P_data * scale

        cr_layer.C.data.copy_(C_data.to(dtype=W.dtype))
        cr_layer.P.data.copy_(P_data.to(dtype=W.dtype))

        if bias is not None and cr_layer.bias is not None:
            cr_layer.bias.data.copy_(bias)

        return cr_layer

Interpolative Decomposition / Column-Row Factorized Linear Layer. W ≈ C @ P where: - C in R^{out_features x rank}: column skeleton from W - P in R^{rank x in_features}: interpolative projection matrix

Forward pass: y = ((x @ P^T) @ C^T) + bias

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Linear | Any, rank: int = 64) ‑> CRLinear: Perform Column Interpolative Decomposition on linear or Conv1D layer weight matrix.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix (out_features, in_features)."""
    return torch.matmul(self.C, self.P)

Reconstructed effective weight matrix (out_features, in_features).

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Sub-quadratic forward execution:
    y = ((x @ P^T) @ C^T) + bias
    """
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)

    x_p = F.linear(x_flat, self.P)
    out = F.linear(x_p, self.C, self.bias)

    return out.view(*orig_shape[:-1], self.out_features)

Sub-quadratic forward execution: y = ((x @ P^T) @ C^T) + bias

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    """Materialize full reconstructed weight matrix."""
    return torch.matmul(self.C, self.P)

Materialize full reconstructed weight matrix.

class CURCompressor (model: nn.Module | None = None, rank: int = 64, min_rank: int = 8, method: str = 'leverage', deterministic: bool = True, max_error: float = 0.2, fallback_to_original: bool = True, max_retries: int = 2, lora_r: int = 0, lora_alpha: float = 1.0, freeze_skeletons: bool = False, target_attention_only: bool = False, target_names: list[str] | None = None, exclude_names: list[str] | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("cur_compressor")
class CURCompressor(BaseCompressor):
    """
    Compressor that replaces eligible Linear and Conv1D layers with CURLinear.

    Args:
        model: PyTorch model.
        rank: Target rank for skeleton column and row selection (default 64).
        min_rank: Minimum allowable rank for compression (default 8).
        method: Selection method ('leverage' or 'energy', default 'leverage').
        deterministic: Deterministic top-k selection if True (default True).
        max_error: Maximum relative Frobenius reconstruction error allowed (default 0.20).
        fallback_to_original: If True, retains original layer when reconstruction error exceeds max_error (default True).
        max_retries: Number of adaptive rank retries when error threshold is exceeded (default 2).
        lora_r: LoRA adapter rank if adding fine-tuning adapters (default 0).
        lora_alpha: LoRA scaling factor (default 1.0).
        freeze_skeletons: If True, freeze C and R parameters during post-compression training (default False).
        target_attention_only: If True, restrict compression to attention projections only.
        target_names: Optional list of layer name substrings to target.
        exclude_names: Optional list of layer name substrings to exclude.
        cfg: Optional configuration object.
    """

    def __init__(
        self,
        model: nn.Module | None = None,
        rank: int = 64,
        min_rank: int = 8,
        method: str = "leverage",
        deterministic: bool = True,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        max_retries: int = 2,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        freeze_skeletons: bool = False,
        target_attention_only: bool = False,
        target_names: list[str] | None = None,
        exclude_names: list[str] | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model if model is not None else nn.Identity(), name="cur", **kwargs)
        c_rank = _cfg_get(cfg, "cur_rank", _cfg_get(cfg, "rank", rank))
        c_min_rank = _cfg_get(cfg, "cur_min_rank", _cfg_get(cfg, "min_rank", min_rank))
        c_method = _cfg_get(cfg, "cur_method", _cfg_get(cfg, "method", method))
        c_det = _cfg_get(cfg, "cur_deterministic", _cfg_get(cfg, "deterministic", deterministic))
        c_max_error = _cfg_get(cfg, "cur_max_error", _cfg_get(cfg, "max_error", max_error))
        c_fallback = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
        c_lora_r = _cfg_get(cfg, "cur_lora_r", _cfg_get(cfg, "lora_r", lora_r))
        c_lora_alpha = _cfg_get(cfg, "cur_lora_alpha", _cfg_get(cfg, "lora_alpha", lora_alpha))
        c_freeze = _cfg_get(cfg, "freeze_skeletons", freeze_skeletons)
        c_target_attn = _cfg_get(
            cfg,
            "cur_target_attention_only",
            _cfg_get(cfg, "target_attention_only", target_attention_only),
        )

        if c_max_error is None or c_max_error == 0.20:
            c_max_error = adaptive_max_error(c_rank, 768, 768, "distilbert")

        self.rank = int(c_rank) if c_rank is not None else 64
        self.min_rank = int(c_min_rank) if c_min_rank is not None else 8
        self.method = str(c_method)
        self.deterministic = bool(c_det)
        self.max_error = float(c_max_error)
        self.fallback_to_original = bool(c_fallback)
        self.max_retries = max(0, int(max_retries))
        self.lora_r = int(c_lora_r or 0)
        self.lora_alpha = float(c_lora_alpha or 1.0)
        self.freeze_skeletons = bool(c_freeze)
        self.target_attention_only = bool(c_target_attn)
        self.cfg = cfg
        self.target_names = target_names or _cfg_get(
            cfg, "cur_target_names", _cfg_get(cfg, "target_layers", None)
        )
        self.exclude_names = exclude_names or _cfg_get(
            cfg,
            "cur_exclude_names",
            [
                "lm_head",
                "classifier",
                "head",
                "score",
                "pooler",
                "embedding",
                "router",
                "gate",
            ],
        )

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        """Identify candidate linear/Conv1D/LoRALinear layers."""
        cls_name = module.__class__.__name__.lower()
        has_base = hasattr(module, "base_layer")
        base_cls = module.base_layer.__class__.__name__.lower() if has_base else ""

        is_linear = isinstance(module, nn.Linear) or "linear" in cls_name
        is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")
        is_lora = has_base and (
            isinstance(module.base_layer, nn.Linear)
            or "conv1d" in base_cls
            or hasattr(module.base_layer, "nf")
        )

        if not (is_linear or is_conv1d or is_lora):
            return False

        if isinstance(module, CURLinear) or (has_base and isinstance(module.base_layer, CURLinear)):
            return False

        name_lower = name.lower()
        if any(skip in name_lower for skip in self.exclude_names):
            return False

        if self.target_attention_only:
            ffn_patterns = [
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            ]
            if any(p in name_lower for p in ffn_patterns):
                return False
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            attn_patterns = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
                "in_proj_qkv",
                "out_proj",
            ]
            if not any(p in name_lower for p in attn_patterns):
                return False

        if self.target_names and not any(t in name for t in self.target_names):
            return False

        return True

    def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
        """Replace eligible Linear and Conv1D layers with CURLinear."""
        if model is not None:
            self.model = model
            self._original_params = self._count_params(self.model)

        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}
        visited_weight_ids: set[int] = set()

        total_orig_decomp_params = 0
        total_comp_decomp_params = 0

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                if self.target_attention_only:
                    is_lin_like = (
                        isinstance(module, nn.Linear)
                        or "conv1d" in module.__class__.__name__.lower()
                        or hasattr(module, "base_layer")
                    )
                    if is_lin_like:
                        for p in module.parameters():
                            p.requires_grad = False
                continue

            target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
            cls_name = module.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(module, "nf") or hasattr(target_mod, "nf")
            has_lora = (hasattr(module, "lora_A") and module.lora_A is not None) or hasattr(
                module, "base_layer"
            )

            # Check tied weights
            if not hasattr(target_mod, "weight") or target_mod.weight is None:
                for p in module.parameters():
                    p.requires_grad = False
                continue
            weight_id = id(target_mod.weight)
            if weight_id in visited_weight_ids:
                logger.info("CURCompressor: Skipping tied weight layer %s", name)
                for p in module.parameters():
                    p.requires_grad = False
                continue

            raw_w = target_mod.weight.data
            orig_w = raw_w.t().float() if is_conv1d else raw_w.float()
            out_f, in_f = orig_w.shape
            max_rank = min(out_f, in_f)

            if max_rank < self.min_rank:
                for p in module.parameters():
                    p.requires_grad = False
                continue

            # Parameter reduction check: r*(in_f + out_f + r) < in_f * out_f
            target_r = min(self.rank, max_rank)
            while target_r >= self.min_rank and (target_r * (in_f + out_f + target_r)) >= 0.85 * (
                in_f * out_f
            ):
                target_r = int(target_r * 0.8)

            if target_r < self.min_rank:
                logger.debug(
                    "CURCompressor: Skipping %s as target rank %d violates parameter reduction constraint",
                    name,
                    target_r,
                )
                for p in module.parameters():
                    p.requires_grad = False
                continue

            cur_r = target_r
            best_cur_layer = None
            best_err = float("inf")
            best_r = cur_r
            norm_orig = torch.norm(orig_w, p="fro").clamp_min(1e-8)
            effective_lora_r = 0 if has_lora else self.lora_r

            for _ in range(self.max_retries + 1):
                try:
                    candidate = CURLinear.from_linear(
                        target_mod,
                        rank=cur_r,
                        method=self.method,
                        deterministic=self.deterministic,
                        lora_r=effective_lora_r,
                        lora_alpha=self.lora_alpha,
                        freeze_skeletons=self.freeze_skeletons,
                    )
                    recon_w = candidate.get_weight().float()
                    err = (torch.norm(orig_w - recon_w, p="fro") / norm_orig).item()

                    if err < best_err:
                        best_err = err
                        best_cur_layer = candidate
                        best_r = cur_r

                    if err <= self.max_error:
                        break

                    next_r = min(int(cur_r * 1.5), max_rank - 1)
                    if next_r <= cur_r or (next_r * (in_f + out_f + next_r)) >= 0.85 * (
                        in_f * out_f
                    ):
                        break
                    cur_r = next_r
                except Exception as ex:
                    logger.warning(
                        "CURCompressor: Failed CUR decomposition on %s at rank %d: %s",
                        name,
                        cur_r,
                        ex,
                    )
                    break

            if best_err > self.max_error:
                logger.warning(
                    "CURCompressor: Layer %s error %.4f exceeded max_error %.4f",
                    name,
                    best_err,
                    self.max_error,
                )
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    # Freeze fallback layers
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

            if best_cur_layer is not None:
                new_layer = best_cur_layer
                compressed_layer = new_layer
                # Preserve existing LoRA adapter weights if the source module has them
                if (
                    hasattr(module, "lora_A")
                    and hasattr(module, "lora_B")
                    and module.lora_A is not None
                    and module.lora_B is not None
                ):
                    src_lA = getattr(module.lora_A, "weight", module.lora_A)
                    src_lB = getattr(module.lora_B, "weight", module.lora_B)
                    if hasattr(new_layer, "lora_A") and new_layer.lora_A is not None:
                        if hasattr(src_lA, "data") and src_lA.shape == new_layer.lora_A.shape:
                            new_layer.lora_A.data.copy_(src_lA.data)
                    if hasattr(new_layer, "lora_B") and new_layer.lora_B is not None:
                        if hasattr(src_lB, "data") and src_lB.shape == new_layer.lora_B.shape:
                            new_layer.lora_B.data.copy_(src_lB.data)

                if hasattr(module, "base_layer"):
                    module.base_layer = compressed_layer  # put compressed layer inside wrapper
                    # module itself (LoRALinear wrapper) remains the replacement
                    new_module = module
                else:
                    new_module = compressed_layer

                if hasattr(module, "base_layer"):
                    for p_name, p in module.named_parameters():
                        if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                            p.requires_grad = True
                else:
                    if name == "":
                        self.model = new_module
                    else:
                        replace_module(self.model, name, new_module)

                self.compressed_layers.append(name)
                self.layer_errors[name] = float(best_err)
                self.layer_ranks[name] = int(best_r)
                visited_weight_ids.add(weight_id)

                total_orig_decomp_params += in_f * out_f
                total_comp_decomp_params += best_r * (in_f + out_f + best_r)

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        avg_err = (
            (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
        )
        effective_param_ratio = (
            total_comp_decomp_params / total_orig_decomp_params
            if total_orig_decomp_params > 0
            else 1.0
        )
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "avg_reconstruction_error": avg_err,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "cur_method": self.method,
                "cur_rank": self.rank,
                "min_rank": self.min_rank,
                "deterministic": self.deterministic,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "target_attention_only": self.target_attention_only,
                "effective_parameter_ratio": round(effective_param_ratio, 4),
            }
        )
        logger.info(
            "Compressor cur finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

    def compress_linear(self, linear: nn.Linear | Any) -> CURLinear:
        """Compress a single linear layer into CURLinear."""
        return CURLinear.from_linear(
            linear,
            rank=self.rank,
            method=self.method,
            deterministic=self.deterministic,
            lora_r=self.lora_r,
            lora_alpha=self.lora_alpha,
            freeze_skeletons=self.freeze_skeletons,
        )

Compressor that replaces eligible Linear and Conv1D layers with CURLinear.

Args
-----=
model: PyTorch model.
rank: Target rank for skeleton column and row selection (default 64).
min_rank: Minimum allowable rank for compression (default 8).
method: Selection method ('leverage' or 'energy', default 'leverage').
deterministic: Deterministic top-k selection if True (default True).
max_error: Maximum relative Frobenius reconstruction error allowed (default 0.20).
fallback_to_original: If True, retains original layer when reconstruction error exceeds max_error (default True).
max_retries: Number of adaptive rank retries when error threshold is exceeded (default 2).
lora_r: LoRA adapter rank if adding fine-tuning adapters (default 0).
lora_alpha: LoRA scaling factor (default 1.0).
freeze_skeletons: If True, freeze C and R parameters during post-compression training (default False).
target_attention_only: If True, restrict compression to attention projections only.
target_names: Optional list of layer name substrings to target.
exclude_names: Optional list of layer name substrings to exclude.
cfg: Optional configuration object.

Ancestors

BaseCompressor

Methods

def compress(self, model: nn.Module | None = None) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
    """Replace eligible Linear and Conv1D layers with CURLinear."""
    if model is not None:
        self.model = model
        self._original_params = self._count_params(self.model)

    self.compressed_layers = []
    self.fallback_layers = []
    self.layer_errors = {}
    self.layer_ranks = {}
    visited_weight_ids: set[int] = set()

    total_orig_decomp_params = 0
    total_comp_decomp_params = 0

    for name, module in list(self.model.named_modules()):
        if not self._is_target_module(name, module):
            if self.target_attention_only:
                is_lin_like = (
                    isinstance(module, nn.Linear)
                    or "conv1d" in module.__class__.__name__.lower()
                    or hasattr(module, "base_layer")
                )
                if is_lin_like:
                    for p in module.parameters():
                        p.requires_grad = False
            continue

        target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
        cls_name = module.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(module, "nf") or hasattr(target_mod, "nf")
        has_lora = (hasattr(module, "lora_A") and module.lora_A is not None) or hasattr(
            module, "base_layer"
        )

        # Check tied weights
        if not hasattr(target_mod, "weight") or target_mod.weight is None:
            for p in module.parameters():
                p.requires_grad = False
            continue
        weight_id = id(target_mod.weight)
        if weight_id in visited_weight_ids:
            logger.info("CURCompressor: Skipping tied weight layer %s", name)
            for p in module.parameters():
                p.requires_grad = False
            continue

        raw_w = target_mod.weight.data
        orig_w = raw_w.t().float() if is_conv1d else raw_w.float()
        out_f, in_f = orig_w.shape
        max_rank = min(out_f, in_f)

        if max_rank < self.min_rank:
            for p in module.parameters():
                p.requires_grad = False
            continue

        # Parameter reduction check: r*(in_f + out_f + r) < in_f * out_f
        target_r = min(self.rank, max_rank)
        while target_r >= self.min_rank and (target_r * (in_f + out_f + target_r)) >= 0.85 * (
            in_f * out_f
        ):
            target_r = int(target_r * 0.8)

        if target_r < self.min_rank:
            logger.debug(
                "CURCompressor: Skipping %s as target rank %d violates parameter reduction constraint",
                name,
                target_r,
            )
            for p in module.parameters():
                p.requires_grad = False
            continue

        cur_r = target_r
        best_cur_layer = None
        best_err = float("inf")
        best_r = cur_r
        norm_orig = torch.norm(orig_w, p="fro").clamp_min(1e-8)
        effective_lora_r = 0 if has_lora else self.lora_r

        for _ in range(self.max_retries + 1):
            try:
                candidate = CURLinear.from_linear(
                    target_mod,
                    rank=cur_r,
                    method=self.method,
                    deterministic=self.deterministic,
                    lora_r=effective_lora_r,
                    lora_alpha=self.lora_alpha,
                    freeze_skeletons=self.freeze_skeletons,
                )
                recon_w = candidate.get_weight().float()
                err = (torch.norm(orig_w - recon_w, p="fro") / norm_orig).item()

                if err < best_err:
                    best_err = err
                    best_cur_layer = candidate
                    best_r = cur_r

                if err <= self.max_error:
                    break

                next_r = min(int(cur_r * 1.5), max_rank - 1)
                if next_r <= cur_r or (next_r * (in_f + out_f + next_r)) >= 0.85 * (
                    in_f * out_f
                ):
                    break
                cur_r = next_r
            except Exception as ex:
                logger.warning(
                    "CURCompressor: Failed CUR decomposition on %s at rank %d: %s",
                    name,
                    cur_r,
                    ex,
                )
                break

        if best_err > self.max_error:
            logger.warning(
                "CURCompressor: Layer %s error %.4f exceeded max_error %.4f",
                name,
                best_err,
                self.max_error,
            )
            if self.fallback_to_original:
                self.fallback_layers.append(name)
                # Freeze fallback layers
                for p in module.parameters():
                    p.requires_grad = False
                continue

        if best_cur_layer is not None:
            new_layer = best_cur_layer
            compressed_layer = new_layer
            # Preserve existing LoRA adapter weights if the source module has them
            if (
                hasattr(module, "lora_A")
                and hasattr(module, "lora_B")
                and module.lora_A is not None
                and module.lora_B is not None
            ):
                src_lA = getattr(module.lora_A, "weight", module.lora_A)
                src_lB = getattr(module.lora_B, "weight", module.lora_B)
                if hasattr(new_layer, "lora_A") and new_layer.lora_A is not None:
                    if hasattr(src_lA, "data") and src_lA.shape == new_layer.lora_A.shape:
                        new_layer.lora_A.data.copy_(src_lA.data)
                if hasattr(new_layer, "lora_B") and new_layer.lora_B is not None:
                    if hasattr(src_lB, "data") and src_lB.shape == new_layer.lora_B.shape:
                        new_layer.lora_B.data.copy_(src_lB.data)

            if hasattr(module, "base_layer"):
                module.base_layer = compressed_layer  # put compressed layer inside wrapper
                # module itself (LoRALinear wrapper) remains the replacement
                new_module = module
            else:
                new_module = compressed_layer

            if hasattr(module, "base_layer"):
                for p_name, p in module.named_parameters():
                    if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                        p.requires_grad = True
            else:
                if name == "":
                    self.model = new_module
                else:
                    replace_module(self.model, name, new_module)

            self.compressed_layers.append(name)
            self.layer_errors[name] = float(best_err)
            self.layer_ranks[name] = int(best_r)
            visited_weight_ids.add(weight_id)

            total_orig_decomp_params += in_f * out_f
            total_comp_decomp_params += best_r * (in_f + out_f + best_r)

    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    avg_err = (
        (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
    )
    effective_param_ratio = (
        total_comp_decomp_params / total_orig_decomp_params
        if total_orig_decomp_params > 0
        else 1.0
    )
    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "layer_errors": self.layer_errors,
            "layer_ranks": self.layer_ranks,
            "avg_reconstruction_error": avg_err,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "cur_method": self.method,
            "cur_rank": self.rank,
            "min_rank": self.min_rank,
            "deterministic": self.deterministic,
            "lora_r": self.lora_r,
            "lora_alpha": self.lora_alpha,
            "target_attention_only": self.target_attention_only,
            "effective_parameter_ratio": round(effective_param_ratio, 4),
        }
    )
    logger.info(
        "Compressor cur finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        self.max_error,
    )
    return self.model, stats

Replace eligible Linear and Conv1D layers with CURLinear.

def compress_linear(self, linear: nn.Linear | Any) ‑> CURLinear

Expand source code

def compress_linear(self, linear: nn.Linear | Any) -> CURLinear:
    """Compress a single linear layer into CURLinear."""
    return CURLinear.from_linear(
        linear,
        rank=self.rank,
        method=self.method,
        deterministic=self.deterministic,
        lora_r=self.lora_r,
        lora_alpha=self.lora_alpha,
        freeze_skeletons=self.freeze_skeletons,
    )

Compress a single linear layer into CURLinear.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class CURLinear (in_features: int, out_features: int, rank: int = 64, bias: bool = True, method: str = 'leverage', lora_r: int = 0, lora_alpha: float = 1.0, freeze_skeletons: bool = False, dtype: torch.dtype | None = None, device: torch.device | None = None)

Expand source code

class CURLinear(nn.Module):
    """
    Linear layer whose weight matrix is factorized via CUR cross-approximation:
    W ≈ C @ U @ R
    - C in R^{out_features x rank}: selected columns of W
    - U in R^{rank x rank}: Moore-Penrose pseudo-inverse core matrix
    - R in R^{rank x in_features}: selected rows of W

    Executes forward pass in factorized form:
    y = ((x @ R^T) @ U^T) @ C^T + bias
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 64,
        bias: bool = True,
        method: str = "leverage",
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        freeze_skeletons: bool = False,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = max(1, min(rank, in_features, out_features))
        self.method = method
        self.lora_r = max(0, int(lora_r))
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / max(1, self.lora_r)) if self.lora_r > 0 else 1.0

        # C: (out_features, rank), U: (rank, rank), R: (rank, in_features)
        self.C = nn.Parameter(torch.empty(out_features, self.rank, dtype=dtype, device=device))
        self.U = nn.Parameter(torch.empty(self.rank, self.rank, dtype=dtype, device=device))
        self.R = nn.Parameter(torch.empty(self.rank, in_features, dtype=dtype, device=device))

        bound = 1.0 / math.sqrt(max(1, in_features))
        nn.init.uniform_(self.C, -bound, bound)
        nn.init.eye_(self.U)
        nn.init.uniform_(self.R, -bound, bound)

        if freeze_skeletons:
            self.C.requires_grad = False
            self.R.requires_grad = False

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty(self.lora_r, in_features, dtype=dtype, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.empty(out_features, self.lora_r, dtype=dtype, device=device)
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix for compatibility."""
        W_core = torch.matmul(torch.matmul(self.C, self.U), self.R)
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            W_core = W_core + (torch.matmul(self.lora_B, self.lora_A) * self.scaling)
        return W_core

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Sub-quadratic forward execution:
        x: (..., in_features)
        y = ((x @ R^T) @ U^T) @ C^T + bias
        """
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)

        # Factorized forward pass:
        # x_r = x @ R^T: (K, rank)
        x_r = F.linear(x_flat, self.R)
        # x_ru = x_r @ U^T: (K, rank)
        x_ru = F.linear(x_r, self.U)
        # out = x_ru @ C^T + bias: (K, out_features)
        out = F.linear(x_ru, self.C, self.bias)

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x_flat, self.lora_A), self.lora_B) * self.scaling
            out = out + lora_out

        return out.view(*orig_shape[:-1], self.out_features)

    def get_weight(self) -> torch.Tensor:
        """Materialize reconstructed weight matrix W = C @ U @ R."""
        return self.weight

    @classmethod
    def from_linear(
        cls,
        linear: nn.Linear | Any,
        rank: int = 64,
        method: str = "leverage",
        deterministic: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        freeze_skeletons: bool = False,
    ) -> CURLinear:
        """
        Construct CURLinear from existing nn.Linear or Conv1D using leverage scores or energy sampling.
        """
        target_mod = getattr(linear, "base_layer", getattr(linear, "linear", linear))
        raw_w = target_mod.weight.data
        bias = target_mod.bias.data if getattr(target_mod, "bias", None) is not None else None
        cls_name = target_mod.__class__.__name__.lower()
        wrapper_name = linear.__class__.__name__.lower()

        is_conv1d = (
            "conv1d" in cls_name
            or "conv1d" in wrapper_name
            or hasattr(target_mod, "nf")
            or hasattr(linear, "nf")
        )
        if is_conv1d:
            W = raw_w.t().contiguous()
        else:
            W = raw_w.contiguous()

        out_features, in_features = W.shape
        r = max(1, min(rank, out_features, in_features))

        cur_layer = cls(
            in_features=in_features,
            out_features=out_features,
            rank=r,
            bias=bias is not None,
            method=method,
            lora_r=lora_r,
            lora_alpha=lora_alpha,
            freeze_skeletons=freeze_skeletons,
            dtype=W.dtype,
            device=W.device,
        )

        W_float = W.float()
        if method == "leverage":
            col_scores, row_scores = compute_leverage_scores(W_float, rank=r)
        elif method == "energy":
            col_scores, row_scores = compute_energy_scores(W_float)
        else:
            col_scores = torch.ones(in_features, device=W.device)
            row_scores = torch.ones(out_features, device=W.device)

        col_idx = select_indices(col_scores, k=r, deterministic=deterministic)
        row_idx = select_indices(row_scores, k=r, deterministic=deterministic)

        # C = W[:, col_idx] (out_features, r)
        # R = W[row_idx, :] (r, in_features)
        C = W_float[:, col_idx]
        R = W_float[row_idx, :]

        # Optimal core matrix: U = C^+ @ W @ R^+
        # This is the exact least-squares solution argmin_U ||W - C U R||_F
        try:
            C_pinv = torch.linalg.pinv(C, rcond=1e-4)
            R_pinv = torch.linalg.pinv(R, rcond=1e-4)
            U = torch.matmul(torch.matmul(C_pinv, W_float), R_pinv)
        except Exception:
            W_int = W_float[row_idx][:, col_idx]
            U = torch.linalg.pinv(W_int, rcond=1e-3)

        cur_layer.C.data.copy_(C.to(dtype=W.dtype))
        cur_layer.U.data.copy_(U.to(dtype=W.dtype))
        cur_layer.R.data.copy_(R.to(dtype=W.dtype))

        if bias is not None and cur_layer.bias is not None:
            cur_layer.bias.data.copy_(bias)

        return cur_layer

Linear layer whose weight matrix is factorized via CUR cross-approximation: W ≈ C @ U @ R - C in R^{out_features x rank}: selected columns of W - U in R^{rank x rank}: Moore-Penrose pseudo-inverse core matrix - R in R^{rank x in_features}: selected rows of W

Executes forward pass in factorized form: y = ((x @ R^T) @ U^T) @ C^T + bias

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Linear | Any, rank: int = 64, method: str = 'leverage', deterministic: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, freeze_skeletons: bool = False) ‑> CURLinear: Construct CURLinear from existing nn.Linear or Conv1D using leverage scores or energy sampling.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix for compatibility."""
    W_core = torch.matmul(torch.matmul(self.C, self.U), self.R)
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        W_core = W_core + (torch.matmul(self.lora_B, self.lora_A) * self.scaling)
    return W_core

Reconstructed effective weight matrix for compatibility.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Sub-quadratic forward execution:
    x: (..., in_features)
    y = ((x @ R^T) @ U^T) @ C^T + bias
    """
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)

    # Factorized forward pass:
    # x_r = x @ R^T: (K, rank)
    x_r = F.linear(x_flat, self.R)
    # x_ru = x_r @ U^T: (K, rank)
    x_ru = F.linear(x_r, self.U)
    # out = x_ru @ C^T + bias: (K, out_features)
    out = F.linear(x_ru, self.C, self.bias)

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x_flat, self.lora_A), self.lora_B) * self.scaling
        out = out + lora_out

    return out.view(*orig_shape[:-1], self.out_features)

Sub-quadratic forward execution: x: (…, in_features) y = ((x @ R^T) @ U^T) @ C^T + bias

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    """Materialize reconstructed weight matrix W = C @ U @ R."""
    return self.weight

Materialize reconstructed weight matrix W = C @ U @ R.

class CURLowRankLinear (in_features: int, out_features: int, rank: int = 64, bias: bool = True, method: str = 'leverage', lora_r: int = 0, lora_alpha: float = 1.0, freeze_skeletons: bool = False, dtype: torch.dtype | None = None, device: torch.device | None = None)

Expand source code

class CURLinear(nn.Module):
    """
    Linear layer whose weight matrix is factorized via CUR cross-approximation:
    W ≈ C @ U @ R
    - C in R^{out_features x rank}: selected columns of W
    - U in R^{rank x rank}: Moore-Penrose pseudo-inverse core matrix
    - R in R^{rank x in_features}: selected rows of W

    Executes forward pass in factorized form:
    y = ((x @ R^T) @ U^T) @ C^T + bias
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 64,
        bias: bool = True,
        method: str = "leverage",
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        freeze_skeletons: bool = False,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = max(1, min(rank, in_features, out_features))
        self.method = method
        self.lora_r = max(0, int(lora_r))
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / max(1, self.lora_r)) if self.lora_r > 0 else 1.0

        # C: (out_features, rank), U: (rank, rank), R: (rank, in_features)
        self.C = nn.Parameter(torch.empty(out_features, self.rank, dtype=dtype, device=device))
        self.U = nn.Parameter(torch.empty(self.rank, self.rank, dtype=dtype, device=device))
        self.R = nn.Parameter(torch.empty(self.rank, in_features, dtype=dtype, device=device))

        bound = 1.0 / math.sqrt(max(1, in_features))
        nn.init.uniform_(self.C, -bound, bound)
        nn.init.eye_(self.U)
        nn.init.uniform_(self.R, -bound, bound)

        if freeze_skeletons:
            self.C.requires_grad = False
            self.R.requires_grad = False

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty(self.lora_r, in_features, dtype=dtype, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.empty(out_features, self.lora_r, dtype=dtype, device=device)
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix for compatibility."""
        W_core = torch.matmul(torch.matmul(self.C, self.U), self.R)
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            W_core = W_core + (torch.matmul(self.lora_B, self.lora_A) * self.scaling)
        return W_core

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Sub-quadratic forward execution:
        x: (..., in_features)
        y = ((x @ R^T) @ U^T) @ C^T + bias
        """
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)

        # Factorized forward pass:
        # x_r = x @ R^T: (K, rank)
        x_r = F.linear(x_flat, self.R)
        # x_ru = x_r @ U^T: (K, rank)
        x_ru = F.linear(x_r, self.U)
        # out = x_ru @ C^T + bias: (K, out_features)
        out = F.linear(x_ru, self.C, self.bias)

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x_flat, self.lora_A), self.lora_B) * self.scaling
            out = out + lora_out

        return out.view(*orig_shape[:-1], self.out_features)

    def get_weight(self) -> torch.Tensor:
        """Materialize reconstructed weight matrix W = C @ U @ R."""
        return self.weight

    @classmethod
    def from_linear(
        cls,
        linear: nn.Linear | Any,
        rank: int = 64,
        method: str = "leverage",
        deterministic: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        freeze_skeletons: bool = False,
    ) -> CURLinear:
        """
        Construct CURLinear from existing nn.Linear or Conv1D using leverage scores or energy sampling.
        """
        target_mod = getattr(linear, "base_layer", getattr(linear, "linear", linear))
        raw_w = target_mod.weight.data
        bias = target_mod.bias.data if getattr(target_mod, "bias", None) is not None else None
        cls_name = target_mod.__class__.__name__.lower()
        wrapper_name = linear.__class__.__name__.lower()

        is_conv1d = (
            "conv1d" in cls_name
            or "conv1d" in wrapper_name
            or hasattr(target_mod, "nf")
            or hasattr(linear, "nf")
        )
        if is_conv1d:
            W = raw_w.t().contiguous()
        else:
            W = raw_w.contiguous()

        out_features, in_features = W.shape
        r = max(1, min(rank, out_features, in_features))

        cur_layer = cls(
            in_features=in_features,
            out_features=out_features,
            rank=r,
            bias=bias is not None,
            method=method,
            lora_r=lora_r,
            lora_alpha=lora_alpha,
            freeze_skeletons=freeze_skeletons,
            dtype=W.dtype,
            device=W.device,
        )

        W_float = W.float()
        if method == "leverage":
            col_scores, row_scores = compute_leverage_scores(W_float, rank=r)
        elif method == "energy":
            col_scores, row_scores = compute_energy_scores(W_float)
        else:
            col_scores = torch.ones(in_features, device=W.device)
            row_scores = torch.ones(out_features, device=W.device)

        col_idx = select_indices(col_scores, k=r, deterministic=deterministic)
        row_idx = select_indices(row_scores, k=r, deterministic=deterministic)

        # C = W[:, col_idx] (out_features, r)
        # R = W[row_idx, :] (r, in_features)
        C = W_float[:, col_idx]
        R = W_float[row_idx, :]

        # Optimal core matrix: U = C^+ @ W @ R^+
        # This is the exact least-squares solution argmin_U ||W - C U R||_F
        try:
            C_pinv = torch.linalg.pinv(C, rcond=1e-4)
            R_pinv = torch.linalg.pinv(R, rcond=1e-4)
            U = torch.matmul(torch.matmul(C_pinv, W_float), R_pinv)
        except Exception:
            W_int = W_float[row_idx][:, col_idx]
            U = torch.linalg.pinv(W_int, rcond=1e-3)

        cur_layer.C.data.copy_(C.to(dtype=W.dtype))
        cur_layer.U.data.copy_(U.to(dtype=W.dtype))
        cur_layer.R.data.copy_(R.to(dtype=W.dtype))

        if bias is not None and cur_layer.bias is not None:
            cur_layer.bias.data.copy_(bias)

        return cur_layer

Executes forward pass in factorized form: y = ((x @ R^T) @ U^T) @ C^T + bias

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Linear | Any, rank: int = 64, method: str = 'leverage', deterministic: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, freeze_skeletons: bool = False) ‑> CURLinear: Construct CURLinear from existing nn.Linear or Conv1D using leverage scores or energy sampling.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix for compatibility."""
    W_core = torch.matmul(torch.matmul(self.C, self.U), self.R)
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        W_core = W_core + (torch.matmul(self.lora_B, self.lora_A) * self.scaling)
    return W_core

Reconstructed effective weight matrix for compatibility.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Sub-quadratic forward execution:
    x: (..., in_features)
    y = ((x @ R^T) @ U^T) @ C^T + bias
    """
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)

    # Factorized forward pass:
    # x_r = x @ R^T: (K, rank)
    x_r = F.linear(x_flat, self.R)
    # x_ru = x_r @ U^T: (K, rank)
    x_ru = F.linear(x_r, self.U)
    # out = x_ru @ C^T + bias: (K, out_features)
    out = F.linear(x_ru, self.C, self.bias)

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x_flat, self.lora_A), self.lora_B) * self.scaling
        out = out + lora_out

    return out.view(*orig_shape[:-1], self.out_features)

Sub-quadratic forward execution: x: (…, in_features) y = ((x @ R^T) @ U^T) @ C^T + bias

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    """Materialize reconstructed weight matrix W = C @ U @ R."""
    return self.weight

Materialize reconstructed weight matrix W = C @ U @ R.

class CURMoECompressor (model: nn.Module | None = None, rank: int = 64, min_rank: int = 8, method: str = 'leverage', deterministic: bool = True, max_error: float = 0.9, fallback_to_original: bool = True, max_retries: int = 2, target_attention_only: bool = False, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("cur_moe_class")
@COMPRESSOR_REGISTRY.register("cur_moe_compressor_class")
class CURMoECompressor(BaseCompressor):
    """
    Compresses Mixture of Experts (MoE) layers and dense FFN layers using CUR cross-approximation.

    Args:
        model: PyTorch model containing MoE or transformer layers.
        rank: Target rank for CUR decomposition.
        min_rank: Minimum allowable rank (default 8).
        method: 'leverage' (statistical leverage scores) or 'energy' (Frobenius energy).
        deterministic: Whether to use deterministic top-k selection (default True).
        max_error: Maximum allowed relative Frobenius error (default 0.90).
        fallback_to_original: Whether to preserve original layer if error > max_error (default True).
        max_retries: Number of adaptive rank escalations (default 2).
        target_attention_only: If True, skip attention projections and compress only FFN/MoE expert layers.
        cfg: Optional configuration object.
    """

    def __init__(
        self,
        model: nn.Module | None = None,
        rank: int = 64,
        min_rank: int = 8,
        method: str = "leverage",
        deterministic: bool = True,
        max_error: float = 0.90,
        fallback_to_original: bool = True,
        max_retries: int = 2,
        target_attention_only: bool = False,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model if model is not None else nn.Identity(), name="cur_moe", **kwargs)
        c_rank = _cfg_get(cfg, "cur_moe_rank", _cfg_get(cfg, "rank", rank))
        c_min_rank = _cfg_get(cfg, "cur_min_rank", _cfg_get(cfg, "min_rank", min_rank))
        c_method = _cfg_get(cfg, "cur_moe_method", _cfg_get(cfg, "method", method))
        c_det = _cfg_get(cfg, "cur_deterministic", _cfg_get(cfg, "deterministic", deterministic))
        c_max_error = _cfg_get(cfg, "cur_moe_max_error", _cfg_get(cfg, "max_error", max_error))
        c_fallback = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
        c_target_attn = _cfg_get(
            cfg,
            "cur_moe_target_attention_only",
            _cfg_get(cfg, "target_attention_only", target_attention_only),
        )

        if c_max_error is None or c_max_error == 0.90:
            c_max_error = adaptive_max_error(c_rank, 768, 768, "distilbert")
            c_max_error = max(c_max_error, 0.95)

        self.rank = int(c_rank) if c_rank is not None else 64
        self.min_rank = int(c_min_rank) if c_min_rank is not None else 8
        self.method = str(c_method)
        self.deterministic = bool(c_det)
        self.max_error = float(c_max_error)
        self.fallback_to_original = bool(c_fallback)
        self.max_retries = max(0, int(max_retries))
        self.target_attention_only = bool(c_target_attn)
        self.cfg = cfg

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.total_orig_decomp_params = 0
        self.total_comp_decomp_params = 0

    def _compress_single_layer(self, name: str, module: nn.Module) -> bool:
        """Compress a single linear or Conv1D layer with Frobenius validation and adaptive retries."""
        target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
        cls_name = target_mod.__class__.__name__.lower()
        wrapper_name = module.__class__.__name__.lower()
        is_conv1d = (
            "conv1d" in cls_name
            or "conv1d" in wrapper_name
            or hasattr(module, "nf")
            or hasattr(target_mod, "nf")
        )
        (hasattr(module, "lora_A") and module.lora_A is not None) or hasattr(module, "base_layer")

        raw_w = target_mod.weight.data
        orig_w = raw_w.t().float() if is_conv1d else raw_w.float()
        norm_orig = torch.norm(orig_w, p="fro").clamp_min(1e-8)
        if norm_orig.item() < 1e-4:
            logger.warning(
                "CURMoECompressor: Layer %s weight Frobenius norm %.6f < 1e-4, skipping",
                name,
                norm_orig.item(),
            )
            return False

        out_f, in_f = orig_w.shape
        max_rank = min(out_f, in_f)

        if max_rank < self.min_rank:
            return False

        # Parameter reduction check: r*(in_f + out_f + r) < in_f * out_f
        target_r = min(self.rank, max_rank)
        while target_r >= self.min_rank and (target_r * (in_f + out_f + target_r)) >= 0.85 * (
            in_f * out_f
        ):
            target_r = int(target_r * 0.8)

        if target_r < self.min_rank:
            return False

        cur_r = target_r
        best_cur_layer = None
        best_err = float("inf")
        best_r = cur_r

        for _ in range(self.max_retries + 1):
            try:
                candidate = CURLinear.from_linear(
                    target_mod,
                    rank=cur_r,
                    method=self.method,
                    deterministic=self.deterministic,
                )
                recon_w = candidate.get_weight().float()
                err = (torch.norm(orig_w - recon_w, p="fro") / norm_orig).item()

                if err < best_err:
                    best_err = err
                    best_cur_layer = candidate
                    best_r = cur_r

                if err <= self.max_error:
                    break

                next_r = min(int(cur_r * 1.5), max_rank - 1)
                if next_r <= cur_r or (next_r * (in_f + out_f + next_r)) >= 0.85 * (in_f * out_f):
                    break
                cur_r = next_r
            except Exception as ex:
                logger.warning("CURMoECompressor: Error on %s with rank %d: %s", name, cur_r, ex)
                break

        if best_err > self.max_error or best_cur_layer is None:
            logger.warning(
                "CURMoECompressor: Layer %s error %.4f exceeded max_error %.4f",
                name,
                best_err,
                self.max_error,
            )
            if self.fallback_to_original:
                self.fallback_layers.append(name)
                # Freeze fallback layers
                for p in module.parameters():
                    p.requires_grad = False
                return False

        new_layer = best_cur_layer
        compressed_layer = new_layer

        has_base_layer = hasattr(module, "base_layer") or (
            hasattr(module, "linear") and (hasattr(module, "lora_A") or hasattr(module, "r"))
        )
        # Preserve existing LoRA adapter weights if the source module has them and not wrapped
        if (
            not has_base_layer
            and hasattr(module, "lora_A")
            and hasattr(module, "lora_B")
            and module.lora_A is not None
            and module.lora_B is not None
        ):
            src_lA = getattr(module.lora_A, "weight", module.lora_A)
            src_lB = getattr(module.lora_B, "weight", module.lora_B)
            dst_lA = getattr(new_layer, "lora_A", None)
            dst_lB = getattr(new_layer, "lora_B", None)
            if (
                dst_lA is not None
                and hasattr(src_lA, "data")
                and hasattr(dst_lA, "data")
                and src_lA.shape == dst_lA.shape
            ):
                dst_lA.data.copy_(src_lA.data)
            if (
                dst_lB is not None
                and hasattr(src_lB, "data")
                and hasattr(dst_lB, "data")
                and src_lB.shape == dst_lB.shape
            ):
                dst_lB.data.copy_(src_lB.data)

        if hasattr(module, "base_layer"):
            module.base_layer = compressed_layer  # put compressed layer inside wrapper
            new_module = module
        elif hasattr(module, "linear") and (hasattr(module, "lora_A") or hasattr(module, "r")):
            module.linear = compressed_layer  # put compressed layer inside framework LoRALinear
            new_module = module
        else:
            new_module = compressed_layer

        if has_base_layer:
            for p_name, p in module.named_parameters():
                if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                    p.requires_grad = True
        else:
            if name == "":
                self.model = new_module
            else:
                replace_module(self.model, name, new_module)

        self.compressed_layers.append(name)
        self.layer_errors[name] = float(best_err)
        self.layer_ranks[name] = int(best_r)
        self.total_orig_decomp_params += in_f * out_f
        self.total_comp_decomp_params += best_r * (in_f + out_f + best_r)
        return True

    def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
        """Compress expert linear layers in MoE architecture, or fallback to FFN layers."""
        if model is not None:
            self.model = model
            self._original_params = self._count_params(self.model)

        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}
        self.total_orig_decomp_params = 0
        self.total_comp_decomp_params = 0
        experts_detected = False

        # Scan for explicit MoE modules (module.experts or expert submodules)
        for name, module in list(self.model.named_modules()):
            name_lower = name.lower()
            if any(
                k in name_lower
                for k in [
                    "router",
                    "gate",
                    "classifier",
                    "score",
                    "lm_head",
                    "pooler",
                    "embedding",
                ]
            ):
                continue

            if hasattr(module, "experts"):
                experts_container = module.experts
                if isinstance(experts_container, (nn.ModuleList, list)):
                    experts_detected = True
                    for exp_idx, expert in enumerate(experts_container):
                        for child_name, child_mod in list(expert.named_children()):
                            child_full_name = (
                                f"{name}.experts.{exp_idx}.{child_name}"
                                if name
                                else f"experts.{exp_idx}.{child_name}"
                            )
                            is_linear_like = (
                                isinstance(child_mod, nn.Linear)
                                or "conv1d" in child_mod.__class__.__name__.lower()
                                or hasattr(child_mod, "nf")
                                or hasattr(child_mod, "base_layer")
                            )
                            if is_linear_like and not isinstance(child_mod, CURLinear):
                                self._compress_single_layer(child_full_name, child_mod)
            elif "expert" in name_lower:
                is_linear_like = (
                    isinstance(module, nn.Linear)
                    or "conv1d" in module.__class__.__name__.lower()
                    or hasattr(module, "nf")
                    or hasattr(module, "base_layer")
                )
                if is_linear_like and not isinstance(module, CURLinear):
                    experts_detected = True
                    self._compress_single_layer(name, module)

        # Fallback to FFN linear layers if no explicit MoE structure was found
        _ATTN_NAMES = frozenset(
            {
                "query",
                "key",
                "value",
                "q_lin",
                "k_lin",
                "v_lin",
                "q_proj",
                "k_proj",
                "v_proj",
                "out_lin",
                "out_proj",
                "o_proj",
                "c_attn",
            }
        )
        if len(self.compressed_layers) == 0:
            visited = set()
            for name, module in list(self.model.named_modules()):
                target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
                if id(target_mod) in visited:
                    continue
                visited.add(id(target_mod))
                name_lower = name.lower()
                if any(
                    skip in name_lower
                    for skip in [
                        "classifier",
                        "score",
                        "lm_head",
                        "pooler",
                        "embedding",
                        "router",
                        "gate",
                    ]
                ):
                    continue

                last_part = name.split(".")[-1]
                is_attn_proj = (last_part in _ATTN_NAMES) or any(
                    k in name_lower for k in ("attn", "attention")
                )

                # In target_attention_only mode (or default MoE behavior), we protect attention projections
                if is_attn_proj:
                    if self.target_attention_only:
                        # Freeze attention layers to protect them during fine-tuning
                        for p in module.parameters():
                            p.requires_grad = False
                    continue

                is_linear_like = (
                    isinstance(module, nn.Linear)
                    or "conv1d" in module.__class__.__name__.lower()
                    or hasattr(module, "nf")
                    or hasattr(module, "base_layer")
                )
                if is_linear_like and not isinstance(module, CURLinear):
                    self._compress_single_layer(name, module)

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        avg_err = (
            (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
        )
        effective_param_ratio = (
            self.total_comp_decomp_params / self.total_orig_decomp_params
            if self.total_orig_decomp_params > 0
            else 1.0
        )
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "avg_reconstruction_error": avg_err,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "cur_method": self.method,
                "cur_moe_rank": self.rank,
                "min_rank": self.min_rank,
                "deterministic": self.deterministic,
                "experts_detected": experts_detected,
                "target_attention_only": self.target_attention_only,
                "effective_parameter_ratio": round(effective_param_ratio, 4),
            }
        )
        logger.info(
            "Compressor cur_moe finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

Compresses Mixture of Experts (MoE) layers and dense FFN layers using CUR cross-approximation.

Args
-----=
model: PyTorch model containing MoE or transformer layers.
rank: Target rank for CUR decomposition.
min_rank: Minimum allowable rank (default 8).
method: 'leverage' (statistical leverage scores) or 'energy' (Frobenius energy).
deterministic: Whether to use deterministic top-k selection (default True).
max_error: Maximum allowed relative Frobenius error (default 0.90).
fallback_to_original: Whether to preserve original layer if error > max_error (default True).
max_retries: Number of adaptive rank escalations (default 2).
target_attention_only: If True, skip attention projections and compress only FFN/MoE expert layers.
cfg: Optional configuration object.

Ancestors

BaseCompressor

Methods

def compress(self, model: nn.Module | None = None) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
    """Compress expert linear layers in MoE architecture, or fallback to FFN layers."""
    if model is not None:
        self.model = model
        self._original_params = self._count_params(self.model)

    self.compressed_layers = []
    self.fallback_layers = []
    self.layer_errors = {}
    self.layer_ranks = {}
    self.total_orig_decomp_params = 0
    self.total_comp_decomp_params = 0
    experts_detected = False

    # Scan for explicit MoE modules (module.experts or expert submodules)
    for name, module in list(self.model.named_modules()):
        name_lower = name.lower()
        if any(
            k in name_lower
            for k in [
                "router",
                "gate",
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
            ]
        ):
            continue

        if hasattr(module, "experts"):
            experts_container = module.experts
            if isinstance(experts_container, (nn.ModuleList, list)):
                experts_detected = True
                for exp_idx, expert in enumerate(experts_container):
                    for child_name, child_mod in list(expert.named_children()):
                        child_full_name = (
                            f"{name}.experts.{exp_idx}.{child_name}"
                            if name
                            else f"experts.{exp_idx}.{child_name}"
                        )
                        is_linear_like = (
                            isinstance(child_mod, nn.Linear)
                            or "conv1d" in child_mod.__class__.__name__.lower()
                            or hasattr(child_mod, "nf")
                            or hasattr(child_mod, "base_layer")
                        )
                        if is_linear_like and not isinstance(child_mod, CURLinear):
                            self._compress_single_layer(child_full_name, child_mod)
        elif "expert" in name_lower:
            is_linear_like = (
                isinstance(module, nn.Linear)
                or "conv1d" in module.__class__.__name__.lower()
                or hasattr(module, "nf")
                or hasattr(module, "base_layer")
            )
            if is_linear_like and not isinstance(module, CURLinear):
                experts_detected = True
                self._compress_single_layer(name, module)

    # Fallback to FFN linear layers if no explicit MoE structure was found
    _ATTN_NAMES = frozenset(
        {
            "query",
            "key",
            "value",
            "q_lin",
            "k_lin",
            "v_lin",
            "q_proj",
            "k_proj",
            "v_proj",
            "out_lin",
            "out_proj",
            "o_proj",
            "c_attn",
        }
    )
    if len(self.compressed_layers) == 0:
        visited = set()
        for name, module in list(self.model.named_modules()):
            target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
            if id(target_mod) in visited:
                continue
            visited.add(id(target_mod))
            name_lower = name.lower()
            if any(
                skip in name_lower
                for skip in [
                    "classifier",
                    "score",
                    "lm_head",
                    "pooler",
                    "embedding",
                    "router",
                    "gate",
                ]
            ):
                continue

            last_part = name.split(".")[-1]
            is_attn_proj = (last_part in _ATTN_NAMES) or any(
                k in name_lower for k in ("attn", "attention")
            )

            # In target_attention_only mode (or default MoE behavior), we protect attention projections
            if is_attn_proj:
                if self.target_attention_only:
                    # Freeze attention layers to protect them during fine-tuning
                    for p in module.parameters():
                        p.requires_grad = False
                continue

            is_linear_like = (
                isinstance(module, nn.Linear)
                or "conv1d" in module.__class__.__name__.lower()
                or hasattr(module, "nf")
                or hasattr(module, "base_layer")
            )
            if is_linear_like and not isinstance(module, CURLinear):
                self._compress_single_layer(name, module)

    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    avg_err = (
        (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
    )
    effective_param_ratio = (
        self.total_comp_decomp_params / self.total_orig_decomp_params
        if self.total_orig_decomp_params > 0
        else 1.0
    )
    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "layer_errors": self.layer_errors,
            "layer_ranks": self.layer_ranks,
            "avg_reconstruction_error": avg_err,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "cur_method": self.method,
            "cur_moe_rank": self.rank,
            "min_rank": self.min_rank,
            "deterministic": self.deterministic,
            "experts_detected": experts_detected,
            "target_attention_only": self.target_attention_only,
            "effective_parameter_ratio": round(effective_param_ratio, 4),
        }
    )
    logger.info(
        "Compressor cur_moe finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        self.max_error,
    )
    return self.model, stats

Compress expert linear layers in MoE architecture, or fallback to FFN layers.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class CacheNotesCompressor (model: nn.Module, task_type: str = 'qa', task_description: str | None = None, use_learned_cpt: bool = True, keep_ratio: float = 0.5, max_error: float = 0.35, fallback_to_original: bool = True, protect_critical: bool = True, target_attention_only: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, skip_attention_qkv: bool = True, recalibrate_ln: bool = True, tokenizer: Any | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, target_modules: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("cachenotes")
@COMPRESSOR_REGISTRY.register("cachenotes_compressor")
class CacheNotesCompressor(BaseCompressor):
    is_custom_attention: bool = True
    """
    CacheNotes Compressor: Decouples static weight compression from KV cache dynamic pruning
    with continuous learnable CPT Generator.

    Production hardening features:
    - Trainable CPTGenerator mapping task descriptions to per-feature importance vectors.
    - Support for LoRALinear wrappers, nn.Linear, and Conv1D (GPT-2).
    - Preserves existing LoRALinear wrappers without creating redundant layers.
    - target_attention_only parameter skipping FFN/MLP and preserving classification heads.
    - Robust config handling (_cfg_get) supporting None, dict, and namespace objects.
    - Fallback freezing: parameters in any fallback layer are frozen (requires_grad=False).
    - Checkpoint saving and loading for CPTGenerator.
    - Comprehensive metric tracking: actual ranks, reconstruction error, and KV cache saving.
    """

    def __init__(
        self,
        model: nn.Module,
        task_type: str = "qa",
        task_description: str | None = None,
        use_learned_cpt: bool = True,
        keep_ratio: float = 0.5,
        max_error: float = 0.35,
        fallback_to_original: bool = True,
        protect_critical: bool = True,
        target_attention_only: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        skip_attention_qkv: bool = True,
        recalibrate_ln: bool = True,
        tokenizer: Any | None = None,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        target_modules: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="cachenotes", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))

        task_type = str(_cfg_get(cfg, "cachenotes_task", _cfg_get(cfg, "task_type", task_type)))
        task_description = str(
            _cfg_get(
                cfg,
                "cachenotes_task_description",
                _cfg_get(cfg, "task_description", task_description or task_type),
            )
        )
        use_learned_cpt = bool(
            _cfg_get(
                cfg,
                "cachenotes_use_learned_cpt",
                _cfg_get(cfg, "use_learned_cpt", use_learned_cpt),
            )
        )
        keep_ratio = float(
            _cfg_get(cfg, "cachenotes_keep_ratio", _cfg_get(cfg, "keep_ratio", keep_ratio))
        )
        max_error = float(
            _cfg_get(cfg, "cachenotes_max_error", _cfg_get(cfg, "max_error", max_error))
        )
        fallback_to_original = bool(
            _cfg_get(
                cfg,
                "cachenotes_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        protect_critical = bool(_cfg_get(cfg, "protect_critical", protect_critical))
        target_attention_only = bool(
            _cfg_get(
                cfg,
                "cachenotes_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        lora_r = int(
            _cfg_get(
                cfg,
                "cachenotes_lora_r",
                _cfg_get(cfg, "lora_r", _cfg_get(cfg, "lora_rank", lora_r)),
            )
        )
        lora_alpha = float(
            _cfg_get(cfg, "cachenotes_lora_alpha", _cfg_get(cfg, "lora_alpha", lora_alpha))
        )
        skip_attention_qkv = bool(
            _cfg_get(
                cfg,
                "cachenotes_skip_attention_qkv",
                _cfg_get(cfg, "skip_attention_qkv", skip_attention_qkv),
            )
        )
        recalibrate_ln = bool(
            _cfg_get(
                cfg,
                "cachenotes_recalibrate_ln",
                _cfg_get(cfg, "recalibrate_ln", recalibrate_ln),
            )
        )

        self.task_type = task_type
        self.task_description = task_description
        self.use_learned_cpt = use_learned_cpt
        self.keep_ratio = max(0.05, min(keep_ratio, 1.0))
        self.max_error = max_error
        self.fallback_to_original = fallback_to_original
        self.protect_critical = protect_critical
        self.target_attention_only = target_attention_only
        self.lora_r = lora_r
        self.lora_alpha = lora_alpha
        self.skip_attention_qkv = skip_attention_qkv
        self.recalibrate_ln = recalibrate_ln
        self.target_modules = target_modules

        self.tokenizer = tokenizer or kwargs.get("tokenizer", getattr(model, "tokenizer", None))
        self.calibration_loader = calibration_loader or kwargs.get("calibration_loader")
        self.calibration_texts = calibration_texts or kwargs.get("calibration_texts")
        self.target_device = (
            torch.device(device)
            if device is not None
            else (next(model.parameters(), torch.tensor(0.0)).device)
        )
        self.comp_config = cfg

        # Initialize learnable CPT Generator
        self.cpt_generator = CPTGenerator().to(self.target_device)
        self.cpt_generator_trained = False

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.actual_ranks: dict[str, int] = {}
        self.stats: dict[str, Any] = {}

    def get_importance(self, d: int, device: torch.device) -> torch.Tensor:
        """
        Retrieves feature importance vector using learned CPT generator if trained/active,
        otherwise falling back to static heuristic.
        """
        if self.use_learned_cpt and self.cpt_generator is not None:
            task_query = self.task_description or self.task_type
            return self.cpt_generator(task_query, d, device=device)
        logger.debug("Falling back to static CPT heuristic for task %s", self.task_type)
        return CompressionPlanningTokens.generate_importance(self.task_type, d, device=device)

    def train_cpt_generator(
        self,
        task_descriptions: list[str] | None = None,
        sample_layers: list[tuple[str, nn.Module]] | None = None,
        epochs: int = 5,
        lr: float = 1e-4,
    ) -> float:
        """
        Train CPT generator on sample layers across diverse task descriptions,
        minimizing reconstruction error after CPT-weighted SVD compression.
        """
        if task_descriptions is None:
            task_descriptions = [
                "question answering: extract concise answer from supporting context passage",
                "code generation: implement efficient python function and classes",
                "summarization: produce concise abstractive summary of input document",
                "mathematical reasoning: solve multi-step theorem and calculus equations",
                "sentiment classification: predict positive or negative customer review sentiment",
            ]

        if sample_layers is None:
            sample_layers = [
                (name, mod)
                for name, mod in self.model.named_modules()
                if self._is_target_module(name, mod)
            ][:4]  # train on up to 4 sample layers for efficiency

        if not sample_layers:
            logger.warning("No sample layers found to train CPT generator.")
            return 0.0

        self.cpt_generator.train()
        for p in self.cpt_generator.parameters():
            p.requires_grad = True

        optimizer = torch.optim.AdamW(self.cpt_generator.parameters(), lr=lr)

        total_loss = 0.0
        steps = 0

        for epoch in range(epochs):
            for task in task_descriptions:
                for name, module in sample_layers:
                    target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
                    cls_name = module.__class__.__name__.lower()
                    is_conv1d = (
                        "conv1d" in cls_name or hasattr(module, "nf") or hasattr(target_mod, "nf")
                    )
                    raw_w = target_mod.weight.data.float().to(self.target_device)
                    W = raw_w.t().contiguous() if is_conv1d else raw_w
                    out_f, in_f = W.shape

                    # Generate learned importance
                    importance = self.cpt_generator(task, out_f, device=W.device)
                    W_cpt = W * importance.unsqueeze(1)

                    r = max(
                        1,
                        min(int(min(out_f, in_f) * self.keep_ratio), min(out_f, in_f)),
                    )
                    try:
                        U, S, Vt = torch.linalg.svd(W_cpt, full_matrices=False)
                        W_lowrank = torch.matmul(U[:, :r] * S[:r].unsqueeze(0), Vt[:r, :])
                        W_approx = W_lowrank / importance.unsqueeze(1)
                        loss = torch.norm(W - W_approx, p="fro") / max(
                            torch.norm(W, p="fro").item(), 1e-6
                        )

                        optimizer.zero_grad()
                        loss.backward()
                        torch.nn.utils.clip_grad_norm_(
                            self.cpt_generator.parameters(), max_norm=1.0
                        )
                        optimizer.step()

                        total_loss += loss.item()
                        steps += 1
                    except Exception as e:
                        logger.debug("CPT generator training step exception: %s", e)

        self.cpt_generator.eval()
        self.cpt_generator_trained = True
        for p in self.cpt_generator.parameters():
            p.requires_grad = False

        avg_loss = total_loss / max(1, steps)
        logger.info(
            "CPT generator trained successfully over %d steps with avg loss %.4f",
            steps,
            avg_loss,
        )
        return avg_loss

    def save_cpt_generator(self, path: str | Path) -> None:
        """Save CPT generator state dict to file."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        torch.save(self.cpt_generator.state_dict(), path)
        logger.info("Saved CPT generator checkpoint to %s", path)

    def load_cpt_generator(self, path: str | Path) -> bool:
        """Load CPT generator state dict from file."""
        path = Path(path)
        if path.exists():
            try:
                state_dict = torch.load(path, map_location=self.target_device)
                self.cpt_generator.load_state_dict(state_dict)
                self.cpt_generator_trained = True
                self.cpt_generator.eval()
                for p in self.cpt_generator.parameters():
                    p.requires_grad = False
                logger.info("Loaded CPT generator checkpoint from %s", path)
                return True
            except Exception as e:
                logger.warning("Failed to load CPT generator from %s: %s", path, e)
        return False

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        if self.target_attention_only:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            name_lower = name.lower()
            if any(
                p in name_lower
                for p in [
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                ]
            ):
                return True
        return False

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        if isinstance(module, CacheNotesLoRALinear):
            return False

        # Do not target inner submodules of LoRA or wrappers directly
        if any(
            name.endswith(sfx)
            for sfx in [".base_layer", ".lora_A", ".lora_B", ".lora_dropout", ".linear"]
        ):
            return False

        target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
        cls_name = module.__class__.__name__.lower()
        target_cls = target_mod.__class__.__name__.lower()

        is_linear_like = (
            isinstance(module, nn.Linear)
            or isinstance(target_mod, nn.Linear)
            or (
                Conv1D is not None
                and (isinstance(module, Conv1D) or isinstance(target_mod, Conv1D))
            )
            or "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(module, "nf")
            or hasattr(target_mod, "nf")
        )
        if not is_linear_like:
            return False

        name_lower = name.lower()
        # Protect classification / output heads
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embed",
                "logits_processor",
                "head",
                "pre_classifier",
            ]
        ):
            return False

        if self._is_gqa_asymmetric(name, module):
            return False

        if self.target_attention_only:
            # Skip FFN / MLP layers
            ffn_patterns = [
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            ]
            if any(p in name_lower for p in ffn_patterns):
                return False
            # In RoBERTa, layer.X.output.dense is FFN, while attention.output.dense is attention
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            # Match attention projections
            attn_patterns = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
                "in_proj_qkv",
                "out_proj",
            ]
            if not any(p in name_lower for p in attn_patterns):
                return False
        else:
            if self.protect_critical and not should_compress_linear(self.model, name):
                return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """Compress linear layers using task-guided CPT-weighted low-rank SVD."""
        # Auto-train learnable CPT generator if active and untrained
        if self.use_learned_cpt and not self.cpt_generator_trained:
            logger.info(
                "CPT generator not yet trained. Running initial task-aware training pass..."
            )
            self.train_cpt_generator(epochs=3, lr=1e-4)

        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.actual_ranks = {}
        seen_weights: set[int] = set()

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
            module.__class__.__name__.lower()
            is_conv1d = (
                (
                    Conv1D is not None
                    and (isinstance(module, Conv1D) or isinstance(target_mod, Conv1D))
                )
                or hasattr(module, "nf")
                or hasattr(target_mod, "nf")
            )
            has_lora = (hasattr(module, "lora_A") and module.lora_A is not None) or hasattr(
                module, "base_layer"
            )

            w_id = id(target_mod.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            raw_w = target_mod.weight.data
            device = raw_w.device
            dtype = raw_w.dtype

            # Standardize shape to [out_features, in_features] in FP32
            W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
            out_f, in_f = W.shape

            # Retrieve CPT importance modulation per output feature (learned or static)
            importance = self.get_importance(out_f, device=device)

            # Weight rows by CPT importance
            W_cpt = W * importance.unsqueeze(1)

            # Compute SVD in CPT-weighted space
            r = max(1, min(int(min(out_f, in_f) * self.keep_ratio), min(out_f, in_f)))
            self.actual_ranks[name] = r
            try:
                U, S, Vt = torch.linalg.svd(W_cpt, full_matrices=False)
                # Low-rank approximation
                W_lowrank = torch.matmul(U[:, :r] * S[:r].unsqueeze(0), Vt[:r, :])
                # Unweight back to physical space: all channels remain active and non-zero
                W_approx = W_lowrank / importance.unsqueeze(1)
                if torch.isnan(W_approx).any() or torch.isinf(W_approx).any():
                    W_approx = W
            except Exception as e:
                logger.warning("SVD failed in CacheNotes on %s: %s", name, e)
                W_approx = W

            w_norm = torch.norm(W, p="fro").item()
            err = float(torch.norm(W - W_approx, p="fro").item() / max(w_norm, 1e-12))
            self.layer_errors[name] = err

            if torch.isnan(W_approx).any() or torch.isinf(W_approx).any() or err > self.max_error:
                logger.warning(
                    "CacheNotes reconstruction error %.4f on %s exceeded threshold %.4f.",
                    err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    logger.info(
                        "Skipping CacheNotes compression for %s (fallback to original).",
                        name,
                    )
                    self.fallback_layers.append(name)
                    # Freeze fallback layers
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

            # Update weights / module
            if has_lora:
                # Layer is already a LoRALinear wrapper: update base layer weights in-place
                if is_conv1d:
                    target_mod.weight.data.copy_(
                        W_approx.t().contiguous().to(dtype=dtype, device=device)
                    )
                else:
                    target_mod.weight.data.copy_(
                        W_approx.contiguous().to(dtype=dtype, device=device)
                    )
                # Keep base weights frozen, but keep LoRA adapters trainable
                target_mod.weight.requires_grad = False
                if getattr(target_mod, "bias", None) is not None and target_mod.bias is not None:
                    target_mod.bias.requires_grad = False
                for p_name, p in module.named_parameters():
                    if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                        p.requires_grad = True
            elif self.lora_r > 0:
                # Not LoRALinear yet, and lora_r > 0: wrap in CacheNotesLoRALinear
                cn_layer = CacheNotesLoRALinear.from_linear(
                    target_mod,
                    weight=W_approx,
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                )
                if name == "":
                    self.model = cn_layer
                else:
                    replace_module(self.model, name, cn_layer)
            else:
                # Pure static compression (lora_r == 0)
                if is_conv1d:
                    target_mod.weight.data.copy_(
                        W_approx.t().contiguous().to(dtype=dtype, device=device)
                    )
                else:
                    target_mod.weight.data.copy_(
                        W_approx.contiguous().to(dtype=dtype, device=device)
                    )

            self.compressed_layers.append(name)

        # LayerNorm Recalibration
        if self.recalibrate_ln and (
            self.calibration_loader is not None or self.calibration_texts is not None
        ):
            try:
                recalibrate_layernorms(
                    self.model,
                    tokenizer=self.tokenizer,
                    calibration_loader=self.calibration_loader,
                    calibration_texts=self.calibration_texts,
                    device=self.target_device,
                )
            except Exception as e:
                logger.warning("CacheNotes LayerNorm recalibration skipped: %s", e)

        # Post-compression forward sanity check
        forward_check_success = False
        try:
            self.model.eval()
            with torch.no_grad():
                d = next(self.model.parameters()).device
                dummy = torch.randint(0, 100, (1, 16), device=d)
                self.model(dummy)
            forward_check_success = True
        except Exception as e:
            logger.debug("CacheNotes post-compression forward check exception: %s", e)

        self._compressed_params = self._count_params(self.model)
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        avg_rank = (
            float(
                sum(self.actual_ranks[k] for k in self.compressed_layers)
                / max(1, len(self.compressed_layers))
            )
            if self.compressed_layers
            else 0.0
        )
        stats = self.get_compression_stats()
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "actual_ranks": self.actual_ranks,
                "avg_rank": avg_rank,
                "avg_reconstruction_error": avg_err,
                "task_type": self.task_type,
                "task_description": self.task_description,
                "use_learned_cpt": self.use_learned_cpt,
                "cpt_generator_trained": self.cpt_generator_trained,
                "keep_ratio": self.keep_ratio,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "target_attention_only": self.target_attention_only,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "forward_check_success": forward_check_success,
                "recalibrate_ln": self.recalibrate_ln,
                "method": "cachenotes_learned_cpt_svd",
                "kv_cache_budget_saving": f"{(1.0 - self.keep_ratio) * 100:.1f}% Footprint Reduction",
            }
        )
        self.stats = stats
        self.model.compression_stats = stats
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: CacheNotes Compressor: Decouples static weight compression from KV cache dynamic pruning with continuous learnable CPT Generator.

Production hardening features: - Trainable CPTGenerator mapping task descriptions to per-feature importance vectors. - Support for LoRALinear wrappers, nn.Linear, and Conv1D (GPT-2). - Preserves existing LoRALinear wrappers without creating redundant layers. - target_attention_only parameter skipping FFN/MLP and preserving classification heads. - Robust config handling (_cfg_get) supporting None, dict, and namespace objects. - Fallback freezing: parameters in any fallback layer are frozen (requires_grad=False). - Checkpoint saving and loading for CPTGenerator. - Comprehensive metric tracking: actual ranks, reconstruction error, and KV cache saving.

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """Compress linear layers using task-guided CPT-weighted low-rank SVD."""
    # Auto-train learnable CPT generator if active and untrained
    if self.use_learned_cpt and not self.cpt_generator_trained:
        logger.info(
            "CPT generator not yet trained. Running initial task-aware training pass..."
        )
        self.train_cpt_generator(epochs=3, lr=1e-4)

    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    self.layer_errors = {}
    self.actual_ranks = {}
    seen_weights: set[int] = set()

    for name, module in list(self.model.named_modules()):
        if not self._is_target_module(name, module):
            continue

        target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
        module.__class__.__name__.lower()
        is_conv1d = (
            (
                Conv1D is not None
                and (isinstance(module, Conv1D) or isinstance(target_mod, Conv1D))
            )
            or hasattr(module, "nf")
            or hasattr(target_mod, "nf")
        )
        has_lora = (hasattr(module, "lora_A") and module.lora_A is not None) or hasattr(
            module, "base_layer"
        )

        w_id = id(target_mod.weight)
        if w_id in seen_weights:
            self.skipped_layers.append(name)
            continue
        seen_weights.add(w_id)

        raw_w = target_mod.weight.data
        device = raw_w.device
        dtype = raw_w.dtype

        # Standardize shape to [out_features, in_features] in FP32
        W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
        out_f, in_f = W.shape

        # Retrieve CPT importance modulation per output feature (learned or static)
        importance = self.get_importance(out_f, device=device)

        # Weight rows by CPT importance
        W_cpt = W * importance.unsqueeze(1)

        # Compute SVD in CPT-weighted space
        r = max(1, min(int(min(out_f, in_f) * self.keep_ratio), min(out_f, in_f)))
        self.actual_ranks[name] = r
        try:
            U, S, Vt = torch.linalg.svd(W_cpt, full_matrices=False)
            # Low-rank approximation
            W_lowrank = torch.matmul(U[:, :r] * S[:r].unsqueeze(0), Vt[:r, :])
            # Unweight back to physical space: all channels remain active and non-zero
            W_approx = W_lowrank / importance.unsqueeze(1)
            if torch.isnan(W_approx).any() or torch.isinf(W_approx).any():
                W_approx = W
        except Exception as e:
            logger.warning("SVD failed in CacheNotes on %s: %s", name, e)
            W_approx = W

        w_norm = torch.norm(W, p="fro").item()
        err = float(torch.norm(W - W_approx, p="fro").item() / max(w_norm, 1e-12))
        self.layer_errors[name] = err

        if torch.isnan(W_approx).any() or torch.isinf(W_approx).any() or err > self.max_error:
            logger.warning(
                "CacheNotes reconstruction error %.4f on %s exceeded threshold %.4f.",
                err,
                name,
                self.max_error,
            )
            if self.fallback_to_original:
                logger.info(
                    "Skipping CacheNotes compression for %s (fallback to original).",
                    name,
                )
                self.fallback_layers.append(name)
                # Freeze fallback layers
                for p in module.parameters():
                    p.requires_grad = False
                continue

        # Update weights / module
        if has_lora:
            # Layer is already a LoRALinear wrapper: update base layer weights in-place
            if is_conv1d:
                target_mod.weight.data.copy_(
                    W_approx.t().contiguous().to(dtype=dtype, device=device)
                )
            else:
                target_mod.weight.data.copy_(
                    W_approx.contiguous().to(dtype=dtype, device=device)
                )
            # Keep base weights frozen, but keep LoRA adapters trainable
            target_mod.weight.requires_grad = False
            if getattr(target_mod, "bias", None) is not None and target_mod.bias is not None:
                target_mod.bias.requires_grad = False
            for p_name, p in module.named_parameters():
                if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                    p.requires_grad = True
        elif self.lora_r > 0:
            # Not LoRALinear yet, and lora_r > 0: wrap in CacheNotesLoRALinear
            cn_layer = CacheNotesLoRALinear.from_linear(
                target_mod,
                weight=W_approx,
                lora_r=self.lora_r,
                lora_alpha=self.lora_alpha,
            )
            if name == "":
                self.model = cn_layer
            else:
                replace_module(self.model, name, cn_layer)
        else:
            # Pure static compression (lora_r == 0)
            if is_conv1d:
                target_mod.weight.data.copy_(
                    W_approx.t().contiguous().to(dtype=dtype, device=device)
                )
            else:
                target_mod.weight.data.copy_(
                    W_approx.contiguous().to(dtype=dtype, device=device)
                )

        self.compressed_layers.append(name)

    # LayerNorm Recalibration
    if self.recalibrate_ln and (
        self.calibration_loader is not None or self.calibration_texts is not None
    ):
        try:
            recalibrate_layernorms(
                self.model,
                tokenizer=self.tokenizer,
                calibration_loader=self.calibration_loader,
                calibration_texts=self.calibration_texts,
                device=self.target_device,
            )
        except Exception as e:
            logger.warning("CacheNotes LayerNorm recalibration skipped: %s", e)

    # Post-compression forward sanity check
    forward_check_success = False
    try:
        self.model.eval()
        with torch.no_grad():
            d = next(self.model.parameters()).device
            dummy = torch.randint(0, 100, (1, 16), device=d)
            self.model(dummy)
        forward_check_success = True
    except Exception as e:
        logger.debug("CacheNotes post-compression forward check exception: %s", e)

    self._compressed_params = self._count_params(self.model)
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )
    avg_rank = (
        float(
            sum(self.actual_ranks[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
        )
        if self.compressed_layers
        else 0.0
    )
    stats = self.get_compression_stats()
    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "fallback_layers": self.fallback_layers,
            "skipped_layers": self.skipped_layers,
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "num_skipped": len(self.skipped_layers),
            "layer_errors": self.layer_errors,
            "actual_ranks": self.actual_ranks,
            "avg_rank": avg_rank,
            "avg_reconstruction_error": avg_err,
            "task_type": self.task_type,
            "task_description": self.task_description,
            "use_learned_cpt": self.use_learned_cpt,
            "cpt_generator_trained": self.cpt_generator_trained,
            "keep_ratio": self.keep_ratio,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "target_attention_only": self.target_attention_only,
            "lora_r": self.lora_r,
            "lora_alpha": self.lora_alpha,
            "forward_check_success": forward_check_success,
            "recalibrate_ln": self.recalibrate_ln,
            "method": "cachenotes_learned_cpt_svd",
            "kv_cache_budget_saving": f"{(1.0 - self.keep_ratio) * 100:.1f}% Footprint Reduction",
        }
    )
    self.stats = stats
    self.model.compression_stats = stats
    return self.model, stats

Compress linear layers using task-guided CPT-weighted low-rank SVD.

def get_importance(self, d: int, device: torch.device) ‑> torch.Tensor

Expand source code

def get_importance(self, d: int, device: torch.device) -> torch.Tensor:
    """
    Retrieves feature importance vector using learned CPT generator if trained/active,
    otherwise falling back to static heuristic.
    """
    if self.use_learned_cpt and self.cpt_generator is not None:
        task_query = self.task_description or self.task_type
        return self.cpt_generator(task_query, d, device=device)
    logger.debug("Falling back to static CPT heuristic for task %s", self.task_type)
    return CompressionPlanningTokens.generate_importance(self.task_type, d, device=device)

Retrieves feature importance vector using learned CPT generator if trained/active, otherwise falling back to static heuristic.

def load_cpt_generator(self, path: str | Path) ‑> bool

Expand source code

def load_cpt_generator(self, path: str | Path) -> bool:
    """Load CPT generator state dict from file."""
    path = Path(path)
    if path.exists():
        try:
            state_dict = torch.load(path, map_location=self.target_device)
            self.cpt_generator.load_state_dict(state_dict)
            self.cpt_generator_trained = True
            self.cpt_generator.eval()
            for p in self.cpt_generator.parameters():
                p.requires_grad = False
            logger.info("Loaded CPT generator checkpoint from %s", path)
            return True
        except Exception as e:
            logger.warning("Failed to load CPT generator from %s: %s", path, e)
    return False

Load CPT generator state dict from file.

def save_cpt_generator(self, path: str | Path) ‑> None

Expand source code

def save_cpt_generator(self, path: str | Path) -> None:
    """Save CPT generator state dict to file."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(self.cpt_generator.state_dict(), path)
    logger.info("Saved CPT generator checkpoint to %s", path)

Save CPT generator state dict to file.

def train_cpt_generator(self, task_descriptions: list[str] | None = None, sample_layers: list[tuple[str, nn.Module]] | None = None, epochs: int = 5, lr: float = 0.0001) ‑> float

Expand source code

def train_cpt_generator(
    self,
    task_descriptions: list[str] | None = None,
    sample_layers: list[tuple[str, nn.Module]] | None = None,
    epochs: int = 5,
    lr: float = 1e-4,
) -> float:
    """
    Train CPT generator on sample layers across diverse task descriptions,
    minimizing reconstruction error after CPT-weighted SVD compression.
    """
    if task_descriptions is None:
        task_descriptions = [
            "question answering: extract concise answer from supporting context passage",
            "code generation: implement efficient python function and classes",
            "summarization: produce concise abstractive summary of input document",
            "mathematical reasoning: solve multi-step theorem and calculus equations",
            "sentiment classification: predict positive or negative customer review sentiment",
        ]

    if sample_layers is None:
        sample_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ][:4]  # train on up to 4 sample layers for efficiency

    if not sample_layers:
        logger.warning("No sample layers found to train CPT generator.")
        return 0.0

    self.cpt_generator.train()
    for p in self.cpt_generator.parameters():
        p.requires_grad = True

    optimizer = torch.optim.AdamW(self.cpt_generator.parameters(), lr=lr)

    total_loss = 0.0
    steps = 0

    for epoch in range(epochs):
        for task in task_descriptions:
            for name, module in sample_layers:
                target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
                cls_name = module.__class__.__name__.lower()
                is_conv1d = (
                    "conv1d" in cls_name or hasattr(module, "nf") or hasattr(target_mod, "nf")
                )
                raw_w = target_mod.weight.data.float().to(self.target_device)
                W = raw_w.t().contiguous() if is_conv1d else raw_w
                out_f, in_f = W.shape

                # Generate learned importance
                importance = self.cpt_generator(task, out_f, device=W.device)
                W_cpt = W * importance.unsqueeze(1)

                r = max(
                    1,
                    min(int(min(out_f, in_f) * self.keep_ratio), min(out_f, in_f)),
                )
                try:
                    U, S, Vt = torch.linalg.svd(W_cpt, full_matrices=False)
                    W_lowrank = torch.matmul(U[:, :r] * S[:r].unsqueeze(0), Vt[:r, :])
                    W_approx = W_lowrank / importance.unsqueeze(1)
                    loss = torch.norm(W - W_approx, p="fro") / max(
                        torch.norm(W, p="fro").item(), 1e-6
                    )

                    optimizer.zero_grad()
                    loss.backward()
                    torch.nn.utils.clip_grad_norm_(
                        self.cpt_generator.parameters(), max_norm=1.0
                    )
                    optimizer.step()

                    total_loss += loss.item()
                    steps += 1
                except Exception as e:
                    logger.debug("CPT generator training step exception: %s", e)

    self.cpt_generator.eval()
    self.cpt_generator_trained = True
    for p in self.cpt_generator.parameters():
        p.requires_grad = False

    avg_loss = total_loss / max(1, steps)
    logger.info(
        "CPT generator trained successfully over %d steps with avg loss %.4f",
        steps,
        avg_loss,
    )
    return avg_loss

Train CPT generator on sample layers across diverse task descriptions, minimizing reconstruction error after CPT-weighted SVD compression.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class CacheNotesKVCache (max_cache_len: int = 128, task_type: str = 'qa', keep_ratio: float = 0.3, num_sinks: int = 4, cpt_generator: CPTGenerator | None = None)

Expand source code

class CacheNotesKVCache:
    """
    Runtime task-guided dynamic KV cache compressor.
    Evaluates token task-relevance dynamically and prunes non-essential context tokens,
    supporting variable sequence lengths during prefill and autoregressive decoding.
    """

    def __init__(
        self,
        max_cache_len: int = 128,
        task_type: str = "qa",
        keep_ratio: float = 0.3,
        num_sinks: int = 4,
        cpt_generator: CPTGenerator | None = None,
    ) -> None:
        self.max_cache_len = int(max_cache_len)
        self.task_type = task_type
        self.keep_ratio = float(keep_ratio)
        self.num_sinks = int(num_sinks)
        self.cpt_generator = cpt_generator
        self.k_cache: torch.Tensor | None = None
        self.v_cache: torch.Tensor | None = None

    def reset(self) -> None:
        """Clear cache state to prevent memory leaks across sequences."""
        self.k_cache = None
        self.v_cache = None

    @torch.no_grad()
    def update_and_compress(
        self,
        new_k: torch.Tensor,
        new_v: torch.Tensor,
        is_new_sequence: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Append new KV tokens and apply task-aware retention when exceeding capacity.
        Shapes: (B, H, N, D).
        Automatically resets cache on new multi-token prompt prefill (N > 1) or when is_new_sequence=True (H6).
        """
        k_4d = new_k
        v_4d = new_v
        if new_k.ndim == 2:
            k_4d = new_k.unsqueeze(0).unsqueeze(0)
            v_4d = new_v.unsqueeze(0).unsqueeze(0)
        elif new_k.ndim == 3:
            k_4d = new_k.unsqueeze(1)
            v_4d = new_v.unsqueeze(1)

        B, H, N, D = k_4d.shape

        # H6: Prevent cross-sequence cache pollution on prompt prefill (N > 1) or explicit flag
        if is_new_sequence or N > 1:
            self.reset()

        if self.k_cache is None:
            self.k_cache = k_4d
            self.v_cache = v_4d
        else:
            if self.k_cache.shape[0] != B:
                self.k_cache = self.k_cache[:1].expand(B, -1, -1, -1).clone()
                self.v_cache = self.v_cache[:1].expand(B, -1, -1, -1).clone()
            self.k_cache = torch.cat([self.k_cache, k_4d], dim=2)
            self.v_cache = torch.cat([self.v_cache, v_4d], dim=2)

        cur_len = self.k_cache.size(2)
        if cur_len <= self.max_cache_len:
            return self.k_cache, self.v_cache

        sinks = min(self.num_sinks, cur_len)
        retained_cap = max(sinks + 1, self.max_cache_len)

        # Task importance weighting across tokens
        if self.cpt_generator is not None:
            cpt_importance = self.cpt_generator(self.task_type, D, device=self.k_cache.device)
        else:
            cpt_importance = CompressionPlanningTokens.generate_importance(
                self.task_type, D, device=self.k_cache.device
            )

        weighted_k = self.k_cache * cpt_importance.view(1, 1, 1, D)
        token_scores = torch.norm(weighted_k.float(), dim=-1)  # (B, H, cur_len)

        # Always preserve initial sink tokens
        token_scores[:, :, :sinks] += 1e6

        _, top_indices = torch.topk(token_scores, retained_cap, dim=-1)
        top_indices, _ = torch.sort(top_indices, dim=-1)

        idx_exp = top_indices.unsqueeze(-1).expand(-1, -1, -1, D)
        self.k_cache = self.k_cache.gather(2, idx_exp).contiguous()
        self.v_cache = self.v_cache.gather(2, idx_exp).contiguous()

        return self.k_cache, self.v_cache

Runtime task-guided dynamic KV cache compressor. Evaluates token task-relevance dynamically and prunes non-essential context tokens, supporting variable sequence lengths during prefill and autoregressive decoding.

Methods

def reset(self) ‑> None

Expand source code

def reset(self) -> None:
    """Clear cache state to prevent memory leaks across sequences."""
    self.k_cache = None
    self.v_cache = None

Clear cache state to prevent memory leaks across sequences.

def update_and_compress(self, new_k: torch.Tensor, new_v: torch.Tensor, is_new_sequence: bool = False) ‑> tuple[torch.Tensor, torch.Tensor]

Expand source code

@torch.no_grad()
def update_and_compress(
    self,
    new_k: torch.Tensor,
    new_v: torch.Tensor,
    is_new_sequence: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Append new KV tokens and apply task-aware retention when exceeding capacity.
    Shapes: (B, H, N, D).
    Automatically resets cache on new multi-token prompt prefill (N > 1) or when is_new_sequence=True (H6).
    """
    k_4d = new_k
    v_4d = new_v
    if new_k.ndim == 2:
        k_4d = new_k.unsqueeze(0).unsqueeze(0)
        v_4d = new_v.unsqueeze(0).unsqueeze(0)
    elif new_k.ndim == 3:
        k_4d = new_k.unsqueeze(1)
        v_4d = new_v.unsqueeze(1)

    B, H, N, D = k_4d.shape

    # H6: Prevent cross-sequence cache pollution on prompt prefill (N > 1) or explicit flag
    if is_new_sequence or N > 1:
        self.reset()

    if self.k_cache is None:
        self.k_cache = k_4d
        self.v_cache = v_4d
    else:
        if self.k_cache.shape[0] != B:
            self.k_cache = self.k_cache[:1].expand(B, -1, -1, -1).clone()
            self.v_cache = self.v_cache[:1].expand(B, -1, -1, -1).clone()
        self.k_cache = torch.cat([self.k_cache, k_4d], dim=2)
        self.v_cache = torch.cat([self.v_cache, v_4d], dim=2)

    cur_len = self.k_cache.size(2)
    if cur_len <= self.max_cache_len:
        return self.k_cache, self.v_cache

    sinks = min(self.num_sinks, cur_len)
    retained_cap = max(sinks + 1, self.max_cache_len)

    # Task importance weighting across tokens
    if self.cpt_generator is not None:
        cpt_importance = self.cpt_generator(self.task_type, D, device=self.k_cache.device)
    else:
        cpt_importance = CompressionPlanningTokens.generate_importance(
            self.task_type, D, device=self.k_cache.device
        )

    weighted_k = self.k_cache * cpt_importance.view(1, 1, 1, D)
    token_scores = torch.norm(weighted_k.float(), dim=-1)  # (B, H, cur_len)

    # Always preserve initial sink tokens
    token_scores[:, :, :sinks] += 1e6

    _, top_indices = torch.topk(token_scores, retained_cap, dim=-1)
    top_indices, _ = torch.sort(top_indices, dim=-1)

    idx_exp = top_indices.unsqueeze(-1).expand(-1, -1, -1, D)
    self.k_cache = self.k_cache.gather(2, idx_exp).contiguous()
    self.v_cache = self.v_cache.gather(2, idx_exp).contiguous()

    return self.k_cache, self.v_cache

Append new KV tokens and apply task-aware retention when exceeding capacity. Shapes: (B, H, N, D). Automatically resets cache on new multi-token prompt prefill (N > 1) or when is_new_sequence=True (H6).

class CacheNotesLoRALinear (in_features: int, out_features: int, weight: torch.Tensor, bias: torch.Tensor | None = None, lora_r: int = 0, lora_alpha: float = 1.0, device: torch.device | None = None, dtype: torch.dtype | None = None)

Expand source code

class CacheNotesLoRALinear(nn.Module):
    """
    Linear layer containing CacheNotes compressed base weights (frozen)
    with optional trainable LoRA adapters (lora_A, lora_B).
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        weight: torch.Tensor,
        bias: torch.Tensor | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features  # Conv1D compatibility
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        target_device = device if device is not None else weight.device
        target_dtype = dtype if dtype is not None else weight.dtype

        self.weight = nn.Parameter(
            weight.to(device=target_device, dtype=target_dtype),
            requires_grad=False,
        )
        if bias is not None:
            self.bias = nn.Parameter(
                bias.to(device=target_device, dtype=target_dtype),
                requires_grad=False,
            )
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty((self.lora_r, in_features), dtype=target_dtype, device=target_device)
            )
            self.lora_B = nn.Parameter(
                torch.zeros(
                    (out_features, self.lora_r),
                    dtype=target_dtype,
                    device=target_device,
                )
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        weight: torch.Tensor,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> CacheNotesLoRALinear:
        cls_name = linear.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(linear, "nf")
        raw_w = linear.weight.data

        if is_conv1d:
            in_f = raw_w.shape[0]
            out_f = raw_w.shape[1]
        else:
            out_f = raw_w.shape[0]
            in_f = raw_w.shape[1]

        bias = (
            linear.bias.detach().clone()
            if getattr(linear, "bias", None) is not None and linear.bias is not None
            else None
        )
        return cls(
            in_features=in_f,
            out_features=out_f,
            weight=weight,
            bias=bias,
            lora_r=lora_r,
            lora_alpha=lora_alpha,
            device=weight.device,
            dtype=weight.dtype,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)
        out = F.linear(
            x_flat,
            self.weight.to(x_flat.dtype),
            self.bias.to(x_flat.dtype) if self.bias is not None else None,
        )

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x_flat.to(self.lora_A.dtype), self.lora_A), self.lora_B)
            out = out + (self.scaling * lora_out).to(out.dtype)

        return out.reshape(*orig_shape[:-1], self.out_features)

Linear layer containing CacheNotes compressed base weights (frozen) with optional trainable LoRA adapters (lora_A, lora_B).

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Module, weight: torch.Tensor, lora_r: int = 0, lora_alpha: float = 1.0) ‑> CacheNotesLoRALinear

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)
    out = F.linear(
        x_flat,
        self.weight.to(x_flat.dtype),
        self.bias.to(x_flat.dtype) if self.bias is not None else None,
    )

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x_flat.to(self.lora_A.dtype), self.lora_A), self.lora_B)
        out = out + (self.scaling * lora_out).to(out.dtype)

    return out.reshape(*orig_shape[:-1], self.out_features)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class CoMeBridge (hidden_size: int, bridge_rank: int = 64, alpha_init: float = 0.1)

Expand source code

class CoMeBridge(nn.Module):
    """
    Learnable low-rank interpolation adapter bridging representation gaps
    between non-consecutive layers.

        W_bridge = I + B @ A     (B: hidden x r, A: r x hidden)
        y        = x + alpha * (W_bridge x - x)
                 = x + alpha * (B @ A) x

    At initialization B = 0, so the bridge is exactly the identity and only
    alpha influences the output; gradient flows to A and alpha from step 1.
    Robustly handles Tensor, tuple, list, dict, and BaseModelOutput returns.
    """

    def __init__(
        self,
        hidden_size: int,
        bridge_rank: int = 64,
        alpha_init: float = 0.1,
    ) -> None:
        super().__init__()
        self.hidden_size = int(hidden_size)
        self.bridge_rank = int(max(1, min(bridge_rank, hidden_size)))
        # Low-rank factors (actual residual low-rank approximation)
        self.A = nn.Parameter(torch.empty(self.bridge_rank, self.hidden_size))
        self.B = nn.Parameter(torch.zeros(self.hidden_size, self.bridge_rank))
        # Optional bias on the low-rank correction (kept zero by default)
        self.correction_bias = nn.Parameter(torch.zeros(self.hidden_size))
        # Learnable mixing coefficient (scalar per bridge)
        self.alpha = nn.Parameter(torch.tensor(float(alpha_init)))
        # Initialize A with kaiming to get good gradient flow
        nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))

    def _bridge_tensor(self, t: torch.Tensor) -> torch.Tensor:
        if t.ndim >= 2 and t.shape[-1] == self.hidden_size:
            A = self.A.to(device=t.device, dtype=t.dtype)
            B = self.B.to(device=t.device, dtype=t.dtype)
            cb = self.correction_bias.to(device=t.device, dtype=t.dtype)
            # residual = B @ (A @ t^T)^T + bias
            low = F.linear(F.linear(t, A), B) + cb
            return t + self.alpha.to(device=t.device, dtype=t.dtype) * low
        return t

    def forward(self, hidden_states: Any, *args: Any, **kwargs: Any) -> Any:
        # Tuple return
        if isinstance(hidden_states, tuple):
            if len(hidden_states) > 0 and isinstance(hidden_states[0], torch.Tensor):
                return (self._bridge_tensor(hidden_states[0]), *hidden_states[1:])
            return hidden_states

        # List return
        if isinstance(hidden_states, list):
            if len(hidden_states) > 0 and isinstance(hidden_states[0], torch.Tensor):
                return [self._bridge_tensor(hidden_states[0])] + hidden_states[1:]
            return hidden_states

        # Dict return
        if isinstance(hidden_states, dict):
            out_dict = dict(hidden_states)
            for k in ("last_hidden_state", "hidden_states", "x"):
                if k in out_dict and isinstance(out_dict[k], torch.Tensor):
                    out_dict[k] = self._bridge_tensor(out_dict[k])
                    return out_dict
            return hidden_states

        # BaseModelOutput or objects with .last_hidden_state / .hidden_states
        modified = False
        if hasattr(hidden_states, "last_hidden_state"):
            t = hidden_states.last_hidden_state
            if isinstance(t, torch.Tensor):
                hidden_states.last_hidden_state = self._bridge_tensor(t)
                modified = True

        if hasattr(hidden_states, "hidden_states"):
            hs = hidden_states.hidden_states
            if isinstance(hs, tuple) and len(hs) > 0 and isinstance(hs[-1], torch.Tensor):
                new_hs = list(hs)
                new_hs[-1] = self._bridge_tensor(new_hs[-1])
                hidden_states.hidden_states = tuple(new_hs)
                modified = True

        if modified:
            return hidden_states

        # Raw tensor
        if isinstance(hidden_states, torch.Tensor):
            return self._bridge_tensor(hidden_states)

        return hidden_states

Learnable low-rank interpolation adapter bridging representation gaps between non-consecutive layers.

W_bridge = I + B @ A     (B: hidden x r, A: r x hidden)
y        = x + alpha * (W_bridge x - x)
         = x + alpha * (B @ A) x

At initialization B = 0, so the bridge is exactly the identity and only alpha influences the output; gradient flows to A and alpha from step 1. Robustly handles Tensor, tuple, list, dict, and BaseModelOutput returns.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Methods

def forward(self, hidden_states: Any, *args: Any, **kwargs: Any) ‑> Any

Expand source code

def forward(self, hidden_states: Any, *args: Any, **kwargs: Any) -> Any:
    # Tuple return
    if isinstance(hidden_states, tuple):
        if len(hidden_states) > 0 and isinstance(hidden_states[0], torch.Tensor):
            return (self._bridge_tensor(hidden_states[0]), *hidden_states[1:])
        return hidden_states

    # List return
    if isinstance(hidden_states, list):
        if len(hidden_states) > 0 and isinstance(hidden_states[0], torch.Tensor):
            return [self._bridge_tensor(hidden_states[0])] + hidden_states[1:]
        return hidden_states

    # Dict return
    if isinstance(hidden_states, dict):
        out_dict = dict(hidden_states)
        for k in ("last_hidden_state", "hidden_states", "x"):
            if k in out_dict and isinstance(out_dict[k], torch.Tensor):
                out_dict[k] = self._bridge_tensor(out_dict[k])
                return out_dict
        return hidden_states

    # BaseModelOutput or objects with .last_hidden_state / .hidden_states
    modified = False
    if hasattr(hidden_states, "last_hidden_state"):
        t = hidden_states.last_hidden_state
        if isinstance(t, torch.Tensor):
            hidden_states.last_hidden_state = self._bridge_tensor(t)
            modified = True

    if hasattr(hidden_states, "hidden_states"):
        hs = hidden_states.hidden_states
        if isinstance(hs, tuple) and len(hs) > 0 and isinstance(hs[-1], torch.Tensor):
            new_hs = list(hs)
            new_hs[-1] = self._bridge_tensor(new_hs[-1])
            hidden_states.hidden_states = tuple(new_hs)
            modified = True

    if modified:
        return hidden_states

    # Raw tensor
    if isinstance(hidden_states, torch.Tensor):
        return self._bridge_tensor(hidden_states)

    return hidden_states

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class CoMeBridgedLayer (layer: nn.Module, bridge: CoMeBridge | None = None)

Expand source code

class CoMeBridgedLayer(nn.Module):
    """Wraps a retained transformer layer with a learnable CoMeBridge."""

    def __init__(self, layer: nn.Module, bridge: CoMeBridge | None = None) -> None:
        super().__init__()
        self.layer = layer
        self.bridge = bridge

    def __getattr__(self, name: str) -> Any:
        try:
            return super().__getattr__(name)
        except AttributeError:
            if "layer" in self.__dict__ or (
                "_modules" in self.__dict__ and "layer" in self._modules
            ):
                return getattr(self.layer, name)
            raise

    def __getitem__(self, key: Any) -> Any:
        return self.layer[key]

    def __contains__(self, key: Any) -> bool:
        return key in self.layer

    def forward(self, hidden_states: Any = None, *args: Any, **kwargs: Any) -> Any:
        if hidden_states is None and "hidden_states" in kwargs:
            hidden_states = kwargs.pop("hidden_states")
        if hidden_states is not None:
            args = (hidden_states,) + args

        attn_mask = kwargs.pop("attention_mask", None)
        if attn_mask is not None:
            if getattr(attn_mask, "dtype", None) == torch.int64:
                attn_mask = attn_mask.to(torch.bool)
            kwargs["attention_mask"] = attn_mask

        out = self.layer(*args, **kwargs)
        if self.bridge is not None:
            out = self.bridge(out)
        return out

Wraps a retained transformer layer with a learnable CoMeBridge.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Methods

def forward(self, hidden_states: Any = None, *args: Any, **kwargs: Any) ‑> Any

Expand source code

def forward(self, hidden_states: Any = None, *args: Any, **kwargs: Any) -> Any:
    if hidden_states is None and "hidden_states" in kwargs:
        hidden_states = kwargs.pop("hidden_states")
    if hidden_states is not None:
        args = (hidden_states,) + args

    attn_mask = kwargs.pop("attention_mask", None)
    if attn_mask is not None:
        if getattr(attn_mask, "dtype", None) == torch.int64:
            attn_mask = attn_mask.to(torch.bool)
        kwargs["attention_mask"] = attn_mask

    out = self.layer(*args, **kwargs)
    if self.bridge is not None:
        out = self.bridge(out)
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class CoMeCompressor (model: nn.Module, tokenizer: Any | None = None, target_layers: int | None = None, sparsity: float = 0.1, bridge_rank: int | None = None, alpha_init: float = 0.1, lora_r: int = 16, lora_alpha: float = 1.0, skip_attention_qkv: bool = True, protect_critical: bool = True, target_attention_only: bool = False, recalibrate_ln: bool = True, allow_forward_failure: bool = False, device: str | torch.device | None = None, target_modules: list[str] | None = None, rank: int | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("come")
@COMPRESSOR_REGISTRY.register("come_compressor")
class CoMeCompressor(BaseCompressor):
    """
    CoMe Compressor: Layer shedding with learnable low-rank interpolation
    adapters (bridges) between non-consecutive retained layers.

    Config parameters (via kwargs or comp_config):
        rank                  : int, optional — SVD-style rank. If provided and
                                bridge_rank/sparsity are not explicitly set,
                                they are derived from rank via safe mappings.
        target_layers         : int, optional — target layer count after shedding.
        sparsity              : float, default 0.10 — fraction of layers to remove.
        bridge_rank           : int, optional — bottleneck rank for the
                                low-rank interpolation bridge (B @ A). If not
                                set, derived from rank (rank // 4) or defaults
                                to 64 for the full hidden size.
        alpha_init            : float, default 0.1 — initial bridge mixing coeff.
        lora_r                : int, default 16 — post-shedding LoRA rank.
        lora_alpha            : float, default 1.0 — LoRA scaling.
        target_attention_only : bool, default False.
        skip_attention_qkv    : bool, default True — protect asymmetric GQA.
        recalibrate_ln        : bool, default True.
        allow_forward_failure : bool, default False.
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any | None = None,
        target_layers: int | None = None,
        sparsity: float = 0.10,
        bridge_rank: int | None = None,
        alpha_init: float = 0.1,
        lora_r: int = 16,
        lora_alpha: float = 1.0,
        skip_attention_qkv: bool = True,
        protect_critical: bool = True,
        target_attention_only: bool = False,
        recalibrate_ln: bool = True,
        allow_forward_failure: bool = False,
        device: str | torch.device | None = None,
        target_modules: list[str] | None = None,
        rank: int | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="come", **kwargs)
        conf = kwargs.get("comp_config", kwargs.get("cfg"))

        # -------- Read config with precedence: explicit kwarg > comp_config > default --------
        rank = _cfg_get(conf, "rank", _cfg_get(conf, "come_rank", rank))
        if rank is not None:
            rank = int(rank)

        target_layers = _cfg_get(
            conf, "come_target_layers", _cfg_get(conf, "target_layers", target_layers)
        )
        if target_layers is not None:
            target_layers = int(target_layers)

        # sparsity: if rank provided and sparsity not explicitly overridden, derive
        sparsity_explicit = _cfg_get(
            conf,
            "come_sparsity",
            _cfg_get(conf, "sparsity", kwargs.get("sparsity")),
        )
        if sparsity_explicit is None and rank is not None:
            sparsity = _rank_to_sparsity(rank)
        elif sparsity_explicit is not None:
            sparsity = float(sparsity_explicit)
        # else: keep the constructor default (0.10)

        # bridge_rank: if rank provided and bridge_rank not explicitly overridden, derive
        bridge_rank_explicit = _cfg_get(
            conf,
            "come_bridge_rank",
            _cfg_get(conf, "bridge_rank", bridge_rank),
        )
        hidden_size_probe = None
        for attr in ("hidden_size", "d_model", "dim"):
            if hasattr(model, "config") and hasattr(model.config, attr):
                hidden_size_probe = int(getattr(model.config, attr))
                break
        if bridge_rank_explicit is None and rank is not None and hidden_size_probe:
            bridge_rank = _rank_to_bridge_rank(rank, hidden_size_probe)
        elif bridge_rank_explicit is not None:
            bridge_rank = int(bridge_rank_explicit)
        else:
            # default: full hidden size (no bottleneck) — safest default
            bridge_rank = hidden_size_probe or 64

        alpha_init = float(
            _cfg_get(conf, "come_alpha_init", _cfg_get(conf, "alpha_init", alpha_init))
        )
        lora_r = int(
            _cfg_get(
                conf,
                "come_lora_r",
                _cfg_get(conf, "lora_r", _cfg_get(conf, "lora_rank", lora_r)),
            )
        )
        lora_alpha = float(
            _cfg_get(conf, "come_lora_alpha", _cfg_get(conf, "lora_alpha", lora_alpha))
        )
        skip_attention_qkv = bool(
            _cfg_get(
                conf,
                "come_skip_attention_qkv",
                _cfg_get(conf, "skip_attention_qkv", skip_attention_qkv),
            )
        )
        protect_critical = bool(_cfg_get(conf, "protect_critical", protect_critical))
        target_attention_only = bool(
            _cfg_get(
                conf,
                "come_target_attention_only",
                _cfg_get(conf, "target_attention_only", target_attention_only),
            )
        )
        recalibrate_ln = bool(
            _cfg_get(
                conf,
                "come_recalibrate_ln",
                _cfg_get(conf, "recalibrate_ln", recalibrate_ln),
            )
        )
        allow_forward_failure = bool(
            _cfg_get(
                conf,
                "allow_forward_failure",
                _cfg_get(kwargs, "allow_forward_failure", allow_forward_failure),
            )
        )

        # -------- Store --------
        self.tokenizer = tokenizer or kwargs.get("tokenizer", getattr(model, "tokenizer", None))
        self.target_layers = target_layers
        self.sparsity = max(0.0, min(sparsity, 0.95))
        self.bridge_rank = int(bridge_rank)
        self.alpha_init = alpha_init
        self.lora_r = lora_r
        self.lora_alpha = lora_alpha
        self.skip_attention_qkv = skip_attention_qkv
        self.protect_critical = protect_critical
        self.target_attention_only = target_attention_only
        self.recalibrate_ln = recalibrate_ln
        self.allow_forward_failure = allow_forward_failure
        self.target_modules = target_modules
        self.rank = rank
        self.target_device = (
            torch.device(device)
            if device is not None
            else (next(model.parameters(), torch.tensor(0.0)).device)
        )
        self.comp_config = conf

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.stats: dict[str, Any] = {}

    # ------------------------------------------------------------------
    # helpers
    # ------------------------------------------------------------------
    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv or self.target_attention_only:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            name_lower = name.lower()
            if any(
                p in name_lower
                for p in [
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                ]
            ):
                return True
        return False

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        if isinstance(module, (CoMeLoRALinear, CoMeBridge, CoMeBridgedLayer)):
            return False
        if ".bridge." in name or name.endswith(".bridge"):
            return False
        if any(
            name.endswith(sfx)
            for sfx in [".base_layer", ".lora_A", ".lora_B", ".lora_dropout", ".linear"]
        ):
            return False

        target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
        cls_name = module.__class__.__name__.lower()
        target_cls = target_mod.__class__.__name__.lower()

        is_linear_like = (
            isinstance(module, nn.Linear)
            or isinstance(target_mod, nn.Linear)
            or (
                Conv1D is not None
                and (isinstance(module, Conv1D) or isinstance(target_mod, Conv1D))
            )
            or "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(module, "nf")
            or hasattr(target_mod, "nf")
        )
        if not is_linear_like:
            return False

        name_lower = name.lower()
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embed",
                "logits_processor",
                "head",
                "pre_classifier",
            ]
        ):
            return False

        if self._is_gqa_asymmetric(name, module):
            return False

        if self.target_attention_only:
            ffn_patterns = [
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            ]
            if any(p in name_lower for p in ffn_patterns):
                return False
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            attn_patterns = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
                "in_proj_qkv",
                "out_proj",
            ]
            if not any(p in name_lower for p in attn_patterns):
                return False
        else:
            if self.protect_critical and not should_compress_linear(self.model, name):
                return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def _get_hidden_size(self, layer: nn.Module) -> int:
        if hasattr(self, "model") and hasattr(self.model, "config"):
            for attr in ("hidden_size", "d_model", "dim"):
                if hasattr(self.model.config, attr):
                    return int(getattr(self.model.config, attr))
        for attr in ("hidden_size", "d_model", "dim", "embed_dim"):
            if hasattr(layer, attr):
                return int(getattr(layer, attr))
        if hasattr(layer, "config"):
            for attr in ("hidden_size", "d_model", "dim"):
                if hasattr(layer.config, attr):
                    return int(getattr(layer.config, attr))
        _OUTPUT_PROJ_KEYS = ("o_proj", "out_proj", "output", "dense")
        for name, m in layer.named_modules():
            if isinstance(m, nn.Linear) and any(k in name for k in _OUTPUT_PROJ_KEYS):
                return m.out_features
        sizes = [m.in_features for _, m in layer.named_modules() if isinstance(m, nn.Linear)]
        if sizes:
            from collections import Counter

            return Counter(sizes).most_common(1)[0][0]
        for m in layer.modules():
            if isinstance(m, nn.LayerNorm):
                ns = m.normalized_shape
                return ns[0] if isinstance(ns, (tuple, list)) else int(ns)
        raise ValueError(f"Cannot determine hidden_size for {type(layer).__name__}")

    def _find_layer_container(
        self,
    ) -> tuple[nn.Module | None, str | None, nn.ModuleList | None]:
        candidates = [
            (
                "roberta.encoder.layer",
                lambda m: getattr(
                    getattr(getattr(m, "roberta", None), "encoder", None), "layer", None
                ),
            ),
            (
                "bert.encoder.layer",
                lambda m: getattr(
                    getattr(getattr(m, "bert", None), "encoder", None), "layer", None
                ),
            ),
            (
                "distilbert.transformer.layer",
                lambda m: getattr(
                    getattr(getattr(m, "distilbert", None), "transformer", None),
                    "layer",
                    None,
                ),
            ),
            (
                "transformer.layer",
                lambda m: getattr(getattr(m, "transformer", None), "layer", None),
            ),
            (
                "transformer.h",
                lambda m: getattr(getattr(m, "transformer", None), "h", None),
            ),
            (
                "transformer.blocks",
                lambda m: getattr(getattr(m, "transformer", None), "blocks", None),
            ),
            (
                "model.layers",
                lambda m: getattr(getattr(m, "model", None), "layers", None),
            ),
            (
                "model.decoder.layers",
                lambda m: getattr(
                    getattr(getattr(m, "model", None), "decoder", None), "layers", None
                ),
            ),
            (
                "model.encoder.layers",
                lambda m: getattr(
                    getattr(getattr(m, "model", None), "encoder", None), "layers", None
                ),
            ),
            (
                "encoder.layer",
                lambda m: getattr(getattr(m, "encoder", None), "layer", None),
            ),
            (
                "decoder.layer",
                lambda m: getattr(getattr(m, "decoder", None), "layer", None),
            ),
            ("layers", lambda m: getattr(m, "layers", None)),
            ("h", lambda m: getattr(m, "h", None)),
            ("blocks", lambda m: getattr(m, "blocks", None)),
            (
                "model.encoder.block",
                lambda m: getattr(
                    getattr(getattr(m, "model", None), "encoder", None), "block", None
                ),
            ),
            (
                "model.decoder.block",
                lambda m: getattr(
                    getattr(getattr(m, "model", None), "decoder", None), "block", None
                ),
            ),
            (
                "encoder.block",
                lambda m: getattr(getattr(m, "encoder", None), "block", None),
            ),
            (
                "decoder.block",
                lambda m: getattr(getattr(m, "decoder", None), "block", None),
            ),
        ]
        for name, getter in candidates:
            res = getter(self.model)
            if isinstance(res, (nn.ModuleList, nn.Sequential, list)) and len(res) > 0:
                parts = name.split(".")
                parent = self.model
                for p in parts[:-1]:
                    parent = getattr(parent, p)
                return parent, parts[-1], res

        best_match = None
        max_score = -1
        for name, module in self.model.named_modules():
            if isinstance(module, (nn.ModuleList, nn.Sequential)) and len(module) >= 2:
                parts = name.split(".")
                last = parts[-1].lower()
                is_layer_name = any(k in last for k in ["layer", "h", "block"])
                score = len(module) + (1000 if is_layer_name else 0)
                if score > max_score:
                    max_score = score
                    parent = self.model
                    for p in parts[:-1]:
                        parent = getattr(parent, p)
                    best_match = (parent, parts[-1], module)

        if best_match is not None:
            return best_match
        return None, None, None

    # ------------------------------------------------------------------
    # compress
    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        parent, attr_name, layers = self._find_layer_container()
        bridges_inserted = 0
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []

        orig_len = 0
        keep_len = 0

        if layers is not None and parent is not None and attr_name is not None:
            orig_len = len(layers)
            keep_len = max(1, int(round(orig_len * (1.0 - self.sparsity))))
            if self.target_layers is not None:
                keep_len = min(orig_len, max(1, self.target_layers))

            if keep_len < orig_len:
                kept_indices = np.round(np.linspace(0, orig_len - 1, keep_len)).astype(int).tolist()
                seen = set()
                kept_indices = [x for x in kept_indices if not (x in seen or seen.add(x))]

                bridged_module_list = nn.ModuleList()
                for i, layer_idx in enumerate(kept_indices):
                    layer = layers[layer_idx]
                    is_gap = i < len(kept_indices) - 1 and kept_indices[i + 1] > layer_idx + 1
                    if is_gap:
                        hidden_size = self._get_hidden_size(layer)
                        # Use hidden_size as cap on bridge_rank
                        effective_bridge_rank = max(1, min(self.bridge_rank, hidden_size))
                        bridge = CoMeBridge(
                            hidden_size=hidden_size,
                            bridge_rank=effective_bridge_rank,
                            alpha_init=self.alpha_init,
                        )
                        p = next(layer.parameters(), None)
                        if p is not None:
                            bridge = bridge.to(device=p.device, dtype=p.dtype)
                        bridged_module_list.append(CoMeBridgedLayer(layer, bridge=bridge))
                        bridges_inserted += 1
                    else:
                        bridged_module_list.append(layer)

                setattr(parent, attr_name, bridged_module_list)
                self.compressed_layers.append(f"{attr_name} (shed {orig_len - keep_len} layers)")
        else:
            # Fallback SVD compression (still supported)
            seen_weights: set[int] = set()
            for name, module in list(self.model.named_modules()):
                if self._is_target_module(name, module):
                    target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
                    w_id = id(target_mod.weight)
                    if w_id in seen_weights:
                        self.skipped_layers.append(name)
                        continue
                    seen_weights.add(w_id)
                    try:
                        raw_w = target_mod.weight.data
                        cls_name = module.__class__.__name__.lower()
                        is_conv1d = (
                            "conv1d" in cls_name
                            or hasattr(module, "nf")
                            or hasattr(target_mod, "nf")
                        )
                        W = (
                            raw_w.t().contiguous().clone().float()
                            if is_conv1d
                            else raw_w.clone().float()
                        )
                        U, S, Vt = torch.linalg.svd(W, full_matrices=False)
                        # Prefer explicit rank, otherwise derive from sparsity
                        if self.rank is not None:
                            r = max(1, min(int(self.rank), len(S)))
                        else:
                            r = max(1, int(len(S) * (1.0 - self.sparsity)))
                        from neuralops.compression_engine.aa_svd_compressor import (
                            SVDLinear,
                        )

                        in_features = W.shape[1]
                        out_features = W.shape[0]
                        svd_layer = SVDLinear(
                            in_features=in_features,
                            out_features=out_features,
                            rank=r,
                            bias=getattr(target_mod, "bias", None) is not None,
                        )
                        svd_layer.set_low_rank_factors(U[:, :r] * S[:r], Vt[:r, :])
                        if getattr(target_mod, "bias", None) is not None:
                            svd_layer.bias.data.copy_(target_mod.bias.data)
                        replace_module(self.model, name, svd_layer)
                        self.compressed_layers.append(name)
                    except Exception as e:
                        logger.warning("CoMe fallback SVD failed on %s: %s", name, e)
                        self.fallback_layers.append(name)
                        for p in module.parameters():
                            p.requires_grad = False

        # Optional post-shedding LoRA
        if self.lora_r > 0:
            seen_weights: set[int] = set()
            for name, module in list(self.model.named_modules()):
                if self._is_target_module(name, module):
                    target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
                    has_lora = (hasattr(module, "lora_A") and module.lora_A is not None) or hasattr(
                        module, "base_layer"
                    )
                    w_id = id(target_mod.weight)
                    if w_id in seen_weights:
                        self.skipped_layers.append(name)
                        continue
                    seen_weights.add(w_id)
                    if has_lora:
                        target_mod.weight.requires_grad = False
                        if (
                            getattr(target_mod, "bias", None) is not None
                            and target_mod.bias is not None
                        ):
                            target_mod.bias.requires_grad = False
                        for p_name, p in module.named_parameters():
                            if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                                p.requires_grad = True
                    else:
                        lora_mod = CoMeLoRALinear.from_linear(
                            target_mod,
                            lora_r=self.lora_r,
                            lora_alpha=self.lora_alpha,
                        )
                        replace_module(self.model, name, lora_mod)

        # LayerNorm recalibration
        if self.recalibrate_ln:
            try:
                recalibrate_layernorms(
                    self.model,
                    tokenizer=self.tokenizer,
                    device=self.target_device,
                )
            except Exception as e:
                logger.warning("CoMe LayerNorm recalibration skipped: %s", e)

        # Forward sanity check
        forward_check_success = False
        try:
            self.model.eval()
            with torch.no_grad():
                d = next(self.model.parameters()).device
                dummy = torch.randint(0, 100, (1, 16), device=d)
                self.model(dummy)
            forward_check_success = True
        except Exception as e:
            logger.error("CoMe post-compression forward sanity check failed: %s", e)
            if not self.allow_forward_failure:
                raise RuntimeError(f"CoMe post-compression forward check failed: {e}") from e

        self._compressed_params = self._count_params(self.model)
        shed_count = orig_len - keep_len if orig_len > 0 else 0
        reduc_ratio = float(shed_count) / orig_len if orig_len > 0 else self.sparsity
        flops_reduction = f"{reduc_ratio * 100:.1f}%"

        bridge_params_count = sum(
            p.numel()
            for m in self.model.modules()
            if isinstance(m, CoMeBridge)
            for p in m.parameters()
        )
        lora_params_count = sum(
            p.numel()
            for m in self.model.modules()
            if isinstance(m, CoMeLoRALinear)
            for p in [m.lora_A, m.lora_B]
            if p is not None
        )

        stats = self.get_compression_stats()
        num_compressed = len(self.compressed_layers)
        num_fallback = len(self.fallback_layers)

        # Write BOTH naming conventions so downstream readers find them
        stats.update(
            {
                "original_layer_count": orig_len,
                "kept_layer_count": keep_len,
                "shed_layer_count": shed_count,
                "layer_reduction_ratio": reduc_ratio,
                "estimated_flops_reduction": flops_reduction,
                "bridges_inserted": bridges_inserted,
                "bridge_rank": self.bridge_rank,
                "rank": self.rank,
                "bridge_parameters": bridge_params_count,
                "lora_parameters": lora_params_count,
                "trainable_parameters": sum(p.numel() for p in self.get_trainable_params()),
                "target_layers": self.target_layers,
                "sparsity": self.sparsity,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "target_attention_only": self.target_attention_only,
                "recalibrate_ln": self.recalibrate_ln,
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                # Primary keys expected by run_quality_sweep.py:
                "num_compressed": num_compressed,
                "num_fallback": num_fallback,
                # Alternate key names used by some readers:
                "num_compressed_layers": num_compressed,
                "num_fallback_layers": num_fallback,
                "compressed_layers_count": num_compressed,
                "fallback_layer_count": num_fallback,
                "num_skipped": len(self.skipped_layers),
                "forward_check_success": forward_check_success,
                "method": "come_layer_shedding_with_bridges",
            }
        )
        self.stats = stats
        self.model.compression_stats = stats
        # Also attach to the compressor instance itself so nested pipelines can recover it
        try:
            self.model._come_stats = stats
        except Exception:
            pass
        return self.model, stats

    def get_trainable_params(self) -> list[nn.Parameter]:
        trainable = []
        for _, param in self.model.named_parameters():
            if param.requires_grad:
                trainable.append(param)
        return trainable

CoMe Compressor: Layer shedding with learnable low-rank interpolation adapters (bridges) between non-consecutive retained layers.

Config parameters (via kwargs or comp_config): rank : int, optional — SVD-style rank. If provided and bridge_rank/sparsity are not explicitly set, they are derived from rank via safe mappings. target_layers : int, optional — target layer count after shedding. sparsity : float, default 0.10 — fraction of layers to remove. bridge_rank : int, optional — bottleneck rank for the low-rank interpolation bridge (B @ A). If not set, derived from rank (rank // 4) or defaults to 64 for the full hidden size. alpha_init : float, default 0.1 — initial bridge mixing coeff. lora_r : int, default 16 — post-shedding LoRA rank. lora_alpha : float, default 1.0 — LoRA scaling. target_attention_only : bool, default False. skip_attention_qkv : bool, default True — protect asymmetric GQA. recalibrate_ln : bool, default True. allow_forward_failure : bool, default False.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def get_trainable_params(self) ‑> list[torch.nn.parameter.Parameter]

Expand source code

def get_trainable_params(self) -> list[nn.Parameter]:
    trainable = []
    for _, param in self.model.named_parameters():
        if param.requires_grad:
            trainable.append(param)
    return trainable

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class CoMeLoRALinear (in_features: int, out_features: int, weight: torch.Tensor, bias: torch.Tensor | None = None, lora_r: int = 0, lora_alpha: float = 1.0, device: torch.device | None = None, dtype: torch.dtype | None = None)

Expand source code

class CoMeLoRALinear(nn.Module):
    """Linear layer: frozen base weights + optional trainable LoRA adapters."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        weight: torch.Tensor,
        bias: torch.Tensor | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features  # Conv1D compatibility
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        target_device = device if device is not None else weight.device
        target_dtype = dtype if dtype is not None else weight.dtype

        self.weight = nn.Parameter(
            weight.to(device=target_device, dtype=target_dtype),
            requires_grad=False,
        )
        if bias is not None:
            self.bias = nn.Parameter(
                bias.to(device=target_device, dtype=target_dtype),
                requires_grad=False,
            )
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty((self.lora_r, in_features), dtype=target_dtype, device=target_device)
            )
            self.lora_B = nn.Parameter(
                torch.zeros(
                    (out_features, self.lora_r),
                    dtype=target_dtype,
                    device=target_device,
                )
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> CoMeLoRALinear:
        target_mod = getattr(linear, "base_layer", getattr(linear, "linear", linear))
        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.detach()

        if is_conv1d:
            in_f = raw_w.shape[0]
            out_f = raw_w.shape[1]
            w = raw_w.t().contiguous().clone()
        else:
            out_f = raw_w.shape[0]
            in_f = raw_w.shape[1]
            w = raw_w.clone()

        bias = (
            target_mod.bias.detach().clone()
            if getattr(target_mod, "bias", None) is not None and target_mod.bias is not None
            else None
        )
        return cls(
            in_features=in_f,
            out_features=out_f,
            weight=w,
            bias=bias,
            lora_r=lora_r,
            lora_alpha=lora_alpha,
            device=w.device,
            dtype=w.dtype,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)
        out = F.linear(
            x_flat,
            self.weight.to(x_flat.dtype),
            self.bias.to(x_flat.dtype) if self.bias is not None else None,
        )
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x_flat.to(self.lora_A.dtype), self.lora_A), self.lora_B)
            out = out + (self.scaling * lora_out).to(out.dtype)
        return out.reshape(*orig_shape[:-1], self.out_features)

Linear layer: frozen base weights + optional trainable LoRA adapters.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Module, lora_r: int = 0, lora_alpha: float = 1.0) ‑> CoMeLoRALinear

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)
    out = F.linear(
        x_flat,
        self.weight.to(x_flat.dtype),
        self.bias.to(x_flat.dtype) if self.bias is not None else None,
    )
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x_flat.to(self.lora_A.dtype), self.lora_A), self.lora_B)
        out = out + (self.scaling * lora_out).to(out.dtype)
    return out.reshape(*orig_shape[:-1], self.out_features)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class CompressedAttentionLayer (config: Any | None = None, rank: int | None = None, rank_qkv: int = 256, rank_out: int = 256, lora_rank: int = 16, original_layer: Any = None, pretrained_attn: nn.Module | None = None, **kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "compressed_attention_layer",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
@ATTENTION_REGISTRY.register(
    "compressed_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class CompressedAttentionLayer(BaseAttention):
    supports_causal: bool = True
    supports_bidirectional: bool = True
    supports_gqa: bool = True
    supports_rope: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        rank_qkv: int = 256,
        rank_out: int = 256,
        lora_rank: int = 16,
        original_layer: Any = None,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = config if config is not None else original_layer
        super().__init__(cfg, rank)
        self.adapter = ModelAdapter(cfg)
        self.hidden_size: int = self.adapter.hidden_size or kwargs.get(
            "hidden_size", kwargs.get("d_model", 768)
        )
        self.num_heads: int = self.adapter.num_heads or kwargs.get(
            "num_heads", kwargs.get("num_attention_heads", 12)
        )
        self.num_kv_heads: int = self.adapter.num_key_value_heads or kwargs.get(
            "num_kv_heads", kwargs.get("num_key_value_heads", self.num_heads)
        )
        self.head_dim: int = self.adapter.head_dim or (self.hidden_size // max(1, self.num_heads))
        self.all_head_size: int = self.num_heads * self.head_dim
        self.kv_head_size: int = self.num_kv_heads * self.head_dim
        self.kv_repeat: int = max(1, self.num_heads // max(1, self.num_kv_heads))
        _mt = getattr(cfg, "model_type", "").lower() if cfg else ""
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        self.is_decoder: bool = bool(
            getattr(cfg, "is_decoder", False) or any(_mt.startswith(p) for p in _decoder_prefixes)
        )
        self.is_nested_encoder: bool = False

        rank_val = rank or kwargs.get("rank_qkv", rank_qkv)
        self.rank_qkv = min(rank_val, self.all_head_size + 2 * self.kv_head_size, self.hidden_size)
        self.rank_out = min(
            rank or kwargs.get("rank_out", rank_out),
            self.all_head_size,
            self.hidden_size,
        )
        self.lora_rank = getattr(cfg, "lora_rank", lora_rank)
        self.lora_scaling = (32.0 / self.lora_rank) if self.lora_rank > 0 else 1.0

        qkv_out_dim = self.all_head_size + 2 * self.kv_head_size

        # Frozen base SVD parameters
        self.Vh_qkv = nn.Parameter(
            torch.randn(self.rank_qkv, self.hidden_size) * 0.02, requires_grad=False
        )
        self.S_qkv_sqrt = nn.Parameter(torch.ones(self.rank_qkv), requires_grad=False)
        self.U_qkv = nn.Parameter(
            torch.randn(qkv_out_dim, self.rank_qkv) * 0.02, requires_grad=False
        )
        self.bias_qkv = nn.Parameter(torch.zeros(qkv_out_dim), requires_grad=False)

        self.Vh_out = nn.Parameter(
            torch.randn(self.rank_out, self.all_head_size) * 0.02, requires_grad=False
        )
        self.S_out_sqrt = nn.Parameter(torch.ones(self.rank_out), requires_grad=False)
        self.U_out = nn.Parameter(
            torch.randn(self.hidden_size, self.rank_out) * 0.02, requires_grad=False
        )
        self.bias_out = nn.Parameter(torch.zeros(self.hidden_size), requires_grad=False)

        # Trainable LoRA residual adapters
        if self.lora_rank > 0:
            self.lora_A_qkv = nn.Parameter(torch.empty(self.lora_rank, self.hidden_size))
            self.lora_B_qkv = nn.Parameter(torch.empty(qkv_out_dim, self.lora_rank))
            self.lora_A_out = nn.Parameter(torch.empty(self.lora_rank, self.all_head_size))
            self.lora_B_out = nn.Parameter(torch.empty(self.hidden_size, self.lora_rank))
            self.reset_lora_parameters()
        else:
            self.register_parameter("lora_A_qkv", None)
            self.register_parameter("lora_B_qkv", None)
            self.register_parameter("lora_A_out", None)
            self.register_parameter("lora_B_out", None)

        dropout_p = (
            getattr(cfg, "dropout", getattr(cfg, "attention_probs_dropout_prob", 0.1))
            if cfg
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)
        self.rotary_emb = getattr(pretrained_attn, "rotary_emb", None)

        src = pretrained_attn or (
            original_layer
            if hasattr(original_layer, "q_lin")
            or hasattr(original_layer, "query")
            or hasattr(original_layer, "c_attn")
            else None
        )
        if src is not None:
            self.copy_weights(src)

    @property
    def qkv_weight(self) -> torch.Tensor:
        """Reconstruct the effective base QKV weight from SVD factors."""
        return (self.U_qkv * self.S_qkv_sqrt.view(1, -1)) @ (
            self.Vh_qkv * self.S_qkv_sqrt.view(-1, 1)
        )

    @property
    def out_weight(self) -> torch.Tensor:
        """Reconstruct the effective base Out weight from SVD factors."""
        return (self.U_out * self.S_out_sqrt.view(1, -1)) @ (
            self.Vh_out * self.S_out_sqrt.view(-1, 1)
        )

    def reset_lora_parameters(self) -> None:
        if self.lora_rank > 0:
            nn.init.kaiming_uniform_(self.lora_A_qkv, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B_qkv)
            nn.init.kaiming_uniform_(self.lora_A_out, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B_out)
            self.lora_A_qkv.requires_grad = True
            self.lora_B_qkv.requires_grad = True
            self.lora_A_out.requires_grad = True
            self.lora_B_out.requires_grad = True

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 1. HuggingFace / standard UniversalAdapter extraction
        for proj in ["q", "k", "v", "o"]:
            try:
                w, b = UniversalAdapter.get_weight_bias(src, proj)
                if w is not None:
                    extracted[proj] = (w, b)
            except Exception as e:
                logger.debug(
                    f"UniversalAdapter failed for {proj}: {e}"
                )  # continue to fallback paths

        # 2. GPT-2 Fused c_attn Conv1D handling
        if hasattr(src, "c_attn") and (
            extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
        ):
            weight = src.c_attn.weight.data
            bias = getattr(src.c_attn, "bias", None)
            bias_data = bias.data if bias is not None else None
            if weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size:
                weight = weight.t().contiguous()
            split_size = weight.shape[0] // 3
            q_w = weight[:split_size, :]
            k_w = weight[split_size : 2 * split_size, :]
            v_w = weight[2 * split_size :, :]
            q_b = bias_data[:split_size] if bias_data is not None else None
            k_b = bias_data[split_size : 2 * split_size] if bias_data is not None else None
            v_b = bias_data[2 * split_size :] if bias_data is not None else None
            extracted["q"] = (q_w, q_b)
            extracted["k"] = (k_w, k_b)
            extracted["v"] = (v_w, v_b)

        if hasattr(src, "c_proj") and extracted["o"][0] is None:
            w_o = src.c_proj.weight.data
            if (
                w_o.dim() == 2
                and w_o.shape[0] != self.hidden_size
                and w_o.shape[1] == self.hidden_size
            ):
                w_o = w_o.t().contiguous()
            b_o = getattr(src.c_proj, "bias", None)
            extracted["o"] = (w_o, b_o.data if b_o is not None else None)

        # 3. RoBERTa / BERT Nested Encoder handling
        if hasattr(src, "self") and (
            extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
        ):
            s = src.self
            if hasattr(s, "query"):
                extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
            if hasattr(s, "key"):
                extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
            if hasattr(s, "value"):
                extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

        if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None:
            dense = src.output.dense
            extracted["o"] = (dense.weight.data, getattr(dense, "bias", None))

        # 4. Modern Decoders (Qwen3.5, LLaMA, Mistral) & DistilBERT
        named_map = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for name in names:
                    if hasattr(src, name):
                        layer = getattr(src, name)
                        if hasattr(layer, "weight"):
                            extracted[proj] = (
                                layer.weight.data,
                                getattr(layer, "bias", None),
                            )
                            break

        # Detect nested encoder architecture
        cls_name = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        # Check output projection attrs FIRST — DistilBERT has out_lin and is NOT nested
        if hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj"):
            self.is_nested_encoder = False
        elif (
            (
                hasattr(src, "self")
                and hasattr(src.self, "query")
                and hasattr(src, "output")
                and hasattr(src.output, "dense")
            )
            or (
                cls_name in ("BertSelfAttention", "RobertaSelfAttention")
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
            )
            and not has_out_proj
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        # Check required weights
        missing = [k for k in ["q", "k", "v"] if extracted[k][0] is None]
        if not self.is_nested_encoder and extracted["o"][0] is None:
            missing.append("o")
        if missing:
            logger.warning(
                f"CompressedAttentionLayer: missing weights for projections {missing}; proceeding with random init"
            )

        # Frobenius norm validation
        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    logger.warning(
                        f"CompressedAttentionLayer: extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f}); continuing"
                    )

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                logger.warning(
                    f"CompressedAttentionLayer: extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f}); continuing"
                )

        # Joint SVD on Q, K, V
        def _orient(w: torch.Tensor, in_dim: int) -> torch.Tensor:
            if w.dim() == 2 and w.shape[1] != in_dim and w.shape[0] == in_dim:
                return w.t().contiguous()
            return w

        with torch.no_grad():
            w_q = _orient(extracted["q"][0].float(), self.hidden_size)
            w_k = _orient(extracted["k"][0].float(), self.hidden_size)
            w_v = _orient(extracted["v"][0].float(), self.hidden_size)
            w_qkv = torch.cat([w_q, w_k, w_v], dim=0)

            U, S, Vh = torch.linalg.svd(w_qkv, full_matrices=False)
            r_qkv = min(self.rank_qkv, U.size(1), Vh.size(0))
            sqrt_s = torch.sqrt(torch.clamp(S[:r_qkv], min=1e-12))

            qkv_out_dim = self.all_head_size + 2 * self.kv_head_size
            self.U_qkv.data[: min(qkv_out_dim, U.size(0)), :r_qkv].copy_(
                U[: min(qkv_out_dim, U.size(0)), :r_qkv]
            )
            self.S_qkv_sqrt.data[:r_qkv].copy_(sqrt_s)
            self.Vh_qkv.data[:r_qkv, : min(self.hidden_size, Vh.size(1))].copy_(
                Vh[:r_qkv, : min(self.hidden_size, Vh.size(1))]
            )

            b_q = (
                extracted["q"][1]
                if extracted["q"][1] is not None
                else torch.zeros(w_q.size(0), device=w_q.device)
            )
            b_k = (
                extracted["k"][1]
                if extracted["k"][1] is not None
                else torch.zeros(w_k.size(0), device=w_k.device)
            )
            b_v = (
                extracted["v"][1]
                if extracted["v"][1] is not None
                else torch.zeros(w_v.size(0), device=w_v.device)
            )
            b_q_val = b_q.data if isinstance(b_q, nn.Parameter) else b_q
            b_k_val = b_k.data if isinstance(b_k, nn.Parameter) else b_k
            b_v_val = b_v.data if isinstance(b_v, nn.Parameter) else b_v
            b_qkv = torch.cat([b_q_val, b_k_val, b_v_val], dim=0)
            self.bias_qkv.data[: min(qkv_out_dim, b_qkv.size(0))].copy_(
                b_qkv[: min(qkv_out_dim, b_qkv.size(0))]
            )

            # Output projection SVD
            if extracted["o"][0] is not None:
                w_o = _orient(extracted["o"][0].float(), self.all_head_size)
                U_o, S_o, Vh_o = torch.linalg.svd(w_o, full_matrices=False)
                r_out = min(self.rank_out, U_o.size(1), Vh_o.size(0))
                sqrt_so = torch.sqrt(torch.clamp(S_o[:r_out], min=1e-12))

                self.U_out.data[: min(self.hidden_size, U_o.size(0)), :r_out].copy_(
                    U_o[: min(self.hidden_size, U_o.size(0)), :r_out]
                )
                self.S_out_sqrt.data[:r_out].copy_(sqrt_so)
                self.Vh_out.data[:r_out, : min(self.all_head_size, Vh_o.size(1))].copy_(
                    Vh_o[:r_out, : min(self.all_head_size, Vh_o.size(1))]
                )

                if extracted["o"][1] is not None:
                    b_o = extracted["o"][1]
                    b_o_data = b_o.data if isinstance(b_o, nn.Parameter) else b_o
                    self.bias_out.data[: min(self.hidden_size, b_o_data.size(0))].copy_(
                        b_o_data[: min(self.hidden_size, b_o_data.size(0))]
                    )
            elif self.is_nested_encoder:
                self.U_out.data.zero_()
                min_d = min(self.hidden_size, self.rank_out)
                self.U_out.data[:min_d, :min_d].fill_diagonal_(1.0)
                self.S_out_sqrt.data.fill_(1.0)
                self.Vh_out.data.zero_()
                self.Vh_out.data[:min_d, :min_d].fill_diagonal_(1.0)
                self.bias_out.data.zero_()

            # Ensure base weights are frozen
            self.Vh_qkv.requires_grad = False
            self.S_qkv_sqrt.requires_grad = False
            self.U_qkv.requires_grad = False
            self.bias_qkv.requires_grad = False
            self.Vh_out.requires_grad = False
            self.S_out_sqrt.requires_grad = False
            self.U_out.requires_grad = False
            self.bias_out.requires_grad = False

            if self.lora_rank > 0:
                self.reset_lora_parameters()

        if hasattr(src, "is_decoder"):
            self.is_decoder = bool(src.is_decoder)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        past_key_value: Any | None = None,
        output_attentions: bool | None = False,
        output_hidden_states: bool | None = False,
        return_dict: bool | None = True,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        if hidden_states.dim() == 2:
            hidden_states = hidden_states.unsqueeze(0)
        B, N, D = hidden_states.shape
        device = hidden_states.device
        dtype = hidden_states.dtype

        # FIX 1: Pre-transpose weight matrices to avoid runtime transposes
        # Vh_qkv is (Rank, D), U_qkv is (Out_Dim, Rank)
        # S_qkv_sqrt is (Rank)

        # Optimize Matmul path
        v_proj = torch.matmul(hidden_states, self.Vh_qkv.t()) * self.S_qkv_sqrt
        qkv = torch.matmul(v_proj, (self.U_qkv * self.S_qkv_sqrt).t()) + self.bias_qkv

        # LoRA residual adapter
        if self.lora_rank > 0 and self.lora_A_qkv is not None and self.lora_B_qkv is not None:
            lora_qkv = (
                F.linear(F.linear(hidden_states, self.lora_A_qkv), self.lora_B_qkv)
                * self.lora_scaling
            )
            qkv = qkv + lora_qkv

        # Split into Q, K, V
        q, k, v = torch.split(
            qkv, [self.all_head_size, self.kv_head_size, self.kv_head_size], dim=-1
        )

        # FIX 2: Remove .contiguous() memory bombs
        q = q.view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, N, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, N, self.num_kv_heads, self.head_dim).transpose(1, 2)

        # Apply RoPE
        position_embeddings = kwargs.get("position_embeddings")
        if (
            position_embeddings is not None
            and isinstance(position_embeddings, (tuple, list))
            and len(position_embeddings) >= 2
        ):
            cos, sin = position_embeddings[0], position_embeddings[1]
            q, k = apply_rotary_pos_emb(q, k, cos, sin, kwargs.get("position_ids"))

        # GQA repeat
        if self.kv_repeat > 1:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        scale = 1.0 / math.sqrt(self.head_dim)

        # Use torch.baddbmm or scaled_dot_product_attention for huge memory savings
        # Fallback to standard matmul if not available or unsupported shapes
        try:
            # PyTorch 2.0+ Flash Attention / Memory Efficient Attention
            # This handles Causal masking automatically and saves massive VRAM
            is_causal = self.is_decoder and (attention_mask is None or attention_mask.dim() != 4)
            attn_out = F.scaled_dot_product_attention(
                q,
                k,
                v,
                attn_mask=attention_mask,
                dropout_p=self.dropout.p if self.training else 0.0,
                is_causal=is_causal,
            )
            probs = None  # SDPA doesn't return probabilities to save memory
        except Exception:
            # Legacy Path (if SDPA fails due to old PyTorch or weird masks)
            scores = torch.matmul(q, k.transpose(-2, -1)) * scale

            mask = attention_mask if attention_mask is not None else kwargs.get("mask")
            if mask is not None:
                if mask.dim() == 4:
                    if mask.dtype == torch.bool:
                        scores = scores.masked_fill(~mask, -10000.0)
                    else:
                        scores = scores + mask.to(scores.dtype)
                else:
                    valid_m = get_token_valid_mask(mask, B, N, device)
                    if valid_m is not None:
                        scores = scores.masked_fill((~valid_m).unsqueeze(1).unsqueeze(2), -10000.0)

            # Explicit Causal Masking
            if self.is_decoder and (mask is None or mask.dim() != 4):
                causal_tril = (
                    torch.tril(torch.ones((N, N), device=device, dtype=torch.bool))
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                scores = scores.masked_fill(~causal_tril, -10000.0)

            # FIX 3: Safe Softmax without extra type casts that create new tensors
            probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(dtype)
            probs = self.dropout(probs)
            attn_out = torch.matmul(probs, v)

        # Transpose back and view (reshape handles non-contiguous memory safely without deep copies)
        attn_out = attn_out.transpose(1, 2).reshape(B, N, self.all_head_size)

        if self.is_nested_encoder:
            out = attn_out
        else:
            # SVD Out forward
            out_v = torch.matmul(attn_out, self.Vh_out.t()) * self.S_out_sqrt
            out = torch.matmul(out_v, (self.U_out * self.S_out_sqrt).t()) + self.bias_out
            if self.lora_rank > 0 and self.lora_A_out is not None and self.lora_B_out is not None:
                lora_out = (
                    F.linear(F.linear(attn_out, self.lora_A_out), self.lora_B_out)
                    * self.lora_scaling
                )
                out = out + lora_out

        return out, (probs if output_attentions and probs is not None else None)

Every custom attention module must inherit from this.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

BaseAttention
torch.nn.modules.module.Module

Subclasses

CompressedAttentionLayer2

Class variables

var is_custom_attention : bool: The type of the None singleton.
var supports_bidirectional : bool: The type of the None singleton.
var supports_causal : bool: The type of the None singleton.
var supports_gqa : bool: The type of the None singleton.
var supports_rope : bool: The type of the None singleton.

Instance variables

prop out_weight : torch.Tensor

Expand source code

@property
def out_weight(self) -> torch.Tensor:
    """Reconstruct the effective base Out weight from SVD factors."""
    return (self.U_out * self.S_out_sqrt.view(1, -1)) @ (
        self.Vh_out * self.S_out_sqrt.view(-1, 1)
    )

Reconstruct the effective base Out weight from SVD factors.

prop qkv_weight : torch.Tensor

Expand source code

@property
def qkv_weight(self) -> torch.Tensor:
    """Reconstruct the effective base QKV weight from SVD factors."""
    return (self.U_qkv * self.S_qkv_sqrt.view(1, -1)) @ (
        self.Vh_qkv * self.S_qkv_sqrt.view(-1, 1)
    )

Reconstruct the effective base QKV weight from SVD factors.

Methods

def reset_lora_parameters(self) ‑> None

Expand source code

def reset_lora_parameters(self) -> None:
    if self.lora_rank > 0:
        nn.init.kaiming_uniform_(self.lora_A_qkv, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B_qkv)
        nn.init.kaiming_uniform_(self.lora_A_out, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B_out)
        self.lora_A_qkv.requires_grad = True
        self.lora_B_qkv.requires_grad = True
        self.lora_A_out.requires_grad = True
        self.lora_B_out.requires_grad = True

Inherited members

BaseAttention:
- copy_weights
- extra_repr
- forward

class CompressionPlanningTokens

Expand source code

class CompressionPlanningTokens:
    """
    Static CPT heuristic importance generator.
    Serves as safe baseline fallback when learnable CPT generator is inactive.
    """

    TASK_KEYWORD_MAP = {
        "qa": [
            "question",
            "answer",
            "context",
            "passage",
            "who",
            "what",
            "where",
            "why",
        ],
        "summarize": ["summary", "abstract", "key", "main", "brief", "conclusion"],
        "code": [
            "function",
            "class",
            "def",
            "return",
            "import",
            "self",
            "error",
            "test",
        ],
        "reasoning": [
            "therefore",
            "because",
            "if",
            "then",
            "thus",
            "proof",
            "step",
            "hence",
        ],
        "math": [
            "solve",
            "equation",
            "matrix",
            "integral",
            "theorem",
            "calculate",
            "equal",
        ],
    }

    @classmethod
    def generate_importance(
        cls,
        task_type: str,
        num_features: int,
        device: torch.device = torch.device("cpu"),
    ) -> torch.Tensor:
        """Generate per-feature importance modulation vector based on static heuristics."""
        task_lower = str(task_type).lower()
        importance = torch.zeros(num_features, device=device)

        # Baseline structure: prompt position bias
        pos_dim = max(1, num_features // 4)
        importance[:pos_dim] = 0.25

        if any(kw in task_lower for kw in ["qa", "question"]):
            importance[pos_dim : 2 * pos_dim] += 0.35
        elif any(kw in task_lower for kw in ["code", "program"]):
            importance[2 * pos_dim : 3 * pos_dim] += 0.40
        elif any(kw in task_lower for kw in ["reason", "math"]):
            importance += 0.20

        return torch.clamp(importance + 0.5, min=0.1, max=2.0)

Static CPT heuristic importance generator. Serves as safe baseline fallback when learnable CPT generator is inactive.

Class variables

var TASK_KEYWORD_MAP: The type of the None singleton.

Static methods

def generate_importance(task_type: str, num_features: int, device: torch.device = device(type='cpu')) ‑> torch.Tensor: Generate per-feature importance modulation vector based on static heuristics.

class Conv1D (nf, nx)

Expand source code

class Conv1D(nn.Module):
    """
    1D-convolutional layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2).

    Basically works like a linear layer but the weights are transposed.

    Args:
        nf (`int`): The number of output features.
        nx (`int`): The number of input features.
    """

    def __init__(self, nf, nx):
        super().__init__()
        self.nf = nf
        self.nx = nx
        self.weight = nn.Parameter(torch.empty(nx, nf))
        self.bias = nn.Parameter(torch.zeros(nf))
        nn.init.normal_(self.weight, std=0.02)

    def __repr__(self) -> str:
        return "Conv1D(nf={nf}, nx={nx})".format(**self.__dict__)

    def forward(self, x):
        size_out = x.size()[:-1] + (self.nf,)
        x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
        x = x.view(size_out)
        return x

1D-convolutional layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2).

Basically works like a linear layer but the weights are transposed.

Args -----= nf (int): The number of output features. nx (int): The number of input features. Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Methods

def forward(self, x) ‑> Callable[..., typing.Any]

Expand source code

def forward(self, x):
    size_out = x.size()[:-1] + (self.nf,)
    x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
    x = x.view(size_out)
    return x

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class CoopQLinear (in_features: int, out_features: int, bias: bool = True, bits: int = 4, symmetric: bool = False, per_channel: bool = True, lora_r: int = 0, lora_alpha: float = 1.0)

Expand source code

class CoopQLinear(nn.Module):
    """
    CoopQ Linear layer with per-channel asymmetric quantization and STE backpropagation.
    Supports both standard nn.Linear and Conv1D weight representations.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        bits: int = 4,
        symmetric: bool = False,
        per_channel: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features  # Conv1D compatibility
        self.bits = max(4, int(bits))
        self.symmetric = symmetric
        self.per_channel = per_channel
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        if bias:
            self.bias = nn.Parameter(torch.empty(out_features))
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(torch.empty(self.lora_r, in_features))
            self.lora_B = nn.Parameter(torch.zeros(out_features, self.lora_r))
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        nn.init.kaiming_uniform_(self.weight, a=5**0.5)
        if self.bias is not None:
            nn.init.zeros_(self.bias)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        bits: int = 4,
        symmetric: bool = False,
        per_channel: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> CoopQLinear:
        cls_name = linear.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(linear, "nf")
        raw_w = linear.weight.data

        if is_conv1d:
            in_f = raw_w.shape[0]
            out_f = raw_w.shape[1]
            w_data = raw_w.t().contiguous()
        else:
            out_f = raw_w.shape[0]
            in_f = raw_w.shape[1]
            w_data = raw_w.contiguous()

        layer = cls(
            in_features=in_f,
            out_features=out_f,
            bias=getattr(linear, "bias", None) is not None and linear.bias is not None,
            bits=bits,
            symmetric=symmetric,
            per_channel=per_channel,
            lora_r=lora_r,
            lora_alpha=lora_alpha,
        )
        layer.weight.data.copy_(w_data)
        if getattr(linear, "bias", None) is not None and linear.bias is not None:
            layer.bias.data.copy_(linear.bias.data)
        return layer

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w_q = quantize_dequantize_coopq(
            self.weight,
            bits=self.bits,
            symmetric=self.symmetric,
            per_channel=self.per_channel,
            use_ste=self.training,
        )
        b_q = None
        if self.bias is not None:
            b_q = quantize_dequantize_coopq(
                self.bias,
                bits=self.bits,
                symmetric=self.symmetric,
                per_channel=False,
                use_ste=self.training,
            )
        out = F.linear(
            x,
            w_q.to(dtype=x.dtype, device=x.device),
            b_q.to(dtype=x.dtype, device=x.device) if b_q is not None else None,
        )
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B)
            out = out + (self.scaling * lora_out).to(out.dtype)
        return out

CoopQ Linear layer with per-channel asymmetric quantization and STE backpropagation. Supports both standard nn.Linear and Conv1D weight representations.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Module, bits: int = 4, symmetric: bool = False, per_channel: bool = True, lora_r: int = 0, lora_alpha: float = 1.0) ‑> CoopQLinear

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    w_q = quantize_dequantize_coopq(
        self.weight,
        bits=self.bits,
        symmetric=self.symmetric,
        per_channel=self.per_channel,
        use_ste=self.training,
    )
    b_q = None
    if self.bias is not None:
        b_q = quantize_dequantize_coopq(
            self.bias,
            bits=self.bits,
            symmetric=self.symmetric,
            per_channel=False,
            use_ste=self.training,
        )
    out = F.linear(
        x,
        w_q.to(dtype=x.dtype, device=x.device),
        b_q.to(dtype=x.dtype, device=x.device) if b_q is not None else None,
    )
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B)
        out = out + (self.scaling * lora_out).to(out.dtype)
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class CoopQQuantizer (model: nn.Module, bits: int = 4, min_bits: int = 4, max_bits: int = 8, symmetric: bool = False, per_channel: bool = True, use_ste: bool = True, recalibrate_ln: bool = True, target_attention_only: bool = False, tokenizer: Any | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

class CoopQQuantizer(BaseCompressor):
    """
    CoopQ: Cooperative Game-Theoretic Bit Allocation Quantization.
    Allocates discrete bit-widths (>= 4 bits) across transformer layers
    proportional to layer sensitivity (Frobenius energy / variance), guaranteeing
    a target mean bit-budget under strict distortion minimization.

    Args:
        model: PyTorch model.
        bits: Target average bit-width (default 4).
        min_bits: Minimum allowable bit-width per layer (floor enforced at >= 4).
        max_bits: Maximum allowable bit-width per layer (default 8).
        symmetric: Whether to use symmetric quantization (default False for asymmetric calibration).
        per_channel: Whether to use per-channel scale factors (default True).
        use_ste: Whether to substitute nn.Linear with CoopQLinear for QAT (default True).
        target_attention_only: If True, only quantize attention projection layers (skip FFN/MLP).
        cfg: Optional configuration dictionary, namespace, or object.
    """

    def __init__(
        self,
        model: nn.Module,
        bits: int = 4,
        min_bits: int = 4,
        max_bits: int = 8,
        symmetric: bool = False,
        per_channel: bool = True,
        use_ste: bool = True,
        recalibrate_ln: bool = True,
        target_attention_only: bool = False,
        tokenizer: Any | None = None,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="coopq", **kwargs)
        # Extract config values with _cfg_get fallbacks
        c_bits = _cfg_get(cfg, "coopq_bits", _cfg_get(cfg, "bits", bits))
        c_min_bits = _cfg_get(cfg, "coopq_min_bits", _cfg_get(cfg, "min_bits", min_bits))
        c_max_bits = _cfg_get(cfg, "coopq_max_bits", _cfg_get(cfg, "max_bits", max_bits))
        c_symmetric = _cfg_get(cfg, "coopq_symmetric", _cfg_get(cfg, "symmetric", symmetric))
        c_per_channel = _cfg_get(
            cfg, "coopq_per_channel", _cfg_get(cfg, "per_channel", per_channel)
        )
        c_use_ste = _cfg_get(cfg, "coopq_use_ste", _cfg_get(cfg, "use_ste", use_ste))
        c_recalibrate_ln = _cfg_get(
            cfg, "coopq_recalibrate_ln", _cfg_get(cfg, "recalibrate_ln", recalibrate_ln)
        )
        c_target_attn = _cfg_get(
            cfg,
            "coopq_target_attention_only",
            _cfg_get(cfg, "target_attention_only", target_attention_only),
        )

        self.min_bits = max(4, int(c_min_bits))
        self.max_bits = max(self.min_bits, int(c_max_bits))
        self.bits = max(self.min_bits, min(int(c_bits), self.max_bits))
        self.symmetric = bool(c_symmetric)
        self.per_channel = bool(c_per_channel)
        self.use_ste = bool(c_use_ste)
        self.recalibrate_ln = bool(c_recalibrate_ln)
        self.target_attention_only = bool(c_target_attn)

        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.target_device = (
            torch.device(device)
            if device is not None
            else (next(model.parameters(), torch.tensor(0.0)).device)
        )
        self.stats: dict[str, Any] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        """Identify candidate linear/Conv1D/LoRALinear modules."""
        cls_name = module.__class__.__name__.lower()
        has_base = hasattr(module, "base_layer")
        base_cls = module.base_layer.__class__.__name__.lower() if has_base else ""

        is_linear = isinstance(module, nn.Linear) or "linear" in cls_name
        is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")
        is_lora = has_base and (
            isinstance(module.base_layer, nn.Linear)
            or "conv1d" in base_cls
            or hasattr(module.base_layer, "nf")
        )

        if not (is_linear or is_conv1d or is_lora):
            return False

        if isinstance(module, CoopQLinear) or (
            has_base and isinstance(module.base_layer, CoopQLinear)
        ):
            return False

        name_lower = name.lower()
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embed",
                "logits_processor",
                "pre_classifier",
                "head",
            ]
        ):
            return False

        if self.target_attention_only:
            ffn_patterns = [
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            ]
            if any(p in name_lower for p in ffn_patterns):
                return False
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            attn_patterns = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
                "in_proj_qkv",
                "out_proj",
            ]
            if not any(p in name_lower for p in attn_patterns):
                return False

        return True

    def _calculate_layer_sensitivities(self) -> dict[str, float]:
        """Compute sensitivity metric per target layer based on normalized Frobenius energy and variance."""
        sensitivities: dict[str, float] = {}
        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
                raw_w = target_mod.weight.data
                cls_name = module.__class__.__name__.lower()
                is_conv1d = (
                    "conv1d" in cls_name or hasattr(module, "nf") or hasattr(target_mod, "nf")
                )
                W = raw_w.t().float() if is_conv1d else raw_w.float()
                frob_norm = torch.norm(W, p="fro").item() / (W.numel() ** 0.5 + 1e-8)
                std_val = torch.std(W).item() if W.numel() > 1 else 0.0
                sensitivities[name] = float(frob_norm + std_val)
        return sensitivities

    def _allocate_bits(self, sensitivities: dict[str, float]) -> dict[str, int]:
        """Allocate discrete integer bit-widths satisfying mean(bits) <= self.bits with floor >= min_bits (>= 4)."""
        if not sensitivities:
            return {}

        total_sens = sum(sensitivities.values()) + 1e-8
        num_layers = len(sensitivities)
        target_total_bits = self.bits * num_layers

        # Allowed bit-levels with floor >= self.min_bits (>= 4)
        allowed_bits = [b for b in [4, 5, 6, 7, 8] if self.min_bits <= b <= self.max_bits]
        if not allowed_bits:
            allowed_bits = [self.min_bits]

        # Continuous allocation proportional to sensitivity
        raw_allocations: dict[str, float] = {}
        mean_sens = total_sens / num_layers
        for name, sens in sensitivities.items():
            sens_ratio = sens / mean_sens
            raw_bits = self.bits + (sens_ratio - 1.0) * (self.max_bits - self.min_bits) / 2.0
            raw_allocations[name] = max(float(self.min_bits), min(float(self.max_bits), raw_bits))

        discrete_allocations: dict[str, int] = {}
        for name, raw_b in raw_allocations.items():
            closest_b = min(allowed_bits, key=lambda b: abs(b - raw_b))
            discrete_allocations[name] = closest_b

        # Adjust sum to match target total bits budget
        current_sum = sum(discrete_allocations.values())
        sorted_names = sorted(sensitivities.keys(), key=lambda k: sensitivities[k], reverse=True)

        if current_sum > target_total_bits:
            for name in reversed(sorted_names):
                curr_b = discrete_allocations[name]
                valid_lower = [b for b in allowed_bits if b < curr_b]
                if valid_lower:
                    discrete_allocations[name] = max(valid_lower)
                    current_sum = sum(discrete_allocations.values())
                    if current_sum <= target_total_bits:
                        break
        elif current_sum < target_total_bits:
            for name in sorted_names:
                curr_b = discrete_allocations[name]
                valid_higher = [b for b in allowed_bits if b > curr_b]
                if valid_higher:
                    discrete_allocations[name] = min(valid_higher)
                    current_sum = sum(discrete_allocations.values())
                    if current_sum >= target_total_bits:
                        break

        # Ensure strict min_bits floor on all layers
        for k in discrete_allocations:
            discrete_allocations[k] = max(self.min_bits, discrete_allocations[k])

        return discrete_allocations

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Quantize linear layers with layer-specific bit allocation (>= 4 bits) and per-channel asymmetric calibration."""
        sensitivities = self._calculate_layer_sensitivities()
        bit_allocations = self._allocate_bits(sensitivities)

        total_bits_weighted = 0
        total_params_linear = 0
        compressed_layers = []
        fallback_layers = []

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                # When target_attention_only is True, unselected FFN/MLP layers must be frozen during fine-tuning
                if self.target_attention_only:
                    is_lin_like = (
                        isinstance(module, nn.Linear)
                        or "conv1d" in module.__class__.__name__.lower()
                        or hasattr(module, "base_layer")
                    )
                    if is_lin_like:
                        for p in module.parameters():
                            p.requires_grad = False
                continue

            if name not in bit_allocations:
                continue

            target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
            cls_name = module.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(module, "nf") or hasattr(target_mod, "nf")
            has_lora = (hasattr(module, "lora_A") and module.lora_A is not None) or hasattr(
                module, "base_layer"
            )

            layer_bits = bit_allocations[name]
            num_params = target_mod.weight.numel()
            total_bits_weighted += layer_bits * num_params
            total_params_linear += num_params

            if self.use_ste:
                src_lora_r = 0
                src_lora_alpha = 1.0
                if (
                    hasattr(module, "lora_A")
                    and hasattr(module, "lora_B")
                    and module.lora_A is not None
                    and module.lora_B is not None
                ):
                    src_lA = getattr(module.lora_A, "weight", module.lora_A)
                    if hasattr(src_lA, "shape") and len(src_lA.shape) == 2:
                        src_lora_r = src_lA.shape[0]
                    src_lora_alpha = getattr(module, "lora_alpha", float(src_lora_r))

                coopq_layer = CoopQLinear.from_linear(
                    target_mod,
                    bits=layer_bits,
                    symmetric=self.symmetric,
                    per_channel=self.per_channel,
                    lora_r=src_lora_r,
                    lora_alpha=src_lora_alpha,
                )
                new_layer = coopq_layer
                compressed_layer = new_layer
                # Preserve existing LoRA adapter weights if the source module has them
                if (
                    hasattr(module, "lora_A")
                    and hasattr(module, "lora_B")
                    and module.lora_A is not None
                    and module.lora_B is not None
                ):
                    src_lA = getattr(module.lora_A, "weight", module.lora_A)
                    src_lB = getattr(module.lora_B, "weight", module.lora_B)
                    if hasattr(new_layer, "lora_A") and new_layer.lora_A is not None:
                        if hasattr(src_lA, "data") and src_lA.shape == new_layer.lora_A.shape:
                            new_layer.lora_A.data.copy_(src_lA.data)
                    if hasattr(new_layer, "lora_B") and new_layer.lora_B is not None:
                        if hasattr(src_lB, "data") and src_lB.shape == new_layer.lora_B.shape:
                            new_layer.lora_B.data.copy_(src_lB.data)

                if hasattr(module, "base_layer"):
                    module.base_layer = compressed_layer  # put compressed layer inside wrapper
                    # module itself (LoRALinear wrapper) remains the replacement
                    new_module = module
                else:
                    new_module = compressed_layer

                if not hasattr(module, "base_layer"):
                    if name == "":
                        self.model = new_module
                    else:
                        replace_module(self.model, name, new_module)
            else:
                raw_w = target_mod.weight.data
                W = raw_w.t().float() if is_conv1d else raw_w.float()
                W_quant = quantize_dequantize_coopq(
                    W,
                    bits=layer_bits,
                    symmetric=self.symmetric,
                    per_channel=self.per_channel,
                    use_ste=False,
                )
                if is_conv1d:
                    target_mod.weight.data.copy_(
                        W_quant.t().contiguous().to(dtype=raw_w.dtype, device=raw_w.device)
                    )
                else:
                    target_mod.weight.data.copy_(
                        W_quant.contiguous().to(dtype=raw_w.dtype, device=raw_w.device)
                    )

                if target_mod.bias is not None:
                    b_quant = quantize_dequantize_coopq(
                        target_mod.bias.data.float(),
                        bits=layer_bits,
                        symmetric=self.symmetric,
                        per_channel=False,
                        use_ste=False,
                    )
                    target_mod.bias.data.copy_(
                        b_quant.to(dtype=target_mod.bias.dtype, device=target_mod.bias.device)
                    )

            # Preserve pre-existing LoRA adapter parameters and trainability
            if has_lora:
                for p_name, p in module.named_parameters():
                    if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                        p.requires_grad = True

            compressed_layers.append(name)

        # LayerNorm Recalibration
        if self.recalibrate_ln:
            recalibrate_layernorms(
                self.model,
                tokenizer=self.tokenizer,
                calibration_loader=self.calibration_loader,
                calibration_texts=self.calibration_texts,
                device=self.target_device,
            )

        avg_bits = (
            total_bits_weighted / total_params_linear
            if total_params_linear > 0
            else float(self.bits)
        )
        effective_compression_rate = 16.0 / max(avg_bits, 1e-3)
        stats = self.get_compression_stats()
        stats.update(
            {
                "target_bits": self.bits,
                "min_bits_floor": self.min_bits,
                "average_allocated_bits": round(avg_bits, 2),
                "layer_bit_allocations": bit_allocations,
                "compressed_layers": compressed_layers,
                "compressed_layer_count": len(compressed_layers),
                "fallback_layers": fallback_layers,
                "fallback_layer_count": len(fallback_layers),
                "symmetric": self.symmetric,
                "per_channel": self.per_channel,
                "use_ste": self.use_ste,
                "target_attention_only": self.target_attention_only,
                "effective_compression_rate": round(effective_compression_rate, 2),
                "recalibrate_ln": self.recalibrate_ln,
                "method": "coopq_asymmetric_ste_bmin4",
                "theoretical_vram_saving": f"{(1.0 - avg_bits / 16.0) * 100:.1f}% vs FP16",
            }
        )
        self.stats = stats
        return self.model, stats

CoopQ: Cooperative Game-Theoretic Bit Allocation Quantization. Allocates discrete bit-widths (>= 4 bits) across transformer layers proportional to layer sensitivity (Frobenius energy / variance), guaranteeing a target mean bit-budget under strict distortion minimization.

Args
-----=
model: PyTorch model.
bits: Target average bit-width (default 4).
min_bits: Minimum allowable bit-width per layer (floor enforced at >= 4).
max_bits: Maximum allowable bit-width per layer (default 8).
symmetric: Whether to use symmetric quantization (default False for asymmetric calibration).
per_channel: Whether to use per-channel scale factors (default True).
use_ste: Whether to substitute nn.Linear with CoopQLinear for QAT (default True).
target_attention_only: If True, only quantize attention projection layers (skip FFN/MLP).
cfg: Optional configuration dictionary, namespace, or object.

Ancestors

BaseCompressor

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Quantize linear layers with layer-specific bit allocation (>= 4 bits) and per-channel asymmetric calibration."""
    sensitivities = self._calculate_layer_sensitivities()
    bit_allocations = self._allocate_bits(sensitivities)

    total_bits_weighted = 0
    total_params_linear = 0
    compressed_layers = []
    fallback_layers = []

    for name, module in list(self.model.named_modules()):
        if not self._is_target_module(name, module):
            # When target_attention_only is True, unselected FFN/MLP layers must be frozen during fine-tuning
            if self.target_attention_only:
                is_lin_like = (
                    isinstance(module, nn.Linear)
                    or "conv1d" in module.__class__.__name__.lower()
                    or hasattr(module, "base_layer")
                )
                if is_lin_like:
                    for p in module.parameters():
                        p.requires_grad = False
            continue

        if name not in bit_allocations:
            continue

        target_mod = getattr(module, "base_layer", getattr(module, "linear", module))
        cls_name = module.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(module, "nf") or hasattr(target_mod, "nf")
        has_lora = (hasattr(module, "lora_A") and module.lora_A is not None) or hasattr(
            module, "base_layer"
        )

        layer_bits = bit_allocations[name]
        num_params = target_mod.weight.numel()
        total_bits_weighted += layer_bits * num_params
        total_params_linear += num_params

        if self.use_ste:
            src_lora_r = 0
            src_lora_alpha = 1.0
            if (
                hasattr(module, "lora_A")
                and hasattr(module, "lora_B")
                and module.lora_A is not None
                and module.lora_B is not None
            ):
                src_lA = getattr(module.lora_A, "weight", module.lora_A)
                if hasattr(src_lA, "shape") and len(src_lA.shape) == 2:
                    src_lora_r = src_lA.shape[0]
                src_lora_alpha = getattr(module, "lora_alpha", float(src_lora_r))

            coopq_layer = CoopQLinear.from_linear(
                target_mod,
                bits=layer_bits,
                symmetric=self.symmetric,
                per_channel=self.per_channel,
                lora_r=src_lora_r,
                lora_alpha=src_lora_alpha,
            )
            new_layer = coopq_layer
            compressed_layer = new_layer
            # Preserve existing LoRA adapter weights if the source module has them
            if (
                hasattr(module, "lora_A")
                and hasattr(module, "lora_B")
                and module.lora_A is not None
                and module.lora_B is not None
            ):
                src_lA = getattr(module.lora_A, "weight", module.lora_A)
                src_lB = getattr(module.lora_B, "weight", module.lora_B)
                if hasattr(new_layer, "lora_A") and new_layer.lora_A is not None:
                    if hasattr(src_lA, "data") and src_lA.shape == new_layer.lora_A.shape:
                        new_layer.lora_A.data.copy_(src_lA.data)
                if hasattr(new_layer, "lora_B") and new_layer.lora_B is not None:
                    if hasattr(src_lB, "data") and src_lB.shape == new_layer.lora_B.shape:
                        new_layer.lora_B.data.copy_(src_lB.data)

            if hasattr(module, "base_layer"):
                module.base_layer = compressed_layer  # put compressed layer inside wrapper
                # module itself (LoRALinear wrapper) remains the replacement
                new_module = module
            else:
                new_module = compressed_layer

            if not hasattr(module, "base_layer"):
                if name == "":
                    self.model = new_module
                else:
                    replace_module(self.model, name, new_module)
        else:
            raw_w = target_mod.weight.data
            W = raw_w.t().float() if is_conv1d else raw_w.float()
            W_quant = quantize_dequantize_coopq(
                W,
                bits=layer_bits,
                symmetric=self.symmetric,
                per_channel=self.per_channel,
                use_ste=False,
            )
            if is_conv1d:
                target_mod.weight.data.copy_(
                    W_quant.t().contiguous().to(dtype=raw_w.dtype, device=raw_w.device)
                )
            else:
                target_mod.weight.data.copy_(
                    W_quant.contiguous().to(dtype=raw_w.dtype, device=raw_w.device)
                )

            if target_mod.bias is not None:
                b_quant = quantize_dequantize_coopq(
                    target_mod.bias.data.float(),
                    bits=layer_bits,
                    symmetric=self.symmetric,
                    per_channel=False,
                    use_ste=False,
                )
                target_mod.bias.data.copy_(
                    b_quant.to(dtype=target_mod.bias.dtype, device=target_mod.bias.device)
                )

        # Preserve pre-existing LoRA adapter parameters and trainability
        if has_lora:
            for p_name, p in module.named_parameters():
                if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                    p.requires_grad = True

        compressed_layers.append(name)

    # LayerNorm Recalibration
    if self.recalibrate_ln:
        recalibrate_layernorms(
            self.model,
            tokenizer=self.tokenizer,
            calibration_loader=self.calibration_loader,
            calibration_texts=self.calibration_texts,
            device=self.target_device,
        )

    avg_bits = (
        total_bits_weighted / total_params_linear
        if total_params_linear > 0
        else float(self.bits)
    )
    effective_compression_rate = 16.0 / max(avg_bits, 1e-3)
    stats = self.get_compression_stats()
    stats.update(
        {
            "target_bits": self.bits,
            "min_bits_floor": self.min_bits,
            "average_allocated_bits": round(avg_bits, 2),
            "layer_bit_allocations": bit_allocations,
            "compressed_layers": compressed_layers,
            "compressed_layer_count": len(compressed_layers),
            "fallback_layers": fallback_layers,
            "fallback_layer_count": len(fallback_layers),
            "symmetric": self.symmetric,
            "per_channel": self.per_channel,
            "use_ste": self.use_ste,
            "target_attention_only": self.target_attention_only,
            "effective_compression_rate": round(effective_compression_rate, 2),
            "recalibrate_ln": self.recalibrate_ln,
            "method": "coopq_asymmetric_ste_bmin4",
            "theoretical_vram_saving": f"{(1.0 - avg_bits / 16.0) * 100:.1f}% vs FP16",
        }
    )
    self.stats = stats
    return self.model, stats

Quantize linear layers with layer-specific bit allocation (>= 4 bits) and per-channel asymmetric calibration.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class CoopQSTEFakeQuantize (*args, **kwargs)

Expand source code

class CoopQSTEFakeQuantize(torch.autograd.Function):
    """
    Straight-Through Estimator (STE) fake-quantization autograd function for CoopQ.
    Forwards per-channel quantized-dequantized weights and passes gradients through on backward pass.
    """

    @staticmethod
    def forward(
        ctx: Any,
        x: torch.Tensor,
        scale: torch.Tensor,
        zero_point: torch.Tensor,
        qmin: float,
        qmax: float,
    ) -> torch.Tensor:
        ctx.save_for_backward(scale, zero_point)
        ctx.qmin = qmin
        ctx.qmax = qmax

        # Quantize and clamp
        q = torch.clamp(torch.round(x / scale) + zero_point, qmin, qmax)
        # De-quantize
        dq = (q - zero_point) * scale
        return dq

    @staticmethod
    def backward(
        ctx: Any, grad_output: torch.Tensor
    ) -> tuple[torch.Tensor, None, None, None, None]:
        # STE passes gradient uninhibited to parameter tensor
        grad_x = grad_output.clone()
        return grad_x, None, None, None, None

Straight-Through Estimator (STE) fake-quantization autograd function for CoopQ. Forwards per-channel quantized-dequantized weights and passes gradients through on backward pass.

Ancestors

torch.autograd.function.Function
torch.autograd.function._SingleLevelFunction
torch._C._FunctionBase
torch.autograd.function.FunctionCtx
torch.autograd.function._HookMixin

Static methods

def backward(ctx: Any, grad_output: torch.Tensor) ‑> tuple[torch.Tensor, None, None, None, None]

Expand source code

@staticmethod
def backward(
    ctx: Any, grad_output: torch.Tensor
) -> tuple[torch.Tensor, None, None, None, None]:
    # STE passes gradient uninhibited to parameter tensor
    grad_x = grad_output.clone()
    return grad_x, None, None, None, None

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the vjp function.)

def forward(ctx: Any, x: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, qmin: float, qmax: float) ‑> torch.Tensor

Expand source code

@staticmethod
def forward(
    ctx: Any,
    x: torch.Tensor,
    scale: torch.Tensor,
    zero_point: torch.Tensor,
    qmin: float,
    qmax: float,
) -> torch.Tensor:
    ctx.save_for_backward(scale, zero_point)
    ctx.qmin = qmin
    ctx.qmax = qmax

    # Quantize and clamp
    q = torch.clamp(torch.round(x / scale) + zero_point, qmin, qmax)
    # De-quantize
    dq = (q - zero_point) * scale
    return dq

Define the forward of the custom autograd Function.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx)::

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See :ref:combining-forward-context for more details

Usage 2 (Separate forward and ctx)::

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass


@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

The forward no longer accepts a ctx argument.
Instead, you must also override the :meth:torch.autograd.Function.setup_context staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.
See :ref:extending-autograd for more details

class DRONECompressor (model: nn.Module, tokenizer: Any | None = None, rank: int = 64, max_error: float = 0.5, fallback_to_original: bool = True, device: str | torch.device | None = None, target_modules: list[str] | None = None, target_attention_only: bool = True, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("drone")
@COMPRESSOR_REGISTRY.register("drone_compressor")
class DRONECompressor(BaseCompressor):
    """
    DRONE Compressor.
    Data-aware SVD with identity weighting (standard data-independent SVD baseline)
    using identical rank adaptation, parameter constraints, and fallback error mechanism as AASVD.
    """

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any | None = None,
        rank: int = 64,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        device: str | torch.device | None = None,
        target_modules: list[str] | None = None,
        target_attention_only: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="drone", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            rank = getattr(cfg, "drone_rank", getattr(cfg, "rank", rank))
            max_error = getattr(cfg, "drone_max_error", getattr(cfg, "max_error", max_error))
            fallback_to_original = getattr(
                cfg,
                "drone_fallback_to_original",
                getattr(cfg, "fallback_to_original", fallback_to_original),
            )
            target_attention_only = getattr(
                cfg,
                "drone_target_attention_only",
                getattr(cfg, "target_attention_only", target_attention_only),
            )

        target_attention_only = kwargs.get(
            "target_attention_only",
            kwargs.get("drone_target_attention_only", target_attention_only),
        )
        self.rank = rank
        if max_error is None or max_error == 0.50:
            max_error = adaptive_max_error(self.rank, 768, 768, "distilbert")

        self.max_error = float(max_error)
        self.fallback_to_original = fallback_to_original
        self.tokenizer = tokenizer
        self.target_modules = target_modules
        self.target_attention_only = target_attention_only
        self.device = (
            torch.device(device) if device is not None else next(model.parameters()).device
        )
        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        if isinstance(module, SVDLinear):
            return False
        if any(
            name.endswith(sfx) for sfx in [".base_layer", ".lora_A", ".lora_B", ".lora_dropout"]
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = module.__class__.__name__.lower()
        target_cls = target_mod.__class__.__name__.lower()

        is_linear_like = (
            isinstance(module, nn.Linear)
            or isinstance(target_mod, nn.Linear)
            or "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(module, "nf")
            or hasattr(target_mod, "nf")
            or "loralinear" in cls_name
            or hasattr(module, "base_layer")
        )
        if not is_linear_like:
            return False
        name_lower = name.lower()
        if any(
            skip in name_lower
            for skip in [
                "lm_head",
                "classifier",
                "head",
                "output_layer",
                "score",
                "pooler",
            ]
        ):
            return False

        if self.target_attention_only:
            ffn_patterns = [
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            ]
            if any(p in name_lower for p in ffn_patterns):
                return False
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            attn_patterns = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(p in name_lower for p in attn_patterns):
                return False
        else:
            if not should_compress_linear(self.model, name):
                return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def compress_linear(
        self,
        layer: nn.Module,
        rank: int | None = None,
    ) -> tuple[SVDLinear | None, float, int]:
        try:
            target_mod = getattr(layer, "base_layer", layer)
            cls_name = layer.__class__.__name__.lower()
            target_cls = target_mod.__class__.__name__.lower()
            is_conv1d = (
                "conv1d" in cls_name
                or "conv1d" in target_cls
                or hasattr(layer, "nf")
                or hasattr(target_mod, "nf")
            )
            weight_tensor = getattr(target_mod, "weight", getattr(layer, "weight", None))
            if weight_tensor is None:
                return None, float("inf"), 0

            raw_w = weight_tensor.data

            if is_conv1d:
                in_feat, out_feat = raw_w.shape[0], raw_w.shape[1]
                W = raw_w.t().contiguous().float()
            else:
                out_feat, in_feat = raw_w.shape[0], raw_w.shape[1]
                W = raw_w.float()

            r = get_dynamic_compression_rank(
                in_feat,
                out_feat,
                model_or_family=self.model,
                default_rank=rank or self.rank,
                comp_config=getattr(self, "comp_config", None),
            )
            max_allowed_rank = max(
                1,
                min(
                    in_feat,
                    out_feat,
                    int(0.75 * (in_feat * out_feat) / (in_feat + out_feat)),
                ),
            )
            r = min(r, max_allowed_rank)
            W_norm = torch.norm(W, p="fro").item()

            W_dev = W.to(device=layer.weight.device, dtype=torch.float32)
            U, S, Vt = torch.linalg.svd(W_dev, full_matrices=False)

            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
            U_r = U[:, :r] * sqrt_S.unsqueeze(0)
            V_r = sqrt_S.unsqueeze(1) * Vt[:r, :]

            W_recon = torch.matmul(U_r, V_r)
            err = float(torch.norm(W - W_recon, p="fro").item() / max(W_norm, 1e-12))

            dtype = weight_tensor.dtype
            device = weight_tensor.device
            bias_tensor = getattr(target_mod, "bias", getattr(layer, "bias", None))

            svd_layer = SVDLinear(
                in_features=in_feat,
                out_features=out_feat,
                rank=r,
                bias=bias_tensor is not None,
                dtype=dtype,
                device=device,
                lora_r=getattr(self, "lora_r", 0),
                lora_alpha=getattr(self, "lora_alpha", 1.0),
            )
            svd_layer.set_low_rank_factors(
                U_r.to(device=device, dtype=dtype),
                V_r.to(device=device, dtype=dtype),
            )
            if bias_tensor is not None and svd_layer.bias is not None:
                svd_layer.bias.data.copy_(bias_tensor.data)

            svd_layer.reconstruction_error = err
            svd_layer.rank = r
            return svd_layer, err, r
        except Exception as e:
            logger.warning("DRONE factorization failed on linear: %s", e)
            return None, float("inf"), 0

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        self.model.eval()
        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}

        for name, module in list(self.model.named_modules()):
            if self._is_target_module(name, module):
                svd_layer, err, r = self.compress_linear(module, rank=self.rank)
                self.layer_errors[name] = err
                self.layer_ranks[name] = r

                if svd_layer is None or err > self.max_error:
                    if self.fallback_to_original:
                        logger.warning(
                            "DRONE reconstruction error %.4f on %s exceeded threshold %.4f. Skipping compression (fallback).",
                            err,
                            name,
                            self.max_error,
                        )
                        self.fallback_layers.append(name)
                        continue
                    else:
                        if svd_layer is None:
                            self.fallback_layers.append(name)
                            continue

                if hasattr(module, "base_layer"):
                    module.base_layer = svd_layer
                else:
                    replace_module(self.model, name, svd_layer)
                self.compressed_layers.append(name)

        stats = {
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "avg_reconstruction_error": (
                float(sum(self.layer_errors.values()) / max(1, len(self.layer_errors)))
                if self.layer_errors
                else 0.0
            ),
            "layer_ranks": self.layer_ranks,
        }
        self.model.compression_stats = stats
        logger.info(
            "Compressor drone finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            stats["avg_reconstruction_error"],
            self.max_error,
        )
        return self.model, stats

DRONE Compressor. Data-aware SVD with identity weighting (standard data-independent SVD baseline) using identical rank adaptation, parameter constraints, and fallback error mechanism as AASVD.

Ancestors

BaseCompressor

Methods

def compress_linear(self, layer: nn.Module, rank: int | None = None) ‑> tuple[SVDLinear | None, float, int]

Expand source code

def compress_linear(
    self,
    layer: nn.Module,
    rank: int | None = None,
) -> tuple[SVDLinear | None, float, int]:
    try:
        target_mod = getattr(layer, "base_layer", layer)
        cls_name = layer.__class__.__name__.lower()
        target_cls = target_mod.__class__.__name__.lower()
        is_conv1d = (
            "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(layer, "nf")
            or hasattr(target_mod, "nf")
        )
        weight_tensor = getattr(target_mod, "weight", getattr(layer, "weight", None))
        if weight_tensor is None:
            return None, float("inf"), 0

        raw_w = weight_tensor.data

        if is_conv1d:
            in_feat, out_feat = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()
        else:
            out_feat, in_feat = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float()

        r = get_dynamic_compression_rank(
            in_feat,
            out_feat,
            model_or_family=self.model,
            default_rank=rank or self.rank,
            comp_config=getattr(self, "comp_config", None),
        )
        max_allowed_rank = max(
            1,
            min(
                in_feat,
                out_feat,
                int(0.75 * (in_feat * out_feat) / (in_feat + out_feat)),
            ),
        )
        r = min(r, max_allowed_rank)
        W_norm = torch.norm(W, p="fro").item()

        W_dev = W.to(device=layer.weight.device, dtype=torch.float32)
        U, S, Vt = torch.linalg.svd(W_dev, full_matrices=False)

        sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
        U_r = U[:, :r] * sqrt_S.unsqueeze(0)
        V_r = sqrt_S.unsqueeze(1) * Vt[:r, :]

        W_recon = torch.matmul(U_r, V_r)
        err = float(torch.norm(W - W_recon, p="fro").item() / max(W_norm, 1e-12))

        dtype = weight_tensor.dtype
        device = weight_tensor.device
        bias_tensor = getattr(target_mod, "bias", getattr(layer, "bias", None))

        svd_layer = SVDLinear(
            in_features=in_feat,
            out_features=out_feat,
            rank=r,
            bias=bias_tensor is not None,
            dtype=dtype,
            device=device,
            lora_r=getattr(self, "lora_r", 0),
            lora_alpha=getattr(self, "lora_alpha", 1.0),
        )
        svd_layer.set_low_rank_factors(
            U_r.to(device=device, dtype=dtype),
            V_r.to(device=device, dtype=dtype),
        )
        if bias_tensor is not None and svd_layer.bias is not None:
            svd_layer.bias.data.copy_(bias_tensor.data)

        svd_layer.reconstruction_error = err
        svd_layer.rank = r
        return svd_layer, err, r
    except Exception as e:
        logger.warning("DRONE factorization failed on linear: %s", e)
        return None, float("inf"), 0

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class DeltaKVCompressor (model: nn.Module, rank: int = 16, sparsity: float = 0.29, max_error: float = 0.35, target_attention_only: bool = False, fallback_to_original: bool = True, min_layer_dim: int = 8, skip_should_compress_filter: bool = False, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("delta_kv_class")
@COMPRESSOR_REGISTRY.register("delta_kv_compressor_class")
class DeltaKVCompressor(BaseCompressor):
    # DKV-FIX-01: this is a compressor, not attention.
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module,
        rank: int = 16,
        sparsity: float = 0.29,
        max_error: float = 0.35,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        min_layer_dim: int = 8,
        skip_should_compress_filter: bool = False,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="delta_kv", **kwargs)
        self.cfg = cfg

        self.rank = int(_cfg_get(cfg, "delta_kv_rank", _cfg_get(cfg, "rank", rank)))
        self.sparsity = float(
            _cfg_get(cfg, "delta_kv_sparsity", _cfg_get(cfg, "sparsity", sparsity))
        )
        self.max_error = float(
            _cfg_get(cfg, "delta_kv_max_error", _cfg_get(cfg, "max_error", max_error))
        )
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "delta_kv_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "delta_kv_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        self.min_layer_dim = max(1, int(_cfg_get(cfg, "min_layer_dim", min_layer_dim)))
        self.skip_should_compress_filter = bool(
            _cfg_get(
                cfg,
                "skip_should_compress_filter",
                skip_should_compress_filter,
            )
        )

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.stats: dict[str, Any] = {}
        self._compressed_params: int = 0

    # ------------------------------------------------------------------
    # Target check (DKV-FIX-07)
    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module, _debug: bool = False) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ("lora_A", "lora_B", "base_layer")):
            if _debug:
                logger.info(f"  [skip] {name}: lora/base_layer part")
            return False

        # Skip already compressed (DKV-FIX-05 companion)
        cls_name = module.__class__.__name__
        if cls_name in (
            "DeltaKVLinear",
            "LSCPrunedLoRALinear",
            "DynamicSVDLinear",
            "SVDLinear",
            "FactorisedLinear",
        ):
            if _debug:
                logger.info(f"  [skip] {name}: already compressed ({cls_name})")
            return False

        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            ]
        ):
            if _debug:
                logger.info(f"  [skip] {name}: skip keyword")
            return False

        target_mod = getattr(module, "base_layer", module)
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = _is_conv1d(target_mod)
        if not (is_linear or is_conv1d):
            if _debug:
                logger.info(f"  [skip] {name}: not Linear/Conv1D")
            return False

        # DKV-FIX-06: bypass framework filter if requested.
        if not self.skip_should_compress_filter:
            try:
                from neuralops.utils.compat import should_compress_linear

                if not should_compress_linear(self.model, name):
                    if _debug:
                        logger.info(f"  [skip] {name}: should_compress_linear=False")
                    return False
            except Exception as e:
                if _debug:
                    logger.info(f"  [skip] {name}: should_compress_linear raised {e}")
                return False

        w = getattr(target_mod, "weight", None)
        if w is None:
            if _debug:
                logger.info(f"  [skip] {name}: weight None")
            return False
        # DKV-FIX-11: min_layer_dim filter.
        if min(w.shape[0], w.shape[1]) < self.min_layer_dim:
            if _debug:
                logger.info(f"  [skip] {name}: dim too small {tuple(w.shape)}")
            return False

        if self.target_attention_only:
            last_part = name.split(".")[-1]
            is_attn = (last_part in _ATTN_NAMES) or any(
                k in name_lower for k in ("attn", "attention")
            )
            if not is_attn:
                if _debug:
                    logger.info(f"  [skip] {name}: target_attention_only filter")
                return False

        if _debug:
            logger.info(f"  [TARGET] {name}: shape={tuple(w.shape)}")
        return True

    # ------------------------------------------------------------------
    # Main compress
    # ------------------------------------------------------------------
    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        # DKV-FIX-09: idempotence
        if getattr(self.model, "_delta_kv_compressed", False):
            logger.warning("DeltaKV: model already compressed; skipping.")
            try:
                stats = self.get_compression_stats()
            except Exception:
                stats = {}
            return self.model, stats

        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}

        total_orig_params = 0
        total_comp_params = 0

        # Diagnostic breakdown if no targets.
        target_list = [
            (name, module)
            for name, module in self.model.named_modules()
            if self._is_target_module(name, module)
        ]
        if not target_list:
            logger.warning("DeltaKV: no compressible layers found. Running diagnostic...")
            n_scanned = 0
            for name, module in self.model.named_modules():
                n_scanned += 1
                self._is_target_module(name, module, _debug=True)
            logger.warning("DeltaKV diagnostic: %d modules scanned, 0 targets.", n_scanned)
            try:
                stats = self.get_compression_stats()
            except Exception:
                stats = {}
            return self.model, stats

        logger.info("DeltaKV: %d target layers identified.", len(target_list))

        # Freeze non-target attention/linear layers when target_attention_only.
        if self.target_attention_only:
            for name, module in self.model.named_modules():
                if any(part in name.split(".") for part in ("lora_A", "lora_B", "base_layer")):
                    continue
                target_mod = getattr(module, "base_layer", module)
                if isinstance(target_mod, nn.Linear) or _is_conv1d(target_mod):
                    if not self._is_target_module(name, module):
                        for p in module.parameters():
                            p.requires_grad = False
                        self.skipped_layers.append(name)

        for name, module in target_list:
            target_mod = getattr(module, "base_layer", module)
            raw_w = target_mod.weight.data
            is_conv1d = _is_conv1d(target_mod)

            if is_conv1d:
                in_f, out_f = raw_w.shape[0], raw_w.shape[1]
                W = raw_w.t().contiguous().float()
            else:
                out_f, in_f = raw_w.shape[0], raw_w.shape[1]
                W = raw_w.float()

            orig_layer_params = out_f * in_f
            total_orig_params += orig_layer_params
            norm_w = torch.norm(W, p="fro").item()

            try:
                U, S, Vt = torch.linalg.svd(W, full_matrices=False)
                # DKV-FIX-03: correct min over 2D shape.
                r = max(1, min(self.rank, out_f, in_f, len(S)))
                L = (U[:, :r] * S[:r].unsqueeze(0)) @ Vt[:r, :]

                residual = W - L
                flat_r = residual.contiguous().abs().reshape(-1)
                keep_n = max(1, min(int(flat_r.numel() * self.sparsity), flat_r.numel()))
                threshold = torch.topk(flat_r, keep_n).values[-1]
                mask = (residual.abs() >= threshold).float()
                residual_sparse = residual * mask

                W_compressed = L + residual_sparse
                err = float(torch.norm(W - W_compressed, p="fro").item() / max(norm_w, 1e-12))
                self.layer_errors[name] = err
                self.layer_ranks[name] = r

                if err > self.max_error and self.fallback_to_original:
                    logger.warning(
                        "DeltaKV: error %.4f on %s > %.4f; keeping original.",
                        err,
                        name,
                        self.max_error,
                    )
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    total_comp_params += orig_layer_params
                    continue

                # Write back in original orientation and dtype.
                if is_conv1d:
                    target_mod.weight.data.copy_(
                        W_compressed.t().contiguous().to(dtype=target_mod.weight.dtype)
                    )
                else:
                    target_mod.weight.data.copy_(W_compressed.to(dtype=target_mod.weight.dtype))

                # DKV-FIX-10: requires_grad handled at module level.
                for p in module.parameters():
                    p.requires_grad = True

                self.compressed_layers.append(name)
                # DKV-FIX-12: honest parameter accounting.
                # Note: the dense weight is retained in-place (no reduction);
                # we report the theoretical storage ratio here for reference.
                comp_layer_params = r * (in_f + out_f) + keep_n
                total_comp_params += comp_layer_params

            except Exception as e:
                logger.warning("DeltaKV: decomposition failed on %s: %s", name, e)
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_params += orig_layer_params

        # DKV-FIX-08: forward sanity check.
        forward_check_success = True
        try:
            with torch.no_grad():
                first_p = next(iter(self.model.parameters()), None)
                dev = first_p.device if first_p is not None else torch.device("cpu")
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                try:
                    self.model(input_ids=dummy, attention_mask=torch.ones_like(dummy))
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception as e2:
                        logger.warning("DeltaKV: forward check failed: %s", e2)
                        forward_check_success = False
        except Exception as e:
            logger.warning("DeltaKV: forward check raised: %s", e)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        effective_ratio = total_comp_params / total_orig_params if total_orig_params > 0 else 1.0
        avg_err = (
            sum(self.layer_errors.values()) / max(1, len(self.layer_errors))
            if self.layer_errors
            else 0.0
        )

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        stats.update(
            {
                "rank": self.rank,
                "sparsity": self.sparsity,
                "max_error": self.max_error,
                "target_attention_only": self.target_attention_only,
                "fallback_to_original": self.fallback_to_original,
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "avg_reconstruction_error": avg_err,
                "effective_parameter_ratio": round(effective_ratio, 4),
                "kv_cache_storage_ratio": f"{self.sparsity * 100:.1f}% of full KV cache",
                "forward_check_success": forward_check_success,
                "method": "delta_kv",
                "note": (
                    "DeltaKV rewrites dense weights in-place (low-rank + sparse residual). "
                    "The stored weight retains its full (out, in) shape; framework-level "
                    "compression_rate reads 1.0. Use effective_parameter_ratio for the "
                    "theoretical storage metric."
                ),
            }
        )
        self.stats = stats

        # DKV-FIX-14: write stats on the model.
        try:
            if not hasattr(self.model, "compression_stats") or self.model.compression_stats is None:
                self.model.compression_stats = {}
            self.model.compression_stats.update(
                {
                    "num_compressed": len(self.compressed_layers),
                    "num_fallback": len(self.fallback_layers),
                    "compression_method": "delta_kv",
                    "effective_parameter_ratio": effective_ratio,
                    "avg_reconstruction_error": avg_err,
                }
            )
        except Exception:
            pass

        self.model._delta_kv_compressed = True
        logger.info(
            "DeltaKV: %d compressed, %d fallback, %d skipped. "
            "Avg err %.4f, effective_param_ratio %.2fx.",
            len(self.compressed_layers),
            len(self.fallback_layers),
            len(self.skipped_layers),
            avg_err,
            effective_ratio,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class DeltaKVStore (max_refs: int = 128, topk: int = 4, compression_ratio: float = 0.29, residual_threshold: float = 0.001)

Expand source code

class DeltaKVStore:
    """
    Optional runtime cache that stores delta residuals relative to retrieved
    historical K/V references. Not used by the weight-only compressor path.
    """

    def __init__(
        self,
        max_refs: int = 128,
        topk: int = 4,
        compression_ratio: float = 0.29,
        residual_threshold: float = 1e-3,
    ) -> None:
        self.max_refs = int(max_refs)
        self.topk = int(topk)
        self.compression_ratio = float(compression_ratio)
        self.residual_threshold = float(residual_threshold)
        self.ref_k: torch.Tensor | None = None
        self.ref_v: torch.Tensor | None = None

    def reset(self) -> None:
        self.ref_k = None
        self.ref_v = None

    @torch.no_grad()
    def update(self, k: torch.Tensor, v: torch.Tensor) -> None:
        orig_k = k.detach()
        orig_v = v.detach()
        if orig_k.ndim == 2:
            orig_k = orig_k.unsqueeze(0).unsqueeze(0)
            orig_v = orig_v.unsqueeze(0).unsqueeze(0)
        elif orig_k.ndim == 3:
            orig_k = orig_k.unsqueeze(1)
            orig_v = orig_v.unsqueeze(1)

        B, H, N, D = orig_k.shape
        if N == 0 or B == 0:
            return

        if self.ref_k is None:
            self.ref_k = orig_k
            self.ref_v = orig_v
        else:
            if self.ref_k.device != orig_k.device or self.ref_k.dtype != orig_k.dtype:
                self.ref_k = self.ref_k.to(device=orig_k.device, dtype=orig_k.dtype)
                self.ref_v = self.ref_v.to(device=orig_v.device, dtype=orig_v.dtype)

            if self.ref_k.shape[0] != B:
                self.ref_k = self.ref_k[:1].expand(B, -1, -1, -1).clone()
                self.ref_v = self.ref_v[:1].expand(B, -1, -1, -1).clone()

            self.ref_k = torch.cat([self.ref_k, orig_k], dim=2)
            self.ref_v = torch.cat([self.ref_v, orig_v], dim=2)

        if self.ref_k.size(2) > self.max_refs:
            self.ref_k = self.ref_k[:, :, -self.max_refs :, :].contiguous()
            self.ref_v = self.ref_v[:, :, -self.max_refs :, :].contiguous()

    @torch.no_grad()
    def compress(
        self,
        k: torch.Tensor,
        v: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        orig_shape = k.shape
        if self.ref_k is None or self.ref_k.size(2) == 0:
            self.update(k, v)
            return k, v

        k_4d = k
        v_4d = v
        if k.ndim == 2:
            k_4d = k.unsqueeze(0).unsqueeze(0)
            v_4d = v.unsqueeze(0).unsqueeze(0)
        elif k.ndim == 3:
            k_4d = k.unsqueeze(1)
            v_4d = v.unsqueeze(1)

        B, H, N, D = k_4d.shape
        ref_b, ref_h, ref_n, ref_d = self.ref_k.shape

        ref_k_cur = self.ref_k.to(device=k.device, dtype=k.dtype)
        ref_v_cur = self.ref_v.to(device=v.device, dtype=v.dtype)

        ref_k_b = ref_k_cur if ref_b == B else ref_k_cur[:1].expand(B, -1, -1, -1)
        ref_v_b = ref_v_cur if ref_b == B else ref_v_cur[:1].expand(B, -1, -1, -1)

        tk = min(self.topk, ref_n)
        if tk == 0:
            self.update(k, v)
            return k, v

        sim = torch.matmul(k_4d.float(), ref_k_b.float().transpose(-2, -1).contiguous()) / (D**0.5)
        _, top_idx = torch.topk(sim, tk, dim=-1)

        top_idx_exp = top_idx.unsqueeze(-1).expand(B, H, N, tk, D)
        ref_k_exp = ref_k_b.unsqueeze(2).expand(B, H, N, ref_n, D)
        ref_v_exp = ref_v_b.unsqueeze(2).expand(B, H, N, ref_n, D)

        rk = ref_k_exp.gather(3, top_idx_exp).mean(dim=3).to(dtype=k_4d.dtype)
        rv = ref_v_exp.gather(3, top_idx_exp).mean(dim=3).to(dtype=v_4d.dtype)

        k_residual = k_4d - rk
        v_residual = v_4d - rv

        k_norm = torch.norm(k_residual.float(), dim=-1, keepdim=True)
        v_norm = torch.norm(v_residual.float(), dim=-1, keepdim=True)

        k_mask = (k_norm > self.residual_threshold).to(dtype=k_4d.dtype)
        v_mask = (v_norm > self.residual_threshold).to(dtype=v_4d.dtype)

        k_hat = rk + k_residual * k_mask
        v_hat = rv + v_residual * v_mask

        self.update(k, v)
        return k_hat.reshape(orig_shape), v_hat.reshape(orig_shape)

Optional runtime cache that stores delta residuals relative to retrieved historical K/V references. Not used by the weight-only compressor path.

Methods

def compress(self, k: torch.Tensor, v: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor]

Expand source code

@torch.no_grad()
def compress(
    self,
    k: torch.Tensor,
    v: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    orig_shape = k.shape
    if self.ref_k is None or self.ref_k.size(2) == 0:
        self.update(k, v)
        return k, v

    k_4d = k
    v_4d = v
    if k.ndim == 2:
        k_4d = k.unsqueeze(0).unsqueeze(0)
        v_4d = v.unsqueeze(0).unsqueeze(0)
    elif k.ndim == 3:
        k_4d = k.unsqueeze(1)
        v_4d = v.unsqueeze(1)

    B, H, N, D = k_4d.shape
    ref_b, ref_h, ref_n, ref_d = self.ref_k.shape

    ref_k_cur = self.ref_k.to(device=k.device, dtype=k.dtype)
    ref_v_cur = self.ref_v.to(device=v.device, dtype=v.dtype)

    ref_k_b = ref_k_cur if ref_b == B else ref_k_cur[:1].expand(B, -1, -1, -1)
    ref_v_b = ref_v_cur if ref_b == B else ref_v_cur[:1].expand(B, -1, -1, -1)

    tk = min(self.topk, ref_n)
    if tk == 0:
        self.update(k, v)
        return k, v

    sim = torch.matmul(k_4d.float(), ref_k_b.float().transpose(-2, -1).contiguous()) / (D**0.5)
    _, top_idx = torch.topk(sim, tk, dim=-1)

    top_idx_exp = top_idx.unsqueeze(-1).expand(B, H, N, tk, D)
    ref_k_exp = ref_k_b.unsqueeze(2).expand(B, H, N, ref_n, D)
    ref_v_exp = ref_v_b.unsqueeze(2).expand(B, H, N, ref_n, D)

    rk = ref_k_exp.gather(3, top_idx_exp).mean(dim=3).to(dtype=k_4d.dtype)
    rv = ref_v_exp.gather(3, top_idx_exp).mean(dim=3).to(dtype=v_4d.dtype)

    k_residual = k_4d - rk
    v_residual = v_4d - rv

    k_norm = torch.norm(k_residual.float(), dim=-1, keepdim=True)
    v_norm = torch.norm(v_residual.float(), dim=-1, keepdim=True)

    k_mask = (k_norm > self.residual_threshold).to(dtype=k_4d.dtype)
    v_mask = (v_norm > self.residual_threshold).to(dtype=v_4d.dtype)

    k_hat = rk + k_residual * k_mask
    v_hat = rv + v_residual * v_mask

    self.update(k, v)
    return k_hat.reshape(orig_shape), v_hat.reshape(orig_shape)

def reset(self) ‑> None

Expand source code

def reset(self) -> None:
    self.ref_k = None
    self.ref_v = None

def update(self, k: torch.Tensor, v: torch.Tensor) ‑> None

Expand source code

@torch.no_grad()
def update(self, k: torch.Tensor, v: torch.Tensor) -> None:
    orig_k = k.detach()
    orig_v = v.detach()
    if orig_k.ndim == 2:
        orig_k = orig_k.unsqueeze(0).unsqueeze(0)
        orig_v = orig_v.unsqueeze(0).unsqueeze(0)
    elif orig_k.ndim == 3:
        orig_k = orig_k.unsqueeze(1)
        orig_v = orig_v.unsqueeze(1)

    B, H, N, D = orig_k.shape
    if N == 0 or B == 0:
        return

    if self.ref_k is None:
        self.ref_k = orig_k
        self.ref_v = orig_v
    else:
        if self.ref_k.device != orig_k.device or self.ref_k.dtype != orig_k.dtype:
            self.ref_k = self.ref_k.to(device=orig_k.device, dtype=orig_k.dtype)
            self.ref_v = self.ref_v.to(device=orig_v.device, dtype=orig_v.dtype)

        if self.ref_k.shape[0] != B:
            self.ref_k = self.ref_k[:1].expand(B, -1, -1, -1).clone()
            self.ref_v = self.ref_v[:1].expand(B, -1, -1, -1).clone()

        self.ref_k = torch.cat([self.ref_k, orig_k], dim=2)
        self.ref_v = torch.cat([self.ref_v, orig_v], dim=2)

    if self.ref_k.size(2) > self.max_refs:
        self.ref_k = self.ref_k[:, :, -self.max_refs :, :].contiguous()
        self.ref_v = self.ref_v[:, :, -self.max_refs :, :].contiguous()

class DirectSVDCompressor (model: nn.Module, rank: int | None = 32, min_rank: int = 16, rank_ratio: float | None = None, max_error: float = 0.35, fallback_to_original: bool = True, protect_critical: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, skip_attention_qkv: bool = True, target_attention_only: bool = False, max_retries: int = 2, target_modules: list[str] | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("direct_svd")
@COMPRESSOR_REGISTRY.register("direct_svd_compressor")
class DirectSVDCompressor(BaseCompressor):
    is_custom_attention: bool = True
    """
    Direct SVD Compressor.

    Performs direct truncated SVD on target linear and Conv1D layers:
        W ≈ U_r * Sigma_r * V_r^T
    Allocates singular energy symmetrically:
        U_factor = U_r * sqrt(Sigma_r)
        V_factor = sqrt(Sigma_r) * V_r^T
    and replaces dense layers with factorized SVDLinear modules.

    Production hardening features:
    - Relative Frobenius reconstruction error validation with configurable max_error.
    - Automatic rank escalation retries when error threshold is exceeded.
    - Safe fallback to original layer if error remains high or decomposition fails.
    - Fallback and non-target layers frozen (requires_grad=False) during recovery.
    - Robust LoRALinear and Conv1D support (with transpose and adapter preservation).
    - target_attention_only parameter: restricts SVD factorization to attention projections.
    - Strict rank safeguard: enforces r < min(in_f, out_f) and parameter reduction r * (in + out) < in * out.
    - Unified critical-layer protection via should_compress_linear.
    - Full FP32 computation for numerical stability.
    - Real metrics recording: original_layer_count, compressed_layer_count, fallback_layer_count, effective_parameter_ratio.
    """

    def __init__(
        self,
        model: nn.Module,
        rank: int | None = 32,
        min_rank: int = 16,
        rank_ratio: float | None = None,
        max_error: float = 0.35,
        fallback_to_original: bool = True,
        protect_critical: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        skip_attention_qkv: bool = True,
        target_attention_only: bool = False,
        max_retries: int = 2,
        target_modules: list[str] | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="direct_svd", **kwargs)
        cfg_obj = cfg if cfg is not None else kwargs.get("comp_config")
        self.cfg = cfg_obj

        rank = _cfg_get(
            cfg_obj,
            "direct_svd_rank",
            _cfg_get(cfg_obj, "svd_rank", _cfg_get(cfg_obj, "rank", rank)),
        )
        min_rank = _cfg_get(cfg_obj, "min_rank", min_rank)
        rank_ratio = _cfg_get(
            cfg_obj,
            "direct_svd_rank_ratio",
            _cfg_get(cfg_obj, "rank_ratio", rank_ratio),
        )
        max_error = _cfg_get(
            cfg_obj, "direct_svd_max_error", _cfg_get(cfg_obj, "max_error", max_error)
        )
        fallback_to_original = _cfg_get(
            cfg_obj,
            "direct_svd_fallback_to_original",
            _cfg_get(cfg_obj, "fallback_to_original", fallback_to_original),
        )
        protect_critical = _cfg_get(cfg_obj, "protect_critical", protect_critical)
        lora_r = _cfg_get(cfg_obj, "lora_r", _cfg_get(cfg_obj, "lora_rank", lora_r))
        lora_alpha = _cfg_get(cfg_obj, "lora_alpha", lora_alpha)
        skip_attention_qkv = _cfg_get(cfg_obj, "skip_attention_qkv", False)
        target_attention_only = _cfg_get(
            cfg_obj,
            "direct_svd_target_attention_only",
            _cfg_get(cfg_obj, "target_attention_only", target_attention_only),
        )

        if "rank" in kwargs and kwargs["rank"] is not None:
            rank = kwargs["rank"]

        self.rank = rank
        self.min_rank = max(1, int(min_rank))
        self.rank_ratio = float(rank_ratio) if rank_ratio is not None else None
        if max_error is None:
            max_error = 0.05

        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.protect_critical = bool(protect_critical)
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.skip_attention_qkv = bool(skip_attention_qkv)
        self.target_attention_only = bool(target_attention_only)
        self.max_retries = int(max_retries)
        self.target_modules = target_modules

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_ranks: dict[str, int] = {}
        self.layer_errors: dict[str, float] = {}

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            name_lower = name.lower()
            if any(
                p in name_lower
                for p in [
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                ]
            ):
                return True
        return False

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
            return False
        if self.protect_critical and any(
            k in name_lower
            for k in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embed",
                "logits_processor",
            ]
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d):
            return False
        if isinstance(target_mod, SVDLinear) or isinstance(module, SVDLinear):
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self._is_gqa_asymmetric(name, target_mod):
            return False

        if self.target_attention_only:
            last_part = name.split(".")[-1]
            is_attn = (last_part in _ATTN_NAMES) or any(
                k in name_lower for k in ("attn", "attention")
            )
            if not is_attn:
                return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def _determine_rank(self, in_f: int, out_f: int, target_r: int | None) -> int:
        """Compute target rank guaranteeing parameter reduction and strict bounds."""
        max_reduction_r = max(1, (in_f * out_f) // (in_f + out_f) - 1)
        max_r = max(1, min(in_f - 1, out_f - 1, max_reduction_r))

        ref_dim = 768.0
        if (
            hasattr(self, "model")
            and hasattr(self.model, "config")
            and self.model.config is not None
        ):
            cfg_obj = self.model.config
            ref_dim = float(
                getattr(
                    cfg_obj,
                    "hidden_size",
                    getattr(cfg_obj, "d_model", getattr(cfg_obj, "n_embd", 768.0)),
                )
            )
        ref_dim = max(1.0, ref_dim)

        if self.rank_ratio is not None:
            r = int(round(min(in_f, out_f) * max(0.01, min(0.99, self.rank_ratio))))
        elif target_r is not None:
            r = int(round(target_r * (min(in_f, out_f) / ref_dim)))
        else:
            r = int(round(32 * (min(in_f, out_f) / ref_dim)))

        # Dynamic architecture-aware rank selection for large models
        r = get_dynamic_compression_rank(
            in_f,
            out_f,
            model_or_family=self.model,
            default_rank=r,
            comp_config=getattr(self, "cfg", None),
        )
        r = max(self.min_rank, r)
        return max(1, min(r, max_r))

    def compress_linear(
        self,
        layer: nn.Module,
        rank: int | None = None,
    ) -> tuple[SVDLinear | None, float, int]:
        """
        Decompose a single linear or Conv1D layer via SVD with symmetric factor distribution.
        Returns: (svd_layer, relative_reconstruction_error, chosen_rank)
        """
        target_mod = getattr(layer, "base_layer", layer)
        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.data

        if is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()  # [out_features, in_features]
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float()  # [out_features, in_features]

        max_reduction_r = max(1, (in_f * out_f) // (in_f + out_f) - 1)
        max_r = max(1, min(in_f - 1, out_f - 1, max_reduction_r))

        current_r = self._determine_rank(in_f, out_f, rank)
        device = raw_w.device
        dtype = raw_w.dtype
        norm_w = torch.norm(W, p="fro").item()

        try:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        except Exception as e:
            logger.warning("DirectSVD SVD computation failed: %s", e)
            return None, 1.0, 0

        best_svd_layer: SVDLinear | None = None
        best_err = 1.0
        chosen_r = current_r

        # Retry with higher rank if error exceeds threshold
        for attempt in range(self.max_retries + 1):
            r = max(1, min(current_r, max_r))
            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
            U_r = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
            V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(dtype=dtype, device=device)

            W_recon = torch.matmul(U_r.float(), V_r.float())
            if norm_w < 1e-12:
                err = 0.0
            else:
                err = float(torch.norm(W - W_recon, p="fro").item() / norm_w)

            best_err = err
            chosen_r = r

            # Determine LoRA parameters to preserve or initialize
            eff_lora_r = self.lora_r
            eff_lora_alpha = self.lora_alpha
            if hasattr(layer, "lora_A") and hasattr(layer, "r"):
                eff_lora_r = getattr(layer, "r", self.lora_r)
                eff_lora_alpha = getattr(layer, "lora_alpha", self.lora_alpha)

            svd_layer = SVDLinear(
                in_features=in_f,
                out_features=out_f,
                rank=r,
                bias=target_mod.bias is not None,
                dtype=dtype,
                device=device,
                lora_r=eff_lora_r,
                lora_alpha=eff_lora_alpha,
            )
            svd_layer.set_low_rank_factors(U_r, V_r)
            if target_mod.bias is not None and svd_layer.bias is not None:
                svd_layer.bias.data.copy_(target_mod.bias.data)

            # If the original layer was a LoRALinear wrapper, copy over existing adapter weights
            if (
                hasattr(layer, "lora_A")
                and hasattr(layer, "lora_B")
                and svd_layer.lora_A is not None
            ):
                try:
                    if (
                        hasattr(layer.lora_A, "weight")
                        and svd_layer.lora_A.weight.shape == layer.lora_A.weight.shape
                    ):
                        svd_layer.lora_A.weight.data.copy_(layer.lora_A.weight.data)
                    if (
                        hasattr(layer.lora_B, "weight")
                        and svd_layer.lora_B.weight.shape == layer.lora_B.weight.shape
                    ):
                        svd_layer.lora_B.weight.data.copy_(layer.lora_B.weight.data)
                except Exception as e:
                    logger.debug("DirectSVD adapter preservation notice: %s", e)

            best_svd_layer = svd_layer

            if err <= self.max_error or r >= max_r:
                break

            if attempt < self.max_retries:
                next_r = min(int(r * 1.5), max_r)
                if next_r <= r:
                    break
                current_r = next_r

        return best_svd_layer, best_err, chosen_r

    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """Compress target linear and Conv1D layers across the model."""
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_ranks = {}
        self.layer_errors = {}
        seen_weights: set[int] = set()

        total_orig_decomp_params = 0
        total_comp_decomp_params = 0

        target_layers = [
            (name, mod)
            for name, mod in list(self.model.named_modules())
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            logger.warning("DirectSVD: No compressible linear layers found.")
            return self.model, self.get_compression_stats()

        # Handle freezing of non-target layers when target_attention_only=True
        if self.target_attention_only:
            for name, module in list(self.model.named_modules()):
                target_mod = getattr(module, "base_layer", module)
                cls_name = target_mod.__class__.__name__.lower()
                is_linear = isinstance(target_mod, nn.Linear)
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                if is_linear or is_conv1d:
                    if not self._is_target_module(name, module):
                        for p in module.parameters():
                            p.requires_grad = False
                        self.skipped_layers.append(name)

        for name, module in target_layers:
            target_mod = getattr(module, "base_layer", module)
            if not hasattr(target_mod, "weight") or target_mod.weight is None:
                continue
            w_id = id(target_mod.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            raw_w = target_mod.weight.data
            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            if is_conv1d:
                in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            else:
                out_f, in_f = raw_w.shape[0], raw_w.shape[1]

            orig_params = out_f * in_f
            total_orig_decomp_params += orig_params

            svd_layer, err, r = self.compress_linear(module, rank=self.rank)
            self.layer_errors[name] = err
            self.layer_ranks[name] = r

            if svd_layer is None or err > self.max_error:
                logger.warning(
                    "DirectSVD reconstruction error %.4f on %s exceeded threshold %.4f.",
                    err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    logger.info(
                        "Skipping DirectSVD compression for %s (fallback to original and frozen).",
                        name,
                    )
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    total_comp_decomp_params += orig_params
                    continue

            replace_module(self.model, name, svd_layer)
            self.compressed_layers.append(name)
            total_comp_decomp_params += r * (in_f + out_f)

        # Post-compression forward sanity check
        forward_check_success = False
        try:
            self.model.eval()
            with torch.no_grad():
                d = next(self.model.parameters()).device
                dummy = torch.randint(0, 100, (1, 16), device=d)
                self.model(dummy)
            forward_check_success = True
        except Exception as e:
            logger.debug("DirectSVD post-compression forward sanity check exception: %s", e)

        effective_ratio = (
            total_comp_decomp_params / total_orig_decomp_params
            if total_orig_decomp_params > 0
            else 1.0
        )
        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "original_layer_count": len(target_layers),
                "layer_ranks": self.layer_ranks,
                "layer_errors": self.layer_errors,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "target_attention_only": self.target_attention_only,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "avg_reconstruction_error": avg_err,
                "effective_parameter_ratio": round(effective_ratio, 4),
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "forward_check_success": forward_check_success,
                "method": "direct_truncated_svd",
            }
        )
        logger.info(
            "Compressor direct_svd finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        self.stats = stats
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: Direct SVD Compressor.

Performs direct truncated SVD on target linear and Conv1D layers: W ≈ U_r * Sigma_r * V_r^T Allocates singular energy symmetrically: U_factor = U_r * sqrt(Sigma_r) V_factor = sqrt(Sigma_r) * V_r^T and replaces dense layers with factorized SVDLinear modules.

Production hardening features: - Relative Frobenius reconstruction error validation with configurable max_error. - Automatic rank escalation retries when error threshold is exceeded. - Safe fallback to original layer if error remains high or decomposition fails. - Fallback and non-target layers frozen (requires_grad=False) during recovery. - Robust LoRALinear and Conv1D support (with transpose and adapter preservation). - target_attention_only parameter: restricts SVD factorization to attention projections. - Strict rank safeguard: enforces r < min(in_f, out_f) and parameter reduction r * (in + out) < in * out. - Unified critical-layer protection via should_compress_linear. - Full FP32 computation for numerical stability. - Real metrics recording: original_layer_count, compressed_layer_count, fallback_layer_count, effective_parameter_ratio.

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """Compress target linear and Conv1D layers across the model."""
    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    self.layer_ranks = {}
    self.layer_errors = {}
    seen_weights: set[int] = set()

    total_orig_decomp_params = 0
    total_comp_decomp_params = 0

    target_layers = [
        (name, mod)
        for name, mod in list(self.model.named_modules())
        if self._is_target_module(name, mod)
    ]

    if not target_layers:
        logger.warning("DirectSVD: No compressible linear layers found.")
        return self.model, self.get_compression_stats()

    # Handle freezing of non-target layers when target_attention_only=True
    if self.target_attention_only:
        for name, module in list(self.model.named_modules()):
            target_mod = getattr(module, "base_layer", module)
            cls_name = target_mod.__class__.__name__.lower()
            is_linear = isinstance(target_mod, nn.Linear)
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            if is_linear or is_conv1d:
                if not self._is_target_module(name, module):
                    for p in module.parameters():
                        p.requires_grad = False
                    self.skipped_layers.append(name)

    for name, module in target_layers:
        target_mod = getattr(module, "base_layer", module)
        if not hasattr(target_mod, "weight") or target_mod.weight is None:
            continue
        w_id = id(target_mod.weight)
        if w_id in seen_weights:
            self.skipped_layers.append(name)
            continue
        seen_weights.add(w_id)

        raw_w = target_mod.weight.data
        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]

        orig_params = out_f * in_f
        total_orig_decomp_params += orig_params

        svd_layer, err, r = self.compress_linear(module, rank=self.rank)
        self.layer_errors[name] = err
        self.layer_ranks[name] = r

        if svd_layer is None or err > self.max_error:
            logger.warning(
                "DirectSVD reconstruction error %.4f on %s exceeded threshold %.4f.",
                err,
                name,
                self.max_error,
            )
            if self.fallback_to_original:
                logger.info(
                    "Skipping DirectSVD compression for %s (fallback to original and frozen).",
                    name,
                )
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_decomp_params += orig_params
                continue

        replace_module(self.model, name, svd_layer)
        self.compressed_layers.append(name)
        total_comp_decomp_params += r * (in_f + out_f)

    # Post-compression forward sanity check
    forward_check_success = False
    try:
        self.model.eval()
        with torch.no_grad():
            d = next(self.model.parameters()).device
            dummy = torch.randint(0, 100, (1, 16), device=d)
            self.model(dummy)
        forward_check_success = True
    except Exception as e:
        logger.debug("DirectSVD post-compression forward sanity check exception: %s", e)

    effective_ratio = (
        total_comp_decomp_params / total_orig_decomp_params
        if total_orig_decomp_params > 0
        else 1.0
    )
    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )
    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "skipped_layers": self.skipped_layers,
            "skipped_layer_count": len(self.skipped_layers),
            "original_layer_count": len(target_layers),
            "layer_ranks": self.layer_ranks,
            "layer_errors": self.layer_errors,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "target_attention_only": self.target_attention_only,
            "lora_r": self.lora_r,
            "lora_alpha": self.lora_alpha,
            "avg_reconstruction_error": avg_err,
            "effective_parameter_ratio": round(effective_ratio, 4),
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "num_skipped": len(self.skipped_layers),
            "forward_check_success": forward_check_success,
            "method": "direct_truncated_svd",
        }
    )
    logger.info(
        "Compressor direct_svd finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        self.max_error,
    )
    self.stats = stats
    return self.model, stats

Compress target linear and Conv1D layers across the model.

def compress_linear(self, layer: nn.Module, rank: int | None = None) ‑> tuple[SVDLinear | None, float, int]

Expand source code

def compress_linear(
    self,
    layer: nn.Module,
    rank: int | None = None,
) -> tuple[SVDLinear | None, float, int]:
    """
    Decompose a single linear or Conv1D layer via SVD with symmetric factor distribution.
    Returns: (svd_layer, relative_reconstruction_error, chosen_rank)
    """
    target_mod = getattr(layer, "base_layer", layer)
    cls_name = target_mod.__class__.__name__.lower()
    is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
    raw_w = target_mod.weight.data

    if is_conv1d:
        in_f, out_f = raw_w.shape[0], raw_w.shape[1]
        W = raw_w.t().contiguous().float()  # [out_features, in_features]
    else:
        out_f, in_f = raw_w.shape[0], raw_w.shape[1]
        W = raw_w.float()  # [out_features, in_features]

    max_reduction_r = max(1, (in_f * out_f) // (in_f + out_f) - 1)
    max_r = max(1, min(in_f - 1, out_f - 1, max_reduction_r))

    current_r = self._determine_rank(in_f, out_f, rank)
    device = raw_w.device
    dtype = raw_w.dtype
    norm_w = torch.norm(W, p="fro").item()

    try:
        U, S, Vt = torch.linalg.svd(W, full_matrices=False)
    except Exception as e:
        logger.warning("DirectSVD SVD computation failed: %s", e)
        return None, 1.0, 0

    best_svd_layer: SVDLinear | None = None
    best_err = 1.0
    chosen_r = current_r

    # Retry with higher rank if error exceeds threshold
    for attempt in range(self.max_retries + 1):
        r = max(1, min(current_r, max_r))
        sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
        U_r = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
        V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(dtype=dtype, device=device)

        W_recon = torch.matmul(U_r.float(), V_r.float())
        if norm_w < 1e-12:
            err = 0.0
        else:
            err = float(torch.norm(W - W_recon, p="fro").item() / norm_w)

        best_err = err
        chosen_r = r

        # Determine LoRA parameters to preserve or initialize
        eff_lora_r = self.lora_r
        eff_lora_alpha = self.lora_alpha
        if hasattr(layer, "lora_A") and hasattr(layer, "r"):
            eff_lora_r = getattr(layer, "r", self.lora_r)
            eff_lora_alpha = getattr(layer, "lora_alpha", self.lora_alpha)

        svd_layer = SVDLinear(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=target_mod.bias is not None,
            dtype=dtype,
            device=device,
            lora_r=eff_lora_r,
            lora_alpha=eff_lora_alpha,
        )
        svd_layer.set_low_rank_factors(U_r, V_r)
        if target_mod.bias is not None and svd_layer.bias is not None:
            svd_layer.bias.data.copy_(target_mod.bias.data)

        # If the original layer was a LoRALinear wrapper, copy over existing adapter weights
        if (
            hasattr(layer, "lora_A")
            and hasattr(layer, "lora_B")
            and svd_layer.lora_A is not None
        ):
            try:
                if (
                    hasattr(layer.lora_A, "weight")
                    and svd_layer.lora_A.weight.shape == layer.lora_A.weight.shape
                ):
                    svd_layer.lora_A.weight.data.copy_(layer.lora_A.weight.data)
                if (
                    hasattr(layer.lora_B, "weight")
                    and svd_layer.lora_B.weight.shape == layer.lora_B.weight.shape
                ):
                    svd_layer.lora_B.weight.data.copy_(layer.lora_B.weight.data)
            except Exception as e:
                logger.debug("DirectSVD adapter preservation notice: %s", e)

        best_svd_layer = svd_layer

        if err <= self.max_error or r >= max_r:
            break

        if attempt < self.max_retries:
            next_r = min(int(r * 1.5), max_r)
            if next_r <= r:
                break
            current_r = next_r

    return best_svd_layer, best_err, chosen_r

Decompose a single linear or Conv1D layer via SVD with symmetric factor distribution. Returns: (svd_layer, relative_reconstruction_error, chosen_rank)

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class DirectSVDEncoderLayer (config: Any | None = None, hidden_size: int | None = None, intermediate_size: int | None = None, num_attention_heads: int | None = None, svd_rank: int = 20, layer_norm_eps: float = 1e-12, hidden_dropout_prob: float = 0.1, lora_r: int = 0, lora_alpha: float = 1.0)

Expand source code

class DirectSVDEncoderLayer(nn.Module):
    """Encoder layer that replaces MHA with SVDSingleHead + standard FFN."""

    def __init__(
        self,
        config: Any | None = None,
        hidden_size: int | None = None,
        intermediate_size: int | None = None,
        num_attention_heads: int | None = None,
        svd_rank: int = 20,
        layer_norm_eps: float = 1e-12,
        hidden_dropout_prob: float = 0.1,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> None:
        super().__init__()
        self.config = config
        self.hidden_size = int(hidden_size or getattr(config, "hidden_size", 768))
        self.intermediate_size = int(
            intermediate_size
            or getattr(config, "intermediate_size", None)
            or (4 * self.hidden_size)
        )
        self.num_attention_heads = int(
            num_attention_heads or getattr(config, "num_attention_heads", 12)
        )
        self.svd_rank = int(getattr(config, "svd_rank", None) or svd_rank)
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)

        self.svd_attention = SVDSingleHead(self.hidden_size, svd_rank=self.svd_rank)
        self.dense_attn = nn.Linear(self.hidden_size, self.hidden_size)
        self.LayerNorm_attn = nn.LayerNorm(self.hidden_size, eps=layer_norm_eps)
        self.dropout_attn = nn.Dropout(hidden_dropout_prob)

        self.dense_intermediate = nn.Linear(self.hidden_size, self.intermediate_size)
        self.intermediate_act = nn.GELU()
        self.dense_output = nn.Linear(self.intermediate_size, self.hidden_size)
        self.LayerNorm_output = nn.LayerNorm(self.hidden_size, eps=layer_norm_eps)
        self.dropout_output = nn.Dropout(hidden_dropout_prob)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        output_attentions: bool = False,
        *args: Any,
        **kwargs: Any,
    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        attn_out = self.svd_attention(hidden_states)
        attn_proj = self.dropout_attn(self.dense_attn(attn_out))
        h = self.LayerNorm_attn(hidden_states + attn_proj)

        inter = self.intermediate_act(self.dense_intermediate(h))
        inter = self.dropout_output(self.dense_output(inter))
        output = self.LayerNorm_output(h + inter)

        if output_attentions:
            return (output, None)
        return output

Encoder layer that replaces MHA with SVDSingleHead + standard FFN.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Methods

def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, head_mask: torch.Tensor | None = None, output_attentions: bool = False, *args: Any, **kwargs: Any) ‑> torch.Tensor | tuple[torch.Tensor, ...]

Expand source code

def forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    head_mask: torch.Tensor | None = None,
    output_attentions: bool = False,
    *args: Any,
    **kwargs: Any,
) -> torch.Tensor | tuple[torch.Tensor, ...]:
    attn_out = self.svd_attention(hidden_states)
    attn_proj = self.dropout_attn(self.dense_attn(attn_out))
    h = self.LayerNorm_attn(hidden_states + attn_proj)

    inter = self.intermediate_act(self.dense_intermediate(h))
    inter = self.dropout_output(self.dense_output(inter))
    output = self.LayerNorm_output(h + inter)

    if output_attentions:
        return (output, None)
    return output

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class DobiSVDCompressor (model: nn.Module, init_rank_ratio: float = 0.5, rank: int | None = None, init_gate_val: float = 2.5, max_error: float = 0.35, fallback_to_original: bool = True, target_attention_only: bool = True, max_retries: int = 2, target_modules: list[str] | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("dobi_class")
@COMPRESSOR_REGISTRY.register("dobi_svd_class")
@COMPRESSOR_REGISTRY.register("dobi_compressor_class")
class DobiSVDCompressor(BaseCompressor):
    """
    Dobi-SVD Compressor: Differentiable SVD with learnable spectral gating.

    Extracts truncated SVD factors from pretrained weights with symmetric
    square-root singular value distribution (U_r * sqrt(Sigma_r), sqrt(Sigma_r) * V_r^T),
    compensates for initial gate value to ensure near-lossless initialization, and
    replaces layers with DobiSVDLinearLayer containing learnable gating parameters.

    Production hardening features:
    - Relative Frobenius reconstruction error validation with configurable max_error.
    - Automatic rank escalation retries when error exceeds threshold.
    - Safe fallback to original module if error remains high.
    - Fallback and non-target layers frozen (requires_grad=False) during recovery.
    - Robust LoRALinear and Conv1D support (with transpose and adapter preservation).
    - target_attention_only parameter: restricts SVD factorization to attention projections.
    - Strict rank bounds guaranteeing parameter reduction.
    - Critical-layer protection via should_compress_linear.
    - Detailed per-layer stats, avg error, and effective parameter ratio.
    """

    def __init__(
        self,
        model: nn.Module,
        init_rank_ratio: float = 0.5,
        rank: int | None = None,
        init_gate_val: float = 2.5,
        max_error: float = 0.35,
        fallback_to_original: bool = True,
        target_attention_only: bool = True,
        max_retries: int = 2,
        target_modules: list[str] | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="dobi", **kwargs)
        cfg_obj = cfg if cfg is not None else kwargs.get("comp_config")
        self.cfg = cfg_obj

        init_rank_ratio = _cfg_get(
            cfg_obj,
            "dobi_init_rank_ratio",
            _cfg_get(cfg_obj, "init_rank_ratio", init_rank_ratio),
        )
        rank = _cfg_get(cfg_obj, "dobi_rank", _cfg_get(cfg_obj, "rank", rank))
        init_gate_val = _cfg_get(cfg_obj, "init_gate_val", init_gate_val)
        max_error = _cfg_get(cfg_obj, "dobi_max_error", _cfg_get(cfg_obj, "max_error", max_error))
        fallback_to_original = _cfg_get(
            cfg_obj,
            "dobi_fallback_to_original",
            _cfg_get(cfg_obj, "fallback_to_original", fallback_to_original),
        )
        target_attention_only = _cfg_get(
            cfg_obj,
            "dobi_target_attention_only",
            _cfg_get(cfg_obj, "target_attention_only", target_attention_only),
        )

        if "rank" in kwargs and kwargs["rank"] is not None:
            rank = kwargs["rank"]

        self.init_rank_ratio = float(init_rank_ratio)
        self.rank = rank
        self.init_gate_val = float(init_gate_val)
        if max_error is None or max_error == 0.35:
            max_error = adaptive_max_error(self.rank, 768, 768, "distilbert")

        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.target_attention_only = bool(target_attention_only)
        self.max_retries = int(max_retries)
        self.target_modules = target_modules

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_ranks: dict[str, int] = {}
        self.layer_errors: dict[str, float] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
            return False
        if any(
            k in name_lower
            for k in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            ]
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d):
            return False
        if isinstance(target_mod, DobiSVDLinearLayer) or isinstance(module, DobiSVDLinearLayer):
            return False
        if not should_compress_linear(self.model, name):
            return False

        if self.target_attention_only:
            last_part = name.split(".")[-1]
            is_attn = (last_part in _ATTN_NAMES) or any(
                k in name_lower for k in ("attn", "attention")
            )
            if not is_attn:
                return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def _determine_rank(self, in_f: int, out_f: int, target_r: int | None) -> int:
        """Compute target rank guaranteeing parameter reduction and physical bounds."""
        max_reduction_r = max(1, (in_f * out_f) // (in_f + out_f) - 1)
        max_r = max(1, min(in_f - 1, out_f - 1, max_reduction_r, 128))

        if target_r is not None:
            r = target_r
        else:
            r = int(round(min(in_f, out_f) * self.init_rank_ratio))

        return max(1, min(r, max_r))

    def compress_linear(
        self,
        layer: nn.Module,
        rank: int | None = None,
    ) -> tuple[DobiSVDLinearLayer | None, float, int]:
        """Decompose a single linear or Conv1D layer into DobiSVDLinearLayer."""
        target_mod = getattr(layer, "base_layer", layer)
        cls_name = target_mod.__class__.__name__.lower()
        raw_w = target_mod.weight.data

        # Determine in_features and out_features accurately for Linear, Conv1D, and fused QKV
        in_f = None
        out_f = None
        if (
            target_mod.bias is not None
            and getattr(target_mod.bias, "ndim", 0) == 1
            and raw_w.ndim == 2
        ):
            bias_len = target_mod.bias.shape[0]
            if raw_w.shape[1] == bias_len and raw_w.shape[0] != bias_len:
                # Conv1D format: [in_features, out_features]
                in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            elif raw_w.shape[0] == bias_len and raw_w.shape[1] != bias_len:
                # Linear format: [out_features, in_features]
                out_f, in_f = raw_w.shape[0], raw_w.shape[1]

        if in_f is None or out_f is None:
            if hasattr(target_mod, "in_features") and hasattr(target_mod, "out_features"):
                in_f, out_f = target_mod.in_features, target_mod.out_features
            elif hasattr(target_mod, "nf") and hasattr(target_mod, "nx"):
                in_f, out_f = target_mod.nx, target_mod.nf
            elif "conv1d" in cls_name or (raw_w.ndim == 2 and raw_w.shape[0] < raw_w.shape[1]):
                in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            else:
                out_f, in_f = raw_w.shape[0], raw_w.shape[1]

        # W must strictly be [out_features, in_features]
        if raw_w.shape == (out_f, in_f):
            W = raw_w.float()
        elif raw_w.shape == (in_f, out_f):
            W = raw_w.t().contiguous().float()
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float()

        assert W.shape == (
            out_f,
            in_f,
        ), f"W shape {W.shape} does not match ({out_f}, {in_f})"

        max_reduction_r = max(1, (in_f * out_f) // (in_f + out_f) - 1)
        max_r = max(1, min(in_f - 1, out_f - 1, max_reduction_r, 128))

        current_r = self._determine_rank(in_f, out_f, rank)
        device = raw_w.device
        dtype = raw_w.dtype
        norm_w = torch.norm(W, p="fro").item()
        if norm_w < 1e-4:
            logger.warning(
                "DobiSVD: weight tensor has near-zero Frobenius norm (%.6e < 1e-4)",
                norm_w,
            )
            return None, 1.0, 0

        try:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        except Exception as e:
            logger.warning("DobiSVD SVD computation failed: %s", e)
            return None, 1.0, 0

        best_layer: DobiSVDLinearLayer | None = None
        best_err = 1.0
        chosen_r = current_r

        # Initial gate scaling factor to achieve exact initial SVD energy
        init_gate_scale = 1.0 / (1.0 / (1.0 + math.exp(-self.init_gate_val)))

        retries = 0 if rank is not None else self.max_retries
        for attempt in range(retries + 1):
            r = max(1, min(current_r, max_r))
            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
            U_r = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
            # Pre-scale V by 1/sigma(alpha) so initial gated output matches SVD reconstruction
            V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :] * init_gate_scale).to(dtype=dtype, device=device)

            dobi_layer = DobiSVDLinearLayer(
                in_features=in_f,
                out_features=out_f,
                rank=r,
                bias=target_mod.bias is not None,
                init_gate_val=self.init_gate_val,
                dtype=dtype,
                device=device,
            )
            dobi_layer.set_low_rank_factors(U_r, V_r)
            assert dobi_layer.U.shape == (
                out_f,
                r,
            ), f"U shape mismatch: {dobi_layer.U.shape} vs ({out_f}, {r})"
            assert dobi_layer.V.shape == (
                r,
                in_f,
            ), f"V shape mismatch: {dobi_layer.V.shape} vs ({r}, {in_f})"

            if target_mod.bias is not None and dobi_layer.bias is not None:
                tb = target_mod.bias.data
                db = dobi_layer.bias.data
                min_b = min(tb.numel(), db.numel())
                db[:min_b].copy_(tb[:min_b])

            W_recon = dobi_layer.weight.float()
            recon_norm = torch.norm(W_recon, p="fro").item()
            if recon_norm < 1e-4:
                logger.warning(
                    "DobiSVD: reconstructed weight Frobenius norm too small (%.6e < 1e-4)",
                    recon_norm,
                )
                continue

            if norm_w < 1e-12:
                err = 0.0
            else:
                err = float(torch.norm(W - W_recon, p="fro").item() / norm_w)

            best_err = err
            chosen_r = r
            best_layer = dobi_layer

            if err <= self.max_error or r >= max_r:
                break

            if attempt < retries:
                next_r = min(int(r * 1.5), max_r)
                if next_r <= r:
                    break
                current_r = next_r

        if best_layer is not None:
            best_layer.reconstruction_error = best_err
            return best_layer, best_err, chosen_r
        return None, best_err, chosen_r

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Compress all target linear and Conv1D layers across the model."""
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_ranks = {}
        self.layer_errors = {}

        total_orig_decomp_params = 0
        total_comp_decomp_params = 0

        target_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            logger.warning("Dobi-SVD: No compressible linear layers found.")
            return self.model, self.get_compression_stats()

        # If target_attention_only=True, freeze non-target linear/conv layers during recovery
        if self.target_attention_only:
            for name, module in list(self.model.named_modules()):
                target_mod = getattr(module, "base_layer", module)
                cls_name = target_mod.__class__.__name__.lower()
                is_linear = isinstance(target_mod, nn.Linear)
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                if is_linear or is_conv1d:
                    if not self._is_target_module(name, module):
                        for p in module.parameters():
                            p.requires_grad = False
                        self.skipped_layers.append(name)

        for name, module in target_layers:
            target_mod = getattr(module, "base_layer", module)
            raw_w = target_mod.weight.data
            cls_name = target_mod.__class__.__name__.lower()
            in_f = None
            out_f = None
            if (
                target_mod.bias is not None
                and getattr(target_mod.bias, "ndim", 0) == 1
                and raw_w.ndim == 2
            ):
                bias_len = target_mod.bias.shape[0]
                if raw_w.shape[1] == bias_len and raw_w.shape[0] != bias_len:
                    in_f, out_f = raw_w.shape[0], raw_w.shape[1]
                elif raw_w.shape[0] == bias_len and raw_w.shape[1] != bias_len:
                    out_f, in_f = raw_w.shape[0], raw_w.shape[1]

            if in_f is None or out_f is None:
                if hasattr(target_mod, "in_features") and hasattr(target_mod, "out_features"):
                    in_f, out_f = target_mod.in_features, target_mod.out_features
                elif hasattr(target_mod, "nf") and hasattr(target_mod, "nx"):
                    in_f, out_f = target_mod.nx, target_mod.nf
                elif "conv1d" in cls_name or (raw_w.ndim == 2 and raw_w.shape[0] < raw_w.shape[1]):
                    in_f, out_f = raw_w.shape[0], raw_w.shape[1]
                else:
                    out_f, in_f = raw_w.shape[0], raw_w.shape[1]

            orig_params = out_f * in_f
            total_orig_decomp_params += orig_params

            new_layer, err, r = self.compress_linear(module, rank=self.rank)
            self.layer_errors[name] = err
            self.layer_ranks[name] = r

            if new_layer is None or err > self.max_error:
                logger.warning(
                    "Dobi-SVD reconstruction error %.4f on %s exceeded threshold %.4f.",
                    err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    logger.info(
                        "Skipping Dobi-SVD compression for %s (fallback to original and frozen).",
                        name,
                    )
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    total_comp_decomp_params += orig_params
                    continue

            # If wrapped in LoRALinear, replace the inner base_layer or linear
            has_base_layer = hasattr(module, "base_layer") or (
                hasattr(module, "linear") and (hasattr(module, "lora_A") or hasattr(module, "r"))
            )
            if hasattr(module, "base_layer"):
                module.base_layer = new_layer
            elif hasattr(module, "linear") and (hasattr(module, "lora_A") or hasattr(module, "r")):
                module.linear = new_layer
            else:
                replace_module(self.model, name, new_layer)

            if has_base_layer:
                for p_name, p in module.named_parameters():
                    if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                        p.requires_grad = True

            self.compressed_layers.append(name)
            # Parameters: U (out*r) + V (r*in) + gate (r)
            comp_params = r * (in_f + out_f) + r
            total_comp_decomp_params += comp_params

        effective_ratio = (
            total_comp_decomp_params / total_orig_decomp_params
            if total_orig_decomp_params > 0
            else 1.0
        )
        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "original_layer_count": len(target_layers),
                "layer_ranks": self.layer_ranks,
                "layer_errors": self.layer_errors,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "target_attention_only": self.target_attention_only,
                "init_rank_ratio": self.init_rank_ratio,
                "init_gate_val": self.init_gate_val,
                "effective_parameter_ratio": round(effective_ratio, 4),
                "avg_reconstruction_error": avg_err,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
            }
        )
        logger.info(
            "Compressor dobi finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

Dobi-SVD Compressor: Differentiable SVD with learnable spectral gating.

Extracts truncated SVD factors from pretrained weights with symmetric square-root singular value distribution (U_r * sqrt(Sigma_r), sqrt(Sigma_r) * V_r^T), compensates for initial gate value to ensure near-lossless initialization, and replaces layers with DobiSVDLinearLayer containing learnable gating parameters.

Production hardening features: - Relative Frobenius reconstruction error validation with configurable max_error. - Automatic rank escalation retries when error exceeds threshold. - Safe fallback to original module if error remains high. - Fallback and non-target layers frozen (requires_grad=False) during recovery. - Robust LoRALinear and Conv1D support (with transpose and adapter preservation). - target_attention_only parameter: restricts SVD factorization to attention projections. - Strict rank bounds guaranteeing parameter reduction. - Critical-layer protection via should_compress_linear. - Detailed per-layer stats, avg error, and effective parameter ratio.

Ancestors

BaseCompressor

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Compress all target linear and Conv1D layers across the model."""
    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    self.layer_ranks = {}
    self.layer_errors = {}

    total_orig_decomp_params = 0
    total_comp_decomp_params = 0

    target_layers = [
        (name, mod)
        for name, mod in self.model.named_modules()
        if self._is_target_module(name, mod)
    ]

    if not target_layers:
        logger.warning("Dobi-SVD: No compressible linear layers found.")
        return self.model, self.get_compression_stats()

    # If target_attention_only=True, freeze non-target linear/conv layers during recovery
    if self.target_attention_only:
        for name, module in list(self.model.named_modules()):
            target_mod = getattr(module, "base_layer", module)
            cls_name = target_mod.__class__.__name__.lower()
            is_linear = isinstance(target_mod, nn.Linear)
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            if is_linear or is_conv1d:
                if not self._is_target_module(name, module):
                    for p in module.parameters():
                        p.requires_grad = False
                    self.skipped_layers.append(name)

    for name, module in target_layers:
        target_mod = getattr(module, "base_layer", module)
        raw_w = target_mod.weight.data
        cls_name = target_mod.__class__.__name__.lower()
        in_f = None
        out_f = None
        if (
            target_mod.bias is not None
            and getattr(target_mod.bias, "ndim", 0) == 1
            and raw_w.ndim == 2
        ):
            bias_len = target_mod.bias.shape[0]
            if raw_w.shape[1] == bias_len and raw_w.shape[0] != bias_len:
                in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            elif raw_w.shape[0] == bias_len and raw_w.shape[1] != bias_len:
                out_f, in_f = raw_w.shape[0], raw_w.shape[1]

        if in_f is None or out_f is None:
            if hasattr(target_mod, "in_features") and hasattr(target_mod, "out_features"):
                in_f, out_f = target_mod.in_features, target_mod.out_features
            elif hasattr(target_mod, "nf") and hasattr(target_mod, "nx"):
                in_f, out_f = target_mod.nx, target_mod.nf
            elif "conv1d" in cls_name or (raw_w.ndim == 2 and raw_w.shape[0] < raw_w.shape[1]):
                in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            else:
                out_f, in_f = raw_w.shape[0], raw_w.shape[1]

        orig_params = out_f * in_f
        total_orig_decomp_params += orig_params

        new_layer, err, r = self.compress_linear(module, rank=self.rank)
        self.layer_errors[name] = err
        self.layer_ranks[name] = r

        if new_layer is None or err > self.max_error:
            logger.warning(
                "Dobi-SVD reconstruction error %.4f on %s exceeded threshold %.4f.",
                err,
                name,
                self.max_error,
            )
            if self.fallback_to_original:
                logger.info(
                    "Skipping Dobi-SVD compression for %s (fallback to original and frozen).",
                    name,
                )
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_decomp_params += orig_params
                continue

        # If wrapped in LoRALinear, replace the inner base_layer or linear
        has_base_layer = hasattr(module, "base_layer") or (
            hasattr(module, "linear") and (hasattr(module, "lora_A") or hasattr(module, "r"))
        )
        if hasattr(module, "base_layer"):
            module.base_layer = new_layer
        elif hasattr(module, "linear") and (hasattr(module, "lora_A") or hasattr(module, "r")):
            module.linear = new_layer
        else:
            replace_module(self.model, name, new_layer)

        if has_base_layer:
            for p_name, p in module.named_parameters():
                if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                    p.requires_grad = True

        self.compressed_layers.append(name)
        # Parameters: U (out*r) + V (r*in) + gate (r)
        comp_params = r * (in_f + out_f) + r
        total_comp_decomp_params += comp_params

    effective_ratio = (
        total_comp_decomp_params / total_orig_decomp_params
        if total_orig_decomp_params > 0
        else 1.0
    )
    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )
    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "skipped_layers": self.skipped_layers,
            "skipped_layer_count": len(self.skipped_layers),
            "original_layer_count": len(target_layers),
            "layer_ranks": self.layer_ranks,
            "layer_errors": self.layer_errors,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "target_attention_only": self.target_attention_only,
            "init_rank_ratio": self.init_rank_ratio,
            "init_gate_val": self.init_gate_val,
            "effective_parameter_ratio": round(effective_ratio, 4),
            "avg_reconstruction_error": avg_err,
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
        }
    )
    logger.info(
        "Compressor dobi finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        self.max_error,
    )
    return self.model, stats

Compress all target linear and Conv1D layers across the model.

def compress_linear(self, layer: nn.Module, rank: int | None = None) ‑> tuple[DobiSVDLinearLayer | None, float, int]

Expand source code

def compress_linear(
    self,
    layer: nn.Module,
    rank: int | None = None,
) -> tuple[DobiSVDLinearLayer | None, float, int]:
    """Decompose a single linear or Conv1D layer into DobiSVDLinearLayer."""
    target_mod = getattr(layer, "base_layer", layer)
    cls_name = target_mod.__class__.__name__.lower()
    raw_w = target_mod.weight.data

    # Determine in_features and out_features accurately for Linear, Conv1D, and fused QKV
    in_f = None
    out_f = None
    if (
        target_mod.bias is not None
        and getattr(target_mod.bias, "ndim", 0) == 1
        and raw_w.ndim == 2
    ):
        bias_len = target_mod.bias.shape[0]
        if raw_w.shape[1] == bias_len and raw_w.shape[0] != bias_len:
            # Conv1D format: [in_features, out_features]
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
        elif raw_w.shape[0] == bias_len and raw_w.shape[1] != bias_len:
            # Linear format: [out_features, in_features]
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]

    if in_f is None or out_f is None:
        if hasattr(target_mod, "in_features") and hasattr(target_mod, "out_features"):
            in_f, out_f = target_mod.in_features, target_mod.out_features
        elif hasattr(target_mod, "nf") and hasattr(target_mod, "nx"):
            in_f, out_f = target_mod.nx, target_mod.nf
        elif "conv1d" in cls_name or (raw_w.ndim == 2 and raw_w.shape[0] < raw_w.shape[1]):
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]

    # W must strictly be [out_features, in_features]
    if raw_w.shape == (out_f, in_f):
        W = raw_w.float()
    elif raw_w.shape == (in_f, out_f):
        W = raw_w.t().contiguous().float()
    else:
        out_f, in_f = raw_w.shape[0], raw_w.shape[1]
        W = raw_w.float()

    assert W.shape == (
        out_f,
        in_f,
    ), f"W shape {W.shape} does not match ({out_f}, {in_f})"

    max_reduction_r = max(1, (in_f * out_f) // (in_f + out_f) - 1)
    max_r = max(1, min(in_f - 1, out_f - 1, max_reduction_r, 128))

    current_r = self._determine_rank(in_f, out_f, rank)
    device = raw_w.device
    dtype = raw_w.dtype
    norm_w = torch.norm(W, p="fro").item()
    if norm_w < 1e-4:
        logger.warning(
            "DobiSVD: weight tensor has near-zero Frobenius norm (%.6e < 1e-4)",
            norm_w,
        )
        return None, 1.0, 0

    try:
        U, S, Vt = torch.linalg.svd(W, full_matrices=False)
    except Exception as e:
        logger.warning("DobiSVD SVD computation failed: %s", e)
        return None, 1.0, 0

    best_layer: DobiSVDLinearLayer | None = None
    best_err = 1.0
    chosen_r = current_r

    # Initial gate scaling factor to achieve exact initial SVD energy
    init_gate_scale = 1.0 / (1.0 / (1.0 + math.exp(-self.init_gate_val)))

    retries = 0 if rank is not None else self.max_retries
    for attempt in range(retries + 1):
        r = max(1, min(current_r, max_r))
        sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
        U_r = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
        # Pre-scale V by 1/sigma(alpha) so initial gated output matches SVD reconstruction
        V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :] * init_gate_scale).to(dtype=dtype, device=device)

        dobi_layer = DobiSVDLinearLayer(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=target_mod.bias is not None,
            init_gate_val=self.init_gate_val,
            dtype=dtype,
            device=device,
        )
        dobi_layer.set_low_rank_factors(U_r, V_r)
        assert dobi_layer.U.shape == (
            out_f,
            r,
        ), f"U shape mismatch: {dobi_layer.U.shape} vs ({out_f}, {r})"
        assert dobi_layer.V.shape == (
            r,
            in_f,
        ), f"V shape mismatch: {dobi_layer.V.shape} vs ({r}, {in_f})"

        if target_mod.bias is not None and dobi_layer.bias is not None:
            tb = target_mod.bias.data
            db = dobi_layer.bias.data
            min_b = min(tb.numel(), db.numel())
            db[:min_b].copy_(tb[:min_b])

        W_recon = dobi_layer.weight.float()
        recon_norm = torch.norm(W_recon, p="fro").item()
        if recon_norm < 1e-4:
            logger.warning(
                "DobiSVD: reconstructed weight Frobenius norm too small (%.6e < 1e-4)",
                recon_norm,
            )
            continue

        if norm_w < 1e-12:
            err = 0.0
        else:
            err = float(torch.norm(W - W_recon, p="fro").item() / norm_w)

        best_err = err
        chosen_r = r
        best_layer = dobi_layer

        if err <= self.max_error or r >= max_r:
            break

        if attempt < retries:
            next_r = min(int(r * 1.5), max_r)
            if next_r <= r:
                break
            current_r = next_r

    if best_layer is not None:
        best_layer.reconstruction_error = best_err
        return best_layer, best_err, chosen_r
    return None, best_err, chosen_r

Decompose a single linear or Conv1D layer into DobiSVDLinearLayer.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class DobiSVDLinearLayer (in_features: int, out_features: int, rank: int, bias: bool = True, init_gate_val: float = 2.5, dtype: torch.dtype = torch.float32, device: torch.device | None = None)

Expand source code

class DobiSVDLinearLayer(nn.Module):
    is_custom_attention: bool = True
    """
    Differentiable SVD Linear Layer (Dobi-SVD).

    Decomposes linear transformations into low-rank factorized parameters
    (U: [out_features, rank], V: [rank, in_features]) with a learnable continuous
    spectral gate vector alpha: [rank], allowing differentiable rank pruning and fine-tuning:
        y = (x @ (V * sigmoid(alpha)).T) @ U.T + bias
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        bias: bool = True,
        init_gate_val: float = 2.5,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = max(1, min(rank, in_features, out_features))

        self.V = nn.Parameter(torch.empty(self.rank, in_features, dtype=dtype, device=device))
        self.U = nn.Parameter(torch.empty(out_features, self.rank, dtype=dtype, device=device))
        # Learnable spectral gate parameter initialized to +2.5 so sigma(alpha) ~ 0.92
        self.gate = nn.Parameter(
            torch.full((self.rank,), fill_value=init_gate_val, dtype=dtype, device=device)
        )

        if bias:
            self.bias = nn.Parameter(torch.empty(out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

        self.reset_parameters()

    def reset_parameters(self) -> None:
        with torch.no_grad():
            if self.U.dtype in (torch.float16, torch.bfloat16):
                u_fp32 = torch.empty_like(self.U, dtype=torch.float32, device="cpu")
                nn.init.orthogonal_(u_fp32)
                self.U.copy_(u_fp32.to(dtype=self.U.dtype, device=self.U.device))
            else:
                nn.init.orthogonal_(self.U)

            if self.V.dtype in (torch.float16, torch.bfloat16):
                v_fp32 = torch.empty_like(self.V, dtype=torch.float32, device="cpu")
                nn.init.orthogonal_(v_fp32)
                self.V.copy_(v_fp32.to(dtype=self.V.dtype, device=self.V.device))
            else:
                nn.init.orthogonal_(self.V)

            if self.bias is not None:
                nn.init.zeros_(self.bias)

    def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
        """Initialize U and V parameters directly from SVD factors."""
        with torch.no_grad():
            self.U.copy_(U_factor)
            self.V.copy_(V_factor)

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix for compatibility."""
        g = torch.sigmoid(self.gate).unsqueeze(1)
        return torch.matmul(self.U, self.V * g)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Differentiable forward pass with spectral gating matching Linear and Conv1D."""
        if x.shape[-1] != self.in_features:
            if x.shape[-1] < self.in_features:
                x = F.pad(x, (0, self.in_features - x.shape[-1]))
            else:
                x = x[..., : self.in_features]
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)

        g = torch.sigmoid(self.gate).unsqueeze(1)  # [rank, 1]
        V_gated = self.V * g  # [rank, in_features]
        latent = F.linear(x_flat, V_gated)  # [..., rank]
        out = F.linear(latent, self.U, self.bias)  # [..., out_features]
        return out.view(*orig_shape[:-1], self.out_features)

    def extra_repr(self) -> str:
        return f"in_features={self.in_features}, out_features={self.out_features}, rank={self.rank}, bias={self.bias is not None}"

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: Differentiable SVD Linear Layer (Dobi-SVD).

Decomposes linear transformations into low-rank factorized parameters (U: [out_features, rank], V: [rank, in_features]) with a learnable continuous spectral gate vector alpha: [rank], allowing differentiable rank pruning and fine-tuning: y = (x @ (V * sigmoid(alpha)).T) @ U.T + bias

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix for compatibility."""
    g = torch.sigmoid(self.gate).unsqueeze(1)
    return torch.matmul(self.U, self.V * g)

Reconstructed effective weight matrix for compatibility.

Methods

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    return f"in_features={self.in_features}, out_features={self.out_features}, rank={self.rank}, bias={self.bias is not None}"

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Differentiable forward pass with spectral gating matching Linear and Conv1D."""
    if x.shape[-1] != self.in_features:
        if x.shape[-1] < self.in_features:
            x = F.pad(x, (0, self.in_features - x.shape[-1]))
        else:
            x = x[..., : self.in_features]
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)

    g = torch.sigmoid(self.gate).unsqueeze(1)  # [rank, 1]
    V_gated = self.V * g  # [rank, in_features]
    latent = F.linear(x_flat, V_gated)  # [..., rank]
    out = F.linear(latent, self.U, self.bias)  # [..., out_features]
    return out.view(*orig_shape[:-1], self.out_features)

Differentiable forward pass with spectral gating matching Linear and Conv1D.

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    with torch.no_grad():
        if self.U.dtype in (torch.float16, torch.bfloat16):
            u_fp32 = torch.empty_like(self.U, dtype=torch.float32, device="cpu")
            nn.init.orthogonal_(u_fp32)
            self.U.copy_(u_fp32.to(dtype=self.U.dtype, device=self.U.device))
        else:
            nn.init.orthogonal_(self.U)

        if self.V.dtype in (torch.float16, torch.bfloat16):
            v_fp32 = torch.empty_like(self.V, dtype=torch.float32, device="cpu")
            nn.init.orthogonal_(v_fp32)
            self.V.copy_(v_fp32.to(dtype=self.V.dtype, device=self.V.device))
        else:
            nn.init.orthogonal_(self.V)

        if self.bias is not None:
            nn.init.zeros_(self.bias)

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) ‑> None

Expand source code

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
    """Initialize U and V parameters directly from SVD factors."""
    with torch.no_grad():
        self.U.copy_(U_factor)
        self.V.copy_(V_factor)

Initialize U and V parameters directly from SVD factors.

class DynamicEntropySVDCompressor (model: nn.Module, min_rank: int = 16, max_rank: int = 128, target_param_ratio: float | None = None, total_rank_budget: int | None = None, target_rank: int | None = None, max_error: float = 0.35, target_attention_only: bool = False, fallback_to_original: bool = True, calibration_loader: Any = None, calibration_texts: list[str] | None = None, target_modules: list[str] | None = None, tokenizer: Any | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

class DynamicEntropySVDCompressor(BaseCompressor):
    """
    Dynamic Entropy SVD Compressor.

    1. Computes singular entropy H(Sigma_l) on layer activation covariances or weights.
    2. Allocates layer ranks within [min_rank, max_rank] subject to parameter/rank budget.
    3. Replaces dense Linear/Conv1D layers with DynamicSVDLinear factorized modules.
    4. Provides relative Frobenius error validation, max_error threshold, and safe fallback.
    5. Freezes fallback and non-target layers during recovery.
    """

    def __init__(
        self,
        model: nn.Module,
        min_rank: int = 16,
        max_rank: int = 128,
        target_param_ratio: float | None = None,
        total_rank_budget: int | None = None,
        target_rank: int | None = None,
        max_error: float = 0.35,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        calibration_loader: Any = None,
        calibration_texts: list[str] | None = None,
        target_modules: list[str] | None = None,
        tokenizer: Any | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="dynamic_entropy_svd", **kwargs)
        cfg_obj = cfg if cfg is not None else kwargs.get("comp_config")
        self.cfg = cfg_obj

        min_rank = _cfg_get(cfg_obj, "dynamic_min_rank", _cfg_get(cfg_obj, "min_rank", min_rank))
        max_rank = _cfg_get(cfg_obj, "dynamic_max_rank", _cfg_get(cfg_obj, "max_rank", max_rank))
        target_rank = _cfg_get(cfg_obj, "dynamic_rank", _cfg_get(cfg_obj, "rank", target_rank))
        total_rank_budget = _cfg_get(cfg_obj, "total_rank_budget", total_rank_budget)
        target_param_ratio = _cfg_get(cfg_obj, "target_param_ratio", target_param_ratio)
        max_error = _cfg_get(
            cfg_obj, "dynamic_max_error", _cfg_get(cfg_obj, "max_error", max_error)
        )
        target_attention_only = _cfg_get(
            cfg_obj,
            "dynamic_target_attention_only",
            _cfg_get(cfg_obj, "target_attention_only", target_attention_only),
        )
        fallback_to_original = _cfg_get(
            cfg_obj,
            "dynamic_fallback_to_original",
            _cfg_get(cfg_obj, "fallback_to_original", fallback_to_original),
        )

        min_rank = max(16, int(min_rank))
        if "rank" in kwargs and kwargs["rank"] is not None:
            target_rank = kwargs["rank"]
            min_rank = max(16, min(min_rank, target_rank))
            max_rank = max(max_rank, target_rank)

        self.min_rank = min_rank
        self.max_rank = max(min_rank, int(max_rank))
        self.target_param_ratio = (
            float(target_param_ratio) if target_param_ratio is not None else None
        )
        self.total_rank_budget = int(total_rank_budget) if total_rank_budget is not None else None
        self.target_rank = int(target_rank) if target_rank is not None else None
        if max_error is None or max_error == 0.35:
            max_error = adaptive_max_error(
                self.target_rank or self.min_rank, 768, 768, "distilbert"
            )

        self.max_error = float(max_error)
        self.target_attention_only = bool(target_attention_only)
        self.fallback_to_original = bool(fallback_to_original)

        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.target_modules = target_modules
        self.tokenizer = (
            tokenizer
            or getattr(model, "tokenizer", None)
            or getattr(getattr(model, "config", None), "tokenizer", None)
        )

        self.layer_entropies: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.layer_errors: dict[str, float] = {}
        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        """Check whether module is a compressible nn.Linear or Conv1D matching target filter."""
        name_lower = name.lower()
        if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
            return False
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            ]
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d):
            return False
        if isinstance(target_mod, DynamicSVDLinear) or isinstance(module, DynamicSVDLinear):
            return False
        if not should_compress_linear(self.model, name):
            return False

        raw_w = target_mod.weight.data
        if min(raw_w.shape[0], raw_w.shape[1]) < 16:
            return False

        if self.target_attention_only:
            last_part = name.split(".")[-1]
            is_attn = (last_part in _ATTN_NAMES) or any(
                k in name_lower for k in ("attn", "attention")
            )
            if not is_attn:
                return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def collect_activation_covariances(
        self,
    ) -> dict[str, torch.Tensor]:
        """
        Collect empirical input activation covariance matrices C_X for each target layer
        using calibration forward passes.
        """
        covariances: dict[str, torch.Tensor] = {}
        sample_counts: dict[str, int] = {}
        hooks = []

        def get_hook(name: str):
            def hook_fn(mod: nn.Module, inp: tuple[Any, ...], out: Any) -> None:
                if not inp or inp[0] is None:
                    return
                x = inp[0].detach()
                if x.ndim >= 2:
                    x_flat = x.reshape(-1, x.shape[-1]).float()
                    # Double-precision covariance accumulation for stability
                    cov = x_flat.t().contiguous() @ x_flat
                    n = x_flat.shape[0]
                    if name not in covariances:
                        covariances[name] = cov
                        sample_counts[name] = n
                    else:
                        covariances[name] += cov
                        sample_counts[name] += n

            return hook_fn

        target_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            return {}

        for name, mod in target_layers:
            target_mod = getattr(mod, "base_layer", mod)
            hooks.append(target_mod.register_forward_hook(get_hook(name)))

        device = next(self.model.parameters()).device
        self.model.eval()

        # Track whether REAL calibration data (not random dummy) was used.
        # If only dummy data is used, the covariance matrix will be rank-deficient
        # (rank <= num_dummy_tokens << hidden_dim), causing the whitening transform
        # to map near-zero eigenvalue directions through inv_sqrt with scale ~100x,
        # which after un-whitening produces reconstruction error near 1.0 on every layer.
        real_calibration_used = False

        try:
            with torch.no_grad():
                if self.calibration_loader is not None:
                    batch_count = 0
                    for batch in self.calibration_loader:
                        if isinstance(batch, dict):
                            batch = {
                                k: v.to(device) if isinstance(v, torch.Tensor) else v
                                for k, v in batch.items()
                            }
                            self.model(**batch)
                        elif isinstance(batch, (list, tuple)):
                            batch = [
                                v.to(device) if isinstance(v, torch.Tensor) else v for v in batch
                            ]
                            self.model(*batch)
                        elif isinstance(batch, torch.Tensor):
                            self.model(batch.to(device))
                        batch_count += 1
                    if batch_count > 0:
                        real_calibration_used = True
                        logger.info(
                            "DynamicEntropySVD: real calibration loader used (%d batches).",
                            batch_count,
                        )
                elif self.calibration_texts is not None and len(self.calibration_texts) > 0:
                    logger.info(
                        "Running calibration forward on %d provided texts.",
                        len(self.calibration_texts),
                    )
                    if self.tokenizer is not None:
                        text_count = 0
                        for text in self.calibration_texts[:32]:
                            try:
                                enc = self.tokenizer(
                                    text,
                                    return_tensors="pt",
                                    truncation=True,
                                    max_length=128,
                                )
                                enc = {k: v.to(device) for k, v in enc.items()}
                                self.model(**enc)
                                text_count += 1
                            except Exception as e:
                                logger.debug("Calibration text pass exception: %s", e)
                                break
                        if text_count > 0:
                            real_calibration_used = True
                            logger.info(
                                "DynamicEntropySVD: real calibration texts used (%d texts).",
                                text_count,
                            )
                    else:
                        logger.info(
                            "No tokenizer available; calibration_texts provided but cannot encode. "
                            "Falling back to weight-only SVD."
                        )
        except Exception as e:
            logger.warning(
                "Activation collection encountered exception (%s); falling back to weight SVD.",
                e,
            )
        finally:
            for h in hooks:
                h.remove()

        # ROOT-CAUSE FIX: If no real calibration data was used, dummy inputs (e.g.
        # 2x32 random tokens) produce a rank-deficient covariance (rank<=64 for
        # hidden_dim=768). The whitening eigenvalue clamp (min=1e-4) then assigns
        # weight 100x to the 704 meaningless directions. After SVD truncation and
        # un-whitening, reconstruction error approaches 1.0 on every layer, causing
        # all 24 layers to fall back and zero compression.
        #
        # Solution: return {} so from_linear uses plain weight-only SVD instead.
        # Expected reconstruction error at rank 64: ~0.15-0.25 (all layers compress).
        if not real_calibration_used:
            logger.warning(
                "DynamicEntropySVD: no real calibration data available "
                "(calibration_loader=None, calibration_texts=None/empty). "
                "Activation-aware whitening DISABLED. Using weight-only SVD for all "
                "layers. Expected reconstruction error at rank=64: ~0.15-0.25."
            )
            return {}

        # Normalize accumulated covariances by total sample count
        normalized_covs: dict[str, torch.Tensor] = {}
        for name, cov in covariances.items():
            cnt = sample_counts.get(name, 1)
            normalized_covs[name] = cov / max(1, cnt)

        logger.info(
            "DynamicEntropySVD: %d activation covariances collected (real data).",
            len(normalized_covs),
        )
        return normalized_covs

    def compute_layer_entropies(self, covariances: dict[str, torch.Tensor]) -> dict[str, float]:
        """
        Compute singular entropy for each target layer using activation covariance
        or fallback weight matrix spectrum.
        """
        entropies: dict[str, float] = {}
        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                target_mod = getattr(module, "base_layer", module)
                raw_w = target_mod.weight.data
                cls_name = target_mod.__class__.__name__.lower()
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                w_mat = raw_w.t().contiguous().float() if is_conv1d else raw_w.float()

                if name in covariances:
                    ent = compute_singular_entropy(covariances[name])
                else:
                    ent = compute_singular_entropy(w_mat)
                entropies[name] = ent
        return entropies

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """
        Execute dynamic entropy-budgeted low-rank SVD compression across all target layers.
        Includes reconstruction error validation, fallback, layer freezing, and parameter tracking.

        Returns:
            Tuple[nn.Module, Dict[str, Any]]: Compressed model and detailed statistics.
        """
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}

        # Step 1: Collect activation covariances
        covariances = self.collect_activation_covariances()

        # Step 2: Compute singular entropies per layer
        self.layer_entropies = self.compute_layer_entropies(covariances)

        if not self.layer_entropies:
            logger.warning("No compressible linear layers found in model.")
            return self.model, self.get_compression_stats()

        # Step 3: Extract layer dimensions
        layer_dims: dict[str, tuple[int, int]] = {}
        total_orig_decomp_params = 0
        total_comp_decomp_params = 0

        for name, module in self.model.named_modules():
            if name in self.layer_entropies:
                target_mod = getattr(module, "base_layer", module)
                raw_w = target_mod.weight.data
                cls_name = target_mod.__class__.__name__.lower()
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                if is_conv1d:
                    in_f, out_f = raw_w.shape[0], raw_w.shape[1]
                else:
                    out_f, in_f = raw_w.shape[0], raw_w.shape[1]
                layer_dims[name] = (out_f, in_f)
                total_orig_decomp_params += out_f * in_f

        # Freeze uncompressed layers when target_attention_only=True
        if self.target_attention_only:
            for name, module in list(self.model.named_modules()):
                target_mod = getattr(module, "base_layer", module)
                cls_name = target_mod.__class__.__name__.lower()
                is_linear = isinstance(target_mod, nn.Linear)
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                if is_linear or is_conv1d:
                    if not self._is_target_module(name, module):
                        for p in module.parameters():
                            p.requires_grad = False
                        self.skipped_layers.append(name)

        # Step 4: Determine total rank budget
        total_rank_budget = self.total_rank_budget
        if total_rank_budget is None and self.target_rank is not None:
            total_rank_budget = self.target_rank * len(self.layer_entropies)

        # Step 5: Dynamic rank allocation
        self.layer_ranks = allocate_entropy_ranks(
            entropies=self.layer_entropies,
            min_rank=self.min_rank,
            max_rank=self.max_rank,
            total_rank_budget=total_rank_budget,
            target_param_ratio=self.target_param_ratio,
            layer_dims=layer_dims,
        )

        # Step 6: Factorize, validate, and replace each target linear layer
        for name, rank in self.layer_ranks.items():
            parent_name, child_name = name.rsplit(".", 1) if "." in name else ("", name)
            if parent_name:
                parent = self.model.get_submodule(parent_name)
            else:
                parent = self.model
            module = getattr(parent, child_name)
            target_mod = getattr(module, "base_layer", module)

            out_f, in_f = layer_dims[name]
            orig_params = out_f * in_f

            cov = covariances.get(name, None)
            factorized_layer = DynamicSVDLinear.from_linear(
                linear=module,
                rank=rank,
                covariance=cov,
            )

            # Compute relative Frobenius reconstruction error
            raw_w = target_mod.weight.data
            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            W = raw_w.t().contiguous().float() if is_conv1d else raw_w.float()
            W_recon = factorized_layer.U.float() @ factorized_layer.V.float()
            norm_w = torch.norm(W, p="fro").item()
            err = float(torch.norm(W - W_recon, p="fro").item() / max(norm_w, 1e-12))
            self.layer_errors[name] = err

            if err > self.max_error and self.fallback_to_original:
                logger.warning(
                    "DynamicEntropySVD error %.4f on %s exceeded threshold %.4f; falling back to original.",
                    err,
                    name,
                    self.max_error,
                )
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_decomp_params += orig_params
                continue

            # If wrapped in LoRALinear, replace the inner base_layer or linear
            has_base_layer = hasattr(module, "base_layer") or (
                hasattr(module, "linear") and (hasattr(module, "lora_A") or hasattr(module, "r"))
            )
            if hasattr(module, "base_layer"):
                module.base_layer = factorized_layer
            elif hasattr(module, "linear") and (hasattr(module, "lora_A") or hasattr(module, "r")):
                module.linear = factorized_layer
            else:
                replace_module(self.model, name, factorized_layer)

            if has_base_layer:
                for p_name, p in module.named_parameters():
                    if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                        p.requires_grad = True

            self.compressed_layers.append(name)
            total_comp_decomp_params += rank * (in_f + out_f)

        # Forward sanity check
        forward_check_success = False
        try:
            self.model.eval()
            with torch.no_grad():
                d = next(self.model.parameters()).device
                dummy = torch.randint(0, 100, (1, 16), device=d)
                self.model(dummy)
            forward_check_success = True
        except Exception as e:
            logger.debug("DynamicEntropySVD forward check exception: %s", e)

        effective_ratio = (
            total_comp_decomp_params / total_orig_decomp_params
            if total_orig_decomp_params > 0
            else 1.0
        )
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )

        stats = self.get_compression_stats()
        stats.update(
            {
                "layer_entropies": self.layer_entropies,
                "layer_ranks": self.layer_ranks,
                "layer_errors": self.layer_errors,
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "original_layer_count": len(layer_dims),
                "max_error": self.max_error,
                "target_attention_only": self.target_attention_only,
                "fallback_to_original": self.fallback_to_original,
                "effective_parameter_ratio": round(effective_ratio, 4),
                "avg_reconstruction_error": avg_err,
                "forward_check_success": forward_check_success,
                "method": "dynamic_entropy_svd",
            }
        )
        self.stats = stats
        logger.info(
            "Compressor dynamic_rank finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

Dynamic Entropy SVD Compressor.

Computes singular entropy H(Sigma_l) on layer activation covariances or weights.
Allocates layer ranks within [min_rank, max_rank] subject to parameter/rank budget.
Replaces dense Linear/Conv1D layers with DynamicSVDLinear factorized modules.
Provides relative Frobenius error validation, max_error threshold, and safe fallback.
Freezes fallback and non-target layers during recovery.

Ancestors

BaseCompressor

Subclasses

neuralops.compression_engine.dynamic_rank.DynamicEntropySVDCompressorRegistered
neuralops.compression_engine.dynamic_rank.EntropyRankAllocatorRegistered

Methods

def collect_activation_covariances(self) ‑> dict[str, torch.Tensor]

Expand source code

def collect_activation_covariances(
    self,
) -> dict[str, torch.Tensor]:
    """
    Collect empirical input activation covariance matrices C_X for each target layer
    using calibration forward passes.
    """
    covariances: dict[str, torch.Tensor] = {}
    sample_counts: dict[str, int] = {}
    hooks = []

    def get_hook(name: str):
        def hook_fn(mod: nn.Module, inp: tuple[Any, ...], out: Any) -> None:
            if not inp or inp[0] is None:
                return
            x = inp[0].detach()
            if x.ndim >= 2:
                x_flat = x.reshape(-1, x.shape[-1]).float()
                # Double-precision covariance accumulation for stability
                cov = x_flat.t().contiguous() @ x_flat
                n = x_flat.shape[0]
                if name not in covariances:
                    covariances[name] = cov
                    sample_counts[name] = n
                else:
                    covariances[name] += cov
                    sample_counts[name] += n

        return hook_fn

    target_layers = [
        (name, mod)
        for name, mod in self.model.named_modules()
        if self._is_target_module(name, mod)
    ]

    if not target_layers:
        return {}

    for name, mod in target_layers:
        target_mod = getattr(mod, "base_layer", mod)
        hooks.append(target_mod.register_forward_hook(get_hook(name)))

    device = next(self.model.parameters()).device
    self.model.eval()

    # Track whether REAL calibration data (not random dummy) was used.
    # If only dummy data is used, the covariance matrix will be rank-deficient
    # (rank <= num_dummy_tokens << hidden_dim), causing the whitening transform
    # to map near-zero eigenvalue directions through inv_sqrt with scale ~100x,
    # which after un-whitening produces reconstruction error near 1.0 on every layer.
    real_calibration_used = False

    try:
        with torch.no_grad():
            if self.calibration_loader is not None:
                batch_count = 0
                for batch in self.calibration_loader:
                    if isinstance(batch, dict):
                        batch = {
                            k: v.to(device) if isinstance(v, torch.Tensor) else v
                            for k, v in batch.items()
                        }
                        self.model(**batch)
                    elif isinstance(batch, (list, tuple)):
                        batch = [
                            v.to(device) if isinstance(v, torch.Tensor) else v for v in batch
                        ]
                        self.model(*batch)
                    elif isinstance(batch, torch.Tensor):
                        self.model(batch.to(device))
                    batch_count += 1
                if batch_count > 0:
                    real_calibration_used = True
                    logger.info(
                        "DynamicEntropySVD: real calibration loader used (%d batches).",
                        batch_count,
                    )
            elif self.calibration_texts is not None and len(self.calibration_texts) > 0:
                logger.info(
                    "Running calibration forward on %d provided texts.",
                    len(self.calibration_texts),
                )
                if self.tokenizer is not None:
                    text_count = 0
                    for text in self.calibration_texts[:32]:
                        try:
                            enc = self.tokenizer(
                                text,
                                return_tensors="pt",
                                truncation=True,
                                max_length=128,
                            )
                            enc = {k: v.to(device) for k, v in enc.items()}
                            self.model(**enc)
                            text_count += 1
                        except Exception as e:
                            logger.debug("Calibration text pass exception: %s", e)
                            break
                    if text_count > 0:
                        real_calibration_used = True
                        logger.info(
                            "DynamicEntropySVD: real calibration texts used (%d texts).",
                            text_count,
                        )
                else:
                    logger.info(
                        "No tokenizer available; calibration_texts provided but cannot encode. "
                        "Falling back to weight-only SVD."
                    )
    except Exception as e:
        logger.warning(
            "Activation collection encountered exception (%s); falling back to weight SVD.",
            e,
        )
    finally:
        for h in hooks:
            h.remove()

    # ROOT-CAUSE FIX: If no real calibration data was used, dummy inputs (e.g.
    # 2x32 random tokens) produce a rank-deficient covariance (rank<=64 for
    # hidden_dim=768). The whitening eigenvalue clamp (min=1e-4) then assigns
    # weight 100x to the 704 meaningless directions. After SVD truncation and
    # un-whitening, reconstruction error approaches 1.0 on every layer, causing
    # all 24 layers to fall back and zero compression.
    #
    # Solution: return {} so from_linear uses plain weight-only SVD instead.
    # Expected reconstruction error at rank 64: ~0.15-0.25 (all layers compress).
    if not real_calibration_used:
        logger.warning(
            "DynamicEntropySVD: no real calibration data available "
            "(calibration_loader=None, calibration_texts=None/empty). "
            "Activation-aware whitening DISABLED. Using weight-only SVD for all "
            "layers. Expected reconstruction error at rank=64: ~0.15-0.25."
        )
        return {}

    # Normalize accumulated covariances by total sample count
    normalized_covs: dict[str, torch.Tensor] = {}
    for name, cov in covariances.items():
        cnt = sample_counts.get(name, 1)
        normalized_covs[name] = cov / max(1, cnt)

    logger.info(
        "DynamicEntropySVD: %d activation covariances collected (real data).",
        len(normalized_covs),
    )
    return normalized_covs

Collect empirical input activation covariance matrices C_X for each target layer using calibration forward passes.

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """
    Execute dynamic entropy-budgeted low-rank SVD compression across all target layers.
    Includes reconstruction error validation, fallback, layer freezing, and parameter tracking.

    Returns:
        Tuple[nn.Module, Dict[str, Any]]: Compressed model and detailed statistics.
    """
    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    self.layer_errors = {}

    # Step 1: Collect activation covariances
    covariances = self.collect_activation_covariances()

    # Step 2: Compute singular entropies per layer
    self.layer_entropies = self.compute_layer_entropies(covariances)

    if not self.layer_entropies:
        logger.warning("No compressible linear layers found in model.")
        return self.model, self.get_compression_stats()

    # Step 3: Extract layer dimensions
    layer_dims: dict[str, tuple[int, int]] = {}
    total_orig_decomp_params = 0
    total_comp_decomp_params = 0

    for name, module in self.model.named_modules():
        if name in self.layer_entropies:
            target_mod = getattr(module, "base_layer", module)
            raw_w = target_mod.weight.data
            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            if is_conv1d:
                in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            else:
                out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            layer_dims[name] = (out_f, in_f)
            total_orig_decomp_params += out_f * in_f

    # Freeze uncompressed layers when target_attention_only=True
    if self.target_attention_only:
        for name, module in list(self.model.named_modules()):
            target_mod = getattr(module, "base_layer", module)
            cls_name = target_mod.__class__.__name__.lower()
            is_linear = isinstance(target_mod, nn.Linear)
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            if is_linear or is_conv1d:
                if not self._is_target_module(name, module):
                    for p in module.parameters():
                        p.requires_grad = False
                    self.skipped_layers.append(name)

    # Step 4: Determine total rank budget
    total_rank_budget = self.total_rank_budget
    if total_rank_budget is None and self.target_rank is not None:
        total_rank_budget = self.target_rank * len(self.layer_entropies)

    # Step 5: Dynamic rank allocation
    self.layer_ranks = allocate_entropy_ranks(
        entropies=self.layer_entropies,
        min_rank=self.min_rank,
        max_rank=self.max_rank,
        total_rank_budget=total_rank_budget,
        target_param_ratio=self.target_param_ratio,
        layer_dims=layer_dims,
    )

    # Step 6: Factorize, validate, and replace each target linear layer
    for name, rank in self.layer_ranks.items():
        parent_name, child_name = name.rsplit(".", 1) if "." in name else ("", name)
        if parent_name:
            parent = self.model.get_submodule(parent_name)
        else:
            parent = self.model
        module = getattr(parent, child_name)
        target_mod = getattr(module, "base_layer", module)

        out_f, in_f = layer_dims[name]
        orig_params = out_f * in_f

        cov = covariances.get(name, None)
        factorized_layer = DynamicSVDLinear.from_linear(
            linear=module,
            rank=rank,
            covariance=cov,
        )

        # Compute relative Frobenius reconstruction error
        raw_w = target_mod.weight.data
        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        W = raw_w.t().contiguous().float() if is_conv1d else raw_w.float()
        W_recon = factorized_layer.U.float() @ factorized_layer.V.float()
        norm_w = torch.norm(W, p="fro").item()
        err = float(torch.norm(W - W_recon, p="fro").item() / max(norm_w, 1e-12))
        self.layer_errors[name] = err

        if err > self.max_error and self.fallback_to_original:
            logger.warning(
                "DynamicEntropySVD error %.4f on %s exceeded threshold %.4f; falling back to original.",
                err,
                name,
                self.max_error,
            )
            self.fallback_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False
            total_comp_decomp_params += orig_params
            continue

        # If wrapped in LoRALinear, replace the inner base_layer or linear
        has_base_layer = hasattr(module, "base_layer") or (
            hasattr(module, "linear") and (hasattr(module, "lora_A") or hasattr(module, "r"))
        )
        if hasattr(module, "base_layer"):
            module.base_layer = factorized_layer
        elif hasattr(module, "linear") and (hasattr(module, "lora_A") or hasattr(module, "r")):
            module.linear = factorized_layer
        else:
            replace_module(self.model, name, factorized_layer)

        if has_base_layer:
            for p_name, p in module.named_parameters():
                if any(t in p_name.lower() for t in ("lora_a", "lora_b", "adapter")):
                    p.requires_grad = True

        self.compressed_layers.append(name)
        total_comp_decomp_params += rank * (in_f + out_f)

    # Forward sanity check
    forward_check_success = False
    try:
        self.model.eval()
        with torch.no_grad():
            d = next(self.model.parameters()).device
            dummy = torch.randint(0, 100, (1, 16), device=d)
            self.model(dummy)
        forward_check_success = True
    except Exception as e:
        logger.debug("DynamicEntropySVD forward check exception: %s", e)

    effective_ratio = (
        total_comp_decomp_params / total_orig_decomp_params
        if total_orig_decomp_params > 0
        else 1.0
    )
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )

    stats = self.get_compression_stats()
    stats.update(
        {
            "layer_entropies": self.layer_entropies,
            "layer_ranks": self.layer_ranks,
            "layer_errors": self.layer_errors,
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "skipped_layers": self.skipped_layers,
            "skipped_layer_count": len(self.skipped_layers),
            "original_layer_count": len(layer_dims),
            "max_error": self.max_error,
            "target_attention_only": self.target_attention_only,
            "fallback_to_original": self.fallback_to_original,
            "effective_parameter_ratio": round(effective_ratio, 4),
            "avg_reconstruction_error": avg_err,
            "forward_check_success": forward_check_success,
            "method": "dynamic_entropy_svd",
        }
    )
    self.stats = stats
    logger.info(
        "Compressor dynamic_rank finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        self.max_error,
    )
    return self.model, stats

Execute dynamic entropy-budgeted low-rank SVD compression across all target layers. Includes reconstruction error validation, fallback, layer freezing, and parameter tracking.

Returns
-----=
Tuple[nn.Module, Dict[str, Any]]: Compressed model and detailed statistics.

def compute_layer_entropies(self, covariances: dict[str, torch.Tensor]) ‑> dict[str, float]

Expand source code

def compute_layer_entropies(self, covariances: dict[str, torch.Tensor]) -> dict[str, float]:
    """
    Compute singular entropy for each target layer using activation covariance
    or fallback weight matrix spectrum.
    """
    entropies: dict[str, float] = {}
    for name, module in self.model.named_modules():
        if self._is_target_module(name, module):
            target_mod = getattr(module, "base_layer", module)
            raw_w = target_mod.weight.data
            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            w_mat = raw_w.t().contiguous().float() if is_conv1d else raw_w.float()

            if name in covariances:
                ent = compute_singular_entropy(covariances[name])
            else:
                ent = compute_singular_entropy(w_mat)
            entropies[name] = ent
    return entropies

Compute singular entropy for each target layer using activation covariance or fallback weight matrix spectrum.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class DynamicEntropySVDCompressorRegistered (model: nn.Module, min_rank: int = 16, max_rank: int = 128, target_param_ratio: float | None = None, total_rank_budget: int | None = None, target_rank: int | None = None, max_error: float = 0.35, target_attention_only: bool = False, fallback_to_original: bool = True, calibration_loader: Any = None, calibration_texts: list[str] | None = None, target_modules: list[str] | None = None, tokenizer: Any | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("dynamic_entropy_svd_class")
@COMPRESSOR_REGISTRY.register("dynamic_rank_class")
@COMPRESSOR_REGISTRY.register("dynamic_entropy_class")
class DynamicEntropySVDCompressorRegistered(DynamicEntropySVDCompressor):
    pass

Dynamic Entropy SVD Compressor.

Computes singular entropy H(Sigma_l) on layer activation covariances or weights.
Allocates layer ranks within [min_rank, max_rank] subject to parameter/rank budget.
Replaces dense Linear/Conv1D layers with DynamicSVDLinear factorized modules.
Provides relative Frobenius error validation, max_error threshold, and safe fallback.
Freezes fallback and non-target layers during recovery.

Ancestors

Inherited members

DynamicEntropySVDCompressor:
- collect_activation_covariances
- compress
- compute_layer_entropies
- get_compression_stats
- summary

class DynamicSVDLinear (in_features: int, out_features: int, rank: int, bias: bool = True, dtype: torch.dtype = torch.float32, device: torch.device | None = None)

Expand source code

class DynamicSVDLinear(nn.Module):
    is_custom_attention: bool = True
    """
    True Low-Rank Factorized Linear Layer executing as:
    y = (x @ V.T) @ U.T + bias
    where V is (rank, in_features) and U is (out_features, rank).
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        bias: bool = True,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = max(1, min(rank, in_features, out_features))

        # Factorized low-rank parameter tensors
        self.V = nn.Parameter(torch.empty(self.rank, in_features, dtype=dtype, device=device))
        self.U = nn.Parameter(torch.empty(out_features, self.rank, dtype=dtype, device=device))

        if bias:
            self.bias = nn.Parameter(torch.empty(out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

        self.reset_parameters()

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix for compatibility."""
        return torch.matmul(self.U, self.V)

    def reset_parameters(self) -> None:
        """Kaiming uniform parameter initialization."""
        nn.init.kaiming_uniform_(self.V, a=math.sqrt(5))
        nn.init.kaiming_uniform_(self.U, a=math.sqrt(5))
        if self.bias is not None:
            fan_in = self.in_features
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            nn.init.uniform_(self.bias, -bound, bound)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass executing factorized low-rank multiplication:
        y = (x @ V^T) @ U^T + bias
        """
        intermediate = F.linear(x, self.V)
        out = F.linear(intermediate, self.U, self.bias)
        return out

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        rank: int,
        covariance: torch.Tensor | None = None,
    ) -> DynamicSVDLinear:
        """
        Construct factorized DynamicSVDLinear directly from dense nn.Linear or Conv1D via SVD.

        Args:
            linear: Source dense linear or Conv1D layer (or LoRALinear wrapper).
            rank: Target rank for low-rank factorization.
            covariance: Optional empirical activation covariance C_X for activation-aware SVD.

        Returns:
            DynamicSVDLinear: Initialized factorized low-rank layer.
        """
        target_mod = getattr(linear, "base_layer", linear)
        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.data

        if is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()  # [out_features, in_features]
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float()  # [out_features, in_features]

        target_rank = max(1, min(rank, in_f, out_f))

        new_layer = cls(
            in_features=in_f,
            out_features=out_f,
            rank=target_rank,
            bias=target_mod.bias is not None,
            dtype=target_mod.weight.dtype,
            device=target_mod.weight.device,
        )

        # Always compute the weight-only SVD baseline first.
        # This serves as both the default path (when covariance=None) and the
        # fallback comparison when the activation-aware path produces higher error.
        U_svd_w, S_svd_w, Vh_svd_w = torch.linalg.svd(W, full_matrices=False)
        sqrt_S_w = torch.sqrt(torch.clamp(S_svd_w[:target_rank], min=1e-12))
        U_w = U_svd_w[:, :target_rank] * sqrt_S_w.unsqueeze(0)  # [out_f, rank]
        V_w = sqrt_S_w.unsqueeze(1) * Vh_svd_w[:target_rank, :]  # [rank, in_f]
        W_rec_w = U_w @ V_w
        norm_W = torch.norm(W, p="fro").clamp(min=1e-12)
        err_weight_only = (torch.norm(W - W_rec_w, p="fro") / norm_W).item()

        # Default to weight-only unless activation-aware is demonstrably better.
        U_r, V_r = U_w, V_w

        if covariance is not None and covariance.numel() > 0 and covariance.shape[0] == in_f:
            # RESEARCH FIX: Activation-aware whitening SVD with eigenvalue clamping lambda_i >= 1e-4
            # NOTE: We only accept the activation-aware result if its reconstruction error is
            # strictly lower than the weight-only baseline. This guards against rank-deficient
            # covariances (e.g. from dummy calibration data) that inflate inv_sqrt by up to 100x
            # and destroy the approximation quality.
            cov = covariance.to(device=W.device, dtype=torch.float32)
            try:
                # Eigen-decomposition with regularized eigenvalue clamping >= 1e-4
                eigvals, eigvecs = torch.linalg.eigh(cov + 1e-4 * torch.eye(in_f, device=W.device))
                eigvals_clamped = torch.clamp(eigvals, min=1e-4)
                cov_sqrt = (
                    eigvecs @ torch.diag(torch.sqrt(eigvals_clamped)) @ eigvecs.t().contiguous()
                )
                cov_inv_sqrt = (
                    eigvecs
                    @ torch.diag(1.0 / torch.sqrt(eigvals_clamped))
                    @ eigvecs.t().contiguous()
                )

                W_scaled = W @ cov_sqrt
                U_svd_a, S_svd_a, Vh_svd_a = torch.linalg.svd(W_scaled, full_matrices=False)

                # Truncate and un-whiten right singular factors
                sqrt_S_a = torch.sqrt(torch.clamp(S_svd_a[:target_rank], min=1e-12))
                U_a = U_svd_a[:, :target_rank] * sqrt_S_a.unsqueeze(0)
                V_a = sqrt_S_a.unsqueeze(1) * (Vh_svd_a[:target_rank, :] @ cov_inv_sqrt)
                # Symmetrically normalize factor magnitudes
                norm_u = torch.norm(U_a, p="fro").clamp_min(1e-8)
                norm_v = torch.norm(V_a, p="fro").clamp_min(1e-8)
                scale = torch.sqrt(norm_u / norm_v)
                U_a = U_a / scale
                V_a = V_a * scale

                W_rec_a = U_a @ V_a
                err_act_aware = (torch.norm(W - W_rec_a, p="fro") / norm_W).item()

                if err_act_aware < err_weight_only:
                    # Activation-aware is genuinely better: accept it
                    U_r, V_r = U_a, V_a
                    logger.debug(
                        "from_linear: activation-aware SVD accepted "
                        "(err_act=%.4f < err_weight=%.4f).",
                        err_act_aware,
                        err_weight_only,
                    )
                else:
                    # Activation-aware is worse (e.g. due to rank-deficient covariance):
                    # stick with the weight-only baseline
                    logger.debug(
                        "from_linear: activation-aware SVD rejected "
                        "(err_act=%.4f >= err_weight=%.4f); using weight-only SVD.",
                        err_act_aware,
                        err_weight_only,
                    )
            except Exception as e:
                logger.warning("Activation-aware SVD failed (%s); using weight-only SVD.", e)
                # U_r, V_r already set to weight-only baseline above

        new_layer.U.data.copy_(U_r.to(dtype=target_mod.weight.dtype))
        new_layer.V.data.copy_(V_r.to(dtype=target_mod.weight.dtype))

        if target_mod.bias is not None and new_layer.bias is not None:
            new_layer.bias.data.copy_(target_mod.bias.data)

        return new_layer

    def extra_repr(self) -> str:
        return f"in_features={self.in_features}, out_features={self.out_features}, rank={self.rank}, bias={self.bias is not None}"

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: True Low-Rank Factorized Linear Layer executing as: y = (x @ V.T) @ U.T + bias where V is (rank, in_features) and U is (out_features, rank).

Static methods

def from_linear(linear: nn.Module, rank: int, covariance: torch.Tensor | None = None) ‑> DynamicSVDLinear

Construct factorized DynamicSVDLinear directly from dense nn.Linear or Conv1D via SVD.

Args
-----=
linear: Source dense linear or Conv1D layer (or LoRALinear wrapper).
rank: Target rank for low-rank factorization.
covariance: Optional empirical activation covariance C_X for activation-aware SVD.
Returns
-----=
DynamicSVDLinear: Initialized factorized low-rank layer.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix for compatibility."""
    return torch.matmul(self.U, self.V)

Reconstructed effective weight matrix for compatibility.

Methods

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    return f"in_features={self.in_features}, out_features={self.out_features}, rank={self.rank}, bias={self.bias is not None}"

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Forward pass executing factorized low-rank multiplication:
    y = (x @ V^T) @ U^T + bias
    """
    intermediate = F.linear(x, self.V)
    out = F.linear(intermediate, self.U, self.bias)
    return out

Forward pass executing factorized low-rank multiplication: y = (x @ V^T) @ U^T + bias

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    """Kaiming uniform parameter initialization."""
    nn.init.kaiming_uniform_(self.V, a=math.sqrt(5))
    nn.init.kaiming_uniform_(self.U, a=math.sqrt(5))
    if self.bias is not None:
        fan_in = self.in_features
        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
        nn.init.uniform_(self.bias, -bound, bound)

Kaiming uniform parameter initialization.

class ENLightenCompressor (model: nn.Module, rank: int = 64, sparsity: float = 0.5, structured_24: bool = True, max_error: float = 0.45, target_attention_only: bool = False, fallback_to_original: bool = True, max_retries: int = 2, target_modules: list[str] | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("enlighten")
@COMPRESSOR_REGISTRY.register("enlighten_compressor")
class ENLightenCompressor(BaseCompressor):
    is_custom_attention: bool = True
    """
    ENLighten Compressor: Low-Rank SVD + 2:4 Structured Sparsity.
    Decomposes weights into low-rank factors plus 2:4 structured sparse residuals,
    preserving Frobenius energy while reducing memory bandwidth and arithmetic complexity.

    Production hardening features:
    - Robust LoRALinear wrapper and Conv1D support (with transpose and adapter preservation).
    - target_attention_only parameter: restricts compression to attention projections and freezes non-target layers.
    - Relative Frobenius reconstruction error validation with configurable max_error.
    - Automatic rank escalation retries when error threshold is exceeded.
    - Safe fallback to original module if error remains high or decomposition is unstable, with fallback freezing.
    - Fully vectorized 2:4 structured pruning eliminating Python loop overhead.
    - Critical-layer protection via should_compress_linear.
    - Real compression metrics recording (effective_parameter_ratio, num_compressed, num_fallback, etc.).
    """

    def __init__(
        self,
        model: nn.Module,
        rank: int = 64,
        sparsity: float = 0.5,
        structured_24: bool = True,
        max_error: float = 0.45,  # raised from 0.35: at rank=64 + 2:4 sparsity, theoretical error is ~0.39
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        max_retries: int = 2,
        target_modules: list[str] | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="enlighten", **kwargs)
        cfg = cfg or kwargs.get("comp_config")
        self.cfg = cfg

        self.rank = int(_cfg_get(cfg, "enlighten_rank", _cfg_get(cfg, "rank", rank)))
        self.sparsity = float(
            _cfg_get(cfg, "enlighten_sparsity", _cfg_get(cfg, "sparsity", sparsity))
        )
        self.structured_24 = bool(
            _cfg_get(
                cfg,
                "enlighten_structured_24",
                _cfg_get(cfg, "structured_24", structured_24),
            )
        )
        self.max_error = float(
            _cfg_get(cfg, "enlighten_max_error", _cfg_get(cfg, "max_error", max_error))
        )
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "enlighten_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "enlighten_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        self.max_retries = int(max_retries)
        self.target_modules = target_modules

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_ranks: dict[str, int] = {}
        self.layer_sparsities: dict[str, float] = {}
        self.layer_errors: dict[str, float] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        """Determine if a module is targeted for ENLighten compression."""
        name_lower = name.lower()
        if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
            return False
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            ]
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d):
            return False

        if not should_compress_linear(self.model, name):
            return False

        if self.target_attention_only:
            last_part = name.split(".")[-1]
            is_attn = (last_part in _ATTN_NAMES) or any(
                k in name_lower for k in ("attn", "attention")
            )
            if not is_attn:
                return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)

        return True

    def _decompose_and_prune(self, W: torch.Tensor, rank: int) -> tuple[torch.Tensor, float]:
        """Perform FP32 SVD, compute residual, apply 2:4 / unstructured pruning, and scale."""
        out_f, in_f = W.shape
        r = max(1, min(rank, out_f, in_f))
        norm_w = torch.linalg.norm(W, ord="fro").item()

        U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        U_r = U[:, :r]
        S_r = S[:r]
        Vt_r = Vt[:r, :]
        L = (U_r * S_r.unsqueeze(0)) @ Vt_r
        residual = W - L

        if self.structured_24:
            residual_sp = vectorized_24_prune(residual)
        else:
            residual_sp = unstructured_prune(residual, self.sparsity)

        W_recon = L + residual_sp

        err = (
            float(torch.linalg.norm(W - W_recon, ord="fro").item() / max(norm_w, 1e-12))
            if norm_w > 1e-12
            else 0.0
        )
        return W_recon, err

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Compress linear and Conv1D layers using low-rank SVD + 2:4 structured sparse residuals."""
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_ranks = {}
        self.layer_sparsities = {}
        self.layer_errors = {}

        total_orig_params = 0
        total_comp_params = 0

        for name, module in list(self.model.named_modules()):
            name_lower = name.lower()
            if any(
                skip in name_lower
                for skip in [
                    "classifier",
                    "score",
                    "lm_head",
                    "pooler",
                    "embedding",
                    "router",
                    "gate",
                ]
            ):
                continue

            if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
                continue

            target_mod = getattr(module, "base_layer", module)
            cls_name = target_mod.__class__.__name__.lower()
            is_linear = isinstance(target_mod, nn.Linear)
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")

            if not (is_linear or is_conv1d):
                continue

            if not hasattr(target_mod, "weight") or target_mod.weight is None:
                continue

            if not self._is_target_module(name, module):
                if self.target_attention_only:
                    if not any(
                        part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]
                    ):
                        for p in module.parameters():
                            p.requires_grad = False
                        self.skipped_layers.append(name)
                continue

            raw_w = target_mod.weight.data
            if is_conv1d:
                in_f, out_f = raw_w.shape[0], raw_w.shape[1]
                W = raw_w.t().contiguous().float()  # [out_features, in_features]
            else:
                out_f, in_f = raw_w.shape[0], raw_w.shape[1]
                W = raw_w.float()  # [out_features, in_features]

            orig_layer_params = out_f * in_f
            total_orig_params += orig_layer_params
            max_r = max(1, min(out_f, in_f))
            current_r = max(1, min(self.rank, max_r))
            best_w_new: torch.Tensor | None = None
            final_err = 1.0

            # Iterative rank retry on error threshold exceedance
            for attempt in range(self.max_retries + 1):
                try:
                    w_cand, err = self._decompose_and_prune(W, current_r)
                    final_err = err
                    best_w_new = w_cand
                    if err <= self.max_error:
                        break
                except Exception as e:
                    logger.debug(
                        "ENLighten decomposition failed on %s (iter %d): %s",
                        name,
                        attempt,
                        e,
                    )
                    break

                if attempt < self.max_retries:
                    next_r = min(int(current_r * 1.5), max_r)
                    if next_r <= current_r:
                        break
                    logger.debug(
                        "ENLighten: Error %.4f > %.4f on %s. Escalating rank from %d to %d.",
                        err,
                        self.max_error,
                        name,
                        current_r,
                        next_r,
                    )
                    current_r = next_r

            self.layer_ranks[name] = current_r
            self.layer_sparsities[name] = 0.5 if self.structured_24 else self.sparsity
            self.layer_errors[name] = final_err

            if best_w_new is None or final_err > self.max_error:
                logger.warning(
                    "ENLighten reconstruction error %.4f on %s exceeded threshold %.4f.",
                    final_err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    logger.info(
                        "Skipping ENLighten compression for %s (fallback to original).",
                        name,
                    )
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    total_comp_params += orig_layer_params
                    continue

            # Assign compressed weight back in original orientation and dtype
            if is_conv1d:
                target_mod.weight.data.copy_(
                    best_w_new.t()
                    .contiguous()
                    .to(dtype=target_mod.weight.dtype, device=target_mod.weight.device)
                )
            else:
                target_mod.weight.data.copy_(
                    best_w_new.to(dtype=target_mod.weight.dtype, device=target_mod.weight.device)
                )

            # Preserve LoRA adapter trainability if layer is a LoRALinear wrapper
            if hasattr(module, "lora_A") and hasattr(module, "lora_B"):
                if hasattr(module.lora_A, "weight") and module.lora_A.weight is not None:
                    module.lora_A.weight.requires_grad = True
                if hasattr(module.lora_B, "weight") and module.lora_B.weight is not None:
                    module.lora_B.weight.requires_grad = True

            self.compressed_layers.append(name)
            # Effective compressed parameter representation:
            # Low-rank factor params (current_r * (in_f + out_f)) + non-zero sparse residual params (0.5 * in_f * out_f for 2:4)
            effective_sparse_frac = 0.5 if self.structured_24 else (1.0 - self.sparsity)
            comp_layer_params = current_r * (in_f + out_f) + int(
                effective_sparse_frac * orig_layer_params
            )
            total_comp_params += comp_layer_params

        self._compressed_params = self._count_params(self.model)
        effective_ratio = total_comp_params / total_orig_params if total_orig_params > 0 else 1.0
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )

        stats = self.get_compression_stats()
        stats.update(
            {
                "rank": self.rank,
                "sparsity": 0.5 if self.structured_24 else self.sparsity,
                "structured_24": self.structured_24,
                "max_error": self.max_error,
                "target_attention_only": self.target_attention_only,
                "fallback_to_original": self.fallback_to_original,
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "layer_ranks": self.layer_ranks,
                "layer_sparsities": self.layer_sparsities,
                "layer_errors": self.layer_errors,
                "avg_reconstruction_error": avg_err,
                "num_compressed": len(self.compressed_layers),
                "num_injected_attention_layers": len(
                    self.compressed_layers
                ),  # field used by framework CSV writer
                "num_fallback": len(self.fallback_layers),
                "effective_parameter_ratio": round(effective_ratio, 4),
            }
        )

        # Register stats directly on the model object so the framework can read
        # num_compressed_layers. ENLighten is a weight-rewriting compressor (no
        # module replacement), so the framework cannot discover compression via
        # module-type inspection — we expose it explicitly here.
        if not hasattr(self.model, "compression_stats") or self.model.compression_stats is None:
            self.model.compression_stats = {}
        self.model.compression_stats.update(
            {
                "num_compressed": len(self.compressed_layers),
                "num_compressed_layers": len(self.compressed_layers),
                "num_fallback_layers": len(self.fallback_layers),
                "compression_method": "enlighten",
                "effective_parameter_ratio": round(effective_ratio, 4),
                "avg_reconstruction_error": avg_err,
            }
        )

        logger.info(
            "ENLighten compression complete: %d compressed, %d fallback. "
            "Avg error: %.4f, Effective Param Ratio: %.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            effective_ratio,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: ENLighten Compressor: Low-Rank SVD + 2:4 Structured Sparsity. Decomposes weights into low-rank factors plus 2:4 structured sparse residuals, preserving Frobenius energy while reducing memory bandwidth and arithmetic complexity.

Production hardening features: - Robust LoRALinear wrapper and Conv1D support (with transpose and adapter preservation). - target_attention_only parameter: restricts compression to attention projections and freezes non-target layers. - Relative Frobenius reconstruction error validation with configurable max_error. - Automatic rank escalation retries when error threshold is exceeded. - Safe fallback to original module if error remains high or decomposition is unstable, with fallback freezing. - Fully vectorized 2:4 structured pruning eliminating Python loop overhead. - Critical-layer protection via should_compress_linear. - Real compression metrics recording (effective_parameter_ratio, num_compressed, num_fallback, etc.).

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Compress linear and Conv1D layers using low-rank SVD + 2:4 structured sparse residuals."""
    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    self.layer_ranks = {}
    self.layer_sparsities = {}
    self.layer_errors = {}

    total_orig_params = 0
    total_comp_params = 0

    for name, module in list(self.model.named_modules()):
        name_lower = name.lower()
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            ]
        ):
            continue

        if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
            continue

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")

        if not (is_linear or is_conv1d):
            continue

        if not hasattr(target_mod, "weight") or target_mod.weight is None:
            continue

        if not self._is_target_module(name, module):
            if self.target_attention_only:
                if not any(
                    part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]
                ):
                    for p in module.parameters():
                        p.requires_grad = False
                    self.skipped_layers.append(name)
            continue

        raw_w = target_mod.weight.data
        if is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()  # [out_features, in_features]
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float()  # [out_features, in_features]

        orig_layer_params = out_f * in_f
        total_orig_params += orig_layer_params
        max_r = max(1, min(out_f, in_f))
        current_r = max(1, min(self.rank, max_r))
        best_w_new: torch.Tensor | None = None
        final_err = 1.0

        # Iterative rank retry on error threshold exceedance
        for attempt in range(self.max_retries + 1):
            try:
                w_cand, err = self._decompose_and_prune(W, current_r)
                final_err = err
                best_w_new = w_cand
                if err <= self.max_error:
                    break
            except Exception as e:
                logger.debug(
                    "ENLighten decomposition failed on %s (iter %d): %s",
                    name,
                    attempt,
                    e,
                )
                break

            if attempt < self.max_retries:
                next_r = min(int(current_r * 1.5), max_r)
                if next_r <= current_r:
                    break
                logger.debug(
                    "ENLighten: Error %.4f > %.4f on %s. Escalating rank from %d to %d.",
                    err,
                    self.max_error,
                    name,
                    current_r,
                    next_r,
                )
                current_r = next_r

        self.layer_ranks[name] = current_r
        self.layer_sparsities[name] = 0.5 if self.structured_24 else self.sparsity
        self.layer_errors[name] = final_err

        if best_w_new is None or final_err > self.max_error:
            logger.warning(
                "ENLighten reconstruction error %.4f on %s exceeded threshold %.4f.",
                final_err,
                name,
                self.max_error,
            )
            if self.fallback_to_original:
                logger.info(
                    "Skipping ENLighten compression for %s (fallback to original).",
                    name,
                )
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_params += orig_layer_params
                continue

        # Assign compressed weight back in original orientation and dtype
        if is_conv1d:
            target_mod.weight.data.copy_(
                best_w_new.t()
                .contiguous()
                .to(dtype=target_mod.weight.dtype, device=target_mod.weight.device)
            )
        else:
            target_mod.weight.data.copy_(
                best_w_new.to(dtype=target_mod.weight.dtype, device=target_mod.weight.device)
            )

        # Preserve LoRA adapter trainability if layer is a LoRALinear wrapper
        if hasattr(module, "lora_A") and hasattr(module, "lora_B"):
            if hasattr(module.lora_A, "weight") and module.lora_A.weight is not None:
                module.lora_A.weight.requires_grad = True
            if hasattr(module.lora_B, "weight") and module.lora_B.weight is not None:
                module.lora_B.weight.requires_grad = True

        self.compressed_layers.append(name)
        # Effective compressed parameter representation:
        # Low-rank factor params (current_r * (in_f + out_f)) + non-zero sparse residual params (0.5 * in_f * out_f for 2:4)
        effective_sparse_frac = 0.5 if self.structured_24 else (1.0 - self.sparsity)
        comp_layer_params = current_r * (in_f + out_f) + int(
            effective_sparse_frac * orig_layer_params
        )
        total_comp_params += comp_layer_params

    self._compressed_params = self._count_params(self.model)
    effective_ratio = total_comp_params / total_orig_params if total_orig_params > 0 else 1.0
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )

    stats = self.get_compression_stats()
    stats.update(
        {
            "rank": self.rank,
            "sparsity": 0.5 if self.structured_24 else self.sparsity,
            "structured_24": self.structured_24,
            "max_error": self.max_error,
            "target_attention_only": self.target_attention_only,
            "fallback_to_original": self.fallback_to_original,
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "skipped_layers": self.skipped_layers,
            "skipped_layer_count": len(self.skipped_layers),
            "layer_ranks": self.layer_ranks,
            "layer_sparsities": self.layer_sparsities,
            "layer_errors": self.layer_errors,
            "avg_reconstruction_error": avg_err,
            "num_compressed": len(self.compressed_layers),
            "num_injected_attention_layers": len(
                self.compressed_layers
            ),  # field used by framework CSV writer
            "num_fallback": len(self.fallback_layers),
            "effective_parameter_ratio": round(effective_ratio, 4),
        }
    )

    # Register stats directly on the model object so the framework can read
    # num_compressed_layers. ENLighten is a weight-rewriting compressor (no
    # module replacement), so the framework cannot discover compression via
    # module-type inspection — we expose it explicitly here.
    if not hasattr(self.model, "compression_stats") or self.model.compression_stats is None:
        self.model.compression_stats = {}
    self.model.compression_stats.update(
        {
            "num_compressed": len(self.compressed_layers),
            "num_compressed_layers": len(self.compressed_layers),
            "num_fallback_layers": len(self.fallback_layers),
            "compression_method": "enlighten",
            "effective_parameter_ratio": round(effective_ratio, 4),
            "avg_reconstruction_error": avg_err,
        }
    )

    logger.info(
        "ENLighten compression complete: %d compressed, %d fallback. "
        "Avg error: %.4f, Effective Param Ratio: %.4f",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        effective_ratio,
    )
    return self.model, stats

Compress linear and Conv1D layers using low-rank SVD + 2:4 structured sparse residuals.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class ERCCompressor (model: nn.Module, allowed_error: float = 0.2, max_error: float | None = None, fallback_to_original: bool = True, target_attention_only: bool = False, rank: int | None = None, max_retries: int = 2, target_modules: list[str] | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("erc")
@COMPRESSOR_REGISTRY.register("erc_svd")
class ERCCompressor(BaseCompressor):
    is_custom_attention: bool = True
    """
    Error-Controlled SVD (ERC-SVD) Compressor.

    Finds the minimal rank r per layer such that
        ||W - W_r||_F / ||W||_F <= allowed_error.
    Factored symmetrically (U_r sqrt(S), sqrt(S) V_r^T) and replaces with SVDLinear.
    """

    def __init__(
        self,
        model: nn.Module,
        allowed_error: float = 0.20,
        max_error: float | None = None,
        fallback_to_original: bool = True,
        target_attention_only: bool = False,
        rank: int | None = None,
        max_retries: int = 2,
        target_modules: list[str] | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="erc", **kwargs)
        cfg = cfg or kwargs.get("comp_config")
        self.cfg = cfg

        c_allowed = _cfg_get(
            cfg, "erc_allowed_error", _cfg_get(cfg, "allowed_error", allowed_error)
        )
        effective_max = _cfg_get(cfg, "erc_max_error", _cfg_get(cfg, "max_error", max_error))
        if effective_max is None and (c_allowed is None or c_allowed <= 0.25):
            effective_max = adaptive_max_error(rank, 768, 768, "distilbert")
        if c_allowed is None or c_allowed <= 0.25:
            c_allowed = adaptive_max_error(rank, 768, 768, "distilbert")

        self.allowed_error = float(c_allowed)
        self.max_error = float(effective_max if effective_max is not None else self.allowed_error)
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "erc_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "erc_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        self.rank = _cfg_get(cfg, "erc_rank", _cfg_get(cfg, "rank", rank))
        if self.rank is not None:
            self.rank = int(self.rank)
        self.max_retries = max(0, int(max_retries))
        self.target_modules = target_modules

        # Pre-LayerNorm encoder detection (e.g., RoBERTa): pre-LN amplifies residual truncation noise.
        # Enforce stricter maximum error threshold (0.15) to prevent below-chance catastrophic collapse.
        model_name_or_cls = (
            getattr(getattr(model, "config", None), "_name_or_path", "") or model.__class__.__name__
        ).lower()
        self.is_pre_layernorm_encoder = "roberta" in model_name_or_cls
        if self.is_pre_layernorm_encoder and self.max_error > 0.30:
            logger.info(
                "ERC-SVD: Detected pre-LayerNorm encoder (%s). Capping max_error from %.2f to 0.30.",
                model_name_or_cls,
                self.max_error,
            )
            self.max_error = min(self.max_error, 0.30)

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_ranks: dict[str, int] = {}
        self.layer_errors: dict[str, float] = {}
        self._compressed_params: int = 0

    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
            return False
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
                "head",
                "output_layer",
            ]
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d):
            return False
        if isinstance(target_mod, SVDLinear):
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)

        if self.target_attention_only:
            last_part = name.split(".")[-1]
            is_attn = (last_part in _ATTN_NAMES) or any(
                k in name_lower for k in ("attn", "attention")
            )
            if not is_attn:
                return False

        return True

    # ------------------------------------------------------------------
    # ER-FIX-02: PEFT-compatible LoRA rank detection
    # ------------------------------------------------------------------
    @staticmethod
    def _detect_lora(module: nn.Module) -> tuple[int, float]:
        """Return (lora_rank, lora_alpha) from an existing LoRA wrapper if present."""
        if not hasattr(module, "lora_A"):
            return 0, 1.0
        rank = 0
        la = getattr(module, "lora_A", None)
        if la is not None:
            if hasattr(la, "weight") and la.weight is not None and la.weight.dim() >= 1:
                rank = int(la.weight.shape[0])
            elif isinstance(la, nn.Parameter) and la.dim() >= 1:
                rank = int(la.shape[0])
        alpha = 1.0
        for attr in ("lora_alpha", "alpha"):
            a = getattr(module, attr, None)
            if a is not None:
                try:
                    alpha = float(a)
                except (TypeError, ValueError):
                    pass
                break
        return rank, alpha

    # ------------------------------------------------------------------
    def compress_linear(
        self,
        module: nn.Module,
    ) -> tuple[SVDLinear | None, float, int]:
        """Decompose a single linear, Conv1D, or LoRALinear layer via error-controlled SVD."""
        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.data

        if is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()
        else:
            out_f, in_f = raw_w.shape
            W = raw_w.float()

        bound_r = (in_f * out_f) // (in_f + out_f)
        max_r = min(in_f, out_f, bound_r)
        if max_r < 16:
            return None, 1.0, 0

        device = raw_w.device
        dtype = raw_w.dtype
        norm_w = torch.norm(W, p="fro").item()

        try:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        except Exception as e:
            logger.warning("ERC-SVD SVD computation failed: %s", e)
            return None, 1.0, 0

        s_squared = S**2
        total_energy = s_squared.sum()
        if total_energy <= 1e-12:
            r_init = 1
        else:
            cum_energy = torch.cumsum(s_squared, dim=0) / total_energy
            target_retention = max(0.0, 1.0 - (self.allowed_error**2))
            idx = torch.where(cum_energy >= target_retention)[0]
            r_init = int(idx[0].item()) + 1 if len(idx) > 0 else max_r

        # ER-FIX-05: clamp r_init to max_r before use.
        r_init = min(r_init, max_r)
        if self.rank is not None:
            r_init = min(r_init, int(self.rank))
        current_r = max(1, min(r_init, max_r))

        best_svd_layer: SVDLinear | None = None
        best_err = 1.0
        chosen_r = current_r

        lora_r, lora_alpha = self._detect_lora(module)

        for attempt in range(self.max_retries + 1):
            r = max(1, min(current_r, max_r))
            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
            U_r = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
            V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(dtype=dtype, device=device)

            W_recon = torch.matmul(U_r.float(), V_r.float())
            if norm_w < 1e-12:
                err = 0.0
            else:
                err = float(torch.norm(W - W_recon, p="fro").item() / norm_w)

            best_err = err
            chosen_r = r

            svd_layer = SVDLinear(
                in_features=in_f,
                out_features=out_f,
                rank=r,
                bias=(
                    getattr(target_mod, "bias", None) is not None and target_mod.bias is not None
                ),
                dtype=dtype,
                device=device,
                lora_r=lora_r,
                lora_alpha=lora_alpha,
            )
            svd_layer.set_low_rank_factors(U_r, V_r)
            if (
                getattr(target_mod, "bias", None) is not None
                and target_mod.bias is not None
                and svd_layer.bias is not None
            ):
                svd_layer.bias.data.copy_(target_mod.bias.data)

            if (
                lora_r > 0
                and getattr(svd_layer, "lora_A", None) is not None
                and getattr(svd_layer, "lora_B", None) is not None
            ):
                try:
                    src_la = getattr(module, "lora_A", None)
                    src_lb = getattr(module, "lora_B", None)
                    la_w = getattr(src_la, "weight", src_la)
                    lb_w = getattr(src_lb, "weight", src_lb)
                    if hasattr(la_w, "data") and hasattr(svd_layer.lora_A, "data"):
                        svd_layer.lora_A.data.copy_(la_w.data)
                    if hasattr(lb_w, "data") and hasattr(svd_layer.lora_B, "data"):
                        svd_layer.lora_B.data.copy_(lb_w.data)
                except Exception as e:
                    logger.debug(
                        "ERC-SVD LoRA copy failed on %s: %s",
                        getattr(module, "_name", "?"),
                        e,
                    )

            best_svd_layer = svd_layer

            if err <= self.max_error or r >= max_r:
                break

            if attempt < self.max_retries:
                next_r = min(int(r * 1.5), max_r)
                if next_r <= r:
                    break
                current_r = next_r

        if best_svd_layer is not None:
            best_svd_layer.reconstruction_error = best_err
        return best_svd_layer, best_err, chosen_r

    # ------------------------------------------------------------------
    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Compress all target linear and Conv1D layers across the model."""
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_ranks = {}
        self.layer_errors = {}

        total_orig_params = 0
        total_comp_params = 0

        if self.target_attention_only:
            for name, module in self.model.named_modules():
                if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
                    continue
                target_mod = getattr(module, "base_layer", module)
                cls_name = target_mod.__class__.__name__.lower()
                is_linear = isinstance(target_mod, nn.Linear)
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                if (is_linear or is_conv1d) and not self._is_target_module(name, module):
                    for p in module.parameters():
                        p.requires_grad = False
                    self.skipped_layers.append(name)

        target_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            logger.warning("ERC-SVD: No compressible linear layers found.")
            return self.model, self.get_compression_stats()

        for name, module in target_layers:
            target_mod = getattr(module, "base_layer", module)
            raw_w = target_mod.weight.data
            orig_params = raw_w.numel()
            total_orig_params += orig_params

            svd_layer, err, r = self.compress_linear(module)
            self.layer_errors[name] = err
            self.layer_ranks[name] = r

            # ER-FIX-01: safe fallback when svd_layer is None.
            if svd_layer is None:
                logger.warning(
                    "ERC-SVD: could not compress %s (rank floor not met or SVD failed).",
                    name,
                )
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_params += orig_params
                continue

            if err > self.max_error:
                logger.warning(
                    "ERC-SVD: reconstruction error %.4f on %s exceeded threshold %.4f.",
                    err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    logger.info("ERC-SVD: fallback to original on %s.", name)
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    total_comp_params += orig_params
                    continue
                # fallback_to_original=False: proceed with high error (warned).

            replace_module(self.model, name, svd_layer)
            self.compressed_layers.append(name)
            comp_params = r * (svd_layer.in_features + svd_layer.out_features)
            total_comp_params += comp_params

        # ER-FIX-03/04: forward sanity check with attention_mask.
        forward_check_success = True
        try:
            with torch.no_grad():
                first_p = next(self.model.parameters(), None)
                dev = first_p.device if first_p is not None else torch.device("cpu")
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                attn = torch.ones_like(dummy)
                try:
                    self.model(input_ids=dummy, attention_mask=attn)
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception as e_inner:
                        logger.debug("ERC-SVD forward check (input_ids) failed: %s", e_inner)
                        first_linear = next(
                            (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                            None,
                        )
                        in_dim = getattr(first_linear, "in_features", 8)
                        self.model(
                            torch.randn(
                                1,
                                8,
                                in_dim,
                                device=dev,
                                dtype=(first_p.dtype if first_p is not None else torch.float32),
                            )
                        )
        except Exception as e:
            logger.debug("ERC-SVD forward sanity check failed: %s", e)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        effective_ratio = (total_comp_params / total_orig_params) if total_orig_params > 0 else 1.0
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "layer_ranks": self.layer_ranks,
                "layer_errors": self.layer_errors,
                "allowed_error": self.allowed_error,
                "max_error": self.max_error,
                "target_attention_only": self.target_attention_only,
                "fallback_to_original": self.fallback_to_original,
                "avg_reconstruction_error": avg_err,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "effective_parameter_ratio": round(effective_ratio, 4),
                "forward_check_success": forward_check_success,
                "method": "erc_svd",
            }
        )
        logger.info(
            "Compressor erc finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: Error-Controlled SVD (ERC-SVD) Compressor.

Finds the minimal rank r per layer such that ||W - W_r||_F / ||W||_F <= allowed_error. Factored symmetrically (U_r sqrt(S), sqrt(S) V_r^T) and replaces with SVDLinear.

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Compress all target linear and Conv1D layers across the model."""
    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    self.layer_ranks = {}
    self.layer_errors = {}

    total_orig_params = 0
    total_comp_params = 0

    if self.target_attention_only:
        for name, module in self.model.named_modules():
            if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
                continue
            target_mod = getattr(module, "base_layer", module)
            cls_name = target_mod.__class__.__name__.lower()
            is_linear = isinstance(target_mod, nn.Linear)
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            if (is_linear or is_conv1d) and not self._is_target_module(name, module):
                for p in module.parameters():
                    p.requires_grad = False
                self.skipped_layers.append(name)

    target_layers = [
        (name, mod)
        for name, mod in self.model.named_modules()
        if self._is_target_module(name, mod)
    ]

    if not target_layers:
        logger.warning("ERC-SVD: No compressible linear layers found.")
        return self.model, self.get_compression_stats()

    for name, module in target_layers:
        target_mod = getattr(module, "base_layer", module)
        raw_w = target_mod.weight.data
        orig_params = raw_w.numel()
        total_orig_params += orig_params

        svd_layer, err, r = self.compress_linear(module)
        self.layer_errors[name] = err
        self.layer_ranks[name] = r

        # ER-FIX-01: safe fallback when svd_layer is None.
        if svd_layer is None:
            logger.warning(
                "ERC-SVD: could not compress %s (rank floor not met or SVD failed).",
                name,
            )
            self.fallback_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False
            total_comp_params += orig_params
            continue

        if err > self.max_error:
            logger.warning(
                "ERC-SVD: reconstruction error %.4f on %s exceeded threshold %.4f.",
                err,
                name,
                self.max_error,
            )
            if self.fallback_to_original:
                logger.info("ERC-SVD: fallback to original on %s.", name)
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_params += orig_params
                continue
            # fallback_to_original=False: proceed with high error (warned).

        replace_module(self.model, name, svd_layer)
        self.compressed_layers.append(name)
        comp_params = r * (svd_layer.in_features + svd_layer.out_features)
        total_comp_params += comp_params

    # ER-FIX-03/04: forward sanity check with attention_mask.
    forward_check_success = True
    try:
        with torch.no_grad():
            first_p = next(self.model.parameters(), None)
            dev = first_p.device if first_p is not None else torch.device("cpu")
            dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
            attn = torch.ones_like(dummy)
            try:
                self.model(input_ids=dummy, attention_mask=attn)
            except Exception:
                try:
                    self.model(dummy)
                except Exception as e_inner:
                    logger.debug("ERC-SVD forward check (input_ids) failed: %s", e_inner)
                    first_linear = next(
                        (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                        None,
                    )
                    in_dim = getattr(first_linear, "in_features", 8)
                    self.model(
                        torch.randn(
                            1,
                            8,
                            in_dim,
                            device=dev,
                            dtype=(first_p.dtype if first_p is not None else torch.float32),
                        )
                    )
    except Exception as e:
        logger.debug("ERC-SVD forward sanity check failed: %s", e)
        forward_check_success = False

    try:
        self._compressed_params = self._count_params(self.model)
    except Exception:
        self._compressed_params = sum(p.numel() for p in self.model.parameters())

    effective_ratio = (total_comp_params / total_orig_params) if total_orig_params > 0 else 1.0
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )
    try:
        stats = self.get_compression_stats()
    except Exception:
        stats = {}
    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "skipped_layers": self.skipped_layers,
            "skipped_layer_count": len(self.skipped_layers),
            "layer_ranks": self.layer_ranks,
            "layer_errors": self.layer_errors,
            "allowed_error": self.allowed_error,
            "max_error": self.max_error,
            "target_attention_only": self.target_attention_only,
            "fallback_to_original": self.fallback_to_original,
            "avg_reconstruction_error": avg_err,
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "num_skipped": len(self.skipped_layers),
            "effective_parameter_ratio": round(effective_ratio, 4),
            "forward_check_success": forward_check_success,
            "method": "erc_svd",
        }
    )
    logger.info(
        "Compressor erc finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        self.max_error,
    )
    return self.model, stats

Compress all target linear and Conv1D layers across the model.

def compress_linear(self, module: nn.Module) ‑> tuple[SVDLinear | None, float, int]

Expand source code

def compress_linear(
    self,
    module: nn.Module,
) -> tuple[SVDLinear | None, float, int]:
    """Decompose a single linear, Conv1D, or LoRALinear layer via error-controlled SVD."""
    target_mod = getattr(module, "base_layer", module)
    cls_name = target_mod.__class__.__name__.lower()
    is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
    raw_w = target_mod.weight.data

    if is_conv1d:
        in_f, out_f = raw_w.shape[0], raw_w.shape[1]
        W = raw_w.t().contiguous().float()
    else:
        out_f, in_f = raw_w.shape
        W = raw_w.float()

    bound_r = (in_f * out_f) // (in_f + out_f)
    max_r = min(in_f, out_f, bound_r)
    if max_r < 16:
        return None, 1.0, 0

    device = raw_w.device
    dtype = raw_w.dtype
    norm_w = torch.norm(W, p="fro").item()

    try:
        U, S, Vt = torch.linalg.svd(W, full_matrices=False)
    except Exception as e:
        logger.warning("ERC-SVD SVD computation failed: %s", e)
        return None, 1.0, 0

    s_squared = S**2
    total_energy = s_squared.sum()
    if total_energy <= 1e-12:
        r_init = 1
    else:
        cum_energy = torch.cumsum(s_squared, dim=0) / total_energy
        target_retention = max(0.0, 1.0 - (self.allowed_error**2))
        idx = torch.where(cum_energy >= target_retention)[0]
        r_init = int(idx[0].item()) + 1 if len(idx) > 0 else max_r

    # ER-FIX-05: clamp r_init to max_r before use.
    r_init = min(r_init, max_r)
    if self.rank is not None:
        r_init = min(r_init, int(self.rank))
    current_r = max(1, min(r_init, max_r))

    best_svd_layer: SVDLinear | None = None
    best_err = 1.0
    chosen_r = current_r

    lora_r, lora_alpha = self._detect_lora(module)

    for attempt in range(self.max_retries + 1):
        r = max(1, min(current_r, max_r))
        sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
        U_r = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
        V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(dtype=dtype, device=device)

        W_recon = torch.matmul(U_r.float(), V_r.float())
        if norm_w < 1e-12:
            err = 0.0
        else:
            err = float(torch.norm(W - W_recon, p="fro").item() / norm_w)

        best_err = err
        chosen_r = r

        svd_layer = SVDLinear(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=(
                getattr(target_mod, "bias", None) is not None and target_mod.bias is not None
            ),
            dtype=dtype,
            device=device,
            lora_r=lora_r,
            lora_alpha=lora_alpha,
        )
        svd_layer.set_low_rank_factors(U_r, V_r)
        if (
            getattr(target_mod, "bias", None) is not None
            and target_mod.bias is not None
            and svd_layer.bias is not None
        ):
            svd_layer.bias.data.copy_(target_mod.bias.data)

        if (
            lora_r > 0
            and getattr(svd_layer, "lora_A", None) is not None
            and getattr(svd_layer, "lora_B", None) is not None
        ):
            try:
                src_la = getattr(module, "lora_A", None)
                src_lb = getattr(module, "lora_B", None)
                la_w = getattr(src_la, "weight", src_la)
                lb_w = getattr(src_lb, "weight", src_lb)
                if hasattr(la_w, "data") and hasattr(svd_layer.lora_A, "data"):
                    svd_layer.lora_A.data.copy_(la_w.data)
                if hasattr(lb_w, "data") and hasattr(svd_layer.lora_B, "data"):
                    svd_layer.lora_B.data.copy_(lb_w.data)
            except Exception as e:
                logger.debug(
                    "ERC-SVD LoRA copy failed on %s: %s",
                    getattr(module, "_name", "?"),
                    e,
                )

        best_svd_layer = svd_layer

        if err <= self.max_error or r >= max_r:
            break

        if attempt < self.max_retries:
            next_r = min(int(r * 1.5), max_r)
            if next_r <= r:
                break
            current_r = next_r

    if best_svd_layer is not None:
        best_svd_layer.reconstruction_error = best_err
    return best_svd_layer, best_err, chosen_r

Decompose a single linear, Conv1D, or LoRALinear layer via error-controlled SVD.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class EntropyBudgetCompressor (model: nn.Module, calibration_loader: Any | None = None, total_budget_ratio: float = 0.7, rank: int | None = None, eps: float = 1e-06, max_error: float = 0.35, target_attention_only: bool = False, fallback_to_original: bool = True, target_modules: list[str] | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("entropy_budget")
@COMPRESSOR_REGISTRY.register("entropy_budget_allocation")
class EntropyBudgetCompressor(BaseCompressor):
    is_custom_attention: bool = True
    """
    Entropy Budget Low-Rank SVD Compressor.

    Computes activation covariance and whiten weights, calculates per-component
    singular entropy contributions, allocates low-rank budgets globally across layers,
    and replaces dense nn.Linear layers with factorized SVDLinear modules with symmetric
    singular value splitting (U * sqrt(Sigma), sqrt(Sigma) * V^T).

    Hardened for:
    - LoRALinear wrappers & Conv1D support.
    - target_attention_only filtering.
    - Relative Frobenius reconstruction error validation & fallback.
    - Freezing of fallback and skipped layers.
    - Real compression metrics recording.
    """

    def __init__(
        self,
        model: nn.Module,
        calibration_loader: Any | None = None,
        total_budget_ratio: float = 0.7,
        rank: int | None = None,
        eps: float = 1e-6,
        max_error: float = 0.35,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        target_modules: list[str] | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="entropy_budget", **kwargs)
        cfg = cfg or kwargs.get("comp_config")
        self.cfg = cfg

        self.calibration_loader = calibration_loader or _cfg_get(cfg, "calibration_loader", None)
        self.total_budget_ratio = float(
            _cfg_get(
                cfg,
                "entropy_budget_total_budget_ratio",
                _cfg_get(
                    cfg,
                    "total_budget_ratio",
                    _cfg_get(cfg, "budget_ratio", total_budget_ratio),
                ),
            )
        )
        self.rank = _cfg_get(
            cfg,
            "entropy_budget_rank",
            _cfg_get(cfg, "rank", _cfg_get(cfg, "target_rank", rank)),
        )
        if self.rank is not None:
            self.rank = int(self.rank)
        self.eps = float(eps)
        c_max_error = _cfg_get(
            cfg, "entropy_budget_max_error", _cfg_get(cfg, "max_error", max_error)
        )
        if c_max_error is None or c_max_error == 0.35:
            c_max_error = adaptive_max_error(self.rank or 64, 768, 768, "distilbert")

        self.max_error = float(c_max_error)
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "entropy_budget_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "entropy_budget_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        self.target_modules = target_modules

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_ranks: dict[str, int] = {}
        self.layer_errors: dict[str, float] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
            return False
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
                "head",
                "output_layer",
            ]
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d):
            return False
        if isinstance(target_mod, SVDLinear):
            return False

        if not should_compress_linear(self.model, name):
            return False

        if hasattr(target_mod, "in_features") and hasattr(target_mod, "out_features"):
            if min(target_mod.in_features, target_mod.out_features) < 16:
                return False
        elif hasattr(target_mod, "weight") and target_mod.weight is not None:
            if min(target_mod.weight.shape) < 16:
                return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)

        if self.target_attention_only:
            last_part = name.split(".")[-1]
            is_attn = (last_part in _ATTN_NAMES) or any(
                k in name_lower for k in ("attn", "attention")
            )
            if not is_attn:
                return False

        return True

    def collect_covariances(self) -> dict[str, torch.Tensor]:
        """Collect empirical input activation covariances with double-precision accumulation."""
        if self.calibration_loader is None:
            return {}

        activations: dict[str, torch.Tensor] = {}
        sample_counts: dict[str, int] = {}
        hooks = []

        def get_hook(name: str):
            def hook(module: nn.Module, inp: tuple[Any, ...], outp: Any) -> None:
                if not inp or inp[0] is None:
                    return
                x = inp[0].detach()
                if x.ndim >= 2:
                    x_flat = x.reshape(-1, x.shape[-1]).float()
                    cov = x_flat.T @ x_flat
                    n = x_flat.shape[0]
                    if name not in activations:
                        activations[name] = cov
                        sample_counts[name] = n
                    else:
                        activations[name] += cov
                        sample_counts[name] += n

            return hook

        target_layers = [
            (name, getattr(mod, "base_layer", mod))
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            return {}

        for name, mod in target_layers:
            hooks.append(mod.register_forward_hook(get_hook(name)))

        device = next(self.model.parameters()).device
        self.model.eval()

        try:
            with torch.no_grad():
                for batch in self.calibration_loader:
                    if isinstance(batch, dict):
                        batch = {
                            k: v.to(device) if isinstance(v, torch.Tensor) else v
                            for k, v in batch.items()
                        }
                        self.model(**batch)
                    elif isinstance(batch, (list, tuple)):
                        batch = [v.to(device) if isinstance(v, torch.Tensor) else v for v in batch]
                        self.model(*batch)
                    elif isinstance(batch, torch.Tensor):
                        self.model(batch.to(device))
        except Exception as e:
            logger.warning(
                "EntropyBudget: Calibration forward pass failed (%s). Falling back to weight spectra.",
                e,
            )
        finally:
            for h in hooks:
                h.remove()

        normalized_covs: dict[str, torch.Tensor] = {}
        for name, cov in activations.items():
            cnt = max(1, sample_counts.get(name, 1))
            normalized_covs[name] = cov / cnt

        return normalized_covs

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Compress model using whitened entropy budget allocation and covariance-aligned singular vectors."""
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_ranks = {}
        self.layer_errors = {}

        total_orig_params = 0
        total_comp_params = 0

        # Freeze skipped layers if target_attention_only is enabled
        if self.target_attention_only:
            for name, module in self.model.named_modules():
                if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
                    continue
                target_mod = getattr(module, "base_layer", module)
                cls_name = target_mod.__class__.__name__.lower()
                is_linear = isinstance(target_mod, nn.Linear)
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                if (is_linear or is_conv1d) and not self._is_target_module(name, module):
                    for p in module.parameters():
                        p.requires_grad = False
                    self.skipped_layers.append(name)

        target_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            logger.warning("EntropyBudget: No compressible linear layers found.")
            return self.model, self.get_compression_stats()

        covs = self.collect_covariances()
        residual_candidates = []
        layer_decomp_info = {}

        # Step 1: Compute covariance-whitened singular decomposition and singular entropy per layer
        for name, module in target_layers:
            target_mod = getattr(module, "base_layer", module)
            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            raw_w = target_mod.weight.data

            if is_conv1d:
                in_f, out_f = raw_w.shape[0], raw_w.shape[1]
                W = raw_w.t().contiguous().float()
            else:
                out_f, in_f = raw_w.shape[0], raw_w.shape[1]
                W = raw_w.float()

            orig_layer_params = out_f * in_f
            total_orig_params += orig_layer_params
            max_r = min(out_f, in_f)
            r_floor = max(1, min(16, max_r))

            cov_inv_sqrt = None
            if name in covs and covs[name].shape == (in_f, in_f):
                cov = covs[name].to(device=W.device, dtype=torch.float32)
                try:
                    eigvals, eigvecs = torch.linalg.eigh(
                        cov + 1e-4 * torch.eye(in_f, device=W.device)
                    )
                    eigvals_clamped = torch.clamp(eigvals, min=1e-4)
                    cov_sqrt = eigvecs @ torch.diag(torch.sqrt(eigvals_clamped)) @ eigvecs.T
                    cov_inv_sqrt = (
                        eigvecs @ torch.diag(1.0 / torch.sqrt(eigvals_clamped)) @ eigvecs.T
                    )
                    M = W @ cov_sqrt
                except Exception:
                    M = W
                    cov_inv_sqrt = None
            else:
                M = W

            U, S, Vt = torch.linalg.svd(M, full_matrices=False)
            layer_decomp_info[name] = (
                U,
                S,
                Vt,
                cov_inv_sqrt,
                max_r,
                r_floor,
                W,
                is_conv1d,
                target_mod,
            )

            total_s = S.sum() + self.eps
            p = torch.clamp(S / total_s, min=self.eps, max=1.0)
            # B18: Shannon entropy in bits using log2
            entropy = -torch.sum(p * torch.log2(p)).item()
            entropy = max(entropy, 1e-4)

            for i in range(r_floor, max_r):
                s_val = S[i].item()
                score = (s_val**2) / entropy
                residual_candidates.append((name, i, s_val, score))

        # Step 2: Allocate ranks
        if self.rank is not None:
            allocated_ranks = {
                name: max(
                    layer_decomp_info[name][5],
                    min(self.rank, layer_decomp_info[name][4]),
                )
                for name, _ in target_layers
            }
        else:
            allocated_ranks = {name: layer_decomp_info[name][5] for name, _ in target_layers}
            base_allocated = sum(allocated_ranks.values())
            total_possible = sum(layer_decomp_info[name][4] for name, _ in target_layers)
            target_total = max(
                base_allocated,
                min(total_possible, int(total_possible * self.total_budget_ratio)),
            )
            residual_budget = target_total - base_allocated

            residual_candidates.sort(key=lambda x: x[3], reverse=True)
            for c in residual_candidates[:residual_budget]:
                allocated_ranks[c[0]] += 1

        # Step 3: Factorized replacement with validation & fallback
        forward_check_success = True

        for name, module in target_layers:
            target_mod = getattr(module, "base_layer", module)
            r = allocated_ranks[name]
            U, S, Vt, cov_inv_sqrt, max_r, _, W, is_conv1d, _ = layer_decomp_info[name]
            out_f, in_f = W.shape
            orig_layer_params = out_f * in_f
            device = target_mod.weight.device
            dtype = target_mod.weight.dtype
            r = max(1, min(r, max_r))

            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
            U_r = U[:, :r] * sqrt_S.unsqueeze(0)
            if cov_inv_sqrt is not None:
                V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]) @ cov_inv_sqrt
                norm_u = torch.norm(U_r, p="fro").clamp_min(1e-8)
                norm_v = torch.norm(V_r, p="fro").clamp_min(1e-8)
                scale = torch.sqrt(norm_u / norm_v)
                U_r = U_r / scale
                V_r = V_r * scale
            else:
                V_r = sqrt_S.unsqueeze(1) * Vt[:r, :]

            # Compute reconstructed weight & Frobenius error
            W_recon = torch.matmul(U_r.float(), V_r.float())
            norm_w = torch.norm(W, p="fro").item()
            err = (
                float(torch.norm(W - W_recon, p="fro").item() / max(norm_w, 1e-12))
                if norm_w > 1e-12
                else 0.0
            )
            self.layer_errors[name] = err
            self.layer_ranks[name] = r

            if err > self.max_error and self.fallback_to_original:
                logger.warning(
                    "EntropyBudget error %.4f on %s > max_error %.4f; falling back.",
                    err,
                    name,
                    self.max_error,
                )
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_params += orig_layer_params
                continue

            U_r = U_r.to(dtype=dtype, device=device)
            V_r = V_r.to(dtype=dtype, device=device)

            lora_r = getattr(module, "r", 0) if hasattr(module, "lora_A") else 0
            lora_alpha = getattr(module, "lora_alpha", 1.0) if hasattr(module, "lora_A") else 1.0

            svd_layer = SVDLinear(
                in_features=in_f,
                out_features=out_f,
                rank=r,
                bias=getattr(target_mod, "bias", None) is not None and target_mod.bias is not None,
                dtype=dtype,
                device=device,
                lora_r=lora_r,
                lora_alpha=lora_alpha,
            )
            svd_layer.set_low_rank_factors(U_r, V_r)
            if (
                getattr(target_mod, "bias", None) is not None
                and target_mod.bias is not None
                and svd_layer.bias is not None
            ):
                svd_layer.bias.data.copy_(target_mod.bias.data)

            # If module was LoRALinear wrapper, copy existing trained adapter weights
            if (
                hasattr(module, "lora_A")
                and hasattr(module, "lora_B")
                and svd_layer.lora_A is not None
                and svd_layer.lora_B is not None
            ):
                if hasattr(module.lora_A, "weight") and module.lora_A.weight is not None:
                    svd_layer.lora_A.data.copy_(module.lora_A.weight.data)
                elif isinstance(module.lora_A, nn.Parameter):
                    svd_layer.lora_A.data.copy_(module.lora_A.data)
                if hasattr(module.lora_B, "weight") and module.lora_B.weight is not None:
                    svd_layer.lora_B.data.copy_(module.lora_B.weight.data)
                elif isinstance(module.lora_B, nn.Parameter):
                    svd_layer.lora_B.data.copy_(module.lora_B.data)

            replace_module(self.model, name, svd_layer)
            self.compressed_layers.append(name)
            comp_layer_params = r * (in_f + out_f)
            total_comp_params += comp_layer_params

        # Sanity check forward pass
        try:
            with torch.no_grad():
                first_p = next(self.model.parameters(), None)
                dev = first_p.device if first_p is not None else torch.device("cpu")
                dummy = torch.randint(0, 100, (1, 8), device=dev)
                try:
                    self.model(dummy)
                except Exception as e:
                    logger.warning(f"copy_weights failed: {e}")
        except Exception:
            forward_check_success = False

        effective_ratio = total_comp_params / total_orig_params if total_orig_params > 0 else 1.0
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )

        stats = self.get_compression_stats()
        stats.update(
            {
                "allocated_ranks": allocated_ranks,
                "layer_ranks": self.layer_ranks,
                "layer_errors": self.layer_errors,
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "avg_reconstruction_error": avg_err,
                "effective_parameter_ratio": round(effective_ratio, 4),
                "total_budget_ratio": self.total_budget_ratio,
                "max_error": self.max_error,
                "target_attention_only": self.target_attention_only,
                "forward_check_success": forward_check_success,
            }
        )
        logger.info(
            "Compressor entropy_budget finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: Entropy Budget Low-Rank SVD Compressor.

Computes activation covariance and whiten weights, calculates per-component singular entropy contributions, allocates low-rank budgets globally across layers, and replaces dense nn.Linear layers with factorized SVDLinear modules with symmetric singular value splitting (U * sqrt(Sigma), sqrt(Sigma) * V^T).

Hardened for: - LoRALinear wrappers & Conv1D support. - target_attention_only filtering. - Relative Frobenius reconstruction error validation & fallback. - Freezing of fallback and skipped layers. - Real compression metrics recording.

Methods

def collect_covariances(self) ‑> dict[str, torch.Tensor]

Expand source code

def collect_covariances(self) -> dict[str, torch.Tensor]:
    """Collect empirical input activation covariances with double-precision accumulation."""
    if self.calibration_loader is None:
        return {}

    activations: dict[str, torch.Tensor] = {}
    sample_counts: dict[str, int] = {}
    hooks = []

    def get_hook(name: str):
        def hook(module: nn.Module, inp: tuple[Any, ...], outp: Any) -> None:
            if not inp or inp[0] is None:
                return
            x = inp[0].detach()
            if x.ndim >= 2:
                x_flat = x.reshape(-1, x.shape[-1]).float()
                cov = x_flat.T @ x_flat
                n = x_flat.shape[0]
                if name not in activations:
                    activations[name] = cov
                    sample_counts[name] = n
                else:
                    activations[name] += cov
                    sample_counts[name] += n

        return hook

    target_layers = [
        (name, getattr(mod, "base_layer", mod))
        for name, mod in self.model.named_modules()
        if self._is_target_module(name, mod)
    ]

    if not target_layers:
        return {}

    for name, mod in target_layers:
        hooks.append(mod.register_forward_hook(get_hook(name)))

    device = next(self.model.parameters()).device
    self.model.eval()

    try:
        with torch.no_grad():
            for batch in self.calibration_loader:
                if isinstance(batch, dict):
                    batch = {
                        k: v.to(device) if isinstance(v, torch.Tensor) else v
                        for k, v in batch.items()
                    }
                    self.model(**batch)
                elif isinstance(batch, (list, tuple)):
                    batch = [v.to(device) if isinstance(v, torch.Tensor) else v for v in batch]
                    self.model(*batch)
                elif isinstance(batch, torch.Tensor):
                    self.model(batch.to(device))
    except Exception as e:
        logger.warning(
            "EntropyBudget: Calibration forward pass failed (%s). Falling back to weight spectra.",
            e,
        )
    finally:
        for h in hooks:
            h.remove()

    normalized_covs: dict[str, torch.Tensor] = {}
    for name, cov in activations.items():
        cnt = max(1, sample_counts.get(name, 1))
        normalized_covs[name] = cov / cnt

    return normalized_covs

Collect empirical input activation covariances with double-precision accumulation.

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Compress model using whitened entropy budget allocation and covariance-aligned singular vectors."""
    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    self.layer_ranks = {}
    self.layer_errors = {}

    total_orig_params = 0
    total_comp_params = 0

    # Freeze skipped layers if target_attention_only is enabled
    if self.target_attention_only:
        for name, module in self.model.named_modules():
            if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
                continue
            target_mod = getattr(module, "base_layer", module)
            cls_name = target_mod.__class__.__name__.lower()
            is_linear = isinstance(target_mod, nn.Linear)
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            if (is_linear or is_conv1d) and not self._is_target_module(name, module):
                for p in module.parameters():
                    p.requires_grad = False
                self.skipped_layers.append(name)

    target_layers = [
        (name, mod)
        for name, mod in self.model.named_modules()
        if self._is_target_module(name, mod)
    ]

    if not target_layers:
        logger.warning("EntropyBudget: No compressible linear layers found.")
        return self.model, self.get_compression_stats()

    covs = self.collect_covariances()
    residual_candidates = []
    layer_decomp_info = {}

    # Step 1: Compute covariance-whitened singular decomposition and singular entropy per layer
    for name, module in target_layers:
        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.data

        if is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float()

        orig_layer_params = out_f * in_f
        total_orig_params += orig_layer_params
        max_r = min(out_f, in_f)
        r_floor = max(1, min(16, max_r))

        cov_inv_sqrt = None
        if name in covs and covs[name].shape == (in_f, in_f):
            cov = covs[name].to(device=W.device, dtype=torch.float32)
            try:
                eigvals, eigvecs = torch.linalg.eigh(
                    cov + 1e-4 * torch.eye(in_f, device=W.device)
                )
                eigvals_clamped = torch.clamp(eigvals, min=1e-4)
                cov_sqrt = eigvecs @ torch.diag(torch.sqrt(eigvals_clamped)) @ eigvecs.T
                cov_inv_sqrt = (
                    eigvecs @ torch.diag(1.0 / torch.sqrt(eigvals_clamped)) @ eigvecs.T
                )
                M = W @ cov_sqrt
            except Exception:
                M = W
                cov_inv_sqrt = None
        else:
            M = W

        U, S, Vt = torch.linalg.svd(M, full_matrices=False)
        layer_decomp_info[name] = (
            U,
            S,
            Vt,
            cov_inv_sqrt,
            max_r,
            r_floor,
            W,
            is_conv1d,
            target_mod,
        )

        total_s = S.sum() + self.eps
        p = torch.clamp(S / total_s, min=self.eps, max=1.0)
        # B18: Shannon entropy in bits using log2
        entropy = -torch.sum(p * torch.log2(p)).item()
        entropy = max(entropy, 1e-4)

        for i in range(r_floor, max_r):
            s_val = S[i].item()
            score = (s_val**2) / entropy
            residual_candidates.append((name, i, s_val, score))

    # Step 2: Allocate ranks
    if self.rank is not None:
        allocated_ranks = {
            name: max(
                layer_decomp_info[name][5],
                min(self.rank, layer_decomp_info[name][4]),
            )
            for name, _ in target_layers
        }
    else:
        allocated_ranks = {name: layer_decomp_info[name][5] for name, _ in target_layers}
        base_allocated = sum(allocated_ranks.values())
        total_possible = sum(layer_decomp_info[name][4] for name, _ in target_layers)
        target_total = max(
            base_allocated,
            min(total_possible, int(total_possible * self.total_budget_ratio)),
        )
        residual_budget = target_total - base_allocated

        residual_candidates.sort(key=lambda x: x[3], reverse=True)
        for c in residual_candidates[:residual_budget]:
            allocated_ranks[c[0]] += 1

    # Step 3: Factorized replacement with validation & fallback
    forward_check_success = True

    for name, module in target_layers:
        target_mod = getattr(module, "base_layer", module)
        r = allocated_ranks[name]
        U, S, Vt, cov_inv_sqrt, max_r, _, W, is_conv1d, _ = layer_decomp_info[name]
        out_f, in_f = W.shape
        orig_layer_params = out_f * in_f
        device = target_mod.weight.device
        dtype = target_mod.weight.dtype
        r = max(1, min(r, max_r))

        sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
        U_r = U[:, :r] * sqrt_S.unsqueeze(0)
        if cov_inv_sqrt is not None:
            V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]) @ cov_inv_sqrt
            norm_u = torch.norm(U_r, p="fro").clamp_min(1e-8)
            norm_v = torch.norm(V_r, p="fro").clamp_min(1e-8)
            scale = torch.sqrt(norm_u / norm_v)
            U_r = U_r / scale
            V_r = V_r * scale
        else:
            V_r = sqrt_S.unsqueeze(1) * Vt[:r, :]

        # Compute reconstructed weight & Frobenius error
        W_recon = torch.matmul(U_r.float(), V_r.float())
        norm_w = torch.norm(W, p="fro").item()
        err = (
            float(torch.norm(W - W_recon, p="fro").item() / max(norm_w, 1e-12))
            if norm_w > 1e-12
            else 0.0
        )
        self.layer_errors[name] = err
        self.layer_ranks[name] = r

        if err > self.max_error and self.fallback_to_original:
            logger.warning(
                "EntropyBudget error %.4f on %s > max_error %.4f; falling back.",
                err,
                name,
                self.max_error,
            )
            self.fallback_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False
            total_comp_params += orig_layer_params
            continue

        U_r = U_r.to(dtype=dtype, device=device)
        V_r = V_r.to(dtype=dtype, device=device)

        lora_r = getattr(module, "r", 0) if hasattr(module, "lora_A") else 0
        lora_alpha = getattr(module, "lora_alpha", 1.0) if hasattr(module, "lora_A") else 1.0

        svd_layer = SVDLinear(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=getattr(target_mod, "bias", None) is not None and target_mod.bias is not None,
            dtype=dtype,
            device=device,
            lora_r=lora_r,
            lora_alpha=lora_alpha,
        )
        svd_layer.set_low_rank_factors(U_r, V_r)
        if (
            getattr(target_mod, "bias", None) is not None
            and target_mod.bias is not None
            and svd_layer.bias is not None
        ):
            svd_layer.bias.data.copy_(target_mod.bias.data)

        # If module was LoRALinear wrapper, copy existing trained adapter weights
        if (
            hasattr(module, "lora_A")
            and hasattr(module, "lora_B")
            and svd_layer.lora_A is not None
            and svd_layer.lora_B is not None
        ):
            if hasattr(module.lora_A, "weight") and module.lora_A.weight is not None:
                svd_layer.lora_A.data.copy_(module.lora_A.weight.data)
            elif isinstance(module.lora_A, nn.Parameter):
                svd_layer.lora_A.data.copy_(module.lora_A.data)
            if hasattr(module.lora_B, "weight") and module.lora_B.weight is not None:
                svd_layer.lora_B.data.copy_(module.lora_B.weight.data)
            elif isinstance(module.lora_B, nn.Parameter):
                svd_layer.lora_B.data.copy_(module.lora_B.data)

        replace_module(self.model, name, svd_layer)
        self.compressed_layers.append(name)
        comp_layer_params = r * (in_f + out_f)
        total_comp_params += comp_layer_params

    # Sanity check forward pass
    try:
        with torch.no_grad():
            first_p = next(self.model.parameters(), None)
            dev = first_p.device if first_p is not None else torch.device("cpu")
            dummy = torch.randint(0, 100, (1, 8), device=dev)
            try:
                self.model(dummy)
            except Exception as e:
                logger.warning(f"copy_weights failed: {e}")
    except Exception:
        forward_check_success = False

    effective_ratio = total_comp_params / total_orig_params if total_orig_params > 0 else 1.0
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )

    stats = self.get_compression_stats()
    stats.update(
        {
            "allocated_ranks": allocated_ranks,
            "layer_ranks": self.layer_ranks,
            "layer_errors": self.layer_errors,
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "skipped_layers": self.skipped_layers,
            "skipped_layer_count": len(self.skipped_layers),
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "num_skipped": len(self.skipped_layers),
            "avg_reconstruction_error": avg_err,
            "effective_parameter_ratio": round(effective_ratio, 4),
            "total_budget_ratio": self.total_budget_ratio,
            "max_error": self.max_error,
            "target_attention_only": self.target_attention_only,
            "forward_check_success": forward_check_success,
        }
    )
    logger.info(
        "Compressor entropy_budget finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        self.max_error,
    )
    return self.model, stats

Compress model using whitened entropy budget allocation and covariance-aligned singular vectors.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class EntropyRankAllocatorRegistered (model: nn.Module, min_rank: int = 16, max_rank: int = 128, target_param_ratio: float | None = None, total_rank_budget: int | None = None, target_rank: int | None = None, max_error: float = 0.35, target_attention_only: bool = False, fallback_to_original: bool = True, calibration_loader: Any = None, calibration_texts: list[str] | None = None, target_modules: list[str] | None = None, tokenizer: Any | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("entropy_rank_allocator_class")
class EntropyRankAllocatorRegistered(DynamicEntropySVDCompressor):
    pass

Dynamic Entropy SVD Compressor.

Computes singular entropy H(Sigma_l) on layer activation covariances or weights.
Allocates layer ranks within [min_rank, max_rank] subject to parameter/rank budget.
Replaces dense Linear/Conv1D layers with DynamicSVDLinear factorized modules.
Provides relative Frobenius error validation, max_error threshold, and safe fallback.
Freezes fallback and non-target layers during recovery.

Ancestors

Inherited members

DynamicEntropySVDCompressor:
- collect_activation_covariances
- compress
- compute_layer_entropies
- get_compression_stats
- summary

class FAQLinear (in_features: int, out_features: int, bias: bool = True, bits: int = 8, symmetric: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, device: torch.device | None = None, dtype: torch.dtype | None = None)

Expand source code

class FAQLinear(nn.Module):
    is_custom_attention: bool = True
    """Linear with FAQ activation-aware quantization, STE backprop and optional LoRA."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        bits: int = 8,
        symmetric: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features  # Conv1D compat
        self.bits = bits
        self.symmetric = symmetric
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        requires_grad = not (self.lora_r > 0)
        self.weight = nn.Parameter(
            torch.empty(out_features, in_features, device=device, dtype=dtype),
            requires_grad=requires_grad,
        )
        if bias:
            self.bias = nn.Parameter(
                torch.empty(out_features, device=device, dtype=dtype),
                requires_grad=requires_grad,
            )
        else:
            self.register_parameter("bias", None)

        self.register_buffer(
            "act_scale",
            torch.ones(in_features, device=device, dtype=dtype),
            persistent=False,  # FQ-FIX-03
        )
        nn.init.kaiming_uniform_(self.weight, a=5**0.5)
        if self.bias is not None:
            nn.init.zeros_(self.bias)

        # FQ-FIX-01: LoRA in base dtype.
        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty((self.lora_r, in_features), dtype=dtype, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.zeros((out_features, self.lora_r), dtype=dtype, device=device)
            )
            with torch.no_grad():
                tmp = self.lora_A.data.float()
                nn.init.kaiming_uniform_(tmp, a=5**0.5)
                self.lora_A.data.copy_(tmp.to(dtype))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    @classmethod
    def from_linear(
        cls,
        module: nn.Linear | Any,
        act_scale: torch.Tensor | None = None,  # FQ-FIX-11: optional
        bits: int = 8,
        symmetric: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> FAQLinear:
        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.data
        if is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous()
        else:
            out_f, in_f = raw_w.shape
            W = raw_w

        eff_lora_r = lora_r
        eff_lora_alpha = lora_alpha
        if hasattr(module, "lora_A") and hasattr(module, "lora_B"):
            r_attr = getattr(module, "r", None)
            if r_attr is not None and int(r_attr) > 0:
                eff_lora_r = int(r_attr)
            else:
                la = getattr(module, "lora_A", None)
                if la is not None and hasattr(la, "weight") and la.weight is not None:
                    eff_lora_r = int(la.weight.shape[0])
            eff_lora_alpha = float(getattr(module, "lora_alpha", eff_lora_alpha))

        layer = cls(
            in_features=in_f,
            out_features=out_f,
            bias=(getattr(target_mod, "bias", None) is not None and target_mod.bias is not None),
            bits=bits,
            symmetric=symmetric,
            lora_r=eff_lora_r,
            lora_alpha=eff_lora_alpha,
            device=W.device,
            dtype=W.dtype,
        )
        with torch.no_grad():
            layer.weight.data.copy_(W)
            if (
                getattr(target_mod, "bias", None) is not None
                and target_mod.bias is not None
                and layer.bias is not None
            ):
                layer.bias.data.copy_(target_mod.bias.data)

            # FQ-FIX-11: act_scale optional and shape-safe.
            if act_scale is not None:
                s = act_scale.detach().view(-1)[:in_f].to(device=W.device, dtype=W.dtype)
                n = min(layer.act_scale.numel(), s.numel())
                layer.act_scale[:n].copy_(s[:n])

            # Copy existing LoRA weights.
            if (
                hasattr(module, "lora_A")
                and hasattr(module, "lora_B")
                and layer.lora_A is not None
                and layer.lora_B is not None
            ):
                try:
                    la = getattr(module.lora_A, "weight", module.lora_A)
                    lb = getattr(module.lora_B, "weight", module.lora_B)
                    if hasattr(la, "data") and la.data.shape == layer.lora_A.data.shape:
                        layer.lora_A.data.copy_(la.data)
                    if hasattr(lb, "data") and lb.data.shape == layer.lora_B.data.shape:
                        layer.lora_B.data.copy_(lb.data)
                except Exception as e:
                    logger.debug("FAQ LoRA copy failed: %s", e)

        return layer

    def get_weight(self) -> torch.Tensor:
        s = (
            torch.clamp(self.act_scale, min=1e-6)
            .unsqueeze(0)
            .to(device=self.weight.device, dtype=self.weight.dtype)
        )
        w_scaled = self.weight * s
        w_q = quantize_dequantize_tensor(
            w_scaled,
            bits=self.bits,
            symmetric=self.symmetric,
            per_channel=True,
            use_ste=False,
        )
        return w_q / s

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        s = (
            torch.clamp(self.act_scale, min=1e-6)
            .unsqueeze(0)
            .to(device=self.weight.device, dtype=self.weight.dtype)
        )
        w_scaled = self.weight * s
        w_q = quantize_dequantize_tensor(
            w_scaled,
            bits=self.bits,
            symmetric=self.symmetric,
            per_channel=True,
            use_ste=self.training,
        )
        w_restored = w_q / s

        b_q = None
        if self.bias is not None:
            # FQ-FIX-07: same STE flag as weight.
            b_q = quantize_dequantize_tensor(
                self.bias,
                bits=self.bits,
                symmetric=self.symmetric,
                per_channel=False,
                use_ste=self.training,
            )

        out = nn.functional.linear(
            x,
            w_restored.to(dtype=x.dtype, device=x.device),
            b_q.to(dtype=x.dtype, device=x.device) if b_q is not None else None,
        )

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            # FQ-FIX-08: LoRA scaling matches PEFT convention.
            lora_out = nn.functional.linear(nn.functional.linear(x, self.lora_A), self.lora_B)
            out = out + self.scaling * lora_out

        return out

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: Linear with FAQ activation-aware quantization, STE backprop and optional LoRA.

Static methods

def from_linear(module: nn.Linear | Any, act_scale: torch.Tensor | None = None, bits: int = 8, symmetric: bool = True, lora_r: int = 0, lora_alpha: float = 1.0) ‑> FAQLinear

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    s = (
        torch.clamp(self.act_scale, min=1e-6)
        .unsqueeze(0)
        .to(device=self.weight.device, dtype=self.weight.dtype)
    )
    w_scaled = self.weight * s
    w_q = quantize_dequantize_tensor(
        w_scaled,
        bits=self.bits,
        symmetric=self.symmetric,
        per_channel=True,
        use_ste=self.training,
    )
    w_restored = w_q / s

    b_q = None
    if self.bias is not None:
        # FQ-FIX-07: same STE flag as weight.
        b_q = quantize_dequantize_tensor(
            self.bias,
            bits=self.bits,
            symmetric=self.symmetric,
            per_channel=False,
            use_ste=self.training,
        )

    out = nn.functional.linear(
        x,
        w_restored.to(dtype=x.dtype, device=x.device),
        b_q.to(dtype=x.dtype, device=x.device) if b_q is not None else None,
    )

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        # FQ-FIX-08: LoRA scaling matches PEFT convention.
        lora_out = nn.functional.linear(nn.functional.linear(x, self.lora_A), self.lora_B)
        out = out + self.scaling * lora_out

    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    s = (
        torch.clamp(self.act_scale, min=1e-6)
        .unsqueeze(0)
        .to(device=self.weight.device, dtype=self.weight.dtype)
    )
    w_scaled = self.weight * s
    w_q = quantize_dequantize_tensor(
        w_scaled,
        bits=self.bits,
        symmetric=self.symmetric,
        per_channel=True,
        use_ste=False,
    )
    return w_q / s

class FAQQuantizer (model: nn.Module, tokenizer: Any | None = None, bits: int = 8, calibration_texts: list[str] | None = None, calibration_data: Any | None = None, alpha: float = 0.5, symmetric: bool = True, use_ste: bool = True, lora_r: int = 16, lora_alpha: float = 1.0, max_error: float = 0.35, target_attention_only: bool = False, fallback_to_original: bool = True, skip_attention_qkv: bool = False, cfg: Any | None = None, num_calibration_batches: int = 8, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("faq")
@COMPRESSOR_REGISTRY.register("faq_quantizer")
class FAQQuantizer(BaseCompressor):
    """
    FAQQuantizer: Feature-Activation-Aware Quantization.

    Collects activation statistics to scale channel-wise quantization step sizes,
    minimizing activation-weighted reconstruction distortion.
    """

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any | None = None,
        bits: int = 8,
        calibration_texts: list[str] | None = None,
        calibration_data: Any | None = None,
        alpha: float = 0.5,
        symmetric: bool = True,
        use_ste: bool = True,
        lora_r: int = 16,
        lora_alpha: float = 1.0,
        max_error: float = 0.35,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        skip_attention_qkv: bool = False,
        cfg: Any | None = None,
        num_calibration_batches: int = 8,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="faq", **kwargs)
        cfg = cfg or kwargs.get("comp_config")
        self.tokenizer = tokenizer
        self.bits = int(_cfg_get(cfg, "faq_bits", _cfg_get(cfg, "bits", bits)))
        self.calibration_texts = calibration_texts
        self.calibration_data = calibration_data
        self.alpha = float(_cfg_get(cfg, "faq_alpha", _cfg_get(cfg, "alpha", alpha)))
        self.symmetric = bool(_cfg_get(cfg, "faq_symmetric", _cfg_get(cfg, "symmetric", symmetric)))
        self.use_ste = bool(_cfg_get(cfg, "faq_use_ste", _cfg_get(cfg, "use_ste", use_ste)))
        self.lora_r = int(_cfg_get(cfg, "lora_r", _cfg_get(cfg, "lora_rank", lora_r)))
        self.lora_alpha = float(_cfg_get(cfg, "lora_alpha", lora_alpha))
        self.max_error = float(
            _cfg_get(cfg, "faq_max_error", _cfg_get(cfg, "max_error", max_error))
        )
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "faq_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "faq_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        self.skip_attention_qkv = bool(_cfg_get(cfg, "skip_attention_qkv", skip_attention_qkv))
        self.cfg = cfg
        self.num_calibration_batches = max(
            1, int(_cfg_get(cfg, "num_calibration_batches", num_calibration_batches))
        )

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self._compressed_params: int = 0  # FQ-FIX-10
        self._injected: bool = False  # FQ-FIX-14

    # ------------------------------------------------------------------
    # Target selection
    # ------------------------------------------------------------------
    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            name_lower = name.lower()
            if any(
                p in name_lower
                for p in [
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                ]
            ):
                return True
        return False

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
            return False
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
                "head",
                "output_layer",
            ]
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear_like = (
            isinstance(target_mod, nn.Linear) or "conv1d" in cls_name or hasattr(target_mod, "nf")
        )
        if not is_linear_like or isinstance(target_mod, FAQLinear):
            return False

        if not should_compress_linear(self.model, name):
            return False

        if self.target_attention_only:
            last_part = name.split(".")[-1]
            is_attn = (last_part in _ATTN_NAMES) or any(
                k in name_lower for k in ("attn", "attention")
            )
            if not is_attn:
                return False

        return True

    # ------------------------------------------------------------------
    # FQ-FIX-12: per-feature E[x^2] accumulation (diagonal of covariance).
    # FQ-FIX-13: DataLoader support.
    # ------------------------------------------------------------------
    @staticmethod
    def _normalize_calibration_batches(data: Any) -> list[Any]:
        """FQ-FIX-13: convert DataLoader / list / tensor / dict into a list of batches."""
        if data is None:
            return []
        if isinstance(data, torch.Tensor):
            return [data]
        if isinstance(data, dict):
            return [data]
        if isinstance(data, (list, tuple)):
            return list(data)
        if hasattr(data, "__iter__"):
            return list(data)
        return [data]

    def _collect_activation_scales(self) -> dict[str, torch.Tensor]:
        """
        FQ-FIX-12: accumulate E[x^2] per input feature.
        This is the diagonal of the input covariance; sufficient for
        per-channel quantization.
        """
        running_sq: dict[str, torch.Tensor] = {}
        sample_counts: dict[str, int] = {}
        hooks = []

        def make_hook(layer_name: str):
            def hook(module, inp, outp):
                if not inp or not isinstance(inp[0], torch.Tensor):
                    return
                x = inp[0].detach()
                target_mod = getattr(module, "base_layer", module)
                cls_name = target_mod.__class__.__name__.lower()
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                in_f = target_mod.weight.shape[0] if is_conv1d else target_mod.weight.shape[1]

                x = x.reshape(-1, in_f).float()
                n = x.shape[0]
                if n == 0:
                    return

                sq = (x * x).sum(dim=0).cpu()  # (in_f,)
                if layer_name not in running_sq:
                    running_sq[layer_name] = sq
                    sample_counts[layer_name] = n
                else:
                    running_sq[layer_name] += sq
                    sample_counts[layer_name] += n

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module) and not self._is_gqa_asymmetric(name, module):
                hooks.append(module.register_forward_hook(make_hook(name)))

        self.model.eval()
        first_p = next(self.model.parameters(), None)
        dev = first_p.device if first_p is not None else torch.device("cpu")

        try:
            with torch.no_grad():
                if self.tokenizer is not None and self.calibration_texts:
                    bs = max(1, self.num_calibration_batches)
                    for idx in range(0, len(self.calibration_texts), bs):
                        sub = self.calibration_texts[idx : idx + bs]
                        if not sub:
                            break
                        try:
                            inputs = self.tokenizer(
                                sub,
                                return_tensors="pt",
                                padding=True,
                                truncation=True,
                                max_length=128,
                            )
                            inputs = {k: v.to(dev) for k, v in inputs.items()}
                            self.model(**inputs)
                        except Exception as e:
                            logger.debug("FAQ calibration batch failed: %s", e)

                elif self.calibration_data is not None:
                    batches = self._normalize_calibration_batches(self.calibration_data)
                    for i, batch in enumerate(batches):
                        if i >= self.num_calibration_batches:
                            break
                        try:
                            if isinstance(batch, torch.Tensor):
                                self.model(batch.to(dev))
                            elif isinstance(batch, dict):
                                b = {
                                    k: v.to(dev)
                                    for k, v in batch.items()
                                    if isinstance(v, torch.Tensor)
                                }
                                self.model(**b)
                            elif isinstance(batch, (list, tuple)):
                                self.model(
                                    *[
                                        b.to(dev) if isinstance(b, torch.Tensor) else b
                                        for b in batch
                                    ]
                                )
                        except Exception as e:
                            logger.debug("FAQ calibration batch %d failed: %s", i, e)
                else:
                    first_linear = next(
                        (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                        None,
                    )
                    if first_linear is not None:
                        dummy = torch.randn(2, 16, first_linear.in_features, device=dev)
                        try:
                            self.model(dummy)
                        except Exception as e:
                            logger.debug("FAQ fallback forward failed: %s", e)
        finally:
            for h in hooks:
                h.remove()

        act_scales: dict[str, torch.Tensor] = {}
        for name, module in self.model.named_modules():
            target_mod = getattr(module, "base_layer", module)
            cls_name = target_mod.__class__.__name__.lower()
            is_linear_like = (
                isinstance(target_mod, nn.Linear)
                or "conv1d" in cls_name
                or hasattr(target_mod, "nf")
            )
            if is_linear_like and hasattr(target_mod, "weight") and target_mod.weight is not None:
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                in_f = target_mod.weight.shape[0] if is_conv1d else target_mod.weight.shape[1]
                if name in running_sq and sample_counts.get(name, 0) > 0:
                    mean_sq = running_sq[name] / float(sample_counts[name])
                    channel_scale = torch.sqrt(torch.clamp(mean_sq, min=1e-6))
                    channel_scale = channel_scale / channel_scale.mean().clamp(min=1e-6)
                    act_scales[name] = channel_scale
                else:
                    act_scales[name] = torch.ones(in_f)
        return act_scales

    # ------------------------------------------------------------------
    # Main compress
    # ------------------------------------------------------------------
    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        # FQ-FIX-14: double-injection guard.
        if self._injected:
            logger.warning("FAQQuantizer.compress called twice; returning current stats.")
            try:
                return self.model, self.get_compression_stats()
            except Exception:
                return self.model, {}

        act_scales = self._collect_activation_scales()
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        seen_weights: set = set()

        # Mark non-target attention as skipped when target_attention_only is set.
        if self.target_attention_only:
            for name, module in self.model.named_modules():
                if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
                    continue
                target_mod = getattr(module, "base_layer", module)
                cls_name = target_mod.__class__.__name__.lower()
                is_linear_like = (
                    isinstance(target_mod, nn.Linear)
                    or "conv1d" in cls_name
                    or hasattr(target_mod, "nf")
                )
                if is_linear_like and not self._is_target_module(name, module):
                    for p in module.parameters():
                        p.requires_grad = False
                    self.skipped_layers.append(name)

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            target_mod = getattr(module, "base_layer", module)
            if not hasattr(target_mod, "weight") or target_mod.weight is None:
                continue

            w_id = id(target_mod.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            raw_w = target_mod.weight.data
            W_orig = raw_w.t().float() if is_conv1d else raw_w.float()
            in_f = W_orig.shape[1]

            s_act = act_scales.get(name, torch.ones(in_f)).to(raw_w.device)
            s_pow = torch.pow(torch.clamp(s_act, min=1e-6), self.alpha).unsqueeze(0)

            try:
                if self.use_ste or self.lora_r > 0:
                    faq_layer = FAQLinear.from_linear(
                        module,
                        act_scale=s_act,
                        bits=self.bits,
                        symmetric=self.symmetric,
                        lora_r=self.lora_r,
                        lora_alpha=self.lora_alpha,
                    )
                    W_recon = faq_layer.get_weight().float()
                    W_norm = torch.norm(W_orig, p="fro").clamp_min(1e-8)
                    err = float(torch.norm(W_orig - W_recon, p="fro").item() / W_norm)
                    self.layer_errors[name] = err

                    if err > self.max_error and self.fallback_to_original:
                        self.fallback_layers.append(name)
                        for p in module.parameters():
                            p.requires_grad = False
                        continue

                    replace_module(self.model, name, faq_layer)
                    self.compressed_layers.append(name)
                else:
                    W_scaled = W_orig * s_pow
                    W_q = quantize_dequantize_tensor(
                        W_scaled,
                        bits=self.bits,
                        symmetric=self.symmetric,
                        per_channel=True,
                        use_ste=False,
                    )
                    W_recon = W_q / s_pow
                    W_norm = torch.norm(W_orig, p="fro").clamp_min(1e-8)
                    err = float(torch.norm(W_orig - W_recon, p="fro").item() / W_norm)
                    self.layer_errors[name] = err

                    if err > self.max_error and self.fallback_to_original:
                        self.fallback_layers.append(name)
                        for p in module.parameters():
                            p.requires_grad = False
                        continue

                    # FQ-FIX-02: no_grad mutation.
                    with torch.no_grad():
                        if is_conv1d:
                            target_mod.weight.data.copy_(W_recon.t().to(dtype=raw_w.dtype))
                        else:
                            target_mod.weight.data.copy_(W_recon.to(dtype=raw_w.dtype))

                        if (
                            getattr(target_mod, "bias", None) is not None
                            and target_mod.bias is not None
                        ):
                            b_q = quantize_dequantize_tensor(
                                target_mod.bias.data,
                                bits=self.bits,
                                symmetric=self.symmetric,
                                per_channel=False,
                                use_ste=False,
                            )
                            target_mod.bias.data.copy_(b_q.to(dtype=target_mod.bias.dtype))

                        if not hasattr(target_mod, "act_scale"):
                            target_mod.register_buffer(
                                "act_scale",
                                s_act.to(dtype=raw_w.dtype, device=raw_w.device),
                                persistent=False,
                            )

                    # Preserve LoRA trainability.
                    if hasattr(module, "lora_A") and hasattr(module, "lora_B"):
                        if hasattr(module.lora_A, "weight") and module.lora_A.weight is not None:
                            module.lora_A.weight.requires_grad = True
                        if hasattr(module.lora_B, "weight") and module.lora_B.weight is not None:
                            module.lora_B.weight.requires_grad = True

                    self.compressed_layers.append(name)
            except Exception as e:
                logger.warning("FAQ quantization failed on %s: %s", name, e)
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False

        # FQ-FIX-04: forward sanity check with attention_mask.
        forward_check_success = True
        try:
            with torch.no_grad():
                dev = next(self.model.parameters(), torch.tensor(0.0)).device
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                attn = torch.ones_like(dummy)
                try:
                    self.model(input_ids=dummy, attention_mask=attn)
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception:
                        first_linear = next(
                            (
                                m
                                for m in self.model.modules()
                                if isinstance(m, (nn.Linear, FAQLinear))
                            ),
                            None,
                        )
                        in_dim = getattr(first_linear, "in_features", 8)
                        first_p = next(self.model.parameters(), None)
                        self.model(
                            torch.randn(
                                1,
                                8,
                                in_dim,
                                device=dev,
                                dtype=(first_p.dtype if first_p is not None else torch.float32),
                            )
                        )
        except Exception as e:
            logger.debug("FAQ forward sanity check failed: %s", e)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        avg_err = (
            (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
        )

        # FQ-FIX-05: effective compression reflects actual quantized layers.
        num_targets = len(self.compressed_layers) + len(self.fallback_layers)
        if num_targets > 0 and self.bits > 0:
            quantized_fraction = len(self.compressed_layers) / num_targets
            fp16_to_int_ratio = 16.0 / float(self.bits)
            effective_compression_rate = 1.0 + quantized_fraction * (fp16_to_int_ratio - 1.0)
        else:
            effective_compression_rate = 1.0

        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "avg_reconstruction_error": avg_err,
                "quantization_bits": self.bits,
                "alpha": self.alpha,
                "symmetric": self.symmetric,
                "use_ste": self.use_ste,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "target_attention_only": self.target_attention_only,
                "forward_check_success": forward_check_success,
                "effective_compression_rate": round(effective_compression_rate, 2),
                "theoretical_vram_saving": f"{(1.0 - self.bits / 16.0) * 100:.1f}% vs FP16",
                "method": "faq_quantization",
            }
        )

        # FQ-FIX-14: mark as injected.
        self._injected = True
        return self.model, stats

FAQQuantizer: Feature-Activation-Aware Quantization.

Collects activation statistics to scale channel-wise quantization step sizes, minimizing activation-weighted reconstruction distortion.

Ancestors

BaseCompressor

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class FLAPCompressor (model: nn.Module, tokenizer: Any | None = None, sparsity: float = 0.2, metric_type: str = 'WIFV', max_error: float = 0.35, target_attention_only: bool = False, fallback_to_original: bool = True, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, structured_24: bool = False, target_modules: list[str] | None = None, cfg: Any | None = None, num_calibration_batches: int = 8, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("flap")
@COMPRESSOR_REGISTRY.register("flap_compressor")
class FLAPCompressor(BaseCompressor):
    """
    FLAP Compressor: Fluctuation-based Adaptive Structured Pruning.

    Production hardening features:
    - Robust LoRALinear and Conv1D support (transpose + adapter preservation).
    - target_attention_only restricts pruning to attention projections.
    - Relative Frobenius reconstruction error validation with configurable max_error.
    - Safe fallback to original module if error exceeds threshold.
    - Real metrics (effective_parameter_ratio, num_compressed, num_fallback, sparsity).
    """

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any | None = None,
        sparsity: float = 0.2,
        metric_type: str = "WIFV",
        max_error: float = 0.35,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        structured_24: bool = False,
        target_modules: list[str] | None = None,
        cfg: Any | None = None,
        num_calibration_batches: int = 8,  # FL-FIX-04
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="flap", **kwargs)
        cfg = cfg or kwargs.get("comp_config")
        self.cfg = cfg

        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        raw_sp = _cfg_get(cfg, "flap_sparsity", _cfg_get(cfg, "sparsity", sparsity))
        self.sparsity = max(0.0, min(float(raw_sp), 0.95))
        self.metric_type = str(
            _cfg_get(cfg, "flap_metric_type", _cfg_get(cfg, "metric_type", metric_type))
        ).upper()
        self.structured_24 = bool(
            _cfg_get(cfg, "flap_structured_24", _cfg_get(cfg, "structured_24", structured_24))
        )
        c_max_error = _cfg_get(cfg, "flap_max_error", _cfg_get(cfg, "max_error", max_error))
        if c_max_error is None or c_max_error == 0.35:
            c_max_error = adaptive_max_error(
                int((1.0 - self.sparsity) * 768), 768, 768, "distilbert"
            )
            c_max_error = max(c_max_error, 0.60)

        self.max_error = float(c_max_error)
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "flap_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "flap_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        self.calibration_loader = calibration_loader or _cfg_get(cfg, "calibration_loader", None)
        self.calibration_texts = calibration_texts
        self.target_modules = target_modules
        self.target_device = (
            torch.device(device)
            if device is not None
            else next(model.parameters(), torch.tensor(0.0)).device
        )
        self.num_calibration_batches = max(
            1, int(_cfg_get(cfg, "num_calibration_batches", num_calibration_batches))
        )

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_sparsities: dict[str, float] = {}
        self._compressed_params: int = 0

    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
            return False
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
                "head",
                "output_layer",
            ]
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d):
            return False

        if not should_compress_linear(self.model, name):
            return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)

        if self.target_attention_only:
            last_part = name.split(".")[-1]
            is_attn = (last_part in _ATTN_NAMES) or any(
                k in name_lower for k in ("attn", "attention")
            )
            if not is_attn:
                return False

        return True

    # ------------------------------------------------------------------
    # FL-FIX-04/05: multiple calibration batches, mean-of-variance accumulation
    # ------------------------------------------------------------------
    def _normalize_calibration_batches(self, calib_inputs: Any) -> list[Any]:
        if calib_inputs is None:
            return []
        if isinstance(calib_inputs, torch.Tensor):
            return [calib_inputs]
        if isinstance(calib_inputs, dict):
            return [calib_inputs]
        if isinstance(calib_inputs, (list, tuple)):
            if len(calib_inputs) == 0:
                return []
            first = calib_inputs[0]
            if isinstance(first, (dict, torch.Tensor)):
                return list(calib_inputs)
            return [calib_inputs]
        return [calib_inputs]

    def _collect_layer_activations(self) -> dict[str, torch.Tensor]:
        """
        Capture per-feature activation variance across calibration batches.
        Returns: {name: variance_tensor of shape (in_features,)}
        """
        self.model.eval()
        variance_sum: dict[str, torch.Tensor] = {}
        batch_count: dict[str, int] = {}
        hooks = []

        def make_hook(name: str):
            def hook(module, inputs, output):
                if not inputs or not isinstance(inputs[0], torch.Tensor):
                    return
                inp = inputs[0].detach().float()
                target_mod = getattr(module, "base_layer", module)
                cls_name = target_mod.__class__.__name__.lower()
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                in_f = target_mod.weight.shape[0] if is_conv1d else target_mod.weight.shape[1]
                flat = inp.reshape(-1, in_f).cpu()
                if flat.size(0) >= 2:
                    var = flat.var(dim=0)
                else:
                    var = torch.zeros(in_f)
                if name not in variance_sum:
                    variance_sum[name] = var
                    batch_count[name] = 1
                else:
                    variance_sum[name] += var
                    batch_count[name] += 1

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                hooks.append(module.register_forward_hook(make_hook(name)))

        try:
            calib_inputs = get_calibration_tensors(
                self.model,
                tokenizer=self.tokenizer,
                calibration_loader=self.calibration_loader,
                calibration_texts=self.calibration_texts,
                device=self.target_device,
            )
            batches = self._normalize_calibration_batches(calib_inputs)
            n = min(len(batches), self.num_calibration_batches)

            with torch.no_grad():
                for i in range(n):
                    b = batches[i]
                    try:
                        if isinstance(b, dict):
                            self.model(**b)
                        elif isinstance(b, torch.Tensor):
                            self.model(b)
                        elif isinstance(b, (list, tuple)):
                            self.model(*b)
                    except Exception as e:
                        logger.debug("FLAP calibration batch %d failed: %s", i, e)
        finally:
            for h in hooks:
                h.remove()

        aggregated: dict[str, torch.Tensor] = {}
        for name, vsum in variance_sum.items():
            aggregated[name] = (vsum / max(1, batch_count[name])).to(self.target_device)
        return aggregated

    # ------------------------------------------------------------------
    def _compute_fluctuation_metric(
        self,
        weight: torch.Tensor,
        activations: torch.Tensor | None,
    ) -> torch.Tensor:
        """Per-column fluctuation saliency score in fp32."""
        d_out, d_in = weight.shape
        device = weight.device

        if activations is not None and activations.numel() > 0:
            inp_var = activations.reshape(-1).float().to(device)
            if inp_var.shape[0] != d_in:
                inp_var = torch.ones(d_in, device=device)
            else:
                inp_var = inp_var.clamp(min=1e-8)
        else:
            inp_var = torch.ones(d_in, device=device)

        W = weight.float()
        if self.metric_type == "IFV":
            scores = inp_var
        elif self.metric_type == "WIFV":
            scores = inp_var * torch.sum(W.pow(2), dim=0)
        elif self.metric_type == "WIFN":
            scores = torch.sum(torch.abs(W), dim=0) * torch.sqrt(inp_var)
        else:
            scores = inp_var * torch.sum(W.pow(2), dim=0)

        return scores

    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        activations_map = self._collect_layer_activations()
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.layer_sparsities = {}

        total_orig_params = 0
        total_comp_params = 0

        if self.target_attention_only:
            for name, module in self.model.named_modules():
                if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
                    continue
                target_mod = getattr(module, "base_layer", module)
                cls_name = target_mod.__class__.__name__.lower()
                is_linear = isinstance(target_mod, nn.Linear)
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                if (is_linear or is_conv1d) and not self._is_target_module(name, module):
                    for p in module.parameters():
                        p.requires_grad = False
                    self.skipped_layers.append(name)

        targets = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not targets:
            logger.warning("FLAP: No compressible layers found.")
            return self.model, self.get_compression_stats()

        for name, module in targets:
            target_mod = getattr(module, "base_layer", module)
            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            raw_w = target_mod.weight.data
            orig_shape = raw_w.shape
            device = raw_w.device
            dtype = raw_w.dtype

            W = raw_w.t().float().contiguous() if is_conv1d else raw_w.float()
            d_out, d_in = W.shape
            orig_layer_params = d_out * d_in
            total_orig_params += orig_layer_params
            norm_w = torch.norm(W, p="fro").item()

            act = activations_map.get(name, None)
            scores = self._compute_fluctuation_metric(W, act)

            # FL-FIX-01: compute actual sparsity from mask.
            if self.structured_24:
                saliency_matrix = scores.unsqueeze(0).expand_as(W)
                W_pruned, mask = apply_24_structured_mask(W, saliency=saliency_matrix, dim=1)
                if mask is not None:
                    mask_f = mask.float()
                    effective_sparsity = 1.0 - (mask_f.sum().item() / max(1, mask_f.numel()))
                else:
                    # Fallback: count zeros.
                    zeros = (W_pruned.abs() < 1e-12).float().sum().item()
                    effective_sparsity = zeros / max(1, W_pruned.numel())
            else:
                n_keep = max(1, int(d_in * (1.0 - self.sparsity)))
                _, keep_idx = torch.topk(scores, k=n_keep, largest=True)
                mask = torch.zeros(d_in, device=device)
                mask[keep_idx] = 1.0
                W_pruned = W * mask.unsqueeze(0)
                effective_sparsity = 1.0 - (n_keep / d_in)

            if norm_w < 1e-12:
                final_err = 0.0
            else:
                final_err = float(torch.norm(W - W_pruned, p="fro").item() / norm_w)

            self.layer_errors[name] = final_err
            self.layer_sparsities[name] = float(effective_sparsity)

            if final_err > self.max_error:
                logger.warning(
                    "FLAP reconstruction error %.4f on %s exceeded threshold %.4f.",
                    final_err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    logger.info("FLAP: fallback to original on %s.", name)
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    total_comp_params += orig_layer_params
                    continue

            # FL-FIX-02: no_grad for weight copy.
            if is_conv1d:
                res_w = W_pruned.t().to(dtype=dtype, device=device).contiguous()
            else:
                res_w = W_pruned.to(dtype=dtype, device=device).contiguous()

            if res_w.shape != orig_shape:
                raise RuntimeError(
                    f"FLAP: shape mismatch on {name}: {tuple(res_w.shape)} vs {tuple(orig_shape)}"
                )

            with torch.no_grad():
                target_mod.weight.data.copy_(res_w)

            # LoRA trainability.
            if hasattr(module, "lora_A") and hasattr(module, "lora_B"):
                if hasattr(module.lora_A, "weight") and module.lora_A.weight is not None:
                    module.lora_A.weight.requires_grad = True
                if hasattr(module.lora_B, "weight") and module.lora_B.weight is not None:
                    module.lora_B.weight.requires_grad = True

            self.compressed_layers.append(name)
            # FL-FIX-03: effective params based on actual sparsity.
            comp_layer_params = int((1.0 - effective_sparsity) * orig_layer_params)
            total_comp_params += max(1, comp_layer_params)

        # FL-FIX-06: forward sanity check with attention_mask.
        forward_check_success = True
        try:
            with torch.no_grad():
                first_p = next(self.model.parameters(), None)
                dev = first_p.device if first_p is not None else torch.device("cpu")
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                attn = torch.ones_like(dummy)
                try:
                    self.model(input_ids=dummy, attention_mask=attn)
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception:
                        first_linear = next(
                            (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                            None,
                        )
                        in_dim = getattr(first_linear, "in_features", 8)
                        self.model(
                            torch.randn(
                                1,
                                8,
                                in_dim,
                                device=dev,
                                dtype=(first_p.dtype if first_p is not None else torch.float32),
                            )
                        )
        except Exception as e:
            logger.debug("FLAP forward sanity check failed: %s", e)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        effective_ratio = (total_comp_params / total_orig_params) if total_orig_params > 0 else 1.0
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "layer_sparsities": self.layer_sparsities,
                "metric_type": self.metric_type,
                "sparsity_target": self.sparsity,
                "structured_24": self.structured_24,
                "max_error": self.max_error,
                "target_attention_only": self.target_attention_only,
                "fallback_to_original": self.fallback_to_original,
                "avg_reconstruction_error": avg_err,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "effective_parameter_ratio": round(effective_ratio, 4),
                "forward_check_success": forward_check_success,
                "method": "flap_fluctuation_pruning",
            }
        )
        logger.info(
            "Compressor flap finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

FLAP Compressor: Fluctuation-based Adaptive Structured Pruning.

Production hardening features: - Robust LoRALinear and Conv1D support (transpose + adapter preservation). - target_attention_only restricts pruning to attention projections. - Relative Frobenius reconstruction error validation with configurable max_error. - Safe fallback to original module if error exceeds threshold. - Real metrics (effective_parameter_ratio, num_compressed, num_fallback, sparsity).

Ancestors

BaseCompressor

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class FLATLLMCompressor (model: nn.Module, tokenizer: Any | None = None, keep_ratio: float = 0.7, damping: float = 0.0001, max_error: float = 0.35, target_attention_only: bool = False, fallback_to_original: bool = True, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, structured_24: bool = False, target_modules: list[str] | None = None, cfg: Any | None = None, num_calibration_batches: int = 8, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("flat_llm")
@COMPRESSOR_REGISTRY.register("flat_llm_compressor")
class FLATLLMCompressor(BaseCompressor):
    """
    FLAT-LLM compressor: Activation-aware output-channel pruning.

    Pipeline:
      1. Collect per-output-channel energy from calibration activations.
      2. Rank channels by energy.
      3. Keep top-k channels (k = round(keep_ratio * out_features)).
      4. Replace Linear with PrunedChannelLinear (factorized storage).

    Optionally applies 2:4 structured sparsity (compute-only; no storage savings
    unless a sparse kernel is used).
    """

    # FLT-FIX-09
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any | None = None,
        keep_ratio: float = 0.7,
        damping: float = 1e-4,
        max_error: float = 0.35,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        structured_24: bool = False,
        target_modules: list[str] | None = None,
        cfg: Any | None = None,
        num_calibration_batches: int = 8,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="flat_llm", **kwargs)
        cfg = cfg or kwargs.get("comp_config")
        self.cfg = cfg

        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        raw_kr = _cfg_get(cfg, "flat_llm_keep_ratio", _cfg_get(cfg, "keep_ratio", keep_ratio))
        self.keep_ratio = max(0.05, min(float(raw_kr), 1.0))
        # FLT-FIX-04
        self.damping = max(
            0.0,
            float(_cfg_get(cfg, "flat_llm_damping", _cfg_get(cfg, "damping", damping))),
        )
        self.structured_24 = bool(
            _cfg_get(
                cfg,
                "flat_llm_structured_24",
                _cfg_get(cfg, "structured_24", structured_24),
            )
        )
        c_max_error = _cfg_get(cfg, "flat_llm_max_error", _cfg_get(cfg, "max_error", max_error))
        if c_max_error is None or c_max_error == 0.35:
            c_max_error = adaptive_max_error(int(self.keep_ratio * 768), 768, 768, "distilbert")
            c_max_error = max(c_max_error, 0.60)

        self.max_error = float(c_max_error)
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "flat_llm_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "flat_llm_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        self.calibration_loader = calibration_loader or _cfg_get(cfg, "calibration_loader", None)
        self.calibration_texts = calibration_texts
        self.target_modules = target_modules
        self.target_device = (
            torch.device(device)
            if device is not None
            else next(model.parameters(), torch.tensor(0.0)).device
        )
        self.num_calibration_batches = max(
            1, int(_cfg_get(cfg, "num_calibration_batches", num_calibration_batches))
        )

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_keep_ratios: dict[str, float] = {}
        self._compressed_params: int = 0

    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
            return False
        if any(
            skip in name_lower
            for skip in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
                "norm",
                "layernorm",
                "rmsnorm",
            )
        ):
            return False

        if isinstance(
            module,
            (
                nn.LayerNorm,
                nn.RMSNorm,
                nn.Embedding,
                nn.BatchNorm1d,
                nn.BatchNorm2d,
                PrunedChannelLinear,
            ),
        ):
            return False

        if self.target_attention_only:
            if not any(attn_kw in name_lower for attn_kw in _ATTN_NAMES):
                return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d):
            return False

        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    # ------------------------------------------------------------------
    # Calibration: per-output-channel energy (FLT-FIX-05)
    # ------------------------------------------------------------------
    def _normalize_calibration_batches(self, calib_inputs: Any) -> list[Any]:
        if calib_inputs is None:
            return []
        if isinstance(calib_inputs, (torch.Tensor, dict)):
            return [calib_inputs]
        if isinstance(calib_inputs, (list, tuple)):
            if len(calib_inputs) == 0:
                return []
            first = calib_inputs[0]
            if isinstance(first, (dict, torch.Tensor)):
                return list(calib_inputs)
            return [calib_inputs]
        return [calib_inputs]

    def _collect_layer_outputs(self) -> dict[str, torch.Tensor]:
        """
        Capture per-output-channel energy across calibration batches.

        Energy metric: E[x_out^2] per output channel, computed as the mean of
        per-batch variance of the output activations. For Conv1D, the axis is
        transposed accordingly.
        """
        self.model.eval()
        energy_sum: dict[str, torch.Tensor] = {}
        batch_count: dict[str, int] = {}
        hooks = []

        def make_hook(layer_name: str):
            def hook(module, inputs, output):
                if not isinstance(output, torch.Tensor):
                    return
                out = output.detach().float()
                target_mod = getattr(module, "base_layer", module)
                cls_name = target_mod.__class__.__name__.lower()
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                out_f = target_mod.weight.shape[1] if is_conv1d else target_mod.weight.shape[0]
                if out.shape[-1] != out_f:
                    return
                flat = out.reshape(-1, out_f).cpu()
                if flat.size(0) >= 2:
                    var = flat.var(dim=0)
                else:
                    var = flat.abs().mean(dim=0)
                if layer_name not in energy_sum:
                    energy_sum[layer_name] = var
                    batch_count[layer_name] = 1
                else:
                    energy_sum[layer_name] += var
                    batch_count[layer_name] += 1

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                hooks.append(module.register_forward_hook(make_hook(name)))

        try:
            calib_inputs = get_calibration_tensors(
                self.model,
                tokenizer=self.tokenizer,
                calibration_loader=self.calibration_loader,
                calibration_texts=self.calibration_texts,
                device=self.target_device,
            )
            batches = self._normalize_calibration_batches(calib_inputs)
            n = min(len(batches), self.num_calibration_batches)

            with torch.no_grad():
                for i in range(n):
                    b = batches[i]
                    try:
                        if isinstance(b, dict):
                            self.model(**b)
                        elif isinstance(b, torch.Tensor):
                            self.model(b)
                        elif isinstance(b, (list, tuple)):
                            self.model(*b)
                    except Exception as e:
                        logger.debug("FLAT-LLM calibration batch %d failed: %s", i, e)
        finally:
            for h in hooks:
                h.remove()

        aggregated: dict[str, torch.Tensor] = {}
        for name, esum in energy_sum.items():
            aggregated[name] = (esum / max(1, batch_count[name])).to(self.target_device)
        return aggregated

    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        from neuralops.core.model_utils import replace_module

        outputs_map = self._collect_layer_outputs()
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.layer_keep_ratios = {}

        # Honest accounting: storage params, not "effective" made-up numbers.
        total_orig_storage = 0
        total_comp_storage = 0

        targets = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not targets:
            logger.warning("FLAT-LLM: no compressible layers found.")
            return self.model, self.get_compression_stats()

        for name, module in targets:
            target_mod = getattr(module, "base_layer", module)
            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            raw_w = target_mod.weight.data
            device = raw_w.device
            dtype = raw_w.dtype

            W = raw_w.t().float().contiguous() if is_conv1d else raw_w.float()
            d_out, d_in = W.shape
            orig_layer_storage = d_out * d_in
            total_orig_storage += orig_layer_storage
            norm_w = torch.norm(W, p="fro").item()

            # Activation-aware saliency
            act = outputs_map.get(name, None)
            if act is not None and act.numel() > 0 and act.shape[0] == d_out:
                saliency = act.float().clamp(min=0.0).to(device)
            else:
                saliency = W.pow(2).sum(dim=1)

            # Apply damping to stabilize near-zero channels
            if self.damping > 0:
                saliency = saliency + self.damping * saliency.mean().clamp_min(1e-12)

            # ---- Decide kept channels ----
            if self.structured_24:
                # 2:4 is a compute-only structured sparsity. For storage, we still
                # keep all output channels; only the mask reduces the values.
                sal_mat = saliency.unsqueeze(1).expand_as(W)
                W_pruned, _mask = apply_24_structured_mask(W, saliency=sal_mat, dim=0)
                # Storage: unchanged (dense)
                kept_idx = torch.arange(d_out, device=device)
                W_kept = W_pruned.contiguous()
                keep_ratio_actual = 1.0
                bias_kept = (
                    target_mod.bias.data.float()
                    if getattr(target_mod, "bias", None) is not None and target_mod.bias is not None
                    else None
                )
                comp_layer_storage = orig_layer_storage
            else:
                keep_dims = max(1, min(int(round(d_out * self.keep_ratio)), d_out))
                _, keep_idx = torch.topk(saliency, k=keep_dims, largest=True)
                keep_idx, _ = torch.sort(keep_idx)  # stable order
                kept_idx = keep_idx.to(device)
                W_kept = W.index_select(0, kept_idx).contiguous()
                keep_ratio_actual = keep_dims / d_out

                # Compute error vs original (with pruned rows treated as zero)
                W_pruned = torch.zeros_like(W)
                W_pruned.index_copy_(0, kept_idx, W_kept)

                # Bias handling: keep only the corresponding entries
                if getattr(target_mod, "bias", None) is not None and target_mod.bias is not None:
                    bias_full = target_mod.bias.data.float()
                    bias_kept = bias_full.index_select(0, kept_idx).contiguous()
                else:
                    bias_kept = None

                # Storage: k * in instead of out * in
                comp_layer_storage = int(keep_dims) * d_in

            # ---- Error check ----
            if norm_w < 1e-12:
                final_err = 0.0
            else:
                if self.structured_24:
                    final_err = float(torch.norm(W - W_pruned, p="fro").item() / norm_w)
                else:
                    W_pruned_full = torch.zeros_like(W)
                    W_pruned_full.index_copy_(0, kept_idx, W_kept)
                    final_err = float(torch.norm(W - W_pruned_full, p="fro").item() / norm_w)

            self.layer_errors[name] = final_err
            self.layer_keep_ratios[name] = float(keep_ratio_actual)

            if final_err > self.max_error:
                logger.warning(
                    "FLAT-LLM error %.4f on %s exceeds %.4f.",
                    final_err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    total_comp_storage += orig_layer_storage
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

            # ---- Build the new module ----
            if self.structured_24:
                # No storage reduction: keep original Linear, just overwrite
                # its weight with the masked values (compute-only sparsity).
                new_module = target_mod
                with torch.no_grad():
                    if is_conv1d:
                        target_mod.weight.data.copy_(W_kept.t().to(dtype=dtype, device=device))
                    else:
                        target_mod.weight.data.copy_(W_kept.to(dtype=dtype, device=device))
            else:
                new_module = PrunedChannelLinear(
                    in_features=d_in,
                    original_out_features=d_out,
                    kept_idx=kept_idx.cpu(),
                    W_kept=W_kept.to(dtype=dtype, device=device),
                    bias_kept=(
                        bias_kept.to(dtype=dtype, device=device) if bias_kept is not None else None
                    ),
                )

            # ---- Splice into the model (preserve LoRA wrappers) ----
            if hasattr(module, "base_layer"):
                module.base_layer = new_module
                # Keep LoRA params trainable
                for pn, p in module.named_parameters():
                    if "lora_" in pn.lower() and not pn.startswith("base_layer"):
                        p.requires_grad = True
            else:
                replace_module(self.model, name, new_module)

            self.compressed_layers.append(name)
            total_comp_storage += comp_layer_storage

        # ---- Forward sanity check (FLT-FIX-07) ----
        forward_check_success = True
        try:
            with torch.no_grad():
                first_p = next(self.model.parameters(), None)
                dev = first_p.device if first_p is not None else torch.device("cpu")
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                attn = torch.ones_like(dummy)
                try:
                    self.model(input_ids=dummy, attention_mask=attn)
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception:
                        first_linear = next(
                            (
                                m
                                for m in self.model.modules()
                                if isinstance(m, (nn.Linear, PrunedChannelLinear))
                            ),
                            None,
                        )
                        in_dim = getattr(first_linear, "in_features", 8)
                        dt = first_p.dtype if first_p is not None else torch.float32
                        self.model(torch.randn(1, 8, in_dim, device=dev, dtype=dt))
        except Exception as e:
            logger.debug("FLAT-LLM forward sanity check failed: %s", e)
            forward_check_success = False

        # Honest storage ratio
        storage_ratio = total_orig_storage / total_comp_storage if total_comp_storage > 0 else 1.0
        effective_ratio = total_comp_storage / total_orig_storage if total_orig_storage > 0 else 1.0

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )

        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "layer_errors": self.layer_errors,
                "layer_keep_ratios": self.layer_keep_ratios,
                "keep_ratio": self.keep_ratio,
                "damping": self.damping,
                "structured_24": self.structured_24,
                "max_error": self.max_error,
                "target_attention_only": self.target_attention_only,
                "fallback_to_original": self.fallback_to_original,
                "avg_reconstruction_error": avg_err,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                # Honest accounting (FLT-FIX-02)
                "orig_layer_storage_params": total_orig_storage,
                "comp_layer_storage_params": total_comp_storage,
                "effective_parameter_ratio": round(effective_ratio, 4),
                "storage_compression_ratio": round(storage_ratio, 3),
                "storage_reduction_pct": round(100.0 * (1.0 - effective_ratio), 2),
                "forward_check_success": forward_check_success,
                "method": "flat_llm_channel_pruning",
            }
        )

        logger.info(
            "Compressor flat_llm finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

FLAT-LLM compressor: Activation-aware output-channel pruning.

Pipeline -----= 1. Collect per-output-channel energy from calibration activations. 2. Rank channels by energy. 3. Keep top-k channels (k = round(keep_ratio * out_features)). 4. Replace Linear with PrunedChannelLinear (factorized storage).

Optionally applies 2:4 structured sparsity (compute-only; no storage savings unless a sparse kernel is used).

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class FRanDICompressor (model: nn.Module | None = None, rank: int = 8, energy_threshold: float | None = None, calibration_loader: Any | None = None, lora_r: int = 0, lora_alpha: float = 1.0, max_error: float = 0.5, target_attention_only: bool = False, fallback_to_original: bool = True, skip_attention_qkv: bool = False, cfg: Any | None = None, target_layers: list[str] | None = None, num_calibration_batches: int = 16, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("frandi_compressor_class")
class FRanDICompressor(BaseCompressor):
    """
    FRanDI (Feature-Regression aligned SVD) Compressor.
    """

    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module | None = None,
        rank: int = 8,
        energy_threshold: float | None = None,
        calibration_loader: Any | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        max_error: float = 0.5,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        skip_attention_qkv: bool = False,
        cfg: Any | None = None,
        target_layers: list[str] | None = None,
        num_calibration_batches: int = 16,
        **kwargs: Any,
    ) -> None:
        super().__init__(model if model is not None else nn.Identity(), name="frandi", **kwargs)
        cfg = cfg or kwargs.get("comp_config")
        self.cfg = cfg

        self.rank = int(_cfg_get(cfg, "frandi_rank", _cfg_get(cfg, "rank", rank)))
        self.energy_threshold = _cfg_get(
            cfg,
            "frandi_energy_threshold",
            _cfg_get(cfg, "energy_threshold", energy_threshold),
        )
        self.calibration_loader = calibration_loader or _cfg_get(cfg, "calibration_loader", None)
        self.lora_r = int(_cfg_get(cfg, "lora_r", _cfg_get(cfg, "lora_rank", lora_r)))
        self.lora_alpha = float(_cfg_get(cfg, "lora_alpha", lora_alpha))
        c_max_error = _cfg_get(cfg, "frandi_max_error", _cfg_get(cfg, "max_error", max_error))
        if c_max_error is None or c_max_error == 0.5:
            c_max_error = adaptive_max_error(self.rank, 768, 768, "distilbert")

        self.max_error = float(c_max_error)
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "frandi_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "frandi_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        self.skip_attention_qkv = bool(_cfg_get(cfg, "skip_attention_qkv", skip_attention_qkv))

        # FR-FIX-12: treat empty list as None.
        tl = target_layers or _cfg_get(cfg, "target_layers", None)
        if tl is not None and len(tl) == 0:
            tl = None
        self.target_layers = tl

        self.num_calibration_batches = max(
            1, int(_cfg_get(cfg, "num_calibration_batches", num_calibration_batches))
        )

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self._compressed_params: int = 0

    # ------------------------------------------------------------------
    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        # FR-FIX-17: GQA guard always active.
        cfg = getattr(self.model, "config", None)
        if cfg is None:
            return False
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is None or num_kv_heads is None:
            return False
        if num_kv_heads >= num_heads:
            return False
        name_lower = name.lower()
        return any(
            p in name_lower
            for p in ("q_proj", "k_proj", "v_proj", "query", "key", "value", "c_attn")
        )

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
            return False
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            ]
        ):
            return False

        if self.target_attention_only:
            if not any(attn_kw in name_lower for attn_kw in _ATTN_NAMES):
                return False

        if self._is_gqa_asymmetric(name, module):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d) or isinstance(target_mod, FRanDILinear):
            return False

        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self.target_layers is not None:
            return any(t in name for t in self.target_layers)
        return True

    # ------------------------------------------------------------------
    # FR-FIX-04: bounded calibration loop, FR-FIX-14: fp32 accumulation
    # ------------------------------------------------------------------
    def _collect_covariances(self) -> dict[str, torch.Tensor]:
        if self.calibration_loader is None or not hasattr(self.model, "named_modules"):
            return {}

        covariances: dict[str, torch.Tensor] = {}
        counts: dict[str, int] = {}
        hooks = []

        def make_hook(layer_name: str):
            def hook(module, inputs, output):
                if len(inputs) > 0 and isinstance(inputs[0], torch.Tensor):
                    x = inputs[0].detach().float()
                    target_mod = getattr(module, "base_layer", module)
                    cls_name = target_mod.__class__.__name__.lower()
                    is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                    in_f = target_mod.weight.shape[0] if is_conv1d else target_mod.weight.shape[1]
                    if x.shape[-1] != in_f:
                        return
                    x_flat = x.reshape(-1, in_f)
                    # FR-FIX-14: accumulate in fp32
                    cov = torch.matmul(x_flat.t().contiguous(), x_flat).float()
                    if layer_name not in covariances:
                        covariances[layer_name] = cov
                        counts[layer_name] = x_flat.size(0)
                    else:
                        covariances[layer_name] = covariances[layer_name] + cov
                        counts[layer_name] += x_flat.size(0)

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                hooks.append(module.register_forward_hook(make_hook(name)))

        device = next(self.model.parameters(), torch.tensor(0.0)).device
        self.model.eval()
        try:
            with torch.no_grad():
                n_seen = 0
                for batch in self.calibration_loader:
                    if n_seen >= self.num_calibration_batches:
                        break
                    try:
                        if isinstance(batch, dict):
                            b = {
                                k: v.to(device)
                                for k, v in batch.items()
                                if isinstance(v, torch.Tensor)
                            }
                            _ = self.model(**b)
                        elif isinstance(batch, (list, tuple)):
                            _ = self.model(batch[0].to(device))
                        elif isinstance(batch, torch.Tensor):
                            _ = self.model(batch.to(device))
                        n_seen += 1
                    except Exception as e:
                        logger.debug("FRanDI calibration batch failed: %s", e)
        finally:
            for h in hooks:
                h.remove()

        # Average
        for name in list(covariances.keys()):
            covariances[name] = covariances[name] / max(1, counts[name])
            covariances[name] = covariances[name].cpu()

        return covariances

    # ------------------------------------------------------------------
    def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
        from neuralops.core.model_utils import replace_module

        if model is not None:
            self.model = model
            try:
                self._original_params = self._count_params(self.model)
            except Exception:
                self._original_params = sum(p.numel() for p in self.model.parameters())

        covs = self._collect_covariances()
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        seen_weights: set = set()

        total_orig_params = 0
        total_comp_params = 0

        # FR-FIX-13: threshold scales with requested rank.
        min_rank_threshold = max(4, self.rank // 2)

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            target_mod = getattr(module, "base_layer", module)
            if not hasattr(target_mod, "weight") or target_mod.weight is None:
                continue

            w_id = id(target_mod.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            raw_w = target_mod.weight.data
            W_target = raw_w.t().contiguous().float() if is_conv1d else raw_w.float()
            out_f, in_f = W_target.shape
            orig_params = in_f * out_f
            total_orig_params += orig_params

            max_rank = min(in_f, out_f)
            bound_r = (in_f * out_f) // max(1, in_f + out_f)
            # FR-FIX-13
            if max_rank < min_rank_threshold or bound_r < max(1, self.rank // 4):
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_params += orig_params
                continue

            cov = covs.get(name, None)
            try:
                frandi_layer = FRanDILinear.from_linear(
                    module,
                    rank=self.rank,
                    energy_threshold=self.energy_threshold,
                    cov_matrix=cov,
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                )
                W_comp = frandi_layer.get_weight().float()
                W_norm = torch.norm(W_target, p="fro").clamp_min(1e-8)
                err = float(torch.norm(W_target - W_comp, p="fro").item() / W_norm)
                self.layer_errors[name] = err

                if err > self.max_error:
                    if self.fallback_to_original:
                        self.fallback_layers.append(name)
                        for p in module.parameters():
                            p.requires_grad = False
                        total_comp_params += orig_params
                        continue

                # Preserve wrapper if present (e.g. LoRALinear)
                if hasattr(module, "base_layer"):
                    module.base_layer = frandi_layer
                else:
                    replace_module(self.model, name, frandi_layer)

                self.compressed_layers.append(name)
                actual_rank = int(getattr(frandi_layer, "rank", self.rank))
                comp_params = actual_rank * (in_f + out_f)
                if getattr(frandi_layer, "lora_r", 0) > 0:
                    comp_params += int(frandi_layer.lora_r) * (in_f + out_f)
                total_comp_params += comp_params
            except Exception as e:
                logger.debug("FRanDI error on %s: %s", name, e)
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_params += orig_params

        # FR-FIX-07: forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                first_p = next(self.model.parameters(), None)
                dev = first_p.device if first_p is not None else torch.device("cpu")
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                attn = torch.ones_like(dummy)
                try:
                    self.model(input_ids=dummy, attention_mask=attn)
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception:
                        first_linear = next(
                            (
                                m
                                for m in self.model.modules()
                                if isinstance(m, (nn.Linear, FRanDILinear))
                            ),
                            None,
                        )
                        in_dim = getattr(first_linear, "in_features", 8)
                        dt = first_p.dtype if first_p is not None else torch.float32
                        dummy_float = torch.randn(1, 8, in_dim, device=dev, dtype=dt)
                        self.model(dummy_float)
        except Exception as e:
            logger.debug("FRanDI forward sanity check failed: %s", e)
            forward_check_success = False

        effective_ratio = total_comp_params / total_orig_params if total_orig_params > 0 else 1.0

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        avg_err = (
            sum(self.layer_errors.values()) / len(self.layer_errors) if self.layer_errors else 0.0
        )
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "avg_reconstruction_error": avg_err,
                "frandi_rank": self.rank,
                "energy_threshold": self.energy_threshold,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "max_error": self.max_error,
                "effective_parameter_ratio": round(effective_ratio, 4),
                "forward_check_success": forward_check_success,
                "method": "frandi",
            }
        )
        logger.info(
            "Compressor frandi finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

FRanDI (Feature-Regression aligned SVD) Compressor.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class FRanDILinear (in_features: int, out_features: int, rank: int = 8, bias: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, dtype: torch.dtype | None = None, device: torch.device | None = None)

Expand source code

class FRanDILinear(nn.Module):
    """
    Feature-Regression Aligned Factorized Linear Layer with optional LoRA.

    W ≈ B @ A
      - A (rank, in_features): input feature-aligned basis
      - B (out_features, rank): output expansion matrix

    Forward:
      y = (x @ A^T) @ B^T + bias
          + (alpha/r) * (x @ lora_A^T @ lora_B^T)
    """

    # FR-FIX-09: NOT an attention module.
    is_custom_attention: bool = False

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 8,
        bias: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = max(1, min(rank, in_features, out_features))
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (lora_alpha / lora_r) if lora_r > 0 else 1.0

        requires_grad_base = not (self.lora_r > 0)

        # FR-FIX-02: init in fp32 for numerical stability, cast to dtype.
        A_init = torch.empty(self.rank, in_features, dtype=torch.float32, device=device)
        nn.init.orthogonal_(A_init)
        self.A = nn.Parameter(A_init.to(dtype=dtype), requires_grad=requires_grad_base)

        self.B = nn.Parameter(
            torch.empty(out_features, self.rank, dtype=dtype, device=device),
            requires_grad=requires_grad_base,
        )
        nn.init.xavier_uniform_(self.B.float())
        with torch.no_grad():
            self.B.data = self.B.data.to(dtype=dtype) if dtype is not None else self.B.data

        if bias:
            self.bias = nn.Parameter(
                torch.zeros(out_features, dtype=dtype, device=device),
                requires_grad=requires_grad_base,
            )
        else:
            self.register_parameter("bias", None)

        # FR-FIX-03: LoRA in the same dtype as base.
        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty((self.lora_r, in_features), dtype=dtype, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.zeros((out_features, self.lora_r), dtype=dtype, device=device)
            )
            with torch.no_grad():
                tmp = self.lora_A.data.float()
                nn.init.kaiming_uniform_(tmp, a=math.sqrt(5))
                self.lora_A.data.copy_(tmp.to(dtype=dtype) if dtype is not None else tmp)
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        # FR-FIX-15: weight cache
        self._cached_weight: torch.Tensor | None = None
        self._cached_version: int = -1

    # ---- cache management ----
    def _current_version(self) -> int:
        v = 0
        for p in (self.A, self.B):
            v += getattr(p, "_version", 0)
        if self.lora_A is not None:
            v += getattr(self.lora_A, "_version", 0)
        if self.lora_B is not None:
            v += getattr(self.lora_B, "_version", 0)
        return v

    @property
    def weight(self) -> torch.Tensor:
        # FR-FIX-15: cache in eval mode.
        if (
            not self.training
            and self._cached_weight is not None
            and self._cached_version == self._current_version()
        ):
            return self._cached_weight
        W = torch.matmul(self.B.float(), self.A.float()).to(self.A.dtype)
        if not self.training:
            self._cached_weight = W
            self._cached_version = self._current_version()
        return W

    def get_weight(self) -> torch.Tensor:
        return self.weight

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)

        x_a = F.linear(x_flat, self.A)  # (N, rank)
        out = F.linear(x_a, self.B, self.bias)  # (N, out)

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x_flat, self.lora_A), self.lora_B)
            out = out + self.scaling * lora_out

        return out.view(*orig_shape[:-1], self.out_features)

    # FR-FIX-16: merge() for PEFT/HF compatibility
    @torch.no_grad()
    def merge(self) -> None:
        """Fold LoRA delta into base factors (approximate, additive)."""
        if self.lora_r <= 0 or self.lora_A is None or self.lora_B is None:
            return
        # Merge into B: B_new = B + (alpha/r) * (lora_B @ lora_A) @ pinv(A)
        # Simpler: absorb into B by solving B_new A = B A + delta_BA
        # We use a least-squares approach via pseudoinverse.
        A_f = self.A.float()
        delta = self.scaling * torch.matmul(self.lora_B.float(), self.lora_A.float())
        # Solve for B_new such that B_new @ A = B @ A + delta:
        # B_new = B + delta @ A^+
        A_pinv = torch.linalg.pinv(A_f)
        B_new = self.B.float() + delta @ A_pinv
        self.B.data.copy_(B_new.to(self.B.dtype))
        # Reset LoRA to zero
        self.lora_A.data.zero_()
        self.lora_B.data.zero_()
        self._cached_weight = None

    def merge_and_unload(self) -> FRanDILinear:
        self.merge()
        return self

    def unmerge(self) -> None:
        # Not implemented (would require tracking original B).
        raise NotImplementedError("FRanDILinear.unmerge is not supported.")

    # ------------------------------------------------------------------
    @staticmethod
    def _select_rank(
        S: torch.Tensor,
        rank: int = 8,
        energy_threshold: float | None = None,
    ) -> int:
        # FR-FIX-11: cap at `rank` even when energy threshold is used.
        max_rank = min(int(rank), len(S))
        if energy_threshold is None:
            return max(1, max_rank)
        total_energy = (S**2).sum().clamp_min(1e-12)
        cum_energy = (S**2).cumsum(dim=0) / total_energy
        idx = torch.where(cum_energy >= energy_threshold)[0]
        if len(idx) > 0:
            return max(1, min(int(idx[0].item()) + 1, max_rank))
        return max(1, max_rank)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Linear | Any,
        rank: int = 8,
        energy_threshold: float | None = None,
        cov_matrix: torch.Tensor | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> FRanDILinear:
        target_mod = getattr(linear, "base_layer", linear)
        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.data

        if is_conv1d:
            in_f = raw_w.shape[0]
            out_f = raw_w.shape[1]
            W_float = raw_w.t().contiguous().float()
        else:
            out_f, in_f = raw_w.shape
            W_float = raw_w.float()

        bias = (
            target_mod.bias.data
            if getattr(target_mod, "bias", None) is not None and target_mod.bias is not None
            else None
        )

        existing_lora_A = getattr(linear, "lora_A", None)
        existing_lora_B = getattr(linear, "lora_B", None)

        B_data: torch.Tensor | None = None
        A_data: torch.Tensor | None = None
        r: int = 1

        # ---- Whitening path ----
        if cov_matrix is not None and cov_matrix.shape == (in_f, in_f):
            try:
                cov_f = cov_matrix.float()
                trace_val = torch.diagonal(cov_f).sum().item()
                reg = max(1e-6, 1e-4 * trace_val / max(in_f, 1))

                # FR-FIX-05: add regularization without materializing eye.
                cov_reg = cov_f.clone()
                idx = torch.arange(in_f, device=cov_reg.device)
                cov_reg[idx, idx] += reg

                L = torch.linalg.cholesky(cov_reg)  # (in, in) lower triangular
                W_tilde = torch.matmul(W_float, L)  # (out, in)

                U, S, Vt = torch.linalg.svd(W_tilde, full_matrices=False)
                r = cls._select_rank(S, rank=rank, energy_threshold=energy_threshold)
                r = max(1, min(r, out_f, in_f))

                sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
                B_data = U[:, :r] * sqrt_S.unsqueeze(0)  # (out, r)
                A_tilde = sqrt_S.unsqueeze(1) * Vt[:r, :]  # (r, in)

                # FR-FIX-10: solve for A such that A @ L = A_tilde.
                # Transposed: L^T @ A^T = A_tilde^T. L^T is upper triangular.
                A_solved = torch.linalg.solve_triangular(
                    L.t().contiguous(),
                    A_tilde.t().contiguous(),
                    upper=True,
                )
                A_data = A_solved.t().contiguous()  # (r, in)
            except Exception as e:
                logger.debug("FRanDI whitening failed: %s; falling back to plain SVD", e)
                B_data = None
                A_data = None

        # ---- Plain SVD path ----
        if B_data is None or A_data is None:
            U, S, Vt = torch.linalg.svd(W_float, full_matrices=False)
            r = cls._select_rank(S, rank=rank, energy_threshold=energy_threshold)
            r = max(1, min(r, out_f, in_f))
            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
            B_data = U[:, :r] * sqrt_S.unsqueeze(0)
            A_data = sqrt_S.unsqueeze(1) * Vt[:r, :]

        # Determine effective LoRA rank
        effective_lora_r = int(lora_r)
        if existing_lora_A is not None and hasattr(existing_lora_A, "weight"):
            try:
                effective_lora_r = int(existing_lora_A.weight.shape[0])
            except (AttributeError, IndexError, TypeError):
                pass

        frandi_layer = cls(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=bias is not None,
            lora_r=effective_lora_r,
            lora_alpha=lora_alpha,
            dtype=raw_w.dtype,
            device=raw_w.device,
        )

        with torch.no_grad():
            frandi_layer.A.data.copy_(A_data[:r, :].to(dtype=raw_w.dtype, device=raw_w.device))
            frandi_layer.B.data.copy_(B_data[:, :r].to(dtype=raw_w.dtype, device=raw_w.device))

            if bias is not None and frandi_layer.bias is not None:
                frandi_layer.bias.data.copy_(bias.to(dtype=raw_w.dtype, device=raw_w.device))

            if (
                existing_lora_A is not None
                and existing_lora_B is not None
                and frandi_layer.lora_A is not None
            ):
                w_A = getattr(existing_lora_A, "weight", existing_lora_A)
                w_B = getattr(existing_lora_B, "weight", existing_lora_B)
                if hasattr(w_A, "data") and hasattr(w_B, "data"):
                    try:
                        if w_A.shape == frandi_layer.lora_A.shape:
                            frandi_layer.lora_A.data.copy_(w_A.data)
                        if w_B.shape == frandi_layer.lora_B.shape:
                            frandi_layer.lora_B.data.copy_(w_B.data)
                        frandi_layer.lora_A.requires_grad = True
                        frandi_layer.lora_B.requires_grad = True
                    except Exception as e:
                        logger.debug("FRanDI LoRA copy failed: %s", e)

        return frandi_layer

Feature-Regression Aligned Factorized Linear Layer with optional LoRA.

W ≈ B @ A - A (rank, in_features): input feature-aligned basis - B (out_features, rank): output expansion matrix

Forward -----= y = (x @ A^T) @ B^T + bias + (alpha/r) * (x @ lora_A^T @ lora_B^T)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Static methods

def from_linear(linear: nn.Linear | Any, rank: int = 8, energy_threshold: float | None = None, cov_matrix: torch.Tensor | None = None, lora_r: int = 0, lora_alpha: float = 1.0) ‑> FRanDILinear

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    # FR-FIX-15: cache in eval mode.
    if (
        not self.training
        and self._cached_weight is not None
        and self._cached_version == self._current_version()
    ):
        return self._cached_weight
    W = torch.matmul(self.B.float(), self.A.float()).to(self.A.dtype)
    if not self.training:
        self._cached_weight = W
        self._cached_version = self._current_version()
    return W

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)

    x_a = F.linear(x_flat, self.A)  # (N, rank)
    out = F.linear(x_a, self.B, self.bias)  # (N, out)

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x_flat, self.lora_A), self.lora_B)
        out = out + self.scaling * lora_out

    return out.view(*orig_shape[:-1], self.out_features)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    return self.weight

def merge(self) ‑> None

Expand source code

@torch.no_grad()
def merge(self) -> None:
    """Fold LoRA delta into base factors (approximate, additive)."""
    if self.lora_r <= 0 or self.lora_A is None or self.lora_B is None:
        return
    # Merge into B: B_new = B + (alpha/r) * (lora_B @ lora_A) @ pinv(A)
    # Simpler: absorb into B by solving B_new A = B A + delta_BA
    # We use a least-squares approach via pseudoinverse.
    A_f = self.A.float()
    delta = self.scaling * torch.matmul(self.lora_B.float(), self.lora_A.float())
    # Solve for B_new such that B_new @ A = B @ A + delta:
    # B_new = B + delta @ A^+
    A_pinv = torch.linalg.pinv(A_f)
    B_new = self.B.float() + delta @ A_pinv
    self.B.data.copy_(B_new.to(self.B.dtype))
    # Reset LoRA to zero
    self.lora_A.data.zero_()
    self.lora_B.data.zero_()
    self._cached_weight = None

Fold LoRA delta into base factors (approximate, additive).

def merge_and_unload(self) ‑> FRanDILinear

Expand source code

def merge_and_unload(self) -> FRanDILinear:
    self.merge()
    return self

def unmerge(self) ‑> None

Expand source code

def unmerge(self) -> None:
    # Not implemented (would require tracking original B).
    raise NotImplementedError("FRanDILinear.unmerge is not supported.")

class FWSVDCompressor (model: nn.Module, tokenizer: Any | None = None, rank: int = 64, max_error: float = 0.5, fallback_to_original: bool = True, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, target_modules: list[str] | None = None, target_attention_only: bool = True, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("fwsvd")
@COMPRESSOR_REGISTRY.register("fwsvd_compressor")
class FWSVDCompressor(BaseCompressor):
    """
    Fisher-Weighted SVD Compressor.
    Projects weights using diagonal Fisher activation weights F_i = E[x_i^2].
    Replaces dense nn.Linear with SVDLinear.
    """

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any | None = None,
        rank: int = 64,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        target_modules: list[str] | None = None,
        target_attention_only: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="fwsvd", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            rank = getattr(cfg, "fwsvd_rank", getattr(cfg, "rank", rank))
            tokenizer = getattr(cfg, "tokenizer", tokenizer)
            calibration_texts = getattr(cfg, "calibration_texts", calibration_texts)
            max_error = getattr(cfg, "fwsvd_max_error", getattr(cfg, "max_error", max_error))
            fallback_to_original = getattr(
                cfg,
                "fwsvd_fallback_to_original",
                getattr(cfg, "fallback_to_original", fallback_to_original),
            )
            target_attention_only = getattr(
                cfg,
                "fwsvd_target_attention_only",
                getattr(cfg, "target_attention_only", target_attention_only),
            )

        target_attention_only = kwargs.get(
            "target_attention_only",
            kwargs.get("fwsvd_target_attention_only", target_attention_only),
        )
        self.rank = rank
        self.max_error = max_error
        self.fallback_to_original = fallback_to_original
        self.tokenizer = tokenizer
        self.calibration_texts = calibration_texts or DEFAULT_CALIBRATION_CORPUS
        self.target_modules = target_modules
        self.target_attention_only = target_attention_only
        self.device = (
            torch.device(device) if device is not None else next(model.parameters()).device
        )
        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        if isinstance(module, SVDLinear):
            return False
        if any(
            name.endswith(sfx) for sfx in [".base_layer", ".lora_A", ".lora_B", ".lora_dropout"]
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = module.__class__.__name__.lower()
        target_cls = target_mod.__class__.__name__.lower()

        is_linear_like = (
            isinstance(module, nn.Linear)
            or isinstance(target_mod, nn.Linear)
            or "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(module, "nf")
            or hasattr(target_mod, "nf")
            or "loralinear" in cls_name
            or hasattr(module, "base_layer")
        )
        if not is_linear_like:
            return False
        name_lower = name.lower()
        if any(
            skip in name_lower
            for skip in [
                "lm_head",
                "classifier",
                "head",
                "output_layer",
                "score",
                "pooler",
            ]
        ):
            return False

        if self.target_attention_only:
            ffn_patterns = [
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            ]
            if any(p in name_lower for p in ffn_patterns):
                return False
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            attn_patterns = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(p in name_lower for p in attn_patterns):
                return False
        else:
            if not should_compress_linear(self.model, name):
                return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def compress_linear(
        self,
        layer: nn.Module,
        rank: int | None = None,
        fisher_diag: torch.Tensor | None = None,
    ) -> tuple[SVDLinear | None, float, int]:
        try:
            target_mod = getattr(layer, "base_layer", layer)
            cls_name = layer.__class__.__name__.lower()
            target_cls = target_mod.__class__.__name__.lower()
            is_conv1d = (
                "conv1d" in cls_name
                or "conv1d" in target_cls
                or hasattr(layer, "nf")
                or hasattr(target_mod, "nf")
            )
            weight_tensor = getattr(target_mod, "weight", getattr(layer, "weight", None))
            if weight_tensor is None:
                return None, float("inf"), 0

            raw_w = weight_tensor.data

            if is_conv1d:
                in_feat, out_feat = raw_w.shape[0], raw_w.shape[1]
                W = raw_w.t().contiguous().float()
            else:
                out_feat, in_feat = raw_w.shape[0], raw_w.shape[1]
                W = raw_w.float()

            r = get_dynamic_compression_rank(
                in_feat,
                out_feat,
                model_or_family=self.model,
                default_rank=rank or self.rank,
                comp_config=getattr(self, "comp_config", None),
            )
            max_allowed_rank = max(
                1,
                min(
                    in_feat,
                    out_feat,
                    int(0.75 * (in_feat * out_feat) / (in_feat + out_feat)),
                ),
            )
            r = min(r, max_allowed_rank)
            W_norm = torch.norm(W, p="fro").item()

            sqrt_f: torch.Tensor | None = None
            inv_sqrt_f: torch.Tensor | None = None

            if fisher_diag is not None and fisher_diag.shape == (in_feat,):
                f_clamped = fisher_diag.to(device=layer.weight.device, dtype=torch.float32).clamp(
                    min=1e-8
                )
                sqrt_f = torch.sqrt(f_clamped)
                inv_sqrt_f = 1.0 / sqrt_f

            if sqrt_f is not None and inv_sqrt_f is not None:
                W_weighted = W * sqrt_f.unsqueeze(0)
            else:
                W_weighted = W

            W_weighted_dev = W_weighted.to(device=layer.weight.device, dtype=torch.float32)
            U, S, Vt = torch.linalg.svd(W_weighted_dev, full_matrices=False)

            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
            U_r = U[:, :r] * sqrt_S.unsqueeze(0)
            V_r = sqrt_S.unsqueeze(1) * Vt[:r, :]

            if inv_sqrt_f is not None:
                V_r = V_r * inv_sqrt_f.unsqueeze(0)

            W_recon = torch.matmul(U_r, V_r)
            err = float(torch.norm(W - W_recon, p="fro").item() / max(W_norm, 1e-12))

            dtype = weight_tensor.dtype
            device = weight_tensor.device
            bias_tensor = getattr(target_mod, "bias", getattr(layer, "bias", None))

            svd_layer = SVDLinear(
                in_features=in_feat,
                out_features=out_feat,
                rank=r,
                bias=bias_tensor is not None,
                dtype=dtype,
                device=device,
                lora_r=getattr(self, "lora_r", 0),
                lora_alpha=getattr(self, "lora_alpha", 1.0),
            )
            svd_layer.set_low_rank_factors(
                U_r.to(device=device, dtype=dtype),
                V_r.to(device=device, dtype=dtype),
            )
            if bias_tensor is not None and svd_layer.bias is not None:
                svd_layer.bias.data.copy_(bias_tensor.data)

            svd_layer.reconstruction_error = err
            svd_layer.rank = r
            return svd_layer, err, r
        except Exception as e:
            logger.warning("FWSVD factorization failed on linear: %s", e)
            return None, float("inf"), 0

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        self.model.eval()
        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}

        fishers: dict[str, torch.Tensor] = {}

        if self.tokenizer is None:
            logger.info("No tokenizer provided for FWSVD; proceeding with standard unweighted SVD.")
        else:
            hooks = []

            def make_hook(name: str):
                def hook(module: nn.Module, inp: tuple[torch.Tensor, ...], outp: Any):
                    if not inp or inp[0] is None:
                        return
                    x = inp[0].detach()
                    dev_type = x.device.type
                    autocast_ctx = (
                        torch.autocast(device_type=dev_type, enabled=False)
                        if hasattr(torch, "autocast") and dev_type in ("cuda", "cpu")
                        else contextlib.nullcontext()
                    )
                    with autocast_ctx:
                        target_mod = getattr(module, "base_layer", module)
                        cls_name = module.__class__.__name__.lower()
                        target_cls = target_mod.__class__.__name__.lower()
                        is_conv1d = (
                            "conv1d" in cls_name
                            or "conv1d" in target_cls
                            or hasattr(module, "nf")
                            or hasattr(target_mod, "nf")
                        )
                        weight_tensor = getattr(
                            target_mod, "weight", getattr(module, "weight", None)
                        )
                        if weight_tensor is None:
                            return
                        in_f = weight_tensor.shape[0] if is_conv1d else weight_tensor.shape[1]
                        x_fp32 = x.float()
                        x_flat = x_fp32.reshape(-1, in_f)
                        mask = x_flat.abs().sum(dim=-1) > 1e-5
                        if mask.sum() == 0:
                            return
                        x_active = x_flat[mask]
                        batch_fisher = torch.mean(x_active**2, dim=0).cpu()

                    if name not in fishers:
                        fishers[name] = batch_fisher
                    else:
                        fishers[name] = 0.5 * (fishers[name] + batch_fisher)

                return hook

            for name, module in self.model.named_modules():
                if self._is_target_module(name, module):
                    hooks.append(module.register_forward_hook(make_hook(name)))

            try:
                inputs = self.tokenizer(
                    self.calibration_texts,
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                    max_length=512,
                )
                inputs = {k: v.to(self.device) for k, v in inputs.items()}
                with torch.no_grad():
                    self.model(**inputs)
            except Exception as e:
                logger.debug("FWSVD calibration forward pass exception: %s", e)

            for h in hooks:
                h.remove()

        for name, module in list(self.model.named_modules()):
            if self._is_target_module(name, module):
                f_diag = fishers.get(name)
                svd_layer, err, r = self.compress_linear(module, rank=self.rank, fisher_diag=f_diag)
                self.layer_errors[name] = err
                self.layer_ranks[name] = r

                if svd_layer is None or err > self.max_error:
                    if self.fallback_to_original:
                        logger.warning(
                            "FWSVD reconstruction error %.4f on %s exceeded threshold %.4f. Skipping compression (fallback).",
                            err,
                            name,
                            self.max_error,
                        )
                        self.fallback_layers.append(name)
                        continue
                    else:
                        if svd_layer is None:
                            self.fallback_layers.append(name)
                            continue

                if hasattr(module, "base_layer"):
                    module.base_layer = svd_layer
                else:
                    replace_module(self.model, name, svd_layer)
                self.compressed_layers.append(name)

        stats = {
            "compressed_layers": len(self.compressed_layers),
            "fallback_layers": len(self.fallback_layers),
            "avg_reconstruction_error": (
                float(sum(self.layer_errors.values()) / max(1, len(self.layer_errors)))
                if self.layer_errors
                else 0.0
            ),
            "layer_ranks": self.layer_ranks,
        }
        self.model.compression_stats = stats
        return self.model, stats

Fisher-Weighted SVD Compressor. Projects weights using diagonal Fisher activation weights F_i = E[x_i^2]. Replaces dense nn.Linear with SVDLinear.

Ancestors

BaseCompressor

Methods

def compress_linear(self, layer: nn.Module, rank: int | None = None, fisher_diag: torch.Tensor | None = None) ‑> tuple[SVDLinear | None, float, int]

Expand source code

def compress_linear(
    self,
    layer: nn.Module,
    rank: int | None = None,
    fisher_diag: torch.Tensor | None = None,
) -> tuple[SVDLinear | None, float, int]:
    try:
        target_mod = getattr(layer, "base_layer", layer)
        cls_name = layer.__class__.__name__.lower()
        target_cls = target_mod.__class__.__name__.lower()
        is_conv1d = (
            "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(layer, "nf")
            or hasattr(target_mod, "nf")
        )
        weight_tensor = getattr(target_mod, "weight", getattr(layer, "weight", None))
        if weight_tensor is None:
            return None, float("inf"), 0

        raw_w = weight_tensor.data

        if is_conv1d:
            in_feat, out_feat = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()
        else:
            out_feat, in_feat = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float()

        r = get_dynamic_compression_rank(
            in_feat,
            out_feat,
            model_or_family=self.model,
            default_rank=rank or self.rank,
            comp_config=getattr(self, "comp_config", None),
        )
        max_allowed_rank = max(
            1,
            min(
                in_feat,
                out_feat,
                int(0.75 * (in_feat * out_feat) / (in_feat + out_feat)),
            ),
        )
        r = min(r, max_allowed_rank)
        W_norm = torch.norm(W, p="fro").item()

        sqrt_f: torch.Tensor | None = None
        inv_sqrt_f: torch.Tensor | None = None

        if fisher_diag is not None and fisher_diag.shape == (in_feat,):
            f_clamped = fisher_diag.to(device=layer.weight.device, dtype=torch.float32).clamp(
                min=1e-8
            )
            sqrt_f = torch.sqrt(f_clamped)
            inv_sqrt_f = 1.0 / sqrt_f

        if sqrt_f is not None and inv_sqrt_f is not None:
            W_weighted = W * sqrt_f.unsqueeze(0)
        else:
            W_weighted = W

        W_weighted_dev = W_weighted.to(device=layer.weight.device, dtype=torch.float32)
        U, S, Vt = torch.linalg.svd(W_weighted_dev, full_matrices=False)

        sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
        U_r = U[:, :r] * sqrt_S.unsqueeze(0)
        V_r = sqrt_S.unsqueeze(1) * Vt[:r, :]

        if inv_sqrt_f is not None:
            V_r = V_r * inv_sqrt_f.unsqueeze(0)

        W_recon = torch.matmul(U_r, V_r)
        err = float(torch.norm(W - W_recon, p="fro").item() / max(W_norm, 1e-12))

        dtype = weight_tensor.dtype
        device = weight_tensor.device
        bias_tensor = getattr(target_mod, "bias", getattr(layer, "bias", None))

        svd_layer = SVDLinear(
            in_features=in_feat,
            out_features=out_feat,
            rank=r,
            bias=bias_tensor is not None,
            dtype=dtype,
            device=device,
            lora_r=getattr(self, "lora_r", 0),
            lora_alpha=getattr(self, "lora_alpha", 1.0),
        )
        svd_layer.set_low_rank_factors(
            U_r.to(device=device, dtype=dtype),
            V_r.to(device=device, dtype=dtype),
        )
        if bias_tensor is not None and svd_layer.bias is not None:
            svd_layer.bias.data.copy_(bias_tensor.data)

        svd_layer.reconstruction_error = err
        svd_layer.rank = r
        return svd_layer, err, r
    except Exception as e:
        logger.warning("FWSVD factorization failed on linear: %s", e)
        return None, float("inf"), 0

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class FusedQKVSVDLinear (in_features: int, out_features: int, rank: int, bias: bool = True, dtype: torch.dtype | None = None, device: torch.device | None = None, lora_r: int = 0, lora_alpha: float = 1.0)

Expand source code

class FusedQKVSVDLinear(nn.Module):
    """
    Compressed replacement for GPT-2's fused c_attn.

    Stores three independent low-rank factors (Q, K, V). Forward concatenates
    them so the output shape matches the original fused projection.

    Supports mean-shift bias absorption: if `_input_mean` is set, the input
    is centered before the low-rank projection and the shift is added back
    through the bias.
    """

    is_custom_attention: bool = False

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        bias: bool = True,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> None:
        super().__init__()
        if out_features % 3 != 0:
            raise ValueError(
                f"FusedQKVSVDLinear requires out_features divisible by 3, got {out_features}"
            )
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.split_size = out_features // 3
        self.rank = int(rank)
        # Conv1D compatible
        self.nf = out_features
        self.nx = in_features
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        if device is None:
            device = torch.device("cpu")
        if dtype is None:
            dtype = torch.float32

        # Q, K, V factor pairs
        self.q_A = nn.Parameter(torch.empty(self.split_size, self.rank, dtype=dtype, device=device))
        self.q_B = nn.Parameter(torch.empty(self.rank, in_features, dtype=dtype, device=device))
        self.k_A = nn.Parameter(torch.empty(self.split_size, self.rank, dtype=dtype, device=device))
        self.k_B = nn.Parameter(torch.empty(self.rank, in_features, dtype=dtype, device=device))
        self.v_A = nn.Parameter(torch.empty(self.split_size, self.rank, dtype=dtype, device=device))
        self.v_B = nn.Parameter(torch.empty(self.rank, in_features, dtype=dtype, device=device))

        for A, B in ((self.q_A, self.q_B), (self.k_A, self.k_B), (self.v_A, self.v_B)):
            nn.init.kaiming_uniform_(A, a=math.sqrt(5))
            nn.init.zeros_(B)

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty(self.lora_r, in_features, dtype=dtype, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.zeros(out_features, self.lora_r, dtype=dtype, device=device)
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        # Mean-shift buffer (persistent so it survives checkpoint save/load)
        self.register_buffer(
            "_input_mean",
            torch.zeros(in_features, dtype=torch.float32, device=device),
            persistent=True,
        )
        self._has_input_mean = False

    # ------------------------------------------------------------------
    def set_fused_factors(self, W_fused: torch.Tensor) -> None:
        """Split a fused (in, out) or (out, in) weight into Q/K/V factors."""
        if W_fused.shape == (self.in_features, self.out_features):
            W = W_fused.t().contiguous().float()
        elif W_fused.shape == (self.out_features, self.in_features):
            W = W_fused.contiguous().float()
        else:
            raise ValueError(
                f"Unexpected fused weight shape {tuple(W_fused.shape)}; "
                f"expected ({self.in_features}, {self.out_features}) or "
                f"({self.out_features}, {self.in_features})"
            )

        S = self.split_size
        splits = [
            (W[:S, :], self.q_A, self.q_B),
            (W[S : 2 * S, :], self.k_A, self.k_B),
            (W[2 * S :, :], self.v_A, self.v_B),
        ]
        for W_part, A_param, B_param in splits:
            U, Svals, Vt = torch.linalg.svd(W_part, full_matrices=False)
            r = min(self.rank, Svals.shape[0])
            sqrt_S = torch.sqrt(torch.clamp(Svals[:r], min=1e-12))
            A_new = U[:, :r] * sqrt_S.unsqueeze(0)
            B_new = sqrt_S.unsqueeze(1) * Vt[:r, :]
            with torch.no_grad():
                A_param.data.zero_()
                B_param.data.zero_()
                A_param.data[:, :r].copy_(A_new.to(device=A_param.device, dtype=A_param.dtype))
                B_param.data[:r, :].copy_(B_new.to(device=B_param.device, dtype=B_param.dtype))

    # ------------------------------------------------------------------
    @staticmethod
    def _forward_factor(x: torch.Tensor, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        h = F.linear(x.to(B.dtype), B)
        return F.linear(h, A)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_dtype = x.dtype

        # Mean-shift centering
        if self._has_input_mean:
            mu = self._input_mean.to(x.device, x.dtype)
            x_c = x - mu
        else:
            x_c = x

        q = self._forward_factor(x_c, self.q_A, self.q_B)
        k = self._forward_factor(x_c, self.k_A, self.k_B)
        v = self._forward_factor(x_c, self.v_A, self.v_B)
        out = torch.cat([q, k, v], dim=-1)

        if self.bias is not None:
            out = out + self.bias.to(out.dtype)

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            a = self.lora_A.to(x.dtype)
            b = self.lora_B.to(x.dtype)
            lora_out = F.linear(F.linear(x, a), b) * self.scaling
            out = out + lora_out

        return out.to(orig_dtype)

    @torch.no_grad()
    def set_input_mean(self, mu: torch.Tensor) -> None:
        """Register the calibration activation mean (input centering)."""
        mu_f = mu.detach().float().to(self._input_mean.device)
        if mu_f.shape[0] != self.in_features:
            raise ValueError(f"input_mean shape {tuple(mu_f.shape)} != ({self.in_features},)")
        self._input_mean.copy_(mu_f)
        self._has_input_mean = True

    # ------------------------------------------------------------------
    @property
    def weight(self) -> torch.Tensor:
        """Reconstruct the fused (in, out) Conv1D weight for compatibility."""
        W_q = self.q_A.float() @ self.q_B.float()
        W_k = self.k_A.float() @ self.k_B.float()
        W_v = self.v_A.float() @ self.v_B.float()
        W_linear = torch.cat([W_q, W_k, W_v], dim=0)
        return W_linear.t().contiguous()

    @weight.setter
    def weight(self, value: torch.Tensor) -> None:
        self.set_fused_factors(value)

Compressed replacement for GPT-2's fused c_attn.

Stores three independent low-rank factors (Q, K, V). Forward concatenates them so the output shape matches the original fused projection.

Supports mean-shift bias absorption: if _input_mean is set, the input is centered before the low-rank projection and the shift is added back through the bias.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstruct the fused (in, out) Conv1D weight for compatibility."""
    W_q = self.q_A.float() @ self.q_B.float()
    W_k = self.k_A.float() @ self.k_B.float()
    W_v = self.v_A.float() @ self.v_B.float()
    W_linear = torch.cat([W_q, W_k, W_v], dim=0)
    return W_linear.t().contiguous()

Reconstruct the fused (in, out) Conv1D weight for compatibility.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_dtype = x.dtype

    # Mean-shift centering
    if self._has_input_mean:
        mu = self._input_mean.to(x.device, x.dtype)
        x_c = x - mu
    else:
        x_c = x

    q = self._forward_factor(x_c, self.q_A, self.q_B)
    k = self._forward_factor(x_c, self.k_A, self.k_B)
    v = self._forward_factor(x_c, self.v_A, self.v_B)
    out = torch.cat([q, k, v], dim=-1)

    if self.bias is not None:
        out = out + self.bias.to(out.dtype)

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        a = self.lora_A.to(x.dtype)
        b = self.lora_B.to(x.dtype)
        lora_out = F.linear(F.linear(x, a), b) * self.scaling
        out = out + lora_out

    return out.to(orig_dtype)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def set_fused_factors(self, W_fused: torch.Tensor) ‑> None

Expand source code

def set_fused_factors(self, W_fused: torch.Tensor) -> None:
    """Split a fused (in, out) or (out, in) weight into Q/K/V factors."""
    if W_fused.shape == (self.in_features, self.out_features):
        W = W_fused.t().contiguous().float()
    elif W_fused.shape == (self.out_features, self.in_features):
        W = W_fused.contiguous().float()
    else:
        raise ValueError(
            f"Unexpected fused weight shape {tuple(W_fused.shape)}; "
            f"expected ({self.in_features}, {self.out_features}) or "
            f"({self.out_features}, {self.in_features})"
        )

    S = self.split_size
    splits = [
        (W[:S, :], self.q_A, self.q_B),
        (W[S : 2 * S, :], self.k_A, self.k_B),
        (W[2 * S :, :], self.v_A, self.v_B),
    ]
    for W_part, A_param, B_param in splits:
        U, Svals, Vt = torch.linalg.svd(W_part, full_matrices=False)
        r = min(self.rank, Svals.shape[0])
        sqrt_S = torch.sqrt(torch.clamp(Svals[:r], min=1e-12))
        A_new = U[:, :r] * sqrt_S.unsqueeze(0)
        B_new = sqrt_S.unsqueeze(1) * Vt[:r, :]
        with torch.no_grad():
            A_param.data.zero_()
            B_param.data.zero_()
            A_param.data[:, :r].copy_(A_new.to(device=A_param.device, dtype=A_param.dtype))
            B_param.data[:r, :].copy_(B_new.to(device=B_param.device, dtype=B_param.dtype))

Split a fused (in, out) or (out, in) weight into Q/K/V factors.

def set_input_mean(self, mu: torch.Tensor) ‑> None

Expand source code

@torch.no_grad()
def set_input_mean(self, mu: torch.Tensor) -> None:
    """Register the calibration activation mean (input centering)."""
    mu_f = mu.detach().float().to(self._input_mean.device)
    if mu_f.shape[0] != self.in_features:
        raise ValueError(f"input_mean shape {tuple(mu_f.shape)} != ({self.in_features},)")
    self._input_mean.copy_(mu_f)
    self._has_input_mean = True

class FusedSVDLinear (in_features: int, out_features: int, rank: int = 32, bias: bool = True, use_rms_norm: bool = False, eps: float = 1e-06, dtype: torch.dtype = torch.float32, device: torch.device | None = None)

Expand source code

class FusedSVDLinear(nn.Module):
    """
    RESEARCH SOTA: Fused SVD-Linear Module.
    Combines low-rank projection $Z = x V^T$, in-SRAM RMSNorm, and up-projection $Y = Z U^T + b$
    into a single kernel invocation, eliminating HBM roundtrips.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 32,
        bias: bool = True,
        use_rms_norm: bool = False,
        eps: float = 1e-6,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = rank
        self.eps = eps
        self.use_rms_norm = use_rms_norm

        # Low-rank factorized parameters
        self.V = nn.Parameter(torch.empty(rank, in_features, dtype=dtype, device=device))
        self.U = nn.Parameter(torch.empty(out_features, rank, dtype=dtype, device=device))

        if use_rms_norm:
            self.norm_weight = nn.Parameter(torch.ones(rank, dtype=dtype, device=device))
        else:
            self.register_parameter("norm_weight", None)

        if bias:
            self.bias = nn.Parameter(torch.empty(out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

        self.reset_parameters()

    def reset_parameters(self) -> None:
        # Symmetrical square-root initialization
        nn.init.orthogonal_(self.U)
        nn.init.orthogonal_(self.V)
        if self.bias is not None:
            nn.init.zeros_(self.bias)
        if self.norm_weight is not None:
            nn.init.ones_(self.norm_weight)

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix for compatibility."""
        return torch.matmul(self.U, self.V)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # RESEARCH FIX: Dispatches fused kernel executing down-projection, RMSNorm, and up-projection
        return fused_svd_linear(
            x=x,
            V=self.V,
            U=self.U,
            norm_weight=self.norm_weight,
            bias=self.bias,
            eps=self.eps,
            use_rms_norm=self.use_rms_norm,
        )

    def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
        """Set precomputed U and V low-rank factors."""
        with torch.no_grad():
            self.U.copy_(U_factor)
            self.V.copy_(V_factor)

    def decompose_from_weight(self, weight: torch.Tensor, bias: torch.Tensor | None = None) -> None:
        """Decompose dense weight matrix into symmetric SVD factors."""
        with torch.no_grad():
            W = weight.float()
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            r = min(self.rank, len(S))
            sqrt_s = torch.sqrt(S[:r])
            self.U.data.zero_()
            self.V.data.zero_()
            self.U.data[:, :r] = (U[:, :r] * sqrt_s).to(self.U.dtype)
            self.V.data[:r, :] = (sqrt_s.unsqueeze(1) * Vt[:r, :]).to(self.V.dtype)
            if bias is not None and self.bias is not None:
                self.bias.data.copy_(bias)

    def get_orthogonality_loss(self) -> torch.Tensor:
        """
        # RESEARCH FIX: Computes Stiefel manifold penalty to prevent singular subspace collapse.
        """
        u_gram = torch.matmul(self.U.T, self.U)
        v_gram = torch.matmul(self.V, self.V.T)
        eye_r = torch.eye(self.rank, device=self.U.device, dtype=self.U.dtype)
        loss_u = torch.norm(u_gram - eye_r, p="fro") ** 2
        loss_v = torch.norm(v_gram - eye_r, p="fro") ** 2
        return (loss_u + loss_v) / (self.rank * self.rank)

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, out_features={self.out_features}, "
            f"rank={self.rank}, rms_norm={self.use_rms_norm}, bias={self.bias is not None}"
        )

RESEARCH SOTA: Fused SVD-Linear Module. Combines low-rank projection $Z = x V^T$, in-SRAM RMSNorm, and up-projection $Y = Z U^T + b$ into a single kernel invocation, eliminating HBM roundtrips.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix for compatibility."""
    return torch.matmul(self.U, self.V)

Reconstructed effective weight matrix for compatibility.

Methods

def decompose_from_weight(self, weight: torch.Tensor, bias: torch.Tensor | None = None) ‑> None

Expand source code

def decompose_from_weight(self, weight: torch.Tensor, bias: torch.Tensor | None = None) -> None:
    """Decompose dense weight matrix into symmetric SVD factors."""
    with torch.no_grad():
        W = weight.float()
        U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        r = min(self.rank, len(S))
        sqrt_s = torch.sqrt(S[:r])
        self.U.data.zero_()
        self.V.data.zero_()
        self.U.data[:, :r] = (U[:, :r] * sqrt_s).to(self.U.dtype)
        self.V.data[:r, :] = (sqrt_s.unsqueeze(1) * Vt[:r, :]).to(self.V.dtype)
        if bias is not None and self.bias is not None:
            self.bias.data.copy_(bias)

Decompose dense weight matrix into symmetric SVD factors.

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    return (
        f"in_features={self.in_features}, out_features={self.out_features}, "
        f"rank={self.rank}, rms_norm={self.use_rms_norm}, bias={self.bias is not None}"
    )

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    # RESEARCH FIX: Dispatches fused kernel executing down-projection, RMSNorm, and up-projection
    return fused_svd_linear(
        x=x,
        V=self.V,
        U=self.U,
        norm_weight=self.norm_weight,
        bias=self.bias,
        eps=self.eps,
        use_rms_norm=self.use_rms_norm,
    )

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def get_orthogonality_loss(self) ‑> torch.Tensor

Expand source code

def get_orthogonality_loss(self) -> torch.Tensor:
    """
    # RESEARCH FIX: Computes Stiefel manifold penalty to prevent singular subspace collapse.
    """
    u_gram = torch.matmul(self.U.T, self.U)
    v_gram = torch.matmul(self.V, self.V.T)
    eye_r = torch.eye(self.rank, device=self.U.device, dtype=self.U.dtype)
    loss_u = torch.norm(u_gram - eye_r, p="fro") ** 2
    loss_v = torch.norm(v_gram - eye_r, p="fro") ** 2
    return (loss_u + loss_v) / (self.rank * self.rank)

RESEARCH FIX: Computes Stiefel manifold penalty to prevent singular subspace collapse.

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    # Symmetrical square-root initialization
    nn.init.orthogonal_(self.U)
    nn.init.orthogonal_(self.V)
    if self.bias is not None:
        nn.init.zeros_(self.bias)
    if self.norm_weight is not None:
        nn.init.ones_(self.norm_weight)

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) ‑> None

Expand source code

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
    """Set precomputed U and V low-rank factors."""
    with torch.no_grad():
        self.U.copy_(U_factor)
        self.V.copy_(V_factor)

Set precomputed U and V low-rank factors.

class GFWSCompressor (model: nn.Module, rank: int = 64, damping: float = 0.01, max_error: float = 0.2, fallback_to_original: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, skip_attention_qkv: bool = True, target_attention_only: bool = False, tokenizer: Any | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, target_modules: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("gfws")
@COMPRESSOR_REGISTRY.register("gfws_d")
class GFWSCompressor(BaseCompressor):
    is_custom_attention: bool = True
    """
    Generalized Fisher / Covariance-Weighted SVD Compressor.
    Whitens weights using empirical input activation covariance square-root Cov(X)^{1/2},
    performs low-rank truncated SVD, and de-whitens to restore optimal compressed weights.

    Production hardening features:
    - Relative Frobenius reconstruction error validation with configurable max_error.
    - Automatic fallback to standard weight SVD if covariance whitening fails or error > max_error.
    - Safe fallback to original weights if even standard SVD exceeds max_error.
    - Robust Conv1D support (GPT-2 c_attn, c_proj) via weight transposition.
    - Unified critical-layer protection via should_compress_linear.
    - Full FP32 computation for numerical stability, casting back to original dtype only when assigning.
    - Trainable LoRA adapters (lora_A, lora_B) for parameter-efficient fine-tuning.
    - Per-layer logging and detailed compression statistics.
    """

    def __init__(
        self,
        model: nn.Module,
        rank: int = 64,
        damping: float = 0.01,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        skip_attention_qkv: bool = True,
        target_attention_only: bool = False,
        tokenizer: Any | None = None,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        target_modules: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="gfws", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            rank = getattr(cfg, "gfws_rank", getattr(cfg, "rank", rank))
            damping = getattr(cfg, "gfws_damping", getattr(cfg, "damping", damping))
            max_error = getattr(cfg, "gfws_max_error", getattr(cfg, "max_error", max_error))
            fallback_to_original = getattr(
                cfg,
                "gfws_fallback_to_original",
                getattr(cfg, "fallback_to_original", fallback_to_original),
            )
            lora_r = getattr(cfg, "lora_r", getattr(cfg, "lora_rank", lora_r))
            lora_alpha = getattr(cfg, "lora_alpha", lora_alpha)
            skip_attention_qkv = getattr(cfg, "skip_attention_qkv", False)

        self.rank = int(rank)
        self.damping = float(damping)
        if max_error is None or max_error == 0.20:
            max_error = adaptive_max_error(self.rank, 768, 768, "distilbert")

        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.skip_attention_qkv = bool(skip_attention_qkv)
        self.target_attention_only = target_attention_only
        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.target_modules = target_modules
        self.target_device = (
            torch.device(device)
            if device is not None
            else next(model.parameters(), torch.tensor(0.0)).device
        )

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.used_covariance: dict[str, bool] = {}

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            name_lower = name.lower()
            if any(
                p in name_lower
                for p in [
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                ]
            ):
                return True
        return False

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        cls_name = module.__class__.__name__.lower()
        is_linear = isinstance(module, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")
        if not (is_linear or is_conv1d):
            return False
        if isinstance(module, GFWSLoRALinear):
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self._is_gqa_asymmetric(name, module):
            return False
        if self.target_attention_only and not any(
            k in name.lower()
            for k in [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "out_proj",
                "o_proj",
                "q_lin",
                "k_lin",
                "v_lin",
                "out_lin",
                "query",
                "key",
                "value",
                "c_attn",
                "c_proj",
            ]
        ):
            return False
        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def _collect_activation_covariances(self) -> dict[str, torch.Tensor]:
        """Collect empirical input activation covariances across target modules."""
        self.model.eval()
        running_cov: dict[str, torch.Tensor] = {}
        sample_counts: dict[str, int] = {}
        hooks = []

        def make_hook(layer_name: str) -> Any:
            def hook(module: nn.Module, inp: tuple[torch.Tensor, ...], outp: Any) -> None:
                if len(inp) > 0 and isinstance(inp[0], torch.Tensor):
                    x = inp[0].detach().float()
                    cls_name = module.__class__.__name__.lower()
                    is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")
                    if is_conv1d:
                        in_f = module.weight.shape[0]
                    else:
                        in_f = module.weight.shape[1]
                    if x.ndim >= 2:
                        flat_x = x.reshape(-1, in_f)
                    else:
                        flat_x = x.unsqueeze(0)
                    n_samples = flat_x.shape[0]
                    if n_samples > 0:
                        batch_cov = torch.matmul(flat_x.T, flat_x).cpu()
                        if layer_name not in running_cov:
                            running_cov[layer_name] = batch_cov
                            sample_counts[layer_name] = n_samples
                        else:
                            running_cov[layer_name] += batch_cov
                            sample_counts[layer_name] += n_samples

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                hooks.append(module.register_forward_hook(make_hook(name)))

        calib_inputs = get_calibration_tensors(
            self.model,
            tokenizer=self.tokenizer,
            calibration_loader=self.calibration_loader,
            calibration_texts=self.calibration_texts,
            device=self.target_device,
        )

        with torch.no_grad():
            try:
                if isinstance(calib_inputs, dict):
                    self.model(**calib_inputs)
                elif isinstance(calib_inputs, torch.Tensor):
                    self.model(calib_inputs)
                elif isinstance(calib_inputs, (list, tuple)):
                    self.model(*calib_inputs)
            except Exception as e:
                logger.debug("GFWS covariance collection forward pass exception: %s", e)

        for h in hooks:
            h.remove()

        covariances: dict[str, torch.Tensor] = {}
        for name, cov_sum in running_cov.items():
            n = float(max(sample_counts.get(name, 1), 1))
            covariances[name] = cov_sum / n

        return covariances

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Compress layers via activation covariance whitening + low-rank SVD, with standard SVD fallback."""
        from neuralops.core.model_utils import replace_module

        covariances = self._collect_activation_covariances()
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}
        self.used_covariance = {}
        seen_weights: set[int] = set()

        target_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        for name, module in target_layers:
            if not hasattr(module, "weight") or module.weight is None:
                continue

            w_id = id(module.weight)
            if w_id in seen_weights:
                for p in module.parameters():
                    p.requires_grad = False
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            cls_name = module.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")
            target_mod = getattr(module, "base_layer", module)
            raw_w = target_mod.weight.data

            if is_conv1d:
                in_f, out_f = raw_w.shape[0], raw_w.shape[1]
                W = raw_w.t().contiguous().float()  # [out_features, in_features]
            else:
                out_f, in_f = raw_w.shape[0], raw_w.shape[1]
                W = raw_w.float()  # [out_features, in_features]

            r = max(1, min(self.rank, out_f, in_f))
            device = target_mod.weight.device
            dtype = target_mod.weight.dtype
            W_norm = torch.norm(W, p="fro").item()

            used_cov = False
            W_approx = None
            final_err = 1.0

            cov = covariances.get(name, None)
            if cov is not None and cov.shape == (in_f, in_f):
                try:
                    cov_dev = cov.to(device=device, dtype=torch.float32)
                    diag_mean = torch.diagonal(cov_dev).mean().item()
                    if not math.isfinite(diag_mean) or diag_mean <= 0.0:
                        diag_mean = 1.0
                    lambda_reg = max(float(self.damping) * diag_mean, 1e-6)
                    cov_reg = cov_dev + lambda_reg * torch.eye(
                        in_f, device=device, dtype=torch.float32
                    )

                    eigvals, eigvecs = torch.linalg.eigh(cov_reg)
                    eigvals_clamped = torch.clamp(eigvals, min=1e-6)

                    # Cov(X)^{1/2} and Cov(X)^{-1/2}
                    cov_sqrt = eigvecs @ torch.diag(torch.sqrt(eigvals_clamped)) @ eigvecs.T
                    cov_inv_sqrt = (
                        eigvecs @ torch.diag(1.0 / torch.sqrt(eigvals_clamped)) @ eigvecs.T
                    )

                    # Whiten weights: M = W * Cov(X)^{1/2}
                    M = torch.matmul(W, cov_sqrt)
                    U, S, Vt = torch.linalg.svd(M, full_matrices=False)

                    # Low-rank approximation of M: U_r * S_r * Vt_r
                    M_r = (U[:, :r] * S[:r].unsqueeze(0)) @ Vt[:r, :]

                    # De-whiten: candidate = M_r * Cov(X)^{-1/2}
                    candidate = torch.matmul(M_r, cov_inv_sqrt)

                    if not (torch.isnan(candidate).any() or torch.isinf(candidate).any()):
                        err = float(torch.norm(W - candidate, p="fro").item() / max(W_norm, 1e-12))
                        if err <= self.max_error:
                            W_approx = candidate
                            final_err = err
                            used_cov = True
                        else:
                            logger.debug(
                                "GFWS covariance de-whitened error (%.4f) > max_error (%.4f) on %s. Trying SVD fallback.",
                                err,
                                self.max_error,
                                name,
                            )
                except Exception as e:
                    logger.debug(
                        "GFWS covariance decomposition exception on layer %s: %s",
                        name,
                        e,
                    )

            # Fallback to standard SVD if covariance whitening was not used or failed
            if W_approx is None:
                try:
                    U, S, Vt = torch.linalg.svd(W, full_matrices=False)
                    candidate = (U[:, :r] * S[:r].unsqueeze(0)) @ Vt[:r, :]
                    if not (torch.isnan(candidate).any() or torch.isinf(candidate).any()):
                        err = float(torch.norm(W - candidate, p="fro").item() / max(W_norm, 1e-12))
                        W_approx = candidate
                        final_err = err
                        used_cov = False
                except Exception as e:
                    logger.warning("GFWS standard SVD failed on %s: %s", name, e)
                    W_approx = None
                    final_err = 1.0

            self.layer_errors[name] = final_err
            self.layer_ranks[name] = r
            self.used_covariance[name] = used_cov

            if W_approx is None or final_err > self.max_error:
                logger.warning(
                    "GFWS reconstruction error %.4f on %s exceeded threshold %.4f.",
                    final_err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    logger.info("Skipping GFWS compression for %s (fallback to original).", name)
                    for p in module.parameters():
                        p.requires_grad = False
                    self.fallback_layers.append(name)
                    continue

            # Replace with GFWSLoRALinear to support optional post-compression LoRA adapters
            new_layer = GFWSLoRALinear.from_linear(
                target_mod,
                weight=W_approx.to(dtype=dtype, device=device),
                lora_r=self.lora_r,
                lora_alpha=self.lora_alpha,
            )
            if hasattr(module, "base_layer"):
                if (
                    hasattr(module, "lora_A")
                    and hasattr(module, "lora_B")
                    and new_layer.lora_A is not None
                    and new_layer.lora_B is not None
                ):
                    with torch.no_grad():
                        src_A = (
                            module.lora_A.weight.data
                            if hasattr(module.lora_A, "weight")
                            else module.lora_A.data
                        )
                        src_B = (
                            module.lora_B.weight.data
                            if hasattr(module.lora_B, "weight")
                            else module.lora_B.data
                        )
                        if (
                            src_A.shape == new_layer.lora_A.shape
                            and src_B.shape == new_layer.lora_B.shape
                        ):
                            new_layer.lora_A.copy_(
                                src_A.to(
                                    new_layer.lora_A.device,
                                    dtype=new_layer.lora_A.dtype,
                                )
                            )
                            new_layer.lora_B.copy_(
                                src_B.to(
                                    new_layer.lora_B.device,
                                    dtype=new_layer.lora_B.dtype,
                                )
                            )
                module.base_layer = new_layer
                new_layer = module
            replace_module(self.model, name, new_layer)
            self.compressed_layers.append(name)

        # Post-compression forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                dev = next(self.model.parameters(), torch.tensor(0.0)).device
                dummy = torch.randint(0, 100, (1, 8), device=dev)
                try:
                    self.model(dummy)
                except Exception:
                    first_linear = next(
                        (
                            m
                            for m in self.model.modules()
                            if isinstance(m, (nn.Linear, GFWSLoRALinear))
                        ),
                        None,
                    )
                    in_dim = getattr(first_linear, "in_features", 8)
                    self.model(torch.randn(1, 8, in_dim, device=dev))
        except Exception as e:
            logger.warning("GFWS post-compression forward check warning: %s", e)
            forward_check_success = False

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        stats.update(
            {
                "rank": self.rank,
                "damping": self.damping,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "used_covariance": self.used_covariance,
                "avg_reconstruction_error": avg_err,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "forward_check_success": forward_check_success,
                "method": "covariance_weighted_svd",
            }
        )
        logger.info(
            "Compressor gfws finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: Generalized Fisher / Covariance-Weighted SVD Compressor. Whitens weights using empirical input activation covariance square-root Cov(X)^{1/2}, performs low-rank truncated SVD, and de-whitens to restore optimal compressed weights.

Production hardening features: - Relative Frobenius reconstruction error validation with configurable max_error. - Automatic fallback to standard weight SVD if covariance whitening fails or error > max_error. - Safe fallback to original weights if even standard SVD exceeds max_error. - Robust Conv1D support (GPT-2 c_attn, c_proj) via weight transposition. - Unified critical-layer protection via should_compress_linear. - Full FP32 computation for numerical stability, casting back to original dtype only when assigning. - Trainable LoRA adapters (lora_A, lora_B) for parameter-efficient fine-tuning. - Per-layer logging and detailed compression statistics.

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Compress layers via activation covariance whitening + low-rank SVD, with standard SVD fallback."""
    from neuralops.core.model_utils import replace_module

    covariances = self._collect_activation_covariances()
    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    self.layer_errors = {}
    self.layer_ranks = {}
    self.used_covariance = {}
    seen_weights: set[int] = set()

    target_layers = [
        (name, mod)
        for name, mod in self.model.named_modules()
        if self._is_target_module(name, mod)
    ]

    for name, module in target_layers:
        if not hasattr(module, "weight") or module.weight is None:
            continue

        w_id = id(module.weight)
        if w_id in seen_weights:
            for p in module.parameters():
                p.requires_grad = False
            self.skipped_layers.append(name)
            continue
        seen_weights.add(w_id)

        cls_name = module.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(module, "nf")
        target_mod = getattr(module, "base_layer", module)
        raw_w = target_mod.weight.data

        if is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()  # [out_features, in_features]
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float()  # [out_features, in_features]

        r = max(1, min(self.rank, out_f, in_f))
        device = target_mod.weight.device
        dtype = target_mod.weight.dtype
        W_norm = torch.norm(W, p="fro").item()

        used_cov = False
        W_approx = None
        final_err = 1.0

        cov = covariances.get(name, None)
        if cov is not None and cov.shape == (in_f, in_f):
            try:
                cov_dev = cov.to(device=device, dtype=torch.float32)
                diag_mean = torch.diagonal(cov_dev).mean().item()
                if not math.isfinite(diag_mean) or diag_mean <= 0.0:
                    diag_mean = 1.0
                lambda_reg = max(float(self.damping) * diag_mean, 1e-6)
                cov_reg = cov_dev + lambda_reg * torch.eye(
                    in_f, device=device, dtype=torch.float32
                )

                eigvals, eigvecs = torch.linalg.eigh(cov_reg)
                eigvals_clamped = torch.clamp(eigvals, min=1e-6)

                # Cov(X)^{1/2} and Cov(X)^{-1/2}
                cov_sqrt = eigvecs @ torch.diag(torch.sqrt(eigvals_clamped)) @ eigvecs.T
                cov_inv_sqrt = (
                    eigvecs @ torch.diag(1.0 / torch.sqrt(eigvals_clamped)) @ eigvecs.T
                )

                # Whiten weights: M = W * Cov(X)^{1/2}
                M = torch.matmul(W, cov_sqrt)
                U, S, Vt = torch.linalg.svd(M, full_matrices=False)

                # Low-rank approximation of M: U_r * S_r * Vt_r
                M_r = (U[:, :r] * S[:r].unsqueeze(0)) @ Vt[:r, :]

                # De-whiten: candidate = M_r * Cov(X)^{-1/2}
                candidate = torch.matmul(M_r, cov_inv_sqrt)

                if not (torch.isnan(candidate).any() or torch.isinf(candidate).any()):
                    err = float(torch.norm(W - candidate, p="fro").item() / max(W_norm, 1e-12))
                    if err <= self.max_error:
                        W_approx = candidate
                        final_err = err
                        used_cov = True
                    else:
                        logger.debug(
                            "GFWS covariance de-whitened error (%.4f) > max_error (%.4f) on %s. Trying SVD fallback.",
                            err,
                            self.max_error,
                            name,
                        )
            except Exception as e:
                logger.debug(
                    "GFWS covariance decomposition exception on layer %s: %s",
                    name,
                    e,
                )

        # Fallback to standard SVD if covariance whitening was not used or failed
        if W_approx is None:
            try:
                U, S, Vt = torch.linalg.svd(W, full_matrices=False)
                candidate = (U[:, :r] * S[:r].unsqueeze(0)) @ Vt[:r, :]
                if not (torch.isnan(candidate).any() or torch.isinf(candidate).any()):
                    err = float(torch.norm(W - candidate, p="fro").item() / max(W_norm, 1e-12))
                    W_approx = candidate
                    final_err = err
                    used_cov = False
            except Exception as e:
                logger.warning("GFWS standard SVD failed on %s: %s", name, e)
                W_approx = None
                final_err = 1.0

        self.layer_errors[name] = final_err
        self.layer_ranks[name] = r
        self.used_covariance[name] = used_cov

        if W_approx is None or final_err > self.max_error:
            logger.warning(
                "GFWS reconstruction error %.4f on %s exceeded threshold %.4f.",
                final_err,
                name,
                self.max_error,
            )
            if self.fallback_to_original:
                logger.info("Skipping GFWS compression for %s (fallback to original).", name)
                for p in module.parameters():
                    p.requires_grad = False
                self.fallback_layers.append(name)
                continue

        # Replace with GFWSLoRALinear to support optional post-compression LoRA adapters
        new_layer = GFWSLoRALinear.from_linear(
            target_mod,
            weight=W_approx.to(dtype=dtype, device=device),
            lora_r=self.lora_r,
            lora_alpha=self.lora_alpha,
        )
        if hasattr(module, "base_layer"):
            if (
                hasattr(module, "lora_A")
                and hasattr(module, "lora_B")
                and new_layer.lora_A is not None
                and new_layer.lora_B is not None
            ):
                with torch.no_grad():
                    src_A = (
                        module.lora_A.weight.data
                        if hasattr(module.lora_A, "weight")
                        else module.lora_A.data
                    )
                    src_B = (
                        module.lora_B.weight.data
                        if hasattr(module.lora_B, "weight")
                        else module.lora_B.data
                    )
                    if (
                        src_A.shape == new_layer.lora_A.shape
                        and src_B.shape == new_layer.lora_B.shape
                    ):
                        new_layer.lora_A.copy_(
                            src_A.to(
                                new_layer.lora_A.device,
                                dtype=new_layer.lora_A.dtype,
                            )
                        )
                        new_layer.lora_B.copy_(
                            src_B.to(
                                new_layer.lora_B.device,
                                dtype=new_layer.lora_B.dtype,
                            )
                        )
            module.base_layer = new_layer
            new_layer = module
        replace_module(self.model, name, new_layer)
        self.compressed_layers.append(name)

    # Post-compression forward sanity check
    forward_check_success = True
    try:
        with torch.no_grad():
            dev = next(self.model.parameters(), torch.tensor(0.0)).device
            dummy = torch.randint(0, 100, (1, 8), device=dev)
            try:
                self.model(dummy)
            except Exception:
                first_linear = next(
                    (
                        m
                        for m in self.model.modules()
                        if isinstance(m, (nn.Linear, GFWSLoRALinear))
                    ),
                    None,
                )
                in_dim = getattr(first_linear, "in_features", 8)
                self.model(torch.randn(1, 8, in_dim, device=dev))
    except Exception as e:
        logger.warning("GFWS post-compression forward check warning: %s", e)
        forward_check_success = False

    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )
    stats.update(
        {
            "rank": self.rank,
            "damping": self.damping,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "compressed_layers": self.compressed_layers,
            "fallback_layers": self.fallback_layers,
            "skipped_layers": self.skipped_layers,
            "layer_errors": self.layer_errors,
            "layer_ranks": self.layer_ranks,
            "used_covariance": self.used_covariance,
            "avg_reconstruction_error": avg_err,
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "num_skipped": len(self.skipped_layers),
            "lora_r": self.lora_r,
            "lora_alpha": self.lora_alpha,
            "forward_check_success": forward_check_success,
            "method": "covariance_weighted_svd",
        }
    )
    logger.info(
        "Compressor gfws finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        self.max_error,
    )
    return self.model, stats

Compress layers via activation covariance whitening + low-rank SVD, with standard SVD fallback.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class GFWSLoRALinear (in_features: int, out_features: int, weight: torch.Tensor, bias: torch.Tensor | None = None, lora_r: int = 0, lora_alpha: float = 1.0, device: torch.device | None = None, dtype: torch.dtype | None = None)

Expand source code

class GFWSLoRALinear(nn.Module):
    """
    Linear layer containing Fisher/Covariance-whitened SVD compressed base weights (frozen)
    with optional trainable LoRA adapters (lora_A, lora_B).
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        weight: torch.Tensor,
        bias: torch.Tensor | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features  # Conv1D compatibility
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        target_device = device if device is not None else weight.device
        target_dtype = dtype if dtype is not None else weight.dtype

        self.weight = nn.Parameter(
            weight.to(device=target_device, dtype=target_dtype),
            requires_grad=False,
        )
        if bias is not None:
            self.bias = nn.Parameter(
                bias.to(device=target_device, dtype=target_dtype),
                requires_grad=False,
            )
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty(
                    (self.lora_r, in_features),
                    dtype=torch.float32,
                    device=target_device,
                )
            )
            self.lora_B = nn.Parameter(
                torch.zeros(
                    (out_features, self.lora_r),
                    dtype=torch.float32,
                    device=target_device,
                )
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        weight: torch.Tensor,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> GFWSLoRALinear:
        bias = linear.bias.detach() if hasattr(linear, "bias") and linear.bias is not None else None
        in_features = getattr(linear, "in_features", weight.shape[1])
        out_features = getattr(linear, "out_features", weight.shape[0])
        return cls(
            in_features=in_features,
            out_features=out_features,
            weight=weight,
            bias=bias,
            lora_r=lora_r,
            lora_alpha=lora_alpha,
            device=weight.device,
            dtype=weight.dtype,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)
        out = F.linear(x_flat, self.weight, self.bias)

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x_flat.float(), self.lora_A), self.lora_B)
            out = out + (self.scaling * lora_out).to(out.dtype)

        return out.reshape(*orig_shape[:-1], self.out_features)

Linear layer containing Fisher/Covariance-whitened SVD compressed base weights (frozen) with optional trainable LoRA adapters (lora_A, lora_B).

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Module, weight: torch.Tensor, lora_r: int = 0, lora_alpha: float = 1.0) ‑> GFWSLoRALinear

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)
    out = F.linear(x_flat, self.weight, self.bias)

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x_flat.float(), self.lora_A), self.lora_B)
        out = out + (self.scaling * lora_out).to(out.dtype)

    return out.reshape(*orig_shape[:-1], self.out_features)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class GreCoDeCCompressor (model: nn.Module | None = None, rank: int = 16, max_iter: int = 15, lora_r: int = 0, lora_alpha: float = 1.0, max_error: float = 0.35, target_attention_only: bool = False, fallback_to_original: bool = True, skip_attention_qkv: bool = False, cfg: Any | None = None, target_layers: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("grecodec_compressor_class")
class GreCoDeCCompressor(BaseCompressor):
    """
    GreCoDeC Compressor.
    Decomposes linear layers using coordinate descent low-rank factorization.
    """

    def __init__(
        self,
        model: nn.Module | None = None,
        rank: int = 16,
        max_iter: int = 15,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        max_error: float = 0.35,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        skip_attention_qkv: bool = False,
        cfg: Any | None = None,
        target_layers: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model if model is not None else nn.Identity(), name="grecodec", **kwargs)
        cfg = cfg or kwargs.get("comp_config")
        self.cfg = cfg

        self.rank = int(_cfg_get(cfg, "grecodec_rank", _cfg_get(cfg, "rank", rank)))
        self.max_iter = int(_cfg_get(cfg, "grecodec_max_iter", _cfg_get(cfg, "max_iter", max_iter)))
        self.lora_r = int(_cfg_get(cfg, "lora_r", _cfg_get(cfg, "lora_rank", lora_r)))
        self.lora_alpha = float(_cfg_get(cfg, "lora_alpha", lora_alpha))
        c_max_error = _cfg_get(cfg, "grecodec_max_error", _cfg_get(cfg, "max_error", max_error))
        if c_max_error is None or c_max_error == 0.35:
            c_max_error = adaptive_max_error(self.rank, 768, 768, "distilbert")

        self.max_error = float(c_max_error)
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "grecodec_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "grecodec_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        self.skip_attention_qkv = bool(_cfg_get(cfg, "skip_attention_qkv", skip_attention_qkv))
        self.target_layers = target_layers or _cfg_get(cfg, "target_layers", None)

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            name_lower = name.lower()
            if any(
                p in name_lower
                for p in [
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                ]
            ):
                return True
        return False

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
            return False
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            ]
        ):
            return False

        if self.target_attention_only:
            if not any(attn_kw in name_lower for attn_kw in _ATTN_NAMES):
                return False

        if self._is_gqa_asymmetric(name, module):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d) or isinstance(target_mod, GreCoDeCLinear):
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self.target_layers is not None:
            return any(t in name for t in self.target_layers)
        return True

    def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
        """Compress linear layers using GreCoDeC decomposition."""
        if model is not None:
            self.model = model
            self._original_params = self._count_params(self.model)

        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        seen_weights: set[int] = set()

        total_orig_params = 0
        total_comp_params = 0

        # Freeze skipped non-target layers if target_attention_only is enabled
        if self.target_attention_only:
            for name, module in self.model.named_modules():
                if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
                    continue
                target_mod = getattr(module, "base_layer", module)
                cls_name = target_mod.__class__.__name__.lower()
                is_linear = isinstance(target_mod, nn.Linear)
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                if (is_linear or is_conv1d) and not self._is_target_module(name, module):
                    for p in module.parameters():
                        p.requires_grad = False
                    self.skipped_layers.append(name)

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            target_mod = getattr(module, "base_layer", module)
            if not hasattr(target_mod, "weight") or target_mod.weight is None:
                continue

            w_id = id(target_mod.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            raw_w = target_mod.weight.data
            W_target = raw_w.t().contiguous().float() if is_conv1d else raw_w.float()
            out_f, in_f = W_target.shape[0], W_target.shape[1]
            orig_params = in_f * out_f
            total_orig_params += orig_params

            max_rank = min(in_f, out_f)
            bound_r = (in_f * out_f) // (in_f + out_f)
            if max_rank < 16 or bound_r < 4:
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_params += orig_params
                continue

            try:
                grecodec_layer = GreCoDeCLinear.from_linear(
                    module,
                    rank=self.rank,
                    max_iter=self.max_iter,
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                )
                W_comp = grecodec_layer.get_weight().float()
                W_norm = torch.norm(W_target, p="fro").clamp_min(1e-8)
                err = float(torch.norm(W_target - W_comp, p="fro").item() / W_norm)
                self.layer_errors[name] = err

                if err > self.max_error:
                    if self.fallback_to_original:
                        self.fallback_layers.append(name)
                        for p in module.parameters():
                            p.requires_grad = False
                        total_comp_params += orig_params
                        continue

                replace_module(self.model, name, grecodec_layer)
                self.compressed_layers.append(name)
                comp_params = grecodec_layer.rank * (in_f + out_f)
                if grecodec_layer.lora_r > 0:
                    comp_params += grecodec_layer.lora_r * (in_f + out_f)
                total_comp_params += comp_params
            except Exception as e:
                logger.debug("GreCoDeC error on %s: %s", name, e)
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_params += orig_params

        # Post-compression forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                first_p = next(self.model.parameters(), None)
                dev = first_p.device if first_p is not None else torch.device("cpu")
                dummy = torch.randint(0, 100, (1, 8), device=dev)
                try:
                    self.model(dummy)
                except Exception:
                    first_linear = next(
                        (
                            m
                            for m in self.model.modules()
                            if isinstance(m, (nn.Linear, GreCoDeCLinear))
                        ),
                        None,
                    )
                    in_dim = getattr(first_linear, "in_features", 8)
                    self.model(torch.randn(1, 8, in_dim, device=dev))
        except Exception:
            forward_check_success = False

        effective_ratio = total_comp_params / total_orig_params if total_orig_params > 0 else 1.0

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        avg_err = (
            (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
        )
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "avg_reconstruction_error": avg_err,
                "grecodec_rank": self.rank,
                "grecodec_max_iter": self.max_iter,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "effective_parameter_ratio": round(effective_ratio, 4),
                "forward_check_success": forward_check_success,
            }
        )
        logger.info(
            "Compressor grecodec finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

GreCoDeC Compressor. Decomposes linear layers using coordinate descent low-rank factorization.

Ancestors

BaseCompressor

Methods

def compress(self, model: nn.Module | None = None) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
    """Compress linear layers using GreCoDeC decomposition."""
    if model is not None:
        self.model = model
        self._original_params = self._count_params(self.model)

    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    self.layer_errors = {}
    seen_weights: set[int] = set()

    total_orig_params = 0
    total_comp_params = 0

    # Freeze skipped non-target layers if target_attention_only is enabled
    if self.target_attention_only:
        for name, module in self.model.named_modules():
            if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
                continue
            target_mod = getattr(module, "base_layer", module)
            cls_name = target_mod.__class__.__name__.lower()
            is_linear = isinstance(target_mod, nn.Linear)
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            if (is_linear or is_conv1d) and not self._is_target_module(name, module):
                for p in module.parameters():
                    p.requires_grad = False
                self.skipped_layers.append(name)

    for name, module in list(self.model.named_modules()):
        if not self._is_target_module(name, module):
            continue

        target_mod = getattr(module, "base_layer", module)
        if not hasattr(target_mod, "weight") or target_mod.weight is None:
            continue

        w_id = id(target_mod.weight)
        if w_id in seen_weights:
            self.skipped_layers.append(name)
            continue
        seen_weights.add(w_id)

        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.data
        W_target = raw_w.t().contiguous().float() if is_conv1d else raw_w.float()
        out_f, in_f = W_target.shape[0], W_target.shape[1]
        orig_params = in_f * out_f
        total_orig_params += orig_params

        max_rank = min(in_f, out_f)
        bound_r = (in_f * out_f) // (in_f + out_f)
        if max_rank < 16 or bound_r < 4:
            self.fallback_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False
            total_comp_params += orig_params
            continue

        try:
            grecodec_layer = GreCoDeCLinear.from_linear(
                module,
                rank=self.rank,
                max_iter=self.max_iter,
                lora_r=self.lora_r,
                lora_alpha=self.lora_alpha,
            )
            W_comp = grecodec_layer.get_weight().float()
            W_norm = torch.norm(W_target, p="fro").clamp_min(1e-8)
            err = float(torch.norm(W_target - W_comp, p="fro").item() / W_norm)
            self.layer_errors[name] = err

            if err > self.max_error:
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    total_comp_params += orig_params
                    continue

            replace_module(self.model, name, grecodec_layer)
            self.compressed_layers.append(name)
            comp_params = grecodec_layer.rank * (in_f + out_f)
            if grecodec_layer.lora_r > 0:
                comp_params += grecodec_layer.lora_r * (in_f + out_f)
            total_comp_params += comp_params
        except Exception as e:
            logger.debug("GreCoDeC error on %s: %s", name, e)
            self.fallback_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False
            total_comp_params += orig_params

    # Post-compression forward sanity check
    forward_check_success = True
    try:
        with torch.no_grad():
            first_p = next(self.model.parameters(), None)
            dev = first_p.device if first_p is not None else torch.device("cpu")
            dummy = torch.randint(0, 100, (1, 8), device=dev)
            try:
                self.model(dummy)
            except Exception:
                first_linear = next(
                    (
                        m
                        for m in self.model.modules()
                        if isinstance(m, (nn.Linear, GreCoDeCLinear))
                    ),
                    None,
                )
                in_dim = getattr(first_linear, "in_features", 8)
                self.model(torch.randn(1, 8, in_dim, device=dev))
    except Exception:
        forward_check_success = False

    effective_ratio = total_comp_params / total_orig_params if total_orig_params > 0 else 1.0

    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    avg_err = (
        (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
    )
    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "skipped_layers": self.skipped_layers,
            "skipped_layer_count": len(self.skipped_layers),
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "num_skipped": len(self.skipped_layers),
            "layer_errors": self.layer_errors,
            "avg_reconstruction_error": avg_err,
            "grecodec_rank": self.rank,
            "grecodec_max_iter": self.max_iter,
            "lora_r": self.lora_r,
            "lora_alpha": self.lora_alpha,
            "effective_parameter_ratio": round(effective_ratio, 4),
            "forward_check_success": forward_check_success,
        }
    )
    logger.info(
        "Compressor grecodec finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        self.max_error,
    )
    return self.model, stats

Compress linear layers using GreCoDeC decomposition.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class GreCoDeCLinear (in_features: int, out_features: int, rank: int = 16, bias: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, dtype: torch.dtype | None = None, device: torch.device | None = None)

Expand source code

class GreCoDeCLinear(nn.Module):
    is_custom_attention: bool = True
    """
    Greedy Coordinate Descent Factorized Linear Layer with optional LoRA adapter.
    W ≈ U @ V
    - U in R^{out_features x rank}
    - V in R^{rank x in_features}

    Forward pass:
    y = ((x @ V^T) @ U^T) + bias + (x @ lora_A^T @ lora_B^T) * (alpha / r)
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 16,
        bias: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = max(1, min(rank, in_features, out_features))
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (lora_alpha / lora_r) if lora_r > 0 else 1.0

        requires_grad = not (self.lora_r > 0)
        self.U = nn.Parameter(
            torch.empty(out_features, self.rank, dtype=dtype, device=device),
            requires_grad=requires_grad,
        )
        self.V = nn.Parameter(
            torch.empty(self.rank, in_features, dtype=dtype, device=device),
            requires_grad=requires_grad,
        )

        bound = 1.0 / math.sqrt(max(1, in_features))
        nn.init.uniform_(self.U, -bound, bound)
        nn.init.uniform_(self.V, -bound, bound)

        if bias:
            self.bias = nn.Parameter(
                torch.zeros(out_features, dtype=dtype, device=device),
                requires_grad=requires_grad,
            )
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty((self.lora_r, in_features), dtype=torch.float32, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.zeros((out_features, self.lora_r), dtype=torch.float32, device=device)
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix for compatibility."""
        return torch.matmul(self.U, self.V)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Sub-quadratic forward execution.
        y = ((x @ V^T) @ U^T) + bias + LoRA
        """
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)

        x_v = F.linear(x_flat, self.V)
        out = F.linear(x_v, self.U, self.bias)

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            out = out + self.scaling * F.linear(
                F.linear(x_flat.float(), self.lora_A), self.lora_B
            ).to(out.dtype)

        out = out
        return out.view(*orig_shape[:-1], self.out_features)

    def get_weight(self) -> torch.Tensor:
        """Materialize full reconstructed weight matrix W = U @ V."""
        return torch.matmul(self.U, self.V)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Linear | Any,
        rank: int = 16,
        max_iter: int = 15,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> GreCoDeCLinear:
        """Construct GreCoDeCLinear by greedy coordinate descent low-rank factorization."""
        target_mod = getattr(linear, "base_layer", linear)
        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.data

        if is_conv1d:
            in_f = raw_w.shape[0]
            out_f = raw_w.shape[1]
            W_float = raw_w.t().contiguous().float()  # [out_features, in_features]
        else:
            out_f = raw_w.shape[0]
            in_f = raw_w.shape[1]
            W_float = raw_w.float()

        bias = (
            target_mod.bias.data
            if getattr(target_mod, "bias", None) is not None and target_mod.bias is not None
            else None
        )
        r = max(1, min(rank, out_f, in_f))

        # Check existing LoRA parameters
        existing_lora_A = getattr(linear, "lora_A", None)
        existing_lora_B = getattr(linear, "lora_B", None)

        # Coordinate descent low-rank approximation with SVD initialization fallback
        U_init, S_init, Vt_init = torch.linalg.svd(W_float, full_matrices=False)
        sqrt_S = torch.sqrt(torch.clamp(S_init[:r], min=1e-12))
        U_mat = (U_init[:, :r] * sqrt_S.unsqueeze(0)).clone()
        V_mat = (sqrt_S.unsqueeze(1) * Vt_init[:r, :]).clone()

        # Iterative coordinate descent refinement
        for _ in range(max_iter):
            # Update U given V: min_U ||W - U V||_F^2 => U = W V^T (V V^T)^{-1}
            VVt = torch.matmul(V_mat, V_mat.t().contiguous())
            VVt_reg = VVt + 1e-6 * torch.eye(r, device=VVt.device, dtype=VVt.dtype)
            try:
                WV_t = torch.matmul(W_float, V_mat.t().contiguous())
                U_mat = torch.linalg.solve(VVt_reg, WV_t.t().contiguous()).t().contiguous()
            except Exception:
                break

            # Update V given U: min_V ||W - U V||_F^2 => V = (U^T U)^{-1} U^T W
            UtU = torch.matmul(U_mat.t().contiguous(), U_mat)
            UtU_reg = UtU + 1e-6 * torch.eye(r, device=UtU.device, dtype=UtU.dtype)
            try:
                UtW = torch.matmul(U_mat.t().contiguous(), W_float)
                V_mat = torch.linalg.solve(UtU_reg, UtW)
            except Exception:
                break

        # Symmetrical norm balancing
        U_norm = torch.norm(U_mat, p=2, dim=0, keepdim=True).clamp_min(1e-8)
        V_norm = torch.norm(V_mat, p=2, dim=1, keepdim=True).clamp_min(1e-8)
        geom = torch.sqrt(U_norm.squeeze(0) * V_norm.squeeze(1))
        U_data = U_mat * (geom / U_norm.squeeze(0)).unsqueeze(0)
        V_data = V_mat * (geom / V_norm.squeeze(1)).unsqueeze(1)

        effective_lora_r = lora_r
        if existing_lora_A is not None and hasattr(existing_lora_A, "weight"):
            effective_lora_r = existing_lora_A.weight.shape[0]

        grecodec_layer = cls(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=bias is not None,
            lora_r=effective_lora_r,
            lora_alpha=lora_alpha,
            dtype=raw_w.dtype,
            device=raw_w.device,
        )

        grecodec_layer.U.data.copy_(U_data.to(dtype=raw_w.dtype, device=raw_w.device))
        grecodec_layer.V.data.copy_(V_data.to(dtype=raw_w.dtype, device=raw_w.device))

        if bias is not None and grecodec_layer.bias is not None:
            grecodec_layer.bias.data.copy_(bias.to(dtype=raw_w.dtype, device=raw_w.device))

        if (
            existing_lora_A is not None
            and existing_lora_B is not None
            and grecodec_layer.lora_A is not None
        ):
            w_A = getattr(existing_lora_A, "weight", existing_lora_A)
            w_B = getattr(existing_lora_B, "weight", existing_lora_B)
            if hasattr(w_A, "data") and hasattr(w_B, "data"):
                grecodec_layer.lora_A.data.copy_(w_A.data)
                grecodec_layer.lora_B.data.copy_(w_B.data)
                grecodec_layer.lora_A.requires_grad = True
                grecodec_layer.lora_B.requires_grad = True

        return grecodec_layer

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: Greedy Coordinate Descent Factorized Linear Layer with optional LoRA adapter. W ≈ U @ V - U in R^{out_features x rank} - V in R^{rank x in_features}

Forward pass: y = ((x @ V^T) @ U^T) + bias + (x @ lora_A^T @ lora_B^T) * (alpha / r)

Static methods

def from_linear(linear: nn.Linear | Any, rank: int = 16, max_iter: int = 15, lora_r: int = 0, lora_alpha: float = 1.0) ‑> GreCoDeCLinear: Construct GreCoDeCLinear by greedy coordinate descent low-rank factorization.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix for compatibility."""
    return torch.matmul(self.U, self.V)

Reconstructed effective weight matrix for compatibility.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Sub-quadratic forward execution.
    y = ((x @ V^T) @ U^T) + bias + LoRA
    """
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)

    x_v = F.linear(x_flat, self.V)
    out = F.linear(x_v, self.U, self.bias)

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        out = out + self.scaling * F.linear(
            F.linear(x_flat.float(), self.lora_A), self.lora_B
        ).to(out.dtype)

    out = out
    return out.view(*orig_shape[:-1], self.out_features)

Sub-quadratic forward execution. y = ((x @ V^T) @ U^T) + bias + LoRA

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    """Materialize full reconstructed weight matrix W = U @ V."""
    return torch.matmul(self.U, self.V)

Materialize full reconstructed weight matrix W = U @ V.

class HashedLinear (in_features: int, out_features: int, n_buckets: int = 65536, bias: bool = True, use_sign: bool = True, lora_rank: int = 0, lora_alpha: float = 1.0, dtype: torch.dtype | None = None, device: torch.device | None = None)

Expand source code

class HashedLinear(nn.Module):
    """
    Linear layer with parameter sharing via a compact hash table.

      W[i, j] = sign[i, j] · weight_table[hash_idx[i, j]]
      y = x @ W^T + bias + (alpha/r) · (x @ lora_A^T) @ lora_B^T

    Storage:
      weight_table:   n_buckets · dtype_size        (the only trained/learned params)
      bias:           out_features · dtype_size
      lora_A, lora_B: (optional) r · (in + out) · dtype_size
      hash_idx:       out · in · 4 bytes  (int32, non-persistent)
      sign_val:       out · in · 1 byte   (int8, non-persistent)

    For compression to actually reduce storage, `out · in · 5` bytes must be
    less than `out · in · dtype_size`. With fp16 (2 bytes) this is almost
    never true unless a sparse storage format is used. Honest accounting is
    therefore critical (see HN-FIX-08).
    """

    # HN-FIX-01
    is_custom_attention: bool = False

    def __init__(
        self,
        in_features: int,
        out_features: int,
        n_buckets: int = 65536,
        bias: bool = True,
        use_sign: bool = True,
        lora_rank: int = 0,
        lora_alpha: float = 1.0,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.nf = self.out_features
        self.nx = self.in_features
        self.n_buckets = max(1, int(n_buckets))
        self.use_sign = bool(use_sign)
        self.lora_rank = max(0, int(lora_rank))
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_rank) if self.lora_rank > 0 else 1.0

        # Learned parameter table
        self.weight_table = nn.Parameter(torch.empty(self.n_buckets, dtype=dtype, device=device))
        bound = 1.0 / math.sqrt(max(1, self.in_features))
        nn.init.uniform_(self.weight_table, -bound, bound)

        # Bias
        if bias:
            self.bias = nn.Parameter(torch.zeros(self.out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

        # Hash buffers (non-persistent, recomputed on load)
        h_idx, s_val = fast_hash_and_sign(
            self.out_features,
            self.in_features,
            self.n_buckets,
            device=device,
        )
        self.register_buffer("hash_idx", h_idx, persistent=False)
        if self.use_sign:
            self.register_buffer("sign_val", s_val, persistent=False)
        else:
            self.register_buffer(
                "sign_val",
                torch.ones_like(s_val, dtype=torch.int8),
                persistent=False,
            )

        # LoRA
        if self.lora_rank > 0:
            self.lora_A = nn.Parameter(
                torch.empty(self.lora_rank, self.in_features, dtype=dtype, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.zeros(self.out_features, self.lora_rank, dtype=dtype, device=device)
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        # Cache for reconstructed weight (HN-FIX-02)
        self._cached_weight: torch.Tensor | None = None
        self._cached_versions: tuple[int, ...] = ()

    # ------------------------------------------------------------------
    def _current_versions(self) -> tuple[int, ...]:
        vs = (getattr(self.weight_table, "_version", 0),)
        if self.use_sign:
            vs += (0,)  # sign is a buffer, never mutated by autograd
        return vs

    def _invalidate_cache(self) -> None:
        self._cached_weight = None
        self._cached_versions = ()

    def get_weight(self) -> torch.Tensor:
        # HN-FIX-02
        if (
            not self.training
            and self._cached_weight is not None
            and self._cached_versions == self._current_versions()
        ):
            return self._cached_weight

        # Gather rows of the table by index and multiply by sign
        w_flat = self.weight_table.index_select(0, self.hash_idx.reshape(-1).long())
        W = w_flat.view(self.out_features, self.in_features)
        if self.use_sign:
            W = W * self.sign_val.to(dtype=W.dtype)
        if not self.training:
            self._cached_weight = W
            self._cached_versions = self._current_versions()
        return W

    @property
    def weight(self) -> torch.Tensor:
        return self.get_weight()

    @property
    def bucket_indices(self) -> torch.Tensor:
        return self.hash_idx

    # ------------------------------------------------------------------
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        W = self.get_weight()
        out = F.linear(x, W, self.bias)

        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            lora = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.scaling
            out = out + lora

        return out

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, out_features={self.out_features}, "
            f"n_buckets={self.n_buckets}, use_sign={self.use_sign}, "
            f"lora_rank={self.lora_rank}"
        )

    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict[str, int]:
        """
        HN-FIX-08: report per-component storage (bytes).
        Used to compute honest compression ratios.
        """

        def _sz(t):
            if t is None:
                return 0
            return t.numel() * t.element_size()

        return {
            "weight_table": _sz(self.weight_table.data),
            "bias": _sz(self.bias.data) if self.bias is not None else 0,
            "lora_A": _sz(self.lora_A.data) if self.lora_A is not None else 0,
            "lora_B": _sz(self.lora_B.data) if self.lora_B is not None else 0,
            "hash_idx": _sz(self.hash_idx),
            "sign_val": _sz(self.sign_val),
        }

    # ------------------------------------------------------------------
    @classmethod
    def from_linear(
        cls,
        layer: nn.Module,
        n_buckets: int = 65536,
        use_sign: bool = True,
        lora_rank: int = 0,
        lora_alpha: float = 1.0,
        fit_lora_residual: bool = True,
        **_ignored: Any,
    ) -> HashedLinear:
        """
        Build a HashedLinear by projecting W onto the hash basis.

        Reconstruction: W_hat[i,j] = sign[i,j] · mean_{bucket}(W · sign)
        which is the least-squares optimum under the HashedNet constraint.

        HN-FIX-07: no energy rescaling (it increases error).
        """
        target_mod = getattr(layer, "base_layer", layer)
        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.data

        # Normalize to (out, in)
        if hasattr(target_mod, "nx") and hasattr(target_mod, "nf"):
            in_f = int(target_mod.nx)
            out_f = int(target_mod.nf)
            W = raw_w.t().contiguous().float()
        elif hasattr(target_mod, "in_features") and hasattr(target_mod, "out_features"):
            in_f = int(target_mod.in_features)
            out_f = int(target_mod.out_features)
            if raw_w.shape == (in_f, out_f):
                W = raw_w.t().contiguous().float()
            else:
                W = raw_w.view(out_f, in_f).float()
        elif is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float()

        # Clamp buckets to a sane range
        n_buckets = max(16, min(int(n_buckets), out_f * in_f))

        bias = (
            target_mod.bias.data
            if getattr(target_mod, "bias", None) is not None and target_mod.bias is not None
            else None
        )

        # Detect existing LoRA
        existing_lora_A = getattr(layer, "lora_A", None)
        existing_lora_B = getattr(layer, "lora_B", None)
        effective_lora_r = max(0, int(lora_rank))
        if existing_lora_A is not None and hasattr(existing_lora_A, "weight"):
            try:
                effective_lora_r = int(existing_lora_A.weight.shape[0])
            except Exception:
                pass

        hashed = cls(
            in_features=in_f,
            out_features=out_f,
            n_buckets=n_buckets,
            bias=bias is not None,
            use_sign=use_sign,
            lora_rank=effective_lora_r,
            lora_alpha=lora_alpha,
            dtype=raw_w.dtype,
            device=raw_w.device,
        )

        # ---- Least-squares fit for weight_table ----
        W_flat = W.flatten()
        h_flat = hashed.hash_idx.reshape(-1).long()
        s_flat = hashed.sign_val.reshape(-1).to(torch.float32)
        if not use_sign:
            s_flat = torch.ones_like(s_flat)

        # accum[b] = sum_{i,j: h(i,j)=b} W[i,j] · sign[i,j]
        accum = torch.zeros(n_buckets, dtype=torch.float32, device=W.device)
        weights_signed = W_flat * s_flat
        accum.scatter_add_(0, h_flat, weights_signed)

        counts = torch.bincount(h_flat, minlength=n_buckets).float()

        # HN-FIX-12: deterministic init for empty buckets
        non_zero = counts > 0
        mean_per_bucket = torch.zeros(n_buckets, dtype=torch.float32, device=W.device)
        mean_per_bucket[non_zero] = accum[non_zero] / counts[non_zero]

        if (~non_zero).any():
            bound = 1.0 / math.sqrt(max(1, in_f))
            gen = torch.Generator(device=W.device).manual_seed(0xCAFEBABE)
            n_empty = int((~non_zero).sum().item())
            mean_per_bucket[~non_zero] = torch.empty(
                n_empty, device=W.device, dtype=torch.float32
            ).uniform_(-bound, bound, generator=gen)

        with torch.no_grad():
            hashed.weight_table.data.copy_(
                mean_per_bucket.to(dtype=raw_w.dtype, device=raw_w.device)
            )
            if bias is not None and hashed.bias is not None:
                hashed.bias.data.copy_(bias.to(dtype=raw_w.dtype, device=raw_w.device))

        # ---- LoRA transfer / fit ----
        if (
            existing_lora_A is not None
            and existing_lora_B is not None
            and hashed.lora_A is not None
        ):
            w_A = getattr(existing_lora_A, "weight", existing_lora_A)
            w_B = getattr(existing_lora_B, "weight", existing_lora_B)
            try:
                if w_A.shape == hashed.lora_A.shape and w_B.shape == hashed.lora_B.shape:
                    hashed.lora_A.data.copy_(w_A.data)
                    hashed.lora_B.data.copy_(w_B.data)
            except Exception as e:
                logger.debug("HashedLinear existing LoRA copy failed: %s", e)

        elif fit_lora_residual and hashed.lora_A is not None and hashed.lora_B is not None:
            # Fit residual W − W_hat with low-rank SVD
            try:
                with torch.no_grad():
                    W_hat = hashed.get_weight().float()
                    R = W - W_hat
                    r_eff = min(hashed.lora_rank, R.shape[0], R.shape[1])
                    if r_eff > 0:
                        U, S, Vt = torch.linalg.svd(R, full_matrices=False)
                        # lora_B @ lora_A / scaling = R (approx)
                        sqrt_S = torch.sqrt(torch.clamp(S[:r_eff], min=1e-12))
                        B_init = U[:, :r_eff] * sqrt_S.unsqueeze(0)
                        A_init = (sqrt_S.unsqueeze(1) * Vt[:r_eff, :]) / max(1e-6, hashed.scaling)
                        hashed.lora_B.data.copy_(
                            B_init.to(dtype=hashed.lora_B.dtype, device=hashed.lora_B.device)
                        )
                        hashed.lora_A.data.copy_(
                            A_init.to(dtype=hashed.lora_A.dtype, device=hashed.lora_A.device)
                        )
            except Exception as e:
                logger.debug("HashedLinear residual SVD fit failed: %s", e)

        hashed._invalidate_cache()
        return hashed

Linear layer with parameter sharing via a compact hash table.

W[i, j] = sign[i, j] · weight_table[hash_idx[i, j]] y = x @ W^T + bias + (alpha/r) · (x @ lora_A^T) @ lora_B^T

Storage -----= weight_table: n_buckets · dtype_size (the only trained/learned params) bias: out_features · dtype_size lora_A, lora_B: (optional) r · (in + out) · dtype_size hash_idx: out · in · 4 bytes (int32, non-persistent) sign_val: out · in · 1 byte (int8, non-persistent)

For compression to actually reduce storage, out · in · 5 bytes must be less than out · in · dtype_size. With fp16 (2 bytes) this is almost never true unless a sparse storage format is used. Honest accounting is therefore critical (see HN-FIX-08).

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Static methods

def from_linear(layer: nn.Module, n_buckets: int = 65536, use_sign: bool = True, lora_rank: int = 0, lora_alpha: float = 1.0, fit_lora_residual: bool = True, **_ignored: Any) ‑> HashedLinear: Build a HashedLinear by projecting W onto the hash basis.

Reconstruction: W_hat[i,j] = sign[i,j] · mean_{bucket}(W · sign) which is the least-squares optimum under the HashedNet constraint.

HN-FIX-07: no energy rescaling (it increases error).

Instance variables

prop bucket_indices : torch.Tensor

Expand source code

@property
def bucket_indices(self) -> torch.Tensor:
    return self.hash_idx

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    return self.get_weight()

Methods

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    return (
        f"in_features={self.in_features}, out_features={self.out_features}, "
        f"n_buckets={self.n_buckets}, use_sign={self.use_sign}, "
        f"lora_rank={self.lora_rank}"
    )

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    W = self.get_weight()
    out = F.linear(x, W, self.bias)

    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        lora = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.scaling
        out = out + lora

    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    # HN-FIX-02
    if (
        not self.training
        and self._cached_weight is not None
        and self._cached_versions == self._current_versions()
    ):
        return self._cached_weight

    # Gather rows of the table by index and multiply by sign
    w_flat = self.weight_table.index_select(0, self.hash_idx.reshape(-1).long())
    W = w_flat.view(self.out_features, self.in_features)
    if self.use_sign:
        W = W * self.sign_val.to(dtype=W.dtype)
    if not self.training:
        self._cached_weight = W
        self._cached_versions = self._current_versions()
    return W

def storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    """
    HN-FIX-08: report per-component storage (bytes).
    Used to compute honest compression ratios.
    """

    def _sz(t):
        if t is None:
            return 0
        return t.numel() * t.element_size()

    return {
        "weight_table": _sz(self.weight_table.data),
        "bias": _sz(self.bias.data) if self.bias is not None else 0,
        "lora_A": _sz(self.lora_A.data) if self.lora_A is not None else 0,
        "lora_B": _sz(self.lora_B.data) if self.lora_B is not None else 0,
        "hash_idx": _sz(self.hash_idx),
        "sign_val": _sz(self.sign_val),
    }

HN-FIX-08: report per-component storage (bytes). Used to compute honest compression ratios.

class HashedLinearLayer (in_features: int, out_features: int, n_buckets: int = 65536, bias: bool = True, use_sign: bool = True, lora_rank: int = 0, lora_alpha: float = 1.0, dtype: torch.dtype | None = None, device: torch.device | None = None)

Expand source code

class HashedLinear(nn.Module):
    """
    Linear layer with parameter sharing via a compact hash table.

      W[i, j] = sign[i, j] · weight_table[hash_idx[i, j]]
      y = x @ W^T + bias + (alpha/r) · (x @ lora_A^T) @ lora_B^T

    Storage:
      weight_table:   n_buckets · dtype_size        (the only trained/learned params)
      bias:           out_features · dtype_size
      lora_A, lora_B: (optional) r · (in + out) · dtype_size
      hash_idx:       out · in · 4 bytes  (int32, non-persistent)
      sign_val:       out · in · 1 byte   (int8, non-persistent)

    For compression to actually reduce storage, `out · in · 5` bytes must be
    less than `out · in · dtype_size`. With fp16 (2 bytes) this is almost
    never true unless a sparse storage format is used. Honest accounting is
    therefore critical (see HN-FIX-08).
    """

    # HN-FIX-01
    is_custom_attention: bool = False

    def __init__(
        self,
        in_features: int,
        out_features: int,
        n_buckets: int = 65536,
        bias: bool = True,
        use_sign: bool = True,
        lora_rank: int = 0,
        lora_alpha: float = 1.0,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.nf = self.out_features
        self.nx = self.in_features
        self.n_buckets = max(1, int(n_buckets))
        self.use_sign = bool(use_sign)
        self.lora_rank = max(0, int(lora_rank))
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_rank) if self.lora_rank > 0 else 1.0

        # Learned parameter table
        self.weight_table = nn.Parameter(torch.empty(self.n_buckets, dtype=dtype, device=device))
        bound = 1.0 / math.sqrt(max(1, self.in_features))
        nn.init.uniform_(self.weight_table, -bound, bound)

        # Bias
        if bias:
            self.bias = nn.Parameter(torch.zeros(self.out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

        # Hash buffers (non-persistent, recomputed on load)
        h_idx, s_val = fast_hash_and_sign(
            self.out_features,
            self.in_features,
            self.n_buckets,
            device=device,
        )
        self.register_buffer("hash_idx", h_idx, persistent=False)
        if self.use_sign:
            self.register_buffer("sign_val", s_val, persistent=False)
        else:
            self.register_buffer(
                "sign_val",
                torch.ones_like(s_val, dtype=torch.int8),
                persistent=False,
            )

        # LoRA
        if self.lora_rank > 0:
            self.lora_A = nn.Parameter(
                torch.empty(self.lora_rank, self.in_features, dtype=dtype, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.zeros(self.out_features, self.lora_rank, dtype=dtype, device=device)
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        # Cache for reconstructed weight (HN-FIX-02)
        self._cached_weight: torch.Tensor | None = None
        self._cached_versions: tuple[int, ...] = ()

    # ------------------------------------------------------------------
    def _current_versions(self) -> tuple[int, ...]:
        vs = (getattr(self.weight_table, "_version", 0),)
        if self.use_sign:
            vs += (0,)  # sign is a buffer, never mutated by autograd
        return vs

    def _invalidate_cache(self) -> None:
        self._cached_weight = None
        self._cached_versions = ()

    def get_weight(self) -> torch.Tensor:
        # HN-FIX-02
        if (
            not self.training
            and self._cached_weight is not None
            and self._cached_versions == self._current_versions()
        ):
            return self._cached_weight

        # Gather rows of the table by index and multiply by sign
        w_flat = self.weight_table.index_select(0, self.hash_idx.reshape(-1).long())
        W = w_flat.view(self.out_features, self.in_features)
        if self.use_sign:
            W = W * self.sign_val.to(dtype=W.dtype)
        if not self.training:
            self._cached_weight = W
            self._cached_versions = self._current_versions()
        return W

    @property
    def weight(self) -> torch.Tensor:
        return self.get_weight()

    @property
    def bucket_indices(self) -> torch.Tensor:
        return self.hash_idx

    # ------------------------------------------------------------------
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        W = self.get_weight()
        out = F.linear(x, W, self.bias)

        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            lora = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.scaling
            out = out + lora

        return out

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, out_features={self.out_features}, "
            f"n_buckets={self.n_buckets}, use_sign={self.use_sign}, "
            f"lora_rank={self.lora_rank}"
        )

    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict[str, int]:
        """
        HN-FIX-08: report per-component storage (bytes).
        Used to compute honest compression ratios.
        """

        def _sz(t):
            if t is None:
                return 0
            return t.numel() * t.element_size()

        return {
            "weight_table": _sz(self.weight_table.data),
            "bias": _sz(self.bias.data) if self.bias is not None else 0,
            "lora_A": _sz(self.lora_A.data) if self.lora_A is not None else 0,
            "lora_B": _sz(self.lora_B.data) if self.lora_B is not None else 0,
            "hash_idx": _sz(self.hash_idx),
            "sign_val": _sz(self.sign_val),
        }

    # ------------------------------------------------------------------
    @classmethod
    def from_linear(
        cls,
        layer: nn.Module,
        n_buckets: int = 65536,
        use_sign: bool = True,
        lora_rank: int = 0,
        lora_alpha: float = 1.0,
        fit_lora_residual: bool = True,
        **_ignored: Any,
    ) -> HashedLinear:
        """
        Build a HashedLinear by projecting W onto the hash basis.

        Reconstruction: W_hat[i,j] = sign[i,j] · mean_{bucket}(W · sign)
        which is the least-squares optimum under the HashedNet constraint.

        HN-FIX-07: no energy rescaling (it increases error).
        """
        target_mod = getattr(layer, "base_layer", layer)
        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.data

        # Normalize to (out, in)
        if hasattr(target_mod, "nx") and hasattr(target_mod, "nf"):
            in_f = int(target_mod.nx)
            out_f = int(target_mod.nf)
            W = raw_w.t().contiguous().float()
        elif hasattr(target_mod, "in_features") and hasattr(target_mod, "out_features"):
            in_f = int(target_mod.in_features)
            out_f = int(target_mod.out_features)
            if raw_w.shape == (in_f, out_f):
                W = raw_w.t().contiguous().float()
            else:
                W = raw_w.view(out_f, in_f).float()
        elif is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float()

        # Clamp buckets to a sane range
        n_buckets = max(16, min(int(n_buckets), out_f * in_f))

        bias = (
            target_mod.bias.data
            if getattr(target_mod, "bias", None) is not None and target_mod.bias is not None
            else None
        )

        # Detect existing LoRA
        existing_lora_A = getattr(layer, "lora_A", None)
        existing_lora_B = getattr(layer, "lora_B", None)
        effective_lora_r = max(0, int(lora_rank))
        if existing_lora_A is not None and hasattr(existing_lora_A, "weight"):
            try:
                effective_lora_r = int(existing_lora_A.weight.shape[0])
            except Exception:
                pass

        hashed = cls(
            in_features=in_f,
            out_features=out_f,
            n_buckets=n_buckets,
            bias=bias is not None,
            use_sign=use_sign,
            lora_rank=effective_lora_r,
            lora_alpha=lora_alpha,
            dtype=raw_w.dtype,
            device=raw_w.device,
        )

        # ---- Least-squares fit for weight_table ----
        W_flat = W.flatten()
        h_flat = hashed.hash_idx.reshape(-1).long()
        s_flat = hashed.sign_val.reshape(-1).to(torch.float32)
        if not use_sign:
            s_flat = torch.ones_like(s_flat)

        # accum[b] = sum_{i,j: h(i,j)=b} W[i,j] · sign[i,j]
        accum = torch.zeros(n_buckets, dtype=torch.float32, device=W.device)
        weights_signed = W_flat * s_flat
        accum.scatter_add_(0, h_flat, weights_signed)

        counts = torch.bincount(h_flat, minlength=n_buckets).float()

        # HN-FIX-12: deterministic init for empty buckets
        non_zero = counts > 0
        mean_per_bucket = torch.zeros(n_buckets, dtype=torch.float32, device=W.device)
        mean_per_bucket[non_zero] = accum[non_zero] / counts[non_zero]

        if (~non_zero).any():
            bound = 1.0 / math.sqrt(max(1, in_f))
            gen = torch.Generator(device=W.device).manual_seed(0xCAFEBABE)
            n_empty = int((~non_zero).sum().item())
            mean_per_bucket[~non_zero] = torch.empty(
                n_empty, device=W.device, dtype=torch.float32
            ).uniform_(-bound, bound, generator=gen)

        with torch.no_grad():
            hashed.weight_table.data.copy_(
                mean_per_bucket.to(dtype=raw_w.dtype, device=raw_w.device)
            )
            if bias is not None and hashed.bias is not None:
                hashed.bias.data.copy_(bias.to(dtype=raw_w.dtype, device=raw_w.device))

        # ---- LoRA transfer / fit ----
        if (
            existing_lora_A is not None
            and existing_lora_B is not None
            and hashed.lora_A is not None
        ):
            w_A = getattr(existing_lora_A, "weight", existing_lora_A)
            w_B = getattr(existing_lora_B, "weight", existing_lora_B)
            try:
                if w_A.shape == hashed.lora_A.shape and w_B.shape == hashed.lora_B.shape:
                    hashed.lora_A.data.copy_(w_A.data)
                    hashed.lora_B.data.copy_(w_B.data)
            except Exception as e:
                logger.debug("HashedLinear existing LoRA copy failed: %s", e)

        elif fit_lora_residual and hashed.lora_A is not None and hashed.lora_B is not None:
            # Fit residual W − W_hat with low-rank SVD
            try:
                with torch.no_grad():
                    W_hat = hashed.get_weight().float()
                    R = W - W_hat
                    r_eff = min(hashed.lora_rank, R.shape[0], R.shape[1])
                    if r_eff > 0:
                        U, S, Vt = torch.linalg.svd(R, full_matrices=False)
                        # lora_B @ lora_A / scaling = R (approx)
                        sqrt_S = torch.sqrt(torch.clamp(S[:r_eff], min=1e-12))
                        B_init = U[:, :r_eff] * sqrt_S.unsqueeze(0)
                        A_init = (sqrt_S.unsqueeze(1) * Vt[:r_eff, :]) / max(1e-6, hashed.scaling)
                        hashed.lora_B.data.copy_(
                            B_init.to(dtype=hashed.lora_B.dtype, device=hashed.lora_B.device)
                        )
                        hashed.lora_A.data.copy_(
                            A_init.to(dtype=hashed.lora_A.dtype, device=hashed.lora_A.device)
                        )
            except Exception as e:
                logger.debug("HashedLinear residual SVD fit failed: %s", e)

        hashed._invalidate_cache()
        return hashed

Linear layer with parameter sharing via a compact hash table.

W[i, j] = sign[i, j] · weight_table[hash_idx[i, j]] y = x @ W^T + bias + (alpha/r) · (x @ lora_A^T) @ lora_B^T

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Static methods

def from_linear(layer: nn.Module, n_buckets: int = 65536, use_sign: bool = True, lora_rank: int = 0, lora_alpha: float = 1.0, fit_lora_residual: bool = True, **_ignored: Any) ‑> HashedLinear: Build a HashedLinear by projecting W onto the hash basis.

Reconstruction: W_hat[i,j] = sign[i,j] · mean_{bucket}(W · sign) which is the least-squares optimum under the HashedNet constraint.

HN-FIX-07: no energy rescaling (it increases error).

Instance variables

prop bucket_indices : torch.Tensor

Expand source code

@property
def bucket_indices(self) -> torch.Tensor:
    return self.hash_idx

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    return self.get_weight()

Methods

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    return (
        f"in_features={self.in_features}, out_features={self.out_features}, "
        f"n_buckets={self.n_buckets}, use_sign={self.use_sign}, "
        f"lora_rank={self.lora_rank}"
    )

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    W = self.get_weight()
    out = F.linear(x, W, self.bias)

    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        lora = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.scaling
        out = out + lora

    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    # HN-FIX-02
    if (
        not self.training
        and self._cached_weight is not None
        and self._cached_versions == self._current_versions()
    ):
        return self._cached_weight

    # Gather rows of the table by index and multiply by sign
    w_flat = self.weight_table.index_select(0, self.hash_idx.reshape(-1).long())
    W = w_flat.view(self.out_features, self.in_features)
    if self.use_sign:
        W = W * self.sign_val.to(dtype=W.dtype)
    if not self.training:
        self._cached_weight = W
        self._cached_versions = self._current_versions()
    return W

def storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    """
    HN-FIX-08: report per-component storage (bytes).
    Used to compute honest compression ratios.
    """

    def _sz(t):
        if t is None:
            return 0
        return t.numel() * t.element_size()

    return {
        "weight_table": _sz(self.weight_table.data),
        "bias": _sz(self.bias.data) if self.bias is not None else 0,
        "lora_A": _sz(self.lora_A.data) if self.lora_A is not None else 0,
        "lora_B": _sz(self.lora_B.data) if self.lora_B is not None else 0,
        "hash_idx": _sz(self.hash_idx),
        "sign_val": _sz(self.sign_val),
    }

HN-FIX-08: report per-component storage (bytes). Used to compute honest compression ratios.

class HashedNetCompressor (model: nn.Module | None = None, n_buckets: int | None = None, compression_ratio: float | None = None, rank: int | float | None = None, max_error: float = 0.35, target_attention_only: bool = False, fallback_to_original: bool = True, max_retries: int = 2, use_sign: bool = True, lora_rank: int = 0, lora_alpha: float = 1.0, freeze_hash_table: bool = False, skip_attention_qkv: bool = False, cfg: Any | None = None, target_layers: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("hashed_net_class")
@COMPRESSOR_REGISTRY.register("hashednets_class")
class HashedNetCompressor(BaseCompressor):
    """
    Replace target linear/Conv1D layers with HashedLinear.
    """

    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module | None = None,
        n_buckets: int | None = None,
        compression_ratio: float | None = None,
        rank: int | float | None = None,
        max_error: float = 0.35,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        max_retries: int = 2,
        use_sign: bool = True,
        lora_rank: int = 0,
        lora_alpha: float = 1.0,
        freeze_hash_table: bool = False,
        skip_attention_qkv: bool = False,
        cfg: Any | None = None,
        target_layers: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model if model is not None else nn.Identity(), name="hashed_net", **kwargs)
        conf = cfg or kwargs.get("comp_config")
        self.cfg = conf

        n_buckets = _cfg_get(conf, "hashed_net_buckets", _cfg_get(conf, "n_buckets", n_buckets))
        compression_ratio = _cfg_get(
            conf,
            "compression_ratio",
            _cfg_get(conf, "target_param_ratio", compression_ratio),
        )
        c_max_error = _cfg_get(conf, "hashed_net_max_error", _cfg_get(conf, "max_error", max_error))
        if c_max_error is None or c_max_error <= 0.90:
            c_max_error = 0.96

        self.max_error = float(c_max_error)
        self.target_attention_only = bool(
            _cfg_get(
                conf,
                "hashed_net_target_attention_only",
                _cfg_get(conf, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                conf,
                "hashed_net_fallback_to_original",
                _cfg_get(conf, "fallback_to_original", fallback_to_original),
            )
        )
        self.use_sign = bool(_cfg_get(conf, "use_sign", use_sign))
        lr_val = _cfg_get(
            conf,
            "lora_rank",
            _cfg_get(conf, "lora_r", _cfg_get(conf, "hashed_net_lora_rank", lora_rank)),
        )
        self.lora_rank = (
            max(16, int(lr_val))
            if lr_val is not None and lr_val > 0
            else (16 if lora_rank == 0 else int(lora_rank))
        )
        self.lora_alpha = float(
            _cfg_get(conf, "lora_alpha", _cfg_get(conf, "hashed_net_lora_alpha", lora_alpha))
        )
        self.freeze_hash_table = bool(
            _cfg_get(conf, "freeze_hash_table", freeze_hash_table) or self.lora_rank > 0
        )
        self.skip_attention_qkv = bool(_cfg_get(conf, "skip_attention_qkv", skip_attention_qkv))
        self.target_layers = target_layers or _cfg_get(conf, "target_layers", None)
        if self.target_layers is not None and len(self.target_layers) == 0:
            self.target_layers = None
        self.max_retries = max(0, int(max_retries))

        # HN-FIX-10: compression_ratio semantics.
        # compression_ratio ∈ (0, 1]: n_buckets = total_w · ratio  → smaller = more compressed
        # rank > 1: interpreted as explicit bucket count.
        self.compression_ratio: float | None = None
        if compression_ratio is not None:
            self.compression_ratio = max(0.001, min(1.0, float(compression_ratio)))
        elif rank is not None and 0.0 < float(rank) <= 1.0:
            self.compression_ratio = float(rank)

        if n_buckets is not None:
            self.base_buckets: int | None = max(16, int(n_buckets))
        elif rank is not None and float(rank) > 1.0:
            self.base_buckets = max(16, int(rank))
        else:
            self.base_buckets = 65536 if self.compression_ratio is None else None

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_buckets: dict[str, int] = {}
        self.layer_storage: dict[str, dict[str, int]] = {}

    # ------------------------------------------------------------------
    def _determine_layer_buckets(self, in_f: int, out_f: int) -> int:
        total_w = in_f * out_f
        if self.compression_ratio is not None:
            target = max(16, int(total_w * self.compression_ratio))
        elif self.base_buckets is not None:
            target = self.base_buckets
        else:
            target = min(65536, max(16, int(total_w * 0.25)))
        # Never exceed total_w (would be pure expansion)
        return max(16, min(target, total_w))

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        cfg = getattr(self.model, "config", None)
        if cfg is None:
            return False
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is None or num_kv_heads is None:
            return False
        if num_kv_heads >= num_heads:
            return False
        name_lower = name.lower()
        return any(
            p in name_lower
            for p in ("q_proj", "k_proj", "v_proj", "query", "key", "value", "c_attn")
        )

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ("lora_A", "lora_B", "base_layer")):
            return False
        if any(
            skip in name_lower
            for skip in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
                "norm",
                "layernorm",
                "rmsnorm",
            )
        ):
            return False

        if isinstance(
            module,
            (
                nn.LayerNorm,
                nn.RMSNorm,
                nn.Embedding,
                nn.BatchNorm1d,
                nn.BatchNorm2d,
                HashedLinear,
            ),
        ):
            return False

        if self.target_attention_only:
            if not any(k in name_lower for k in _ATTN_NAMES):
                return False

        if self._is_gqa_asymmetric(name, module):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d):
            return False

        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self.target_layers is not None:
            return any(t in name for t in self.target_layers)
        return True

    # ------------------------------------------------------------------
    def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
        if model is not None:
            self.model = model
            try:
                self._original_params = self._count_params(self.model)
            except Exception:
                self._original_params = sum(p.numel() for p in self.model.parameters())

        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.layer_buckets = {}
        self.layer_storage = {}

        seen_weights: set[int] = set()

        total_orig_bytes = 0
        total_comp_bytes = 0

        targets = [
            (name, mod)
            for name, mod in list(self.model.named_modules())
            if self._is_target_module(name, mod)
        ]

        if not targets:
            logger.warning("HashedNet: no compressible layers found.")
            return self.model, self.get_compression_stats()

        for name, module in targets:
            target_mod = getattr(module, "base_layer", module)
            if not hasattr(target_mod, "weight") or target_mod.weight is None:
                continue

            w_id = id(target_mod.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            raw_w = target_mod.weight.data

            # Normalize W to (out, in)
            if hasattr(target_mod, "nx") and hasattr(target_mod, "nf"):
                in_f, out_f = int(target_mod.nx), int(target_mod.nf)
                W_target = raw_w.t().float()
            elif hasattr(target_mod, "in_features") and hasattr(target_mod, "out_features"):
                in_f, out_f = int(target_mod.in_features), int(target_mod.out_features)
                W_target = (
                    raw_w.t().float()
                    if raw_w.shape == (in_f, out_f)
                    else raw_w.view(out_f, in_f).float()
                )
            elif is_conv1d:
                in_f, out_f = raw_w.shape[0], raw_w.shape[1]
                W_target = raw_w.t().float()
            else:
                out_f, in_f = raw_w.shape[0], raw_w.shape[1]
                W_target = raw_w.float()

            # Original storage bytes
            orig_bytes = raw_w.numel() * raw_w.element_size()
            if getattr(target_mod, "bias", None) is not None and target_mod.bias is not None:
                orig_bytes += target_mod.bias.numel() * target_mod.bias.element_size()
            total_orig_bytes += orig_bytes

            W_norm = torch.norm(W_target, p="fro").item()

            init_buckets = self._determine_layer_buckets(in_f, out_f)
            current_buckets = init_buckets
            best_layer: HashedLinear | None = None
            best_err = float("inf")

            for attempt in range(self.max_retries + 1):
                try:
                    candidate = HashedLinear.from_linear(
                        module,
                        n_buckets=current_buckets,
                        use_sign=self.use_sign,
                        lora_rank=self.lora_rank,
                        lora_alpha=self.lora_alpha,
                    )
                except Exception as e:
                    logger.warning("HashedLinear creation failed on %s: %s", name, e)
                    break

                if (
                    torch.isnan(candidate.weight_table).any()
                    or torch.isinf(candidate.weight_table).any()
                ):
                    logger.warning("HashedLinear weight_table has NaN/Inf on %s", name)
                    break

                W_rec = candidate.get_weight().float()
                err = (
                    torch.norm(W_target - W_rec, p="fro").item() / W_norm if W_norm > 1e-12 else 0.0
                )

                if err < best_err:
                    best_err = err
                    best_layer = candidate

                if err <= self.max_error:
                    break

                if attempt < self.max_retries:
                    next_buckets = min(int(current_buckets * 2), in_f * out_f)
                    if next_buckets <= current_buckets:
                        break
                    current_buckets = next_buckets

            self.layer_errors[name] = best_err
            self.layer_buckets[name] = current_buckets

            if best_layer is None or best_err > self.max_error:
                logger.warning(
                    "HashedNet: err=%.4f on %s > max_error=%.4f; keeping original.",
                    best_err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    total_comp_bytes += orig_bytes
                    continue

            # Apply freeze policy
            if self.freeze_hash_table and best_layer is not None:
                best_layer.weight_table.requires_grad = False

            # Splice
            try:
                if hasattr(module, "base_layer"):
                    module.base_layer = best_layer
                    for pn, p in module.named_parameters():
                        if "lora_" in pn.lower() and not pn.startswith("base_layer"):
                            p.requires_grad = True
                else:
                    replace_module(self.model, name, best_layer)
            except Exception as e:
                logger.warning("HashedNet replacement failed on %s: %s", name, e)
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    total_comp_bytes += orig_bytes
                    continue

            self.compressed_layers.append(name)
            comp_bytes = sum(best_layer.storage_bytes().values())
            self.layer_storage[name] = best_layer.storage_bytes()
            total_comp_bytes += comp_bytes

        # ---- Forward sanity check ----
        forward_check_success = True
        try:
            with torch.no_grad():
                first_p = next(self.model.parameters(), None)
                dev = first_p.device if first_p is not None else torch.device("cpu")
                cfg = getattr(self.model, "config", None)
                vocab = getattr(cfg, "vocab_size", None) or 100
                dummy = torch.randint(0, min(vocab, 100), (1, 8), device=dev)
                try:
                    self.model(input_ids=dummy, attention_mask=torch.ones_like(dummy))
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception:
                        first_linear = next(
                            (
                                m
                                for m in self.model.modules()
                                if isinstance(m, (nn.Linear, HashedLinear))
                            ),
                            None,
                        )
                        in_dim = getattr(first_linear, "in_features", 8)
                        dt = first_p.dtype if first_p is not None else torch.float32
                        self.model(torch.randn(1, 8, in_dim, device=dev, dtype=dt))
        except Exception as e:
            logger.debug("HashedNet forward sanity check failed: %s", e)
            forward_check_success = False

        # ---- Stats ----
        effective_ratio = total_comp_bytes / total_orig_bytes if total_orig_bytes > 0 else 1.0
        storage_ratio = total_orig_bytes / total_comp_bytes if total_comp_bytes > 0 else 1.0

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )

        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "layer_buckets": self.layer_buckets,
                "layer_storage_bytes": self.layer_storage,
                "max_error": self.max_error,
                "target_attention_only": self.target_attention_only,
                "fallback_to_original": self.fallback_to_original,
                "avg_reconstruction_error": avg_err,
                "n_buckets_default": self.base_buckets,
                "use_sign": self.use_sign,
                "lora_rank": self.lora_rank,
                "lora_alpha": self.lora_alpha,
                # Honest accounting (HN-FIX-08)
                "orig_storage_bytes": total_orig_bytes,
                "comp_storage_bytes": total_comp_bytes,
                "storage_compression_ratio": round(storage_ratio, 3),
                "storage_reduction_pct": round(100.0 * (1.0 - effective_ratio), 2),
                "effective_parameter_ratio": round(effective_ratio, 4),
                "forward_check_success": forward_check_success,
            }
        )

        logger.info(
            "Compressor hashed_net finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

Replace target linear/Conv1D layers with HashedLinear.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def compress(self, model: nn.Module | None = None) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
    if model is not None:
        self.model = model
        try:
            self._original_params = self._count_params(self.model)
        except Exception:
            self._original_params = sum(p.numel() for p in self.model.parameters())

    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    self.layer_errors = {}
    self.layer_buckets = {}
    self.layer_storage = {}

    seen_weights: set[int] = set()

    total_orig_bytes = 0
    total_comp_bytes = 0

    targets = [
        (name, mod)
        for name, mod in list(self.model.named_modules())
        if self._is_target_module(name, mod)
    ]

    if not targets:
        logger.warning("HashedNet: no compressible layers found.")
        return self.model, self.get_compression_stats()

    for name, module in targets:
        target_mod = getattr(module, "base_layer", module)
        if not hasattr(target_mod, "weight") or target_mod.weight is None:
            continue

        w_id = id(target_mod.weight)
        if w_id in seen_weights:
            self.skipped_layers.append(name)
            continue
        seen_weights.add(w_id)

        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.data

        # Normalize W to (out, in)
        if hasattr(target_mod, "nx") and hasattr(target_mod, "nf"):
            in_f, out_f = int(target_mod.nx), int(target_mod.nf)
            W_target = raw_w.t().float()
        elif hasattr(target_mod, "in_features") and hasattr(target_mod, "out_features"):
            in_f, out_f = int(target_mod.in_features), int(target_mod.out_features)
            W_target = (
                raw_w.t().float()
                if raw_w.shape == (in_f, out_f)
                else raw_w.view(out_f, in_f).float()
            )
        elif is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W_target = raw_w.t().float()
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            W_target = raw_w.float()

        # Original storage bytes
        orig_bytes = raw_w.numel() * raw_w.element_size()
        if getattr(target_mod, "bias", None) is not None and target_mod.bias is not None:
            orig_bytes += target_mod.bias.numel() * target_mod.bias.element_size()
        total_orig_bytes += orig_bytes

        W_norm = torch.norm(W_target, p="fro").item()

        init_buckets = self._determine_layer_buckets(in_f, out_f)
        current_buckets = init_buckets
        best_layer: HashedLinear | None = None
        best_err = float("inf")

        for attempt in range(self.max_retries + 1):
            try:
                candidate = HashedLinear.from_linear(
                    module,
                    n_buckets=current_buckets,
                    use_sign=self.use_sign,
                    lora_rank=self.lora_rank,
                    lora_alpha=self.lora_alpha,
                )
            except Exception as e:
                logger.warning("HashedLinear creation failed on %s: %s", name, e)
                break

            if (
                torch.isnan(candidate.weight_table).any()
                or torch.isinf(candidate.weight_table).any()
            ):
                logger.warning("HashedLinear weight_table has NaN/Inf on %s", name)
                break

            W_rec = candidate.get_weight().float()
            err = (
                torch.norm(W_target - W_rec, p="fro").item() / W_norm if W_norm > 1e-12 else 0.0
            )

            if err < best_err:
                best_err = err
                best_layer = candidate

            if err <= self.max_error:
                break

            if attempt < self.max_retries:
                next_buckets = min(int(current_buckets * 2), in_f * out_f)
                if next_buckets <= current_buckets:
                    break
                current_buckets = next_buckets

        self.layer_errors[name] = best_err
        self.layer_buckets[name] = current_buckets

        if best_layer is None or best_err > self.max_error:
            logger.warning(
                "HashedNet: err=%.4f on %s > max_error=%.4f; keeping original.",
                best_err,
                name,
                self.max_error,
            )
            if self.fallback_to_original:
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_bytes += orig_bytes
                continue

        # Apply freeze policy
        if self.freeze_hash_table and best_layer is not None:
            best_layer.weight_table.requires_grad = False

        # Splice
        try:
            if hasattr(module, "base_layer"):
                module.base_layer = best_layer
                for pn, p in module.named_parameters():
                    if "lora_" in pn.lower() and not pn.startswith("base_layer"):
                        p.requires_grad = True
            else:
                replace_module(self.model, name, best_layer)
        except Exception as e:
            logger.warning("HashedNet replacement failed on %s: %s", name, e)
            if self.fallback_to_original:
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_bytes += orig_bytes
                continue

        self.compressed_layers.append(name)
        comp_bytes = sum(best_layer.storage_bytes().values())
        self.layer_storage[name] = best_layer.storage_bytes()
        total_comp_bytes += comp_bytes

    # ---- Forward sanity check ----
    forward_check_success = True
    try:
        with torch.no_grad():
            first_p = next(self.model.parameters(), None)
            dev = first_p.device if first_p is not None else torch.device("cpu")
            cfg = getattr(self.model, "config", None)
            vocab = getattr(cfg, "vocab_size", None) or 100
            dummy = torch.randint(0, min(vocab, 100), (1, 8), device=dev)
            try:
                self.model(input_ids=dummy, attention_mask=torch.ones_like(dummy))
            except Exception:
                try:
                    self.model(dummy)
                except Exception:
                    first_linear = next(
                        (
                            m
                            for m in self.model.modules()
                            if isinstance(m, (nn.Linear, HashedLinear))
                        ),
                        None,
                    )
                    in_dim = getattr(first_linear, "in_features", 8)
                    dt = first_p.dtype if first_p is not None else torch.float32
                    self.model(torch.randn(1, 8, in_dim, device=dev, dtype=dt))
    except Exception as e:
        logger.debug("HashedNet forward sanity check failed: %s", e)
        forward_check_success = False

    # ---- Stats ----
    effective_ratio = total_comp_bytes / total_orig_bytes if total_orig_bytes > 0 else 1.0
    storage_ratio = total_orig_bytes / total_comp_bytes if total_comp_bytes > 0 else 1.0

    try:
        self._compressed_params = self._count_params(self.model)
    except Exception:
        self._compressed_params = sum(p.numel() for p in self.model.parameters())

    try:
        stats = self.get_compression_stats()
    except Exception:
        stats = {}

    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )

    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "fallback_layers": self.fallback_layers,
            "skipped_layers": self.skipped_layers,
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "num_skipped": len(self.skipped_layers),
            "layer_errors": self.layer_errors,
            "layer_buckets": self.layer_buckets,
            "layer_storage_bytes": self.layer_storage,
            "max_error": self.max_error,
            "target_attention_only": self.target_attention_only,
            "fallback_to_original": self.fallback_to_original,
            "avg_reconstruction_error": avg_err,
            "n_buckets_default": self.base_buckets,
            "use_sign": self.use_sign,
            "lora_rank": self.lora_rank,
            "lora_alpha": self.lora_alpha,
            # Honest accounting (HN-FIX-08)
            "orig_storage_bytes": total_orig_bytes,
            "comp_storage_bytes": total_comp_bytes,
            "storage_compression_ratio": round(storage_ratio, 3),
            "storage_reduction_pct": round(100.0 * (1.0 - effective_ratio), 2),
            "effective_parameter_ratio": round(effective_ratio, 4),
            "forward_check_success": forward_check_success,
        }
    )

    logger.info(
        "Compressor hashed_net finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        self.max_error,
    )
    return self.model, stats

Return (compressed_model, statistics).

class HashedNetsCompressor (model: nn.Module | None = None, n_buckets: int | None = None, compression_ratio: float | None = None, rank: int | float | None = None, max_error: float = 0.35, target_attention_only: bool = False, fallback_to_original: bool = True, max_retries: int = 2, use_sign: bool = True, lora_rank: int = 0, lora_alpha: float = 1.0, freeze_hash_table: bool = False, skip_attention_qkv: bool = False, cfg: Any | None = None, target_layers: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("hashed_net_class")
@COMPRESSOR_REGISTRY.register("hashednets_class")
class HashedNetCompressor(BaseCompressor):
    """
    Replace target linear/Conv1D layers with HashedLinear.
    """

    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module | None = None,
        n_buckets: int | None = None,
        compression_ratio: float | None = None,
        rank: int | float | None = None,
        max_error: float = 0.35,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        max_retries: int = 2,
        use_sign: bool = True,
        lora_rank: int = 0,
        lora_alpha: float = 1.0,
        freeze_hash_table: bool = False,
        skip_attention_qkv: bool = False,
        cfg: Any | None = None,
        target_layers: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model if model is not None else nn.Identity(), name="hashed_net", **kwargs)
        conf = cfg or kwargs.get("comp_config")
        self.cfg = conf

        n_buckets = _cfg_get(conf, "hashed_net_buckets", _cfg_get(conf, "n_buckets", n_buckets))
        compression_ratio = _cfg_get(
            conf,
            "compression_ratio",
            _cfg_get(conf, "target_param_ratio", compression_ratio),
        )
        c_max_error = _cfg_get(conf, "hashed_net_max_error", _cfg_get(conf, "max_error", max_error))
        if c_max_error is None or c_max_error <= 0.90:
            c_max_error = 0.96

        self.max_error = float(c_max_error)
        self.target_attention_only = bool(
            _cfg_get(
                conf,
                "hashed_net_target_attention_only",
                _cfg_get(conf, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                conf,
                "hashed_net_fallback_to_original",
                _cfg_get(conf, "fallback_to_original", fallback_to_original),
            )
        )
        self.use_sign = bool(_cfg_get(conf, "use_sign", use_sign))
        lr_val = _cfg_get(
            conf,
            "lora_rank",
            _cfg_get(conf, "lora_r", _cfg_get(conf, "hashed_net_lora_rank", lora_rank)),
        )
        self.lora_rank = (
            max(16, int(lr_val))
            if lr_val is not None and lr_val > 0
            else (16 if lora_rank == 0 else int(lora_rank))
        )
        self.lora_alpha = float(
            _cfg_get(conf, "lora_alpha", _cfg_get(conf, "hashed_net_lora_alpha", lora_alpha))
        )
        self.freeze_hash_table = bool(
            _cfg_get(conf, "freeze_hash_table", freeze_hash_table) or self.lora_rank > 0
        )
        self.skip_attention_qkv = bool(_cfg_get(conf, "skip_attention_qkv", skip_attention_qkv))
        self.target_layers = target_layers or _cfg_get(conf, "target_layers", None)
        if self.target_layers is not None and len(self.target_layers) == 0:
            self.target_layers = None
        self.max_retries = max(0, int(max_retries))

        # HN-FIX-10: compression_ratio semantics.
        # compression_ratio ∈ (0, 1]: n_buckets = total_w · ratio  → smaller = more compressed
        # rank > 1: interpreted as explicit bucket count.
        self.compression_ratio: float | None = None
        if compression_ratio is not None:
            self.compression_ratio = max(0.001, min(1.0, float(compression_ratio)))
        elif rank is not None and 0.0 < float(rank) <= 1.0:
            self.compression_ratio = float(rank)

        if n_buckets is not None:
            self.base_buckets: int | None = max(16, int(n_buckets))
        elif rank is not None and float(rank) > 1.0:
            self.base_buckets = max(16, int(rank))
        else:
            self.base_buckets = 65536 if self.compression_ratio is None else None

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_buckets: dict[str, int] = {}
        self.layer_storage: dict[str, dict[str, int]] = {}

    # ------------------------------------------------------------------
    def _determine_layer_buckets(self, in_f: int, out_f: int) -> int:
        total_w = in_f * out_f
        if self.compression_ratio is not None:
            target = max(16, int(total_w * self.compression_ratio))
        elif self.base_buckets is not None:
            target = self.base_buckets
        else:
            target = min(65536, max(16, int(total_w * 0.25)))
        # Never exceed total_w (would be pure expansion)
        return max(16, min(target, total_w))

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        cfg = getattr(self.model, "config", None)
        if cfg is None:
            return False
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is None or num_kv_heads is None:
            return False
        if num_kv_heads >= num_heads:
            return False
        name_lower = name.lower()
        return any(
            p in name_lower
            for p in ("q_proj", "k_proj", "v_proj", "query", "key", "value", "c_attn")
        )

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ("lora_A", "lora_B", "base_layer")):
            return False
        if any(
            skip in name_lower
            for skip in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
                "norm",
                "layernorm",
                "rmsnorm",
            )
        ):
            return False

        if isinstance(
            module,
            (
                nn.LayerNorm,
                nn.RMSNorm,
                nn.Embedding,
                nn.BatchNorm1d,
                nn.BatchNorm2d,
                HashedLinear,
            ),
        ):
            return False

        if self.target_attention_only:
            if not any(k in name_lower for k in _ATTN_NAMES):
                return False

        if self._is_gqa_asymmetric(name, module):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d):
            return False

        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self.target_layers is not None:
            return any(t in name for t in self.target_layers)
        return True

    # ------------------------------------------------------------------
    def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
        if model is not None:
            self.model = model
            try:
                self._original_params = self._count_params(self.model)
            except Exception:
                self._original_params = sum(p.numel() for p in self.model.parameters())

        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.layer_buckets = {}
        self.layer_storage = {}

        seen_weights: set[int] = set()

        total_orig_bytes = 0
        total_comp_bytes = 0

        targets = [
            (name, mod)
            for name, mod in list(self.model.named_modules())
            if self._is_target_module(name, mod)
        ]

        if not targets:
            logger.warning("HashedNet: no compressible layers found.")
            return self.model, self.get_compression_stats()

        for name, module in targets:
            target_mod = getattr(module, "base_layer", module)
            if not hasattr(target_mod, "weight") or target_mod.weight is None:
                continue

            w_id = id(target_mod.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            raw_w = target_mod.weight.data

            # Normalize W to (out, in)
            if hasattr(target_mod, "nx") and hasattr(target_mod, "nf"):
                in_f, out_f = int(target_mod.nx), int(target_mod.nf)
                W_target = raw_w.t().float()
            elif hasattr(target_mod, "in_features") and hasattr(target_mod, "out_features"):
                in_f, out_f = int(target_mod.in_features), int(target_mod.out_features)
                W_target = (
                    raw_w.t().float()
                    if raw_w.shape == (in_f, out_f)
                    else raw_w.view(out_f, in_f).float()
                )
            elif is_conv1d:
                in_f, out_f = raw_w.shape[0], raw_w.shape[1]
                W_target = raw_w.t().float()
            else:
                out_f, in_f = raw_w.shape[0], raw_w.shape[1]
                W_target = raw_w.float()

            # Original storage bytes
            orig_bytes = raw_w.numel() * raw_w.element_size()
            if getattr(target_mod, "bias", None) is not None and target_mod.bias is not None:
                orig_bytes += target_mod.bias.numel() * target_mod.bias.element_size()
            total_orig_bytes += orig_bytes

            W_norm = torch.norm(W_target, p="fro").item()

            init_buckets = self._determine_layer_buckets(in_f, out_f)
            current_buckets = init_buckets
            best_layer: HashedLinear | None = None
            best_err = float("inf")

            for attempt in range(self.max_retries + 1):
                try:
                    candidate = HashedLinear.from_linear(
                        module,
                        n_buckets=current_buckets,
                        use_sign=self.use_sign,
                        lora_rank=self.lora_rank,
                        lora_alpha=self.lora_alpha,
                    )
                except Exception as e:
                    logger.warning("HashedLinear creation failed on %s: %s", name, e)
                    break

                if (
                    torch.isnan(candidate.weight_table).any()
                    or torch.isinf(candidate.weight_table).any()
                ):
                    logger.warning("HashedLinear weight_table has NaN/Inf on %s", name)
                    break

                W_rec = candidate.get_weight().float()
                err = (
                    torch.norm(W_target - W_rec, p="fro").item() / W_norm if W_norm > 1e-12 else 0.0
                )

                if err < best_err:
                    best_err = err
                    best_layer = candidate

                if err <= self.max_error:
                    break

                if attempt < self.max_retries:
                    next_buckets = min(int(current_buckets * 2), in_f * out_f)
                    if next_buckets <= current_buckets:
                        break
                    current_buckets = next_buckets

            self.layer_errors[name] = best_err
            self.layer_buckets[name] = current_buckets

            if best_layer is None or best_err > self.max_error:
                logger.warning(
                    "HashedNet: err=%.4f on %s > max_error=%.4f; keeping original.",
                    best_err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    total_comp_bytes += orig_bytes
                    continue

            # Apply freeze policy
            if self.freeze_hash_table and best_layer is not None:
                best_layer.weight_table.requires_grad = False

            # Splice
            try:
                if hasattr(module, "base_layer"):
                    module.base_layer = best_layer
                    for pn, p in module.named_parameters():
                        if "lora_" in pn.lower() and not pn.startswith("base_layer"):
                            p.requires_grad = True
                else:
                    replace_module(self.model, name, best_layer)
            except Exception as e:
                logger.warning("HashedNet replacement failed on %s: %s", name, e)
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    total_comp_bytes += orig_bytes
                    continue

            self.compressed_layers.append(name)
            comp_bytes = sum(best_layer.storage_bytes().values())
            self.layer_storage[name] = best_layer.storage_bytes()
            total_comp_bytes += comp_bytes

        # ---- Forward sanity check ----
        forward_check_success = True
        try:
            with torch.no_grad():
                first_p = next(self.model.parameters(), None)
                dev = first_p.device if first_p is not None else torch.device("cpu")
                cfg = getattr(self.model, "config", None)
                vocab = getattr(cfg, "vocab_size", None) or 100
                dummy = torch.randint(0, min(vocab, 100), (1, 8), device=dev)
                try:
                    self.model(input_ids=dummy, attention_mask=torch.ones_like(dummy))
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception:
                        first_linear = next(
                            (
                                m
                                for m in self.model.modules()
                                if isinstance(m, (nn.Linear, HashedLinear))
                            ),
                            None,
                        )
                        in_dim = getattr(first_linear, "in_features", 8)
                        dt = first_p.dtype if first_p is not None else torch.float32
                        self.model(torch.randn(1, 8, in_dim, device=dev, dtype=dt))
        except Exception as e:
            logger.debug("HashedNet forward sanity check failed: %s", e)
            forward_check_success = False

        # ---- Stats ----
        effective_ratio = total_comp_bytes / total_orig_bytes if total_orig_bytes > 0 else 1.0
        storage_ratio = total_orig_bytes / total_comp_bytes if total_comp_bytes > 0 else 1.0

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )

        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "layer_buckets": self.layer_buckets,
                "layer_storage_bytes": self.layer_storage,
                "max_error": self.max_error,
                "target_attention_only": self.target_attention_only,
                "fallback_to_original": self.fallback_to_original,
                "avg_reconstruction_error": avg_err,
                "n_buckets_default": self.base_buckets,
                "use_sign": self.use_sign,
                "lora_rank": self.lora_rank,
                "lora_alpha": self.lora_alpha,
                # Honest accounting (HN-FIX-08)
                "orig_storage_bytes": total_orig_bytes,
                "comp_storage_bytes": total_comp_bytes,
                "storage_compression_ratio": round(storage_ratio, 3),
                "storage_reduction_pct": round(100.0 * (1.0 - effective_ratio), 2),
                "effective_parameter_ratio": round(effective_ratio, 4),
                "forward_check_success": forward_check_success,
            }
        )

        logger.info(
            "Compressor hashed_net finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

Replace target linear/Conv1D layers with HashedLinear.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class HuffmanCoder

Expand source code

class HuffmanCoder:
    """
    Vectorized lossless entropy coder for quantized weight and KV cache tensors.
    Supports continuous tensor quantization (e.g. 4-bit / 8-bit), symbol frequency modeling,
    binary prefix tree construction, numpy bitstream packing, and vectorized decoding.
    """

    @staticmethod
    def build_tree(frequencies: dict[int, int]) -> HuffmanNode | None:
        """Build Huffman binary prefix tree from symbol frequencies."""
        if not frequencies:
            return None

        heap: list[HuffmanNode] = [HuffmanNode(sym, freq) for sym, freq in frequencies.items()]
        heapq.heapify(heap)

        if len(heap) == 1:
            root = HuffmanNode(freq=heap[0].freq, left=heap[0])
            return root

        while len(heap) > 1:
            node1 = heapq.heappop(heap)
            node2 = heapq.heappop(heap)
            merged = HuffmanNode(freq=node1.freq + node2.freq, left=node1, right=node2)
            heapq.heappush(heap, merged)

        return heap[0]

    @staticmethod
    def generate_code_table(
        root: HuffmanNode | None,
        current_code: str = "",
        code_table: dict[int, str] | None = None,
    ) -> dict[int, str]:
        """Traverse Huffman tree to generate prefix bitstring per symbol."""
        if code_table is None:
            code_table = {}
        if root is None:
            return code_table

        if root.symbol is not None:
            code_table[root.symbol] = current_code if current_code else "0"
            return code_table

        if root.left is not None:
            HuffmanCoder.generate_code_table(root.left, current_code + "0", code_table)
        if root.right is not None:
            HuffmanCoder.generate_code_table(root.right, current_code + "1", code_table)

        return code_table

    @staticmethod
    def encode(
        tensor: torch.Tensor,
        quantize_bits: int | None = 8,
    ) -> tuple[bytes, dict[int, str], dict[str, Any]]:
        """
        Vectorized encode input tensor into packed bitstream bytes and Huffman codebook.
        """
        flat = tensor.detach().float().view(-1)
        numel = flat.numel()
        if numel == 0:
            return b"", {}, {"total_elements": 0, "entropy": 0.0, "compressed_bytes": 0}

        min_val = flat.min().item()
        max_val = flat.max().item()

        # Handle high quantize_bits properly: skip quantization / identity mapping
        if quantize_bits is None or quantize_bits >= 16:
            if tensor.dtype == torch.bfloat16:
                raw_bytes = tensor.detach().cpu().view(torch.int16).numpy().tobytes()
            else:
                raw_bytes = tensor.detach().cpu().numpy().tobytes()
            metadata = {
                "scale": 1.0,
                "min_val": 0.0,
                "orig_shape": list(tensor.shape),
                "orig_dtype": str(tensor.dtype),
                "total_elements": numel,
                "entropy": 0.0,
                "pad_len": 0,
                "compressed_bytes": len(raw_bytes),
                "is_raw": True,
            }
            return raw_bytes, {}, metadata

        # Quantization
        num_levels = 2 ** int(quantize_bits)
        scale = (max_val - min_val) / max(num_levels - 1, 1)
        scale = max(scale, 1e-12)
        quantized = torch.clamp(torch.round((flat - min_val) / scale), 0, num_levels - 1).long()

        q_np = quantized.cpu().numpy()
        symbols, counts = np.unique(q_np, return_counts=True)
        frequencies = {int(s): int(c) for s, c in zip(symbols, counts)}

        root = HuffmanCoder.build_tree(frequencies)
        code_table = HuffmanCoder.generate_code_table(root)

        probs = counts / numel
        entropy = -float(np.sum(probs * np.log2(probs + 1e-12)))

        # Vectorized bitstream packing using LUT mapping and numpy.packbits
        max_sym = int(symbols.max())
        lut = [code_table.get(i, "0") for i in range(max_sym + 1)]
        bit_chunks = [lut[v] for v in q_np]
        bitstream_str = "".join(bit_chunks)
        pad_len = (8 - len(bitstream_str) % 8) % 8
        padded = bitstream_str + "0" * pad_len

        # Fast bit packing using np.packbits
        bit_arr = np.frombuffer(padded.encode("ascii"), dtype=np.uint8) - ord("0")
        packed_np = np.packbits(bit_arr)
        byte_array = packed_np.tobytes()

        metadata = {
            "scale": scale,
            "min_val": min_val,
            "orig_shape": list(tensor.shape),
            "orig_dtype": str(tensor.dtype),
            "total_elements": numel,
            "entropy": round(entropy, 3),
            "pad_len": pad_len,
            "compressed_bytes": len(byte_array),
            "is_raw": False,
        }
        return byte_array, code_table, metadata

    @staticmethod
    def decode(
        packed_bytes: bytes,
        code_table: dict[int, str],
        metadata: dict[str, Any],
        device: torch.device | None = None,
    ) -> torch.Tensor:
        """Vectorized decode bitstream back into reconstructed tensor."""
        total_elements = metadata.get("total_elements", 0)
        if total_elements == 0:
            return torch.empty(metadata.get("orig_shape", [0]), device=device)

        if metadata.get("is_raw", False):
            orig_dtype_str = metadata.get("orig_dtype", "torch.float32")
            target_dtype = getattr(torch, orig_dtype_str.replace("torch.", ""), torch.float32)
            if target_dtype == torch.bfloat16:
                arr = np.frombuffer(packed_bytes, dtype=np.int16)
                tensor_out = (
                    torch.from_numpy(arr.copy())
                    .view(torch.bfloat16)
                    .reshape(metadata.get("orig_shape"))
                )
            else:
                dtype_str = orig_dtype_str.replace("torch.", "")
                np_dtype = getattr(np, dtype_str if hasattr(np, dtype_str) else "float32")
                arr = np.frombuffer(packed_bytes, dtype=np_dtype)
                tensor_out = (
                    torch.from_numpy(arr.copy())
                    .reshape(metadata.get("orig_shape"))
                    .to(dtype=target_dtype)
                )
            return tensor_out.to(device=device)

        scale = metadata.get("scale", 1.0)
        min_val = metadata.get("min_val", 0.0)
        shape = metadata.get("orig_shape", [total_elements])

        # Vectorized unpacking of bits
        pad_len = metadata.get("pad_len", 0)
        raw_arr = np.frombuffer(packed_bytes, dtype=np.uint8)
        bits = np.unpackbits(raw_arr)
        if pad_len > 0:
            bits = bits[:-pad_len]

        # Fast flat-array binary trie traversal
        node_left = [-1]
        node_right = [-1]
        node_sym = [-1]

        for sym, code in code_table.items():
            curr = 0
            for b in code:
                if b == "0":
                    if node_left[curr] == -1:
                        node_left[curr] = len(node_left)
                        node_left.append(-1)
                        node_right.append(-1)
                        node_sym.append(-1)
                    curr = node_left[curr]
                else:
                    if node_right[curr] == -1:
                        node_right[curr] = len(node_right)
                        node_left.append(-1)
                        node_right.append(-1)
                        node_sym.append(-1)
                    curr = node_right[curr]
            node_sym[curr] = sym

        nl = np.array(node_left, dtype=np.int32)
        nr = np.array(node_right, dtype=np.int32)
        ns = np.array(node_sym, dtype=np.int32)

        decoded = np.empty(total_elements, dtype=np.float32)
        curr = 0
        idx = 0

        # Fast traversal loop
        for b in bits:
            curr = nr[curr] if b else nl[curr]
            sym = ns[curr]
            if sym != -1:
                decoded[idx] = sym
                idx += 1
                curr = 0
                if idx == total_elements:
                    break

        decoded = decoded * scale + min_val
        return torch.from_numpy(decoded).reshape(shape).to(device=device)

Vectorized lossless entropy coder for quantized weight and KV cache tensors. Supports continuous tensor quantization (e.g. 4-bit / 8-bit), symbol frequency modeling, binary prefix tree construction, numpy bitstream packing, and vectorized decoding.

Static methods

def build_tree(frequencies: dict[int, int]) ‑> HuffmanNode | None

Expand source code

@staticmethod
def build_tree(frequencies: dict[int, int]) -> HuffmanNode | None:
    """Build Huffman binary prefix tree from symbol frequencies."""
    if not frequencies:
        return None

    heap: list[HuffmanNode] = [HuffmanNode(sym, freq) for sym, freq in frequencies.items()]
    heapq.heapify(heap)

    if len(heap) == 1:
        root = HuffmanNode(freq=heap[0].freq, left=heap[0])
        return root

    while len(heap) > 1:
        node1 = heapq.heappop(heap)
        node2 = heapq.heappop(heap)
        merged = HuffmanNode(freq=node1.freq + node2.freq, left=node1, right=node2)
        heapq.heappush(heap, merged)

    return heap[0]

Build Huffman binary prefix tree from symbol frequencies.

def decode(packed_bytes: bytes, code_table: dict[int, str], metadata: dict[str, Any], device: torch.device | None = None) ‑> torch.Tensor

Expand source code

@staticmethod
def decode(
    packed_bytes: bytes,
    code_table: dict[int, str],
    metadata: dict[str, Any],
    device: torch.device | None = None,
) -> torch.Tensor:
    """Vectorized decode bitstream back into reconstructed tensor."""
    total_elements = metadata.get("total_elements", 0)
    if total_elements == 0:
        return torch.empty(metadata.get("orig_shape", [0]), device=device)

    if metadata.get("is_raw", False):
        orig_dtype_str = metadata.get("orig_dtype", "torch.float32")
        target_dtype = getattr(torch, orig_dtype_str.replace("torch.", ""), torch.float32)
        if target_dtype == torch.bfloat16:
            arr = np.frombuffer(packed_bytes, dtype=np.int16)
            tensor_out = (
                torch.from_numpy(arr.copy())
                .view(torch.bfloat16)
                .reshape(metadata.get("orig_shape"))
            )
        else:
            dtype_str = orig_dtype_str.replace("torch.", "")
            np_dtype = getattr(np, dtype_str if hasattr(np, dtype_str) else "float32")
            arr = np.frombuffer(packed_bytes, dtype=np_dtype)
            tensor_out = (
                torch.from_numpy(arr.copy())
                .reshape(metadata.get("orig_shape"))
                .to(dtype=target_dtype)
            )
        return tensor_out.to(device=device)

    scale = metadata.get("scale", 1.0)
    min_val = metadata.get("min_val", 0.0)
    shape = metadata.get("orig_shape", [total_elements])

    # Vectorized unpacking of bits
    pad_len = metadata.get("pad_len", 0)
    raw_arr = np.frombuffer(packed_bytes, dtype=np.uint8)
    bits = np.unpackbits(raw_arr)
    if pad_len > 0:
        bits = bits[:-pad_len]

    # Fast flat-array binary trie traversal
    node_left = [-1]
    node_right = [-1]
    node_sym = [-1]

    for sym, code in code_table.items():
        curr = 0
        for b in code:
            if b == "0":
                if node_left[curr] == -1:
                    node_left[curr] = len(node_left)
                    node_left.append(-1)
                    node_right.append(-1)
                    node_sym.append(-1)
                curr = node_left[curr]
            else:
                if node_right[curr] == -1:
                    node_right[curr] = len(node_right)
                    node_left.append(-1)
                    node_right.append(-1)
                    node_sym.append(-1)
                curr = node_right[curr]
        node_sym[curr] = sym

    nl = np.array(node_left, dtype=np.int32)
    nr = np.array(node_right, dtype=np.int32)
    ns = np.array(node_sym, dtype=np.int32)

    decoded = np.empty(total_elements, dtype=np.float32)
    curr = 0
    idx = 0

    # Fast traversal loop
    for b in bits:
        curr = nr[curr] if b else nl[curr]
        sym = ns[curr]
        if sym != -1:
            decoded[idx] = sym
            idx += 1
            curr = 0
            if idx == total_elements:
                break

    decoded = decoded * scale + min_val
    return torch.from_numpy(decoded).reshape(shape).to(device=device)

Vectorized decode bitstream back into reconstructed tensor.

def encode(tensor: torch.Tensor, quantize_bits: int | None = 8) ‑> tuple[bytes, dict[int, str], dict[str, typing.Any]]

Expand source code

@staticmethod
def encode(
    tensor: torch.Tensor,
    quantize_bits: int | None = 8,
) -> tuple[bytes, dict[int, str], dict[str, Any]]:
    """
    Vectorized encode input tensor into packed bitstream bytes and Huffman codebook.
    """
    flat = tensor.detach().float().view(-1)
    numel = flat.numel()
    if numel == 0:
        return b"", {}, {"total_elements": 0, "entropy": 0.0, "compressed_bytes": 0}

    min_val = flat.min().item()
    max_val = flat.max().item()

    # Handle high quantize_bits properly: skip quantization / identity mapping
    if quantize_bits is None or quantize_bits >= 16:
        if tensor.dtype == torch.bfloat16:
            raw_bytes = tensor.detach().cpu().view(torch.int16).numpy().tobytes()
        else:
            raw_bytes = tensor.detach().cpu().numpy().tobytes()
        metadata = {
            "scale": 1.0,
            "min_val": 0.0,
            "orig_shape": list(tensor.shape),
            "orig_dtype": str(tensor.dtype),
            "total_elements": numel,
            "entropy": 0.0,
            "pad_len": 0,
            "compressed_bytes": len(raw_bytes),
            "is_raw": True,
        }
        return raw_bytes, {}, metadata

    # Quantization
    num_levels = 2 ** int(quantize_bits)
    scale = (max_val - min_val) / max(num_levels - 1, 1)
    scale = max(scale, 1e-12)
    quantized = torch.clamp(torch.round((flat - min_val) / scale), 0, num_levels - 1).long()

    q_np = quantized.cpu().numpy()
    symbols, counts = np.unique(q_np, return_counts=True)
    frequencies = {int(s): int(c) for s, c in zip(symbols, counts)}

    root = HuffmanCoder.build_tree(frequencies)
    code_table = HuffmanCoder.generate_code_table(root)

    probs = counts / numel
    entropy = -float(np.sum(probs * np.log2(probs + 1e-12)))

    # Vectorized bitstream packing using LUT mapping and numpy.packbits
    max_sym = int(symbols.max())
    lut = [code_table.get(i, "0") for i in range(max_sym + 1)]
    bit_chunks = [lut[v] for v in q_np]
    bitstream_str = "".join(bit_chunks)
    pad_len = (8 - len(bitstream_str) % 8) % 8
    padded = bitstream_str + "0" * pad_len

    # Fast bit packing using np.packbits
    bit_arr = np.frombuffer(padded.encode("ascii"), dtype=np.uint8) - ord("0")
    packed_np = np.packbits(bit_arr)
    byte_array = packed_np.tobytes()

    metadata = {
        "scale": scale,
        "min_val": min_val,
        "orig_shape": list(tensor.shape),
        "orig_dtype": str(tensor.dtype),
        "total_elements": numel,
        "entropy": round(entropy, 3),
        "pad_len": pad_len,
        "compressed_bytes": len(byte_array),
        "is_raw": False,
    }
    return byte_array, code_table, metadata

Vectorized encode input tensor into packed bitstream bytes and Huffman codebook.

def generate_code_table(root: HuffmanNode | None, current_code: str = '', code_table: dict[int, str] | None = None) ‑> dict[int, str]

Expand source code

@staticmethod
def generate_code_table(
    root: HuffmanNode | None,
    current_code: str = "",
    code_table: dict[int, str] | None = None,
) -> dict[int, str]:
    """Traverse Huffman tree to generate prefix bitstring per symbol."""
    if code_table is None:
        code_table = {}
    if root is None:
        return code_table

    if root.symbol is not None:
        code_table[root.symbol] = current_code if current_code else "0"
        return code_table

    if root.left is not None:
        HuffmanCoder.generate_code_table(root.left, current_code + "0", code_table)
    if root.right is not None:
        HuffmanCoder.generate_code_table(root.right, current_code + "1", code_table)

    return code_table

Traverse Huffman tree to generate prefix bitstring per symbol.

class HuffmanCompressor (model: nn.Module, quantize_bits: int = 8, max_error: float = 0.2, target_attention_only: bool = False, fallback_to_original: bool = True, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("huffman")
@COMPRESSOR_REGISTRY.register("huffman_coder")
class HuffmanCompressor(BaseCompressor):
    """
    HuffmanCompressor: Model-wide entropy compression profiler and simulator.
    """

    def __init__(
        self,
        model: nn.Module,
        quantize_bits: int = 8,
        max_error: float = 0.20,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="huffman", **kwargs)
        cfg = cfg or kwargs.get("comp_config")
        self.cfg = cfg

        self.quantize_bits = int(
            _cfg_get(
                cfg,
                "huffman_quantize_bits",
                _cfg_get(cfg, "quantize_bits", _cfg_get(cfg, "bits", quantize_bits)),
            )
        )
        self.max_error = float(
            _cfg_get(cfg, "huffman_max_error", _cfg_get(cfg, "max_error", max_error))
        )
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "huffman_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "huffman_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_lower = name.lower()
        parts = name.split(".")
        if any(part in parts for part in ["lora_A", "lora_B", "base_layer"]):
            return False
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            ]
        ):
            return False

        if self.target_attention_only:
            if not any(attn_kw in name_lower for attn_kw in _ATTN_NAMES):
                return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d):
            return False
        if not should_compress_linear(self.model, name):
            return False
        return True

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Profile and apply Huffman entropy coding across all linear layers."""
        first_p = next(self.model.parameters(), None)
        _model_device = first_p.device if first_p is not None else torch.device("cpu")

        total_orig_bytes = 0
        total_comp_bytes = 0
        total_entropy = 0.0

        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        seen_weights: set[int] = set()

        # Freeze skipped non-target layers if target_attention_only is enabled
        if self.target_attention_only:
            for name, module in self.model.named_modules():
                if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
                    continue
                target_mod = getattr(module, "base_layer", module)
                cls_name = target_mod.__class__.__name__.lower()
                is_linear = isinstance(target_mod, nn.Linear)
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                if (is_linear or is_conv1d) and not self._is_target_module(name, module):
                    _freeze_base_only(module)
                    self.skipped_layers.append(name)

        for name, module in self.model.named_modules():
            if not self._is_target_module(name, module):
                continue

            target_mod = getattr(module, "base_layer", module)
            if not hasattr(target_mod, "weight") or target_mod.weight is None:
                continue

            w_id = id(target_mod.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            raw_w = target_mod.weight.data
            orig_shape = raw_w.shape
            orig_device = raw_w.device
            orig_dtype = raw_w.dtype

            # Byte accounting based on actual element_size()
            elem_bytes = raw_w.element_size() if hasattr(raw_w, "element_size") else 2
            layer_orig_bytes = raw_w.numel() * elem_bytes

            W_orig = raw_w.t().float() if is_conv1d else raw_w.float()
            norm_w = torch.norm(W_orig, p="fro").item()

            try:
                packed_b, code_map, meta = HuffmanCoder.encode(
                    W_orig, quantize_bits=self.quantize_bits
                )
                W_rec = HuffmanCoder.decode(packed_b, code_map, meta, device=orig_device).to(
                    dtype=torch.float32
                )

                if norm_w < 1e-12:
                    err = 0.0
                else:
                    err = float(torch.norm(W_orig - W_rec, p="fro").item() / norm_w)

                self.layer_errors[name] = err

                d_out, d_in = W_orig.shape
                fam = getattr(getattr(self.model, "config", None), "model_type", None)
                eff_max_err = max(self.max_error, adaptive_max_error(None, d_in, d_out, family=fam))

                if err > eff_max_err:
                    logger.warning(
                        "Huffman error %.4f on %s exceeded max_error %.4f (base max_err=%.4f)",
                        err,
                        name,
                        eff_max_err,
                        self.max_error,
                    )
                    if self.fallback_to_original:
                        self.fallback_layers.append(name)
                        _freeze_base_only(module)
                        total_orig_bytes += layer_orig_bytes
                        total_comp_bytes += layer_orig_bytes
                        continue

                # Apply reconstructed weight
                if is_conv1d:
                    res_w = W_rec.t().to(dtype=orig_dtype, device=orig_device)
                else:
                    res_w = W_rec.to(dtype=orig_dtype, device=orig_device)

                assert res_w.shape == orig_shape, (
                    f"Shape mismatch on {name}: {res_w.shape} vs {orig_shape}"
                )
                target_mod.weight.data.copy_(res_w)

                # Preserve LoRA adapter trainability
                for lora_attr in ("lora_A", "lora_B", "lora_a", "lora_b"):
                    if hasattr(module, lora_attr):
                        l_mod = getattr(module, lora_attr)
                        if (
                            isinstance(l_mod, nn.Module)
                            and hasattr(l_mod, "weight")
                            and l_mod.weight is not None
                        ):
                            l_mod.weight.requires_grad = True
                        elif isinstance(l_mod, nn.Parameter):
                            l_mod.requires_grad = True

                self.compressed_layers.append(name)
                total_orig_bytes += layer_orig_bytes
                total_comp_bytes += meta["compressed_bytes"]
                total_entropy += meta["entropy"]

            except Exception as e:
                logger.debug("Huffman exception on %s: %s", name, e)
                self.fallback_layers.append(name)
                _freeze_base_only(module)
                total_orig_bytes += layer_orig_bytes
                total_comp_bytes += layer_orig_bytes

        # Ensure all parameters on target device
        self.model.to(device=_model_device)

        forward_check_success = None
        try:
            with torch.no_grad():
                dev = next(self.model.parameters(), torch.tensor(0.0)).device
                tokenizer = getattr(self.model, "tokenizer", None) or getattr(
                    getattr(self.model, "config", None), "tokenizer", None
                )
                if tokenizer is not None:
                    tok_out = tokenizer("Hello world", return_tensors="pt")
                    tok_out = {k: v.to(dev) for k, v in tok_out.items()}
                    self.model(**tok_out)
                    forward_check_success = True
                else:
                    # Look for token embeddings
                    has_embed = any(
                        any(kw in n.lower() for kw in ("embed_tokens", "wte", "word_embeddings"))
                        for n, _ in self.model.named_modules()
                    )
                    if has_embed:
                        dummy_ids = torch.randint(0, 100, (1, 8), device=dev)
                        dummy_mask = torch.ones((1, 8), device=dev)
                        try:
                            self.model(input_ids=dummy_ids, attention_mask=dummy_mask)
                        except TypeError:
                            self.model(dummy_ids)
                        forward_check_success = True
                    else:
                        first_linear = next(
                            (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                            None,
                        )
                        if first_linear is not None:
                            in_dim = getattr(first_linear, "in_features", 8)
                            self.model(torch.randn(1, 8, in_dim, device=dev))
                            forward_check_success = True
                        else:
                            forward_check_success = None
        except Exception as fwd_exc:
            logger.debug("Huffman forward check failed: %s", fwd_exc)
            forward_check_success = None

        stats = self.get_compression_stats()
        avg_err = (
            (sum(self.layer_errors.values()) / max(1, len(self.layer_errors)))
            if self.layer_errors
            else 0.0
        )
        stats["quantize_bits"] = self.quantize_bits
        stats["avg_entropy_bits"] = round(total_entropy / max(len(self.compressed_layers), 1), 2)
        stats["total_compressed_bytes"] = total_comp_bytes
        stats["huffman_compression_ratio"] = round(total_orig_bytes / max(total_comp_bytes, 1), 2)
        stats["theoretical_compression_ratio"] = stats["huffman_compression_ratio"]
        stats["num_compressed"] = len(self.compressed_layers)
        stats["num_fallback"] = len(self.fallback_layers)
        stats["num_skipped"] = len(self.skipped_layers)
        stats["compressed_layers"] = self.compressed_layers
        stats["fallback_layers"] = self.fallback_layers
        stats["skipped_layers"] = self.skipped_layers
        stats["avg_reconstruction_error"] = avg_err
        stats["forward_check_success"] = forward_check_success
        logger.info(
            "Compressor huffman finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

HuffmanCompressor: Model-wide entropy compression profiler and simulator.

Ancestors

BaseCompressor

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Profile and apply Huffman entropy coding across all linear layers."""
    first_p = next(self.model.parameters(), None)
    _model_device = first_p.device if first_p is not None else torch.device("cpu")

    total_orig_bytes = 0
    total_comp_bytes = 0
    total_entropy = 0.0

    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    self.layer_errors = {}
    seen_weights: set[int] = set()

    # Freeze skipped non-target layers if target_attention_only is enabled
    if self.target_attention_only:
        for name, module in self.model.named_modules():
            if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
                continue
            target_mod = getattr(module, "base_layer", module)
            cls_name = target_mod.__class__.__name__.lower()
            is_linear = isinstance(target_mod, nn.Linear)
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            if (is_linear or is_conv1d) and not self._is_target_module(name, module):
                _freeze_base_only(module)
                self.skipped_layers.append(name)

    for name, module in self.model.named_modules():
        if not self._is_target_module(name, module):
            continue

        target_mod = getattr(module, "base_layer", module)
        if not hasattr(target_mod, "weight") or target_mod.weight is None:
            continue

        w_id = id(target_mod.weight)
        if w_id in seen_weights:
            self.skipped_layers.append(name)
            continue
        seen_weights.add(w_id)

        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.data
        orig_shape = raw_w.shape
        orig_device = raw_w.device
        orig_dtype = raw_w.dtype

        # Byte accounting based on actual element_size()
        elem_bytes = raw_w.element_size() if hasattr(raw_w, "element_size") else 2
        layer_orig_bytes = raw_w.numel() * elem_bytes

        W_orig = raw_w.t().float() if is_conv1d else raw_w.float()
        norm_w = torch.norm(W_orig, p="fro").item()

        try:
            packed_b, code_map, meta = HuffmanCoder.encode(
                W_orig, quantize_bits=self.quantize_bits
            )
            W_rec = HuffmanCoder.decode(packed_b, code_map, meta, device=orig_device).to(
                dtype=torch.float32
            )

            if norm_w < 1e-12:
                err = 0.0
            else:
                err = float(torch.norm(W_orig - W_rec, p="fro").item() / norm_w)

            self.layer_errors[name] = err

            d_out, d_in = W_orig.shape
            fam = getattr(getattr(self.model, "config", None), "model_type", None)
            eff_max_err = max(self.max_error, adaptive_max_error(None, d_in, d_out, family=fam))

            if err > eff_max_err:
                logger.warning(
                    "Huffman error %.4f on %s exceeded max_error %.4f (base max_err=%.4f)",
                    err,
                    name,
                    eff_max_err,
                    self.max_error,
                )
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    _freeze_base_only(module)
                    total_orig_bytes += layer_orig_bytes
                    total_comp_bytes += layer_orig_bytes
                    continue

            # Apply reconstructed weight
            if is_conv1d:
                res_w = W_rec.t().to(dtype=orig_dtype, device=orig_device)
            else:
                res_w = W_rec.to(dtype=orig_dtype, device=orig_device)

            assert res_w.shape == orig_shape, (
                f"Shape mismatch on {name}: {res_w.shape} vs {orig_shape}"
            )
            target_mod.weight.data.copy_(res_w)

            # Preserve LoRA adapter trainability
            for lora_attr in ("lora_A", "lora_B", "lora_a", "lora_b"):
                if hasattr(module, lora_attr):
                    l_mod = getattr(module, lora_attr)
                    if (
                        isinstance(l_mod, nn.Module)
                        and hasattr(l_mod, "weight")
                        and l_mod.weight is not None
                    ):
                        l_mod.weight.requires_grad = True
                    elif isinstance(l_mod, nn.Parameter):
                        l_mod.requires_grad = True

            self.compressed_layers.append(name)
            total_orig_bytes += layer_orig_bytes
            total_comp_bytes += meta["compressed_bytes"]
            total_entropy += meta["entropy"]

        except Exception as e:
            logger.debug("Huffman exception on %s: %s", name, e)
            self.fallback_layers.append(name)
            _freeze_base_only(module)
            total_orig_bytes += layer_orig_bytes
            total_comp_bytes += layer_orig_bytes

    # Ensure all parameters on target device
    self.model.to(device=_model_device)

    forward_check_success = None
    try:
        with torch.no_grad():
            dev = next(self.model.parameters(), torch.tensor(0.0)).device
            tokenizer = getattr(self.model, "tokenizer", None) or getattr(
                getattr(self.model, "config", None), "tokenizer", None
            )
            if tokenizer is not None:
                tok_out = tokenizer("Hello world", return_tensors="pt")
                tok_out = {k: v.to(dev) for k, v in tok_out.items()}
                self.model(**tok_out)
                forward_check_success = True
            else:
                # Look for token embeddings
                has_embed = any(
                    any(kw in n.lower() for kw in ("embed_tokens", "wte", "word_embeddings"))
                    for n, _ in self.model.named_modules()
                )
                if has_embed:
                    dummy_ids = torch.randint(0, 100, (1, 8), device=dev)
                    dummy_mask = torch.ones((1, 8), device=dev)
                    try:
                        self.model(input_ids=dummy_ids, attention_mask=dummy_mask)
                    except TypeError:
                        self.model(dummy_ids)
                    forward_check_success = True
                else:
                    first_linear = next(
                        (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                        None,
                    )
                    if first_linear is not None:
                        in_dim = getattr(first_linear, "in_features", 8)
                        self.model(torch.randn(1, 8, in_dim, device=dev))
                        forward_check_success = True
                    else:
                        forward_check_success = None
    except Exception as fwd_exc:
        logger.debug("Huffman forward check failed: %s", fwd_exc)
        forward_check_success = None

    stats = self.get_compression_stats()
    avg_err = (
        (sum(self.layer_errors.values()) / max(1, len(self.layer_errors)))
        if self.layer_errors
        else 0.0
    )
    stats["quantize_bits"] = self.quantize_bits
    stats["avg_entropy_bits"] = round(total_entropy / max(len(self.compressed_layers), 1), 2)
    stats["total_compressed_bytes"] = total_comp_bytes
    stats["huffman_compression_ratio"] = round(total_orig_bytes / max(total_comp_bytes, 1), 2)
    stats["theoretical_compression_ratio"] = stats["huffman_compression_ratio"]
    stats["num_compressed"] = len(self.compressed_layers)
    stats["num_fallback"] = len(self.fallback_layers)
    stats["num_skipped"] = len(self.skipped_layers)
    stats["compressed_layers"] = self.compressed_layers
    stats["fallback_layers"] = self.fallback_layers
    stats["skipped_layers"] = self.skipped_layers
    stats["avg_reconstruction_error"] = avg_err
    stats["forward_check_success"] = forward_check_success
    logger.info(
        "Compressor huffman finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        self.max_error,
    )
    return self.model, stats

Profile and apply Huffman entropy coding across all linear layers.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class HuffmanNode (symbol: int | None = None, freq: int = 0, left: HuffmanNode | None = None, right: HuffmanNode | None = None)

Expand source code

class HuffmanNode:
    """Node in the Huffman prefix tree."""

    def __init__(
        self,
        symbol: int | None = None,
        freq: int = 0,
        left: HuffmanNode | None = None,
        right: HuffmanNode | None = None,
    ) -> None:
        self.symbol = symbol
        self.freq = freq
        self.left = left
        self.right = right

    def __lt__(self, other: HuffmanNode) -> bool:
        return self.freq < other.freq

Node in the Huffman prefix tree.

class HybridKVCache (max_cache_len: int = 128, num_sinks: int = 4, recent_window: int = 32, budget: float = 0.3, static_ratio: float = 0.7, static_heads_mask: torch.Tensor | None = None)

Expand source code

class HybridKVCache:
    """
    Dynamic KV Cache Manager for HybridKV.
    Maintains separate retention policies for static vs dynamic attention heads:
    - Static heads: pinned attention sinks (first S tokens) + global heavy hitters
    - Dynamic heads: recent local sliding window (last W tokens) + chunked key retention
    Preserves all model stationary weights and strictly compresses dynamic KV cache tensors.
    Supports variable sequence lengths across prefill and token-by-token decoding.
    """

    def __init__(
        self,
        max_cache_len: int = 128,
        num_sinks: int = 4,
        recent_window: int = 32,
        budget: float = 0.3,
        static_ratio: float = 0.7,
        static_heads_mask: torch.Tensor | None = None,
    ) -> None:
        self.budget = float(budget)
        self.max_cache_len = (
            max(8, int(max_cache_len * self.budget)) if self.budget > 0 else int(max_cache_len)
        )
        self.num_sinks = int(num_sinks)
        self.recent_window = int(recent_window)
        self.static_ratio = float(static_ratio)
        self.static_heads_mask = static_heads_mask
        self.k_cache: torch.Tensor | None = None
        self.v_cache: torch.Tensor | None = None

    def reset(self) -> None:
        """Clear cache state to prevent memory leaks."""
        self.k_cache = None
        self.v_cache = None

    @torch.no_grad()
    def update_and_compress(
        self,
        new_k: torch.Tensor,
        new_v: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Append new KV representations and apply hybrid eviction if cache exceeds capacity.
        Shapes: (B, H, N, D) or (B, N, D) or (N, D).
        """
        orig_ndim = new_k.ndim
        k_4d = new_k
        v_4d = new_v
        if orig_ndim == 2:
            k_4d = new_k.unsqueeze(0).unsqueeze(0)
            v_4d = new_v.unsqueeze(0).unsqueeze(0)
        elif orig_ndim == 3:
            k_4d = new_k.unsqueeze(1)
            v_4d = new_v.unsqueeze(1)

        B, H, N, D = k_4d.shape

        if self.k_cache is None:
            self.k_cache = k_4d
            self.v_cache = v_4d
        else:
            if self.k_cache.device != k_4d.device or self.k_cache.dtype != k_4d.dtype:
                self.k_cache = self.k_cache.to(device=k_4d.device, dtype=k_4d.dtype)
                self.v_cache = self.v_cache.to(device=v_4d.device, dtype=v_4d.dtype)
            if self.k_cache.shape[0] != B:
                self.k_cache = self.k_cache[:1].expand(B, -1, -1, -1).clone()
                self.v_cache = self.v_cache[:1].expand(B, -1, -1, -1).clone()
            self.k_cache = torch.cat([self.k_cache, k_4d], dim=2)
            self.v_cache = torch.cat([self.v_cache, v_4d], dim=2)

        cur_len = self.k_cache.size(2)
        if cur_len <= self.max_cache_len:
            out_k, out_v = self.k_cache, self.v_cache
            if orig_ndim == 2:
                return out_k.squeeze(0).squeeze(0), out_v.squeeze(0).squeeze(0)
            elif orig_ndim == 3:
                return out_k.squeeze(1), out_v.squeeze(1)
            return out_k, out_v

        # Eviction per head according to static vs dynamic classification
        target_len = self.max_cache_len
        head_k_list: list[torch.Tensor] = []
        head_v_list: list[torch.Tensor] = []

        is_static = (
            self.static_heads_mask
            if self.static_heads_mask is not None
            else torch.ones(H, dtype=torch.bool, device=self.k_cache.device)
        )

        for h_idx in range(H):
            k_h = self.k_cache[:, h_idx, :, :]  # (B, cur_len, D)
            v_h = self.v_cache[:, h_idx, :, :]

            if is_static[h_idx]:
                # Static head policy: Sink tokens + Heavy Hitters via L2 Key norm
                sinks_k = k_h[:, : self.num_sinks, :]
                sinks_v = v_h[:, : self.num_sinks, :]
                remaining_k = k_h[:, self.num_sinks :, :]
                remaining_v = v_h[:, self.num_sinks :, :]

                rem_budget = max(0, target_len - self.num_sinks)
                if remaining_k.size(1) > rem_budget and rem_budget > 0:
                    k_norms = torch.norm(remaining_k, p=2, dim=-1).mean(dim=0)
                    _, topk_indices = torch.topk(k_norms, k=rem_budget, largest=True)
                    topk_indices, _ = torch.sort(topk_indices)
                    kept_k = remaining_k[:, topk_indices, :]
                    kept_v = remaining_v[:, topk_indices, :]
                    evicted_k = torch.cat([sinks_k, kept_k], dim=1)
                    evicted_v = torch.cat([sinks_v, kept_v], dim=1)
                else:
                    evicted_k = sinks_k
                    evicted_v = sinks_v
            else:
                # Dynamic head policy: Sliding local window strictly
                window_budget = min(self.recent_window, target_len)
                evicted_k = k_h[:, -window_budget:, :]
                evicted_v = v_h[:, -window_budget:, :]

            head_k_list.append(evicted_k.unsqueeze(1))
            head_v_list.append(evicted_v.unsqueeze(1))

        self.k_cache = torch.cat(head_k_list, dim=1)
        self.v_cache = torch.cat(head_v_list, dim=1)

        out_k, out_v = self.k_cache, self.v_cache
        if orig_ndim == 2:
            return out_k.squeeze(0).squeeze(0), out_v.squeeze(0).squeeze(0)
        elif orig_ndim == 3:
            return out_k.squeeze(1), out_v.squeeze(1)
        return out_k, out_v

Dynamic KV Cache Manager for HybridKV. Maintains separate retention policies for static vs dynamic attention heads: - Static heads: pinned attention sinks (first S tokens) + global heavy hitters - Dynamic heads: recent local sliding window (last W tokens) + chunked key retention Preserves all model stationary weights and strictly compresses dynamic KV cache tensors. Supports variable sequence lengths across prefill and token-by-token decoding.

Methods

def reset(self) ‑> None

Expand source code

def reset(self) -> None:
    """Clear cache state to prevent memory leaks."""
    self.k_cache = None
    self.v_cache = None

Clear cache state to prevent memory leaks.

def update_and_compress(self, new_k: torch.Tensor, new_v: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor]

Expand source code

@torch.no_grad()
def update_and_compress(
    self,
    new_k: torch.Tensor,
    new_v: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Append new KV representations and apply hybrid eviction if cache exceeds capacity.
    Shapes: (B, H, N, D) or (B, N, D) or (N, D).
    """
    orig_ndim = new_k.ndim
    k_4d = new_k
    v_4d = new_v
    if orig_ndim == 2:
        k_4d = new_k.unsqueeze(0).unsqueeze(0)
        v_4d = new_v.unsqueeze(0).unsqueeze(0)
    elif orig_ndim == 3:
        k_4d = new_k.unsqueeze(1)
        v_4d = new_v.unsqueeze(1)

    B, H, N, D = k_4d.shape

    if self.k_cache is None:
        self.k_cache = k_4d
        self.v_cache = v_4d
    else:
        if self.k_cache.device != k_4d.device or self.k_cache.dtype != k_4d.dtype:
            self.k_cache = self.k_cache.to(device=k_4d.device, dtype=k_4d.dtype)
            self.v_cache = self.v_cache.to(device=v_4d.device, dtype=v_4d.dtype)
        if self.k_cache.shape[0] != B:
            self.k_cache = self.k_cache[:1].expand(B, -1, -1, -1).clone()
            self.v_cache = self.v_cache[:1].expand(B, -1, -1, -1).clone()
        self.k_cache = torch.cat([self.k_cache, k_4d], dim=2)
        self.v_cache = torch.cat([self.v_cache, v_4d], dim=2)

    cur_len = self.k_cache.size(2)
    if cur_len <= self.max_cache_len:
        out_k, out_v = self.k_cache, self.v_cache
        if orig_ndim == 2:
            return out_k.squeeze(0).squeeze(0), out_v.squeeze(0).squeeze(0)
        elif orig_ndim == 3:
            return out_k.squeeze(1), out_v.squeeze(1)
        return out_k, out_v

    # Eviction per head according to static vs dynamic classification
    target_len = self.max_cache_len
    head_k_list: list[torch.Tensor] = []
    head_v_list: list[torch.Tensor] = []

    is_static = (
        self.static_heads_mask
        if self.static_heads_mask is not None
        else torch.ones(H, dtype=torch.bool, device=self.k_cache.device)
    )

    for h_idx in range(H):
        k_h = self.k_cache[:, h_idx, :, :]  # (B, cur_len, D)
        v_h = self.v_cache[:, h_idx, :, :]

        if is_static[h_idx]:
            # Static head policy: Sink tokens + Heavy Hitters via L2 Key norm
            sinks_k = k_h[:, : self.num_sinks, :]
            sinks_v = v_h[:, : self.num_sinks, :]
            remaining_k = k_h[:, self.num_sinks :, :]
            remaining_v = v_h[:, self.num_sinks :, :]

            rem_budget = max(0, target_len - self.num_sinks)
            if remaining_k.size(1) > rem_budget and rem_budget > 0:
                k_norms = torch.norm(remaining_k, p=2, dim=-1).mean(dim=0)
                _, topk_indices = torch.topk(k_norms, k=rem_budget, largest=True)
                topk_indices, _ = torch.sort(topk_indices)
                kept_k = remaining_k[:, topk_indices, :]
                kept_v = remaining_v[:, topk_indices, :]
                evicted_k = torch.cat([sinks_k, kept_k], dim=1)
                evicted_v = torch.cat([sinks_v, kept_v], dim=1)
            else:
                evicted_k = sinks_k
                evicted_v = sinks_v
        else:
            # Dynamic head policy: Sliding local window strictly
            window_budget = min(self.recent_window, target_len)
            evicted_k = k_h[:, -window_budget:, :]
            evicted_v = v_h[:, -window_budget:, :]

        head_k_list.append(evicted_k.unsqueeze(1))
        head_v_list.append(evicted_v.unsqueeze(1))

    self.k_cache = torch.cat(head_k_list, dim=1)
    self.v_cache = torch.cat(head_v_list, dim=1)

    out_k, out_v = self.k_cache, self.v_cache
    if orig_ndim == 2:
        return out_k.squeeze(0).squeeze(0), out_v.squeeze(0).squeeze(0)
    elif orig_ndim == 3:
        return out_k.squeeze(1), out_v.squeeze(1)
    return out_k, out_v

Append new KV representations and apply hybrid eviction if cache exceeds capacity. Shapes: (B, H, N, D) or (B, N, D) or (N, D).

class HybridKVCompressor (model: nn.Module, budget: float = 0.3, static_ratio: float = 0.7, max_cache_len: int = 128, num_sinks: int = 4, recent_window: int = 32, target_attention_only: bool = True, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("hybridkv_class")
@COMPRESSOR_REGISTRY.register("hybridkv_compressor")
class HybridKVCompressor(BaseCompressor):
    """
    HybridKV Compressor: Classifies projection heads into static/dynamic types
    and attaches HybridKVCache instances to attention layers without modifying model weights.
    """

    def __init__(
        self,
        model: nn.Module,
        budget: float = 0.3,
        static_ratio: float = 0.7,
        max_cache_len: int = 128,
        num_sinks: int = 4,
        recent_window: int = 32,
        target_attention_only: bool = True,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="hybridkv", **kwargs)
        cfg = cfg or kwargs.get("comp_config")
        self.cfg = cfg

        self.budget = float(_cfg_get(cfg, "hybridkv_budget", _cfg_get(cfg, "budget", budget)))
        self.static_ratio = float(
            _cfg_get(
                cfg,
                "hybridkv_static_ratio",
                _cfg_get(cfg, "static_ratio", static_ratio),
            )
        )
        self.max_cache_len = int(
            _cfg_get(
                cfg,
                "hybridkv_max_cache_len",
                _cfg_get(cfg, "max_cache_len", max_cache_len),
            )
        )
        self.num_sinks = int(
            _cfg_get(cfg, "hybridkv_num_sinks", _cfg_get(cfg, "num_sinks", num_sinks))
        )
        self.recent_window = int(
            _cfg_get(
                cfg,
                "hybridkv_recent_window",
                _cfg_get(cfg, "recent_window", recent_window),
            )
        )
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "hybridkv_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )

    def _classify_head(self, weight: torch.Tensor) -> str:
        """Classify attention projection weight into static vs dynamic based on variance."""
        var = torch.var(weight.float(), dim=-1).mean().item()
        return "static" if var < 0.1 else "dynamic"

    def _is_attention_proj(self, name: str, module: nn.Module) -> bool:
        """Check if module corresponds to attention projection."""
        name_lower = name.lower()
        if any(part in name.split(".") for part in ["lora_A", "lora_B"]):
            return False
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            ]
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d):
            return False

        if self.target_attention_only:
            return any(kw in name_lower for kw in _ATTN_NAMES)
        return True

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Classify heads and configure dynamic KV caches without mutating model weights."""
        head_stats: dict[str, Any] = {}
        total_heads = 0
        static_heads_count = 0
        dynamic_heads_count = 0

        for name, module in self.model.named_modules():
            if self._is_attention_proj(name, module):
                target_mod = getattr(module, "base_layer", module)
                W = target_mod.weight.data
                cls_name = target_mod.__class__.__name__.lower()
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                out_f = W.shape[1] if is_conv1d else W.shape[0]

                head_dim = 64 if out_f % 64 == 0 else (128 if out_f % 128 == 0 else 32)
                num_heads = max(1, out_f // head_dim)
                head_type = self._classify_head(W)

                num_static = max(1, int(num_heads * self.static_ratio))
                num_dynamic = max(0, num_heads - num_static)
                static_mask = torch.zeros(num_heads, dtype=torch.bool, device=W.device)
                static_mask[:num_static] = True

                kv_cache = HybridKVCache(
                    max_cache_len=self.max_cache_len,
                    num_sinks=self.num_sinks,
                    recent_window=self.recent_window,
                    budget=self.budget,
                    static_ratio=self.static_ratio,
                    static_heads_mask=static_mask,
                )
                # Attach KV cache manager without mutating stationary weights
                module.hybrid_kv_cache = kv_cache
                if hasattr(module, "base_layer"):
                    module.base_layer.hybrid_kv_cache = kv_cache

                head_stats[name] = {
                    "head_type": head_type,
                    "num_heads": num_heads,
                    "static_heads": num_static,
                    "dynamic_heads": num_dynamic,
                }
                total_heads += num_heads
                static_heads_count += num_static
                dynamic_heads_count += num_dynamic
            elif hasattr(module, "weight") or list(module.parameters(recurse=False)):
                # Freeze fallback/skipped layers to prevent unintended gradient flow
                for p in module.parameters(recurse=False):
                    p.requires_grad = False

        # Forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                dev = next(self.model.parameters(), torch.tensor(0.0)).device
                dummy = torch.randint(0, 100, (1, 8), device=dev)
                try:
                    self.model(dummy)
                except Exception:
                    first_linear = next(
                        (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                        None,
                    )
                    in_dim = getattr(first_linear, "in_features", 8)
                    self.model(torch.randn(1, 8, in_dim, device=dev))
        except Exception:
            forward_check_success = False

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        stats["budget"] = self.budget
        stats["static_ratio"] = self.static_ratio
        stats["max_cache_len"] = self.max_cache_len
        stats["num_sinks"] = self.num_sinks
        stats["recent_window"] = self.recent_window
        stats["total_attention_heads"] = total_heads
        stats["static_heads_count"] = static_heads_count
        stats["dynamic_heads_count"] = dynamic_heads_count
        stats["head_classifications"] = head_stats
        stats["attached_caches"] = len(head_stats)
        stats["kv_cache_budget_saving"] = f"{(1.0 - self.budget) * 100:.1f}% Footprint Reduction"
        stats["forward_check_success"] = forward_check_success
        return self.model, stats

HybridKV Compressor: Classifies projection heads into static/dynamic types and attaches HybridKVCache instances to attention layers without modifying model weights.

Ancestors

BaseCompressor

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Classify heads and configure dynamic KV caches without mutating model weights."""
    head_stats: dict[str, Any] = {}
    total_heads = 0
    static_heads_count = 0
    dynamic_heads_count = 0

    for name, module in self.model.named_modules():
        if self._is_attention_proj(name, module):
            target_mod = getattr(module, "base_layer", module)
            W = target_mod.weight.data
            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            out_f = W.shape[1] if is_conv1d else W.shape[0]

            head_dim = 64 if out_f % 64 == 0 else (128 if out_f % 128 == 0 else 32)
            num_heads = max(1, out_f // head_dim)
            head_type = self._classify_head(W)

            num_static = max(1, int(num_heads * self.static_ratio))
            num_dynamic = max(0, num_heads - num_static)
            static_mask = torch.zeros(num_heads, dtype=torch.bool, device=W.device)
            static_mask[:num_static] = True

            kv_cache = HybridKVCache(
                max_cache_len=self.max_cache_len,
                num_sinks=self.num_sinks,
                recent_window=self.recent_window,
                budget=self.budget,
                static_ratio=self.static_ratio,
                static_heads_mask=static_mask,
            )
            # Attach KV cache manager without mutating stationary weights
            module.hybrid_kv_cache = kv_cache
            if hasattr(module, "base_layer"):
                module.base_layer.hybrid_kv_cache = kv_cache

            head_stats[name] = {
                "head_type": head_type,
                "num_heads": num_heads,
                "static_heads": num_static,
                "dynamic_heads": num_dynamic,
            }
            total_heads += num_heads
            static_heads_count += num_static
            dynamic_heads_count += num_dynamic
        elif hasattr(module, "weight") or list(module.parameters(recurse=False)):
            # Freeze fallback/skipped layers to prevent unintended gradient flow
            for p in module.parameters(recurse=False):
                p.requires_grad = False

    # Forward sanity check
    forward_check_success = True
    try:
        with torch.no_grad():
            dev = next(self.model.parameters(), torch.tensor(0.0)).device
            dummy = torch.randint(0, 100, (1, 8), device=dev)
            try:
                self.model(dummy)
            except Exception:
                first_linear = next(
                    (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                    None,
                )
                in_dim = getattr(first_linear, "in_features", 8)
                self.model(torch.randn(1, 8, in_dim, device=dev))
    except Exception:
        forward_check_success = False

    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    stats["budget"] = self.budget
    stats["static_ratio"] = self.static_ratio
    stats["max_cache_len"] = self.max_cache_len
    stats["num_sinks"] = self.num_sinks
    stats["recent_window"] = self.recent_window
    stats["total_attention_heads"] = total_heads
    stats["static_heads_count"] = static_heads_count
    stats["dynamic_heads_count"] = dynamic_heads_count
    stats["head_classifications"] = head_stats
    stats["attached_caches"] = len(head_stats)
    stats["kv_cache_budget_saving"] = f"{(1.0 - self.budget) * 100:.1f}% Footprint Reduction"
    stats["forward_check_success"] = forward_check_success
    return self.model, stats

Classify heads and configure dynamic KV caches without mutating model weights.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class HybridSVDCompressor (model: nn.Module, svd_rank: int = 128, use_full_layer_replacement: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, max_error: float = 0.85, target_attention_only: bool = False, fallback_to_original: bool = True, cfg: Any | None = None, target_modules: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("hybrid_svd")
@COMPRESSOR_REGISTRY.register("hybrid_svd_layer")
@COMPRESSOR_REGISTRY.register("direct_svd_encoder")
class HybridSVDCompressor(BaseCompressor):
    """
    Weight-based SVD compressor for Linear / Conv1D layers.
    """

    def __init__(
        self,
        model: nn.Module,
        svd_rank: int = 128,
        use_full_layer_replacement: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        max_error: float = 0.85,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        cfg: Any | None = None,
        target_modules: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="hybrid_svd", **kwargs)
        cfg = cfg or kwargs.get("comp_config")
        self.cfg = cfg

        self.svd_rank = int(
            _cfg_get(
                cfg,
                "hybrid_svd_rank",
                _cfg_get(cfg, "svd_rank", _cfg_get(cfg, "rank", svd_rank)),
            )
        )
        self.use_full_layer_replacement = bool(
            _cfg_get(cfg, "use_full_layer_replacement", use_full_layer_replacement)
        )
        self.lora_r = int(_cfg_get(cfg, "lora_r", _cfg_get(cfg, "lora_rank", lora_r)))
        self.lora_alpha = float(_cfg_get(cfg, "lora_alpha", lora_alpha))
        self.max_error = float(
            _cfg_get(cfg, "hybrid_svd_max_error", _cfg_get(cfg, "max_error", max_error))
        )
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "hybrid_svd_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "hybrid_svd_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        self.target_modules = target_modules or _cfg_get(cfg, "target_modules", None)

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self._compressed_params: int = 0

    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
            return False
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            ]
        ):
            return False

        if self.target_attention_only:
            if not any(attn_kw in name_lower for attn_kw in _ATTN_NAMES):
                return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d) or isinstance(target_mod, SVDLinear):
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    # ------------------------------------------------------------------
    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Decompose linear layers via SVD and replace with SVDLinear."""
        replaced_count = 0
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        seen_weights: set = set()

        total_orig_params = 0
        total_comp_params = 0

        # Freeze non-target layers if attention-only.
        if self.target_attention_only:
            for name, module in self.model.named_modules():
                if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
                    continue
                target_mod = getattr(module, "base_layer", module)
                cls_name = target_mod.__class__.__name__.lower()
                is_linear = isinstance(target_mod, nn.Linear)
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                if (is_linear or is_conv1d) and not self._is_target_module(name, module):
                    for p in module.parameters():
                        p.requires_grad = False
                    self.skipped_layers.append(name)

        # Fallback: linear-layer SVD.
        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            target_mod = getattr(module, "base_layer", module)
            if not hasattr(target_mod, "weight") or target_mod.weight is None:
                continue

            w_id = id(target_mod.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            raw_w = target_mod.weight.data
            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            W = raw_w.t().contiguous().float() if is_conv1d else raw_w.float()
            out_f, in_f = W.shape
            orig_params = in_f * out_f
            total_orig_params += orig_params

            max_rank = min(in_f, out_f)
            bound_r = (in_f * out_f) // (in_f + out_f)
            if max_rank < 16 or bound_r < 4:
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_params += orig_params
                continue

            r = min(self.svd_rank, bound_r, max_rank)

            try:
                U, S, Vt = torch.linalg.svd(W, full_matrices=False)
                for escalation_step in range(4):
                    sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
                    U_r = (U[:, :r] * sqrt_S.unsqueeze(0)).to(
                        dtype=raw_w.dtype, device=raw_w.device
                    )
                    V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(
                        dtype=raw_w.dtype, device=raw_w.device
                    )

                    W_recon = torch.matmul(U[:, :r] * S[:r].unsqueeze(0), Vt[:r, :])
                    err = (
                        torch.norm(W - W_recon, p="fro") / torch.norm(W, p="fro").clamp_min(1e-8)
                    ).item()
                    self.layer_errors[name] = float(err)

                    if err > self.max_error and r < min(bound_r, max_rank):
                        r_new = min(int(r * 1.5), bound_r, max_rank)
                        if r_new > r:
                            logger.info(
                                "HybridSVD: %s err=%.4f at rank=%d, escalating to rank=%d",
                                name,
                                err,
                                r,
                                r_new,
                            )
                            r = r_new
                            continue

                    if err > self.max_error:
                        logger.warning(
                            "HybridSVD: layer %s reconstruction error %.4f > max_error %.4f (final rank=%d)",
                            name,
                            err,
                            self.max_error,
                            r,
                        )
                        if self.fallback_to_original:
                            self.fallback_layers.append(name)
                            for p in module.parameters():
                                p.requires_grad = False
                            total_comp_params += orig_params
                            break
                    break

                if name in self.fallback_layers:
                    continue

                # Detect existing LoRA rank.
                existing_lora_A = getattr(module, "lora_A", None)
                existing_lora_B = getattr(module, "lora_B", None)
                effective_lora_r = self.lora_r
                if (
                    existing_lora_A is not None
                    and hasattr(existing_lora_A, "weight")
                    and hasattr(existing_lora_A.weight, "shape")
                ):
                    try:
                        effective_lora_r = int(existing_lora_A.weight.shape[0])
                    except (AttributeError, IndexError, TypeError):
                        pass

                svd_layer = SVDLinear(
                    in_f,
                    out_f,
                    r,
                    bias=(getattr(target_mod, "bias", None) is not None),
                    dtype=raw_w.dtype,
                    device=raw_w.device,
                    lora_r=effective_lora_r,
                    lora_alpha=self.lora_alpha,
                )
                svd_layer.set_low_rank_factors(U_r, V_r)
                if getattr(target_mod, "bias", None) is not None and target_mod.bias is not None:
                    svd_layer.bias.data.copy_(target_mod.bias.data)

                if (
                    existing_lora_A is not None
                    and existing_lora_B is not None
                    and getattr(svd_layer, "lora_A", None) is not None
                ):
                    w_A = getattr(existing_lora_A, "weight", existing_lora_A)
                    w_B = getattr(existing_lora_B, "weight", existing_lora_B)
                    if hasattr(w_A, "data") and hasattr(w_B, "data"):
                        try:
                            svd_layer.lora_A.data.copy_(w_A.data)
                            svd_layer.lora_B.data.copy_(w_B.data)
                            svd_layer.lora_A.requires_grad = True
                            svd_layer.lora_B.requires_grad = True
                        except Exception as e:
                            logger.debug("LoRA copy failed for %s: %s", name, e)

                replace_module(self.model, name, svd_layer)
                replaced_count += 1
                self.compressed_layers.append(name)

                # FIX: use actual svd_layer.rank for parameter accounting.
                actual_rank = int(getattr(svd_layer, "rank", r))
                comp_params = actual_rank * (in_f + out_f)
                if getattr(svd_layer, "lora_r", 0) > 0:
                    comp_params += int(svd_layer.lora_r) * (in_f + out_f)
                total_comp_params += comp_params

            except Exception as e:
                logger.warning("HybridSVD compression exception on %s: %s", name, e)
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_params += orig_params

        # Post-compression forward sanity check (single attempt).
        forward_check_success = True
        try:
            self.model.eval()
            with torch.no_grad():
                first_p = next(self.model.parameters(), None)
                dev = first_p.device if first_p is not None else torch.device("cpu")
                cfg = getattr(self.model, "config", None)
                vocab_size = getattr(cfg, "vocab_size", 1000)
                max_tok = min(100, vocab_size - 1) if vocab_size > 10 else 2
                min_tok = min(2, max_tok)
                dummy = torch.randint(min_tok, max_tok + 1, (1, 8), device=dev, dtype=torch.long)
                attn_mask = torch.ones_like(dummy)
                try:
                    self.model(input_ids=dummy, attention_mask=attn_mask)
                except Exception:
                    # Try hidden-states form.
                    first_linear = next(
                        (m for m in self.model.modules() if isinstance(m, (nn.Linear, SVDLinear))),
                        None,
                    )
                    in_dim = getattr(first_linear, "in_features", getattr(cfg, "hidden_size", 8))
                    dummy_float = torch.randn(
                        1,
                        8,
                        in_dim,
                        device=dev,
                        dtype=first_p.dtype if first_p is not None else torch.float32,
                    )
                    self.model(dummy_float)
        except Exception as e:
            logger.warning("HybridSVD post-compression forward check warning: %s", e)
            forward_check_success = False

        effective_ratio = (total_comp_params / total_orig_params) if total_orig_params > 0 else 1.0

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        stats = self.get_compression_stats()
        avg_err = (
            (sum(self.layer_errors.values()) / max(1, len(self.layer_errors)))
            if self.layer_errors
            else 0.0
        )
        stats.update(
            {
                "replaced_layers_count": replaced_count,
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "svd_rank": self.svd_rank,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "max_error": self.max_error,
                "target_attention_only": self.target_attention_only,
                "fallback_to_original": self.fallback_to_original,
                "avg_reconstruction_error": avg_err,
                "effective_parameter_ratio": round(effective_ratio, 4),
                "forward_check_success": forward_check_success,
                "method": "hybrid_svd",
            }
        )
        return self.model, stats

Weight-based SVD compressor for Linear / Conv1D layers.

Ancestors

BaseCompressor

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Decompose linear layers via SVD and replace with SVDLinear."""
    replaced_count = 0
    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    self.layer_errors = {}
    seen_weights: set = set()

    total_orig_params = 0
    total_comp_params = 0

    # Freeze non-target layers if attention-only.
    if self.target_attention_only:
        for name, module in self.model.named_modules():
            if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
                continue
            target_mod = getattr(module, "base_layer", module)
            cls_name = target_mod.__class__.__name__.lower()
            is_linear = isinstance(target_mod, nn.Linear)
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            if (is_linear or is_conv1d) and not self._is_target_module(name, module):
                for p in module.parameters():
                    p.requires_grad = False
                self.skipped_layers.append(name)

    # Fallback: linear-layer SVD.
    for name, module in list(self.model.named_modules()):
        if not self._is_target_module(name, module):
            continue

        target_mod = getattr(module, "base_layer", module)
        if not hasattr(target_mod, "weight") or target_mod.weight is None:
            continue

        w_id = id(target_mod.weight)
        if w_id in seen_weights:
            self.skipped_layers.append(name)
            continue
        seen_weights.add(w_id)

        raw_w = target_mod.weight.data
        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        W = raw_w.t().contiguous().float() if is_conv1d else raw_w.float()
        out_f, in_f = W.shape
        orig_params = in_f * out_f
        total_orig_params += orig_params

        max_rank = min(in_f, out_f)
        bound_r = (in_f * out_f) // (in_f + out_f)
        if max_rank < 16 or bound_r < 4:
            self.fallback_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False
            total_comp_params += orig_params
            continue

        r = min(self.svd_rank, bound_r, max_rank)

        try:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            for escalation_step in range(4):
                sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
                U_r = (U[:, :r] * sqrt_S.unsqueeze(0)).to(
                    dtype=raw_w.dtype, device=raw_w.device
                )
                V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(
                    dtype=raw_w.dtype, device=raw_w.device
                )

                W_recon = torch.matmul(U[:, :r] * S[:r].unsqueeze(0), Vt[:r, :])
                err = (
                    torch.norm(W - W_recon, p="fro") / torch.norm(W, p="fro").clamp_min(1e-8)
                ).item()
                self.layer_errors[name] = float(err)

                if err > self.max_error and r < min(bound_r, max_rank):
                    r_new = min(int(r * 1.5), bound_r, max_rank)
                    if r_new > r:
                        logger.info(
                            "HybridSVD: %s err=%.4f at rank=%d, escalating to rank=%d",
                            name,
                            err,
                            r,
                            r_new,
                        )
                        r = r_new
                        continue

                if err > self.max_error:
                    logger.warning(
                        "HybridSVD: layer %s reconstruction error %.4f > max_error %.4f (final rank=%d)",
                        name,
                        err,
                        self.max_error,
                        r,
                    )
                    if self.fallback_to_original:
                        self.fallback_layers.append(name)
                        for p in module.parameters():
                            p.requires_grad = False
                        total_comp_params += orig_params
                        break
                break

            if name in self.fallback_layers:
                continue

            # Detect existing LoRA rank.
            existing_lora_A = getattr(module, "lora_A", None)
            existing_lora_B = getattr(module, "lora_B", None)
            effective_lora_r = self.lora_r
            if (
                existing_lora_A is not None
                and hasattr(existing_lora_A, "weight")
                and hasattr(existing_lora_A.weight, "shape")
            ):
                try:
                    effective_lora_r = int(existing_lora_A.weight.shape[0])
                except (AttributeError, IndexError, TypeError):
                    pass

            svd_layer = SVDLinear(
                in_f,
                out_f,
                r,
                bias=(getattr(target_mod, "bias", None) is not None),
                dtype=raw_w.dtype,
                device=raw_w.device,
                lora_r=effective_lora_r,
                lora_alpha=self.lora_alpha,
            )
            svd_layer.set_low_rank_factors(U_r, V_r)
            if getattr(target_mod, "bias", None) is not None and target_mod.bias is not None:
                svd_layer.bias.data.copy_(target_mod.bias.data)

            if (
                existing_lora_A is not None
                and existing_lora_B is not None
                and getattr(svd_layer, "lora_A", None) is not None
            ):
                w_A = getattr(existing_lora_A, "weight", existing_lora_A)
                w_B = getattr(existing_lora_B, "weight", existing_lora_B)
                if hasattr(w_A, "data") and hasattr(w_B, "data"):
                    try:
                        svd_layer.lora_A.data.copy_(w_A.data)
                        svd_layer.lora_B.data.copy_(w_B.data)
                        svd_layer.lora_A.requires_grad = True
                        svd_layer.lora_B.requires_grad = True
                    except Exception as e:
                        logger.debug("LoRA copy failed for %s: %s", name, e)

            replace_module(self.model, name, svd_layer)
            replaced_count += 1
            self.compressed_layers.append(name)

            # FIX: use actual svd_layer.rank for parameter accounting.
            actual_rank = int(getattr(svd_layer, "rank", r))
            comp_params = actual_rank * (in_f + out_f)
            if getattr(svd_layer, "lora_r", 0) > 0:
                comp_params += int(svd_layer.lora_r) * (in_f + out_f)
            total_comp_params += comp_params

        except Exception as e:
            logger.warning("HybridSVD compression exception on %s: %s", name, e)
            self.fallback_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False
            total_comp_params += orig_params

    # Post-compression forward sanity check (single attempt).
    forward_check_success = True
    try:
        self.model.eval()
        with torch.no_grad():
            first_p = next(self.model.parameters(), None)
            dev = first_p.device if first_p is not None else torch.device("cpu")
            cfg = getattr(self.model, "config", None)
            vocab_size = getattr(cfg, "vocab_size", 1000)
            max_tok = min(100, vocab_size - 1) if vocab_size > 10 else 2
            min_tok = min(2, max_tok)
            dummy = torch.randint(min_tok, max_tok + 1, (1, 8), device=dev, dtype=torch.long)
            attn_mask = torch.ones_like(dummy)
            try:
                self.model(input_ids=dummy, attention_mask=attn_mask)
            except Exception:
                # Try hidden-states form.
                first_linear = next(
                    (m for m in self.model.modules() if isinstance(m, (nn.Linear, SVDLinear))),
                    None,
                )
                in_dim = getattr(first_linear, "in_features", getattr(cfg, "hidden_size", 8))
                dummy_float = torch.randn(
                    1,
                    8,
                    in_dim,
                    device=dev,
                    dtype=first_p.dtype if first_p is not None else torch.float32,
                )
                self.model(dummy_float)
    except Exception as e:
        logger.warning("HybridSVD post-compression forward check warning: %s", e)
        forward_check_success = False

    effective_ratio = (total_comp_params / total_orig_params) if total_orig_params > 0 else 1.0

    try:
        self._compressed_params = self._count_params(self.model)
    except Exception:
        self._compressed_params = sum(p.numel() for p in self.model.parameters())

    stats = self.get_compression_stats()
    avg_err = (
        (sum(self.layer_errors.values()) / max(1, len(self.layer_errors)))
        if self.layer_errors
        else 0.0
    )
    stats.update(
        {
            "replaced_layers_count": replaced_count,
            "compressed_layers": self.compressed_layers,
            "fallback_layers": self.fallback_layers,
            "skipped_layers": self.skipped_layers,
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "num_skipped": len(self.skipped_layers),
            "layer_errors": self.layer_errors,
            "svd_rank": self.svd_rank,
            "lora_r": self.lora_r,
            "lora_alpha": self.lora_alpha,
            "max_error": self.max_error,
            "target_attention_only": self.target_attention_only,
            "fallback_to_original": self.fallback_to_original,
            "avg_reconstruction_error": avg_err,
            "effective_parameter_ratio": round(effective_ratio, 4),
            "forward_check_success": forward_check_success,
            "method": "hybrid_svd",
        }
    )
    return self.model, stats

Decompose linear layers via SVD and replace with SVDLinear.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class KMeans (n_clusters=8, *, init='k-means++', n_init='auto', max_iter=300, tol=0.0001, verbose=0, random_state=None, copy_x=True, algorithm='lloyd')

Expand source code

class KMeans(_BaseKMeans):
    """K-Means clustering.

    Read more in the :ref:`User Guide <k_means>`.

    Parameters
    ----------

    n_clusters : int, default=8
        The number of clusters to form as well as the number of
        centroids to generate.

        For an example of how to choose an optimal value for `n_clusters` refer to
        :ref:`sphx_glr_auto_examples_cluster_plot_kmeans_silhouette_analysis.py`.

    init : {'k-means++', 'random'}, callable or array-like of shape \
            (n_clusters, n_features), default='k-means++'
        Method for initialization:

        * 'k-means++' : selects initial cluster centroids using sampling \
            based on an empirical probability distribution of the points' \
            contribution to the overall inertia. This technique speeds up \
            convergence. The algorithm implemented is "greedy k-means++". It \
            differs from the vanilla k-means++ by making several trials at \
            each sampling step and choosing the best centroid among them.

        * 'random': choose `n_clusters` observations (rows) at random from \
        data for the initial centroids.

        * If an array is passed, it should be of shape (n_clusters, n_features)\
        and gives the initial centers.

        * If a callable is passed, it should take arguments X, n_clusters and a\
        random state and return an initialization.

        For an example of how to use the different `init` strategies, see
        :ref:`sphx_glr_auto_examples_cluster_plot_kmeans_digits.py`.

        For an evaluation of the impact of initialization, see the example
        :ref:`sphx_glr_auto_examples_cluster_plot_kmeans_stability_low_dim_dense.py`.

    n_init : 'auto' or int, default='auto'
        Number of times the k-means algorithm is run with different centroid
        seeds. The final results is the best output of `n_init` consecutive runs
        in terms of inertia. Several runs are recommended for sparse
        high-dimensional problems (see :ref:`kmeans_sparse_high_dim`).

        When `n_init='auto'`, the number of runs depends on the value of init:
        10 if using `init='random'` or `init` is a callable;
        1 if using `init='k-means++'` or `init` is an array-like.

        .. versionadded:: 1.2
           Added 'auto' option for `n_init`.

        .. versionchanged:: 1.4
           Default value for `n_init` changed to `'auto'`.

    max_iter : int, default=300
        Maximum number of iterations of the k-means algorithm for a
        single run.

    tol : float, default=1e-4
        Relative tolerance with regards to Frobenius norm of the difference
        in the cluster centers of two consecutive iterations to declare
        convergence.

    verbose : int, default=0
        Verbosity mode.

    random_state : int, RandomState instance or None, default=None
        Determines random number generation for centroid initialization. Use
        an int to make the randomness deterministic.
        See :term:`Glossary <random_state>`.

    copy_x : bool, default=True
        When pre-computing distances it is more numerically accurate to center
        the data first. If copy_x is True (default), then the original data is
        not modified. If False, the original data is modified, and put back
        before the function returns, but small numerical differences may be
        introduced by subtracting and then adding the data mean. Note that if
        the original data is not C-contiguous, a copy will be made even if
        copy_x is False. If the original data is sparse, but not in CSR format,
        a copy will be made even if copy_x is False.

    algorithm : {"lloyd", "elkan"}, default="lloyd"
        K-means algorithm to use. The classical EM-style algorithm is `"lloyd"`.
        The `"elkan"` variation can be more efficient on some datasets with
        well-defined clusters, by using the triangle inequality. However it's
        more memory intensive due to the allocation of an extra array of shape
        `(n_samples, n_clusters)`.

        .. versionchanged:: 0.18
            Added Elkan algorithm

        .. versionchanged:: 1.1
            Renamed "full" to "lloyd", and deprecated "auto" and "full".
            Changed "auto" to use "lloyd" instead of "elkan".

    Attributes
    ----------
    cluster_centers_ : ndarray of shape (n_clusters, n_features)
        Coordinates of cluster centers. If the algorithm stops before fully
        converging (see ``tol`` and ``max_iter``), these will not be
        consistent with ``labels_``.

    labels_ : ndarray of shape (n_samples,)
        Labels of each point

    inertia_ : float
        Sum of squared distances of samples to their closest cluster center,
        weighted by the sample weights if provided.

    n_iter_ : int
        Number of iterations run.

    n_features_in_ : int
        Number of features seen during :term:`fit`.

        .. versionadded:: 0.24

    feature_names_in_ : ndarray of shape (`n_features_in_`,)
        Names of features seen during :term:`fit`. Defined only when `X`
        has feature names that are all strings.

        .. versionadded:: 1.0

    See Also
    --------
    MiniBatchKMeans : Alternative online implementation that does incremental
        updates of the centers positions using mini-batches.
        For large scale learning (say n_samples > 10k) MiniBatchKMeans is
        probably much faster than the default batch implementation.

    Notes
    -----
    The k-means problem is solved using either Lloyd's or Elkan's algorithm.

    The average complexity is given by O(k n T), where n is the number of
    samples and T is the number of iteration.

    The worst case complexity is given by O(n^(k+2/p)) with
    n = n_samples, p = n_features.
    Refer to :doi:`"How slow is the k-means method?" D. Arthur and S. Vassilvitskii -
    SoCG2006.<10.1145/1137856.1137880>` for more details.

    In practice, the k-means algorithm is very fast (one of the fastest
    clustering algorithms available), but it falls in local minima. That's why
    it can be useful to restart it several times.

    If the algorithm stops before fully converging (because of ``tol`` or
    ``max_iter``), ``labels_`` and ``cluster_centers_`` will not be consistent,
    i.e. the ``cluster_centers_`` will not be the means of the points in each
    cluster. Also, the estimator will reassign ``labels_`` after the last
    iteration to make ``labels_`` consistent with ``predict`` on the training
    set.

    Examples
    --------

    >>> from sklearn.cluster import KMeans
    >>> import numpy as np
    >>> X = np.array([[1, 2], [1, 4], [1, 0],
    ...               [10, 2], [10, 4], [10, 0]])
    >>> kmeans = KMeans(n_clusters=2, random_state=0, n_init="auto").fit(X)
    >>> kmeans.labels_
    array([1, 1, 1, 0, 0, 0], dtype=int32)
    >>> kmeans.predict([[0, 0], [12, 3]])
    array([1, 0], dtype=int32)
    >>> kmeans.cluster_centers_
    array([[10.,  2.],
           [ 1.,  2.]])

    For examples of common problems with K-Means and how to address them see
    :ref:`sphx_glr_auto_examples_cluster_plot_kmeans_assumptions.py`.

    For a demonstration of how K-Means can be used to cluster text documents see
    :ref:`sphx_glr_auto_examples_text_plot_document_clustering.py`.

    For a comparison between K-Means and MiniBatchKMeans refer to example
    :ref:`sphx_glr_auto_examples_cluster_plot_mini_batch_kmeans.py`.

    For a comparison between K-Means and BisectingKMeans refer to example
    :ref:`sphx_glr_auto_examples_cluster_plot_bisect_kmeans.py`.
    """

    _parameter_constraints: dict = {
        **_BaseKMeans._parameter_constraints,
        "copy_x": ["boolean"],
        "algorithm": [StrOptions({"lloyd", "elkan"})],
    }

    def __init__(
        self,
        n_clusters=8,
        *,
        init="k-means++",
        n_init="auto",
        max_iter=300,
        tol=1e-4,
        verbose=0,
        random_state=None,
        copy_x=True,
        algorithm="lloyd",
    ):
        super().__init__(
            n_clusters=n_clusters,
            init=init,
            n_init=n_init,
            max_iter=max_iter,
            tol=tol,
            verbose=verbose,
            random_state=random_state,
        )

        self.copy_x = copy_x
        self.algorithm = algorithm

    def _check_params_vs_input(self, X):
        super()._check_params_vs_input(X, default_n_init=10)

        self._algorithm = self.algorithm
        if self._algorithm == "elkan" and self.n_clusters == 1:
            warnings.warn(
                (
                    "algorithm='elkan' doesn't make sense for a single "
                    "cluster. Using 'lloyd' instead."
                ),
                RuntimeWarning,
            )
            self._algorithm = "lloyd"

    def _warn_mkl_vcomp(self, n_active_threads):
        """Warn when vcomp and mkl are both present"""
        warnings.warn(
            "KMeans is known to have a memory leak on Windows "
            "with MKL, when there are less chunks than available "
            "threads. You can avoid it by setting the environment"
            f" variable OMP_NUM_THREADS={n_active_threads}."
        )

    @_fit_context(prefer_skip_nested_validation=True)
    def fit(self, X, y=None, sample_weight=None):
        """Compute k-means clustering.

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape (n_samples, n_features)
            Training instances to cluster. It must be noted that the data
            will be converted to C ordering, which will cause a memory
            copy if the given data is not C-contiguous.
            If a sparse matrix is passed, a copy will be made if it's not in
            CSR format.

        y : Ignored
            Not used, present here for API consistency by convention.

        sample_weight : array-like of shape (n_samples,), default=None
            The weights for each observation in X. If None, all observations
            are assigned equal weight. `sample_weight` is not used during
            initialization if `init` is a callable or a user provided array.

            .. versionadded:: 0.20

        Returns
        -------
        self : object
            Fitted estimator.
        """
        X = validate_data(
            self,
            X,
            accept_sparse="csr",
            dtype=[np.float64, np.float32],
            order="C",
            copy=self.copy_x,
            accept_large_sparse=False,
        )

        self._check_params_vs_input(X)

        random_state = check_random_state(self.random_state)
        sample_weight = _check_sample_weight(sample_weight, X, dtype=X.dtype)
        self._n_threads = _openmp_effective_n_threads()

        # Validate init array
        init = self.init
        init_is_array_like = _is_arraylike_not_scalar(init)
        if init_is_array_like:
            init = check_array(init, dtype=X.dtype, copy=True, order="C")
            self._validate_center_shape(X, init)

        # subtract of mean of x for more accurate distance computations
        if not sp.issparse(X):
            X_mean = X.mean(axis=0)
            # The copy was already done above
            X -= X_mean

            if init_is_array_like:
                init -= X_mean

        # precompute squared norms of data points
        x_squared_norms = row_norms(X, squared=True)

        if self._algorithm == "elkan":
            kmeans_single = _kmeans_single_elkan
        else:
            kmeans_single = _kmeans_single_lloyd
            self._check_mkl_vcomp(X, X.shape[0])

        best_inertia, best_labels = None, None

        for i in range(self._n_init):
            # Initialize centers
            centers_init = self._init_centroids(
                X,
                x_squared_norms=x_squared_norms,
                init=init,
                random_state=random_state,
                sample_weight=sample_weight,
            )
            if self.verbose:
                print("Initialization complete")

            # run a k-means once
            labels, inertia, centers, n_iter_ = kmeans_single(
                X,
                sample_weight,
                centers_init,
                max_iter=self.max_iter,
                verbose=self.verbose,
                tol=self._tol,
                n_threads=self._n_threads,
            )

            # determine if these results are the best so far
            # we chose a new run if it has a better inertia and the clustering is
            # different from the best so far (it's possible that the inertia is
            # slightly better even if the clustering is the same with potentially
            # permuted labels, due to rounding errors)
            if best_inertia is None or (
                inertia < best_inertia
                and not _is_same_clustering(labels, best_labels, self.n_clusters)
            ):
                best_labels = labels
                best_centers = centers
                best_inertia = inertia
                best_n_iter = n_iter_

        if not sp.issparse(X):
            if not self.copy_x:
                X += X_mean
            best_centers += X_mean

        distinct_clusters = len(set(best_labels))
        if distinct_clusters < self.n_clusters:
            warnings.warn(
                "Number of distinct clusters ({}) found smaller than "
                "n_clusters ({}). Possibly due to duplicate points "
                "in X.".format(distinct_clusters, self.n_clusters),
                ConvergenceWarning,
                stacklevel=2,
            )

        self.cluster_centers_ = best_centers
        self._n_features_out = self.cluster_centers_.shape[0]
        self.labels_ = best_labels
        self.inertia_ = best_inertia
        self.n_iter_ = best_n_iter
        return self

K-Means clustering.

Read more in the :ref:User Guide <k_means>.

Parameters

n_clusters : int, default=8

The number of clusters to form as well as the number of centroids to generate.

For an example of how to choose an optimal value for n_clusters refer to :ref:sphx_glr_auto_examples_cluster_plot_kmeans_silhouette_analysis.py.

init : {'k-means++', 'random'}, callable or array-like of shape (n_clusters, n_features), default='k-means++'

Method for initialization:

'k-means++' : selects initial cluster centroids using sampling based on an empirical probability distribution of the points' contribution to the overall inertia. This technique speeds up convergence. The algorithm implemented is "greedy k-means++". It differs from the vanilla k-means++ by making several trials at each sampling step and choosing the best centroid among them.
'random': choose n_clusters observations (rows) at random from data for the initial centroids.
If an array is passed, it should be of shape (n_clusters, n_features) and gives the initial centers.
If a callable is passed, it should take arguments X, n_clusters and a random state and return an initialization.

For an example of how to use the different init strategies, see :ref:sphx_glr_auto_examples_cluster_plot_kmeans_digits.py.

For an evaluation of the impact of initialization, see the example :ref:sphx_glr_auto_examples_cluster_plot_kmeans_stability_low_dim_dense.py.

n_init : 'auto' or int, default='auto'

Number of times the k-means algorithm is run with different centroid seeds. The final results is the best output of n_init consecutive runs in terms of inertia. Several runs are recommended for sparse high-dimensional problems (see :ref:kmeans_sparse_high_dim).

When n_init='auto', the number of runs depends on the value of init: 10 if using init='random' or init is a callable; 1 if using init='k-means++' or init is an array-like.

Added in version: 1.2

Added 'auto' option for n_init.

Changed in version: 1.4

Default value for n_init changed to 'auto'.

max_iter : int, default=300

Maximum number of iterations of the k-means algorithm for a single run.

tol : float, default=1e-4

Relative tolerance with regards to Frobenius norm of the difference in the cluster centers of two consecutive iterations to declare convergence.

verbose : int, default=0

Verbosity mode.

random_state : int, RandomState instance or None, default=None

Determines random number generation for centroid initialization. Use an int to make the randomness deterministic. See :term:Glossary <random_state>.

copy_x : bool, default=True

When pre-computing distances it is more numerically accurate to center the data first. If copy_x is True (default), then the original data is not modified. If False, the original data is modified, and put back before the function returns, but small numerical differences may be introduced by subtracting and then adding the data mean. Note that if the original data is not C-contiguous, a copy will be made even if copy_x is False. If the original data is sparse, but not in CSR format, a copy will be made even if copy_x is False.

algorithm : {"lloyd", "elkan"}, default="lloyd"

K-means algorithm to use. The classical EM-style algorithm is "lloyd". The "elkan" variation can be more efficient on some datasets with well-defined clusters, by using the triangle inequality. However it's more memory intensive due to the allocation of an extra array of shape (n_samples, n_clusters).

Changed in version: 0.18

Added Elkan algorithm

Changed in version: 1.1

Renamed "full" to "lloyd", and deprecated "auto" and "full". Changed "auto" to use "lloyd" instead of "elkan".

Attributes

cluster_centers_ : ndarray of shape (n_clusters, n_features): Coordinates of cluster centers. If the algorithm stops before fully converging (see tol and max_iter), these will not be consistent with labels_.
labels_ : ndarray of shape (n_samples,): Labels of each point
inertia_ : float: Sum of squared distances of samples to their closest cluster center, weighted by the sample weights if provided.
n_iter_ : int: Number of iterations run.
n_features_in_ : int: Number of features seen during :term:fit.

Added in version: 0.24
feature_names_in_ : ndarray of shape (n_features_in_,): Names of features seen during :term:fit. Defined only when X has feature names that are all strings.

Added in version: 1.0

Notes

The k-means problem is solved using either Lloyd's or Elkan's algorithm.

The average complexity is given by O(k n T), where n is the number of samples and T is the number of iteration.

The worst case complexity is given by O(n^(k+2/p)) with n = n_samples, p = n_features. Refer to :doi:"How slow is the k-means method?" D. Arthur and S. Vassilvitskii - SoCG2006.<10.1145/1137856.1137880> for more details.

In practice, the k-means algorithm is very fast (one of the fastest clustering algorithms available), but it falls in local minima. That's why it can be useful to restart it several times.

If the algorithm stops before fully converging (because of tol or max_iter), labels_ and cluster_centers_ will not be consistent, i.e. the cluster_centers_ will not be the means of the points in each cluster. Also, the estimator will reassign labels_ after the last iteration to make labels_ consistent with predict on the training set.

Examples

>>> from sklearn.cluster import KMeans
>>> import numpy as np
>>> X = np.array([[1, 2], [1, 4], [1, 0],
...               [10, 2], [10, 4], [10, 0]])
>>> kmeans = KMeans(n_clusters=2, random_state=0, n_init="auto").fit(X)
>>> kmeans.labels_
array([1, 1, 1, 0, 0, 0], dtype=int32)
>>> kmeans.predict([[0, 0], [12, 3]])
array([1, 0], dtype=int32)
>>> kmeans.cluster_centers_
array([[10.,  2.],
       [ 1.,  2.]])

For examples of common problems with K-Means and how to address them see :ref:sphx_glr_auto_examples_cluster_plot_kmeans_assumptions.py.

For a demonstration of how K-Means can be used to cluster text documents see :ref:sphx_glr_auto_examples_text_plot_document_clustering.py.

For a comparison between K-Means and MiniBatchKMeans refer to example :ref:sphx_glr_auto_examples_cluster_plot_mini_batch_kmeans.py.

For a comparison between K-Means and BisectingKMeans refer to example :ref:sphx_glr_auto_examples_cluster_plot_bisect_kmeans.py.

Ancestors

sklearn.cluster._kmeans._BaseKMeans
sklearn.base.ClassNamePrefixFeaturesOutMixin
sklearn.base.TransformerMixin
sklearn.utils._set_output._SetOutputMixin
sklearn.base.ClusterMixin
sklearn.base.BaseEstimator
sklearn.utils._repr_html.base.ReprHTMLMixin
sklearn.utils._repr_html.base._HTMLDocumentationLinkMixin
sklearn.utils._metadata_requests._MetadataRequester
abc.ABC

Methods

def fit(self, X, y=None, sample_weight=None)

Expand source code

@_fit_context(prefer_skip_nested_validation=True)
def fit(self, X, y=None, sample_weight=None):
    """Compute k-means clustering.

    Parameters
    ----------
    X : {array-like, sparse matrix} of shape (n_samples, n_features)
        Training instances to cluster. It must be noted that the data
        will be converted to C ordering, which will cause a memory
        copy if the given data is not C-contiguous.
        If a sparse matrix is passed, a copy will be made if it's not in
        CSR format.

    y : Ignored
        Not used, present here for API consistency by convention.

    sample_weight : array-like of shape (n_samples,), default=None
        The weights for each observation in X. If None, all observations
        are assigned equal weight. `sample_weight` is not used during
        initialization if `init` is a callable or a user provided array.

        .. versionadded:: 0.20

    Returns
    -------
    self : object
        Fitted estimator.
    """
    X = validate_data(
        self,
        X,
        accept_sparse="csr",
        dtype=[np.float64, np.float32],
        order="C",
        copy=self.copy_x,
        accept_large_sparse=False,
    )

    self._check_params_vs_input(X)

    random_state = check_random_state(self.random_state)
    sample_weight = _check_sample_weight(sample_weight, X, dtype=X.dtype)
    self._n_threads = _openmp_effective_n_threads()

    # Validate init array
    init = self.init
    init_is_array_like = _is_arraylike_not_scalar(init)
    if init_is_array_like:
        init = check_array(init, dtype=X.dtype, copy=True, order="C")
        self._validate_center_shape(X, init)

    # subtract of mean of x for more accurate distance computations
    if not sp.issparse(X):
        X_mean = X.mean(axis=0)
        # The copy was already done above
        X -= X_mean

        if init_is_array_like:
            init -= X_mean

    # precompute squared norms of data points
    x_squared_norms = row_norms(X, squared=True)

    if self._algorithm == "elkan":
        kmeans_single = _kmeans_single_elkan
    else:
        kmeans_single = _kmeans_single_lloyd
        self._check_mkl_vcomp(X, X.shape[0])

    best_inertia, best_labels = None, None

    for i in range(self._n_init):
        # Initialize centers
        centers_init = self._init_centroids(
            X,
            x_squared_norms=x_squared_norms,
            init=init,
            random_state=random_state,
            sample_weight=sample_weight,
        )
        if self.verbose:
            print("Initialization complete")

        # run a k-means once
        labels, inertia, centers, n_iter_ = kmeans_single(
            X,
            sample_weight,
            centers_init,
            max_iter=self.max_iter,
            verbose=self.verbose,
            tol=self._tol,
            n_threads=self._n_threads,
        )

        # determine if these results are the best so far
        # we chose a new run if it has a better inertia and the clustering is
        # different from the best so far (it's possible that the inertia is
        # slightly better even if the clustering is the same with potentially
        # permuted labels, due to rounding errors)
        if best_inertia is None or (
            inertia < best_inertia
            and not _is_same_clustering(labels, best_labels, self.n_clusters)
        ):
            best_labels = labels
            best_centers = centers
            best_inertia = inertia
            best_n_iter = n_iter_

    if not sp.issparse(X):
        if not self.copy_x:
            X += X_mean
        best_centers += X_mean

    distinct_clusters = len(set(best_labels))
    if distinct_clusters < self.n_clusters:
        warnings.warn(
            "Number of distinct clusters ({}) found smaller than "
            "n_clusters ({}). Possibly due to duplicate points "
            "in X.".format(distinct_clusters, self.n_clusters),
            ConvergenceWarning,
            stacklevel=2,
        )

    self.cluster_centers_ = best_centers
    self._n_features_out = self.cluster_centers_.shape[0]
    self.labels_ = best_labels
    self.inertia_ = best_inertia
    self.n_iter_ = best_n_iter
    return self

Compute k-means clustering.

Parameters

X : {array-like, sparse matrix} of shape (n_samples, n_features): Training instances to cluster. It must be noted that the data will be converted to C ordering, which will cause a memory copy if the given data is not C-contiguous. If a sparse matrix is passed, a copy will be made if it's not in CSR format.
y : Ignored: Not used, present here for API consistency by convention.
sample_weight : array-like of shape (n_samples,), default=None: The weights for each observation in X. If None, all observations are assigned equal weight. sample_weight is not used during initialization if init is a callable or a user provided array.

Added in version: 0.20

Returns

self : object: Fitted estimator.

def set_fit_request(self: sklearn.cluster._kmeans.KMeans, *, sample_weight: bool | str | None = '$UNCHANGED$') ‑> sklearn.cluster._kmeans.KMeans

Expand source code

def func(*args, **kw):
    """Updates the `_metadata_request` attribute of the consumer (`instance`)
    for the parameters provided as `**kw`.

    This docstring is overwritten below.
    See REQUESTER_DOC for expected functionality.
    """
    if not _routing_enabled():
        raise RuntimeError(
            "This method is only available when metadata routing is enabled."
            " You can enable it using"
            " sklearn.set_config(enable_metadata_routing=True)."
        )

    if self.validate_keys and (set(kw) - set(self.keys)):
        raise TypeError(
            f"Unexpected args: {set(kw) - set(self.keys)} in {self.name}. "
            f"Accepted arguments are: {set(self.keys)}"
        )

    # This makes it possible to use the decorated method as an unbound method,
    # for instance when monkeypatching.
    # https://github.com/scikit-learn/scikit-learn/issues/28632
    if instance is None:
        _instance = args[0]
        args = args[1:]
    else:
        _instance = instance

    # Replicating python's behavior when positional args are given other than
    # `self`, and `self` is only allowed if this method is unbound.
    if args:
        raise TypeError(
            f"set_{self.name}_request() takes 0 positional argument but"
            f" {len(args)} were given"
        )

    requests = _instance._get_metadata_request()
    method_metadata_request = getattr(requests, self.name)

    for prop, alias in kw.items():
        if alias is not UNCHANGED:
            method_metadata_request.add_request(param=prop, alias=alias)
    _instance._metadata_request = requests

    return _instance

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a :term:meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see :func:sklearn.set_config). Please check the :ref:User Guide <metadata_routing> on how the routing mechanism works.

The options for each parameter are:

True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to fit.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version: 1.3

Parameters

sample_weight : str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for sample_weight parameter in fit.

Returns

self : object: The updated object.

def set_score_request(self: sklearn.cluster._kmeans.KMeans, *, sample_weight: bool | str | None = '$UNCHANGED$') ‑> sklearn.cluster._kmeans.KMeans

Expand source code

def func(*args, **kw):
    """Updates the `_metadata_request` attribute of the consumer (`instance`)
    for the parameters provided as `**kw`.

    This docstring is overwritten below.
    See REQUESTER_DOC for expected functionality.
    """
    if not _routing_enabled():
        raise RuntimeError(
            "This method is only available when metadata routing is enabled."
            " You can enable it using"
            " sklearn.set_config(enable_metadata_routing=True)."
        )

    if self.validate_keys and (set(kw) - set(self.keys)):
        raise TypeError(
            f"Unexpected args: {set(kw) - set(self.keys)} in {self.name}. "
            f"Accepted arguments are: {set(self.keys)}"
        )

    # This makes it possible to use the decorated method as an unbound method,
    # for instance when monkeypatching.
    # https://github.com/scikit-learn/scikit-learn/issues/28632
    if instance is None:
        _instance = args[0]
        args = args[1:]
    else:
        _instance = instance

    # Replicating python's behavior when positional args are given other than
    # `self`, and `self` is only allowed if this method is unbound.
    if args:
        raise TypeError(
            f"set_{self.name}_request() takes 0 positional argument but"
            f" {len(args)} were given"
        )

    requests = _instance._get_metadata_request()
    method_metadata_request = getattr(requests, self.name)

    for prop, alias in kw.items():
        if alias is not UNCHANGED:
            method_metadata_request.add_request(param=prop, alias=alias)
    _instance._metadata_request = requests

    return _instance

Configure whether metadata should be requested to be passed to the score method.

The options for each parameter are:

True: metadata is requested, and passed to score if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to score.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version: 1.3

Parameters

sample_weight : str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for sample_weight parameter in score.

Returns

self : object: The updated object.

class KMeansCompressor (model: nn.Module | None = None, n_clusters: int = 256, bits: int | None = None, max_error: float = 0.35, target_attention_only: bool = False, fallback_to_original: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, skip_attention_qkv: bool = False, skip_should_compress_filter: bool = False, min_layer_dim: int = 8, kmeans_seed: int = 51966, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("kmeans_class")
@COMPRESSOR_REGISTRY.register("kmeans_quantize_class")
@COMPRESSOR_REGISTRY.register("kmeans_compressor_class")
class KMeansCompressor(BaseCompressor):
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module | None = None,
        n_clusters: int = 256,
        bits: int | None = None,
        max_error: float = 0.35,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        skip_attention_qkv: bool = False,
        skip_should_compress_filter: bool = False,
        min_layer_dim: int = 8,
        kmeans_seed: int = 0xCAFE,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model if model is not None else nn.Identity(), name="kmeans", **kwargs)
        cfg = cfg or kwargs.get("comp_config")
        self.cfg = cfg

        parsed_bits = _cfg_get(cfg, "kmeans_bits", _cfg_get(cfg, "bits", bits))
        if parsed_bits is not None:
            self.n_clusters = 2 ** int(parsed_bits)
        else:
            self.n_clusters = int(
                _cfg_get(cfg, "kmeans_clusters", _cfg_get(cfg, "n_clusters", n_clusters))
            )
        self.bits = int(math.ceil(math.log2(max(self.n_clusters, 2))))

        self.max_error = float(
            _cfg_get(cfg, "kmeans_max_error", _cfg_get(cfg, "max_error", max_error))
        )
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "kmeans_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "kmeans_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        self.lora_r = int(_cfg_get(cfg, "lora_r", _cfg_get(cfg, "lora_rank", lora_r)))
        self.lora_alpha = float(_cfg_get(cfg, "lora_alpha", lora_alpha))
        self.skip_attention_qkv = bool(_cfg_get(cfg, "skip_attention_qkv", skip_attention_qkv))
        self.skip_should_compress_filter = bool(
            _cfg_get(
                cfg,
                "skip_should_compress_filter",
                skip_should_compress_filter,
            )
        )
        self.min_layer_dim = max(1, int(_cfg_get(cfg, "min_layer_dim", min_layer_dim)))
        self.kmeans_seed = int(_cfg_get(cfg, "kmeans_seed", kmeans_seed))

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_storage: dict[str, dict[str, int]] = {}
        self.stats: dict[str, Any] = {}
        self._compressed_params: int = 0

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            if any(
                p in name.lower()
                for p in (
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                )
            ):
                return True
        return False

    def _is_target_module(self, name: str, module: nn.Module, _debug: bool = False) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ("lora_A", "lora_B", "base_layer")):
            if _debug:
                logger.info(f"  [skip] {name}: lora/base_layer")
            return False
        if any(
            skip in name_lower
            for skip in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            )
        ):
            if _debug:
                logger.info(f"  [skip] {name}: skip keyword")
            return False

        if self.target_attention_only:
            if not any(k in name_lower for k in _ATTN_NAMES):
                if _debug:
                    logger.info(f"  [skip] {name}: not attention")
                return False

        if self._is_gqa_asymmetric(name, module):
            if _debug:
                logger.info(f"  [skip] {name}: GQA asymmetric")
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_lower = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_lower or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d):
            if _debug:
                logger.info(f"  [skip] {name}: not linear-like")
            return False
        if isinstance(target_mod, KMeansLinear):
            if _debug:
                logger.info(f"  [skip] {name}: already KMeansLinear")
            return False

        w = getattr(target_mod, "weight", None)
        if w is None:
            if _debug:
                logger.info(f"  [skip] {name}: weight None")
            return False
        if min(w.shape[0], w.shape[1]) < self.min_layer_dim:
            if _debug:
                logger.info(f"  [skip] {name}: dim too small {tuple(w.shape)}")
            return False

        if not self.skip_should_compress_filter:
            try:
                if not should_compress_linear(self.model, name):
                    if _debug:
                        logger.info(f"  [skip] {name}: should_compress_linear=False")
                    return False
            except Exception as e:
                if _debug:
                    logger.info(f"  [skip] {name}: should_compress_linear raised {e}")
                return False

        if _debug:
            logger.info(f"  [TARGET] {name}: shape={tuple(w.shape)}")
        return True

    # ------------------------------------------------------------------
    def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
        # Idempotence
        if getattr(self.model, "_kmeans_compressed", False):
            logger.warning("KMeansCompressor: model already compressed; skipping.")
            try:
                return self.model, self.get_compression_stats()
            except Exception:
                return self.model, {}

        if model is not None:
            self.model = model
            try:
                self._original_params = self._count_params(self.model)
            except Exception:
                self._original_params = sum(p.numel() for p in self.model.parameters())

        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.layer_storage = {}
        seen_weights: set[int] = set()

        total_orig_bytes = 0
        total_comp_bytes = 0

        target_list = [
            (name, module)
            for name, module in self.model.named_modules()
            if self._is_target_module(name, module)
        ]

        if not target_list:
            logger.warning("KMeansCompressor: no compressible layers found.")
            n = 0
            for name, module in self.model.named_modules():
                n += 1
                self._is_target_module(name, module, _debug=True)
            logger.warning("KMeansCompressor diagnostic: %d modules scanned, 0 targets.", n)
            try:
                return self.model, self.get_compression_stats()
            except Exception:
                return self.model, {}

        logger.info("KMeansCompressor: %d target layers identified.", len(target_list))

        for name, module in target_list:
            target_mod = getattr(module, "base_layer", module)
            if not hasattr(target_mod, "weight") or target_mod.weight is None:
                continue

            w_id = id(target_mod.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            cls_lower = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_lower or hasattr(target_mod, "nf")
            orig_w = target_mod.weight.data
            W_target = orig_w.t().contiguous().float() if is_conv1d else orig_w.float()
            norm_orig = torch.norm(W_target, p="fro").clamp_min(1e-8)
            _out_f, _in_f = W_target.shape[0], W_target.shape[1]

            # Original bytes
            orig_bytes = orig_w.numel() * orig_w.element_size()
            bias_holder = getattr(target_mod, "bias", None)
            if bias_holder is not None:
                orig_bytes += bias_holder.numel() * bias_holder.element_size()
            total_orig_bytes += orig_bytes

            new_layer: KMeansLinear | None = None
            err = 1.0
            try:
                new_layer = KMeansLinear(
                    module,
                    n_clusters=self.n_clusters,
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                    kmeans_seed=self.kmeans_seed,
                )
                recon_w = new_layer.get_weight().float()
                err = float(torch.norm(W_target - recon_w, p="fro").item() / norm_orig.item())
            except Exception as ex:
                logger.warning("KMeansCompressor: failed on %s: %s", name, ex)
                new_layer = None
                err = 1.0

            self.layer_errors[name] = float(err)

            if new_layer is None or err > self.max_error:
                logger.warning(
                    "KMeansCompressor: error %.4f on %s > %.4f. %s",
                    err,
                    name,
                    self.max_error,
                    ("Falling back." if self.fallback_to_original else "Using compressed."),
                )
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_bytes += orig_bytes
                continue

            # KM-FIX-18: preserve LoRA wrapper
            try:
                if hasattr(module, "base_layer"):
                    module.base_layer = new_layer
                    # Keep LoRA param trainable
                    for pn, p in module.named_parameters():
                        if "lora_" in pn.lower() and not pn.startswith("base_layer"):
                            p.requires_grad = True
                else:
                    replace_module(self.model, name, new_layer)
                self.compressed_layers.append(name)
            except Exception as e:
                logger.warning("KMeansCompressor: replacement failed on %s: %s", name, e)
                self.fallback_layers.append(name)
                total_comp_bytes += orig_bytes
                continue

            # Honest byte accounting
            st = new_layer.storage_bytes()
            self.layer_storage[name] = st
            total_comp_bytes += sum(st.values())

        # Forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                dev = _get_device(self.model)
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                try:
                    self.model(input_ids=dummy, attention_mask=torch.ones_like(dummy))
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception as e2:
                        logger.warning("KMeans forward check failed: %s", e2)
                        forward_check_success = False
        except Exception as e:
            logger.warning("KMeans forward check raised: %s", e)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        storage_ratio = total_orig_bytes / total_comp_bytes if total_comp_bytes > 0 else 1.0
        effective_ratio = total_comp_bytes / total_orig_bytes if total_orig_bytes > 0 else 1.0
        avg_err = (
            sum(self.layer_errors.values()) / max(1, len(self.layer_errors))
            if self.layer_errors
            else 0.0
        )

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "layer_storage_bytes": self.layer_storage,
                "avg_reconstruction_error": avg_err,
                "max_error": self.max_error,
                "target_attention_only": self.target_attention_only,
                "fallback_to_original": self.fallback_to_original,
                "n_clusters": self.n_clusters,
                "effective_bits_per_weight": self.bits,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                # Honest accounting
                "orig_storage_bytes": total_orig_bytes,
                "comp_storage_bytes": total_comp_bytes,
                "storage_compression_ratio": round(storage_ratio, 3),
                "storage_reduction_pct": round(100.0 * (1.0 - effective_ratio), 2),
                "effective_parameter_ratio": round(effective_ratio, 4),
                "forward_check_success": forward_check_success,
                "method": "kmeans",
                "note": (
                    "Storage accounting is honest (bytes). indices are stored as "
                    f"{_dtype_for_clusters(self.n_clusters)} to minimize overhead. "
                    "The framework's parameter counter uses model.parameters(), "
                    "which does NOT include the `indices` buffer."
                ),
            }
        )
        self.stats = stats

        try:
            if not hasattr(self.model, "compression_stats") or self.model.compression_stats is None:
                self.model.compression_stats = {}
            self.model.compression_stats.update(
                {
                    "num_compressed": len(self.compressed_layers),
                    "num_fallback": len(self.fallback_layers),
                    "compression_method": "kmeans",
                    "effective_parameter_ratio": effective_ratio,
                    "storage_compression_ratio": storage_ratio,
                }
            )
        except Exception:
            pass

        self.model._kmeans_compressed = True
        logger.info(
            "KMeans: %d compressed, %d fallback, avg_err=%.4f, storage=%.2fx.",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            storage_ratio,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class KMeansLinear (linear: nn.Module, n_clusters: int = 256, lora_r: int = 0, lora_alpha: float = 1.0, kmeans_seed: int = 51966)

Expand source code

class KMeansLinear(nn.Module):
    """
    Linear layer with weight matrix quantized via 1D k-means codebook.

    Storage:
      - centroids:  (C,)         fp16/fp32
      - indices:    (out, in)    int8/int16/int32  (smallest that fits C)
      - bias:       (out,)       fp16/fp32
      - lora_A/B:   (r, in) / (out, r)   optional

    Forward:
      W_hat = centroids[indices]
      y = F.linear(x, W_hat) + bias + scaling · (x @ A^T) @ B^T
    """

    is_custom_attention: bool = False
    is_custom_linear: bool = True

    def __init__(
        self,
        linear: nn.Module,
        n_clusters: int = 256,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        kmeans_seed: int = 0xCAFE,
    ) -> None:
        super().__init__()
        target_mod = getattr(linear, "base_layer", linear)
        cls_lower = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_lower or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.data

        if is_conv1d:
            W = raw_w.t().contiguous().float()
            self.in_features = raw_w.shape[0]
            self.out_features = raw_w.shape[1]
        else:
            W = raw_w.float()
            self.in_features = int(getattr(target_mod, "in_features", W.shape[1]))
            self.out_features = int(getattr(target_mod, "out_features", W.shape[0]))

        self.n_clusters = int(min(n_clusters, W.numel()))
        self._index_dtype = _dtype_for_clusters(self.n_clusters)

        # Detect existing LoRA
        detected_r, detected_alpha, src_A, src_B = _detect_lora(linear)
        self.lora_r = int(detected_r) if detected_r > 0 else int(lora_r)
        self.lora_alpha = float(detected_alpha) if detected_r > 0 else float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        # Fit centroids and indices
        centroids, indices = fast_kmeans_1d(
            W,
            n_clusters=self.n_clusters,
            seed=kmeans_seed,
        )
        # KM-FIX-23: centroids frozen if LoRA exists (pure quantizer mode).
        centroids_trainable = self.lora_r == 0
        self.centroids = nn.Parameter(
            centroids.to(dtype=raw_w.dtype),
            requires_grad=centroids_trainable,
        )
        # Store indices as int8/int16/int32
        self.register_buffer("indices", indices.to(self._index_dtype))

        # Bias
        bias_holder = getattr(target_mod, "bias", None)
        if bias_holder is not None:
            self.bias = nn.Parameter(bias_holder.data.clone(), requires_grad=False)
        else:
            self.register_parameter("bias", None)

        # LoRA branch
        if self.lora_r > 0:
            dev = raw_w.device
            dt = raw_w.dtype
            self.lora_A = nn.Parameter(
                torch.empty((self.lora_r, self.in_features), dtype=dt, device=dev)
            )
            self.lora_B = nn.Parameter(
                torch.zeros((self.out_features, self.lora_r), dtype=dt, device=dev)
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))

            if src_A is not None and src_B is not None:
                if src_A.shape == self.lora_A.shape:
                    self.lora_A.data.copy_(src_A.to(self.lora_A.device, dtype=self.lora_A.dtype))
                if src_B.shape == self.lora_B.shape:
                    self.lora_B.data.copy_(src_B.to(self.lora_B.device, dtype=self.lora_B.dtype))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        # KM-FIX-20: cache in eval
        self._cached_weight: torch.Tensor | None = None

    # ------------------------------------------------------------------
    def get_weight(self) -> torch.Tensor:
        # KM-FIX-20: cache in eval mode
        if not self.training and self._cached_weight is not None:
            return self._cached_weight
        idx_long = self.indices.long().reshape(-1)
        W = self.centroids.index_select(0, idx_long).reshape(self.out_features, self.in_features)
        if not self.training:
            self._cached_weight = W
        return W

    @property
    def weight(self) -> torch.Tensor:
        return self.get_weight()

    def train(self, mode: bool = True) -> KMeansLinear:
        super().train(mode)
        if mode:
            self._cached_weight = None
        return self

    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict[str, int]:
        """KM-FIX-07/22: honest storage accounting."""

        def _sz(t):
            if t is None:
                return 0
            return t.numel() * t.element_size()

        return {
            "centroids": _sz(self.centroids.data),
            "indices": _sz(self.indices),
            "bias": _sz(self.bias.data) if self.bias is not None else 0,
            "lora_A": _sz(self.lora_A.data) if self.lora_A is not None else 0,
            "lora_B": _sz(self.lora_B.data) if self.lora_B is not None else 0,
        }

    # ------------------------------------------------------------------
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Bias cast only when dtype/device mismatch
        if self.bias is not None:
            b = self.bias
            if b.dtype != x.dtype or b.device != x.device:
                b = b.to(dtype=x.dtype, device=x.device)
        else:
            b = None

        if self.training:
            out = KMeansSTEFunction.apply(self.centroids, self.indices, x, b)
        else:
            W = self.get_weight().to(dtype=x.dtype, device=x.device)
            out = F.linear(x, W, b)

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B)
            out = out + self.scaling * lora_out.to(out.dtype)

        return out

Linear layer with weight matrix quantized via 1D k-means codebook.

Storage -----= - centroids: (C,) fp16/fp32 - indices: (out, in) int8/int16/int32 (smallest that fits C) - bias: (out,) fp16/fp32 - lora_A/B: (r, in) / (out, r) optional

Forward -----= W_hat = centroids[indices] y = F.linear(x, W_hat) + bias + scaling · (x @ A^T) @ B^T

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.
var is_custom_linear : bool: The type of the None singleton.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    return self.get_weight()

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    # Bias cast only when dtype/device mismatch
    if self.bias is not None:
        b = self.bias
        if b.dtype != x.dtype or b.device != x.device:
            b = b.to(dtype=x.dtype, device=x.device)
    else:
        b = None

    if self.training:
        out = KMeansSTEFunction.apply(self.centroids, self.indices, x, b)
    else:
        W = self.get_weight().to(dtype=x.dtype, device=x.device)
        out = F.linear(x, W, b)

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B)
        out = out + self.scaling * lora_out.to(out.dtype)

    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    # KM-FIX-20: cache in eval mode
    if not self.training and self._cached_weight is not None:
        return self._cached_weight
    idx_long = self.indices.long().reshape(-1)
    W = self.centroids.index_select(0, idx_long).reshape(self.out_features, self.in_features)
    if not self.training:
        self._cached_weight = W
    return W

def storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    """KM-FIX-07/22: honest storage accounting."""

    def _sz(t):
        if t is None:
            return 0
        return t.numel() * t.element_size()

    return {
        "centroids": _sz(self.centroids.data),
        "indices": _sz(self.indices),
        "bias": _sz(self.bias.data) if self.bias is not None else 0,
        "lora_A": _sz(self.lora_A.data) if self.lora_A is not None else 0,
        "lora_B": _sz(self.lora_B.data) if self.lora_B is not None else 0,
    }

KM-FIX-07/22: honest storage accounting.

def train(self, mode: bool = True) ‑> KMeansLinear

Expand source code

def train(self, mode: bool = True) -> KMeansLinear:
    super().train(mode)
    if mode:
        self._cached_weight = None
    return self

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. :class:Dropout, :class:BatchNorm, etc.

Args
-----=
mode : bool: whether to set training mode (True) or evaluation mode (False). Default: True.
Returns
-----=
Module: self

class KMeansSTEFunction (*args, **kwargs)

Expand source code

class KMeansSTEFunction(torch.autograd.Function):
    """
    STE for codebook lookup.

    Forward:  out = F.linear(x, centroids[indices], bias)
    Backward: gradient to centroids via scatter-add; gradient to x flows
              through the same effective W.
    """

    @staticmethod
    def forward(ctx, centroids, indices, x, bias):
        # KM-FIX-17: safe indexing
        idx_long = indices.long()
        W = centroids.index_select(0, idx_long.reshape(-1)).reshape(
            indices.shape[0], indices.shape[1]
        )
        W = W.to(dtype=x.dtype)
        ctx.save_for_backward(centroids, idx_long, x, bias)
        return F.linear(x, W, bias)

    @staticmethod
    def backward(ctx, grad_output):
        centroids, idx_long, x, bias = ctx.saved_tensors
        dtype = grad_output.dtype
        device = grad_output.device

        # Reconstruct W
        W = (
            centroids.index_select(0, idx_long.reshape(-1))
            .reshape(idx_long.shape[0], idx_long.shape[1])
            .to(dtype=dtype, device=device)
        )

        x_cast = x.to(dtype=dtype, device=device)

        grad_x = grad_output.matmul(W) if ctx.needs_input_grad[2] else None

        grad_bias = None
        if bias is not None and ctx.needs_input_grad[3]:
            grad_bias = grad_output.sum(dim=list(range(grad_output.ndim - 1)))

        grad_centroids = None
        if ctx.needs_input_grad[0]:
            if x_cast.ndim == 2:
                grad_W = grad_output.t().contiguous().matmul(x_cast)
            else:
                go = grad_output.reshape(-1, grad_output.shape[-1])
                xf = x_cast.reshape(-1, x_cast.shape[-1])
                grad_W = go.t().contiguous().matmul(xf)

            grad_centroids = torch.zeros_like(centroids)
            flat_idx = idx_long.reshape(-1)
            flat_grad = grad_W.reshape(-1).to(dtype=centroids.dtype, device=centroids.device)
            grad_centroids.scatter_add_(0, flat_idx, flat_grad)

        return grad_centroids, None, grad_x, grad_bias

STE for codebook lookup.

Forward: out = F.linear(x, centroids[indices], bias) Backward: gradient to centroids via scatter-add; gradient to x flows through the same effective W.

Ancestors

torch.autograd.function.Function
torch.autograd.function._SingleLevelFunction
torch._C._FunctionBase
torch.autograd.function.FunctionCtx
torch.autograd.function._HookMixin

Static methods

def backward(ctx, grad_output)

Expand source code

@staticmethod
def backward(ctx, grad_output):
    centroids, idx_long, x, bias = ctx.saved_tensors
    dtype = grad_output.dtype
    device = grad_output.device

    # Reconstruct W
    W = (
        centroids.index_select(0, idx_long.reshape(-1))
        .reshape(idx_long.shape[0], idx_long.shape[1])
        .to(dtype=dtype, device=device)
    )

    x_cast = x.to(dtype=dtype, device=device)

    grad_x = grad_output.matmul(W) if ctx.needs_input_grad[2] else None

    grad_bias = None
    if bias is not None and ctx.needs_input_grad[3]:
        grad_bias = grad_output.sum(dim=list(range(grad_output.ndim - 1)))

    grad_centroids = None
    if ctx.needs_input_grad[0]:
        if x_cast.ndim == 2:
            grad_W = grad_output.t().contiguous().matmul(x_cast)
        else:
            go = grad_output.reshape(-1, grad_output.shape[-1])
            xf = x_cast.reshape(-1, x_cast.shape[-1])
            grad_W = go.t().contiguous().matmul(xf)

        grad_centroids = torch.zeros_like(centroids)
        flat_idx = idx_long.reshape(-1)
        flat_grad = grad_W.reshape(-1).to(dtype=centroids.dtype, device=centroids.device)
        grad_centroids.scatter_add_(0, flat_idx, flat_grad)

    return grad_centroids, None, grad_x, grad_bias

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the vjp function.)

def forward(ctx, centroids, indices, x, bias)

Expand source code

@staticmethod
def forward(ctx, centroids, indices, x, bias):
    # KM-FIX-17: safe indexing
    idx_long = indices.long()
    W = centroids.index_select(0, idx_long.reshape(-1)).reshape(
        indices.shape[0], indices.shape[1]
    )
    W = W.to(dtype=x.dtype)
    ctx.save_for_backward(centroids, idx_long, x, bias)
    return F.linear(x, W, bias)

Define the forward of the custom autograd Function.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx)::

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See :ref:combining-forward-context for more details

Usage 2 (Separate forward and ctx)::

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass


@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

The forward no longer accepts a ctx argument.
Instead, you must also override the :meth:torch.autograd.Function.setup_context staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.
See :ref:extending-autograd for more details

class KVTCCompressor (model: nn.Module, bits: int = 4, pca_energy: float = 0.95, max_error: float = 0.5, fallback_to_original: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, target_attention_only: bool = False, skip_attention_qkv: bool = True, min_rank: int = 4, target_modules: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("kvtc")
@COMPRESSOR_REGISTRY.register("kvtc_compressor")
@COMPRESSOR_REGISTRY.register("kvtc_quantizer")
@COMPRESSOR_REGISTRY.register("pca_quant")
@COMPRESSOR_REGISTRY.register("pca_quant_compressor")
class KVTCCompressor(BaseCompressor):
    """
    KVTC: KV Cache Transform Coding compressor (Static Weight PCA + Quantization).
    Compresses linear projection weights using PCA decorrelation,
    percentile clipping, uniform integer quantization, and entropy coding estimation.
    """

    def __init__(
        self,
        model: nn.Module,
        bits: int = 4,
        pca_energy: float = 0.95,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        target_attention_only: bool = False,
        skip_attention_qkv: bool = True,
        min_rank: int = 4,
        target_modules: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="kvtc", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg", getattr(model, "config", None)))
        self.cfg = cfg
        self.bits = int(_cfg_get(cfg, "kvtc_bits", _cfg_get(cfg, "bits", bits)))
        self.pca_energy = float(
            _cfg_get(cfg, "kvtc_pca_energy", _cfg_get(cfg, "pca_energy", pca_energy))
        )
        self.max_error = float(
            _cfg_get(cfg, "kvtc_max_error", _cfg_get(cfg, "max_error", max_error))
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "kvtc_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        self.lora_r = int(_cfg_get(cfg, "lora_r", lora_r))
        self.lora_alpha = float(_cfg_get(cfg, "lora_alpha", lora_alpha))
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "kvtc_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.skip_attention_qkv = bool(_cfg_get(cfg, "skip_attention_qkv", skip_attention_qkv))
        self.min_rank = max(1, int(_cfg_get(cfg, "min_rank", min_rank)))
        self.target_modules = target_modules

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.stats: dict[str, Any] = {}

    @staticmethod
    def _percentile_clip(W: torch.Tensor, low: float = 0.001, high: float = 0.999) -> torch.Tensor:
        """Clip extreme outliers from weight tensor before computing quantization scale."""
        if W.numel() == 0:
            return W
        q_low = torch.quantile(W.float(), low)
        q_high = torch.quantile(W.float(), high)
        return torch.clamp(W, min=q_low, max=q_high)

    def _quantize_to_int(
        self, W: torch.Tensor, bits: int
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Quantize continuous tensor to true integer format (torch.int8 or torch.int16).
        Uses signed symmetric quantization around 0:
            int_w in [-2^(bits-1), 2^(bits-1) - 1]
            scale = max(|W|) / (2^(bits-1) - 1)
            W_dequant ≈ int_w * scale + zero_point (zero_point = 0.0)
        """
        qmax = (2 ** (bits - 1)) - 1
        qmin = -(2 ** (bits - 1))

        max_abs = torch.clamp(W.abs().max(), min=1e-6)
        scale = (max_abs / float(qmax)).item()
        zero_point = 0.0

        q_float = torch.round(W / scale)
        q_clamped = torch.clamp(q_float, qmin, qmax)

        scale_t = torch.tensor(scale, dtype=torch.float32, device=W.device)
        zp_t = torch.tensor(zero_point, dtype=torch.float32, device=W.device)

        if bits <= 8:
            int_w = q_clamped.to(torch.int8)
        else:
            int_w = q_clamped.to(torch.int16)

        return int_w, scale_t, zp_t

    def _dequantize_from_int(
        self, int_w: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor
    ) -> torch.Tensor:
        """Dequantize integer weight back to floating point."""
        return int_w.to(dtype=scale.dtype) * scale + zero_point

    def _estimate_actual_bytes(self, model: nn.Module) -> int:
        """
        Estimate physical memory storage in bytes across all parameters and buffers,
        taking into account actual tensor element size (int8=1 byte, fp16=2 bytes, fp32=4 bytes).
        Uses tensor storage data_ptr() to avoid duplicate counting of shared weights.
        """
        total_bytes = 0
        seen_ptrs: set[int] = set()

        for p in model.parameters():
            ptr = p.data_ptr()
            if ptr not in seen_ptrs:
                seen_ptrs.add(ptr)
                total_bytes += p.numel() * p.element_size()

        for b in model.buffers():
            ptr = b.data_ptr()
            if ptr not in seen_ptrs:
                seen_ptrs.add(ptr)
                total_bytes += b.numel() * b.element_size()

        return total_bytes

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            name_lower = name.lower()
            if any(
                p in name_lower
                for p in [
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                ]
            ):
                return True
        return False

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
            return False
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            ]
        ):
            return False

        if self.target_attention_only:
            if not any(attn_kw in name_lower for attn_kw in _ATTN_NAMES):
                return False

        if self._is_gqa_asymmetric(name, module):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear_like = (
            isinstance(target_mod, nn.Linear) or "conv1d" in cls_name or hasattr(target_mod, "nf")
        )
        if not is_linear_like:
            return False
        if isinstance(target_mod, KVTCQuantizedLoRALinear):
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """Compress linear layers using PCA decorrelation and true integer quantization."""
        total_entropy = 0.0
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        layer_details: dict[str, Any] = {}
        seen_weights: set[int] = set()

        total_orig_params = 0
        total_comp_params = 0

        # Physical byte tracking before compression
        initial_actual_bytes = self._estimate_actual_bytes(self.model)

        # Freeze non-target linear/conv1d layers when target_attention_only is True
        if self.target_attention_only:
            for name, module in self.model.named_modules():
                target_mod = getattr(module, "base_layer", module)
                cls_name = target_mod.__class__.__name__.lower()
                is_linear = isinstance(target_mod, nn.Linear)
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                if (is_linear or is_conv1d) and not self._is_target_module(name, module):
                    for p in module.parameters():
                        p.requires_grad = False
                    self.skipped_layers.append(name)

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            target_mod = getattr(module, "base_layer", module)
            if not hasattr(target_mod, "weight") or target_mod.weight is None:
                continue

            w_id = id(target_mod.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            raw_w = target_mod.weight.data
            orig_params = raw_w.numel()
            if getattr(target_mod, "bias", None) is not None:
                orig_params += target_mod.bias.numel()
            total_orig_params += orig_params

            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            dtype = raw_w.dtype

            W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
            d_out, d_in = W.shape
            min_dim = min(d_out, d_in)

            try:
                # Step 1: PCA decorrelation via SVD in FP32
                U, S, Vt = torch.linalg.svd(W, full_matrices=False)
                total_energy = (S**2).sum()
                if total_energy > 0:
                    cum_energy = torch.cumsum(S**2, dim=0) / total_energy
                    idx = torch.where(cum_energy >= self.pca_energy)[0]
                    rank = int(idx[0].item()) + 1 if len(idx) > 0 else min_dim
                else:
                    rank = min_dim

                effective_min_rank = min(self.min_rank, min_dim)
                rank = max(effective_min_rank, min(rank, min_dim))

                # Reconstruct low-rank approximation
                U_r = U[:, :rank]
                S_r = S[:rank]
                Vt_r = Vt[:rank, :]
                W_lowrank = (U_r * S_r.unsqueeze(0)) @ Vt_r

                # Step 2: Percentile outlier clipping before quantization
                W_clipped = self._percentile_clip(W_lowrank, low=0.001, high=0.999)

                # Step 3: Integer quantization
                int_w, scale, zp = self._quantize_to_int(W_clipped, bits=self.bits)
                W_quant = self._dequantize_from_int(int_w, scale, zp)
                W_quant = torch.nan_to_num(W_quant, nan=0.0)

                # Reconstruction error check against original W
                norm_orig = torch.norm(W, p="fro").item()
                rel_err = (torch.norm(W - W_quant, p="fro") / max(norm_orig, 1e-12)).item()
                self.layer_errors[name] = float(rel_err)

                if not math.isfinite(rel_err) or rel_err > self.max_error:
                    logger.warning(
                        "KVTC reconstruction error %.4f on %s exceeded max_error %.4f%s",
                        rel_err,
                        name,
                        self.max_error,
                        (
                            ". Keeping original layer."
                            if self.fallback_to_original
                            else ". Proceeding with quantized weight."
                        ),
                    )
                    if self.fallback_to_original:
                        self.fallback_layers.append(name)
                        for p in module.parameters():
                            p.requires_grad = False
                        total_comp_params += orig_params
                        continue

                # Step 4: Empirical Shannon entropy calculation on discrete integer symbols
                entropy = SimpleEntropyCoder.compute_entropy(int_w)
                total_entropy += entropy

                new_layer = KVTCQuantizedLoRALinear.from_linear(
                    module,
                    int_weight=int_w,
                    scale=scale,
                    zero_point=zp,
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                    compute_dtype=dtype,
                )

                # Effective compressed parameters count for reporting
                quant_equiv_params = int(orig_params * (self.bits / 16.0))
                lora_params = (
                    (new_layer.lora_A.numel() + new_layer.lora_B.numel())
                    if (new_layer.lora_A is not None and new_layer.lora_B is not None)
                    else 0
                )
                total_comp_params += quant_equiv_params + lora_params

                replace_module(self.model, name, new_layer)
                self.compressed_layers.append(name)
                layer_details[name] = {
                    "rank": rank,
                    "rel_error": rel_err,
                    "entropy": entropy,
                    "shape": tuple(int_w.shape),
                }
            except Exception as e:
                logger.warning(
                    "KVTC compression exception on %s: %s. Keeping original layer.",
                    name,
                    e,
                )
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_params += orig_params
                self.layer_errors[name] = 1.0

        # Physical byte tracking after compression
        final_actual_bytes = self._estimate_actual_bytes(self.model)

        # Real post-compression forward sanity check
        forward_check_success = False
        forward_check_error = None
        try:
            p0 = next(self.model.parameters(), None)
            if p0 is None:
                p0 = next(self.model.buffers(), None)
            dev = p0.device if p0 is not None else torch.device("cpu")
            self.model.eval()

            # Try text classification / sequence classification format: input_ids (1, 8)
            with torch.no_grad():
                test_passed = False
                try:
                    dummy_ids = torch.randint(0, 100, (1, 8), device=dev)
                    _ = self.model(dummy_ids)
                    test_passed = True
                except Exception:
                    pass

                if not test_passed:
                    # Try continuous hidden tensor (1, 8, d_model)
                    dummy_x = torch.randn(1, 8, 64, device=dev)
                    _ = self.model(dummy_x)
                    test_passed = True

            forward_check_success = test_passed
        except Exception as e:
            forward_check_error = str(e)
            logger.error("KVTC post-compression forward pass failed: %s", e)
            if not kwargs.get("ignore_forward_check", False):
                raise RuntimeError(
                    f"KVTC produced a model that fails the forward pass sanity check: {e}"
                ) from e

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        avg_entropy = total_entropy / max(len(self.compressed_layers), 1)
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        param_ratio = (
            round(total_comp_params / max(total_orig_params, 1), 4)
            if total_orig_params > 0
            else 1.0
        )

        bytes_reduction_pct = (
            (initial_actual_bytes - final_actual_bytes) / float(initial_actual_bytes) * 100.0
            if initial_actual_bytes > 0
            else 0.0
        )

        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "bits": self.bits,
                "pca_energy": self.pca_energy,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "target_attention_only": self.target_attention_only,
                "skip_attention_qkv": self.skip_attention_qkv,
                "avg_entropy_bits": round(avg_entropy, 2),
                "entropy_lower_bound_bits": round(avg_entropy, 2),
                "entropy_coding_note": "Theoretical lower bound; actual bitstream not implemented.",
                "avg_reconstruction_error": avg_err,
                "effective_parameter_ratio": param_ratio,
                "initial_actual_bytes": initial_actual_bytes,
                "final_actual_bytes": final_actual_bytes,
                "actual_bytes_reduction_pct": round(bytes_reduction_pct, 2),
                "theoretical_entropy_compression_ratio": round(16.0 / max(avg_entropy, 1e-4), 2),
                "theoretical_vram_saving": f"{(1.0 - self.bits / 16.0) * 100:.1f}% vs FP16",
                "layer_errors": self.layer_errors,
                "layer_details": layer_details,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "forward_check_success": forward_check_success,
                "forward_check_error": forward_check_error,
                "method": "kvtc_pca_quant",
            }
        )
        self.stats = stats
        return self.model, stats

KVTC: KV Cache Transform Coding compressor (Static Weight PCA + Quantization). Compresses linear projection weights using PCA decorrelation, percentile clipping, uniform integer quantization, and entropy coding estimation.

Ancestors

BaseCompressor

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """Compress linear layers using PCA decorrelation and true integer quantization."""
    total_entropy = 0.0
    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    self.layer_errors = {}
    layer_details: dict[str, Any] = {}
    seen_weights: set[int] = set()

    total_orig_params = 0
    total_comp_params = 0

    # Physical byte tracking before compression
    initial_actual_bytes = self._estimate_actual_bytes(self.model)

    # Freeze non-target linear/conv1d layers when target_attention_only is True
    if self.target_attention_only:
        for name, module in self.model.named_modules():
            target_mod = getattr(module, "base_layer", module)
            cls_name = target_mod.__class__.__name__.lower()
            is_linear = isinstance(target_mod, nn.Linear)
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            if (is_linear or is_conv1d) and not self._is_target_module(name, module):
                for p in module.parameters():
                    p.requires_grad = False
                self.skipped_layers.append(name)

    for name, module in list(self.model.named_modules()):
        if not self._is_target_module(name, module):
            continue

        target_mod = getattr(module, "base_layer", module)
        if not hasattr(target_mod, "weight") or target_mod.weight is None:
            continue

        w_id = id(target_mod.weight)
        if w_id in seen_weights:
            self.skipped_layers.append(name)
            continue
        seen_weights.add(w_id)

        raw_w = target_mod.weight.data
        orig_params = raw_w.numel()
        if getattr(target_mod, "bias", None) is not None:
            orig_params += target_mod.bias.numel()
        total_orig_params += orig_params

        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        dtype = raw_w.dtype

        W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
        d_out, d_in = W.shape
        min_dim = min(d_out, d_in)

        try:
            # Step 1: PCA decorrelation via SVD in FP32
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            total_energy = (S**2).sum()
            if total_energy > 0:
                cum_energy = torch.cumsum(S**2, dim=0) / total_energy
                idx = torch.where(cum_energy >= self.pca_energy)[0]
                rank = int(idx[0].item()) + 1 if len(idx) > 0 else min_dim
            else:
                rank = min_dim

            effective_min_rank = min(self.min_rank, min_dim)
            rank = max(effective_min_rank, min(rank, min_dim))

            # Reconstruct low-rank approximation
            U_r = U[:, :rank]
            S_r = S[:rank]
            Vt_r = Vt[:rank, :]
            W_lowrank = (U_r * S_r.unsqueeze(0)) @ Vt_r

            # Step 2: Percentile outlier clipping before quantization
            W_clipped = self._percentile_clip(W_lowrank, low=0.001, high=0.999)

            # Step 3: Integer quantization
            int_w, scale, zp = self._quantize_to_int(W_clipped, bits=self.bits)
            W_quant = self._dequantize_from_int(int_w, scale, zp)
            W_quant = torch.nan_to_num(W_quant, nan=0.0)

            # Reconstruction error check against original W
            norm_orig = torch.norm(W, p="fro").item()
            rel_err = (torch.norm(W - W_quant, p="fro") / max(norm_orig, 1e-12)).item()
            self.layer_errors[name] = float(rel_err)

            if not math.isfinite(rel_err) or rel_err > self.max_error:
                logger.warning(
                    "KVTC reconstruction error %.4f on %s exceeded max_error %.4f%s",
                    rel_err,
                    name,
                    self.max_error,
                    (
                        ". Keeping original layer."
                        if self.fallback_to_original
                        else ". Proceeding with quantized weight."
                    ),
                )
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    total_comp_params += orig_params
                    continue

            # Step 4: Empirical Shannon entropy calculation on discrete integer symbols
            entropy = SimpleEntropyCoder.compute_entropy(int_w)
            total_entropy += entropy

            new_layer = KVTCQuantizedLoRALinear.from_linear(
                module,
                int_weight=int_w,
                scale=scale,
                zero_point=zp,
                lora_r=self.lora_r,
                lora_alpha=self.lora_alpha,
                compute_dtype=dtype,
            )

            # Effective compressed parameters count for reporting
            quant_equiv_params = int(orig_params * (self.bits / 16.0))
            lora_params = (
                (new_layer.lora_A.numel() + new_layer.lora_B.numel())
                if (new_layer.lora_A is not None and new_layer.lora_B is not None)
                else 0
            )
            total_comp_params += quant_equiv_params + lora_params

            replace_module(self.model, name, new_layer)
            self.compressed_layers.append(name)
            layer_details[name] = {
                "rank": rank,
                "rel_error": rel_err,
                "entropy": entropy,
                "shape": tuple(int_w.shape),
            }
        except Exception as e:
            logger.warning(
                "KVTC compression exception on %s: %s. Keeping original layer.",
                name,
                e,
            )
            self.fallback_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False
            total_comp_params += orig_params
            self.layer_errors[name] = 1.0

    # Physical byte tracking after compression
    final_actual_bytes = self._estimate_actual_bytes(self.model)

    # Real post-compression forward sanity check
    forward_check_success = False
    forward_check_error = None
    try:
        p0 = next(self.model.parameters(), None)
        if p0 is None:
            p0 = next(self.model.buffers(), None)
        dev = p0.device if p0 is not None else torch.device("cpu")
        self.model.eval()

        # Try text classification / sequence classification format: input_ids (1, 8)
        with torch.no_grad():
            test_passed = False
            try:
                dummy_ids = torch.randint(0, 100, (1, 8), device=dev)
                _ = self.model(dummy_ids)
                test_passed = True
            except Exception:
                pass

            if not test_passed:
                # Try continuous hidden tensor (1, 8, d_model)
                dummy_x = torch.randn(1, 8, 64, device=dev)
                _ = self.model(dummy_x)
                test_passed = True

        forward_check_success = test_passed
    except Exception as e:
        forward_check_error = str(e)
        logger.error("KVTC post-compression forward pass failed: %s", e)
        if not kwargs.get("ignore_forward_check", False):
            raise RuntimeError(
                f"KVTC produced a model that fails the forward pass sanity check: {e}"
            ) from e

    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    avg_entropy = total_entropy / max(len(self.compressed_layers), 1)
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )
    param_ratio = (
        round(total_comp_params / max(total_orig_params, 1), 4)
        if total_orig_params > 0
        else 1.0
    )

    bytes_reduction_pct = (
        (initial_actual_bytes - final_actual_bytes) / float(initial_actual_bytes) * 100.0
        if initial_actual_bytes > 0
        else 0.0
    )

    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "fallback_layers": self.fallback_layers,
            "skipped_layers": self.skipped_layers,
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "num_skipped": len(self.skipped_layers),
            "bits": self.bits,
            "pca_energy": self.pca_energy,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "target_attention_only": self.target_attention_only,
            "skip_attention_qkv": self.skip_attention_qkv,
            "avg_entropy_bits": round(avg_entropy, 2),
            "entropy_lower_bound_bits": round(avg_entropy, 2),
            "entropy_coding_note": "Theoretical lower bound; actual bitstream not implemented.",
            "avg_reconstruction_error": avg_err,
            "effective_parameter_ratio": param_ratio,
            "initial_actual_bytes": initial_actual_bytes,
            "final_actual_bytes": final_actual_bytes,
            "actual_bytes_reduction_pct": round(bytes_reduction_pct, 2),
            "theoretical_entropy_compression_ratio": round(16.0 / max(avg_entropy, 1e-4), 2),
            "theoretical_vram_saving": f"{(1.0 - self.bits / 16.0) * 100:.1f}% vs FP16",
            "layer_errors": self.layer_errors,
            "layer_details": layer_details,
            "lora_r": self.lora_r,
            "lora_alpha": self.lora_alpha,
            "forward_check_success": forward_check_success,
            "forward_check_error": forward_check_error,
            "method": "kvtc_pca_quant",
        }
    )
    self.stats = stats
    return self.model, stats

Compress linear layers using PCA decorrelation and true integer quantization.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class KVTCQuantizedLoRALinear (in_features: int, out_features: int, int_weight: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, bias: torch.Tensor | None = None, lora_r: int = 0, lora_alpha: float = 1.0, device: torch.device | None = None, compute_dtype: torch.dtype = torch.float32)

Expand source code

class KVTCQuantizedLoRALinear(nn.Module):
    """
    Linear layer containing true integer quantized base weights (int8/int16 buffer)
    with scale and zero_point buffers, dequantized on the fly during forward(),
    plus optional trainable LoRA adapters (lora_A, lora_B).
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        int_weight: torch.Tensor,
        scale: torch.Tensor,
        zero_point: torch.Tensor,
        bias: torch.Tensor | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        device: torch.device | None = None,
        compute_dtype: torch.dtype = torch.float32,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features  # Conv1D compatibility
        self.compute_dtype = compute_dtype
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        target_device = device if device is not None else int_weight.device

        # Store true integer weights as buffer (not float Parameter)
        self.register_buffer("int_weight", int_weight.to(device=target_device))
        self.register_buffer("scale", scale.to(device=target_device, dtype=torch.float32))
        self.register_buffer("zero_point", zero_point.to(device=target_device, dtype=torch.float32))

        if bias is not None:
            self.bias = nn.Parameter(
                bias.to(device=target_device, dtype=compute_dtype),
                requires_grad=False,
            )
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty(
                    (self.lora_r, in_features),
                    dtype=compute_dtype,
                    device=target_device,
                )
            )
            self.lora_B = nn.Parameter(
                torch.zeros(
                    (out_features, self.lora_r),
                    dtype=compute_dtype,
                    device=target_device,
                )
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    @property
    def weight(self) -> torch.Tensor:
        """Backward-compatibility property returning on-the-fly dequantized weight."""
        return self.int_weight.to(dtype=self.compute_dtype) * self.scale + self.zero_point

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        int_weight: torch.Tensor,
        scale: torch.Tensor,
        zero_point: torch.Tensor,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        compute_dtype: torch.dtype | None = None,
    ) -> KVTCQuantizedLoRALinear:
        target_mod = getattr(linear, "base_layer", linear)
        bias = (
            target_mod.bias.detach()
            if hasattr(target_mod, "bias") and target_mod.bias is not None
            else None
        )
        in_features = getattr(target_mod, "in_features", int_weight.shape[1])
        out_features = getattr(target_mod, "out_features", int_weight.shape[0])

        existing_lora_A = getattr(linear, "lora_A", None)
        existing_lora_B = getattr(linear, "lora_B", None)
        effective_lora_r = lora_r
        effective_lora_alpha = lora_alpha
        if existing_lora_A is not None and existing_lora_B is not None:
            w_A = getattr(existing_lora_A, "weight", existing_lora_A)
            effective_lora_r = w_A.shape[0]
            if hasattr(linear, "scaling") and linear.scaling:
                effective_lora_alpha = linear.scaling * effective_lora_r
            elif hasattr(linear, "lora_alpha"):
                effective_lora_alpha = linear.lora_alpha

        if compute_dtype is None:
            if hasattr(target_mod, "weight") and target_mod.weight is not None:
                compute_dtype = target_mod.weight.dtype
            else:
                compute_dtype = torch.float32

        new_layer = cls(
            in_features=in_features,
            out_features=out_features,
            int_weight=int_weight,
            scale=scale,
            zero_point=zero_point,
            bias=bias,
            lora_r=effective_lora_r,
            lora_alpha=effective_lora_alpha,
            device=int_weight.device,
            compute_dtype=compute_dtype,
        )

        if existing_lora_A is not None and existing_lora_B is not None:
            w_A = getattr(existing_lora_A, "weight", existing_lora_A)
            w_B = getattr(existing_lora_B, "weight", existing_lora_B)
            if (
                hasattr(w_A, "data")
                and hasattr(w_B, "data")
                and new_layer.lora_A is not None
                and new_layer.lora_B is not None
            ):
                new_layer.lora_A.data.copy_(
                    w_A.data.to(dtype=new_layer.lora_A.dtype, device=new_layer.lora_A.device)
                )
                new_layer.lora_B.data.copy_(
                    w_B.data.to(dtype=new_layer.lora_B.dtype, device=new_layer.lora_B.device)
                )
                new_layer.lora_A.requires_grad = True
                new_layer.lora_B.requires_grad = True

        return new_layer

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w_eff = self.int_weight.to(dtype=x.dtype) * self.scale.to(
            dtype=x.dtype
        ) + self.zero_point.to(dtype=x.dtype)
        bias = self.bias.to(dtype=x.dtype) if self.bias is not None else None
        out = F.linear(x, w_eff, bias)
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(
                F.linear(x.to(dtype=self.lora_A.dtype), self.lora_A), self.lora_B
            ).to(dtype=out.dtype)
            out = out + self.scaling * lora_out
        return out

Linear layer containing true integer quantized base weights (int8/int16 buffer) with scale and zero_point buffers, dequantized on the fly during forward(), plus optional trainable LoRA adapters (lora_A, lora_B).

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Module, int_weight: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, lora_r: int = 0, lora_alpha: float = 1.0, compute_dtype: torch.dtype | None = None) ‑> KVTCQuantizedLoRALinear

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Backward-compatibility property returning on-the-fly dequantized weight."""
    return self.int_weight.to(dtype=self.compute_dtype) * self.scale + self.zero_point

Backward-compatibility property returning on-the-fly dequantized weight.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    w_eff = self.int_weight.to(dtype=x.dtype) * self.scale.to(
        dtype=x.dtype
    ) + self.zero_point.to(dtype=x.dtype)
    bias = self.bias.to(dtype=x.dtype) if self.bias is not None else None
    out = F.linear(x, w_eff, bias)
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(
            F.linear(x.to(dtype=self.lora_A.dtype), self.lora_A), self.lora_B
        ).to(dtype=out.dtype)
        out = out + self.scaling * lora_out
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class KVTCStore (bits: int = 4, pca_energy: float = 0.95)

Expand source code

class KVTCStore:
    is_custom_attention: bool = True
    """
    Runtime KV Cache Transform Coding engine.
    Applies PCA decorrelation across head feature dimensions and adaptive quantization,
    supporting dynamic sequence lengths across prefill and decoding phases.
    """

    def __init__(self, bits: int = 4, pca_energy: float = 0.95) -> None:
        self.bits = int(bits)
        self.pca_energy = float(pca_energy)
        self.pca_basis_k: torch.Tensor | None = None
        self.pca_basis_v: torch.Tensor | None = None

    def reset(self) -> None:
        """Reset PCA basis to prevent cross-sequence memory leaks."""
        self.pca_basis_k = None
        self.pca_basis_v = None

    @torch.no_grad()
    def calibrate_basis(self, k: torch.Tensor, v: torch.Tensor) -> None:
        """Calibrate PCA projection basis from sample K/V tensor (B, H, N, D)."""
        k_flat = k.detach().float().reshape(-1, k.shape[-1])
        v_flat = v.detach().float().reshape(-1, v.shape[-1])

        # SVD on centered channels in FP32
        _, _, Vt_k = torch.linalg.svd(
            k_flat - k_flat.mean(dim=0, keepdim=True), full_matrices=False
        )
        _, _, Vt_v = torch.linalg.svd(
            v_flat - v_flat.mean(dim=0, keepdim=True), full_matrices=False
        )

        self.pca_basis_k = Vt_k.t().contiguous().to(device=k.device, dtype=k.dtype)
        self.pca_basis_v = Vt_v.t().contiguous().to(device=v.device, dtype=v.dtype)

    @torch.no_grad()
    def compress_kv_cache(
        self, k: torch.Tensor, v: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
        """
        Transform and quantize dynamic KV cache tensors.
        Returns reconstructed (k_hat, v_hat) and coding statistics.
        """
        orig_shape = k.shape
        D = orig_shape[-1]

        if self.pca_basis_k is None or self.pca_basis_v is None:
            self.calibrate_basis(k, v)

        k_flat = k.reshape(-1, D)
        v_flat = v.reshape(-1, D)

        k_trans = torch.matmul(k_flat, self.pca_basis_k)
        v_trans = torch.matmul(v_flat, self.pca_basis_v)

        qmax = (2 ** (self.bits - 1)) - 1
        scale_k = torch.clamp(k_trans.abs().max() / float(qmax), min=1e-6)
        scale_v = torch.clamp(v_trans.abs().max() / float(qmax), min=1e-6)

        q_k = torch.clamp(torch.round(k_trans / scale_k), -qmax, qmax)
        q_v = torch.clamp(torch.round(v_trans / scale_v), -qmax, qmax)

        h_k = SimpleEntropyCoder.compute_entropy(q_k)
        h_v = SimpleEntropyCoder.compute_entropy(q_v)

        k_hat = torch.matmul(q_k * scale_k, self.pca_basis_k.t().contiguous()).reshape(orig_shape)
        v_hat = torch.matmul(q_v * scale_v, self.pca_basis_v.t().contiguous()).reshape(orig_shape)

        mean_h = (h_k + h_v) / 2.0
        stats: dict[str, Any] = {
            "entropy_k_bits": h_k,
            "entropy_v_bits": h_v,
            "mean_entropy_bits": mean_h,
            "entropy_lower_bound_bits": mean_h,
            "note": "Theoretical lower bound; actual bitstream not implemented.",
            "theoretical_compression_ratio": 16.0 / max(mean_h, 1e-4),
        }
        return k_hat, v_hat, stats

Class variables

var is_custom_attention : bool: Runtime KV Cache Transform Coding engine. Applies PCA decorrelation across head feature dimensions and adaptive quantization, supporting dynamic sequence lengths across prefill and decoding phases.

Methods

def calibrate_basis(self, k: torch.Tensor, v: torch.Tensor) ‑> None

Expand source code

@torch.no_grad()
def calibrate_basis(self, k: torch.Tensor, v: torch.Tensor) -> None:
    """Calibrate PCA projection basis from sample K/V tensor (B, H, N, D)."""
    k_flat = k.detach().float().reshape(-1, k.shape[-1])
    v_flat = v.detach().float().reshape(-1, v.shape[-1])

    # SVD on centered channels in FP32
    _, _, Vt_k = torch.linalg.svd(
        k_flat - k_flat.mean(dim=0, keepdim=True), full_matrices=False
    )
    _, _, Vt_v = torch.linalg.svd(
        v_flat - v_flat.mean(dim=0, keepdim=True), full_matrices=False
    )

    self.pca_basis_k = Vt_k.t().contiguous().to(device=k.device, dtype=k.dtype)
    self.pca_basis_v = Vt_v.t().contiguous().to(device=v.device, dtype=v.dtype)

Calibrate PCA projection basis from sample K/V tensor (B, H, N, D).

def compress_kv_cache(self, k: torch.Tensor, v: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor, dict[str, typing.Any]]

Expand source code

@torch.no_grad()
def compress_kv_cache(
    self, k: torch.Tensor, v: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
    """
    Transform and quantize dynamic KV cache tensors.
    Returns reconstructed (k_hat, v_hat) and coding statistics.
    """
    orig_shape = k.shape
    D = orig_shape[-1]

    if self.pca_basis_k is None or self.pca_basis_v is None:
        self.calibrate_basis(k, v)

    k_flat = k.reshape(-1, D)
    v_flat = v.reshape(-1, D)

    k_trans = torch.matmul(k_flat, self.pca_basis_k)
    v_trans = torch.matmul(v_flat, self.pca_basis_v)

    qmax = (2 ** (self.bits - 1)) - 1
    scale_k = torch.clamp(k_trans.abs().max() / float(qmax), min=1e-6)
    scale_v = torch.clamp(v_trans.abs().max() / float(qmax), min=1e-6)

    q_k = torch.clamp(torch.round(k_trans / scale_k), -qmax, qmax)
    q_v = torch.clamp(torch.round(v_trans / scale_v), -qmax, qmax)

    h_k = SimpleEntropyCoder.compute_entropy(q_k)
    h_v = SimpleEntropyCoder.compute_entropy(q_v)

    k_hat = torch.matmul(q_k * scale_k, self.pca_basis_k.t().contiguous()).reshape(orig_shape)
    v_hat = torch.matmul(q_v * scale_v, self.pca_basis_v.t().contiguous()).reshape(orig_shape)

    mean_h = (h_k + h_v) / 2.0
    stats: dict[str, Any] = {
        "entropy_k_bits": h_k,
        "entropy_v_bits": h_v,
        "mean_entropy_bits": mean_h,
        "entropy_lower_bound_bits": mean_h,
        "note": "Theoretical lower bound; actual bitstream not implemented.",
        "theoretical_compression_ratio": 16.0 / max(mean_h, 1e-4),
    }
    return k_hat, v_hat, stats

Transform and quantize dynamic KV cache tensors. Returns reconstructed (k_hat, v_hat) and coding statistics.

def reset(self) ‑> None

Expand source code

def reset(self) -> None:
    """Reset PCA basis to prevent cross-sequence memory leaks."""
    self.pca_basis_k = None
    self.pca_basis_v = None

Reset PCA basis to prevent cross-sequence memory leaks.

class KroneckerCompressor (model: nn.Module | None = None, rank: int | None = None, max_error: float = 0.8, fallback_to_original: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, target_attention_only: bool = False, skip_attention_qkv: bool = False, skip_should_compress_filter: bool = False, min_layer_dim: int = 8, cfg: Any | None = None, target_layers: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("kronecker_class")
@COMPRESSOR_REGISTRY.register("kronecker_compressor_class")
class KroneckerCompressor(BaseCompressor):
    # KC-FIX-01
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module | None = None,
        rank: int | None = None,
        max_error: float = 0.80,
        fallback_to_original: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        target_attention_only: bool = False,
        skip_attention_qkv: bool = False,
        skip_should_compress_filter: bool = False,
        min_layer_dim: int = 8,
        cfg: Any | None = None,
        target_layers: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(
            model if model is not None else nn.Identity(),
            name="kronecker",
            **kwargs,
        )
        conf = cfg or kwargs.get("comp_config")

        rank = _cfg_get(conf, "kronecker_rank", _cfg_get(conf, "rank", rank))
        max_error = _cfg_get(conf, "kronecker_max_error", _cfg_get(conf, "max_error", max_error))
        fallback_to_original = _cfg_get(
            conf,
            "kronecker_fallback_to_original",
            _cfg_get(conf, "fallback_to_original", fallback_to_original),
        )
        lora_r = _cfg_get(conf, "lora_r", _cfg_get(conf, "lora_rank", lora_r))
        lora_alpha = _cfg_get(conf, "lora_alpha", lora_alpha)
        target_attention_only = _cfg_get(
            conf,
            "kronecker_target_attention_only",
            _cfg_get(conf, "target_attention_only", target_attention_only),
        )
        skip_attention_qkv = _cfg_get(conf, "skip_attention_qkv", skip_attention_qkv)
        skip_should_compress_filter = _cfg_get(
            conf,
            "skip_should_compress_filter",
            skip_should_compress_filter,
        )
        min_layer_dim = _cfg_get(conf, "min_layer_dim", min_layer_dim)
        target_layers = _cfg_get(conf, "target_layers", target_layers)

        self.rank = int(rank) if rank is not None and rank > 0 else None
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.target_attention_only = bool(target_attention_only)
        self.skip_attention_qkv = bool(skip_attention_qkv)
        self.skip_should_compress_filter = bool(skip_should_compress_filter)
        self.min_layer_dim = max(1, int(min_layer_dim))
        self.target_layers = target_layers
        self.cfg = conf

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int | None] = {}
        self.layer_modes: dict[str, str] = {}
        self.stats: dict[str, Any] = {}
        self._compressed_params: int = 0

    # ------------------------------------------------------------------
    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            if any(
                p in name.lower()
                for p in (
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                )
            ):
                return True
        return False

    # KC-FIX-12: debug-enriched target check.
    def _is_target_module(self, name: str, module: nn.Module, _debug: bool = False) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ("lora_A", "lora_B", "base_layer")):
            if _debug:
                logger.info(f"  [skip] {name}: lora/base_layer part")
            return False
        if any(
            skip in name_lower
            for skip in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            )
        ):
            if _debug:
                logger.info(f"  [skip] {name}: skip keyword")
            return False

        if self.target_attention_only:
            if not any(k in name_lower for k in _ATTN_NAMES):
                if _debug:
                    logger.info(f"  [skip] {name}: not attention")
                return False

        if self._is_gqa_asymmetric(name, module):
            if _debug:
                logger.info(f"  [skip] {name}: GQA asymmetric")
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_lower = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_lower or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d):
            if _debug:
                logger.info(f"  [skip] {name}: not linear-like")
            return False
        if isinstance(target_mod, KroneckerLinear):
            if _debug:
                logger.info(f"  [skip] {name}: already KroneckerLinear")
            return False

        w = getattr(target_mod, "weight", None)
        if w is None:
            if _debug:
                logger.info(f"  [skip] {name}: weight None")
            return False

        # KC-FIX-11: filter tiny layers.
        if min(w.shape[0], w.shape[1]) < self.min_layer_dim:
            if _debug:
                logger.info(f"  [skip] {name}: dim too small {tuple(w.shape)}")
            return False

        # KC-FIX-14: guard against framework errors.
        if not self.skip_should_compress_filter:
            try:
                if not should_compress_linear(self.model, name):
                    if _debug:
                        logger.info(f"  [skip] {name}: should_compress_linear=False")
                    return False
            except Exception as e:
                if _debug:
                    logger.info(f"  [skip] {name}: should_compress_linear raised {e}")
                return False

        # KC-FIX-13: empty target_layers means "no targets".
        if self.target_layers is not None:
            ok = any(t in name for t in self.target_layers)
            if not ok and _debug:
                logger.info(f"  [skip] {name}: not in target_layers")
            return ok

        if _debug:
            logger.info(f"  [TARGET] {name}: shape={tuple(w.shape)}")
        return True

    # ------------------------------------------------------------------
    def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
        # KC-FIX-09: idempotence.
        if getattr(self.model, "_kronecker_compressed", False):
            logger.warning("KroneckerCompressor: model already compressed; skipping.")
            try:
                return self.model, self.get_compression_stats()
            except Exception:
                return self.model, {}

        if model is not None:
            self.model = model
            try:
                self._original_params = self._count_params(self.model)
            except Exception:
                self._original_params = sum(p.numel() for p in self.model.parameters())

        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}
        self.layer_modes = {}
        seen_weights: set[int] = set()

        total_orig_params = 0
        total_comp_params = 0

        # Freeze non-targets when attention_only.
        if self.target_attention_only:
            for name, module in self.model.named_modules():
                target_mod = getattr(module, "base_layer", module)
                cls_lower = target_mod.__class__.__name__.lower()
                if (
                    isinstance(target_mod, nn.Linear)
                    or "conv1d" in cls_lower
                    or hasattr(target_mod, "nf")
                ):
                    if not self._is_target_module(name, module):
                        for p in module.parameters():
                            p.requires_grad = False
                        self.skipped_layers.append(name)

        targets = [
            (name, mod)
            for name, mod in list(self.model.named_modules())
            if self._is_target_module(name, mod)
        ]

        if not targets:
            # KC-FIX-12: diagnostic breakdown.
            logger.warning("KroneckerCompressor: no targets found. Diagnostic...")
            n_scanned = 0
            for name, module in self.model.named_modules():
                n_scanned += 1
                self._is_target_module(name, module, _debug=True)
            logger.warning(
                "KroneckerCompressor diagnostic: %d modules scanned, 0 targets.",
                n_scanned,
            )
            try:
                return self.model, self.get_compression_stats()
            except Exception:
                return self.model, {}

        logger.info("KroneckerCompressor: %d target layers identified.", len(targets))

        for name, module in targets:
            target_mod = getattr(module, "base_layer", module)
            if not hasattr(target_mod, "weight") or target_mod.weight is None:
                continue

            w_id = id(target_mod.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            cls_lower = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_lower or hasattr(target_mod, "nf")
            raw_w = target_mod.weight.data

            orig_layer_params = raw_w.numel()
            bias_mod = getattr(target_mod, "bias", None)
            if bias_mod is not None:
                orig_layer_params += bias_mod.numel()
            total_orig_params += orig_layer_params

            W_target = raw_w.t().float() if is_conv1d else raw_w.float()
            W_norm = torch.norm(W_target, p="fro").item()

            new_layer: KroneckerLinear | None = None
            err = 1.0
            try:
                new_layer = KroneckerLinear.from_linear(
                    module,
                    rank=self.rank,
                    max_error=self.max_error,
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                )
                W_comp = new_layer.get_weight().float()
                err = float(torch.norm(W_target - W_comp, p="fro").item() / max(W_norm, 1e-12))
            except Exception as e:
                logger.warning("KroneckerLinear factorization failed on %s: %s", name, e)
                new_layer = None
                err = 1.0

            self.layer_errors[name] = err

            # KC-FIX-15: only set mode/rank after decision.
            if new_layer is None or err > self.max_error:
                logger.warning(
                    "Kronecker: error %.4f on %s > %.4f. %s",
                    err,
                    name,
                    self.max_error,
                    ("Falling back." if self.fallback_to_original else "Keeping original weight."),
                )
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    self.layer_ranks[name] = None
                    self.layer_modes[name] = "fallback"
                    for p in module.parameters():
                        p.requires_grad = False
                    # KC-FIX-06: use orig_layer_params (no double-count).
                    total_comp_params += orig_layer_params
                    continue
                else:
                    self.fallback_layers.append(name)
                    self.layer_ranks[name] = None
                    self.layer_modes[name] = "fallback"
                    total_comp_params += orig_layer_params
                    continue

            # Accepted.
            mode = "kronecker" if getattr(new_layer, "can_decompose", False) else "svd"
            self.layer_modes[name] = mode
            self.layer_ranks[name] = getattr(new_layer, "rank", None)

            comp_layer_params = sum(p.numel() for p in new_layer.parameters())
            total_comp_params += comp_layer_params

            try:
                replace_module(self.model, name, new_layer)
                self.compressed_layers.append(name)
            except Exception as e:
                logger.warning("Kronecker: replace_module failed for %s: %s", name, e)
                self.fallback_layers.append(name)

        # KC-FIX-08: robust forward check.
        forward_check_success = True
        try:
            with torch.no_grad():
                dev = _get_device(self.model)
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                try:
                    self.model(input_ids=dummy, attention_mask=torch.ones_like(dummy))
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception as e2:
                        logger.warning("Kronecker forward check failed: %s", e2)
                        forward_check_success = False
        except Exception as e:
            logger.warning("Kronecker forward check raised: %s", e)
            forward_check_success = False

        # KC-FIX-03: safe param count.
        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        # KC-FIX-05: avg_err only over compressed layers.
        compressed_errs = [
            self.layer_errors[n] for n in self.compressed_layers if n in self.layer_errors
        ]
        avg_err = sum(compressed_errs) / len(compressed_errs) if compressed_errs else 0.0
        param_ratio = (
            total_comp_params / max(total_orig_params, 1) if total_orig_params > 0 else 1.0
        )

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "layer_modes": self.layer_modes,
                "rank": self.rank,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "target_attention_only": self.target_attention_only,
                "forward_check_success": forward_check_success,
                "avg_reconstruction_error": avg_err,
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layer_count": len(self.fallback_layers),
                "effective_parameter_ratio": round(param_ratio, 4),
                "method": "kronecker",
            }
        )
        self.stats = stats

        # KC-FIX-09: write stats onto model.
        try:
            if not hasattr(self.model, "compression_stats") or self.model.compression_stats is None:
                self.model.compression_stats = {}
            self.model.compression_stats.update(
                {
                    "num_compressed": len(self.compressed_layers),
                    "num_fallback": len(self.fallback_layers),
                    "compression_method": "kronecker",
                    "effective_parameter_ratio": param_ratio,
                    "avg_reconstruction_error": avg_err,
                }
            )
        except Exception:
            pass

        self.model._kronecker_compressed = True
        logger.info(
            "Kronecker: %d compressed, %d fallback. param_ratio=%.4f, avg_err=%.4f.",
            len(self.compressed_layers),
            len(self.fallback_layers),
            param_ratio,
            avg_err,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

Expand source code

class KroneckerLinear(nn.Module):
    # KL-FIX-01: this is a linear layer, not an attention layer.
    is_custom_attention: bool = False
    is_custom_linear: bool = True

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        rank: int | None = None,
        m1: int | None = None,
        n1: int | None = None,
        m2: int | None = None,
        n2: int | None = None,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> None:
        super().__init__()
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.nf = self.out_features  # Conv1D compatibility
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        # ---- Factor dimensions ----
        if m1 is not None and m2 is not None and n1 is not None and n2 is not None:
            if m1 * m2 != self.out_features or n1 * n2 != self.in_features:
                raise ValueError(
                    f"Factor dims mismatch: {m1}*{m2}!={self.out_features} "
                    f"or {n1}*{n2}!={self.in_features}"
                )
            self.m1, self.m2 = int(m1), int(m2)
            self.n1, self.n2 = int(n1), int(n2)
        else:
            self.n1, self.n2 = factorize_dim(self.in_features)
            self.m1, self.m2 = factorize_dim(self.out_features)

        # Decomposable when both factorizations are non-trivial OR dims are tiny.
        self.can_decompose = (
            (self.n1 * self.n2 == self.in_features)
            and (self.m1 * self.m2 == self.out_features)
            and (self.n1 > 1 or self.in_features <= 3)
            and (self.m1 > 1 or self.out_features <= 3)
        )

        # ---- Parameter allocation ----
        requires_grad = not (self.lora_r > 0)
        self._init_base_parameters(
            bias=bias,
            dtype=dtype,
            device=device,
            requires_grad=requires_grad,
            rank=rank,
        )

        # ---- LoRA parameters ----
        if self.lora_r > 0:
            # KL-FIX-12: keep LoRA in the same dtype as base parameters.
            param_dtype = (
                next(self.parameters()).dtype
                if any(True for _ in self.parameters())
                else torch.float32
            )
            param_device = (
                next(self.parameters()).device
                if any(True for _ in self.parameters())
                else torch.device("cpu")
            )
            self.lora_A = nn.Parameter(
                torch.empty(
                    (self.lora_r, self.in_features),
                    dtype=param_dtype,
                    device=param_device,
                )
            )
            self.lora_B = nn.Parameter(
                torch.zeros(
                    (self.out_features, self.lora_r),
                    dtype=param_dtype,
                    device=param_device,
                )
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    # ------------------------------------------------------------------
    # KL-FIX-09: clean parameter initialization
    # ------------------------------------------------------------------
    def _init_base_parameters(
        self,
        bias: bool,
        dtype: torch.dtype | None,
        device: torch.device | None,
        requires_grad: bool,
        rank: int | None,
    ) -> None:
        if self.can_decompose:
            # Kronecker factors A (m1, n1), B (m2, n2)
            self.A = nn.Parameter(
                torch.empty(self.m1, self.n1, dtype=dtype, device=device),
                requires_grad=requires_grad,
            )
            self.B = nn.Parameter(
                torch.empty(self.m2, self.n2, dtype=dtype, device=device),
                requires_grad=requires_grad,
            )
            fan_in = max(self.in_features, 1)
            bound = 1.0 / math.sqrt(fan_in)
            nn.init.uniform_(self.A, -bound, bound)
            nn.init.uniform_(self.B, -bound, bound)

            # Ensure U / V don't exist.
            self.register_parameter("U", None)
            self.register_parameter("V", None)

            if bias:
                self.bias = nn.Parameter(
                    torch.empty(self.out_features, dtype=dtype, device=device),
                    requires_grad=requires_grad,
                )
                nn.init.uniform_(self.bias, -bound, bound)
            else:
                self.register_parameter("bias", None)
        else:
            effective_rank = rank or max(1, min(self.in_features, self.out_features) // 4)
            self.rank = min(effective_rank, self.in_features, self.out_features)
            self.U = nn.Parameter(
                torch.empty(self.out_features, self.rank, dtype=dtype, device=device),
                requires_grad=requires_grad,
            )
            self.V = nn.Parameter(
                torch.empty(self.rank, self.in_features, dtype=dtype, device=device),
                requires_grad=requires_grad,
            )
            nn.init.xavier_uniform_(self.U)
            nn.init.xavier_uniform_(self.V)

            # Ensure A / B don't exist.
            self.register_parameter("A", None)
            self.register_parameter("B", None)

            if bias:
                self.bias = nn.Parameter(
                    torch.zeros(self.out_features, dtype=dtype, device=device),
                    requires_grad=requires_grad,
                )
            else:
                self.register_parameter("bias", None)

    # ------------------------------------------------------------------
    # KL-FIX-02 / KL-FIX-07: cache lifecycle
    # ------------------------------------------------------------------
    def _load_from_state_dict(self, *args, **kwargs):
        # KL-FIX-07: state_dict load invalidates weight cache.
        self._cached_weight = None
        super()._load_from_state_dict(*args, **kwargs)

    def _apply(self, fn, *args, **kwargs):
        # _apply is called for .to(), .cuda(), .float(), .half(), etc.
        # Invalidate cache whenever parameters move.
        self._cached_weight = None
        return super()._apply(fn, *args, **kwargs)

    def _invalidate_cache(self) -> None:
        self._cached_weight = None

    # ------------------------------------------------------------------
    # Weight accessors (KL-FIX-10: no stale caching)
    # ------------------------------------------------------------------
    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix (out_features, in_features)."""
        return self.get_weight()

    def get_weight(self) -> torch.Tensor:
        if self.can_decompose and self.A is not None and self.B is not None:
            return torch.kron(self.A, self.B)
        if self.U is not None and self.V is not None:
            return torch.matmul(self.U, self.V)
        # Fallback: should never happen for a properly initialized layer.
        raise RuntimeError("KroneckerLinear has no valid parameter set.")

    # ------------------------------------------------------------------
    # Forward
    # ------------------------------------------------------------------
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.can_decompose and self.A is not None and self.B is not None:
            out = self._forward_kron(x)
        else:
            out = self._forward_svd(x)

        # KL-FIX-12: LoRA in native dtype (no explicit fp32 cast).
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B)
            out = out + self.scaling * lora_out.to(out.dtype)
        return out

    def _forward_kron(self, x: torch.Tensor) -> torch.Tensor:
        """(A ⊗ B) x computed without materializing the full weight."""
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)  # (K, n1 * n2)
        K = x_flat.shape[0]

        # Step 1: contract with B over n2  ->  (K * n1, m2)
        x_b = torch.matmul(
            x_flat.reshape(K * self.n1, self.n2),
            self.B.t().contiguous(),
        )
        # Step 2: (K, m2, n1) -> (K * m2, n1) @ A^T  ->  (K * m2, m1)
        x_b = (
            x_b.reshape(K, self.n1, self.m2)
            .transpose(1, 2)
            .contiguous()
            .reshape(K * self.m2, self.n1)
        )
        x_a = torch.matmul(x_b, self.A.t().contiguous())
        # Step 3: back to (K, m1 * m2)
        out = (
            x_a.reshape(K, self.m2, self.m1)
            .transpose(1, 2)
            .contiguous()
            .reshape(K, self.out_features)
        )

        if self.bias is not None:
            out = out + self.bias
        return out.reshape(*orig_shape[:-1], self.out_features)

    def _forward_svd(self, x: torch.Tensor) -> torch.Tensor:
        return F.linear(F.linear(x, self.V), self.U, self.bias)

    # ------------------------------------------------------------------
    # SVD fitting (KL-FIX-03, KL-FIX-05, KL-FIX-09, KL-FIX-11)
    # ------------------------------------------------------------------
    def _fit_svd(
        self,
        weight: torch.Tensor,
        bias: torch.Tensor | None = None,
        rank: int | None = None,
        max_error: float = 0.20,
    ) -> None:
        """Fit low-rank SVD factors ensuring parameter reduction."""
        in_f, out_f = self.in_features, self.out_features
        target_dtype = weight.dtype
        target_device = weight.device
        W_norm = torch.norm(weight, p="fro").item()

        # Parameter reduction constraint: r * (in + out) < in * out
        max_allowed_rank = max(
            1,
            min(in_f, out_f, int(0.75 * (in_f * out_f) / max(in_f + out_f, 1))),
        )
        base_r = rank if rank is not None and rank > 0 else max(1, min(in_f, out_f) // 4)
        base_r = min(base_r, max_allowed_rank)

        try:
            U, S, Vt = torch.linalg.svd(weight.float(), full_matrices=False)
            candidate_ranks = [base_r]
            if base_r * 2 <= max_allowed_rank:
                candidate_ranks.append(base_r * 2)
            if base_r * 4 <= max_allowed_rank:
                candidate_ranks.append(base_r * 4)

            best_r = base_r
            best_U_r = None
            best_V_r = None
            for cand_r in candidate_ranks:
                U_r = U[:, :cand_r] * torch.sqrt(S[:cand_r].unsqueeze(0))
                V_r = torch.sqrt(S[:cand_r].unsqueeze(1)) * Vt[:cand_r, :]
                W_cand = torch.matmul(U_r, V_r)
                err = float(
                    torch.norm(weight.float() - W_cand, p="fro").item() / max(W_norm, 1e-12)
                )
                best_r = cand_r
                best_U_r = U_r
                best_V_r = V_r
                if err <= max_error:
                    break
        except Exception:
            best_r = base_r
            best_U_r = torch.zeros(out_f, best_r, device=target_device, dtype=torch.float32)
            best_V_r = torch.zeros(best_r, in_f, device=target_device, dtype=torch.float32)

        self.can_decompose = False
        # KL-FIX-05: clean removal of Kronecker parameters.
        self.register_parameter("A", None)
        self.register_parameter("B", None)

        self.rank = best_r
        # KL-FIX-09: clean registration of SVD parameters.
        self.U = nn.Parameter(best_U_r.to(dtype=target_dtype, device=target_device))
        self.V = nn.Parameter(best_V_r.to(dtype=target_dtype, device=target_device))

        # KL-FIX-03 + KL-FIX-11: unconditional, dtype-matched bias.
        if bias is not None:
            self.bias = nn.Parameter(bias.clone().to(dtype=target_dtype, device=target_device))
        else:
            self.register_parameter("bias", None)

        # KL-FIX-10: invalidate cache.
        self._invalidate_cache()

    # ------------------------------------------------------------------
    # Factory (KL-FIX-04, KL-FIX-08)
    # ------------------------------------------------------------------
    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        rank: int | None = None,
        max_error: float = 0.20,
        m1: int | None = None,
        n1: int | None = None,
        m2: int | None = None,
        n2: int | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> KroneckerLinear:
        """Build a KroneckerLinear from a pretrained Linear/Conv1D/LoRA wrapper."""
        target_mod = getattr(linear, "base_layer", linear)
        cls_lower = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_lower or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.data

        if is_conv1d:
            in_features = raw_w.shape[0]
            out_features = raw_w.shape[1]
            weight = raw_w.t().contiguous().float()
        else:
            out_features = raw_w.shape[0]
            in_features = raw_w.shape[1]
            weight = raw_w.float()

        bias_holder = getattr(target_mod, "bias", None)
        bias = bias_holder.data.clone() if bias_holder is not None else None
        target_dtype = raw_w.dtype
        target_device = raw_w.device

        # KL-FIX-04: robust LoRA detection.
        existing_lora_A = getattr(linear, "lora_A", None)
        existing_lora_B = getattr(linear, "lora_B", None)
        effective_lora_r = lora_r
        effective_lora_alpha = lora_alpha
        if existing_lora_A is not None and existing_lora_B is not None:
            w_A = getattr(existing_lora_A, "weight", existing_lora_A)
            if isinstance(w_A, torch.Tensor) and w_A.dim() >= 1:
                effective_lora_r = int(w_A.shape[0])
            scaling = getattr(linear, "scaling", None)
            if isinstance(scaling, (int, float)) and scaling > 0:
                effective_lora_alpha = float(scaling) * effective_lora_r
            elif hasattr(linear, "lora_alpha"):
                effective_lora_alpha = float(linear.lora_alpha)

        kron_layer = cls(
            in_features=in_features,
            out_features=out_features,
            bias=bias is not None,
            rank=rank,
            m1=m1,
            n1=n1,
            m2=m2,
            n2=n2,
            dtype=target_dtype,
            device=target_device,
            lora_r=effective_lora_r,
            lora_alpha=effective_lora_alpha,
        )

        # Transfer existing LoRA if shapes match.
        if (
            existing_lora_A is not None
            and existing_lora_B is not None
            and kron_layer.lora_A is not None
            and kron_layer.lora_B is not None
        ):
            w_A = getattr(existing_lora_A, "weight", existing_lora_A)
            w_B = getattr(existing_lora_B, "weight", existing_lora_B)
            if isinstance(w_A, torch.Tensor) and w_A.shape == kron_layer.lora_A.shape:
                kron_layer.lora_A.data.copy_(
                    w_A.to(kron_layer.lora_A.device, dtype=kron_layer.lora_A.dtype)
                )
            if isinstance(w_B, torch.Tensor) and w_B.shape == kron_layer.lora_B.shape:
                kron_layer.lora_B.data.copy_(
                    w_B.to(kron_layer.lora_B.device, dtype=kron_layer.lora_B.dtype)
                )

        W_norm = torch.norm(weight, p="fro").item()
        err = 1.0

        if kron_layer.can_decompose:
            m1_, m2_ = kron_layer.m1, kron_layer.m2
            n1_, n2_ = kron_layer.n1, kron_layer.n2
            try:
                # Rearrangement for Nearest Kronecker Product.
                W_tensor = weight.view(m1_, m2_, n1_, n2_).permute(0, 2, 1, 3).contiguous()
                W_reordered = W_tensor.view(m1_ * n1_, m2_ * n2_).float()

                U, S, Vt = torch.linalg.svd(W_reordered, full_matrices=False)
                sigma1 = S[0].item() if S.numel() > 0 else 1.0
                scale = math.sqrt(max(sigma1, 1e-12))

                u1 = U[:, 0] * scale
                v1 = Vt[0, :] * scale

                kron_layer.A.data.copy_(
                    u1.view(m1_, n1_).to(dtype=target_dtype, device=target_device)
                )
                kron_layer.B.data.copy_(
                    v1.view(m2_, n2_).to(dtype=target_dtype, device=target_device)
                )

                W_recon = torch.kron(kron_layer.A, kron_layer.B).float()
                err = float(torch.norm(weight - W_recon, p="fro").item() / max(W_norm, 1e-12))
            except Exception as e:
                import logging

                logging.getLogger(__name__).debug(
                    "KroneckerLinear NKP failed: %s; falling back to SVD.", e
                )
                err = 1.0

            if err > max_error:
                kron_layer._fit_svd(weight, bias, rank=rank, max_error=max_error)
        else:
            kron_layer._fit_svd(weight, bias, rank=rank, max_error=max_error)

        if bias is not None and kron_layer.bias is not None:
            kron_layer.bias.data.copy_(
                bias.to(dtype=kron_layer.bias.dtype, device=kron_layer.bias.device)
            )

        # KL-FIX-08: forward sanity check.
        try:
            with torch.no_grad():
                dummy = torch.randn(2, in_features, device=target_device, dtype=target_dtype)
                _ = kron_layer(dummy)
        except Exception as e:
            import logging

            logging.getLogger(__name__).warning(
                "KroneckerLinear forward sanity check failed: %s", e
            )

        return kron_layer

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.
var is_custom_linear : bool: The type of the None singleton.

Static methods

def from_linear(linear: nn.Module, rank: int | None = None, max_error: float = 0.2, m1: int | None = None, n1: int | None = None, m2: int | None = None, n2: int | None = None, lora_r: int = 0, lora_alpha: float = 1.0) ‑> neuralops.compression_engine.kronecker_compressor.KroneckerLinear: Build a KroneckerLinear from a pretrained Linear/Conv1D/LoRA wrapper.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix (out_features, in_features)."""
    return self.get_weight()

Reconstructed effective weight matrix (out_features, in_features).

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    if self.can_decompose and self.A is not None and self.B is not None:
        out = self._forward_kron(x)
    else:
        out = self._forward_svd(x)

    # KL-FIX-12: LoRA in native dtype (no explicit fp32 cast).
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B)
        out = out + self.scaling * lora_out.to(out.dtype)
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    if self.can_decompose and self.A is not None and self.B is not None:
        return torch.kron(self.A, self.B)
    if self.U is not None and self.V is not None:
        return torch.matmul(self.U, self.V)
    # Fallback: should never happen for a properly initialized layer.
    raise RuntimeError("KroneckerLinear has no valid parameter set.")

class KroneckerTTLoRACompressor (model: torch.nn.modules.module.Module | None = None, rank: int = 8, gamma: float = 1.0, lora_r: int = 4, lora_alpha: float = 8.0, max_error: float = 0.35, target_attention_only: bool = False, fallback_to_original: bool = True, skip_attention_qkv: bool = False, skip_should_compress_filter: bool = False, min_layer_dim: int = 8, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("kronecker_tt_lora_class")
@COMPRESSOR_REGISTRY.register("kronecker_tt_lora_compressor_class")
class KroneckerTTLoRACompressor(BaseCompressor):
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module | None = None,
        rank: int = 8,
        gamma: float = 1.0,
        lora_r: int = 4,
        lora_alpha: float = 8.0,
        max_error: float = 0.35,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        skip_attention_qkv: bool = False,
        skip_should_compress_filter: bool = False,
        min_layer_dim: int = 8,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(
            model if model is not None else nn.Identity(),
            name="kronecker_tt_lora",
            **kwargs,
        )
        self.rank = int(_cfg_get(cfg, "tt_rank", _cfg_get(cfg, "rank", rank)))
        self.gamma = float(_cfg_get(cfg, "gamma", gamma))
        self.lora_r = int(_cfg_get(cfg, "lora_r", lora_r))
        self.lora_alpha = float(_cfg_get(cfg, "lora_alpha", lora_alpha))
        self.max_error = float(
            _cfg_get(
                cfg,
                "kronecker_tt_lora_max_error",
                _cfg_get(cfg, "max_error", max_error),
            )
        )
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "kronecker_tt_lora_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "kronecker_tt_lora_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )
        self.skip_attention_qkv = bool(_cfg_get(cfg, "skip_attention_qkv", skip_attention_qkv))
        self.skip_should_compress_filter = bool(
            _cfg_get(
                cfg,
                "skip_should_compress_filter",
                skip_should_compress_filter,
            )
        )
        self.min_layer_dim = max(1, int(_cfg_get(cfg, "min_layer_dim", min_layer_dim)))
        self.cfg = cfg

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.stats: dict[str, Any] = {}
        self._compressed_params: int = 0

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            if any(
                p in name.lower()
                for p in (
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                )
            ):
                return True
        return False

    # KTL-FIX-09: debug-enriched target check.
    def _is_target_module(self, name: str, module: nn.Module, _debug: bool = False) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ("lora_A", "lora_B", "base_layer")):
            if _debug:
                logger.info(f"  [skip] {name}: lora/base_layer")
            return False
        if any(
            skip in name_lower
            for skip in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            )
        ):
            if _debug:
                logger.info(f"  [skip] {name}: skip keyword")
            return False

        if self.target_attention_only:
            if not any(k in name_lower for k in _ATTN_NAMES):
                if _debug:
                    logger.info(f"  [skip] {name}: not attention")
                return False

        if self._is_gqa_asymmetric(name, module):
            if _debug:
                logger.info(f"  [skip] {name}: GQA asymmetric")
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_lower = target_mod.__class__.__name__.lower()
        is_linear_like = (
            isinstance(target_mod, nn.Linear) or "conv1d" in cls_lower or hasattr(target_mod, "nf")
        )
        if not is_linear_like:
            if _debug:
                logger.info(f"  [skip] {name}: not linear-like")
            return False
        if isinstance(target_mod, KroneckerTTLoRALayer):
            if _debug:
                logger.info(f"  [skip] {name}: already KroneckerTTLoRALayer")
            return False

        # KTL-FIX-12
        if not self.skip_should_compress_filter:
            try:
                if not should_compress_linear(self.model, name):
                    if _debug:
                        logger.info(f"  [skip] {name}: should_compress_linear=False")
                    return False
            except Exception as e:
                if _debug:
                    logger.info(f"  [skip] {name}: should_compress_linear raised {e}")
                return False

        w = getattr(target_mod, "weight", None)
        if w is None:
            if _debug:
                logger.info(f"  [skip] {name}: weight None")
            return False
        if min(w.shape[0], w.shape[1]) < self.min_layer_dim:
            if _debug:
                logger.info(f"  [skip] {name}: dim too small {tuple(w.shape)}")
            return False

        if _debug:
            logger.info(f"  [TARGET] {name}: shape={tuple(w.shape)}")
        return True

    # ------------------------------------------------------------------
    def compress_linear(self, linear_layer: nn.Module) -> KroneckerTTLoRALayer:
        r = max(1, min(self.rank, 32))
        r_inner = max(1, r // 2)
        return KroneckerTTLoRALayer.from_linear(
            base_layer=linear_layer,
            tt_ranks=(1, r_inner, r_inner, 1),
            gamma=self.gamma,
            lora_r=self.lora_r,
            lora_alpha=self.lora_alpha,
        )

    # ------------------------------------------------------------------
    def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
        # KTL-FIX-08: idempotence.
        if getattr(self.model, "_kronecker_tt_lora_compressed", False):
            logger.warning("KroneckerTTLoRA: model already compressed; skipping.")
            try:
                return self.model, self.get_compression_stats()
            except Exception:
                return self.model, {}

        if model is not None:
            self.model = model
            try:
                self._original_params = self._count_params(self.model)
            except Exception:
                self._original_params = sum(p.numel() for p in self.model.parameters())

        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        seen_weights: set[int] = set()

        total_orig_params = 0
        total_comp_params = 0

        target_list = [
            (name, module)
            for name, module in self.model.named_modules()
            if self._is_target_module(name, module)
        ]
        if not target_list:
            logger.warning("KroneckerTTLoRA: no target layers found. Diagnostic...")
            n = 0
            for name, module in self.model.named_modules():
                n += 1
                self._is_target_module(name, module, _debug=True)
            logger.warning("KroneckerTTLoRA diagnostic: %d modules scanned, 0 targets.", n)
            try:
                return self.model, self.get_compression_stats()
            except Exception:
                return self.model, {}

        logger.info("KroneckerTTLoRA: %d target layers identified.", len(target_list))

        # Freeze non-targets when attention_only.
        if self.target_attention_only:
            for name, module in self.model.named_modules():
                target_mod = getattr(module, "base_layer", module)
                if (
                    isinstance(target_mod, nn.Linear)
                    or "conv1d" in target_mod.__class__.__name__.lower()
                    or hasattr(target_mod, "nf")
                ):
                    if not self._is_target_module(name, module):
                        for p in module.parameters():
                            p.requires_grad = False
                        self.skipped_layers.append(name)

        for name, module in target_list:
            target_mod = getattr(module, "base_layer", module)
            if not hasattr(target_mod, "weight") or target_mod.weight is None:
                continue

            w_id = id(target_mod.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            raw_w = target_mod.weight.data
            orig_params = raw_w.numel()
            if getattr(target_mod, "bias", None) is not None:
                orig_params += target_mod.bias.numel()
            total_orig_params += orig_params

            try:
                adapted = self.compress_linear(module)
            except Exception as e:
                logger.warning("KroneckerTTLoRA: compress_linear failed for %s: %s", name, e)
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_comp_params += orig_params
                continue

            # Error check
            cls_lower = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_lower or hasattr(target_mod, "nf")
            W_orig = raw_w.t().contiguous().float() if is_conv1d else raw_w.float()
            delta_w = adapted.get_delta_w().float()
            # Note: delta_w is the adapter delta, not a reconstruction of the base weight.
            # Compute the *relative adapter magnitude* as a stability metric.
            base_norm = torch.norm(W_orig, p="fro").item()
            adapter_norm = torch.norm(delta_w, p="fro").item()
            err = adapter_norm / max(base_norm, 1e-12)
            self.layer_errors[name] = float(err)

            if err > self.max_error:
                logger.warning(
                    "KroneckerTTLoRA: adapter delta norm ratio %.4f on %s > %.4f.",
                    err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    total_comp_params += orig_params
                    continue

            # Sum adapter params only (trainable).
            adapter_params = sum(
                p.numel()
                for p in (
                    adapted.A1,
                    adapted.A2,
                    adapted.G1,
                    adapted.G2,
                    adapted.G3,
                    adapted.lora_A,
                    adapted.lora_B,
                )
                if p is not None
            )
            # KTL-FIX-13: total stored params = base + adapter.
            total_comp_params += orig_params + adapter_params

            try:
                replace_module(self.model, name, adapted)
                self.compressed_layers.append(name)
            except Exception as e:
                logger.warning("KroneckerTTLoRA: replace_module failed for %s: %s", name, e)
                self.fallback_layers.append(name)

        # KTL-FIX-02: actual forward sanity check.
        forward_check_success = True
        try:
            with torch.no_grad():
                first_p = next(iter(self.model.parameters()), None)
                dev = first_p.device if first_p is not None else torch.device("cpu")
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                try:
                    self.model(input_ids=dummy, attention_mask=torch.ones_like(dummy))
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception as e2:
                        logger.warning("KroneckerTTLoRA: forward check failed: %s", e2)
                        forward_check_success = False
        except Exception as e:
            logger.warning("KroneckerTTLoRA: forward check raised: %s", e)
            forward_check_success = False

        # KTL-FIX-05: safe param count.
        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        # KTL-FIX-07: honest stats.
        param_ratio = (
            total_comp_params / max(total_orig_params, 1) if total_orig_params > 0 else 1.0
        )
        avg_err = (
            sum(self.layer_errors.values()) / max(1, len(self.layer_errors))
            if self.layer_errors
            else 0.0
        )

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "tt_rank": self.rank,
                "gamma": self.gamma,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "max_error": self.max_error,
                "target_attention_only": self.target_attention_only,
                "avg_adapter_delta_ratio": avg_err,
                "effective_parameter_ratio": round(param_ratio, 4),
                "forward_check_success": forward_check_success,
                "method": "kronecker_tt_lora",
                # KTL-FIX-07 + KTL-FIX-13: explicit note.
                "note": (
                    "KroneckerTTLoRA is a PEFT-style adapter. The original weight is "
                    "kept frozen and a low-rank Kronecker+TT update is added. "
                    "Total stored parameters increase by adapter size. "
                    "`effective_parameter_ratio` reflects (base + adapter) / base. "
                    "Trainable parameters are only the adapter (A1, A2, G*, lora_*)."
                ),
            }
        )
        self.stats = stats

        try:
            if not hasattr(self.model, "compression_stats") or self.model.compression_stats is None:
                self.model.compression_stats = {}
            self.model.compression_stats.update(
                {
                    "num_compressed": len(self.compressed_layers),
                    "num_fallback": len(self.fallback_layers),
                    "compression_method": "kronecker_tt_lora",
                    "effective_parameter_ratio": param_ratio,
                }
            )
        except Exception:
            pass

        self.model._kronecker_tt_lora_compressed = True
        logger.info(
            "KroneckerTTLoRA: %d compressed, %d fallback. avg_adapter_delta=%.4f, "
            "param_ratio=%.4f.",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            param_ratio,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def compress_linear(self, linear_layer: torch.nn.modules.module.Module) ‑> KroneckerTTLoRALayer

Expand source code

def compress_linear(self, linear_layer: nn.Module) -> KroneckerTTLoRALayer:
    r = max(1, min(self.rank, 32))
    r_inner = max(1, r // 2)
    return KroneckerTTLoRALayer.from_linear(
        base_layer=linear_layer,
        tt_ranks=(1, r_inner, r_inner, 1),
        gamma=self.gamma,
        lora_r=self.lora_r,
        lora_alpha=self.lora_alpha,
    )

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class KroneckerTTLoRALayer (base_layer: torch.nn.modules.linear.Linear, m1: int, n1: int, m2: int, n2: int, tt_shapes: tuple[int, int, int], tt_ranks: tuple[int, int, int, int] = (1, 4, 4, 1), gamma: float = 1.0, lora_r: int = 4, lora_alpha: float = 8.0, r_min: int = 2, r_max: int = 32, entropy_threshold: float = 0.75)

Expand source code

class KroneckerTTLoRALayer(nn.Module):
    # KTL-FIX-01: this is a linear layer, not attention.
    is_custom_attention: bool = False
    is_custom_linear: bool = True

    def __init__(
        self,
        base_layer: nn.Linear,
        m1: int,
        n1: int,
        m2: int,
        n2: int,
        tt_shapes: tuple[int, int, int],
        tt_ranks: tuple[int, int, int, int] = (1, 4, 4, 1),
        gamma: float = 1.0,
        lora_r: int = 4,
        lora_alpha: float = 8.0,
        r_min: int = 2,
        r_max: int = 32,
        entropy_threshold: float = 0.75,
    ) -> None:
        super().__init__()
        self.base_layer = base_layer
        self.base_layer.weight.requires_grad = False
        if getattr(base_layer, "bias", None) is not None:
            base_layer.bias.requires_grad = False

        self.in_features = base_layer.in_features
        self.out_features = base_layer.out_features

        self.padded_out_features = m1 * m2
        self.padded_in_features = n1 * n2
        # KTL-FIX-14: informative errors.
        if self.padded_out_features < self.out_features:
            raise ValueError(
                f"m1*m2 ({self.padded_out_features}) must be >= out_features ({self.out_features})"
            )
        if self.padded_in_features < self.in_features:
            raise ValueError(
                f"n1*n2 ({self.padded_in_features}) must be >= in_features ({self.in_features})"
            )

        tt_prod = tt_shapes[0] * tt_shapes[1] * tt_shapes[2]
        expected_prod = self.padded_out_features * self.padded_in_features
        if tt_prod != expected_prod:
            raise ValueError(
                f"tt_shapes product ({tt_prod}) must equal "
                f"padded_out * padded_in ({expected_prod}); "
                f"got tt_shapes={tt_shapes}, padded_out={self.padded_out_features}, "
                f"padded_in={self.padded_in_features}"
            )
        if tt_ranks[0] != 1 or tt_ranks[3] != 1:
            raise ValueError(f"Boundary TT ranks must be 1; got {tt_ranks}")

        self.m1, self.n1 = m1, n1
        self.m2, self.n2 = m2, n2
        self.tt_shapes = tuple(tt_shapes)
        self.tt_ranks = list(tt_ranks)
        self.gamma = float(gamma)
        self.lora_r = max(1, int(lora_r))
        self.lora_alpha = float(lora_alpha)
        self.lora_scaling = self.lora_alpha / float(self.lora_r)
        self.r_min = int(r_min)
        self.r_max = int(r_max)
        self.entropy_threshold = float(entropy_threshold)

        device = base_layer.weight.device
        dtype = base_layer.weight.dtype

        self.A1 = nn.Parameter(_init_orthogonal_spectral((m1, n1), (m1, n1), device, dtype))
        self.A2 = nn.Parameter(_init_orthogonal_spectral((m2, n2), (m2, n2), device, dtype))

        r1, r2 = self.tt_ranks[1], self.tt_ranks[2]
        d1, d2, d3 = self.tt_shapes

        self.G1 = nn.Parameter(_init_orthogonal_spectral((1, d1, r1), (d1, r1), device, dtype))
        self.G2 = nn.Parameter(
            _init_orthogonal_spectral((r1, d2, r2), (r1 * d2, r2), device, dtype)
        )
        self.G3 = nn.Parameter(_init_orthogonal_spectral((r2, d3, 1), (r2, d3), device, dtype))

        self.lora_A = nn.Parameter(
            torch.zeros((self.lora_r, self.in_features), dtype=dtype, device=device)
        )
        self.lora_B = nn.Parameter(
            torch.zeros((self.out_features, self.lora_r), dtype=dtype, device=device)
        )
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)

    # ------------------------------------------------------------------
    @classmethod
    def factorize_dimensions(cls, in_f: int, out_f: int) -> dict[str, Any]:
        """Auto-factorize dims and pick 3-factor TT shape."""

        def best_factor_pair(dim: int) -> tuple[int, int]:
            for d in range(int(math.isqrt(dim)), 0, -1):
                if dim % d == 0:
                    return d, dim // d
            return 1, dim

        m1, m2 = best_factor_pair(out_f)
        n1, n2 = best_factor_pair(in_f)

        total_padded = (m1 * m2) * (n1 * n2)
        c = int(round(total_padded ** (1.0 / 3.0)))
        best: tuple[int, int, int] | None = None
        best_spread = float("inf")
        for d in range(max(1, c - 64), c + 65):
            if total_padded % d != 0:
                continue
            rem = total_padded // d
            for d2 in range(1, int(math.isqrt(rem)) + 1):
                if rem % d2 == 0:
                    d3 = rem // d2
                    spread = max(d, d2, d3) - min(d, d2, d3)
                    if spread < best_spread:
                        best = (d, d2, d3)
                        best_spread = spread
                    break
        tt_shapes = best if best is not None else (m1, m2, n1 * n2)

        return {"m1": m1, "n1": n1, "m2": m2, "n2": n2, "tt_shapes": tt_shapes}

    @classmethod
    def from_linear(
        cls,
        base_layer: nn.Module,
        m1: int | None = None,
        n1: int | None = None,
        m2: int | None = None,
        n2: int | None = None,
        tt_shapes: tuple[int, int, int] | None = None,
        tt_ranks: tuple[int, int, int, int] = (1, 4, 4, 1),
        gamma: float = 1.0,
        lora_r: int = 4,
        lora_alpha: float = 8.0,
        **kwargs: Any,
    ) -> "KroneckerTTLoRALayer":
        """Build a KroneckerTTLoRALayer from Linear / Conv1D / LoRA wrapper."""
        target_mod = getattr(base_layer, "base_layer", base_layer)
        cls_lower = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_lower or hasattr(target_mod, "nf")

        raw_w = target_mod.weight.data
        bias_holder = getattr(target_mod, "bias", None)
        has_bias = bias_holder is not None

        if is_conv1d:
            in_features = raw_w.shape[0]
            out_features = raw_w.shape[1]
            lin = nn.Linear(
                in_features,
                out_features,
                bias=has_bias,
                device=raw_w.device,
                dtype=raw_w.dtype,
            )
            lin.weight.data.copy_(raw_w.t().contiguous())
            if has_bias:
                lin.bias.data.copy_(bias_holder.data)
            effective_base = lin
        elif isinstance(target_mod, nn.Linear):
            effective_base = target_mod
        else:
            # KTL-FIX-11: fallback for exotic modules with explicit dims.
            in_features = getattr(target_mod, "in_features", raw_w.shape[1])
            out_features = getattr(target_mod, "out_features", raw_w.shape[0])
            lin = nn.Linear(
                in_features,
                out_features,
                bias=has_bias,
                device=raw_w.device,
                dtype=raw_w.dtype,
            )
            lin.weight.data.copy_(raw_w)
            if has_bias:
                lin.bias.data.copy_(bias_holder.data)
            effective_base = lin

        dims = cls.factorize_dimensions(effective_base.in_features, effective_base.out_features)

        # Detect existing LoRA.
        existing_lora_A = getattr(base_layer, "lora_A", None)
        existing_lora_B = getattr(base_layer, "lora_B", None)
        eff_lora_r = lora_r
        eff_lora_alpha = lora_alpha
        if existing_lora_A is not None and existing_lora_B is not None:
            w_A = getattr(existing_lora_A, "weight", existing_lora_A)
            if isinstance(w_A, torch.Tensor) and w_A.dim() >= 1:
                eff_lora_r = int(w_A.shape[0])
            scaling = getattr(base_layer, "scaling", None)
            if isinstance(scaling, (int, float)) and scaling:
                eff_lora_alpha = float(scaling) * eff_lora_r
            elif hasattr(base_layer, "lora_alpha"):
                eff_lora_alpha = float(base_layer.lora_alpha)

        layer = cls(
            base_layer=effective_base,
            m1=m1 or dims["m1"],
            n1=n1 or dims["n1"],
            m2=m2 or dims["m2"],
            n2=n2 or dims["n2"],
            tt_shapes=tt_shapes or dims["tt_shapes"],
            tt_ranks=tt_ranks,
            gamma=gamma,
            lora_r=eff_lora_r,
            lora_alpha=eff_lora_alpha,
            **kwargs,
        )

        if existing_lora_A is not None and existing_lora_B is not None:
            src_A = getattr(existing_lora_A, "weight", existing_lora_A)
            src_B = getattr(existing_lora_B, "weight", existing_lora_B)
            if isinstance(src_A, torch.Tensor) and src_A.shape == layer.lora_A.shape:
                layer.lora_A.data.copy_(src_A.to(layer.lora_A.device, dtype=layer.lora_A.dtype))
            if isinstance(src_B, torch.Tensor) and src_B.shape == layer.lora_B.shape:
                layer.lora_B.data.copy_(src_B.to(layer.lora_B.device, dtype=layer.lora_B.dtype))

        return layer

    # ------------------------------------------------------------------
    # TT and delta_w construction (KTL-FIX-03)
    # ------------------------------------------------------------------
    def get_tt_tensor(self) -> torch.Tensor:
        d1, d2, d3 = self.tt_shapes
        r1, r2 = self.tt_ranks[1], self.tt_ranks[2]
        g1 = self.G1.view(d1, r1)
        g2 = self.G2.view(r1, d2 * r2)
        core12 = torch.matmul(g1, g2).view(d1 * d2, r2)
        g3 = self.G3.view(r2, d3)
        tt = torch.matmul(core12, g3)
        return tt.view(self.padded_out_features, self.padded_in_features)

    def get_delta_w(self) -> torch.Tensor:
        kron = torch.kron(self.A1, self.A2)
        tt_tensor = self.get_tt_tensor()
        delta_padded = self.gamma * (kron * tt_tensor)
        return delta_padded[: self.out_features, : self.in_features]

    @property
    def weight(self) -> torch.Tensor:
        # KTL-FIX-15: consistent with forward (bias handled in forward, not here).
        W = self.base_layer.weight + self.get_delta_w()
        if self.lora_A is not None and self.lora_B is not None:
            W = W + self.lora_scaling * torch.matmul(self.lora_B, self.lora_A)
        return W

    @property
    def bias(self) -> torch.Tensor | None:
        return getattr(self.base_layer, "bias", None)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # KTL-FIX-03: compute delta_w once.
        base_out = F.linear(x, self.base_layer.weight, getattr(self.base_layer, "bias", None))
        delta_w = self.get_delta_w()
        kron_tt_out = F.linear(x, delta_w)
        lora_bypass = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling
        return base_out + kron_tt_out + lora_bypass

    # ------------------------------------------------------------------
    # Dynamic rank adaptation
    # ------------------------------------------------------------------
    def adapt_tt_ranks(self, delta_r: int = 1) -> None:
        with torch.no_grad():
            delta_w = self.get_delta_w().float()
            try:
                _, S, _ = torch.linalg.svd(delta_w, full_matrices=False)
            except Exception:
                return
            total = S.sum().item()
            if total > 1e-9:
                p = S / total
                p = p[p > 1e-9]
                h_sv = float(-torch.sum(p * torch.log(p)).item())
            else:
                h_sv = 0.0

            if h_sv > self.entropy_threshold:
                direction = 1
            elif h_sv < self.entropy_threshold * 0.5:
                direction = -1
            else:
                direction = 0

            if direction == 0:
                return
            new_r1 = max(self.r_min, min(self.r_max, self.tt_ranks[1] + direction * delta_r))
            new_r2 = max(self.r_min, min(self.r_max, self.tt_ranks[2] + direction * delta_r))
            if new_r1 != self.tt_ranks[1] or new_r2 != self.tt_ranks[2]:
                self._resize_cores(new_r1, new_r2)

    def _resize_cores(self, new_r1: int, new_r2: int) -> None:
        # KTL-FIX-06: validate inputs.
        if new_r1 < 1 or new_r2 < 1:
            return
        old_r1, old_r2 = self.tt_ranks[1], self.tt_ranks[2]
        d1, d2, d3 = self.tt_shapes

        new_G1 = _init_orthogonal_spectral(
            (1, d1, new_r1), (d1, new_r1), self.G1.device, self.G1.dtype
        )
        new_G2 = _init_orthogonal_spectral(
            (new_r1, d2, new_r2), (new_r1 * d2, new_r2), self.G2.device, self.G2.dtype
        )
        new_G3 = _init_orthogonal_spectral(
            (new_r2, d3, 1), (new_r2, d3), self.G3.device, self.G3.dtype
        )

        min_r1 = min(old_r1, new_r1)
        min_r2 = min(old_r2, new_r2)
        try:
            new_G1[:, :, :min_r1] = self.G1.data[:, :, :min_r1]
            new_G2[:min_r1, :, :min_r2] = self.G2.data[:min_r1, :, :min_r2]
            new_G3[:min_r2, :, :] = self.G3.data[:min_r2, :, :]
        except Exception as e:
            logger.warning("KTL: _resize_cores copy failed: %s", e)
            return

        self.G1 = nn.Parameter(new_G1)
        self.G2 = nn.Parameter(new_G2)
        self.G3 = nn.Parameter(new_G3)
        self.tt_ranks[1] = new_r1
        self.tt_ranks[2] = new_r2

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.
var is_custom_linear : bool: The type of the None singleton.

Static methods

def factorize_dimensions(in_f: int, out_f: int) ‑> dict[str, typing.Any]: Auto-factorize dims and pick 3-factor TT shape.
def from_linear(base_layer: torch.nn.modules.module.Module, m1: int | None = None, n1: int | None = None, m2: int | None = None, n2: int | None = None, tt_shapes: tuple[int, int, int] | None = None, tt_ranks: tuple[int, int, int, int] = (1, 4, 4, 1), gamma: float = 1.0, lora_r: int = 4, lora_alpha: float = 8.0, **kwargs: Any) ‑> KroneckerTTLoRALayer: Build a KroneckerTTLoRALayer from Linear / Conv1D / LoRA wrapper.

Instance variables

prop bias : torch.Tensor | None

Expand source code

@property
def bias(self) -> torch.Tensor | None:
    return getattr(self.base_layer, "bias", None)

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    # KTL-FIX-15: consistent with forward (bias handled in forward, not here).
    W = self.base_layer.weight + self.get_delta_w()
    if self.lora_A is not None and self.lora_B is not None:
        W = W + self.lora_scaling * torch.matmul(self.lora_B, self.lora_A)
    return W

Methods

def adapt_tt_ranks(self, delta_r: int = 1) ‑> None

Expand source code

def adapt_tt_ranks(self, delta_r: int = 1) -> None:
    with torch.no_grad():
        delta_w = self.get_delta_w().float()
        try:
            _, S, _ = torch.linalg.svd(delta_w, full_matrices=False)
        except Exception:
            return
        total = S.sum().item()
        if total > 1e-9:
            p = S / total
            p = p[p > 1e-9]
            h_sv = float(-torch.sum(p * torch.log(p)).item())
        else:
            h_sv = 0.0

        if h_sv > self.entropy_threshold:
            direction = 1
        elif h_sv < self.entropy_threshold * 0.5:
            direction = -1
        else:
            direction = 0

        if direction == 0:
            return
        new_r1 = max(self.r_min, min(self.r_max, self.tt_ranks[1] + direction * delta_r))
        new_r2 = max(self.r_min, min(self.r_max, self.tt_ranks[2] + direction * delta_r))
        if new_r1 != self.tt_ranks[1] or new_r2 != self.tt_ranks[2]:
            self._resize_cores(new_r1, new_r2)

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    # KTL-FIX-03: compute delta_w once.
    base_out = F.linear(x, self.base_layer.weight, getattr(self.base_layer, "bias", None))
    delta_w = self.get_delta_w()
    kron_tt_out = F.linear(x, delta_w)
    lora_bypass = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling
    return base_out + kron_tt_out + lora_bypass

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def get_delta_w(self) ‑> torch.Tensor

Expand source code

def get_delta_w(self) -> torch.Tensor:
    kron = torch.kron(self.A1, self.A2)
    tt_tensor = self.get_tt_tensor()
    delta_padded = self.gamma * (kron * tt_tensor)
    return delta_padded[: self.out_features, : self.in_features]

def get_tt_tensor(self) ‑> torch.Tensor

Expand source code

def get_tt_tensor(self) -> torch.Tensor:
    d1, d2, d3 = self.tt_shapes
    r1, r2 = self.tt_ranks[1], self.tt_ranks[2]
    g1 = self.G1.view(d1, r1)
    g2 = self.G2.view(r1, d2 * r2)
    core12 = torch.matmul(g1, g2).view(d1 * d2, r2)
    g3 = self.G3.view(r2, d3)
    tt = torch.matmul(core12, g3)
    return tt.view(self.padded_out_features, self.padded_in_features)

class LIFTRCompressor (model: nn.Module, replace_attention: bool = False, gelu_coefficient: float = 0.8, validation_calibration: bool = False, calibration_threshold: float = 0.2, target_attention_only: bool = False, skip_should_compress_filter: bool = False, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("liftr_class")
@COMPRESSOR_REGISTRY.register("liftr_compressor_class")
class LIFTRCompressor(BaseCompressor):
    # LIFTR-FIX-01: this is a compressor, not an attention.
    is_custom_attention: bool = False

    CRITICAL_MODULE_KEYWORDS = [
        "classifier",
        "score",
        "lm_head",
        "pooler",
        "embed",
        "wte",
        "wpe",
        "shared",
        "head",
        "logits",
    ]

    def __init__(
        self,
        model: nn.Module,
        replace_attention: bool = False,
        gelu_coefficient: float = 0.8,
        validation_calibration: bool = False,
        calibration_threshold: float = 0.20,
        target_attention_only: bool = False,
        skip_should_compress_filter: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="liftr", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        self.cfg = cfg

        self.replace_attention = bool(
            _cfg_get(
                cfg,
                "liftr_replace_attention",
                _cfg_get(cfg, "replace_attention", replace_attention),
            )
        )
        self.gelu_coefficient = float(
            _cfg_get(
                cfg,
                "liftr_gelu_coefficient",
                _cfg_get(cfg, "gelu_coefficient", gelu_coefficient),
            )
        )
        self.validation_calibration = bool(
            _cfg_get(
                cfg,
                "validation_calibration",
                validation_calibration,
            )
        )
        self.calibration_threshold = float(
            _cfg_get(
                cfg,
                "calibration_threshold",
                calibration_threshold,
            )
        )
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "liftr_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.skip_should_compress_filter = bool(
            _cfg_get(
                cfg,
                "skip_should_compress_filter",
                skip_should_compress_filter,
            )
        )

        self.linearized_activations_count = 0
        self.linearized_attentions_count = 0
        self.replaced_modules: list[str] = []
        self.skipped_layers: list[str] = []
        self.stats: dict[str, Any] = {}
        self._compressed_params: int = 0

    def _is_critical(self, name: str) -> bool:
        low = name.lower()
        if any(part in name.split(".") for part in ("lora_A", "lora_B")):
            return True
        return any(kw in low for kw in self.CRITICAL_MODULE_KEYWORDS)

    def _is_already_compressed(self, module: nn.Module) -> bool:
        """LIFTR-FIX-06: skip already-compressed module types."""
        cls_name = module.__class__.__name__
        return cls_name in (
            "LSCPrunedLoRALinear",
            "DynamicSVDLinear",
            "SVDLinear",
            "FactorisedLinear",
            "FactorizedLinear",
            "LowRankPlusSparseLinear",
            "_LinearizedAttentionWrapper",
        )

    def _linearize_module(
        self,
        module: nn.Module,
        prefix: str = "",
        _depth: int = 0,
    ) -> None:
        """Recursively replaces activation functions and (optionally) attention blocks."""
        if _depth > 16:
            return

        for name, child in list(module.named_children()):
            full_name = f"{prefix}.{name}" if prefix else name

            if self._is_critical(full_name):
                continue

            if self._is_already_compressed(child):
                if hasattr(child, "lora_A") or child.__class__.__name__.lower() == "loralinear":
                    continue
                if isinstance(child, _LinearizedAttentionWrapper):
                    continue

            # Skip LoRA-wrapped modules.
            if hasattr(child, "lora_A") or child.__class__.__name__.lower() == "loralinear":
                continue

            is_attn_layer = any(k in full_name.lower() for k in _ATTN_NAMES)

            # --- Activation linearization ---
            if isinstance(child, nn.ReLU):
                if self.target_attention_only and not is_attn_layer:
                    self.skipped_layers.append(full_name)
                    continue
                setattr(module, name, nn.Identity())
                self.linearized_activations_count += 1
                self.replaced_modules.append(full_name)
                continue

            if isinstance(child, _SILU_CLASSES):
                if self.target_attention_only and not is_attn_layer:
                    self.skipped_layers.append(full_name)
                    continue
                setattr(module, name, _LinearizedGELU(alpha=self.gelu_coefficient))
                self.linearized_activations_count += 1
                self.replaced_modules.append(full_name)
                continue

            # --- Attention linearization ---
            if self.replace_attention and is_attn_layer:
                if "norm" in name.lower():
                    self._linearize_module(child, prefix=full_name, _depth=_depth + 1)
                    continue
                if isinstance(child, (nn.Linear, nn.LayerNorm, nn.RMSNorm, nn.Embedding)):
                    self._linearize_module(child, prefix=full_name, _depth=_depth + 1)
                    continue

                # BERT / RoBERTa nested structure: linearize inner self-attention.
                if hasattr(child, "self") and hasattr(child, "output"):
                    self._linearize_module(child, prefix=full_name, _depth=_depth + 1)
                    continue

                if isinstance(child, _LinearizedAttentionWrapper):
                    continue

                if not hasattr(child, "_liftr_skip") and not isinstance(
                    child, _LinearizedAttentionWrapper
                ):
                    try:
                        wrapped = _LinearizedAttentionWrapper(child, layer_name=full_name)
                        setattr(module, name, wrapped)
                        self.linearized_attentions_count += 1
                        self.replaced_modules.append(full_name)
                    except Exception as e:
                        logger.warning("LIFTR: failed to wrap %s: %s", full_name, e)
                continue

            # --- Recurse ---
            self._linearize_module(child, prefix=full_name, _depth=_depth + 1)

    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        # LIFTR-FIX-15: idempotence
        if getattr(self.model, "_liftr_compressed", False):
            logger.warning("LIFTR: model already compressed; skipping.")
            try:
                stats = self.get_compression_stats()
            except Exception:
                stats = {}
            return self.model, stats

        # LIFTR-FIX-11: safe device resolution.
        first_param = next(iter(self.model.parameters()), None)
        device = first_param.device if first_param is not None else torch.device("cpu")

        # LIFTR-FIX-03: dummy_in defined at function scope.
        dummy_in: torch.Tensor | None = None
        orig_out: torch.Tensor | None = None

        if self.validation_calibration:
            try:
                self.model.eval()
                dummy_in = torch.randint(0, 100, (1, 8), device=device, dtype=torch.long)
                with torch.no_grad():
                    res = self.model(dummy_in)
                    if hasattr(res, "logits"):
                        orig_out = res.logits.detach()
                    elif isinstance(res, (tuple, list)):
                        orig_out = res[0].detach()
                    else:
                        orig_out = res.detach()
            except Exception as e:
                logger.debug("LIFTR baseline calibration failed: %s", e)
                dummy_in = None
                orig_out = None

        # Apply linearization.
        self._linearize_module(self.model)

        # Freeze parameters in skipped modules.
        if self.skipped_layers:
            skipped_set = set(self.skipped_layers)
            for name, module in self.model.named_modules():
                if name in skipped_set:
                    for p in module.parameters(recurse=False):
                        p.requires_grad = False

        # Post-linearization validation.
        validation_passed = True
        rel_diff = 0.0
        if self.validation_calibration and orig_out is not None and dummy_in is not None:
            try:
                with torch.no_grad():
                    res_new = self.model(dummy_in)
                    if hasattr(res_new, "logits"):
                        new_out = res_new.logits
                    elif isinstance(res_new, (tuple, list)):
                        new_out = res_new[0]
                    else:
                        new_out = res_new
                    rel_diff = (
                        torch.norm(orig_out.float() - new_out.float())
                        / (torch.norm(orig_out.float()) + 1e-8)
                    ).item()
                    if rel_diff > self.calibration_threshold:
                        logger.warning(
                            "LIFTR: validation diff %.4f > threshold %.4f.",
                            rel_diff,
                            self.calibration_threshold,
                        )
                        validation_passed = False
            except Exception as e:
                logger.warning("LIFTR: post-linearization validation failed: %s", e)
                validation_passed = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        stats.update(
            {
                "linearized_activations": self.linearized_activations_count,
                "linearized_attentions": self.linearized_attentions_count,
                "replace_attention": self.replace_attention,
                "replaced_modules": self.replaced_modules,
                "gelu_coefficient": self.gelu_coefficient,
                "target_attention_only": self.target_attention_only,
                "validation_calibration": self.validation_calibration,
                "validation_passed": validation_passed,
                "calibration_rel_diff": rel_diff,
                # LIFTR-FIX-08: explicit note.
                "effective_parameter_ratio": 1.0,
                "flops_reduction_pct": 0.0,
                "method": "liftr_linearization",
                "note": (
                    "LIFTR performs functional linearization (activation & attention). "
                    "It does NOT reduce parameters; original attention submodules are "
                    "retained as `self.original_attn` for safe delegation."
                ),
            }
        )
        self.stats = stats
        self.model._liftr_compressed = True
        logger.info(
            "LIFTR: %d activations linearized, %d attentions wrapped.",
            self.linearized_activations_count,
            self.linearized_attentions_count,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var CRITICAL_MODULE_KEYWORDS: The type of the None singleton.
var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class LSCCompressor (model: nn.Module, rank: int = 32, energy_ratio: float | None = None, max_error: float = 0.85, fallback_to_original: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, skip_attention_qkv: bool = False, target_attention_only: bool = False, target_modules: list[str] | None = None, min_layer_dim: int = 8, skip_should_compress_filter: bool = False, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("lsc_class")
@COMPRESSOR_REGISTRY.register("lsc_compressor_class")
class LSCCompressor(BaseCompressor):
    # LSC-FIX-01: this is a compressor, not an attention layer.
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module,
        rank: int = 32,
        energy_ratio: float | None = None,
        max_error: float = 0.85,
        fallback_to_original: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        skip_attention_qkv: bool = False,
        target_attention_only: bool = False,
        target_modules: list[str] | None = None,
        min_layer_dim: int = 8,
        skip_should_compress_filter: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="lsc", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            rank = _cfg_get(cfg, "lsc_rank", _cfg_get(cfg, "rank", rank))
            energy_ratio = _cfg_get(
                cfg, "lsc_energy_ratio", _cfg_get(cfg, "energy_ratio", energy_ratio)
            )
            max_error = _cfg_get(cfg, "lsc_max_error", _cfg_get(cfg, "max_error", max_error))
            fallback_to_original = _cfg_get(
                cfg,
                "lsc_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
            lora_r = _cfg_get(cfg, "lora_r", lora_r)
            lora_alpha = _cfg_get(cfg, "lora_alpha", lora_alpha)
            skip_attention_qkv = _cfg_get(cfg, "skip_attention_qkv", skip_attention_qkv)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)
            min_layer_dim = _cfg_get(cfg, "min_layer_dim", min_layer_dim)
            skip_should_compress_filter = _cfg_get(
                cfg,
                "skip_should_compress_filter",
                skip_should_compress_filter,
            )

        self.rank = max(1, int(rank))
        self.energy_ratio = float(energy_ratio) if energy_ratio is not None else None
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.skip_attention_qkv = bool(skip_attention_qkv)
        self.target_attention_only = bool(target_attention_only)
        self.target_modules = target_modules
        self.min_layer_dim = max(1, int(min_layer_dim))
        self.skip_should_compress_filter = bool(skip_should_compress_filter)

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.stats: dict[str, Any] = {}
        self._compressed_params: int = 0

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            if any(
                p in name.lower()
                for p in [
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                ]
            ):
                return True
        return False

    # ------------------------------------------------------------------
    # Target check with debug
    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module, _debug: bool = False) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ("lora_A", "lora_B", "base_layer")):
            if _debug:
                logger.info(f"  [skip] {name}: lora/base_layer part")
            return False

        cls_name = module.__class__.__name__
        if cls_name in (
            "LSCPrunedLoRALinear",
            "DynamicSVDLinear",
            "SVDLinear",
            "FactorisedLinear",
        ):
            if _debug:
                logger.info(f"  [skip] {name}: already compressed ({cls_name})")
            return False

        base_mod = getattr(module, "base_layer", module)
        cls_lower = base_mod.__class__.__name__.lower()
        is_linear_like = (
            isinstance(base_mod, nn.Linear) or "conv1d" in cls_lower or hasattr(base_mod, "nf")
        )
        if not is_linear_like:
            if _debug:
                logger.info(f"  [skip] {name}: not linear-like")
            return False

        if not self.skip_should_compress_filter:
            try:
                if not should_compress_linear(self.model, name):
                    if _debug:
                        logger.info(f"  [skip] {name}: should_compress_linear=False")
                    return False
            except Exception as e:
                if _debug:
                    logger.info(f"  [skip] {name}: should_compress_linear raised {e}")
                return False

        w = getattr(base_mod, "weight", None)
        if w is None:
            if _debug:
                logger.info(f"  [skip] {name}: weight None")
            return False
        if min(w.shape[0], w.shape[1]) < self.min_layer_dim:
            if _debug:
                logger.info(f"  [skip] {name}: dim too small {tuple(w.shape)}")
            return False

        if self._is_gqa_asymmetric(name, module):
            if _debug:
                logger.info(f"  [skip] {name}: GQA asymmetric")
            return False

        if self.target_attention_only:
            if not any(k in name_lower for k in _ATTN_NAMES):
                if _debug:
                    logger.info(f"  [skip] {name}: not attention")
                return False

        if self.target_modules is not None:
            ok = any(t in name for t in self.target_modules)
            if not ok and _debug:
                logger.info(f"  [skip] {name}: not in target_modules")
            return ok

        if _debug:
            logger.info(f"  [TARGET] {name}: shape={tuple(w.shape)}")
        return True

    # ------------------------------------------------------------------
    # Main compress
    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        layer_details: dict[str, Any] = {}

        seen_weights: set[int] = set()
        seen_param_ids: set[int] = set()
        total_orig_params = 0
        total_comp_params = 0

        # LSC-FIX-12: deduplicated parameter accounting.
        for p in self.model.parameters():
            if id(p) not in seen_param_ids:
                seen_param_ids.add(id(p))
                total_orig_params += p.numel()

        count_param_ids: set[int] = set()

        def _add_params(mod: nn.Module) -> None:
            nonlocal total_comp_params
            for p in mod.parameters(recurse=False):
                if id(p) not in count_param_ids:
                    count_param_ids.add(id(p))
                    total_comp_params += p.numel()

        # LSC-FIX-06: diagnostic breakdown if no targets
        target_list = [
            (name, module)
            for name, module in self.model.named_modules()
            if self._is_target_module(name, module)
        ]
        if not target_list:
            logger.warning("LSC: no compressible layers found. Running diagnostic...")
            n_scanned = 0
            for name, module in self.model.named_modules():
                n_scanned += 1
                self._is_target_module(name, module, _debug=True)
            logger.warning("LSC diagnostic: %d modules scanned, 0 targets.", n_scanned)
            try:
                stats = self.get_compression_stats()
            except Exception:
                stats = {}
            return self.model, stats

        logger.info("LSC: %d target layers identified.", len(target_list))

        for name, module in target_list:
            target_layer = getattr(module, "base_layer", module)
            if not hasattr(target_layer, "weight") or target_layer.weight is None:
                _add_params(module)
                continue

            w_id = id(target_layer.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                _add_params(module)
                continue
            seen_weights.add(w_id)

            raw_w = target_layer.weight.data
            cls_name = target_layer.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_layer, "nf")
            device = raw_w.device
            dtype = raw_w.dtype

            # FP32 for numerical stability
            W = raw_w.t().clone().float() if is_conv1d else raw_w.clone().float()
            d_out, d_in = W.shape
            total_coeffs = d_out * d_in

            try:
                W_fft = torch.fft.fft2(W)
                magnitude = torch.abs(W_fft)
                flat_mag = magnitude.reshape(-1)

                # LSC-FIX-08: guard against all-zero magnitude
                if flat_mag.max().item() <= 1e-12:
                    logger.warning(
                        "LSC: all-zero magnitude for %s; keeping original.",
                        name,
                    )
                    self.fallback_layers.append(name)
                    for p in module.parameters(recurse=False):
                        p.requires_grad = False
                    _add_params(module)
                    continue

                if self.energy_ratio is not None:
                    sorted_mag, _ = torch.sort(flat_mag, descending=True)
                    total_energy = (sorted_mag**2).sum().clamp(min=1e-12)
                    cum_energy = torch.cumsum(sorted_mag**2, dim=0) / total_energy
                    idx = torch.where(cum_energy >= self.energy_ratio)[0]
                    k = idx[0].item() + 1 if len(idx) > 0 else total_coeffs
                    k = max(1, min(k, total_coeffs))
                else:
                    k = max(1, min(self.rank * self.rank, total_coeffs))

                thr = torch.topk(flat_mag, k=k, largest=True).values[-1]
                spectral_mask = (magnitude >= thr).float()

                W_comp = torch.fft.ifft2(W_fft * spectral_mask).real
                W_comp = torch.nan_to_num(W_comp, nan=0.0, posinf=0.0, neginf=0.0)

                norm_orig = torch.norm(W, p="fro").item()
                rel_err = torch.norm(W - W_comp, p="fro").item() / max(norm_orig, 1e-12)
                self.layer_errors[name] = float(rel_err)

                fam = getattr(getattr(self.model, "config", None), "model_type", None)
                eff_max_err = max(
                    self.max_error,
                    adaptive_max_error(self.rank, d_in, d_out, family=fam),
                )
                if fam in ("bert", "distilbert", "roberta"):
                    eff_max_err = max(eff_max_err, 0.999)

                if not math.isfinite(rel_err) or rel_err > eff_max_err:
                    logger.warning(
                        "LSC: err %.4f on %s > %.4f (base max_err=%.4f). %s",
                        rel_err,
                        name,
                        eff_max_err,
                        self.max_error,
                        (
                            "Keeping original."
                            if self.fallback_to_original
                            else "Proceeding with filtered weights."
                        ),
                    )
                    if self.fallback_to_original:
                        self.fallback_layers.append(name)
                        for p in module.parameters(recurse=False):
                            p.requires_grad = False
                        _add_params(module)
                        continue

                new_w = W_comp.t().contiguous() if is_conv1d else W_comp
                expected_shape = target_layer.weight.shape
                if new_w.shape != expected_shape:
                    logger.warning(
                        "LSC: shape mismatch for %s: expected %s, got %s; skipping.",
                        name,
                        tuple(expected_shape),
                        tuple(new_w.shape),
                    )
                    for p in module.parameters(recurse=False):
                        p.requires_grad = False
                    _add_params(module)
                    continue

                new_layer = LSCPrunedLoRALinear.from_linear(
                    module,
                    weight=new_w.to(dtype=dtype, device=device),
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                )
                replace_module(self.model, name, new_layer)
                self.compressed_layers.append(name)
                _add_params(new_layer)
                layer_details[name] = {
                    "k_coeffs": int(k),
                    "rel_error": float(rel_err),
                    "shape": tuple(new_w.shape),
                }
            except Exception as e:
                logger.warning(
                    "LSC: exception on %s: %s. Keeping original layer.",
                    name,
                    e,
                )
                self.fallback_layers.append(name)
                self.layer_errors[name] = 1.0
                for p in module.parameters(recurse=False):
                    p.requires_grad = False
                _add_params(module)

        # LSC-FIX-04 + LSC-FIX-10: robust forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                first_param = next(iter(self.model.parameters()), None)
                dev = first_param.device if first_param is not None else torch.device("cpu")
                dummy_ids = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                dummy_mask = torch.ones_like(dummy_ids)
                try:
                    self.model(input_ids=dummy_ids, attention_mask=dummy_mask)
                except Exception:
                    try:
                        self.model(dummy_ids)
                    except Exception as e2:
                        logger.warning("LSC: forward check failed: %s", e2)
                        forward_check_success = False
        except Exception as e:
            logger.warning("LSC: forward check raised: %s", e)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        eff_ratio = total_comp_params / max(1, total_orig_params)

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "rank": self.rank,
                "energy_ratio": self.energy_ratio,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "avg_reconstruction_error": avg_err,
                "layer_errors": self.layer_errors,
                "layer_details": layer_details,
                "effective_parameter_ratio": eff_ratio,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "forward_check_success": forward_check_success,
                "method": "lsc_fourier_spectral",
                "note": (
                    "LSC performs in-place spectral filtering of weights. "
                    "The stored weight is still dense (out x in); "
                    "framework-level compression_rate reads 1.0. "
                    "Use effective_parameter_ratio for the true storage metric."
                ),
            }
        )
        self.stats = stats

        # LSC-FIX-13: write stats onto the model for framework visibility.
        try:
            if not hasattr(self.model, "compression_stats") or self.model.compression_stats is None:
                self.model.compression_stats = {}
            self.model.compression_stats.update(
                {
                    "num_compressed": len(self.compressed_layers),
                    "num_fallback": len(self.fallback_layers),
                    "compression_method": "lsc_fourier_spectral",
                    "effective_parameter_ratio": eff_ratio,
                    "avg_reconstruction_error": avg_err,
                }
            )
        except Exception:
            pass

        logger.info(
            "Compressor lsc finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        logger.info(
            "LSC: %d compressed, %d fallback, %d skipped. "
            "Avg err %.4f, effective_param_ratio %.2fx.",
            len(self.compressed_layers),
            len(self.fallback_layers),
            len(self.skipped_layers),
            avg_err,
            eff_ratio,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class LSCPrunedLoRALinear (in_features: int, out_features: int, weight: torch.Tensor, bias: torch.Tensor | None = None, lora_r: int = 0, lora_alpha: float = 1.0, device: torch.device | None = None, dtype: torch.dtype | None = None)

Expand source code

class LSCPrunedLoRALinear(nn.Module):
    # LSC-FIX-07: this is a linear layer, not an attention module.
    is_custom_attention: bool = False
    is_custom_linear: bool = True

    def __init__(
        self,
        in_features: int,
        out_features: int,
        weight: torch.Tensor,
        bias: torch.Tensor | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features  # Conv1D compatibility
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0
        self.rank = self.lora_r  # Alias for detection

        target_device = device if device is not None else weight.device
        target_dtype = dtype if dtype is not None else weight.dtype

        self.weight = nn.Parameter(
            weight.to(device=target_device, dtype=target_dtype),
            requires_grad=False,
        )
        if bias is not None:
            self.bias = nn.Parameter(
                bias.to(device=target_device, dtype=target_dtype),
                requires_grad=False,
            )
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            # LSC-FIX-03: LoRA parameters in base weight's dtype for forward consistency.
            self.lora_A = nn.Parameter(
                torch.empty((self.lora_r, in_features), dtype=target_dtype, device=target_device)
            )
            self.lora_B = nn.Parameter(
                torch.zeros(
                    (out_features, self.lora_r),
                    dtype=target_dtype,
                    device=target_device,
                )
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        weight: torch.Tensor,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> LSCPrunedLoRALinear:
        base_layer = getattr(linear, "base_layer", linear)
        bias = (
            base_layer.bias.detach()
            if hasattr(base_layer, "bias") and base_layer.bias is not None
            else None
        )
        in_features = getattr(base_layer, "in_features", weight.shape[1])
        out_features = getattr(base_layer, "out_features", weight.shape[0])

        existing_lora_r = _detect_lora_rank(linear)
        eff_lora_r = existing_lora_r if existing_lora_r > 0 else lora_r
        eff_lora_alpha = _detect_lora_alpha(linear) if existing_lora_r > 0 else lora_alpha

        layer = cls(
            in_features=in_features,
            out_features=out_features,
            weight=weight,
            bias=bias,
            lora_r=eff_lora_r,
            lora_alpha=eff_lora_alpha,
            device=weight.device,
            dtype=weight.dtype,
        )

        # Preserve existing LoRA if shapes match.
        if (
            hasattr(linear, "lora_A")
            and hasattr(linear, "lora_B")
            and layer.lora_A is not None
            and layer.lora_B is not None
        ):
            with torch.no_grad():
                src_A = getattr(linear.lora_A, "weight", linear.lora_A)
                src_B = getattr(linear.lora_B, "weight", linear.lora_B)
                if isinstance(src_A, torch.Tensor) and src_A.shape == layer.lora_A.shape:
                    layer.lora_A.data.copy_(src_A.to(layer.lora_A.device, dtype=layer.lora_A.dtype))
                if isinstance(src_B, torch.Tensor) and src_B.shape == layer.lora_B.shape:
                    layer.lora_B.data.copy_(src_B.to(layer.lora_B.device, dtype=layer.lora_B.dtype))

        if layer.lora_A is not None:
            layer.lora_A.requires_grad = True
        if layer.lora_B is not None:
            layer.lora_B.requires_grad = True

        return layer

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = F.linear(x, self.weight, self.bias)
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            # Preserve dtype consistency: do the LoRA computation in the
            # base weight's dtype (input dtype may differ).
            lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B)
            out = out + self.scaling * lora_out.to(out.dtype)
        return out

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.
var is_custom_linear : bool: The type of the None singleton.

Static methods

def from_linear(linear: nn.Module, weight: torch.Tensor, lora_r: int = 0, lora_alpha: float = 1.0) ‑> LSCPrunedLoRALinear

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    out = F.linear(x, self.weight, self.bias)
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        # Preserve dtype consistency: do the LoRA computation in the
        # base weight's dtype (input dtype may differ).
        lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B)
        out = out + self.scaling * lora_out.to(out.dtype)
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class LiLLaMACompressor (model: nn.Module, tokenizer: Any | None = None, rank: int = 32, calibration_texts: list[str] | None = None, calibration_loader: Any | None = None, device: str | torch.device | None = None, max_error: float = 0.2, fallback_to_original: bool = True, skip_attention: bool = False, target_attention_only: bool = False, target_modules: list[str] | None = None, min_layer_dim: int = 8, skip_should_compress_filter: bool = False, max_calibration_batches: int = 8, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("lillama_class")
@COMPRESSOR_REGISTRY.register("lillama_compressor_class")
@COMPRESSOR_REGISTRY.register("lillama_distiller_class")
class LiLLaMADistiller(BaseCompressor):
    # LL-FIX-01: this is a local distillation / weight-rewriter, not standard factorized attention.
    is_custom_attention: bool = False
    category: str = "weight_rewriting"

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any | None = None,
        rank: int = 32,
        calibration_texts: list[str] | None = None,
        calibration_loader: Any | None = None,
        device: str | torch.device | None = None,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        skip_attention: bool = False,
        target_attention_only: bool = False,
        target_modules: list[str] | None = None,
        min_layer_dim: int = 8,
        skip_should_compress_filter: bool = False,
        max_calibration_batches: int = 8,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="lillama", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            rank = _cfg_get(cfg, "lillama_rank", _cfg_get(cfg, "rank", rank))
            max_error = _cfg_get(cfg, "max_error", max_error)
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            skip_attention = _cfg_get(cfg, "skip_attention", skip_attention)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)
            min_layer_dim = _cfg_get(cfg, "min_layer_dim", min_layer_dim)
            skip_should_compress_filter = _cfg_get(
                cfg,
                "skip_should_compress_filter",
                skip_should_compress_filter,
            )
            max_calibration_batches = _cfg_get(
                cfg, "max_calibration_batches", max_calibration_batches
            )

        # LL-FIX-11: tokenizer resolution.
        self.tokenizer = (
            tokenizer
            or getattr(model, "tokenizer", None)
            or getattr(getattr(model, "config", None), "tokenizer", None)
        )
        self.rank = max(1, int(rank))
        self.calibration_texts = calibration_texts
        self.calibration_loader = calibration_loader
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.skip_attention = bool(skip_attention)
        self.target_attention_only = bool(target_attention_only)
        self.target_modules = target_modules
        self.min_layer_dim = max(1, int(min_layer_dim))
        self.skip_should_compress_filter = bool(skip_should_compress_filter)
        self.max_calibration_batches = max(1, int(max_calibration_batches))

        # LL-FIX-03: safe device.
        if device is not None:
            self.target_device = torch.device(device)
        else:
            self.target_device = _get_device(model)

        # LL-FIX-15: initialize all lists in __init__.
        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_details: dict[str, Any] = {}
        self.stats: dict[str, Any] = {}
        self._compressed_params: int = 0

    # ------------------------------------------------------------------
    # Target check (LL-FIX-06)
    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module, _debug: bool = False) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ("lora_A", "lora_B", "base_layer")):
            if _debug:
                logger.info(f"  [skip] {name}: lora/base_layer part")
            return False

        # Skip already-compressed.
        cls_name = module.__class__.__name__
        if cls_name in (
            "LSCPrunedLoRALinear",
            "DynamicSVDLinear",
            "SVDLinear",
            "FactorisedLinear",
            "LowRankPlusSparseLinear",
        ):
            if _debug:
                logger.info(f"  [skip] {name}: already compressed ({cls_name})")
            return False

        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            ]
        ):
            if _debug:
                logger.info(f"  [skip] {name}: skip keyword")
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_lower = target_mod.__class__.__name__.lower()
        is_linear_like = (
            isinstance(target_mod, nn.Linear) or "conv1d" in cls_lower or hasattr(target_mod, "nf")
        )
        if not is_linear_like:
            if _debug:
                logger.info(f"  [skip] {name}: not linear-like")
            return False

        # LL-FIX-09: allow bypassing framework filter.
        if not self.skip_should_compress_filter:
            try:
                if not should_compress_linear(self.model, name):
                    if _debug:
                        logger.info(f"  [skip] {name}: should_compress_linear=False")
                    return False
            except Exception as e:
                if _debug:
                    logger.info(f"  [skip] {name}: should_compress_linear raised {e}")
                return False

        # LL-FIX-10: filter tiny layers.
        w = getattr(target_mod, "weight", None)
        if w is None:
            if _debug:
                logger.info(f"  [skip] {name}: weight None")
            return False
        if min(w.shape[0], w.shape[1]) < self.min_layer_dim:
            if _debug:
                logger.info(f"  [skip] {name}: dim too small {tuple(w.shape)}")
            return False

        if self.target_attention_only:
            is_attn = any(k in name_lower for k in _ATTN_NAMES)
            if not is_attn:
                if _debug:
                    logger.info(f"  [skip] {name}: not attention")
                return False

        if self.skip_attention:
            attn_keywords = {
                "query",
                "key",
                "value",
                "q_lin",
                "k_lin",
                "v_lin",
                "q_proj",
                "k_proj",
                "v_proj",
                "out_lin",
                "out_proj",
                "dense",
                "o_proj",
                "c_attn",
            }
            last = name.split(".")[-1].lower()
            if last in attn_keywords or any(k in name_lower for k in ("self_attn", "attn.c_attn")):
                if _debug:
                    logger.info(f"  [skip] {name}: skip_attention")
                return False

        if self.target_modules is not None:
            ok = any(t in name for t in self.target_modules)
            if not ok and _debug:
                logger.info(f"  [skip] {name}: not in target_modules")
            return ok

        if _debug:
            logger.info(f"  [TARGET] {name}: shape={tuple(w.shape)}")
        return True

    # ------------------------------------------------------------------
    # IO activation collection
    # ------------------------------------------------------------------
    def _collect_io_activations(
        self,
    ) -> dict[str, tuple[torch.Tensor, torch.Tensor]]:
        """Collect input/output activations across target layers."""
        self.model.eval()
        activations: dict[str, dict[str, list[torch.Tensor]]] = {}
        hooks = []

        def make_hook(layer_name: str):
            def hook(module: nn.Module, inputs: tuple[Any, ...], output: Any):
                inp = inputs[0].detach() if inputs and isinstance(inputs[0], torch.Tensor) else None
                if isinstance(output, torch.Tensor):
                    out = output.detach()
                elif (
                    isinstance(output, (tuple, list))
                    and output
                    and isinstance(output[0], torch.Tensor)
                ):
                    out = output[0].detach()
                else:
                    out = None
                if inp is not None and out is not None:
                    if layer_name not in activations:
                        activations[layer_name] = {"in": [], "out": []}
                    activations[layer_name]["in"].append(inp.cpu())
                    activations[layer_name]["out"].append(out.cpu())

            return hook

        target_list = [
            (name, module)
            for name, module in self.model.named_modules()
            if self._is_target_module(name, module)
        ]
        if not target_list:
            return {}

        for name, module in target_list:
            hooks.append(module.register_forward_hook(make_hook(name)))

        try:
            calib_inputs = get_calibration_tensors(
                self.model,
                tokenizer=self.tokenizer,
                calibration_loader=self.calibration_loader,
                calibration_texts=self.calibration_texts,
                device=self.target_device,
            )
            batches = self._normalize_batches(calib_inputs)
            n_batches = min(len(batches), self.max_calibration_batches)
            with torch.no_grad():
                for i in range(n_batches):
                    batch = batches[i]
                    try:
                        if isinstance(batch, dict):
                            batch = {
                                k: (v.to(self.target_device) if isinstance(v, torch.Tensor) else v)
                                for k, v in batch.items()
                            }
                            self.model(**batch)
                        elif isinstance(batch, torch.Tensor):
                            self.model(batch.to(self.target_device))
                        elif isinstance(batch, (list, tuple)):
                            batch = [
                                (v.to(self.target_device) if isinstance(v, torch.Tensor) else v)
                                for v in batch
                            ]
                            self.model(*batch)
                    except Exception as e:
                        logger.debug("Calibration batch %d failed: %s", i, e)
        except Exception as e:
            logger.warning("LiLLaMA: calibration forward raised: %s", e)
        finally:
            for h in hooks:
                h.remove()

        aggregated: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
        MAX_ACTIVATION_SAMPLES = 512
        for name, data in activations.items():
            if data["in"] and data["out"]:
                try:
                    cat_in = torch.cat(data["in"], dim=0)
                    cat_out = torch.cat(data["out"], dim=0)
                    if cat_in.size(0) > MAX_ACTIVATION_SAMPLES:
                        cat_in = cat_in[:MAX_ACTIVATION_SAMPLES]
                        cat_out = cat_out[:MAX_ACTIVATION_SAMPLES]
                    aggregated[name] = (cat_in, cat_out)
                except Exception as e:
                    logger.debug("Activation concat failed for %s: %s", name, e)
        return aggregated

    @staticmethod
    def _normalize_batches(calib_inputs: Any) -> list[Any]:
        if calib_inputs is None:
            return []
        if isinstance(calib_inputs, torch.Tensor):
            return [calib_inputs]
        if isinstance(calib_inputs, dict):
            return [calib_inputs]
        if isinstance(calib_inputs, (list, tuple)):
            if len(calib_inputs) == 0:
                return []
            first = calib_inputs[0]
            if isinstance(first, (dict, torch.Tensor)):
                return list(calib_inputs)
            return [calib_inputs]
        return [calib_inputs]

    # ------------------------------------------------------------------
    # SVD factorization helper
    # ------------------------------------------------------------------
    def _factorize(
        self,
        W: torch.Tensor,
        rank: int,
    ) -> tuple[torch.Tensor, int]:
        """Return (W_recon, rank_used)."""
        try:
            U, S, Vh = torch.linalg.svd(W, full_matrices=False)
            total_energy = (S**2).sum().clamp(min=1e-12)
            cum_ratio = torch.cumsum(S**2, 0) / total_energy
            r = min(rank, len(S))
            cum_e = cum_ratio[r - 1].item() if r > 0 else 0.0
            if cum_e < 0.95:
                idx = torch.where(cum_ratio >= 0.95)[0]
                if len(idx) > 0:
                    r = min(int(idx[0].item()) + 1, len(S))
            W_recon = (U[:, :r] * S[:r].unsqueeze(0)) @ Vh[:r, :]
            return W_recon, r
        except Exception as e:
            logger.debug("LiLLaMA SVD factorization failed: %s", e)
            return W, 0

    # ------------------------------------------------------------------
    # Main compress
    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        # LL-FIX-08: idempotence.
        if getattr(self.model, "_lillama_compressed", False):
            logger.warning("LiLLaMA: model already compressed; skipping.")
            try:
                return self.model, self.get_compression_stats()
            except Exception:
                return self.model, {}

        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_details = {}

        activations_map = self._collect_io_activations()
        logger.info("LiLLaMA: %d target activation sets collected.", len(activations_map))

        # LL-FIX-06: diagnostic if no targets.
        target_list = [
            (name, module)
            for name, module in self.model.named_modules()
            if self._is_target_module(name, module)
        ]
        if not target_list:
            logger.warning("LiLLaMA: no compressible layers found. Running diagnostic...")
            n_scanned = 0
            for name, module in self.model.named_modules():
                n_scanned += 1
                self._is_target_module(name, module, _debug=True)
            logger.warning("LiLLaMA diagnostic: %d modules scanned, 0 targets.", n_scanned)
            try:
                return self.model, self.get_compression_stats()
            except Exception:
                return self.model, {}

        logger.info("LiLLaMA: %d target layers identified.", len(target_list))

        total_orig_params = 0
        total_comp_params = 0

        # Deduplicated parameter accounting.
        seen_param_ids: set = set()
        count_param_ids: set = set()

        for p in self.model.parameters():
            if id(p) not in seen_param_ids:
                seen_param_ids.add(id(p))
                total_orig_params += p.numel()

        def _count_params_of(mod: nn.Module):
            nonlocal total_comp_params
            for p in mod.parameters(recurse=False):
                if id(p) not in count_param_ids:
                    count_param_ids.add(id(p))
                    total_comp_params += p.numel()

        for name, module in target_list:
            target_layer = getattr(module, "base_layer", module)
            raw_w = target_layer.weight.data
            is_conv1d = _is_conv1d(target_layer)
            device = raw_w.device
            dtype = target_layer.weight.dtype

            W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
            d_out, d_in = W.shape
            used_cross_cov = False
            r_used = self.rank

            # LL-FIX-04: initialize W_new to W (safe default).
            W_new = W

            if name in activations_map:
                inp, out = activations_map[name]
                try:
                    inp_flat = inp.reshape(-1, inp.shape[-1]).float()
                    out_flat = out.reshape(-1, out.shape[-1]).float()
                    if inp_flat.shape[-1] == d_in and out_flat.shape[-1] == d_out:
                        C = torch.matmul(out_flat.t().contiguous(), inp_flat) / max(
                            len(inp_flat), 1
                        )
                        W_new, r_used = self._factorize(C, self.rank)
                        used_cross_cov = True
                    else:
                        W_new, r_used = self._factorize(W, self.rank)
                except Exception as e:
                    logger.debug("Cross-covariance factorization failed for %s: %s", name, e)
                    W_new, r_used = self._factorize(W, self.rank)
            else:
                W_new, r_used = self._factorize(W, self.rank)

            norm_orig = torch.norm(W, p="fro")
            rel_err = (torch.norm(W - W_new, p="fro") / (norm_orig + 1e-8)).item()
            self.layer_errors[name] = rel_err

            if rel_err > self.max_error:
                logger.warning(
                    "LiLLaMA: err %.4f on %s > %.4f. %s",
                    rel_err,
                    name,
                    self.max_error,
                    (
                        "Keeping original."
                        if self.fallback_to_original
                        else "Proceeding with distilled weights."
                    ),
                )
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    for p in module.parameters(recurse=False):
                        p.requires_grad = False
                    _count_params_of(module)
                    continue

            # Shape integrity check.
            final_w = W_new.t().contiguous() if is_conv1d else W_new
            if final_w.shape != target_layer.weight.shape:
                logger.warning(
                    "LiLLaMA: shape mismatch for %s: expected %s, got %s; skipping.",
                    name,
                    tuple(target_layer.weight.shape),
                    tuple(final_w.shape),
                )
                _count_params_of(module)
                continue

            target_layer.weight.data.copy_(final_w.to(dtype=dtype, device=device))
            target_layer.weight.requires_grad = False
            if getattr(target_layer, "bias", None) is not None:
                target_layer.bias.requires_grad = False

            # Preserve LoRA adapter trainability.
            if hasattr(module, "lora_A") and hasattr(module, "lora_B"):
                for attr in ("lora_A", "lora_B"):
                    p = getattr(module, attr, None)
                    if p is None:
                        continue
                    if hasattr(p, "weight"):
                        p.weight.requires_grad = True
                    elif isinstance(p, nn.Parameter):
                        p.requires_grad = True

            self.compressed_layers.append(name)
            _count_params_of(module)

            self.layer_details[name] = {
                "rank": int(r_used),
                "rel_error": float(rel_err),
                "used_cross_cov": bool(used_cross_cov),
                "shape": tuple(final_w.shape),
            }

        # LL-FIX-07: forward sanity check.
        forward_check_success = True
        try:
            with torch.no_grad():
                dummy = torch.randint(0, 100, (1, 8), device=self.target_device, dtype=torch.long)
                try:
                    self.model(input_ids=dummy, attention_mask=torch.ones_like(dummy))
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception as e2:
                        logger.warning("LiLLaMA: forward check failed: %s", e2)
                        forward_check_success = False
        except Exception as e:
            logger.warning("LiLLaMA: forward check raised: %s", e)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        # LL-FIX-04: avg error only over compressed layers.
        compressed_errs = [
            self.layer_errors[n] for n in self.compressed_layers if n in self.layer_errors
        ]
        avg_err = sum(compressed_errs) / len(compressed_errs) if compressed_errs else 0.0
        eff_ratio = total_comp_params / max(1, total_orig_params)

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        stats.update(
            {
                "distilled_layers": len(self.compressed_layers),
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "rank": self.rank,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "avg_reconstruction_error": avg_err,
                "layer_errors": self.layer_errors,
                "layer_details": self.layer_details,
                "effective_parameter_ratio": eff_ratio,
                "forward_check_success": forward_check_success,
                "method": "lillama_cross_covariance_svd",
                # LL-FIX-12: honest note.
                "note": (
                    "LiLLaMA rewrites weights in-place with the distilled reconstruction. "
                    "The dense weight keeps its full (out, in) shape; framework-level "
                    "compression_rate reads 1.0. Use effective_parameter_ratio for the "
                    "theoretical storage metric."
                ),
            }
        )
        self.stats = stats

        # LL-FIX-13: write stats onto model.
        try:
            if not hasattr(self.model, "compression_stats") or self.model.compression_stats is None:
                self.model.compression_stats = {}
            self.model.compression_stats.update(
                {
                    "num_compressed": len(self.compressed_layers),
                    "num_fallback": len(self.fallback_layers),
                    "compression_method": "lillama_cross_covariance_svd",
                    "effective_parameter_ratio": eff_ratio,
                    "avg_reconstruction_error": avg_err,
                }
            )
        except Exception:
            pass

        self.model._lillama_compressed = True
        logger.info(
            "LiLLaMA: %d distilled, %d fallback. Avg err %.4f, eff_param_ratio %.2fx.",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            eff_ratio,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var category : str: The type of the None singleton.
var is_custom_attention : bool: The type of the None singleton.

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    # LL-FIX-08: idempotence.
    if getattr(self.model, "_lillama_compressed", False):
        logger.warning("LiLLaMA: model already compressed; skipping.")
        try:
            return self.model, self.get_compression_stats()
        except Exception:
            return self.model, {}

    self.compressed_layers = []
    self.fallback_layers = []
    self.layer_errors = {}
    self.layer_details = {}

    activations_map = self._collect_io_activations()
    logger.info("LiLLaMA: %d target activation sets collected.", len(activations_map))

    # LL-FIX-06: diagnostic if no targets.
    target_list = [
        (name, module)
        for name, module in self.model.named_modules()
        if self._is_target_module(name, module)
    ]
    if not target_list:
        logger.warning("LiLLaMA: no compressible layers found. Running diagnostic...")
        n_scanned = 0
        for name, module in self.model.named_modules():
            n_scanned += 1
            self._is_target_module(name, module, _debug=True)
        logger.warning("LiLLaMA diagnostic: %d modules scanned, 0 targets.", n_scanned)
        try:
            return self.model, self.get_compression_stats()
        except Exception:
            return self.model, {}

    logger.info("LiLLaMA: %d target layers identified.", len(target_list))

    total_orig_params = 0
    total_comp_params = 0

    # Deduplicated parameter accounting.
    seen_param_ids: set = set()
    count_param_ids: set = set()

    for p in self.model.parameters():
        if id(p) not in seen_param_ids:
            seen_param_ids.add(id(p))
            total_orig_params += p.numel()

    def _count_params_of(mod: nn.Module):
        nonlocal total_comp_params
        for p in mod.parameters(recurse=False):
            if id(p) not in count_param_ids:
                count_param_ids.add(id(p))
                total_comp_params += p.numel()

    for name, module in target_list:
        target_layer = getattr(module, "base_layer", module)
        raw_w = target_layer.weight.data
        is_conv1d = _is_conv1d(target_layer)
        device = raw_w.device
        dtype = target_layer.weight.dtype

        W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
        d_out, d_in = W.shape
        used_cross_cov = False
        r_used = self.rank

        # LL-FIX-04: initialize W_new to W (safe default).
        W_new = W

        if name in activations_map:
            inp, out = activations_map[name]
            try:
                inp_flat = inp.reshape(-1, inp.shape[-1]).float()
                out_flat = out.reshape(-1, out.shape[-1]).float()
                if inp_flat.shape[-1] == d_in and out_flat.shape[-1] == d_out:
                    C = torch.matmul(out_flat.t().contiguous(), inp_flat) / max(
                        len(inp_flat), 1
                    )
                    W_new, r_used = self._factorize(C, self.rank)
                    used_cross_cov = True
                else:
                    W_new, r_used = self._factorize(W, self.rank)
            except Exception as e:
                logger.debug("Cross-covariance factorization failed for %s: %s", name, e)
                W_new, r_used = self._factorize(W, self.rank)
        else:
            W_new, r_used = self._factorize(W, self.rank)

        norm_orig = torch.norm(W, p="fro")
        rel_err = (torch.norm(W - W_new, p="fro") / (norm_orig + 1e-8)).item()
        self.layer_errors[name] = rel_err

        if rel_err > self.max_error:
            logger.warning(
                "LiLLaMA: err %.4f on %s > %.4f. %s",
                rel_err,
                name,
                self.max_error,
                (
                    "Keeping original."
                    if self.fallback_to_original
                    else "Proceeding with distilled weights."
                ),
            )
            if self.fallback_to_original:
                self.fallback_layers.append(name)
                for p in module.parameters(recurse=False):
                    p.requires_grad = False
                _count_params_of(module)
                continue

        # Shape integrity check.
        final_w = W_new.t().contiguous() if is_conv1d else W_new
        if final_w.shape != target_layer.weight.shape:
            logger.warning(
                "LiLLaMA: shape mismatch for %s: expected %s, got %s; skipping.",
                name,
                tuple(target_layer.weight.shape),
                tuple(final_w.shape),
            )
            _count_params_of(module)
            continue

        target_layer.weight.data.copy_(final_w.to(dtype=dtype, device=device))
        target_layer.weight.requires_grad = False
        if getattr(target_layer, "bias", None) is not None:
            target_layer.bias.requires_grad = False

        # Preserve LoRA adapter trainability.
        if hasattr(module, "lora_A") and hasattr(module, "lora_B"):
            for attr in ("lora_A", "lora_B"):
                p = getattr(module, attr, None)
                if p is None:
                    continue
                if hasattr(p, "weight"):
                    p.weight.requires_grad = True
                elif isinstance(p, nn.Parameter):
                    p.requires_grad = True

        self.compressed_layers.append(name)
        _count_params_of(module)

        self.layer_details[name] = {
            "rank": int(r_used),
            "rel_error": float(rel_err),
            "used_cross_cov": bool(used_cross_cov),
            "shape": tuple(final_w.shape),
        }

    # LL-FIX-07: forward sanity check.
    forward_check_success = True
    try:
        with torch.no_grad():
            dummy = torch.randint(0, 100, (1, 8), device=self.target_device, dtype=torch.long)
            try:
                self.model(input_ids=dummy, attention_mask=torch.ones_like(dummy))
            except Exception:
                try:
                    self.model(dummy)
                except Exception as e2:
                    logger.warning("LiLLaMA: forward check failed: %s", e2)
                    forward_check_success = False
    except Exception as e:
        logger.warning("LiLLaMA: forward check raised: %s", e)
        forward_check_success = False

    try:
        self._compressed_params = self._count_params(self.model)
    except Exception:
        self._compressed_params = sum(p.numel() for p in self.model.parameters())

    # LL-FIX-04: avg error only over compressed layers.
    compressed_errs = [
        self.layer_errors[n] for n in self.compressed_layers if n in self.layer_errors
    ]
    avg_err = sum(compressed_errs) / len(compressed_errs) if compressed_errs else 0.0
    eff_ratio = total_comp_params / max(1, total_orig_params)

    try:
        stats = self.get_compression_stats()
    except Exception:
        stats = {}

    stats.update(
        {
            "distilled_layers": len(self.compressed_layers),
            "compressed_layers": self.compressed_layers,
            "fallback_layers": self.fallback_layers,
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "rank": self.rank,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "avg_reconstruction_error": avg_err,
            "layer_errors": self.layer_errors,
            "layer_details": self.layer_details,
            "effective_parameter_ratio": eff_ratio,
            "forward_check_success": forward_check_success,
            "method": "lillama_cross_covariance_svd",
            # LL-FIX-12: honest note.
            "note": (
                "LiLLaMA rewrites weights in-place with the distilled reconstruction. "
                "The dense weight keeps its full (out, in) shape; framework-level "
                "compression_rate reads 1.0. Use effective_parameter_ratio for the "
                "theoretical storage metric."
            ),
        }
    )
    self.stats = stats

    # LL-FIX-13: write stats onto model.
    try:
        if not hasattr(self.model, "compression_stats") or self.model.compression_stats is None:
            self.model.compression_stats = {}
        self.model.compression_stats.update(
            {
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "compression_method": "lillama_cross_covariance_svd",
                "effective_parameter_ratio": eff_ratio,
                "avg_reconstruction_error": avg_err,
            }
        )
    except Exception:
        pass

    self.model._lillama_compressed = True
    logger.info(
        "LiLLaMA: %d distilled, %d fallback. Avg err %.4f, eff_param_ratio %.2fx.",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        eff_ratio,
    )
    return self.model, stats

Return (compressed_model, statistics).

class LiLLaMADistiller (model: nn.Module, tokenizer: Any | None = None, rank: int = 32, calibration_texts: list[str] | None = None, calibration_loader: Any | None = None, device: str | torch.device | None = None, max_error: float = 0.2, fallback_to_original: bool = True, skip_attention: bool = False, target_attention_only: bool = False, target_modules: list[str] | None = None, min_layer_dim: int = 8, skip_should_compress_filter: bool = False, max_calibration_batches: int = 8, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("lillama_class")
@COMPRESSOR_REGISTRY.register("lillama_compressor_class")
@COMPRESSOR_REGISTRY.register("lillama_distiller_class")
class LiLLaMADistiller(BaseCompressor):
    # LL-FIX-01: this is a local distillation / weight-rewriter, not standard factorized attention.
    is_custom_attention: bool = False
    category: str = "weight_rewriting"

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any | None = None,
        rank: int = 32,
        calibration_texts: list[str] | None = None,
        calibration_loader: Any | None = None,
        device: str | torch.device | None = None,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        skip_attention: bool = False,
        target_attention_only: bool = False,
        target_modules: list[str] | None = None,
        min_layer_dim: int = 8,
        skip_should_compress_filter: bool = False,
        max_calibration_batches: int = 8,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="lillama", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            rank = _cfg_get(cfg, "lillama_rank", _cfg_get(cfg, "rank", rank))
            max_error = _cfg_get(cfg, "max_error", max_error)
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            skip_attention = _cfg_get(cfg, "skip_attention", skip_attention)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)
            min_layer_dim = _cfg_get(cfg, "min_layer_dim", min_layer_dim)
            skip_should_compress_filter = _cfg_get(
                cfg,
                "skip_should_compress_filter",
                skip_should_compress_filter,
            )
            max_calibration_batches = _cfg_get(
                cfg, "max_calibration_batches", max_calibration_batches
            )

        # LL-FIX-11: tokenizer resolution.
        self.tokenizer = (
            tokenizer
            or getattr(model, "tokenizer", None)
            or getattr(getattr(model, "config", None), "tokenizer", None)
        )
        self.rank = max(1, int(rank))
        self.calibration_texts = calibration_texts
        self.calibration_loader = calibration_loader
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.skip_attention = bool(skip_attention)
        self.target_attention_only = bool(target_attention_only)
        self.target_modules = target_modules
        self.min_layer_dim = max(1, int(min_layer_dim))
        self.skip_should_compress_filter = bool(skip_should_compress_filter)
        self.max_calibration_batches = max(1, int(max_calibration_batches))

        # LL-FIX-03: safe device.
        if device is not None:
            self.target_device = torch.device(device)
        else:
            self.target_device = _get_device(model)

        # LL-FIX-15: initialize all lists in __init__.
        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_details: dict[str, Any] = {}
        self.stats: dict[str, Any] = {}
        self._compressed_params: int = 0

    # ------------------------------------------------------------------
    # Target check (LL-FIX-06)
    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module, _debug: bool = False) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ("lora_A", "lora_B", "base_layer")):
            if _debug:
                logger.info(f"  [skip] {name}: lora/base_layer part")
            return False

        # Skip already-compressed.
        cls_name = module.__class__.__name__
        if cls_name in (
            "LSCPrunedLoRALinear",
            "DynamicSVDLinear",
            "SVDLinear",
            "FactorisedLinear",
            "LowRankPlusSparseLinear",
        ):
            if _debug:
                logger.info(f"  [skip] {name}: already compressed ({cls_name})")
            return False

        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            ]
        ):
            if _debug:
                logger.info(f"  [skip] {name}: skip keyword")
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_lower = target_mod.__class__.__name__.lower()
        is_linear_like = (
            isinstance(target_mod, nn.Linear) or "conv1d" in cls_lower or hasattr(target_mod, "nf")
        )
        if not is_linear_like:
            if _debug:
                logger.info(f"  [skip] {name}: not linear-like")
            return False

        # LL-FIX-09: allow bypassing framework filter.
        if not self.skip_should_compress_filter:
            try:
                if not should_compress_linear(self.model, name):
                    if _debug:
                        logger.info(f"  [skip] {name}: should_compress_linear=False")
                    return False
            except Exception as e:
                if _debug:
                    logger.info(f"  [skip] {name}: should_compress_linear raised {e}")
                return False

        # LL-FIX-10: filter tiny layers.
        w = getattr(target_mod, "weight", None)
        if w is None:
            if _debug:
                logger.info(f"  [skip] {name}: weight None")
            return False
        if min(w.shape[0], w.shape[1]) < self.min_layer_dim:
            if _debug:
                logger.info(f"  [skip] {name}: dim too small {tuple(w.shape)}")
            return False

        if self.target_attention_only:
            is_attn = any(k in name_lower for k in _ATTN_NAMES)
            if not is_attn:
                if _debug:
                    logger.info(f"  [skip] {name}: not attention")
                return False

        if self.skip_attention:
            attn_keywords = {
                "query",
                "key",
                "value",
                "q_lin",
                "k_lin",
                "v_lin",
                "q_proj",
                "k_proj",
                "v_proj",
                "out_lin",
                "out_proj",
                "dense",
                "o_proj",
                "c_attn",
            }
            last = name.split(".")[-1].lower()
            if last in attn_keywords or any(k in name_lower for k in ("self_attn", "attn.c_attn")):
                if _debug:
                    logger.info(f"  [skip] {name}: skip_attention")
                return False

        if self.target_modules is not None:
            ok = any(t in name for t in self.target_modules)
            if not ok and _debug:
                logger.info(f"  [skip] {name}: not in target_modules")
            return ok

        if _debug:
            logger.info(f"  [TARGET] {name}: shape={tuple(w.shape)}")
        return True

    # ------------------------------------------------------------------
    # IO activation collection
    # ------------------------------------------------------------------
    def _collect_io_activations(
        self,
    ) -> dict[str, tuple[torch.Tensor, torch.Tensor]]:
        """Collect input/output activations across target layers."""
        self.model.eval()
        activations: dict[str, dict[str, list[torch.Tensor]]] = {}
        hooks = []

        def make_hook(layer_name: str):
            def hook(module: nn.Module, inputs: tuple[Any, ...], output: Any):
                inp = inputs[0].detach() if inputs and isinstance(inputs[0], torch.Tensor) else None
                if isinstance(output, torch.Tensor):
                    out = output.detach()
                elif (
                    isinstance(output, (tuple, list))
                    and output
                    and isinstance(output[0], torch.Tensor)
                ):
                    out = output[0].detach()
                else:
                    out = None
                if inp is not None and out is not None:
                    if layer_name not in activations:
                        activations[layer_name] = {"in": [], "out": []}
                    activations[layer_name]["in"].append(inp.cpu())
                    activations[layer_name]["out"].append(out.cpu())

            return hook

        target_list = [
            (name, module)
            for name, module in self.model.named_modules()
            if self._is_target_module(name, module)
        ]
        if not target_list:
            return {}

        for name, module in target_list:
            hooks.append(module.register_forward_hook(make_hook(name)))

        try:
            calib_inputs = get_calibration_tensors(
                self.model,
                tokenizer=self.tokenizer,
                calibration_loader=self.calibration_loader,
                calibration_texts=self.calibration_texts,
                device=self.target_device,
            )
            batches = self._normalize_batches(calib_inputs)
            n_batches = min(len(batches), self.max_calibration_batches)
            with torch.no_grad():
                for i in range(n_batches):
                    batch = batches[i]
                    try:
                        if isinstance(batch, dict):
                            batch = {
                                k: (v.to(self.target_device) if isinstance(v, torch.Tensor) else v)
                                for k, v in batch.items()
                            }
                            self.model(**batch)
                        elif isinstance(batch, torch.Tensor):
                            self.model(batch.to(self.target_device))
                        elif isinstance(batch, (list, tuple)):
                            batch = [
                                (v.to(self.target_device) if isinstance(v, torch.Tensor) else v)
                                for v in batch
                            ]
                            self.model(*batch)
                    except Exception as e:
                        logger.debug("Calibration batch %d failed: %s", i, e)
        except Exception as e:
            logger.warning("LiLLaMA: calibration forward raised: %s", e)
        finally:
            for h in hooks:
                h.remove()

        aggregated: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
        MAX_ACTIVATION_SAMPLES = 512
        for name, data in activations.items():
            if data["in"] and data["out"]:
                try:
                    cat_in = torch.cat(data["in"], dim=0)
                    cat_out = torch.cat(data["out"], dim=0)
                    if cat_in.size(0) > MAX_ACTIVATION_SAMPLES:
                        cat_in = cat_in[:MAX_ACTIVATION_SAMPLES]
                        cat_out = cat_out[:MAX_ACTIVATION_SAMPLES]
                    aggregated[name] = (cat_in, cat_out)
                except Exception as e:
                    logger.debug("Activation concat failed for %s: %s", name, e)
        return aggregated

    @staticmethod
    def _normalize_batches(calib_inputs: Any) -> list[Any]:
        if calib_inputs is None:
            return []
        if isinstance(calib_inputs, torch.Tensor):
            return [calib_inputs]
        if isinstance(calib_inputs, dict):
            return [calib_inputs]
        if isinstance(calib_inputs, (list, tuple)):
            if len(calib_inputs) == 0:
                return []
            first = calib_inputs[0]
            if isinstance(first, (dict, torch.Tensor)):
                return list(calib_inputs)
            return [calib_inputs]
        return [calib_inputs]

    # ------------------------------------------------------------------
    # SVD factorization helper
    # ------------------------------------------------------------------
    def _factorize(
        self,
        W: torch.Tensor,
        rank: int,
    ) -> tuple[torch.Tensor, int]:
        """Return (W_recon, rank_used)."""
        try:
            U, S, Vh = torch.linalg.svd(W, full_matrices=False)
            total_energy = (S**2).sum().clamp(min=1e-12)
            cum_ratio = torch.cumsum(S**2, 0) / total_energy
            r = min(rank, len(S))
            cum_e = cum_ratio[r - 1].item() if r > 0 else 0.0
            if cum_e < 0.95:
                idx = torch.where(cum_ratio >= 0.95)[0]
                if len(idx) > 0:
                    r = min(int(idx[0].item()) + 1, len(S))
            W_recon = (U[:, :r] * S[:r].unsqueeze(0)) @ Vh[:r, :]
            return W_recon, r
        except Exception as e:
            logger.debug("LiLLaMA SVD factorization failed: %s", e)
            return W, 0

    # ------------------------------------------------------------------
    # Main compress
    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        # LL-FIX-08: idempotence.
        if getattr(self.model, "_lillama_compressed", False):
            logger.warning("LiLLaMA: model already compressed; skipping.")
            try:
                return self.model, self.get_compression_stats()
            except Exception:
                return self.model, {}

        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        self.layer_details = {}

        activations_map = self._collect_io_activations()
        logger.info("LiLLaMA: %d target activation sets collected.", len(activations_map))

        # LL-FIX-06: diagnostic if no targets.
        target_list = [
            (name, module)
            for name, module in self.model.named_modules()
            if self._is_target_module(name, module)
        ]
        if not target_list:
            logger.warning("LiLLaMA: no compressible layers found. Running diagnostic...")
            n_scanned = 0
            for name, module in self.model.named_modules():
                n_scanned += 1
                self._is_target_module(name, module, _debug=True)
            logger.warning("LiLLaMA diagnostic: %d modules scanned, 0 targets.", n_scanned)
            try:
                return self.model, self.get_compression_stats()
            except Exception:
                return self.model, {}

        logger.info("LiLLaMA: %d target layers identified.", len(target_list))

        total_orig_params = 0
        total_comp_params = 0

        # Deduplicated parameter accounting.
        seen_param_ids: set = set()
        count_param_ids: set = set()

        for p in self.model.parameters():
            if id(p) not in seen_param_ids:
                seen_param_ids.add(id(p))
                total_orig_params += p.numel()

        def _count_params_of(mod: nn.Module):
            nonlocal total_comp_params
            for p in mod.parameters(recurse=False):
                if id(p) not in count_param_ids:
                    count_param_ids.add(id(p))
                    total_comp_params += p.numel()

        for name, module in target_list:
            target_layer = getattr(module, "base_layer", module)
            raw_w = target_layer.weight.data
            is_conv1d = _is_conv1d(target_layer)
            device = raw_w.device
            dtype = target_layer.weight.dtype

            W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
            d_out, d_in = W.shape
            used_cross_cov = False
            r_used = self.rank

            # LL-FIX-04: initialize W_new to W (safe default).
            W_new = W

            if name in activations_map:
                inp, out = activations_map[name]
                try:
                    inp_flat = inp.reshape(-1, inp.shape[-1]).float()
                    out_flat = out.reshape(-1, out.shape[-1]).float()
                    if inp_flat.shape[-1] == d_in and out_flat.shape[-1] == d_out:
                        C = torch.matmul(out_flat.t().contiguous(), inp_flat) / max(
                            len(inp_flat), 1
                        )
                        W_new, r_used = self._factorize(C, self.rank)
                        used_cross_cov = True
                    else:
                        W_new, r_used = self._factorize(W, self.rank)
                except Exception as e:
                    logger.debug("Cross-covariance factorization failed for %s: %s", name, e)
                    W_new, r_used = self._factorize(W, self.rank)
            else:
                W_new, r_used = self._factorize(W, self.rank)

            norm_orig = torch.norm(W, p="fro")
            rel_err = (torch.norm(W - W_new, p="fro") / (norm_orig + 1e-8)).item()
            self.layer_errors[name] = rel_err

            if rel_err > self.max_error:
                logger.warning(
                    "LiLLaMA: err %.4f on %s > %.4f. %s",
                    rel_err,
                    name,
                    self.max_error,
                    (
                        "Keeping original."
                        if self.fallback_to_original
                        else "Proceeding with distilled weights."
                    ),
                )
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    for p in module.parameters(recurse=False):
                        p.requires_grad = False
                    _count_params_of(module)
                    continue

            # Shape integrity check.
            final_w = W_new.t().contiguous() if is_conv1d else W_new
            if final_w.shape != target_layer.weight.shape:
                logger.warning(
                    "LiLLaMA: shape mismatch for %s: expected %s, got %s; skipping.",
                    name,
                    tuple(target_layer.weight.shape),
                    tuple(final_w.shape),
                )
                _count_params_of(module)
                continue

            target_layer.weight.data.copy_(final_w.to(dtype=dtype, device=device))
            target_layer.weight.requires_grad = False
            if getattr(target_layer, "bias", None) is not None:
                target_layer.bias.requires_grad = False

            # Preserve LoRA adapter trainability.
            if hasattr(module, "lora_A") and hasattr(module, "lora_B"):
                for attr in ("lora_A", "lora_B"):
                    p = getattr(module, attr, None)
                    if p is None:
                        continue
                    if hasattr(p, "weight"):
                        p.weight.requires_grad = True
                    elif isinstance(p, nn.Parameter):
                        p.requires_grad = True

            self.compressed_layers.append(name)
            _count_params_of(module)

            self.layer_details[name] = {
                "rank": int(r_used),
                "rel_error": float(rel_err),
                "used_cross_cov": bool(used_cross_cov),
                "shape": tuple(final_w.shape),
            }

        # LL-FIX-07: forward sanity check.
        forward_check_success = True
        try:
            with torch.no_grad():
                dummy = torch.randint(0, 100, (1, 8), device=self.target_device, dtype=torch.long)
                try:
                    self.model(input_ids=dummy, attention_mask=torch.ones_like(dummy))
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception as e2:
                        logger.warning("LiLLaMA: forward check failed: %s", e2)
                        forward_check_success = False
        except Exception as e:
            logger.warning("LiLLaMA: forward check raised: %s", e)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        # LL-FIX-04: avg error only over compressed layers.
        compressed_errs = [
            self.layer_errors[n] for n in self.compressed_layers if n in self.layer_errors
        ]
        avg_err = sum(compressed_errs) / len(compressed_errs) if compressed_errs else 0.0
        eff_ratio = total_comp_params / max(1, total_orig_params)

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        stats.update(
            {
                "distilled_layers": len(self.compressed_layers),
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "rank": self.rank,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "avg_reconstruction_error": avg_err,
                "layer_errors": self.layer_errors,
                "layer_details": self.layer_details,
                "effective_parameter_ratio": eff_ratio,
                "forward_check_success": forward_check_success,
                "method": "lillama_cross_covariance_svd",
                # LL-FIX-12: honest note.
                "note": (
                    "LiLLaMA rewrites weights in-place with the distilled reconstruction. "
                    "The dense weight keeps its full (out, in) shape; framework-level "
                    "compression_rate reads 1.0. Use effective_parameter_ratio for the "
                    "theoretical storage metric."
                ),
            }
        )
        self.stats = stats

        # LL-FIX-13: write stats onto model.
        try:
            if not hasattr(self.model, "compression_stats") or self.model.compression_stats is None:
                self.model.compression_stats = {}
            self.model.compression_stats.update(
                {
                    "num_compressed": len(self.compressed_layers),
                    "num_fallback": len(self.fallback_layers),
                    "compression_method": "lillama_cross_covariance_svd",
                    "effective_parameter_ratio": eff_ratio,
                    "avg_reconstruction_error": avg_err,
                }
            )
        except Exception:
            pass

        self.model._lillama_compressed = True
        logger.info(
            "LiLLaMA: %d distilled, %d fallback. Avg err %.4f, eff_param_ratio %.2fx.",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            eff_ratio,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var category : str: The type of the None singleton.
var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class LoRALinear (in_features: int, out_features: int, bias: bool = True, lora_rank: int = 16, lora_alpha: float = 32.0, freeze_base: bool = False, device: torch.device | None = None, dtype: torch.dtype | None = None)

Expand source code

class LoRALinear(nn.Module):
    is_custom_attention: bool = False

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        freeze_base: bool = False,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        factory_kwargs = {"device": device, "dtype": dtype}
        self.in_features = in_features
        self.out_features = out_features
        self.lora_rank = lora_rank
        self.lora_alpha = lora_alpha
        self.lora_scaling = (lora_alpha / lora_rank) if lora_rank > 0 else 1.0
        self.freeze_base = freeze_base

        self.weight = nn.Parameter(torch.empty((out_features, in_features), **factory_kwargs))
        if bias:
            self.bias = nn.Parameter(torch.empty(out_features, **factory_kwargs))
        else:
            self.register_parameter("bias", None)

        if self.lora_rank > 0:
            self.lora_A = nn.Parameter(torch.empty((lora_rank, in_features), **factory_kwargs))
            self.lora_B = nn.Parameter(torch.empty((out_features, lora_rank), **factory_kwargs))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        self.reset_parameters()
        if freeze_base:
            self.freeze_base_weights()

    def reset_parameters(self) -> None:
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            nn.init.uniform_(self.bias, -bound, bound)
        if self.lora_rank > 0 and self.lora_A is not None:
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            self.lora_A.requires_grad = True
        if self.lora_rank > 0 and self.lora_B is not None:
            nn.init.zeros_(self.lora_B)
            self.lora_B.requires_grad = True

    def freeze_base_weights(self) -> None:
        self.weight.requires_grad = False
        if self.bias is not None:
            self.bias.requires_grad = False
        if self.lora_rank > 0 and self.lora_A is not None:
            self.lora_A.requires_grad = True
        if self.lora_rank > 0 and self.lora_B is not None:
            self.lora_B.requires_grad = True

    def unfreeze_base_weights(self) -> None:
        self.weight.requires_grad = True
        if self.bias is not None:
            self.bias.requires_grad = True

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = F.linear(x, self.weight, self.bias)
        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling
            out = out + lora_out
        return out

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    out = F.linear(x, self.weight, self.bias)
    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling
        out = out + lora_out
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    self.weight.requires_grad = False
    if self.bias is not None:
        self.bias.requires_grad = False
    if self.lora_rank > 0 and self.lora_A is not None:
        self.lora_A.requires_grad = True
    if self.lora_rank > 0 and self.lora_B is not None:
        self.lora_B.requires_grad = True

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
    if self.bias is not None:
        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
        nn.init.uniform_(self.bias, -bound, bound)
    if self.lora_rank > 0 and self.lora_A is not None:
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        self.lora_A.requires_grad = True
    if self.lora_rank > 0 and self.lora_B is not None:
        nn.init.zeros_(self.lora_B)
        self.lora_B.requires_grad = True

def unfreeze_base_weights(self) ‑> None

Expand source code

def unfreeze_base_weights(self) -> None:
    self.weight.requires_grad = True
    if self.bias is not None:
        self.bias.requires_grad = True

class LowRankCompressor (model: nn.Module, rank: int = 8, energy_threshold: float | None = None, max_error: float = 0.5, target_attention_only: bool = False, fallback_to_original: bool = True, target_modules: list[str] | None = None, **kwargs: Any)

Expand source code

class LowRankCompressor(BaseCompressor):
    is_custom_attention: bool = True
    """
    SoKA Low-Rank Linear Compressor.
    Replaces dense nn.Linear and Conv1D layers with factorized LowRankLinear layers.
    """

    def __init__(
        self,
        model: nn.Module,
        rank: int = 8,
        energy_threshold: float | None = None,
        max_error: float = 0.5,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        target_modules: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="soka", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            rank = _cfg_get(cfg, "soka_rank", _cfg_get(cfg, "rank", rank))
            energy_threshold = _cfg_get(
                cfg, "soka_energy", _cfg_get(cfg, "energy_threshold", energy_threshold)
            )
            max_error = _cfg_get(cfg, "max_error", max_error)
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)

        self.rank = max(1, int(rank))
        self.energy_threshold = energy_threshold
        self.max_error = float(max_error)
        self.target_attention_only = bool(target_attention_only)
        self.fallback_to_original = bool(fallback_to_original)
        self.target_modules = target_modules
        self.compressed_layers: list[dict[str, Any]] = []
        self.layer_errors: dict[str, float] = {}
        self.stats: dict[str, Any] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        if any(k in name_l for k in ["classifier", "score", "lm_head", "pooler", "embedding"]):
            return False

        if self.target_attention_only:
            attn_keywords = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(k in name_l for k in attn_keywords):
                return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__
        is_linear_like = (
            isinstance(base_mod, nn.Linear)
            or "conv1d" in cls_name.lower()
            or hasattr(base_mod, "nf")
        )
        if not is_linear_like:
            return False
        if isinstance(base_mod, LowRankLinear):
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    def _select_rank(self, S: torch.Tensor, in_f: int, out_f: int) -> int:
        """Determines rank dynamically based on singular energy or fixed target, capped for true reduction."""
        param_cap = max(1, (in_f * out_f) // (in_f + out_f) - 1)
        max_r = min(len(S), param_cap)

        if self.energy_threshold is None:
            return max(1, min(self.rank, max_r))
        total_energy = torch.sum(S**2)
        if total_energy <= 0:
            return max(1, min(self.rank, max_r))
        cum_energy = torch.cumsum(S**2, dim=0) / total_energy
        idx = torch.where(cum_energy >= self.energy_threshold)[0]
        r = idx[0].item() if len(idx) > 0 else len(S)
        return max(1, min(r, max_r))

    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """Replaces target Linear and Conv1D modules with LowRankLinear."""
        self.compressed_layers = []
        self.layer_errors = {}

        linear_layers: list[tuple[str, nn.Module]] = []
        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                linear_layers.append((name, module))

        if not linear_layers:
            logger.warning("SoKA: No compressible linear layers found.")
            return self.model, self.get_compression_stats()

        for name, module in linear_layers:
            base_mod = getattr(module, "base_layer", module)
            raw_w = base_mod.weight.data
            cls_name = base_mod.__class__.__name__
            is_conv1d = (
                "conv1d" in cls_name.lower()
                or hasattr(base_mod, "nf")
                or (
                    raw_w.ndim == 2
                    and hasattr(base_mod, "in_features")
                    and base_mod.in_features == raw_w.shape[0]
                    and raw_w.shape[0] < raw_w.shape[1]
                )
            )

            W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
            device = raw_w.device
            dtype = raw_w.dtype
            d_out, d_in = W.shape

            try:
                U, S, Vh = torch.linalg.svd(W, full_matrices=False)
                r = self._select_rank(S, d_in, d_out)
                sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))

                new_layer = LowRankLinear(
                    in_features=d_in,
                    out_features=d_out,
                    rank=r,
                    bias=(base_mod.bias is not None),
                    device=device,
                    dtype=dtype,
                )

                # Assign factorized weights: W ≈ (U_r sqrt(S)) @ (sqrt(S) V_r^T)
                # U: [d_out, r], V: [d_in, r]
                new_layer.U.data = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
                new_layer.V.data = (Vh[:r, :].T * sqrt_S.unsqueeze(0)).to(
                    dtype=dtype, device=device
                )

                if base_mod.bias is not None:
                    new_layer.bias.data.copy_(base_mod.bias.data)

                # Frobenius reconstruction error check
                W_rec = new_layer.weight.float()
                norm_orig = torch.norm(W, p="fro")
                rel_err = (torch.norm(W - W_rec, p="fro") / (norm_orig + 1e-8)).item()
                self.layer_errors[name] = rel_err

                if self.fallback_to_original and rel_err > self.max_error:
                    logger.warning(
                        "SoKA reconstruction error %.4f on %s exceeded threshold %.4f. Skipping.",
                        rel_err,
                        name,
                        self.max_error,
                    )
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

                compressed_layer = new_layer
                # Preserve existing LoRA adapter weights if the source module has them
                if (
                    hasattr(module, "lora_A")
                    and hasattr(module, "lora_B")
                    and module.lora_A is not None
                    and module.lora_B is not None
                ):
                    src_lA = getattr(module.lora_A, "weight", module.lora_A)
                    src_lB = getattr(module.lora_B, "weight", module.lora_B)
                    if hasattr(new_layer, "lora_A") and new_layer.lora_A is not None:
                        if hasattr(src_lA, "data") and src_lA.shape == new_layer.lora_A.shape:
                            new_layer.lora_A.data.copy_(src_lA.data)
                    if hasattr(new_layer, "lora_B") and new_layer.lora_B is not None:
                        if hasattr(src_lB, "data") and src_lB.shape == new_layer.lora_B.shape:
                            new_layer.lora_B.data.copy_(src_lB.data)

                if hasattr(module, "base_layer"):
                    module.base_layer = compressed_layer  # put compressed layer inside wrapper
                    # module itself (LoRALinear wrapper) remains the replacement
                    new_module = module
                else:
                    new_module = compressed_layer
                if not hasattr(module, "base_layer"):
                    if name == "":
                        self.model = new_module
                    else:
                        replace_module(self.model, name, new_module)
                self.compressed_layers.append(
                    {
                        "name": name,
                        "original_shape": (d_out, d_in),
                        "rank": r,
                        "error": rel_err,
                    }
                )
            except Exception as e:
                logger.debug("SVD decomposition error in %s: %s", name, e)
                for p in module.parameters():
                    p.requires_grad = False

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        stats.update(
            {
                "total_layers": len(self.compressed_layers),
                "method": "soka_low_rank_svd",
                "fixed_rank": self.rank,
                "energy_threshold": self.energy_threshold,
                "layer_errors": self.layer_errors,
                "target_attention_only": self.target_attention_only,
                "fallback_to_original": self.fallback_to_original,
                "max_error": self.max_error,
            }
        )
        self.stats = stats
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: SoKA Low-Rank Linear Compressor. Replaces dense nn.Linear and Conv1D layers with factorized LowRankLinear layers.

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """Replaces target Linear and Conv1D modules with LowRankLinear."""
    self.compressed_layers = []
    self.layer_errors = {}

    linear_layers: list[tuple[str, nn.Module]] = []
    for name, module in self.model.named_modules():
        if self._is_target_module(name, module):
            linear_layers.append((name, module))

    if not linear_layers:
        logger.warning("SoKA: No compressible linear layers found.")
        return self.model, self.get_compression_stats()

    for name, module in linear_layers:
        base_mod = getattr(module, "base_layer", module)
        raw_w = base_mod.weight.data
        cls_name = base_mod.__class__.__name__
        is_conv1d = (
            "conv1d" in cls_name.lower()
            or hasattr(base_mod, "nf")
            or (
                raw_w.ndim == 2
                and hasattr(base_mod, "in_features")
                and base_mod.in_features == raw_w.shape[0]
                and raw_w.shape[0] < raw_w.shape[1]
            )
        )

        W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
        device = raw_w.device
        dtype = raw_w.dtype
        d_out, d_in = W.shape

        try:
            U, S, Vh = torch.linalg.svd(W, full_matrices=False)
            r = self._select_rank(S, d_in, d_out)
            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))

            new_layer = LowRankLinear(
                in_features=d_in,
                out_features=d_out,
                rank=r,
                bias=(base_mod.bias is not None),
                device=device,
                dtype=dtype,
            )

            # Assign factorized weights: W ≈ (U_r sqrt(S)) @ (sqrt(S) V_r^T)
            # U: [d_out, r], V: [d_in, r]
            new_layer.U.data = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
            new_layer.V.data = (Vh[:r, :].T * sqrt_S.unsqueeze(0)).to(
                dtype=dtype, device=device
            )

            if base_mod.bias is not None:
                new_layer.bias.data.copy_(base_mod.bias.data)

            # Frobenius reconstruction error check
            W_rec = new_layer.weight.float()
            norm_orig = torch.norm(W, p="fro")
            rel_err = (torch.norm(W - W_rec, p="fro") / (norm_orig + 1e-8)).item()
            self.layer_errors[name] = rel_err

            if self.fallback_to_original and rel_err > self.max_error:
                logger.warning(
                    "SoKA reconstruction error %.4f on %s exceeded threshold %.4f. Skipping.",
                    rel_err,
                    name,
                    self.max_error,
                )
                for p in module.parameters():
                    p.requires_grad = False
                continue

            compressed_layer = new_layer
            # Preserve existing LoRA adapter weights if the source module has them
            if (
                hasattr(module, "lora_A")
                and hasattr(module, "lora_B")
                and module.lora_A is not None
                and module.lora_B is not None
            ):
                src_lA = getattr(module.lora_A, "weight", module.lora_A)
                src_lB = getattr(module.lora_B, "weight", module.lora_B)
                if hasattr(new_layer, "lora_A") and new_layer.lora_A is not None:
                    if hasattr(src_lA, "data") and src_lA.shape == new_layer.lora_A.shape:
                        new_layer.lora_A.data.copy_(src_lA.data)
                if hasattr(new_layer, "lora_B") and new_layer.lora_B is not None:
                    if hasattr(src_lB, "data") and src_lB.shape == new_layer.lora_B.shape:
                        new_layer.lora_B.data.copy_(src_lB.data)

            if hasattr(module, "base_layer"):
                module.base_layer = compressed_layer  # put compressed layer inside wrapper
                # module itself (LoRALinear wrapper) remains the replacement
                new_module = module
            else:
                new_module = compressed_layer
            if not hasattr(module, "base_layer"):
                if name == "":
                    self.model = new_module
                else:
                    replace_module(self.model, name, new_module)
            self.compressed_layers.append(
                {
                    "name": name,
                    "original_shape": (d_out, d_in),
                    "rank": r,
                    "error": rel_err,
                }
            )
        except Exception as e:
            logger.debug("SVD decomposition error in %s: %s", name, e)
            for p in module.parameters():
                p.requires_grad = False

    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    stats.update(
        {
            "total_layers": len(self.compressed_layers),
            "method": "soka_low_rank_svd",
            "fixed_rank": self.rank,
            "energy_threshold": self.energy_threshold,
            "layer_errors": self.layer_errors,
            "target_attention_only": self.target_attention_only,
            "fallback_to_original": self.fallback_to_original,
            "max_error": self.max_error,
        }
    )
    self.stats = stats
    return self.model, stats

Replaces target Linear and Conv1D modules with LowRankLinear.

class SoKACompressor (model: nn.Module, rank: int = 8, energy_threshold: float | None = None, max_error: float = 0.5, target_attention_only: bool = False, fallback_to_original: bool = True, target_modules: list[str] | None = None, **kwargs: Any)

Expand source code

class LowRankCompressor(BaseCompressor):
    is_custom_attention: bool = True
    """
    SoKA Low-Rank Linear Compressor.
    Replaces dense nn.Linear and Conv1D layers with factorized LowRankLinear layers.
    """

    def __init__(
        self,
        model: nn.Module,
        rank: int = 8,
        energy_threshold: float | None = None,
        max_error: float = 0.5,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        target_modules: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="soka", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            rank = _cfg_get(cfg, "soka_rank", _cfg_get(cfg, "rank", rank))
            energy_threshold = _cfg_get(
                cfg, "soka_energy", _cfg_get(cfg, "energy_threshold", energy_threshold)
            )
            max_error = _cfg_get(cfg, "max_error", max_error)
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)

        self.rank = max(1, int(rank))
        self.energy_threshold = energy_threshold
        self.max_error = float(max_error)
        self.target_attention_only = bool(target_attention_only)
        self.fallback_to_original = bool(fallback_to_original)
        self.target_modules = target_modules
        self.compressed_layers: list[dict[str, Any]] = []
        self.layer_errors: dict[str, float] = {}
        self.stats: dict[str, Any] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        if any(k in name_l for k in ["classifier", "score", "lm_head", "pooler", "embedding"]):
            return False

        if self.target_attention_only:
            attn_keywords = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(k in name_l for k in attn_keywords):
                return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__
        is_linear_like = (
            isinstance(base_mod, nn.Linear)
            or "conv1d" in cls_name.lower()
            or hasattr(base_mod, "nf")
        )
        if not is_linear_like:
            return False
        if isinstance(base_mod, LowRankLinear):
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    def _select_rank(self, S: torch.Tensor, in_f: int, out_f: int) -> int:
        """Determines rank dynamically based on singular energy or fixed target, capped for true reduction."""
        param_cap = max(1, (in_f * out_f) // (in_f + out_f) - 1)
        max_r = min(len(S), param_cap)

        if self.energy_threshold is None:
            return max(1, min(self.rank, max_r))
        total_energy = torch.sum(S**2)
        if total_energy <= 0:
            return max(1, min(self.rank, max_r))
        cum_energy = torch.cumsum(S**2, dim=0) / total_energy
        idx = torch.where(cum_energy >= self.energy_threshold)[0]
        r = idx[0].item() if len(idx) > 0 else len(S)
        return max(1, min(r, max_r))

    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """Replaces target Linear and Conv1D modules with LowRankLinear."""
        self.compressed_layers = []
        self.layer_errors = {}

        linear_layers: list[tuple[str, nn.Module]] = []
        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                linear_layers.append((name, module))

        if not linear_layers:
            logger.warning("SoKA: No compressible linear layers found.")
            return self.model, self.get_compression_stats()

        for name, module in linear_layers:
            base_mod = getattr(module, "base_layer", module)
            raw_w = base_mod.weight.data
            cls_name = base_mod.__class__.__name__
            is_conv1d = (
                "conv1d" in cls_name.lower()
                or hasattr(base_mod, "nf")
                or (
                    raw_w.ndim == 2
                    and hasattr(base_mod, "in_features")
                    and base_mod.in_features == raw_w.shape[0]
                    and raw_w.shape[0] < raw_w.shape[1]
                )
            )

            W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
            device = raw_w.device
            dtype = raw_w.dtype
            d_out, d_in = W.shape

            try:
                U, S, Vh = torch.linalg.svd(W, full_matrices=False)
                r = self._select_rank(S, d_in, d_out)
                sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))

                new_layer = LowRankLinear(
                    in_features=d_in,
                    out_features=d_out,
                    rank=r,
                    bias=(base_mod.bias is not None),
                    device=device,
                    dtype=dtype,
                )

                # Assign factorized weights: W ≈ (U_r sqrt(S)) @ (sqrt(S) V_r^T)
                # U: [d_out, r], V: [d_in, r]
                new_layer.U.data = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
                new_layer.V.data = (Vh[:r, :].T * sqrt_S.unsqueeze(0)).to(
                    dtype=dtype, device=device
                )

                if base_mod.bias is not None:
                    new_layer.bias.data.copy_(base_mod.bias.data)

                # Frobenius reconstruction error check
                W_rec = new_layer.weight.float()
                norm_orig = torch.norm(W, p="fro")
                rel_err = (torch.norm(W - W_rec, p="fro") / (norm_orig + 1e-8)).item()
                self.layer_errors[name] = rel_err

                if self.fallback_to_original and rel_err > self.max_error:
                    logger.warning(
                        "SoKA reconstruction error %.4f on %s exceeded threshold %.4f. Skipping.",
                        rel_err,
                        name,
                        self.max_error,
                    )
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

                compressed_layer = new_layer
                # Preserve existing LoRA adapter weights if the source module has them
                if (
                    hasattr(module, "lora_A")
                    and hasattr(module, "lora_B")
                    and module.lora_A is not None
                    and module.lora_B is not None
                ):
                    src_lA = getattr(module.lora_A, "weight", module.lora_A)
                    src_lB = getattr(module.lora_B, "weight", module.lora_B)
                    if hasattr(new_layer, "lora_A") and new_layer.lora_A is not None:
                        if hasattr(src_lA, "data") and src_lA.shape == new_layer.lora_A.shape:
                            new_layer.lora_A.data.copy_(src_lA.data)
                    if hasattr(new_layer, "lora_B") and new_layer.lora_B is not None:
                        if hasattr(src_lB, "data") and src_lB.shape == new_layer.lora_B.shape:
                            new_layer.lora_B.data.copy_(src_lB.data)

                if hasattr(module, "base_layer"):
                    module.base_layer = compressed_layer  # put compressed layer inside wrapper
                    # module itself (LoRALinear wrapper) remains the replacement
                    new_module = module
                else:
                    new_module = compressed_layer
                if not hasattr(module, "base_layer"):
                    if name == "":
                        self.model = new_module
                    else:
                        replace_module(self.model, name, new_module)
                self.compressed_layers.append(
                    {
                        "name": name,
                        "original_shape": (d_out, d_in),
                        "rank": r,
                        "error": rel_err,
                    }
                )
            except Exception as e:
                logger.debug("SVD decomposition error in %s: %s", name, e)
                for p in module.parameters():
                    p.requires_grad = False

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        stats.update(
            {
                "total_layers": len(self.compressed_layers),
                "method": "soka_low_rank_svd",
                "fixed_rank": self.rank,
                "energy_threshold": self.energy_threshold,
                "layer_errors": self.layer_errors,
                "target_attention_only": self.target_attention_only,
                "fallback_to_original": self.fallback_to_original,
                "max_error": self.max_error,
            }
        )
        self.stats = stats
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: SoKA Low-Rank Linear Compressor. Replaces dense nn.Linear and Conv1D layers with factorized LowRankLinear layers.

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """Replaces target Linear and Conv1D modules with LowRankLinear."""
    self.compressed_layers = []
    self.layer_errors = {}

    linear_layers: list[tuple[str, nn.Module]] = []
    for name, module in self.model.named_modules():
        if self._is_target_module(name, module):
            linear_layers.append((name, module))

    if not linear_layers:
        logger.warning("SoKA: No compressible linear layers found.")
        return self.model, self.get_compression_stats()

    for name, module in linear_layers:
        base_mod = getattr(module, "base_layer", module)
        raw_w = base_mod.weight.data
        cls_name = base_mod.__class__.__name__
        is_conv1d = (
            "conv1d" in cls_name.lower()
            or hasattr(base_mod, "nf")
            or (
                raw_w.ndim == 2
                and hasattr(base_mod, "in_features")
                and base_mod.in_features == raw_w.shape[0]
                and raw_w.shape[0] < raw_w.shape[1]
            )
        )

        W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
        device = raw_w.device
        dtype = raw_w.dtype
        d_out, d_in = W.shape

        try:
            U, S, Vh = torch.linalg.svd(W, full_matrices=False)
            r = self._select_rank(S, d_in, d_out)
            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))

            new_layer = LowRankLinear(
                in_features=d_in,
                out_features=d_out,
                rank=r,
                bias=(base_mod.bias is not None),
                device=device,
                dtype=dtype,
            )

            # Assign factorized weights: W ≈ (U_r sqrt(S)) @ (sqrt(S) V_r^T)
            # U: [d_out, r], V: [d_in, r]
            new_layer.U.data = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
            new_layer.V.data = (Vh[:r, :].T * sqrt_S.unsqueeze(0)).to(
                dtype=dtype, device=device
            )

            if base_mod.bias is not None:
                new_layer.bias.data.copy_(base_mod.bias.data)

            # Frobenius reconstruction error check
            W_rec = new_layer.weight.float()
            norm_orig = torch.norm(W, p="fro")
            rel_err = (torch.norm(W - W_rec, p="fro") / (norm_orig + 1e-8)).item()
            self.layer_errors[name] = rel_err

            if self.fallback_to_original and rel_err > self.max_error:
                logger.warning(
                    "SoKA reconstruction error %.4f on %s exceeded threshold %.4f. Skipping.",
                    rel_err,
                    name,
                    self.max_error,
                )
                for p in module.parameters():
                    p.requires_grad = False
                continue

            compressed_layer = new_layer
            # Preserve existing LoRA adapter weights if the source module has them
            if (
                hasattr(module, "lora_A")
                and hasattr(module, "lora_B")
                and module.lora_A is not None
                and module.lora_B is not None
            ):
                src_lA = getattr(module.lora_A, "weight", module.lora_A)
                src_lB = getattr(module.lora_B, "weight", module.lora_B)
                if hasattr(new_layer, "lora_A") and new_layer.lora_A is not None:
                    if hasattr(src_lA, "data") and src_lA.shape == new_layer.lora_A.shape:
                        new_layer.lora_A.data.copy_(src_lA.data)
                if hasattr(new_layer, "lora_B") and new_layer.lora_B is not None:
                    if hasattr(src_lB, "data") and src_lB.shape == new_layer.lora_B.shape:
                        new_layer.lora_B.data.copy_(src_lB.data)

            if hasattr(module, "base_layer"):
                module.base_layer = compressed_layer  # put compressed layer inside wrapper
                # module itself (LoRALinear wrapper) remains the replacement
                new_module = module
            else:
                new_module = compressed_layer
            if not hasattr(module, "base_layer"):
                if name == "":
                    self.model = new_module
                else:
                    replace_module(self.model, name, new_module)
            self.compressed_layers.append(
                {
                    "name": name,
                    "original_shape": (d_out, d_in),
                    "rank": r,
                    "error": rel_err,
                }
            )
        except Exception as e:
            logger.debug("SVD decomposition error in %s: %s", name, e)
            for p in module.parameters():
                p.requires_grad = False

    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    stats.update(
        {
            "total_layers": len(self.compressed_layers),
            "method": "soka_low_rank_svd",
            "fixed_rank": self.rank,
            "energy_threshold": self.energy_threshold,
            "layer_errors": self.layer_errors,
            "target_attention_only": self.target_attention_only,
            "fallback_to_original": self.fallback_to_original,
            "max_error": self.max_error,
        }
    )
    self.stats = stats
    return self.model, stats

Replaces target Linear and Conv1D modules with LowRankLinear.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class LowRankLinear (in_features: int, out_features: int, rank: int, bias: bool = True, device: torch.device | None = None, dtype: torch.dtype | None = None)

Expand source code

class LowRankLinear(nn.Module):
    """
    Factorized Linear Layer executing as y = (x @ V) @ U^T + b.
    Reduces parameter storage from (d_in * d_out) to r * (d_in + d_out).
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        bias: bool = True,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = max(1, min(rank, in_features, out_features))

        # V: [in_features, rank] -> maps x to bottleneck
        # U: [out_features, rank] -> maps bottleneck to output
        self.V = nn.Parameter(torch.empty((in_features, self.rank), device=device, dtype=dtype))
        self.U = nn.Parameter(torch.empty((out_features, self.rank), device=device, dtype=dtype))

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, device=device, dtype=dtype))
        else:
            self.register_parameter("bias", None)

        nn.init.normal_(self.V, std=0.02)
        nn.init.normal_(self.U, std=0.02)

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix for compatibility."""
        return torch.matmul(self.U, self.V.t().contiguous())

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [..., in_features]
        bottleneck = torch.matmul(x, self.V)  # [..., rank]
        out = torch.matmul(bottleneck, self.U.T)  # [..., out_features]
        if self.bias is not None:
            out = out + self.bias
        return out

Factorized Linear Layer executing as y = (x @ V) @ U^T + b. Reduces parameter storage from (d_in * d_out) to r * (d_in + d_out).

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix for compatibility."""
    return torch.matmul(self.U, self.V.t().contiguous())

Reconstructed effective weight matrix for compatibility.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    # x: [..., in_features]
    bottleneck = torch.matmul(x, self.V)  # [..., rank]
    out = torch.matmul(bottleneck, self.U.T)  # [..., out_features]
    if self.bias is not None:
        out = out + self.bias
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class MGIECompressor (model: nn.Module, rank: int | None = None, max_error: float = 0.5, fallback_to_original: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, skip_attention_qkv: bool = False, target_attention_only: bool = False, min_rank: int = 4, target_modules: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("mgie")
@COMPRESSOR_REGISTRY.register("mgie_svd")
class MGIECompressor(BaseCompressor):
    """
    MGIE-SVD Compressor (Maximum Information Entropy SVD).
    """

    # MGIE-FIX-01
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module,
        rank: int | None = None,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        skip_attention_qkv: bool = False,
        target_attention_only: bool = False,
        min_rank: int = 4,
        target_modules: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="mgie", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))

        if cfg is not None:
            rank = _cfg_get(cfg, "mgie_rank", _cfg_get(cfg, "rank", rank))
            max_error = _cfg_get(cfg, "mgie_max_error", _cfg_get(cfg, "max_error", max_error))
            fallback_to_original = _cfg_get(
                cfg,
                "mgie_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
            lora_r = _cfg_get(cfg, "lora_r", lora_r)
            lora_alpha = _cfg_get(cfg, "lora_alpha", lora_alpha)
            skip_attention_qkv = _cfg_get(cfg, "skip_attention_qkv", skip_attention_qkv)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)
            min_rank = _cfg_get(cfg, "min_rank", min_rank)

        if "rank" in kwargs and kwargs["rank"] is not None:
            rank = kwargs["rank"]

        self.rank = rank
        if max_error is None or max_error == 0.50:
            max_error = adaptive_max_error(self.rank, 768, 768, "distilbert")

        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.skip_attention_qkv = bool(skip_attention_qkv)
        self.target_attention_only = bool(target_attention_only)
        self.min_rank = max(1, int(min_rank))
        self.target_modules = target_modules

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.stats: dict[str, Any] = {}

    # ------------------------------------------------------------------
    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            name_lower = name.lower()
            if any(
                p in name_lower
                for p in (
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                )
            ):
                return True
        return False

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        # MGIE-FIX-07: skip already-compressed SVDLinear
        if isinstance(module, SVDLinear):
            return False
        name_l = name.lower()
        if any(
            k in name_l
            for k in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "layernorm",
                "rmsnorm",
                "batchnorm",
            )
        ):
            return False

        cls_name = module.__class__.__name__.lower()
        is_linear_like = (
            isinstance(module, nn.Linear)
            or "conv1d" in cls_name
            or hasattr(module, "nf")
            or hasattr(module, "base_layer")
        )
        if not is_linear_like:
            return False

        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self._is_gqa_asymmetric(name, module):
            return False

        if self.target_attention_only:
            if not any(k in name_l for k in _ATTN_NAMES):
                return False

        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    # ------------------------------------------------------------------
    def _compute_entropy_rank(self, S: torch.Tensor, max_rank: int) -> int:
        """
        Compute rank from singular entropy of squared singular values
        (energy distribution).
        """
        if max_rank <= 1:
            return 1

        energy = S**2
        total_energy = energy.sum()
        if total_energy <= 1e-12:
            return max(1, min(self.min_rank, max_rank))

        p = torch.clamp(energy / total_energy, min=1e-12, max=1.0)
        entropy = -torch.sum(p * torch.log(p)).item()
        max_entropy = math.log(float(max_rank))

        normalized_h = entropy / max(1e-6, max_entropy)
        allocated_rank = max(
            self.min_rank,
            int(round(normalized_h * max_rank)),
        )

        if self.rank is not None:
            allocated_rank = min(allocated_rank, self.rank)

        return max(1, min(allocated_rank, max_rank))

    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}

        seen_weights: set[int] = set()

        # MGIE-FIX-04: compute original parameter count properly.
        total_orig_params = sum(p.numel() for p in self.model.parameters())

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            target_layer = getattr(module, "base_layer", module)
            if not hasattr(target_layer, "weight") or target_layer.weight is None:
                continue

            w_id = id(target_layer.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            raw_w = target_layer.weight.data
            cls_name = target_layer.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_layer, "nf")

            W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
            d_out, d_in = W.shape
            max_r = min(d_out, d_in)
            device = raw_w.device
            dtype = raw_w.dtype

            try:
                U, S, Vt = torch.linalg.svd(W, full_matrices=False)
                r = self._compute_entropy_rank(S, max_r)

                sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
                U_r = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
                V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(dtype=dtype, device=device)

                W_approx = torch.nan_to_num(U_r @ V_r, nan=0.0)
                w_norm = torch.norm(W, p="fro").clamp_min(1e-12).item()
                err = (torch.norm(W.to(device) - W_approx, p="fro") / w_norm).item()
                self.layer_errors[name] = float(err)
                self.layer_ranks[name] = r

                if not math.isfinite(err) or err > self.max_error:
                    logger.warning(
                        "MGIE: reconstruction error %.4f on %s exceeded %.4f.",
                        err,
                        name,
                        self.max_error,
                    )
                    if self.fallback_to_original:
                        self.fallback_layers.append(name)
                        for p in module.parameters():
                            p.requires_grad = False
                        continue

                existing_lora_r = getattr(module, "r", getattr(module, "lora_r", 0))
                eff_lora_r = existing_lora_r if existing_lora_r > 0 else self.lora_r
                eff_lora_alpha = getattr(
                    module,
                    "lora_alpha",
                    getattr(module, "alpha", self.lora_alpha),
                )

                svd_layer = SVDLinear(
                    in_features=d_in,
                    out_features=d_out,
                    rank=r,
                    bias=target_layer.bias is not None,
                    dtype=dtype,
                    device=device,
                    lora_r=eff_lora_r,
                    lora_alpha=eff_lora_alpha,
                )

                # MGIE-FIX-08: verify API availability
                if not hasattr(svd_layer, "set_low_rank_factors"):
                    raise AttributeError(
                        "SVDLinear does not implement set_low_rank_factors; "
                        "MGIE requires an updated aa_svd_compressor."
                    )
                svd_layer.set_low_rank_factors(U_r, V_r)

                # MGIE-FIX-05: bias copy under no_grad
                if target_layer.bias is not None and svd_layer.bias is not None:
                    with torch.no_grad():
                        svd_layer.bias.data.copy_(
                            target_layer.bias.data.to(
                                device=svd_layer.bias.device,
                                dtype=svd_layer.bias.dtype,
                            )
                        )

                # Preserve pre-existing LoRA adapter parameters if wrapped
                if (
                    hasattr(module, "lora_A")
                    and hasattr(module, "lora_B")
                    and svd_layer.lora_A is not None
                    and svd_layer.lora_B is not None
                ):
                    with torch.no_grad():
                        src_A = (
                            module.lora_A.weight.data
                            if hasattr(module.lora_A, "weight")
                            else module.lora_A.data
                        )
                        src_B = (
                            module.lora_B.weight.data
                            if hasattr(module.lora_B, "weight")
                            else module.lora_B.data
                        )
                        if (
                            src_A.shape == svd_layer.lora_A.shape
                            and src_B.shape == svd_layer.lora_B.shape
                        ):
                            svd_layer.lora_A.copy_(
                                src_A.to(
                                    svd_layer.lora_A.device,
                                    dtype=svd_layer.lora_A.dtype,
                                )
                            )
                            svd_layer.lora_B.copy_(
                                src_B.to(
                                    svd_layer.lora_B.device,
                                    dtype=svd_layer.lora_B.dtype,
                                )
                            )
                    svd_layer.lora_A.requires_grad = True
                    svd_layer.lora_B.requires_grad = True

                # MGIE-FIX-02: preserve LoRA wrapper.
                if hasattr(module, "base_layer"):
                    module.base_layer = svd_layer
                else:
                    replace_module(self.model, name, svd_layer)

                self.compressed_layers.append(name)

            except Exception as e:
                logger.warning("MGIE: failed on %s: %s", name, e)
                self.fallback_layers.append(name)
                self.layer_errors[name] = 1.0
                for p in module.parameters():
                    p.requires_grad = False

        # --------------------------------------------------------------
        # Forward sanity check (MGIE-FIX-03)
        # --------------------------------------------------------------
        forward_check_success = True
        try:
            with torch.no_grad():
                dev = next(self.model.parameters(), torch.tensor(0.0)).device
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                try:
                    self.model(input_ids=dummy)
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception:
                        cfg = getattr(self.model, "config", None)
                        hidden = getattr(cfg, "hidden_size", None) if cfg else None
                        if hidden is None:
                            first_lin = next(
                                (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                                None,
                            )
                            hidden = getattr(first_lin, "in_features", 8)
                        dt = next(self.model.parameters()).dtype
                        dummy_feat = torch.randn(1, 8, hidden, device=dev, dtype=dt)
                        self.model(dummy_feat)
        except Exception as e:
            logger.warning("MGIE forward sanity check failed: %s", e)
            forward_check_success = False

        # --------------------------------------------------------------
        # Stats
        # --------------------------------------------------------------
        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        # MGIE-FIX-04: accurate ratio
        eff_ratio = self._compressed_params / max(1, total_orig_params)

        avg_err = (
            sum(self.layer_errors.values()) / max(1, len(self.layer_errors))
            if self.layer_errors
            else 0.0
        )

        stats = self.get_compression_stats()
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "avg_reconstruction_error": avg_err,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "total_params_before": total_orig_params,
                "total_params_after": self._compressed_params,
                "effective_parameter_ratio": round(eff_ratio, 4),
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "forward_check_success": forward_check_success,
                "method": "mgie_svd",
            }
        )
        self.stats = stats
        logger.info(
            "Compressor mgie finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

MGIE-SVD Compressor (Maximum Information Entropy SVD).

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class MGIESVDBudgetCompressor (model: nn.Module, total_rank_budget: int = 1000, rank: int | None = None, max_error: float = 0.5, fallback_to_original: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, skip_attention_qkv: bool = False, target_attention_only: bool = False, min_rank: int = 4, target_modules: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("mgie_budget")
class MGIESVDBudgetCompressor(BaseCompressor):
    """
    MGIE-SVD Budget Compressor: Global Entropy-Proportional Rank Allocation.
    """

    # MGIE-BUDGET-FIX-01
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module,
        total_rank_budget: int = 1000,
        rank: int | None = None,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        skip_attention_qkv: bool = False,
        target_attention_only: bool = False,
        min_rank: int = 4,
        target_modules: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="mgie_budget", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))

        if cfg is not None:
            total_rank_budget = _cfg_get(
                cfg,
                "total_rank_budget",
                _cfg_get(cfg, "rank_budget", total_rank_budget),
            )
            rank = _cfg_get(cfg, "mgie_rank", _cfg_get(cfg, "rank", rank))
            max_error = _cfg_get(cfg, "mgie_max_error", _cfg_get(cfg, "max_error", max_error))
            fallback_to_original = _cfg_get(
                cfg,
                "mgie_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
            lora_r = _cfg_get(cfg, "lora_r", lora_r)
            lora_alpha = _cfg_get(cfg, "lora_alpha", lora_alpha)
            skip_attention_qkv = _cfg_get(cfg, "skip_attention_qkv", skip_attention_qkv)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)
            min_rank = _cfg_get(cfg, "min_rank", min_rank)

        if "rank" in kwargs and kwargs["rank"] is not None:
            rank = kwargs["rank"]

        self.total_rank_budget = max(1, int(total_rank_budget))
        self.rank = rank
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.skip_attention_qkv = bool(skip_attention_qkv)
        self.target_attention_only = bool(target_attention_only)
        self.min_rank = max(1, int(min_rank))
        self.target_modules = target_modules

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.stats: dict[str, Any] = {}

    # ------------------------------------------------------------------
    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            name_lower = name.lower()
            if any(
                p in name_lower
                for p in (
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                )
            ):
                return True
        return False

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        # MGIE-BUDGET-FIX-07: skip already-compressed SVDLinear
        if isinstance(module, SVDLinear):
            return False
        name_l = name.lower()
        if any(
            k in name_l
            for k in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "layernorm",
                "rmsnorm",
                "batchnorm",
            )
        ):
            return False

        cls_name = module.__class__.__name__.lower()
        is_linear_like = (
            isinstance(module, nn.Linear)
            or "conv1d" in cls_name
            or hasattr(module, "nf")
            or hasattr(module, "base_layer")
        )
        if not is_linear_like:
            return False

        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self._is_gqa_asymmetric(name, module):
            return False

        if self.target_attention_only:
            if not any(k in name_l for k in _ATTN_NAMES):
                return False

        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    # ------------------------------------------------------------------
    @staticmethod
    def _compute_entropy(
        W: torch.Tensor,
    ) -> tuple[float, torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Compute singular entropy on squared singular values (energy distribution).
        Returns (entropy, U, S, Vt) in FP32.
        """
        U, S, Vt = torch.linalg.svd(W.float(), full_matrices=False)
        energy = S**2
        total_energy = energy.sum()
        if total_energy <= 1e-12:
            return 0.0, U, S, Vt
        p = torch.clamp(energy / total_energy, min=1e-12, max=1.0)
        entropy = -torch.sum(p * torch.log(p)).item()
        return max(0.0, float(entropy)), U, S, Vt

    # ------------------------------------------------------------------
    def _allocate_ranks(
        self,
        entropies: list[float],
        max_ranks: list[int],
    ) -> list[int]:
        """
        Allocate ranks proportional to entropy, then enforce total_rank_budget
        (MGIE-BUDGET-FIX-06).
        """
        n = len(entropies)
        if n == 0:
            return []

        # Fallback if all entropies ≈ 0.
        total_ent = sum(max(e, 1e-4) for e in entropies)
        if total_ent <= 1e-9:
            per = max(1, self.total_rank_budget // max(1, n))
            return [max(1, min(per, mr)) for mr in max_ranks]

        # Initial allocation.
        ranks = []
        for e, mr in zip(entropies, max_ranks):
            frac = max(e, 1e-4) / total_ent
            r = int(round(frac * self.total_rank_budget))
            r = max(self.min_rank, min(r, mr))
            ranks.append(r)

        # Enforce budget: reduce the largest ranks first.
        # We never go below min_rank (or 1 if max_rank < min_rank).
        total = sum(ranks)
        while total > self.total_rank_budget:
            # pick the layer with the largest current rank that can be reduced
            candidates = [
                i for i in range(n) if ranks[i] > max(1, min(self.min_rank, max_ranks[i]))
            ]
            if not candidates:
                break
            idx = max(candidates, key=lambda i: ranks[i])
            ranks[idx] -= 1
            total = sum(ranks)

        # If under budget, try to increase ranks of layers with largest
        # entropy-per-bit deficit (simple heuristic: increase the layer with
        # the largest max_rank - rank, weighted by entropy).
        if total < self.total_rank_budget:
            order = sorted(
                range(n),
                key=lambda i: entropies[i] / max(1, ranks[i]),
                reverse=True,
            )
            while total < self.total_rank_budget:
                progressed = False
                for i in order:
                    if ranks[i] < max_ranks[i]:
                        ranks[i] += 1
                        total += 1
                        progressed = True
                        if total >= self.total_rank_budget:
                            break
                if not progressed:
                    break

        # Final clamp.
        ranks = [max(1, min(r, max_ranks[i])) for i, r in enumerate(ranks)]
        return ranks

    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}

        seen_weights: set[int] = set()

        total_orig_params = sum(p.numel() for p in self.model.parameters())

        # ------------------------------------------------------------------
        # 1. Collect target layers.
        # ------------------------------------------------------------------
        target_layers: list[tuple[str, nn.Module, nn.Module]] = []
        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue
            base = getattr(module, "base_layer", module)
            if not hasattr(base, "weight") or base.weight is None:
                continue
            w_id = id(base.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)
            target_layers.append((name, module, base))

        if not target_layers:
            logger.warning("MGIE-Budget: no compressible linear layers found.")
            stats = self.get_compression_stats()
            stats.update(
                {
                    "num_compressed": 0,
                    "num_fallback": 0,
                    "num_skipped": len(self.skipped_layers),
                    "method": "mgie_budget",
                }
            )
            self.stats = stats
            return self.model, stats

        # ------------------------------------------------------------------
        # 2. Compute per-layer entropy and SVD.
        # ------------------------------------------------------------------
        layer_data: list[dict[str, Any]] = []
        entropies: list[float] = []
        max_ranks: list[int] = []

        for name, module, target_layer in target_layers:
            raw_w = target_layer.weight.data
            cls_name = target_layer.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_layer, "nf")

            W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
            d_out, d_in = W.shape
            max_r = min(d_out, d_in)

            ent, U, S, Vt = self._compute_entropy(W)
            layer_data.append(
                {
                    "name": name,
                    "module": module,
                    "target_layer": target_layer,
                    "W": W,
                    "U": U,
                    "S": S,
                    "Vt": Vt,
                    "is_conv1d": is_conv1d,
                    "d_in": d_in,
                    "d_out": d_out,
                    "max_r": max_r,
                }
            )
            entropies.append(ent)
            max_ranks.append(max_r)

        # ------------------------------------------------------------------
        # 3. Allocate ranks (with budget enforcement).
        # ------------------------------------------------------------------
        if self.rank is not None:
            # Explicit rank overrides budget.
            for info in layer_data:
                info["rank"] = max(1, min(self.rank, info["max_r"]))
        else:
            ranks = self._allocate_ranks(entropies, max_ranks)
            for info, r in zip(layer_data, ranks):
                info["rank"] = max(1, min(r, info["max_r"]))

        # ------------------------------------------------------------------
        # 4. Apply compression.
        # ------------------------------------------------------------------
        for info in layer_data:
            name = info["name"]
            module = info["module"]
            target_layer = info["target_layer"]
            W = info["W"]
            U = info["U"]
            S = info["S"]
            Vt = info["Vt"]
            r = info["rank"]
            d_in = info["d_in"]
            d_out = info["d_out"]
            device = target_layer.weight.device
            dtype = target_layer.weight.dtype

            try:
                sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
                U_r = (U[:, :r] * sqrt_S.unsqueeze(0)).to(device=device, dtype=dtype)
                V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(device=device, dtype=dtype)

                W_approx = torch.nan_to_num(U_r @ V_r, nan=0.0)
                w_norm = torch.norm(W, p="fro").clamp_min(1e-12).item()
                err = (torch.norm(W.to(device) - W_approx, p="fro") / w_norm).item()
                self.layer_errors[name] = float(err)
                self.layer_ranks[name] = r

                if not math.isfinite(err) or err > self.max_error:
                    logger.warning(
                        "MGIE-Budget: reconstruction error %.4f on %s exceeded %.4f.",
                        err,
                        name,
                        self.max_error,
                    )
                    if self.fallback_to_original:
                        self.fallback_layers.append(name)
                        for p in module.parameters():
                            p.requires_grad = False
                        continue

                existing_lora_r = getattr(module, "r", getattr(module, "lora_r", 0))
                eff_lora_r = existing_lora_r if existing_lora_r > 0 else self.lora_r
                # MGIE-BUDGET-FIX-08: correct lora_alpha retrieval
                eff_lora_alpha = getattr(module, "lora_alpha", None)
                if eff_lora_alpha is None:
                    eff_lora_alpha = getattr(module, "alpha", self.lora_alpha)
                eff_lora_alpha = float(eff_lora_alpha)

                svd_layer = SVDLinear(
                    in_features=d_in,
                    out_features=d_out,
                    rank=r,
                    bias=(target_layer.bias is not None),
                    dtype=dtype,
                    device=device,
                    lora_r=eff_lora_r,
                    lora_alpha=eff_lora_alpha,
                )

                # MGIE-BUDGET-FIX-09: verify API availability
                if not hasattr(svd_layer, "set_low_rank_factors"):
                    raise AttributeError(
                        "SVDLinear does not implement set_low_rank_factors; "
                        "MGIE-Budget requires an updated aa_svd_compressor."
                    )
                svd_layer.set_low_rank_factors(U_r, V_r)

                # MGIE-BUDGET-FIX-05: bias copy under no_grad.
                if target_layer.bias is not None and svd_layer.bias is not None:
                    with torch.no_grad():
                        svd_layer.bias.data.copy_(
                            target_layer.bias.data.to(
                                device=svd_layer.bias.device,
                                dtype=svd_layer.bias.dtype,
                            )
                        )

                # Preserve pre-existing LoRA adapter parameters if wrapped.
                if (
                    hasattr(module, "lora_A")
                    and hasattr(module, "lora_B")
                    and svd_layer.lora_A is not None
                    and svd_layer.lora_B is not None
                ):
                    with torch.no_grad():
                        src_A = (
                            module.lora_A.weight.data
                            if hasattr(module.lora_A, "weight")
                            else module.lora_A.data
                        )
                        src_B = (
                            module.lora_B.weight.data
                            if hasattr(module.lora_B, "weight")
                            else module.lora_B.data
                        )
                        if (
                            src_A.shape == svd_layer.lora_A.shape
                            and src_B.shape == svd_layer.lora_B.shape
                        ):
                            svd_layer.lora_A.copy_(
                                src_A.to(
                                    svd_layer.lora_A.device,
                                    dtype=svd_layer.lora_A.dtype,
                                )
                            )
                            svd_layer.lora_B.copy_(
                                src_B.to(
                                    svd_layer.lora_B.device,
                                    dtype=svd_layer.lora_B.dtype,
                                )
                            )
                    svd_layer.lora_A.requires_grad = True
                    svd_layer.lora_B.requires_grad = True

                # MGIE-BUDGET-FIX-02: preserve LoRA wrapper.
                if hasattr(module, "base_layer"):
                    module.base_layer = svd_layer
                else:
                    replace_module(self.model, name, svd_layer)

                self.compressed_layers.append(name)

            except Exception as e:
                logger.warning("MGIE-Budget: failed on %s: %s", name, e)
                self.fallback_layers.append(name)
                self.layer_errors[name] = 1.0
                for p in module.parameters():
                    p.requires_grad = False

        # ------------------------------------------------------------------
        # 5. Forward sanity check (MGIE-BUDGET-FIX-03)
        # ------------------------------------------------------------------
        forward_check_success = True
        try:
            with torch.no_grad():
                dev = next(self.model.parameters(), torch.tensor(0.0)).device
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                try:
                    self.model(input_ids=dummy)
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception:
                        cfg = getattr(self.model, "config", None)
                        hidden = getattr(cfg, "hidden_size", None) if cfg else None
                        if hidden is None:
                            first_lin = next(
                                (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                                None,
                            )
                            hidden = getattr(first_lin, "in_features", 8)
                        dt = next(self.model.parameters()).dtype
                        dummy_feat = torch.randn(1, 8, hidden, device=dev, dtype=dt)
                        self.model(dummy_feat)
        except Exception as e:
            logger.warning("MGIE-Budget forward sanity check failed: %s", e)
            forward_check_success = False

        # ------------------------------------------------------------------
        # 6. Stats (MGIE-BUDGET-FIX-04)
        # ------------------------------------------------------------------
        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        eff_ratio = self._compressed_params / max(1, total_orig_params)

        avg_err = (
            sum(self.layer_errors.values()) / max(1, len(self.layer_errors))
            if self.layer_errors
            else 0.0
        )

        allocated_ranks = [self.layer_ranks.get(name, 0) for name in self.layer_ranks]
        total_allocated = sum(allocated_ranks)

        stats = self.get_compression_stats()
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "avg_reconstruction_error": avg_err,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "total_rank_budget": self.total_rank_budget,
                "total_allocated_rank": total_allocated,
                "rank_budget_utilization_pct": round(
                    100.0 * total_allocated / max(1, self.total_rank_budget), 2
                ),
                "total_params_before": total_orig_params,
                "total_params_after": self._compressed_params,
                "effective_parameter_ratio": round(eff_ratio, 4),
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "forward_check_success": forward_check_success,
                "method": "mgie_budget",
            }
        )
        self.stats = stats
        return self.model, stats

MGIE-SVD Budget Compressor: Global Entropy-Proportional Rank Allocation.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class MLoRQCompressor (model: nn.Module, total_memory_budget_mb: float = 500, bits_candidates: list[int] | None = None, max_error: float = 0.5, fallback_to_original: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, skip_attention_qkv: bool = False, target_attention_only: bool = False, min_rank: int = 4, target_modules: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("mlorq")
@COMPRESSOR_REGISTRY.register("mlorq_compressor")
class MLoRQCompressor(BaseCompressor):
    """
    MLoRQ Compressor: Multi-Layer Optimal Rank and Quantization under Memory Budget.
    """

    is_custom_attention: bool = False  # MLR-FIX-01

    def __init__(
        self,
        model: nn.Module,
        total_memory_budget_mb: float = 500,
        bits_candidates: list[int] | None = None,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        skip_attention_qkv: bool = False,
        target_attention_only: bool = False,
        min_rank: int = 4,
        target_modules: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="mlorq", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            total_memory_budget_mb = _cfg_get(
                cfg,
                "mlorq_budget_mb",
                _cfg_get(cfg, "total_memory_budget_mb", total_memory_budget_mb),
            )
            bits_candidates = _cfg_get(
                cfg,
                "mlorq_bits_candidates",
                _cfg_get(cfg, "bits_candidates", bits_candidates),
            )
            max_error = _cfg_get(cfg, "mlorq_max_error", _cfg_get(cfg, "max_error", max_error))
            fallback_to_original = _cfg_get(
                cfg,
                "mlorq_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
            lora_r = _cfg_get(cfg, "lora_r", lora_r)
            lora_alpha = _cfg_get(cfg, "lora_alpha", lora_alpha)
            skip_attention_qkv = _cfg_get(cfg, "skip_attention_qkv", skip_attention_qkv)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)
            min_rank = _cfg_get(cfg, "min_rank", min_rank)

        self.total_budget_bits = float(total_memory_budget_mb) * 8 * 1024 * 1024  # bits
        self.bits_candidates = sorted(set(int(b) for b in (bits_candidates or [4, 8])))
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.skip_attention_qkv = bool(skip_attention_qkv)
        self.target_attention_only = bool(target_attention_only)
        self.min_rank = max(1, int(min_rank))
        self.target_modules = target_modules

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.layer_bits: dict[str, int] = {}
        self.stats: dict[str, Any] = {}

    # ------------------------------------------------------------------
    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            name_lower = name.lower()
            if any(
                p in name_lower
                for p in (
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                )
            ):
                return True
        return False

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        if any(
            k in name_l
            for k in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "layernorm",
                "rmsnorm",
                "batchnorm",
            )
        ):
            return False
        if isinstance(module, (nn.LayerNorm, nn.RMSNorm, nn.Embedding, MLoRQLoRALinear)):
            return False

        cls_name = module.__class__.__name__.lower()
        is_linear_like = (
            isinstance(module, nn.Linear)
            or "conv1d" in cls_name
            or hasattr(module, "nf")
            or hasattr(module, "base_layer")
        )
        if not is_linear_like:
            return False

        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self._is_gqa_asymmetric(name, module):
            return False

        if self.target_attention_only and not any(k in name_l for k in _ATTN_NAMES):
            return False

        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    # ------------------------------------------------------------------
    @staticmethod
    def _quantize_float(t: torch.Tensor, bits: int) -> torch.Tensor:
        """Non-packed quantization for error estimation."""
        if bits >= 32:
            return t.float()
        tf = t.float()
        max_abs = tf.abs().max().clamp_min(1e-8)
        if bits >= 8:
            scale = max_abs / 127.0
            return torch.round(tf / scale).clamp(-128, 127) * scale
        qmax = (1 << (bits - 1)) - 1
        scale = max_abs / float(qmax)
        return torch.round(tf / scale).clamp(-qmax - 1, qmax) * scale

    # ------------------------------------------------------------------
    def _generate_candidates(self, W: torch.Tensor) -> list[dict[str, Any]]:
        """Generate (rank, bits, error, cost_bits, U, S, Vt) candidates for a weight."""
        out_f, in_f = W.shape
        max_rank = min(out_f, in_f)

        # Randomized SVD for large matrices (MLR-FIX-09)
        q = min(max_rank + 10, max_rank)
        if W.numel() > 1_000_000:
            try:
                U, S, V = torch.svd_lowrank(W.float(), q=q, niter=4)
                Vt = V.t().contiguous()
            except Exception:
                U, S, Vt = torch.linalg.svd(W.float(), full_matrices=False)
        else:
            U, S, Vt = torch.linalg.svd(W.float(), full_matrices=False)

        # Rank list (unique, <= max_rank)
        ranks = set()
        ranks.add(min(self.min_rank, max_rank))
        for frac in (8, 4, 2):
            r = max(self.min_rank, max_rank // frac)
            if r <= max_rank:
                ranks.add(r)
        ranks.add(max_rank)
        ranks = sorted(r for r in ranks if r >= 1 and r <= max_rank)

        candidates = []
        for bits in self.bits_candidates:
            for r in ranks:
                if r > U.shape[1]:
                    continue
                U_r = U[:, :r]
                S_r = S[:r]
                Vt_r = Vt[:r, :]

                U_q = self._quantize_float(U_r, bits)
                S_q = self._quantize_float(S_r, bits)
                Vt_q = self._quantize_float(Vt_r, bits)

                W_rec = U_q @ torch.diag(S_q) @ Vt_q
                w_norm = torch.norm(W, p="fro").clamp_min(1e-12).item()
                err = float((torch.norm(W - W_rec, p="fro") / w_norm).item())

                # cost in bits for packed storage
                cost_bits = (out_f * r + r * in_f + r) * bits
                candidates.append(
                    {
                        "rank": r,
                        "bits": bits,
                        "error": err,
                        "cost_bits": cost_bits,
                        "U": U_r,
                        "S": S_r,
                        "Vt": Vt_r,
                    }
                )
        candidates.sort(key=lambda c: c["cost_bits"])
        return candidates

    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}
        self.layer_bits = {}
        seen_weights: set[int] = set()

        # ---------- 1. Collect target layers ----------
        layer_infos: list[dict[str, Any]] = []
        total_orig_bits = 0

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            base = getattr(module, "base_layer", module)
            w = getattr(base, "weight", None)
            if w is None:
                continue
            w_id = id(w)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            raw_w = w.data
            orig_bits = raw_w.numel() * raw_w.element_size() * 8
            total_orig_bits += orig_bits

            cls_name = base.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(base, "nf")
            W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()

            layer_infos.append(
                {
                    "name": name,
                    "module": module,
                    "base": base,
                    "W": W,
                    "shape": W.shape,
                    "is_conv1d": is_conv1d,
                    "device": raw_w.device,
                    "dtype": raw_w.dtype,
                    "orig_bits": orig_bits,
                }
            )

        if not layer_infos:
            logger.warning("MLoRQ: no target layers found.")
            stats = self.get_compression_stats()
            stats.update({"num_compressed": 0, "method": "mlorq"})
            self.stats = stats
            return self.model, stats

        # ---------- 2. Generate candidates ----------
        for info in layer_infos:
            cands = self._generate_candidates(info["W"])
            # Filter by max_error
            valid = [c for c in cands if c["error"] <= self.max_error]
            if not valid:
                # fallback original
                info["candidates"] = []
                info["chosen"] = None
                info["current_cost"] = info["orig_bits"]
                info["current_error"] = 0.0
            else:
                info["candidates"] = valid
                info["chosen"] = valid[0]  # cheapest valid
                info["current_cost"] = valid[0]["cost_bits"]
                info["current_error"] = valid[0]["error"]

        # ---------- 3. Greedy budget allocation (MLR-FIX-04) ----------
        total_cost = sum(info["current_cost"] for info in layer_infos)
        budget = self.total_budget_bits

        # Downgrade if over budget
        while total_cost > budget:
            best_info = None
            best_ratio = float("inf")
            best_new = None
            for info in layer_infos:
                if info["chosen"] is None:
                    continue
                # find cheaper candidate than current
                cheaper = [c for c in info["candidates"] if c["cost_bits"] < info["current_cost"]]
                if not cheaper:
                    continue
                # choose the cheapest cheaper candidate? Actually we want the one with smallest error increase per bit saved.
                # We can evaluate each cheaper candidate, but for simplicity pick the next cheaper (max cost < current).
                next_cheaper = max(cheaper, key=lambda c: c["cost_bits"])
                delta_cost = info["current_cost"] - next_cheaper["cost_bits"]
                delta_err = next_cheaper["error"] - info["current_error"]
                if delta_cost <= 0:
                    continue
                ratio = delta_err / delta_cost
                if ratio < best_ratio:
                    best_ratio = ratio
                    best_info = info
                    best_new = next_cheaper
            if best_info is None:
                break
            # apply downgrade
            best_info["chosen"] = best_new
            best_info["current_cost"] = best_new["cost_bits"]
            best_info["current_error"] = best_new["error"]
            total_cost = sum(info["current_cost"] for info in layer_infos)

        # Upgrade if under budget
        while total_cost < budget:
            best_info = None
            best_ratio = -1.0
            best_new = None
            for info in layer_infos:
                if info["chosen"] is None:
                    continue
                # find more expensive candidate
                more = [c for c in info["candidates"] if c["cost_bits"] > info["current_cost"]]
                if not more:
                    continue
                next_more = min(more, key=lambda c: c["cost_bits"])
                delta_cost = next_more["cost_bits"] - info["current_cost"]
                delta_err = info["current_error"] - next_more["error"]
                if delta_cost <= 0:
                    continue
                if total_cost + delta_cost > budget:
                    continue
                ratio = delta_err / delta_cost
                if ratio > best_ratio:
                    best_ratio = ratio
                    best_info = info
                    best_new = next_more
            if best_info is None:
                break
            best_info["chosen"] = best_new
            best_info["current_cost"] = best_new["cost_bits"]
            best_info["current_error"] = best_new["error"]
            total_cost = sum(info["current_cost"] for info in layer_infos)

        # ---------- 4. Apply chosen candidates ----------
        total_comp_bytes = 0
        for info in layer_infos:
            name = info["name"]
            module = info["module"]
            chosen = info["chosen"]

            if chosen is None:
                # fallback original
                self.fallback_layers.append(name)
                total_comp_bytes += info["orig_bits"] // 8
                continue

            try:
                new_layer = MLoRQLoRALinear.from_linear(
                    module,
                    U=chosen["U"],
                    S=chosen["S"],
                    Vt=chosen["Vt"],
                    bits=chosen["bits"],
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                )
                # MLR-FIX-05: preserve wrapper
                if hasattr(module, "base_layer"):
                    module.base_layer = new_layer
                else:
                    replace_module(self.model, name, new_layer)

                self.compressed_layers.append(name)
                self.layer_errors[name] = float(chosen["error"])
                self.layer_ranks[name] = int(chosen["rank"])
                self.layer_bits[name] = int(chosen["bits"])
                total_comp_bytes += new_layer.storage_bytes()
            except Exception as e:
                logger.warning("MLoRQ: failed on %s: %s", name, e)
                self.fallback_layers.append(name)
                total_comp_bytes += info["orig_bits"] // 8

        # ---------- 5. Forward sanity check (MLR-FIX-06) ----------
        forward_check_success = True
        try:
            with torch.no_grad():
                dev = next(self.model.parameters(), torch.tensor(0.0)).device
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                try:
                    self.model(input_ids=dummy)
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception:
                        cfg = getattr(self.model, "config", None)
                        hidden = getattr(cfg, "hidden_size", None) if cfg else None
                        if hidden is None:
                            first_lin = next(
                                (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                                None,
                            )
                            hidden = getattr(first_lin, "in_features", 8)
                        dt = next(self.model.parameters()).dtype
                        dummy_feat = torch.randn(1, 8, hidden, device=dev, dtype=dt)
                        self.model(dummy_feat)
        except Exception as e:
            logger.warning("MLoRQ forward sanity check failed: %s", e)
            forward_check_success = False

        # ---------- 6. Stats ----------
        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        avg_err = (
            sum(self.layer_errors.values()) / max(1, len(self.layer_errors))
            if self.layer_errors
            else 0.0
        )
        storage_ratio = (
            total_orig_bits / 8 / max(1, total_comp_bytes) if total_comp_bytes > 0 else 1.0
        )
        stats = self.get_compression_stats()
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "layer_bits": self.layer_bits,
                "avg_reconstruction_error": avg_err,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "total_budget_bits": self.total_budget_bits,
                "total_used_bits": total_cost,
                "budget_utilization_pct": round(
                    100.0 * total_cost / max(1, self.total_budget_bits), 2
                ),
                "orig_storage_bytes": total_orig_bits // 8,
                "comp_storage_bytes": total_comp_bytes,
                "storage_compression_ratio": round(storage_ratio, 3),
                "storage_reduction_pct": round(
                    100.0 * (1.0 - total_comp_bytes / max(1, total_orig_bits // 8)), 2
                ),
                "forward_check_success": forward_check_success,
                "method": "mlorq",
            }
        )
        self.stats = stats
        return self.model, stats

MLoRQ Compressor: Multi-Layer Optimal Rank and Quantization under Memory Budget.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class MLoRQLoRALinear (in_features: int, out_features: int, rank: int, bits: int, U_packed: torch.Tensor, U_scale: torch.Tensor, U_shape: torch.Size, S_packed: torch.Tensor, S_scale: torch.Tensor, S_shape: torch.Size, Vt_packed: torch.Tensor, Vt_scale: torch.Tensor, Vt_shape: torch.Size, bias: torch.Tensor | None = None, lora_r: int = 0, lora_alpha: float = 1.0, device: torch.device | None = None, dtype: torch.dtype | None = None)

Expand source code

class MLoRQLoRALinear(nn.Module):
    """
    Low-rank quantized linear layer:

        y = (x @ Vt) * S @ U^T + bias + LoRA(x)

    where U, S, Vt are stored as PACKED int8 tensors with per-tensor scales.
    """

    is_custom_attention: bool = False  # MLR-FIX-01

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        bits: int,
        U_packed: torch.Tensor,
        U_scale: torch.Tensor,
        U_shape: torch.Size,
        S_packed: torch.Tensor,
        S_scale: torch.Tensor,
        S_shape: torch.Size,
        Vt_packed: torch.Tensor,
        Vt_scale: torch.Tensor,
        Vt_shape: torch.Size,
        bias: torch.Tensor | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.nf = self.out_features
        self.rank = int(rank)
        self.bits = int(bits)
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        target_device = device if device is not None else U_packed.device
        target_dtype = dtype if dtype is not None else torch.float32

        self.register_buffer("U_packed", U_packed.to(target_device), persistent=True)
        self.register_buffer(
            "U_scale", U_scale.to(target_device, dtype=torch.float32), persistent=True
        )
        self.register_buffer("S_packed", S_packed.to(target_device), persistent=True)
        self.register_buffer(
            "S_scale", S_scale.to(target_device, dtype=torch.float32), persistent=True
        )
        self.register_buffer("Vt_packed", Vt_packed.to(target_device), persistent=True)
        self.register_buffer(
            "Vt_scale", Vt_scale.to(target_device, dtype=torch.float32), persistent=True
        )

        self.U_shape = tuple(U_shape)
        self.S_shape = tuple(S_shape)
        self.Vt_shape = tuple(Vt_shape)

        if bias is not None:
            self.bias = nn.Parameter(
                bias.to(device=target_device, dtype=target_dtype),
                requires_grad=False,
            )
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty(
                    (self.lora_r, self.in_features),
                    dtype=torch.float32,
                    device=target_device,
                )
            )
            self.lora_B = nn.Parameter(
                torch.zeros(
                    (self.out_features, self.lora_r),
                    dtype=torch.float32,
                    device=target_device,
                )
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    # ------------------------------------------------------------------
    def _dequantize(self, dtype: torch.dtype, device: torch.device):
        U = _dequantize_packed(self.U_packed, self.bits, self.U_scale, self.U_shape, device, dtype)
        S = _dequantize_packed(self.S_packed, self.bits, self.S_scale, self.S_shape, device, dtype)
        Vt = _dequantize_packed(
            self.Vt_packed, self.bits, self.Vt_scale, self.Vt_shape, device, dtype
        )
        return U, S, Vt

    @property
    def weight(self) -> torch.Tensor:
        dtype = self.U_scale.dtype
        device = self.U_scale.device
        U, S, Vt = self._dequantize(dtype, device)
        return (U * S.unsqueeze(0)) @ Vt

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        U, S, Vt = self._dequantize(x.dtype, x.device)
        # (B, in) @ (in, r) -> (B, r)
        x_v = F.linear(x, Vt)
        x_v = x_v * S
        out = F.linear(x_v, U)
        if self.bias is not None:
            out = out + self.bias.to(dtype=x.dtype, device=x.device)
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            out = out + self.scaling * F.linear(F.linear(x.float(), self.lora_A), self.lora_B).to(
                out.dtype
            )
        return out

    def storage_bytes(self) -> int:
        total = 0
        for name, buf in self.named_buffers():
            if buf is not None:
                total += buf.numel() * buf.element_size()
        if self.bias is not None:
            total += self.bias.numel() * self.bias.element_size()
        if self.lora_A is not None:
            total += self.lora_A.numel() * self.lora_A.element_size()
        if self.lora_B is not None:
            total += self.lora_B.numel() * self.lora_B.element_size()
        return total

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        U: torch.Tensor,
        S: torch.Tensor,
        Vt: torch.Tensor,
        bits: int,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> MLoRQLoRALinear:
        base_layer = getattr(linear, "base_layer", linear)
        bias = base_layer.bias.detach() if getattr(base_layer, "bias", None) is not None else None
        in_features = int(getattr(base_layer, "in_features", Vt.shape[1]))
        out_features = int(getattr(base_layer, "out_features", U.shape[0]))
        rank = int(min(U.shape[1], Vt.shape[0], S.numel()))

        # quantize + pack
        U_packed, U_scale, U_shape = _quantize_packed(U, bits)
        S_packed, S_scale, S_shape = _quantize_packed(S, bits)
        Vt_packed, Vt_scale, Vt_shape = _quantize_packed(Vt, bits)

        existing_lora_r = getattr(linear, "r", getattr(linear, "lora_r", 0))
        eff_lora_r = existing_lora_r if existing_lora_r > 0 else lora_r
        eff_lora_alpha = getattr(linear, "lora_alpha", getattr(linear, "alpha", lora_alpha))

        layer = cls(
            in_features=in_features,
            out_features=out_features,
            rank=rank,
            bits=bits,
            U_packed=U_packed,
            U_scale=U_scale,
            U_shape=U_shape,
            S_packed=S_packed,
            S_scale=S_scale,
            S_shape=S_shape,
            Vt_packed=Vt_packed,
            Vt_scale=Vt_scale,
            Vt_shape=Vt_shape,
            bias=bias,
            lora_r=eff_lora_r,
            lora_alpha=eff_lora_alpha,
            device=U.device,
            dtype=U.dtype,
        )

        # Preserve existing LoRA weights if present
        if (
            hasattr(linear, "lora_A")
            and hasattr(linear, "lora_B")
            and layer.lora_A is not None
            and layer.lora_B is not None
        ):
            with torch.no_grad():
                src_A = (
                    linear.lora_A.weight.data
                    if hasattr(linear.lora_A, "weight")
                    else linear.lora_A.data
                )
                src_B = (
                    linear.lora_B.weight.data
                    if hasattr(linear.lora_B, "weight")
                    else linear.lora_B.data
                )
                if src_A.shape == layer.lora_A.shape and src_B.shape == layer.lora_B.shape:
                    layer.lora_A.copy_(src_A.to(layer.lora_A.device, dtype=layer.lora_A.dtype))
                    layer.lora_B.copy_(src_B.to(layer.lora_B.device, dtype=layer.lora_B.dtype))
            layer.lora_A.requires_grad = True
            layer.lora_B.requires_grad = True

        return layer

Low-rank quantized linear layer:

y = (x @ Vt) * S @ U^T + bias + LoRA(x)

where U, S, Vt are stored as PACKED int8 tensors with per-tensor scales.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Static methods

def from_linear(linear: nn.Module, U: torch.Tensor, S: torch.Tensor, Vt: torch.Tensor, bits: int, lora_r: int = 0, lora_alpha: float = 1.0) ‑> MLoRQLoRALinear

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    dtype = self.U_scale.dtype
    device = self.U_scale.device
    U, S, Vt = self._dequantize(dtype, device)
    return (U * S.unsqueeze(0)) @ Vt

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    U, S, Vt = self._dequantize(x.dtype, x.device)
    # (B, in) @ (in, r) -> (B, r)
    x_v = F.linear(x, Vt)
    x_v = x_v * S
    out = F.linear(x_v, U)
    if self.bias is not None:
        out = out + self.bias.to(dtype=x.dtype, device=x.device)
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        out = out + self.scaling * F.linear(F.linear(x.float(), self.lora_A), self.lora_B).to(
            out.dtype
        )
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def storage_bytes(self) ‑> int

Expand source code

def storage_bytes(self) -> int:
    total = 0
    for name, buf in self.named_buffers():
        if buf is not None:
            total += buf.numel() * buf.element_size()
    if self.bias is not None:
        total += self.bias.numel() * self.bias.element_size()
    if self.lora_A is not None:
        total += self.lora_A.numel() * self.lora_A.element_size()
    if self.lora_B is not None:
        total += self.lora_B.numel() * self.lora_B.element_size()
    return total

class MambaShedderCompressor (model: nn.Module, tokenizer: Any | None = None, sparsity: float = 0.2, shed_layers: int = 0, max_error: float = 0.2, fallback_to_original: bool = True, validate_after_shed: bool = True, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, structured_24: bool = False, target_attention_only: bool = False, target_modules: list[str] | None = None, max_calibration_batches: int = 8, min_layer_dim: int = 8, skip_should_compress_filter: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("mamba_shedder_class")
@COMPRESSOR_REGISTRY.register("mamba_shedder_compressor_class")
class MambaShedderCompressor(BaseCompressor):
    is_custom_attention: bool = False  # MS-FIX-01

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any | None = None,
        sparsity: float = 0.2,
        shed_layers: int = 0,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        validate_after_shed: bool = True,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        structured_24: bool = False,
        target_attention_only: bool = False,
        target_modules: list[str] | None = None,
        max_calibration_batches: int = 8,
        min_layer_dim: int = 8,
        skip_should_compress_filter: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="mamba_shedder", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))

        if cfg is not None:
            sparsity = _cfg_get(cfg, "mamba_shedder_sparsity", _cfg_get(cfg, "sparsity", sparsity))
            shed_layers = _cfg_get(
                cfg,
                "mamba_shedder_shed_layers",
                _cfg_get(cfg, "shed_layers", shed_layers),
            )
            structured_24 = _cfg_get(
                cfg,
                "mamba_shedder_structured_24",
                _cfg_get(cfg, "structured_24", structured_24),
            )
            max_error = _cfg_get(cfg, "max_error", max_error)
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)
            max_calibration_batches = _cfg_get(
                cfg, "max_calibration_batches", max_calibration_batches
            )
            min_layer_dim = _cfg_get(cfg, "min_layer_dim", min_layer_dim)
            skip_should_compress_filter = _cfg_get(
                cfg,
                "skip_should_compress_filter",
                skip_should_compress_filter,
            )
            lora_r = _cfg_get(cfg, "lora_r", lora_r)
            lora_alpha = _cfg_get(cfg, "lora_alpha", lora_alpha)

        self.tokenizer = (
            tokenizer
            or getattr(model, "tokenizer", None)
            or getattr(getattr(model, "config", None), "tokenizer", None)
        )

        self.sparsity = max(0.0, min(float(sparsity), 0.95))
        self.shed_layers = max(0, int(shed_layers))
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.validate_after_shed = bool(validate_after_shed)
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.structured_24 = bool(structured_24)
        self.target_attention_only = bool(target_attention_only)
        self.target_modules = target_modules
        self.max_calibration_batches = max(1, int(max_calibration_batches))
        self.min_layer_dim = max(1, int(min_layer_dim))
        self.skip_should_compress_filter = bool(skip_should_compress_filter)
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)

        if device is not None:
            self.target_device = torch.device(device)
        else:
            first_param = next(iter(model.parameters()), None)
            self.target_device = (
                first_param.device if first_param is not None else torch.device("cpu")
            )

        self.layer_errors: dict[str, float] = {}
        self.stats: dict[str, Any] = {}
        self._compressed_params: int = 0

    # ------------------------------------------------------------------
    # Target check (MS-FIX-16: cached)
    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module, _debug: bool = False) -> bool:
        name_lower = name.lower()
        if any(part in name.split(".") for part in ("lora_A", "lora_B", "base_layer")):
            if _debug:
                logger.info(f"  [skip] {name}: lora/base_layer part")
            return False

        cls_name = module.__class__.__name__
        if cls_name in (
            "DynamicSVDLinear",
            "SVDLinear",
            "SVDLinearLayer",
            "FactorisedLinear",
            "FactorizedLinear",
            "LowRankPlusSparseLinear",
            "MambaShedderPrunedLinear",
        ):
            if _debug:
                logger.info(f"  [skip] {name}: already compressed ({cls_name})")
            return False

        if not self.skip_should_compress_filter:
            try:
                if not should_compress_linear(self.model, name):
                    if _debug:
                        logger.info(f"  [skip] {name}: should_compress_linear=False")
                    return False
            except Exception as e:
                if _debug:
                    logger.info(f"  [skip] {name}: should_compress_linear raised {e}")
                return False

        base_mod = getattr(module, "base_layer", module)
        cls_lower = base_mod.__class__.__name__.lower()
        is_linear_like = (
            isinstance(base_mod, nn.Linear) or "conv1d" in cls_lower or hasattr(base_mod, "nf")
        )
        if not is_linear_like:
            if _debug:
                logger.info(f"  [skip] {name}: not linear-like")
            return False

        w = getattr(base_mod, "weight", None)
        if w is None:
            if _debug:
                logger.info(f"  [skip] {name}: weight None")
            return False
        if min(w.shape[0], w.shape[1]) < self.min_layer_dim:
            if _debug:
                logger.info(f"  [skip] {name}: dim too small {tuple(w.shape)}")
            return False

        if self.target_attention_only and not any(k in name_lower for k in _ATTN_NAMES):
            if _debug:
                logger.info(f"  [skip] {name}: not attention")
            return False

        if self.target_modules is not None:
            ok = any(t in name for t in self.target_modules)
            if not ok and _debug:
                logger.info(f"  [skip] {name}: not in target_modules")
            return ok

        if _debug:
            logger.info(f"  [TARGET] {name}: shape={tuple(w.shape)}")
        return True

    # ------------------------------------------------------------------
    # Activation collection
    # ------------------------------------------------------------------
    def _collect_activations(self, target_names: list[str]) -> dict[str, torch.Tensor]:
        self.model.eval()
        accum: dict[str, list[torch.Tensor]] = {}
        hooks = []

        def make_hook(layer_name: str):
            def hook(module, inputs, output):
                if inputs and isinstance(inputs[0], torch.Tensor):
                    inp = inputs[0].detach()
                    if inp.dim() >= 2 and inp.shape[0] * inp.shape[1] > 32768:
                        inp = inp.reshape(-1, inp.shape[-1])[:32768].unsqueeze(0)
                    accum.setdefault(layer_name, []).append(inp.cpu())

            return hook

        target_set = set(target_names)
        for name, module in self.model.named_modules():
            if name in target_set:
                hooks.append(module.register_forward_hook(make_hook(name)))

        try:
            calib_inputs = get_calibration_tensors(
                self.model,
                tokenizer=self.tokenizer,
                calibration_loader=self.calibration_loader,
                calibration_texts=self.calibration_texts,
                device=self.target_device,
            )
            batches = self._normalize_calibration_batches(calib_inputs)
            n_batches = min(len(batches), self.max_calibration_batches)

            with torch.no_grad():
                for i in range(n_batches):
                    batch = batches[i]
                    try:
                        if isinstance(batch, dict):
                            self.model(**batch)
                        elif isinstance(batch, torch.Tensor):
                            self.model(batch)
                        elif isinstance(batch, (list, tuple)):
                            self.model(*batch)
                    except Exception as e:
                        logger.debug("Calibration batch %d failed: %s", i, e)
        finally:
            for h in hooks:
                h.remove()

        aggregated: dict[str, torch.Tensor] = {}
        for name, act_list in accum.items():
            if act_list:
                try:
                    aggregated[name] = torch.cat(act_list, dim=0)
                except Exception as e:
                    logger.debug("Activation concat failed for %s: %s", name, e)
        return aggregated

    @staticmethod
    def _normalize_calibration_batches(calib_inputs: Any) -> list[Any]:
        if calib_inputs is None:
            return []
        if isinstance(calib_inputs, (list, tuple)):
            if len(calib_inputs) == 0:
                return []
            first = calib_inputs[0]
            if isinstance(first, (dict, torch.Tensor)):
                return list(calib_inputs)
            if isinstance(first, (list, tuple)):
                return [calib_inputs]
            return [calib_inputs]
        if isinstance(calib_inputs, (dict, torch.Tensor)):
            return [calib_inputs]
        return [calib_inputs]

    # ------------------------------------------------------------------
    # Layer shedding
    # ------------------------------------------------------------------
    def _apply_layer_shedding(self) -> int:
        if self.shed_layers <= 0:
            return 0
        if getattr(self.model, "_mamba_shed_done", False):
            logger.info("MambaShedder: layer shedding already applied; skipping.")
            return 0

        target_container = None
        container_parent = None
        container_attr = None

        candidates = [
            "backbone.layers",
            "transformer.layer",
            "transformer.h",
            "model.layers",
            "encoder.layer",
            "layers",
            "h",
        ]

        def _resolve(path: str):
            parts = path.split(".")
            parent = self.model
            for p in parts[:-1]:
                parent = getattr(parent, p, None)
                if parent is None:
                    return None, None
            container = getattr(parent, parts[-1], None)
            if isinstance(container, (nn.ModuleList, list)):
                return parent, parts[-1]
            return None, None

        for path in candidates:
            parent, attr = _resolve(path)
            if parent is None:
                continue
            container = getattr(parent, attr)
            if len(container) > self.shed_layers:
                target_container = container
                container_parent = parent
                container_attr = attr
                break

        if target_container is None:
            # Generic search: longest module list.
            best_len = 0
            for name, mod in self.model.named_modules():
                if not isinstance(mod, nn.ModuleList):
                    continue
                if len(mod) <= self.shed_layers or len(mod) <= best_len:
                    continue
                parent, attr = _resolve(name)
                if parent is None:
                    continue
                target_container = mod
                container_parent = parent
                container_attr = attr
                best_len = len(mod)

        if target_container is None or len(target_container) <= self.shed_layers:
            logger.warning("MambaShedder: no suitable ModuleList found.")
            return 0

        orig_len = len(target_container)
        keep_len = orig_len - self.shed_layers

        # Always keep first and last (per Mamba-Shedder paper).
        keep_set = {0, orig_len - 1}
        n_middle = keep_len - len(keep_set)
        if n_middle > 0 and orig_len > 2:
            middle = list(range(1, orig_len - 1))
            rng = random.Random(42)
            middle_kept = sorted(rng.sample(middle, min(n_middle, len(middle))))
            keep_set.update(middle_kept)
        else:
            keep_set = set(range(keep_len))

        kept_indices = sorted(keep_set)[:keep_len]
        old_container = getattr(container_parent, container_attr)
        new_layers = nn.ModuleList([target_container[i] for i in kept_indices])
        setattr(container_parent, container_attr, new_layers)

        # Validation with revert
        if self.validate_after_shed:
            try:
                self.model.eval()
                dummy = torch.randint(0, 100, (1, 8), device=self.target_device)
                with torch.no_grad():
                    try:
                        self.model(input_ids=dummy)
                    except Exception:
                        try:
                            self.model(dummy)
                        except Exception:
                            cfg = getattr(self.model, "config", None)
                            hidden = getattr(cfg, "hidden_size", None) if cfg else None
                            if hidden is None:
                                first_lin = next(
                                    (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                                    None,
                                )
                                hidden = getattr(first_lin, "in_features", 8)
                            dt = next(self.model.parameters()).dtype
                            self.model(
                                torch.randn(1, 8, hidden, device=self.target_device, dtype=dt)
                            )
            except Exception as e:
                logger.warning("MambaShedder: shedding validation failed (%s); reverting.", e)
                setattr(container_parent, container_attr, old_container)
                return 0

        shed_done = orig_len - len(new_layers)
        logger.info(
            "MambaShedder: shed %d layers (%d → %d).",
            shed_done,
            orig_len,
            len(new_layers),
        )
        self.model._mamba_shed_done = True
        return shed_done

    # ------------------------------------------------------------------
    # Main compress
    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        from neuralops.core.model_utils import replace_module

        shed_done = self._apply_layer_shedding()

        # Single-pass target discovery (MS-FIX-16)
        targets = [
            (name, module)
            for name, module in self.model.named_modules()
            if self._is_target_module(name, module)
        ]

        if not targets:
            logger.warning("MambaShedder: no target layers found.")
            n = 0
            for name, module in self.model.named_modules():
                n += 1
                self._is_target_module(name, module, _debug=True)
            logger.warning("MambaShedder diagnostic: %d modules scanned, 0 targets.", n)

        activations_map = self._collect_activations([n for n, _ in targets])

        pruned_layers_count = 0
        self.layer_errors = {}
        layer_details: dict[str, Any] = {}

        total_orig_bytes = 0
        total_comp_bytes = 0
        seen_weight_ids: set[int] = set()

        for name, module in targets:
            base_mod = getattr(module, "base_layer", module)
            w = getattr(base_mod, "weight", None)
            if w is None:
                continue
            w_id = id(w)
            if w_id in seen_weight_ids:
                continue
            seen_weight_ids.add(w_id)

            raw_w = w.data
            device = raw_w.device
            dtype = raw_w.dtype
            cls_name = base_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(base_mod, "nf")

            W = raw_w.t().clone().float() if is_conv1d else raw_w.clone().float()
            d_out, d_in = W.shape

            orig_bytes = raw_w.numel() * raw_w.element_size()
            if getattr(base_mod, "bias", None) is not None:
                orig_bytes += base_mod.bias.numel() * base_mod.bias.element_size()
            total_orig_bytes += orig_bytes

            # Sensitivity: Fisher or weight-norm.
            act = activations_map.get(name, None)
            used_act_sensitivity = False
            if act is not None and act.numel() > 0:
                try:
                    act_flat = act.reshape(-1, act.shape[-1]).float().to(device)
                    if act_flat.shape[-1] == d_in:
                        fisher = (act_flat**2).mean(dim=0).clamp(min=1e-8)
                        sens = fisher.sqrt()
                        used_act_sensitivity = True
                    else:
                        sens = torch.norm(W, dim=0)
                except Exception as e:
                    logger.debug("Fisher computation failed for %s: %s", name, e)
                    sens = torch.norm(W, dim=0)
            else:
                sens = torch.norm(W, dim=0)

            # Decide kept output channels (or 2:4 mask).
            if self.structured_24:
                # 2:4 → no storage savings (dense storage kept). Fallback to
                # in-place weight modification.
                saliency = sens.unsqueeze(0).expand_as(W)
                try:
                    W_pruned, _ = apply_24_structured_mask(W, saliency=saliency, dim=1)
                except Exception as e:
                    logger.warning("2:4 pruning failed for %s: %s", name, e)
                    total_comp_bytes += orig_bytes
                    continue

                norm_orig = torch.norm(W, p="fro").clamp_min(1e-12)
                rel_err = (torch.norm(W - W_pruned, p="fro") / norm_orig).item()
                self.layer_errors[name] = rel_err

                if rel_err > self.max_error and self.fallback_to_original:
                    total_comp_bytes += orig_bytes
                    continue

                final_w = W_pruned.t().contiguous() if is_conv1d else W_pruned
                if final_w.shape != raw_w.shape:
                    total_comp_bytes += orig_bytes
                    continue

                with torch.no_grad():
                    base_mod.weight.data.copy_(final_w.to(dtype=dtype, device=device))
                    base_mod.weight.requires_grad = False
                    if getattr(base_mod, "bias", None) is not None:
                        base_mod.bias.requires_grad = False

                pruned_layers_count += 1
                total_comp_bytes += orig_bytes  # no real storage change
                layer_details[name] = {
                    "rel_error": rel_err,
                    "used_act_sensitivity": used_act_sensitivity,
                    "storage_bytes_before": orig_bytes,
                    "storage_bytes_after": orig_bytes,
                    "note": "2:4 sparsity: no storage savings",
                }
            else:
                # Channel pruning → real storage savings.
                n_keep = max(1, int(d_out * (1.0 - self.sparsity)))
                _, keep_idx = torch.topk(sens, k=n_keep, largest=True)
                keep_idx, _ = torch.sort(keep_idx)
                keep_idx = keep_idx.to(device)

                W_kept = W.index_select(0, keep_idx).contiguous()
                W_pruned_full = torch.zeros_like(W)
                W_pruned_full.index_copy_(0, keep_idx, W_kept)

                norm_orig = torch.norm(W, p="fro").clamp_min(1e-12)
                rel_err = (torch.norm(W - W_pruned_full, p="fro") / norm_orig).item()
                self.layer_errors[name] = rel_err

                if rel_err > self.max_error and self.fallback_to_original:
                    logger.warning(
                        "MambaShedder: err %.4f on %s > %.4f; keeping original.",
                        rel_err,
                        name,
                        self.max_error,
                    )
                    total_comp_bytes += orig_bytes
                    continue

                bias_kept = None
                if getattr(base_mod, "bias", None) is not None:
                    b_full = base_mod.bias.data.float()
                    bias_kept = (
                        b_full.index_select(0, keep_idx).contiguous().to(dtype=dtype, device=device)
                    )

                # Existing LoRA
                lora_A_src = getattr(module, "lora_A", None)
                lora_B_src = getattr(module, "lora_B", None)
                eff_lora_r = self.lora_r
                if lora_A_src is not None and hasattr(lora_A_src, "weight"):
                    eff_lora_r = int(lora_A_src.weight.shape[0])

                new_layer = MambaShedderPrunedLinear(
                    in_features=d_in,
                    original_out_features=d_out,
                    W_kept=W_kept.to(dtype=dtype, device=device),
                    keep_idx=keep_idx.cpu(),
                    bias_kept=bias_kept,
                    lora_r=eff_lora_r,
                    lora_alpha=self.lora_alpha,
                )

                # Preserve LoRA weights if shapes match.
                if (
                    lora_A_src is not None
                    and lora_B_src is not None
                    and new_layer.lora_A is not None
                ):
                    with torch.no_grad():
                        src_A = (
                            lora_A_src.weight.data
                            if hasattr(lora_A_src, "weight")
                            else lora_A_src.data
                        )
                        src_B = (
                            lora_B_src.weight.data
                            if hasattr(lora_B_src, "weight")
                            else lora_B_src.data
                        )
                        if (
                            src_A.shape == new_layer.lora_A.shape
                            and src_B.shape == new_layer.lora_B.shape
                        ):
                            new_layer.lora_A.copy_(
                                src_A.to(
                                    new_layer.lora_A.device,
                                    dtype=new_layer.lora_A.dtype,
                                )
                            )
                            new_layer.lora_B.copy_(
                                src_B.to(
                                    new_layer.lora_B.device,
                                    dtype=new_layer.lora_B.dtype,
                                )
                            )

                # Splice with wrapper preservation (MS-FIX-18)
                if hasattr(module, "base_layer"):
                    module.base_layer = new_layer
                else:
                    replace_module(self.model, name, new_layer)

                pruned_layers_count += 1
                total_comp_bytes += new_layer.storage_bytes()
                layer_details[name] = {
                    "rel_error": rel_err,
                    "used_act_sensitivity": used_act_sensitivity,
                    "storage_bytes_before": orig_bytes,
                    "storage_bytes_after": new_layer.storage_bytes(),
                    "keep_ratio": n_keep / d_out,
                }

        # ------------------------------------------------------------------
        # Forward sanity check (MS-FIX-17)
        # ------------------------------------------------------------------
        forward_check_success = True
        try:
            with torch.no_grad():
                dev = next(self.model.parameters(), torch.tensor(0.0)).device
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                try:
                    self.model(input_ids=dummy)
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception:
                        cfg = getattr(self.model, "config", None)
                        hidden = getattr(cfg, "hidden_size", None) if cfg else None
                        if hidden is None:
                            first_lin = next(
                                (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                                None,
                            )
                            hidden = getattr(first_lin, "in_features", 8)
                        dt = next(self.model.parameters()).dtype
                        self.model(torch.randn(1, 8, hidden, device=dev, dtype=dt))
        except Exception as e:
            logger.warning("MambaShedder forward sanity check failed: %s", e)
            forward_check_success = False

        # ------------------------------------------------------------------
        # Stats
        # ------------------------------------------------------------------
        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        avg_err = (
            sum(self.layer_errors.values()) / max(1, len(self.layer_errors))
            if self.layer_errors
            else 0.0
        )

        storage_ratio = total_orig_bytes / max(1, total_comp_bytes) if total_comp_bytes > 0 else 1.0
        effective_ratio = (
            total_comp_bytes / max(1, total_orig_bytes) if total_orig_bytes > 0 else 1.0
        )

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        stats.update(
            {
                "pruned_layers": pruned_layers_count,
                "shed_layers": shed_done,
                "sparsity_target": self.sparsity,
                "structured_24": self.structured_24,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "avg_reconstruction_error": avg_err,
                "layer_errors": self.layer_errors,
                "layer_details": layer_details,
                "orig_storage_bytes": total_orig_bytes,
                "comp_storage_bytes": total_comp_bytes,
                "storage_compression_ratio": round(storage_ratio, 3),
                "storage_reduction_pct": round(100.0 * (1.0 - effective_ratio), 2),
                "effective_parameter_ratio": round(effective_ratio, 4),
                "forward_check_success": forward_check_success,
                "method": "mamba_shedder_sensitivity",
                "note": (
                    "Channel pruning now replaces target layers with "
                    "MambaShedderPrunedLinear (real storage reduction). "
                    "2:4 sparsity stores dense weights, so its storage "
                    "compression ratio stays at 1.0."
                ),
            }
        )
        self.stats = stats

        logger.info(
            "MambaShedder: %d pruned, %d shed, avg_err=%.4f, storage=%.2fx.",
            pruned_layers_count,
            shed_done,
            avg_err,
            storage_ratio,
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class MambaShedderPrunedLinear (in_features: int, original_out_features: int, W_kept: torch.Tensor, keep_idx: torch.Tensor, bias_kept: torch.Tensor | None = None, lora_r: int = 0, lora_alpha: float = 1.0)

Expand source code

class MambaShedderPrunedLinear(nn.Module):
    """
    Linear layer with pruned output channels.

    Stores only the kept rows of W:
      - W_kept:    (k, in_features)
      - bias_kept: (k,) or None
      - keep_idx:  (k,) int64 (output positions of kept rows)
      - original_out_features: int

    Forward:
      h = F.linear(x, W_kept, bias_kept)              # (..., k)
      out = zeros(..., original_out) then scatter to keep_idx
      + LoRA(x) if present.

    Storage: k * in instead of out * in.
    """

    is_custom_attention: bool = False

    def __init__(
        self,
        in_features: int,
        original_out_features: int,
        W_kept: torch.Tensor,
        keep_idx: torch.Tensor,
        bias_kept: torch.Tensor | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> None:
        super().__init__()
        self.in_features = int(in_features)
        self.original_out_features = int(original_out_features)
        self.out_features = self.original_out_features
        self.nf = self.out_features
        self.num_kept = int(W_kept.shape[0])
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        self.register_buffer("keep_idx", keep_idx.long(), persistent=True)
        self.W_kept = nn.Parameter(W_kept, requires_grad=False)

        if bias_kept is not None:
            self.bias_kept = nn.Parameter(bias_kept, requires_grad=False)
        else:
            self.register_parameter("bias_kept", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty(
                    (self.lora_r, self.in_features),
                    dtype=torch.float32,
                    device=W_kept.device,
                )
            )
            self.lora_B = nn.Parameter(
                torch.zeros(
                    (self.original_out_features, self.lora_r),
                    dtype=torch.float32,
                    device=W_kept.device,
                )
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    @property
    def weight(self) -> torch.Tensor:
        """Reconstruct (original_out, in) weight with zeros in pruned rows."""
        W = self.W_kept.new_zeros(self.original_out_features, self.in_features)
        W.index_copy_(0, self.keep_idx, self.W_kept)
        return W

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = F.linear(x, self.W_kept, self.bias_kept)  # (..., k)
        out = torch.zeros(
            *x.shape[:-1],
            self.original_out_features,
            dtype=h.dtype,
            device=h.device,
        )
        out.scatter_(-1, self.keep_idx.expand(*x.shape[:-1], -1), h)
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x.float(), self.lora_A), self.lora_B)
            out = out + self.scaling * lora_out.to(out.dtype)
        return out

    def storage_bytes(self) -> int:
        total = self.W_kept.numel() * self.W_kept.element_size()
        total += self.keep_idx.numel() * self.keep_idx.element_size()
        if self.bias_kept is not None:
            total += self.bias_kept.numel() * self.bias_kept.element_size()
        if self.lora_A is not None:
            total += self.lora_A.numel() * self.lora_A.element_size()
        if self.lora_B is not None:
            total += self.lora_B.numel() * self.lora_B.element_size()
        return total

Linear layer with pruned output channels.

Stores only the kept rows of W: - W_kept: (k, in_features) - bias_kept: (k,) or None - keep_idx: (k,) int64 (output positions of kept rows) - original_out_features: int

Forward -----= h = F.linear(x, W_kept, bias_kept) # (…, k) out = zeros(…, original_out) then scatter to keep_idx + LoRA(x) if present.

Storage: k * in instead of out * in.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstruct (original_out, in) weight with zeros in pruned rows."""
    W = self.W_kept.new_zeros(self.original_out_features, self.in_features)
    W.index_copy_(0, self.keep_idx, self.W_kept)
    return W

Reconstruct (original_out, in) weight with zeros in pruned rows.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    h = F.linear(x, self.W_kept, self.bias_kept)  # (..., k)
    out = torch.zeros(
        *x.shape[:-1],
        self.original_out_features,
        dtype=h.dtype,
        device=h.device,
    )
    out.scatter_(-1, self.keep_idx.expand(*x.shape[:-1], -1), h)
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x.float(), self.lora_A), self.lora_B)
        out = out + self.scaling * lora_out.to(out.dtype)
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def storage_bytes(self) ‑> int

Expand source code

def storage_bytes(self) -> int:
    total = self.W_kept.numel() * self.W_kept.element_size()
    total += self.keep_idx.numel() * self.keep_idx.element_size()
    if self.bias_kept is not None:
        total += self.bias_kept.numel() * self.bias_kept.element_size()
    if self.lora_A is not None:
        total += self.lora_A.numel() * self.lora_A.element_size()
    if self.lora_B is not None:
        total += self.lora_B.numel() * self.lora_B.element_size()
    return total

class MiniBatchKMeans (n_clusters=8, *, init='k-means++', max_iter=100, batch_size=1024, verbose=0, compute_labels=True, random_state=None, tol=0.0, max_no_improvement=10, init_size=None, n_init='auto', reassignment_ratio=0.01)

Expand source code

class MiniBatchKMeans(_BaseKMeans):
    """
    Mini-Batch K-Means clustering.

    Read more in the :ref:`User Guide <mini_batch_kmeans>`.

    Parameters
    ----------

    n_clusters : int, default=8
        The number of clusters to form as well as the number of
        centroids to generate.

    init : {'k-means++', 'random'}, callable or array-like of shape \
            (n_clusters, n_features), default='k-means++'
        Method for initialization:

        'k-means++' : selects initial cluster centroids using sampling based on
        an empirical probability distribution of the points' contribution to the
        overall inertia. This technique speeds up convergence. The algorithm
        implemented is "greedy k-means++". It differs from the vanilla k-means++
        by making several trials at each sampling step and choosing the best centroid
        among them.

        'random': choose `n_clusters` observations (rows) at random from data
        for the initial centroids.

        If an array is passed, it should be of shape (n_clusters, n_features)
        and gives the initial centers.

        If a callable is passed, it should take arguments X, n_clusters and a
        random state and return an initialization.

        For an evaluation of the impact of initialization, see the example
        :ref:`sphx_glr_auto_examples_cluster_plot_kmeans_stability_low_dim_dense.py`.

    max_iter : int, default=100
        Maximum number of iterations over the complete dataset before
        stopping independently of any early stopping criterion heuristics.

    batch_size : int, default=1024
        Size of the mini batches.
        For faster computations, you can set `batch_size > 256 * number_of_cores`
        to enable :ref:`parallelism <lower-level-parallelism-with-openmp>`
        on all cores.

        .. versionchanged:: 1.0
           `batch_size` default changed from 100 to 1024.

    verbose : int, default=0
        Verbosity mode.

    compute_labels : bool, default=True
        Compute label assignment and inertia for the complete dataset
        once the minibatch optimization has converged in fit.

    random_state : int, RandomState instance or None, default=None
        Determines random number generation for centroid initialization and
        random reassignment. Use an int to make the randomness deterministic.
        See :term:`Glossary <random_state>`.

    tol : float, default=0.0
        Control early stopping based on the relative center changes as
        measured by a smoothed, variance-normalized of the mean center
        squared position changes. This early stopping heuristics is
        closer to the one used for the batch variant of the algorithms
        but induces a slight computational and memory overhead over the
        inertia heuristic.

        To disable convergence detection based on normalized center
        change, set tol to 0.0 (default).

    max_no_improvement : int, default=10
        Control early stopping based on the consecutive number of mini
        batches that does not yield an improvement on the smoothed inertia.

        To disable convergence detection based on inertia, set
        max_no_improvement to None.

    init_size : int, default=None
        Number of samples to randomly sample for speeding up the
        initialization (sometimes at the expense of accuracy): the
        only algorithm is initialized by running a batch KMeans on a
        random subset of the data. This needs to be larger than n_clusters.

        If `None`, the heuristic is `init_size = 3 * batch_size` if
        `3 * batch_size < n_clusters`, else `init_size = 3 * n_clusters`.

    n_init : 'auto' or int, default="auto"
        Number of random initializations that are tried.
        In contrast to KMeans, the algorithm is only run once, using the best of
        the `n_init` initializations as measured by inertia. Several runs are
        recommended for sparse high-dimensional problems (see
        :ref:`kmeans_sparse_high_dim`).

        When `n_init='auto'`, the number of runs depends on the value of init:
        3 if using `init='random'` or `init` is a callable;
        1 if using `init='k-means++'` or `init` is an array-like.

        .. versionadded:: 1.2
           Added 'auto' option for `n_init`.

        .. versionchanged:: 1.4
           Default value for `n_init` changed to `'auto'` in version.

    reassignment_ratio : float, default=0.01
        Control the fraction of the maximum number of counts for a center to
        be reassigned. A higher value means that low count centers are more
        easily reassigned, which means that the model will take longer to
        converge, but should converge in a better clustering. However, too high
        a value may cause convergence issues, especially with a small batch
        size.

    Attributes
    ----------

    cluster_centers_ : ndarray of shape (n_clusters, n_features)
        Coordinates of cluster centers.

    labels_ : ndarray of shape (n_samples,)
        Labels of each point (if compute_labels is set to True).

    inertia_ : float
        The value of the inertia criterion associated with the chosen
        partition if compute_labels is set to True. If compute_labels is set to
        False, it's an approximation of the inertia based on an exponentially
        weighted average of the batch inertiae.
        The inertia is defined as the sum of square distances of samples to
        their cluster center, weighted by the sample weights if provided.

    n_iter_ : int
        Number of iterations over the full dataset.

    n_steps_ : int
        Number of minibatches processed.

        .. versionadded:: 1.0

    n_features_in_ : int
        Number of features seen during :term:`fit`.

        .. versionadded:: 0.24

    feature_names_in_ : ndarray of shape (`n_features_in_`,)
        Names of features seen during :term:`fit`. Defined only when `X`
        has feature names that are all strings.

        .. versionadded:: 1.0

    See Also
    --------
    KMeans : The classic implementation of the clustering method based on the
        Lloyd's algorithm. It consumes the whole set of input data at each
        iteration.

    Notes
    -----
    See https://www.eecs.tufts.edu/~dsculley/papers/fastkmeans.pdf

    When there are too few points in the dataset, some centers may be
    duplicated, which means that a proper clustering in terms of the number
    of requesting clusters and the number of returned clusters will not
    always match. One solution is to set `reassignment_ratio=0`, which
    prevents reassignments of clusters that are too small.

    See :ref:`sphx_glr_auto_examples_cluster_plot_birch_vs_minibatchkmeans.py` for a
    comparison with :class:`~sklearn.cluster.BIRCH`.

    Examples
    --------
    >>> from sklearn.cluster import MiniBatchKMeans
    >>> import numpy as np
    >>> X = np.array([[1, 2], [1, 4], [1, 0],
    ...               [4, 2], [4, 0], [4, 4],
    ...               [4, 5], [0, 1], [2, 2],
    ...               [3, 2], [5, 5], [1, -1]])
    >>> # manually fit on batches
    >>> kmeans = MiniBatchKMeans(n_clusters=2,
    ...                          random_state=0,
    ...                          batch_size=6,
    ...                          n_init="auto")
    >>> kmeans = kmeans.partial_fit(X[0:6,:])
    >>> kmeans = kmeans.partial_fit(X[6:12,:])
    >>> kmeans.cluster_centers_
    array([[3.375, 3.  ],
           [0.75 , 0.5 ]])
    >>> kmeans.predict([[0, 0], [4, 4]])
    array([1, 0], dtype=int32)
    >>> # fit on the whole data
    >>> kmeans = MiniBatchKMeans(n_clusters=2,
    ...                          random_state=0,
    ...                          batch_size=6,
    ...                          max_iter=10,
    ...                          n_init="auto").fit(X)
    >>> kmeans.cluster_centers_
    array([[3.55102041, 2.48979592],
           [1.06896552, 1.        ]])
    >>> kmeans.predict([[0, 0], [4, 4]])
    array([1, 0], dtype=int32)

    For a comparison of Mini-Batch K-Means clustering with other clustering algorithms,
    see :ref:`sphx_glr_auto_examples_cluster_plot_cluster_comparison.py`
    """

    _parameter_constraints: dict = {
        **_BaseKMeans._parameter_constraints,
        "batch_size": [Interval(Integral, 1, None, closed="left")],
        "compute_labels": ["boolean"],
        "max_no_improvement": [Interval(Integral, 0, None, closed="left"), None],
        "init_size": [Interval(Integral, 1, None, closed="left"), None],
        "reassignment_ratio": [Interval(Real, 0, None, closed="left")],
    }

    def __init__(
        self,
        n_clusters=8,
        *,
        init="k-means++",
        max_iter=100,
        batch_size=1024,
        verbose=0,
        compute_labels=True,
        random_state=None,
        tol=0.0,
        max_no_improvement=10,
        init_size=None,
        n_init="auto",
        reassignment_ratio=0.01,
    ):
        super().__init__(
            n_clusters=n_clusters,
            init=init,
            max_iter=max_iter,
            verbose=verbose,
            random_state=random_state,
            tol=tol,
            n_init=n_init,
        )

        self.max_no_improvement = max_no_improvement
        self.batch_size = batch_size
        self.compute_labels = compute_labels
        self.init_size = init_size
        self.reassignment_ratio = reassignment_ratio

    def _check_params_vs_input(self, X):
        super()._check_params_vs_input(X, default_n_init=3)

        self._batch_size = min(self.batch_size, X.shape[0])

        # init_size
        self._init_size = self.init_size
        if self._init_size is None:
            self._init_size = 3 * self._batch_size
            if self._init_size < self.n_clusters:
                self._init_size = 3 * self.n_clusters
        elif self._init_size < self.n_clusters:
            warnings.warn(
                (
                    f"init_size={self._init_size} should be larger than "
                    f"n_clusters={self.n_clusters}. Setting it to "
                    "min(3*n_clusters, n_samples)"
                ),
                RuntimeWarning,
                stacklevel=2,
            )
            self._init_size = 3 * self.n_clusters
        self._init_size = min(self._init_size, X.shape[0])

        # reassignment_ratio
        if self.reassignment_ratio < 0:
            raise ValueError(
                "reassignment_ratio should be >= 0, got "
                f"{self.reassignment_ratio} instead."
            )

    def _warn_mkl_vcomp(self, n_active_threads):
        """Warn when vcomp and mkl are both present"""
        warnings.warn(
            "MiniBatchKMeans is known to have a memory leak on "
            "Windows with MKL, when there are less chunks than "
            "available threads. You can prevent it by setting "
            f"batch_size >= {self._n_threads * CHUNK_SIZE} or by "
            "setting the environment variable "
            f"OMP_NUM_THREADS={n_active_threads}"
        )

    def _mini_batch_convergence(
        self, step, n_steps, n_samples, centers_squared_diff, batch_inertia
    ):
        """Helper function to encapsulate the early stopping logic"""
        # Normalize inertia to be able to compare values when
        # batch_size changes
        batch_inertia /= self._batch_size

        # count steps starting from 1 for user friendly verbose mode.
        step = step + 1

        # Ignore first iteration because it's inertia from initialization.
        if step == 1:
            if self.verbose:
                print(
                    f"Minibatch step {step}/{n_steps}: mean batch "
                    f"inertia: {batch_inertia}"
                )
            return False

        # Compute an Exponentially Weighted Average of the inertia to
        # monitor the convergence while discarding minibatch-local stochastic
        # variability: https://en.wikipedia.org/wiki/Moving_average
        if self._ewa_inertia is None:
            self._ewa_inertia = batch_inertia
        else:
            alpha = self._batch_size * 2.0 / (n_samples + 1)
            alpha = min(alpha, 1)
            self._ewa_inertia = self._ewa_inertia * (1 - alpha) + batch_inertia * alpha

        # Log progress to be able to monitor convergence
        if self.verbose:
            print(
                f"Minibatch step {step}/{n_steps}: mean batch inertia: "
                f"{batch_inertia}, ewa inertia: {self._ewa_inertia}"
            )

        # Early stopping based on absolute tolerance on squared change of
        # centers position
        if self._tol > 0.0 and centers_squared_diff <= self._tol:
            if self.verbose:
                print(f"Converged (small centers change) at step {step}/{n_steps}")
            return True

        # Early stopping heuristic due to lack of improvement on smoothed
        # inertia
        if self._ewa_inertia_min is None or self._ewa_inertia < self._ewa_inertia_min:
            self._no_improvement = 0
            self._ewa_inertia_min = self._ewa_inertia
        else:
            self._no_improvement += 1

        if (
            self.max_no_improvement is not None
            and self._no_improvement >= self.max_no_improvement
        ):
            if self.verbose:
                print(
                    "Converged (lack of improvement in inertia) at step "
                    f"{step}/{n_steps}"
                )
            return True

        return False

    def _random_reassign(self):
        """Check if a random reassignment needs to be done.

        Do random reassignments each time 10 * n_clusters samples have been
        processed.

        If there are empty clusters we always want to reassign.
        """
        self._n_since_last_reassign += self._batch_size
        if (self._counts == 0).any() or self._n_since_last_reassign >= (
            10 * self.n_clusters
        ):
            self._n_since_last_reassign = 0
            return True
        return False

    @_fit_context(prefer_skip_nested_validation=True)
    def fit(self, X, y=None, sample_weight=None):
        """Compute the centroids on X by chunking it into mini-batches.

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape (n_samples, n_features)
            Training instances to cluster. It must be noted that the data
            will be converted to C ordering, which will cause a memory copy
            if the given data is not C-contiguous.
            If a sparse matrix is passed, a copy will be made if it's not in
            CSR format.

        y : Ignored
            Not used, present here for API consistency by convention.

        sample_weight : array-like of shape (n_samples,), default=None
            The weights for each observation in X. If None, all observations
            are assigned equal weight. `sample_weight` is not used during
            initialization if `init` is a callable or a user provided array.

            .. versionadded:: 0.20

        Returns
        -------
        self : object
            Fitted estimator.
        """
        X = validate_data(
            self,
            X,
            accept_sparse="csr",
            dtype=[np.float64, np.float32],
            order="C",
            accept_large_sparse=False,
        )

        self._check_params_vs_input(X)
        random_state = check_random_state(self.random_state)
        sample_weight = _check_sample_weight(sample_weight, X, dtype=X.dtype)
        self._n_threads = _openmp_effective_n_threads()
        n_samples, n_features = X.shape

        # Validate init array
        init = self.init
        if _is_arraylike_not_scalar(init):
            init = check_array(init, dtype=X.dtype, copy=True, order="C")
            self._validate_center_shape(X, init)

        self._check_mkl_vcomp(X, self._batch_size)

        # precompute squared norms of data points
        x_squared_norms = row_norms(X, squared=True)

        # Validation set for the init
        validation_indices = random_state.randint(0, n_samples, self._init_size)
        X_valid = X[validation_indices]
        sample_weight_valid = sample_weight[validation_indices]

        # perform several inits with random subsets
        best_inertia = None
        for init_idx in range(self._n_init):
            if self.verbose:
                print(f"Init {init_idx + 1}/{self._n_init} with method {init}")

            # Initialize the centers using only a fraction of the data as we
            # expect n_samples to be very large when using MiniBatchKMeans.
            cluster_centers = self._init_centroids(
                X,
                x_squared_norms=x_squared_norms,
                init=init,
                random_state=random_state,
                init_size=self._init_size,
                sample_weight=sample_weight,
            )

            # Compute inertia on a validation set.
            _, inertia = _labels_inertia_threadpool_limit(
                X_valid,
                sample_weight_valid,
                cluster_centers,
                n_threads=self._n_threads,
            )

            if self.verbose:
                print(f"Inertia for init {init_idx + 1}/{self._n_init}: {inertia}")
            if best_inertia is None or inertia < best_inertia:
                init_centers = cluster_centers
                best_inertia = inertia

        centers = init_centers
        centers_new = np.empty_like(centers)

        # Initialize counts
        self._counts = np.zeros(self.n_clusters, dtype=X.dtype)

        # Attributes to monitor the convergence
        self._ewa_inertia = None
        self._ewa_inertia_min = None
        self._no_improvement = 0

        # Initialize number of samples seen since last reassignment
        self._n_since_last_reassign = 0

        n_steps = (self.max_iter * n_samples) // self._batch_size

        with _get_threadpool_controller().limit(limits=1, user_api="blas"):
            # Perform the iterative optimization until convergence
            for i in range(n_steps):
                # Sample a minibatch from the full dataset
                minibatch_indices = random_state.randint(0, n_samples, self._batch_size)

                # Perform the actual update step on the minibatch data
                batch_inertia = _mini_batch_step(
                    X=X[minibatch_indices],
                    sample_weight=sample_weight[minibatch_indices],
                    centers=centers,
                    centers_new=centers_new,
                    weight_sums=self._counts,
                    random_state=random_state,
                    random_reassign=self._random_reassign(),
                    reassignment_ratio=self.reassignment_ratio,
                    verbose=self.verbose,
                    n_threads=self._n_threads,
                )

                if self._tol > 0.0:
                    centers_squared_diff = np.sum((centers_new - centers) ** 2)
                else:
                    centers_squared_diff = 0

                centers, centers_new = centers_new, centers

                # Monitor convergence and do early stopping if necessary
                if self._mini_batch_convergence(
                    i, n_steps, n_samples, centers_squared_diff, batch_inertia
                ):
                    break

        self.cluster_centers_ = centers
        self._n_features_out = self.cluster_centers_.shape[0]

        self.n_steps_ = i + 1
        self.n_iter_ = int(np.ceil(((i + 1) * self._batch_size) / n_samples))

        if self.compute_labels:
            self.labels_, self.inertia_ = _labels_inertia_threadpool_limit(
                X,
                sample_weight,
                self.cluster_centers_,
                n_threads=self._n_threads,
            )
        else:
            self.inertia_ = self._ewa_inertia * n_samples

        return self

    @_fit_context(prefer_skip_nested_validation=True)
    def partial_fit(self, X, y=None, sample_weight=None):
        """Update k means estimate on a single mini-batch X.

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape (n_samples, n_features)
            Training instances to cluster. It must be noted that the data
            will be converted to C ordering, which will cause a memory copy
            if the given data is not C-contiguous.
            If a sparse matrix is passed, a copy will be made if it's not in
            CSR format.

        y : Ignored
            Not used, present here for API consistency by convention.

        sample_weight : array-like of shape (n_samples,), default=None
            The weights for each observation in X. If None, all observations
            are assigned equal weight. `sample_weight` is not used during
            initialization if `init` is a callable or a user provided array.

        Returns
        -------
        self : object
            Return updated estimator.
        """
        has_centers = hasattr(self, "cluster_centers_")

        X = validate_data(
            self,
            X,
            accept_sparse="csr",
            dtype=[np.float64, np.float32],
            order="C",
            accept_large_sparse=False,
            reset=not has_centers,
        )

        self._random_state = getattr(
            self, "_random_state", check_random_state(self.random_state)
        )
        sample_weight = _check_sample_weight(sample_weight, X, dtype=X.dtype)
        self.n_steps_ = getattr(self, "n_steps_", 0)

        # precompute squared norms of data points
        x_squared_norms = row_norms(X, squared=True)

        if not has_centers:
            # this instance has not been fitted yet (fit or partial_fit)
            self._check_params_vs_input(X)
            self._n_threads = _openmp_effective_n_threads()

            # Validate init array
            init = self.init
            if _is_arraylike_not_scalar(init):
                init = check_array(init, dtype=X.dtype, copy=True, order="C")
                self._validate_center_shape(X, init)

            self._check_mkl_vcomp(X, X.shape[0])

            # initialize the cluster centers
            self.cluster_centers_ = self._init_centroids(
                X,
                x_squared_norms=x_squared_norms,
                init=init,
                random_state=self._random_state,
                init_size=self._init_size,
                sample_weight=sample_weight,
            )

            # Initialize counts
            self._counts = np.zeros(self.n_clusters, dtype=X.dtype)

            # Initialize number of samples seen since last reassignment
            self._n_since_last_reassign = 0

        with _get_threadpool_controller().limit(limits=1, user_api="blas"):
            _mini_batch_step(
                X,
                sample_weight=sample_weight,
                centers=self.cluster_centers_,
                centers_new=self.cluster_centers_,
                weight_sums=self._counts,
                random_state=self._random_state,
                random_reassign=self._random_reassign(),
                reassignment_ratio=self.reassignment_ratio,
                verbose=self.verbose,
                n_threads=self._n_threads,
            )

        if self.compute_labels:
            self.labels_, self.inertia_ = _labels_inertia_threadpool_limit(
                X,
                sample_weight,
                self.cluster_centers_,
                n_threads=self._n_threads,
            )

        self.n_steps_ += 1
        self._n_features_out = self.cluster_centers_.shape[0]

        return self

Mini-Batch K-Means clustering.

Read more in the :ref:User Guide <mini_batch_kmeans>.

Parameters

n_clusters : int, default=8

The number of clusters to form as well as the number of centroids to generate.

init : {'k-means++', 'random'}, callable or array-like of shape (n_clusters, n_features), default='k-means++'

Method for initialization:

'k-means++' : selects initial cluster centroids using sampling based on an empirical probability distribution of the points' contribution to the overall inertia. This technique speeds up convergence. The algorithm implemented is "greedy k-means++". It differs from the vanilla k-means++ by making several trials at each sampling step and choosing the best centroid among them.

'random': choose n_clusters observations (rows) at random from data for the initial centroids.

If an array is passed, it should be of shape (n_clusters, n_features) and gives the initial centers.

If a callable is passed, it should take arguments X, n_clusters and a random state and return an initialization.

For an evaluation of the impact of initialization, see the example :ref:sphx_glr_auto_examples_cluster_plot_kmeans_stability_low_dim_dense.py.

max_iter : int, default=100

Maximum number of iterations over the complete dataset before stopping independently of any early stopping criterion heuristics.

batch_size : int, default=1024

Size of the mini batches. For faster computations, you can set batch_size > 256 * number_of_cores to enable :ref:parallelism <lower-level-parallelism-with-openmp> on all cores.

Changed in version: 1.0

batch_size default changed from 100 to 1024.

verbose : int, default=0

Verbosity mode.

compute_labels : bool, default=True

Compute label assignment and inertia for the complete dataset once the minibatch optimization has converged in fit.

random_state : int, RandomState instance or None, default=None

Determines random number generation for centroid initialization and random reassignment. Use an int to make the randomness deterministic. See :term:Glossary <random_state>.

tol : float, default=0.0

Control early stopping based on the relative center changes as measured by a smoothed, variance-normalized of the mean center squared position changes. This early stopping heuristics is closer to the one used for the batch variant of the algorithms but induces a slight computational and memory overhead over the inertia heuristic.

To disable convergence detection based on normalized center change, set tol to 0.0 (default).

max_no_improvement : int, default=10

Control early stopping based on the consecutive number of mini batches that does not yield an improvement on the smoothed inertia.

To disable convergence detection based on inertia, set max_no_improvement to None.

init_size : int, default=None

Number of samples to randomly sample for speeding up the initialization (sometimes at the expense of accuracy): the only algorithm is initialized by running a batch KMeans on a random subset of the data. This needs to be larger than n_clusters.

If None, the heuristic is init_size = 3 * batch_size if 3 * batch_size < n_clusters, else init_size = 3 * n_clusters.

n_init : 'auto' or int, default="auto"

Number of random initializations that are tried. In contrast to KMeans, the algorithm is only run once, using the best of the n_init initializations as measured by inertia. Several runs are recommended for sparse high-dimensional problems (see :ref:kmeans_sparse_high_dim).

When n_init='auto', the number of runs depends on the value of init: 3 if using init='random' or init is a callable; 1 if using init='k-means++' or init is an array-like.

Added in version: 1.2

Added 'auto' option for n_init.

Changed in version: 1.4

Default value for n_init changed to 'auto' in version.

reassignment_ratio : float, default=0.01

Control the fraction of the maximum number of counts for a center to be reassigned. A higher value means that low count centers are more easily reassigned, which means that the model will take longer to converge, but should converge in a better clustering. However, too high a value may cause convergence issues, especially with a small batch size.

Attributes

cluster_centers_ : ndarray of shape (n_clusters, n_features): Coordinates of cluster centers.
labels_ : ndarray of shape (n_samples,): Labels of each point (if compute_labels is set to True).
inertia_ : float: The value of the inertia criterion associated with the chosen partition if compute_labels is set to True. If compute_labels is set to False, it's an approximation of the inertia based on an exponentially weighted average of the batch inertiae. The inertia is defined as the sum of square distances of samples to their cluster center, weighted by the sample weights if provided.
n_iter_ : int: Number of iterations over the full dataset.
n_steps_ : int: Number of minibatches processed.

Added in version: 1.0
n_features_in_ : int: Number of features seen during :term:fit.

Added in version: 0.24
feature_names_in_ : ndarray of shape (n_features_in_,): Names of features seen during :term:fit. Defined only when X has feature names that are all strings.

Added in version: 1.0

Notes

See https://www.eecs.tufts.edu/~dsculley/papers/fastkmeans.pdf

When there are too few points in the dataset, some centers may be duplicated, which means that a proper clustering in terms of the number of requesting clusters and the number of returned clusters will not always match. One solution is to set reassignment_ratio=0, which prevents reassignments of clusters that are too small.

See :ref:sphx_glr_auto_examples_cluster_plot_birch_vs_minibatchkmeans.py for a comparison with :class:~sklearn.cluster.BIRCH.

Examples

>>> from sklearn.cluster import MiniBatchKMeans
>>> import numpy as np
>>> X = np.array([[1, 2], [1, 4], [1, 0],
...               [4, 2], [4, 0], [4, 4],
...               [4, 5], [0, 1], [2, 2],
...               [3, 2], [5, 5], [1, -1]])
>>> # manually fit on batches
>>> kmeans = MiniBatchKMeans(n_clusters=2,
...                          random_state=0,
...                          batch_size=6,
...                          n_init="auto")
>>> kmeans = kmeans.partial_fit(X[0:6,:])
>>> kmeans = kmeans.partial_fit(X[6:12,:])
>>> kmeans.cluster_centers_
array([[3.375, 3.  ],
       [0.75 , 0.5 ]])
>>> kmeans.predict([[0, 0], [4, 4]])
array([1, 0], dtype=int32)
>>> # fit on the whole data
>>> kmeans = MiniBatchKMeans(n_clusters=2,
...                          random_state=0,
...                          batch_size=6,
...                          max_iter=10,
...                          n_init="auto").fit(X)
>>> kmeans.cluster_centers_
array([[3.55102041, 2.48979592],
       [1.06896552, 1.        ]])
>>> kmeans.predict([[0, 0], [4, 4]])
array([1, 0], dtype=int32)

For a comparison of Mini-Batch K-Means clustering with other clustering algorithms, see :ref:sphx_glr_auto_examples_cluster_plot_cluster_comparison.py

Ancestors

sklearn.cluster._kmeans._BaseKMeans
sklearn.base.ClassNamePrefixFeaturesOutMixin
sklearn.base.TransformerMixin
sklearn.utils._set_output._SetOutputMixin
sklearn.base.ClusterMixin
sklearn.base.BaseEstimator
sklearn.utils._repr_html.base.ReprHTMLMixin
sklearn.utils._repr_html.base._HTMLDocumentationLinkMixin
sklearn.utils._metadata_requests._MetadataRequester
abc.ABC

Methods

def fit(self, X, y=None, sample_weight=None)

Expand source code

@_fit_context(prefer_skip_nested_validation=True)
def fit(self, X, y=None, sample_weight=None):
    """Compute the centroids on X by chunking it into mini-batches.

    Parameters
    ----------
    X : {array-like, sparse matrix} of shape (n_samples, n_features)
        Training instances to cluster. It must be noted that the data
        will be converted to C ordering, which will cause a memory copy
        if the given data is not C-contiguous.
        If a sparse matrix is passed, a copy will be made if it's not in
        CSR format.

    y : Ignored
        Not used, present here for API consistency by convention.

    sample_weight : array-like of shape (n_samples,), default=None
        The weights for each observation in X. If None, all observations
        are assigned equal weight. `sample_weight` is not used during
        initialization if `init` is a callable or a user provided array.

        .. versionadded:: 0.20

    Returns
    -------
    self : object
        Fitted estimator.
    """
    X = validate_data(
        self,
        X,
        accept_sparse="csr",
        dtype=[np.float64, np.float32],
        order="C",
        accept_large_sparse=False,
    )

    self._check_params_vs_input(X)
    random_state = check_random_state(self.random_state)
    sample_weight = _check_sample_weight(sample_weight, X, dtype=X.dtype)
    self._n_threads = _openmp_effective_n_threads()
    n_samples, n_features = X.shape

    # Validate init array
    init = self.init
    if _is_arraylike_not_scalar(init):
        init = check_array(init, dtype=X.dtype, copy=True, order="C")
        self._validate_center_shape(X, init)

    self._check_mkl_vcomp(X, self._batch_size)

    # precompute squared norms of data points
    x_squared_norms = row_norms(X, squared=True)

    # Validation set for the init
    validation_indices = random_state.randint(0, n_samples, self._init_size)
    X_valid = X[validation_indices]
    sample_weight_valid = sample_weight[validation_indices]

    # perform several inits with random subsets
    best_inertia = None
    for init_idx in range(self._n_init):
        if self.verbose:
            print(f"Init {init_idx + 1}/{self._n_init} with method {init}")

        # Initialize the centers using only a fraction of the data as we
        # expect n_samples to be very large when using MiniBatchKMeans.
        cluster_centers = self._init_centroids(
            X,
            x_squared_norms=x_squared_norms,
            init=init,
            random_state=random_state,
            init_size=self._init_size,
            sample_weight=sample_weight,
        )

        # Compute inertia on a validation set.
        _, inertia = _labels_inertia_threadpool_limit(
            X_valid,
            sample_weight_valid,
            cluster_centers,
            n_threads=self._n_threads,
        )

        if self.verbose:
            print(f"Inertia for init {init_idx + 1}/{self._n_init}: {inertia}")
        if best_inertia is None or inertia < best_inertia:
            init_centers = cluster_centers
            best_inertia = inertia

    centers = init_centers
    centers_new = np.empty_like(centers)

    # Initialize counts
    self._counts = np.zeros(self.n_clusters, dtype=X.dtype)

    # Attributes to monitor the convergence
    self._ewa_inertia = None
    self._ewa_inertia_min = None
    self._no_improvement = 0

    # Initialize number of samples seen since last reassignment
    self._n_since_last_reassign = 0

    n_steps = (self.max_iter * n_samples) // self._batch_size

    with _get_threadpool_controller().limit(limits=1, user_api="blas"):
        # Perform the iterative optimization until convergence
        for i in range(n_steps):
            # Sample a minibatch from the full dataset
            minibatch_indices = random_state.randint(0, n_samples, self._batch_size)

            # Perform the actual update step on the minibatch data
            batch_inertia = _mini_batch_step(
                X=X[minibatch_indices],
                sample_weight=sample_weight[minibatch_indices],
                centers=centers,
                centers_new=centers_new,
                weight_sums=self._counts,
                random_state=random_state,
                random_reassign=self._random_reassign(),
                reassignment_ratio=self.reassignment_ratio,
                verbose=self.verbose,
                n_threads=self._n_threads,
            )

            if self._tol > 0.0:
                centers_squared_diff = np.sum((centers_new - centers) ** 2)
            else:
                centers_squared_diff = 0

            centers, centers_new = centers_new, centers

            # Monitor convergence and do early stopping if necessary
            if self._mini_batch_convergence(
                i, n_steps, n_samples, centers_squared_diff, batch_inertia
            ):
                break

    self.cluster_centers_ = centers
    self._n_features_out = self.cluster_centers_.shape[0]

    self.n_steps_ = i + 1
    self.n_iter_ = int(np.ceil(((i + 1) * self._batch_size) / n_samples))

    if self.compute_labels:
        self.labels_, self.inertia_ = _labels_inertia_threadpool_limit(
            X,
            sample_weight,
            self.cluster_centers_,
            n_threads=self._n_threads,
        )
    else:
        self.inertia_ = self._ewa_inertia * n_samples

    return self

Compute the centroids on X by chunking it into mini-batches.

Parameters

X : {array-like, sparse matrix} of shape (n_samples, n_features): Training instances to cluster. It must be noted that the data will be converted to C ordering, which will cause a memory copy if the given data is not C-contiguous. If a sparse matrix is passed, a copy will be made if it's not in CSR format.
y : Ignored: Not used, present here for API consistency by convention.
sample_weight : array-like of shape (n_samples,), default=None: The weights for each observation in X. If None, all observations are assigned equal weight. sample_weight is not used during initialization if init is a callable or a user provided array.

Added in version: 0.20

Returns

self : object: Fitted estimator.

def partial_fit(self, X, y=None, sample_weight=None)

Expand source code

@_fit_context(prefer_skip_nested_validation=True)
def partial_fit(self, X, y=None, sample_weight=None):
    """Update k means estimate on a single mini-batch X.

    Parameters
    ----------
    X : {array-like, sparse matrix} of shape (n_samples, n_features)
        Training instances to cluster. It must be noted that the data
        will be converted to C ordering, which will cause a memory copy
        if the given data is not C-contiguous.
        If a sparse matrix is passed, a copy will be made if it's not in
        CSR format.

    y : Ignored
        Not used, present here for API consistency by convention.

    sample_weight : array-like of shape (n_samples,), default=None
        The weights for each observation in X. If None, all observations
        are assigned equal weight. `sample_weight` is not used during
        initialization if `init` is a callable or a user provided array.

    Returns
    -------
    self : object
        Return updated estimator.
    """
    has_centers = hasattr(self, "cluster_centers_")

    X = validate_data(
        self,
        X,
        accept_sparse="csr",
        dtype=[np.float64, np.float32],
        order="C",
        accept_large_sparse=False,
        reset=not has_centers,
    )

    self._random_state = getattr(
        self, "_random_state", check_random_state(self.random_state)
    )
    sample_weight = _check_sample_weight(sample_weight, X, dtype=X.dtype)
    self.n_steps_ = getattr(self, "n_steps_", 0)

    # precompute squared norms of data points
    x_squared_norms = row_norms(X, squared=True)

    if not has_centers:
        # this instance has not been fitted yet (fit or partial_fit)
        self._check_params_vs_input(X)
        self._n_threads = _openmp_effective_n_threads()

        # Validate init array
        init = self.init
        if _is_arraylike_not_scalar(init):
            init = check_array(init, dtype=X.dtype, copy=True, order="C")
            self._validate_center_shape(X, init)

        self._check_mkl_vcomp(X, X.shape[0])

        # initialize the cluster centers
        self.cluster_centers_ = self._init_centroids(
            X,
            x_squared_norms=x_squared_norms,
            init=init,
            random_state=self._random_state,
            init_size=self._init_size,
            sample_weight=sample_weight,
        )

        # Initialize counts
        self._counts = np.zeros(self.n_clusters, dtype=X.dtype)

        # Initialize number of samples seen since last reassignment
        self._n_since_last_reassign = 0

    with _get_threadpool_controller().limit(limits=1, user_api="blas"):
        _mini_batch_step(
            X,
            sample_weight=sample_weight,
            centers=self.cluster_centers_,
            centers_new=self.cluster_centers_,
            weight_sums=self._counts,
            random_state=self._random_state,
            random_reassign=self._random_reassign(),
            reassignment_ratio=self.reassignment_ratio,
            verbose=self.verbose,
            n_threads=self._n_threads,
        )

    if self.compute_labels:
        self.labels_, self.inertia_ = _labels_inertia_threadpool_limit(
            X,
            sample_weight,
            self.cluster_centers_,
            n_threads=self._n_threads,
        )

    self.n_steps_ += 1
    self._n_features_out = self.cluster_centers_.shape[0]

    return self

Update k means estimate on a single mini-batch X.

Parameters

X : {array-like, sparse matrix} of shape (n_samples, n_features): Training instances to cluster. It must be noted that the data will be converted to C ordering, which will cause a memory copy if the given data is not C-contiguous. If a sparse matrix is passed, a copy will be made if it's not in CSR format.
y : Ignored: Not used, present here for API consistency by convention.
sample_weight : array-like of shape (n_samples,), default=None: The weights for each observation in X. If None, all observations are assigned equal weight. sample_weight is not used during initialization if init is a callable or a user provided array.

Returns

self : object: Return updated estimator.

def set_fit_request(self: sklearn.cluster._kmeans.MiniBatchKMeans, *, sample_weight: bool | str | None = '$UNCHANGED$') ‑> sklearn.cluster._kmeans.MiniBatchKMeans

Expand source code

def func(*args, **kw):
    """Updates the `_metadata_request` attribute of the consumer (`instance`)
    for the parameters provided as `**kw`.

    This docstring is overwritten below.
    See REQUESTER_DOC for expected functionality.
    """
    if not _routing_enabled():
        raise RuntimeError(
            "This method is only available when metadata routing is enabled."
            " You can enable it using"
            " sklearn.set_config(enable_metadata_routing=True)."
        )

    if self.validate_keys and (set(kw) - set(self.keys)):
        raise TypeError(
            f"Unexpected args: {set(kw) - set(self.keys)} in {self.name}. "
            f"Accepted arguments are: {set(self.keys)}"
        )

    # This makes it possible to use the decorated method as an unbound method,
    # for instance when monkeypatching.
    # https://github.com/scikit-learn/scikit-learn/issues/28632
    if instance is None:
        _instance = args[0]
        args = args[1:]
    else:
        _instance = instance

    # Replicating python's behavior when positional args are given other than
    # `self`, and `self` is only allowed if this method is unbound.
    if args:
        raise TypeError(
            f"set_{self.name}_request() takes 0 positional argument but"
            f" {len(args)} were given"
        )

    requests = _instance._get_metadata_request()
    method_metadata_request = getattr(requests, self.name)

    for prop, alias in kw.items():
        if alias is not UNCHANGED:
            method_metadata_request.add_request(param=prop, alias=alias)
    _instance._metadata_request = requests

    return _instance

Configure whether metadata should be requested to be passed to the fit method.

The options for each parameter are:

True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to fit.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version: 1.3

Parameters

sample_weight : str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for sample_weight parameter in fit.

Returns

self : object: The updated object.

def set_partial_fit_request(self: sklearn.cluster._kmeans.MiniBatchKMeans, *, sample_weight: bool | str | None = '$UNCHANGED$') ‑> sklearn.cluster._kmeans.MiniBatchKMeans

Expand source code

def func(*args, **kw):
    """Updates the `_metadata_request` attribute of the consumer (`instance`)
    for the parameters provided as `**kw`.

    This docstring is overwritten below.
    See REQUESTER_DOC for expected functionality.
    """
    if not _routing_enabled():
        raise RuntimeError(
            "This method is only available when metadata routing is enabled."
            " You can enable it using"
            " sklearn.set_config(enable_metadata_routing=True)."
        )

    if self.validate_keys and (set(kw) - set(self.keys)):
        raise TypeError(
            f"Unexpected args: {set(kw) - set(self.keys)} in {self.name}. "
            f"Accepted arguments are: {set(self.keys)}"
        )

    # This makes it possible to use the decorated method as an unbound method,
    # for instance when monkeypatching.
    # https://github.com/scikit-learn/scikit-learn/issues/28632
    if instance is None:
        _instance = args[0]
        args = args[1:]
    else:
        _instance = instance

    # Replicating python's behavior when positional args are given other than
    # `self`, and `self` is only allowed if this method is unbound.
    if args:
        raise TypeError(
            f"set_{self.name}_request() takes 0 positional argument but"
            f" {len(args)} were given"
        )

    requests = _instance._get_metadata_request()
    method_metadata_request = getattr(requests, self.name)

    for prop, alias in kw.items():
        if alias is not UNCHANGED:
            method_metadata_request.add_request(param=prop, alias=alias)
    _instance._metadata_request = requests

    return _instance

Configure whether metadata should be requested to be passed to the partial_fit method.

The options for each parameter are:

True: metadata is requested, and passed to partial_fit if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to partial_fit.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version: 1.3

Parameters

sample_weight : str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for sample_weight parameter in partial_fit.

Returns

self : object: The updated object.

def set_score_request(self: sklearn.cluster._kmeans.MiniBatchKMeans, *, sample_weight: bool | str | None = '$UNCHANGED$') ‑> sklearn.cluster._kmeans.MiniBatchKMeans

Expand source code

def func(*args, **kw):
    """Updates the `_metadata_request` attribute of the consumer (`instance`)
    for the parameters provided as `**kw`.

    This docstring is overwritten below.
    See REQUESTER_DOC for expected functionality.
    """
    if not _routing_enabled():
        raise RuntimeError(
            "This method is only available when metadata routing is enabled."
            " You can enable it using"
            " sklearn.set_config(enable_metadata_routing=True)."
        )

    if self.validate_keys and (set(kw) - set(self.keys)):
        raise TypeError(
            f"Unexpected args: {set(kw) - set(self.keys)} in {self.name}. "
            f"Accepted arguments are: {set(self.keys)}"
        )

    # This makes it possible to use the decorated method as an unbound method,
    # for instance when monkeypatching.
    # https://github.com/scikit-learn/scikit-learn/issues/28632
    if instance is None:
        _instance = args[0]
        args = args[1:]
    else:
        _instance = instance

    # Replicating python's behavior when positional args are given other than
    # `self`, and `self` is only allowed if this method is unbound.
    if args:
        raise TypeError(
            f"set_{self.name}_request() takes 0 positional argument but"
            f" {len(args)} were given"
        )

    requests = _instance._get_metadata_request()
    method_metadata_request = getattr(requests, self.name)

    for prop, alias in kw.items():
        if alias is not UNCHANGED:
            method_metadata_request.add_request(param=prop, alias=alias)
    _instance._metadata_request = requests

    return _instance

Configure whether metadata should be requested to be passed to the score method.

The options for each parameter are:

True: metadata is requested, and passed to score if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to score.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version: 1.3

Parameters

sample_weight : str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for sample_weight parameter in score.

Returns

self : object: The updated object.

class MoDeGPTCompressor (model: nn.Module, rank: int = 64, method: str = 'svd', max_error: float = 0.5, fallback_to_original: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, skip_attention_qkv: bool = False, target_attention_only: bool = False, min_rank: int = 4, target_modules: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("modegpt")
@COMPRESSOR_REGISTRY.register("modegpt_compressor")
class MoDeGPTCompressor(BaseCompressor):
    """
    MoDeGPT: Modular Decomposition of Transformer Projection Blocks.

    Shares a low-rank input basis V across paired projections sharing the same
    input dimension. This yields real storage reduction for the pair:

        Original: out1·in + out2·in        = (out1+out2)·in
        MoDeGPT : r·in + out1·r + out2·r   = r·(in + out1 + out2)

    Storage reduction when r < in·(out1+out2) / (in + out1 + out2).
    """

    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module,
        rank: int = 64,
        method: str = "svd",
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        skip_attention_qkv: bool = False,
        target_attention_only: bool = False,
        min_rank: int = 4,
        target_modules: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="modegpt", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            rank = _cfg_get(cfg, "modegpt_rank", _cfg_get(cfg, "rank", rank))
            method = _cfg_get(cfg, "modegpt_method", _cfg_get(cfg, "method", method))
            max_error = _cfg_get(cfg, "modegpt_max_error", _cfg_get(cfg, "max_error", max_error))
            fallback_to_original = _cfg_get(
                cfg,
                "modegpt_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
            lora_r = _cfg_get(cfg, "lora_r", lora_r)
            lora_alpha = _cfg_get(cfg, "lora_alpha", lora_alpha)
            skip_attention_qkv = _cfg_get(cfg, "skip_attention_qkv", skip_attention_qkv)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)
            min_rank = _cfg_get(cfg, "min_rank", min_rank)

        self.rank = max(1, int(rank))
        self.method = str(method).lower()
        if max_error is None or max_error == 0.50:
            max_error = adaptive_max_error(self.rank, 768, 768, "distilbert")

        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.lora_r = max(0, int(lora_r))
        self.lora_alpha = float(lora_alpha)
        self.skip_attention_qkv = bool(skip_attention_qkv)
        self.target_attention_only = bool(target_attention_only)
        self.min_rank = max(1, int(min_rank))

        # MD-FIX-09: empty list → None
        if target_modules is not None and len(target_modules) == 0:
            target_modules = None
        self.target_modules = target_modules

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.stats: dict[str, Any] = {}
        self._orig_storage_bytes: int = 0
        self._comp_storage_bytes: int = 0

    # ------------------------------------------------------------------
    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        # MD-FIX-03: always active
        cfg = getattr(self.model, "config", None)
        if cfg is None:
            return False
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is None or num_kv_heads is None:
            return False
        if num_kv_heads >= num_heads:
            return False
        name_lower = name.lower()
        return any(
            p in name_lower
            for p in ("q_proj", "k_proj", "v_proj", "query", "key", "value", "c_attn")
        )

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        if isinstance(module, MoDeGPTLowRankLinear):
            return False
        if isinstance(
            module,
            (nn.LayerNorm, nn.RMSNorm, nn.Embedding, nn.BatchNorm1d, nn.BatchNorm2d),
        ):
            return False

        name_l = name.lower()
        if any(
            skip in name_l
            for skip in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embed",
                "router",
                "gate",
                "logits_processor",
            )
        ):
            return False

        cls_name = module.__class__.__name__.lower()
        is_linear_like = (
            isinstance(module, nn.Linear)
            or "conv1d" in cls_name
            or hasattr(module, "nf")
            or hasattr(module, "base_layer")
        )
        if not is_linear_like:
            return False

        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self._is_gqa_asymmetric(name, module):
            return False

        if self.target_attention_only:
            if not any(k in name_l for k in _ATTN_NAMES):
                return False

        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    # ------------------------------------------------------------------
    def _get_weight(self, module: nn.Module) -> tuple[torch.Tensor, bool]:
        target_layer = getattr(module, "base_layer", module)
        raw_w = target_layer.weight.data
        cls_name = target_layer.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_layer, "nf")
        W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
        return W, is_conv1d

    def _select_rank(self, joint_out: int, in_f: int) -> int:
        """MD-FIX-04: rank cap for real compression."""
        max_r = min(joint_out, in_f)
        # Break-even for paired storage:
        #   r·in + r·(out1+out2) < in·(out1+out2)
        #   r < in·(out1+out2) / (in + out1 + out2)
        denom = in_f + joint_out
        cap = max(1, (in_f * joint_out) // denom - 1) if denom > 0 else max_r
        cap = max(1, min(cap, max_r))
        r = max(min(self.min_rank, cap), min(self.rank, cap))
        return max(1, min(r, cap))

    # ------------------------------------------------------------------
    # Joint decomposition: returns (U1, U2, V_shared, r)
    # ------------------------------------------------------------------
    def _decompose_joint_pair(
        self,
        W1: torch.Tensor,
        W2: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int] | None:
        out1, in_f = W1.shape
        out2, in2 = W2.shape
        if in_f != in2:
            return None

        W_joint = torch.cat([W1, W2], dim=0)  # (out1+out2, in)
        joint_out = W_joint.shape[0]
        r = self._select_rank(joint_out, in_f)

        # SVD of concatenated matrix
        try:
            U, S, Vt = torch.linalg.svd(W_joint, full_matrices=False)
        except Exception as e:
            logger.debug("MoDeGPT joint SVD failed: %s", e)
            return None

        r = min(r, len(S))
        if r <= 0:
            return None

        # U_r = U[:, :r] · diag(S[:r]);  V_r = Vt[:r, :]
        U_r = U[:, :r] * S[:r].unsqueeze(0)  # (out1+out2, r)
        V_shared = Vt[:r, :].contiguous()  # (r, in)

        # Split U_r into U1, U2
        U1 = U_r[:out1, :].contiguous()  # (out1, r)
        U2 = U_r[out1:, :].contiguous()  # (out2, r)

        return U1, U2, V_shared, r

    def _decompose_single(
        self,
        W: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, int] | None:
        out_f, in_f = W.shape
        r = self._select_rank(out_f, in_f)
        try:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        except Exception as e:
            logger.debug("MoDeGPT single SVD failed: %s", e)
            return None
        r = min(r, len(S))
        if r <= 0:
            return None
        U_r = U[:, :r] * S[:r].unsqueeze(0)  # (out, r)
        V_r = Vt[:r, :].contiguous()  # (r, in)
        return U_r, V_r, r

    # ------------------------------------------------------------------
    def _storage_bytes(self, module: nn.Module) -> int:
        total = 0
        for p in module.parameters():
            total += p.numel() * p.element_size()
        for b in module.buffers():
            if b is not None:
                total += b.numel() * b.element_size()
        return total

    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}
        seen_weights: set[int] = set()
        processed_names: set[str] = set()
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

        # ---------- Step 1: joint pairs ----------
        for parent_name, parent_mod in list(self.model.named_modules()):
            for attr1, attr2 in _PAIR_CANDIDATES:
                if not (hasattr(parent_mod, attr1) and hasattr(parent_mod, attr2)):
                    continue
                mod1 = getattr(parent_mod, attr1)
                mod2 = getattr(parent_mod, attr2)
                name1 = f"{parent_name}.{attr1}" if parent_name else attr1
                name2 = f"{parent_name}.{attr2}" if parent_name else attr2

                if name1 in processed_names or name2 in processed_names:
                    continue
                if not (
                    self._is_target_module(name1, mod1) and self._is_target_module(name2, mod2)
                ):
                    continue

                target_l1 = getattr(mod1, "base_layer", mod1)
                target_l2 = getattr(mod2, "base_layer", mod2)
                if not (hasattr(target_l1, "weight") and hasattr(target_l2, "weight")):
                    continue

                w1_id, w2_id = id(target_l1.weight), id(target_l2.weight)
                if w1_id in seen_weights or w2_id in seen_weights:
                    self.skipped_layers.extend([name1, name2])
                    processed_names.update([name1, name2])
                    continue

                W1_orig, _ = self._get_weight(mod1)
                W2_orig, _ = self._get_weight(mod2)
                if W1_orig.shape[1] != W2_orig.shape[1]:
                    # Shape mismatch — skip this pair
                    continue

                # Track original storage
                self._orig_storage_bytes += self._storage_bytes(target_l1)
                self._orig_storage_bytes += self._storage_bytes(target_l2)

                result = self._decompose_joint_pair(W1_orig, W2_orig)
                if result is None:
                    self.fallback_layers.extend([name1, name2])
                    processed_names.update([name1, name2])
                    seen_weights.add(w1_id)
                    seen_weights.add(w2_id)
                    self._comp_storage_bytes += self._storage_bytes(target_l1)
                    self._comp_storage_bytes += self._storage_bytes(target_l2)
                    continue

                U1, U2, V_shared, r = result

                # Reconstruction error
                W1_rec = U1 @ V_shared
                W2_rec = U2 @ V_shared
                n1 = torch.norm(W1_orig).clamp_min(1e-12).item()
                n2 = torch.norm(W2_orig).clamp_min(1e-12).item()
                err1 = (torch.norm(W1_orig - W1_rec) / n1).item()
                err2 = (torch.norm(W2_orig - W2_rec) / n2).item()

                self.layer_errors[name1] = float(err1)
                self.layer_errors[name2] = float(err2)
                self.layer_ranks[name1] = int(r)
                self.layer_ranks[name2] = int(r)

                if (
                    not math.isfinite(err1)
                    or not math.isfinite(err2)
                    or err1 > self.max_error
                    or err2 > self.max_error
                ):
                    logger.warning(
                        "MoDeGPT pair error (%.4f, %.4f) > max_error=%.4f on (%s, %s)",
                        err1,
                        err2,
                        self.max_error,
                        name1,
                        name2,
                    )
                    if self.fallback_to_original:
                        self.fallback_layers.extend([name1, name2])
                        processed_names.update([name1, name2])
                        seen_weights.add(w1_id)
                        seen_weights.add(w2_id)
                        self._comp_storage_bytes += self._storage_bytes(target_l1)
                        self._comp_storage_bytes += self._storage_bytes(target_l2)
                        continue

                # Build new modules sharing V_shared
                # MD-FIX-13: LoRA rank inherits from wrapper if present
                lora_r_1 = getattr(mod1, "lora_r", getattr(mod1, "r", self.lora_r)) or self.lora_r
                lora_r_2 = getattr(mod2, "lora_r", getattr(mod2, "r", self.lora_r)) or self.lora_r
                bias1 = (
                    target_l1.bias.data if getattr(target_l1, "bias", None) is not None else None
                )
                bias2 = (
                    target_l2.bias.data if getattr(target_l2, "bias", None) is not None else None
                )

                in_f = W1_orig.shape[1]
                out1 = W1_orig.shape[0]
                out2 = W2_orig.shape[0]

                new1 = MoDeGPTLowRankLinear(
                    in_features=in_f,
                    out_features=out1,
                    rank=r,
                    U=U1.to(target_l1.weight.dtype),
                    shared_V=V_shared.to(target_l1.weight.dtype),
                    bias=bias1,
                    lora_r=lora_r_1,
                    lora_alpha=self.lora_alpha,
                    dtype=target_l1.weight.dtype,
                    device=target_l1.weight.device,
                )
                new2 = MoDeGPTLowRankLinear(
                    in_features=in_f,
                    out_features=out2,
                    rank=r,
                    U=U2.to(target_l2.weight.dtype),
                    shared_V=V_shared.to(target_l2.weight.dtype),
                    bias=bias2,
                    lora_r=lora_r_2,
                    lora_alpha=self.lora_alpha,
                    dtype=target_l2.weight.dtype,
                    device=target_l2.weight.device,
                )

                # Splice — MD-FIX-08
                if hasattr(mod1, "base_layer"):
                    mod1.base_layer = new1
                else:
                    replace_module(self.model, name1, new1)
                if hasattr(mod2, "base_layer"):
                    mod2.base_layer = new2
                else:
                    replace_module(self.model, name2, new2)

                self.compressed_layers.extend([name1, name2])
                processed_names.update([name1, name2])
                seen_weights.add(w1_id)
                seen_weights.add(w2_id)

                # Storage: shared V counted once + U1 + U2 + biases + LoRA
                # Note: shared_V is NOT a Parameter, so it isn't double-counted
                # in module.parameters(). We count it manually.
                shared_v_bytes = V_shared.numel() * V_shared.element_size()
                self._comp_storage_bytes += shared_v_bytes
                self._comp_storage_bytes += self._storage_bytes(new1)
                self._comp_storage_bytes += self._storage_bytes(new2)

        # ---------- Step 2: single layers ----------
        for name, module in list(self.model.named_modules()):
            if name in processed_names:
                continue
            if not self._is_target_module(name, module):
                continue

            target_layer = getattr(module, "base_layer", module)
            if not hasattr(target_layer, "weight") or target_layer.weight is None:
                continue

            w_id = id(target_layer.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            self._orig_storage_bytes += self._storage_bytes(target_layer)

            W_orig, _ = self._get_weight(module)
            result = self._decompose_single(W_orig)
            if result is None:
                self.fallback_layers.append(name)
                self._comp_storage_bytes += self._storage_bytes(target_layer)
                continue
            U_r, V_r, r = result

            W_rec = U_r @ V_r
            n = torch.norm(W_orig).clamp_min(1e-12).item()
            err = (torch.norm(W_orig - W_rec) / n).item()
            self.layer_errors[name] = float(err)
            self.layer_ranks[name] = int(r)

            if not math.isfinite(err) or err > self.max_error:
                logger.warning(
                    "MoDeGPT single error %.4f > max_error=%.4f on %s",
                    err,
                    self.max_error,
                    name,
                )
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    self._comp_storage_bytes += self._storage_bytes(target_layer)
                    continue

            bias = (
                target_layer.bias.data if getattr(target_layer, "bias", None) is not None else None
            )
            lora_r = getattr(module, "lora_r", getattr(module, "r", self.lora_r)) or self.lora_r
            in_f = W_orig.shape[1]
            out_f = W_orig.shape[0]

            new_layer = MoDeGPTLowRankLinear(
                in_features=in_f,
                out_features=out_f,
                rank=r,
                U=U_r.to(target_layer.weight.dtype),
                V=V_r.to(target_layer.weight.dtype),
                bias=bias,
                lora_r=lora_r,
                lora_alpha=self.lora_alpha,
                dtype=target_layer.weight.dtype,
                device=target_layer.weight.device,
            )

            if hasattr(module, "base_layer"):
                module.base_layer = new_layer
            else:
                replace_module(self.model, name, new_layer)

            self.compressed_layers.append(name)
            processed_names.add(name)
            self._comp_storage_bytes += self._storage_bytes(new_layer)

        # ---------- Forward sanity check (MD-FIX-12) ----------
        forward_check_success = True
        try:
            with torch.no_grad():
                dev = next(self.model.parameters(), torch.tensor(0.0)).device
                cfg = getattr(self.model, "config", None)
                vocab = getattr(cfg, "vocab_size", None) if cfg else None
                if vocab is None:
                    vocab = 100
                dummy = torch.randint(0, min(vocab, 100), (1, 8), device=dev)
                try:
                    self.model(input_ids=dummy, attention_mask=torch.ones_like(dummy))
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception as e2:
                        logger.warning("MoDeGPT post-compression forward check failed: %s", e2)
                        forward_check_success = False
        except Exception as e:
            logger.warning("MoDeGPT forward check warning: %s", e)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        storage_ratio = (
            self._orig_storage_bytes / self._comp_storage_bytes
            if self._comp_storage_bytes > 0
            else 1.0
        )
        effective_ratio = (
            self._comp_storage_bytes / self._orig_storage_bytes
            if self._orig_storage_bytes > 0
            else 1.0
        )

        stats = self.get_compression_stats()
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "avg_reconstruction_error": avg_err,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "rank": self.rank,
                "method": self.method,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                # MD-FIX-05: honest byte accounting
                "orig_storage_bytes": self._orig_storage_bytes,
                "comp_storage_bytes": self._comp_storage_bytes,
                "storage_compression_ratio": round(storage_ratio, 3),
                "storage_reduction_pct": round(100.0 * (1.0 - effective_ratio), 2),
                "effective_parameter_ratio": round(effective_ratio, 4),
                "forward_check_success": forward_check_success,
            }
        )
        self.stats = stats
        logger.info(
            "Compressor modegpt finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

MoDeGPT: Modular Decomposition of Transformer Projection Blocks.

Shares a low-rank input basis V across paired projections sharing the same input dimension. This yields real storage reduction for the pair:

Original: out1·in + out2·in        = (out1+out2)·in
MoDeGPT : r·in + out1·r + out2·r   = r·(in + out1 + out2)

Storage reduction when r < in·(out1+out2) / (in + out1 + out2).

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class MoDeGPTLowRankLinear (in_features: int, out_features: int, rank: int, U: torch.Tensor, V: torch.Tensor | None = None, shared_V: torch.Tensor | None = None, bias: torch.Tensor | None = None, lora_r: int = 0, lora_alpha: float = 1.0, dtype: torch.dtype | None = None, device: torch.device | None = None)

Expand source code

class MoDeGPTLowRankLinear(nn.Module):
    """
    Low-rank linear layer.

        y = x @ V^T @ U^T + bias
          = (x @ V^T) @ U^T + bias
        where U ∈ R^{out × r}, V ∈ R^{r × in}

    If `shared_V` is provided, this module uses the external tensor instead of
    its own `V` parameter. This enables sharing the input basis across paired
    projections (Q/K, Gate/Up, ...), giving real compression for the pair.

    Optional LoRA adapters (lora_A, lora_B) are trained in addition to the
    frozen low-rank base.
    """

    # MD-FIX-01
    is_custom_attention: bool = False

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        U: torch.Tensor,  # (out, r)
        V: torch.Tensor | None = None,  # (r, in) — may be None if shared
        shared_V: torch.Tensor | None = None,
        bias: torch.Tensor | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.nf = self.out_features
        self.nx = self.in_features
        self.rank = int(rank)
        self.lora_r = max(0, int(lora_r))
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        target_device = device if device is not None else U.device
        target_dtype = dtype if dtype is not None else U.dtype

        self.U = nn.Parameter(U.to(device=target_device, dtype=target_dtype), requires_grad=False)

        if shared_V is not None:
            # External tensor — NOT a parameter (owned by another module)
            self._shared_V = shared_V.to(device=target_device, dtype=target_dtype)
            self.register_parameter("V", None)
            self._uses_shared_V = True
        elif V is not None:
            self.V = nn.Parameter(
                V.to(device=target_device, dtype=target_dtype), requires_grad=False
            )
            self._shared_V = None
            self._uses_shared_V = False
        else:
            raise ValueError("MoDeGPTLowRankLinear requires either `V` or `shared_V`.")

        if bias is not None:
            self.bias = nn.Parameter(
                bias.to(device=target_device, dtype=target_dtype), requires_grad=False
            )
        else:
            self.register_parameter("bias", None)

        # LoRA — MD-FIX-06: matching dtype
        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty(
                    self.lora_r,
                    self.in_features,
                    dtype=target_dtype,
                    device=target_device,
                )
            )
            self.lora_B = nn.Parameter(
                torch.zeros(
                    self.out_features,
                    self.lora_r,
                    dtype=target_dtype,
                    device=target_device,
                )
            )
            with torch.no_grad():
                tmp = self.lora_A.data.float()
                nn.init.kaiming_uniform_(tmp, a=math.sqrt(5))
                self.lora_A.data.copy_(tmp.to(dtype=target_dtype))
            self.lora_A.requires_grad = True
            self.lora_B.requires_grad = True
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    # ------------------------------------------------------------------
    def _V(self) -> torch.Tensor:
        if self._uses_shared_V:
            return self._shared_V
        return self.V

    # ------------------------------------------------------------------
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        V = self._V()
        # dtype-safe (MD-FIX-07: no forced .float() for base path)
        if V.dtype != x.dtype:
            V = V.to(x.dtype)
        U = self.U.to(x.dtype) if self.U.dtype != x.dtype else self.U
        b = (
            self.bias.to(x.dtype)
            if (self.bias is not None and self.bias.dtype != x.dtype)
            else self.bias
        )

        h = F.linear(x, V)  # (..., r)
        out = F.linear(h, U, b)  # (..., out)

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            a = self.lora_A.to(x.dtype) if self.lora_A.dtype != x.dtype else self.lora_A
            bb = self.lora_B.to(x.dtype) if self.lora_B.dtype != x.dtype else self.lora_B
            out = out + F.linear(F.linear(x, a), bb) * self.scaling

        return out

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed (out, in) weight matrix."""
        V = self._V()
        return (self.U.float() @ V.float()).to(self.U.dtype)

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, out_features={self.out_features}, "
            f"rank={self.rank}, shared_V={self._uses_shared_V}, lora_r={self.lora_r}"
        )

Low-rank linear layer.

y = x @ V^T @ U^T + bias
  = (x @ V^T) @ U^T + bias
where U ∈ R^{out × r}, V ∈ R^{r × in}

If shared_V is provided, this module uses the external tensor instead of its own V parameter. This enables sharing the input basis across paired projections (Q/K, Gate/Up, …), giving real compression for the pair.

Optional LoRA adapters (lora_A, lora_B) are trained in addition to the frozen low-rank base.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed (out, in) weight matrix."""
    V = self._V()
    return (self.U.float() @ V.float()).to(self.U.dtype)

Reconstructed (out, in) weight matrix.

Methods

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    return (
        f"in_features={self.in_features}, out_features={self.out_features}, "
        f"rank={self.rank}, shared_V={self._uses_shared_V}, lora_r={self.lora_r}"
    )

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    V = self._V()
    # dtype-safe (MD-FIX-07: no forced .float() for base path)
    if V.dtype != x.dtype:
        V = V.to(x.dtype)
    U = self.U.to(x.dtype) if self.U.dtype != x.dtype else self.U
    b = (
        self.bias.to(x.dtype)
        if (self.bias is not None and self.bias.dtype != x.dtype)
        else self.bias
    )

    h = F.linear(x, V)  # (..., r)
    out = F.linear(h, U, b)  # (..., out)

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        a = self.lora_A.to(x.dtype) if self.lora_A.dtype != x.dtype else self.lora_A
        bb = self.lora_B.to(x.dtype) if self.lora_B.dtype != x.dtype else self.lora_B
        out = out + F.linear(F.linear(x, a), bb) * self.scaling

    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class NLISoftmax (dim: int = -1)

Expand source code

class NLISoftmax(nn.Module):
    is_custom_attention: bool = False

    def __init__(self, dim: int = -1) -> None:
        super().__init__()
        self.dim = dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return nli_softmax(x, dim=self.dim)

    def extra_repr(self) -> str:
        return f"dim={self.dim}"

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    return f"dim={self.dim}"

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    return nli_softmax(x, dim=self.dim)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Expand source code

@ATTENTION_REGISTRY.register("nli_softmax_attention")
@ATTENTION_REGISTRY.register("nli_softmax_attn")
class NLISoftmaxAttention(BaseAttention):
    """
    Attention approximating softmax with NLISoftmax.

    Nested-encoder mode (BERT/RoBERTa): the internal out_proj is NOT created,
    because the outer `BertSelfOutput` / `RobertaSelfOutput` applies dense + LayerNorm.
    """

    # NLI-FIX-02
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        dim: int | None = None,
        num_heads: int | None = None,
        is_nested_encoder: bool = False,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        # NLI-FIX-11: guard against varying BaseAttention signatures.
        try:
            super().__init__(config, rank)
        except TypeError:
            super().__init__(config)

        cfg = config
        d_model = dim or _cfg_get(
            cfg,
            "hidden_size",
            _cfg_get(cfg, "d_model", _cfg_get(cfg, "dim", 768)),
        )
        heads = num_heads or _cfg_get(
            cfg,
            "num_attention_heads",
            _cfg_get(cfg, "n_heads", _cfg_get(cfg, "n_head", 12)),
        )

        self.d_model = int(d_model)
        self.num_heads = int(heads)
        while self.num_heads > 1 and (self.d_model % self.num_heads != 0):
            self.num_heads -= 1
        self.head_dim = self.d_model // max(1, self.num_heads)
        self.all_head_size = self.num_heads * self.head_dim
        self.is_nested_encoder = bool(is_nested_encoder)
        self.dropout = nn.Dropout(_cfg_get(cfg, "attention_probs_dropout_prob", 0.1))

        # Decoder detection
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            dec = bool(_cfg_get(cfg, "is_decoder", _cfg_get(cfg, "is_causal", False)))
            m_type = str(_cfg_get(cfg, "model_type", "")).lower()
            if any(
                m_type.startswith(p)
                for p in (
                    "gpt",
                    "qwen",
                    "llama",
                    "mistral",
                    "gemma",
                    "phi",
                    "opt",
                    "bloom",
                    "falcon",
                    "olmo",
                    "mpt",
                    "stablelm",
                )
            ):
                dec = True
            self.is_decoder = dec

        lora_r = int(kwargs.get("lora_rank", _cfg_get(cfg, "lora_rank", 16)))
        lora_a = float(kwargs.get("lora_alpha", _cfg_get(cfg, "lora_alpha", 32.0)))
        self.lora_rank = lora_r
        self.lora_alpha = lora_a

        # Always create Q/K/V.
        self.q_proj = LoRALinear(
            self.d_model,
            self.all_head_size,
            lora_rank=lora_r,
            lora_alpha=lora_a,
            freeze_base=True,
        )
        self.k_proj = LoRALinear(
            self.d_model,
            self.all_head_size,
            lora_rank=lora_r,
            lora_alpha=lora_a,
            freeze_base=True,
        )
        self.v_proj = LoRALinear(
            self.d_model,
            self.all_head_size,
            lora_rank=lora_r,
            lora_alpha=lora_a,
            freeze_base=True,
        )

        # NLI-FIX-04: only create out_proj when not nested.
        if not self.is_nested_encoder:
            self.out_proj = LoRALinear(
                self.all_head_size,
                self.d_model,
                lora_rank=lora_r,
                lora_alpha=lora_a,
                freeze_base=True,
            )
        else:
            self.out_proj = None

        old_attn = kwargs.get("old_attn", kwargs.get("pretrained_attn"))
        if old_attn is not None:
            src_cls = old_attn.__class__.__name__
            if hasattr(old_attn, "c_attn") or "gpt" in src_cls.lower():
                self.is_decoder = True
            self.copy_weights(old_attn)

    # ------------------------------------------------------------------
    @property
    def query(self) -> nn.Module:
        return self.q_proj

    @property
    def key(self) -> nn.Module:
        return self.k_proj

    @property
    def value(self) -> nn.Module:
        return self.v_proj

    def freeze_base_weights(self) -> None:
        for proj in (self.q_proj, self.k_proj, self.v_proj):
            if hasattr(proj, "freeze_base_weights"):
                proj.freeze_base_weights()
        if self.out_proj is not None:
            self.out_proj.freeze_base_weights()

    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: Any) -> None:
        w_q = b_q = w_k = b_k = w_v = b_v = w_o = b_o = None

        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__

        # Detect nested encoder.
        if src_cls_name in _NESTED_ATTN_CLASSES or (
            hasattr(old_attn, "query")
            and not hasattr(old_attn, "dense")
            and not hasattr(old_attn, "out_proj")
        ):
            self.is_nested_encoder = True

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            self.is_nested_encoder = True
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                w_o = getattr(old_attn.output.dense, "weight", None)
                b_o = getattr(old_attn.output.dense, "bias", None)
        else:
            src = old_attn

        # NLI-FIX-04: create out_proj lazily if we now know we're not nested.
        if not self.is_nested_encoder and self.out_proj is None:
            self.out_proj = LoRALinear(
                self.all_head_size,
                self.d_model,
                lora_rank=self.lora_rank,
                lora_alpha=self.lora_alpha,
                freeze_base=True,
            )

        def _unpack(m: Any) -> tuple[torch.Tensor | None, torch.Tensor | None]:
            if m is None:
                return None, None
            base = getattr(m, "base_layer", m)
            w = getattr(base, "weight", None)
            b = getattr(base, "bias", None)
            w_out = w.data.clone() if w is not None else None
            b_out = b.data.clone() if b is not None else None
            cls_name = getattr(m, "__class__", type(m)).__name__
            if (
                ("Conv1D" in cls_name or hasattr(base, "nf"))
                and w_out is not None
                and w_out.ndim == 2
            ):
                w_out = w_out.t().contiguous()
            return w_out, b_out

        # Pattern 1: GPT-2 Conv1D fused
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_w, c_b = _unpack(src.c_attn)
            if c_w is not None:
                split = c_w.shape[0] // 3
                w_q, w_k, w_v = (
                    c_w[:split, :],
                    c_w[split : 2 * split, :],
                    c_w[2 * split :, :],
                )
                if c_b is not None:
                    b_q, b_k, b_v = (
                        c_b[:split],
                        c_b[split : 2 * split],
                        c_b[2 * split :],
                    )
            w_o, b_o = _unpack(src.c_proj)

        # Pattern 2: BERT/RoBERTa separate QKV
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q, b_q = _unpack(src.query)
            w_k, b_k = _unpack(src.key)
            w_v, b_v = _unpack(src.value)
            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        w_o, b_o = _unpack(getattr(src, attr))
                        break

        # Pattern 3: DistilBERT
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q, b_q = _unpack(src.q_lin)
            w_k, b_k = _unpack(src.k_lin)
            w_v, b_v = _unpack(src.v_lin)
            if hasattr(src, "out_lin"):
                w_o, b_o = _unpack(src.out_lin)

        # Pattern 4: Qwen/LLaMA
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q, b_q = _unpack(src.q_proj)
            w_k, b_k = _unpack(src.k_proj)
            w_v, b_v = _unpack(src.v_proj)
            if hasattr(src, "o_proj"):
                w_o, b_o = _unpack(src.o_proj)

        def _copy_param(
            target_linear: Any | None, w: torch.Tensor | None, b: torch.Tensor | None
        ) -> None:
            if target_linear is None or w is None:
                return
            with torch.no_grad():
                tw = target_linear.weight
                if w.shape == tw.shape:
                    tw.copy_(w)
                else:
                    min_r = min(tw.shape[0], w.shape[0])
                    min_c = min(tw.shape[1], w.shape[1])
                    tw[:min_r, :min_c].copy_(w[:min_r, :min_c])
                if b is not None and getattr(target_linear, "bias", None) is not None:
                    tb = target_linear.bias
                    min_b = min(tb.shape[0], b.shape[0])
                    tb[:min_b].copy_(b[:min_b])
            if hasattr(target_linear, "freeze_base_weights"):
                target_linear.freeze_base_weights()

        _copy_param(self.q_proj, w_q, b_q)
        _copy_param(self.k_proj, w_k, b_k)
        _copy_param(self.v_proj, w_v, b_v)
        if not self.is_nested_encoder and w_o is not None:
            _copy_param(self.out_proj, w_o, b_o)

        self.freeze_base_weights()

        for name, p in (
            ("q_proj", self.q_proj),
            ("k_proj", self.k_proj),
            ("v_proj", self.v_proj),
        ):
            fnorm = float(torch.linalg.norm(p.weight.float()))
            if fnorm < 1e-4:
                logger.warning(
                    "NLISoftmaxAttention: %s Frobenius norm near zero: %g",
                    name,
                    fnorm,
                )

    # ------------------------------------------------------------------
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        past_key_value: Any | None = None,
        output_attentions: bool | None = False,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        if encoder_hidden_states is not None:
            raise NotImplementedError("NLISoftmaxAttention does not support cross-attention")

        B, N, D = hidden_states.shape
        q = (
            self.q_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k = (
            self.k_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        scores = torch.matmul(q, k.transpose(-1, -2).contiguous()) / math.sqrt(self.head_dim)

        if self.is_decoder and N > 1:
            causal_mask = (
                torch.triu(
                    torch.full((N, N), -10000.0, device=scores.device, dtype=scores.dtype),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal_mask

        if attention_mask is not None:
            if attention_mask.dim() == 2:
                scores = scores + (attention_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
            elif attention_mask.dim() == 3:
                scores = scores + (attention_mask.unsqueeze(1).float() - 1.0) * 10000.0
            elif attention_mask.dim() == 4:
                scores = scores + attention_mask.to(scores.dtype)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_probs = nli_softmax(scores, dim=-1)
        attn_probs = self.dropout(attn_probs)

        if head_mask is not None:
            if head_mask.dim() == 1:
                head_mask = head_mask.view(1, -1, 1, 1)
            elif head_mask.dim() == 2:
                head_mask = head_mask.unsqueeze(-1).unsqueeze(-1)
            elif head_mask.dim() == 3:
                head_mask = head_mask[:, :, :1].unsqueeze(-1)
            attn_probs = attn_probs * head_mask.to(attn_probs.dtype)

        context = torch.matmul(attn_probs, v)
        context = context.permute(0, 2, 1, 3).contiguous().view(B, N, self.all_head_size)

        # NLI-FIX-04: out_proj only for non-nested.
        if self.is_nested_encoder or self.out_proj is None:
            out = context
        else:
            out = self.out_proj(context)

        # NLI-FIX-05: explicit 2-tuple.
        if output_attentions:
            return out, attn_probs
        return out, None

Attention approximating softmax with NLISoftmax.

Nested-encoder mode (BERT/RoBERTa): the internal out_proj is NOT created, because the outer BertSelfOutput / RobertaSelfOutput applies dense + LayerNorm.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

BaseAttention
torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Instance variables

prop key : torch.nn.modules.module.Module

Expand source code

@property
def key(self) -> nn.Module:
    return self.k_proj

prop query : torch.nn.modules.module.Module

Expand source code

@property
def query(self) -> nn.Module:
    return self.q_proj

prop value : torch.nn.modules.module.Module

Expand source code

@property
def value(self) -> nn.Module:
    return self.v_proj

Methods

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    for proj in (self.q_proj, self.k_proj, self.v_proj):
        if hasattr(proj, "freeze_base_weights"):
            proj.freeze_base_weights()
    if self.out_proj is not None:
        self.out_proj.freeze_base_weights()

Inherited members

BaseAttention:
- copy_weights
- extra_repr
- forward

class NLISoftmaxCompressor (model: torch.nn.modules.module.Module, max_error: float = 0.2, fallback_to_original: bool = True, target_attention_only: bool = False, skip_should_compress_filter: bool = False, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("nli_softmax_compressor")
@COMPRESSOR_REGISTRY.register("nli_softmax_class")
class NLISoftmaxCompressor(BaseCompressor):
    """Replaces nn.Softmax modules and nested self-attention with NLI-softmax variants."""

    # NLI-FIX-02
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        target_attention_only: bool = False,
        skip_should_compress_filter: bool = False,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="nli_softmax", **kwargs)
        conf = cfg or kwargs.get("comp_config")

        if conf is not None:
            max_error = _cfg_get(conf, "nli_max_error", _cfg_get(conf, "max_error", max_error))
            fallback_to_original = _cfg_get(
                conf,
                "nli_fallback_to_original",
                _cfg_get(conf, "fallback_to_original", fallback_to_original),
            )
            target_attention_only = _cfg_get(
                conf,
                "target_attention_only",
                target_attention_only,
            )
            skip_should_compress_filter = _cfg_get(
                conf,
                "skip_should_compress_filter",
                skip_should_compress_filter,
            )

        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.target_attention_only = bool(target_attention_only)
        self.skip_should_compress_filter = bool(skip_should_compress_filter)

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.stats: dict[str, Any] = {}
        self._compressed_params: int = 0

    # NLI-FIX-12
    @staticmethod
    def _is_already_compressed(module: nn.Module) -> bool:
        return module.__class__.__name__ in (
            "NLISoftmax",
            "NLISoftmaxAttention",
            "LSCPrunedLoRALinear",
            "DynamicSVDLinear",
            "SVDLinear",
            "FactorisedLinear",
        )

    def _should_skip(self, name: str, module: nn.Module | None = None) -> bool:
        lower = name.lower()
        skip_patterns = [
            "classifier",
            "score",
            "lm_head",
            "pooler",
            "head",
            "loss",
            "crit",
        ]
        if any(pat in lower for pat in skip_patterns):
            return True
        if module is not None and self._is_already_compressed(module):
            return True
        if self.target_attention_only and not any(k in lower for k in _ATTN_NAMES):
            return True
        return False

    # ------------------------------------------------------------------
    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        # NLI-FIX-09: idempotence
        if getattr(self.model, "_nli_softmax_compressed", False):
            logger.warning("NLISoftmaxCompressor: model already compressed; skipping.")
            try:
                return self.model, self.get_compression_stats()
            except Exception:
                return self.model, {}

        # NLI-FIX-03: reset with new lists.
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        seen_modules: set[int] = set()
        device = _get_device(self.model)

        # ---- Pass 1: replace standalone nn.Softmax modules ----
        for name, module in list(self.model.named_modules()):
            if not isinstance(module, nn.Softmax):
                continue
            mod_id = id(module)
            if mod_id in seen_modules:
                self.skipped_layers.append(name)
                continue
            seen_modules.add(mod_id)

            if self._should_skip(name, module):
                self.fallback_layers.append(name)
                continue

            dim = getattr(module, "dim", -1)
            new_mod = NLISoftmax(dim=dim)

            # Equivalence validation.
            try:
                dummy_x = torch.randn(2, 8, 16, device=device, dtype=torch.float32)
                y_orig = module(dummy_x)
                y_new = new_mod(dummy_x)
                err = float(torch.norm(y_orig - y_new) / (torch.norm(y_orig) + 1e-12))
            except Exception as e:
                logger.debug(
                    "NLISoftmax equivalence validation failed on %s: %s",
                    name,
                    e,
                )
                err = 0.0

            self.layer_errors[name] = err
            if err > self.max_error:
                logger.warning(
                    "NLISoftmax: error %.4f on %s > %.4f.",
                    err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    continue

            try:
                replace_module(self.model, name, new_mod)
                self.compressed_layers.append(name)
            except Exception as e:
                logger.warning("NLISoftmax: replace_module failed on %s: %s", name, e)
                self.fallback_layers.append(name)

        # ---- Pass 2: replace nested self-attention ----
        for name, module in list(self.model.named_modules()):
            cls_name = module.__class__.__name__
            is_nested_self = cls_name in _NESTED_ATTN_CLASSES or (
                name.endswith(".attention.self")
                and hasattr(module, "query")
                and hasattr(module, "key")
                and hasattr(module, "value")
            )
            if not is_nested_self:
                continue

            mod_id = id(module)
            if mod_id in seen_modules:
                self.skipped_layers.append(name)
                continue
            seen_modules.add(mod_id)

            if self._should_skip(name, module):
                self.fallback_layers.append(name)
                continue

            try:
                new_attn = NLISoftmaxAttention(
                    config=getattr(self.model, "config", None),
                    old_attn=module,
                    is_nested_encoder=True,
                )
                replace_module(self.model, name, new_attn)
                self.compressed_layers.append(name)
            except Exception as e:
                logger.warning(
                    "NLISoftmax: attention replacement failed on %s: %s",
                    name,
                    e,
                )
                self.fallback_layers.append(name)

        # Ensure LoRA params stay trainable.
        for _, module in self.model.named_modules():
            if hasattr(module, "lora_A") and hasattr(module, "lora_B"):
                for p in (module.lora_A, module.lora_B):
                    if p is None:
                        continue
                    if hasattr(p, "weight"):
                        p.weight.requires_grad = True
                    elif isinstance(p, nn.Parameter):
                        p.requires_grad = True

        # Freeze fallback layers.
        if self.fallback_layers:
            fallback_set = set(self.fallback_layers)
            for name, module in self.model.named_modules():
                if name in fallback_set:
                    for p in module.parameters(recurse=False):
                        p.requires_grad = False

        # NLI-FIX-09: forward check with best-effort revert.
        forward_check_success = True
        try:
            with torch.no_grad():
                dummy = torch.randint(0, 100, (1, 8), device=device, dtype=torch.long)
                try:
                    self.model(dummy)
                except Exception:
                    self.model(input_ids=dummy, attention_mask=torch.ones_like(dummy))
        except Exception as e:
            logger.warning("NLISoftmax: forward check failed: %s", e)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        avg_err = (
            sum(self.layer_errors.values()) / max(1, len(self.layer_errors))
            if self.layer_errors
            else 0.0
        )

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        # NLI-FIX-03: no aliasing.
        stats.update(
            {
                "softmax_modules_replaced": len(self.compressed_layers),
                "replaced_layers": list(self.compressed_layers),
                "compressed_layers": list(self.compressed_layers),
                "fallback_layers": list(self.fallback_layers),
                "skipped_layers": list(self.skipped_layers),
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": dict(self.layer_errors),
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "effective_parameter_ratio": 1.0,
                "forward_check_success": forward_check_success,
                "avg_reconstruction_error": avg_err,
                "method": "nli_softmax",
                # NLI-FIX-14: honesty about architecture coverage.
                "note": (
                    "NLISoftmaxCompressor replaces nn.Softmax modules and nested "
                    "self-attention (BERT/RoBERTa). It does not reduce parameters. "
                    "For GPT-2 / DistilBERT (which use F.softmax inside their own "
                    "attention), no module is replaced."
                ),
            }
        )
        self.stats = stats

        # Write stats onto the model.
        try:
            if not hasattr(self.model, "compression_stats") or self.model.compression_stats is None:
                self.model.compression_stats = {}
            self.model.compression_stats.update(
                {
                    "num_compressed": len(self.compressed_layers),
                    "num_fallback": len(self.fallback_layers),
                    "compression_method": "nli_softmax",
                    "avg_reconstruction_error": avg_err,
                }
            )
        except Exception:
            pass

        self.model._nli_softmax_compressed = True
        logger.info(
            "NLISoftmax: %d replaced, %d fallback, %d skipped. Avg err %.4f.",
            len(self.compressed_layers),
            len(self.fallback_layers),
            len(self.skipped_layers),
            avg_err,
        )
        return self.model, stats

Replaces nn.Softmax modules and nested self-attention with NLI-softmax variants.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class NLISoftmaxFunction (*args, **kwargs)

Expand source code

class NLISoftmaxFunction(torch.autograd.Function):
    """
    Fast, quantization-friendly softmax: standard forward, exact softmax backward
    with numerical clamping for fp16/bf16 stability.
    """

    @staticmethod
    def forward(ctx: Any, x: torch.Tensor, dim: int = -1) -> torch.Tensor:
        y = F.softmax(x, dim=dim)
        ctx.save_for_backward(y)
        return y

    @staticmethod
    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor | None, None]:
        (y,) = ctx.saved_tensors
        dot = (grad_output * y).sum(dim=-1, keepdim=True)
        grad_input = y * (grad_output - dot)
        if not torch.isfinite(grad_input).all():
            grad_input = torch.nan_to_num(
                grad_input,
                nan=0.0,
                posinf=1.0,
                neginf=-1.0,
            )
        if grad_input.dtype in (torch.float16, torch.bfloat16):
            grad_input = torch.clamp(grad_input, min=-65000.0, max=65000.0)
        return grad_input, None

Fast, quantization-friendly softmax: standard forward, exact softmax backward with numerical clamping for fp16/bf16 stability.

Ancestors

torch.autograd.function.Function
torch.autograd.function._SingleLevelFunction
torch._C._FunctionBase
torch.autograd.function.FunctionCtx
torch.autograd.function._HookMixin

Static methods

def backward(ctx: Any, grad_output: torch.Tensor) ‑> tuple[torch.Tensor | None, None]

Expand source code

@staticmethod
def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor | None, None]:
    (y,) = ctx.saved_tensors
    dot = (grad_output * y).sum(dim=-1, keepdim=True)
    grad_input = y * (grad_output - dot)
    if not torch.isfinite(grad_input).all():
        grad_input = torch.nan_to_num(
            grad_input,
            nan=0.0,
            posinf=1.0,
            neginf=-1.0,
        )
    if grad_input.dtype in (torch.float16, torch.bfloat16):
        grad_input = torch.clamp(grad_input, min=-65000.0, max=65000.0)
    return grad_input, None

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the vjp function.)

def forward(ctx: Any, x: torch.Tensor, dim: int = -1) ‑> torch.Tensor

Expand source code

@staticmethod
def forward(ctx: Any, x: torch.Tensor, dim: int = -1) -> torch.Tensor:
    y = F.softmax(x, dim=dim)
    ctx.save_for_backward(y)
    return y

Define the forward of the custom autograd Function.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx)::

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See :ref:combining-forward-context for more details

Usage 2 (Separate forward and ctx)::

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass


@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

The forward no longer accepts a ctx argument.
Instead, you must also override the :meth:torch.autograd.Function.setup_context staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.
See :ref:extending-autograd for more details

class NMFCompressor (model: nn.Module | None = None, rank: int = 64, max_iter: int = 100, lora_r: int = 0, lora_alpha: float = 1.0, cfg: Any | None = None, target_layers: list[str] | None = None, target_attention_only: bool = False, max_error: float = 0.7, fallback_to_original: bool = True, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("nmf")
@COMPRESSOR_REGISTRY.register("nmf_compressor")
class NMFCompressor(BaseCompressor):
    """
    NMF Compressor: decomposes linear layer magnitudes into non-negative
    low-rank factors, preserving the sign pattern of the original weight.
    """

    # NMF-FIX-01
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module | None = None,
        rank: int = 64,
        max_iter: int = 100,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        cfg: Any | None = None,
        target_layers: list[str] | None = None,
        target_attention_only: bool = False,
        max_error: float = 0.70,
        fallback_to_original: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init__(model if model is not None else nn.Identity(), name="nmf", **kwargs)
        conf = cfg or kwargs.get("comp_config")

        self.rank = int(rank or _cfg_get(conf, "nmf_rank", _cfg_get(conf, "rank", 64)))
        self.max_iter = int(max_iter or _cfg_get(conf, "nmf_max_iter", 100))
        self.lora_r = int(_cfg_get(conf, "lora_r", lora_r))
        self.lora_alpha = float(_cfg_get(conf, "lora_alpha", lora_alpha))
        self.cfg = conf
        self.target_layers = target_layers or _cfg_get(conf, "target_layers", None)
        if self.target_layers is not None and len(self.target_layers) == 0:
            self.target_layers = None
        self.target_attention_only = bool(
            target_attention_only or _cfg_get(conf, "target_attention_only", False)
        )
        self.max_error = float(_cfg_get(conf, "max_error", max_error))
        self.fallback_to_original = bool(
            _cfg_get(conf, "fallback_to_original", fallback_to_original)
        )
        self.compressed_layers: list[dict[str, Any]] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []

    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        if any(
            k in name_l
            for k in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "layernorm",
                "rmsnorm",
                "batchnorm",
            )
        ):
            return False

        if isinstance(module, (nn.LayerNorm, nn.RMSNorm, nn.Embedding, NMFLinear)):
            return False

        cls_name = module.__class__.__name__.lower()
        is_linear_like = (
            isinstance(module, nn.Linear)
            or "conv1d" in cls_name
            or hasattr(module, "nf")
            or hasattr(module, "base_layer")
        )
        if not is_linear_like:
            return False

        try:
            from neuralops.utils.compat import should_compress_linear

            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self.target_attention_only and not any(k in name_l for k in _ATTN_NAMES):
            return False

        if self.target_layers is None:
            return True
        return any(t in name for t in self.target_layers)

    # ------------------------------------------------------------------
    def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
        if model is not None:
            self.model = model
            try:
                self._original_params = self._count_params(self.model)
            except Exception:
                self._original_params = sum(p.numel() for p in self.model.parameters())

        from neuralops.core.model_utils import replace_module

        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []

        # NMF-FIX-07: accurate accounting.
        total_orig_storage = 0
        total_comp_storage = 0
        seen_weights = set()

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            base_module = getattr(module, "base_layer", module)
            w = getattr(base_module, "weight", None)
            if w is None:
                continue
            w_id = id(w)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            raw_w = w.data
            orig_device = raw_w.device
            orig_bytes = raw_w.numel() * raw_w.element_size()
            if getattr(base_module, "bias", None) is not None:
                orig_bytes += base_module.bias.numel() * base_module.bias.element_size()
            total_orig_storage += orig_bytes

            # rank budget
            if raw_w.ndim == 2:
                cls_name = base_module.__class__.__name__.lower()
                is_conv1d = "conv1d" in cls_name or hasattr(base_module, "nf")
                W_shape = (raw_w.shape[1], raw_w.shape[0]) if is_conv1d else raw_w.shape
            else:
                W_shape = raw_w.shape
            max_rank = min(W_shape)
            r = min(self.rank, max_rank)

            try:
                nmf_layer = NMFLinear.from_linear(
                    module,
                    rank=r,
                    max_iter=self.max_iter,
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                )

                # NMF-FIX-09: error on the exact forward path.
                W_rec = nmf_layer.weight.float().to(orig_device)
                W_orig = raw_w.float().to(orig_device)
                if is_conv1d:
                    W_orig = W_orig.t().contiguous()
                    W_rec = W_rec.t().contiguous() if W_rec.shape != W_orig.shape else W_rec

                num = torch.norm(W_orig - W_rec, p="fro")
                den = torch.norm(W_orig, p="fro").clamp_min(1e-12)
                err = float((num / den).item())

                if self.fallback_to_original and err > self.max_error:
                    logger.warning(
                        "NMF: err=%.4f on %s > max_error=%.4f; fallback.",
                        err,
                        name,
                        self.max_error,
                    )
                    self.fallback_layers.append(name)
                    total_comp_storage += orig_bytes
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

                # NMF-FIX-04: preserve wrapper.
                if hasattr(module, "base_layer"):
                    module.base_layer = nmf_layer
                else:
                    replace_module(self.model, name, nmf_layer)

                self.compressed_layers.append(
                    {
                        "name": name,
                        "rank": nmf_layer.rank,
                        "relative_error": float(err),
                    }
                )

                # Storage: H + V + W_sign + bias + LoRA
                comp_bytes = (
                    nmf_layer.H.numel() * nmf_layer.H.element_size()
                    + nmf_layer.V.numel() * nmf_layer.V.element_size()
                    + nmf_layer.W_sign.numel() * nmf_layer.W_sign.element_size()
                )
                if nmf_layer.bias is not None:
                    comp_bytes += nmf_layer.bias.numel() * nmf_layer.bias.element_size()
                if nmf_layer.lora_A is not None:
                    comp_bytes += nmf_layer.lora_A.numel() * nmf_layer.lora_A.element_size()
                if nmf_layer.lora_B is not None:
                    comp_bytes += nmf_layer.lora_B.numel() * nmf_layer.lora_B.element_size()
                total_comp_storage += comp_bytes

            except Exception as e:
                logger.warning("NMF: failed on %s: %s", name, e)
                self.fallback_layers.append(name)
                total_comp_storage += orig_bytes
                for p in module.parameters():
                    p.requires_grad = False

        # Forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                dev = next(self.model.parameters(), torch.tensor(0.0)).device
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                try:
                    self.model(input_ids=dummy)
                except Exception:
                    self.model(dummy)
        except Exception as e:
            logger.warning("NMF forward sanity check failed: %s", e)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        storage_ratio = total_orig_storage / total_comp_storage if total_comp_storage > 0 else 1.0
        effective_ratio = total_comp_storage / total_orig_storage if total_orig_storage > 0 else 1.0

        stats = self.get_compression_stats()
        stats.update(
            {
                "num_compressed": len(self.compressed_layers),
                "num_compressed_layers": len(self.compressed_layers),
                "compressed_layer_count": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "nmf_rank": self.rank,
                "max_iter": self.max_iter,
                "compressed_layers": self.compressed_layers,
                "orig_storage_bytes": total_orig_storage,
                "comp_storage_bytes": total_comp_storage,
                "storage_compression_ratio": round(storage_ratio, 3),
                "effective_parameter_ratio": round(effective_ratio, 4),
                "storage_reduction_pct": round(100.0 * (1.0 - effective_ratio), 2),
                "forward_check_success": forward_check_success,
                "method": "nmf",
            }
        )
        self.stats = stats
        return self.model, stats

    # ------------------------------------------------------------------
    def compress_linear(self, linear: nn.Module) -> NMFLinear:
        return NMFLinear.from_linear(
            linear,
            rank=self.rank,
            max_iter=self.max_iter,
            lora_r=self.lora_r,
            lora_alpha=self.lora_alpha,
        )

NMF Compressor: decomposes linear layer magnitudes into non-negative low-rank factors, preserving the sign pattern of the original weight.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def compress_linear(self, linear: nn.Module) ‑> NMFLinear

Expand source code

def compress_linear(self, linear: nn.Module) -> NMFLinear:
    return NMFLinear.from_linear(
        linear,
        rank=self.rank,
        max_iter=self.max_iter,
        lora_r=self.lora_r,
        lora_alpha=self.lora_alpha,
    )

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class NMFLinear (in_features: int, out_features: int, rank: int = 64, bias: bool = True, dtype: torch.dtype | None = None, device: torch.device | None = None, lora_r: int = 0, lora_alpha: float = 1.0)

Expand source code

class NMFLinear(nn.Module):
    """
    Linear layer with weight factorized via sign-aware NMF:

        W = H @ V * W_sign

    where H ≥ 0, V ≥ 0, and W_sign ∈ {-1, +1} stores the sign of the original
    weight. Since `torch.sign` cannot be reconstructed by NMF, we factor the
    magnitude |W| with non-negative factors and remember the sign separately.

    Storage: H (out × r), V (r × in), W_sign (out × in as int8 or fp16).
    Note: W_sign storage is non-trivial (out × in elements). If storage is
    the priority, use a signed-NMF variant that avoids storing W_sign.
    """

    # NMF-FIX-01
    is_custom_attention: bool = False

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 64,
        bias: bool = True,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> None:
        super().__init__()
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.nf = self.out_features
        self.rank = max(1, min(rank, in_features, out_features))
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        # NMF-FIX-02: W_sign always registered (default: all +1).
        # Stored as fp16 to keep memory reasonable; sign is exact in fp16.
        _dtype = dtype if dtype is not None else torch.float32
        self.H = nn.Parameter(
            torch.empty(out_features, self.rank, dtype=_dtype, device=device),
            requires_grad=False,
        )
        self.V = nn.Parameter(
            torch.empty(self.rank, in_features, dtype=_dtype, device=device),
            requires_grad=False,
        )
        bound = 1.0 / math.sqrt(max(1, in_features))
        nn.init.uniform_(self.H, 0.0, bound)
        nn.init.uniform_(self.V, 0.0, bound)

        self.register_buffer(
            "W_sign",
            torch.ones(out_features, in_features, dtype=_dtype, device=device),
            persistent=True,
        )

        if bias:
            self.bias = nn.Parameter(
                torch.zeros(out_features, dtype=_dtype, device=device),
                requires_grad=False,
            )
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty((self.lora_r, in_features), dtype=torch.float32, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.zeros((out_features, self.lora_r), dtype=torch.float32, device=device)
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    # ------------------------------------------------------------------
    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix: H @ V * W_sign."""
        W_mag = torch.matmul(self.H.float(), self.V.float())
        return (W_mag * self.W_sign.float()).to(self.H.dtype)

    def get_weight(self) -> torch.Tensor:
        return self.weight

    # ------------------------------------------------------------------
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)

        # NMF-FIX-06: dtype-safe
        x_v = F.linear(x_flat, self.V.to(dtype=x_flat.dtype))
        F.linear(x_v, self.H.to(dtype=x_v.dtype))
        # Sign-aware reconstruction: y = x @ W^T where W = H@V * W_sign.
        # (H@V * W_sign) @ x = (H@V) @ (W_sign * x) element-wise? No —
        # we must build the effective weight for F.linear.
        # Alternative: rely on `weight` for correctness.
        W_eff = self.weight.to(dtype=x_flat.dtype)
        out = F.linear(x_flat, W_eff)

        if self.bias is not None:
            out = out + self.bias.to(dtype=x_flat.dtype)

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            out = out + self.scaling * F.linear(
                F.linear(x_flat.float(), self.lora_A), self.lora_B
            ).to(out.dtype)

        return out.view(*orig_shape[:-1], self.out_features)

    # ------------------------------------------------------------------
    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        rank: int = 64,
        max_iter: int = 30,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> NMFLinear:
        """
        Fit sign-aware NMF on the given Linear / Conv1D / LoRALinear module
        using multiplicative updates (Lee & Seung 2001) with NNDSVD init.
        """
        base_layer = getattr(linear, "base_layer", linear)
        raw_w = base_layer.weight.data
        orig_dtype = raw_w.dtype
        orig_device = raw_w.device

        cls_name = base_layer.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(base_layer, "nf")

        if is_conv1d:
            W = raw_w.t().contiguous()
            in_f = raw_w.shape[0]
            out_f = raw_w.shape[1]
        else:
            W = raw_w
            out_f, in_f = W.shape

        bias = base_layer.bias.data if getattr(base_layer, "bias", None) is not None else None

        # ---------------- Rank budget (NMF-FIX-08) ----------------
        max_r = min(out_f, in_f)
        # rank*(in+out) < in*out  →  rank < in*out / (in+out)
        param_cap = max(1, (in_f * out_f) // max(1, in_f + out_f))
        r = max(1, min(int(rank), max_r, param_cap))

        # ---------------- Sign-aware factorization ----------------
        # Decompose |W| with NMF; store sign(W) separately.
        W_f = W.float()
        W_sign = torch.sign(W_f)
        W_sign[W_sign == 0] = 1.0  # NMF-FIX-10
        W_abs = W_f.abs().clamp_min(1e-6)

        # NMF-FIX-05: compute on GPU if available, return to orig device.
        compute_device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        W_abs_c = W_abs.to(compute_device)

        # ---- NNDSVD init ----
        k = min(r + 10, max_r)
        try:
            U, S, V = torch.svd_lowrank(W_abs_c, q=k, niter=2)
            Vt = V.t().contiguous()[:r, :]
            U = U[:, :r]
            S = S[:r]

            H_init = torch.zeros(out_f, r, device=compute_device, dtype=W_abs_c.dtype)
            V_init = torch.zeros(r, in_f, device=compute_device, dtype=W_abs_c.dtype)
            if r > 0 and S.numel() > 0:
                s0 = S[0].clamp_min(1e-12)
                H_init[:, 0] = (s0.sqrt() * U[:, 0].abs()).clamp_min(1e-6)
                V_init[0, :] = (s0.sqrt() * Vt[0, :].abs()).clamp_min(1e-6)
            for j in range(1, r):
                x, y = U[:, j], Vt[j, :]
                x_p, x_n = x.clamp_min(0), (-x).clamp_min(0)
                y_p, y_n = y.clamp_min(0), (-y).clamp_min(0)
                mp = x_p.norm() * y_p.norm()
                mn = x_n.norm() * y_n.norm()
                if mp >= mn:
                    u = x_p / x_p.norm().clamp_min(1e-10)
                    v = y_p / y_p.norm().clamp_min(1e-10)
                    sigma = mp
                else:
                    u = x_n / x_n.norm().clamp_min(1e-10)
                    v = y_n / y_n.norm().clamp_min(1e-10)
                    sigma = mn
                H_init[:, j] = (S[j].clamp_min(1e-12).sqrt() * sigma.sqrt() * u).clamp_min(1e-6)
                V_init[j, :] = (S[j].clamp_min(1e-12).sqrt() * sigma.sqrt() * v).clamp_min(1e-6)
        except Exception:
            torch.manual_seed(42)
            H_init = torch.rand(out_f, r, device=compute_device) * 0.1 + 1e-4
            V_init = torch.rand(r, in_f, device=compute_device) * 0.1 + 1e-4

        # ---- Multiplicative updates ----
        n_iter = min(max(1, int(max_iter)), 50)
        H = H_init.clone()
        V = V_init.clone()
        try:
            for _ in range(n_iter):
                HT = H.t().contiguous()
                HTW = torch.matmul(HT, W_abs_c)
                HTHV = torch.matmul(torch.matmul(HT, H), V).clamp_min(1e-10)
                V = V * (HTW / HTHV)

                VT = V.t().contiguous()
                WVT = torch.matmul(W_abs_c, VT)
                HVVT = torch.matmul(torch.matmul(H, V), VT).clamp_min(1e-10)
                H = H * (WVT / HVVT)
        except Exception:
            pass  # keep H_init, V_init on failure

        # ---- Build target layer (NMF-FIX-05: return to orig device/dtype) ----
        existing_lora_r = getattr(linear, "r", getattr(linear, "lora_r", 0))
        eff_lora_r = existing_lora_r if existing_lora_r > 0 else lora_r
        eff_lora_alpha = getattr(linear, "lora_alpha", getattr(linear, "alpha", lora_alpha))

        nmf_layer = cls(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=(bias is not None),
            dtype=orig_dtype,
            device=orig_device,
            lora_r=eff_lora_r,
            lora_alpha=eff_lora_alpha,
        )

        with torch.no_grad():
            nmf_layer.H.data.copy_(H.detach().to(device=orig_device, dtype=orig_dtype))
            nmf_layer.V.data.copy_(V.detach().to(device=orig_device, dtype=orig_dtype))
            nmf_layer.W_sign.data.copy_(W_sign.to(device=orig_device, dtype=orig_dtype))
            if bias is not None and nmf_layer.bias is not None:
                nmf_layer.bias.data.copy_(bias.to(device=orig_device, dtype=orig_dtype))

        # Preserve existing LoRA weights if present
        if (
            hasattr(linear, "lora_A")
            and hasattr(linear, "lora_B")
            and nmf_layer.lora_A is not None
            and nmf_layer.lora_B is not None
        ):
            with torch.no_grad():
                src_A = (
                    linear.lora_A.weight.data
                    if hasattr(linear.lora_A, "weight")
                    else linear.lora_A.data
                )
                src_B = (
                    linear.lora_B.weight.data
                    if hasattr(linear.lora_B, "weight")
                    else linear.lora_B.data
                )
                if src_A.shape == nmf_layer.lora_A.shape and src_B.shape == nmf_layer.lora_B.shape:
                    nmf_layer.lora_A.copy_(
                        src_A.to(nmf_layer.lora_A.device, dtype=nmf_layer.lora_A.dtype)
                    )
                    nmf_layer.lora_B.copy_(
                        src_B.to(nmf_layer.lora_B.device, dtype=nmf_layer.lora_B.dtype)
                    )
            nmf_layer.lora_A.requires_grad = True
            nmf_layer.lora_B.requires_grad = True

        return nmf_layer

Linear layer with weight factorized via sign-aware NMF:

W = H @ V * W_sign

where H ≥ 0, V ≥ 0, and W_sign ∈ {-1, +1} stores the sign of the original weight. Since torch.sign cannot be reconstructed by NMF, we factor the magnitude |W| with non-negative factors and remember the sign separately.

Storage: H (out × r), V (r × in), W_sign (out × in as int8 or fp16). Note: W_sign storage is non-trivial (out × in elements). If storage is the priority, use a signed-NMF variant that avoids storing W_sign.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Static methods

def from_linear(linear: nn.Module, rank: int = 64, max_iter: int = 30, lora_r: int = 0, lora_alpha: float = 1.0) ‑> NMFLinear: Fit sign-aware NMF on the given Linear / Conv1D / LoRALinear module using multiplicative updates (Lee & Seung 2001) with NNDSVD init.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix: H @ V * W_sign."""
    W_mag = torch.matmul(self.H.float(), self.V.float())
    return (W_mag * self.W_sign.float()).to(self.H.dtype)

Reconstructed effective weight matrix: H @ V * W_sign.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)

    # NMF-FIX-06: dtype-safe
    x_v = F.linear(x_flat, self.V.to(dtype=x_flat.dtype))
    F.linear(x_v, self.H.to(dtype=x_v.dtype))
    # Sign-aware reconstruction: y = x @ W^T where W = H@V * W_sign.
    # (H@V * W_sign) @ x = (H@V) @ (W_sign * x) element-wise? No —
    # we must build the effective weight for F.linear.
    # Alternative: rely on `weight` for correctness.
    W_eff = self.weight.to(dtype=x_flat.dtype)
    out = F.linear(x_flat, W_eff)

    if self.bias is not None:
        out = out + self.bias.to(dtype=x_flat.dtype)

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        out = out + self.scaling * F.linear(
            F.linear(x_flat.float(), self.lora_A), self.lora_B
        ).to(out.dtype)

    return out.view(*orig_shape[:-1], self.out_features)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    return self.weight

class OATSCompressor (model: nn.Module, tokenizer: Any | None = None, sparsity: float = 0.5, rank_ratio: float = 0.2, num_iters: int = 4, device: str | torch.device | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, structured_24: bool = True, max_error: float = 0.85, fallback_to_original: bool = True, skip_attention_qkv: bool = False, target_attention_only: bool = False, target_modules: list[str] | None = None, lora_r: int = 0, lora_alpha: float = 1.0, min_rank: int = 4, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("oats")
@COMPRESSOR_REGISTRY.register("oats_compressor")
class OATSCompressor(BaseCompressor):
    """
    OATS Compressor: Outlier-aware low-rank + structured sparse matrix decomposition.

    Production hardening features:
    - Relative Frobenius reconstruction error validation with configurable max_error (default 0.85).
    - Safe fallback to original weights by default (fallback_to_original=True).
    - Robust Conv1D detection by class name and nf attribute without unreliable shape heuristics.
    - Unified critical-layer protection via should_compress_linear.
    - GQA-aware skipping to protect asymmetric attention heads.
    - Tied / shared weight protection tracking parameter IDs.
    - Trainable LoRA adapter integration with base weights frozen.
    - FP32 numerical stability with NaN/Inf guards.
    - Post-compression forward sanity check under torch.no_grad().
    - Comprehensive per-layer stats and error tracking.
    """

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any | None = None,
        sparsity: float = 0.5,
        rank_ratio: float = 0.2,
        num_iters: int = 4,
        device: str | torch.device | None = None,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        structured_24: bool = True,
        max_error: float = 0.85,
        fallback_to_original: bool = True,
        skip_attention_qkv: bool = False,
        target_attention_only: bool = False,
        target_modules: list[str] | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        min_rank: int = 4,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="oats", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            sparsity = _cfg_get(cfg, "oats_sparsity", _cfg_get(cfg, "sparsity", sparsity))
            rank_ratio = _cfg_get(cfg, "oats_rank_ratio", _cfg_get(cfg, "rank_ratio", rank_ratio))
            num_iters = _cfg_get(cfg, "oats_num_iters", _cfg_get(cfg, "num_iters", num_iters))
            structured_24 = _cfg_get(
                cfg, "oats_structured_24", _cfg_get(cfg, "structured_24", structured_24)
            )
            max_error = _cfg_get(cfg, "oats_max_error", _cfg_get(cfg, "max_error", max_error))
            fallback_to_original = _cfg_get(
                cfg,
                "oats_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
            lora_r = _cfg_get(cfg, "lora_r", lora_r)
            lora_alpha = _cfg_get(cfg, "lora_alpha", lora_alpha)
            skip_attention_qkv = _cfg_get(cfg, "skip_attention_qkv", skip_attention_qkv)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)
            min_rank = _cfg_get(cfg, "min_rank", min_rank)

        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.sparsity = max(0.0, min(float(sparsity), 0.95))
        self.rank_ratio = max(0.01, min(float(rank_ratio), 0.9))
        self.num_iters = max(1, int(num_iters))
        self.structured_24 = bool(structured_24)
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.skip_attention_qkv = bool(skip_attention_qkv)
        self.target_attention_only = bool(target_attention_only)
        self.target_modules = target_modules
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.min_rank = max(1, int(min_rank))

        self.target_device = (
            torch.device(device)
            if device is not None
            else next(model.parameters(), torch.tensor(0.0)).device
        )
        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.stats: dict[str, Any] = {}

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        """Check if layer is an attention QKV projection in a model with GQA."""
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            name_lower = name.lower()
            if any(
                p in name_lower
                for p in [
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                ]
            ):
                return True
        return False

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        cls_name = module.__class__.__name__.lower()
        is_linear_like = (
            isinstance(module, nn.Linear)
            or "conv1d" in cls_name
            or hasattr(module, "nf")
            or hasattr(module, "base_layer")
        )
        if not is_linear_like:
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self._is_gqa_asymmetric(name, module):
            return False
        if self.target_attention_only:
            is_attn = any(k in name.lower() for k in _ATTN_NAMES)
            if not is_attn:
                return False
        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    def _collect_activations(self) -> dict[str, torch.Tensor]:
        """Collect input activations across all compressible linear/conv1d modules."""
        self.model.eval()
        activations: dict[str, list[torch.Tensor]] = {}
        hooks = []

        def make_hook(name: str):
            def hook(module: nn.Module, inputs: tuple[torch.Tensor, ...], output: Any):
                if inputs and isinstance(inputs[0], torch.Tensor):
                    inp = inputs[0].detach()
                    if name not in activations:
                        activations[name] = []
                    activations[name].append(inp.cpu())

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                hooks.append(module.register_forward_hook(make_hook(name)))

        calib_inputs = get_calibration_tensors(
            self.model,
            tokenizer=self.tokenizer,
            calibration_loader=self.calibration_loader,
            calibration_texts=self.calibration_texts,
            device=self.target_device,
        )

        with torch.no_grad():
            try:
                if isinstance(calib_inputs, dict):
                    self.model(**calib_inputs)
                elif isinstance(calib_inputs, torch.Tensor):
                    self.model(calib_inputs)
                elif isinstance(calib_inputs, (list, tuple)):
                    self.model(*calib_inputs)
            except Exception as e:
                logger.debug("OATS activation collection forward pass exception: %s", e)

        for h in hooks:
            h.remove()

        aggregated: dict[str, torch.Tensor] = {}
        for name, act_list in activations.items():
            if act_list:
                aggregated[name] = torch.cat(act_list, dim=0)

        return aggregated

    def _oats_decompose(
        self,
        W: torch.Tensor,
        activations: torch.Tensor | None,
    ) -> tuple[torch.Tensor, int]:
        """Decomposes weight into outlier low-rank L + sparse S under metric M."""
        d_out, d_in = W.shape
        device = W.device

        # Compute M
        if activations is not None and activations.numel() > 0:
            X = activations.reshape(-1, activations.shape[-1]).float().to(device)
            if X.shape[-1] == d_in:
                M = torch.matmul(X.T, X) / max(X.shape[0], 1)
                diag_mean = torch.diagonal(M).mean().item()
                reg = max(1e-4 * diag_mean, 1e-5)
                M = M + reg * torch.eye(d_in, device=device)
            else:
                M = torch.eye(d_in, device=device)
        else:
            M = torch.eye(d_in, device=device)

        try:
            L_M = torch.linalg.cholesky(M)
            L_M_inv = torch.linalg.inv(L_M)
            if torch.isnan(L_M).any() or torch.isnan(L_M_inv).any():
                raise ValueError("Cholesky produced NaNs")
        except Exception:
            L_M = torch.eye(d_in, device=device)
            L_M_inv = torch.eye(d_in, device=device)

        # Scale weights: W_scaled = W @ L_M
        W_scaled = torch.matmul(W, L_M)
        rank = max(
            min(self.min_rank, min(d_out, d_in)),
            min(int(self.rank_ratio * min(d_out, d_in)), min(d_out, d_in)),
        )

        # Initialize L via truncated SVD
        try:
            U, S, Vh = torch.linalg.svd(W_scaled, full_matrices=False)
            r = min(rank, len(S))
            L = torch.matmul(U[:, :r] * S[:r].unsqueeze(0), Vh[:r, :])
        except Exception:
            r = rank
            L = torch.zeros_like(W_scaled)

        # Initialize S
        R = W_scaled - L
        if self.structured_24:
            S_mat, _ = apply_24_structured_mask(R, saliency=torch.abs(R), dim=1)
        else:
            flat = R.abs().reshape(-1)
            k = max(1, int(flat.numel() * (1.0 - self.sparsity)))
            thr = torch.topk(flat, k, largest=True).values[-1]
            mask = (R.abs() >= thr).float()
            S_mat = R * mask

        # Alternating minimization iterations
        for _ in range(self.num_iters):
            try:
                U, S, Vh = torch.linalg.svd(W_scaled - S_mat, full_matrices=False)
                r = min(rank, len(S))
                L = torch.matmul(U[:, :r] * S[:r].unsqueeze(0), Vh[:r, :])
            except Exception as e:
                logger.warning(
                    "OATS SVD alternating iteration failed: %s; keeping previous decomposition",
                    e,
                )
                break

            R = W_scaled - L
            if self.structured_24:
                S_mat, _ = apply_24_structured_mask(R, saliency=torch.abs(R), dim=1)
            else:
                flat = R.abs().reshape(-1)
                k = max(1, int(flat.numel() * (1.0 - self.sparsity)))
                thr = torch.topk(flat, k, largest=True).values[-1]
                mask = (R.abs() >= thr).float()
                S_mat = R * mask

        # De-scale: W_comp = (L + S_mat) @ L_M^-1
        try:
            W_comp = torch.matmul(L + S_mat, L_M_inv)
        except Exception as e:
            logger.warning("OATS de-scaling failed: %s; returning original weight", e)
            return W, rank

        if torch.isnan(W_comp).any() or torch.isinf(W_comp).any():
            logger.warning("OATS decomposition produced NaNs or Infs; returning original weight")
            return W, rank
        return W_comp, r

    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """Executes OATS decomposition across compressible Linear and Conv1D modules."""
        activations_map = self._collect_activations()
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}
        seen_weights: set[int] = set()
        total_orig_params = 0
        total_comp_params = 0

        for name, module in list(self.model.named_modules()):
            if self._is_target_module(name, module):
                base_module = getattr(module, "base_layer", module)
                if not hasattr(base_module, "weight") or base_module.weight is None:
                    continue

                w_id = id(base_module.weight)
                if w_id in seen_weights:
                    logger.info("Skipping %s due to shared/tied weight.", name)
                    self.skipped_layers.append(name)
                    continue
                seen_weights.add(w_id)

                for p in module.parameters():
                    total_orig_params += p.numel()

                raw_w = base_module.weight.data
                cls_name = base_module.__class__.__name__.lower()
                is_conv1d = "conv1d" in cls_name or hasattr(base_module, "nf")

                W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
                device = raw_w.device
                dtype = raw_w.dtype

                act = activations_map.get(name, None)
                try:
                    W_comp, r = self._oats_decompose(W, act)
                    w_norm = torch.norm(W, p="fro").item()
                    err = (torch.norm(W - W_comp, p="fro") / max(w_norm, 1e-12)).item()
                    self.layer_errors[name] = float(err)
                    self.layer_ranks[name] = r

                    if not math.isfinite(err) or err > self.max_error:
                        logger.warning(
                            "OATS reconstruction error %.4f on %s exceeded threshold %.4f.",
                            err,
                            name,
                            self.max_error,
                        )
                        if self.fallback_to_original:
                            logger.info(
                                "Skipping OATS compression for %s (fallback to original).",
                                name,
                            )
                            self.fallback_layers.append(name)
                            for p in module.parameters():
                                p.requires_grad = False
                                total_comp_params += p.numel()
                            continue

                    W_target = W_comp
                    compressed_layer = OATSLoRALinear.from_linear(
                        module,
                        weight=W_target.to(dtype=dtype, device=device),
                        lora_r=self.lora_r,
                        lora_alpha=self.lora_alpha,
                    )
                    replace_module(self.model, name, compressed_layer)
                    self.compressed_layers.append(name)
                    for p in compressed_layer.parameters():
                        total_comp_params += p.numel()
                except Exception as e:
                    logger.warning("OATS decomposition exception in %s: %s", name, e)
                    self.fallback_layers.append(name)
                    self.layer_errors[name] = 1.0
                    for p in module.parameters():
                        p.requires_grad = False
                        total_comp_params += p.numel()

        # Post-compression forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                dev = next(self.model.parameters(), torch.tensor(0.0)).device
                dummy = torch.randint(0, 100, (1, 8), device=dev)
                try:
                    self.model(dummy)
                except Exception:
                    self.model(dummy.float())
        except Exception as e:
            logger.warning("OATS post-compression forward check warning: %s", e)
            forward_check_success = False

        self._compressed_params = self._count_params(self.model)
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        stats = self.get_compression_stats()
        eff_ratio = total_comp_params / max(1, total_orig_params)
        stats.update(
            {
                "decomposed_layers": len(self.compressed_layers),
                "sparsity_target": self.sparsity,
                "rank_ratio": self.rank_ratio,
                "structured_24": self.structured_24,
                "num_iters": self.num_iters,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "avg_reconstruction_error": avg_err,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "effective_parameter_ratio": eff_ratio,
                "forward_check_success": forward_check_success,
                "method": "oats_outlier_decomposition",
            }
        )
        self.stats = stats
        return self.model, stats

OATS Compressor: Outlier-aware low-rank + structured sparse matrix decomposition.

Production hardening features: - Relative Frobenius reconstruction error validation with configurable max_error (default 0.85). - Safe fallback to original weights by default (fallback_to_original=True). - Robust Conv1D detection by class name and nf attribute without unreliable shape heuristics. - Unified critical-layer protection via should_compress_linear. - GQA-aware skipping to protect asymmetric attention heads. - Tied / shared weight protection tracking parameter IDs. - Trainable LoRA adapter integration with base weights frozen. - FP32 numerical stability with NaN/Inf guards. - Post-compression forward sanity check under torch.no_grad(). - Comprehensive per-layer stats and error tracking.

Ancestors

BaseCompressor

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """Executes OATS decomposition across compressible Linear and Conv1D modules."""
    activations_map = self._collect_activations()
    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    self.layer_errors = {}
    self.layer_ranks = {}
    seen_weights: set[int] = set()
    total_orig_params = 0
    total_comp_params = 0

    for name, module in list(self.model.named_modules()):
        if self._is_target_module(name, module):
            base_module = getattr(module, "base_layer", module)
            if not hasattr(base_module, "weight") or base_module.weight is None:
                continue

            w_id = id(base_module.weight)
            if w_id in seen_weights:
                logger.info("Skipping %s due to shared/tied weight.", name)
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            for p in module.parameters():
                total_orig_params += p.numel()

            raw_w = base_module.weight.data
            cls_name = base_module.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(base_module, "nf")

            W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
            device = raw_w.device
            dtype = raw_w.dtype

            act = activations_map.get(name, None)
            try:
                W_comp, r = self._oats_decompose(W, act)
                w_norm = torch.norm(W, p="fro").item()
                err = (torch.norm(W - W_comp, p="fro") / max(w_norm, 1e-12)).item()
                self.layer_errors[name] = float(err)
                self.layer_ranks[name] = r

                if not math.isfinite(err) or err > self.max_error:
                    logger.warning(
                        "OATS reconstruction error %.4f on %s exceeded threshold %.4f.",
                        err,
                        name,
                        self.max_error,
                    )
                    if self.fallback_to_original:
                        logger.info(
                            "Skipping OATS compression for %s (fallback to original).",
                            name,
                        )
                        self.fallback_layers.append(name)
                        for p in module.parameters():
                            p.requires_grad = False
                            total_comp_params += p.numel()
                        continue

                W_target = W_comp
                compressed_layer = OATSLoRALinear.from_linear(
                    module,
                    weight=W_target.to(dtype=dtype, device=device),
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                )
                replace_module(self.model, name, compressed_layer)
                self.compressed_layers.append(name)
                for p in compressed_layer.parameters():
                    total_comp_params += p.numel()
            except Exception as e:
                logger.warning("OATS decomposition exception in %s: %s", name, e)
                self.fallback_layers.append(name)
                self.layer_errors[name] = 1.0
                for p in module.parameters():
                    p.requires_grad = False
                    total_comp_params += p.numel()

    # Post-compression forward sanity check
    forward_check_success = True
    try:
        with torch.no_grad():
            dev = next(self.model.parameters(), torch.tensor(0.0)).device
            dummy = torch.randint(0, 100, (1, 8), device=dev)
            try:
                self.model(dummy)
            except Exception:
                self.model(dummy.float())
    except Exception as e:
        logger.warning("OATS post-compression forward check warning: %s", e)
        forward_check_success = False

    self._compressed_params = self._count_params(self.model)
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )
    stats = self.get_compression_stats()
    eff_ratio = total_comp_params / max(1, total_orig_params)
    stats.update(
        {
            "decomposed_layers": len(self.compressed_layers),
            "sparsity_target": self.sparsity,
            "rank_ratio": self.rank_ratio,
            "structured_24": self.structured_24,
            "num_iters": self.num_iters,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "compressed_layers": self.compressed_layers,
            "fallback_layers": self.fallback_layers,
            "skipped_layers": self.skipped_layers,
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "num_skipped": len(self.skipped_layers),
            "layer_errors": self.layer_errors,
            "layer_ranks": self.layer_ranks,
            "avg_reconstruction_error": avg_err,
            "lora_r": self.lora_r,
            "lora_alpha": self.lora_alpha,
            "effective_parameter_ratio": eff_ratio,
            "forward_check_success": forward_check_success,
            "method": "oats_outlier_decomposition",
        }
    )
    self.stats = stats
    return self.model, stats

Executes OATS decomposition across compressible Linear and Conv1D modules.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class OATSLoRALinear (in_features: int, out_features: int, weight: torch.Tensor, bias: torch.Tensor | None = None, lora_r: int = 0, lora_alpha: float = 1.0, device: torch.device | None = None, dtype: torch.dtype | None = None)

Expand source code

class OATSLoRALinear(nn.Module):
    is_custom_attention: bool = True
    """
    Linear layer containing OATS decomposed base weights (frozen)
    with optional trainable LoRA adapters (lora_A, lora_B).
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        weight: torch.Tensor,
        bias: torch.Tensor | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features  # Conv1D compatibility
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        target_device = device if device is not None else weight.device
        target_dtype = dtype if dtype is not None else weight.dtype

        # Strictly ensure nn.Linear weight orientation (out_features, in_features) for F.linear.
        # Conv1D stores weight as (in_features, out_features) = (nx, nf).
        if weight.shape == (in_features, out_features) and in_features != out_features:
            weight = weight.t().contiguous()

        self.weight = nn.Parameter(
            weight.to(device=target_device, dtype=target_dtype),
            requires_grad=False,
        )
        if bias is not None:
            self.bias = nn.Parameter(
                bias.to(device=target_device, dtype=target_dtype),
                requires_grad=False,
            )
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty(
                    (self.lora_r, in_features),
                    dtype=torch.float32,
                    device=target_device,
                )
            )
            self.lora_B = nn.Parameter(
                torch.zeros(
                    (out_features, self.lora_r),
                    dtype=torch.float32,
                    device=target_device,
                )
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        weight: torch.Tensor | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        **kwargs: Any,
    ) -> OATSLoRALinear:
        base_layer = getattr(linear, "base_layer", linear)
        if weight is None and hasattr(base_layer, "weight"):
            weight = base_layer.weight.detach()
        if "rank" in kwargs and lora_r == 0:
            lora_r = kwargs["rank"]
        bias = (
            base_layer.bias.detach()
            if hasattr(base_layer, "bias") and base_layer.bias is not None
            else None
        )
        in_features = getattr(base_layer, "in_features", getattr(base_layer, "nx", None))
        out_features = getattr(base_layer, "out_features", getattr(base_layer, "nf", None))
        if in_features is None and out_features is None:
            out_features, in_features = weight.shape
        elif in_features is None:
            in_features = weight.shape[1] if weight.shape[0] == out_features else weight.shape[0]
        elif out_features is None:
            out_features = weight.shape[0] if weight.shape[1] == in_features else weight.shape[1]

        # Strictly orient weight as (out_features, in_features) for F.linear
        if weight.shape == (in_features, out_features) and in_features != out_features:
            weight = weight.t().contiguous()

        existing_lora_r = getattr(linear, "r", getattr(linear, "lora_r", 0))
        eff_lora_r = existing_lora_r if existing_lora_r > 0 else lora_r
        eff_lora_alpha = getattr(linear, "lora_alpha", getattr(linear, "alpha", lora_alpha))

        layer = cls(
            in_features=in_features,
            out_features=out_features,
            weight=weight,
            bias=bias,
            lora_r=eff_lora_r,
            lora_alpha=eff_lora_alpha,
            device=weight.device,
            dtype=weight.dtype,
        )

        # Preserve existing LoRA weights if present
        if (
            hasattr(linear, "lora_A")
            and hasattr(linear, "lora_B")
            and layer.lora_A is not None
            and layer.lora_B is not None
        ):
            with torch.no_grad():
                src_A = (
                    linear.lora_A.weight.data
                    if hasattr(linear.lora_A, "weight")
                    else linear.lora_A.data
                )
                src_B = (
                    linear.lora_B.weight.data
                    if hasattr(linear.lora_B, "weight")
                    else linear.lora_B.data
                )
                if src_A.shape == layer.lora_A.shape and src_B.shape == layer.lora_B.shape:
                    layer.lora_A.copy_(src_A.to(layer.lora_A.device, dtype=layer.lora_A.dtype))
                    layer.lora_B.copy_(src_B.to(layer.lora_B.device, dtype=layer.lora_B.dtype))

        if layer.lora_A is not None:
            layer.lora_A.requires_grad = True
        if layer.lora_B is not None:
            layer.lora_B.requires_grad = True

        return layer

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = F.linear(x, self.weight, self.bias)
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            out = out + self.scaling * F.linear(F.linear(x.float(), self.lora_A), self.lora_B).to(
                out.dtype
            )
        return out

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: Linear layer containing OATS decomposed base weights (frozen) with optional trainable LoRA adapters (lora_A, lora_B).

Static methods

def from_linear(linear: nn.Module, weight: torch.Tensor | None = None, lora_r: int = 0, lora_alpha: float = 1.0, **kwargs: Any) ‑> OATSLoRALinear

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    out = F.linear(x, self.weight, self.bias)
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        out = out + self.scaling * F.linear(F.linear(x.float(), self.lora_A), self.lora_B).to(
            out.dtype
        )
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class OptimalSubspaceLoRALinear (in_features: int, out_features: int, weight: torch.Tensor, bias: torch.Tensor | None = None, lora_r: int = 0, lora_alpha: float = 1.0, device: torch.device | None = None, dtype: torch.dtype | None = None)

Expand source code

class OptimalSubspaceLoRALinear(nn.Module):
    """
    Frozen optimal-subspace-projected base weight + optional trainable LoRA.
    """

    is_custom_attention: bool = False  # OSP-FIX-01

    def __init__(
        self,
        in_features: int,
        out_features: int,
        weight: torch.Tensor,
        bias: torch.Tensor | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features
        self.nx = in_features
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        target_device = device if device is not None else weight.device
        target_dtype = dtype if dtype is not None else weight.dtype

        # Normalize to (out, in)
        if (
            weight.shape[0] == in_features
            and weight.shape[1] == out_features
            and in_features != out_features
        ):
            weight = weight.t().contiguous()

        self.weight = nn.Parameter(
            weight.to(device=target_device, dtype=target_dtype),
            requires_grad=False,
        )
        if bias is not None:
            self.bias = nn.Parameter(
                bias.to(device=target_device, dtype=target_dtype),
                requires_grad=False,
            )
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty(
                    (self.lora_r, in_features),
                    dtype=torch.float32,
                    device=target_device,
                )
            )
            self.lora_B = nn.Parameter(
                torch.zeros(
                    (out_features, self.lora_r),
                    dtype=torch.float32,
                    device=target_device,
                )
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        weight: torch.Tensor,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> OptimalSubspaceLoRALinear:
        base_layer = getattr(linear, "base_layer", linear)
        bias = (
            base_layer.bias.detach()
            if hasattr(base_layer, "bias") and base_layer.bias is not None
            else None
        )
        in_features = getattr(base_layer, "in_features", getattr(base_layer, "nx", None))
        out_features = getattr(base_layer, "out_features", getattr(base_layer, "nf", None))
        if in_features is None or out_features is None:
            out_features = out_features or weight.shape[0]
            in_features = in_features or weight.shape[1]

        existing_lora_r = getattr(linear, "r", getattr(linear, "lora_r", 0))
        eff_lora_r = existing_lora_r if existing_lora_r > 0 else lora_r
        eff_lora_alpha = getattr(linear, "lora_alpha", getattr(linear, "alpha", lora_alpha))

        layer = cls(
            in_features=in_features,
            out_features=out_features,
            weight=weight,
            bias=bias,
            lora_r=eff_lora_r,
            lora_alpha=eff_lora_alpha,
            device=weight.device,
            dtype=weight.dtype,
        )

        # Preserve existing LoRA weights if present
        if (
            hasattr(linear, "lora_A")
            and hasattr(linear, "lora_B")
            and layer.lora_A is not None
            and layer.lora_B is not None
        ):
            with torch.no_grad():
                src_A = (
                    linear.lora_A.weight.data
                    if hasattr(linear.lora_A, "weight")
                    else linear.lora_A.data
                )
                src_B = (
                    linear.lora_B.weight.data
                    if hasattr(linear.lora_B, "weight")
                    else linear.lora_B.data
                )
                if src_A.shape == layer.lora_A.shape and src_B.shape == layer.lora_B.shape:
                    layer.lora_A.copy_(src_A.to(layer.lora_A.device, dtype=layer.lora_A.dtype))
                    layer.lora_B.copy_(src_B.to(layer.lora_B.device, dtype=layer.lora_B.dtype))
            layer.lora_A.requires_grad = True
            layer.lora_B.requires_grad = True

        return layer

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = F.linear(x, self.weight, self.bias)
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            out = out + self.scaling * F.linear(F.linear(x.float(), self.lora_A), self.lora_B).to(
                out.dtype
            )
        return out

Frozen optimal-subspace-projected base weight + optional trainable LoRA.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Static methods

def from_linear(linear: nn.Module, weight: torch.Tensor, lora_r: int = 0, lora_alpha: float = 1.0) ‑> OptimalSubspaceLoRALinear

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    out = F.linear(x, self.weight, self.bias)
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        out = out + self.scaling * F.linear(F.linear(x.float(), self.lora_A), self.lora_B).to(
            out.dtype
        )
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class OptimalSubspaceProjectionCompressor (model: nn.Module, calibration_loader: Any | None = None, rank_ratio: float = 0.5, power_iters: int = 2, tokenizer: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, max_error: float = 0.5, fallback_to_original: bool = True, target_attention_only: bool = False, target_modules: list[str] | None = None, lora_r: int = 0, lora_alpha: float = 1.0, skip_attention_qkv: bool = True, min_rank: int = 4, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("optimal_subspace_projection")
class OptimalSubspaceProjectionCompressor(BaseCompressor):
    """
    Optimal Subspace Projection Compressor via activation-aware randomized SVD.
    """

    # OSP-FIX-01: this is a compressor, not attention.
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module,
        calibration_loader: Any | None = None,
        rank_ratio: float = 0.5,
        power_iters: int = 2,
        tokenizer: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        target_attention_only: bool = False,
        target_modules: list[str] | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        skip_attention_qkv: bool = True,
        min_rank: int = 4,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="optimal_subspace_projection", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))

        if cfg is not None:
            rank_ratio = _cfg_get(
                cfg,
                "optimal_subspace_rank_ratio",
                _cfg_get(cfg, "rank_ratio", rank_ratio),
            )
            power_iters = _cfg_get(
                cfg,
                "optimal_subspace_power_iters",
                _cfg_get(cfg, "power_iters", power_iters),
            )
            max_error = _cfg_get(
                cfg, "optimal_subspace_max_error", _cfg_get(cfg, "max_error", max_error)
            )
            fallback_to_original = _cfg_get(
                cfg,
                "optimal_subspace_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
            lora_r = _cfg_get(cfg, "lora_r", lora_r)
            lora_alpha = _cfg_get(cfg, "lora_alpha", lora_alpha)
            skip_attention_qkv = _cfg_get(cfg, "skip_attention_qkv", skip_attention_qkv)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)
            min_rank = _cfg_get(cfg, "min_rank", min_rank)

        self.calibration_loader = calibration_loader
        self.rank_ratio = max(0.01, min(float(rank_ratio), 1.0))
        self.power_iters = max(0, int(power_iters))
        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.calibration_texts = calibration_texts
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.target_attention_only = bool(target_attention_only)
        self.target_modules = target_modules
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.skip_attention_qkv = bool(skip_attention_qkv)
        self.min_rank = max(1, int(min_rank))

        self.target_device = (
            torch.device(device)
            if device is not None
            else next(model.parameters(), torch.tensor(0.0)).device
        )
        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.stats: dict[str, Any] = {}

    # ------------------------------------------------------------------
    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            name_lower = name.lower()
            if any(
                p in name_lower
                for p in (
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                )
            ):
                return True
        return False

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        # OSP-FIX-07: skip already-compressed modules.
        if isinstance(module, OptimalSubspaceLoRALinear):
            return False
        # Skip nested lora params.
        if any(part in name.split(".") for part in ("lora_A", "lora_B")):
            return False

        cls_name = module.__class__.__name__.lower()
        is_linear_like = (
            isinstance(module, nn.Linear)
            or "conv1d" in cls_name
            or hasattr(module, "nf")
            or hasattr(module, "base_layer")
        )
        if not is_linear_like:
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self._is_gqa_asymmetric(name, module):
            return False
        if self.target_attention_only:
            if not any(k in name.lower() for k in _ATTN_NAMES):
                return False
        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    # ------------------------------------------------------------------
    @staticmethod
    def _compute_rank(d_out: int, d_in: int, rank_ratio: float, min_rank: int) -> int:
        """OSP-FIX-04: clean, correct rank formula."""
        max_rank = min(d_out, d_in)
        target = max(1, int(round(max_rank * rank_ratio)))
        rank = max(target, min_rank)
        rank = min(rank, max_rank)
        return rank

    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        # --------------------------------------------------------------
        # 1. Collect activation covariances (dedup by base weight id).
        # --------------------------------------------------------------
        act_covs: dict[int, torch.Tensor] = {}
        weight_id_to_scan_id: dict[int, int] = {}  # base weight id → cov key
        hooks = []

        def hook_fn(module: nn.Module, inputs: tuple[torch.Tensor, ...], output: Any) -> None:
            if not inputs or not isinstance(inputs[0], torch.Tensor):
                return
            base = getattr(module, "base_layer", module)
            w = getattr(base, "weight", None)
            if w is None:
                return
            w_id = id(w)
            key = weight_id_to_scan_id.get(w_id, w_id)

            x = inputs[0].detach()
            x_flat = x.reshape(-1, x.shape[-1]).float()
            cov = torch.matmul(x_flat.T, x_flat) / max(x_flat.shape[0], 1)
            if key not in act_covs:
                act_covs[key] = cov.cpu()
            else:
                act_covs[key] = act_covs[key] + cov.cpu()

        # OSP-FIX-06: register one hook per unique base weight.
        registered_ids: set[int] = set()
        for name, module in self.model.named_modules():
            if not self._is_target_module(name, module):
                continue
            base = getattr(module, "base_layer", module)
            w = getattr(base, "weight", None)
            if w is None:
                continue
            w_id = id(w)
            if w_id in registered_ids:
                continue
            registered_ids.add(w_id)
            weight_id_to_scan_id[w_id] = w_id
            hooks.append(module.register_forward_hook(hook_fn))

        calib_inputs = get_calibration_tensors(
            self.model,
            tokenizer=self.tokenizer,
            calibration_loader=self.calibration_loader,
            calibration_texts=self.calibration_texts,
            device=self.target_device,
        )

        with torch.no_grad():
            try:
                if isinstance(calib_inputs, dict):
                    self.model(**calib_inputs)
                elif isinstance(calib_inputs, torch.Tensor):
                    self.model(calib_inputs)
                elif isinstance(calib_inputs, (list, tuple)):
                    self.model(*calib_inputs)
            except Exception as e:
                logger.debug("OSP activation collection failed: %s", e)

        for h in hooks:
            h.remove()

        # --------------------------------------------------------------
        # 2. Compress each target layer.
        # --------------------------------------------------------------
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.layer_ranks = {}
        seen_weights: set[int] = set()
        total_orig_params = 0
        total_comp_params = 0

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            base_module = getattr(module, "base_layer", module)
            if not hasattr(base_module, "weight") or base_module.weight is None:
                continue

            w_id = id(base_module.weight)
            if w_id in seen_weights:
                logger.info("OSP: skipping %s (shared/tied weight)", name)
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            # OSP-FIX-03: count only BASE parameters.
            base_param_count = sum(p.numel() for p in base_module.parameters())
            total_orig_params += base_param_count

            raw_w = base_module.weight.data
            cls_name = base_module.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(base_module, "nf")

            W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
            d_out, d_in = W.shape
            device = raw_w.device
            dtype = raw_w.dtype

            # ---------------- Cholesky preconditioner ----------------
            C = act_covs.get(w_id)
            L = torch.eye(d_in, device=device)
            L_inv = torch.eye(d_in, device=device)

            if C is not None and C.shape[-1] == d_in:
                try:
                    C = C.to(device).float()
                    diag_mean = torch.diagonal(C).mean().item()
                    reg = max(1e-4 * max(diag_mean, 0.0), 1e-5)
                    C_reg = C + reg * torch.eye(d_in, device=device)
                    L = torch.linalg.cholesky(C_reg)
                    L_inv = torch.linalg.inv(L)
                    if not torch.isfinite(L).all() or not torch.isfinite(L_inv).all():
                        raise ValueError("Cholesky produced non-finite values")
                except Exception as e:
                    logger.debug("OSP: Cholesky fallback for %s: %s", name, e)
                    L = torch.eye(d_in, device=device)
                    L_inv = torch.eye(d_in, device=device)

            # ---------------- Randomized subspace iteration --------
            M = torch.matmul(W, L)
            rank = self._compute_rank(d_out, d_in, self.rank_ratio, self.min_rank)

            # OSP-FIX-09: Omega on the same device as M.
            Omega = torch.randn(M.shape[1], rank, device=M.device, dtype=M.dtype)
            Y = torch.matmul(M, Omega)
            for _ in range(self.power_iters):
                Y = torch.matmul(M, torch.matmul(M.T, Y))

            try:
                Q, _ = torch.linalg.qr(Y)
                B = torch.matmul(Q.T, M)
                U, S, Vt = torch.linalg.svd(B, full_matrices=False)
                r = min(rank, len(S))

                # OSP-FIX-10: empty rank → fallback.
                if r < 1:
                    raise ValueError("Empty SVD rank")

                U_r = U[:, :r]
                S_r = S[:r]
                Vt_r = Vt[:r, :]
                new_U = torch.matmul(Q, U_r)  # (d_out, r)
                new_V = torch.matmul(Vt_r, L_inv)  # (r, d_in)
                new_weight = torch.matmul(new_U * S_r.unsqueeze(0), new_V)  # (d_out, d_in)
                new_weight = torch.nan_to_num(new_weight, nan=0.0, posinf=0.0, neginf=0.0)

                # ---------------- Error check ----------------
                w_norm = torch.norm(W, p="fro").clamp_min(1e-12).item()
                err = (torch.norm(W - new_weight, p="fro") / w_norm).item()
                self.layer_errors[name] = float(err)
                self.layer_ranks[name] = r

                if not math.isfinite(err) or err > self.max_error:
                    logger.warning(
                        "OSP: reconstruction error %.4f on %s exceeded %.4f.",
                        err,
                        name,
                        self.max_error,
                    )
                    if self.fallback_to_original:
                        self.fallback_layers.append(name)
                        total_comp_params += base_param_count  # OSP-FIX-08
                        for p in module.parameters():
                            p.requires_grad = False
                        continue

                # ---------------- Commit ----------------
                W_target = new_weight.to(dtype=dtype, device=device)  # OSP-FIX-11
                compressed_layer = OptimalSubspaceLoRALinear.from_linear(
                    module,
                    weight=W_target,
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                )

                # OSP-FIX-02: preserve LoRA wrapper.
                if hasattr(module, "base_layer"):
                    module.base_layer = compressed_layer
                else:
                    replace_module(self.model, name, compressed_layer)

                self.compressed_layers.append(name)
                total_comp_params += sum(p.numel() for p in compressed_layer.parameters())

            except Exception as e:
                logger.warning("OSP: SVD/decomposition failed on %s: %s", name, e)
                self.fallback_layers.append(name)
                self.layer_errors[name] = 1.0
                total_comp_params += base_param_count
                for p in module.parameters():
                    p.requires_grad = False

        # --------------------------------------------------------------
        # 3. Forward sanity check.
        # --------------------------------------------------------------
        forward_check_success = True
        try:
            with torch.no_grad():
                dev = next(self.model.parameters(), torch.tensor(0.0)).device
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                # OSP-FIX-05: no `.float()` on int64 input ids.
                try:
                    self.model(input_ids=dummy)
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception:
                        # Try embedding-level call as last resort.
                        cfg = getattr(self.model, "config", None)
                        hidden = getattr(cfg, "hidden_size", None) if cfg else None
                        if hidden is None:
                            first_lin = next(
                                (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                                None,
                            )
                            hidden = getattr(first_lin, "in_features", 8)
                        dummy_feat = torch.randn(
                            1,
                            8,
                            hidden,
                            device=dev,
                            dtype=next(self.model.parameters()).dtype,
                        )
                        self.model(dummy_feat)
        except Exception as e:
            logger.warning("OSP forward sanity check failed: %s", e)
            forward_check_success = False

        # --------------------------------------------------------------
        # 4. Stats.
        # --------------------------------------------------------------
        self._compressed_params = self._count_params(self.model)
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        eff_ratio = total_comp_params / max(1, total_orig_params) if total_orig_params > 0 else 1.0
        stats = self.get_compression_stats()
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "avg_reconstruction_error": avg_err,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "rank_ratio": self.rank_ratio,
                "power_iters": self.power_iters,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "effective_parameter_ratio": round(eff_ratio, 4),
                "forward_check_success": forward_check_success,
                "method": "optimal_subspace_projection",
            }
        )
        self.stats = stats
        return self.model, stats

Optimal Subspace Projection Compressor via activation-aware randomized SVD.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class Path (*args)

Expand source code

class Path(PurePath):
    """PurePath subclass that can make system calls.

    Path represents a filesystem path but unlike PurePath, also offers
    methods to do system calls on path objects. Depending on your system,
    instantiating a Path will return either a PosixPath or a WindowsPath
    object. You can also instantiate a PosixPath or WindowsPath directly,
    but cannot instantiate a WindowsPath on a POSIX system or vice versa.
    """
    __slots__ = ('_info',)

    def __new__(cls, *args, **kwargs):
        if cls is Path:
            cls = WindowsPath if os.name == 'nt' else PosixPath
        return object.__new__(cls)

    @property
    def info(self):
        """
        A PathInfo object that exposes the file type and other file attributes
        of this path.
        """
        try:
            return self._info
        except AttributeError:
            self._info = PathInfo(self)
            return self._info

    def stat(self, *, follow_symlinks=True):
        """
        Return the result of the stat() system call on this path, like
        os.stat() does.
        """
        return os.stat(self, follow_symlinks=follow_symlinks)

    def lstat(self):
        """
        Like stat(), except if the path points to a symlink, the symlink's
        status information is returned, rather than its target's.
        """
        return os.lstat(self)

    def exists(self, *, follow_symlinks=True):
        """
        Whether this path exists.

        This method normally follows symlinks; to check whether a symlink exists,
        add the argument follow_symlinks=False.
        """
        if follow_symlinks:
            return os.path.exists(self)
        return os.path.lexists(self)

    def is_dir(self, *, follow_symlinks=True):
        """
        Whether this path is a directory.
        """
        if follow_symlinks:
            return os.path.isdir(self)
        try:
            return S_ISDIR(self.stat(follow_symlinks=follow_symlinks).st_mode)
        except (OSError, ValueError):
            return False

    def is_file(self, *, follow_symlinks=True):
        """
        Whether this path is a regular file (also True for symlinks pointing
        to regular files).
        """
        if follow_symlinks:
            return os.path.isfile(self)
        try:
            return S_ISREG(self.stat(follow_symlinks=follow_symlinks).st_mode)
        except (OSError, ValueError):
            return False

    def is_mount(self):
        """
        Check if this path is a mount point
        """
        return os.path.ismount(self)

    def is_symlink(self):
        """
        Whether this path is a symbolic link.
        """
        return os.path.islink(self)

    def is_junction(self):
        """
        Whether this path is a junction.
        """
        return os.path.isjunction(self)

    def is_block_device(self):
        """
        Whether this path is a block device.
        """
        try:
            return S_ISBLK(self.stat().st_mode)
        except (OSError, ValueError):
            return False

    def is_char_device(self):
        """
        Whether this path is a character device.
        """
        try:
            return S_ISCHR(self.stat().st_mode)
        except (OSError, ValueError):
            return False

    def is_fifo(self):
        """
        Whether this path is a FIFO.
        """
        try:
            return S_ISFIFO(self.stat().st_mode)
        except (OSError, ValueError):
            return False

    def is_socket(self):
        """
        Whether this path is a socket.
        """
        try:
            return S_ISSOCK(self.stat().st_mode)
        except (OSError, ValueError):
            return False

    def samefile(self, other_path):
        """Return whether other_path is the same or not as this file
        (as returned by os.path.samefile()).
        """
        st = self.stat()
        try:
            other_st = other_path.stat()
        except AttributeError:
            other_st = self.with_segments(other_path).stat()
        return (st.st_ino == other_st.st_ino and
                st.st_dev == other_st.st_dev)

    def open(self, mode='r', buffering=-1, encoding=None,
             errors=None, newline=None):
        """
        Open the file pointed to by this path and return a file object, as
        the built-in open() function does.
        """
        if "b" not in mode:
            encoding = io.text_encoding(encoding)
        return io.open(self, mode, buffering, encoding, errors, newline)

    def read_bytes(self):
        """
        Open the file in bytes mode, read it, and close the file.
        """
        with self.open(mode='rb', buffering=0) as f:
            return f.read()

    def read_text(self, encoding=None, errors=None, newline=None):
        """
        Open the file in text mode, read it, and close the file.
        """
        # Call io.text_encoding() here to ensure any warning is raised at an
        # appropriate stack level.
        encoding = io.text_encoding(encoding)
        with self.open(mode='r', encoding=encoding, errors=errors, newline=newline) as f:
            return f.read()

    def write_bytes(self, data):
        """
        Open the file in bytes mode, write to it, and close the file.
        """
        # type-check for the buffer interface before truncating the file
        view = memoryview(data)
        with self.open(mode='wb') as f:
            return f.write(view)

    def write_text(self, data, encoding=None, errors=None, newline=None):
        """
        Open the file in text mode, write to it, and close the file.
        """
        # Call io.text_encoding() here to ensure any warning is raised at an
        # appropriate stack level.
        encoding = io.text_encoding(encoding)
        if not isinstance(data, str):
            raise TypeError('data must be str, not %s' %
                            data.__class__.__name__)
        with self.open(mode='w', encoding=encoding, errors=errors, newline=newline) as f:
            return f.write(data)

    _remove_leading_dot = operator.itemgetter(slice(2, None))
    _remove_trailing_slash = operator.itemgetter(slice(-1))

    def _filter_trailing_slash(self, paths):
        sep = self.parser.sep
        anchor_len = len(self.anchor)
        for path_str in paths:
            if len(path_str) > anchor_len and path_str[-1] == sep:
                path_str = path_str[:-1]
            yield path_str

    def _from_dir_entry(self, dir_entry, path_str):
        path = self.with_segments(path_str)
        path._str = path_str
        path._info = DirEntryInfo(dir_entry)
        return path

    def iterdir(self):
        """Yield path objects of the directory contents.

        The children are yielded in arbitrary order, and the
        special entries '.' and '..' are not included.
        """
        root_dir = str(self)
        with os.scandir(root_dir) as scandir_it:
            entries = list(scandir_it)
        if root_dir == '.':
            return (self._from_dir_entry(e, e.name) for e in entries)
        else:
            return (self._from_dir_entry(e, e.path) for e in entries)

    def glob(self, pattern, *, case_sensitive=None, recurse_symlinks=False):
        """Iterate over this subtree and yield all existing files (of any
        kind, including directories) matching the given relative pattern.
        """
        sys.audit("pathlib.Path.glob", self, pattern)
        if case_sensitive is None:
            case_sensitive = self.parser is posixpath
            case_pedantic = False
        else:
            # The user has expressed a case sensitivity choice, but we don't
            # know the case sensitivity of the underlying filesystem, so we
            # must use scandir() for everything, including non-wildcard parts.
            case_pedantic = True
        parts = self._parse_pattern(pattern)
        recursive = True if recurse_symlinks else _no_recurse_symlinks
        globber = _StringGlobber(self.parser.sep, case_sensitive, case_pedantic, recursive)
        select = globber.selector(parts[::-1])
        root = str(self)
        paths = select(self.parser.join(root, ''))

        # Normalize results
        if root == '.':
            paths = map(self._remove_leading_dot, paths)
        if parts[-1] == '':
            paths = map(self._remove_trailing_slash, paths)
        elif parts[-1] == '**':
            paths = self._filter_trailing_slash(paths)
        paths = map(self._from_parsed_string, paths)
        return paths

    def rglob(self, pattern, *, case_sensitive=None, recurse_symlinks=False):
        """Recursively yield all existing files (of any kind, including
        directories) matching the given relative pattern, anywhere in
        this subtree.
        """
        sys.audit("pathlib.Path.rglob", self, pattern)
        pattern = self.parser.join('**', pattern)
        return self.glob(pattern, case_sensitive=case_sensitive, recurse_symlinks=recurse_symlinks)

    def walk(self, top_down=True, on_error=None, follow_symlinks=False):
        """Walk the directory tree from this directory, similar to os.walk()."""
        sys.audit("pathlib.Path.walk", self, on_error, follow_symlinks)
        root_dir = str(self)
        if not follow_symlinks:
            follow_symlinks = os._walk_symlinks_as_files
        results = os.walk(root_dir, top_down, on_error, follow_symlinks)
        for path_str, dirnames, filenames in results:
            if root_dir == '.':
                path_str = path_str[2:]
            yield self._from_parsed_string(path_str), dirnames, filenames

    def absolute(self):
        """Return an absolute version of this path
        No normalization or symlink resolution is performed.

        Use resolve() to resolve symlinks and remove '..' segments.
        """
        if self.is_absolute():
            return self
        if self.root:
            drive = os.path.splitroot(os.getcwd())[0]
            return self._from_parsed_parts(drive, self.root, self._tail)
        if self.drive:
            # There is a CWD on each drive-letter drive.
            cwd = os.path.abspath(self.drive)
        else:
            cwd = os.getcwd()
        if not self._tail:
            # Fast path for "empty" paths, e.g. Path("."), Path("") or Path().
            # We pass only one argument to with_segments() to avoid the cost
            # of joining, and we exploit the fact that getcwd() returns a
            # fully-normalized string by storing it in _str. This is used to
            # implement Path.cwd().
            return self._from_parsed_string(cwd)
        drive, root, rel = os.path.splitroot(cwd)
        if not rel:
            return self._from_parsed_parts(drive, root, self._tail)
        tail = rel.split(self.parser.sep)
        tail.extend(self._tail)
        return self._from_parsed_parts(drive, root, tail)

    @classmethod
    def cwd(cls):
        """Return a new path pointing to the current working directory."""
        cwd = os.getcwd()
        path = cls(cwd)
        path._str = cwd  # getcwd() returns a normalized path
        return path

    def resolve(self, strict=False):
        """
        Make the path absolute, resolving all symlinks on the way and also
        normalizing it.
        """

        return self.with_segments(os.path.realpath(self, strict=strict))

    if pwd:
        def owner(self, *, follow_symlinks=True):
            """
            Return the login name of the file owner.
            """
            uid = self.stat(follow_symlinks=follow_symlinks).st_uid
            return pwd.getpwuid(uid).pw_name
    else:
        def owner(self, *, follow_symlinks=True):
            """
            Return the login name of the file owner.
            """
            f = f"{type(self).__name__}.owner()"
            raise UnsupportedOperation(f"{f} is unsupported on this system")

    if grp:
        def group(self, *, follow_symlinks=True):
            """
            Return the group name of the file gid.
            """
            gid = self.stat(follow_symlinks=follow_symlinks).st_gid
            return grp.getgrgid(gid).gr_name
    else:
        def group(self, *, follow_symlinks=True):
            """
            Return the group name of the file gid.
            """
            f = f"{type(self).__name__}.group()"
            raise UnsupportedOperation(f"{f} is unsupported on this system")

    if hasattr(os, "readlink"):
        def readlink(self):
            """
            Return the path to which the symbolic link points.
            """
            return self.with_segments(os.readlink(self))
    else:
        def readlink(self):
            """
            Return the path to which the symbolic link points.
            """
            f = f"{type(self).__name__}.readlink()"
            raise UnsupportedOperation(f"{f} is unsupported on this system")

    def touch(self, mode=0o666, exist_ok=True):
        """
        Create this file with the given access mode, if it doesn't exist.
        """

        if exist_ok:
            # First try to bump modification time
            # Implementation note: GNU touch uses the UTIME_NOW option of
            # the utimensat() / futimens() functions.
            try:
                os.utime(self, None)
            except OSError:
                # Avoid exception chaining
                pass
            else:
                return
        flags = os.O_CREAT | os.O_WRONLY
        if not exist_ok:
            flags |= os.O_EXCL
        fd = os.open(self, flags, mode)
        os.close(fd)

    def mkdir(self, mode=0o777, parents=False, exist_ok=False):
        """
        Create a new directory at this given path.
        """
        try:
            os.mkdir(self, mode)
        except FileNotFoundError:
            if not parents or self.parent == self:
                raise
            self.parent.mkdir(parents=True, exist_ok=True)
            self.mkdir(mode, parents=False, exist_ok=exist_ok)
        except OSError:
            # Cannot rely on checking for EEXIST, since the operating system
            # could give priority to other errors like EACCES or EROFS
            if not exist_ok or not self.is_dir():
                raise

    def chmod(self, mode, *, follow_symlinks=True):
        """
        Change the permissions of the path, like os.chmod().
        """
        os.chmod(self, mode, follow_symlinks=follow_symlinks)

    def lchmod(self, mode):
        """
        Like chmod(), except if the path points to a symlink, the symlink's
        permissions are changed, rather than its target's.
        """
        self.chmod(mode, follow_symlinks=False)

    def unlink(self, missing_ok=False):
        """
        Remove this file or link.
        If the path is a directory, use rmdir() instead.
        """
        try:
            os.unlink(self)
        except FileNotFoundError:
            if not missing_ok:
                raise

    def rmdir(self):
        """
        Remove this directory.  The directory must be empty.
        """
        os.rmdir(self)

    def _delete(self):
        """
        Delete this file or directory (including all sub-directories).
        """
        if self.is_symlink() or self.is_junction():
            self.unlink()
        elif self.is_dir():
            # Lazy import to improve module import time
            import shutil
            shutil.rmtree(self)
        else:
            self.unlink()

    def rename(self, target):
        """
        Rename this path to the target path.

        The target path may be absolute or relative. Relative paths are
        interpreted relative to the current working directory, *not* the
        directory of the Path object.

        Returns the new Path instance pointing to the target path.
        """
        os.rename(self, target)
        if not hasattr(target, 'with_segments'):
            target = self.with_segments(target)
        return target

    def replace(self, target):
        """
        Rename this path to the target path, overwriting if that path exists.

        The target path may be absolute or relative. Relative paths are
        interpreted relative to the current working directory, *not* the
        directory of the Path object.

        Returns the new Path instance pointing to the target path.
        """
        os.replace(self, target)
        if not hasattr(target, 'with_segments'):
            target = self.with_segments(target)
        return target

    def copy(self, target, **kwargs):
        """
        Recursively copy this file or directory tree to the given destination.
        """
        if not hasattr(target, 'with_segments'):
            target = self.with_segments(target)
        ensure_distinct_paths(self, target)
        target._copy_from(self, **kwargs)
        return target.joinpath()  # Empty join to ensure fresh metadata.

    def copy_into(self, target_dir, **kwargs):
        """
        Copy this file or directory tree into the given existing directory.
        """
        name = self.name
        if not name:
            raise ValueError(f"{self!r} has an empty name")
        elif hasattr(target_dir, 'with_segments'):
            target = target_dir / name
        else:
            target = self.with_segments(target_dir, name)
        return self.copy(target, **kwargs)

    def _copy_from(self, source, follow_symlinks=True, preserve_metadata=False):
        """
        Recursively copy the given path to this path.
        """
        if not follow_symlinks and source.info.is_symlink():
            self._copy_from_symlink(source, preserve_metadata)
        elif source.info.is_dir():
            children = source.iterdir()
            os.mkdir(self)
            for child in children:
                self.joinpath(child.name)._copy_from(
                    child, follow_symlinks, preserve_metadata)
            if preserve_metadata:
                copy_info(source.info, self)
        else:
            self._copy_from_file(source, preserve_metadata)

    def _copy_from_file(self, source, preserve_metadata=False):
        ensure_different_files(source, self)
        with magic_open(source, 'rb') as source_f:
            with open(self, 'wb') as target_f:
                copyfileobj(source_f, target_f)
        if preserve_metadata:
            copy_info(source.info, self)

    if copyfile2:
        # Use fast OS routine for local file copying where available.
        _copy_from_file_fallback = _copy_from_file
        def _copy_from_file(self, source, preserve_metadata=False):
            try:
                source = os.fspath(source)
            except TypeError:
                pass
            else:
                copyfile2(source, str(self))
                return
            self._copy_from_file_fallback(source, preserve_metadata)

    if os.name == 'nt':
        # If a directory-symlink is copied *before* its target, then
        # os.symlink() incorrectly creates a file-symlink on Windows. Avoid
        # this by passing *target_is_dir* to os.symlink() on Windows.
        def _copy_from_symlink(self, source, preserve_metadata=False):
            os.symlink(str(source.readlink()), self, source.info.is_dir())
            if preserve_metadata:
                copy_info(source.info, self, follow_symlinks=False)
    else:
        def _copy_from_symlink(self, source, preserve_metadata=False):
            os.symlink(str(source.readlink()), self)
            if preserve_metadata:
                copy_info(source.info, self, follow_symlinks=False)

    def move(self, target):
        """
        Recursively move this file or directory tree to the given destination.
        """
        # Use os.replace() if the target is os.PathLike and on the same FS.
        try:
            target = self.with_segments(target)
        except TypeError:
            pass
        else:
            ensure_different_files(self, target)
            try:
                os.replace(self, target)
            except OSError as err:
                if err.errno != EXDEV:
                    raise
            else:
                return target.joinpath()  # Empty join to ensure fresh metadata.
        # Fall back to copy+delete.
        target = self.copy(target, follow_symlinks=False, preserve_metadata=True)
        self._delete()
        return target

    def move_into(self, target_dir):
        """
        Move this file or directory tree into the given existing directory.
        """
        name = self.name
        if not name:
            raise ValueError(f"{self!r} has an empty name")
        elif hasattr(target_dir, 'with_segments'):
            target = target_dir / name
        else:
            target = self.with_segments(target_dir, name)
        return self.move(target)

    if hasattr(os, "symlink"):
        def symlink_to(self, target, target_is_directory=False):
            """
            Make this path a symlink pointing to the target path.
            Note the order of arguments (link, target) is the reverse of os.symlink.
            """
            os.symlink(target, self, target_is_directory)
    else:
        def symlink_to(self, target, target_is_directory=False):
            """
            Make this path a symlink pointing to the target path.
            Note the order of arguments (link, target) is the reverse of os.symlink.
            """
            f = f"{type(self).__name__}.symlink_to()"
            raise UnsupportedOperation(f"{f} is unsupported on this system")

    if hasattr(os, "link"):
        def hardlink_to(self, target):
            """
            Make this path a hard link pointing to the same file as *target*.

            Note the order of arguments (self, target) is the reverse of os.link's.
            """
            os.link(target, self)
    else:
        def hardlink_to(self, target):
            """
            Make this path a hard link pointing to the same file as *target*.

            Note the order of arguments (self, target) is the reverse of os.link's.
            """
            f = f"{type(self).__name__}.hardlink_to()"
            raise UnsupportedOperation(f"{f} is unsupported on this system")

    def expanduser(self):
        """ Return a new path with expanded ~ and ~user constructs
        (as returned by os.path.expanduser)
        """
        if (not (self.drive or self.root) and
            self._tail and self._tail[0][:1] == '~'):
            homedir = os.path.expanduser(self._tail[0])
            if homedir[:1] == "~":
                raise RuntimeError("Could not determine home directory.")
            drv, root, tail = self._parse_path(homedir)
            return self._from_parsed_parts(drv, root, tail + self._tail[1:])

        return self

    @classmethod
    def home(cls):
        """Return a new path pointing to expanduser('~').
        """
        homedir = os.path.expanduser("~")
        if homedir == "~":
            raise RuntimeError("Could not determine home directory.")
        return cls(homedir)

    def as_uri(self):
        """Return the path as a URI."""
        if not self.is_absolute():
            raise ValueError("relative paths can't be expressed as file URIs")
        from urllib.request import pathname2url
        return pathname2url(str(self), add_scheme=True)

    @classmethod
    def from_uri(cls, uri):
        """Return a new path from the given 'file' URI."""
        from urllib.error import URLError
        from urllib.request import url2pathname
        try:
            path = cls(url2pathname(uri, require_scheme=True))
        except URLError as exc:
            raise ValueError(exc.reason) from None
        if not path.is_absolute():
            raise ValueError(f"URI is not absolute: {uri!r}")
        return path

PurePath subclass that can make system calls.

Path represents a filesystem path but unlike PurePath, also offers methods to do system calls on path objects. Depending on your system, instantiating a Path will return either a PosixPath or a WindowsPath object. You can also instantiate a PosixPath or WindowsPath directly, but cannot instantiate a WindowsPath on a POSIX system or vice versa.

Ancestors

pathlib.PurePath

Subclasses

pathlib.PosixPath
pathlib.WindowsPath

Static methods

def cwd(): Return a new path pointing to the current working directory.
def from_uri(uri): Return a new path from the given 'file' URI.
def home(): Return a new path pointing to expanduser('~').

Instance variables

prop info

Expand source code

@property
def info(self):
    """
    A PathInfo object that exposes the file type and other file attributes
    of this path.
    """
    try:
        return self._info
    except AttributeError:
        self._info = PathInfo(self)
        return self._info

A PathInfo object that exposes the file type and other file attributes of this path.

Methods

def absolute(self)

Expand source code

def absolute(self):
    """Return an absolute version of this path
    No normalization or symlink resolution is performed.

    Use resolve() to resolve symlinks and remove '..' segments.
    """
    if self.is_absolute():
        return self
    if self.root:
        drive = os.path.splitroot(os.getcwd())[0]
        return self._from_parsed_parts(drive, self.root, self._tail)
    if self.drive:
        # There is a CWD on each drive-letter drive.
        cwd = os.path.abspath(self.drive)
    else:
        cwd = os.getcwd()
    if not self._tail:
        # Fast path for "empty" paths, e.g. Path("."), Path("") or Path().
        # We pass only one argument to with_segments() to avoid the cost
        # of joining, and we exploit the fact that getcwd() returns a
        # fully-normalized string by storing it in _str. This is used to
        # implement Path.cwd().
        return self._from_parsed_string(cwd)
    drive, root, rel = os.path.splitroot(cwd)
    if not rel:
        return self._from_parsed_parts(drive, root, self._tail)
    tail = rel.split(self.parser.sep)
    tail.extend(self._tail)
    return self._from_parsed_parts(drive, root, tail)

Return an absolute version of this path No normalization or symlink resolution is performed.

Use resolve() to resolve symlinks and remove '..' segments.

def as_uri(self)

Expand source code

def as_uri(self):
    """Return the path as a URI."""
    if not self.is_absolute():
        raise ValueError("relative paths can't be expressed as file URIs")
    from urllib.request import pathname2url
    return pathname2url(str(self), add_scheme=True)

Return the path as a URI.

def chmod(self, mode, *, follow_symlinks=True)

Expand source code

def chmod(self, mode, *, follow_symlinks=True):
    """
    Change the permissions of the path, like os.chmod().
    """
    os.chmod(self, mode, follow_symlinks=follow_symlinks)

Change the permissions of the path, like os.chmod().

def copy(self, target, **kwargs)

Expand source code

def copy(self, target, **kwargs):
    """
    Recursively copy this file or directory tree to the given destination.
    """
    if not hasattr(target, 'with_segments'):
        target = self.with_segments(target)
    ensure_distinct_paths(self, target)
    target._copy_from(self, **kwargs)
    return target.joinpath()  # Empty join to ensure fresh metadata.

Recursively copy this file or directory tree to the given destination.

def copy_into(self, target_dir, **kwargs)

Expand source code

def copy_into(self, target_dir, **kwargs):
    """
    Copy this file or directory tree into the given existing directory.
    """
    name = self.name
    if not name:
        raise ValueError(f"{self!r} has an empty name")
    elif hasattr(target_dir, 'with_segments'):
        target = target_dir / name
    else:
        target = self.with_segments(target_dir, name)
    return self.copy(target, **kwargs)

Copy this file or directory tree into the given existing directory.

def exists(self, *, follow_symlinks=True)

Expand source code

def exists(self, *, follow_symlinks=True):
    """
    Whether this path exists.

    This method normally follows symlinks; to check whether a symlink exists,
    add the argument follow_symlinks=False.
    """
    if follow_symlinks:
        return os.path.exists(self)
    return os.path.lexists(self)

Whether this path exists.

This method normally follows symlinks; to check whether a symlink exists, add the argument follow_symlinks=False.

def expanduser(self)

Expand source code

def expanduser(self):
    """ Return a new path with expanded ~ and ~user constructs
    (as returned by os.path.expanduser)
    """
    if (not (self.drive or self.root) and
        self._tail and self._tail[0][:1] == '~'):
        homedir = os.path.expanduser(self._tail[0])
        if homedir[:1] == "~":
            raise RuntimeError("Could not determine home directory.")
        drv, root, tail = self._parse_path(homedir)
        return self._from_parsed_parts(drv, root, tail + self._tail[1:])

    return self

Return a new path with expanded ~ and ~user constructs (as returned by os.path.expanduser)

def glob(self, pattern, *, case_sensitive=None, recurse_symlinks=False)

Expand source code

def glob(self, pattern, *, case_sensitive=None, recurse_symlinks=False):
    """Iterate over this subtree and yield all existing files (of any
    kind, including directories) matching the given relative pattern.
    """
    sys.audit("pathlib.Path.glob", self, pattern)
    if case_sensitive is None:
        case_sensitive = self.parser is posixpath
        case_pedantic = False
    else:
        # The user has expressed a case sensitivity choice, but we don't
        # know the case sensitivity of the underlying filesystem, so we
        # must use scandir() for everything, including non-wildcard parts.
        case_pedantic = True
    parts = self._parse_pattern(pattern)
    recursive = True if recurse_symlinks else _no_recurse_symlinks
    globber = _StringGlobber(self.parser.sep, case_sensitive, case_pedantic, recursive)
    select = globber.selector(parts[::-1])
    root = str(self)
    paths = select(self.parser.join(root, ''))

    # Normalize results
    if root == '.':
        paths = map(self._remove_leading_dot, paths)
    if parts[-1] == '':
        paths = map(self._remove_trailing_slash, paths)
    elif parts[-1] == '**':
        paths = self._filter_trailing_slash(paths)
    paths = map(self._from_parsed_string, paths)
    return paths

Iterate over this subtree and yield all existing files (of any kind, including directories) matching the given relative pattern.

def group(self, *, follow_symlinks=True)

Expand source code

def group(self, *, follow_symlinks=True):
    """
    Return the group name of the file gid.
    """
    gid = self.stat(follow_symlinks=follow_symlinks).st_gid
    return grp.getgrgid(gid).gr_name

Return the group name of the file gid.

def hardlink_to(self, target)

Expand source code

def hardlink_to(self, target):
    """
    Make this path a hard link pointing to the same file as *target*.

    Note the order of arguments (self, target) is the reverse of os.link's.
    """
    os.link(target, self)

Make this path a hard link pointing to the same file as target.

Note the order of arguments (self, target) is the reverse of os.link's.

def is_block_device(self)

Expand source code

def is_block_device(self):
    """
    Whether this path is a block device.
    """
    try:
        return S_ISBLK(self.stat().st_mode)
    except (OSError, ValueError):
        return False

Whether this path is a block device.

def is_char_device(self)

Expand source code

def is_char_device(self):
    """
    Whether this path is a character device.
    """
    try:
        return S_ISCHR(self.stat().st_mode)
    except (OSError, ValueError):
        return False

Whether this path is a character device.

def is_dir(self, *, follow_symlinks=True)

Expand source code

def is_dir(self, *, follow_symlinks=True):
    """
    Whether this path is a directory.
    """
    if follow_symlinks:
        return os.path.isdir(self)
    try:
        return S_ISDIR(self.stat(follow_symlinks=follow_symlinks).st_mode)
    except (OSError, ValueError):
        return False

Whether this path is a directory.

def is_fifo(self)

Expand source code

def is_fifo(self):
    """
    Whether this path is a FIFO.
    """
    try:
        return S_ISFIFO(self.stat().st_mode)
    except (OSError, ValueError):
        return False

Whether this path is a FIFO.

def is_file(self, *, follow_symlinks=True)

Expand source code

def is_file(self, *, follow_symlinks=True):
    """
    Whether this path is a regular file (also True for symlinks pointing
    to regular files).
    """
    if follow_symlinks:
        return os.path.isfile(self)
    try:
        return S_ISREG(self.stat(follow_symlinks=follow_symlinks).st_mode)
    except (OSError, ValueError):
        return False

Whether this path is a regular file (also True for symlinks pointing to regular files).

def is_junction(self)

Expand source code

def is_junction(self):
    """
    Whether this path is a junction.
    """
    return os.path.isjunction(self)

Whether this path is a junction.

def is_mount(self)

Expand source code

def is_mount(self):
    """
    Check if this path is a mount point
    """
    return os.path.ismount(self)

Check if this path is a mount point

def is_socket(self)

Expand source code

def is_socket(self):
    """
    Whether this path is a socket.
    """
    try:
        return S_ISSOCK(self.stat().st_mode)
    except (OSError, ValueError):
        return False

Whether this path is a socket.

def is_symlink(self)

Expand source code

def is_symlink(self):
    """
    Whether this path is a symbolic link.
    """
    return os.path.islink(self)

Whether this path is a symbolic link.

def iterdir(self)

Expand source code

def iterdir(self):
    """Yield path objects of the directory contents.

    The children are yielded in arbitrary order, and the
    special entries '.' and '..' are not included.
    """
    root_dir = str(self)
    with os.scandir(root_dir) as scandir_it:
        entries = list(scandir_it)
    if root_dir == '.':
        return (self._from_dir_entry(e, e.name) for e in entries)
    else:
        return (self._from_dir_entry(e, e.path) for e in entries)

Yield path objects of the directory contents.

The children are yielded in arbitrary order, and the special entries '.' and '..' are not included.

def lchmod(self, mode)

Expand source code

def lchmod(self, mode):
    """
    Like chmod(), except if the path points to a symlink, the symlink's
    permissions are changed, rather than its target's.
    """
    self.chmod(mode, follow_symlinks=False)

Like chmod(), except if the path points to a symlink, the symlink's permissions are changed, rather than its target's.

def lstat(self)

Expand source code

def lstat(self):
    """
    Like stat(), except if the path points to a symlink, the symlink's
    status information is returned, rather than its target's.
    """
    return os.lstat(self)

Like stat(), except if the path points to a symlink, the symlink's status information is returned, rather than its target's.

def mkdir(self, mode=511, parents=False, exist_ok=False)

Expand source code

def mkdir(self, mode=0o777, parents=False, exist_ok=False):
    """
    Create a new directory at this given path.
    """
    try:
        os.mkdir(self, mode)
    except FileNotFoundError:
        if not parents or self.parent == self:
            raise
        self.parent.mkdir(parents=True, exist_ok=True)
        self.mkdir(mode, parents=False, exist_ok=exist_ok)
    except OSError:
        # Cannot rely on checking for EEXIST, since the operating system
        # could give priority to other errors like EACCES or EROFS
        if not exist_ok or not self.is_dir():
            raise

Create a new directory at this given path.

def move(self, target)

Expand source code

def move(self, target):
    """
    Recursively move this file or directory tree to the given destination.
    """
    # Use os.replace() if the target is os.PathLike and on the same FS.
    try:
        target = self.with_segments(target)
    except TypeError:
        pass
    else:
        ensure_different_files(self, target)
        try:
            os.replace(self, target)
        except OSError as err:
            if err.errno != EXDEV:
                raise
        else:
            return target.joinpath()  # Empty join to ensure fresh metadata.
    # Fall back to copy+delete.
    target = self.copy(target, follow_symlinks=False, preserve_metadata=True)
    self._delete()
    return target

Recursively move this file or directory tree to the given destination.

def move_into(self, target_dir)

Expand source code

def move_into(self, target_dir):
    """
    Move this file or directory tree into the given existing directory.
    """
    name = self.name
    if not name:
        raise ValueError(f"{self!r} has an empty name")
    elif hasattr(target_dir, 'with_segments'):
        target = target_dir / name
    else:
        target = self.with_segments(target_dir, name)
    return self.move(target)

Move this file or directory tree into the given existing directory.

def open(self, mode='r', buffering=-1, encoding=None, errors=None, newline=None)

Expand source code

def open(self, mode='r', buffering=-1, encoding=None,
         errors=None, newline=None):
    """
    Open the file pointed to by this path and return a file object, as
    the built-in open() function does.
    """
    if "b" not in mode:
        encoding = io.text_encoding(encoding)
    return io.open(self, mode, buffering, encoding, errors, newline)

Open the file pointed to by this path and return a file object, as the built-in open() function does.

def owner(self, *, follow_symlinks=True)

Expand source code

def owner(self, *, follow_symlinks=True):
    """
    Return the login name of the file owner.
    """
    uid = self.stat(follow_symlinks=follow_symlinks).st_uid
    return pwd.getpwuid(uid).pw_name

Return the login name of the file owner.

def read_bytes(self)

Expand source code

def read_bytes(self):
    """
    Open the file in bytes mode, read it, and close the file.
    """
    with self.open(mode='rb', buffering=0) as f:
        return f.read()

Open the file in bytes mode, read it, and close the file.

def read_text(self, encoding=None, errors=None, newline=None)

Expand source code

def read_text(self, encoding=None, errors=None, newline=None):
    """
    Open the file in text mode, read it, and close the file.
    """
    # Call io.text_encoding() here to ensure any warning is raised at an
    # appropriate stack level.
    encoding = io.text_encoding(encoding)
    with self.open(mode='r', encoding=encoding, errors=errors, newline=newline) as f:
        return f.read()

Open the file in text mode, read it, and close the file.

def readlink(self)

Expand source code

def readlink(self):
    """
    Return the path to which the symbolic link points.
    """
    return self.with_segments(os.readlink(self))

Return the path to which the symbolic link points.

def rename(self, target)

Expand source code

def rename(self, target):
    """
    Rename this path to the target path.

    The target path may be absolute or relative. Relative paths are
    interpreted relative to the current working directory, *not* the
    directory of the Path object.

    Returns the new Path instance pointing to the target path.
    """
    os.rename(self, target)
    if not hasattr(target, 'with_segments'):
        target = self.with_segments(target)
    return target

Rename this path to the target path.

The target path may be absolute or relative. Relative paths are interpreted relative to the current working directory, not the directory of the Path object.

Returns the new Path instance pointing to the target path.

def replace(self, target)

Expand source code

def replace(self, target):
    """
    Rename this path to the target path, overwriting if that path exists.

    The target path may be absolute or relative. Relative paths are
    interpreted relative to the current working directory, *not* the
    directory of the Path object.

    Returns the new Path instance pointing to the target path.
    """
    os.replace(self, target)
    if not hasattr(target, 'with_segments'):
        target = self.with_segments(target)
    return target

Rename this path to the target path, overwriting if that path exists.

The target path may be absolute or relative. Relative paths are interpreted relative to the current working directory, not the directory of the Path object.

Returns the new Path instance pointing to the target path.

def resolve(self, strict=False)

Expand source code

def resolve(self, strict=False):
    """
    Make the path absolute, resolving all symlinks on the way and also
    normalizing it.
    """

    return self.with_segments(os.path.realpath(self, strict=strict))

Make the path absolute, resolving all symlinks on the way and also normalizing it.

def rglob(self, pattern, *, case_sensitive=None, recurse_symlinks=False)

Expand source code

def rglob(self, pattern, *, case_sensitive=None, recurse_symlinks=False):
    """Recursively yield all existing files (of any kind, including
    directories) matching the given relative pattern, anywhere in
    this subtree.
    """
    sys.audit("pathlib.Path.rglob", self, pattern)
    pattern = self.parser.join('**', pattern)
    return self.glob(pattern, case_sensitive=case_sensitive, recurse_symlinks=recurse_symlinks)

Recursively yield all existing files (of any kind, including directories) matching the given relative pattern, anywhere in this subtree.

def rmdir(self)

Expand source code

def rmdir(self):
    """
    Remove this directory.  The directory must be empty.
    """
    os.rmdir(self)

Remove this directory. The directory must be empty.

def samefile(self, other_path)

Expand source code

def samefile(self, other_path):
    """Return whether other_path is the same or not as this file
    (as returned by os.path.samefile()).
    """
    st = self.stat()
    try:
        other_st = other_path.stat()
    except AttributeError:
        other_st = self.with_segments(other_path).stat()
    return (st.st_ino == other_st.st_ino and
            st.st_dev == other_st.st_dev)

Return whether other_path is the same or not as this file (as returned by os.path.samefile()).

def stat(self, *, follow_symlinks=True)

Expand source code

def stat(self, *, follow_symlinks=True):
    """
    Return the result of the stat() system call on this path, like
    os.stat() does.
    """
    return os.stat(self, follow_symlinks=follow_symlinks)

Return the result of the stat() system call on this path, like os.stat() does.

def symlink_to(self, target, target_is_directory=False)

Expand source code

def symlink_to(self, target, target_is_directory=False):
    """
    Make this path a symlink pointing to the target path.
    Note the order of arguments (link, target) is the reverse of os.symlink.
    """
    os.symlink(target, self, target_is_directory)

Make this path a symlink pointing to the target path. Note the order of arguments (link, target) is the reverse of os.symlink.

def touch(self, mode=438, exist_ok=True)

Expand source code

def touch(self, mode=0o666, exist_ok=True):
    """
    Create this file with the given access mode, if it doesn't exist.
    """

    if exist_ok:
        # First try to bump modification time
        # Implementation note: GNU touch uses the UTIME_NOW option of
        # the utimensat() / futimens() functions.
        try:
            os.utime(self, None)
        except OSError:
            # Avoid exception chaining
            pass
        else:
            return
    flags = os.O_CREAT | os.O_WRONLY
    if not exist_ok:
        flags |= os.O_EXCL
    fd = os.open(self, flags, mode)
    os.close(fd)

Create this file with the given access mode, if it doesn't exist.

def unlink(self, missing_ok=False)

Expand source code

def unlink(self, missing_ok=False):
    """
    Remove this file or link.
    If the path is a directory, use rmdir() instead.
    """
    try:
        os.unlink(self)
    except FileNotFoundError:
        if not missing_ok:
            raise

Remove this file or link. If the path is a directory, use rmdir() instead.

def walk(self, top_down=True, on_error=None, follow_symlinks=False)

Expand source code

def walk(self, top_down=True, on_error=None, follow_symlinks=False):
    """Walk the directory tree from this directory, similar to os.walk()."""
    sys.audit("pathlib.Path.walk", self, on_error, follow_symlinks)
    root_dir = str(self)
    if not follow_symlinks:
        follow_symlinks = os._walk_symlinks_as_files
    results = os.walk(root_dir, top_down, on_error, follow_symlinks)
    for path_str, dirnames, filenames in results:
        if root_dir == '.':
            path_str = path_str[2:]
        yield self._from_parsed_string(path_str), dirnames, filenames

Walk the directory tree from this directory, similar to os.walk().

def write_bytes(self, data)

Expand source code

def write_bytes(self, data):
    """
    Open the file in bytes mode, write to it, and close the file.
    """
    # type-check for the buffer interface before truncating the file
    view = memoryview(data)
    with self.open(mode='wb') as f:
        return f.write(view)

Open the file in bytes mode, write to it, and close the file.

def write_text(self, data, encoding=None, errors=None, newline=None)

Expand source code

def write_text(self, data, encoding=None, errors=None, newline=None):
    """
    Open the file in text mode, write to it, and close the file.
    """
    # Call io.text_encoding() here to ensure any warning is raised at an
    # appropriate stack level.
    encoding = io.text_encoding(encoding)
    if not isinstance(data, str):
        raise TypeError('data must be str, not %s' %
                        data.__class__.__name__)
    with self.open(mode='w', encoding=encoding, errors=errors, newline=newline) as f:
        return f.write(data)

Open the file in text mode, write to it, and close the file.

class ProxSparsePrunedLoRALinear (in_features: int, out_features: int, pruned_weight: torch.Tensor, mask: torch.Tensor, bias: torch.Tensor | None = None, lora_r: int = 0, lora_alpha: float = 1.0, device: torch.device | None = None, dtype: torch.dtype | None = None)

Expand source code

class ProxSparsePrunedLoRALinear(nn.Module):
    """
    Frozen pruned base weight + trainable LoRA adapters.

    NOTE: `weight` is the ALREADY-PRUNED weight matrix. Forward does not
    re-apply the mask. `mask` is kept only as a compact bool buffer for
    downstream analysis / merge operations.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        pruned_weight: torch.Tensor,
        mask: torch.Tensor,
        bias: torch.Tensor | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        target_device = device if device is not None else pruned_weight.device
        target_dtype = dtype if dtype is not None else pruned_weight.dtype

        # PS-FIX-02: weight is already masked; do not multiply again.
        self.weight = nn.Parameter(
            pruned_weight.to(device=target_device, dtype=target_dtype),
            requires_grad=False,
        )
        # PS-FIX-03: bool buffer (1 byte per element, non-persistent).
        self.register_buffer(
            "mask",
            mask.bool().to(device=target_device),
            persistent=False,
        )

        if bias is not None:
            self.bias = nn.Parameter(
                bias.to(device=target_device, dtype=target_dtype),
                requires_grad=False,
            )
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty(
                    (self.lora_r, in_features),
                    dtype=torch.float32,
                    device=target_device,
                )
            )
            self.lora_B = nn.Parameter(
                torch.zeros(
                    (out_features, self.lora_r),
                    dtype=torch.float32,
                    device=target_device,
                )
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        pruned_weight: torch.Tensor,
        mask: torch.Tensor,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> ProxSparsePrunedLoRALinear:
        base_mod = getattr(linear, "base_layer", linear)

        # PS-FIX-09: robust in/out detection for Linear and Conv1D.
        cls_name = base_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(base_mod, "nf")

        if is_conv1d:
            in_features = int(getattr(base_mod, "nx", pruned_weight.shape[1]))
            out_features = int(getattr(base_mod, "nf", pruned_weight.shape[0]))
        else:
            in_features = int(getattr(base_mod, "in_features", pruned_weight.shape[1]))
            out_features = int(getattr(base_mod, "out_features", pruned_weight.shape[0]))

        bias = None
        if getattr(base_mod, "bias", None) is not None:
            bias = base_mod.bias.detach()

        instance = cls(
            in_features=in_features,
            out_features=out_features,
            pruned_weight=pruned_weight,
            mask=mask,
            bias=bias,
            lora_r=lora_r,
            lora_alpha=lora_alpha,
            device=pruned_weight.device,
            dtype=pruned_weight.dtype,
        )

        # Preserve existing source LoRA parameters if shapes match.
        if hasattr(linear, "lora_A") and hasattr(linear, "lora_B"):
            src_A = getattr(linear.lora_A, "weight", linear.lora_A)
            src_B = getattr(linear.lora_B, "weight", linear.lora_B)
            if (
                src_A is not None
                and src_B is not None
                and instance.lora_A is not None
                and instance.lora_B is not None
            ):
                if src_A.shape == instance.lora_A.shape:
                    instance.lora_A.data.copy_(src_A.data)
                if src_B.shape == instance.lora_B.shape:
                    instance.lora_B.data.copy_(src_B.data)
        return instance

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # PS-FIX-02: no re-multiplication by mask; weight is already masked.
        w = self.weight.to(dtype=x.dtype, device=x.device)
        # PS-FIX-06: bias cast includes device.
        b = None
        if self.bias is not None:
            b = self.bias.to(dtype=x.dtype, device=x.device)

        out = F.linear(x, w, b)

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_delta = F.linear(F.linear(x.float(), self.lora_A), self.lora_B) * self.scaling
            out = out + lora_delta.to(out.dtype)
        return out

Frozen pruned base weight + trainable LoRA adapters.

NOTE: weight is the ALREADY-PRUNED weight matrix. Forward does not re-apply the mask. mask is kept only as a compact bool buffer for downstream analysis / merge operations.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Module, pruned_weight: torch.Tensor, mask: torch.Tensor, lora_r: int = 0, lora_alpha: float = 1.0) ‑> ProxSparsePrunedLoRALinear

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    # PS-FIX-02: no re-multiplication by mask; weight is already masked.
    w = self.weight.to(dtype=x.dtype, device=x.device)
    # PS-FIX-06: bias cast includes device.
    b = None
    if self.bias is not None:
        b = self.bias.to(dtype=x.dtype, device=x.device)

    out = F.linear(x, w, b)

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_delta = F.linear(F.linear(x.float(), self.lora_A), self.lora_B) * self.scaling
        out = out + lora_delta.to(out.dtype)
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class ProxSparseCompressor (model: nn.Module, sparsity: float = 0.5, lambda_prox: float = 0.01, steps: int = 5, learning_rate: float = 0.5, structured_24: bool = False, block_size: tuple[int, int] = (1, 1), skip_attention_qkv: bool = True, use_lora: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, max_error: float | None = None, fallback_to_original: bool = True, min_numel: int = 16, target_modules: list[str] | None = None, target_attention_only: bool = False, device: str | torch.device | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("proxsparse")
@COMPRESSOR_REGISTRY.register("proxsparse_pruner")
@COMPRESSOR_REGISTRY.register("proxsparse_compressor")
class ProxSparsePruner(BaseCompressor):
    """
    ProxSparse: proximal gradient sparsity with structured / 2:4 / percentile
    constraint, optional LoRA fine-tuning, and safe fallback.
    """

    def __init__(
        self,
        model: nn.Module,
        sparsity: float = 0.5,
        lambda_prox: float = 0.01,
        steps: int = 5,
        learning_rate: float = 0.5,  # PS-FIX-04
        structured_24: bool = False,
        block_size: tuple[int, int] = (1, 1),
        skip_attention_qkv: bool = True,
        use_lora: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        max_error: float | None = None,
        fallback_to_original: bool = True,
        min_numel: int = 16,
        target_modules: list[str] | None = None,
        target_attention_only: bool = False,
        device: str | torch.device | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="proxsparse", **kwargs)
        cfg_obj = cfg or kwargs.get("comp_config")

        self.sparsity = max(
            0.0,
            min(
                float(
                    sparsity
                    if sparsity is not None
                    else _cfg_get(
                        cfg_obj,
                        "proxsparse_sparsity",
                        _cfg_get(cfg_obj, "sparsity", 0.5),
                    )
                ),
                0.95,
            ),
        )
        self.lambda_prox = max(
            0.0,
            float(
                lambda_prox
                if lambda_prox is not None
                else _cfg_get(cfg_obj, "proxsparse_lambda", _cfg_get(cfg_obj, "lambda_prox", 0.01))
            ),
        )
        self.steps = max(
            1,
            int(
                steps
                if steps is not None
                else _cfg_get(cfg_obj, "proxsparse_steps", _cfg_get(cfg_obj, "steps", 5))
            ),
        )
        # PS-FIX-04
        self.learning_rate = max(
            1e-4,
            float(
                learning_rate
                if learning_rate is not None
                else _cfg_get(cfg_obj, "proxsparse_lr", _cfg_get(cfg_obj, "learning_rate", 0.5))
            ),
        )
        self.structured_24 = bool(
            structured_24
            if structured_24 is not None
            else _cfg_get(
                cfg_obj,
                "proxsparse_structured_24",
                _cfg_get(cfg_obj, "structured_24", False),
            )
        )
        self.block_size = block_size or _cfg_get(cfg_obj, "block_size", (1, 1))
        self.target_attention_only = bool(
            target_attention_only or _cfg_get(cfg_obj, "target_attention_only", False)
        )
        self.skip_attention_qkv = (
            bool(
                skip_attention_qkv
                if skip_attention_qkv is not None
                else _cfg_get(cfg_obj, "skip_attention_qkv", True)
            )
            and not self.target_attention_only
        )
        self.use_lora = bool(use_lora or _cfg_get(cfg_obj, "use_lora", False) or (lora_r > 0))
        self.lora_r = int(lora_r or _cfg_get(cfg_obj, "lora_r", 0))
        self.lora_alpha = float(lora_alpha or _cfg_get(cfg_obj, "lora_alpha", 1.0))
        self.min_numel = max(1, int(min_numel or _cfg_get(cfg_obj, "min_numel", 16)))
        self.target_modules = target_modules

        if max_error is not None:
            self.max_error: float | None = float(max_error)
        elif _cfg_get(cfg_obj, "max_error", None) is not None:
            self.max_error = float(_cfg_get(cfg_obj, "max_error"))
        else:
            self.max_error = None

        self.fallback_to_original = bool(
            fallback_to_original
            if fallback_to_original is not None
            else _cfg_get(cfg_obj, "fallback_to_original", True)
        )

        # Device
        if device is not None:
            try:
                self.target_device = torch.device(device)
            except Exception:
                self.target_device = (
                    next(model.parameters()).device
                    if list(model.parameters())
                    else torch.device("cpu")
                )
        else:
            self.target_device = (
                next(model.parameters()).device if list(model.parameters()) else torch.device("cpu")
            )

        self.visited_modules: set[int] = set()
        self.compressed_layers: list[str] = []
        self.protected_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.failed_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_actual_sparsity: dict[str, float] = {}
        self.stats: dict[str, Any] = {}

    # ------------------------------------------------------------------
    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if self.target_attention_only:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        is_gqa = num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads

        name_l = name.lower()
        is_attn_qkv = any(
            p in name_l
            for p in (
                "q_proj",
                "k_proj",
                "v_proj",
                "c_attn",
                "query",
                "key",
                "value",
                "q_lin",
                "k_lin",
                "v_lin",
            )
        )
        return (is_gqa or self.skip_attention_qkv) and is_attn_qkv

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        if any(
            k in name_l
            for k in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embed",
                "embedding",
                "cls_head",
                "head",
                "logits_processor",
            )
        ):
            self.protected_layers.append(name)
            return False

        if self.target_attention_only:
            attn_keywords = (
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
                "dense",
            )
            if not any(k in name_l for k in attn_keywords):
                self.skipped_layers.append(name)
                return False

        if self._is_gqa_asymmetric(name, module):
            self.skipped_layers.append(name)
            return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__
        raw_w = getattr(base_mod, "weight", None)
        if raw_w is None or not isinstance(raw_w, torch.Tensor):
            return False

        is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_mod, "nf")
        is_linear_like = isinstance(base_mod, nn.Linear) or is_conv1d
        if not is_linear_like or isinstance(module, ProxSparsePrunedLoRALinear):
            return False

        if raw_w.numel() < self.min_numel:
            self.skipped_layers.append(name)
            return False

        if not should_compress_linear(self.model, name):
            self.protected_layers.append(name)
            return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    # ------------------------------------------------------------------
    def _soft_threshold(self, tensor: torch.Tensor, threshold: float) -> torch.Tensor:
        """prox_{a ||.||_1}(x) = sign(x) * max(|x| - a, 0)."""
        return torch.sign(tensor) * torch.clamp(torch.abs(tensor) - threshold, min=0.0)

    def _project_sparse(self, tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Hard projection onto sparse constraint S. Returns (pruned, mask)."""
        if self.structured_24:
            pruned, mask = apply_24_structured_mask(tensor, saliency=torch.abs(tensor), dim=1)
            return pruned, (mask > 0)
        elif self.block_size != (1, 1):
            pruned, mask = apply_block_mask(
                tensor, block_size=self.block_size, sparsity=self.sparsity
            )
            return pruned, (mask > 0)
        else:
            flat = tensor.abs().reshape(-1)
            k = max(1, int(flat.numel() * (1.0 - self.sparsity)))
            thr = torch.topk(flat, k, largest=True).values[-1]
            mask = tensor.abs() >= thr
            return tensor * mask.float(), mask

    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        self.compressed_layers.clear()
        self.visited_modules.clear()
        self.protected_layers.clear()
        self.skipped_layers.clear()
        self.fallback_layers.clear()
        self.failed_layers.clear()
        self.layer_errors.clear()
        self.layer_actual_sparsity.clear()

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            mod_id = id(module)
            if mod_id in self.visited_modules:
                self.skipped_layers.append(name)
                continue
            self.visited_modules.add(mod_id)

            try:
                base_mod = getattr(module, "base_layer", module)
                raw_w = base_mod.weight.data
                cls_name = base_mod.__class__.__name__
                is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_mod, "nf")
                device = raw_w.device
                dtype = raw_w.dtype

                # Always operate in (out, in) orientation.
                W_0 = raw_w.t().contiguous().float() if is_conv1d else raw_w.contiguous().float()

                # Early NaN/Inf guard on source.
                if not torch.isfinite(W_0).all():
                    logger.warning("ProxSparse: %s has non-finite weights; skipping.", name)
                    self.failed_layers.append(name)
                    continue

                # ------------------------------------------------
                # PS-FIX-01: Correct proximal gradient iteration.
                # ------------------------------------------------
                eta = self.learning_rate
                W = W_0.clone()
                mask = None
                for _ in range(self.steps):
                    grad_W = W - W_0  # grad of 0.5*||W - W_0||^2
                    W_grad = W - eta * grad_W  # gradient descent step
                    W_prox = self._soft_threshold(W_grad, eta * self.lambda_prox)
                    W, mask = self._project_sparse(W_prox)

                # Final projection (mask from last iteration already applied).
                W_final = torch.nan_to_num(W, nan=0.0, posinf=0.0, neginf=0.0)

                w_norm = torch.norm(W_0, p="fro").clamp_min(1e-12).item()
                err = (torch.norm(W_0 - W_final, p="fro") / w_norm).item()
                self.layer_errors[name] = float(err)

                # Actual sparsity in the final weight.
                actual_zero = (W_final.abs() < 1e-12).float().mean().item()
                self.layer_actual_sparsity[name] = float(actual_zero)

                is_invalid = torch.isnan(W_final).any() or torch.isinf(W_final).any()
                if self.max_error is not None and err > self.max_error:
                    logger.warning(
                        "ProxSparse: reconstruction error %.4f on %s exceeded %.4f.",
                        err,
                        name,
                        self.max_error,
                    )
                    is_invalid = True

                if is_invalid and self.fallback_to_original:
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

                # ---- Commit ----
                if self.use_lora and self.lora_r > 0:
                    lora_layer = ProxSparsePrunedLoRALinear.from_linear(
                        linear=module,
                        pruned_weight=W_final.to(dtype=dtype, device=device),
                        mask=mask.to(device=device),
                        lora_r=self.lora_r,
                        lora_alpha=self.lora_alpha,
                    )
                    replace_module(self.model, name, lora_layer)
                else:
                    final_w = W_final.t().contiguous() if is_conv1d else W_final.contiguous()
                    base_mod.weight.data.copy_(final_w.to(dtype=dtype, device=device))

                self.compressed_layers.append(name)

            except Exception as e_comp:
                logger.warning("ProxSparse: failed on %s: %s", name, e_comp)
                self.failed_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False

        # Forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                dummy = torch.randint(0, 100, (1, 8), device=self.target_device)
                try:
                    self.model(input_ids=dummy)
                except Exception:
                    self.model(dummy)
        except Exception as ex_sanity:
            logger.warning("ProxSparse forward sanity check failed: %s", ex_sanity)
            forward_check_success = False

        self._compressed_params = self._count_params(self.model)
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )

        # PS-FIX-08: effective_parameter_ratio separate from compression_ratio.
        effective_ratio = 1.0
        if self.compressed_layers:
            ratios = []
            for n in self.compressed_layers:
                actual_sp = self.layer_actual_sparsity.get(n, 0.0)
                # Non-zeros in the pruned base + LoRA params (if any)
                ratio = 1.0 - actual_sp
                if self.use_lora and self.lora_r > 0:
                    # Very rough estimate; LoRA adds lora_r*(in+out) params.
                    ratio += 0.02  # approximate placeholder; real accounting per layer
                ratios.append(ratio)
            effective_ratio = sum(ratios) / len(ratios)

        stats = self.get_compression_stats()
        stats.update(
            {
                "effective_parameter_ratio": round(effective_ratio, 4),
                "pruned_layers": self.compressed_layers,
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "protected_layers": self.protected_layers,
                "protected_layer_count": len(self.protected_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "failed_layers": self.failed_layers,
                "failed_layer_count": len(self.failed_layers),
                "layer_errors": self.layer_errors,
                "layer_actual_sparsity": self.layer_actual_sparsity,
                "avg_reconstruction_error": avg_err,
                "forward_check_success": forward_check_success,
                "sparsity_target": self.sparsity,
                "lambda_prox": self.lambda_prox,
                "learning_rate": self.learning_rate,
                "structured_24": self.structured_24,
                "steps": self.steps,
                "skip_attention_qkv": self.skip_attention_qkv,
                "use_lora": self.use_lora,
                "lora_r": self.lora_r,
                "method": "proxsparse_regularized",
            }
        )
        self.stats = stats
        return self.model, stats

ProxSparse: proximal gradient sparsity with structured / 2:4 / percentile constraint, optional LoRA fine-tuning, and safe fallback.

Ancestors

BaseCompressor

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    self.compressed_layers.clear()
    self.visited_modules.clear()
    self.protected_layers.clear()
    self.skipped_layers.clear()
    self.fallback_layers.clear()
    self.failed_layers.clear()
    self.layer_errors.clear()
    self.layer_actual_sparsity.clear()

    for name, module in list(self.model.named_modules()):
        if not self._is_target_module(name, module):
            continue

        mod_id = id(module)
        if mod_id in self.visited_modules:
            self.skipped_layers.append(name)
            continue
        self.visited_modules.add(mod_id)

        try:
            base_mod = getattr(module, "base_layer", module)
            raw_w = base_mod.weight.data
            cls_name = base_mod.__class__.__name__
            is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_mod, "nf")
            device = raw_w.device
            dtype = raw_w.dtype

            # Always operate in (out, in) orientation.
            W_0 = raw_w.t().contiguous().float() if is_conv1d else raw_w.contiguous().float()

            # Early NaN/Inf guard on source.
            if not torch.isfinite(W_0).all():
                logger.warning("ProxSparse: %s has non-finite weights; skipping.", name)
                self.failed_layers.append(name)
                continue

            # ------------------------------------------------
            # PS-FIX-01: Correct proximal gradient iteration.
            # ------------------------------------------------
            eta = self.learning_rate
            W = W_0.clone()
            mask = None
            for _ in range(self.steps):
                grad_W = W - W_0  # grad of 0.5*||W - W_0||^2
                W_grad = W - eta * grad_W  # gradient descent step
                W_prox = self._soft_threshold(W_grad, eta * self.lambda_prox)
                W, mask = self._project_sparse(W_prox)

            # Final projection (mask from last iteration already applied).
            W_final = torch.nan_to_num(W, nan=0.0, posinf=0.0, neginf=0.0)

            w_norm = torch.norm(W_0, p="fro").clamp_min(1e-12).item()
            err = (torch.norm(W_0 - W_final, p="fro") / w_norm).item()
            self.layer_errors[name] = float(err)

            # Actual sparsity in the final weight.
            actual_zero = (W_final.abs() < 1e-12).float().mean().item()
            self.layer_actual_sparsity[name] = float(actual_zero)

            is_invalid = torch.isnan(W_final).any() or torch.isinf(W_final).any()
            if self.max_error is not None and err > self.max_error:
                logger.warning(
                    "ProxSparse: reconstruction error %.4f on %s exceeded %.4f.",
                    err,
                    name,
                    self.max_error,
                )
                is_invalid = True

            if is_invalid and self.fallback_to_original:
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                continue

            # ---- Commit ----
            if self.use_lora and self.lora_r > 0:
                lora_layer = ProxSparsePrunedLoRALinear.from_linear(
                    linear=module,
                    pruned_weight=W_final.to(dtype=dtype, device=device),
                    mask=mask.to(device=device),
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                )
                replace_module(self.model, name, lora_layer)
            else:
                final_w = W_final.t().contiguous() if is_conv1d else W_final.contiguous()
                base_mod.weight.data.copy_(final_w.to(dtype=dtype, device=device))

            self.compressed_layers.append(name)

        except Exception as e_comp:
            logger.warning("ProxSparse: failed on %s: %s", name, e_comp)
            self.failed_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False

    # Forward sanity check
    forward_check_success = True
    try:
        with torch.no_grad():
            dummy = torch.randint(0, 100, (1, 8), device=self.target_device)
            try:
                self.model(input_ids=dummy)
            except Exception:
                self.model(dummy)
    except Exception as ex_sanity:
        logger.warning("ProxSparse forward sanity check failed: %s", ex_sanity)
        forward_check_success = False

    self._compressed_params = self._count_params(self.model)
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )

    # PS-FIX-08: effective_parameter_ratio separate from compression_ratio.
    effective_ratio = 1.0
    if self.compressed_layers:
        ratios = []
        for n in self.compressed_layers:
            actual_sp = self.layer_actual_sparsity.get(n, 0.0)
            # Non-zeros in the pruned base + LoRA params (if any)
            ratio = 1.0 - actual_sp
            if self.use_lora and self.lora_r > 0:
                # Very rough estimate; LoRA adds lora_r*(in+out) params.
                ratio += 0.02  # approximate placeholder; real accounting per layer
            ratios.append(ratio)
        effective_ratio = sum(ratios) / len(ratios)

    stats = self.get_compression_stats()
    stats.update(
        {
            "effective_parameter_ratio": round(effective_ratio, 4),
            "pruned_layers": self.compressed_layers,
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "protected_layers": self.protected_layers,
            "protected_layer_count": len(self.protected_layers),
            "skipped_layers": self.skipped_layers,
            "skipped_layer_count": len(self.skipped_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "failed_layers": self.failed_layers,
            "failed_layer_count": len(self.failed_layers),
            "layer_errors": self.layer_errors,
            "layer_actual_sparsity": self.layer_actual_sparsity,
            "avg_reconstruction_error": avg_err,
            "forward_check_success": forward_check_success,
            "sparsity_target": self.sparsity,
            "lambda_prox": self.lambda_prox,
            "learning_rate": self.learning_rate,
            "structured_24": self.structured_24,
            "steps": self.steps,
            "skip_attention_qkv": self.skip_attention_qkv,
            "use_lora": self.use_lora,
            "lora_r": self.lora_r,
            "method": "proxsparse_regularized",
        }
    )
    self.stats = stats
    return self.model, stats

Return (compressed_model, statistics).

class ProxSparsePruner (model: nn.Module, sparsity: float = 0.5, lambda_prox: float = 0.01, steps: int = 5, learning_rate: float = 0.5, structured_24: bool = False, block_size: tuple[int, int] = (1, 1), skip_attention_qkv: bool = True, use_lora: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, max_error: float | None = None, fallback_to_original: bool = True, min_numel: int = 16, target_modules: list[str] | None = None, target_attention_only: bool = False, device: str | torch.device | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("proxsparse")
@COMPRESSOR_REGISTRY.register("proxsparse_pruner")
@COMPRESSOR_REGISTRY.register("proxsparse_compressor")
class ProxSparsePruner(BaseCompressor):
    """
    ProxSparse: proximal gradient sparsity with structured / 2:4 / percentile
    constraint, optional LoRA fine-tuning, and safe fallback.
    """

    def __init__(
        self,
        model: nn.Module,
        sparsity: float = 0.5,
        lambda_prox: float = 0.01,
        steps: int = 5,
        learning_rate: float = 0.5,  # PS-FIX-04
        structured_24: bool = False,
        block_size: tuple[int, int] = (1, 1),
        skip_attention_qkv: bool = True,
        use_lora: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        max_error: float | None = None,
        fallback_to_original: bool = True,
        min_numel: int = 16,
        target_modules: list[str] | None = None,
        target_attention_only: bool = False,
        device: str | torch.device | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="proxsparse", **kwargs)
        cfg_obj = cfg or kwargs.get("comp_config")

        self.sparsity = max(
            0.0,
            min(
                float(
                    sparsity
                    if sparsity is not None
                    else _cfg_get(
                        cfg_obj,
                        "proxsparse_sparsity",
                        _cfg_get(cfg_obj, "sparsity", 0.5),
                    )
                ),
                0.95,
            ),
        )
        self.lambda_prox = max(
            0.0,
            float(
                lambda_prox
                if lambda_prox is not None
                else _cfg_get(cfg_obj, "proxsparse_lambda", _cfg_get(cfg_obj, "lambda_prox", 0.01))
            ),
        )
        self.steps = max(
            1,
            int(
                steps
                if steps is not None
                else _cfg_get(cfg_obj, "proxsparse_steps", _cfg_get(cfg_obj, "steps", 5))
            ),
        )
        # PS-FIX-04
        self.learning_rate = max(
            1e-4,
            float(
                learning_rate
                if learning_rate is not None
                else _cfg_get(cfg_obj, "proxsparse_lr", _cfg_get(cfg_obj, "learning_rate", 0.5))
            ),
        )
        self.structured_24 = bool(
            structured_24
            if structured_24 is not None
            else _cfg_get(
                cfg_obj,
                "proxsparse_structured_24",
                _cfg_get(cfg_obj, "structured_24", False),
            )
        )
        self.block_size = block_size or _cfg_get(cfg_obj, "block_size", (1, 1))
        self.target_attention_only = bool(
            target_attention_only or _cfg_get(cfg_obj, "target_attention_only", False)
        )
        self.skip_attention_qkv = (
            bool(
                skip_attention_qkv
                if skip_attention_qkv is not None
                else _cfg_get(cfg_obj, "skip_attention_qkv", True)
            )
            and not self.target_attention_only
        )
        self.use_lora = bool(use_lora or _cfg_get(cfg_obj, "use_lora", False) or (lora_r > 0))
        self.lora_r = int(lora_r or _cfg_get(cfg_obj, "lora_r", 0))
        self.lora_alpha = float(lora_alpha or _cfg_get(cfg_obj, "lora_alpha", 1.0))
        self.min_numel = max(1, int(min_numel or _cfg_get(cfg_obj, "min_numel", 16)))
        self.target_modules = target_modules

        if max_error is not None:
            self.max_error: float | None = float(max_error)
        elif _cfg_get(cfg_obj, "max_error", None) is not None:
            self.max_error = float(_cfg_get(cfg_obj, "max_error"))
        else:
            self.max_error = None

        self.fallback_to_original = bool(
            fallback_to_original
            if fallback_to_original is not None
            else _cfg_get(cfg_obj, "fallback_to_original", True)
        )

        # Device
        if device is not None:
            try:
                self.target_device = torch.device(device)
            except Exception:
                self.target_device = (
                    next(model.parameters()).device
                    if list(model.parameters())
                    else torch.device("cpu")
                )
        else:
            self.target_device = (
                next(model.parameters()).device if list(model.parameters()) else torch.device("cpu")
            )

        self.visited_modules: set[int] = set()
        self.compressed_layers: list[str] = []
        self.protected_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.failed_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_actual_sparsity: dict[str, float] = {}
        self.stats: dict[str, Any] = {}

    # ------------------------------------------------------------------
    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if self.target_attention_only:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        is_gqa = num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads

        name_l = name.lower()
        is_attn_qkv = any(
            p in name_l
            for p in (
                "q_proj",
                "k_proj",
                "v_proj",
                "c_attn",
                "query",
                "key",
                "value",
                "q_lin",
                "k_lin",
                "v_lin",
            )
        )
        return (is_gqa or self.skip_attention_qkv) and is_attn_qkv

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        if any(
            k in name_l
            for k in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embed",
                "embedding",
                "cls_head",
                "head",
                "logits_processor",
            )
        ):
            self.protected_layers.append(name)
            return False

        if self.target_attention_only:
            attn_keywords = (
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
                "dense",
            )
            if not any(k in name_l for k in attn_keywords):
                self.skipped_layers.append(name)
                return False

        if self._is_gqa_asymmetric(name, module):
            self.skipped_layers.append(name)
            return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__
        raw_w = getattr(base_mod, "weight", None)
        if raw_w is None or not isinstance(raw_w, torch.Tensor):
            return False

        is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_mod, "nf")
        is_linear_like = isinstance(base_mod, nn.Linear) or is_conv1d
        if not is_linear_like or isinstance(module, ProxSparsePrunedLoRALinear):
            return False

        if raw_w.numel() < self.min_numel:
            self.skipped_layers.append(name)
            return False

        if not should_compress_linear(self.model, name):
            self.protected_layers.append(name)
            return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    # ------------------------------------------------------------------
    def _soft_threshold(self, tensor: torch.Tensor, threshold: float) -> torch.Tensor:
        """prox_{a ||.||_1}(x) = sign(x) * max(|x| - a, 0)."""
        return torch.sign(tensor) * torch.clamp(torch.abs(tensor) - threshold, min=0.0)

    def _project_sparse(self, tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Hard projection onto sparse constraint S. Returns (pruned, mask)."""
        if self.structured_24:
            pruned, mask = apply_24_structured_mask(tensor, saliency=torch.abs(tensor), dim=1)
            return pruned, (mask > 0)
        elif self.block_size != (1, 1):
            pruned, mask = apply_block_mask(
                tensor, block_size=self.block_size, sparsity=self.sparsity
            )
            return pruned, (mask > 0)
        else:
            flat = tensor.abs().reshape(-1)
            k = max(1, int(flat.numel() * (1.0 - self.sparsity)))
            thr = torch.topk(flat, k, largest=True).values[-1]
            mask = tensor.abs() >= thr
            return tensor * mask.float(), mask

    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        self.compressed_layers.clear()
        self.visited_modules.clear()
        self.protected_layers.clear()
        self.skipped_layers.clear()
        self.fallback_layers.clear()
        self.failed_layers.clear()
        self.layer_errors.clear()
        self.layer_actual_sparsity.clear()

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            mod_id = id(module)
            if mod_id in self.visited_modules:
                self.skipped_layers.append(name)
                continue
            self.visited_modules.add(mod_id)

            try:
                base_mod = getattr(module, "base_layer", module)
                raw_w = base_mod.weight.data
                cls_name = base_mod.__class__.__name__
                is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_mod, "nf")
                device = raw_w.device
                dtype = raw_w.dtype

                # Always operate in (out, in) orientation.
                W_0 = raw_w.t().contiguous().float() if is_conv1d else raw_w.contiguous().float()

                # Early NaN/Inf guard on source.
                if not torch.isfinite(W_0).all():
                    logger.warning("ProxSparse: %s has non-finite weights; skipping.", name)
                    self.failed_layers.append(name)
                    continue

                # ------------------------------------------------
                # PS-FIX-01: Correct proximal gradient iteration.
                # ------------------------------------------------
                eta = self.learning_rate
                W = W_0.clone()
                mask = None
                for _ in range(self.steps):
                    grad_W = W - W_0  # grad of 0.5*||W - W_0||^2
                    W_grad = W - eta * grad_W  # gradient descent step
                    W_prox = self._soft_threshold(W_grad, eta * self.lambda_prox)
                    W, mask = self._project_sparse(W_prox)

                # Final projection (mask from last iteration already applied).
                W_final = torch.nan_to_num(W, nan=0.0, posinf=0.0, neginf=0.0)

                w_norm = torch.norm(W_0, p="fro").clamp_min(1e-12).item()
                err = (torch.norm(W_0 - W_final, p="fro") / w_norm).item()
                self.layer_errors[name] = float(err)

                # Actual sparsity in the final weight.
                actual_zero = (W_final.abs() < 1e-12).float().mean().item()
                self.layer_actual_sparsity[name] = float(actual_zero)

                is_invalid = torch.isnan(W_final).any() or torch.isinf(W_final).any()
                if self.max_error is not None and err > self.max_error:
                    logger.warning(
                        "ProxSparse: reconstruction error %.4f on %s exceeded %.4f.",
                        err,
                        name,
                        self.max_error,
                    )
                    is_invalid = True

                if is_invalid and self.fallback_to_original:
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

                # ---- Commit ----
                if self.use_lora and self.lora_r > 0:
                    lora_layer = ProxSparsePrunedLoRALinear.from_linear(
                        linear=module,
                        pruned_weight=W_final.to(dtype=dtype, device=device),
                        mask=mask.to(device=device),
                        lora_r=self.lora_r,
                        lora_alpha=self.lora_alpha,
                    )
                    replace_module(self.model, name, lora_layer)
                else:
                    final_w = W_final.t().contiguous() if is_conv1d else W_final.contiguous()
                    base_mod.weight.data.copy_(final_w.to(dtype=dtype, device=device))

                self.compressed_layers.append(name)

            except Exception as e_comp:
                logger.warning("ProxSparse: failed on %s: %s", name, e_comp)
                self.failed_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False

        # Forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                dummy = torch.randint(0, 100, (1, 8), device=self.target_device)
                try:
                    self.model(input_ids=dummy)
                except Exception:
                    self.model(dummy)
        except Exception as ex_sanity:
            logger.warning("ProxSparse forward sanity check failed: %s", ex_sanity)
            forward_check_success = False

        self._compressed_params = self._count_params(self.model)
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )

        # PS-FIX-08: effective_parameter_ratio separate from compression_ratio.
        effective_ratio = 1.0
        if self.compressed_layers:
            ratios = []
            for n in self.compressed_layers:
                actual_sp = self.layer_actual_sparsity.get(n, 0.0)
                # Non-zeros in the pruned base + LoRA params (if any)
                ratio = 1.0 - actual_sp
                if self.use_lora and self.lora_r > 0:
                    # Very rough estimate; LoRA adds lora_r*(in+out) params.
                    ratio += 0.02  # approximate placeholder; real accounting per layer
                ratios.append(ratio)
            effective_ratio = sum(ratios) / len(ratios)

        stats = self.get_compression_stats()
        stats.update(
            {
                "effective_parameter_ratio": round(effective_ratio, 4),
                "pruned_layers": self.compressed_layers,
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "protected_layers": self.protected_layers,
                "protected_layer_count": len(self.protected_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "failed_layers": self.failed_layers,
                "failed_layer_count": len(self.failed_layers),
                "layer_errors": self.layer_errors,
                "layer_actual_sparsity": self.layer_actual_sparsity,
                "avg_reconstruction_error": avg_err,
                "forward_check_success": forward_check_success,
                "sparsity_target": self.sparsity,
                "lambda_prox": self.lambda_prox,
                "learning_rate": self.learning_rate,
                "structured_24": self.structured_24,
                "steps": self.steps,
                "skip_attention_qkv": self.skip_attention_qkv,
                "use_lora": self.use_lora,
                "lora_r": self.lora_r,
                "method": "proxsparse_regularized",
            }
        )
        self.stats = stats
        return self.model, stats

ProxSparse: proximal gradient sparsity with structured / 2:4 / percentile constraint, optional LoRA fine-tuning, and safe fallback.

Ancestors

BaseCompressor

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class PrunedChannelLinear (in_features: int, original_out_features: int, kept_idx: torch.Tensor, W_kept: torch.Tensor, bias_kept: torch.Tensor | None = None)

Expand source code

class PrunedChannelLinear(nn.Module):
    """
    Linear layer with pruned output channels.

    Stores only the kept rows of the effective weight matrix:
      - W_kept:      (k, in_features)  fp16/fp32   -- only kept rows
      - bias_kept:   (k,)              optional
      - kept_idx:    (k,)              int64  -- output positions of kept rows
      - original_out_features: int

    Forward:
      h = F.linear(x, W_kept, bias_kept)          # (..., k)
      out = zeros(..., original_out_features)
      out.scatter_(-1, kept_idx, h)               # (..., original_out)
    """

    is_custom_attention = False

    def __init__(
        self,
        in_features: int,
        original_out_features: int,
        kept_idx: torch.Tensor,  # (k,) int64
        W_kept: torch.Tensor,  # (k, in_features)
        bias_kept: torch.Tensor | None = None,
    ) -> None:
        super().__init__()
        self.in_features = int(in_features)
        self.original_out_features = int(original_out_features)
        self.out_features = self.original_out_features  # for API compatibility
        self.num_kept = int(W_kept.shape[0])

        self.register_buffer("kept_idx", kept_idx.long(), persistent=True)
        self.W_kept = nn.Parameter(W_kept, requires_grad=False)

        if bias_kept is not None:
            self.bias_kept = nn.Parameter(bias_kept, requires_grad=False)
        else:
            self.register_parameter("bias_kept", None)

        self._cached_weight: torch.Tensor | None = None

    @property
    def weight(self) -> torch.Tensor:
        """Reconstruct effective (original_out, in) weight with zeros in pruned rows."""
        if self._cached_weight is not None:
            return self._cached_weight
        W = self.W_kept.new_zeros(self.original_out_features, self.in_features)
        W[self.kept_idx] = self.W_kept
        self._cached_weight = W
        return W

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = F.linear(x, self.W_kept, self.bias_kept)  # (..., k)
        out = torch.zeros(
            *x.shape[:-1],
            self.original_out_features,
            dtype=h.dtype,
            device=h.device,
        )
        out.scatter_(-1, self.kept_idx.expand(*x.shape[:-1], -1), h)
        return out

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, original_out_features={self.original_out_features}, "
            f"num_kept={self.num_kept}"
        )

Linear layer with pruned output channels.

Stores only the kept rows of the effective weight matrix: - W_kept: (k, in_features) fp16/fp32 – only kept rows - bias_kept: (k,) optional - kept_idx: (k,) int64 – output positions of kept rows - original_out_features: int

Forward -----= h = F.linear(x, W_kept, bias_kept) # (…, k) out = zeros(…, original_out_features) out.scatter_(-1, kept_idx, h) # (…, original_out)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention: The type of the None singleton.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstruct effective (original_out, in) weight with zeros in pruned rows."""
    if self._cached_weight is not None:
        return self._cached_weight
    W = self.W_kept.new_zeros(self.original_out_features, self.in_features)
    W[self.kept_idx] = self.W_kept
    self._cached_weight = W
    return W

Reconstruct effective (original_out, in) weight with zeros in pruned rows.

Methods

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    return (
        f"in_features={self.in_features}, original_out_features={self.original_out_features}, "
        f"num_kept={self.num_kept}"
    )

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    h = F.linear(x, self.W_kept, self.bias_kept)  # (..., k)
    out = torch.zeros(
        *x.shape[:-1],
        self.original_out_features,
        dtype=h.dtype,
        device=h.device,
    )
    out.scatter_(-1, self.kept_idx.expand(*x.shape[:-1], -1), h)
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class ROASTCompressor (model: nn.Module, n_buckets: int = 1024, tile_size: int = 32, skip_attention_qkv: bool = True, use_lora: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, target_modules: list[str] | None = None, target_attention_only: bool = False, max_error: float = 0.5, fallback_to_original: bool = True, min_numel: int = 16, device: str | torch.device | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("roast")
@COMPRESSOR_REGISTRY.register("roast_compressor")
@COMPRESSOR_REGISTRY.register("roast_hashed_linear")
class ROASTCompressor(BaseCompressor):
    """
    Replace eligible linear and Conv1D layers with ROASTLinear (hardware-aware weight sharing).
    """

    def __init__(
        self,
        model: nn.Module,
        n_buckets: int = 1024,
        tile_size: int = 32,
        skip_attention_qkv: bool = True,
        use_lora: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        target_modules: list[str] | None = None,
        target_attention_only: bool = False,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        min_numel: int = 16,
        device: str | torch.device | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="roast", **kwargs)
        cfg_obj = cfg or kwargs.get("comp_config")

        self.n_buckets = max(
            1,
            int(n_buckets if n_buckets is not None else _cfg_get(cfg_obj, "n_buckets", 1024)),
        )
        self.tile_size = max(
            1,
            int(tile_size if tile_size is not None else _cfg_get(cfg_obj, "tile_size", 32)),
        )
        self.target_attention_only = bool(
            target_attention_only
            or _cfg_get(
                cfg_obj,
                "target_attention_only",
                kwargs.get("target_attention_only", False),
            )
        )
        self.skip_attention_qkv = (
            bool(
                skip_attention_qkv
                if skip_attention_qkv is not None
                else _cfg_get(cfg_obj, "skip_attention_qkv", True)
            )
            and not self.target_attention_only
        )
        self.use_lora = bool(use_lora or _cfg_get(cfg_obj, "use_lora", False) or (lora_r > 0))
        self.lora_r = int(lora_r or _cfg_get(cfg_obj, "lora_r", 0))
        self.lora_alpha = float(lora_alpha or _cfg_get(cfg_obj, "lora_alpha", 1.0))
        self.target_modules = target_modules
        self.max_error = float(
            max_error if max_error is not None else _cfg_get(cfg_obj, "max_error", 0.50)
        )
        self.fallback_to_original = bool(
            fallback_to_original
            if fallback_to_original is not None
            else _cfg_get(cfg_obj, "fallback_to_original", True)
        )
        self.min_numel = max(
            1,
            int(min_numel if min_numel is not None else _cfg_get(cfg_obj, "min_numel", 16)),
        )

        # Target device derivation
        if device is not None:
            try:
                self.target_device = torch.device(device)
            except Exception:
                self.target_device = (
                    next(model.parameters()).device
                    if list(model.parameters())
                    else torch.device("cpu")
                )
        else:
            self.target_device = (
                next(model.parameters()).device if list(model.parameters()) else torch.device("cpu")
            )

        # Diagnostics & Tracking
        self.visited_modules: set[int] = set()
        self.compressed_layers: list[str] = []
        self.protected_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.failed_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        """Check if layer belongs to an attention projection in a GQA architecture."""
        if self.target_attention_only:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        is_gqa = num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads

        name_l = name.lower()
        is_attn_qkv = any(
            p in name_l
            for p in [
                "q_proj",
                "k_proj",
                "v_proj",
                "c_attn",
                "query",
                "key",
                "value",
                "q_lin",
                "k_lin",
                "v_lin",
                "self_attn.q",
                "self_attn.k",
                "self_attn.v",
            ]
        )
        return (is_gqa or self.skip_attention_qkv) and is_attn_qkv

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        if any(
            k in name_l
            for k in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embed",
                "embedding",
                "cls_head",
                "head",
                "logits_processor",
            ]
        ):
            self.protected_layers.append(name)
            return False

        if self.target_attention_only:
            attn_keywords = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
                "dense",
            ]
            if not any(k in name_l for k in attn_keywords):
                self.skipped_layers.append(name)
                return False

        if self._is_gqa_asymmetric(name, module):
            self.skipped_layers.append(name)
            return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__
        raw_w = getattr(base_mod, "weight", None)
        if raw_w is None or not isinstance(raw_w, torch.Tensor):
            return False

        is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_mod, "nf")
        is_linear_like = isinstance(base_mod, nn.Linear) or is_conv1d
        if not is_linear_like or isinstance(module, ROASTLinear):
            return False

        if raw_w.numel() < self.min_numel:
            self.skipped_layers.append(name)
            return False

        if not should_compress_linear(self.model, name):
            self.protected_layers.append(name)
            return False

        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """Compress the target modules using ROAST tiling with LoRA support."""
        self.compressed_layers.clear()
        self.visited_modules.clear()
        self.protected_layers.clear()
        self.skipped_layers.clear()
        self.fallback_layers.clear()
        self.failed_layers.clear()
        self.layer_errors.clear()

        for name, module in list(self.model.named_modules()):
            if self._is_target_module(name, module):
                mod_id = id(module)
                if mod_id in self.visited_modules:
                    logger.debug("ROAST: Skipping tied module %s", name)
                    self.skipped_layers.append(name)
                    continue
                self.visited_modules.add(mod_id)

                try:
                    base_mod = getattr(module, "base_layer", module)
                    raw_w = base_mod.weight.data
                    cls_name = base_mod.__class__.__name__
                    is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_mod, "nf")
                    W_orig = (
                        raw_w.t().contiguous().float() if is_conv1d else raw_w.contiguous().float()
                    )

                    new_layer = ROASTLinear.from_linear(
                        module,
                        n_buckets=self.n_buckets,
                        tile_size=self.tile_size,
                        lora_r=self.lora_r if self.use_lora else 0,
                        lora_alpha=self.lora_alpha,
                    )
                    W_rec = new_layer.weight.float()
                    norm_orig = torch.norm(W_orig, p="fro").clamp_min(1e-8)
                    err = (torch.norm(W_orig - W_rec, p="fro") / norm_orig).item()
                    self.layer_errors[name] = float(err)

                    is_invalid = torch.isnan(W_rec).any() or torch.isinf(W_rec).any()
                    if err > self.max_error:
                        logger.warning(
                            "ROAST: reconstruction error %.4f on %s exceeded threshold %.4f.",
                            err,
                            name,
                            self.max_error,
                        )
                        is_invalid = True

                    if is_invalid and self.fallback_to_original:
                        self.fallback_layers.append(name)
                        for p in module.parameters():
                            p.requires_grad = False
                        continue

                    replace_module(self.model, name, new_layer)
                    self.compressed_layers.append(name)

                except Exception as e_roast:
                    logger.warning("ROAST: Failed to process %s: %s", name, e_roast)
                    self.failed_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False

        # Post-compression forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                dummy_input = torch.randint(0, 100, (1, 8), device=self.target_device)
                self.model(dummy_input)
        except Exception as ex_sanity:
            logger.warning("ROAST post-compression forward sanity check warning: %s", ex_sanity)
            forward_check_success = False

        self._compressed_params = self._count_params(self.model)
        avg_err = (
            (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
        )
        stats = self.get_compression_stats()
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "protected_layers": self.protected_layers,
                "protected_layer_count": len(self.protected_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "failed_layers": self.failed_layers,
                "failed_layer_count": len(self.failed_layers),
                "layer_errors": self.layer_errors,
                "avg_reconstruction_error": avg_err,
                "forward_check_success": forward_check_success,
                "target_attention_only": self.target_attention_only,
                "n_buckets": self.n_buckets,
                "tile_size": self.tile_size,
                "skip_attention_qkv": self.skip_attention_qkv,
                "use_lora": self.use_lora,
                "lora_r": self.lora_r,
                "method": "roast",
            }
        )
        self.stats = stats
        return self.model, stats

Replace eligible linear and Conv1D layers with ROASTLinear (hardware-aware weight sharing).

Ancestors

BaseCompressor

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """Compress the target modules using ROAST tiling with LoRA support."""
    self.compressed_layers.clear()
    self.visited_modules.clear()
    self.protected_layers.clear()
    self.skipped_layers.clear()
    self.fallback_layers.clear()
    self.failed_layers.clear()
    self.layer_errors.clear()

    for name, module in list(self.model.named_modules()):
        if self._is_target_module(name, module):
            mod_id = id(module)
            if mod_id in self.visited_modules:
                logger.debug("ROAST: Skipping tied module %s", name)
                self.skipped_layers.append(name)
                continue
            self.visited_modules.add(mod_id)

            try:
                base_mod = getattr(module, "base_layer", module)
                raw_w = base_mod.weight.data
                cls_name = base_mod.__class__.__name__
                is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_mod, "nf")
                W_orig = (
                    raw_w.t().contiguous().float() if is_conv1d else raw_w.contiguous().float()
                )

                new_layer = ROASTLinear.from_linear(
                    module,
                    n_buckets=self.n_buckets,
                    tile_size=self.tile_size,
                    lora_r=self.lora_r if self.use_lora else 0,
                    lora_alpha=self.lora_alpha,
                )
                W_rec = new_layer.weight.float()
                norm_orig = torch.norm(W_orig, p="fro").clamp_min(1e-8)
                err = (torch.norm(W_orig - W_rec, p="fro") / norm_orig).item()
                self.layer_errors[name] = float(err)

                is_invalid = torch.isnan(W_rec).any() or torch.isinf(W_rec).any()
                if err > self.max_error:
                    logger.warning(
                        "ROAST: reconstruction error %.4f on %s exceeded threshold %.4f.",
                        err,
                        name,
                        self.max_error,
                    )
                    is_invalid = True

                if is_invalid and self.fallback_to_original:
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

                replace_module(self.model, name, new_layer)
                self.compressed_layers.append(name)

            except Exception as e_roast:
                logger.warning("ROAST: Failed to process %s: %s", name, e_roast)
                self.failed_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False

    # Post-compression forward sanity check
    forward_check_success = True
    try:
        with torch.no_grad():
            dummy_input = torch.randint(0, 100, (1, 8), device=self.target_device)
            self.model(dummy_input)
    except Exception as ex_sanity:
        logger.warning("ROAST post-compression forward sanity check warning: %s", ex_sanity)
        forward_check_success = False

    self._compressed_params = self._count_params(self.model)
    avg_err = (
        (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
    )
    stats = self.get_compression_stats()
    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "protected_layers": self.protected_layers,
            "protected_layer_count": len(self.protected_layers),
            "skipped_layers": self.skipped_layers,
            "skipped_layer_count": len(self.skipped_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "failed_layers": self.failed_layers,
            "failed_layer_count": len(self.failed_layers),
            "layer_errors": self.layer_errors,
            "avg_reconstruction_error": avg_err,
            "forward_check_success": forward_check_success,
            "target_attention_only": self.target_attention_only,
            "n_buckets": self.n_buckets,
            "tile_size": self.tile_size,
            "skip_attention_qkv": self.skip_attention_qkv,
            "use_lora": self.use_lora,
            "lora_r": self.lora_r,
            "method": "roast",
        }
    )
    self.stats = stats
    return self.model, stats

Compress the target modules using ROAST tiling with LoRA support.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class ROASTLinear (in_features: int, out_features: int, n_buckets: int = 1024, tile_size: int = 32, bias: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, device: torch.device | None = None, dtype: torch.dtype | None = None)

Expand source code

class ROASTLinear(nn.Module):
    """
    Linear layer with weight sharing based on tiles (blocks) rather than
    individual scalar weights. Each tile of size (tile_size, tile_size) shares
    a single weight from a small table, with optional trainable LoRA adapters.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        n_buckets: int = 1024,
        tile_size: int = 32,
        bias: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.tile_size = max(1, int(tile_size))
        self.n_buckets = max(1, int(n_buckets))
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        # Number of tiles along each dimension (padded)
        self.tiles_out = (out_features + self.tile_size - 1) // self.tile_size
        self.tiles_in = (in_features + self.tile_size - 1) // self.tile_size

        # Weight table: frozen if LoRA is active
        requires_grad = not (self.lora_r > 0)
        self.weight_table = nn.Parameter(
            torch.randn(self.n_buckets, device=device, dtype=dtype),
            requires_grad=requires_grad,
        )

        if bias:
            self.bias = nn.Parameter(
                torch.zeros(out_features, device=device, dtype=dtype),
                requires_grad=requires_grad,
            )
        else:
            self.register_parameter("bias", None)

        # Precompute tile indices with high-entropy bitwise mixing (B11, J1)
        tile_idx_i = (
            torch.arange(self.tiles_out, device=device).unsqueeze(1).expand(-1, self.tiles_in)
        )
        tile_idx_j = (
            torch.arange(self.tiles_in, device=device).unsqueeze(0).expand(self.tiles_out, -1)
        )
        h = (tile_idx_i.long() * 0x9E3779B1) ^ (tile_idx_j.long() * 0x85EBCA77)
        h = h ^ (h >> 16)
        h = h * 0x7FEB352D
        h = (h ^ (h >> 15)) & 0x7FFFFFFF
        self.register_buffer(
            "tile_hash_idx",
            (h % self.n_buckets).long(),
            persistent=False,
        )

        self._cached_weight: torch.Tensor | None = None
        self._cache_stale: bool = True

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty((self.lora_r, in_features), dtype=torch.float32, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.zeros((out_features, self.lora_r), dtype=torch.float32, device=device)
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    def _invalidate_cache(self) -> None:
        self._cached_weight = None
        self._cache_stale = True

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix (out_features, in_features) with evaluation caching."""
        if self._cached_weight is not None and not self._cache_stale and not self.training:
            return self._cached_weight
        W_tiles = self.weight_table[self.tile_hash_idx]  # (tiles_out, tiles_in)
        W_full = torch.repeat_interleave(
            torch.repeat_interleave(W_tiles, self.tile_size, dim=0),
            self.tile_size,
            dim=1,
        )
        W_res = W_full[: self.out_features, : self.in_features]
        if not self.training:
            self._cached_weight = W_res
            self._cache_stale = False
        return W_res

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Vectorized forward pass without per-element python loops."""
        W = self.weight.to(device=x.device, dtype=x.dtype)
        out = F.linear(x, W, self.bias.to(dtype=x.dtype) if self.bias is not None else None)
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x.float(), self.lora_A), self.lora_B) * self.scaling
            out = out + lora_out.to(out.dtype)
        return out

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        n_buckets: int = 1024,
        tile_size: int = 32,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> ROASTLinear:
        """Initialize from a pretrained nn.Linear, Conv1D, or LoRALinear base_layer."""
        base_mod = getattr(linear, "base_layer", linear)
        raw_w = base_mod.weight.data
        cls_name = base_mod.__class__.__name__
        is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_mod, "nf")

        if is_conv1d:
            W = raw_w.t().contiguous().float()
            out_f, in_f = W.shape
        else:
            W = raw_w.contiguous().float()
            out_f, in_f = W.shape

        bias = getattr(base_mod, "bias", None) is not None
        roast = cls(
            in_features=in_f,
            out_features=out_f,
            n_buckets=n_buckets,
            tile_size=tile_size,
            bias=bias,
            lora_r=lora_r,
            lora_alpha=lora_alpha,
            device=raw_w.device,
            dtype=raw_w.dtype,
        )
        if bias and base_mod.bias is not None:
            roast.bias.data.copy_(base_mod.bias.data)

        # Preserve existing source LoRA weights if present
        if (
            hasattr(linear, "lora_A")
            and hasattr(linear, "lora_B")
            and linear.lora_A is not None
            and linear.lora_B is not None
        ):
            src_A = getattr(linear.lora_A, "weight", linear.lora_A)
            src_B = getattr(linear.lora_B, "weight", linear.lora_B)
            if roast.lora_A is not None and roast.lora_B is not None:
                if src_A.shape == roast.lora_A.shape:
                    roast.lora_A.data.copy_(src_A.data)
                if src_B.shape == roast.lora_B.shape:
                    roast.lora_B.data.copy_(src_B.data)

        ts = max(1, tile_size)
        tiles_out = (out_f + ts - 1) // ts
        tiles_in = (in_f + ts - 1) // ts

        # Vectorized computation of per-tile means and scatter into weight_table
        with torch.no_grad():
            pad_out = tiles_out * ts - out_f
            pad_in = tiles_in * ts - in_f
            if pad_out > 0 or pad_in > 0:
                W_padded = F.pad(W, (0, pad_in, 0, pad_out), mode="constant", value=0.0)
                mask = F.pad(
                    torch.ones_like(W),
                    (0, pad_in, 0, pad_out),
                    mode="constant",
                    value=0.0,
                )
                W_blocks = (
                    W_padded.view(tiles_out, ts, tiles_in, ts)
                    .permute(0, 2, 1, 3)
                    .contiguous()
                    .view(tiles_out, tiles_in, -1)
                )
                mask_blocks = (
                    mask.view(tiles_out, ts, tiles_in, ts)
                    .permute(0, 2, 1, 3)
                    .contiguous()
                    .view(tiles_out, tiles_in, -1)
                )
                tile_means = W_blocks.sum(dim=-1) / mask_blocks.sum(dim=-1).clamp(min=1.0)
            else:
                tile_means = (
                    W.view(tiles_out, ts, tiles_in, ts)
                    .permute(0, 2, 1, 3)
                    .contiguous()
                    .view(tiles_out, tiles_in, -1)
                    .mean(dim=-1)
                )

            flat_idx = roast.tile_hash_idx.view(-1).to(device=roast.weight_table.device)
            flat_means = tile_means.view(-1).to(
                device=roast.weight_table.device, dtype=roast.weight_table.dtype
            )
            roast.weight_table.data.scatter_(0, flat_idx, flat_means)

        return roast

Linear layer with weight sharing based on tiles (blocks) rather than individual scalar weights. Each tile of size (tile_size, tile_size) shares a single weight from a small table, with optional trainable LoRA adapters.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Module, n_buckets: int = 1024, tile_size: int = 32, lora_r: int = 0, lora_alpha: float = 1.0) ‑> ROASTLinear: Initialize from a pretrained nn.Linear, Conv1D, or LoRALinear base_layer.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix (out_features, in_features) with evaluation caching."""
    if self._cached_weight is not None and not self._cache_stale and not self.training:
        return self._cached_weight
    W_tiles = self.weight_table[self.tile_hash_idx]  # (tiles_out, tiles_in)
    W_full = torch.repeat_interleave(
        torch.repeat_interleave(W_tiles, self.tile_size, dim=0),
        self.tile_size,
        dim=1,
    )
    W_res = W_full[: self.out_features, : self.in_features]
    if not self.training:
        self._cached_weight = W_res
        self._cache_stale = False
    return W_res

Reconstructed effective weight matrix (out_features, in_features) with evaluation caching.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Vectorized forward pass without per-element python loops."""
    W = self.weight.to(device=x.device, dtype=x.dtype)
    out = F.linear(x, W, self.bias.to(dtype=x.dtype) if self.bias is not None else None)
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x.float(), self.lora_A), self.lora_B) * self.scaling
        out = out + lora_out.to(out.dtype)
    return out

Vectorized forward pass without per-element python loops.

class RPCACompressor (model: nn.Module, rank: int = 64, mu: float = 0.001, max_iter: int = 5, min_rank: int = 8, skip_attention_qkv: bool = False, target_attention_only: bool = False, use_lora: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, target_modules: list[str] | None = None, max_error: float = 0.5, fallback_to_original: bool = True, min_numel: int = 16, device: str | torch.device | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("rpca")
@COMPRESSOR_REGISTRY.register("rpca_compressor")
class RPCACompressor(BaseCompressor):
    is_custom_attention: bool = True
    """
    Robust Principal Component Analysis (RPCA) Low-Rank Compressor.

    Decomposes weight matrices W into low-rank component L and sparse outlier matrix S
    via Principal Component Pursuit / alternating soft-thresholding:
        min ||L||_* + lambda ||S||_1  s.t.  W = L + S
    Factorizes the low-rank component L into symmetric factors (U * sqrt(Sigma), sqrt(Sigma) * V^T)
    and replaces nn.Linear or Conv1D with factorized SVDLinear modules with trainable LoRA fine-tuning adapters.
    """

    def __init__(
        self,
        model: nn.Module,
        rank: int = 64,
        mu: float = 1e-3,
        max_iter: int = 5,
        min_rank: int = 8,
        skip_attention_qkv: bool = False,
        target_attention_only: bool = False,
        use_lora: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        target_modules: list[str] | None = None,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        min_numel: int = 16,
        device: str | torch.device | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="rpca", **kwargs)
        cfg_obj = cfg or kwargs.get("comp_config")

        self.rank = int(rank or _cfg_get(cfg_obj, "rpca_rank", _cfg_get(cfg_obj, "rank", 64)))
        self.mu = float(mu or _cfg_get(cfg_obj, "rpca_mu", _cfg_get(cfg_obj, "mu", 1e-3)))
        self.max_iter = int(
            max_iter or _cfg_get(cfg_obj, "rpca_max_iter", _cfg_get(cfg_obj, "max_iter", 5))
        )
        self.min_rank = max(1, int(min_rank or _cfg_get(cfg_obj, "min_rank", 8)))
        self.skip_attention_qkv = bool(
            skip_attention_qkv
            if skip_attention_qkv is not None
            else _cfg_get(cfg_obj, "skip_attention_qkv", False)
        )
        self.target_attention_only = bool(
            target_attention_only
            if target_attention_only is not None
            else _cfg_get(cfg_obj, "target_attention_only", False)
        )
        self.use_lora = bool(use_lora or _cfg_get(cfg_obj, "use_lora", False) or (lora_r > 0))
        self.lora_r = int(lora_r or _cfg_get(cfg_obj, "lora_r", 0))
        self.lora_alpha = float(lora_alpha or _cfg_get(cfg_obj, "lora_alpha", 1.0))
        self.target_modules = target_modules
        self.max_error = float(max_error or _cfg_get(cfg_obj, "max_error", 0.50))
        self.fallback_to_original = bool(
            fallback_to_original
            if fallback_to_original is not None
            else _cfg_get(cfg_obj, "fallback_to_original", True)
        )
        self.min_numel = max(1, int(min_numel or _cfg_get(cfg_obj, "min_numel", 16)))

        # Target device derivation
        if device is not None:
            try:
                self.target_device = torch.device(device)
            except Exception:
                self.target_device = (
                    next(model.parameters()).device
                    if list(model.parameters())
                    else torch.device("cpu")
                )
        else:
            self.target_device = (
                next(model.parameters()).device if list(model.parameters()) else torch.device("cpu")
            )

        # Diagnostics & Tracking
        self.visited_modules: set[int] = set()
        self.compressed_layers: list[str] = []
        self.protected_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.failed_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.replaced_layers: list[dict[str, Any]] = []

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        """Check if layer belongs to an attention projection in a GQA architecture."""
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        is_gqa = num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads

        name_l = name.lower()
        is_attn_qkv = any(
            p in name_l
            for p in [
                "q_proj",
                "k_proj",
                "v_proj",
                "c_attn",
                "query",
                "key",
                "value",
                "q_lin",
                "k_lin",
                "v_lin",
                "self_attn.q",
                "self_attn.k",
                "self_attn.v",
            ]
        )
        return (is_gqa or self.skip_attention_qkv) and is_attn_qkv

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        # 1. Protection for critical / classification / head / pooler / embedding layers
        if any(
            k in name_l
            for k in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embed",
                "embedding",
                "cls_head",
                "head",
                "logits_processor",
            ]
        ):
            self.protected_layers.append(name)
            return False

        # Target attention only filter
        if self.target_attention_only:
            attn_keywords = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(k in name_l for k in attn_keywords):
                return False

        # 2. GQA-aware skipping
        if self._is_gqa_asymmetric(name, module):
            self.skipped_layers.append(name)
            return False

        # 3. Check linear or Conv1D nature
        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__
        raw_w = getattr(base_mod, "weight", None)
        if raw_w is None or not isinstance(raw_w, torch.Tensor):
            return False

        is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_mod, "nf")
        is_linear_like = isinstance(base_mod, nn.Linear) or is_conv1d
        if not is_linear_like or isinstance(base_mod, SVDLinear):
            return False

        # 4. Tiny layer guard
        if raw_w.numel() < self.min_numel:
            self.skipped_layers.append(name)
            return False

        # 5. Respect framework compatibility rules
        if not should_compress_linear(self.model, name):
            self.protected_layers.append(name)
            return False

        # 6. User-defined explicit target modules
        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def compress_linear(
        self,
        layer: nn.Module,
        rank: int,
    ) -> SVDLinear | None:
        """Decompose a single linear layer via RPCA and return factorized SVDLinear with LoRA."""
        base_layer = getattr(layer, "base_layer", layer)
        raw_w = base_layer.weight.data
        cls_name = base_layer.__class__.__name__
        is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_layer, "nf")

        if is_conv1d:
            W = raw_w.t().contiguous().float()
            out_f, in_f = W.shape
        else:
            W = raw_w.contiguous().float()
            out_f = getattr(base_layer, "out_features", W.shape[0])
            in_f = getattr(base_layer, "in_features", W.shape[1])

        max_r = min(out_f, in_f)
        param_cap = max(1, (in_f * out_f) // (in_f + out_f) - 1)
        r = min(rank, max_r - 1, param_cap)
        r = max(min(self.min_rank, max_r - 1), r)
        if r <= 0 or r >= max_r:
            logger.debug(
                "RPCA: rank %d >= max_r %d; skipping layer for no parameter reduction",
                r,
                max_r,
            )
            return None

        device = base_layer.weight.device
        dtype = base_layer.weight.dtype

        L = W.clone()
        S_sparse = torch.zeros_like(W)

        # Inexact ALM / PCP Iterations in FP32 with try/except protection
        try:
            for _ in range(self.max_iter):
                # Low-rank projection via truncated SVD
                U, S_vals, Vt = torch.linalg.svd(L, full_matrices=False)
                sqrt_S = torch.sqrt(torch.clamp(S_vals[:r], min=1e-12))
                U_r = U[:, :r] * sqrt_S.unsqueeze(0)
                V_r = sqrt_S.unsqueeze(1) * Vt[:r, :]
                L_recon = U_r @ V_r

                # Sparse soft-thresholding update on residual
                resid = W - L_recon
                S_sparse = torch.sign(resid) * torch.clamp(torch.abs(resid) - self.mu, min=0.0)

                # Update low-rank target
                L = W - S_sparse

            # Final factorization of L
            U, S_vals, Vt = torch.linalg.svd(L, full_matrices=False)
        except Exception as e_svd:
            logger.warning(
                "RPCA decomposition failed: %s; falling back to standard SVD on W.",
                e_svd,
            )
            try:
                U, S_vals, Vt = torch.linalg.svd(W, full_matrices=False)
            except Exception as e_fb:
                logger.warning("RPCA standard SVD fallback failed: %s", e_fb)
                return None

        sqrt_S = torch.sqrt(torch.clamp(S_vals[:r], min=1e-12))
        U_final = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
        V_final = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(dtype=dtype, device=device)

        # LoRA rank selection
        has_existing_lora = hasattr(layer, "lora_A") and hasattr(layer, "lora_B")
        effective_lora_r = self.lora_r
        if has_existing_lora and getattr(layer, "r", 0) > 0:
            effective_lora_r = getattr(layer, "r", self.lora_r)

        svd_layer = SVDLinear(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=base_layer.bias is not None,
            dtype=dtype,
            device=device,
            lora_r=effective_lora_r if (self.use_lora or has_existing_lora) else 0,
            lora_alpha=self.lora_alpha,
        )
        svd_layer.set_low_rank_factors(U_final, V_final)
        if base_layer.bias is not None:
            svd_layer.bias.data.copy_(base_layer.bias.data)

        # Preserve existing LoRA adapters
        if has_existing_lora and svd_layer.use_lora:
            try:
                if hasattr(layer.lora_A, "weight"):
                    svd_layer.lora_A.weight.data.copy_(layer.lora_A.weight.data)
                if hasattr(layer.lora_B, "weight"):
                    svd_layer.lora_B.weight.data.copy_(layer.lora_B.weight.data)
            except Exception as e_lora:
                logger.warning("Could not copy existing LoRA adapters: %s", e_lora)

        # Frobenius reconstruction error check
        W_rec = torch.matmul(U_final.float(), V_final.float())
        norm_orig = torch.norm(W, p="fro").clamp_min(1e-8)
        rel_err = (torch.norm(W - W_rec, p="fro") / norm_orig).item()

        is_invalid = torch.isnan(W_rec).any() or torch.isinf(W_rec).any()
        if rel_err > self.max_error:
            logger.warning(
                "RPCA: reconstruction error %.4f exceeded threshold %.4f.",
                rel_err,
                self.max_error,
            )
            is_invalid = True

        if is_invalid and self.fallback_to_original:
            return None

        svd_layer.reconstruction_error = rel_err
        return svd_layer

    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """Compress all linear layers in the model using RPCA decomposition."""
        self.layer_errors = {}
        self.replaced_layers = []
        self.visited_modules.clear()
        self.compressed_layers.clear()
        self.protected_layers.clear()
        self.skipped_layers.clear()
        self.fallback_layers.clear()
        self.failed_layers.clear()

        target_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            logger.warning("RPCA: No compressible linear layers found.")
            return self.model, self.get_compression_stats()

        for name, module in target_layers:
            mod_id = id(module)
            if mod_id in self.visited_modules:
                logger.debug("RPCA: Skipping tied module %s", name)
                self.skipped_layers.append(name)
                continue
            self.visited_modules.add(mod_id)

            try:
                svd_layer = self.compress_linear(module, rank=self.rank)
                if svd_layer is not None:
                    err = getattr(svd_layer, "reconstruction_error", 0.0)
                    self.layer_errors[name] = err
                    self.replaced_layers.append(
                        {"name": name, "rank": svd_layer.rank, "error": err}
                    )
                    replace_module(self.model, name, svd_layer)
                    self.compressed_layers.append(name)
                else:
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
            except Exception as e_comp:
                logger.warning("Failed to compress %s: %s", name, e_comp)
                self.failed_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False

        # Post-compression forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                dummy_input = torch.randint(0, 100, (1, 8), device=self.target_device)
                self.model(dummy_input)
        except Exception as ex_sanity:
            logger.warning("RPCA post-compression forward sanity check warning: %s", ex_sanity)
            forward_check_success = False

        self._compressed_params = self._count_params(self.model)
        avg_err = (
            (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
        )
        stats = self.get_compression_stats()
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "protected_layers": self.protected_layers,
                "protected_layer_count": len(self.protected_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "failed_layers": self.failed_layers,
                "failed_layer_count": len(self.failed_layers),
                "layer_errors": self.layer_errors,
                "avg_reconstruction_error": avg_err,
                "forward_check_success": forward_check_success,
                "rank": self.rank,
                "min_rank": self.min_rank,
                "mu": self.mu,
                "max_iter": self.max_iter,
                "skip_attention_qkv": self.skip_attention_qkv,
                "target_attention_only": self.target_attention_only,
                "use_lora": self.use_lora,
                "lora_r": self.lora_r,
                "method": "rpca",
            }
        )
        self.stats = stats
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: Robust Principal Component Analysis (RPCA) Low-Rank Compressor.

Decomposes weight matrices W into low-rank component L and sparse outlier matrix S via Principal Component Pursuit / alternating soft-thresholding: min ||L||_* + lambda ||S||_1 s.t. W = L + S Factorizes the low-rank component L into symmetric factors (U * sqrt(Sigma), sqrt(Sigma) * V^T) and replaces nn.Linear or Conv1D with factorized SVDLinear modules with trainable LoRA fine-tuning adapters.

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """Compress all linear layers in the model using RPCA decomposition."""
    self.layer_errors = {}
    self.replaced_layers = []
    self.visited_modules.clear()
    self.compressed_layers.clear()
    self.protected_layers.clear()
    self.skipped_layers.clear()
    self.fallback_layers.clear()
    self.failed_layers.clear()

    target_layers = [
        (name, mod)
        for name, mod in self.model.named_modules()
        if self._is_target_module(name, mod)
    ]

    if not target_layers:
        logger.warning("RPCA: No compressible linear layers found.")
        return self.model, self.get_compression_stats()

    for name, module in target_layers:
        mod_id = id(module)
        if mod_id in self.visited_modules:
            logger.debug("RPCA: Skipping tied module %s", name)
            self.skipped_layers.append(name)
            continue
        self.visited_modules.add(mod_id)

        try:
            svd_layer = self.compress_linear(module, rank=self.rank)
            if svd_layer is not None:
                err = getattr(svd_layer, "reconstruction_error", 0.0)
                self.layer_errors[name] = err
                self.replaced_layers.append(
                    {"name": name, "rank": svd_layer.rank, "error": err}
                )
                replace_module(self.model, name, svd_layer)
                self.compressed_layers.append(name)
            else:
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
        except Exception as e_comp:
            logger.warning("Failed to compress %s: %s", name, e_comp)
            self.failed_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False

    # Post-compression forward sanity check
    forward_check_success = True
    try:
        with torch.no_grad():
            dummy_input = torch.randint(0, 100, (1, 8), device=self.target_device)
            self.model(dummy_input)
    except Exception as ex_sanity:
        logger.warning("RPCA post-compression forward sanity check warning: %s", ex_sanity)
        forward_check_success = False

    self._compressed_params = self._count_params(self.model)
    avg_err = (
        (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
    )
    stats = self.get_compression_stats()
    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "protected_layers": self.protected_layers,
            "protected_layer_count": len(self.protected_layers),
            "skipped_layers": self.skipped_layers,
            "skipped_layer_count": len(self.skipped_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "failed_layers": self.failed_layers,
            "failed_layer_count": len(self.failed_layers),
            "layer_errors": self.layer_errors,
            "avg_reconstruction_error": avg_err,
            "forward_check_success": forward_check_success,
            "rank": self.rank,
            "min_rank": self.min_rank,
            "mu": self.mu,
            "max_iter": self.max_iter,
            "skip_attention_qkv": self.skip_attention_qkv,
            "target_attention_only": self.target_attention_only,
            "use_lora": self.use_lora,
            "lora_r": self.lora_r,
            "method": "rpca",
        }
    )
    self.stats = stats
    return self.model, stats

Compress all linear layers in the model using RPCA decomposition.

def compress_linear(self, layer: nn.Module, rank: int) ‑> SVDLinear | None

Expand source code

def compress_linear(
    self,
    layer: nn.Module,
    rank: int,
) -> SVDLinear | None:
    """Decompose a single linear layer via RPCA and return factorized SVDLinear with LoRA."""
    base_layer = getattr(layer, "base_layer", layer)
    raw_w = base_layer.weight.data
    cls_name = base_layer.__class__.__name__
    is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_layer, "nf")

    if is_conv1d:
        W = raw_w.t().contiguous().float()
        out_f, in_f = W.shape
    else:
        W = raw_w.contiguous().float()
        out_f = getattr(base_layer, "out_features", W.shape[0])
        in_f = getattr(base_layer, "in_features", W.shape[1])

    max_r = min(out_f, in_f)
    param_cap = max(1, (in_f * out_f) // (in_f + out_f) - 1)
    r = min(rank, max_r - 1, param_cap)
    r = max(min(self.min_rank, max_r - 1), r)
    if r <= 0 or r >= max_r:
        logger.debug(
            "RPCA: rank %d >= max_r %d; skipping layer for no parameter reduction",
            r,
            max_r,
        )
        return None

    device = base_layer.weight.device
    dtype = base_layer.weight.dtype

    L = W.clone()
    S_sparse = torch.zeros_like(W)

    # Inexact ALM / PCP Iterations in FP32 with try/except protection
    try:
        for _ in range(self.max_iter):
            # Low-rank projection via truncated SVD
            U, S_vals, Vt = torch.linalg.svd(L, full_matrices=False)
            sqrt_S = torch.sqrt(torch.clamp(S_vals[:r], min=1e-12))
            U_r = U[:, :r] * sqrt_S.unsqueeze(0)
            V_r = sqrt_S.unsqueeze(1) * Vt[:r, :]
            L_recon = U_r @ V_r

            # Sparse soft-thresholding update on residual
            resid = W - L_recon
            S_sparse = torch.sign(resid) * torch.clamp(torch.abs(resid) - self.mu, min=0.0)

            # Update low-rank target
            L = W - S_sparse

        # Final factorization of L
        U, S_vals, Vt = torch.linalg.svd(L, full_matrices=False)
    except Exception as e_svd:
        logger.warning(
            "RPCA decomposition failed: %s; falling back to standard SVD on W.",
            e_svd,
        )
        try:
            U, S_vals, Vt = torch.linalg.svd(W, full_matrices=False)
        except Exception as e_fb:
            logger.warning("RPCA standard SVD fallback failed: %s", e_fb)
            return None

    sqrt_S = torch.sqrt(torch.clamp(S_vals[:r], min=1e-12))
    U_final = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
    V_final = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(dtype=dtype, device=device)

    # LoRA rank selection
    has_existing_lora = hasattr(layer, "lora_A") and hasattr(layer, "lora_B")
    effective_lora_r = self.lora_r
    if has_existing_lora and getattr(layer, "r", 0) > 0:
        effective_lora_r = getattr(layer, "r", self.lora_r)

    svd_layer = SVDLinear(
        in_features=in_f,
        out_features=out_f,
        rank=r,
        bias=base_layer.bias is not None,
        dtype=dtype,
        device=device,
        lora_r=effective_lora_r if (self.use_lora or has_existing_lora) else 0,
        lora_alpha=self.lora_alpha,
    )
    svd_layer.set_low_rank_factors(U_final, V_final)
    if base_layer.bias is not None:
        svd_layer.bias.data.copy_(base_layer.bias.data)

    # Preserve existing LoRA adapters
    if has_existing_lora and svd_layer.use_lora:
        try:
            if hasattr(layer.lora_A, "weight"):
                svd_layer.lora_A.weight.data.copy_(layer.lora_A.weight.data)
            if hasattr(layer.lora_B, "weight"):
                svd_layer.lora_B.weight.data.copy_(layer.lora_B.weight.data)
        except Exception as e_lora:
            logger.warning("Could not copy existing LoRA adapters: %s", e_lora)

    # Frobenius reconstruction error check
    W_rec = torch.matmul(U_final.float(), V_final.float())
    norm_orig = torch.norm(W, p="fro").clamp_min(1e-8)
    rel_err = (torch.norm(W - W_rec, p="fro") / norm_orig).item()

    is_invalid = torch.isnan(W_rec).any() or torch.isinf(W_rec).any()
    if rel_err > self.max_error:
        logger.warning(
            "RPCA: reconstruction error %.4f exceeded threshold %.4f.",
            rel_err,
            self.max_error,
        )
        is_invalid = True

    if is_invalid and self.fallback_to_original:
        return None

    svd_layer.reconstruction_error = rel_err
    return svd_layer

Decompose a single linear layer via RPCA and return factorized SVDLinear with LoRA.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class ReplaceMeCompressor (model: nn.Module, remove_layers: int = 2, use_lora: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, fallback_to_original: bool = True, device: str | torch.device | None = None, config: Any | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("replaceme")
@COMPRESSOR_REGISTRY.register("replaceme_compressor")
class ReplaceMeCompressor(BaseCompressor):
    """
    Replaces the trailing N layers of a Transformer with a lightweight
    LayerNorm + Linear residual block.
    """

    def __init__(
        self,
        model: nn.Module,
        remove_layers: int = 2,
        use_lora: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        fallback_to_original: bool = True,
        device: str | torch.device | None = None,
        config: Any | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="replaceme", **kwargs)
        cfg_obj = cfg or config or kwargs.get("comp_config")

        self.remove_layers = max(
            1,
            int(
                remove_layers
                if remove_layers is not None
                else _cfg_get(cfg_obj, "remove_layers", 2)
            ),
        )
        self.use_lora = bool(use_lora or _cfg_get(cfg_obj, "use_lora", False) or (lora_r > 0))
        self.lora_r = int(lora_r or _cfg_get(cfg_obj, "lora_r", 0))
        self.lora_alpha = float(lora_alpha or _cfg_get(cfg_obj, "lora_alpha", 1.0))
        self.fallback_to_original = bool(
            fallback_to_original
            if fallback_to_original is not None
            else _cfg_get(cfg_obj, "fallback_to_original", True)
        )

        # Device resolution
        if device is not None:
            try:
                self.target_device = torch.device(device)
            except Exception:
                self.target_device = (
                    next(model.parameters()).device
                    if list(model.parameters())
                    else torch.device("cpu")
                )
        else:
            self.target_device = (
                next(model.parameters()).device if list(model.parameters()) else torch.device("cpu")
            )

    # ------------------------------------------------------------------
    def _find_transformer_layers(self) -> nn.ModuleList | None:
        """Locate the transformer layer list across common architectures."""
        # Ordered list of attribute paths (most specific first).
        candidates = [
            ("distilbert", "transformer", "layer"),
            ("roberta", "encoder", "layer"),
            ("bert", "encoder", "layer"),
            ("encoder", "layer"),
            ("transformer", "h"),  # GPT-2
            ("transformer", "layer"),
            ("model", "layers"),  # LLaMA / Qwen
        ]
        for path in candidates:
            curr = self.model
            ok = True
            for attr in path:
                if hasattr(curr, attr):
                    curr = getattr(curr, attr)
                else:
                    ok = False
                    break
            if ok and isinstance(curr, (nn.ModuleList, list)) and len(curr) > 0:
                return curr

        # Fallback: scan for the first long ModuleList whose elements look like
        # transformer layers.
        for _, module in self.model.named_modules():
            if isinstance(module, nn.ModuleList) and len(module) >= 2:
                first_cls = module[0].__class__.__name__.lower()
                if any(k in first_cls for k in ("layer", "block", "decoderlayer", "encoderlayer")):
                    return module
        return None

    # ------------------------------------------------------------------
    def _infer_hidden_size(self, layers: nn.ModuleList) -> int:
        # Try config first.
        cfg = getattr(self.model, "config", None)
        if cfg is not None:
            for key in ("hidden_size", "d_model", "dim", "n_embd"):
                v = _cfg_get(cfg, key, None)
                if isinstance(v, int) and v > 0:
                    return int(v)
        # Fallback: inspect first layer's parameters.
        try:
            first_layer = layers[0]
            for p in first_layer.parameters():
                if p.dim() == 2:
                    return int(p.shape[-1])
        except Exception:
            pass
        return 768  # last-resort default

    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        layers = self._find_transformer_layers()
        layers_replaced = 0
        fallback_layers: list[int] = []

        if layers is None:
            logger.warning("ReplaceMe: could not locate transformer layer list.")
            stats = self.get_compression_stats()
            stats.update(
                {
                    "layers_removed": 0,
                    "layers_replaced": 0,
                    "fallback_layers": [],
                    "fallback_count": 0,
                    "target_removed": self.remove_layers,
                    "hidden_size": None,
                    "forward_check_success": False,
                    "use_lora": self.use_lora,
                    "lora_r": self.lora_r,
                    "method": "replaceme",
                }
            )
            self.stats = stats
            return self.model, stats

        hidden_size = self._infer_hidden_size(layers)
        num_layers = len(layers)
        n_remove = min(self.remove_layers, num_layers)

        target_dtype = (
            next(self.model.parameters()).dtype if list(self.model.parameters()) else torch.float32
        )

        for i in range(num_layers - n_remove, num_layers):
            old_layer = layers[i]

            # Detect whether the original layer returns a tuple.
            try:
                dummy_h = torch.randn(
                    1,
                    4,
                    hidden_size,
                    device=self.target_device,
                    dtype=target_dtype,
                )
                with torch.no_grad():
                    old_out = old_layer(dummy_h)
                return_tuple = isinstance(old_out, (tuple, list))
            except Exception as e:
                logger.debug(
                    "ReplaceMe: could not probe layer %d output type (%s); "
                    "defaulting to tuple=True.",
                    i,
                    e,
                )
                return_tuple = True

            new_block = ReplacementBlock(
                hidden_size,
                lora_r=self.lora_r if self.use_lora else 0,
                lora_alpha=self.lora_alpha,
                return_tuple=return_tuple,
                device=self.target_device,
                dtype=target_dtype,
            )
            layers[i] = new_block
            layers_replaced += 1

        # Forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                dummy_input = torch.randint(
                    0,
                    100,
                    (1, 8),
                    device=self.target_device,
                )
                try:
                    self.model(input_ids=dummy_input)
                except Exception:
                    # Some non-HF models expect positional arg.
                    self.model(dummy_input)
        except Exception as e:
            logger.warning("ReplaceMe forward sanity check failed: %s", e)
            forward_check_success = False

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        stats.update(
            {
                "layers_removed": layers_replaced,
                "layers_replaced": layers_replaced,
                "fallback_layers": fallback_layers,  # RM-FIX-03: list of indices
                "fallback_count": len(fallback_layers),
                "target_removed": self.remove_layers,
                "hidden_size": hidden_size,
                "forward_check_success": forward_check_success,
                "use_lora": self.use_lora,
                "lora_r": self.lora_r,
                "method": "replaceme",
            }
        )
        self.stats = stats
        return self.model, stats

Replaces the trailing N layers of a Transformer with a lightweight LayerNorm + Linear residual block.

Ancestors

BaseCompressor

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class ReplacementBlock (hidden_size: int, lora_r: int = 0, lora_alpha: float = 1.0, return_tuple: bool = True, device: torch.device | None = None, dtype: torch.dtype | None = None)

Expand source code

class ReplacementBlock(nn.Module):
    """
    Drop-in replacement for a transformer layer.

    Preserves tensor shapes and maintains signal propagation via a residual
    connection with optional LoRA fine-tuning adapters.
    """

    def __init__(
        self,
        hidden_size: int,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        return_tuple: bool = True,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.return_tuple = bool(return_tuple)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        self.ln = nn.LayerNorm(hidden_size, device=device, dtype=dtype)
        self.linear = nn.Linear(hidden_size, hidden_size, device=device, dtype=dtype)
        # Zero-init → block starts as near-identity
        nn.init.zeros_(self.linear.weight)
        if self.linear.bias is not None:
            nn.init.zeros_(self.linear.bias)

        if self.lora_r > 0:
            self.linear.weight.requires_grad = False
            if self.linear.bias is not None:
                self.linear.bias.requires_grad = False
            self.lora_A = nn.Parameter(
                torch.empty((self.lora_r, hidden_size), dtype=torch.float32, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.zeros((hidden_size, self.lora_r), dtype=torch.float32, device=device)
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    def forward(
        self,
        hidden_states: torch.Tensor | tuple[torch.Tensor, ...] | Any,
        *args: Any,
        **kwargs: Any,
    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        # Extract primary tensor
        if isinstance(hidden_states, (tuple, list)):
            primary = hidden_states[0]
        elif hasattr(hidden_states, "last_hidden_state"):
            primary = hidden_states.last_hidden_state
        else:
            primary = hidden_states

        normed = self.ln(primary)
        delta = self.linear(normed)
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_delta = F.linear(F.linear(normed.float(), self.lora_A), self.lora_B) * self.scaling
            delta = delta + lora_delta.to(delta.dtype)

        out = primary + delta

        # RM-FIX-02: return correct tuple shape
        if self.return_tuple:
            return (out,)
        return out

Drop-in replacement for a transformer layer.

Preserves tensor shapes and maintains signal propagation via a residual connection with optional LoRA fine-tuning adapters.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Methods

def forward(self, hidden_states: torch.Tensor | tuple[torch.Tensor, ...] | Any, *args: Any, **kwargs: Any) ‑> torch.Tensor | tuple[torch.Tensor, ...]

Expand source code

def forward(
    self,
    hidden_states: torch.Tensor | tuple[torch.Tensor, ...] | Any,
    *args: Any,
    **kwargs: Any,
) -> torch.Tensor | tuple[torch.Tensor, ...]:
    # Extract primary tensor
    if isinstance(hidden_states, (tuple, list)):
        primary = hidden_states[0]
    elif hasattr(hidden_states, "last_hidden_state"):
        primary = hidden_states.last_hidden_state
    else:
        primary = hidden_states

    normed = self.ln(primary)
    delta = self.linear(normed)
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_delta = F.linear(F.linear(normed.float(), self.lora_A), self.lora_B) * self.scaling
        delta = delta + lora_delta.to(delta.dtype)

    out = primary + delta

    # RM-FIX-02: return correct tuple shape
    if self.return_tuple:
        return (out,)
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class ResSVDCompressor (model: nn.Module, base_rank: int = 32, residual_rank: int = 8, scaling: float = 1.0, factorize_base: bool = True, freeze_base: bool = True, target_modules: list[str] | None = None, **kwargs: Any)

Expand source code

class ResSVDCompressor(BaseCompressor):
    """
    Spectral Residual LoRA (Res-SVD) Compressor.

    Replaces standard Linear layers with ResSVDLinear modules, truncating base weights
    at rank r (frozen) and initializing adapter matrices A, B directly from the residual error tail.
    """

    def __init__(
        self,
        model: nn.Module,
        base_rank: int = 32,
        residual_rank: int = 8,
        scaling: float = 1.0,
        factorize_base: bool = True,
        freeze_base: bool = True,
        target_modules: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="res_svd", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            base_rank = getattr(cfg, "base_rank", getattr(cfg, "rank", base_rank))
            residual_rank = getattr(cfg, "residual_rank", residual_rank)
            scaling = getattr(cfg, "scaling", getattr(cfg, "lora_alpha", scaling))
            factorize_base = getattr(cfg, "factorize_base", factorize_base)
            freeze_base = getattr(cfg, "freeze_base", freeze_base)

        if "rank" in kwargs and kwargs["rank"] is not None:
            base_rank = kwargs["rank"]

        self.base_rank = base_rank
        self.residual_rank = residual_rank
        self.scaling = scaling
        self.factorize_base = factorize_base
        self.freeze_base = freeze_base
        self.target_modules = target_modules
        self.replaced_layers: list[str] = []

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        """Check whether module is an eligible nn.Linear layer matching filter."""
        if not isinstance(module, nn.Linear):
            return False
        if isinstance(module, ResSVDLinear):
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """
        Compress target linear layers into ResSVDLinear modules.

        Returns:
            Tuple[nn.Module, Dict[str, Any]]: Compressed model and compression statistics.
        """
        self.replaced_layers = []

        # Find target linear layers
        target_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            logger.warning("No compressible linear layers found in model.")
            return self.model, self.get_compression_stats()

        for name, old_linear in target_layers:
            res_layer = ResSVDLinear.from_linear(
                linear=old_linear,
                base_rank=self.base_rank,
                residual_rank=self.residual_rank,
                scaling=self.scaling,
                factorize_base=self.factorize_base,
                freeze_base=self.freeze_base,
            )
            replace_module(self.model, name, res_layer)
            self.replaced_layers.append(name)

        # Calculate statistics
        total_params = self._count_params(self.model)
        trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        frozen_params = total_params - trainable_params

        stats = self.get_compression_stats()
        stats["trainable_parameters"] = trainable_params
        stats["frozen_parameters"] = frozen_params
        stats["trainable_percentage"] = (
            (trainable_params / total_params * 100.0) if total_params > 0 else 0.0
        )
        stats["base_rank"] = self.base_rank
        stats["residual_rank"] = self.residual_rank
        stats["replaced_layers"] = self.replaced_layers

        return self.model, stats

Spectral Residual LoRA (Res-SVD) Compressor.

Replaces standard Linear layers with ResSVDLinear modules, truncating base weights at rank r (frozen) and initializing adapter matrices A, B directly from the residual error tail.

Ancestors

BaseCompressor

Subclasses

neuralops.compression_engine.res_svd.ResSVDCompressorRegistered
SpectralResidualLoRA

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """
    Compress target linear layers into ResSVDLinear modules.

    Returns:
        Tuple[nn.Module, Dict[str, Any]]: Compressed model and compression statistics.
    """
    self.replaced_layers = []

    # Find target linear layers
    target_layers = [
        (name, mod)
        for name, mod in self.model.named_modules()
        if self._is_target_module(name, mod)
    ]

    if not target_layers:
        logger.warning("No compressible linear layers found in model.")
        return self.model, self.get_compression_stats()

    for name, old_linear in target_layers:
        res_layer = ResSVDLinear.from_linear(
            linear=old_linear,
            base_rank=self.base_rank,
            residual_rank=self.residual_rank,
            scaling=self.scaling,
            factorize_base=self.factorize_base,
            freeze_base=self.freeze_base,
        )
        replace_module(self.model, name, res_layer)
        self.replaced_layers.append(name)

    # Calculate statistics
    total_params = self._count_params(self.model)
    trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
    frozen_params = total_params - trainable_params

    stats = self.get_compression_stats()
    stats["trainable_parameters"] = trainable_params
    stats["frozen_parameters"] = frozen_params
    stats["trainable_percentage"] = (
        (trainable_params / total_params * 100.0) if total_params > 0 else 0.0
    )
    stats["base_rank"] = self.base_rank
    stats["residual_rank"] = self.residual_rank
    stats["replaced_layers"] = self.replaced_layers

    return self.model, stats

Compress target linear layers into ResSVDLinear modules.

Returns
-----=
Tuple[nn.Module, Dict[str, Any]]: Compressed model and compression statistics.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class ResSVDCompressorRegistered (model: nn.Module, base_rank: int = 32, residual_rank: int = 8, scaling: float = 1.0, factorize_base: bool = True, freeze_base: bool = True, target_modules: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("res_svd")
class ResSVDCompressorRegistered(ResSVDCompressor):
    pass

Spectral Residual LoRA (Res-SVD) Compressor.

Replaces standard Linear layers with ResSVDLinear modules, truncating base weights at rank r (frozen) and initializing adapter matrices A, B directly from the residual error tail.

Ancestors

Inherited members

ResSVDCompressor:
- compress
- get_compression_stats
- summary

class ResSVDLinear (in_features: int, out_features: int, base_rank: int = 32, residual_rank: int = 8, scaling: float = 1.0, factorize_base: bool = True, bias: bool = True, dtype: torch.dtype = torch.float32, device: torch.device | None = None)

Expand source code

class ResSVDLinear(nn.Module):
    """
    Spectral Residual LoRA (Res-SVD) Linear Layer.

    Executes forward pass as:
    y = x @ W_base^T + scaling * (x @ A^T) @ B^T + bias
    where W_base is either factorized as (base_V, base_U) with requires_grad=False,
    and (lora_A, lora_B) are trainable adapter matrices initialized from the residual tail.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        base_rank: int = 32,
        residual_rank: int = 8,
        scaling: float = 1.0,
        factorize_base: bool = True,
        bias: bool = True,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        max_possible_rank = min(in_features, out_features)

        # Bounds validation
        self.base_rank = max(1, min(base_rank, max_possible_rank))
        remaining_rank = max(0, max_possible_rank - self.base_rank)
        self.residual_rank = max(0, min(residual_rank, remaining_rank))
        self.scaling = float(scaling)
        self.factorize_base = factorize_base

        # Base parameters (Frozen)
        if self.factorize_base:
            self.base_V = nn.Parameter(
                torch.empty(self.base_rank, in_features, dtype=dtype, device=device),
                requires_grad=False,
            )
            self.base_U = nn.Parameter(
                torch.empty(out_features, self.base_rank, dtype=dtype, device=device),
                requires_grad=False,
            )
        else:
            self.base_weight = nn.Parameter(
                torch.empty(out_features, in_features, dtype=dtype, device=device),
                requires_grad=False,
            )

        # Trainable Residual LoRA Adapters
        if self.residual_rank > 0:
            self.lora_A = nn.Parameter(
                torch.empty(self.residual_rank, in_features, dtype=dtype, device=device),
                requires_grad=True,
            )
            self.lora_B = nn.Parameter(
                torch.empty(out_features, self.residual_rank, dtype=dtype, device=device),
                requires_grad=True,
            )
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        if bias:
            self.bias = nn.Parameter(torch.empty(out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

        self.reset_parameters()

    def reset_parameters(self) -> None:
        """Kaiming uniform parameter initialization for base and zero for adapters."""
        if self.factorize_base:
            nn.init.kaiming_uniform_(self.base_V, a=math.sqrt(5))
            nn.init.kaiming_uniform_(self.base_U, a=math.sqrt(5))
        else:
            nn.init.kaiming_uniform_(self.base_weight, a=math.sqrt(5))

        if self.residual_rank > 0:
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)

        if self.bias is not None:
            fan_in = self.in_features
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            nn.init.uniform_(self.bias, -bound, bound)

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix for compatibility."""
        if self.factorize_base:
            base = torch.matmul(self.base_U, self.base_V)
        else:
            base = self.base_weight
        if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            res = self.scaling * torch.matmul(self.lora_B, self.lora_A)
            return base + res
        return base

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward execution:
        y = Base(x) + scaling * (x @ A^T) @ B^T + bias
        """
        # RESEARCH FIX: Two-stage factorized forward for frozen base to minimize memory footprint
        if self.factorize_base:
            base_intermediate = F.linear(x, self.base_V)
            base_out = F.linear(base_intermediate, self.base_U)
        else:
            base_out = F.linear(x, self.base_weight)

        # Trainable spectral residual LoRA adapter forward
        if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            res_intermediate = F.linear(x, self.lora_A)
            res_out = self.scaling * F.linear(res_intermediate, self.lora_B)
            out = base_out + res_out
        else:
            out = base_out

        if self.bias is not None:
            out = out + self.bias

        return out

    @classmethod
    def from_linear(
        cls,
        linear: nn.Linear,
        base_rank: int = 32,
        residual_rank: int = 8,
        scaling: float = 1.0,
        factorize_base: bool = True,
        freeze_base: bool = True,
    ) -> ResSVDLinear:
        """
        Construct ResSVDLinear by decomposing a pretrained nn.Linear weight:
        1. Base factors U_r, V_r from the dominant top-r singular components.
        2. Residual adapter factors A, B directly initialized from the residual tail (r+1 to r+k).

        Args:
            linear: Source dense linear layer.
            base_rank: Rank r for the frozen base matrix.
            residual_rank: Rank k for the trainable residual adapter.
            scaling: Multiplier alpha / k for adapter output.
            factorize_base: Whether to store base weights as factorized U, V matrices.
            freeze_base: Whether to freeze base factors (requires_grad=False).

        Returns:
            ResSVDLinear: Spectral residual LoRA module.
        """
        in_f = linear.in_features
        out_f = linear.out_features
        max_rank = min(in_f, out_f)

        r = max(1, min(base_rank, max_rank))
        rem = max(0, max_rank - r)
        k = max(0, min(residual_rank, rem))

        new_layer = cls(
            in_features=in_f,
            out_features=out_f,
            base_rank=r,
            residual_rank=k,
            scaling=scaling,
            factorize_base=factorize_base,
            bias=linear.bias is not None,
            dtype=linear.weight.dtype,
            device=linear.weight.device,
        )

        W = linear.weight.data.float()  # (out_f, in_f)

        # SVD decomposition of source weight
        U, S, Vh = torch.linalg.svd(W, full_matrices=False)

        # 1. Base Truncation (Dominant Top-r Singular Components)
        U_base = U[:, :r] * torch.sqrt(S[:r]).unsqueeze(0)  # (out_f, r)
        V_base = torch.sqrt(S[:r]).unsqueeze(1) * Vh[:r, :]  # (r, in_f)

        if factorize_base:
            new_layer.base_U.data.copy_(U_base.to(dtype=linear.weight.dtype))
            new_layer.base_V.data.copy_(V_base.to(dtype=linear.weight.dtype))
            new_layer.base_U.requires_grad = not freeze_base
            new_layer.base_V.requires_grad = not freeze_base
        else:
            W_base = (U[:, :r] * S[:r].unsqueeze(0)) @ Vh[:r, :]
            new_layer.base_weight.data.copy_(W_base.to(dtype=linear.weight.dtype))
            new_layer.base_weight.requires_grad = not freeze_base

        # 2. Residual Tail Initialization (Singular directions r to r+k)
        # RESEARCH FIX: Analytical spectral initialization of LoRA adapter weights (A, B) directly
        # from the residual tail E = W - U_r V_r minimizes initial reconstruction loss to zero-shot optimality.
        if k > 0:
            S_tail = S[r : r + k]
            U_tail = U[:, r : r + k]
            Vh_tail = Vh[r : r + k, :]

            # Scale S_tail by adapter scaling factor gamma: (B @ A) * gamma = E_tail
            scale_factor = max(1e-8, float(scaling))
            S_tail_scaled = S_tail / scale_factor

            sqrt_S = torch.sqrt(torch.clamp(S_tail_scaled, min=0.0))
            lora_A_init = sqrt_S.unsqueeze(1) * Vh_tail  # (k, in_f)
            lora_B_init = U_tail * sqrt_S.unsqueeze(0)  # (out_f, k)

            new_layer.lora_A.data.copy_(lora_A_init.to(dtype=linear.weight.dtype))
            new_layer.lora_B.data.copy_(lora_B_init.to(dtype=linear.weight.dtype))
            new_layer.lora_A.requires_grad = True
            new_layer.lora_B.requires_grad = True

        # Copy bias if present
        if linear.bias is not None:
            new_layer.bias.data.copy_(linear.bias.data)

        return new_layer

    def merge_weights(self) -> nn.Linear:
        """
        Merge frozen base weights and trainable residual adapter weights into a dense nn.Linear.

        Returns:
            nn.Linear: Standard dense linear layer with merged weights.
        """
        device = next(self.parameters()).device
        dtype = self.base_U.dtype if self.factorize_base else self.base_weight.dtype

        merged_linear = nn.Linear(
            self.in_features,
            self.out_features,
            bias=self.bias is not None,
            device=device,
            dtype=dtype,
        )

        with torch.no_grad():
            if self.factorize_base:
                W_base = self.base_U.float() @ self.base_V.float()
            else:
                W_base = self.base_weight.float()

            if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
                delta_W = self.scaling * (self.lora_B.float() @ self.lora_A.float())
                W_total = W_base + delta_W
            else:
                W_total = W_base

            merged_linear.weight.data.copy_(W_total.to(dtype=dtype))
            if self.bias is not None:
                merged_linear.bias.data.copy_(self.bias.data)

        return merged_linear

    def get_residual_error_norm(self, original_weight: torch.Tensor | None = None) -> float:
        """
        Calculate relative Frobenius norm reconstruction error:
        || W_orig - (W_base + scaling * B @ A) ||_F / || W_orig ||_F
        """
        with torch.no_grad():
            if self.factorize_base:
                W_approx = self.base_U.float() @ self.base_V.float()
            else:
                W_approx = self.base_weight.float()

            if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
                W_approx = W_approx + self.scaling * (self.lora_B.float() @ self.lora_A.float())

            if original_weight is not None:
                W_orig = original_weight.float().to(W_approx.device)
                error = torch.norm(W_orig - W_approx, p="fro") / (
                    torch.norm(W_orig, p="fro") + 1e-12
                )
                return float(error.item())

            return 0.0

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, out_features={self.out_features}, "
            f"base_rank={self.base_rank}, residual_rank={self.residual_rank}, "
            f"scaling={self.scaling}, factorize_base={self.factorize_base}"
        )

Spectral Residual LoRA (Res-SVD) Linear Layer.

Executes forward pass as: y = x @ W_base^T + scaling * (x @ A^T) @ B^T + bias where W_base is either factorized as (base_V, base_U) with requires_grad=False, and (lora_A, lora_B) are trainable adapter matrices initialized from the residual tail.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Linear, base_rank: int = 32, residual_rank: int = 8, scaling: float = 1.0, factorize_base: bool = True, freeze_base: bool = True) ‑> ResSVDLinear

Construct ResSVDLinear by decomposing a pretrained nn.Linear weight: 1. Base factors U_r, V_r from the dominant top-r singular components. 2. Residual adapter factors A, B directly initialized from the residual tail (r+1 to r+k).

Args
-----=
linear: Source dense linear layer.
base_rank: Rank r for the frozen base matrix.
residual_rank: Rank k for the trainable residual adapter.
scaling: Multiplier alpha / k for adapter output.
factorize_base: Whether to store base weights as factorized U, V matrices.
freeze_base: Whether to freeze base factors (requires_grad=False).
Returns
-----=
ResSVDLinear: Spectral residual LoRA module.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix for compatibility."""
    if self.factorize_base:
        base = torch.matmul(self.base_U, self.base_V)
    else:
        base = self.base_weight
    if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        res = self.scaling * torch.matmul(self.lora_B, self.lora_A)
        return base + res
    return base

Reconstructed effective weight matrix for compatibility.

Methods

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    return (
        f"in_features={self.in_features}, out_features={self.out_features}, "
        f"base_rank={self.base_rank}, residual_rank={self.residual_rank}, "
        f"scaling={self.scaling}, factorize_base={self.factorize_base}"
    )

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Forward execution:
    y = Base(x) + scaling * (x @ A^T) @ B^T + bias
    """
    # RESEARCH FIX: Two-stage factorized forward for frozen base to minimize memory footprint
    if self.factorize_base:
        base_intermediate = F.linear(x, self.base_V)
        base_out = F.linear(base_intermediate, self.base_U)
    else:
        base_out = F.linear(x, self.base_weight)

    # Trainable spectral residual LoRA adapter forward
    if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        res_intermediate = F.linear(x, self.lora_A)
        res_out = self.scaling * F.linear(res_intermediate, self.lora_B)
        out = base_out + res_out
    else:
        out = base_out

    if self.bias is not None:
        out = out + self.bias

    return out

Forward execution: y = Base(x) + scaling * (x @ A^T) @ B^T + bias

def get_residual_error_norm(self, original_weight: torch.Tensor | None = None) ‑> float

Expand source code

def get_residual_error_norm(self, original_weight: torch.Tensor | None = None) -> float:
    """
    Calculate relative Frobenius norm reconstruction error:
    || W_orig - (W_base + scaling * B @ A) ||_F / || W_orig ||_F
    """
    with torch.no_grad():
        if self.factorize_base:
            W_approx = self.base_U.float() @ self.base_V.float()
        else:
            W_approx = self.base_weight.float()

        if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            W_approx = W_approx + self.scaling * (self.lora_B.float() @ self.lora_A.float())

        if original_weight is not None:
            W_orig = original_weight.float().to(W_approx.device)
            error = torch.norm(W_orig - W_approx, p="fro") / (
                torch.norm(W_orig, p="fro") + 1e-12
            )
            return float(error.item())

        return 0.0

Calculate relative Frobenius norm reconstruction error: || W_orig - (W_base + scaling * B @ A) ||_F / || W_orig ||_F

def merge_weights(self) ‑> torch.nn.modules.linear.Linear

Expand source code

def merge_weights(self) -> nn.Linear:
    """
    Merge frozen base weights and trainable residual adapter weights into a dense nn.Linear.

    Returns:
        nn.Linear: Standard dense linear layer with merged weights.
    """
    device = next(self.parameters()).device
    dtype = self.base_U.dtype if self.factorize_base else self.base_weight.dtype

    merged_linear = nn.Linear(
        self.in_features,
        self.out_features,
        bias=self.bias is not None,
        device=device,
        dtype=dtype,
    )

    with torch.no_grad():
        if self.factorize_base:
            W_base = self.base_U.float() @ self.base_V.float()
        else:
            W_base = self.base_weight.float()

        if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            delta_W = self.scaling * (self.lora_B.float() @ self.lora_A.float())
            W_total = W_base + delta_W
        else:
            W_total = W_base

        merged_linear.weight.data.copy_(W_total.to(dtype=dtype))
        if self.bias is not None:
            merged_linear.bias.data.copy_(self.bias.data)

    return merged_linear

Merge frozen base weights and trainable residual adapter weights into a dense nn.Linear.

Returns
-----=
nn.Linear: Standard dense linear layer with merged weights.

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    """Kaiming uniform parameter initialization for base and zero for adapters."""
    if self.factorize_base:
        nn.init.kaiming_uniform_(self.base_V, a=math.sqrt(5))
        nn.init.kaiming_uniform_(self.base_U, a=math.sqrt(5))
    else:
        nn.init.kaiming_uniform_(self.base_weight, a=math.sqrt(5))

    if self.residual_rank > 0:
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)

    if self.bias is not None:
        fan_in = self.in_features
        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
        nn.init.uniform_(self.bias, -bound, bound)

Kaiming uniform parameter initialization for base and zero for adapters.

class RotationalSliceGPTCompressor (model: nn.Module | None = None, r_ratio: float = 0.4, lambda_reg: float = 0.01, slice_ratio: float | None = None, quant_bits: int = 4, max_error: float = 0.2, fallback_to_original: bool = True, tokenizer: Any | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, target_modules: list[str] | None = None, target_attention_only: bool = False, seed: int = 42, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("rotational_slicegpt_int4")
@COMPRESSOR_REGISTRY.register("rotational_slicegpt")
class RotationalSliceGPTINT4Compressor(BaseCompressor):
    """
    Rotational-SliceGPT-INT4 hybrid compression engine.
    """

    # RS-FIX-04: not an attention module.
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module | None = None,
        r_ratio: float = 0.4,  # RS-FIX-03: lower default
        lambda_reg: float = 1e-2,
        slice_ratio: float | None = None,
        quant_bits: int = 4,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        tokenizer: Any | None = None,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        target_modules: list[str] | None = None,
        target_attention_only: bool = False,
        seed: int = 42,  # RS-FIX-05
        **kwargs: Any,
    ) -> None:
        # Remove duplicates that arrive through kwargs/comp_config.
        for k in (
            "max_error",
            "fallback_to_original",
            "r_ratio",
            "quant_bits",
            "lambda_reg",
        ):
            kwargs.pop(k, None)
        target_model = model if model is not None else nn.Identity()
        super().__init__(target_model, name="rotational_slicegpt_int4", **kwargs)

        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        r_ratio = _cfg_get(cfg, "rotational_slicegpt_r_ratio", _cfg_get(cfg, "r_ratio", r_ratio))
        lambda_reg = _cfg_get(cfg, "lambda_reg", lambda_reg)
        quant_bits = _cfg_get(cfg, "quant_bits", quant_bits)
        max_error = _cfg_get(cfg, "max_error", max_error)
        fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
        target_attention_only = bool(
            target_attention_only or _cfg_get(cfg, "target_attention_only", False)
        )

        self.r_ratio = float(slice_ratio if slice_ratio is not None else r_ratio)
        if self.r_ratio >= 0.5:
            logger.warning(
                "RotationalSliceGPT: r_ratio=%.2f ≥ 0.5 will INCREASE parameter count. "
                "Consider r_ratio ≤ 0.4 for real savings.",
                self.r_ratio,
            )
        self.lambda_reg = float(lambda_reg)
        self.quant_bits = int(quant_bits)
        self.quantizer = AstroQuantizer(bits=self.quant_bits)
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.target_modules = target_modules
        self.target_attention_only = target_attention_only
        self.seed = int(seed)

        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.target_device = (
            torch.device(device)
            if device is not None
            else (
                next(target_model.parameters(), torch.tensor(0.0)).device
                if model is not None
                else torch.device("cpu")
            )
        )
        self.layer_errors: dict[str, float] = {}
        self.stats: dict[str, Any] = {}
        self._orig_bytes = 0
        self._comp_bytes = 0

    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        if any(
            k in name_l
            for k in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "layernorm",
                "rmsnorm",
                "batchnorm",
            )
        ):
            return False

        if isinstance(
            module,
            (
                nn.LayerNorm,
                nn.RMSNorm,
                nn.Embedding,
                nn.BatchNorm1d,
                nn.BatchNorm2d,
                RotationalSliceGPTINT4Linear,
            ),
        ):
            return False

        if self.target_attention_only:
            attn_kw = (
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            )
            if not any(k in name_l for k in attn_kw):
                return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__.lower()
        is_linear_like = (
            isinstance(base_mod, nn.Linear) or "conv1d" in cls_name or hasattr(base_mod, "nf")
        )
        if not is_linear_like:
            return False

        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    # ------------------------------------------------------------------
    def get_orthogonal_matrix(self, d_in: int, device: torch.device) -> torch.Tensor:
        """Deterministic QR-based orthogonal Q (RS-FIX-05)."""
        gen = torch.Generator(device=device).manual_seed(self.seed)
        H = torch.randn(d_in, d_in, device=device, dtype=torch.float32, generator=gen)
        Q, _ = torch.linalg.qr(H)
        return Q

    # ------------------------------------------------------------------
    def compress_matrix(
        self,
        W: torch.Tensor,
        X: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, torch.Tensor]:
        """
        Returns:
            W_quant_indices:  (D_out, r) int8
            W_quant_centroids:(C,)       fp32
            Q_P_r:            (D_in, r)  fp32
            rel_error:        float
        """
        D_out, D_in = W.shape
        r = max(1, min(int(D_in * self.r_ratio), D_in))
        device = W.device

        # 1. Orthogonal rotation
        Q = self.get_orthogonal_matrix(D_in, device)
        W_fp32 = W.detach().float()
        X_fp32 = X.detach().float().to(device)

        W_tilde = W_fp32 @ Q  # (D_out, D_in)
        X_tilde = Q.T @ X_fp32  # (D_in, N)

        # 2. Variance-based slicing (RS-FIX-06: vectorized)
        if X_tilde.shape[1] > 1:
            channel_variances = X_tilde.var(dim=1)  # (D_in,)
        else:
            channel_variances = torch.ones(D_in, device=device)
        top_r_indices = torch.argsort(channel_variances, descending=True)[:r]

        # P_r as a (D_in, r) selector matrix via scatter
        P_r = torch.zeros(D_in, r, device=device, dtype=torch.float32)
        P_r[top_r_indices, torch.arange(r, device=device)] = 1.0

        W_sliced = W_tilde @ P_r  # (D_out, r)
        X_sliced = P_r.T @ X_tilde  # (r, N)

        # 3. Second-order Hessian-based INT4 quantization
        H = 2.0 * (X_sliced @ X_sliced.T) / max(X_sliced.shape[1], 1) + self.lambda_reg * torch.eye(
            r, device=device
        )
        try:
            H_inv = torch.linalg.inv(H)
        except Exception:
            H_inv = torch.linalg.pinv(H)
        H_inv = torch.nan_to_num(H_inv, nan=0.0)

        # RS-FIX-10: quantize the entire column of residuals at once.
        W_quant = W_sliced.clone()
        # Full-matrix OBS: quantize each column then propagate the error.
        for j in range(r):
            w_j = W_quant[:, j].clone()
            idx_j, cents_j = self.quantizer.quantize(w_j)
            q_j = AstroQuantizer.dequantize(idx_j, cents_j).to(w_j.dtype)
            e_j = w_j - q_j
            W_quant[:, j] = q_j

            diag_inv = float(H_inv[j, j].item())
            if not math.isfinite(diag_inv) or diag_inv <= 1e-8:
                diag_inv = 1e-5

            if j < r - 1:
                h_slice = torch.nan_to_num(H_inv[j, j + 1 : r] / diag_inv, nan=0.0)
                W_quant[:, j + 1 : r] -= torch.outer(e_j, h_slice)

        # Re-quantize the final W_quant column-wise to get indices + centroids
        # (OBS updates can change values slightly, so we requantize once).
        W_quant.t().contiguous()  # (r, D_out) → quantize per row
        _W_idx_list, _W_cents_list = [], []
        per_col_centroids = []
        W_indices = torch.empty_like(W_quant, dtype=torch.int8)
        for j in range(r):
            idx_j, cents_j = self.quantizer.quantize(W_quant[:, j])
            W_indices[:, j] = idx_j
            per_col_centroids.append(cents_j)
        # Use a single shared centroid table = average of per-column tables.
        W_centroids = torch.stack(per_col_centroids, dim=0).mean(dim=0)

        # Reconstruct and measure error (using the shared table for honesty).
        W_deq = AstroQuantizer.dequantize(W_indices, W_centroids).to(W_quant.dtype)
        W_recon = (W_deq @ P_r.T) @ Q.T  # (D_out, D_in)
        norm_orig = W_fp32.norm(p="fro").clamp_min(1e-8)
        rel_error = float((W_fp32 - W_recon).norm(p="fro") / norm_orig)

        Q_P_r = Q @ P_r  # (D_in, r) — RS-FIX-08
        return W_indices, W_centroids, Q_P_r, rel_error

    compress_tensor = compress_matrix  # backward-compat alias

    # ------------------------------------------------------------------
    def compress(self, *args: Any, **kwargs: Any) -> Any:
        # Dual-mode: .compress(W, X) → tuple, .compress() → (model, stats)
        if (
            len(args) >= 2
            and isinstance(args[0], torch.Tensor)
            and isinstance(args[1], torch.Tensor)
        ):
            return self.compress_matrix(args[0], args[1])

        activations_map = self._collect_layer_activations()
        compressed_layers_count = 0
        fallback_layers: list[str] = []
        self.layer_errors = {}
        layer_shapes: dict[str, tuple[int, int]] = {}
        self._orig_bytes = 0
        self._comp_bytes = 0

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            base_mod = getattr(module, "base_layer", module)
            raw_w = base_mod.weight.data
            cls_name = base_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(base_mod, "nf")
            device = raw_w.device
            dtype = raw_w.dtype

            W = raw_w.t().contiguous().float() if is_conv1d else raw_w.clone().float()
            d_out, d_in = W.shape
            layer_shapes[name] = (d_out, d_in)

            # Build X for slicing (RS-FIX-07: per-channel mean)
            act = activations_map.get(name, None)
            if act is not None and act.numel() > 0:
                act_flat = act.reshape(-1, act.shape[-1]).float()
                if act_flat.shape[-1] == d_in:
                    # Synthesize a small sample matrix from the mean vector.
                    mean_act = act_flat.mean(dim=0, keepdim=True)  # (1, D_in)
                    # Repeat with small jitter to give variance-based slicing signal.
                    X = mean_act.T + 0.01 * torch.randn(d_in, 32, device=device)
                else:
                    X = torch.randn(d_in, 32, device=device)
            else:
                X = torch.randn(d_in, 32, device=device)

            try:
                W_idx, W_cents, Q_P_r, rel_error = self.compress_matrix(W, X)
            except Exception as e:
                logger.warning("RotationalSliceGPT failed on %s: %s", name, e)
                fallback_layers.append(name)
                continue

            self.layer_errors[name] = rel_error

            if rel_error > self.max_error:
                logger.warning(
                    "RotationalSliceGPT reconstruction error %.4f on %s exceeded %.4f",
                    rel_error,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

            # Cast buffers to the original dtype & device
            W_idx = W_idx.to(device=device)
            W_cents = W_cents.to(device=device, dtype=torch.float32)
            Q_P_r_dev = Q_P_r.to(device=device, dtype=dtype)

            bias = base_mod.bias.detach() if getattr(base_mod, "bias", None) is not None else None

            new_linear = RotationalSliceGPTINT4Linear(
                W_indices=W_idx,
                W_centroids=W_cents,
                Q_P_r=Q_P_r_dev,
                bias=bias,
            )

            # Splice into model
            if hasattr(module, "base_layer"):
                # Preserve any LoRA params
                if hasattr(module, "lora_A") and module.lora_A is not None:
                    if not hasattr(new_linear, "lora_A"):
                        new_linear.lora_A = module.lora_A
                if hasattr(module, "lora_B") and module.lora_B is not None:
                    if not hasattr(new_linear, "lora_B"):
                        new_linear.lora_B = module.lora_B
                module.base_layer = new_linear
            else:
                replace_module(self.model, name, new_linear)

            # RS-FIX-09: honest byte accounting
            orig_bytes = raw_w.numel() * raw_w.element_size()
            if base_mod.bias is not None:
                orig_bytes += base_mod.bias.numel() * base_mod.bias.element_size()
            # comp: int8 indices + fp32 centroid table + Q_P_r + bias
            comp_bytes = W_idx.numel() * W_idx.element_size()
            comp_bytes += W_cents.numel() * W_cents.element_size()
            comp_bytes += Q_P_r_dev.numel() * Q_P_r_dev.element_size()
            if bias is not None:
                comp_bytes += bias.numel() * bias.element_size()

            self._orig_bytes += orig_bytes
            self._comp_bytes += comp_bytes
            compressed_layers_count += 1

        # Forward sanity check
        forward_check_success = True
        try:
            device = next(self.model.parameters()).device
            with torch.no_grad():
                cfg = getattr(self.model, "config", None)
                vocab = getattr(cfg, "vocab_size", None) or 100
                dummy = torch.randint(0, min(vocab, 100), (1, 8), device=device)
                self.model(dummy)
        except Exception as ex:
            logger.warning("RotationalSliceGPT forward sanity check: %s", ex)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        stats = self.get_compression_stats()
        avg_err = float(np.mean(list(self.layer_errors.values()))) if self.layer_errors else 0.0
        storage_ratio = self._orig_bytes / self._comp_bytes if self._comp_bytes > 0 else 1.0

        stats.update(
            {
                "compressed_layers": compressed_layers_count,
                "fallback_layers": fallback_layers,
                "r_ratio": self.r_ratio,
                "quant_bits": self.quant_bits,
                "lambda_reg": self.lambda_reg,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "layer_errors": self.layer_errors,
                "layer_shapes": layer_shapes,
                "avg_reconstruction_error": avg_err,
                "target_attention_only": self.target_attention_only,
                "forward_check_success": forward_check_success,
                "orig_storage_bytes": self._orig_bytes,
                "comp_storage_bytes": self._comp_bytes,
                "storage_compression_ratio": round(storage_ratio, 3),
                "storage_reduction_pct": round(
                    100.0 * (1.0 - self._comp_bytes / max(1, self._orig_bytes)), 2
                ),
                "method": "rotational_slicegpt_int4",
            }
        )
        self.stats = stats
        logger.info(
            "RotationalSliceGPT: %d layers compressed, %d fallback, "
            "avg_err=%.4f, storage_ratio=%.2fx, bits=%d, r_ratio=%.2f",
            compressed_layers_count,
            len(fallback_layers),
            avg_err,
            storage_ratio,
            self.quant_bits,
            self.r_ratio,
        )
        return self.model, stats

    # ------------------------------------------------------------------
    def _collect_layer_activations(self) -> dict[str, torch.Tensor]:
        """
        RS-FIX-07: store per-channel mean |x| instead of full activation
        tensors to avoid OOM on large models.
        """
        self.model.eval()
        activations: dict[str, list[torch.Tensor]] = {}
        hooks = []

        def make_hook(layer_name: str):
            def hook(module, inputs, output):
                if inputs and isinstance(inputs[0], torch.Tensor):
                    x = inputs[0].detach()
                    if x.dim() >= 2:
                        col_mean = x.abs().mean(dim=tuple(range(x.dim() - 1)))
                    else:
                        col_mean = x.abs()
                    activations.setdefault(layer_name, []).append(col_mean.cpu())

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                target = getattr(module, "base_layer", module)
                hooks.append(target.register_forward_hook(make_hook(name)))

        try:
            with torch.no_grad():
                calib_inputs = get_calibration_tensors(
                    self.model,
                    tokenizer=self.tokenizer,
                    calibration_loader=self.calibration_loader,
                    calibration_texts=self.calibration_texts,
                    device=self.target_device,
                )
                try:
                    if isinstance(calib_inputs, dict):
                        self.model(**calib_inputs)
                    elif isinstance(calib_inputs, torch.Tensor):
                        self.model(calib_inputs)
                    elif isinstance(calib_inputs, (list, tuple)):
                        self.model(*calib_inputs)
                except Exception as e:
                    logger.debug("Activation capture failed: %s", e)
        finally:
            for h in hooks:
                h.remove()

        aggregated: dict[str, torch.Tensor] = {}
        for name, lst in activations.items():
            if lst:
                aggregated[name] = torch.stack([t.float() for t in lst], dim=0).mean(dim=0)
        return aggregated

    # ------------------------------------------------------------------
    def compress_linear(
        self,
        linear_layer: nn.Linear,
        X: torch.Tensor | None = None,
    ) -> RotationalSliceGPTINT4Linear:
        """Convenience wrapper for compressing a single nn.Linear."""
        W = linear_layer.weight.data.float()
        device = W.device
        if X is None:
            X = torch.randn(W.shape[1], 128, device=device)
        elif X.dim() == 3:
            X = X.reshape(-1, W.shape[1]).T
        elif X.dim() == 2 and X.shape[0] != W.shape[1]:
            X = X.T
        W_idx, W_cents, Q_P_r, _ = self.compress_matrix(W, X)
        bias = linear_layer.bias.data if getattr(linear_layer, "bias", None) is not None else None
        return RotationalSliceGPTINT4Linear(
            W_indices=W_idx.to(device),
            W_centroids=W_cents.to(device),
            Q_P_r=Q_P_r.to(device=device, dtype=W.dtype),
            bias=bias,
        )

Rotational-SliceGPT-INT4 hybrid compression engine.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def compress(self, *args: Any, **kwargs: Any) ‑> Any

Expand source code

def compress(self, *args: Any, **kwargs: Any) -> Any:
    # Dual-mode: .compress(W, X) → tuple, .compress() → (model, stats)
    if (
        len(args) >= 2
        and isinstance(args[0], torch.Tensor)
        and isinstance(args[1], torch.Tensor)
    ):
        return self.compress_matrix(args[0], args[1])

    activations_map = self._collect_layer_activations()
    compressed_layers_count = 0
    fallback_layers: list[str] = []
    self.layer_errors = {}
    layer_shapes: dict[str, tuple[int, int]] = {}
    self._orig_bytes = 0
    self._comp_bytes = 0

    for name, module in list(self.model.named_modules()):
        if not self._is_target_module(name, module):
            continue

        base_mod = getattr(module, "base_layer", module)
        raw_w = base_mod.weight.data
        cls_name = base_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(base_mod, "nf")
        device = raw_w.device
        dtype = raw_w.dtype

        W = raw_w.t().contiguous().float() if is_conv1d else raw_w.clone().float()
        d_out, d_in = W.shape
        layer_shapes[name] = (d_out, d_in)

        # Build X for slicing (RS-FIX-07: per-channel mean)
        act = activations_map.get(name, None)
        if act is not None and act.numel() > 0:
            act_flat = act.reshape(-1, act.shape[-1]).float()
            if act_flat.shape[-1] == d_in:
                # Synthesize a small sample matrix from the mean vector.
                mean_act = act_flat.mean(dim=0, keepdim=True)  # (1, D_in)
                # Repeat with small jitter to give variance-based slicing signal.
                X = mean_act.T + 0.01 * torch.randn(d_in, 32, device=device)
            else:
                X = torch.randn(d_in, 32, device=device)
        else:
            X = torch.randn(d_in, 32, device=device)

        try:
            W_idx, W_cents, Q_P_r, rel_error = self.compress_matrix(W, X)
        except Exception as e:
            logger.warning("RotationalSliceGPT failed on %s: %s", name, e)
            fallback_layers.append(name)
            continue

        self.layer_errors[name] = rel_error

        if rel_error > self.max_error:
            logger.warning(
                "RotationalSliceGPT reconstruction error %.4f on %s exceeded %.4f",
                rel_error,
                name,
                self.max_error,
            )
            if self.fallback_to_original:
                fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                continue

        # Cast buffers to the original dtype & device
        W_idx = W_idx.to(device=device)
        W_cents = W_cents.to(device=device, dtype=torch.float32)
        Q_P_r_dev = Q_P_r.to(device=device, dtype=dtype)

        bias = base_mod.bias.detach() if getattr(base_mod, "bias", None) is not None else None

        new_linear = RotationalSliceGPTINT4Linear(
            W_indices=W_idx,
            W_centroids=W_cents,
            Q_P_r=Q_P_r_dev,
            bias=bias,
        )

        # Splice into model
        if hasattr(module, "base_layer"):
            # Preserve any LoRA params
            if hasattr(module, "lora_A") and module.lora_A is not None:
                if not hasattr(new_linear, "lora_A"):
                    new_linear.lora_A = module.lora_A
            if hasattr(module, "lora_B") and module.lora_B is not None:
                if not hasattr(new_linear, "lora_B"):
                    new_linear.lora_B = module.lora_B
            module.base_layer = new_linear
        else:
            replace_module(self.model, name, new_linear)

        # RS-FIX-09: honest byte accounting
        orig_bytes = raw_w.numel() * raw_w.element_size()
        if base_mod.bias is not None:
            orig_bytes += base_mod.bias.numel() * base_mod.bias.element_size()
        # comp: int8 indices + fp32 centroid table + Q_P_r + bias
        comp_bytes = W_idx.numel() * W_idx.element_size()
        comp_bytes += W_cents.numel() * W_cents.element_size()
        comp_bytes += Q_P_r_dev.numel() * Q_P_r_dev.element_size()
        if bias is not None:
            comp_bytes += bias.numel() * bias.element_size()

        self._orig_bytes += orig_bytes
        self._comp_bytes += comp_bytes
        compressed_layers_count += 1

    # Forward sanity check
    forward_check_success = True
    try:
        device = next(self.model.parameters()).device
        with torch.no_grad():
            cfg = getattr(self.model, "config", None)
            vocab = getattr(cfg, "vocab_size", None) or 100
            dummy = torch.randint(0, min(vocab, 100), (1, 8), device=device)
            self.model(dummy)
    except Exception as ex:
        logger.warning("RotationalSliceGPT forward sanity check: %s", ex)
        forward_check_success = False

    try:
        self._compressed_params = self._count_params(self.model)
    except Exception:
        self._compressed_params = sum(p.numel() for p in self.model.parameters())

    stats = self.get_compression_stats()
    avg_err = float(np.mean(list(self.layer_errors.values()))) if self.layer_errors else 0.0
    storage_ratio = self._orig_bytes / self._comp_bytes if self._comp_bytes > 0 else 1.0

    stats.update(
        {
            "compressed_layers": compressed_layers_count,
            "fallback_layers": fallback_layers,
            "r_ratio": self.r_ratio,
            "quant_bits": self.quant_bits,
            "lambda_reg": self.lambda_reg,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "layer_errors": self.layer_errors,
            "layer_shapes": layer_shapes,
            "avg_reconstruction_error": avg_err,
            "target_attention_only": self.target_attention_only,
            "forward_check_success": forward_check_success,
            "orig_storage_bytes": self._orig_bytes,
            "comp_storage_bytes": self._comp_bytes,
            "storage_compression_ratio": round(storage_ratio, 3),
            "storage_reduction_pct": round(
                100.0 * (1.0 - self._comp_bytes / max(1, self._orig_bytes)), 2
            ),
            "method": "rotational_slicegpt_int4",
        }
    )
    self.stats = stats
    logger.info(
        "RotationalSliceGPT: %d layers compressed, %d fallback, "
        "avg_err=%.4f, storage_ratio=%.2fx, bits=%d, r_ratio=%.2f",
        compressed_layers_count,
        len(fallback_layers),
        avg_err,
        storage_ratio,
        self.quant_bits,
        self.r_ratio,
    )
    return self.model, stats

Return (compressed_model, statistics).

def compress_linear(self, linear_layer: nn.Linear, X: torch.Tensor | None = None) ‑> RotationalSliceGPTINT4Linear

Expand source code

def compress_linear(
    self,
    linear_layer: nn.Linear,
    X: torch.Tensor | None = None,
) -> RotationalSliceGPTINT4Linear:
    """Convenience wrapper for compressing a single nn.Linear."""
    W = linear_layer.weight.data.float()
    device = W.device
    if X is None:
        X = torch.randn(W.shape[1], 128, device=device)
    elif X.dim() == 3:
        X = X.reshape(-1, W.shape[1]).T
    elif X.dim() == 2 and X.shape[0] != W.shape[1]:
        X = X.T
    W_idx, W_cents, Q_P_r, _ = self.compress_matrix(W, X)
    bias = linear_layer.bias.data if getattr(linear_layer, "bias", None) is not None else None
    return RotationalSliceGPTINT4Linear(
        W_indices=W_idx.to(device),
        W_centroids=W_cents.to(device),
        Q_P_r=Q_P_r.to(device=device, dtype=W.dtype),
        bias=bias,
    )

Convenience wrapper for compressing a single nn.Linear.

def compress_matrix(self, W: torch.Tensor, X: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, torch.Tensor]

Expand source code

def compress_matrix(
    self,
    W: torch.Tensor,
    X: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, torch.Tensor]:
    """
    Returns:
        W_quant_indices:  (D_out, r) int8
        W_quant_centroids:(C,)       fp32
        Q_P_r:            (D_in, r)  fp32
        rel_error:        float
    """
    D_out, D_in = W.shape
    r = max(1, min(int(D_in * self.r_ratio), D_in))
    device = W.device

    # 1. Orthogonal rotation
    Q = self.get_orthogonal_matrix(D_in, device)
    W_fp32 = W.detach().float()
    X_fp32 = X.detach().float().to(device)

    W_tilde = W_fp32 @ Q  # (D_out, D_in)
    X_tilde = Q.T @ X_fp32  # (D_in, N)

    # 2. Variance-based slicing (RS-FIX-06: vectorized)
    if X_tilde.shape[1] > 1:
        channel_variances = X_tilde.var(dim=1)  # (D_in,)
    else:
        channel_variances = torch.ones(D_in, device=device)
    top_r_indices = torch.argsort(channel_variances, descending=True)[:r]

    # P_r as a (D_in, r) selector matrix via scatter
    P_r = torch.zeros(D_in, r, device=device, dtype=torch.float32)
    P_r[top_r_indices, torch.arange(r, device=device)] = 1.0

    W_sliced = W_tilde @ P_r  # (D_out, r)
    X_sliced = P_r.T @ X_tilde  # (r, N)

    # 3. Second-order Hessian-based INT4 quantization
    H = 2.0 * (X_sliced @ X_sliced.T) / max(X_sliced.shape[1], 1) + self.lambda_reg * torch.eye(
        r, device=device
    )
    try:
        H_inv = torch.linalg.inv(H)
    except Exception:
        H_inv = torch.linalg.pinv(H)
    H_inv = torch.nan_to_num(H_inv, nan=0.0)

    # RS-FIX-10: quantize the entire column of residuals at once.
    W_quant = W_sliced.clone()
    # Full-matrix OBS: quantize each column then propagate the error.
    for j in range(r):
        w_j = W_quant[:, j].clone()
        idx_j, cents_j = self.quantizer.quantize(w_j)
        q_j = AstroQuantizer.dequantize(idx_j, cents_j).to(w_j.dtype)
        e_j = w_j - q_j
        W_quant[:, j] = q_j

        diag_inv = float(H_inv[j, j].item())
        if not math.isfinite(diag_inv) or diag_inv <= 1e-8:
            diag_inv = 1e-5

        if j < r - 1:
            h_slice = torch.nan_to_num(H_inv[j, j + 1 : r] / diag_inv, nan=0.0)
            W_quant[:, j + 1 : r] -= torch.outer(e_j, h_slice)

    # Re-quantize the final W_quant column-wise to get indices + centroids
    # (OBS updates can change values slightly, so we requantize once).
    W_quant.t().contiguous()  # (r, D_out) → quantize per row
    _W_idx_list, _W_cents_list = [], []
    per_col_centroids = []
    W_indices = torch.empty_like(W_quant, dtype=torch.int8)
    for j in range(r):
        idx_j, cents_j = self.quantizer.quantize(W_quant[:, j])
        W_indices[:, j] = idx_j
        per_col_centroids.append(cents_j)
    # Use a single shared centroid table = average of per-column tables.
    W_centroids = torch.stack(per_col_centroids, dim=0).mean(dim=0)

    # Reconstruct and measure error (using the shared table for honesty).
    W_deq = AstroQuantizer.dequantize(W_indices, W_centroids).to(W_quant.dtype)
    W_recon = (W_deq @ P_r.T) @ Q.T  # (D_out, D_in)
    norm_orig = W_fp32.norm(p="fro").clamp_min(1e-8)
    rel_error = float((W_fp32 - W_recon).norm(p="fro") / norm_orig)

    Q_P_r = Q @ P_r  # (D_in, r) — RS-FIX-08
    return W_indices, W_centroids, Q_P_r, rel_error

Returns

-----=

W_quant_indices

(D_out, r) int8

W_quant_centroids:(C,) fp32

Q_P_r

   (D_in, r)  fp32

rel_error

float

def compress_tensor(self, W: torch.Tensor, X: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, torch.Tensor]

Expand source code

def compress_matrix(
    self,
    W: torch.Tensor,
    X: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, torch.Tensor]:
    """
    Returns:
        W_quant_indices:  (D_out, r) int8
        W_quant_centroids:(C,)       fp32
        Q_P_r:            (D_in, r)  fp32
        rel_error:        float
    """
    D_out, D_in = W.shape
    r = max(1, min(int(D_in * self.r_ratio), D_in))
    device = W.device

    # 1. Orthogonal rotation
    Q = self.get_orthogonal_matrix(D_in, device)
    W_fp32 = W.detach().float()
    X_fp32 = X.detach().float().to(device)

    W_tilde = W_fp32 @ Q  # (D_out, D_in)
    X_tilde = Q.T @ X_fp32  # (D_in, N)

    # 2. Variance-based slicing (RS-FIX-06: vectorized)
    if X_tilde.shape[1] > 1:
        channel_variances = X_tilde.var(dim=1)  # (D_in,)
    else:
        channel_variances = torch.ones(D_in, device=device)
    top_r_indices = torch.argsort(channel_variances, descending=True)[:r]

    # P_r as a (D_in, r) selector matrix via scatter
    P_r = torch.zeros(D_in, r, device=device, dtype=torch.float32)
    P_r[top_r_indices, torch.arange(r, device=device)] = 1.0

    W_sliced = W_tilde @ P_r  # (D_out, r)
    X_sliced = P_r.T @ X_tilde  # (r, N)

    # 3. Second-order Hessian-based INT4 quantization
    H = 2.0 * (X_sliced @ X_sliced.T) / max(X_sliced.shape[1], 1) + self.lambda_reg * torch.eye(
        r, device=device
    )
    try:
        H_inv = torch.linalg.inv(H)
    except Exception:
        H_inv = torch.linalg.pinv(H)
    H_inv = torch.nan_to_num(H_inv, nan=0.0)

    # RS-FIX-10: quantize the entire column of residuals at once.
    W_quant = W_sliced.clone()
    # Full-matrix OBS: quantize each column then propagate the error.
    for j in range(r):
        w_j = W_quant[:, j].clone()
        idx_j, cents_j = self.quantizer.quantize(w_j)
        q_j = AstroQuantizer.dequantize(idx_j, cents_j).to(w_j.dtype)
        e_j = w_j - q_j
        W_quant[:, j] = q_j

        diag_inv = float(H_inv[j, j].item())
        if not math.isfinite(diag_inv) or diag_inv <= 1e-8:
            diag_inv = 1e-5

        if j < r - 1:
            h_slice = torch.nan_to_num(H_inv[j, j + 1 : r] / diag_inv, nan=0.0)
            W_quant[:, j + 1 : r] -= torch.outer(e_j, h_slice)

    # Re-quantize the final W_quant column-wise to get indices + centroids
    # (OBS updates can change values slightly, so we requantize once).
    W_quant.t().contiguous()  # (r, D_out) → quantize per row
    _W_idx_list, _W_cents_list = [], []
    per_col_centroids = []
    W_indices = torch.empty_like(W_quant, dtype=torch.int8)
    for j in range(r):
        idx_j, cents_j = self.quantizer.quantize(W_quant[:, j])
        W_indices[:, j] = idx_j
        per_col_centroids.append(cents_j)
    # Use a single shared centroid table = average of per-column tables.
    W_centroids = torch.stack(per_col_centroids, dim=0).mean(dim=0)

    # Reconstruct and measure error (using the shared table for honesty).
    W_deq = AstroQuantizer.dequantize(W_indices, W_centroids).to(W_quant.dtype)
    W_recon = (W_deq @ P_r.T) @ Q.T  # (D_out, D_in)
    norm_orig = W_fp32.norm(p="fro").clamp_min(1e-8)
    rel_error = float((W_fp32 - W_recon).norm(p="fro") / norm_orig)

    Q_P_r = Q @ P_r  # (D_in, r) — RS-FIX-08
    return W_indices, W_centroids, Q_P_r, rel_error

Returns

-----=

W_quant_indices

(D_out, r) int8

W_quant_centroids:(C,) fp32

Q_P_r

   (D_in, r)  fp32

rel_error

float

def get_orthogonal_matrix(self, d_in: int, device: torch.device) ‑> torch.Tensor

Expand source code

def get_orthogonal_matrix(self, d_in: int, device: torch.device) -> torch.Tensor:
    """Deterministic QR-based orthogonal Q (RS-FIX-05)."""
    gen = torch.Generator(device=device).manual_seed(self.seed)
    H = torch.randn(d_in, d_in, device=device, dtype=torch.float32, generator=gen)
    Q, _ = torch.linalg.qr(H)
    return Q

Deterministic QR-based orthogonal Q (RS-FIX-05).

class RotationalSliceGPTINT4Compressor (model: nn.Module | None = None, r_ratio: float = 0.4, lambda_reg: float = 0.01, slice_ratio: float | None = None, quant_bits: int = 4, max_error: float = 0.2, fallback_to_original: bool = True, tokenizer: Any | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, target_modules: list[str] | None = None, target_attention_only: bool = False, seed: int = 42, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("rotational_slicegpt_int4")
@COMPRESSOR_REGISTRY.register("rotational_slicegpt")
class RotationalSliceGPTINT4Compressor(BaseCompressor):
    """
    Rotational-SliceGPT-INT4 hybrid compression engine.
    """

    # RS-FIX-04: not an attention module.
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module | None = None,
        r_ratio: float = 0.4,  # RS-FIX-03: lower default
        lambda_reg: float = 1e-2,
        slice_ratio: float | None = None,
        quant_bits: int = 4,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        tokenizer: Any | None = None,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        target_modules: list[str] | None = None,
        target_attention_only: bool = False,
        seed: int = 42,  # RS-FIX-05
        **kwargs: Any,
    ) -> None:
        # Remove duplicates that arrive through kwargs/comp_config.
        for k in (
            "max_error",
            "fallback_to_original",
            "r_ratio",
            "quant_bits",
            "lambda_reg",
        ):
            kwargs.pop(k, None)
        target_model = model if model is not None else nn.Identity()
        super().__init__(target_model, name="rotational_slicegpt_int4", **kwargs)

        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        r_ratio = _cfg_get(cfg, "rotational_slicegpt_r_ratio", _cfg_get(cfg, "r_ratio", r_ratio))
        lambda_reg = _cfg_get(cfg, "lambda_reg", lambda_reg)
        quant_bits = _cfg_get(cfg, "quant_bits", quant_bits)
        max_error = _cfg_get(cfg, "max_error", max_error)
        fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
        target_attention_only = bool(
            target_attention_only or _cfg_get(cfg, "target_attention_only", False)
        )

        self.r_ratio = float(slice_ratio if slice_ratio is not None else r_ratio)
        if self.r_ratio >= 0.5:
            logger.warning(
                "RotationalSliceGPT: r_ratio=%.2f ≥ 0.5 will INCREASE parameter count. "
                "Consider r_ratio ≤ 0.4 for real savings.",
                self.r_ratio,
            )
        self.lambda_reg = float(lambda_reg)
        self.quant_bits = int(quant_bits)
        self.quantizer = AstroQuantizer(bits=self.quant_bits)
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.target_modules = target_modules
        self.target_attention_only = target_attention_only
        self.seed = int(seed)

        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.target_device = (
            torch.device(device)
            if device is not None
            else (
                next(target_model.parameters(), torch.tensor(0.0)).device
                if model is not None
                else torch.device("cpu")
            )
        )
        self.layer_errors: dict[str, float] = {}
        self.stats: dict[str, Any] = {}
        self._orig_bytes = 0
        self._comp_bytes = 0

    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        if any(
            k in name_l
            for k in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "layernorm",
                "rmsnorm",
                "batchnorm",
            )
        ):
            return False

        if isinstance(
            module,
            (
                nn.LayerNorm,
                nn.RMSNorm,
                nn.Embedding,
                nn.BatchNorm1d,
                nn.BatchNorm2d,
                RotationalSliceGPTINT4Linear,
            ),
        ):
            return False

        if self.target_attention_only:
            attn_kw = (
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            )
            if not any(k in name_l for k in attn_kw):
                return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__.lower()
        is_linear_like = (
            isinstance(base_mod, nn.Linear) or "conv1d" in cls_name or hasattr(base_mod, "nf")
        )
        if not is_linear_like:
            return False

        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    # ------------------------------------------------------------------
    def get_orthogonal_matrix(self, d_in: int, device: torch.device) -> torch.Tensor:
        """Deterministic QR-based orthogonal Q (RS-FIX-05)."""
        gen = torch.Generator(device=device).manual_seed(self.seed)
        H = torch.randn(d_in, d_in, device=device, dtype=torch.float32, generator=gen)
        Q, _ = torch.linalg.qr(H)
        return Q

    # ------------------------------------------------------------------
    def compress_matrix(
        self,
        W: torch.Tensor,
        X: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, torch.Tensor]:
        """
        Returns:
            W_quant_indices:  (D_out, r) int8
            W_quant_centroids:(C,)       fp32
            Q_P_r:            (D_in, r)  fp32
            rel_error:        float
        """
        D_out, D_in = W.shape
        r = max(1, min(int(D_in * self.r_ratio), D_in))
        device = W.device

        # 1. Orthogonal rotation
        Q = self.get_orthogonal_matrix(D_in, device)
        W_fp32 = W.detach().float()
        X_fp32 = X.detach().float().to(device)

        W_tilde = W_fp32 @ Q  # (D_out, D_in)
        X_tilde = Q.T @ X_fp32  # (D_in, N)

        # 2. Variance-based slicing (RS-FIX-06: vectorized)
        if X_tilde.shape[1] > 1:
            channel_variances = X_tilde.var(dim=1)  # (D_in,)
        else:
            channel_variances = torch.ones(D_in, device=device)
        top_r_indices = torch.argsort(channel_variances, descending=True)[:r]

        # P_r as a (D_in, r) selector matrix via scatter
        P_r = torch.zeros(D_in, r, device=device, dtype=torch.float32)
        P_r[top_r_indices, torch.arange(r, device=device)] = 1.0

        W_sliced = W_tilde @ P_r  # (D_out, r)
        X_sliced = P_r.T @ X_tilde  # (r, N)

        # 3. Second-order Hessian-based INT4 quantization
        H = 2.0 * (X_sliced @ X_sliced.T) / max(X_sliced.shape[1], 1) + self.lambda_reg * torch.eye(
            r, device=device
        )
        try:
            H_inv = torch.linalg.inv(H)
        except Exception:
            H_inv = torch.linalg.pinv(H)
        H_inv = torch.nan_to_num(H_inv, nan=0.0)

        # RS-FIX-10: quantize the entire column of residuals at once.
        W_quant = W_sliced.clone()
        # Full-matrix OBS: quantize each column then propagate the error.
        for j in range(r):
            w_j = W_quant[:, j].clone()
            idx_j, cents_j = self.quantizer.quantize(w_j)
            q_j = AstroQuantizer.dequantize(idx_j, cents_j).to(w_j.dtype)
            e_j = w_j - q_j
            W_quant[:, j] = q_j

            diag_inv = float(H_inv[j, j].item())
            if not math.isfinite(diag_inv) or diag_inv <= 1e-8:
                diag_inv = 1e-5

            if j < r - 1:
                h_slice = torch.nan_to_num(H_inv[j, j + 1 : r] / diag_inv, nan=0.0)
                W_quant[:, j + 1 : r] -= torch.outer(e_j, h_slice)

        # Re-quantize the final W_quant column-wise to get indices + centroids
        # (OBS updates can change values slightly, so we requantize once).
        W_quant.t().contiguous()  # (r, D_out) → quantize per row
        _W_idx_list, _W_cents_list = [], []
        per_col_centroids = []
        W_indices = torch.empty_like(W_quant, dtype=torch.int8)
        for j in range(r):
            idx_j, cents_j = self.quantizer.quantize(W_quant[:, j])
            W_indices[:, j] = idx_j
            per_col_centroids.append(cents_j)
        # Use a single shared centroid table = average of per-column tables.
        W_centroids = torch.stack(per_col_centroids, dim=0).mean(dim=0)

        # Reconstruct and measure error (using the shared table for honesty).
        W_deq = AstroQuantizer.dequantize(W_indices, W_centroids).to(W_quant.dtype)
        W_recon = (W_deq @ P_r.T) @ Q.T  # (D_out, D_in)
        norm_orig = W_fp32.norm(p="fro").clamp_min(1e-8)
        rel_error = float((W_fp32 - W_recon).norm(p="fro") / norm_orig)

        Q_P_r = Q @ P_r  # (D_in, r) — RS-FIX-08
        return W_indices, W_centroids, Q_P_r, rel_error

    compress_tensor = compress_matrix  # backward-compat alias

    # ------------------------------------------------------------------
    def compress(self, *args: Any, **kwargs: Any) -> Any:
        # Dual-mode: .compress(W, X) → tuple, .compress() → (model, stats)
        if (
            len(args) >= 2
            and isinstance(args[0], torch.Tensor)
            and isinstance(args[1], torch.Tensor)
        ):
            return self.compress_matrix(args[0], args[1])

        activations_map = self._collect_layer_activations()
        compressed_layers_count = 0
        fallback_layers: list[str] = []
        self.layer_errors = {}
        layer_shapes: dict[str, tuple[int, int]] = {}
        self._orig_bytes = 0
        self._comp_bytes = 0

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            base_mod = getattr(module, "base_layer", module)
            raw_w = base_mod.weight.data
            cls_name = base_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(base_mod, "nf")
            device = raw_w.device
            dtype = raw_w.dtype

            W = raw_w.t().contiguous().float() if is_conv1d else raw_w.clone().float()
            d_out, d_in = W.shape
            layer_shapes[name] = (d_out, d_in)

            # Build X for slicing (RS-FIX-07: per-channel mean)
            act = activations_map.get(name, None)
            if act is not None and act.numel() > 0:
                act_flat = act.reshape(-1, act.shape[-1]).float()
                if act_flat.shape[-1] == d_in:
                    # Synthesize a small sample matrix from the mean vector.
                    mean_act = act_flat.mean(dim=0, keepdim=True)  # (1, D_in)
                    # Repeat with small jitter to give variance-based slicing signal.
                    X = mean_act.T + 0.01 * torch.randn(d_in, 32, device=device)
                else:
                    X = torch.randn(d_in, 32, device=device)
            else:
                X = torch.randn(d_in, 32, device=device)

            try:
                W_idx, W_cents, Q_P_r, rel_error = self.compress_matrix(W, X)
            except Exception as e:
                logger.warning("RotationalSliceGPT failed on %s: %s", name, e)
                fallback_layers.append(name)
                continue

            self.layer_errors[name] = rel_error

            if rel_error > self.max_error:
                logger.warning(
                    "RotationalSliceGPT reconstruction error %.4f on %s exceeded %.4f",
                    rel_error,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

            # Cast buffers to the original dtype & device
            W_idx = W_idx.to(device=device)
            W_cents = W_cents.to(device=device, dtype=torch.float32)
            Q_P_r_dev = Q_P_r.to(device=device, dtype=dtype)

            bias = base_mod.bias.detach() if getattr(base_mod, "bias", None) is not None else None

            new_linear = RotationalSliceGPTINT4Linear(
                W_indices=W_idx,
                W_centroids=W_cents,
                Q_P_r=Q_P_r_dev,
                bias=bias,
            )

            # Splice into model
            if hasattr(module, "base_layer"):
                # Preserve any LoRA params
                if hasattr(module, "lora_A") and module.lora_A is not None:
                    if not hasattr(new_linear, "lora_A"):
                        new_linear.lora_A = module.lora_A
                if hasattr(module, "lora_B") and module.lora_B is not None:
                    if not hasattr(new_linear, "lora_B"):
                        new_linear.lora_B = module.lora_B
                module.base_layer = new_linear
            else:
                replace_module(self.model, name, new_linear)

            # RS-FIX-09: honest byte accounting
            orig_bytes = raw_w.numel() * raw_w.element_size()
            if base_mod.bias is not None:
                orig_bytes += base_mod.bias.numel() * base_mod.bias.element_size()
            # comp: int8 indices + fp32 centroid table + Q_P_r + bias
            comp_bytes = W_idx.numel() * W_idx.element_size()
            comp_bytes += W_cents.numel() * W_cents.element_size()
            comp_bytes += Q_P_r_dev.numel() * Q_P_r_dev.element_size()
            if bias is not None:
                comp_bytes += bias.numel() * bias.element_size()

            self._orig_bytes += orig_bytes
            self._comp_bytes += comp_bytes
            compressed_layers_count += 1

        # Forward sanity check
        forward_check_success = True
        try:
            device = next(self.model.parameters()).device
            with torch.no_grad():
                cfg = getattr(self.model, "config", None)
                vocab = getattr(cfg, "vocab_size", None) or 100
                dummy = torch.randint(0, min(vocab, 100), (1, 8), device=device)
                self.model(dummy)
        except Exception as ex:
            logger.warning("RotationalSliceGPT forward sanity check: %s", ex)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        stats = self.get_compression_stats()
        avg_err = float(np.mean(list(self.layer_errors.values()))) if self.layer_errors else 0.0
        storage_ratio = self._orig_bytes / self._comp_bytes if self._comp_bytes > 0 else 1.0

        stats.update(
            {
                "compressed_layers": compressed_layers_count,
                "fallback_layers": fallback_layers,
                "r_ratio": self.r_ratio,
                "quant_bits": self.quant_bits,
                "lambda_reg": self.lambda_reg,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "layer_errors": self.layer_errors,
                "layer_shapes": layer_shapes,
                "avg_reconstruction_error": avg_err,
                "target_attention_only": self.target_attention_only,
                "forward_check_success": forward_check_success,
                "orig_storage_bytes": self._orig_bytes,
                "comp_storage_bytes": self._comp_bytes,
                "storage_compression_ratio": round(storage_ratio, 3),
                "storage_reduction_pct": round(
                    100.0 * (1.0 - self._comp_bytes / max(1, self._orig_bytes)), 2
                ),
                "method": "rotational_slicegpt_int4",
            }
        )
        self.stats = stats
        logger.info(
            "RotationalSliceGPT: %d layers compressed, %d fallback, "
            "avg_err=%.4f, storage_ratio=%.2fx, bits=%d, r_ratio=%.2f",
            compressed_layers_count,
            len(fallback_layers),
            avg_err,
            storage_ratio,
            self.quant_bits,
            self.r_ratio,
        )
        return self.model, stats

    # ------------------------------------------------------------------
    def _collect_layer_activations(self) -> dict[str, torch.Tensor]:
        """
        RS-FIX-07: store per-channel mean |x| instead of full activation
        tensors to avoid OOM on large models.
        """
        self.model.eval()
        activations: dict[str, list[torch.Tensor]] = {}
        hooks = []

        def make_hook(layer_name: str):
            def hook(module, inputs, output):
                if inputs and isinstance(inputs[0], torch.Tensor):
                    x = inputs[0].detach()
                    if x.dim() >= 2:
                        col_mean = x.abs().mean(dim=tuple(range(x.dim() - 1)))
                    else:
                        col_mean = x.abs()
                    activations.setdefault(layer_name, []).append(col_mean.cpu())

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                target = getattr(module, "base_layer", module)
                hooks.append(target.register_forward_hook(make_hook(name)))

        try:
            with torch.no_grad():
                calib_inputs = get_calibration_tensors(
                    self.model,
                    tokenizer=self.tokenizer,
                    calibration_loader=self.calibration_loader,
                    calibration_texts=self.calibration_texts,
                    device=self.target_device,
                )
                try:
                    if isinstance(calib_inputs, dict):
                        self.model(**calib_inputs)
                    elif isinstance(calib_inputs, torch.Tensor):
                        self.model(calib_inputs)
                    elif isinstance(calib_inputs, (list, tuple)):
                        self.model(*calib_inputs)
                except Exception as e:
                    logger.debug("Activation capture failed: %s", e)
        finally:
            for h in hooks:
                h.remove()

        aggregated: dict[str, torch.Tensor] = {}
        for name, lst in activations.items():
            if lst:
                aggregated[name] = torch.stack([t.float() for t in lst], dim=0).mean(dim=0)
        return aggregated

    # ------------------------------------------------------------------
    def compress_linear(
        self,
        linear_layer: nn.Linear,
        X: torch.Tensor | None = None,
    ) -> RotationalSliceGPTINT4Linear:
        """Convenience wrapper for compressing a single nn.Linear."""
        W = linear_layer.weight.data.float()
        device = W.device
        if X is None:
            X = torch.randn(W.shape[1], 128, device=device)
        elif X.dim() == 3:
            X = X.reshape(-1, W.shape[1]).T
        elif X.dim() == 2 and X.shape[0] != W.shape[1]:
            X = X.T
        W_idx, W_cents, Q_P_r, _ = self.compress_matrix(W, X)
        bias = linear_layer.bias.data if getattr(linear_layer, "bias", None) is not None else None
        return RotationalSliceGPTINT4Linear(
            W_indices=W_idx.to(device),
            W_centroids=W_cents.to(device),
            Q_P_r=Q_P_r.to(device=device, dtype=W.dtype),
            bias=bias,
        )

Rotational-SliceGPT-INT4 hybrid compression engine.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def compress(self, *args: Any, **kwargs: Any) ‑> Any

Expand source code

def compress(self, *args: Any, **kwargs: Any) -> Any:
    # Dual-mode: .compress(W, X) → tuple, .compress() → (model, stats)
    if (
        len(args) >= 2
        and isinstance(args[0], torch.Tensor)
        and isinstance(args[1], torch.Tensor)
    ):
        return self.compress_matrix(args[0], args[1])

    activations_map = self._collect_layer_activations()
    compressed_layers_count = 0
    fallback_layers: list[str] = []
    self.layer_errors = {}
    layer_shapes: dict[str, tuple[int, int]] = {}
    self._orig_bytes = 0
    self._comp_bytes = 0

    for name, module in list(self.model.named_modules()):
        if not self._is_target_module(name, module):
            continue

        base_mod = getattr(module, "base_layer", module)
        raw_w = base_mod.weight.data
        cls_name = base_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(base_mod, "nf")
        device = raw_w.device
        dtype = raw_w.dtype

        W = raw_w.t().contiguous().float() if is_conv1d else raw_w.clone().float()
        d_out, d_in = W.shape
        layer_shapes[name] = (d_out, d_in)

        # Build X for slicing (RS-FIX-07: per-channel mean)
        act = activations_map.get(name, None)
        if act is not None and act.numel() > 0:
            act_flat = act.reshape(-1, act.shape[-1]).float()
            if act_flat.shape[-1] == d_in:
                # Synthesize a small sample matrix from the mean vector.
                mean_act = act_flat.mean(dim=0, keepdim=True)  # (1, D_in)
                # Repeat with small jitter to give variance-based slicing signal.
                X = mean_act.T + 0.01 * torch.randn(d_in, 32, device=device)
            else:
                X = torch.randn(d_in, 32, device=device)
        else:
            X = torch.randn(d_in, 32, device=device)

        try:
            W_idx, W_cents, Q_P_r, rel_error = self.compress_matrix(W, X)
        except Exception as e:
            logger.warning("RotationalSliceGPT failed on %s: %s", name, e)
            fallback_layers.append(name)
            continue

        self.layer_errors[name] = rel_error

        if rel_error > self.max_error:
            logger.warning(
                "RotationalSliceGPT reconstruction error %.4f on %s exceeded %.4f",
                rel_error,
                name,
                self.max_error,
            )
            if self.fallback_to_original:
                fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                continue

        # Cast buffers to the original dtype & device
        W_idx = W_idx.to(device=device)
        W_cents = W_cents.to(device=device, dtype=torch.float32)
        Q_P_r_dev = Q_P_r.to(device=device, dtype=dtype)

        bias = base_mod.bias.detach() if getattr(base_mod, "bias", None) is not None else None

        new_linear = RotationalSliceGPTINT4Linear(
            W_indices=W_idx,
            W_centroids=W_cents,
            Q_P_r=Q_P_r_dev,
            bias=bias,
        )

        # Splice into model
        if hasattr(module, "base_layer"):
            # Preserve any LoRA params
            if hasattr(module, "lora_A") and module.lora_A is not None:
                if not hasattr(new_linear, "lora_A"):
                    new_linear.lora_A = module.lora_A
            if hasattr(module, "lora_B") and module.lora_B is not None:
                if not hasattr(new_linear, "lora_B"):
                    new_linear.lora_B = module.lora_B
            module.base_layer = new_linear
        else:
            replace_module(self.model, name, new_linear)

        # RS-FIX-09: honest byte accounting
        orig_bytes = raw_w.numel() * raw_w.element_size()
        if base_mod.bias is not None:
            orig_bytes += base_mod.bias.numel() * base_mod.bias.element_size()
        # comp: int8 indices + fp32 centroid table + Q_P_r + bias
        comp_bytes = W_idx.numel() * W_idx.element_size()
        comp_bytes += W_cents.numel() * W_cents.element_size()
        comp_bytes += Q_P_r_dev.numel() * Q_P_r_dev.element_size()
        if bias is not None:
            comp_bytes += bias.numel() * bias.element_size()

        self._orig_bytes += orig_bytes
        self._comp_bytes += comp_bytes
        compressed_layers_count += 1

    # Forward sanity check
    forward_check_success = True
    try:
        device = next(self.model.parameters()).device
        with torch.no_grad():
            cfg = getattr(self.model, "config", None)
            vocab = getattr(cfg, "vocab_size", None) or 100
            dummy = torch.randint(0, min(vocab, 100), (1, 8), device=device)
            self.model(dummy)
    except Exception as ex:
        logger.warning("RotationalSliceGPT forward sanity check: %s", ex)
        forward_check_success = False

    try:
        self._compressed_params = self._count_params(self.model)
    except Exception:
        self._compressed_params = sum(p.numel() for p in self.model.parameters())

    stats = self.get_compression_stats()
    avg_err = float(np.mean(list(self.layer_errors.values()))) if self.layer_errors else 0.0
    storage_ratio = self._orig_bytes / self._comp_bytes if self._comp_bytes > 0 else 1.0

    stats.update(
        {
            "compressed_layers": compressed_layers_count,
            "fallback_layers": fallback_layers,
            "r_ratio": self.r_ratio,
            "quant_bits": self.quant_bits,
            "lambda_reg": self.lambda_reg,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "layer_errors": self.layer_errors,
            "layer_shapes": layer_shapes,
            "avg_reconstruction_error": avg_err,
            "target_attention_only": self.target_attention_only,
            "forward_check_success": forward_check_success,
            "orig_storage_bytes": self._orig_bytes,
            "comp_storage_bytes": self._comp_bytes,
            "storage_compression_ratio": round(storage_ratio, 3),
            "storage_reduction_pct": round(
                100.0 * (1.0 - self._comp_bytes / max(1, self._orig_bytes)), 2
            ),
            "method": "rotational_slicegpt_int4",
        }
    )
    self.stats = stats
    logger.info(
        "RotationalSliceGPT: %d layers compressed, %d fallback, "
        "avg_err=%.4f, storage_ratio=%.2fx, bits=%d, r_ratio=%.2f",
        compressed_layers_count,
        len(fallback_layers),
        avg_err,
        storage_ratio,
        self.quant_bits,
        self.r_ratio,
    )
    return self.model, stats

Return (compressed_model, statistics).

def compress_linear(self, linear_layer: nn.Linear, X: torch.Tensor | None = None) ‑> RotationalSliceGPTINT4Linear

Expand source code

def compress_linear(
    self,
    linear_layer: nn.Linear,
    X: torch.Tensor | None = None,
) -> RotationalSliceGPTINT4Linear:
    """Convenience wrapper for compressing a single nn.Linear."""
    W = linear_layer.weight.data.float()
    device = W.device
    if X is None:
        X = torch.randn(W.shape[1], 128, device=device)
    elif X.dim() == 3:
        X = X.reshape(-1, W.shape[1]).T
    elif X.dim() == 2 and X.shape[0] != W.shape[1]:
        X = X.T
    W_idx, W_cents, Q_P_r, _ = self.compress_matrix(W, X)
    bias = linear_layer.bias.data if getattr(linear_layer, "bias", None) is not None else None
    return RotationalSliceGPTINT4Linear(
        W_indices=W_idx.to(device),
        W_centroids=W_cents.to(device),
        Q_P_r=Q_P_r.to(device=device, dtype=W.dtype),
        bias=bias,
    )

Convenience wrapper for compressing a single nn.Linear.

def compress_matrix(self, W: torch.Tensor, X: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, torch.Tensor]

Expand source code

def compress_matrix(
    self,
    W: torch.Tensor,
    X: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, torch.Tensor]:
    """
    Returns:
        W_quant_indices:  (D_out, r) int8
        W_quant_centroids:(C,)       fp32
        Q_P_r:            (D_in, r)  fp32
        rel_error:        float
    """
    D_out, D_in = W.shape
    r = max(1, min(int(D_in * self.r_ratio), D_in))
    device = W.device

    # 1. Orthogonal rotation
    Q = self.get_orthogonal_matrix(D_in, device)
    W_fp32 = W.detach().float()
    X_fp32 = X.detach().float().to(device)

    W_tilde = W_fp32 @ Q  # (D_out, D_in)
    X_tilde = Q.T @ X_fp32  # (D_in, N)

    # 2. Variance-based slicing (RS-FIX-06: vectorized)
    if X_tilde.shape[1] > 1:
        channel_variances = X_tilde.var(dim=1)  # (D_in,)
    else:
        channel_variances = torch.ones(D_in, device=device)
    top_r_indices = torch.argsort(channel_variances, descending=True)[:r]

    # P_r as a (D_in, r) selector matrix via scatter
    P_r = torch.zeros(D_in, r, device=device, dtype=torch.float32)
    P_r[top_r_indices, torch.arange(r, device=device)] = 1.0

    W_sliced = W_tilde @ P_r  # (D_out, r)
    X_sliced = P_r.T @ X_tilde  # (r, N)

    # 3. Second-order Hessian-based INT4 quantization
    H = 2.0 * (X_sliced @ X_sliced.T) / max(X_sliced.shape[1], 1) + self.lambda_reg * torch.eye(
        r, device=device
    )
    try:
        H_inv = torch.linalg.inv(H)
    except Exception:
        H_inv = torch.linalg.pinv(H)
    H_inv = torch.nan_to_num(H_inv, nan=0.0)

    # RS-FIX-10: quantize the entire column of residuals at once.
    W_quant = W_sliced.clone()
    # Full-matrix OBS: quantize each column then propagate the error.
    for j in range(r):
        w_j = W_quant[:, j].clone()
        idx_j, cents_j = self.quantizer.quantize(w_j)
        q_j = AstroQuantizer.dequantize(idx_j, cents_j).to(w_j.dtype)
        e_j = w_j - q_j
        W_quant[:, j] = q_j

        diag_inv = float(H_inv[j, j].item())
        if not math.isfinite(diag_inv) or diag_inv <= 1e-8:
            diag_inv = 1e-5

        if j < r - 1:
            h_slice = torch.nan_to_num(H_inv[j, j + 1 : r] / diag_inv, nan=0.0)
            W_quant[:, j + 1 : r] -= torch.outer(e_j, h_slice)

    # Re-quantize the final W_quant column-wise to get indices + centroids
    # (OBS updates can change values slightly, so we requantize once).
    W_quant.t().contiguous()  # (r, D_out) → quantize per row
    _W_idx_list, _W_cents_list = [], []
    per_col_centroids = []
    W_indices = torch.empty_like(W_quant, dtype=torch.int8)
    for j in range(r):
        idx_j, cents_j = self.quantizer.quantize(W_quant[:, j])
        W_indices[:, j] = idx_j
        per_col_centroids.append(cents_j)
    # Use a single shared centroid table = average of per-column tables.
    W_centroids = torch.stack(per_col_centroids, dim=0).mean(dim=0)

    # Reconstruct and measure error (using the shared table for honesty).
    W_deq = AstroQuantizer.dequantize(W_indices, W_centroids).to(W_quant.dtype)
    W_recon = (W_deq @ P_r.T) @ Q.T  # (D_out, D_in)
    norm_orig = W_fp32.norm(p="fro").clamp_min(1e-8)
    rel_error = float((W_fp32 - W_recon).norm(p="fro") / norm_orig)

    Q_P_r = Q @ P_r  # (D_in, r) — RS-FIX-08
    return W_indices, W_centroids, Q_P_r, rel_error

Returns

-----=

W_quant_indices

(D_out, r) int8

W_quant_centroids:(C,) fp32

Q_P_r

   (D_in, r)  fp32

rel_error

float

def compress_tensor(self, W: torch.Tensor, X: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, torch.Tensor]

Expand source code

def compress_matrix(
    self,
    W: torch.Tensor,
    X: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, torch.Tensor]:
    """
    Returns:
        W_quant_indices:  (D_out, r) int8
        W_quant_centroids:(C,)       fp32
        Q_P_r:            (D_in, r)  fp32
        rel_error:        float
    """
    D_out, D_in = W.shape
    r = max(1, min(int(D_in * self.r_ratio), D_in))
    device = W.device

    # 1. Orthogonal rotation
    Q = self.get_orthogonal_matrix(D_in, device)
    W_fp32 = W.detach().float()
    X_fp32 = X.detach().float().to(device)

    W_tilde = W_fp32 @ Q  # (D_out, D_in)
    X_tilde = Q.T @ X_fp32  # (D_in, N)

    # 2. Variance-based slicing (RS-FIX-06: vectorized)
    if X_tilde.shape[1] > 1:
        channel_variances = X_tilde.var(dim=1)  # (D_in,)
    else:
        channel_variances = torch.ones(D_in, device=device)
    top_r_indices = torch.argsort(channel_variances, descending=True)[:r]

    # P_r as a (D_in, r) selector matrix via scatter
    P_r = torch.zeros(D_in, r, device=device, dtype=torch.float32)
    P_r[top_r_indices, torch.arange(r, device=device)] = 1.0

    W_sliced = W_tilde @ P_r  # (D_out, r)
    X_sliced = P_r.T @ X_tilde  # (r, N)

    # 3. Second-order Hessian-based INT4 quantization
    H = 2.0 * (X_sliced @ X_sliced.T) / max(X_sliced.shape[1], 1) + self.lambda_reg * torch.eye(
        r, device=device
    )
    try:
        H_inv = torch.linalg.inv(H)
    except Exception:
        H_inv = torch.linalg.pinv(H)
    H_inv = torch.nan_to_num(H_inv, nan=0.0)

    # RS-FIX-10: quantize the entire column of residuals at once.
    W_quant = W_sliced.clone()
    # Full-matrix OBS: quantize each column then propagate the error.
    for j in range(r):
        w_j = W_quant[:, j].clone()
        idx_j, cents_j = self.quantizer.quantize(w_j)
        q_j = AstroQuantizer.dequantize(idx_j, cents_j).to(w_j.dtype)
        e_j = w_j - q_j
        W_quant[:, j] = q_j

        diag_inv = float(H_inv[j, j].item())
        if not math.isfinite(diag_inv) or diag_inv <= 1e-8:
            diag_inv = 1e-5

        if j < r - 1:
            h_slice = torch.nan_to_num(H_inv[j, j + 1 : r] / diag_inv, nan=0.0)
            W_quant[:, j + 1 : r] -= torch.outer(e_j, h_slice)

    # Re-quantize the final W_quant column-wise to get indices + centroids
    # (OBS updates can change values slightly, so we requantize once).
    W_quant.t().contiguous()  # (r, D_out) → quantize per row
    _W_idx_list, _W_cents_list = [], []
    per_col_centroids = []
    W_indices = torch.empty_like(W_quant, dtype=torch.int8)
    for j in range(r):
        idx_j, cents_j = self.quantizer.quantize(W_quant[:, j])
        W_indices[:, j] = idx_j
        per_col_centroids.append(cents_j)
    # Use a single shared centroid table = average of per-column tables.
    W_centroids = torch.stack(per_col_centroids, dim=0).mean(dim=0)

    # Reconstruct and measure error (using the shared table for honesty).
    W_deq = AstroQuantizer.dequantize(W_indices, W_centroids).to(W_quant.dtype)
    W_recon = (W_deq @ P_r.T) @ Q.T  # (D_out, D_in)
    norm_orig = W_fp32.norm(p="fro").clamp_min(1e-8)
    rel_error = float((W_fp32 - W_recon).norm(p="fro") / norm_orig)

    Q_P_r = Q @ P_r  # (D_in, r) — RS-FIX-08
    return W_indices, W_centroids, Q_P_r, rel_error

Returns

-----=

W_quant_indices

(D_out, r) int8

W_quant_centroids:(C,) fp32

Q_P_r

   (D_in, r)  fp32

rel_error

float

def get_orthogonal_matrix(self, d_in: int, device: torch.device) ‑> torch.Tensor

Expand source code

def get_orthogonal_matrix(self, d_in: int, device: torch.device) -> torch.Tensor:
    """Deterministic QR-based orthogonal Q (RS-FIX-05)."""
    gen = torch.Generator(device=device).manual_seed(self.seed)
    H = torch.randn(d_in, d_in, device=device, dtype=torch.float32, generator=gen)
    Q, _ = torch.linalg.qr(H)
    return Q

Deterministic QR-based orthogonal Q (RS-FIX-05).

class RotationalSliceGPTInt4Compressor (model: nn.Module | None = None, r_ratio: float = 0.4, lambda_reg: float = 0.01, slice_ratio: float | None = None, quant_bits: int = 4, max_error: float = 0.2, fallback_to_original: bool = True, tokenizer: Any | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: str | torch.device | None = None, target_modules: list[str] | None = None, target_attention_only: bool = False, seed: int = 42, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("rotational_slicegpt_int4")
@COMPRESSOR_REGISTRY.register("rotational_slicegpt")
class RotationalSliceGPTINT4Compressor(BaseCompressor):
    """
    Rotational-SliceGPT-INT4 hybrid compression engine.
    """

    # RS-FIX-04: not an attention module.
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module | None = None,
        r_ratio: float = 0.4,  # RS-FIX-03: lower default
        lambda_reg: float = 1e-2,
        slice_ratio: float | None = None,
        quant_bits: int = 4,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        tokenizer: Any | None = None,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        device: str | torch.device | None = None,
        target_modules: list[str] | None = None,
        target_attention_only: bool = False,
        seed: int = 42,  # RS-FIX-05
        **kwargs: Any,
    ) -> None:
        # Remove duplicates that arrive through kwargs/comp_config.
        for k in (
            "max_error",
            "fallback_to_original",
            "r_ratio",
            "quant_bits",
            "lambda_reg",
        ):
            kwargs.pop(k, None)
        target_model = model if model is not None else nn.Identity()
        super().__init__(target_model, name="rotational_slicegpt_int4", **kwargs)

        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        r_ratio = _cfg_get(cfg, "rotational_slicegpt_r_ratio", _cfg_get(cfg, "r_ratio", r_ratio))
        lambda_reg = _cfg_get(cfg, "lambda_reg", lambda_reg)
        quant_bits = _cfg_get(cfg, "quant_bits", quant_bits)
        max_error = _cfg_get(cfg, "max_error", max_error)
        fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
        target_attention_only = bool(
            target_attention_only or _cfg_get(cfg, "target_attention_only", False)
        )

        self.r_ratio = float(slice_ratio if slice_ratio is not None else r_ratio)
        if self.r_ratio >= 0.5:
            logger.warning(
                "RotationalSliceGPT: r_ratio=%.2f ≥ 0.5 will INCREASE parameter count. "
                "Consider r_ratio ≤ 0.4 for real savings.",
                self.r_ratio,
            )
        self.lambda_reg = float(lambda_reg)
        self.quant_bits = int(quant_bits)
        self.quantizer = AstroQuantizer(bits=self.quant_bits)
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.target_modules = target_modules
        self.target_attention_only = target_attention_only
        self.seed = int(seed)

        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.target_device = (
            torch.device(device)
            if device is not None
            else (
                next(target_model.parameters(), torch.tensor(0.0)).device
                if model is not None
                else torch.device("cpu")
            )
        )
        self.layer_errors: dict[str, float] = {}
        self.stats: dict[str, Any] = {}
        self._orig_bytes = 0
        self._comp_bytes = 0

    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        if any(
            k in name_l
            for k in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "layernorm",
                "rmsnorm",
                "batchnorm",
            )
        ):
            return False

        if isinstance(
            module,
            (
                nn.LayerNorm,
                nn.RMSNorm,
                nn.Embedding,
                nn.BatchNorm1d,
                nn.BatchNorm2d,
                RotationalSliceGPTINT4Linear,
            ),
        ):
            return False

        if self.target_attention_only:
            attn_kw = (
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            )
            if not any(k in name_l for k in attn_kw):
                return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__.lower()
        is_linear_like = (
            isinstance(base_mod, nn.Linear) or "conv1d" in cls_name or hasattr(base_mod, "nf")
        )
        if not is_linear_like:
            return False

        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    # ------------------------------------------------------------------
    def get_orthogonal_matrix(self, d_in: int, device: torch.device) -> torch.Tensor:
        """Deterministic QR-based orthogonal Q (RS-FIX-05)."""
        gen = torch.Generator(device=device).manual_seed(self.seed)
        H = torch.randn(d_in, d_in, device=device, dtype=torch.float32, generator=gen)
        Q, _ = torch.linalg.qr(H)
        return Q

    # ------------------------------------------------------------------
    def compress_matrix(
        self,
        W: torch.Tensor,
        X: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, torch.Tensor]:
        """
        Returns:
            W_quant_indices:  (D_out, r) int8
            W_quant_centroids:(C,)       fp32
            Q_P_r:            (D_in, r)  fp32
            rel_error:        float
        """
        D_out, D_in = W.shape
        r = max(1, min(int(D_in * self.r_ratio), D_in))
        device = W.device

        # 1. Orthogonal rotation
        Q = self.get_orthogonal_matrix(D_in, device)
        W_fp32 = W.detach().float()
        X_fp32 = X.detach().float().to(device)

        W_tilde = W_fp32 @ Q  # (D_out, D_in)
        X_tilde = Q.T @ X_fp32  # (D_in, N)

        # 2. Variance-based slicing (RS-FIX-06: vectorized)
        if X_tilde.shape[1] > 1:
            channel_variances = X_tilde.var(dim=1)  # (D_in,)
        else:
            channel_variances = torch.ones(D_in, device=device)
        top_r_indices = torch.argsort(channel_variances, descending=True)[:r]

        # P_r as a (D_in, r) selector matrix via scatter
        P_r = torch.zeros(D_in, r, device=device, dtype=torch.float32)
        P_r[top_r_indices, torch.arange(r, device=device)] = 1.0

        W_sliced = W_tilde @ P_r  # (D_out, r)
        X_sliced = P_r.T @ X_tilde  # (r, N)

        # 3. Second-order Hessian-based INT4 quantization
        H = 2.0 * (X_sliced @ X_sliced.T) / max(X_sliced.shape[1], 1) + self.lambda_reg * torch.eye(
            r, device=device
        )
        try:
            H_inv = torch.linalg.inv(H)
        except Exception:
            H_inv = torch.linalg.pinv(H)
        H_inv = torch.nan_to_num(H_inv, nan=0.0)

        # RS-FIX-10: quantize the entire column of residuals at once.
        W_quant = W_sliced.clone()
        # Full-matrix OBS: quantize each column then propagate the error.
        for j in range(r):
            w_j = W_quant[:, j].clone()
            idx_j, cents_j = self.quantizer.quantize(w_j)
            q_j = AstroQuantizer.dequantize(idx_j, cents_j).to(w_j.dtype)
            e_j = w_j - q_j
            W_quant[:, j] = q_j

            diag_inv = float(H_inv[j, j].item())
            if not math.isfinite(diag_inv) or diag_inv <= 1e-8:
                diag_inv = 1e-5

            if j < r - 1:
                h_slice = torch.nan_to_num(H_inv[j, j + 1 : r] / diag_inv, nan=0.0)
                W_quant[:, j + 1 : r] -= torch.outer(e_j, h_slice)

        # Re-quantize the final W_quant column-wise to get indices + centroids
        # (OBS updates can change values slightly, so we requantize once).
        W_quant.t().contiguous()  # (r, D_out) → quantize per row
        _W_idx_list, _W_cents_list = [], []
        per_col_centroids = []
        W_indices = torch.empty_like(W_quant, dtype=torch.int8)
        for j in range(r):
            idx_j, cents_j = self.quantizer.quantize(W_quant[:, j])
            W_indices[:, j] = idx_j
            per_col_centroids.append(cents_j)
        # Use a single shared centroid table = average of per-column tables.
        W_centroids = torch.stack(per_col_centroids, dim=0).mean(dim=0)

        # Reconstruct and measure error (using the shared table for honesty).
        W_deq = AstroQuantizer.dequantize(W_indices, W_centroids).to(W_quant.dtype)
        W_recon = (W_deq @ P_r.T) @ Q.T  # (D_out, D_in)
        norm_orig = W_fp32.norm(p="fro").clamp_min(1e-8)
        rel_error = float((W_fp32 - W_recon).norm(p="fro") / norm_orig)

        Q_P_r = Q @ P_r  # (D_in, r) — RS-FIX-08
        return W_indices, W_centroids, Q_P_r, rel_error

    compress_tensor = compress_matrix  # backward-compat alias

    # ------------------------------------------------------------------
    def compress(self, *args: Any, **kwargs: Any) -> Any:
        # Dual-mode: .compress(W, X) → tuple, .compress() → (model, stats)
        if (
            len(args) >= 2
            and isinstance(args[0], torch.Tensor)
            and isinstance(args[1], torch.Tensor)
        ):
            return self.compress_matrix(args[0], args[1])

        activations_map = self._collect_layer_activations()
        compressed_layers_count = 0
        fallback_layers: list[str] = []
        self.layer_errors = {}
        layer_shapes: dict[str, tuple[int, int]] = {}
        self._orig_bytes = 0
        self._comp_bytes = 0

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            base_mod = getattr(module, "base_layer", module)
            raw_w = base_mod.weight.data
            cls_name = base_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(base_mod, "nf")
            device = raw_w.device
            dtype = raw_w.dtype

            W = raw_w.t().contiguous().float() if is_conv1d else raw_w.clone().float()
            d_out, d_in = W.shape
            layer_shapes[name] = (d_out, d_in)

            # Build X for slicing (RS-FIX-07: per-channel mean)
            act = activations_map.get(name, None)
            if act is not None and act.numel() > 0:
                act_flat = act.reshape(-1, act.shape[-1]).float()
                if act_flat.shape[-1] == d_in:
                    # Synthesize a small sample matrix from the mean vector.
                    mean_act = act_flat.mean(dim=0, keepdim=True)  # (1, D_in)
                    # Repeat with small jitter to give variance-based slicing signal.
                    X = mean_act.T + 0.01 * torch.randn(d_in, 32, device=device)
                else:
                    X = torch.randn(d_in, 32, device=device)
            else:
                X = torch.randn(d_in, 32, device=device)

            try:
                W_idx, W_cents, Q_P_r, rel_error = self.compress_matrix(W, X)
            except Exception as e:
                logger.warning("RotationalSliceGPT failed on %s: %s", name, e)
                fallback_layers.append(name)
                continue

            self.layer_errors[name] = rel_error

            if rel_error > self.max_error:
                logger.warning(
                    "RotationalSliceGPT reconstruction error %.4f on %s exceeded %.4f",
                    rel_error,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

            # Cast buffers to the original dtype & device
            W_idx = W_idx.to(device=device)
            W_cents = W_cents.to(device=device, dtype=torch.float32)
            Q_P_r_dev = Q_P_r.to(device=device, dtype=dtype)

            bias = base_mod.bias.detach() if getattr(base_mod, "bias", None) is not None else None

            new_linear = RotationalSliceGPTINT4Linear(
                W_indices=W_idx,
                W_centroids=W_cents,
                Q_P_r=Q_P_r_dev,
                bias=bias,
            )

            # Splice into model
            if hasattr(module, "base_layer"):
                # Preserve any LoRA params
                if hasattr(module, "lora_A") and module.lora_A is not None:
                    if not hasattr(new_linear, "lora_A"):
                        new_linear.lora_A = module.lora_A
                if hasattr(module, "lora_B") and module.lora_B is not None:
                    if not hasattr(new_linear, "lora_B"):
                        new_linear.lora_B = module.lora_B
                module.base_layer = new_linear
            else:
                replace_module(self.model, name, new_linear)

            # RS-FIX-09: honest byte accounting
            orig_bytes = raw_w.numel() * raw_w.element_size()
            if base_mod.bias is not None:
                orig_bytes += base_mod.bias.numel() * base_mod.bias.element_size()
            # comp: int8 indices + fp32 centroid table + Q_P_r + bias
            comp_bytes = W_idx.numel() * W_idx.element_size()
            comp_bytes += W_cents.numel() * W_cents.element_size()
            comp_bytes += Q_P_r_dev.numel() * Q_P_r_dev.element_size()
            if bias is not None:
                comp_bytes += bias.numel() * bias.element_size()

            self._orig_bytes += orig_bytes
            self._comp_bytes += comp_bytes
            compressed_layers_count += 1

        # Forward sanity check
        forward_check_success = True
        try:
            device = next(self.model.parameters()).device
            with torch.no_grad():
                cfg = getattr(self.model, "config", None)
                vocab = getattr(cfg, "vocab_size", None) or 100
                dummy = torch.randint(0, min(vocab, 100), (1, 8), device=device)
                self.model(dummy)
        except Exception as ex:
            logger.warning("RotationalSliceGPT forward sanity check: %s", ex)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        stats = self.get_compression_stats()
        avg_err = float(np.mean(list(self.layer_errors.values()))) if self.layer_errors else 0.0
        storage_ratio = self._orig_bytes / self._comp_bytes if self._comp_bytes > 0 else 1.0

        stats.update(
            {
                "compressed_layers": compressed_layers_count,
                "fallback_layers": fallback_layers,
                "r_ratio": self.r_ratio,
                "quant_bits": self.quant_bits,
                "lambda_reg": self.lambda_reg,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "layer_errors": self.layer_errors,
                "layer_shapes": layer_shapes,
                "avg_reconstruction_error": avg_err,
                "target_attention_only": self.target_attention_only,
                "forward_check_success": forward_check_success,
                "orig_storage_bytes": self._orig_bytes,
                "comp_storage_bytes": self._comp_bytes,
                "storage_compression_ratio": round(storage_ratio, 3),
                "storage_reduction_pct": round(
                    100.0 * (1.0 - self._comp_bytes / max(1, self._orig_bytes)), 2
                ),
                "method": "rotational_slicegpt_int4",
            }
        )
        self.stats = stats
        logger.info(
            "RotationalSliceGPT: %d layers compressed, %d fallback, "
            "avg_err=%.4f, storage_ratio=%.2fx, bits=%d, r_ratio=%.2f",
            compressed_layers_count,
            len(fallback_layers),
            avg_err,
            storage_ratio,
            self.quant_bits,
            self.r_ratio,
        )
        return self.model, stats

    # ------------------------------------------------------------------
    def _collect_layer_activations(self) -> dict[str, torch.Tensor]:
        """
        RS-FIX-07: store per-channel mean |x| instead of full activation
        tensors to avoid OOM on large models.
        """
        self.model.eval()
        activations: dict[str, list[torch.Tensor]] = {}
        hooks = []

        def make_hook(layer_name: str):
            def hook(module, inputs, output):
                if inputs and isinstance(inputs[0], torch.Tensor):
                    x = inputs[0].detach()
                    if x.dim() >= 2:
                        col_mean = x.abs().mean(dim=tuple(range(x.dim() - 1)))
                    else:
                        col_mean = x.abs()
                    activations.setdefault(layer_name, []).append(col_mean.cpu())

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                target = getattr(module, "base_layer", module)
                hooks.append(target.register_forward_hook(make_hook(name)))

        try:
            with torch.no_grad():
                calib_inputs = get_calibration_tensors(
                    self.model,
                    tokenizer=self.tokenizer,
                    calibration_loader=self.calibration_loader,
                    calibration_texts=self.calibration_texts,
                    device=self.target_device,
                )
                try:
                    if isinstance(calib_inputs, dict):
                        self.model(**calib_inputs)
                    elif isinstance(calib_inputs, torch.Tensor):
                        self.model(calib_inputs)
                    elif isinstance(calib_inputs, (list, tuple)):
                        self.model(*calib_inputs)
                except Exception as e:
                    logger.debug("Activation capture failed: %s", e)
        finally:
            for h in hooks:
                h.remove()

        aggregated: dict[str, torch.Tensor] = {}
        for name, lst in activations.items():
            if lst:
                aggregated[name] = torch.stack([t.float() for t in lst], dim=0).mean(dim=0)
        return aggregated

    # ------------------------------------------------------------------
    def compress_linear(
        self,
        linear_layer: nn.Linear,
        X: torch.Tensor | None = None,
    ) -> RotationalSliceGPTINT4Linear:
        """Convenience wrapper for compressing a single nn.Linear."""
        W = linear_layer.weight.data.float()
        device = W.device
        if X is None:
            X = torch.randn(W.shape[1], 128, device=device)
        elif X.dim() == 3:
            X = X.reshape(-1, W.shape[1]).T
        elif X.dim() == 2 and X.shape[0] != W.shape[1]:
            X = X.T
        W_idx, W_cents, Q_P_r, _ = self.compress_matrix(W, X)
        bias = linear_layer.bias.data if getattr(linear_layer, "bias", None) is not None else None
        return RotationalSliceGPTINT4Linear(
            W_indices=W_idx.to(device),
            W_centroids=W_cents.to(device),
            Q_P_r=Q_P_r.to(device=device, dtype=W.dtype),
            bias=bias,
        )

Rotational-SliceGPT-INT4 hybrid compression engine.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def compress_linear(self, linear_layer: nn.Linear, X: torch.Tensor | None = None) ‑> RotationalSliceGPTINT4Linear

Expand source code

def compress_linear(
    self,
    linear_layer: nn.Linear,
    X: torch.Tensor | None = None,
) -> RotationalSliceGPTINT4Linear:
    """Convenience wrapper for compressing a single nn.Linear."""
    W = linear_layer.weight.data.float()
    device = W.device
    if X is None:
        X = torch.randn(W.shape[1], 128, device=device)
    elif X.dim() == 3:
        X = X.reshape(-1, W.shape[1]).T
    elif X.dim() == 2 and X.shape[0] != W.shape[1]:
        X = X.T
    W_idx, W_cents, Q_P_r, _ = self.compress_matrix(W, X)
    bias = linear_layer.bias.data if getattr(linear_layer, "bias", None) is not None else None
    return RotationalSliceGPTINT4Linear(
        W_indices=W_idx.to(device),
        W_centroids=W_cents.to(device),
        Q_P_r=Q_P_r.to(device=device, dtype=W.dtype),
        bias=bias,
    )

Convenience wrapper for compressing a single nn.Linear.

def compress_matrix(self, W: torch.Tensor, X: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, torch.Tensor]

Expand source code

def compress_matrix(
    self,
    W: torch.Tensor,
    X: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, torch.Tensor]:
    """
    Returns:
        W_quant_indices:  (D_out, r) int8
        W_quant_centroids:(C,)       fp32
        Q_P_r:            (D_in, r)  fp32
        rel_error:        float
    """
    D_out, D_in = W.shape
    r = max(1, min(int(D_in * self.r_ratio), D_in))
    device = W.device

    # 1. Orthogonal rotation
    Q = self.get_orthogonal_matrix(D_in, device)
    W_fp32 = W.detach().float()
    X_fp32 = X.detach().float().to(device)

    W_tilde = W_fp32 @ Q  # (D_out, D_in)
    X_tilde = Q.T @ X_fp32  # (D_in, N)

    # 2. Variance-based slicing (RS-FIX-06: vectorized)
    if X_tilde.shape[1] > 1:
        channel_variances = X_tilde.var(dim=1)  # (D_in,)
    else:
        channel_variances = torch.ones(D_in, device=device)
    top_r_indices = torch.argsort(channel_variances, descending=True)[:r]

    # P_r as a (D_in, r) selector matrix via scatter
    P_r = torch.zeros(D_in, r, device=device, dtype=torch.float32)
    P_r[top_r_indices, torch.arange(r, device=device)] = 1.0

    W_sliced = W_tilde @ P_r  # (D_out, r)
    X_sliced = P_r.T @ X_tilde  # (r, N)

    # 3. Second-order Hessian-based INT4 quantization
    H = 2.0 * (X_sliced @ X_sliced.T) / max(X_sliced.shape[1], 1) + self.lambda_reg * torch.eye(
        r, device=device
    )
    try:
        H_inv = torch.linalg.inv(H)
    except Exception:
        H_inv = torch.linalg.pinv(H)
    H_inv = torch.nan_to_num(H_inv, nan=0.0)

    # RS-FIX-10: quantize the entire column of residuals at once.
    W_quant = W_sliced.clone()
    # Full-matrix OBS: quantize each column then propagate the error.
    for j in range(r):
        w_j = W_quant[:, j].clone()
        idx_j, cents_j = self.quantizer.quantize(w_j)
        q_j = AstroQuantizer.dequantize(idx_j, cents_j).to(w_j.dtype)
        e_j = w_j - q_j
        W_quant[:, j] = q_j

        diag_inv = float(H_inv[j, j].item())
        if not math.isfinite(diag_inv) or diag_inv <= 1e-8:
            diag_inv = 1e-5

        if j < r - 1:
            h_slice = torch.nan_to_num(H_inv[j, j + 1 : r] / diag_inv, nan=0.0)
            W_quant[:, j + 1 : r] -= torch.outer(e_j, h_slice)

    # Re-quantize the final W_quant column-wise to get indices + centroids
    # (OBS updates can change values slightly, so we requantize once).
    W_quant.t().contiguous()  # (r, D_out) → quantize per row
    _W_idx_list, _W_cents_list = [], []
    per_col_centroids = []
    W_indices = torch.empty_like(W_quant, dtype=torch.int8)
    for j in range(r):
        idx_j, cents_j = self.quantizer.quantize(W_quant[:, j])
        W_indices[:, j] = idx_j
        per_col_centroids.append(cents_j)
    # Use a single shared centroid table = average of per-column tables.
    W_centroids = torch.stack(per_col_centroids, dim=0).mean(dim=0)

    # Reconstruct and measure error (using the shared table for honesty).
    W_deq = AstroQuantizer.dequantize(W_indices, W_centroids).to(W_quant.dtype)
    W_recon = (W_deq @ P_r.T) @ Q.T  # (D_out, D_in)
    norm_orig = W_fp32.norm(p="fro").clamp_min(1e-8)
    rel_error = float((W_fp32 - W_recon).norm(p="fro") / norm_orig)

    Q_P_r = Q @ P_r  # (D_in, r) — RS-FIX-08
    return W_indices, W_centroids, Q_P_r, rel_error

Returns

-----=

W_quant_indices

(D_out, r) int8

W_quant_centroids:(C,) fp32

Q_P_r

   (D_in, r)  fp32

rel_error

float

def compress_tensor(self, W: torch.Tensor, X: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, torch.Tensor]

Expand source code

def compress_matrix(
    self,
    W: torch.Tensor,
    X: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, torch.Tensor]:
    """
    Returns:
        W_quant_indices:  (D_out, r) int8
        W_quant_centroids:(C,)       fp32
        Q_P_r:            (D_in, r)  fp32
        rel_error:        float
    """
    D_out, D_in = W.shape
    r = max(1, min(int(D_in * self.r_ratio), D_in))
    device = W.device

    # 1. Orthogonal rotation
    Q = self.get_orthogonal_matrix(D_in, device)
    W_fp32 = W.detach().float()
    X_fp32 = X.detach().float().to(device)

    W_tilde = W_fp32 @ Q  # (D_out, D_in)
    X_tilde = Q.T @ X_fp32  # (D_in, N)

    # 2. Variance-based slicing (RS-FIX-06: vectorized)
    if X_tilde.shape[1] > 1:
        channel_variances = X_tilde.var(dim=1)  # (D_in,)
    else:
        channel_variances = torch.ones(D_in, device=device)
    top_r_indices = torch.argsort(channel_variances, descending=True)[:r]

    # P_r as a (D_in, r) selector matrix via scatter
    P_r = torch.zeros(D_in, r, device=device, dtype=torch.float32)
    P_r[top_r_indices, torch.arange(r, device=device)] = 1.0

    W_sliced = W_tilde @ P_r  # (D_out, r)
    X_sliced = P_r.T @ X_tilde  # (r, N)

    # 3. Second-order Hessian-based INT4 quantization
    H = 2.0 * (X_sliced @ X_sliced.T) / max(X_sliced.shape[1], 1) + self.lambda_reg * torch.eye(
        r, device=device
    )
    try:
        H_inv = torch.linalg.inv(H)
    except Exception:
        H_inv = torch.linalg.pinv(H)
    H_inv = torch.nan_to_num(H_inv, nan=0.0)

    # RS-FIX-10: quantize the entire column of residuals at once.
    W_quant = W_sliced.clone()
    # Full-matrix OBS: quantize each column then propagate the error.
    for j in range(r):
        w_j = W_quant[:, j].clone()
        idx_j, cents_j = self.quantizer.quantize(w_j)
        q_j = AstroQuantizer.dequantize(idx_j, cents_j).to(w_j.dtype)
        e_j = w_j - q_j
        W_quant[:, j] = q_j

        diag_inv = float(H_inv[j, j].item())
        if not math.isfinite(diag_inv) or diag_inv <= 1e-8:
            diag_inv = 1e-5

        if j < r - 1:
            h_slice = torch.nan_to_num(H_inv[j, j + 1 : r] / diag_inv, nan=0.0)
            W_quant[:, j + 1 : r] -= torch.outer(e_j, h_slice)

    # Re-quantize the final W_quant column-wise to get indices + centroids
    # (OBS updates can change values slightly, so we requantize once).
    W_quant.t().contiguous()  # (r, D_out) → quantize per row
    _W_idx_list, _W_cents_list = [], []
    per_col_centroids = []
    W_indices = torch.empty_like(W_quant, dtype=torch.int8)
    for j in range(r):
        idx_j, cents_j = self.quantizer.quantize(W_quant[:, j])
        W_indices[:, j] = idx_j
        per_col_centroids.append(cents_j)
    # Use a single shared centroid table = average of per-column tables.
    W_centroids = torch.stack(per_col_centroids, dim=0).mean(dim=0)

    # Reconstruct and measure error (using the shared table for honesty).
    W_deq = AstroQuantizer.dequantize(W_indices, W_centroids).to(W_quant.dtype)
    W_recon = (W_deq @ P_r.T) @ Q.T  # (D_out, D_in)
    norm_orig = W_fp32.norm(p="fro").clamp_min(1e-8)
    rel_error = float((W_fp32 - W_recon).norm(p="fro") / norm_orig)

    Q_P_r = Q @ P_r  # (D_in, r) — RS-FIX-08
    return W_indices, W_centroids, Q_P_r, rel_error

Returns

-----=

W_quant_indices

(D_out, r) int8

W_quant_centroids:(C,) fp32

Q_P_r

   (D_in, r)  fp32

rel_error

float

def get_orthogonal_matrix(self, d_in: int, device: torch.device) ‑> torch.Tensor

Expand source code

def get_orthogonal_matrix(self, d_in: int, device: torch.device) -> torch.Tensor:
    """Deterministic QR-based orthogonal Q (RS-FIX-05)."""
    gen = torch.Generator(device=device).manual_seed(self.seed)
    H = torch.randn(d_in, d_in, device=device, dtype=torch.float32, generator=gen)
    Q, _ = torch.linalg.qr(H)
    return Q

Deterministic QR-based orthogonal Q (RS-FIX-05).

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class RotationalSliceGPTINT4Linear (W_indices: torch.Tensor, W_centroids: torch.Tensor, Q_P_r: torch.Tensor, bias: torch.Tensor | None = None)

Expand source code

class RotationalSliceGPTINT4Linear(nn.Module):
    """
    Linear layer storing:
      - Q_P_r:        (D_in, r) fp16/32 — the projection matrix
      - W_indices:    (D_out, r) int8   — quantized weight in sliced space
      - W_centroids:  (C,)       fp32   — per-tensor centroid table
      - bias:         (D_out,)   or None

    Forward: y = (x @ Q_P_r) → dequantize(W_indices) → @ .T → + bias
    """

    is_custom_attention: bool = False

    def __init__(
        self,
        W_indices: torch.Tensor,  # (D_out, r) int8
        W_centroids: torch.Tensor,  # (C,) fp32
        Q_P_r: torch.Tensor,  # (D_in, r) float
        bias: torch.Tensor | None = None,
    ) -> None:
        super().__init__()
        self.register_buffer("Q_P_r", Q_P_r)
        self.register_buffer("W_indices", W_indices)
        self.register_buffer("W_centroids", W_centroids)

        if bias is not None:
            self.bias = nn.Parameter(bias.clone(), requires_grad=False)
        else:
            self.register_parameter("bias", None)

    def _dequantize_W(self) -> torch.Tensor:
        return AstroQuantizer.dequantize(self.W_indices, self.W_centroids).to(self.Q_P_r.dtype)

    @property
    def weight(self) -> torch.Tensor:
        """Effective weight matrix (out_features, in_features)."""
        W_q = self._dequantize_W()  # (D_out, r)
        return torch.matmul(W_q, self.Q_P_r.T)  # (D_out, D_in)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (..., D_in)
        x_proj = torch.matmul(x, self.Q_P_r)  # (..., r)
        W_q = self._dequantize_W()  # (D_out, r)
        y = torch.matmul(x_proj, W_q.T)  # (..., D_out)
        if self.bias is not None:
            y = y + self.bias
        return y

Linear layer storing: - Q_P_r: (D_in, r) fp16/32 — the projection matrix - W_indices: (D_out, r) int8 — quantized weight in sliced space - W_centroids: (C,) fp32 — per-tensor centroid table - bias: (D_out,) or None

Forward: y = (x @ Q_P_r) → dequantize(W_indices) → @ .T → + bias

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Effective weight matrix (out_features, in_features)."""
    W_q = self._dequantize_W()  # (D_out, r)
    return torch.matmul(W_q, self.Q_P_r.T)  # (D_out, D_in)

Effective weight matrix (out_features, in_features).

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    # x: (..., D_in)
    x_proj = torch.matmul(x, self.Q_P_r)  # (..., r)
    W_q = self._dequantize_W()  # (D_out, r)
    y = torch.matmul(x_proj, W_q.T)  # (..., D_out)
    if self.bias is not None:
        y = y + self.bias
    return y

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class SAESCompressor (model: nn.Module, rank: int | None = None, rank_ratio: float = 0.5, correcting_steps: int = 2, target_modules: list[str] | None = None, target_attention_only: bool = False, max_error: float = 0.6, fallback_to_original: bool = True, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("saes_svd_v2")
@COMPRESSOR_REGISTRY.register("saes_svd_compressor")
@COMPRESSOR_REGISTRY.register("saes_svd")
@COMPRESSOR_REGISTRY.register("saes")
class SAESSVDCompressorV2(BaseCompressor):
    is_custom_attention: bool = True
    """
    SAES-SVD v2 Compressor with iterative spectrum-aware residual correction,
    Frobenius validation, and safe fallback.
    """

    def __init__(
        self,
        model: nn.Module,
        rank: int | None = None,
        rank_ratio: float = 0.5,
        correcting_steps: int = 2,
        target_modules: list[str] | None = None,
        target_attention_only: bool = False,
        max_error: float = 0.60,
        fallback_to_original: bool = True,  # Fix 1: Default changed to True
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="saes_svd_v2", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            rank = _cfg_get(cfg, "saes_rank", _cfg_get(cfg, "rank", rank))
            rank_ratio = _cfg_get(cfg, "saes_rank_ratio", _cfg_get(cfg, "rank_ratio", rank_ratio))
            correcting_steps = _cfg_get(cfg, "correcting_steps", correcting_steps)
            max_error = _cfg_get(cfg, "max_error", max_error)
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)

        if "rank" in kwargs and kwargs["rank"] is not None:
            rank = kwargs["rank"]
        if "fallback_to_original" in kwargs:
            fallback_to_original = bool(kwargs["fallback_to_original"])

        self.rank = rank
        self.rank_ratio = float(rank_ratio)
        self.correcting_steps = int(correcting_steps)
        self.target_modules = target_modules
        self.target_attention_only = bool(target_attention_only)
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.replaced_layers: list[dict[str, Any]] = []

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        """
        Fix 6: Unified critical-layer protection delegated to should_compress_linear.
        """
        name_l = name.lower()
        if self.target_attention_only:
            attn_keywords = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(k in name_l for k in attn_keywords):
                return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__.lower()
        # Fix 3: Reliable Conv1D check using class name and hasattr(nf)
        is_linear_like = (
            isinstance(base_mod, nn.Linear) or "conv1d" in cls_name or hasattr(base_mod, "nf")
        )
        if not is_linear_like:
            return False
        if isinstance(base_mod, SVDLinear):
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    def compress_linear(
        self,
        layer: nn.Module,
        rank: int | None = None,
        name: str = "",
    ) -> SVDLinear | None:
        """Decompose a single linear or Conv1D layer with spectrum-aware error suppression."""
        base_layer = getattr(layer, "base_layer", layer)
        raw_w = base_layer.weight.data
        cls_name = base_layer.__class__.__name__.lower()
        # Fix 3: Conv1D detection by class name or nf attribute
        is_conv1d = "conv1d" in cls_name or hasattr(base_layer, "nf")

        if is_conv1d:
            W = raw_w.t().contiguous().float()
            out_f = W.shape[0]
            in_f = W.shape[1]
        else:
            W = raw_w.float()
            out_f = getattr(base_layer, "out_features", W.shape[0])
            in_f = getattr(base_layer, "in_features", W.shape[1])

        max_r = min(out_f, in_f)
        # Cap rank to guarantee parameter reduction
        param_cap = max(1, (in_f * out_f) // (in_f + out_f) - 1)
        effective_max_r = min(max_r, param_cap)

        if rank is not None:
            r = max(1, min(rank, effective_max_r))
        else:
            r = max(1, min(int(round(max_r * self.rank_ratio)), effective_max_r))

        device = base_layer.weight.device
        dtype = base_layer.weight.dtype

        # Compute full SVD in FP32
        try:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        except Exception as e:
            logger.warning("SAES-SVD SVD failure for %s: %s", name, e)
            return None

        # Initial symmetric square-root singular value allocation
        sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
        A = U[:, :r] * sqrt_S.unsqueeze(0)  # [out_features, r]
        B = sqrt_S.unsqueeze(1) * Vt[:r, :]  # [r, in_features]

        # Iterative residual error suppression
        eps_eye = 1e-6 * torch.eye(r, device=W.device, dtype=torch.float32)
        for _ in range(self.correcting_steps):
            E = W - A @ B
            # Update A towards residual gradient
            BBt = B @ B.t().contiguous() + eps_eye
            A = A + E @ B.t().contiguous() @ torch.linalg.pinv(BBt)

            # Update B towards residual gradient
            AtA = A.t().contiguous() @ A + eps_eye
            B = B + torch.linalg.pinv(AtA) @ A.t().contiguous() @ (W - A @ B)

        # Compute relative Frobenius reconstruction error
        W_rec = A @ B
        norm_orig = torch.norm(W, p="fro")
        rel_err = (torch.norm(W - W_rec, p="fro") / (norm_orig + 1e-8)).item()

        # Fix 2: Always log a warning when reconstruction error exceeds max_error
        if rel_err > self.max_error:
            logger.warning(
                "SAES-SVD reconstruction error %.4f on %s exceeded threshold %.4f%s",
                rel_err,
                name,
                self.max_error,
                (
                    ". Keeping original layer."
                    if self.fallback_to_original
                    else ". Proceeding with compressed layer."
                ),
            )
            if self.fallback_to_original:
                return None

        has_bias = getattr(base_layer, "bias", None) is not None
        has_lora = hasattr(layer, "lora_A") and hasattr(layer, "lora_B")
        lora_r = getattr(layer, "r", 0) if has_lora else 0

        svd_layer = SVDLinear(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=has_bias,
            dtype=dtype,
            device=device,
            lora_r=lora_r,
        )
        svd_layer.set_low_rank_factors(
            U_factor=A.to(dtype=dtype, device=device),
            V_factor=B.to(dtype=dtype, device=device),
        )
        # Fix 5: Ensure bias is copied correctly for both Linear and Conv1D
        if has_bias and svd_layer.bias is not None:
            svd_layer.bias.data.copy_(base_layer.bias.data.to(dtype=dtype, device=device))

        if has_lora and svd_layer.use_lora:
            try:
                if hasattr(layer.lora_A, "weight"):
                    svd_layer.lora_A.weight.data.copy_(layer.lora_A.weight.data)
                if hasattr(layer.lora_B, "weight"):
                    svd_layer.lora_B.weight.data.copy_(layer.lora_B.weight.data)
            except Exception as e_lora:
                logger.warning("Could not preserve LoRA adapters in SAES-SVD: %s", e_lora)

        # Store metadata on module
        svd_layer.reconstruction_error = rel_err
        svd_layer.rank = r
        return svd_layer

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Compress all linear and Conv1D layers in the model."""
        self.layer_errors = {}
        self.layer_ranks = {}
        self.replaced_layers = []

        target_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            logger.warning("SAES-SVD v2: No compressible linear layers found.")
            return self.model, self.get_compression_stats()

        for name, module in target_layers:
            svd_layer = self.compress_linear(module, rank=self.rank, name=name)
            if svd_layer is not None:
                err = getattr(svd_layer, "reconstruction_error", 0.0)
                r = getattr(svd_layer, "rank", self.rank)
                self.layer_errors[name] = err
                self.layer_ranks[name] = r
                # Fix 4: Include per-layer rank and error in returned stats
                self.replaced_layers.append({"name": name, "rank": r, "error": err})
                replace_module(self.model, name, svd_layer)
            else:
                for p in module.parameters():
                    p.requires_grad = False

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        # Fix 4 & Fix 6: Detailed per-layer stats, errors, ranks, and fallback status
        stats.update(
            {
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "replaced_layers": self.replaced_layers,
                "compressed_layers": len(self.replaced_layers),
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "correcting_steps": self.correcting_steps,
                "target_attention_only": self.target_attention_only,
            }
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: SAES-SVD v2 Compressor with iterative spectrum-aware residual correction, Frobenius validation, and safe fallback.

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Compress all linear and Conv1D layers in the model."""
    self.layer_errors = {}
    self.layer_ranks = {}
    self.replaced_layers = []

    target_layers = [
        (name, mod)
        for name, mod in self.model.named_modules()
        if self._is_target_module(name, mod)
    ]

    if not target_layers:
        logger.warning("SAES-SVD v2: No compressible linear layers found.")
        return self.model, self.get_compression_stats()

    for name, module in target_layers:
        svd_layer = self.compress_linear(module, rank=self.rank, name=name)
        if svd_layer is not None:
            err = getattr(svd_layer, "reconstruction_error", 0.0)
            r = getattr(svd_layer, "rank", self.rank)
            self.layer_errors[name] = err
            self.layer_ranks[name] = r
            # Fix 4: Include per-layer rank and error in returned stats
            self.replaced_layers.append({"name": name, "rank": r, "error": err})
            replace_module(self.model, name, svd_layer)
        else:
            for p in module.parameters():
                p.requires_grad = False

    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    # Fix 4 & Fix 6: Detailed per-layer stats, errors, ranks, and fallback status
    stats.update(
        {
            "layer_errors": self.layer_errors,
            "layer_ranks": self.layer_ranks,
            "replaced_layers": self.replaced_layers,
            "compressed_layers": len(self.replaced_layers),
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "correcting_steps": self.correcting_steps,
            "target_attention_only": self.target_attention_only,
        }
    )
    return self.model, stats

Compress all linear and Conv1D layers in the model.

def compress_linear(self, layer: nn.Module, rank: int | None = None, name: str = '') ‑> SVDLinear | None

Expand source code

def compress_linear(
    self,
    layer: nn.Module,
    rank: int | None = None,
    name: str = "",
) -> SVDLinear | None:
    """Decompose a single linear or Conv1D layer with spectrum-aware error suppression."""
    base_layer = getattr(layer, "base_layer", layer)
    raw_w = base_layer.weight.data
    cls_name = base_layer.__class__.__name__.lower()
    # Fix 3: Conv1D detection by class name or nf attribute
    is_conv1d = "conv1d" in cls_name or hasattr(base_layer, "nf")

    if is_conv1d:
        W = raw_w.t().contiguous().float()
        out_f = W.shape[0]
        in_f = W.shape[1]
    else:
        W = raw_w.float()
        out_f = getattr(base_layer, "out_features", W.shape[0])
        in_f = getattr(base_layer, "in_features", W.shape[1])

    max_r = min(out_f, in_f)
    # Cap rank to guarantee parameter reduction
    param_cap = max(1, (in_f * out_f) // (in_f + out_f) - 1)
    effective_max_r = min(max_r, param_cap)

    if rank is not None:
        r = max(1, min(rank, effective_max_r))
    else:
        r = max(1, min(int(round(max_r * self.rank_ratio)), effective_max_r))

    device = base_layer.weight.device
    dtype = base_layer.weight.dtype

    # Compute full SVD in FP32
    try:
        U, S, Vt = torch.linalg.svd(W, full_matrices=False)
    except Exception as e:
        logger.warning("SAES-SVD SVD failure for %s: %s", name, e)
        return None

    # Initial symmetric square-root singular value allocation
    sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
    A = U[:, :r] * sqrt_S.unsqueeze(0)  # [out_features, r]
    B = sqrt_S.unsqueeze(1) * Vt[:r, :]  # [r, in_features]

    # Iterative residual error suppression
    eps_eye = 1e-6 * torch.eye(r, device=W.device, dtype=torch.float32)
    for _ in range(self.correcting_steps):
        E = W - A @ B
        # Update A towards residual gradient
        BBt = B @ B.t().contiguous() + eps_eye
        A = A + E @ B.t().contiguous() @ torch.linalg.pinv(BBt)

        # Update B towards residual gradient
        AtA = A.t().contiguous() @ A + eps_eye
        B = B + torch.linalg.pinv(AtA) @ A.t().contiguous() @ (W - A @ B)

    # Compute relative Frobenius reconstruction error
    W_rec = A @ B
    norm_orig = torch.norm(W, p="fro")
    rel_err = (torch.norm(W - W_rec, p="fro") / (norm_orig + 1e-8)).item()

    # Fix 2: Always log a warning when reconstruction error exceeds max_error
    if rel_err > self.max_error:
        logger.warning(
            "SAES-SVD reconstruction error %.4f on %s exceeded threshold %.4f%s",
            rel_err,
            name,
            self.max_error,
            (
                ". Keeping original layer."
                if self.fallback_to_original
                else ". Proceeding with compressed layer."
            ),
        )
        if self.fallback_to_original:
            return None

    has_bias = getattr(base_layer, "bias", None) is not None
    has_lora = hasattr(layer, "lora_A") and hasattr(layer, "lora_B")
    lora_r = getattr(layer, "r", 0) if has_lora else 0

    svd_layer = SVDLinear(
        in_features=in_f,
        out_features=out_f,
        rank=r,
        bias=has_bias,
        dtype=dtype,
        device=device,
        lora_r=lora_r,
    )
    svd_layer.set_low_rank_factors(
        U_factor=A.to(dtype=dtype, device=device),
        V_factor=B.to(dtype=dtype, device=device),
    )
    # Fix 5: Ensure bias is copied correctly for both Linear and Conv1D
    if has_bias and svd_layer.bias is not None:
        svd_layer.bias.data.copy_(base_layer.bias.data.to(dtype=dtype, device=device))

    if has_lora and svd_layer.use_lora:
        try:
            if hasattr(layer.lora_A, "weight"):
                svd_layer.lora_A.weight.data.copy_(layer.lora_A.weight.data)
            if hasattr(layer.lora_B, "weight"):
                svd_layer.lora_B.weight.data.copy_(layer.lora_B.weight.data)
        except Exception as e_lora:
            logger.warning("Could not preserve LoRA adapters in SAES-SVD: %s", e_lora)

    # Store metadata on module
    svd_layer.reconstruction_error = rel_err
    svd_layer.rank = r
    return svd_layer

Decompose a single linear or Conv1D layer with spectrum-aware error suppression.

class SAESSVDCompressorV2 (model: nn.Module, rank: int | None = None, rank_ratio: float = 0.5, correcting_steps: int = 2, target_modules: list[str] | None = None, target_attention_only: bool = False, max_error: float = 0.6, fallback_to_original: bool = True, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("saes_svd_v2")
@COMPRESSOR_REGISTRY.register("saes_svd_compressor")
@COMPRESSOR_REGISTRY.register("saes_svd")
@COMPRESSOR_REGISTRY.register("saes")
class SAESSVDCompressorV2(BaseCompressor):
    is_custom_attention: bool = True
    """
    SAES-SVD v2 Compressor with iterative spectrum-aware residual correction,
    Frobenius validation, and safe fallback.
    """

    def __init__(
        self,
        model: nn.Module,
        rank: int | None = None,
        rank_ratio: float = 0.5,
        correcting_steps: int = 2,
        target_modules: list[str] | None = None,
        target_attention_only: bool = False,
        max_error: float = 0.60,
        fallback_to_original: bool = True,  # Fix 1: Default changed to True
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="saes_svd_v2", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            rank = _cfg_get(cfg, "saes_rank", _cfg_get(cfg, "rank", rank))
            rank_ratio = _cfg_get(cfg, "saes_rank_ratio", _cfg_get(cfg, "rank_ratio", rank_ratio))
            correcting_steps = _cfg_get(cfg, "correcting_steps", correcting_steps)
            max_error = _cfg_get(cfg, "max_error", max_error)
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)

        if "rank" in kwargs and kwargs["rank"] is not None:
            rank = kwargs["rank"]
        if "fallback_to_original" in kwargs:
            fallback_to_original = bool(kwargs["fallback_to_original"])

        self.rank = rank
        self.rank_ratio = float(rank_ratio)
        self.correcting_steps = int(correcting_steps)
        self.target_modules = target_modules
        self.target_attention_only = bool(target_attention_only)
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.replaced_layers: list[dict[str, Any]] = []

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        """
        Fix 6: Unified critical-layer protection delegated to should_compress_linear.
        """
        name_l = name.lower()
        if self.target_attention_only:
            attn_keywords = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(k in name_l for k in attn_keywords):
                return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__.lower()
        # Fix 3: Reliable Conv1D check using class name and hasattr(nf)
        is_linear_like = (
            isinstance(base_mod, nn.Linear) or "conv1d" in cls_name or hasattr(base_mod, "nf")
        )
        if not is_linear_like:
            return False
        if isinstance(base_mod, SVDLinear):
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    def compress_linear(
        self,
        layer: nn.Module,
        rank: int | None = None,
        name: str = "",
    ) -> SVDLinear | None:
        """Decompose a single linear or Conv1D layer with spectrum-aware error suppression."""
        base_layer = getattr(layer, "base_layer", layer)
        raw_w = base_layer.weight.data
        cls_name = base_layer.__class__.__name__.lower()
        # Fix 3: Conv1D detection by class name or nf attribute
        is_conv1d = "conv1d" in cls_name or hasattr(base_layer, "nf")

        if is_conv1d:
            W = raw_w.t().contiguous().float()
            out_f = W.shape[0]
            in_f = W.shape[1]
        else:
            W = raw_w.float()
            out_f = getattr(base_layer, "out_features", W.shape[0])
            in_f = getattr(base_layer, "in_features", W.shape[1])

        max_r = min(out_f, in_f)
        # Cap rank to guarantee parameter reduction
        param_cap = max(1, (in_f * out_f) // (in_f + out_f) - 1)
        effective_max_r = min(max_r, param_cap)

        if rank is not None:
            r = max(1, min(rank, effective_max_r))
        else:
            r = max(1, min(int(round(max_r * self.rank_ratio)), effective_max_r))

        device = base_layer.weight.device
        dtype = base_layer.weight.dtype

        # Compute full SVD in FP32
        try:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        except Exception as e:
            logger.warning("SAES-SVD SVD failure for %s: %s", name, e)
            return None

        # Initial symmetric square-root singular value allocation
        sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
        A = U[:, :r] * sqrt_S.unsqueeze(0)  # [out_features, r]
        B = sqrt_S.unsqueeze(1) * Vt[:r, :]  # [r, in_features]

        # Iterative residual error suppression
        eps_eye = 1e-6 * torch.eye(r, device=W.device, dtype=torch.float32)
        for _ in range(self.correcting_steps):
            E = W - A @ B
            # Update A towards residual gradient
            BBt = B @ B.t().contiguous() + eps_eye
            A = A + E @ B.t().contiguous() @ torch.linalg.pinv(BBt)

            # Update B towards residual gradient
            AtA = A.t().contiguous() @ A + eps_eye
            B = B + torch.linalg.pinv(AtA) @ A.t().contiguous() @ (W - A @ B)

        # Compute relative Frobenius reconstruction error
        W_rec = A @ B
        norm_orig = torch.norm(W, p="fro")
        rel_err = (torch.norm(W - W_rec, p="fro") / (norm_orig + 1e-8)).item()

        # Fix 2: Always log a warning when reconstruction error exceeds max_error
        if rel_err > self.max_error:
            logger.warning(
                "SAES-SVD reconstruction error %.4f on %s exceeded threshold %.4f%s",
                rel_err,
                name,
                self.max_error,
                (
                    ". Keeping original layer."
                    if self.fallback_to_original
                    else ". Proceeding with compressed layer."
                ),
            )
            if self.fallback_to_original:
                return None

        has_bias = getattr(base_layer, "bias", None) is not None
        has_lora = hasattr(layer, "lora_A") and hasattr(layer, "lora_B")
        lora_r = getattr(layer, "r", 0) if has_lora else 0

        svd_layer = SVDLinear(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=has_bias,
            dtype=dtype,
            device=device,
            lora_r=lora_r,
        )
        svd_layer.set_low_rank_factors(
            U_factor=A.to(dtype=dtype, device=device),
            V_factor=B.to(dtype=dtype, device=device),
        )
        # Fix 5: Ensure bias is copied correctly for both Linear and Conv1D
        if has_bias and svd_layer.bias is not None:
            svd_layer.bias.data.copy_(base_layer.bias.data.to(dtype=dtype, device=device))

        if has_lora and svd_layer.use_lora:
            try:
                if hasattr(layer.lora_A, "weight"):
                    svd_layer.lora_A.weight.data.copy_(layer.lora_A.weight.data)
                if hasattr(layer.lora_B, "weight"):
                    svd_layer.lora_B.weight.data.copy_(layer.lora_B.weight.data)
            except Exception as e_lora:
                logger.warning("Could not preserve LoRA adapters in SAES-SVD: %s", e_lora)

        # Store metadata on module
        svd_layer.reconstruction_error = rel_err
        svd_layer.rank = r
        return svd_layer

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Compress all linear and Conv1D layers in the model."""
        self.layer_errors = {}
        self.layer_ranks = {}
        self.replaced_layers = []

        target_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            logger.warning("SAES-SVD v2: No compressible linear layers found.")
            return self.model, self.get_compression_stats()

        for name, module in target_layers:
            svd_layer = self.compress_linear(module, rank=self.rank, name=name)
            if svd_layer is not None:
                err = getattr(svd_layer, "reconstruction_error", 0.0)
                r = getattr(svd_layer, "rank", self.rank)
                self.layer_errors[name] = err
                self.layer_ranks[name] = r
                # Fix 4: Include per-layer rank and error in returned stats
                self.replaced_layers.append({"name": name, "rank": r, "error": err})
                replace_module(self.model, name, svd_layer)
            else:
                for p in module.parameters():
                    p.requires_grad = False

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        # Fix 4 & Fix 6: Detailed per-layer stats, errors, ranks, and fallback status
        stats.update(
            {
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "replaced_layers": self.replaced_layers,
                "compressed_layers": len(self.replaced_layers),
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "correcting_steps": self.correcting_steps,
                "target_attention_only": self.target_attention_only,
            }
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: SAES-SVD v2 Compressor with iterative spectrum-aware residual correction, Frobenius validation, and safe fallback.

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Compress all linear and Conv1D layers in the model."""
    self.layer_errors = {}
    self.layer_ranks = {}
    self.replaced_layers = []

    target_layers = [
        (name, mod)
        for name, mod in self.model.named_modules()
        if self._is_target_module(name, mod)
    ]

    if not target_layers:
        logger.warning("SAES-SVD v2: No compressible linear layers found.")
        return self.model, self.get_compression_stats()

    for name, module in target_layers:
        svd_layer = self.compress_linear(module, rank=self.rank, name=name)
        if svd_layer is not None:
            err = getattr(svd_layer, "reconstruction_error", 0.0)
            r = getattr(svd_layer, "rank", self.rank)
            self.layer_errors[name] = err
            self.layer_ranks[name] = r
            # Fix 4: Include per-layer rank and error in returned stats
            self.replaced_layers.append({"name": name, "rank": r, "error": err})
            replace_module(self.model, name, svd_layer)
        else:
            for p in module.parameters():
                p.requires_grad = False

    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    # Fix 4 & Fix 6: Detailed per-layer stats, errors, ranks, and fallback status
    stats.update(
        {
            "layer_errors": self.layer_errors,
            "layer_ranks": self.layer_ranks,
            "replaced_layers": self.replaced_layers,
            "compressed_layers": len(self.replaced_layers),
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "correcting_steps": self.correcting_steps,
            "target_attention_only": self.target_attention_only,
        }
    )
    return self.model, stats

Compress all linear and Conv1D layers in the model.

def compress_linear(self, layer: nn.Module, rank: int | None = None, name: str = '') ‑> SVDLinear | None

Expand source code

def compress_linear(
    self,
    layer: nn.Module,
    rank: int | None = None,
    name: str = "",
) -> SVDLinear | None:
    """Decompose a single linear or Conv1D layer with spectrum-aware error suppression."""
    base_layer = getattr(layer, "base_layer", layer)
    raw_w = base_layer.weight.data
    cls_name = base_layer.__class__.__name__.lower()
    # Fix 3: Conv1D detection by class name or nf attribute
    is_conv1d = "conv1d" in cls_name or hasattr(base_layer, "nf")

    if is_conv1d:
        W = raw_w.t().contiguous().float()
        out_f = W.shape[0]
        in_f = W.shape[1]
    else:
        W = raw_w.float()
        out_f = getattr(base_layer, "out_features", W.shape[0])
        in_f = getattr(base_layer, "in_features", W.shape[1])

    max_r = min(out_f, in_f)
    # Cap rank to guarantee parameter reduction
    param_cap = max(1, (in_f * out_f) // (in_f + out_f) - 1)
    effective_max_r = min(max_r, param_cap)

    if rank is not None:
        r = max(1, min(rank, effective_max_r))
    else:
        r = max(1, min(int(round(max_r * self.rank_ratio)), effective_max_r))

    device = base_layer.weight.device
    dtype = base_layer.weight.dtype

    # Compute full SVD in FP32
    try:
        U, S, Vt = torch.linalg.svd(W, full_matrices=False)
    except Exception as e:
        logger.warning("SAES-SVD SVD failure for %s: %s", name, e)
        return None

    # Initial symmetric square-root singular value allocation
    sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
    A = U[:, :r] * sqrt_S.unsqueeze(0)  # [out_features, r]
    B = sqrt_S.unsqueeze(1) * Vt[:r, :]  # [r, in_features]

    # Iterative residual error suppression
    eps_eye = 1e-6 * torch.eye(r, device=W.device, dtype=torch.float32)
    for _ in range(self.correcting_steps):
        E = W - A @ B
        # Update A towards residual gradient
        BBt = B @ B.t().contiguous() + eps_eye
        A = A + E @ B.t().contiguous() @ torch.linalg.pinv(BBt)

        # Update B towards residual gradient
        AtA = A.t().contiguous() @ A + eps_eye
        B = B + torch.linalg.pinv(AtA) @ A.t().contiguous() @ (W - A @ B)

    # Compute relative Frobenius reconstruction error
    W_rec = A @ B
    norm_orig = torch.norm(W, p="fro")
    rel_err = (torch.norm(W - W_rec, p="fro") / (norm_orig + 1e-8)).item()

    # Fix 2: Always log a warning when reconstruction error exceeds max_error
    if rel_err > self.max_error:
        logger.warning(
            "SAES-SVD reconstruction error %.4f on %s exceeded threshold %.4f%s",
            rel_err,
            name,
            self.max_error,
            (
                ". Keeping original layer."
                if self.fallback_to_original
                else ". Proceeding with compressed layer."
            ),
        )
        if self.fallback_to_original:
            return None

    has_bias = getattr(base_layer, "bias", None) is not None
    has_lora = hasattr(layer, "lora_A") and hasattr(layer, "lora_B")
    lora_r = getattr(layer, "r", 0) if has_lora else 0

    svd_layer = SVDLinear(
        in_features=in_f,
        out_features=out_f,
        rank=r,
        bias=has_bias,
        dtype=dtype,
        device=device,
        lora_r=lora_r,
    )
    svd_layer.set_low_rank_factors(
        U_factor=A.to(dtype=dtype, device=device),
        V_factor=B.to(dtype=dtype, device=device),
    )
    # Fix 5: Ensure bias is copied correctly for both Linear and Conv1D
    if has_bias and svd_layer.bias is not None:
        svd_layer.bias.data.copy_(base_layer.bias.data.to(dtype=dtype, device=device))

    if has_lora and svd_layer.use_lora:
        try:
            if hasattr(layer.lora_A, "weight"):
                svd_layer.lora_A.weight.data.copy_(layer.lora_A.weight.data)
            if hasattr(layer.lora_B, "weight"):
                svd_layer.lora_B.weight.data.copy_(layer.lora_B.weight.data)
        except Exception as e_lora:
            logger.warning("Could not preserve LoRA adapters in SAES-SVD: %s", e_lora)

    # Store metadata on module
    svd_layer.reconstruction_error = rel_err
    svd_layer.rank = r
    return svd_layer

Decompose a single linear or Conv1D layer with spectrum-aware error suppression.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class SALRCompressor (model: nn.Module, rank: int = 64, sparsity: float = 0.5, min_rank: int = 8, skip_attention_qkv: bool = False, target_attention_only: bool = False, use_lora: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, target_modules: list[str] | None = None, max_error: float = 0.5, fallback_to_original: bool = True, min_numel: int = 16, device: str | torch.device | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("salr")
class SALRCompressor(BaseCompressor):
    is_custom_attention: bool = True
    """
    SALR Compressor: Structure-Aware Low-Rank and Sparse Pruning.

    Factorizes dense nn.Linear or Conv1D layers into symmetric low-rank factors
    (U * sqrt(Sigma), sqrt(Sigma) * V^T) and applies structured sparse pruning
    to trailing singular components in the bottleneck with LoRA adapters for recovery.
    """

    def __init__(
        self,
        model: nn.Module,
        rank: int = 64,
        sparsity: float = 0.5,
        min_rank: int = 8,
        skip_attention_qkv: bool = False,
        target_attention_only: bool = False,
        use_lora: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        target_modules: list[str] | None = None,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        min_numel: int = 16,
        device: str | torch.device | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="salr", **kwargs)
        cfg_obj = cfg or kwargs.get("comp_config")

        if "rank" in kwargs and kwargs["rank"] is not None:
            rank = kwargs["rank"]
        elif cfg_obj is not None:
            rank = _cfg_get(cfg_obj, "salr_rank", _cfg_get(cfg_obj, "rank", rank))

        self.rank = int(rank)
        self.sparsity = max(
            0.0,
            min(
                float(
                    sparsity
                    or _cfg_get(cfg_obj, "salr_sparsity", _cfg_get(cfg_obj, "sparsity", 0.5))
                ),
                0.90,
            ),
        )
        self.min_rank = max(1, int(min_rank or _cfg_get(cfg_obj, "min_rank", 8)))
        self.skip_attention_qkv = bool(
            skip_attention_qkv
            if skip_attention_qkv is not None
            else _cfg_get(cfg_obj, "skip_attention_qkv", False)
        )
        self.target_attention_only = bool(
            target_attention_only
            if target_attention_only is not None
            else _cfg_get(cfg_obj, "target_attention_only", False)
        )
        self.use_lora = bool(use_lora or _cfg_get(cfg_obj, "use_lora", False) or (lora_r > 0))
        self.lora_r = int(lora_r or _cfg_get(cfg_obj, "lora_r", 0))
        self.lora_alpha = float(lora_alpha or _cfg_get(cfg_obj, "lora_alpha", 1.0))
        self.target_modules = target_modules
        self.max_error = float(max_error or _cfg_get(cfg_obj, "max_error", 0.50))
        self.fallback_to_original = bool(
            fallback_to_original
            if fallback_to_original is not None
            else _cfg_get(cfg_obj, "fallback_to_original", True)
        )
        self.min_numel = max(1, int(min_numel or _cfg_get(cfg_obj, "min_numel", 16)))

        # Target device derivation
        if device is not None:
            try:
                self.target_device = torch.device(device)
            except Exception:
                self.target_device = (
                    next(model.parameters()).device
                    if list(model.parameters())
                    else torch.device("cpu")
                )
        else:
            self.target_device = (
                next(model.parameters()).device if list(model.parameters()) else torch.device("cpu")
            )

        # Diagnostics & Tracking
        self.visited_modules: set[int] = set()
        self.compressed_layers: list[str] = []
        self.protected_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.failed_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.replaced_layers: list[dict[str, Any]] = []

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        """Check if layer belongs to an attention projection in a GQA architecture."""
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        is_gqa = num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads

        name_l = name.lower()
        is_attn_qkv = any(
            p in name_l
            for p in [
                "q_proj",
                "k_proj",
                "v_proj",
                "c_attn",
                "query",
                "key",
                "value",
                "q_lin",
                "k_lin",
                "v_lin",
                "self_attn.q",
                "self_attn.k",
                "self_attn.v",
            ]
        )
        return (is_gqa or self.skip_attention_qkv) and is_attn_qkv

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        # 1. Protection for critical / classification / head / pooler / embedding layers
        if any(
            k in name_l
            for k in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embed",
                "embedding",
                "cls_head",
                "head",
                "logits_processor",
            ]
        ):
            self.protected_layers.append(name)
            return False

        # Target attention only filter
        if self.target_attention_only:
            attn_keywords = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(k in name_l for k in attn_keywords):
                return False

        # 2. GQA-aware skipping
        if self._is_gqa_asymmetric(name, module):
            self.skipped_layers.append(name)
            return False

        # 3. Check linear or Conv1D nature
        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__
        raw_w = getattr(base_mod, "weight", None)
        if raw_w is None or not isinstance(raw_w, torch.Tensor):
            return False

        is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_mod, "nf")
        is_linear_like = isinstance(base_mod, nn.Linear) or is_conv1d
        if not is_linear_like or isinstance(base_mod, SALRLayer):
            return False

        # 4. Tiny layer guard
        if raw_w.numel() < self.min_numel:
            self.skipped_layers.append(name)
            return False

        # 5. Respect framework compatibility rules
        if not should_compress_linear(self.model, name):
            self.protected_layers.append(name)
            return False

        # 6. User-defined explicit target modules
        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def compress_linear(
        self,
        layer: nn.Module,
        rank: int,
    ) -> SALRLayer | None:
        """Decompose a single linear or Conv1D layer into SALRLayer with validation."""
        base_layer = getattr(layer, "base_layer", layer)
        raw_w = base_layer.weight.data
        cls_name = base_layer.__class__.__name__
        is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_layer, "nf")

        if is_conv1d:
            W = raw_w.t().contiguous().float()
            out_f, in_f = W.shape
        else:
            W = raw_w.contiguous().float()
            out_f = getattr(base_layer, "out_features", W.shape[0])
            in_f = getattr(base_layer, "in_features", W.shape[1])

        max_r = min(out_f, in_f)
        param_cap = max(1, (in_f * out_f) // (in_f + out_f) - 1)
        r = max(self.min_rank, min(rank, max_r - 1, param_cap))
        if r >= max_r:
            logger.debug(
                "SALR: rank %d >= max_r %d; skipping layer for no parameter reduction",
                r,
                max_r,
            )
            return None

        device = base_layer.weight.device
        dtype = base_layer.weight.dtype

        # SVD in FP32
        try:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        except Exception as e:
            logger.warning("SALR SVD decomposition failure: %s", e)
            return None

        sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
        U_r = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
        V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(dtype=dtype, device=device)

        has_existing_lora = hasattr(layer, "lora_A") and hasattr(layer, "lora_B")
        effective_lora_r = self.lora_r
        if has_existing_lora and getattr(layer, "r", 0) > 0:
            effective_lora_r = getattr(layer, "r", self.lora_r)

        salr_layer = SALRLayer(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            sparsity=self.sparsity,
            bias=base_layer.bias is not None,
            dtype=dtype,
            device=device,
            lora_r=effective_lora_r if (self.use_lora or has_existing_lora) else 0,
            lora_alpha=self.lora_alpha,
        )
        salr_layer.set_low_rank_factors(U_r, V_r)
        if base_layer.bias is not None:
            salr_layer.bias.data.copy_(base_layer.bias.data)

        # Preserve existing LoRA adapters
        if has_existing_lora and salr_layer.lora_r > 0:
            try:
                if hasattr(layer.lora_A, "weight") and salr_layer.lora_A is not None:
                    salr_layer.lora_A.data.copy_(layer.lora_A.weight.data)
                elif (
                    hasattr(layer, "lora_A")
                    and isinstance(layer.lora_A, torch.Tensor)
                    and salr_layer.lora_A is not None
                ):
                    salr_layer.lora_A.data.copy_(layer.lora_A.data)
                if hasattr(layer.lora_B, "weight") and salr_layer.lora_B is not None:
                    salr_layer.lora_B.data.copy_(layer.lora_B.weight.data)
                elif (
                    hasattr(layer, "lora_B")
                    and isinstance(layer.lora_B, torch.Tensor)
                    and salr_layer.lora_B is not None
                ):
                    salr_layer.lora_B.data.copy_(layer.lora_B.data)
            except Exception as e_lora:
                logger.warning("Could not copy existing LoRA adapters: %s", e_lora)

        # Frobenius reconstruction error check
        W_rec = salr_layer.weight.float()
        norm_orig = torch.norm(W, p="fro").clamp_min(1e-8)
        rel_err = (torch.norm(W - W_rec, p="fro") / norm_orig).item()

        is_invalid = torch.isnan(W_rec).any() or torch.isinf(W_rec).any()
        if rel_err > self.max_error:
            logger.warning(
                "SALR: reconstruction error %.4f exceeded threshold %.4f.",
                rel_err,
                self.max_error,
            )
            is_invalid = True

        if is_invalid and self.fallback_to_original:
            return None

        salr_layer.reconstruction_error = rel_err
        return salr_layer

    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """Compress all linear layers in the model."""
        self.layer_errors = {}
        self.replaced_layers = []
        self.visited_modules.clear()
        self.compressed_layers.clear()
        self.protected_layers.clear()
        self.skipped_layers.clear()
        self.fallback_layers.clear()
        self.failed_layers.clear()

        target_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            logger.warning("SALR: No compressible linear layers found.")
            return self.model, self.get_compression_stats()

        for name, module in target_layers:
            mod_id = id(module)
            if mod_id in self.visited_modules:
                logger.debug("SALR: Skipping tied module %s", name)
                self.skipped_layers.append(name)
                continue
            self.visited_modules.add(mod_id)

            try:
                new_layer = self.compress_linear(module, rank=self.rank)
                if new_layer is not None:
                    err = getattr(new_layer, "reconstruction_error", 0.0)
                    self.layer_errors[name] = err
                    self.replaced_layers.append(
                        {"name": name, "rank": new_layer.rank, "error": err}
                    )
                    replace_module(self.model, name, new_layer)
                    self.compressed_layers.append(name)
                else:
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
            except Exception as e_comp:
                logger.warning("Failed to compress %s: %s", name, e_comp)
                self.failed_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False

        # Post-compression forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                dummy_input = torch.randint(0, 100, (1, 8), device=self.target_device)
                self.model(dummy_input)
        except Exception as ex_sanity:
            logger.warning("SALR post-compression forward sanity check warning: %s", ex_sanity)
            forward_check_success = False

        self._compressed_params = self._count_params(self.model)
        avg_err = (
            (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
        )
        stats = self.get_compression_stats()
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "protected_layers": self.protected_layers,
                "protected_layer_count": len(self.protected_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "failed_layers": self.failed_layers,
                "failed_layer_count": len(self.failed_layers),
                "layer_errors": self.layer_errors,
                "avg_reconstruction_error": avg_err,
                "forward_check_success": forward_check_success,
                "rank": self.rank,
                "min_rank": self.min_rank,
                "sparsity": self.sparsity,
                "skip_attention_qkv": self.skip_attention_qkv,
                "target_attention_only": self.target_attention_only,
                "use_lora": self.use_lora,
                "lora_r": self.lora_r,
                "method": "salr",
            }
        )
        self.stats = stats
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: SALR Compressor: Structure-Aware Low-Rank and Sparse Pruning.

Factorizes dense nn.Linear or Conv1D layers into symmetric low-rank factors (U * sqrt(Sigma), sqrt(Sigma) * V^T) and applies structured sparse pruning to trailing singular components in the bottleneck with LoRA adapters for recovery.

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """Compress all linear layers in the model."""
    self.layer_errors = {}
    self.replaced_layers = []
    self.visited_modules.clear()
    self.compressed_layers.clear()
    self.protected_layers.clear()
    self.skipped_layers.clear()
    self.fallback_layers.clear()
    self.failed_layers.clear()

    target_layers = [
        (name, mod)
        for name, mod in self.model.named_modules()
        if self._is_target_module(name, mod)
    ]

    if not target_layers:
        logger.warning("SALR: No compressible linear layers found.")
        return self.model, self.get_compression_stats()

    for name, module in target_layers:
        mod_id = id(module)
        if mod_id in self.visited_modules:
            logger.debug("SALR: Skipping tied module %s", name)
            self.skipped_layers.append(name)
            continue
        self.visited_modules.add(mod_id)

        try:
            new_layer = self.compress_linear(module, rank=self.rank)
            if new_layer is not None:
                err = getattr(new_layer, "reconstruction_error", 0.0)
                self.layer_errors[name] = err
                self.replaced_layers.append(
                    {"name": name, "rank": new_layer.rank, "error": err}
                )
                replace_module(self.model, name, new_layer)
                self.compressed_layers.append(name)
            else:
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
        except Exception as e_comp:
            logger.warning("Failed to compress %s: %s", name, e_comp)
            self.failed_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False

    # Post-compression forward sanity check
    forward_check_success = True
    try:
        with torch.no_grad():
            dummy_input = torch.randint(0, 100, (1, 8), device=self.target_device)
            self.model(dummy_input)
    except Exception as ex_sanity:
        logger.warning("SALR post-compression forward sanity check warning: %s", ex_sanity)
        forward_check_success = False

    self._compressed_params = self._count_params(self.model)
    avg_err = (
        (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
    )
    stats = self.get_compression_stats()
    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "protected_layers": self.protected_layers,
            "protected_layer_count": len(self.protected_layers),
            "skipped_layers": self.skipped_layers,
            "skipped_layer_count": len(self.skipped_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "failed_layers": self.failed_layers,
            "failed_layer_count": len(self.failed_layers),
            "layer_errors": self.layer_errors,
            "avg_reconstruction_error": avg_err,
            "forward_check_success": forward_check_success,
            "rank": self.rank,
            "min_rank": self.min_rank,
            "sparsity": self.sparsity,
            "skip_attention_qkv": self.skip_attention_qkv,
            "target_attention_only": self.target_attention_only,
            "use_lora": self.use_lora,
            "lora_r": self.lora_r,
            "method": "salr",
        }
    )
    self.stats = stats
    return self.model, stats

Compress all linear layers in the model.

def compress_linear(self, layer: nn.Module, rank: int) ‑> SALRLayer | None

Expand source code

def compress_linear(
    self,
    layer: nn.Module,
    rank: int,
) -> SALRLayer | None:
    """Decompose a single linear or Conv1D layer into SALRLayer with validation."""
    base_layer = getattr(layer, "base_layer", layer)
    raw_w = base_layer.weight.data
    cls_name = base_layer.__class__.__name__
    is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_layer, "nf")

    if is_conv1d:
        W = raw_w.t().contiguous().float()
        out_f, in_f = W.shape
    else:
        W = raw_w.contiguous().float()
        out_f = getattr(base_layer, "out_features", W.shape[0])
        in_f = getattr(base_layer, "in_features", W.shape[1])

    max_r = min(out_f, in_f)
    param_cap = max(1, (in_f * out_f) // (in_f + out_f) - 1)
    r = max(self.min_rank, min(rank, max_r - 1, param_cap))
    if r >= max_r:
        logger.debug(
            "SALR: rank %d >= max_r %d; skipping layer for no parameter reduction",
            r,
            max_r,
        )
        return None

    device = base_layer.weight.device
    dtype = base_layer.weight.dtype

    # SVD in FP32
    try:
        U, S, Vt = torch.linalg.svd(W, full_matrices=False)
    except Exception as e:
        logger.warning("SALR SVD decomposition failure: %s", e)
        return None

    sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
    U_r = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
    V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(dtype=dtype, device=device)

    has_existing_lora = hasattr(layer, "lora_A") and hasattr(layer, "lora_B")
    effective_lora_r = self.lora_r
    if has_existing_lora and getattr(layer, "r", 0) > 0:
        effective_lora_r = getattr(layer, "r", self.lora_r)

    salr_layer = SALRLayer(
        in_features=in_f,
        out_features=out_f,
        rank=r,
        sparsity=self.sparsity,
        bias=base_layer.bias is not None,
        dtype=dtype,
        device=device,
        lora_r=effective_lora_r if (self.use_lora or has_existing_lora) else 0,
        lora_alpha=self.lora_alpha,
    )
    salr_layer.set_low_rank_factors(U_r, V_r)
    if base_layer.bias is not None:
        salr_layer.bias.data.copy_(base_layer.bias.data)

    # Preserve existing LoRA adapters
    if has_existing_lora and salr_layer.lora_r > 0:
        try:
            if hasattr(layer.lora_A, "weight") and salr_layer.lora_A is not None:
                salr_layer.lora_A.data.copy_(layer.lora_A.weight.data)
            elif (
                hasattr(layer, "lora_A")
                and isinstance(layer.lora_A, torch.Tensor)
                and salr_layer.lora_A is not None
            ):
                salr_layer.lora_A.data.copy_(layer.lora_A.data)
            if hasattr(layer.lora_B, "weight") and salr_layer.lora_B is not None:
                salr_layer.lora_B.data.copy_(layer.lora_B.weight.data)
            elif (
                hasattr(layer, "lora_B")
                and isinstance(layer.lora_B, torch.Tensor)
                and salr_layer.lora_B is not None
            ):
                salr_layer.lora_B.data.copy_(layer.lora_B.data)
        except Exception as e_lora:
            logger.warning("Could not copy existing LoRA adapters: %s", e_lora)

    # Frobenius reconstruction error check
    W_rec = salr_layer.weight.float()
    norm_orig = torch.norm(W, p="fro").clamp_min(1e-8)
    rel_err = (torch.norm(W - W_rec, p="fro") / norm_orig).item()

    is_invalid = torch.isnan(W_rec).any() or torch.isinf(W_rec).any()
    if rel_err > self.max_error:
        logger.warning(
            "SALR: reconstruction error %.4f exceeded threshold %.4f.",
            rel_err,
            self.max_error,
        )
        is_invalid = True

    if is_invalid and self.fallback_to_original:
        return None

    salr_layer.reconstruction_error = rel_err
    return salr_layer

Decompose a single linear or Conv1D layer into SALRLayer with validation.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class SALRLayer (in_features: int, out_features: int, rank: int, sparsity: float = 0.5, bias: bool = True, dtype: torch.dtype = torch.float32, device: torch.device | None = None, lora_r: int = 0, lora_alpha: float = 1.0)

Expand source code

class SALRLayer(nn.Module):
    """
    Structure-Aware Low-Rank + Sparse Layer (SALR).

    Combines truncated SVD rank reduction with sparse singular component selection
    and optional trainable LoRA fine-tuning adapters:
        y = (x @ (V * mask).T) @ U.T + bias + (lora_B @ lora_A @ x) * scaling
    where U: [out_features, rank], V: [rank, in_features], mask: [rank, 1].
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        sparsity: float = 0.5,
        bias: bool = True,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = max(1, min(rank, in_features, out_features))
        self.sparsity = max(0.0, min(float(sparsity), 0.90))
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        requires_grad = not (self.lora_r > 0)
        self.V = nn.Parameter(
            torch.empty(self.rank, in_features, dtype=dtype, device=device),
            requires_grad=requires_grad,
        )
        self.U = nn.Parameter(
            torch.empty(out_features, self.rank, dtype=dtype, device=device),
            requires_grad=requires_grad,
        )

        if bias:
            self.bias = nn.Parameter(
                torch.empty(out_features, dtype=dtype, device=device),
                requires_grad=requires_grad,
            )
        else:
            self.register_parameter("bias", None)

        # 1D binary mask on singular bottleneck components: [rank, 1]
        mask = torch.ones(self.rank, 1, dtype=torch.float32, device=device)
        if 0.0 < self.sparsity < 1.0:
            num_zeros = int(round(self.rank * self.sparsity))
            if num_zeros > 0:
                mask[-num_zeros:, :] = 0.0
        self.register_buffer("mask", mask)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty((self.lora_r, in_features), dtype=torch.float32, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.zeros((out_features, self.lora_r), dtype=torch.float32, device=device)
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        self.reset_parameters()

    def reset_parameters(self) -> None:
        nn.init.orthogonal_(self.U)
        nn.init.orthogonal_(self.V)
        if self.bias is not None:
            nn.init.zeros_(self.bias)

    def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
        with torch.no_grad():
            self.U.copy_(U_factor)
            self.V.copy_(V_factor)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Two-stage forward pass with sparse bottleneck mask and optional LoRA branch."""
        V_masked = self.V * self.mask.to(dtype=self.V.dtype)
        latent = F.linear(x, V_masked)  # [..., rank]
        out = F.linear(latent, self.U, self.bias)  # [..., out_features]

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x.float(), self.lora_A), self.lora_B) * self.scaling
            out = out + lora_out.to(out.dtype)
        return out

    @property
    def weight(self) -> torch.Tensor:
        """Effective reconstructed weight matrix."""
        V_masked = self.V * self.mask.to(dtype=self.V.dtype)
        return torch.matmul(self.U, V_masked)

    def extra_repr(self) -> str:
        return f"in_features={self.in_features}, out_features={self.out_features}, rank={self.rank}, sparsity={self.sparsity}, lora_r={self.lora_r}"

Structure-Aware Low-Rank + Sparse Layer (SALR).

Combines truncated SVD rank reduction with sparse singular component selection and optional trainable LoRA fine-tuning adapters: y = (x @ (V * mask).T) @ U.T + bias + (lora_B @ lora_A @ x) * scaling where U: [out_features, rank], V: [rank, in_features], mask: [rank, 1].

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Effective reconstructed weight matrix."""
    V_masked = self.V * self.mask.to(dtype=self.V.dtype)
    return torch.matmul(self.U, V_masked)

Effective reconstructed weight matrix.

Methods

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    return f"in_features={self.in_features}, out_features={self.out_features}, rank={self.rank}, sparsity={self.sparsity}, lora_r={self.lora_r}"

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Two-stage forward pass with sparse bottleneck mask and optional LoRA branch."""
    V_masked = self.V * self.mask.to(dtype=self.V.dtype)
    latent = F.linear(x, V_masked)  # [..., rank]
    out = F.linear(latent, self.U, self.bias)  # [..., out_features]

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x.float(), self.lora_A), self.lora_B) * self.scaling
        out = out + lora_out.to(out.dtype)
    return out

Two-stage forward pass with sparse bottleneck mask and optional LoRA branch.

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    nn.init.orthogonal_(self.U)
    nn.init.orthogonal_(self.V)
    if self.bias is not None:
        nn.init.zeros_(self.bias)

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) ‑> None

Expand source code

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
    with torch.no_grad():
        self.U.copy_(U_factor)
        self.V.copy_(V_factor)

class SDRMCompressor (model: nn.Module, total_param_ratio: float = 0.7, min_rank: int = 8, skip_attention_qkv: bool = False, target_attention_only: bool = False, use_lora: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, target_modules: list[str] | None = None, max_error: float = 0.5, fallback_to_original: bool = True, min_numel: int = 16, device: str | torch.device | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("sdrm_class")
@COMPRESSOR_REGISTRY.register("sdrm_compressor_class")
class SDRMCompressor(BaseCompressor):
    is_custom_attention: bool = True
    """
    SDRM: Scaled Dimension Rank Modulation.
    Assigns layer rank proportional to parameter budget for edge inference and factorizes
    target layers into low-rank SVDLinear modules with optional LoRA fine-tuning adapters.
    """

    def __init__(
        self,
        model: nn.Module,
        total_param_ratio: float = 0.7,
        min_rank: int = 8,
        skip_attention_qkv: bool = False,
        target_attention_only: bool = False,
        use_lora: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        target_modules: list[str] | None = None,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        min_numel: int = 16,
        device: str | torch.device | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="sdrm", **kwargs)
        cfg_obj = cfg or kwargs.get("comp_config")

        self.total_param_ratio = max(
            0.05,
            min(
                float(
                    total_param_ratio
                    or _cfg_get(cfg_obj, "total_param_ratio", _cfg_get(cfg_obj, "ratio", 0.7))
                ),
                1.0,
            ),
        )
        self.min_rank = max(1, int(min_rank or _cfg_get(cfg_obj, "min_rank", 8)))
        self.skip_attention_qkv = bool(
            skip_attention_qkv
            if skip_attention_qkv is not None
            else _cfg_get(cfg_obj, "skip_attention_qkv", False)
        )
        self.target_attention_only = bool(
            target_attention_only
            if target_attention_only is not None
            else _cfg_get(cfg_obj, "target_attention_only", False)
        )
        self.use_lora = bool(use_lora or _cfg_get(cfg_obj, "use_lora", False) or (lora_r > 0))
        self.lora_r = int(lora_r or _cfg_get(cfg_obj, "lora_r", 0))
        self.lora_alpha = float(lora_alpha or _cfg_get(cfg_obj, "lora_alpha", 1.0))
        self.target_modules = target_modules
        self.max_error = float(max_error or _cfg_get(cfg_obj, "max_error", 0.50))
        self.fallback_to_original = bool(
            fallback_to_original
            if fallback_to_original is not None
            else _cfg_get(cfg_obj, "fallback_to_original", True)
        )
        self.min_numel = max(1, int(min_numel or _cfg_get(cfg_obj, "min_numel", 16)))
        self.protect_output_proj: bool = bool(_cfg_get(cfg_obj, "protect_output_proj", True))

        # Target device derivation
        if device is not None:
            try:
                self.target_device = torch.device(device)
            except Exception:
                self.target_device = (
                    next(model.parameters()).device
                    if list(model.parameters())
                    else torch.device("cpu")
                )
        else:
            self.target_device = (
                next(model.parameters()).device if list(model.parameters()) else torch.device("cpu")
            )

        # Diagnostics & Tracking
        self.visited_modules: set[int] = set()
        self.compressed_layers: list[str] = []
        self.protected_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.failed_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.replaced_layers: list[dict[str, Any]] = []

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        """Check if layer belongs to an attention projection in a GQA architecture."""
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(
            cfg, "num_attention_heads", getattr(module, "num_attention_heads", None)
        )
        num_kv_heads = getattr(
            cfg, "num_key_value_heads", getattr(module, "num_key_value_heads", None)
        )
        is_gqa = num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads

        name_l = name.lower()
        is_attn_qkv = any(
            p in name_l
            for p in [
                "q_proj",
                "k_proj",
                "v_proj",
                "c_attn",
                "query",
                "key",
                "value",
                "q_lin",
                "k_lin",
                "v_lin",
                "self_attn.q",
                "self_attn.k",
                "self_attn.v",
            ]
        )
        return (is_gqa or self.skip_attention_qkv) and is_attn_qkv

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        # 1. Protection for critical / classification / head / pooler / embedding layers
        PROTECTED_ALWAYS = [
            "classifier",
            "score",
            "lm_head",
            "pooler",
            "embedding",
            "cls_head",
            "head",
            "logits_processor",
        ]
        ATTENTION_OUTPUT = ["out_lin", "out_proj", "o_proj", "c_proj", "output.dense"]

        if any(k in name_l for k in PROTECTED_ALWAYS):
            self.protected_layers.append(name)
            return False

        if self.protect_output_proj and any(k in name_l for k in ATTENTION_OUTPUT):
            self.protected_layers.append(name)
            return False

        # Target attention only filter
        if self.target_attention_only:
            attn_keywords = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(k in name_l for k in attn_keywords):
                return False

        # 2. GQA-aware skipping
        if self._is_gqa_asymmetric(name, module):
            self.skipped_layers.append(name)
            return False

        # 3. Check linear or Conv1D nature
        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__
        raw_w = getattr(base_mod, "weight", None)
        if raw_w is None or not isinstance(raw_w, torch.Tensor):
            return False

        is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_mod, "nf")
        is_linear_like = isinstance(base_mod, nn.Linear) or is_conv1d
        if not is_linear_like or isinstance(base_mod, SVDLinear):
            return False

        # 4. Tiny layer guard
        if raw_w.numel() < self.min_numel:
            self.skipped_layers.append(name)
            return False

        # 5. Respect framework compatibility rules
        if not should_compress_linear(self.model, name):
            self.protected_layers.append(name)
            return False

        # 6. User-defined explicit target modules
        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def compress_linear(
        self,
        layer: nn.Module,
        rank: int,
    ) -> SVDLinear | None:
        """Decompose a single linear or Conv1D layer into SVDLinear in FP32 with validation."""
        base_layer = getattr(layer, "base_layer", layer)
        raw_w = base_layer.weight.data
        cls_name = base_layer.__class__.__name__
        is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_layer, "nf")

        if is_conv1d:
            W = raw_w.t().contiguous().float()
            out_f, in_f = W.shape
        else:
            W = raw_w.contiguous().float()
            out_f = getattr(base_layer, "out_features", W.shape[0])
            in_f = getattr(base_layer, "in_features", W.shape[1])

        norm_orig = torch.norm(W, p="fro").item()
        if norm_orig < 1e-4:
            return None

        max_r = min(out_f, in_f)
        param_cap = max(1, (in_f * out_f) // (in_f + out_f) - 1)

        if self.min_rank > max_r - 1:
            logger.warning(
                "SDRM: min_rank=%d exceeds max_r=%d for %s; skipping",
                self.min_rank,
                max_r - 1,
                cls_name,
            )
            return None

        min_rank_effective = min(self.min_rank, max_r - 1, param_cap)
        r = max(min_rank_effective, min(rank, max_r - 1, param_cap))

        if r >= max_r:
            logger.debug(
                "SDRM: rank %d >= max_r %d; skipping layer for no parameter reduction",
                r,
                max_r,
            )
            return None

        device = base_layer.weight.device
        dtype = base_layer.weight.dtype

        try:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        except Exception as e:
            logger.warning("SDRM SVD decomposition failure: %s", e)
            return None

        sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
        U_r = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
        V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(dtype=dtype, device=device)

        has_base_layer = hasattr(layer, "base_layer")
        has_existing_lora = hasattr(layer, "lora_A") and hasattr(layer, "lora_B")
        effective_lora_r = (
            0
            if has_base_layer
            else (
                getattr(layer, "r", self.lora_r)
                if (has_existing_lora and getattr(layer, "r", 0) > 0)
                else self.lora_r
            )
        )

        svd_layer = SVDLinear(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=base_layer.bias is not None,
            dtype=dtype,
            device=device,
            lora_r=effective_lora_r if (self.use_lora or has_existing_lora) else 0,
            lora_alpha=self.lora_alpha,
        )
        svd_layer.set_low_rank_factors(U_r, V_r)
        if base_layer.bias is not None:
            svd_layer.bias.data.copy_(base_layer.bias.data)

        # Preserve existing LoRA adapters when replacing full module
        if not has_base_layer and has_existing_lora and svd_layer.use_lora:
            try:
                src_lA = getattr(layer.lora_A, "weight", layer.lora_A)
                src_lB = getattr(layer.lora_B, "weight", layer.lora_B)
                dst_lA = getattr(svd_layer.lora_A, "weight", svd_layer.lora_A)
                dst_lB = getattr(svd_layer.lora_B, "weight", svd_layer.lora_B)
                if (
                    hasattr(src_lA, "data")
                    and hasattr(dst_lA, "data")
                    and src_lA.shape == dst_lA.shape
                ):
                    dst_lA.data.copy_(src_lA.data)
                if (
                    hasattr(src_lB, "data")
                    and hasattr(dst_lB, "data")
                    and src_lB.shape == dst_lB.shape
                ):
                    dst_lB.data.copy_(src_lB.data)
            except Exception as e_lora:
                logger.warning("Could not copy existing LoRA adapters in SDRM: %s", e_lora)

        # Frobenius reconstruction error check
        W_rec = torch.matmul(U_r.float(), V_r.float())
        norm_orig_clamped = torch.norm(W, p="fro").clamp_min(1e-8)
        rel_err = (torch.norm(W - W_rec, p="fro") / norm_orig_clamped).item()

        is_invalid = torch.isnan(W_rec).any() or torch.isinf(W_rec).any()
        if rel_err > self.max_error:
            logger.warning(
                "SDRM: reconstruction error %.4f exceeded threshold %.4f.",
                rel_err,
                self.max_error,
            )
            is_invalid = True

        if is_invalid and self.fallback_to_original:
            return None

        svd_layer.reconstruction_error = rel_err
        return svd_layer

    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """Compress linear layers using budget-proportional rank assignment."""
        self.layer_errors = {}
        self.replaced_layers = []
        self.visited_modules.clear()
        self.compressed_layers.clear()
        self.protected_layers.clear()
        self.skipped_layers.clear()
        self.fallback_layers.clear()
        self.failed_layers.clear()

        target_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            logger.warning("SDRM: No compressible linear layers found.")
            return self.model, self.get_compression_stats()

        for name, module in target_layers:
            mod_id = id(module)
            if mod_id in self.visited_modules:
                logger.debug("SDRM: Skipping tied module %s", name)
                self.skipped_layers.append(name)
                continue
            self.visited_modules.add(mod_id)

            try:
                base_m = getattr(module, "base_layer", module)
                raw_w = base_m.weight.data
                cls_name = base_m.__class__.__name__
                is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_m, "nf")
                out_f = (
                    raw_w.shape[1] if is_conv1d else getattr(base_m, "out_features", raw_w.shape[0])
                )
                in_f = (
                    raw_w.shape[0] if is_conv1d else getattr(base_m, "in_features", raw_w.shape[1])
                )

                params = in_f * out_f
                allowed = params * self.total_param_ratio
                rank = max(self.min_rank, int(allowed / max(in_f + out_f, 1)))

                svd_layer = self.compress_linear(module, rank=rank)
                if svd_layer is not None:
                    err = getattr(svd_layer, "reconstruction_error", 0.0)
                    self.layer_errors[name] = err
                    self.replaced_layers.append(
                        {"name": name, "rank": svd_layer.rank, "error": err}
                    )
                    if hasattr(module, "base_layer"):
                        module.base_layer = svd_layer
                        if hasattr(module, "lora_A") and module.lora_A is not None:
                            for p in (
                                module.lora_A.parameters()
                                if hasattr(module.lora_A, "parameters")
                                else [module.lora_A]
                            ):
                                p.requires_grad = True
                        if hasattr(module, "lora_B") and module.lora_B is not None:
                            for p in (
                                module.lora_B.parameters()
                                if hasattr(module.lora_B, "parameters")
                                else [module.lora_B]
                            ):
                                p.requires_grad = True
                    else:
                        replace_module(self.model, name, svd_layer)
                    self.compressed_layers.append(name)
                else:
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
            except Exception as e_comp:
                logger.warning("Failed to compress %s: %s", name, e_comp)
                self.failed_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False

        # Post-compression forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                dummy_input = torch.randint(0, 100, (1, 8), device=self.target_device)
                self.model(dummy_input)
        except Exception as ex_sanity:
            logger.warning("SDRM post-compression forward sanity check warning: %s", ex_sanity)
            forward_check_success = False

        self._compressed_params = self._count_params(self.model)
        avg_err = (
            (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
        )
        stats = self.get_compression_stats()
        stats.update(
            {
                "total_param_ratio": self.total_param_ratio,
                "compressed_layers": self.compressed_layers,
                "compressed_layer_count": len(self.compressed_layers),
                "protected_layers": self.protected_layers,
                "protected_layer_count": len(self.protected_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "failed_layers": self.failed_layers,
                "failed_layer_count": len(self.failed_layers),
                "layer_errors": self.layer_errors,
                "avg_reconstruction_error": avg_err,
                "forward_check_success": forward_check_success,
                "min_rank": self.min_rank,
                "skip_attention_qkv": self.skip_attention_qkv,
                "target_attention_only": self.target_attention_only,
                "use_lora": self.use_lora,
                "lora_r": self.lora_r,
                "method": "sdrm",
            }
        )
        self.stats = stats
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: SDRM: Scaled Dimension Rank Modulation. Assigns layer rank proportional to parameter budget for edge inference and factorizes target layers into low-rank SVDLinear modules with optional LoRA fine-tuning adapters.

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """Compress linear layers using budget-proportional rank assignment."""
    self.layer_errors = {}
    self.replaced_layers = []
    self.visited_modules.clear()
    self.compressed_layers.clear()
    self.protected_layers.clear()
    self.skipped_layers.clear()
    self.fallback_layers.clear()
    self.failed_layers.clear()

    target_layers = [
        (name, mod)
        for name, mod in self.model.named_modules()
        if self._is_target_module(name, mod)
    ]

    if not target_layers:
        logger.warning("SDRM: No compressible linear layers found.")
        return self.model, self.get_compression_stats()

    for name, module in target_layers:
        mod_id = id(module)
        if mod_id in self.visited_modules:
            logger.debug("SDRM: Skipping tied module %s", name)
            self.skipped_layers.append(name)
            continue
        self.visited_modules.add(mod_id)

        try:
            base_m = getattr(module, "base_layer", module)
            raw_w = base_m.weight.data
            cls_name = base_m.__class__.__name__
            is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_m, "nf")
            out_f = (
                raw_w.shape[1] if is_conv1d else getattr(base_m, "out_features", raw_w.shape[0])
            )
            in_f = (
                raw_w.shape[0] if is_conv1d else getattr(base_m, "in_features", raw_w.shape[1])
            )

            params = in_f * out_f
            allowed = params * self.total_param_ratio
            rank = max(self.min_rank, int(allowed / max(in_f + out_f, 1)))

            svd_layer = self.compress_linear(module, rank=rank)
            if svd_layer is not None:
                err = getattr(svd_layer, "reconstruction_error", 0.0)
                self.layer_errors[name] = err
                self.replaced_layers.append(
                    {"name": name, "rank": svd_layer.rank, "error": err}
                )
                if hasattr(module, "base_layer"):
                    module.base_layer = svd_layer
                    if hasattr(module, "lora_A") and module.lora_A is not None:
                        for p in (
                            module.lora_A.parameters()
                            if hasattr(module.lora_A, "parameters")
                            else [module.lora_A]
                        ):
                            p.requires_grad = True
                    if hasattr(module, "lora_B") and module.lora_B is not None:
                        for p in (
                            module.lora_B.parameters()
                            if hasattr(module.lora_B, "parameters")
                            else [module.lora_B]
                        ):
                            p.requires_grad = True
                else:
                    replace_module(self.model, name, svd_layer)
                self.compressed_layers.append(name)
            else:
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
        except Exception as e_comp:
            logger.warning("Failed to compress %s: %s", name, e_comp)
            self.failed_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False

    # Post-compression forward sanity check
    forward_check_success = True
    try:
        with torch.no_grad():
            dummy_input = torch.randint(0, 100, (1, 8), device=self.target_device)
            self.model(dummy_input)
    except Exception as ex_sanity:
        logger.warning("SDRM post-compression forward sanity check warning: %s", ex_sanity)
        forward_check_success = False

    self._compressed_params = self._count_params(self.model)
    avg_err = (
        (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
    )
    stats = self.get_compression_stats()
    stats.update(
        {
            "total_param_ratio": self.total_param_ratio,
            "compressed_layers": self.compressed_layers,
            "compressed_layer_count": len(self.compressed_layers),
            "protected_layers": self.protected_layers,
            "protected_layer_count": len(self.protected_layers),
            "skipped_layers": self.skipped_layers,
            "skipped_layer_count": len(self.skipped_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "failed_layers": self.failed_layers,
            "failed_layer_count": len(self.failed_layers),
            "layer_errors": self.layer_errors,
            "avg_reconstruction_error": avg_err,
            "forward_check_success": forward_check_success,
            "min_rank": self.min_rank,
            "skip_attention_qkv": self.skip_attention_qkv,
            "target_attention_only": self.target_attention_only,
            "use_lora": self.use_lora,
            "lora_r": self.lora_r,
            "method": "sdrm",
        }
    )
    self.stats = stats
    return self.model, stats

Compress linear layers using budget-proportional rank assignment.

def compress_linear(self, layer: nn.Module, rank: int) ‑> SVDLinear | None

Expand source code

def compress_linear(
    self,
    layer: nn.Module,
    rank: int,
) -> SVDLinear | None:
    """Decompose a single linear or Conv1D layer into SVDLinear in FP32 with validation."""
    base_layer = getattr(layer, "base_layer", layer)
    raw_w = base_layer.weight.data
    cls_name = base_layer.__class__.__name__
    is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_layer, "nf")

    if is_conv1d:
        W = raw_w.t().contiguous().float()
        out_f, in_f = W.shape
    else:
        W = raw_w.contiguous().float()
        out_f = getattr(base_layer, "out_features", W.shape[0])
        in_f = getattr(base_layer, "in_features", W.shape[1])

    norm_orig = torch.norm(W, p="fro").item()
    if norm_orig < 1e-4:
        return None

    max_r = min(out_f, in_f)
    param_cap = max(1, (in_f * out_f) // (in_f + out_f) - 1)

    if self.min_rank > max_r - 1:
        logger.warning(
            "SDRM: min_rank=%d exceeds max_r=%d for %s; skipping",
            self.min_rank,
            max_r - 1,
            cls_name,
        )
        return None

    min_rank_effective = min(self.min_rank, max_r - 1, param_cap)
    r = max(min_rank_effective, min(rank, max_r - 1, param_cap))

    if r >= max_r:
        logger.debug(
            "SDRM: rank %d >= max_r %d; skipping layer for no parameter reduction",
            r,
            max_r,
        )
        return None

    device = base_layer.weight.device
    dtype = base_layer.weight.dtype

    try:
        U, S, Vt = torch.linalg.svd(W, full_matrices=False)
    except Exception as e:
        logger.warning("SDRM SVD decomposition failure: %s", e)
        return None

    sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
    U_r = (U[:, :r] * sqrt_S.unsqueeze(0)).to(dtype=dtype, device=device)
    V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(dtype=dtype, device=device)

    has_base_layer = hasattr(layer, "base_layer")
    has_existing_lora = hasattr(layer, "lora_A") and hasattr(layer, "lora_B")
    effective_lora_r = (
        0
        if has_base_layer
        else (
            getattr(layer, "r", self.lora_r)
            if (has_existing_lora and getattr(layer, "r", 0) > 0)
            else self.lora_r
        )
    )

    svd_layer = SVDLinear(
        in_features=in_f,
        out_features=out_f,
        rank=r,
        bias=base_layer.bias is not None,
        dtype=dtype,
        device=device,
        lora_r=effective_lora_r if (self.use_lora or has_existing_lora) else 0,
        lora_alpha=self.lora_alpha,
    )
    svd_layer.set_low_rank_factors(U_r, V_r)
    if base_layer.bias is not None:
        svd_layer.bias.data.copy_(base_layer.bias.data)

    # Preserve existing LoRA adapters when replacing full module
    if not has_base_layer and has_existing_lora and svd_layer.use_lora:
        try:
            src_lA = getattr(layer.lora_A, "weight", layer.lora_A)
            src_lB = getattr(layer.lora_B, "weight", layer.lora_B)
            dst_lA = getattr(svd_layer.lora_A, "weight", svd_layer.lora_A)
            dst_lB = getattr(svd_layer.lora_B, "weight", svd_layer.lora_B)
            if (
                hasattr(src_lA, "data")
                and hasattr(dst_lA, "data")
                and src_lA.shape == dst_lA.shape
            ):
                dst_lA.data.copy_(src_lA.data)
            if (
                hasattr(src_lB, "data")
                and hasattr(dst_lB, "data")
                and src_lB.shape == dst_lB.shape
            ):
                dst_lB.data.copy_(src_lB.data)
        except Exception as e_lora:
            logger.warning("Could not copy existing LoRA adapters in SDRM: %s", e_lora)

    # Frobenius reconstruction error check
    W_rec = torch.matmul(U_r.float(), V_r.float())
    norm_orig_clamped = torch.norm(W, p="fro").clamp_min(1e-8)
    rel_err = (torch.norm(W - W_rec, p="fro") / norm_orig_clamped).item()

    is_invalid = torch.isnan(W_rec).any() or torch.isinf(W_rec).any()
    if rel_err > self.max_error:
        logger.warning(
            "SDRM: reconstruction error %.4f exceeded threshold %.4f.",
            rel_err,
            self.max_error,
        )
        is_invalid = True

    if is_invalid and self.fallback_to_original:
        return None

    svd_layer.reconstruction_error = rel_err
    return svd_layer

Decompose a single linear or Conv1D layer into SVDLinear in FP32 with validation.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class SLiMCompressor (model: nn.Module, bits: int = 4, sparsity: float = 0.5, rank: int = 8, structured_24: bool = True, target_attention_only: bool = False, max_error: float = 0.5, fallback_to_original: bool = True, target_modules: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("slim")
@COMPRESSOR_REGISTRY.register("slim_compressor")
class SLiMCompressor(BaseCompressor):
    is_custom_attention: bool = True
    """
    SLiM Compressor: One-shot joint b-bit quantization, 2:4 structured sparsity, and low-rank residual recovery.
    """

    def __init__(
        self,
        model: nn.Module,
        bits: int = 4,
        sparsity: float = 0.5,
        rank: int = 8,
        structured_24: bool = True,
        target_attention_only: bool = False,
        max_error: float = 0.5,
        fallback_to_original: bool = True,  # Fix 1: Default changed to True
        target_modules: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="slim", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            bits = _cfg_get(cfg, "slim_bits", _cfg_get(cfg, "bits", bits))
            sparsity = _cfg_get(cfg, "slim_sparsity", _cfg_get(cfg, "sparsity", sparsity))
            rank = _cfg_get(cfg, "slim_rank", _cfg_get(cfg, "rank", rank))
            structured_24 = _cfg_get(
                cfg, "slim_structured_24", _cfg_get(cfg, "structured_24", structured_24)
            )
            max_error = _cfg_get(cfg, "max_error", max_error)
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)

        if "fallback_to_original" in kwargs:
            fallback_to_original = bool(kwargs["fallback_to_original"])

        self.bits = int(bits)
        self.sparsity = max(0.0, min(float(sparsity), 0.95))
        self.rank = max(1, int(rank))
        self.structured_24 = structured_24
        self.target_attention_only = bool(target_attention_only)
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.target_modules = target_modules
        self.layer_errors: dict[str, float] = {}
        self.stats: dict[str, Any] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        """
        Fix 4: Unified critical-layer protection delegated to should_compress_linear.
        Eliminates duplicate keyword matching.
        """
        name_l = name.lower()
        if self.target_attention_only:
            attn_keywords = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(k in name_l for k in attn_keywords):
                return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__.lower()
        # Fix 2: Reliable Conv1D check using class name and hasattr(nf)
        is_linear_like = (
            isinstance(base_mod, nn.Linear) or "conv1d" in cls_name or hasattr(base_mod, "nf")
        )
        if not is_linear_like:
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    def _quantize(self, W: torch.Tensor) -> torch.Tensor:
        """Quantizes weight to b-bit representation (biases remain unquantized)."""
        if self.bits >= 32:
            return W
        min_val = W.min()
        max_val = W.max()
        if (max_val - min_val).abs() < 1e-6:
            max_val = min_val + 1e-6
        scale = (max_val - min_val) / (2.0**self.bits - 1.0)
        W_q = torch.round((W - min_val) / scale) * scale + min_val
        return W_q

    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """Executes SLiM compression pipeline across all Linear and Conv1D layers."""
        self.layer_errors = {}
        pruned_layers_count = 0
        total_svd_energy = 0.0

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                base_mod = getattr(module, "base_layer", module)
                raw_w = base_mod.weight.data
                cls_name = base_mod.__class__.__name__.lower()
                # Fix 2: Check for Conv1D via class name and nf attribute
                is_conv1d = "conv1d" in cls_name or hasattr(base_mod, "nf")
                device = raw_w.device
                dtype = raw_w.dtype

                W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()

                # 1. Quantize (Weights only; biases are preserved in full precision to avoid drift)
                W_q = self._quantize(W)

                # 2. Apply 2:4 structured sparsity or magnitude sparsity
                if self.structured_24:
                    W_sp, _ = apply_24_structured_mask(W_q, saliency=torch.abs(W_q), dim=1)
                else:
                    flat = W_q.abs().reshape(-1)
                    k = max(1, int(flat.numel() * (1.0 - self.sparsity)))
                    thr = torch.topk(flat, k, largest=True).values[-1]
                    mask = (W_q.abs() >= thr).float()
                    W_sp = W_q * mask

                # 3. Low-Rank Residual Correction
                residual = W - W_sp
                try:
                    U, S, Vh = torch.linalg.svd(residual, full_matrices=False)
                    r = min(self.rank, len(S))
                    if r > 0:
                        correction = torch.matmul(U[:, :r] * S[:r].unsqueeze(0), Vh[:r, :])
                        total_svd_energy += torch.norm(correction).item()
                        W_final = W_sp + correction
                    else:
                        W_final = W_sp
                except Exception as e:
                    logger.debug("SVD residual correction exception in %s: %s", name, e)
                    W_final = W_sp

                # Relative Frobenius error check
                norm_orig = torch.norm(W, p="fro")
                rel_err = (torch.norm(W - W_final, p="fro") / (norm_orig + 1e-8)).item()
                self.layer_errors[name] = rel_err

                # Fix 3: Always log a warning when reconstruction error exceeds max_error
                if rel_err > self.max_error:
                    logger.warning(
                        "SLiM reconstruction error %.4f on %s exceeded threshold %.4f%s",
                        rel_err,
                        name,
                        self.max_error,
                        (
                            ". Skipping (fallback to original)."
                            if self.fallback_to_original
                            else ". Proceeding with compressed weight."
                        ),
                    )
                    if self.fallback_to_original:
                        logger.warning(
                            "Fallback triggered for %s. Layer is kept original and NOT wrapped in LoRA. It will not be fine-tuned.",
                            name,
                        )
                        for p in module.parameters():
                            p.requires_grad = False
                        continue

                final_w = W_final.t().contiguous() if is_conv1d else W_final
                base_mod.weight.data = final_w.to(dtype=dtype, device=device)
                compressed_layer = base_mod
                # Preserve existing LoRA adapter weights if the source module has them
                if (
                    hasattr(module, "lora_A")
                    and hasattr(module, "lora_B")
                    and module.lora_A is not None
                    and module.lora_B is not None
                ):
                    src_lA = getattr(module.lora_A, "weight", module.lora_A)
                    src_lB = getattr(module.lora_B, "weight", module.lora_B)
                    if hasattr(compressed_layer, "lora_A") and compressed_layer.lora_A is not None:
                        if (
                            hasattr(src_lA, "data")
                            and src_lA.shape == compressed_layer.lora_A.shape
                        ):
                            compressed_layer.lora_A.data.copy_(src_lA.data)
                    if hasattr(compressed_layer, "lora_B") and compressed_layer.lora_B is not None:
                        if (
                            hasattr(src_lB, "data")
                            and src_lB.shape == compressed_layer.lora_B.shape
                        ):
                            compressed_layer.lora_B.data.copy_(src_lB.data)

                # Determine if module is wrapped (LoRALinear) or standalone
                is_wrapped = hasattr(module, "base_layer")
                if is_wrapped:
                    module.base_layer = compressed_layer
                    new_module = module
                else:
                    new_module = compressed_layer
                    replace_module(self.model, name, new_module)
                pruned_layers_count += 1

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        # Fix 6: Store max_error and fallback_to_original in stats
        stats.update(
            {
                "pruned_layers": pruned_layers_count,
                "quant_bits": self.bits,
                "rank": self.rank,
                "structured_24": self.structured_24,
                "total_svd_energy": total_svd_energy,
                "layer_errors": self.layer_errors,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "target_attention_only": self.target_attention_only,
                "biases_quantized": False,  # Fix 5: Explicitly document bias policy
                "method": "slim_quant_sparse_lowrank",
            }
        )
        self.stats = stats
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: SLiM Compressor: One-shot joint b-bit quantization, 2:4 structured sparsity, and low-rank residual recovery.

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """Executes SLiM compression pipeline across all Linear and Conv1D layers."""
    self.layer_errors = {}
    pruned_layers_count = 0
    total_svd_energy = 0.0

    for name, module in self.model.named_modules():
        if self._is_target_module(name, module):
            base_mod = getattr(module, "base_layer", module)
            raw_w = base_mod.weight.data
            cls_name = base_mod.__class__.__name__.lower()
            # Fix 2: Check for Conv1D via class name and nf attribute
            is_conv1d = "conv1d" in cls_name or hasattr(base_mod, "nf")
            device = raw_w.device
            dtype = raw_w.dtype

            W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()

            # 1. Quantize (Weights only; biases are preserved in full precision to avoid drift)
            W_q = self._quantize(W)

            # 2. Apply 2:4 structured sparsity or magnitude sparsity
            if self.structured_24:
                W_sp, _ = apply_24_structured_mask(W_q, saliency=torch.abs(W_q), dim=1)
            else:
                flat = W_q.abs().reshape(-1)
                k = max(1, int(flat.numel() * (1.0 - self.sparsity)))
                thr = torch.topk(flat, k, largest=True).values[-1]
                mask = (W_q.abs() >= thr).float()
                W_sp = W_q * mask

            # 3. Low-Rank Residual Correction
            residual = W - W_sp
            try:
                U, S, Vh = torch.linalg.svd(residual, full_matrices=False)
                r = min(self.rank, len(S))
                if r > 0:
                    correction = torch.matmul(U[:, :r] * S[:r].unsqueeze(0), Vh[:r, :])
                    total_svd_energy += torch.norm(correction).item()
                    W_final = W_sp + correction
                else:
                    W_final = W_sp
            except Exception as e:
                logger.debug("SVD residual correction exception in %s: %s", name, e)
                W_final = W_sp

            # Relative Frobenius error check
            norm_orig = torch.norm(W, p="fro")
            rel_err = (torch.norm(W - W_final, p="fro") / (norm_orig + 1e-8)).item()
            self.layer_errors[name] = rel_err

            # Fix 3: Always log a warning when reconstruction error exceeds max_error
            if rel_err > self.max_error:
                logger.warning(
                    "SLiM reconstruction error %.4f on %s exceeded threshold %.4f%s",
                    rel_err,
                    name,
                    self.max_error,
                    (
                        ". Skipping (fallback to original)."
                        if self.fallback_to_original
                        else ". Proceeding with compressed weight."
                    ),
                )
                if self.fallback_to_original:
                    logger.warning(
                        "Fallback triggered for %s. Layer is kept original and NOT wrapped in LoRA. It will not be fine-tuned.",
                        name,
                    )
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

            final_w = W_final.t().contiguous() if is_conv1d else W_final
            base_mod.weight.data = final_w.to(dtype=dtype, device=device)
            compressed_layer = base_mod
            # Preserve existing LoRA adapter weights if the source module has them
            if (
                hasattr(module, "lora_A")
                and hasattr(module, "lora_B")
                and module.lora_A is not None
                and module.lora_B is not None
            ):
                src_lA = getattr(module.lora_A, "weight", module.lora_A)
                src_lB = getattr(module.lora_B, "weight", module.lora_B)
                if hasattr(compressed_layer, "lora_A") and compressed_layer.lora_A is not None:
                    if (
                        hasattr(src_lA, "data")
                        and src_lA.shape == compressed_layer.lora_A.shape
                    ):
                        compressed_layer.lora_A.data.copy_(src_lA.data)
                if hasattr(compressed_layer, "lora_B") and compressed_layer.lora_B is not None:
                    if (
                        hasattr(src_lB, "data")
                        and src_lB.shape == compressed_layer.lora_B.shape
                    ):
                        compressed_layer.lora_B.data.copy_(src_lB.data)

            # Determine if module is wrapped (LoRALinear) or standalone
            is_wrapped = hasattr(module, "base_layer")
            if is_wrapped:
                module.base_layer = compressed_layer
                new_module = module
            else:
                new_module = compressed_layer
                replace_module(self.model, name, new_module)
            pruned_layers_count += 1

    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    # Fix 6: Store max_error and fallback_to_original in stats
    stats.update(
        {
            "pruned_layers": pruned_layers_count,
            "quant_bits": self.bits,
            "rank": self.rank,
            "structured_24": self.structured_24,
            "total_svd_energy": total_svd_energy,
            "layer_errors": self.layer_errors,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "target_attention_only": self.target_attention_only,
            "biases_quantized": False,  # Fix 5: Explicitly document bias policy
            "method": "slim_quant_sparse_lowrank",
        }
    )
    self.stats = stats
    return self.model, stats

Executes SLiM compression pipeline across all Linear and Conv1D layers.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class STEFakeQuantize (*args, **kwargs)

Expand source code

class STEFakeQuantize(torch.autograd.Function):
    """
    Straight-Through Estimator (STE) fake-quantization autograd function.
    Clamps inputs to [qmin, qmax] in discrete integer space, de-quantizes back,
    and passes gradients through uninhibited across all channels.
    """

    @staticmethod
    def forward(
        ctx: Any,
        x: torch.Tensor,
        scale: torch.Tensor,
        zero_point: torch.Tensor,
        qmin: float | torch.Tensor,
        qmax: float | torch.Tensor,
    ) -> torch.Tensor:
        ctx.save_for_backward(x, scale, zero_point)
        # Quantize, clamp to discrete range, and dequantize
        q = torch.clamp(torch.round(x / scale) + zero_point, qmin, qmax)
        dq = (q - zero_point) * scale
        return dq

    @staticmethod
    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor | None, ...]:
        # Straight-Through Estimator passes gradient unmodified
        grad_x = grad_output.clone()
        return grad_x, None, None, None, None

Straight-Through Estimator (STE) fake-quantization autograd function. Clamps inputs to [qmin, qmax] in discrete integer space, de-quantizes back, and passes gradients through uninhibited across all channels.

Ancestors

torch.autograd.function.Function
torch.autograd.function._SingleLevelFunction
torch._C._FunctionBase
torch.autograd.function.FunctionCtx
torch.autograd.function._HookMixin

Static methods

def backward(ctx: Any, grad_output: torch.Tensor) ‑> tuple[torch.Tensor | None, ...]

Expand source code

@staticmethod
def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor | None, ...]:
    # Straight-Through Estimator passes gradient unmodified
    grad_x = grad_output.clone()
    return grad_x, None, None, None, None

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the vjp function.)

def forward(ctx: Any, x: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, qmin: float | torch.Tensor, qmax: float | torch.Tensor) ‑> torch.Tensor

Expand source code

@staticmethod
def forward(
    ctx: Any,
    x: torch.Tensor,
    scale: torch.Tensor,
    zero_point: torch.Tensor,
    qmin: float | torch.Tensor,
    qmax: float | torch.Tensor,
) -> torch.Tensor:
    ctx.save_for_backward(x, scale, zero_point)
    # Quantize, clamp to discrete range, and dequantize
    q = torch.clamp(torch.round(x / scale) + zero_point, qmin, qmax)
    dq = (q - zero_point) * scale
    return dq

Define the forward of the custom autograd Function.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx)::

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See :ref:combining-forward-context for more details

Usage 2 (Separate forward and ctx)::

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass


@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

The forward no longer accepts a ctx argument.
Instead, you must also override the :meth:torch.autograd.Function.setup_context staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.
See :ref:extending-autograd for more details

class SVDLLMCompressor (model: nn.Module, rank_ratio: float = 0.2, rank: int | None = None, lora_rank: int = 16, lora_alpha: float = 32.0, max_error: float = 0.5, fallback_to_original: bool = False, calibration_texts: list[str] | None = None, tokenizer: Any | None = None, device: str | torch.device | None = None, target_attention_only: bool = True, recovery_steps: int = 50, learning_rate: float = 0.0002, weight_decay: float = 0.01, batch_size: int = 8, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("svd_llm")
@COMPRESSOR_REGISTRY.register("svd_llm_compressor")
class SVDLLMCompressor(BaseCompressor):
    """
    SVD-LLM (ICLR 2025) Compressor with Cholesky activation whitening,
    sequential LoRA recovery fine-tuning, and fixed rank parameter ratios.
    """

    def __init__(
        self,
        model: nn.Module,
        rank_ratio: float = 0.20,
        rank: int | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        max_error: float = 0.50,
        fallback_to_original: bool = False,
        calibration_texts: list[str] | None = None,
        tokenizer: Any | None = None,
        device: str | torch.device | None = None,
        target_attention_only: bool = True,
        recovery_steps: int = 50,
        learning_rate: float = 2e-4,
        weight_decay: float = 0.01,
        batch_size: int = 8,
        **kwargs: Any,
    ) -> None:
        super().__init__(model=model, name="svd_llm", **kwargs)
        self.rank_ratio = rank_ratio
        self.explicit_rank = rank
        self.lora_rank = lora_rank
        self.lora_alpha = lora_alpha
        self.max_error = max_error
        self.fallback_to_original = fallback_to_original
        self.target_attention_only = target_attention_only
        self.recovery_steps = recovery_steps
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.batch_size = batch_size

        if device is not None:
            self.device = torch.device(device)
        else:
            p = next(model.parameters(), None)
            self.device = p.device if p is not None else torch.device("cpu")

        self.tokenizer = tokenizer
        if calibration_texts is not None:
            self.calibration_texts = calibration_texts
        else:
            self.calibration_texts = self._load_calibration_data(n_samples=256)

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}

    def _load_calibration_data(self, n_samples: int = 256) -> list[str]:
        """
        Load calibration samples from WikiText-2 (or Salesforce/wikitext fallback),
        following SVD-LLM (ICLR 2025) paper protocol. Falls back to default corpus if offline.
        """
        try:
            from datasets import load_dataset

            try:
                ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
            except Exception:
                ds = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="train")
            texts = [t.strip() for t in ds["text"] if t and len(t.strip()) > 50]
            if len(texts) >= n_samples:
                logger.info(
                    "Loaded %d calibration samples from WikiText-2 for SVD-LLM.",
                    n_samples,
                )
                return texts[:n_samples]
            elif texts:
                return texts
        except Exception as e:
            logger.warning(
                "Could not download WikiText-2 calibration dataset (%s); using default corpus.",
                e,
            )
        return DEFAULT_CALIBRATION_CORPUS

    def compute_rank_from_ratio(self, m: int, n: int, ratio: float) -> int:
        """
        Compute low rank r from parameter retention ratio R:
        Original params = m * n
        Target params = R * m * n
        Factorized params = r * (m + n)
        r = (R * m * n) / (m + n)
        """
        if self.explicit_rank is not None and self.explicit_rank > 0:
            return min(self.explicit_rank, min(m, n))
        r_calc = (ratio * m * n) / (m + n)
        r = int(round(r_calc))
        return max(1, min(r, min(m, n)))

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        if isinstance(module, SVDLLMLinear):
            return False
        if any(
            name.endswith(sfx) for sfx in [".base_layer", ".lora_A", ".lora_B", ".lora_dropout"]
        ):
            return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = module.__class__.__name__.lower()
        target_cls = target_mod.__class__.__name__.lower()

        is_linear_like = (
            isinstance(module, nn.Linear)
            or isinstance(target_mod, nn.Linear)
            or "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(module, "nf")
            or hasattr(target_mod, "nf")
        )
        if not is_linear_like:
            return False

        name_lower = name.lower()
        if any(
            skip in name_lower
            for skip in [
                "lm_head",
                "classifier",
                "head",
                "output_layer",
                "score",
                "pooler",
            ]
        ):
            return False

        if self.target_attention_only:
            ffn_patterns = [
                "ffn",
                "mlp",
                "lin1",
                "lin2",
                "intermediate",
                "gate_proj",
                "up_proj",
                "down_proj",
                "c_fc",
            ]
            if any(p in name_lower for p in ffn_patterns):
                return False
            if "output.dense" in name_lower and "attention" not in name_lower:
                return False
            attn_patterns = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "out_lin",
                "q_lin",
                "k_lin",
                "v_lin",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(p in name_lower for p in attn_patterns):
                return False
        else:
            if not should_compress_linear(self.model, name):
                return False

        return True

    def _collect_covariances(self) -> dict[str, torch.Tensor]:
        """
        Collect empirical input activation covariances C = (1/N) X^T X via forward hooks.
        Uses streaming accumulation to avoid storing large activation buffers in memory.
        """
        covariances: dict[str, torch.Tensor] = {}
        counts: dict[str, int] = {}
        hooks = []

        def make_hook(name: str):
            def hook(module: nn.Module, inp: tuple[torch.Tensor, ...], outp: Any):
                if not inp or inp[0] is None:
                    return
                x = inp[0].detach()
                dev_type = x.device.type
                autocast_ctx = (
                    torch.autocast(device_type=dev_type, enabled=False)
                    if hasattr(torch, "autocast") and dev_type in ("cuda", "cpu")
                    else contextlib.nullcontext()
                )
                with autocast_ctx:
                    target_mod = getattr(module, "base_layer", module)
                    cls_name = module.__class__.__name__.lower()
                    target_cls = target_mod.__class__.__name__.lower()
                    is_conv1d = (
                        "conv1d" in cls_name
                        or "conv1d" in target_cls
                        or hasattr(module, "nf")
                        or hasattr(target_mod, "nf")
                    )
                    weight_tensor = getattr(target_mod, "weight", getattr(module, "weight", None))
                    if weight_tensor is None:
                        return
                    in_f = weight_tensor.shape[0] if is_conv1d else weight_tensor.shape[1]
                    x_fp32 = x.float()
                    x_flat = x_fp32.reshape(-1, in_f)
                    mask = x_flat.abs().sum(dim=-1) > 1e-5
                    if mask.sum() == 0:
                        return
                    x_active = x_flat[mask]
                    batch_cov = torch.matmul(x_active.T, x_active).cpu()
                    n_samples = x_active.size(0)

                if name not in covariances:
                    covariances[name] = batch_cov
                    counts[name] = n_samples
                else:
                    covariances[name] += batch_cov
                    counts[name] += n_samples

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                hooks.append(module.register_forward_hook(make_hook(name)))

        if self.tokenizer is not None and self.calibration_texts:
            try:
                inputs = self.tokenizer(
                    self.calibration_texts,
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                    max_length=512,
                )
                inputs = {k: v.to(self.device) for k, v in inputs.items()}
                with torch.no_grad():
                    self.model(**inputs)
            except Exception as e:
                logger.warning("SVD-LLM calibration forward pass exception: %s", e)

        for h in hooks:
            h.remove()

        final_cov: dict[str, torch.Tensor] = {}
        for name, cov in covariances.items():
            n = max(1, counts.get(name, 1))
            final_cov[name] = cov / float(n)

        return final_cov

    def _robust_cholesky(self, C: torch.Tensor, n: int) -> torch.Tensor:
        """
        Compute Cholesky factor S where S = cholesky(C + eps * I).
        Paper specification: eps = 1e-6 * mean(diag(C))
        Uses adaptive regularisation if C is singular or ill-conditioned.
        """
        mean_diag = torch.mean(torch.diag(C)).item()
        for exp in [-6, -5, -4, -3, -2, -1, 0]:
            eps = max(10**exp * mean_diag, 10**exp)
            try:
                S = torch.linalg.cholesky(C + eps * torch.eye(n, device=C.device, dtype=C.dtype))
                return S
            except Exception:
                continue

        # Eigen-fallback for extreme singularity
        eigvals, eigvecs = torch.linalg.eigh(C)
        eigvals = torch.clamp(eigvals, min=1e-6)
        return eigvecs @ torch.diag(torch.sqrt(eigvals))

    def _cholesky_whiten_and_decompose(
        self,
        W: torch.Tensor,
        C: torch.Tensor | None,
        target_rank: int,
    ) -> tuple[torch.Tensor, torch.Tensor, float]:
        """
        Full SVD-LLM core algorithm (Section 3.1):
        1. S = cholesky(C + eps * I)
        2. W_white = W @ S
        3. U, Sigma, Vt = svd(W_white)
        4. Truncate to rank r
        5. Vt_unw = Vt_r @ inv(S)
        6. W_u = U_r * sqrt(Sigma_r), W_v = sqrt(Sigma_r) * Vt_unw
        7. Relative Frobenius error: ||W - W_u @ W_v||_F / ||W||_F
        """
        m, n = W.shape
        W_fp32 = W.float()

        if C is not None and C.shape == (n, n):
            C_fp32 = C.float().to(device=W.device)
            C_sym = 0.5 * (C_fp32 + C_fp32.T)
            S = self._robust_cholesky(C_sym, n)
            try:
                S_inv = torch.linalg.pinv(S)
            except Exception:
                S_inv = torch.inverse(S + 1e-5 * torch.eye(n, device=S.device, dtype=S.dtype))
            W_white = W_fp32 @ S
        else:
            # Identity whitening (standard SVD fallback if covariance unavailable)
            S = torch.eye(n, device=W.device, dtype=W_fp32.dtype)
            S_inv = torch.eye(n, device=W.device, dtype=W_fp32.dtype)
            W_white = W_fp32

        U, Sigma, Vt = torch.linalg.svd(W_white, full_matrices=False)
        r = min(target_rank, min(m, n), U.size(1), Vt.size(0))

        U_r = U[:, :r]
        Sigma_r = Sigma[:r]
        Vt_r = Vt[:r, :]

        Vt_unw = Vt_r @ S_inv
        sqrt_sigma = torch.sqrt(torch.clamp(Sigma_r, min=1e-12))
        W_u = U_r * sqrt_sigma.unsqueeze(0)
        W_v = sqrt_sigma.unsqueeze(1) * Vt_unw

        W_rec = W_u @ W_v
        orig_norm = torch.norm(W_fp32, p="fro").clamp(min=1e-12)
        error = (torch.norm(W_fp32 - W_rec, p="fro") / orig_norm).item()

        return W_u, W_v, error

    def _compute_recovery_loss(
        self, batch_ids: torch.Tensor, batch_mask: torch.Tensor
    ) -> torch.Tensor | None:
        """
        Compute recovery training loss across both CausalLM and Sequence Classification architectures.
        - CausalLM (GPT-2, LLaMA): cross-entropy autoregressive loss with labels=batch_ids.
        - Sequence Classification (DistilBERT, RoBERTa): unsupervised entropy minimization over predicted logits.
        """
        try:
            # 1. CausalLM path
            if (
                hasattr(self.model, "lm_head")
                or "causallm" in self.model.__class__.__name__.lower()
            ):
                outputs = self.model(
                    input_ids=batch_ids, attention_mask=batch_mask, labels=batch_ids
                )
                loss = outputs.loss if hasattr(outputs, "loss") else None
                if loss is None and isinstance(outputs, tuple):
                    loss = outputs[0]
                return loss

            # 2. Sequence Classification / generic encoder path: entropy minimization
            outputs = self.model(input_ids=batch_ids, attention_mask=batch_mask)
            logits = getattr(outputs, "logits", None)
            if logits is None and isinstance(outputs, tuple):
                logits = outputs[0]
            if logits is not None and isinstance(logits, torch.Tensor) and logits.ndim >= 2:
                probs = torch.softmax(logits.float(), dim=-1)
                entropy = -(probs * torch.log(probs + 1e-10)).sum(dim=-1).mean()
                return entropy
            return None
        except Exception as e:
            logger.debug("Error computing recovery loss: %s", e)
            return None

    def _sequential_lora_recovery(self) -> None:
        """
        Execute 2-Stage Sequential LoRA fine-tuning (Section 3.2, Eq. 7):
        Stage 1: Update W_u with LoRA (W_v frozen)
        Stage 2: Update W_v with LoRA (W_u frozen)
        """
        if self.recovery_steps <= 0 or self.lora_rank <= 0:
            return
        if self.tokenizer is None or not self.calibration_texts:
            return

        logger.info(
            "[SVD-LLM] Starting 2-Stage Sequential LoRA Recovery (%d steps/stage)...",
            self.recovery_steps,
        )
        device = self.device

        # Step 0: Freeze ALL model parameters (embeddings, layernorms, heads, and base weights)
        for p in self.model.parameters():
            p.requires_grad = False

        # Prepare calibration DataLoader
        encodings = self.tokenizer(
            self.calibration_texts,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors="pt",
        )
        dataset = torch.utils.data.TensorDataset(
            encodings["input_ids"], encodings["attention_mask"]
        )
        loader = torch.utils.data.DataLoader(
            dataset, batch_size=min(self.batch_size, len(dataset)), shuffle=True
        )

        for stage, stage_name in [
            ("stage1", "Stage 1 (Update W_u)"),
            ("stage2", "Stage 2 (Update W_v)"),
        ]:
            logger.info(f"[SVD-LLM Sequential LoRA] {stage_name}...")
            # Set active LoRA adapters
            for m in self.model.modules():
                if isinstance(m, SVDLLMLinear):
                    m.set_active_lora_stage(stage)

            trainable_params = [p for p in self.model.parameters() if p.requires_grad]
            if not trainable_params:
                continue

            optimizer = torch.optim.AdamW(
                trainable_params, lr=self.learning_rate, weight_decay=self.weight_decay
            )
            step = 0
            self.model.train()

            while step < self.recovery_steps:
                for batch_ids, batch_mask in loader:
                    if step >= self.recovery_steps:
                        break
                    batch_ids = batch_ids.to(device)
                    batch_mask = batch_mask.to(device)

                    optimizer.zero_grad()
                    loss = self._compute_recovery_loss(batch_ids, batch_mask)

                    if loss is not None and not torch.isnan(loss):
                        loss.backward()
                        torch.nn.utils.clip_grad_norm_(trainable_params, 1.0)
                        optimizer.step()
                    step += 1

        # Freeze all adapters after recovery
        for p in self.model.parameters():
            p.requires_grad = False
        for m in self.model.modules():
            if isinstance(m, SVDLLMLinear):
                m.set_active_lora_stage(None)
        self.model.eval()
        logger.info("[SVD-LLM] Sequential LoRA Recovery Complete.")

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """
        Compress the target model using full SVD-LLM Cholesky whitening and structural replacement.
        """
        covariances = self._collect_covariances()
        target_layers: list[tuple[str, nn.Module]] = []
        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                target_layers.append((name, module))

        logger.info(
            "Found %d linear layers targeted for SVD-LLM compression.",
            len(target_layers),
        )

        for name, layer in target_layers:
            target_mod = getattr(layer, "base_layer", layer)
            cls_name = layer.__class__.__name__.lower()
            target_cls = target_mod.__class__.__name__.lower()
            is_conv1d = (
                "conv1d" in cls_name
                or "conv1d" in target_cls
                or hasattr(layer, "nf")
                or hasattr(target_mod, "nf")
            )

            weight_tensor = getattr(target_mod, "weight", getattr(layer, "weight", None))
            if weight_tensor is None:
                continue

            bias_tensor = getattr(target_mod, "bias", getattr(layer, "bias", None))
            W = weight_tensor.data.clone()
            if is_conv1d:
                W = W.t().contiguous()

            m, n = W.shape
            r = self.compute_rank_from_ratio(m, n, self.rank_ratio)

            # Check parameter reduction condition
            if r * (m + n) >= m * n:
                self.fallback_layers.append(name)
                continue

            C = covariances.get(name, None)
            try:
                W_u, W_v, error = self._cholesky_whiten_and_decompose(W, C, target_rank=r)
            except Exception as decomp_err:
                logger.warning("SVD-LLM decomposition failed for layer %s: %s", name, decomp_err)
                self.fallback_layers.append(name)
                continue

            self.layer_errors[name] = error
            self.layer_ranks[name] = r

            # In faithful SVD-LLM (ICLR 2025), fallback is not used; all target layers are compressed.
            if self.fallback_to_original and error > self.max_error:
                logger.info(
                    "SVD-LLM layer %s error %.4f > max_error %.4f; skipping (fallback).",
                    name,
                    error,
                    self.max_error,
                )
                self.fallback_layers.append(name)
                continue

            svd_layer = SVDLLMLinear(
                in_features=n,
                out_features=m,
                rank=r,
                bias=bias_tensor is not None,
                dtype=weight_tensor.dtype,
                device=weight_tensor.device,
                lora_rank=self.lora_rank,
                lora_alpha=self.lora_alpha,
            )
            svd_layer.set_factors(W_u, W_v, bias_tensor)

            replace_module(self.model, name, svd_layer)
            self.compressed_layers.append(name)

        # Run sequential LoRA fine-tuning
        self._sequential_lora_recovery()

        stats = {
            "num_compressed_layers": len(self.compressed_layers),
            "num_fallback_layers": len(self.fallback_layers),
            "compressed_layers": self.compressed_layers,
            "fallback_layers": self.fallback_layers,
            "layer_errors": self.layer_errors,
            "layer_ranks": self.layer_ranks,
            "mean_error": float(sum(self.layer_errors.values()) / max(1, len(self.layer_errors))),
        }
        self.model.compression_stats = stats
        return self.model, stats

SVD-LLM (ICLR 2025) Compressor with Cholesky activation whitening, sequential LoRA recovery fine-tuning, and fixed rank parameter ratios.

Ancestors

BaseCompressor

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """
    Compress the target model using full SVD-LLM Cholesky whitening and structural replacement.
    """
    covariances = self._collect_covariances()
    target_layers: list[tuple[str, nn.Module]] = []
    for name, module in self.model.named_modules():
        if self._is_target_module(name, module):
            target_layers.append((name, module))

    logger.info(
        "Found %d linear layers targeted for SVD-LLM compression.",
        len(target_layers),
    )

    for name, layer in target_layers:
        target_mod = getattr(layer, "base_layer", layer)
        cls_name = layer.__class__.__name__.lower()
        target_cls = target_mod.__class__.__name__.lower()
        is_conv1d = (
            "conv1d" in cls_name
            or "conv1d" in target_cls
            or hasattr(layer, "nf")
            or hasattr(target_mod, "nf")
        )

        weight_tensor = getattr(target_mod, "weight", getattr(layer, "weight", None))
        if weight_tensor is None:
            continue

        bias_tensor = getattr(target_mod, "bias", getattr(layer, "bias", None))
        W = weight_tensor.data.clone()
        if is_conv1d:
            W = W.t().contiguous()

        m, n = W.shape
        r = self.compute_rank_from_ratio(m, n, self.rank_ratio)

        # Check parameter reduction condition
        if r * (m + n) >= m * n:
            self.fallback_layers.append(name)
            continue

        C = covariances.get(name, None)
        try:
            W_u, W_v, error = self._cholesky_whiten_and_decompose(W, C, target_rank=r)
        except Exception as decomp_err:
            logger.warning("SVD-LLM decomposition failed for layer %s: %s", name, decomp_err)
            self.fallback_layers.append(name)
            continue

        self.layer_errors[name] = error
        self.layer_ranks[name] = r

        # In faithful SVD-LLM (ICLR 2025), fallback is not used; all target layers are compressed.
        if self.fallback_to_original and error > self.max_error:
            logger.info(
                "SVD-LLM layer %s error %.4f > max_error %.4f; skipping (fallback).",
                name,
                error,
                self.max_error,
            )
            self.fallback_layers.append(name)
            continue

        svd_layer = SVDLLMLinear(
            in_features=n,
            out_features=m,
            rank=r,
            bias=bias_tensor is not None,
            dtype=weight_tensor.dtype,
            device=weight_tensor.device,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )
        svd_layer.set_factors(W_u, W_v, bias_tensor)

        replace_module(self.model, name, svd_layer)
        self.compressed_layers.append(name)

    # Run sequential LoRA fine-tuning
    self._sequential_lora_recovery()

    stats = {
        "num_compressed_layers": len(self.compressed_layers),
        "num_fallback_layers": len(self.fallback_layers),
        "compressed_layers": self.compressed_layers,
        "fallback_layers": self.fallback_layers,
        "layer_errors": self.layer_errors,
        "layer_ranks": self.layer_ranks,
        "mean_error": float(sum(self.layer_errors.values()) / max(1, len(self.layer_errors))),
    }
    self.model.compression_stats = stats
    return self.model, stats

Compress the target model using full SVD-LLM Cholesky whitening and structural replacement.

def compute_rank_from_ratio(self, m: int, n: int, ratio: float) ‑> int

Expand source code

def compute_rank_from_ratio(self, m: int, n: int, ratio: float) -> int:
    """
    Compute low rank r from parameter retention ratio R:
    Original params = m * n
    Target params = R * m * n
    Factorized params = r * (m + n)
    r = (R * m * n) / (m + n)
    """
    if self.explicit_rank is not None and self.explicit_rank > 0:
        return min(self.explicit_rank, min(m, n))
    r_calc = (ratio * m * n) / (m + n)
    r = int(round(r_calc))
    return max(1, min(r, min(m, n)))

Compute low rank r from parameter retention ratio R: Original params = m * n Target params = R * m * n Factorized params = r * (m + n) r = (R * m * n) / (m + n)

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class SVDLLMLinear (in_features: int, out_features: int, rank: int, bias: bool = True, dtype: torch.dtype = torch.float32, device: torch.device | None = None, lora_rank: int = 16, lora_alpha: float = 32.0)

Expand source code

class SVDLLMLinear(nn.Module):
    """
    SVD-LLM Factorized Linear Layer executing as:
    y = (x @ W_v.T + lora_v(x)) @ (W_u.T) + lora_u(x @ W_v.T) + bias

    Specifically, with sequential LoRA support:
    - W_u ∈ R^(d_out × r), W_v ∈ R^(r × d_in)
    - LoRA for W_u: A_u ∈ R^(r_L × r), B_u ∈ R^(d_out × r_L)
    - LoRA for W_v: A_v ∈ R^(r_L × d_in), B_v ∈ R^(r × r_L)
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        bias: bool = True,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = max(1, min(rank, in_features, out_features))
        self.lora_rank = lora_rank
        self.lora_alpha = lora_alpha
        self.scaling = lora_alpha / max(1, lora_rank)

        # Explicit factor buffers: W_u and W_v
        self.W_u = nn.Parameter(torch.empty(out_features, self.rank, device=device, dtype=dtype))
        self.W_v = nn.Parameter(torch.empty(self.rank, in_features, device=device, dtype=dtype))
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, device=device, dtype=dtype))
        else:
            self.register_parameter("bias", None)

        # LoRA adapters for W_u and W_v
        if self.lora_rank > 0:
            # LoRA for W_u: input dim = r, output dim = out_features
            self.lora_A_u = nn.Parameter(
                torch.empty(lora_rank, self.rank, device=device, dtype=dtype)
            )
            self.lora_B_u = nn.Parameter(
                torch.empty(out_features, lora_rank, device=device, dtype=dtype)
            )

            # LoRA for W_v: input dim = in_features, output dim = r
            self.lora_A_v = nn.Parameter(
                torch.empty(lora_rank, in_features, device=device, dtype=dtype)
            )
            self.lora_B_v = nn.Parameter(
                torch.empty(self.rank, lora_rank, device=device, dtype=dtype)
            )

            self.reset_lora_parameters()
        else:
            self.register_parameter("lora_A_u", None)
            self.register_parameter("lora_B_u", None)
            self.register_parameter("lora_A_v", None)
            self.register_parameter("lora_B_v", None)

        self.reset_base_parameters()

    def reset_base_parameters(self) -> None:
        with torch.no_grad():
            nn.init.orthogonal_(self.W_u)
            nn.init.orthogonal_(self.W_v)
            if self.bias is not None:
                nn.init.zeros_(self.bias)

    def reset_lora_parameters(self) -> None:
        with torch.no_grad():
            if self.lora_rank > 0:
                nn.init.kaiming_uniform_(self.lora_A_u, a=math.sqrt(5))
                nn.init.zeros_(self.lora_B_u)
                nn.init.kaiming_uniform_(self.lora_A_v, a=math.sqrt(5))
                nn.init.zeros_(self.lora_B_v)

    def set_factors(
        self, W_u: torch.Tensor, W_v: torch.Tensor, bias: torch.Tensor | None = None
    ) -> None:
        with torch.no_grad():
            r = min(self.rank, W_u.size(1), W_v.size(0))
            self.W_u.data.zero_()
            self.W_v.data.zero_()
            self.W_u.data[:, :r].copy_(W_u[:, :r].to(device=self.W_u.device, dtype=self.W_u.dtype))
            self.W_v.data[:r, :].copy_(W_v[:r, :].to(device=self.W_v.device, dtype=self.W_v.dtype))
            if bias is not None and self.bias is not None:
                self.bias.data.copy_(bias.to(device=self.bias.device, dtype=self.bias.dtype))

    def set_active_lora_stage(self, stage: str | None) -> None:
        """
        stage == 'stage1': train A_u, B_u only (W_v frozen)
        stage == 'stage2': train A_v, B_v only (W_u frozen)
        stage == 'both': train all LoRA adapters
        stage is None: freeze all LoRA adapters
        """
        self.W_u.requires_grad = False
        self.W_v.requires_grad = False
        if self.bias is not None:
            self.bias.requires_grad = False

        if self.lora_rank == 0:
            return

        if stage == "stage1":
            self.lora_A_u.requires_grad = True
            self.lora_B_u.requires_grad = True
            self.lora_A_v.requires_grad = False
            self.lora_B_v.requires_grad = False
        elif stage == "stage2":
            self.lora_A_u.requires_grad = False
            self.lora_B_u.requires_grad = False
            self.lora_A_v.requires_grad = True
            self.lora_B_v.requires_grad = True
        elif stage == "both":
            self.lora_A_u.requires_grad = True
            self.lora_B_u.requires_grad = True
            self.lora_A_v.requires_grad = True
            self.lora_B_v.requires_grad = True
        else:
            self.lora_A_u.requires_grad = False
            self.lora_B_u.requires_grad = False
            self.lora_A_v.requires_grad = False
            self.lora_B_v.requires_grad = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 1. Compute latent representation through W_v
        # latent = x @ W_v.T + (optional LoRA_v)
        latent = F.linear(x, self.W_v)
        if self.lora_rank > 0 and self.lora_A_v is not None and self.lora_B_v is not None:
            lora_v = (
                F.linear(F.linear(x.to(self.lora_A_v.dtype), self.lora_A_v), self.lora_B_v)
                * self.scaling
            )
            latent = latent + lora_v.to(latent.dtype)

        # 2. Compute output through W_u
        # out = latent @ W_u.T + (optional LoRA_u)
        out = F.linear(latent, self.W_u)
        if self.lora_rank > 0 and self.lora_A_u is not None and self.lora_B_u is not None:
            lora_u = (
                F.linear(
                    F.linear(latent.to(self.lora_A_u.dtype), self.lora_A_u),
                    self.lora_B_u,
                )
                * self.scaling
            )
            out = out + lora_u.to(out.dtype)

        if self.bias is not None:
            out = out + self.bias
        return out

SVD-LLM Factorized Linear Layer executing as: y = (x @ W_v.T + lora_v(x)) @ (W_u.T) + lora_u(x @ W_v.T) + bias

Specifically, with sequential LoRA support: - W_u ∈ R^(d_out × r), W_v ∈ R^(r × d_in) - LoRA for W_u: A_u ∈ R^(r_L × r), B_u ∈ R^(d_out × r_L) - LoRA for W_v: A_v ∈ R^(r_L × d_in), B_v ∈ R^(r × r_L)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    # 1. Compute latent representation through W_v
    # latent = x @ W_v.T + (optional LoRA_v)
    latent = F.linear(x, self.W_v)
    if self.lora_rank > 0 and self.lora_A_v is not None and self.lora_B_v is not None:
        lora_v = (
            F.linear(F.linear(x.to(self.lora_A_v.dtype), self.lora_A_v), self.lora_B_v)
            * self.scaling
        )
        latent = latent + lora_v.to(latent.dtype)

    # 2. Compute output through W_u
    # out = latent @ W_u.T + (optional LoRA_u)
    out = F.linear(latent, self.W_u)
    if self.lora_rank > 0 and self.lora_A_u is not None and self.lora_B_u is not None:
        lora_u = (
            F.linear(
                F.linear(latent.to(self.lora_A_u.dtype), self.lora_A_u),
                self.lora_B_u,
            )
            * self.scaling
        )
        out = out + lora_u.to(out.dtype)

    if self.bias is not None:
        out = out + self.bias
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def reset_base_parameters(self) ‑> None

Expand source code

def reset_base_parameters(self) -> None:
    with torch.no_grad():
        nn.init.orthogonal_(self.W_u)
        nn.init.orthogonal_(self.W_v)
        if self.bias is not None:
            nn.init.zeros_(self.bias)

def reset_lora_parameters(self) ‑> None

Expand source code

def reset_lora_parameters(self) -> None:
    with torch.no_grad():
        if self.lora_rank > 0:
            nn.init.kaiming_uniform_(self.lora_A_u, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B_u)
            nn.init.kaiming_uniform_(self.lora_A_v, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B_v)

def set_active_lora_stage(self, stage: str | None) ‑> None

Expand source code

def set_active_lora_stage(self, stage: str | None) -> None:
    """
    stage == 'stage1': train A_u, B_u only (W_v frozen)
    stage == 'stage2': train A_v, B_v only (W_u frozen)
    stage == 'both': train all LoRA adapters
    stage is None: freeze all LoRA adapters
    """
    self.W_u.requires_grad = False
    self.W_v.requires_grad = False
    if self.bias is not None:
        self.bias.requires_grad = False

    if self.lora_rank == 0:
        return

    if stage == "stage1":
        self.lora_A_u.requires_grad = True
        self.lora_B_u.requires_grad = True
        self.lora_A_v.requires_grad = False
        self.lora_B_v.requires_grad = False
    elif stage == "stage2":
        self.lora_A_u.requires_grad = False
        self.lora_B_u.requires_grad = False
        self.lora_A_v.requires_grad = True
        self.lora_B_v.requires_grad = True
    elif stage == "both":
        self.lora_A_u.requires_grad = True
        self.lora_B_u.requires_grad = True
        self.lora_A_v.requires_grad = True
        self.lora_B_v.requires_grad = True
    else:
        self.lora_A_u.requires_grad = False
        self.lora_B_u.requires_grad = False
        self.lora_A_v.requires_grad = False
        self.lora_B_v.requires_grad = False

stage == 'stage1': train A_u, B_u only (W_v frozen) stage == 'stage2': train A_v, B_v only (W_u frozen) stage == 'both': train all LoRA adapters stage is None: freeze all LoRA adapters

def set_factors(self, W_u: torch.Tensor, W_v: torch.Tensor, bias: torch.Tensor | None = None) ‑> None

Expand source code

def set_factors(
    self, W_u: torch.Tensor, W_v: torch.Tensor, bias: torch.Tensor | None = None
) -> None:
    with torch.no_grad():
        r = min(self.rank, W_u.size(1), W_v.size(0))
        self.W_u.data.zero_()
        self.W_v.data.zero_()
        self.W_u.data[:, :r].copy_(W_u[:, :r].to(device=self.W_u.device, dtype=self.W_u.dtype))
        self.W_v.data[:r, :].copy_(W_v[:r, :].to(device=self.W_v.device, dtype=self.W_v.dtype))
        if bias is not None and self.bias is not None:
            self.bias.data.copy_(bias.to(device=self.bias.device, dtype=self.bias.dtype))

class SVDLinear (in_features: int, out_features: int, rank: int, bias: bool = True, use_fused_kernel: bool = False, dtype: torch.dtype = torch.float32, device: torch.device | None = None, lora_rank: int = 0, lora_alpha: float = 1.0, lora_r: int | None = None, freeze_base: bool = False)

Expand source code

class SVDLinear(nn.Module):
    """
    Low-rank factorized linear layer:
        y = (x @ V.T) @ U.T + LoRA(x) * scaling + b
    Parameter count: r·(d_in + d_out) + LoRA params.
    """

    # AASVD-FIX-01
    is_custom_attention: bool = False

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        bias: bool = True,
        use_fused_kernel: bool = False,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
        lora_rank: int = 0,
        lora_alpha: float = 1.0,
        lora_r: int | None = None,
        freeze_base: bool = False,
    ) -> None:
        super().__init__()
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.rank = max(1, min(rank, self.in_features, self.out_features))
        self.use_fused_kernel = bool(use_fused_kernel)

        # Accept both lora_rank and legacy lora_r
        eff_lora = int(lora_r) if (lora_r is not None and lora_r > 0) else max(0, int(lora_rank))
        self.lora_rank = eff_lora
        self.lora_r = eff_lora
        self.lora_alpha = float(lora_alpha)
        self.lora_scaling = (self.lora_alpha / self.lora_rank) if self.lora_rank > 0 else 1.0
        self.scaling = self.lora_scaling

        requires_grad = not (self.lora_rank > 0 or freeze_base)
        fk = {"dtype": dtype, "device": device}

        self.V = nn.Parameter(
            torch.empty(self.rank, self.in_features, **fk), requires_grad=requires_grad
        )
        self.U = nn.Parameter(
            torch.empty(self.out_features, self.rank, **fk), requires_grad=requires_grad
        )

        if bias:
            self.bias = nn.Parameter(
                torch.empty(self.out_features, **fk), requires_grad=requires_grad
            )
        else:
            self.register_parameter("bias", None)

        if self.lora_rank > 0:
            self.lora_A = nn.Parameter(torch.empty(self.lora_rank, self.in_features, **fk))
            self.lora_B = nn.Parameter(torch.zeros(self.out_features, self.lora_rank, **fk))
            with torch.no_grad():
                nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        # AASVD-FIX-07: input mean is not tracked; remove the buffer.
        self._input_mean: torch.Tensor | None = None

        self.reset_parameters()

    # ------------------------------------------------------------------
    @property
    def weight(self) -> torch.Tensor:
        return self.get_weight()

    def get_weight(self) -> torch.Tensor:
        w = torch.matmul(self.U, self.V)
        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            w = w + torch.matmul(self.lora_B, self.lora_A) * self.lora_scaling
        return w

    @property
    def use_lora(self) -> bool:
        return bool(self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None)

    def freeze_base_weights(self) -> None:
        if self.V.requires_grad:
            self.V.requires_grad = False
        if self.U.requires_grad:
            self.U.requires_grad = False
        if self.bias is not None and self.bias.requires_grad:
            self.bias.requires_grad = False
        if self.lora_A is not None:
            self.lora_A.requires_grad = True
        if self.lora_B is not None:
            self.lora_B.requires_grad = True

    def reset_parameters(self) -> None:
        with torch.no_grad():
            for param in (self.V, self.U):
                if param.dtype in (torch.float16, torch.bfloat16):
                    tmp = torch.empty_like(param, dtype=torch.float32, device="cpu")
                    _orthogonal_init(tmp)
                    param.copy_(tmp.to(dtype=param.dtype, device=param.device))
                else:
                    _orthogonal_init(param)
            if self.bias is not None:
                nn.init.zeros_(self.bias)
            if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
                nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
                nn.init.zeros_(self.lora_B)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._input_mean is not None:
            x = x - self._input_mean.to(x.device, x.dtype)

        if self.use_fused_kernel and (self.lora_rank == 0 or self.lora_A is None):
            try:
                from neuralops.attentions.fused_svd_kernel import fused_svd_linear

                return fused_svd_linear(x=x, V=self.V, U=self.U, bias=self.bias, use_rms_norm=False)
            except Exception as e:
                logger.debug("SVDLinear fused kernel failed (%s); using F.linear.", e)

        latent = F.linear(x, self.V)
        out = F.linear(latent, self.U)

        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            a = self.lora_A.to(x.dtype) if self.lora_A.dtype != x.dtype else self.lora_A
            b = self.lora_B.to(x.dtype) if self.lora_B.dtype != x.dtype else self.lora_B
            lora_latent = F.linear(x, a)
            out = out + F.linear(lora_latent, b) * self.lora_scaling

        if self.bias is not None:
            out = out + self.bias.to(out.dtype)
        return out

    def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
        with torch.no_grad():
            r = min(self.rank, U_factor.size(1), V_factor.size(0))
            self.U.data.zero_()
            self.V.data.zero_()
            self.U.data[:, :r].copy_(U_factor[:, :r].to(self.U.device, self.U.dtype))
            self.V.data[:r, :].copy_(V_factor[:r, :].to(self.V.device, self.V.dtype))

Low-rank factorized linear layer: y = (x @ V.T) @ U.T + LoRA(x) * scaling + b Parameter count: r·(d_in + d_out) + LoRA params.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Instance variables

prop use_lora : bool

Expand source code

@property
def use_lora(self) -> bool:
    return bool(self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None)

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    return self.get_weight()

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    if self._input_mean is not None:
        x = x - self._input_mean.to(x.device, x.dtype)

    if self.use_fused_kernel and (self.lora_rank == 0 or self.lora_A is None):
        try:
            from neuralops.attentions.fused_svd_kernel import fused_svd_linear

            return fused_svd_linear(x=x, V=self.V, U=self.U, bias=self.bias, use_rms_norm=False)
        except Exception as e:
            logger.debug("SVDLinear fused kernel failed (%s); using F.linear.", e)

    latent = F.linear(x, self.V)
    out = F.linear(latent, self.U)

    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        a = self.lora_A.to(x.dtype) if self.lora_A.dtype != x.dtype else self.lora_A
        b = self.lora_B.to(x.dtype) if self.lora_B.dtype != x.dtype else self.lora_B
        lora_latent = F.linear(x, a)
        out = out + F.linear(lora_latent, b) * self.lora_scaling

    if self.bias is not None:
        out = out + self.bias.to(out.dtype)
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    if self.V.requires_grad:
        self.V.requires_grad = False
    if self.U.requires_grad:
        self.U.requires_grad = False
    if self.bias is not None and self.bias.requires_grad:
        self.bias.requires_grad = False
    if self.lora_A is not None:
        self.lora_A.requires_grad = True
    if self.lora_B is not None:
        self.lora_B.requires_grad = True

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    w = torch.matmul(self.U, self.V)
    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        w = w + torch.matmul(self.lora_B, self.lora_A) * self.lora_scaling
    return w

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    with torch.no_grad():
        for param in (self.V, self.U):
            if param.dtype in (torch.float16, torch.bfloat16):
                tmp = torch.empty_like(param, dtype=torch.float32, device="cpu")
                _orthogonal_init(tmp)
                param.copy_(tmp.to(dtype=param.dtype, device=param.device))
            else:
                _orthogonal_init(param)
        if self.bias is not None:
            nn.init.zeros_(self.bias)
        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) ‑> None

Expand source code

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
    with torch.no_grad():
        r = min(self.rank, U_factor.size(1), V_factor.size(0))
        self.U.data.zero_()
        self.V.data.zero_()
        self.U.data[:, :r].copy_(U_factor[:, :r].to(self.U.device, self.U.dtype))
        self.V.data[:r, :].copy_(V_factor[:r, :].to(self.V.device, self.V.dtype))

class AdaptiveSVDLinear (in_features: int, out_features: int, rank: int, bias: bool = True, use_fused_kernel: bool = False, dtype: torch.dtype = torch.float32, device: torch.device | None = None, lora_rank: int = 0, lora_alpha: float = 1.0, lora_r: int | None = None, freeze_base: bool = False)

Expand source code

class SVDLinear(nn.Module):
    """
    Low-rank factorized linear layer:
        y = (x @ V.T) @ U.T + LoRA(x) * scaling + b
    Parameter count: r·(d_in + d_out) + LoRA params.
    """

    # AASVD-FIX-01
    is_custom_attention: bool = False

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        bias: bool = True,
        use_fused_kernel: bool = False,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
        lora_rank: int = 0,
        lora_alpha: float = 1.0,
        lora_r: int | None = None,
        freeze_base: bool = False,
    ) -> None:
        super().__init__()
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.rank = max(1, min(rank, self.in_features, self.out_features))
        self.use_fused_kernel = bool(use_fused_kernel)

        # Accept both lora_rank and legacy lora_r
        eff_lora = int(lora_r) if (lora_r is not None and lora_r > 0) else max(0, int(lora_rank))
        self.lora_rank = eff_lora
        self.lora_r = eff_lora
        self.lora_alpha = float(lora_alpha)
        self.lora_scaling = (self.lora_alpha / self.lora_rank) if self.lora_rank > 0 else 1.0
        self.scaling = self.lora_scaling

        requires_grad = not (self.lora_rank > 0 or freeze_base)
        fk = {"dtype": dtype, "device": device}

        self.V = nn.Parameter(
            torch.empty(self.rank, self.in_features, **fk), requires_grad=requires_grad
        )
        self.U = nn.Parameter(
            torch.empty(self.out_features, self.rank, **fk), requires_grad=requires_grad
        )

        if bias:
            self.bias = nn.Parameter(
                torch.empty(self.out_features, **fk), requires_grad=requires_grad
            )
        else:
            self.register_parameter("bias", None)

        if self.lora_rank > 0:
            self.lora_A = nn.Parameter(torch.empty(self.lora_rank, self.in_features, **fk))
            self.lora_B = nn.Parameter(torch.zeros(self.out_features, self.lora_rank, **fk))
            with torch.no_grad():
                nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        # AASVD-FIX-07: input mean is not tracked; remove the buffer.
        self._input_mean: torch.Tensor | None = None

        self.reset_parameters()

    # ------------------------------------------------------------------
    @property
    def weight(self) -> torch.Tensor:
        return self.get_weight()

    def get_weight(self) -> torch.Tensor:
        w = torch.matmul(self.U, self.V)
        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            w = w + torch.matmul(self.lora_B, self.lora_A) * self.lora_scaling
        return w

    @property
    def use_lora(self) -> bool:
        return bool(self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None)

    def freeze_base_weights(self) -> None:
        if self.V.requires_grad:
            self.V.requires_grad = False
        if self.U.requires_grad:
            self.U.requires_grad = False
        if self.bias is not None and self.bias.requires_grad:
            self.bias.requires_grad = False
        if self.lora_A is not None:
            self.lora_A.requires_grad = True
        if self.lora_B is not None:
            self.lora_B.requires_grad = True

    def reset_parameters(self) -> None:
        with torch.no_grad():
            for param in (self.V, self.U):
                if param.dtype in (torch.float16, torch.bfloat16):
                    tmp = torch.empty_like(param, dtype=torch.float32, device="cpu")
                    _orthogonal_init(tmp)
                    param.copy_(tmp.to(dtype=param.dtype, device=param.device))
                else:
                    _orthogonal_init(param)
            if self.bias is not None:
                nn.init.zeros_(self.bias)
            if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
                nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
                nn.init.zeros_(self.lora_B)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._input_mean is not None:
            x = x - self._input_mean.to(x.device, x.dtype)

        if self.use_fused_kernel and (self.lora_rank == 0 or self.lora_A is None):
            try:
                from neuralops.attentions.fused_svd_kernel import fused_svd_linear

                return fused_svd_linear(x=x, V=self.V, U=self.U, bias=self.bias, use_rms_norm=False)
            except Exception as e:
                logger.debug("SVDLinear fused kernel failed (%s); using F.linear.", e)

        latent = F.linear(x, self.V)
        out = F.linear(latent, self.U)

        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            a = self.lora_A.to(x.dtype) if self.lora_A.dtype != x.dtype else self.lora_A
            b = self.lora_B.to(x.dtype) if self.lora_B.dtype != x.dtype else self.lora_B
            lora_latent = F.linear(x, a)
            out = out + F.linear(lora_latent, b) * self.lora_scaling

        if self.bias is not None:
            out = out + self.bias.to(out.dtype)
        return out

    def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
        with torch.no_grad():
            r = min(self.rank, U_factor.size(1), V_factor.size(0))
            self.U.data.zero_()
            self.V.data.zero_()
            self.U.data[:, :r].copy_(U_factor[:, :r].to(self.U.device, self.U.dtype))
            self.V.data[:r, :].copy_(V_factor[:r, :].to(self.V.device, self.V.dtype))

Low-rank factorized linear layer: y = (x @ V.T) @ U.T + LoRA(x) * scaling + b Parameter count: r·(d_in + d_out) + LoRA params.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Instance variables

prop use_lora : bool

Expand source code

@property
def use_lora(self) -> bool:
    return bool(self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None)

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    return self.get_weight()

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    if self._input_mean is not None:
        x = x - self._input_mean.to(x.device, x.dtype)

    if self.use_fused_kernel and (self.lora_rank == 0 or self.lora_A is None):
        try:
            from neuralops.attentions.fused_svd_kernel import fused_svd_linear

            return fused_svd_linear(x=x, V=self.V, U=self.U, bias=self.bias, use_rms_norm=False)
        except Exception as e:
            logger.debug("SVDLinear fused kernel failed (%s); using F.linear.", e)

    latent = F.linear(x, self.V)
    out = F.linear(latent, self.U)

    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        a = self.lora_A.to(x.dtype) if self.lora_A.dtype != x.dtype else self.lora_A
        b = self.lora_B.to(x.dtype) if self.lora_B.dtype != x.dtype else self.lora_B
        lora_latent = F.linear(x, a)
        out = out + F.linear(lora_latent, b) * self.lora_scaling

    if self.bias is not None:
        out = out + self.bias.to(out.dtype)
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    if self.V.requires_grad:
        self.V.requires_grad = False
    if self.U.requires_grad:
        self.U.requires_grad = False
    if self.bias is not None and self.bias.requires_grad:
        self.bias.requires_grad = False
    if self.lora_A is not None:
        self.lora_A.requires_grad = True
    if self.lora_B is not None:
        self.lora_B.requires_grad = True

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    w = torch.matmul(self.U, self.V)
    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        w = w + torch.matmul(self.lora_B, self.lora_A) * self.lora_scaling
    return w

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    with torch.no_grad():
        for param in (self.V, self.U):
            if param.dtype in (torch.float16, torch.bfloat16):
                tmp = torch.empty_like(param, dtype=torch.float32, device="cpu")
                _orthogonal_init(tmp)
                param.copy_(tmp.to(dtype=param.dtype, device=param.device))
            else:
                _orthogonal_init(param)
        if self.bias is not None:
            nn.init.zeros_(self.bias)
        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) ‑> None

Expand source code

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
    with torch.no_grad():
        r = min(self.rank, U_factor.size(1), V_factor.size(0))
        self.U.data.zero_()
        self.V.data.zero_()
        self.U.data[:, :r].copy_(U_factor[:, :r].to(self.U.device, self.U.dtype))
        self.V.data[:r, :].copy_(V_factor[:r, :].to(self.V.device, self.V.dtype))

class SVDLinearCompressor (model: nn.Module, rank: int = 192, max_error: float = 0.2, fallback_to_original: bool = True, protect_critical: bool = True, target_attention_only: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, skip_attention_qkv: bool = False, target_modules: list[str] | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("svd_linear")
@COMPRESSOR_REGISTRY.register("svd_linear_compressor")
class SVDLinearCompressor(BaseCompressor):
    """Compresses model by replacing Linear/Conv1D layers with SVDLinearLayer."""

    def __init__(
        self,
        model: nn.Module,
        rank: int = 192,
        max_error: float = 0.20,
        fallback_to_original: bool = True,
        protect_critical: bool = True,
        target_attention_only: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        skip_attention_qkv: bool = False,
        target_modules: list[str] | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="svd_linear", **kwargs)
        conf = cfg or kwargs.get("comp_config")
        if conf is not None:
            rank = _cfg_get(conf, "svd_rank", _cfg_get(conf, "rank", rank))
            max_error = _cfg_get(conf, "max_error", max_error)
            fallback_to_original = _cfg_get(conf, "fallback_to_original", fallback_to_original)
            protect_critical = _cfg_get(conf, "protect_critical", protect_critical)
            target_attention_only = _cfg_get(conf, "target_attention_only", target_attention_only)
            lora_r = _cfg_get(conf, "lora_r", _cfg_get(conf, "lora_rank", lora_r))
            lora_alpha = _cfg_get(conf, "lora_alpha", lora_alpha)
            skip_attention_qkv = _cfg_get(conf, "skip_attention_qkv", skip_attention_qkv)

        self.rank = int(rank)
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.protect_critical = bool(protect_critical)
        self.target_attention_only = bool(target_attention_only)
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.skip_attention_qkv = bool(skip_attention_qkv)
        self.target_modules = target_modules
        self.layer_errors: dict[str, float] = {}
        self.replaced_layers: list[dict[str, Any]] = []
        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.stats: dict[str, Any] = {}

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        if not self.skip_attention_qkv:
            return False
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads:
            name_lower = name.lower()
            if any(
                p in name_lower
                for p in [
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "query",
                    "key",
                    "value",
                    "c_attn",
                ]
            ):
                return True
        return False

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        if any(
            k in name_l
            for k in [
                "out_lin",
                "out_proj",
                "o_proj",
                "c_proj",
                "output.dense",
                "self.output",
            ]
        ):
            return False
        if self.protect_critical and any(
            k in name_l
            for k in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embed",
                "logits_processor",
            ]
        ):
            return False

        if self.target_attention_only:
            attn_keywords = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(k in name_l for k in attn_keywords):
                return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__.lower()
        is_linear_like = (
            isinstance(base_mod, nn.Linear) or "conv1d" in cls_name or hasattr(base_mod, "nf")
        )
        if not is_linear_like:
            return False
        if isinstance(base_mod, SVDLinearLayer):
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self._is_gqa_asymmetric(name, module):
            return False
        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """Compress linear layers with SVD factorization."""
        self.layer_errors = {}
        self.replaced_layers = []
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        seen_weights: set[int] = set()

        target_layers = [
            (name, mod)
            for name, mod in list(self.model.named_modules())
            if self._is_target_module(name, mod)
        ]

        if not target_layers:
            logger.warning("SVDLinear: No compressible linear layers found.")
            return self.model, self.get_compression_stats()

        for name, module in target_layers:
            base_mod = getattr(module, "base_layer", module)
            if not hasattr(base_mod, "weight") or base_mod.weight is None:
                continue
            w_id = id(base_mod.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            try:
                new_layer = SVDLinearLayer(
                    module,
                    rank=self.rank,
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                )
                err = new_layer.reconstruction_error
                self.layer_errors[name] = err

                if err > self.max_error:
                    logger.warning(
                        "SVDLinear reconstruction error %.4f on %s exceeded threshold %.4f.",
                        err,
                        name,
                        self.max_error,
                    )
                    if self.fallback_to_original:
                        logger.info(
                            "Skipping SVDLinear compression for %s (fallback to original).",
                            name,
                        )
                        self.fallback_layers.append(name)
                        for p in module.parameters():
                            p.requires_grad = False
                        continue

                self.replaced_layers.append({"name": name, "rank": new_layer.rank, "error": err})
                self.compressed_layers.append(name)
                replace_module(self.model, name, new_layer)
            except Exception as e:
                logger.warning("Failed to compress %s with SVDLinear: %s", name, e)
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False

        # Post-compression forward sanity check
        forward_check_success = False
        try:
            self.model.eval()
            with torch.no_grad():
                d = next(self.model.parameters()).device
                dummy = torch.randint(0, 100, (1, 16), device=d)
                self.model(dummy)
            forward_check_success = True
        except Exception as e:
            logger.debug("SVDLinear post-compression forward sanity check exception: %s", e)

        self._compressed_params = self._count_params(self.model)
        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        stats = self.get_compression_stats()
        stats.update(
            {
                "layer_errors": self.layer_errors,
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "avg_reconstruction_error": avg_err,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "target_attention_only": self.target_attention_only,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "forward_check_success": forward_check_success,
                "method": "svd_linear_factorization",
            }
        )
        self.stats = stats
        return self.model, stats

Compresses model by replacing Linear/Conv1D layers with SVDLinearLayer.

Ancestors

BaseCompressor

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """Compress linear layers with SVD factorization."""
    self.layer_errors = {}
    self.replaced_layers = []
    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    seen_weights: set[int] = set()

    target_layers = [
        (name, mod)
        for name, mod in list(self.model.named_modules())
        if self._is_target_module(name, mod)
    ]

    if not target_layers:
        logger.warning("SVDLinear: No compressible linear layers found.")
        return self.model, self.get_compression_stats()

    for name, module in target_layers:
        base_mod = getattr(module, "base_layer", module)
        if not hasattr(base_mod, "weight") or base_mod.weight is None:
            continue
        w_id = id(base_mod.weight)
        if w_id in seen_weights:
            self.skipped_layers.append(name)
            continue
        seen_weights.add(w_id)

        try:
            new_layer = SVDLinearLayer(
                module,
                rank=self.rank,
                lora_r=self.lora_r,
                lora_alpha=self.lora_alpha,
            )
            err = new_layer.reconstruction_error
            self.layer_errors[name] = err

            if err > self.max_error:
                logger.warning(
                    "SVDLinear reconstruction error %.4f on %s exceeded threshold %.4f.",
                    err,
                    name,
                    self.max_error,
                )
                if self.fallback_to_original:
                    logger.info(
                        "Skipping SVDLinear compression for %s (fallback to original).",
                        name,
                    )
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

            self.replaced_layers.append({"name": name, "rank": new_layer.rank, "error": err})
            self.compressed_layers.append(name)
            replace_module(self.model, name, new_layer)
        except Exception as e:
            logger.warning("Failed to compress %s with SVDLinear: %s", name, e)
            self.fallback_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False

    # Post-compression forward sanity check
    forward_check_success = False
    try:
        self.model.eval()
        with torch.no_grad():
            d = next(self.model.parameters()).device
            dummy = torch.randint(0, 100, (1, 16), device=d)
            self.model(dummy)
        forward_check_success = True
    except Exception as e:
        logger.debug("SVDLinear post-compression forward sanity check exception: %s", e)

    self._compressed_params = self._count_params(self.model)
    avg_err = (
        sum(self.layer_errors[k] for k in self.compressed_layers)
        / max(1, len(self.compressed_layers))
        if self.compressed_layers
        else 0.0
    )
    stats = self.get_compression_stats()
    stats.update(
        {
            "layer_errors": self.layer_errors,
            "compressed_layers": self.compressed_layers,
            "fallback_layers": self.fallback_layers,
            "skipped_layers": self.skipped_layers,
            "num_compressed": len(self.compressed_layers),
            "num_fallback": len(self.fallback_layers),
            "num_skipped": len(self.skipped_layers),
            "avg_reconstruction_error": avg_err,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "target_attention_only": self.target_attention_only,
            "lora_r": self.lora_r,
            "lora_alpha": self.lora_alpha,
            "forward_check_success": forward_check_success,
            "method": "svd_linear_factorization",
        }
    )
    self.stats = stats
    return self.model, stats

Compress linear layers with SVD factorization.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class SVDLinearLayer (linear_layer: nn.Module, rank: int, use_fused_kernel: bool = False, use_rms_norm: bool = False, eps: float = 1e-06, lora_r: int = 0, lora_alpha: float = 1.0)

Expand source code

class SVDLinearLayer(nn.Module):
    is_custom_attention: bool = True
    """
    Replace a linear layer with factorized U @ V, keeping top-`rank` components.
    Optionally dispatches fused in-SRAM kernel to eliminate memory roundtrips.
    Includes optional trainable LoRA adapters (lora_A, lora_B).
    """

    def __init__(
        self,
        linear_layer: nn.Module,
        rank: int,
        use_fused_kernel: bool = False,
        use_rms_norm: bool = False,
        eps: float = 1e-6,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> None:
        super().__init__()
        base_layer = getattr(linear_layer, "base_layer", linear_layer)
        raw_w = base_layer.weight.data
        cls_name = base_layer.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(base_layer, "nf")

        if is_conv1d:
            W = raw_w.t().contiguous().float()
            self.out_features = W.shape[0]
            self.in_features = W.shape[1]
        else:
            W = raw_w.float()
            self.out_features = getattr(base_layer, "out_features", W.shape[0])
            self.in_features = getattr(base_layer, "in_features", W.shape[1])

        self.nf = self.out_features  # Conv1D compatibility
        param_cap = max(
            1,
            (self.in_features * self.out_features) // (self.in_features + self.out_features) - 1,
        )
        self.rank: int = max(1, min(rank, self.in_features, self.out_features, param_cap))
        self.use_fused_kernel: bool = use_fused_kernel
        self.use_rms_norm: bool = use_rms_norm
        self.eps: float = eps

        has_existing_lora = hasattr(linear_layer, "lora_A") and hasattr(linear_layer, "lora_B")
        effective_lora_r = lora_r
        if has_existing_lora and getattr(linear_layer, "r", 0) > 0:
            effective_lora_r = getattr(linear_layer, "r", lora_r)

        self.lora_r = int(effective_lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        device = base_layer.weight.device
        dtype = base_layer.weight.dtype

        with torch.no_grad():
            U, S, Vh = torch.linalg.svd(W, full_matrices=False)
            sqrt_s = torch.sqrt(torch.clamp(S[: self.rank], min=1e-12))
            U_weighted = (U[:, : self.rank] * sqrt_s.unsqueeze(0)).to(device=device, dtype=dtype)
            V_weighted = (sqrt_s.unsqueeze(1) * Vh[: self.rank, :]).to(device=device, dtype=dtype)

        requires_grad = not (self.lora_r > 0)
        self.U = nn.Parameter(U_weighted, requires_grad=requires_grad)
        self.V = nn.Parameter(V_weighted, requires_grad=requires_grad)

        if use_rms_norm:
            self.norm_weight = nn.Parameter(torch.ones(self.rank, dtype=dtype, device=device))
        else:
            self.register_parameter("norm_weight", None)

        if base_layer.bias is not None:
            self.bias = nn.Parameter(base_layer.bias.data.clone(), requires_grad=requires_grad)
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty((self.lora_r, self.in_features), dtype=torch.float32, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.zeros((self.out_features, self.lora_r), dtype=torch.float32, device=device)
            )
            if has_existing_lora:
                try:
                    if hasattr(linear_layer.lora_A, "weight"):
                        self.lora_A.data.copy_(linear_layer.lora_A.weight.data)
                    if hasattr(linear_layer.lora_B, "weight"):
                        self.lora_B.data.copy_(linear_layer.lora_B.weight.data)
                except Exception:
                    nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            else:
                nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        # Frobenius reconstruction error
        W_rec = torch.matmul(U_weighted.float(), V_weighted.float())
        norm_orig = torch.norm(W, p="fro")
        self.reconstruction_error = (torch.norm(W - W_rec, p="fro") / (norm_orig + 1e-8)).item()

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix for compatibility."""
        return torch.matmul(self.U, self.V)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass with optional fused kernel acceleration and LoRA branch."""
        if self.use_fused_kernel and self.lora_r <= 0 and fused_svd_linear is not None:
            return fused_svd_linear(
                x=x,
                V=self.V,
                U=self.U,
                norm_weight=self.norm_weight if self.use_rms_norm else None,
                bias=self.bias,
                eps=self.eps,
                use_rms_norm=self.use_rms_norm,
            )

        latent = F.linear(x, self.V)
        out = F.linear(latent, self.U, self.bias)

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            orig_shape = x.shape
            x_flat = x.reshape(-1, self.in_features)
            lora_out = F.linear(F.linear(x_flat.float(), self.lora_A), self.lora_B)
            out = out + (self.scaling * lora_out).reshape(*orig_shape[:-1], self.out_features).to(
                out.dtype
            )

        return out

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: Replace a linear layer with factorized U @ V, keeping top-rank components. Optionally dispatches fused in-SRAM kernel to eliminate memory roundtrips. Includes optional trainable LoRA adapters (lora_A, lora_B).

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix for compatibility."""
    return torch.matmul(self.U, self.V)

Reconstructed effective weight matrix for compatibility.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass with optional fused kernel acceleration and LoRA branch."""
    if self.use_fused_kernel and self.lora_r <= 0 and fused_svd_linear is not None:
        return fused_svd_linear(
            x=x,
            V=self.V,
            U=self.U,
            norm_weight=self.norm_weight if self.use_rms_norm else None,
            bias=self.bias,
            eps=self.eps,
            use_rms_norm=self.use_rms_norm,
        )

    latent = F.linear(x, self.V)
    out = F.linear(latent, self.U, self.bias)

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)
        lora_out = F.linear(F.linear(x_flat.float(), self.lora_A), self.lora_B)
        out = out + (self.scaling * lora_out).reshape(*orig_shape[:-1], self.out_features).to(
            out.dtype
        )

    return out

Forward pass with optional fused kernel acceleration and LoRA branch.

class SVDSingleHead (hidden_size: int, svd_rank: int = 20)

Expand source code

class SVDSingleHead(nn.Module):
    """
    DEPRECATED single-head SVD attention.

    Performing SVD on live activations is O(N^2 d), numerically fragile, and
    theoretically incorrect as an attention replacement. Use the weight-based
    HybridSVDCompressor instead.

    Kept strictly for backwards compatibility. First use emits a warning.
    """

    is_custom_attention: bool = True
    _warned: bool = False

    def __init__(self, hidden_size: int, svd_rank: int = 20) -> None:
        super().__init__()
        if not SVDSingleHead._warned:
            warnings.warn(
                "SVDSingleHead is deprecated (SVD on live activations is "
                "computationally prohibitive and incorrect). Use "
                "HybridSVDCompressor for weight-based SVD instead.",
                DeprecationWarning,
                stacklevel=2,
            )
            SVDSingleHead._warned = True
        self.hidden_size = hidden_size
        self.svd_rank = int(svd_rank)
        self._pos_encoding: torch.Tensor | None = None

    def _get_pos_encoding(
        self, seq_len: int, device: torch.device, dtype: torch.dtype
    ) -> torch.Tensor:
        if (
            self._pos_encoding is None
            or self._pos_encoding.shape[1] != seq_len
            or self._pos_encoding.device != device
        ):
            pe = torch.zeros(seq_len, self.hidden_size, device=device)
            position = torch.arange(0, seq_len, dtype=torch.float, device=device).unsqueeze(1)
            div_term = torch.exp(
                torch.arange(0, self.hidden_size, 2, dtype=torch.float, device=device)
                * (-math.log(10000.0) / self.hidden_size)
            )
            pe[:, 0::2] = torch.sin(position * div_term)
            pe[:, 1::2] = torch.cos(position * div_term)
            self._pos_encoding = pe.unsqueeze(0).to(dtype)
        return self._pos_encoding

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        B, L, _ = hidden_states.shape
        pe = self._get_pos_encoding(L, hidden_states.device, hidden_states.dtype)
        x = hidden_states + pe[:, :L, :]

        x_flat = F.normalize(x.reshape(-1, self.hidden_size).float(), p=2, dim=-1)
        try:
            U, S, Vh = torch.linalg.svd(x_flat, full_matrices=False)
            k = min(self.svd_rank, len(S))
            weighted_U = U[:, :k] * S[:k].view(1, -1)
            output_flat = torch.matmul(weighted_U, Vh[:k, :])
            return output_flat.reshape(B, L, self.hidden_size).to(hidden_states.dtype)
        except Exception:
            return x

DEPRECATED single-head SVD attention.

Performing SVD on live activations is O(N^2 d), numerically fragile, and theoretically incorrect as an attention replacement. Use the weight-based HybridSVDCompressor instead.

Kept strictly for backwards compatibility. First use emits a warning.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def forward(self, hidden_states: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    B, L, _ = hidden_states.shape
    pe = self._get_pos_encoding(L, hidden_states.device, hidden_states.dtype)
    x = hidden_states + pe[:, :L, :]

    x_flat = F.normalize(x.reshape(-1, self.hidden_size).float(), p=2, dim=-1)
    try:
        U, S, Vh = torch.linalg.svd(x_flat, full_matrices=False)
        k = min(self.svd_rank, len(S))
        weighted_U = U[:, :k] * S[:k].view(1, -1)
        output_flat = torch.matmul(weighted_U, Vh[:k, :])
        return output_flat.reshape(B, L, self.hidden_size).to(hidden_states.dtype)
    except Exception:
        return x

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class SharedSubspace (V_k: torch.Tensor, V_v: torch.Tensor, mean_k: torch.Tensor, mean_v: torch.Tensor, bit_alloc_k: torch.Tensor, bit_alloc_v: torch.Tensor, group_id: int = 0, layer_indices: list[int] = <factory>)

Expand source code

@dataclass
class SharedSubspace:
    """Shared low-rank subspace for a group of consecutive transformer layers."""

    V_k: torch.Tensor
    V_v: torch.Tensor
    mean_k: torch.Tensor
    mean_v: torch.Tensor
    bit_alloc_k: torch.Tensor
    bit_alloc_v: torch.Tensor
    group_id: int = 0
    layer_indices: list[int] = field(default_factory=list)

Shared low-rank subspace for a group of consecutive transformer layers.

Instance variables

var V_k : torch.Tensor: The type of the None singleton.
var V_v : torch.Tensor: The type of the None singleton.
var bit_alloc_k : torch.Tensor: The type of the None singleton.
var bit_alloc_v : torch.Tensor: The type of the None singleton.
var group_id : int: The type of the None singleton.
var layer_indices : list[int]: The type of the None singleton.
var mean_k : torch.Tensor: The type of the None singleton.
var mean_v : torch.Tensor: The type of the None singleton.

class SimpleEntropyCoder

Expand source code

class SimpleEntropyCoder:
    """Lightweight entropy coder using empirical symbol frequency estimation."""

    @staticmethod
    def compute_entropy(tensor: torch.Tensor) -> float:
        """Compute empirical Shannon entropy in bits per symbol."""
        if tensor.numel() == 0:
            return 0.0
        unique_vals, counts = torch.unique(tensor, return_counts=True)
        probs = counts.float() / float(tensor.numel())
        entropy = -torch.sum(probs * torch.log2(torch.clamp(probs, min=1e-12))).item()
        return float(entropy)

    @staticmethod
    def compute_compression_ratio(tensor: torch.Tensor, uncompressed_bits: int = 16) -> float:
        """Estimate theoretical compression ratio relative to FP16/FP32 baseline."""
        entropy = SimpleEntropyCoder.compute_entropy(tensor)
        return float(uncompressed_bits / max(entropy, 1e-4))

Lightweight entropy coder using empirical symbol frequency estimation.

Static methods

def compute_compression_ratio(tensor: torch.Tensor, uncompressed_bits: int = 16) ‑> float

Expand source code

@staticmethod
def compute_compression_ratio(tensor: torch.Tensor, uncompressed_bits: int = 16) -> float:
    """Estimate theoretical compression ratio relative to FP16/FP32 baseline."""
    entropy = SimpleEntropyCoder.compute_entropy(tensor)
    return float(uncompressed_bits / max(entropy, 1e-4))

Estimate theoretical compression ratio relative to FP16/FP32 baseline.

def compute_entropy(tensor: torch.Tensor) ‑> float

Expand source code

@staticmethod
def compute_entropy(tensor: torch.Tensor) -> float:
    """Compute empirical Shannon entropy in bits per symbol."""
    if tensor.numel() == 0:
        return 0.0
    unique_vals, counts = torch.unique(tensor, return_counts=True)
    probs = counts.float() / float(tensor.numel())
    entropy = -torch.sum(probs * torch.log2(torch.clamp(probs, min=1e-12))).item()
    return float(entropy)

Compute empirical Shannon entropy in bits per symbol.

class SliceGPTCompressor (model: nn.Module, tokenizer: Any | None = None, sparsity: float = 0.6, non_uniform: bool = False, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, target_attention_only: bool = False, device: str | torch.device | None = None, max_error: float = 0.5, fallback_to_original: bool = True, target_modules: list[str] | None = None, use_lora: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("slicegpt")
@COMPRESSOR_REGISTRY.register("slicegpt_compressor")
class SliceGPTCompressor(BaseCompressor):
    """
    SliceGPT Compression Strategy (real low-rank implementation).

    Performs orthogonal rotation of each target weight matrix using the
    eigendecomposition of its input activation covariance, then truncates
    to the top-r principal coordinates. The truncated weight is stored as a
    rank-r SVDLinear factorization, achieving genuine parameter reduction.
    """

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any | None = None,
        sparsity: float = 0.6,
        non_uniform: bool = False,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        target_attention_only: bool = False,
        device: str | torch.device | None = None,
        max_error: float = 0.5,
        fallback_to_original: bool = True,
        target_modules: list[str] | None = None,
        use_lora: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="slicegpt", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            sparsity = _cfg_get(cfg, "slicegpt_sparsity", _cfg_get(cfg, "sparsity", sparsity))
            non_uniform = _cfg_get(cfg, "non_uniform", non_uniform)
            max_error = _cfg_get(cfg, "max_error", max_error)
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)
            use_lora = _cfg_get(cfg, "use_lora", use_lora)
            lora_r = _cfg_get(cfg, "lora_r", lora_r)
            lora_alpha = _cfg_get(cfg, "lora_alpha", lora_alpha)

        if "fallback_to_original" in kwargs:
            fallback_to_original = bool(kwargs["fallback_to_original"])
        if "use_lora" in kwargs:
            use_lora = bool(kwargs["use_lora"])
        if "lora_r" in kwargs:
            lora_r = int(kwargs["lora_r"])

        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.sparsity = max(0.0, min(float(sparsity), 0.95))
        self.non_uniform = non_uniform
        self.target_attention_only = bool(target_attention_only)
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.target_device = (
            torch.device(device)
            if device is not None
            else next(model.parameters(), torch.tensor(0.0)).device
        )
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.target_modules = target_modules
        self.use_lora = bool(use_lora) or (lora_r > 0)
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.stats: dict[str, Any] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        """
        Unified critical-layer protection delegated to should_compress_linear.
        """
        name_l = name.lower()
        if self.target_attention_only:
            attn_keywords = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(k in name_l for k in attn_keywords):
                return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__.lower()
        is_linear_like = (
            isinstance(base_mod, nn.Linear) or "conv1d" in cls_name or hasattr(base_mod, "nf")
        )
        if not is_linear_like:
            return False
        # Skip layers that are already factorized
        if SVDLinear is not None and isinstance(base_mod, SVDLinear):
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    def _collect_layer_activations(self) -> dict[str, torch.Tensor]:
        """Collect input activations across target linear layers."""
        self.model.eval()
        activations: dict[str, list[torch.Tensor]] = {}
        hooks = []

        def make_hook(name: str):
            def hook(module: nn.Module, inputs: tuple[torch.Tensor, ...], output: Any):
                if inputs and isinstance(inputs[0], torch.Tensor):
                    inp = inputs[0].detach()
                    activations.setdefault(name, []).append(inp.cpu())

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                hooks.append(module.register_forward_hook(make_hook(name)))

        calib_inputs = get_calibration_tensors(
            self.model,
            tokenizer=self.tokenizer,
            calibration_loader=self.calibration_loader,
            calibration_texts=self.calibration_texts,
            device=self.target_device,
        )

        with torch.no_grad():
            try:
                if isinstance(calib_inputs, dict):
                    self.model(**calib_inputs)
                elif isinstance(calib_inputs, torch.Tensor):
                    self.model(calib_inputs)
                elif isinstance(calib_inputs, (list, tuple)):
                    self.model(*calib_inputs)
            except Exception as e:
                logger.debug(
                    "Forward pass for activation collection had exception: %s. "
                    "Using synthetic linear inputs.",
                    e,
                )

        for h in hooks:
            h.remove()

        aggregated: dict[str, torch.Tensor] = {}
        for name, act_list in activations.items():
            if act_list:
                aggregated[name] = torch.cat(act_list, dim=0)
        return aggregated

    def _compute_orthogonal_rotation(
        self,
        X: torch.Tensor | None,
        d_in: int,
        device: torch.device | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Compute orthogonal rotation matrix Q via eigen-decomposition of activation covariance."""
        if device is None:
            device = X.device if X is not None else torch.device("cpu")
        if X is not None and X.numel() > 0:
            X_flat = X.reshape(-1, X.shape[-1]).float().to(device)
            if X_flat.shape[-1] == d_in and X_flat.shape[0] > 1:
                cov = torch.matmul(X_flat.T, X_flat) / X_flat.shape[0]
                cov = cov + 1e-4 * torch.eye(d_in, device=device)
                try:
                    eigvals, Q = torch.linalg.eigh(cov)
                    idx = torch.argsort(eigvals, descending=True)
                    return Q[:, idx], eigvals[idx]
                except Exception as e:
                    logger.debug(
                        "eigh covariance failed (%s). Falling back to QR random basis.",
                        e,
                    )

        rand_mat = torch.randn(d_in, d_in, device=device)
        Q, _ = torch.linalg.qr(rand_mat)
        return Q, torch.ones(d_in, device=device)

    @staticmethod
    def _detect_lora(module: nn.Module) -> tuple[int, float]:
        """Detect existing LoRA rank and alpha from a wrapped module, if any."""
        if not hasattr(module, "lora_A") or module.lora_A is None:
            return 0, 1.0
        rank = 0
        la = getattr(module, "lora_A", None)
        if la is not None:
            if hasattr(la, "weight") and la.weight is not None and la.weight.dim() >= 1:
                rank = int(la.weight.shape[0])
            elif isinstance(la, nn.Parameter) and la.dim() >= 1:
                rank = int(la.shape[0])
        alpha = 1.0
        for attr in ("lora_alpha", "alpha"):
            a = getattr(module, attr, None)
            if a is not None:
                try:
                    alpha = float(a)
                except (TypeError, ValueError):
                    pass
                break
        return rank, alpha

    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """
        Applies SliceGPT rotation + rank-r truncation and replaces each weight
        with a real SVDLinear low-rank factorization.
        """
        if SVDLinear is None:
            logger.error(
                "SliceGPT: SVDLinear could not be imported (%s). Cannot perform "
                "real compression; returning model unchanged.",
                _SVD_IMPORT_ERROR,
            )
            stats = self.get_compression_stats()
            stats.update(
                {
                    "is_noop": True,
                    "noop_reason": "svd_linear_import_failed",
                    "method": "slicegpt_unavailable",
                }
            )
            return self.model, stats

        self.layer_errors = {}
        self.layer_ranks = {}
        activations_map = self._collect_layer_activations()
        sliced_layers_count = 0
        total_params_before = 0
        total_params_after = 0

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            base_mod = getattr(module, "base_layer", module)
            raw_w = base_mod.weight.data
            cls_name = base_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(base_mod, "nf")
            device = raw_w.device
            dtype = raw_w.dtype

            if is_conv1d:
                W = raw_w.t().float()
                d_out, d_in = W.shape
            else:
                W = raw_w.float()
                d_out, d_in = W.shape

            total_params_before += d_out * d_in

            act = activations_map.get(name, None)
            Q, _ = self._compute_orthogonal_rotation(act, d_in, device=device)
            Q = Q.to(device=device, dtype=torch.float32)

            # Rotate weight into the PCA basis
            W_pca = torch.matmul(W, Q)

            # Determine number of principal coordinates to retain
            if self.non_uniform and any(k in name for k in ("query", "key", "q_lin", "k_lin")):
                head_sparsity = max(0.0, self.sparsity * 0.75)
                keep = max(1, int(d_in * (1.0 - head_sparsity)))
            elif self.non_uniform and any(k in name for k in ("mlp", "ffn", "dense")):
                mlp_sparsity = min(0.9, self.sparsity * 1.15)
                keep = max(1, int(d_in * (1.0 - mlp_sparsity)))
            else:
                keep = max(1, int(d_in * (1.0 - self.sparsity)))

            # Rank r cannot exceed min(d_out, d_in)
            effective_keep = min(keep, d_out, d_in)

            # Reconstruct truncated dense weight for error measurement only
            W_trunc = W_pca.clone()
            if effective_keep < d_in:
                W_trunc[:, effective_keep:] = 0.0
            W_reconstructed = torch.matmul(W_trunc, Q.T)

            norm_orig = torch.norm(W, p="fro").clamp_min(1e-8)
            rel_err = (torch.norm(W - W_reconstructed, p="fro") / norm_orig).item()
            self.layer_errors[name] = float(rel_err)

            # Fallback on high reconstruction error
            if rel_err > self.max_error:
                logger.warning(
                    "SliceGPT reconstruction error %.4f on %s exceeded threshold %.4f%s",
                    rel_err,
                    name,
                    self.max_error,
                    (
                        ". Skipping (fallback to original)."
                        if self.fallback_to_original
                        else ". Proceeding with sliced weight."
                    ),
                )
                if self.fallback_to_original:
                    logger.warning(
                        "Fallback triggered for %s. Layer is kept original and NOT wrapped in LoRA. "
                        "It will not be fine-tuned.",
                        name,
                    )
                    for p in module.parameters():
                        p.requires_grad = False
                    total_params_after += d_out * d_in
                    continue

            # === REAL LOW-RANK FACTORIZATION ===
            # W_reconstructed = W_pca[:, :keep] @ Q[:, :keep]^T
            # Factor:
            #   U_factor  = W_pca[:, :keep]           shape (d_out, keep)
            #   Vt_factor = Q[:, :keep].T             shape (keep, d_in)
            U_factor = W_pca[:, :effective_keep].contiguous().to(dtype=dtype, device=device)
            Vt_factor = Q[:, :effective_keep].T.contiguous().to(dtype=dtype, device=device)

            # Parameter count of the SVDLinear factorization
            factorized_params = effective_keep * (d_out + d_in)
            # Only apply if it actually reduces parameter count
            if factorized_params >= d_out * d_in:
                logger.info(
                    "SliceGPT: keep=%d on %s would not reduce params (%d >= %d). Keeping dense.",
                    effective_keep,
                    name,
                    factorized_params,
                    d_out * d_in,
                )
                total_params_after += d_out * d_in
                continue

            # Detect existing LoRA adapters so we can preserve their weights
            existing_lora_r, existing_lora_alpha = self._detect_lora(module)
            if self.use_lora and self.lora_r > 0 and existing_lora_r == 0:
                # Apply new LoRA on top of the slicegpt factorization
                used_lora_r = self.lora_r
                used_lora_alpha = self.lora_alpha
            else:
                used_lora_r = existing_lora_r
                used_lora_alpha = existing_lora_alpha

            # Build the SVDLinear factorization
            svd_layer = SVDLinear(
                in_features=d_in,
                out_features=d_out,
                rank=effective_keep,
                bias=(getattr(base_mod, "bias", None) is not None and base_mod.bias is not None),
                dtype=dtype,
                device=device,
                lora_r=used_lora_r,
                lora_alpha=used_lora_alpha,
            )
            svd_layer.set_low_rank_factors(U_factor, Vt_factor)

            # Copy bias
            if (
                getattr(base_mod, "bias", None) is not None
                and base_mod.bias is not None
                and getattr(svd_layer, "bias", None) is not None
            ):
                svd_layer.bias.data.copy_(base_mod.bias.data)

            # Copy existing LoRA adapters if they were present on the source module
            if existing_lora_r > 0 and getattr(svd_layer, "lora_A", None) is not None:
                try:
                    src_la = getattr(module.lora_A, "weight", module.lora_A)
                    src_lb = getattr(module.lora_B, "weight", module.lora_B)
                    la_target = svd_layer.lora_A
                    lb_target = svd_layer.lora_B
                    if (
                        hasattr(src_la, "data")
                        and la_target is not None
                        and src_la.shape == la_target.shape
                    ):
                        la_target.data.copy_(src_la.data)
                    if (
                        hasattr(src_lb, "data")
                        and lb_target is not None
                        and src_lb.shape == lb_target.shape
                    ):
                        lb_target.data.copy_(src_lb.data)
                except Exception as e_lora:
                    logger.debug("SliceGPT LoRA adapter copy skipped on %s: %s", name, e_lora)

            # Replace the module
            is_wrapped = hasattr(module, "base_layer")
            if is_wrapped:
                module.base_layer = svd_layer
            else:
                replace_module(self.model, name, svd_layer)

            self.layer_ranks[name] = effective_keep
            sliced_layers_count += 1
            total_params_after += factorized_params

        # Post-compression forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                dummy = torch.randint(0, 100, (1, 8), device=self.target_device)
                self.model(dummy)
        except Exception as e:
            logger.warning("SliceGPT post-compression forward check warning: %s", e)
            forward_check_success = False

        self._compressed_params = self._count_params(self.model)
        avg_err = (
            sum(self.layer_errors.values()) / len(self.layer_errors) if self.layer_errors else 0.0
        )
        compression_rate = (
            total_params_before / max(1, total_params_after) if total_params_after > 0 else 1.0
        )

        stats = self.get_compression_stats()
        stats.update(
            {
                "sliced_layers": sliced_layers_count,
                "sparsity_target": self.sparsity,
                "non_uniform": self.non_uniform,
                "layer_errors": self.layer_errors,
                "layer_ranks": self.layer_ranks,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "target_attention_only": self.target_attention_only,
                "use_lora": self.use_lora,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "total_params_before": total_params_before,
                "total_params_after": total_params_after,
                "compression_rate": round(compression_rate, 4),
                "avg_reconstruction_error": avg_err,
                "forward_check_success": forward_check_success,
                "is_noop": (total_params_after >= total_params_before),
                "noop_reason": (
                    "no_factorizable_layers_or_fallback_all"
                    if total_params_after >= total_params_before
                    else None
                ),
                "method": "slicegpt_real_low_rank_factorization",
            }
        )
        self.stats = stats

        logger.info(
            "SliceGPT (real): %d layers factorized, avg recon error %.4f, "
            "params %d -> %d (%.2fx compression)",
            sliced_layers_count,
            avg_err,
            total_params_before,
            total_params_after,
            compression_rate,
        )
        return self.model, stats

SliceGPT Compression Strategy (real low-rank implementation).

Performs orthogonal rotation of each target weight matrix using the eigendecomposition of its input activation covariance, then truncates to the top-r principal coordinates. The truncated weight is stored as a rank-r SVDLinear factorization, achieving genuine parameter reduction.

Ancestors

BaseCompressor

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """
    Applies SliceGPT rotation + rank-r truncation and replaces each weight
    with a real SVDLinear low-rank factorization.
    """
    if SVDLinear is None:
        logger.error(
            "SliceGPT: SVDLinear could not be imported (%s). Cannot perform "
            "real compression; returning model unchanged.",
            _SVD_IMPORT_ERROR,
        )
        stats = self.get_compression_stats()
        stats.update(
            {
                "is_noop": True,
                "noop_reason": "svd_linear_import_failed",
                "method": "slicegpt_unavailable",
            }
        )
        return self.model, stats

    self.layer_errors = {}
    self.layer_ranks = {}
    activations_map = self._collect_layer_activations()
    sliced_layers_count = 0
    total_params_before = 0
    total_params_after = 0

    for name, module in list(self.model.named_modules()):
        if not self._is_target_module(name, module):
            continue

        base_mod = getattr(module, "base_layer", module)
        raw_w = base_mod.weight.data
        cls_name = base_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(base_mod, "nf")
        device = raw_w.device
        dtype = raw_w.dtype

        if is_conv1d:
            W = raw_w.t().float()
            d_out, d_in = W.shape
        else:
            W = raw_w.float()
            d_out, d_in = W.shape

        total_params_before += d_out * d_in

        act = activations_map.get(name, None)
        Q, _ = self._compute_orthogonal_rotation(act, d_in, device=device)
        Q = Q.to(device=device, dtype=torch.float32)

        # Rotate weight into the PCA basis
        W_pca = torch.matmul(W, Q)

        # Determine number of principal coordinates to retain
        if self.non_uniform and any(k in name for k in ("query", "key", "q_lin", "k_lin")):
            head_sparsity = max(0.0, self.sparsity * 0.75)
            keep = max(1, int(d_in * (1.0 - head_sparsity)))
        elif self.non_uniform and any(k in name for k in ("mlp", "ffn", "dense")):
            mlp_sparsity = min(0.9, self.sparsity * 1.15)
            keep = max(1, int(d_in * (1.0 - mlp_sparsity)))
        else:
            keep = max(1, int(d_in * (1.0 - self.sparsity)))

        # Rank r cannot exceed min(d_out, d_in)
        effective_keep = min(keep, d_out, d_in)

        # Reconstruct truncated dense weight for error measurement only
        W_trunc = W_pca.clone()
        if effective_keep < d_in:
            W_trunc[:, effective_keep:] = 0.0
        W_reconstructed = torch.matmul(W_trunc, Q.T)

        norm_orig = torch.norm(W, p="fro").clamp_min(1e-8)
        rel_err = (torch.norm(W - W_reconstructed, p="fro") / norm_orig).item()
        self.layer_errors[name] = float(rel_err)

        # Fallback on high reconstruction error
        if rel_err > self.max_error:
            logger.warning(
                "SliceGPT reconstruction error %.4f on %s exceeded threshold %.4f%s",
                rel_err,
                name,
                self.max_error,
                (
                    ". Skipping (fallback to original)."
                    if self.fallback_to_original
                    else ". Proceeding with sliced weight."
                ),
            )
            if self.fallback_to_original:
                logger.warning(
                    "Fallback triggered for %s. Layer is kept original and NOT wrapped in LoRA. "
                    "It will not be fine-tuned.",
                    name,
                )
                for p in module.parameters():
                    p.requires_grad = False
                total_params_after += d_out * d_in
                continue

        # === REAL LOW-RANK FACTORIZATION ===
        # W_reconstructed = W_pca[:, :keep] @ Q[:, :keep]^T
        # Factor:
        #   U_factor  = W_pca[:, :keep]           shape (d_out, keep)
        #   Vt_factor = Q[:, :keep].T             shape (keep, d_in)
        U_factor = W_pca[:, :effective_keep].contiguous().to(dtype=dtype, device=device)
        Vt_factor = Q[:, :effective_keep].T.contiguous().to(dtype=dtype, device=device)

        # Parameter count of the SVDLinear factorization
        factorized_params = effective_keep * (d_out + d_in)
        # Only apply if it actually reduces parameter count
        if factorized_params >= d_out * d_in:
            logger.info(
                "SliceGPT: keep=%d on %s would not reduce params (%d >= %d). Keeping dense.",
                effective_keep,
                name,
                factorized_params,
                d_out * d_in,
            )
            total_params_after += d_out * d_in
            continue

        # Detect existing LoRA adapters so we can preserve their weights
        existing_lora_r, existing_lora_alpha = self._detect_lora(module)
        if self.use_lora and self.lora_r > 0 and existing_lora_r == 0:
            # Apply new LoRA on top of the slicegpt factorization
            used_lora_r = self.lora_r
            used_lora_alpha = self.lora_alpha
        else:
            used_lora_r = existing_lora_r
            used_lora_alpha = existing_lora_alpha

        # Build the SVDLinear factorization
        svd_layer = SVDLinear(
            in_features=d_in,
            out_features=d_out,
            rank=effective_keep,
            bias=(getattr(base_mod, "bias", None) is not None and base_mod.bias is not None),
            dtype=dtype,
            device=device,
            lora_r=used_lora_r,
            lora_alpha=used_lora_alpha,
        )
        svd_layer.set_low_rank_factors(U_factor, Vt_factor)

        # Copy bias
        if (
            getattr(base_mod, "bias", None) is not None
            and base_mod.bias is not None
            and getattr(svd_layer, "bias", None) is not None
        ):
            svd_layer.bias.data.copy_(base_mod.bias.data)

        # Copy existing LoRA adapters if they were present on the source module
        if existing_lora_r > 0 and getattr(svd_layer, "lora_A", None) is not None:
            try:
                src_la = getattr(module.lora_A, "weight", module.lora_A)
                src_lb = getattr(module.lora_B, "weight", module.lora_B)
                la_target = svd_layer.lora_A
                lb_target = svd_layer.lora_B
                if (
                    hasattr(src_la, "data")
                    and la_target is not None
                    and src_la.shape == la_target.shape
                ):
                    la_target.data.copy_(src_la.data)
                if (
                    hasattr(src_lb, "data")
                    and lb_target is not None
                    and src_lb.shape == lb_target.shape
                ):
                    lb_target.data.copy_(src_lb.data)
            except Exception as e_lora:
                logger.debug("SliceGPT LoRA adapter copy skipped on %s: %s", name, e_lora)

        # Replace the module
        is_wrapped = hasattr(module, "base_layer")
        if is_wrapped:
            module.base_layer = svd_layer
        else:
            replace_module(self.model, name, svd_layer)

        self.layer_ranks[name] = effective_keep
        sliced_layers_count += 1
        total_params_after += factorized_params

    # Post-compression forward sanity check
    forward_check_success = True
    try:
        with torch.no_grad():
            dummy = torch.randint(0, 100, (1, 8), device=self.target_device)
            self.model(dummy)
    except Exception as e:
        logger.warning("SliceGPT post-compression forward check warning: %s", e)
        forward_check_success = False

    self._compressed_params = self._count_params(self.model)
    avg_err = (
        sum(self.layer_errors.values()) / len(self.layer_errors) if self.layer_errors else 0.0
    )
    compression_rate = (
        total_params_before / max(1, total_params_after) if total_params_after > 0 else 1.0
    )

    stats = self.get_compression_stats()
    stats.update(
        {
            "sliced_layers": sliced_layers_count,
            "sparsity_target": self.sparsity,
            "non_uniform": self.non_uniform,
            "layer_errors": self.layer_errors,
            "layer_ranks": self.layer_ranks,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "target_attention_only": self.target_attention_only,
            "use_lora": self.use_lora,
            "lora_r": self.lora_r,
            "lora_alpha": self.lora_alpha,
            "total_params_before": total_params_before,
            "total_params_after": total_params_after,
            "compression_rate": round(compression_rate, 4),
            "avg_reconstruction_error": avg_err,
            "forward_check_success": forward_check_success,
            "is_noop": (total_params_after >= total_params_before),
            "noop_reason": (
                "no_factorizable_layers_or_fallback_all"
                if total_params_after >= total_params_before
                else None
            ),
            "method": "slicegpt_real_low_rank_factorization",
        }
    )
    self.stats = stats

    logger.info(
        "SliceGPT (real): %d layers factorized, avg recon error %.4f, "
        "params %d -> %d (%.2fx compression)",
        sliced_layers_count,
        avg_err,
        total_params_before,
        total_params_after,
        compression_rate,
    )
    return self.model, stats

Applies SliceGPT rotation + rank-r truncation and replaces each weight with a real SVDLinear low-rank factorization.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class SliceGPTPCACompressor (model: nn.Module, calibration_loader: Any | None = None, sparsity: float = 0.3, tokenizer: Any | None = None, calibration_texts: list[str] | None = None, skip_attention_qkv: bool = False, target_attention_only: bool = False, use_lora: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, max_error: float = 0.5, fallback_to_original: bool = True, min_numel: int = 16, target_modules: list[str] | None = None, device: str | torch.device | None = None, cfg: Any | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("slicegpt_pca")
class SliceGPTPCACompressor(BaseCompressor):
    """
    Production-hardened SliceGPT PCA-based compressor that performs activation-aware
    principal coordinate projection and dimension slicing on Linear and Conv1D layers.

    Replaces each sliced dense weight with a real SVDLinear low-rank factorization
    so that parameter count actually drops from d_out*d_in to (d_out+d_in)*keep.
    """

    def __init__(
        self,
        model: nn.Module,
        calibration_loader: Any | None = None,
        sparsity: float = 0.3,
        tokenizer: Any | None = None,
        calibration_texts: list[str] | None = None,
        skip_attention_qkv: bool = False,
        target_attention_only: bool = False,
        use_lora: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        min_numel: int = 16,
        target_modules: list[str] | None = None,
        device: str | torch.device | None = None,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="slicegpt_pca", **kwargs)
        cfg_obj = cfg or kwargs.get("comp_config")

        self.calibration_loader = calibration_loader
        self.sparsity = max(
            0.0,
            min(
                float(
                    sparsity
                    or _cfg_get(
                        cfg_obj,
                        "slicegpt_pca_sparsity",
                        _cfg_get(cfg_obj, "sparsity", 0.3),
                    )
                ),
                0.95,
            ),
        )
        self.tokenizer = (
            tokenizer or getattr(model, "tokenizer", None) or _cfg_get(cfg_obj, "tokenizer", None)
        )
        self.calibration_texts = calibration_texts or _cfg_get(cfg_obj, "calibration_texts", None)
        self.skip_attention_qkv = bool(
            skip_attention_qkv
            if skip_attention_qkv is not None
            else _cfg_get(cfg_obj, "skip_attention_qkv", False)
        )
        self.target_attention_only = bool(
            target_attention_only
            if target_attention_only is not None
            else _cfg_get(cfg_obj, "target_attention_only", False)
        )
        self.use_lora = bool(use_lora or _cfg_get(cfg_obj, "use_lora", False) or (lora_r > 0))
        self.lora_r = int(lora_r or _cfg_get(cfg_obj, "lora_r", 0))
        self.lora_alpha = float(lora_alpha or _cfg_get(cfg_obj, "lora_alpha", 1.0))
        self.max_error = float(max_error or _cfg_get(cfg_obj, "max_error", 0.50))
        self.fallback_to_original = bool(
            fallback_to_original
            if fallback_to_original is not None
            else _cfg_get(cfg_obj, "fallback_to_original", True)
        )
        self.min_numel = max(1, int(min_numel or _cfg_get(cfg_obj, "min_numel", 16)))
        self.target_modules = target_modules

        # Target device derivation
        if device is not None:
            try:
                self.target_device = torch.device(device)
            except Exception:
                self.target_device = (
                    next(model.parameters()).device
                    if list(model.parameters())
                    else torch.device("cpu")
                )
        else:
            self.target_device = (
                next(model.parameters()).device if list(model.parameters()) else torch.device("cpu")
            )

        # Diagnostics & Tracking
        self.visited_modules: set[int] = set()
        self.sliced_layers: list[str] = []
        self.protected_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.failed_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.layer_ranks: dict[str, int] = {}
        self.stats: dict[str, Any] = {}

    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        """Check if layer belongs to an attention projection in a GQA architecture."""
        cfg = getattr(self.model, "config", None)
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        is_gqa = num_heads is not None and num_kv_heads is not None and num_kv_heads < num_heads

        name_l = name.lower()
        is_attn_qkv = any(
            p in name_l
            for p in [
                "q_proj",
                "k_proj",
                "v_proj",
                "c_attn",
                "query",
                "key",
                "value",
                "q_lin",
                "k_lin",
                "v_lin",
                "self_attn.q",
                "self_attn.k",
                "self_attn.v",
            ]
        )
        return (is_gqa or self.skip_attention_qkv) and is_attn_qkv

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        # 1. Protection for critical / pooler / classification / embedding / CLS layers
        if any(
            k in name_l
            for k in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embed",
                "embedding",
                "cls_head",
                "head",
                "logits_processor",
            ]
        ):
            self.protected_layers.append(name)
            return False

        # Target attention only filter
        if self.target_attention_only:
            attn_keywords = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(k in name_l for k in attn_keywords):
                return False

        # 2. GQA-aware skipping
        if self._is_gqa_asymmetric(name, module):
            self.skipped_layers.append(name)
            return False

        # 3. Check linear or Conv1D nature
        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__
        raw_w = getattr(base_mod, "weight", None)
        if raw_w is None or not isinstance(raw_w, torch.Tensor):
            return False

        is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_mod, "nf")
        is_linear_like = isinstance(base_mod, nn.Linear) or is_conv1d
        if not is_linear_like:
            return False

        # Skip already-factorized layers
        if SVDLinear is not None and isinstance(base_mod, SVDLinear):
            return False

        # 4. Tiny layer guard
        if raw_w.numel() < self.min_numel:
            self.skipped_layers.append(name)
            return False

        # 5. Respect framework compatibility rules
        if not should_compress_linear(self.model, name):
            self.protected_layers.append(name)
            return False

        # 6. User-defined explicit target modules
        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        return True

    def _collect_layer_activations(self) -> dict[str, torch.Tensor]:
        """Capture streaming empirical covariance across target layers to avoid memory spikes."""
        self.model.eval()
        sum_x: dict[str, torch.Tensor] = {}
        sum_xx: dict[str, torch.Tensor] = {}
        counts: dict[str, int] = {}
        hooks = []

        def make_hook(layer_name: str):
            def hook(module: nn.Module, inputs: tuple[torch.Tensor, ...], output: Any):
                if inputs and isinstance(inputs[0], torch.Tensor):
                    inp = inputs[0].detach()
                    x_flat = inp.reshape(-1, inp.shape[-1]).float()
                    n = x_flat.shape[0]
                    s_x = x_flat.sum(dim=0, keepdim=True).cpu()
                    s_xx = torch.matmul(x_flat.T, x_flat).cpu()
                    if layer_name not in sum_x:
                        sum_x[layer_name] = s_x
                        sum_xx[layer_name] = s_xx
                        counts[layer_name] = n
                    else:
                        sum_x[layer_name] += s_x
                        sum_xx[layer_name] += s_xx
                        counts[layer_name] += n

            return hook

        try:
            for name, module in self.model.named_modules():
                if self._is_target_module(name, module):
                    hooks.append(module.register_forward_hook(make_hook(name)))

            calib_inputs = get_calibration_tensors(
                self.model,
                tokenizer=self.tokenizer,
                calibration_loader=self.calibration_loader,
                calibration_texts=self.calibration_texts,
            )

            with torch.no_grad():
                try:
                    if isinstance(calib_inputs, dict):
                        self.model(
                            **{
                                k: (v.to(self.target_device) if isinstance(v, torch.Tensor) else v)
                                for k, v in calib_inputs.items()
                            }
                        )
                    elif isinstance(calib_inputs, torch.Tensor):
                        self.model(calib_inputs.to(self.target_device))
                    elif isinstance(calib_inputs, (list, tuple)):
                        self.model(
                            *[
                                (t.to(self.target_device) if isinstance(t, torch.Tensor) else t)
                                for t in calib_inputs
                            ]
                        )
                except Exception as ex_calib:
                    logger.debug(
                        "SliceGPTPCA primary calibration forward failed (%s); trying synthetic fallback.",
                        ex_calib,
                    )
                    try:
                        dummy_ids = torch.randint(0, 100, (1, 16), device=self.target_device)
                        self.model(input_ids=dummy_ids)
                    except Exception as ex_synth:
                        logger.warning("SliceGPTPCA calibration forward pass skipped: %s", ex_synth)
        finally:
            for h in hooks:
                try:
                    h.remove()
                except Exception as e:
                    # Fix 5: corrected error message (was "copy_weights failed")
                    logger.warning("SliceGPTPCA hook removal failed: %s", e)

        cov_map: dict[str, torch.Tensor] = {}
        for name, n in counts.items():
            if n > 1:
                mu = sum_x[name] / float(n)
                cov = (sum_xx[name] - float(n) * torch.matmul(mu.T, mu)) / float(max(1, n - 1))
                cov_map[name] = cov
        return cov_map

    @staticmethod
    def _detect_lora(module: nn.Module) -> tuple[int, float]:
        """Detect existing LoRA rank and alpha from a wrapped module, if any."""
        if not hasattr(module, "lora_A") or module.lora_A is None:
            return 0, 1.0
        rank = 0
        la = getattr(module, "lora_A", None)
        if la is not None:
            if hasattr(la, "weight") and la.weight is not None and la.weight.dim() >= 1:
                rank = int(la.weight.shape[0])
            elif isinstance(la, nn.Parameter) and la.dim() >= 1:
                rank = int(la.shape[0])
        alpha = 1.0
        for attr in ("lora_alpha", "alpha"):
            a = getattr(module, attr, None)
            if a is not None:
                try:
                    alpha = float(a)
                except (TypeError, ValueError):
                    pass
                break
        return rank, alpha

    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """
        Compresses the model using PCA-based principal subspace projection and
        real low-rank factorization (SVDLinear).

        Returns (compressed_model, stats).
        """
        if SVDLinear is None:
            logger.error(
                "SliceGPTPCA: SVDLinear could not be imported (%s). Cannot perform "
                "real compression; returning model unchanged.",
                _SVD_IMPORT_ERROR,
            )
            stats = self.get_compression_stats()
            stats.update(
                {
                    "is_noop": True,
                    "noop_reason": "svd_linear_import_failed",
                    "method": "slicegpt_pca_unavailable",
                }
            )
            return self.model, stats

        self.layer_errors = {}
        self.layer_ranks = {}
        self.visited_modules.clear()
        self.sliced_layers.clear()
        self.protected_layers.clear()
        self.skipped_layers.clear()
        self.fallback_layers.clear()
        self.failed_layers.clear()

        activations_map = self._collect_layer_activations()
        sliced_layers_count = 0
        pruned_features_count = 0
        total_params_before = 0
        total_params_after = 0

        target_items = [
            (n, m) for n, m in self.model.named_modules() if self._is_target_module(n, m)
        ]

        for name, module in target_items:
            # Tied weight protection
            mod_id = id(module)
            if mod_id in self.visited_modules:
                logger.debug("SliceGPTPCA: Skipping tied module %s", name)
                self.skipped_layers.append(name)
                continue
            self.visited_modules.add(mod_id)

            try:
                base_mod = getattr(module, "base_layer", module)
                raw_w = base_mod.weight.data
                cls_name = base_mod.__class__.__name__
                is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_mod, "nf")
                device = raw_w.device
                dtype = raw_w.dtype

                # FP32 weight extraction
                if is_conv1d:
                    W = raw_w.t().contiguous().float()
                else:
                    W = raw_w.contiguous().float()
                d_out, d_in = W.shape

                total_params_before += d_out * d_in

                cov_mat = activations_map.get(name, None)
                if cov_mat is not None and cov_mat.shape[0] == d_in and cov_mat.shape[1] == d_in:
                    cov = cov_mat.to(device=device, dtype=torch.float32) + 1e-5 * torch.eye(
                        d_in, device=device
                    )
                    try:
                        eigvals, V = torch.linalg.eigh(cov)
                        idx = torch.argsort(eigvals, descending=True)
                        eigvals = eigvals[idx]
                        V = V[:, idx]
                        eye_diff = torch.norm(
                            torch.matmul(V.T, V) - torch.eye(d_in), p="fro"
                        ).item()
                        if eye_diff > 1e-2 or torch.isnan(V).any():
                            raise ValueError(f"Eigh basis non-orthogonal (diff={eye_diff:.4f})")
                    except Exception as e_eigh:
                        logger.debug(
                            "SliceGPTPCA: eigh failed on %s (%s); fallback to QR.",
                            name,
                            e_eigh,
                        )
                        H = torch.randn(d_in, d_in)
                        V, _ = torch.linalg.qr(H)
                        eigvals = torch.ones(d_in)
                else:
                    H = torch.randn(d_in, d_in)
                    V, _ = torch.linalg.qr(H)
                    eigvals = torch.ones(d_in)

                V = V.to(device=device, dtype=torch.float32)
                W = W.to(device=device, dtype=torch.float32)

                # Compute cumulative variance
                pos_eigvals = torch.clamp(eigvals.to(device=device), min=0.0)
                total_var = torch.sum(pos_eigvals).item()
                if total_var > 0:
                    cum_var = torch.cumsum(pos_eigvals, dim=0) / total_var
                    target_var = 1.0 - self.sparsity
                    keep = int((cum_var < target_var).sum().item()) + 1
                    keep = max(1, min(keep, d_in))
                else:
                    keep = max(1, int(d_in * (1.0 - self.sparsity)))

                # Rank cannot exceed min(d_out, d_in)
                effective_keep = min(keep, d_out, d_in)

                # Rotate weight into PCA coordinate space and truncate
                W_pca = torch.matmul(W, V)
                W_trunc = W_pca.clone()
                if effective_keep < d_in:
                    W_trunc[:, effective_keep:] = 0.0

                # Reconstruct truncated weight for error measurement
                W_reconstructed = torch.matmul(W_trunc, V.T)

                # Relative Frobenius error validation
                norm_orig = torch.norm(W, p="fro").clamp_min(1e-8)
                rel_err = (torch.norm(W - W_reconstructed, p="fro") / norm_orig).item()
                self.layer_errors[name] = float(rel_err)

                is_invalid = (
                    torch.isnan(W_reconstructed).any() or torch.isinf(W_reconstructed).any()
                )
                if rel_err > self.max_error:
                    logger.warning(
                        "SliceGPTPCA: reconstruction error %.4f on %s exceeded threshold %.4f.",
                        rel_err,
                        name,
                        self.max_error,
                    )
                    is_invalid = True

                if is_invalid and self.fallback_to_original:
                    logger.warning(
                        "Fallback triggered for %s. Layer is kept original and NOT wrapped in LoRA. "
                        "It will not be fine-tuned.",
                        name,
                    )
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    total_params_after += d_out * d_in
                    continue

                # === REAL LOW-RANK FACTORIZATION ===
                # W_reconstructed = W_pca[:, :keep] @ V[:, :keep].T
                # Factor:
                #   U_factor  = W_pca[:, :keep]      shape (d_out, keep)
                #   Vt_factor = V[:, :keep].T        shape (keep, d_in)
                U_factor = W_pca[:, :effective_keep].contiguous().to(dtype=dtype, device=device)
                Vt_factor = V[:, :effective_keep].T.contiguous().to(dtype=dtype, device=device)

                factorized_params = effective_keep * (d_out + d_in)
                # Only apply if it actually reduces parameter count
                if factorized_params >= d_out * d_in:
                    logger.info(
                        "SliceGPTPCA: keep=%d on %s would not reduce params (%d >= %d). Keeping dense.",
                        effective_keep,
                        name,
                        factorized_params,
                        d_out * d_in,
                    )
                    total_params_after += d_out * d_in
                    continue

                # Detect existing LoRA adapters so we can preserve their weights
                existing_lora_r, existing_lora_alpha = self._detect_lora(module)
                if self.use_lora and self.lora_r > 0 and existing_lora_r == 0:
                    used_lora_r = self.lora_r
                    used_lora_alpha = self.lora_alpha
                else:
                    used_lora_r = existing_lora_r
                    used_lora_alpha = existing_lora_alpha

                # Build the SVDLinear factorization
                svd_layer = SVDLinear(
                    in_features=d_in,
                    out_features=d_out,
                    rank=effective_keep,
                    bias=(
                        getattr(base_mod, "bias", None) is not None and base_mod.bias is not None
                    ),
                    dtype=dtype,
                    device=device,
                    lora_r=used_lora_r,
                    lora_alpha=used_lora_alpha,
                )
                svd_layer.set_low_rank_factors(U_factor, Vt_factor)

                # Copy bias
                if (
                    getattr(base_mod, "bias", None) is not None
                    and base_mod.bias is not None
                    and getattr(svd_layer, "bias", None) is not None
                ):
                    svd_layer.bias.data.copy_(base_mod.bias.data)

                # Copy existing LoRA adapters if present
                if existing_lora_r > 0 and getattr(svd_layer, "lora_A", None) is not None:
                    try:
                        src_la = getattr(module.lora_A, "weight", module.lora_A)
                        src_lb = getattr(module.lora_B, "weight", module.lora_B)
                        la_target = svd_layer.lora_A
                        lb_target = svd_layer.lora_B
                        if (
                            hasattr(src_la, "data")
                            and la_target is not None
                            and src_la.shape == la_target.shape
                        ):
                            la_target.data.copy_(src_la.data)
                        if (
                            hasattr(src_lb, "data")
                            and lb_target is not None
                            and src_lb.shape == lb_target.shape
                        ):
                            lb_target.data.copy_(src_lb.data)
                    except Exception as e_lora:
                        logger.debug("SliceGPTPCA LoRA copy skipped on %s: %s", name, e_lora)

                # Replace module
                is_wrapped = hasattr(module, "base_layer")
                if is_wrapped:
                    module.base_layer = svd_layer
                else:
                    replace_module(self.model, name, svd_layer)

                self.layer_ranks[name] = effective_keep
                self.sliced_layers.append(name)
                sliced_layers_count += 1
                if effective_keep < d_in:
                    pruned_features_count += d_in - effective_keep
                total_params_after += factorized_params

            except Exception as e_layer:
                logger.warning("SliceGPTPCA: Failed to compress %s: %s", name, e_layer)
                self.failed_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False

        # Post-compression forward sanity check
        forward_check_success = True
        try:
            with torch.no_grad():
                dummy_input = torch.randint(0, 100, (1, 8), device=self.target_device)
                self.model(dummy_input)
        except Exception as ex_sanity:
            logger.warning(
                "SliceGPTPCA post-compression forward sanity check warning: %s",
                ex_sanity,
            )
            forward_check_success = False

        self._compressed_params = self._count_params(self.model)
        avg_err = (
            (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
        )
        compression_rate = (
            total_params_before / max(1, total_params_after) if total_params_after > 0 else 1.0
        )

        stats = self.get_compression_stats()
        stats.update(
            {
                "sliced_layers": self.sliced_layers,
                "sliced_layer_count": sliced_layers_count,
                "layer_ranks": self.layer_ranks,
                "pruned_features": pruned_features_count,
                "protected_layers": self.protected_layers,
                "protected_layer_count": len(self.protected_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "failed_layers": self.failed_layers,
                "failed_layer_count": len(self.failed_layers),
                "sparsity_target": self.sparsity,
                "layer_errors": self.layer_errors,
                "avg_reconstruction_error": avg_err,
                "forward_check_success": forward_check_success,
                "skip_attention_qkv": self.skip_attention_qkv,
                "target_attention_only": self.target_attention_only,
                "use_lora": self.use_lora,
                "lora_r": self.lora_r,
                "total_params_before": total_params_before,
                "total_params_after": total_params_after,
                "compression_rate": round(compression_rate, 4),
                "num_compressed": sliced_layers_count,
                "num_compressed_layers": sliced_layers_count,
                "num_fallback": len(self.fallback_layers),
                "num_fallback_layers": len(self.fallback_layers),
                "is_noop": (total_params_after >= total_params_before),
                "noop_reason": (
                    "no_factorizable_layers_or_fallback_all"
                    if total_params_after >= total_params_before
                    else None
                ),
                "method": "slicegpt_pca_real_low_rank_factorization",
            }
        )
        self.stats = stats

        logger.info(
            "SliceGPTPCA (real): %d layers factorized, avg recon error %.4f, "
            "params %d -> %d (%.2fx compression)",
            sliced_layers_count,
            avg_err,
            total_params_before,
            total_params_after,
            compression_rate,
        )
        return self.model, stats

Production-hardened SliceGPT PCA-based compressor that performs activation-aware principal coordinate projection and dimension slicing on Linear and Conv1D layers.

Replaces each sliced dense weight with a real SVDLinear low-rank factorization so that parameter count actually drops from d_outd_in to (d_out+d_in)keep.

Ancestors

BaseCompressor

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """
    Compresses the model using PCA-based principal subspace projection and
    real low-rank factorization (SVDLinear).

    Returns (compressed_model, stats).
    """
    if SVDLinear is None:
        logger.error(
            "SliceGPTPCA: SVDLinear could not be imported (%s). Cannot perform "
            "real compression; returning model unchanged.",
            _SVD_IMPORT_ERROR,
        )
        stats = self.get_compression_stats()
        stats.update(
            {
                "is_noop": True,
                "noop_reason": "svd_linear_import_failed",
                "method": "slicegpt_pca_unavailable",
            }
        )
        return self.model, stats

    self.layer_errors = {}
    self.layer_ranks = {}
    self.visited_modules.clear()
    self.sliced_layers.clear()
    self.protected_layers.clear()
    self.skipped_layers.clear()
    self.fallback_layers.clear()
    self.failed_layers.clear()

    activations_map = self._collect_layer_activations()
    sliced_layers_count = 0
    pruned_features_count = 0
    total_params_before = 0
    total_params_after = 0

    target_items = [
        (n, m) for n, m in self.model.named_modules() if self._is_target_module(n, m)
    ]

    for name, module in target_items:
        # Tied weight protection
        mod_id = id(module)
        if mod_id in self.visited_modules:
            logger.debug("SliceGPTPCA: Skipping tied module %s", name)
            self.skipped_layers.append(name)
            continue
        self.visited_modules.add(mod_id)

        try:
            base_mod = getattr(module, "base_layer", module)
            raw_w = base_mod.weight.data
            cls_name = base_mod.__class__.__name__
            is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_mod, "nf")
            device = raw_w.device
            dtype = raw_w.dtype

            # FP32 weight extraction
            if is_conv1d:
                W = raw_w.t().contiguous().float()
            else:
                W = raw_w.contiguous().float()
            d_out, d_in = W.shape

            total_params_before += d_out * d_in

            cov_mat = activations_map.get(name, None)
            if cov_mat is not None and cov_mat.shape[0] == d_in and cov_mat.shape[1] == d_in:
                cov = cov_mat.to(device=device, dtype=torch.float32) + 1e-5 * torch.eye(
                    d_in, device=device
                )
                try:
                    eigvals, V = torch.linalg.eigh(cov)
                    idx = torch.argsort(eigvals, descending=True)
                    eigvals = eigvals[idx]
                    V = V[:, idx]
                    eye_diff = torch.norm(
                        torch.matmul(V.T, V) - torch.eye(d_in), p="fro"
                    ).item()
                    if eye_diff > 1e-2 or torch.isnan(V).any():
                        raise ValueError(f"Eigh basis non-orthogonal (diff={eye_diff:.4f})")
                except Exception as e_eigh:
                    logger.debug(
                        "SliceGPTPCA: eigh failed on %s (%s); fallback to QR.",
                        name,
                        e_eigh,
                    )
                    H = torch.randn(d_in, d_in)
                    V, _ = torch.linalg.qr(H)
                    eigvals = torch.ones(d_in)
            else:
                H = torch.randn(d_in, d_in)
                V, _ = torch.linalg.qr(H)
                eigvals = torch.ones(d_in)

            V = V.to(device=device, dtype=torch.float32)
            W = W.to(device=device, dtype=torch.float32)

            # Compute cumulative variance
            pos_eigvals = torch.clamp(eigvals.to(device=device), min=0.0)
            total_var = torch.sum(pos_eigvals).item()
            if total_var > 0:
                cum_var = torch.cumsum(pos_eigvals, dim=0) / total_var
                target_var = 1.0 - self.sparsity
                keep = int((cum_var < target_var).sum().item()) + 1
                keep = max(1, min(keep, d_in))
            else:
                keep = max(1, int(d_in * (1.0 - self.sparsity)))

            # Rank cannot exceed min(d_out, d_in)
            effective_keep = min(keep, d_out, d_in)

            # Rotate weight into PCA coordinate space and truncate
            W_pca = torch.matmul(W, V)
            W_trunc = W_pca.clone()
            if effective_keep < d_in:
                W_trunc[:, effective_keep:] = 0.0

            # Reconstruct truncated weight for error measurement
            W_reconstructed = torch.matmul(W_trunc, V.T)

            # Relative Frobenius error validation
            norm_orig = torch.norm(W, p="fro").clamp_min(1e-8)
            rel_err = (torch.norm(W - W_reconstructed, p="fro") / norm_orig).item()
            self.layer_errors[name] = float(rel_err)

            is_invalid = (
                torch.isnan(W_reconstructed).any() or torch.isinf(W_reconstructed).any()
            )
            if rel_err > self.max_error:
                logger.warning(
                    "SliceGPTPCA: reconstruction error %.4f on %s exceeded threshold %.4f.",
                    rel_err,
                    name,
                    self.max_error,
                )
                is_invalid = True

            if is_invalid and self.fallback_to_original:
                logger.warning(
                    "Fallback triggered for %s. Layer is kept original and NOT wrapped in LoRA. "
                    "It will not be fine-tuned.",
                    name,
                )
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                total_params_after += d_out * d_in
                continue

            # === REAL LOW-RANK FACTORIZATION ===
            # W_reconstructed = W_pca[:, :keep] @ V[:, :keep].T
            # Factor:
            #   U_factor  = W_pca[:, :keep]      shape (d_out, keep)
            #   Vt_factor = V[:, :keep].T        shape (keep, d_in)
            U_factor = W_pca[:, :effective_keep].contiguous().to(dtype=dtype, device=device)
            Vt_factor = V[:, :effective_keep].T.contiguous().to(dtype=dtype, device=device)

            factorized_params = effective_keep * (d_out + d_in)
            # Only apply if it actually reduces parameter count
            if factorized_params >= d_out * d_in:
                logger.info(
                    "SliceGPTPCA: keep=%d on %s would not reduce params (%d >= %d). Keeping dense.",
                    effective_keep,
                    name,
                    factorized_params,
                    d_out * d_in,
                )
                total_params_after += d_out * d_in
                continue

            # Detect existing LoRA adapters so we can preserve their weights
            existing_lora_r, existing_lora_alpha = self._detect_lora(module)
            if self.use_lora and self.lora_r > 0 and existing_lora_r == 0:
                used_lora_r = self.lora_r
                used_lora_alpha = self.lora_alpha
            else:
                used_lora_r = existing_lora_r
                used_lora_alpha = existing_lora_alpha

            # Build the SVDLinear factorization
            svd_layer = SVDLinear(
                in_features=d_in,
                out_features=d_out,
                rank=effective_keep,
                bias=(
                    getattr(base_mod, "bias", None) is not None and base_mod.bias is not None
                ),
                dtype=dtype,
                device=device,
                lora_r=used_lora_r,
                lora_alpha=used_lora_alpha,
            )
            svd_layer.set_low_rank_factors(U_factor, Vt_factor)

            # Copy bias
            if (
                getattr(base_mod, "bias", None) is not None
                and base_mod.bias is not None
                and getattr(svd_layer, "bias", None) is not None
            ):
                svd_layer.bias.data.copy_(base_mod.bias.data)

            # Copy existing LoRA adapters if present
            if existing_lora_r > 0 and getattr(svd_layer, "lora_A", None) is not None:
                try:
                    src_la = getattr(module.lora_A, "weight", module.lora_A)
                    src_lb = getattr(module.lora_B, "weight", module.lora_B)
                    la_target = svd_layer.lora_A
                    lb_target = svd_layer.lora_B
                    if (
                        hasattr(src_la, "data")
                        and la_target is not None
                        and src_la.shape == la_target.shape
                    ):
                        la_target.data.copy_(src_la.data)
                    if (
                        hasattr(src_lb, "data")
                        and lb_target is not None
                        and src_lb.shape == lb_target.shape
                    ):
                        lb_target.data.copy_(src_lb.data)
                except Exception as e_lora:
                    logger.debug("SliceGPTPCA LoRA copy skipped on %s: %s", name, e_lora)

            # Replace module
            is_wrapped = hasattr(module, "base_layer")
            if is_wrapped:
                module.base_layer = svd_layer
            else:
                replace_module(self.model, name, svd_layer)

            self.layer_ranks[name] = effective_keep
            self.sliced_layers.append(name)
            sliced_layers_count += 1
            if effective_keep < d_in:
                pruned_features_count += d_in - effective_keep
            total_params_after += factorized_params

        except Exception as e_layer:
            logger.warning("SliceGPTPCA: Failed to compress %s: %s", name, e_layer)
            self.failed_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False

    # Post-compression forward sanity check
    forward_check_success = True
    try:
        with torch.no_grad():
            dummy_input = torch.randint(0, 100, (1, 8), device=self.target_device)
            self.model(dummy_input)
    except Exception as ex_sanity:
        logger.warning(
            "SliceGPTPCA post-compression forward sanity check warning: %s",
            ex_sanity,
        )
        forward_check_success = False

    self._compressed_params = self._count_params(self.model)
    avg_err = (
        (sum(self.layer_errors.values()) / len(self.layer_errors)) if self.layer_errors else 0.0
    )
    compression_rate = (
        total_params_before / max(1, total_params_after) if total_params_after > 0 else 1.0
    )

    stats = self.get_compression_stats()
    stats.update(
        {
            "sliced_layers": self.sliced_layers,
            "sliced_layer_count": sliced_layers_count,
            "layer_ranks": self.layer_ranks,
            "pruned_features": pruned_features_count,
            "protected_layers": self.protected_layers,
            "protected_layer_count": len(self.protected_layers),
            "skipped_layers": self.skipped_layers,
            "skipped_layer_count": len(self.skipped_layers),
            "fallback_layers": self.fallback_layers,
            "fallback_layer_count": len(self.fallback_layers),
            "failed_layers": self.failed_layers,
            "failed_layer_count": len(self.failed_layers),
            "sparsity_target": self.sparsity,
            "layer_errors": self.layer_errors,
            "avg_reconstruction_error": avg_err,
            "forward_check_success": forward_check_success,
            "skip_attention_qkv": self.skip_attention_qkv,
            "target_attention_only": self.target_attention_only,
            "use_lora": self.use_lora,
            "lora_r": self.lora_r,
            "total_params_before": total_params_before,
            "total_params_after": total_params_after,
            "compression_rate": round(compression_rate, 4),
            "num_compressed": sliced_layers_count,
            "num_compressed_layers": sliced_layers_count,
            "num_fallback": len(self.fallback_layers),
            "num_fallback_layers": len(self.fallback_layers),
            "is_noop": (total_params_after >= total_params_before),
            "noop_reason": (
                "no_factorizable_layers_or_fallback_all"
                if total_params_after >= total_params_before
                else None
            ),
            "method": "slicegpt_pca_real_low_rank_factorization",
        }
    )
    self.stats = stats

    logger.info(
        "SliceGPTPCA (real): %d layers factorized, avg recon error %.4f, "
        "params %d -> %d (%.2fx compression)",
        sliced_layers_count,
        avg_err,
        total_params_before,
        total_params_after,
        compression_rate,
    )
    return self.model, stats

Compresses the model using PCA-based principal subspace projection and real low-rank factorization (SVDLinear).

Returns (compressed_model, stats).

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class SliceGPTPCALoRALinear (in_features: int, out_features: int, weight: torch.Tensor, bias: torch.Tensor | None = None, lora_r: int = 0, lora_alpha: float = 1.0, device: torch.device | None = None, dtype: torch.dtype | None = None)

Expand source code

class SliceGPTPCALoRALinear(nn.Module):
    """
    Linear layer containing sliced/PCA-projected base weights (frozen)
    with trainable LoRA adapters (lora_A, lora_B) for post-compression fine-tuning.
    Base weights are frozen with requires_grad=False; initial LoRA contribution is zero.

    NOTE: This wrapper stores the base weight densely. For real parameter
    reduction, the compressor now uses SVDLinear instead. This class is kept
    for backward compatibility only.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        weight: torch.Tensor,
        bias: torch.Tensor | None = None,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        target_device = device if device is not None else weight.device
        target_dtype = dtype if dtype is not None else weight.dtype

        self.weight = nn.Parameter(
            weight.to(device=target_device, dtype=target_dtype),
            requires_grad=False,
        )
        if bias is not None:
            self.bias = nn.Parameter(
                bias.to(device=target_device, dtype=target_dtype),
                requires_grad=False,
            )
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty(
                    (self.lora_r, in_features),
                    device=target_device,
                    dtype=torch.float32,
                )
            )
            self.lora_B = nn.Parameter(
                torch.zeros(
                    (out_features, self.lora_r),
                    device=target_device,
                    dtype=torch.float32,
                )
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        sliced_weight: torch.Tensor,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> SliceGPTPCALoRALinear:
        base_mod = getattr(linear, "base_layer", linear)
        bias = (
            base_mod.bias.detach()
            if hasattr(base_mod, "bias") and base_mod.bias is not None
            else None
        )
        in_features = getattr(base_mod, "in_features", sliced_weight.shape[1])
        out_features = getattr(base_mod, "out_features", sliced_weight.shape[0])

        has_existing_lora = hasattr(linear, "lora_A") and hasattr(linear, "lora_B")
        effective_lora_r = lora_r
        if has_existing_lora and getattr(linear, "r", 0) > 0:
            effective_lora_r = getattr(linear, "r", lora_r)

        layer = cls(
            in_features=in_features,
            out_features=out_features,
            weight=sliced_weight,
            bias=bias,
            lora_r=effective_lora_r,
            lora_alpha=lora_alpha,
            device=sliced_weight.device,
            dtype=sliced_weight.dtype,
        )

        if has_existing_lora and layer.lora_r > 0:
            try:
                if hasattr(linear.lora_A, "weight") and layer.lora_A is not None:
                    layer.lora_A.data.copy_(linear.lora_A.weight.data)
                if hasattr(linear.lora_B, "weight") and layer.lora_B is not None:
                    layer.lora_B.data.copy_(linear.lora_B.weight.data)
            except Exception as e_lora:
                logger.warning("Could not copy existing LoRA adapters in SliceGPTPCA: %s", e_lora)
        return layer

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = F.linear(
            x,
            self.weight.to(x.dtype),
            self.bias.to(x.dtype) if self.bias is not None else None,
        )
        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_delta = (
                F.linear(
                    F.linear(x.float(), self.lora_A),
                    self.lora_B,
                )
                * self.scaling
            )
            out = out + lora_delta.to(out.dtype)
        return out

Linear layer containing sliced/PCA-projected base weights (frozen) with trainable LoRA adapters (lora_A, lora_B) for post-compression fine-tuning. Base weights are frozen with requires_grad=False; initial LoRA contribution is zero.

NOTE: This wrapper stores the base weight densely. For real parameter reduction, the compressor now uses SVDLinear instead. This class is kept for backward compatibility only.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Module, sliced_weight: torch.Tensor, lora_r: int = 0, lora_alpha: float = 1.0) ‑> SliceGPTPCALoRALinear

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    out = F.linear(
        x,
        self.weight.to(x.dtype),
        self.bias.to(x.dtype) if self.bias is not None else None,
    )
    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_delta = (
            F.linear(
                F.linear(x.float(), self.lora_A),
                self.lora_B,
            )
            * self.scaling
        )
        out = out + lora_delta.to(out.dtype)
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class SparseGPTCompressor (model: nn.Module, tokenizer: Any | None = None, sparsity: float = 0.5, prunen: int = 2, prunem: int = 4, block_size: int = 128, percdamp: float = 0.01, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, target_attention_only: bool = False, device: str | torch.device | None = None, structured_24: bool = False, max_error: float = 0.5, fallback_to_original: bool = True, target_modules: list[str] | None = None, **kwargs: Any)

Expand source code

class SparseGPTCompressor(BaseCompressor):
    """
    In-memory PyTorch SparseGPT Compressor.

    Performs layer-by-layer second-order pruning with Optimal Brain Surgeon (OBS) weight
    compensation, supporting unstructured sparsity, block sparsity, and strict 2:4 structured sparsity.
    """

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any | None = None,
        sparsity: float = 0.5,
        prunen: int = 2,
        prunem: int = 4,
        block_size: int = 128,
        percdamp: float = 0.01,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        target_attention_only: bool = False,
        device: str | torch.device | None = None,
        structured_24: bool = False,
        max_error: float = 0.5,
        fallback_to_original: bool = True,
        target_modules: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="sparsegpt", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            sparsity = _cfg_get(cfg, "sparsegpt_sparsity", _cfg_get(cfg, "sparsity", sparsity))
            structured_24 = _cfg_get(
                cfg,
                "sparsegpt_structured_24",
                _cfg_get(cfg, "structured_24", structured_24),
            )
            max_error = _cfg_get(cfg, "max_error", max_error)
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)

        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.sparsity = max(0.0, min(float(sparsity), 0.95))
        self.prunen = prunen
        self.prunem = prunem
        self.block_size = block_size
        self.percdamp = percdamp
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.target_attention_only = bool(target_attention_only)
        self.structured_24 = structured_24 or (prunen == 2 and prunem == 4 and self.sparsity == 0.5)
        if self.structured_24:
            self.sparsity = 0.5
            self.prunen = 2
            self.prunem = 4
        self.target_device = (
            torch.device(device)
            if device is not None
            else next(model.parameters(), torch.tensor(0.0)).device
        )
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.target_modules = target_modules
        self.layer_errors: dict[str, float] = {}
        self.stats: dict[str, Any] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        if any(k in name_l for k in ["classifier", "score", "lm_head", "pooler", "embedding"]):
            return False

        if self.target_attention_only:
            attn_keywords = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(k in name_l for k in attn_keywords):
                return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__
        is_linear_like = (
            isinstance(base_mod, nn.Linear)
            or "conv1d" in cls_name.lower()
            or hasattr(base_mod, "nf")
        )
        if not is_linear_like:
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    def _collect_layer_activations(self) -> dict[str, torch.Tensor]:
        """Capture layer inputs via forward hooks."""
        self.model.eval()
        activations: dict[str, list[torch.Tensor]] = {}
        hooks = []

        def make_hook(name: str):
            def hook(module: nn.Module, inputs: tuple[torch.Tensor, ...], output: Any):
                if inputs and isinstance(inputs[0], torch.Tensor):
                    inp = inputs[0].detach()
                    if name not in activations:
                        activations[name] = []
                    activations[name].append(inp.cpu())

            return hook

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                hooks.append(module.register_forward_hook(make_hook(name)))

        calib_inputs = get_calibration_tensors(
            self.model,
            tokenizer=self.tokenizer,
            calibration_loader=self.calibration_loader,
            calibration_texts=self.calibration_texts,
            device=self.target_device,
        )

        with torch.no_grad():
            try:
                if isinstance(calib_inputs, dict):
                    self.model(**calib_inputs)
                elif isinstance(calib_inputs, torch.Tensor):
                    self.model(calib_inputs)
                elif isinstance(calib_inputs, (list, tuple)):
                    self.model(*calib_inputs)
            except Exception as e:
                logger.debug("SparseGPT activation collection forward pass exception: %s", e)

        for h in hooks:
            h.remove()

        aggregated: dict[str, torch.Tensor] = {}
        for name, act_list in activations.items():
            if act_list:
                aggregated[name] = torch.cat(act_list, dim=0)

        return aggregated

    def _compress_linear_module(
        self,
        name: str,
        module: nn.Module,
        activations: torch.Tensor | None,
    ) -> bool:
        """Executes second-order SparseGPT pruning on a single Linear or Conv1D layer."""
        base_mod = getattr(module, "base_layer", module)
        raw_w = base_mod.weight.data
        cls_name = base_mod.__class__.__name__
        is_conv1d = (
            "conv1d" in cls_name.lower()
            or hasattr(base_mod, "nf")
            or (
                raw_w.ndim == 2
                and hasattr(base_mod, "in_features")
                and base_mod.in_features == raw_w.shape[0]
                and raw_w.shape[0] < raw_w.shape[1]
            )
        )

        W = raw_w.t().clone().float() if is_conv1d else raw_w.clone().float()
        d_out, d_in = W.shape
        device = W.device

        # Compute empirical Hessian and inverse
        if activations is not None and activations.numel() > 0:
            act_flat = activations.reshape(-1, activations.shape[-1]).float()
            if act_flat.shape[-1] == d_in:
                _, H_inv = compute_empirical_hessian(act_flat, damping=self.percdamp, device=device)
            else:
                2.0 * torch.eye(d_in, device=device)
                H_inv = 0.5 * torch.eye(d_in, device=device)
        else:
            2.0 * torch.eye(d_in, device=device)
            H_inv = 0.5 * torch.eye(d_in, device=device)

        H_inv = H_inv.to(device=device, dtype=torch.float32)

        # Iterate over columns in blocks
        W_pruned = W.clone()
        block_cols = min(self.block_size, d_in)

        for i1 in range(0, d_in, block_cols):
            i2 = min(i1 + block_cols, d_in)
            count = i2 - i1

            W_block_init = W_pruned[:, i1:i2].clone()
            W_block = W_block_init.clone()
            H_inv_block = H_inv[i1:i2, i1:i2].clone()

            for j in range(count):
                w_col = W_block[:, j].clone()
                diag_inv = max(H_inv_block[j, j].item(), 1e-8)
                saliency_col = (w_col**2) / (2.0 * diag_inv)

                # N:M structured or unstructured pruning with OBS error compensation (B10, H25)
                if self.structured_24 or (self.prunem > 0 and self.prunen > 0):
                    M = max(1, self.prunem)
                    N = max(1, min(self.prunen, M))
                    # At the end of every M-group or block boundary, apply N:M mask across the group
                    if (j + 1) % M == 0 or (j + 1 == count):
                        start_c = (j // M) * M
                        sub_w = W_block[:, start_c : j + 1]
                        sub_h_diag = torch.diagonal(
                            H_inv_block[start_c : j + 1, start_c : j + 1]
                        ).clamp(min=1e-8)
                        sub_saliency = (sub_w**2) / (2.0 * sub_h_diag.unsqueeze(0))
                        k_keep = min(N, sub_w.shape[1])
                        _, top_idx = torch.topk(sub_saliency, k=k_keep, dim=1, largest=True)
                        sub_mask = torch.zeros_like(sub_w)
                        sub_mask.scatter_(1, top_idx, 1.0)
                        pruned_sub = sub_w * sub_mask
                        err_sub = sub_w - pruned_sub
                        W_block[:, start_c : j + 1] = pruned_sub

                        # Optimal Brain Surgeon error compensation update
                        if j < count - 1:
                            err_col = err_sub[:, -1]
                            h_slice = H_inv_block[j, j + 1 : count] / diag_inv
                            W_block[:, j + 1 : count] -= torch.outer(err_col, h_slice)
                else:
                    k_prune = int(d_out * self.sparsity)
                    if k_prune > 0:
                        _, prune_idx = torch.topk(saliency_col, k=k_prune, largest=False)
                        mask_col = torch.ones(d_out, device=device)
                        mask_col[prune_idx] = 0.0
                        pruned_w_col = w_col * mask_col
                        err = w_col - pruned_w_col
                        W_block[:, j] = pruned_w_col

                        # Optimal Brain Surgeon error compensation update
                        if j < count - 1:
                            h_slice = H_inv_block[j, j + 1 : count] / diag_inv
                            W_block[:, j + 1 : count] -= torch.outer(err, h_slice)

            W_pruned[:, i1:i2] = W_block

            # Update remaining columns outside the current block
            if i2 < d_in:
                err_block = W_block_init - W_block
                h_cross = H_inv[i1:i2, i2:]  # [count, d_in - i2]
                inv_diag_block = torch.linalg.pinv(H_inv_block)
                update = torch.matmul(err_block, torch.matmul(inv_diag_block, h_cross))
                W_pruned[:, i2:] -= update

        # Relative Frobenius error check
        norm_orig = torch.norm(W, p="fro")
        rel_err = (torch.norm(W - W_pruned, p="fro") / (norm_orig + 1e-8)).item()
        self.layer_errors[name] = rel_err

        if self.fallback_to_original and rel_err > self.max_error:
            logger.warning(
                "SparseGPT reconstruction error %.4f on %s exceeded threshold %.4f. Skipping.",
                rel_err,
                name,
                self.max_error,
            )
            logger.warning(
                "Fallback triggered for %s. Layer is kept original and NOT wrapped in LoRA. It will not be fine-tuned.",
                name,
            )
            for p in module.parameters():
                p.requires_grad = False
            return False

        final_w = W_pruned.t() if is_conv1d else W_pruned
        base_mod.weight.data = final_w.to(
            dtype=base_mod.weight.dtype, device=base_mod.weight.device
        )
        compressed_layer = base_mod
        # Preserve existing LoRA adapter weights if the source module has them
        if (
            hasattr(module, "lora_A")
            and hasattr(module, "lora_B")
            and module.lora_A is not None
            and module.lora_B is not None
        ):
            src_lA = getattr(module.lora_A, "weight", module.lora_A)
            src_lB = getattr(module.lora_B, "weight", module.lora_B)
            if hasattr(compressed_layer, "lora_A") and compressed_layer.lora_A is not None:
                if hasattr(src_lA, "data") and src_lA.shape == compressed_layer.lora_A.shape:
                    compressed_layer.lora_A.data.copy_(src_lA.data)
            if hasattr(compressed_layer, "lora_B") and compressed_layer.lora_B is not None:
                if hasattr(src_lB, "data") and src_lB.shape == compressed_layer.lora_B.shape:
                    compressed_layer.lora_B.data.copy_(src_lB.data)

        # Determine if module is wrapped (LoRALinear) or standalone
        is_wrapped = hasattr(module, "base_layer")
        if is_wrapped:
            module.base_layer = compressed_layer
            new_module = module
        else:
            new_module = compressed_layer
            replace_module(self.model, name, new_module)
        return True

    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """Executes SparseGPT compression across all target linear layers."""
        self.layer_errors = {}
        activations_map = self._collect_layer_activations()
        pruned_layers_count = 0

        for name, module in self.model.named_modules():
            if self._is_target_module(name, module):
                act = activations_map.get(name, None)
                if self._compress_linear_module(name, module, act):
                    pruned_layers_count += 1

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        stats.update(
            {
                "pruned_layers": pruned_layers_count,
                "sparsity_target": self.sparsity,
                "structured_24": self.structured_24,
                "layer_errors": self.layer_errors,
                "target_attention_only": self.target_attention_only,
                "fallback_to_original": self.fallback_to_original,
                "max_error": self.max_error,
                "method": "sparsegpt_second_order",
            }
        )
        self.stats = stats
        return self.model, stats

In-memory PyTorch SparseGPT Compressor.

Performs layer-by-layer second-order pruning with Optimal Brain Surgeon (OBS) weight compensation, supporting unstructured sparsity, block sparsity, and strict 2:4 structured sparsity.

Ancestors

BaseCompressor

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """Executes SparseGPT compression across all target linear layers."""
    self.layer_errors = {}
    activations_map = self._collect_layer_activations()
    pruned_layers_count = 0

    for name, module in self.model.named_modules():
        if self._is_target_module(name, module):
            act = activations_map.get(name, None)
            if self._compress_linear_module(name, module, act):
                pruned_layers_count += 1

    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    stats.update(
        {
            "pruned_layers": pruned_layers_count,
            "sparsity_target": self.sparsity,
            "structured_24": self.structured_24,
            "layer_errors": self.layer_errors,
            "target_attention_only": self.target_attention_only,
            "fallback_to_original": self.fallback_to_original,
            "max_error": self.max_error,
            "method": "sparsegpt_second_order",
        }
    )
    self.stats = stats
    return self.model, stats

Executes SparseGPT compression across all target linear layers.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class SpectralResidualLoRA (*args: Any, **kwargs: Any)

Expand source code

class SpectralResidualLoRA(ResSVDCompressor):
    """PEFT-focused wrapper for Spectral Residual LoRA."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.name = "spectral_residual_lora"

PEFT-focused wrapper for Spectral Residual LoRA.

Ancestors

Subclasses

neuralops.compression_engine.res_svd.SpectralResidualLoRARegistered

Inherited members

ResSVDCompressor:
- compress
- get_compression_stats
- summary

class SpectralResidualLoRARegistered (*args: Any, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("spectral_residual_lora")
class SpectralResidualLoRARegistered(SpectralResidualLoRA):
    pass

PEFT-focused wrapper for Spectral Residual LoRA.

Ancestors

Inherited members

SpectralResidualLoRA:
- compress
- get_compression_stats
- summary

class SubspaceLinear (in_features: int, out_features: int, rank: int = 64, bias: bool = True, lora_r: int = 0, lora_alpha: float = 1.0, dtype: torch.dtype | None = None, device: torch.device | None = None)

Expand source code

class SubspaceLinear(nn.Module):
    """
    Subspace Projection Factorized Linear Layer with optional LoRA adapter.

    W ≈ B @ A, where:
      A (rank, in_features):  orthonormal / projection basis
      B (out_features, rank): coefficient expansion matrix

    Forward:
      y = ((x @ A^T) @ B^T) + bias + (lora_B @ lora_A @ x) * scaling
    """

    # SP-FIX-01
    is_custom_attention: bool = False

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 64,
        bias: bool = True,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.nf = self.out_features  # Conv1D compatibility
        self.nx = self.in_features  # Conv1D compatibility
        self.rank = max(1, min(rank, in_features, out_features))
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)
        self.scaling = (self.lora_alpha / self.lora_r) if self.lora_r > 0 else 1.0

        self.A = nn.Parameter(torch.empty(self.rank, in_features, dtype=dtype, device=device))
        self.B = nn.Parameter(torch.empty(out_features, self.rank, dtype=dtype, device=device))

        # Orthogonal init for A; xavier for B. Do it in fp32 for stability.
        with torch.no_grad():
            if dtype in (torch.float16, torch.bfloat16):
                A_fp32 = torch.empty(self.rank, in_features, dtype=torch.float32)
                nn.init.orthogonal_(A_fp32)
                self.A.data.copy_(A_fp32.to(dtype=dtype, device=device))
            else:
                nn.init.orthogonal_(self.A)

            B_fp32 = self.B.data.float()
            nn.init.xavier_uniform_(B_fp32)
            self.B.data.copy_(B_fp32.to(dtype=dtype, device=device))

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

        if self.lora_r > 0:
            self.lora_A = nn.Parameter(
                torch.empty((self.lora_r, in_features), dtype=torch.float32, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.zeros((out_features, self.lora_r), dtype=torch.float32, device=device)
            )
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

    # ------------------------------------------------------------------
    @property
    def weight(self) -> torch.Tensor:
        return torch.matmul(self.B.float(), self.A.float()).to(self.A.dtype)

    def get_weight(self) -> torch.Tensor:
        return self.weight

    # ------------------------------------------------------------------
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)

        # SP-FIX-03: dtype-safe
        A = self.A if self.A.dtype == x_flat.dtype else self.A.to(x_flat.dtype)
        B = self.B if self.B.dtype == x_flat.dtype else self.B.to(x_flat.dtype)
        b = None
        if self.bias is not None:
            b = self.bias if self.bias.dtype == x_flat.dtype else self.bias.to(x_flat.dtype)

        x_proj = F.linear(x_flat, A)  # (K, rank)
        out = F.linear(x_proj, B, b)  # (K, out)

        if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x_flat.float(), self.lora_A), self.lora_B) * self.scaling
            out = out + lora_out.to(out.dtype)

        return out.view(*orig_shape[:-1], self.out_features)

    # ------------------------------------------------------------------
    @classmethod
    def from_linear(
        cls,
        linear: nn.Module,
        rank: int = 64,
        method: str = "svd",
        use_lora: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
    ) -> SubspaceLinear:
        """
        Project linear or Conv1D weight matrix onto top-rank subspace.
        """
        base_mod = getattr(linear, "base_layer", linear)
        raw_w = base_mod.weight.data
        cls_name = base_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(base_mod, "nf")

        if is_conv1d:
            W = raw_w.t().contiguous()
            out_f, in_f = W.shape
        else:
            W = raw_w
            out_f, in_f = W.shape

        bias = base_mod.bias.data if getattr(base_mod, "bias", None) is not None else None

        # ---------------- Rank budget (SP-FIX-07) ----------------
        max_r = min(out_f, in_f)
        param_cap = max(1, (in_f * out_f) // max(1, in_f + out_f))
        r = max(1, min(int(rank), max_r, param_cap))

        # ---------------- Existing LoRA detection ----------------
        has_existing_lora = (
            hasattr(linear, "lora_A")
            and hasattr(linear, "lora_B")
            and _get_lora_param(linear.lora_A) is not None
            and _get_lora_param(linear.lora_B) is not None
        )
        existing_r = getattr(linear, "r", getattr(linear, "lora_r", 0))
        effective_lora_r = (
            int(existing_r) if (has_existing_lora and existing_r > 0) else int(lora_r)
        )
        effective_lora_alpha = float(
            getattr(linear, "lora_alpha", getattr(linear, "alpha", lora_alpha))
        )

        subspace_layer = cls(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=(bias is not None),
            lora_r=effective_lora_r if (use_lora or has_existing_lora) else 0,
            lora_alpha=effective_lora_alpha,
            dtype=W.dtype,
            device=W.device,
        )

        # ---------------- SVD factorization ----------------
        U, S, Vt = torch.linalg.svd(W.float(), full_matrices=False)
        A_data = Vt[:r, :]  # (r, in_f)
        B_data = U[:, :r] * S[:r].unsqueeze(0)  # (out_f, r)

        with torch.no_grad():
            subspace_layer.A.data.copy_(A_data.to(dtype=W.dtype, device=W.device))
            subspace_layer.B.data.copy_(B_data.to(dtype=W.dtype, device=W.device))

            # SP-FIX-11: bias under no_grad
            if bias is not None and subspace_layer.bias is not None:
                subspace_layer.bias.data.copy_(bias.to(dtype=W.dtype, device=W.device))

        # ---------------- LoRA transfer (SP-FIX-04) ----------------
        if (
            has_existing_lora
            and subspace_layer.lora_A is not None
            and subspace_layer.lora_B is not None
        ):
            try:
                src_A = _get_lora_param(linear.lora_A)
                src_B = _get_lora_param(linear.lora_B)
                if src_A is not None and src_B is not None:
                    if src_A.shape == subspace_layer.lora_A.shape:
                        subspace_layer.lora_A.data.copy_(
                            src_A.to(
                                subspace_layer.lora_A.device,
                                dtype=subspace_layer.lora_A.dtype,
                            )
                        )
                    if src_B.shape == subspace_layer.lora_B.shape:
                        subspace_layer.lora_B.data.copy_(
                            src_B.to(
                                subspace_layer.lora_B.device,
                                dtype=subspace_layer.lora_B.dtype,
                            )
                        )
            except Exception as e:
                logger.debug("SubspaceLinear LoRA transfer failed: %s", e)

        return subspace_layer

Subspace Projection Factorized Linear Layer with optional LoRA adapter.

W ≈ B @ A, where: A (rank, in_features): orthonormal / projection basis B (out_features, rank): coefficient expansion matrix

Forward -----= y = ((x @ A^T) @ B^T) + bias + (lora_B @ lora_A @ x) * scaling

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Static methods

def from_linear(linear: nn.Module, rank: int = 64, method: str = 'svd', use_lora: bool = False, lora_r: int = 0, lora_alpha: float = 1.0) ‑> SubspaceLinear: Project linear or Conv1D weight matrix onto top-rank subspace.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    return torch.matmul(self.B.float(), self.A.float()).to(self.A.dtype)

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)

    # SP-FIX-03: dtype-safe
    A = self.A if self.A.dtype == x_flat.dtype else self.A.to(x_flat.dtype)
    B = self.B if self.B.dtype == x_flat.dtype else self.B.to(x_flat.dtype)
    b = None
    if self.bias is not None:
        b = self.bias if self.bias.dtype == x_flat.dtype else self.bias.to(x_flat.dtype)

    x_proj = F.linear(x_flat, A)  # (K, rank)
    out = F.linear(x_proj, B, b)  # (K, out)

    if self.lora_r > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x_flat.float(), self.lora_A), self.lora_B) * self.scaling
        out = out + lora_out.to(out.dtype)

    return out.view(*orig_shape[:-1], self.out_features)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    return self.weight

class SubspaceProjectionCompressor (model: nn.Module | None = None, rank: int = 64, method: str = 'svd', cfg: Any | None = None, target_layers: list[str] | None = None, target_attention_only: bool = False, use_lora: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, max_error: float = 0.5, fallback_to_original: bool = True, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("subspace_projection")
@COMPRESSOR_REGISTRY.register("subspace_projection_compressor")
class SubspaceProjectionCompressor(BaseCompressor):
    """
    Subspace Projection Compression Engine.
    """

    # SP-FIX-01
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module | None = None,
        rank: int = 64,
        method: str = "svd",
        cfg: Any | None = None,
        target_layers: list[str] | None = None,
        target_attention_only: bool = False,
        use_lora: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init__(
            model if model is not None else nn.Identity(),
            name="subspace_projection",
            **kwargs,
        )
        cfg_obj = cfg or kwargs.get("comp_config")

        self.rank = int(rank or _cfg_get(cfg_obj, "subspace_rank", _cfg_get(cfg_obj, "rank", 64)))
        self.method = str(method or _cfg_get(cfg_obj, "subspace_method", "svd"))
        if self.method != "svd":
            logger.warning(
                "SubspaceProjection: method='%s' is not implemented; using 'svd'.",
                self.method,
            )
            self.method = "svd"

        self.cfg = cfg_obj
        tl = target_layers or _cfg_get(cfg_obj, "target_layers", None)
        if tl is not None and len(tl) == 0:
            tl = None
        self.target_layers = tl

        self.target_attention_only = bool(
            target_attention_only or _cfg_get(cfg_obj, "target_attention_only", False)
        )
        self.use_lora = bool(use_lora or _cfg_get(cfg_obj, "use_lora", False) or (lora_r > 0))
        self.lora_r = int(lora_r or _cfg_get(cfg_obj, "lora_r", 0))
        self.lora_alpha = float(lora_alpha or _cfg_get(cfg_obj, "lora_alpha", 1.0))
        self.max_error = float(_cfg_get(cfg_obj, "max_error", max_error))
        self.fallback_to_original = bool(
            _cfg_get(cfg_obj, "fallback_to_original", fallback_to_original)
        )

        self.compressed_layers: list[dict[str, Any]] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.stats: dict[str, Any] = {}

    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        if any(
            k in name_l
            for k in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "layernorm",
                "rmsnorm",
                "batchnorm",
            )
        ):
            return False

        if isinstance(module, SubspaceLinear):
            return False

        if self.target_attention_only:
            attn_keywords = (
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            )
            if not any(k in name_l for k in attn_keywords):
                return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__.lower()
        is_linear_like = (
            isinstance(base_mod, nn.Linear) or "conv1d" in cls_name or hasattr(base_mod, "nf")
        )
        if not is_linear_like:
            return False

        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self.target_layers is None:
            return True
        return any(t in name for t in self.target_layers)

    # ------------------------------------------------------------------
    def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
        from neuralops.core.model_utils import replace_module

        if model is not None:
            self.model = model
            try:
                self._original_params = self._count_params(self.model)
            except Exception:
                self._original_params = sum(p.numel() for p in self.model.parameters())

        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}

        total_orig_bytes = 0
        total_comp_bytes = 0
        seen_weights: set[int] = set()

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            base_mod = getattr(module, "base_layer", module)
            w = getattr(base_mod, "weight", None)
            if w is None:
                continue
            w_id = id(w)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            raw_w = w.data
            cls_name = base_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(base_mod, "nf")

            W_orig = raw_w.t().contiguous().float() if is_conv1d else raw_w.float()
            max_rank = min(W_orig.shape[0], W_orig.shape[1])
            r = min(self.rank, max_rank)

            # storage before
            orig_bytes = raw_w.numel() * raw_w.element_size()
            if getattr(base_mod, "bias", None) is not None:
                orig_bytes += base_mod.bias.numel() * base_mod.bias.element_size()
            total_orig_bytes += orig_bytes

            try:
                new_layer = SubspaceLinear.from_linear(
                    module,
                    rank=r,
                    method=self.method,
                    use_lora=self.use_lora,
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                )
            except Exception as e:
                logger.warning("Subspace projection failed on %s: %s", name, e)
                self.fallback_layers.append(name)
                self.layer_errors[name] = 1.0
                total_comp_bytes += orig_bytes
                for p in module.parameters():
                    p.requires_grad = False
                continue

            # error on forward path
            W_rec = new_layer.weight.float()
            if is_conv1d:
                W_orig_cmp = W_orig
                W_rec_cmp = W_rec.t().contiguous() if W_rec.shape != W_orig.shape else W_rec
            else:
                W_orig_cmp = W_orig
                W_rec_cmp = W_rec

            denom = torch.norm(W_orig_cmp, p="fro").clamp_min(1e-12)
            err = float((torch.norm(W_orig_cmp - W_rec_cmp, p="fro") / denom).item())
            self.layer_errors[name] = err

            if self.fallback_to_original and err > self.max_error:
                logger.warning(
                    "SubspaceProjection: err=%.4f on %s > max_error=%.4f; keeping original.",
                    err,
                    name,
                    self.max_error,
                )
                self.fallback_layers.append(name)
                total_comp_bytes += orig_bytes
                for p in module.parameters():
                    p.requires_grad = False
                continue

            # SP-FIX-10: preserve LoRA wrapper
            try:
                if hasattr(module, "base_layer"):
                    module.base_layer = new_layer
                    for pn, p in module.named_parameters():
                        if "lora_" in pn.lower() and not pn.startswith("base_layer"):
                            p.requires_grad = True
                else:
                    replace_module(self.model, name, new_layer)
            except Exception as e:
                logger.warning("SubspaceProjection replacement failed on %s: %s", name, e)
                self.fallback_layers.append(name)
                total_comp_bytes += orig_bytes
                continue

            self.compressed_layers.append(
                {
                    "name": name,
                    "rank": new_layer.rank,
                    "relative_error": float(err),
                }
            )

            # storage after
            comp_bytes = (
                new_layer.A.numel() * new_layer.A.element_size()
                + new_layer.B.numel() * new_layer.B.element_size()
            )
            if new_layer.bias is not None:
                comp_bytes += new_layer.bias.numel() * new_layer.bias.element_size()
            if new_layer.lora_A is not None:
                comp_bytes += new_layer.lora_A.numel() * new_layer.lora_A.element_size()
            if new_layer.lora_B is not None:
                comp_bytes += new_layer.lora_B.numel() * new_layer.lora_B.element_size()
            total_comp_bytes += comp_bytes

        # ------------------------------------------------------------------
        # SP-FIX-08: Forward sanity check
        # ------------------------------------------------------------------
        forward_check_success = True
        try:
            with torch.no_grad():
                dev = next(self.model.parameters(), torch.tensor(0.0)).device
                dummy = torch.randint(0, 100, (1, 8), device=dev, dtype=torch.long)
                try:
                    self.model(input_ids=dummy)
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception:
                        cfg = getattr(self.model, "config", None)
                        hidden = getattr(cfg, "hidden_size", None) if cfg else None
                        if hidden is None:
                            first_lin = next(
                                (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                                None,
                            )
                            hidden = getattr(first_lin, "in_features", 8)
                        dt = next(self.model.parameters()).dtype
                        self.model(torch.randn(1, 8, hidden, device=dev, dtype=dt))
        except Exception as e:
            logger.warning("SubspaceProjection forward sanity check failed: %s", e)
            forward_check_success = False

        # ------------------------------------------------------------------
        # Stats (SP-FIX-09)
        # ------------------------------------------------------------------
        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        avg_err = (
            sum(self.layer_errors.values()) / max(1, len(self.layer_errors))
            if self.layer_errors
            else 0.0
        )
        storage_ratio = total_orig_bytes / max(1, total_comp_bytes) if total_comp_bytes > 0 else 1.0
        effective_ratio = (
            total_comp_bytes / max(1, total_orig_bytes) if total_orig_bytes > 0 else 1.0
        )

        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        stats.update(
            {
                "compressed_layer_count": len(self.compressed_layers),
                "num_compressed": len(self.compressed_layers),
                "fallback_layers": self.fallback_layers,
                "fallback_layer_count": len(self.fallback_layers),
                "skipped_layers": self.skipped_layers,
                "skipped_layer_count": len(self.skipped_layers),
                "compressed_layers": self.compressed_layers,
                "layer_errors": self.layer_errors,
                "avg_reconstruction_error": avg_err,
                "subspace_rank": self.rank,
                "subspace_method": self.method,
                "target_attention_only": self.target_attention_only,
                "fallback_to_original": self.fallback_to_original,
                "max_error": self.max_error,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                "orig_storage_bytes": total_orig_bytes,
                "comp_storage_bytes": total_comp_bytes,
                "storage_compression_ratio": round(storage_ratio, 3),
                "effective_parameter_ratio": round(effective_ratio, 4),
                "storage_reduction_pct": round(100.0 * (1.0 - effective_ratio), 2),
                "forward_check_success": forward_check_success,
                "method": "subspace_projection",
            }
        )
        self.stats = stats
        return self.model, stats

    # ------------------------------------------------------------------
    def compress_linear(self, linear: nn.Module) -> SubspaceLinear:
        return SubspaceLinear.from_linear(
            linear,
            rank=self.rank,
            method=self.method,
            use_lora=self.use_lora,
            lora_r=self.lora_r,
            lora_alpha=self.lora_alpha,
        )

Subspace Projection Compression Engine.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def compress_linear(self, linear: nn.Module) ‑> SubspaceLinear

Expand source code

def compress_linear(self, linear: nn.Module) -> SubspaceLinear:
    return SubspaceLinear.from_linear(
        linear,
        rank=self.rank,
        method=self.method,
        use_lora=self.use_lora,
        lora_r=self.lora_r,
        lora_alpha=self.lora_alpha,
    )

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class SwiftSVDActivationCompressor (model: nn.Module, tokenizer: Any | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, rank: int = 64, max_error: float = 0.5, fallback_to_original: bool = True, target_attention_only: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, skip_attention_qkv: bool = False, min_rank: int = 4, target_modules: list[str] | None = None, num_calibration_batches: int = 8, max_calibration_tokens: int = 32768, use_fp64_cross_cov: bool = True, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("swift_svd_act")
@COMPRESSOR_REGISTRY.register("swift_svd")
@COMPRESSOR_REGISTRY.register("swift_svd_compressor")
class SwiftSVDActivationCompressor(BaseCompressor):
    """
    Swift-SVD Activation-Aware Compressor.

    Captures empirical input-output cross-covariances across calibration data,
    derives activation-aligned subspaces, projects weights onto the principal
    activation subspace, and replaces dense nn.Linear / Conv1D layers with
    factorized SVDLinear modules.
    """

    # SW-FIX-09
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any | None = None,
        calibration_loader: Any | None = None,
        calibration_texts: list[str] | None = None,
        rank: int = 64,
        max_error: float = 0.50,
        fallback_to_original: bool = True,
        target_attention_only: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        skip_attention_qkv: bool = False,
        min_rank: int = 4,
        target_modules: list[str] | None = None,
        num_calibration_batches: int = 8,
        max_calibration_tokens: int = 32768,
        use_fp64_cross_cov: bool = True,  # SW-FIX-16
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="swift_svd_act", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            rank = _cfg_get(cfg, "swift_svd_act_rank", _cfg_get(cfg, "rank", rank))
            tokenizer = _cfg_get(cfg, "tokenizer", tokenizer)
            calibration_loader = _cfg_get(cfg, "calibration_loader", calibration_loader)
            calibration_texts = _cfg_get(cfg, "calibration_texts", calibration_texts)
            max_error = _cfg_get(cfg, "max_error", max_error)
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)
            lora_r = _cfg_get(cfg, "lora_r", lora_r)
            lora_alpha = _cfg_get(cfg, "lora_alpha", lora_alpha)
            skip_attention_qkv = _cfg_get(cfg, "skip_attention_qkv", skip_attention_qkv)
            min_rank = _cfg_get(cfg, "min_rank", min_rank)
            num_calibration_batches = _cfg_get(
                cfg, "num_calibration_batches", num_calibration_batches
            )
            max_calibration_tokens = _cfg_get(cfg, "max_calibration_tokens", max_calibration_tokens)
            use_fp64_cross_cov = _cfg_get(cfg, "use_fp64_cross_cov", use_fp64_cross_cov)

        if "rank" in kwargs and kwargs["rank"] is not None:
            rank = kwargs["rank"]

        self.tokenizer = tokenizer
        self.calibration_loader = calibration_loader
        self.calibration_texts = calibration_texts
        self.rank = int(rank)
        self.max_error = float(max_error)
        self.fallback_to_original = bool(fallback_to_original)
        self.target_attention_only = bool(target_attention_only)
        self.lora_r = max(0, int(lora_r))
        self.lora_alpha = float(lora_alpha)
        self.skip_attention_qkv = bool(skip_attention_qkv)
        self.min_rank = max(1, int(min_rank))
        self.use_fp64_cross_cov = bool(use_fp64_cross_cov)

        # SW-FIX-12: empty list → None
        if target_modules is not None and len(target_modules) == 0:
            target_modules = None
        self.target_modules = target_modules

        self.num_calibration_batches = max(1, int(num_calibration_batches))
        self.max_calibration_tokens = max(1, int(max_calibration_tokens))

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}
        self.replaced_layers: list[dict[str, Any]] = []
        self.stats: dict[str, Any] = {}
        self._compressed_params: int = 0
        self._orig_storage_bytes: int = 0
        self._comp_storage_bytes: int = 0

    # ------------------------------------------------------------------
    # Target selection
    # ------------------------------------------------------------------
    def _is_gqa_asymmetric(self, name: str, module: nn.Module) -> bool:
        # SW-FIX-11: always active, not gated by skip_attention_qkv.
        cfg = getattr(self.model, "config", None)
        if cfg is None:
            return False
        num_heads = getattr(cfg, "num_attention_heads", None)
        num_kv_heads = getattr(cfg, "num_key_value_heads", None)
        if num_heads is None or num_kv_heads is None:
            return False
        if num_kv_heads >= num_heads:
            return False
        name_lower = name.lower()
        return any(
            p in name_lower
            for p in ("q_proj", "k_proj", "v_proj", "query", "key", "value", "c_attn")
        )

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        if any(
            k in name_l
            for k in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embed",
                "logits_processor",
            )
        ):
            return False

        if isinstance(
            module,
            (nn.LayerNorm, nn.RMSNorm, nn.Embedding, nn.BatchNorm1d, nn.BatchNorm2d),
        ):
            return False

        if self.target_attention_only:
            attn_keywords = (
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            )
            if not any(k in name_l for k in attn_keywords):
                return False

        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__.lower()
        is_linear_like = (
            isinstance(base_mod, nn.Linear) or "conv1d" in cls_name or hasattr(base_mod, "nf")
        )
        if not is_linear_like:
            return False
        if isinstance(base_mod, SVDLinear):
            return False

        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self._is_gqa_asymmetric(name, module):
            return False

        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    # ------------------------------------------------------------------
    # Calibration batch normalization (SW-FIX-02)
    # ------------------------------------------------------------------
    @staticmethod
    def _normalize_calibration_batches(calib_inputs: Any) -> list[Any]:
        """
        Convert whatever `get_calibration_tensors` returns into a list of batches.

        SW-FIX-02: A single tuple like `(input_ids, attention_mask)` is treated
        as ONE batch, not as two batches. Detection rule: a batch is a tuple/
        list ONLY IF its elements are themselves tensors/dicts AND the tuple
        length is small (≤ 3). Otherwise, the top-level container is a list of
        batches.
        """
        if calib_inputs is None:
            return []

        # Single tensor / dict / small tuple → single batch
        if isinstance(calib_inputs, torch.Tensor):
            return [calib_inputs]
        if isinstance(calib_inputs, dict):
            return [calib_inputs]

        if isinstance(calib_inputs, (list, tuple)):
            if len(calib_inputs) == 0:
                return []

            # Case A: top-level list/tuple where every element is a batch
            # (dict, tensor, or tuple-of-tensors)
            all_batches = all(
                isinstance(x, (dict, torch.Tensor, list, tuple)) for x in calib_inputs
            )
            if all_batches:
                # Distinguish "list of batches" from "single batch as tuple".
                # Heuristic: single batch if the outer tuple contains only
                # tensors (e.g., (input_ids, attention_mask)).
                if (
                    all(isinstance(x, torch.Tensor) for x in calib_inputs)
                    and len(calib_inputs) <= 3
                ):
                    return [calib_inputs]
                # Otherwise, it's a list of batches.
                return list(calib_inputs)

        return [calib_inputs]

    # ------------------------------------------------------------------
    # Activation collection (SW-FIX-04, SW-FIX-07, SW-FIX-16)
    # ------------------------------------------------------------------
    def _collect_activations(self) -> dict[str, tuple[torch.Tensor, int]]:
        """
        Collect activations and accumulate cross-covariance C = E[Y^T X] per layer.

        Returns {name: (cross_cov_fp32, n_tokens)}.
        cross_cov shape: [out_features, in_features] in fp32 on CPU.
        """
        accum: dict[str, dict[str, Any]] = {}
        hooks = []
        compute_dtype = torch.float64 if self.use_fp64_cross_cov else torch.float32

        def make_hook(layer_name: str):
            def hook(module: nn.Module, inp: tuple[Any, ...], outp: Any) -> None:
                if not inp or inp[0] is None:
                    return
                x_in = inp[0].detach()
                if isinstance(outp, torch.Tensor):
                    x_out = outp.detach()
                elif (
                    isinstance(outp, (tuple, list))
                    and len(outp) > 0
                    and isinstance(outp[0], torch.Tensor)
                ):
                    x_out = outp[0].detach()
                else:
                    return

                in_f = x_in.shape[-1]
                out_f = x_out.shape[-1]

                x_in_2d = x_in.reshape(-1, in_f).float()
                x_out_2d = x_out.reshape(-1, out_f).float()

                n = x_in_2d.shape[0]
                if n > self.max_calibration_tokens:
                    x_in_2d = x_in_2d[: self.max_calibration_tokens]
                    x_out_2d = x_out_2d[: self.max_calibration_tokens]
                    n = self.max_calibration_tokens

                if n == 0:
                    return

                if layer_name not in accum:
                    accum[layer_name] = {
                        "cov": torch.zeros(out_f, in_f, dtype=compute_dtype),
                        "n": 0,
                        "in_f": in_f,
                        "out_f": out_f,
                    }
                elif accum[layer_name]["in_f"] != in_f or accum[layer_name]["out_f"] != out_f:
                    # Shape changed across batches → skip
                    return

                with torch.no_grad():
                    cov_local = (x_out_2d.to(compute_dtype).t() @ x_in_2d.to(compute_dtype)).cpu()
                accum[layer_name]["cov"] += cov_local
                accum[layer_name]["n"] += n

            return hook

        target_layers = [
            (name, mod)
            for name, mod in self.model.named_modules()
            if self._is_target_module(name, mod)
        ]
        if not target_layers:
            return {}

        for name, mod in target_layers:
            hooks.append(mod.register_forward_hook(make_hook(name)))

        device = next(self.model.parameters(), torch.tensor(0.0)).device
        self.model.eval()

        try:
            calib_inputs = get_calibration_tensors(
                self.model,
                tokenizer=self.tokenizer,
                calibration_loader=self.calibration_loader,
                calibration_texts=self.calibration_texts,
                device=device,
            )
            batches = self._normalize_calibration_batches(calib_inputs)
            n_batches = min(len(batches), self.num_calibration_batches)

            if n_batches == 0:
                logger.warning("SwiftSVDActivation: no calibration batches available.")
                return {}

            with torch.no_grad():
                for i in range(n_batches):
                    batch = batches[i]
                    try:
                        if isinstance(batch, dict):
                            b_dev = {
                                k: v.to(device) if isinstance(v, torch.Tensor) else v
                                for k, v in batch.items()
                            }
                            self.model(**b_dev)
                        elif isinstance(batch, torch.Tensor):
                            self.model(batch.to(device))
                        elif isinstance(batch, (list, tuple)):
                            b_dev = [
                                v.to(device) if isinstance(v, torch.Tensor) else v for v in batch
                            ]
                            self.model(*b_dev)
                    except Exception as e:
                        logger.debug("Calibration batch %d forward failed: %s", i, e)
        finally:
            for h in hooks:
                h.remove()

        aggregated: dict[str, tuple[torch.Tensor, int]] = {}
        for name, data in accum.items():
            n = max(1, data["n"])
            cov = (data["cov"] / n).float()
            aggregated[name] = (cov, int(n))
        return aggregated

    # ------------------------------------------------------------------
    # LoRA rank detection (SW-FIX-03)
    # ------------------------------------------------------------------
    @staticmethod
    def _detect_lora_rank(module: nn.Module) -> int:
        for attr in ("r", "lora_r", "rank"):
            val = getattr(module, attr, None)
            if isinstance(val, int) and val > 0:
                return val
        la = getattr(module, "lora_A", None)
        if la is not None:
            if hasattr(la, "out_features"):
                try:
                    return int(la.out_features)
                except Exception:
                    pass
            if hasattr(la, "weight") and la.weight is not None and la.weight.dim() >= 1:
                return int(la.weight.shape[0])
            if isinstance(la, torch.Tensor) and la.dim() >= 1:
                return int(la.shape[0])
        return 0

    # ------------------------------------------------------------------
    # Rank selection (SW-FIX-01, SW-FIX-13)
    # ------------------------------------------------------------------
    def _select_rank(self, in_f: int, out_f: int) -> int:
        """
        SW-FIX-13: clear, monotonic rank selection.

        1. max_r = min(in_f, out_f)
        2. Dense break-even rank (where param count of SVD == param count of dense):
             2r(in + out) < in · out  →  r < in·out / (2(in+out))
           We use `param_cap = (in·out) // (in + out) - 1` (a slightly stricter bound
           often used in the codebase) and clamp to max_r.
        3. requested_r = clamp(self.rank, self.min_rank, param_cap)
        """
        max_r = min(in_f, out_f)
        if in_f * out_f > 16:
            param_cap = max(1, (in_f * out_f) // (in_f + out_f) - 1)
        else:
            param_cap = max_r
        param_cap = max(1, min(param_cap, max_r))

        r_lower = min(self.min_rank, max_r, param_cap)
        r_upper = min(self.rank, max_r, param_cap)

        # SW-FIX-13: monotone selection — start from r_lower, aim for r_upper.
        r = max(r_lower, r_upper)
        r = max(1, min(r, param_cap))
        return r

    # ------------------------------------------------------------------
    # Per-layer compression
    # ------------------------------------------------------------------
    def compress_linear(
        self,
        layer: nn.Module,
        act_pair: tuple[torch.Tensor, int] | None = None,
    ) -> SVDLinear | None:
        base_layer = getattr(layer, "base_layer", layer)
        raw_w = base_layer.weight.data
        cls_name = base_layer.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(base_layer, "nf")

        W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
        out_f, in_f = W.shape

        r = self._select_rank(in_f, out_f)
        device = raw_w.device
        dtype = raw_w.dtype
        A: torch.Tensor | None = None
        B: torch.Tensor | None = None

        # ---- Activation subspace projection ----
        if act_pair is not None:
            try:
                cov, _n = act_pair
                if cov.shape[0] == out_f and cov.shape[1] == in_f:
                    C = cov.to(device=device, dtype=torch.float32)
                    U_c, _, Vt_c = torch.linalg.svd(C, full_matrices=False)
                    r_act = min(r, U_c.shape[1], Vt_c.shape[0])
                    if r_act > 0:
                        U_act = U_c[:, :r_act]
                        V_act = Vt_c[:r_act, :].t().contiguous()

                        # M = U_act^T W V_act
                        M = U_act.t().contiguous() @ W @ V_act
                        U_m, S_m, Vt_m = torch.linalg.svd(M, full_matrices=False)
                        r_act = min(r_act, len(S_m))
                        if r_act > 0:
                            sqrt_S_m = torch.sqrt(torch.clamp(S_m[:r_act], min=1e-12))
                            A = (U_act @ U_m[:, :r_act]) * sqrt_S_m.unsqueeze(0)
                            B = sqrt_S_m.unsqueeze(1) @ (Vt_m[:r_act, :] @ V_act.t().contiguous())
                            r = r_act
            except Exception as e:
                logger.debug(
                    "Activation subspace projection failed (%s). Falling back to direct SVD.",
                    e,
                )
                A, B = None, None

        # ---- Direct SVD fallback ----
        if A is None or B is None:
            try:
                U, S, Vt = torch.linalg.svd(W, full_matrices=False)
                r = min(r, len(S))
                sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
                A = U[:, :r] * sqrt_S.unsqueeze(0)
                B = sqrt_S.unsqueeze(1) @ Vt[:r, :]
            except Exception as e:
                logger.warning("SwiftSVDActivation SVD failure: %s", e)
                return None

        # ---- Error check ----
        W_rec = A @ B
        norm_orig = torch.norm(W, p="fro").item()
        rel_err = (torch.norm(W - W_rec) / max(norm_orig, 1e-12)).item()

        if not math.isfinite(rel_err) or rel_err > self.max_error:
            # Try direct SVD (may have better conditioning)
            try:
                U, S, Vt = torch.linalg.svd(W, full_matrices=False)
                r_d = min(r, len(S))
                sqrt_S = torch.sqrt(torch.clamp(S[:r_d], min=1e-12))
                A_d = U[:, :r_d] * sqrt_S.unsqueeze(0)
                B_d = sqrt_S.unsqueeze(1) @ Vt[:r_d, :]
                W_rec_d = A_d @ B_d
                err_d = (torch.norm(W - W_rec_d) / max(norm_orig, 1e-12)).item()
                if err_d < rel_err:
                    A, B, r, rel_err = A_d, B_d, r_d, err_d
            except Exception as e:
                logger.debug("Direct SVD fallback failed: %s", e)

        # SW-FIX-14: preserve original on high error
        if not math.isfinite(rel_err) or rel_err > self.max_error:
            logger.warning(
                "SwiftSVDActivation reconstruction error %.4f exceeded threshold %.4f on %s.",
                rel_err,
                self.max_error,
                getattr(base_layer, "_name", "?"),
            )
            return None

        # ---- Build SVDLinear ----
        effective_lora_r = self.lora_r
        detected = self._detect_lora_rank(layer)
        if detected > 0:
            effective_lora_r = detected

        svd_layer = SVDLinear(
            in_features=in_f,
            out_features=out_f,
            rank=r,
            bias=base_layer.bias is not None,
            dtype=dtype,
            device=device,
            lora_r=effective_lora_r,
            lora_alpha=self.lora_alpha,
        )
        svd_layer.set_low_rank_factors(
            U_factor=A.to(dtype=dtype, device=device),
            V_factor=B.to(dtype=dtype, device=device),
        )
        if base_layer.bias is not None and svd_layer.bias is not None:
            svd_layer.bias.data.copy_(base_layer.bias.data)

        # ---- LoRA transfer ----
        if effective_lora_r > 0 and getattr(svd_layer, "use_lora", False):
            try:
                src_lA = getattr(layer, "lora_A", None)
                src_lB = getattr(layer, "lora_B", None)
                if src_lA is not None and src_lB is not None:
                    src_lA_w = getattr(src_lA, "weight", src_lA)
                    src_lB_w = getattr(src_lB, "weight", src_lB)
                    dst_lA_w = getattr(svd_layer.lora_A, "weight", svd_layer.lora_A)
                    dst_lB_w = getattr(svd_layer.lora_B, "weight", svd_layer.lora_B)
                    if src_lA_w.shape == dst_lA_w.shape:
                        dst_lA_w.data.copy_(src_lA_w.data)
                    if src_lB_w.shape == dst_lB_w.shape:
                        dst_lB_w.data.copy_(src_lB_w.data)
            except Exception as e:
                logger.debug("LoRA transfer failed: %s", e)

        svd_layer.reconstruction_error = rel_err
        return svd_layer

    # ------------------------------------------------------------------
    # Weight-sharing detection
    # ------------------------------------------------------------------
    @staticmethod
    def _get_weight_tensor(module: nn.Module) -> torch.Tensor | None:
        base = getattr(module, "base_layer", module)
        return getattr(base, "weight", None)

    # ------------------------------------------------------------------
    # Main compress
    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        self.replaced_layers = []
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0
        seen_weights: set[int] = set()

        target_layers: list[tuple[str, nn.Module]] = []
        for name, mod in list(self.model.named_modules()):
            if not self._is_target_module(name, mod):
                continue
            w = self._get_weight_tensor(mod)
            if w is None:
                continue
            w_id = id(w)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)
            target_layers.append((name, mod))

        if not target_layers:
            logger.warning("SwiftSVDActivation: no compressible linear layers found.")
            return self.model, self.get_compression_stats()

        activations = self._collect_activations()

        for name, module in target_layers:
            act_pair = activations.get(name, None)
            base_layer = getattr(module, "base_layer", module)
            raw_w = base_layer.weight.data

            # Original bytes
            orig_bytes = raw_w.numel() * raw_w.element_size()
            if getattr(base_layer, "bias", None) is not None:
                orig_bytes += base_layer.bias.numel() * base_layer.bias.element_size()
            self._orig_storage_bytes += orig_bytes

            svd_layer = self.compress_linear(module, act_pair=act_pair)
            if svd_layer is not None:
                err = getattr(svd_layer, "reconstruction_error", 0.0)
                self.layer_errors[name] = err
                self.replaced_layers.append(
                    {
                        "name": name,
                        "rank": svd_layer.rank,
                        "error": err,
                    }
                )

                if hasattr(module, "base_layer"):
                    module.base_layer = svd_layer
                    for attr in ("lora_A", "lora_B"):
                        sub = getattr(module, attr, None)
                        if sub is None:
                            continue
                        params = sub.parameters() if hasattr(sub, "parameters") else iter([sub])
                        for p in params:
                            if hasattr(p, "requires_grad"):
                                p.requires_grad = True
                else:
                    replace_module(self.model, name, svd_layer)

                self.compressed_layers.append(name)
                # Honest byte count
                comp_bytes = 0
                for p in svd_layer.parameters():
                    comp_bytes += p.numel() * p.element_size()
                for b in svd_layer.buffers():
                    if b is not None:
                        comp_bytes += b.numel() * b.element_size()
                self._comp_storage_bytes += comp_bytes
            else:
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                # Original storage retained
                self._comp_storage_bytes += orig_bytes

        # ---- Forward sanity check (SW-FIX-15) ----
        forward_check_success = True
        try:
            with torch.no_grad():
                dev = next(self.model.parameters(), torch.tensor(0.0)).device
                cfg = getattr(self.model, "config", None)
                vocab = getattr(cfg, "vocab_size", None) if cfg else None
                if vocab is None:
                    vocab = 100
                dummy = torch.randint(0, min(vocab, 100), (1, 8), device=dev)
                try:
                    self.model(input_ids=dummy, attention_mask=torch.ones_like(dummy))
                except Exception:
                    try:
                        self.model(dummy)
                    except Exception as e2:
                        logger.warning("Post-compression forward check failed: %s", e2)
                        forward_check_success = False
        except Exception as e:
            logger.warning("Post-compression forward check warning: %s", e)
            forward_check_success = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        avg_err = (
            sum(self.layer_errors[k] for k in self.compressed_layers)
            / max(1, len(self.compressed_layers))
            if self.compressed_layers
            else 0.0
        )
        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        storage_ratio = (
            self._orig_storage_bytes / self._comp_storage_bytes
            if self._comp_storage_bytes > 0
            else 1.0
        )
        effective_ratio = (
            self._comp_storage_bytes / self._orig_storage_bytes
            if self._orig_storage_bytes > 0
            else 1.0
        )

        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": self.fallback_layers,
                "skipped_layers": self.skipped_layers,
                "num_compressed": len(self.compressed_layers),
                "num_fallback": len(self.fallback_layers),
                "num_skipped": len(self.skipped_layers),
                "layer_errors": self.layer_errors,
                "avg_reconstruction_error": avg_err,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "lora_r": self.lora_r,
                "lora_alpha": self.lora_alpha,
                # SW-FIX-10: honest accounting
                "orig_storage_bytes": self._orig_storage_bytes,
                "comp_storage_bytes": self._comp_storage_bytes,
                "storage_compression_ratio": round(storage_ratio, 3),
                "storage_reduction_pct": round(100.0 * (1.0 - effective_ratio), 2),
                "effective_parameter_ratio": round(effective_ratio, 4),
                "forward_check_success": forward_check_success,
                "method": "swift_svd_act",
            }
        )
        self.stats = stats
        logger.info(
            "Swift-SVD-ACT: %d compressed, %d fallback, avg_err=%.4f, storage_ratio=%.2fx",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            storage_ratio,
        )
        return self.model, stats

Swift-SVD Activation-Aware Compressor.

Captures empirical input-output cross-covariances across calibration data, derives activation-aligned subspaces, projects weights onto the principal activation subspace, and replaces dense nn.Linear / Conv1D layers with factorized SVDLinear modules.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def compress_linear(self, layer: nn.Module, act_pair: tuple[torch.Tensor, int] | None = None) ‑> SVDLinear | None

Expand source code

def compress_linear(
    self,
    layer: nn.Module,
    act_pair: tuple[torch.Tensor, int] | None = None,
) -> SVDLinear | None:
    base_layer = getattr(layer, "base_layer", layer)
    raw_w = base_layer.weight.data
    cls_name = base_layer.__class__.__name__.lower()
    is_conv1d = "conv1d" in cls_name or hasattr(base_layer, "nf")

    W = raw_w.t().contiguous().clone().float() if is_conv1d else raw_w.clone().float()
    out_f, in_f = W.shape

    r = self._select_rank(in_f, out_f)
    device = raw_w.device
    dtype = raw_w.dtype
    A: torch.Tensor | None = None
    B: torch.Tensor | None = None

    # ---- Activation subspace projection ----
    if act_pair is not None:
        try:
            cov, _n = act_pair
            if cov.shape[0] == out_f and cov.shape[1] == in_f:
                C = cov.to(device=device, dtype=torch.float32)
                U_c, _, Vt_c = torch.linalg.svd(C, full_matrices=False)
                r_act = min(r, U_c.shape[1], Vt_c.shape[0])
                if r_act > 0:
                    U_act = U_c[:, :r_act]
                    V_act = Vt_c[:r_act, :].t().contiguous()

                    # M = U_act^T W V_act
                    M = U_act.t().contiguous() @ W @ V_act
                    U_m, S_m, Vt_m = torch.linalg.svd(M, full_matrices=False)
                    r_act = min(r_act, len(S_m))
                    if r_act > 0:
                        sqrt_S_m = torch.sqrt(torch.clamp(S_m[:r_act], min=1e-12))
                        A = (U_act @ U_m[:, :r_act]) * sqrt_S_m.unsqueeze(0)
                        B = sqrt_S_m.unsqueeze(1) @ (Vt_m[:r_act, :] @ V_act.t().contiguous())
                        r = r_act
        except Exception as e:
            logger.debug(
                "Activation subspace projection failed (%s). Falling back to direct SVD.",
                e,
            )
            A, B = None, None

    # ---- Direct SVD fallback ----
    if A is None or B is None:
        try:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            r = min(r, len(S))
            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
            A = U[:, :r] * sqrt_S.unsqueeze(0)
            B = sqrt_S.unsqueeze(1) @ Vt[:r, :]
        except Exception as e:
            logger.warning("SwiftSVDActivation SVD failure: %s", e)
            return None

    # ---- Error check ----
    W_rec = A @ B
    norm_orig = torch.norm(W, p="fro").item()
    rel_err = (torch.norm(W - W_rec) / max(norm_orig, 1e-12)).item()

    if not math.isfinite(rel_err) or rel_err > self.max_error:
        # Try direct SVD (may have better conditioning)
        try:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            r_d = min(r, len(S))
            sqrt_S = torch.sqrt(torch.clamp(S[:r_d], min=1e-12))
            A_d = U[:, :r_d] * sqrt_S.unsqueeze(0)
            B_d = sqrt_S.unsqueeze(1) @ Vt[:r_d, :]
            W_rec_d = A_d @ B_d
            err_d = (torch.norm(W - W_rec_d) / max(norm_orig, 1e-12)).item()
            if err_d < rel_err:
                A, B, r, rel_err = A_d, B_d, r_d, err_d
        except Exception as e:
            logger.debug("Direct SVD fallback failed: %s", e)

    # SW-FIX-14: preserve original on high error
    if not math.isfinite(rel_err) or rel_err > self.max_error:
        logger.warning(
            "SwiftSVDActivation reconstruction error %.4f exceeded threshold %.4f on %s.",
            rel_err,
            self.max_error,
            getattr(base_layer, "_name", "?"),
        )
        return None

    # ---- Build SVDLinear ----
    effective_lora_r = self.lora_r
    detected = self._detect_lora_rank(layer)
    if detected > 0:
        effective_lora_r = detected

    svd_layer = SVDLinear(
        in_features=in_f,
        out_features=out_f,
        rank=r,
        bias=base_layer.bias is not None,
        dtype=dtype,
        device=device,
        lora_r=effective_lora_r,
        lora_alpha=self.lora_alpha,
    )
    svd_layer.set_low_rank_factors(
        U_factor=A.to(dtype=dtype, device=device),
        V_factor=B.to(dtype=dtype, device=device),
    )
    if base_layer.bias is not None and svd_layer.bias is not None:
        svd_layer.bias.data.copy_(base_layer.bias.data)

    # ---- LoRA transfer ----
    if effective_lora_r > 0 and getattr(svd_layer, "use_lora", False):
        try:
            src_lA = getattr(layer, "lora_A", None)
            src_lB = getattr(layer, "lora_B", None)
            if src_lA is not None and src_lB is not None:
                src_lA_w = getattr(src_lA, "weight", src_lA)
                src_lB_w = getattr(src_lB, "weight", src_lB)
                dst_lA_w = getattr(svd_layer.lora_A, "weight", svd_layer.lora_A)
                dst_lB_w = getattr(svd_layer.lora_B, "weight", svd_layer.lora_B)
                if src_lA_w.shape == dst_lA_w.shape:
                    dst_lA_w.data.copy_(src_lA_w.data)
                if src_lB_w.shape == dst_lB_w.shape:
                    dst_lB_w.data.copy_(src_lB_w.data)
        except Exception as e:
            logger.debug("LoRA transfer failed: %s", e)

    svd_layer.reconstruction_error = rel_err
    return svd_layer

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class SwiftSVDCompressor (model: nn.Module, rank: int | None = None, rank_ratio: float = 0.5, target_modules: list[str] | None = None, max_error: float = 0.6, max_retries: int = 3, fallback_to_original: bool = True, validation_calibration: bool = False, calibration_threshold: float = 0.1, target_attention_only: bool = False, lora_r: int = 0, lora_alpha: float = 1.0, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("swift_svd_class")
class SwiftSVDCompressor(BaseCompressor):
    is_custom_attention: bool = True
    """
    Swift-SVD Compressor: Fast Closed-Form Truncated SVD.

    Decomposes dense weight matrices into rank-r factorized layers with symmetric
    square-root singular value splitting (U_r * sqrt(Sigma_r), sqrt(Sigma_r) * V_r^T)
    and replaces nn.Linear/Conv1D with factorized SVDLinear modules.
    """

    def __init__(
        self,
        model: nn.Module,
        rank: int | None = None,
        rank_ratio: float = 0.5,
        target_modules: list[str] | None = None,
        max_error: float = 0.60,
        max_retries: int = 3,
        fallback_to_original: bool = True,
        validation_calibration: bool = False,
        calibration_threshold: float = 0.10,
        target_attention_only: bool = False,
        lora_r: int = 0,
        lora_alpha: float = 1.0,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="swift_svd", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            rank = _cfg_get(cfg, "swift_svd_rank", _cfg_get(cfg, "rank", rank))
            rank_ratio = _cfg_get(
                cfg, "swift_svd_rank_ratio", _cfg_get(cfg, "rank_ratio", rank_ratio)
            )
            max_error = _cfg_get(
                cfg, "max_error", _cfg_get(cfg, "max_reconstruction_error", max_error)
            )
            max_retries = _cfg_get(cfg, "max_retries", max_retries)
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            validation_calibration = _cfg_get(cfg, "validation_calibration", validation_calibration)
            calibration_threshold = _cfg_get(cfg, "calibration_threshold", calibration_threshold)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)
            lora_r = _cfg_get(cfg, "lora_r", lora_r)
            lora_alpha = _cfg_get(cfg, "lora_alpha", lora_alpha)

        if "rank" in kwargs and kwargs["rank"] is not None:
            rank = kwargs["rank"]
        if "max_error" in kwargs and kwargs["max_error"] is not None:
            max_error = float(kwargs["max_error"])
        if "max_retries" in kwargs and kwargs["max_retries"] is not None:
            max_retries = int(kwargs["max_retries"])
        if "fallback_to_original" in kwargs:
            fallback_to_original = bool(kwargs["fallback_to_original"])
        if "validation_calibration" in kwargs:
            validation_calibration = bool(kwargs["validation_calibration"])
        if "calibration_threshold" in kwargs:
            calibration_threshold = float(kwargs["calibration_threshold"])
        if "target_attention_only" in kwargs:
            target_attention_only = bool(kwargs["target_attention_only"])
        if "lora_r" in kwargs:
            lora_r = int(kwargs["lora_r"])
        if "lora_alpha" in kwargs:
            lora_alpha = float(kwargs["lora_alpha"])

        self.rank = rank
        self.rank_ratio = float(rank_ratio)
        self.target_modules = target_modules
        self.max_error = float(max_error)
        self.max_retries = max(0, int(max_retries))
        self.fallback_to_original = bool(fallback_to_original)
        self.validation_calibration = bool(validation_calibration)
        self.calibration_threshold = float(calibration_threshold)
        self.target_attention_only = bool(target_attention_only)
        self.lora_r = int(lora_r)
        self.lora_alpha = float(lora_alpha)

        self.compressed_layers: list[dict[str, Any]] = []
        # SV-FIX-02: initialize in __init__.
        self._compressed_params: int = 0
        self.stats: dict[str, Any] = {}

    # ------------------------------------------------------------------
    # Target selection
    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        if self.target_attention_only:
            attn_keywords = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            if not any(k in name.lower() for k in attn_keywords):
                return False

        base_layer = getattr(module, "base_layer", module)
        cls_name = base_layer.__class__.__name__.lower()
        is_linear_like = (
            isinstance(base_layer, nn.Linear)
            or "conv1d" in cls_name
            or hasattr(base_layer, "nf")
            or (
                hasattr(base_layer, "weight")
                and hasattr(base_layer.weight, "ndim")
                and base_layer.weight.ndim == 2
                and not isinstance(base_layer, (nn.Embedding, nn.LayerNorm, nn.RMSNorm))
            )
        )
        if not is_linear_like:
            return False
        if isinstance(module, SVDLinear) or isinstance(base_layer, SVDLinear):
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self.target_modules is None:
            return True
        return any(t in name for t in self.target_modules)

    # ------------------------------------------------------------------
    # Per-layer compression
    # ------------------------------------------------------------------
    def compress_linear(
        self,
        layer: nn.Module,
        rank: int | None = None,
        name: str = "",
    ) -> tuple[SVDLinear | None, float, int]:
        """
        Decompose a single linear or Conv1D layer via truncated SVD.
        Returns (svd_layer_or_None, relative_frobenius_error, final_rank).
        """
        base_layer = getattr(layer, "base_layer", layer)
        device = base_layer.weight.device
        dtype = base_layer.weight.dtype

        cls_name = base_layer.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(base_layer, "nf")

        raw_w = base_layer.weight.data
        # Always move weight to a canonical (out_features, in_features) form.
        if is_conv1d:
            W = raw_w.t().contiguous().float()
            out_f, in_f = W.shape[0], W.shape[1]
        else:
            W = raw_w.float()
            out_f, in_f = W.shape[0], W.shape[1]

        # Param budget: r * (in_f + out_f) < in_f * out_f
        max_param_r = max(1, (in_f * out_f) // (in_f + out_f) - 1)
        max_r = min(out_f, in_f, max_param_r)
        if max_r < 1:
            return None, 1.0, 0

        # Target rank
        if rank is not None:
            target_r = max(1, min(int(rank), max_r))
        else:
            target_r = max(1, min(int(round(min(out_f, in_f) * self.rank_ratio)), max_r))

        # SV-FIX-01: single try/except with explicit CUDA OOM handling.
        U: torch.Tensor | None = None
        S: torch.Tensor | None = None
        Vt: torch.Tensor | None = None
        try:
            if W.is_cuda:
                U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            else:
                W_dev = W.to(device) if device.type == "cuda" else W
                U, S, Vt = torch.linalg.svd(W_dev, full_matrices=False)
        except torch.cuda.OutOfMemoryError:
            try:
                torch.cuda.empty_cache()
                U, S, Vt = torch.linalg.svd(W.cpu(), full_matrices=False)
                U, S, Vt = U.to(device), S.to(device), Vt.to(device)
            except Exception as e2:
                logger.warning("SwiftSVD SVD (OOM fallback) failed for %s: %s", name, e2)
                return None, 1.0, 0
        except Exception as e:
            logger.warning("SwiftSVD SVD failed for %s: %s", name, e)
            return None, 1.0, 0
        norm_orig = torch.linalg.norm(W).item() + 1e-8
        cur_r = target_r
        best_factors: tuple[torch.Tensor, torch.Tensor] | None = None
        best_err = float("inf")
        best_r = cur_r

        # Iterative rank-step retries
        attempts = self.max_retries + 1
        for attempt in range(attempts):
            sqrt_S = torch.sqrt(torch.clamp(S[:cur_r], min=1e-12))
            U_r = U[:, :cur_r] * sqrt_S.unsqueeze(0)
            V_r = sqrt_S.unsqueeze(1) * Vt[:cur_r, :]

            W_rec = torch.matmul(U_r, V_r)
            rel_err = torch.linalg.norm(W - W_rec).item() / norm_orig

            if rel_err < best_err:
                best_err = rel_err
                best_factors = (U_r, V_r)
                best_r = cur_r

            if rel_err <= self.max_error or cur_r >= max_r:
                break

            remaining = attempts - attempt - 1
            if remaining <= 0:
                break
            step = max(4, (max_r - cur_r) // max(1, remaining))
            cur_r = min(cur_r + step, max_r)

        if best_err > self.max_error:
            logger.warning(
                "SwiftSVD: %s error (%.4f) > max_error (%.4f) at rank %d. %s",
                name,
                best_err,
                self.max_error,
                best_r,
                (
                    "Keeping original."
                    if self.fallback_to_original
                    else "Using compressed layer anyway."
                ),
            )
            if self.fallback_to_original:
                return None, best_err, best_r

        if best_factors is None:
            return None, best_err, best_r

        U_final, V_final = best_factors
        U_final = U_final.to(dtype=dtype, device=device)
        V_final = V_final.to(dtype=dtype, device=device)

        has_bias = getattr(base_layer, "bias", None) is not None
        has_base_layer = hasattr(layer, "base_layer")
        has_existing_lora = hasattr(layer, "lora_A") and hasattr(layer, "lora_B")

        if has_base_layer:
            # Wrapped: SVD becomes the new base, existing LoRA stays on outer.
            effective_lora_r = 0
        elif has_existing_lora:
            detected = self._detect_lora_rank(layer)
            effective_lora_r = detected if detected > 0 else self.lora_r
        else:
            effective_lora_r = self.lora_r

        svd_layer = SVDLinear(
            in_features=in_f,
            out_features=out_f,
            rank=best_r,
            bias=has_bias,
            dtype=dtype,
            device=device,
            lora_r=effective_lora_r,
            lora_alpha=self.lora_alpha,
        )
        svd_layer.set_low_rank_factors(U_final, V_final)

        if has_bias and svd_layer.bias is not None and base_layer.bias is not None:
            svd_layer.bias.data.copy_(base_layer.bias.data.to(dtype=dtype, device=device))

        # SV-FIX-04: robust LoRA copy via data-shape check.
        if not has_base_layer and has_existing_lora and getattr(svd_layer, "use_lora", False):
            try:
                src_lA = getattr(layer, "lora_A", None)
                src_lB = getattr(layer, "lora_B", None)
                if src_lA is not None and src_lB is not None:
                    src_lA_w = getattr(src_lA, "weight", src_lA)
                    src_lB_w = getattr(src_lB, "weight", src_lB)
                    dst_lA_w = getattr(svd_layer.lora_A, "weight", svd_layer.lora_A)
                    dst_lB_w = getattr(svd_layer.lora_B, "weight", svd_layer.lora_B)
                    if src_lA_w.shape == dst_lA_w.shape:
                        dst_lA_w.data.copy_(src_lA_w.data)
                    if src_lB_w.shape == dst_lB_w.shape:
                        dst_lB_w.data.copy_(src_lB_w.data)
            except Exception as e:
                logger.warning("SwiftSVD LoRA copy warning for %s: %s", name, e)

        # Dimensional integrity check
        if svd_layer.in_features != in_f or svd_layer.out_features != out_f:
            logger.error(
                "SwiftSVD: shape mismatch for %s: expected (in=%d, out=%d), got (in=%d, out=%d)",
                name,
                in_f,
                out_f,
                svd_layer.in_features,
                svd_layer.out_features,
            )
            return None, best_err, best_r

        # Optional calibration validation
        if self.validation_calibration:
            try:
                with torch.no_grad():
                    dummy_x = torch.randn(1, in_f, device=device, dtype=dtype)
                    orig_out = base_layer(dummy_x)
                    svd_out = svd_layer(dummy_x)
                    rel_diff = (
                        torch.norm(orig_out.float() - svd_out.float())
                        / (torch.norm(orig_out.float()) + 1e-8)
                    ).item()
                    if rel_diff > self.calibration_threshold:
                        logger.warning(
                            "SwiftSVD: %s calibration diff=%.4f > %.4f. Keeping original.",
                            name,
                            rel_diff,
                            self.calibration_threshold,
                        )
                        return None, rel_diff, best_r
            except Exception as ce:
                logger.warning("SwiftSVD: %s calibration error %s. Proceeding.", name, ce)

        return svd_layer, best_err, best_r

    # ------------------------------------------------------------------
    # LoRA rank detection
    # ------------------------------------------------------------------
    @staticmethod
    def _detect_lora_rank(module: nn.Module) -> int:
        for attr in ("r", "lora_r", "rank"):
            val = getattr(module, attr, None)
            if isinstance(val, int) and val > 0:
                return val
        la = getattr(module, "lora_A", None)
        if la is not None:
            if hasattr(la, "out_features"):
                try:
                    return int(la.out_features)
                except Exception:
                    pass
            if hasattr(la, "weight") and la.weight is not None and la.weight.dim() >= 1:
                return int(la.weight.shape[0])
        return 0

    # ------------------------------------------------------------------
    # Main compress
    # ------------------------------------------------------------------
    def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
        """Compress all eligible linear and Conv1D layers."""
        # SV-FIX-07: deduplicate shared weights.
        seen_weights: set = set()
        target_layers: list[tuple[str, nn.Module]] = []
        for name, mod in self.model.named_modules():
            if not self._is_target_module(name, mod):
                continue
            base = getattr(mod, "base_layer", mod)
            w = getattr(base, "weight", None)
            if w is None:
                continue
            w_id = id(w)
            if w_id in seen_weights:
                continue
            seen_weights.add(w_id)
            target_layers.append((name, mod))

        if not target_layers:
            logger.warning("SwiftSVD: No compressible linear layers found.")
            try:
                stats = self.get_compression_stats()
            except Exception:
                stats = {}
            return self.model, stats

        self.compressed_layers = []
        fallback_layers: list[str] = []
        total_error = 0.0

        for name, module in target_layers:
            svd_layer, rel_err, final_rank = self.compress_linear(module, rank=self.rank, name=name)

            if svd_layer is not None:
                # SV-FIX-03: raw_shape from base_layer, correctly oriented.
                base_mod = getattr(module, "base_layer", module)
                cls_name = base_mod.__class__.__name__.lower()
                is_conv1d = "conv1d" in cls_name or hasattr(base_mod, "nf")
                if is_conv1d:
                    expected_in = base_mod.weight.shape[0]
                    expected_out = base_mod.weight.shape[1]
                else:
                    expected_in = getattr(base_mod, "in_features", base_mod.weight.shape[1])
                    expected_out = getattr(base_mod, "out_features", base_mod.weight.shape[0])

                if svd_layer.in_features != expected_in or svd_layer.out_features != expected_out:
                    logger.error(
                        "SwiftSVD: skipping %s due to shape mismatch "
                        "(expected in=%d, out=%d; got in=%d, out=%d)",
                        name,
                        expected_in,
                        expected_out,
                        svd_layer.in_features,
                        svd_layer.out_features,
                    )
                    fallback_layers.append(name)
                    continue

                if hasattr(module, "base_layer"):
                    module.base_layer = svd_layer
                    for attr in ("lora_A", "lora_B"):
                        sub = getattr(module, attr, None)
                        if sub is None:
                            continue
                        params = sub.parameters() if hasattr(sub, "parameters") else iter([sub])
                        for p in params:
                            p.requires_grad = True
                else:
                    replace_module(self.model, name, svd_layer)

                self.compressed_layers.append(
                    {
                        "name": name,
                        "original_shape": (expected_out, expected_in),
                        "rank": final_rank,
                        "rel_error": float(rel_err),
                    }
                )
                total_error += rel_err
            else:
                fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                logger.info("SwiftSVD: kept original layer for %s.", name)

        # SV-FIX-02: safe param count.
        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        # SV-FIX-05: guard stats.
        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}

        avg_err = (total_error / len(self.compressed_layers)) if self.compressed_layers else 0.0
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": fallback_layers,
                "num_compressed_layers": len(self.compressed_layers),
                "num_fallback": len(fallback_layers),
                "avg_reconstruction_error": avg_err,
                "rank": self.rank,
                "rank_ratio": self.rank_ratio,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "param_reduction_pct": stats.get("param_reduction_pct", 0.0),
                "method": "swift_svd",
            }
        )
        self.stats = stats
        logger.info(
            "SwiftSVD: %d layers compressed, avg error %.4f, param reduction %.2f%%",
            len(self.compressed_layers),
            avg_err,
            stats.get("param_reduction_pct", 0.0),
        )
        return self.model, stats

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: Swift-SVD Compressor: Fast Closed-Form Truncated SVD.

Decomposes dense weight matrices into rank-r factorized layers with symmetric square-root singular value splitting (U_r * sqrt(Sigma_r), sqrt(Sigma_r) * V_r^T) and replaces nn.Linear/Conv1D with factorized SVDLinear modules.

Methods

def compress(self, **kwargs: Any) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, **kwargs: Any) -> tuple[nn.Module, dict[str, Any]]:
    """Compress all eligible linear and Conv1D layers."""
    # SV-FIX-07: deduplicate shared weights.
    seen_weights: set = set()
    target_layers: list[tuple[str, nn.Module]] = []
    for name, mod in self.model.named_modules():
        if not self._is_target_module(name, mod):
            continue
        base = getattr(mod, "base_layer", mod)
        w = getattr(base, "weight", None)
        if w is None:
            continue
        w_id = id(w)
        if w_id in seen_weights:
            continue
        seen_weights.add(w_id)
        target_layers.append((name, mod))

    if not target_layers:
        logger.warning("SwiftSVD: No compressible linear layers found.")
        try:
            stats = self.get_compression_stats()
        except Exception:
            stats = {}
        return self.model, stats

    self.compressed_layers = []
    fallback_layers: list[str] = []
    total_error = 0.0

    for name, module in target_layers:
        svd_layer, rel_err, final_rank = self.compress_linear(module, rank=self.rank, name=name)

        if svd_layer is not None:
            # SV-FIX-03: raw_shape from base_layer, correctly oriented.
            base_mod = getattr(module, "base_layer", module)
            cls_name = base_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(base_mod, "nf")
            if is_conv1d:
                expected_in = base_mod.weight.shape[0]
                expected_out = base_mod.weight.shape[1]
            else:
                expected_in = getattr(base_mod, "in_features", base_mod.weight.shape[1])
                expected_out = getattr(base_mod, "out_features", base_mod.weight.shape[0])

            if svd_layer.in_features != expected_in or svd_layer.out_features != expected_out:
                logger.error(
                    "SwiftSVD: skipping %s due to shape mismatch "
                    "(expected in=%d, out=%d; got in=%d, out=%d)",
                    name,
                    expected_in,
                    expected_out,
                    svd_layer.in_features,
                    svd_layer.out_features,
                )
                fallback_layers.append(name)
                continue

            if hasattr(module, "base_layer"):
                module.base_layer = svd_layer
                for attr in ("lora_A", "lora_B"):
                    sub = getattr(module, attr, None)
                    if sub is None:
                        continue
                    params = sub.parameters() if hasattr(sub, "parameters") else iter([sub])
                    for p in params:
                        p.requires_grad = True
            else:
                replace_module(self.model, name, svd_layer)

            self.compressed_layers.append(
                {
                    "name": name,
                    "original_shape": (expected_out, expected_in),
                    "rank": final_rank,
                    "rel_error": float(rel_err),
                }
            )
            total_error += rel_err
        else:
            fallback_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False
            logger.info("SwiftSVD: kept original layer for %s.", name)

    # SV-FIX-02: safe param count.
    try:
        self._compressed_params = self._count_params(self.model)
    except Exception:
        self._compressed_params = sum(p.numel() for p in self.model.parameters())

    # SV-FIX-05: guard stats.
    try:
        stats = self.get_compression_stats()
    except Exception:
        stats = {}

    avg_err = (total_error / len(self.compressed_layers)) if self.compressed_layers else 0.0
    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "fallback_layers": fallback_layers,
            "num_compressed_layers": len(self.compressed_layers),
            "num_fallback": len(fallback_layers),
            "avg_reconstruction_error": avg_err,
            "rank": self.rank,
            "rank_ratio": self.rank_ratio,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "param_reduction_pct": stats.get("param_reduction_pct", 0.0),
            "method": "swift_svd",
        }
    )
    self.stats = stats
    logger.info(
        "SwiftSVD: %d layers compressed, avg error %.4f, param reduction %.2f%%",
        len(self.compressed_layers),
        avg_err,
        stats.get("param_reduction_pct", 0.0),
    )
    return self.model, stats

Compress all eligible linear and Conv1D layers.

def compress_linear(self, layer: nn.Module, rank: int | None = None, name: str = '') ‑> tuple[SVDLinear | None, float, int]

Expand source code

def compress_linear(
    self,
    layer: nn.Module,
    rank: int | None = None,
    name: str = "",
) -> tuple[SVDLinear | None, float, int]:
    """
    Decompose a single linear or Conv1D layer via truncated SVD.
    Returns (svd_layer_or_None, relative_frobenius_error, final_rank).
    """
    base_layer = getattr(layer, "base_layer", layer)
    device = base_layer.weight.device
    dtype = base_layer.weight.dtype

    cls_name = base_layer.__class__.__name__.lower()
    is_conv1d = "conv1d" in cls_name or hasattr(base_layer, "nf")

    raw_w = base_layer.weight.data
    # Always move weight to a canonical (out_features, in_features) form.
    if is_conv1d:
        W = raw_w.t().contiguous().float()
        out_f, in_f = W.shape[0], W.shape[1]
    else:
        W = raw_w.float()
        out_f, in_f = W.shape[0], W.shape[1]

    # Param budget: r * (in_f + out_f) < in_f * out_f
    max_param_r = max(1, (in_f * out_f) // (in_f + out_f) - 1)
    max_r = min(out_f, in_f, max_param_r)
    if max_r < 1:
        return None, 1.0, 0

    # Target rank
    if rank is not None:
        target_r = max(1, min(int(rank), max_r))
    else:
        target_r = max(1, min(int(round(min(out_f, in_f) * self.rank_ratio)), max_r))

    # SV-FIX-01: single try/except with explicit CUDA OOM handling.
    U: torch.Tensor | None = None
    S: torch.Tensor | None = None
    Vt: torch.Tensor | None = None
    try:
        if W.is_cuda:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        else:
            W_dev = W.to(device) if device.type == "cuda" else W
            U, S, Vt = torch.linalg.svd(W_dev, full_matrices=False)
    except torch.cuda.OutOfMemoryError:
        try:
            torch.cuda.empty_cache()
            U, S, Vt = torch.linalg.svd(W.cpu(), full_matrices=False)
            U, S, Vt = U.to(device), S.to(device), Vt.to(device)
        except Exception as e2:
            logger.warning("SwiftSVD SVD (OOM fallback) failed for %s: %s", name, e2)
            return None, 1.0, 0
    except Exception as e:
        logger.warning("SwiftSVD SVD failed for %s: %s", name, e)
        return None, 1.0, 0
    norm_orig = torch.linalg.norm(W).item() + 1e-8
    cur_r = target_r
    best_factors: tuple[torch.Tensor, torch.Tensor] | None = None
    best_err = float("inf")
    best_r = cur_r

    # Iterative rank-step retries
    attempts = self.max_retries + 1
    for attempt in range(attempts):
        sqrt_S = torch.sqrt(torch.clamp(S[:cur_r], min=1e-12))
        U_r = U[:, :cur_r] * sqrt_S.unsqueeze(0)
        V_r = sqrt_S.unsqueeze(1) * Vt[:cur_r, :]

        W_rec = torch.matmul(U_r, V_r)
        rel_err = torch.linalg.norm(W - W_rec).item() / norm_orig

        if rel_err < best_err:
            best_err = rel_err
            best_factors = (U_r, V_r)
            best_r = cur_r

        if rel_err <= self.max_error or cur_r >= max_r:
            break

        remaining = attempts - attempt - 1
        if remaining <= 0:
            break
        step = max(4, (max_r - cur_r) // max(1, remaining))
        cur_r = min(cur_r + step, max_r)

    if best_err > self.max_error:
        logger.warning(
            "SwiftSVD: %s error (%.4f) > max_error (%.4f) at rank %d. %s",
            name,
            best_err,
            self.max_error,
            best_r,
            (
                "Keeping original."
                if self.fallback_to_original
                else "Using compressed layer anyway."
            ),
        )
        if self.fallback_to_original:
            return None, best_err, best_r

    if best_factors is None:
        return None, best_err, best_r

    U_final, V_final = best_factors
    U_final = U_final.to(dtype=dtype, device=device)
    V_final = V_final.to(dtype=dtype, device=device)

    has_bias = getattr(base_layer, "bias", None) is not None
    has_base_layer = hasattr(layer, "base_layer")
    has_existing_lora = hasattr(layer, "lora_A") and hasattr(layer, "lora_B")

    if has_base_layer:
        # Wrapped: SVD becomes the new base, existing LoRA stays on outer.
        effective_lora_r = 0
    elif has_existing_lora:
        detected = self._detect_lora_rank(layer)
        effective_lora_r = detected if detected > 0 else self.lora_r
    else:
        effective_lora_r = self.lora_r

    svd_layer = SVDLinear(
        in_features=in_f,
        out_features=out_f,
        rank=best_r,
        bias=has_bias,
        dtype=dtype,
        device=device,
        lora_r=effective_lora_r,
        lora_alpha=self.lora_alpha,
    )
    svd_layer.set_low_rank_factors(U_final, V_final)

    if has_bias and svd_layer.bias is not None and base_layer.bias is not None:
        svd_layer.bias.data.copy_(base_layer.bias.data.to(dtype=dtype, device=device))

    # SV-FIX-04: robust LoRA copy via data-shape check.
    if not has_base_layer and has_existing_lora and getattr(svd_layer, "use_lora", False):
        try:
            src_lA = getattr(layer, "lora_A", None)
            src_lB = getattr(layer, "lora_B", None)
            if src_lA is not None and src_lB is not None:
                src_lA_w = getattr(src_lA, "weight", src_lA)
                src_lB_w = getattr(src_lB, "weight", src_lB)
                dst_lA_w = getattr(svd_layer.lora_A, "weight", svd_layer.lora_A)
                dst_lB_w = getattr(svd_layer.lora_B, "weight", svd_layer.lora_B)
                if src_lA_w.shape == dst_lA_w.shape:
                    dst_lA_w.data.copy_(src_lA_w.data)
                if src_lB_w.shape == dst_lB_w.shape:
                    dst_lB_w.data.copy_(src_lB_w.data)
        except Exception as e:
            logger.warning("SwiftSVD LoRA copy warning for %s: %s", name, e)

    # Dimensional integrity check
    if svd_layer.in_features != in_f or svd_layer.out_features != out_f:
        logger.error(
            "SwiftSVD: shape mismatch for %s: expected (in=%d, out=%d), got (in=%d, out=%d)",
            name,
            in_f,
            out_f,
            svd_layer.in_features,
            svd_layer.out_features,
        )
        return None, best_err, best_r

    # Optional calibration validation
    if self.validation_calibration:
        try:
            with torch.no_grad():
                dummy_x = torch.randn(1, in_f, device=device, dtype=dtype)
                orig_out = base_layer(dummy_x)
                svd_out = svd_layer(dummy_x)
                rel_diff = (
                    torch.norm(orig_out.float() - svd_out.float())
                    / (torch.norm(orig_out.float()) + 1e-8)
                ).item()
                if rel_diff > self.calibration_threshold:
                    logger.warning(
                        "SwiftSVD: %s calibration diff=%.4f > %.4f. Keeping original.",
                        name,
                        rel_diff,
                        self.calibration_threshold,
                    )
                    return None, rel_diff, best_r
        except Exception as ce:
            logger.warning("SwiftSVD: %s calibration error %s. Proceeding.", name, ce)

    return svd_layer, best_err, best_r

Decompose a single linear or Conv1D layer via truncated SVD. Returns (svd_layer_or_None, relative_frobenius_error, final_rank).

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class TTCompressor (model: nn.Module | None = None, rank: int | None = None, tt_rank: int = 16, num_factors: int = 2, cfg: Any | None = None, target_layers: list[str] | None = None, max_error: float = 0.2, energy_threshold: float = 0.995, fallback_policy: str = 'svd', target_attention_only: bool = False, strict_rank: bool = False, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("tt")
@COMPRESSOR_REGISTRY.register("tt_linear")
@COMPRESSOR_REGISTRY.register("tt_compressor")
@COMPRESSOR_REGISTRY.register("tt_linear_compressor")
class TTCompressor(BaseCompressor):
    """
    Replaces Linear/Conv1D with TTLinearLayer when TT succeeds; otherwise
    falls back to SVDLinearLayer. Layers that neither can compress within
    error budget remain untouched.

    Args:
        tt_rank:           TT rank bound.
        num_factors:       Number of TT factors per dimension (default 2).
        max_error:         Relative Frobenius error threshold (default 0.20).
        energy_threshold:  Energy retained per TT core (default 0.995).
        fallback_policy:   'svd' | 'original' | 'none'. Default: 'svd'.
        target_attention_only: Only compress attention projections.
        strict_rank:       Do not exceed tt_rank.
    """

    # TT-FIX-01
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module | None = None,
        rank: int | None = None,
        tt_rank: int = 16,
        num_factors: int = 2,
        cfg: Any | None = None,
        target_layers: list[str] | None = None,
        max_error: float = 0.20,
        energy_threshold: float = 0.995,
        fallback_policy: str = "svd",  # 'svd' | 'original' | 'none'
        target_attention_only: bool = False,
        strict_rank: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(model if model is not None else nn.Identity(), name="tt", **kwargs)
        comp_cfg = kwargs.get("comp_config", cfg)
        if comp_cfg is not None:
            rank = _cfg_get(comp_cfg, "tt_rank", _cfg_get(comp_cfg, "rank", rank))
            num_factors = _cfg_get(comp_cfg, "num_factors", num_factors)
            max_error = _cfg_get(comp_cfg, "max_error", max_error)
            energy_threshold = _cfg_get(comp_cfg, "energy_threshold", energy_threshold)
            fallback_policy = _cfg_get(comp_cfg, "fallback_policy", fallback_policy)
            target_attention_only = _cfg_get(
                comp_cfg, "target_attention_only", target_attention_only
            )
            strict_rank = _cfg_get(comp_cfg, "strict_rank", strict_rank)

        # kwargs overrides
        for k, caster in (
            ("rank", int),
            ("tt_rank", int),
            ("num_factors", int),
            ("max_error", float),
            ("energy_threshold", float),
            ("strict_rank", bool),
            ("target_attention_only", bool),
        ):
            if k in kwargs and kwargs[k] is not None:
                if k == "rank":
                    rank = caster(kwargs[k])
                else:
                    locals_dict = None  # noqa
                    # direct assignment below
        if "max_error" in kwargs and kwargs["max_error"] is not None:
            max_error = float(kwargs["max_error"])
        if "num_factors" in kwargs and kwargs["num_factors"] is not None:
            num_factors = int(kwargs["num_factors"])
        if "energy_threshold" in kwargs and kwargs["energy_threshold"] is not None:
            energy_threshold = float(kwargs["energy_threshold"])
        if "fallback_policy" in kwargs and kwargs["fallback_policy"] is not None:
            fallback_policy = str(kwargs["fallback_policy"])
        if "target_attention_only" in kwargs:
            target_attention_only = bool(kwargs["target_attention_only"])
        if "strict_rank" in kwargs and kwargs["strict_rank"] is not None:
            strict_rank = bool(kwargs["strict_rank"])
        if "tt_rank" in kwargs and kwargs["tt_rank"] is not None:
            tt_rank = int(kwargs["tt_rank"])
        elif rank is not None:
            tt_rank = int(rank)

        self.tt_rank = max(1, int(tt_rank))
        self.num_factors = max(1, int(num_factors))
        self.cfg = comp_cfg
        self.target_layers = target_layers or _cfg_get(comp_cfg, "target_layers", None)
        self.max_error = float(max_error)
        self.energy_threshold = float(energy_threshold)
        self.fallback_policy = str(fallback_policy).lower()
        if self.fallback_policy not in ("svd", "original", "none"):
            self.fallback_policy = "svd"
        self.target_attention_only = bool(target_attention_only)
        self.strict_rank = bool(strict_rank)
        self.compressed_layers: list[dict[str, Any]] = []
        self.tt_layers: list[str] = []
        self.svd_layers: list[str] = []

    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_l = name.lower()
        if any(
            k in name_l
            for k in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "layernorm",
                "rmsnorm",
                "batchnorm",
            )
        ):
            return False

        if isinstance(
            module,
            (
                nn.LayerNorm,
                nn.RMSNorm,
                nn.Embedding,
                nn.BatchNorm1d,
                nn.BatchNorm2d,
                TTLinearLayer,
                SVDLinearLayer,
            ),
        ):
            return False

        if self.target_attention_only:
            attn_kw = (
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            )
            if not any(k in name_l for k in attn_kw):
                return False

        base_layer = getattr(module, "base_layer", module)
        cls_name = base_layer.__class__.__name__.lower()
        is_linear_like = (
            isinstance(base_layer, nn.Linear)
            or "conv1d" in cls_name
            or hasattr(base_layer, "nf")
            or (
                hasattr(base_layer, "weight")
                and hasattr(base_layer.weight, "ndim")
                and base_layer.weight.ndim == 2
            )
        )
        if not is_linear_like:
            return False

        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass

        if self.target_layers and not any(t in name for t in self.target_layers):
            return False
        return True

    # ------------------------------------------------------------------
    def compress(self, model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
        if model is not None:
            self.model = model
            self._original_params = self._count_params(self.model)

        self.compressed_layers = []
        self.tt_layers = []
        self.svd_layers = []
        fallback_layers: list[str] = []
        total_error = 0.0
        orig_bytes = 0
        comp_bytes = 0

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            base_layer = getattr(module, "base_layer", module)
            raw_w = base_layer.weight.data
            cls_name = base_layer.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(base_layer, "nf")

            if hasattr(base_layer, "nx") and hasattr(base_layer, "nf"):
                in_f, out_f = int(base_layer.nx), int(base_layer.nf)
            elif hasattr(base_layer, "in_features") and hasattr(base_layer, "out_features"):
                in_f, out_f = int(base_layer.in_features), int(base_layer.out_features)
            elif is_conv1d:
                in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            else:
                out_f, in_f = raw_w.shape[0], raw_w.shape[1]

            # 1) Try TT
            tt_layer = TTLinearLayer.from_linear(
                module,
                tt_rank=self.tt_rank,
                num_factors=self.num_factors,
                max_error=self.max_error,
                energy_threshold=self.energy_threshold,
                strict_rank=self.strict_rank,
            )

            chosen_layer = None
            method = None
            rel_err = None

            if tt_layer is not None:
                rel_err = getattr(tt_layer, "reconstruction_error", None)
                chosen_layer = tt_layer
                method = "tt"

            # 2) If TT failed, try SVD fallback
            if chosen_layer is None and self.fallback_policy in ("svd", "original"):
                svd_layer = _try_svd_fallback(
                    module,
                    self.max_error,
                    target_rank=self.tt_rank * 4,
                )
                if svd_layer is not None:
                    chosen_layer = svd_layer
                    method = "svd"
                    rel_err = getattr(svd_layer, "reconstruction_error", None)

            # 3) If both failed, decide fallback
            if chosen_layer is None:
                if self.fallback_policy == "original":
                    fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    continue
                else:
                    fallback_layers.append(name)
                    continue

            # Splice into model
            if hasattr(module, "base_layer"):
                module.base_layer = chosen_layer
            else:
                replace_module(self.model, name, chosen_layer)

            # Byte accounting (TT-FIX-08)
            o_bytes = raw_w.numel() * raw_w.element_size()
            if getattr(base_layer, "bias", None) is not None and base_layer.bias is not None:
                o_bytes += base_layer.bias.numel() * base_layer.bias.element_size()
            c_bytes = sum(p.numel() * p.element_size() for p in chosen_layer.parameters())
            c_bytes += sum(
                b.numel() * b.element_size() for b in chosen_layer.buffers() if b is not None
            )

            orig_bytes += o_bytes
            comp_bytes += c_bytes
            if method == "tt":
                self.tt_layers.append(name)
            elif method == "svd":
                self.svd_layers.append(name)

            self.compressed_layers.append(
                {
                    "name": name,
                    "shape": (in_f, out_f),
                    "method": method,
                    "rel_error": float(rel_err) if rel_err is not None else None,
                }
            )
            if rel_err is not None:
                total_error += float(rel_err)

        # Forward sanity check
        forward_ok = True
        try:
            self.model.eval()
            with torch.no_grad():
                device = next(self.model.parameters()).device
                cfg = getattr(self.model, "config", None)
                vocab = getattr(cfg, "vocab_size", None) or 100
                dummy = torch.randint(0, min(vocab, 100), (1, 8), device=device)
                self.model(dummy)
        except Exception as e:
            logger.warning("TT forward sanity failed: %s", e)
            forward_ok = False

        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        stats = self.get_compression_stats()
        avg_err = total_error / max(1, len(self.compressed_layers))
        storage_ratio = orig_bytes / comp_bytes if comp_bytes > 0 else 1.0

        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "num_compressed_layers": len(self.compressed_layers),
                "fallback_layers": fallback_layers,
                "num_fallback": len(fallback_layers),
                "tt_layers": self.tt_layers,
                "svd_layers": self.svd_layers,
                "num_tt_layers": len(self.tt_layers),
                "num_svd_layers": len(self.svd_layers),
                "tt_rank": self.tt_rank,
                "num_factors": self.num_factors,
                "avg_reconstruction_error": avg_err,
                "max_error": self.max_error,
                "energy_threshold": self.energy_threshold,
                "fallback_policy": self.fallback_policy,
                "forward_check_success": forward_ok,
                "orig_storage_bytes": orig_bytes,
                "comp_storage_bytes": comp_bytes,
                "storage_compression_ratio": round(storage_ratio, 3),
                "storage_reduction_pct": round(100.0 * (1.0 - comp_bytes / max(1, orig_bytes)), 2),
                "method": "tensor_train_with_svd_fallback",
            }
        )
        logger.info(
            "TT compressor: %d TT + %d SVD layers, %d fallback, avg_err=%.4f, storage_ratio=%.2fx",
            len(self.tt_layers),
            len(self.svd_layers),
            len(fallback_layers),
            avg_err,
            storage_ratio,
        )
        return self.model, stats

    # ------------------------------------------------------------------
    def compress_linear(self, linear_layer: nn.Module) -> nn.Module | None:
        """Compress a single layer to TT (or SVD if TT fails)."""
        tt = TTLinearLayer.from_linear(
            linear_layer,
            tt_rank=self.tt_rank,
            num_factors=self.num_factors,
            max_error=self.max_error,
            energy_threshold=self.energy_threshold,
            strict_rank=self.strict_rank,
        )
        if tt is not None:
            return tt
        if self.fallback_policy in ("svd", "original"):
            return _try_svd_fallback(linear_layer, self.max_error)
        return None

Replaces Linear/Conv1D with TTLinearLayer when TT succeeds; otherwise falls back to SVDLinearLayer. Layers that neither can compress within error budget remain untouched.

Args

-----=

tt_rank

  TT rank bound.

num_factors

Number of TT factors per dimension (default 2).

max_error

Relative Frobenius error threshold (default 0.20).

energy_threshold

Energy retained per TT core (default 0.995).

fallback_policy

'svd' | 'original' | 'none'. Default: 'svd'.

target_attention_only

Only compress attention projections.

strict_rank

Do not exceed tt_rank.

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Methods

def compress_linear(self, linear_layer: nn.Module) ‑> torch.nn.modules.module.Module | None

Expand source code

def compress_linear(self, linear_layer: nn.Module) -> nn.Module | None:
    """Compress a single layer to TT (or SVD if TT fails)."""
    tt = TTLinearLayer.from_linear(
        linear_layer,
        tt_rank=self.tt_rank,
        num_factors=self.num_factors,
        max_error=self.max_error,
        energy_threshold=self.energy_threshold,
        strict_rank=self.strict_rank,
    )
    if tt is not None:
        return tt
    if self.fallback_policy in ("svd", "original"):
        return _try_svd_fallback(linear_layer, self.max_error)
    return None

Compress a single layer to TT (or SVD if TT fails).

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class TTLinear (in_features: int, out_features: int, tt_rank: int = 16, bias: bool = True, shape: tuple[tuple[int, ...], tuple[int, ...]] | None = None, dtype: torch.dtype | None = None, device: torch.device | None = None)

Expand source code

class TTLinearLayer(nn.Module):
    """
    Tensor-Train (TT) linear layer.

    Weight W ∈ R^{out × in} is reshaped to a multi-dim tensor of shape
    (q_0, q_1, ..., p_0, p_1, ...) with q_dims × p_dims = out × in, then
    factorized into TT cores G_1 ... G_d of shape (r_{k-1}, p_k, q_k, r_k),
    with boundary ranks r_0 = r_d = 1.

    Forward contracts TT cores with x directly, without materializing W
    for d == 2; for larger d, an optimized generic path is used.
    """

    # TT-FIX-01
    is_custom_attention: bool = False

    @staticmethod
    def _max_rank_for_compression(
        out_features: int,
        in_features: int,
        p_dims: tuple[int, ...],
        q_dims: tuple[int, ...],
        safety_margin: float = 0.9,
    ) -> int:
        """
        Max TT-rank so TT parameter count stays below dense count.
        Param count ≈ r · Σ_i (p_i q_i) for uniform rank.
        """
        if not p_dims or not q_dims or len(p_dims) != len(q_dims):
            return 1
        sum_pq = sum(p * q for p, q in zip(p_dims, q_dims))
        if sum_pq <= 0:
            return 1
        breakeven = (out_features * in_features) / sum_pq
        return max(1, int(breakeven * safety_margin))

    def __init__(
        self,
        in_features: int,
        out_features: int,
        tt_rank: int = 16,
        bias: bool = True,
        shape: tuple[tuple[int, ...], tuple[int, ...]] | None = None,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features
        self.nx = in_features
        self.tt_rank = int(tt_rank)

        if shape is None:
            self.p_dims, self.q_dims = factorize_dimension_pair(in_features, out_features, 2)
        else:
            self.p_dims, self.q_dims = shape
            if (
                int(np.prod(self.p_dims)) != in_features
                or int(np.prod(self.q_dims)) != out_features
            ):
                raise ValueError(
                    f"Shape mismatch: p={self.p_dims} prod={np.prod(self.p_dims)} in={in_features}, "
                    f"q={self.q_dims} prod={np.prod(self.q_dims)} out={out_features}"
                )

        self.d = len(self.p_dims)
        self.cores = nn.ParameterList()

        r_left = 1
        for i in range(self.d - 1):
            p_i, q_i = self.p_dims[i], self.q_dims[i]
            r_right = min(self.tt_rank, max(1, p_i * q_i * r_left))
            core = nn.Parameter(
                torch.empty(r_left, p_i, q_i, r_right, dtype=dtype, device=device),
                requires_grad=False,
            )
            bound = 1.0 / math.sqrt(max(1, r_left * p_i))
            nn.init.uniform_(core, -bound, bound)
            self.cores.append(core)
            r_left = r_right

        core_last = nn.Parameter(
            torch.empty(r_left, self.p_dims[-1], self.q_dims[-1], 1, dtype=dtype, device=device),
            requires_grad=False,
        )
        nn.init.uniform_(
            core_last,
            -1.0 / math.sqrt(max(1, r_left * self.p_dims[-1])),
            1.0 / math.sqrt(max(1, r_left * self.p_dims[-1])),
        )
        self.cores.append(core_last)

        if bias:
            self.bias = nn.Parameter(
                torch.zeros(out_features, dtype=dtype, device=device),
                requires_grad=False,
            )
        else:
            self.register_parameter("bias", None)

        self._cached_weight: torch.Tensor | None = None
        self._cached_version: int = -1

    # ---- cache management (TT-FIX-06) ----
    def _invalidate_cache(self) -> None:
        self._cached_weight = None
        self._cached_version = -1

    def _current_version(self) -> int:
        # Sum of _version counters for all cores + bias. Invalidates on any
        # in-place mutation of any parameter.
        v = sum(getattr(c, "_version", 0) for c in self.cores)
        if self.bias is not None:
            v += getattr(self.bias, "_version", 0)
        return v

    def materialize_weight(self) -> torch.Tensor:
        """Reconstruct W ∈ R^{out × in} by contracting TT-cores."""
        if (
            not self.training
            and self._cached_weight is not None
            and self._cached_version == self._current_version()
        ):
            return self._cached_weight

        res = self.cores[0][0]  # (p0, q0, r1)
        for i in range(1, self.d):
            core = self.cores[i]  # (r_i, p_i, q_i, r_{i+1})
            res = torch.tensordot(res, core, dims=([-1], [0]))
        res = res[..., 0]  # drop trailing r=1

        q_idx = [2 * k + 1 for k in range(self.d)]
        p_idx = [2 * k for k in range(self.d)]
        res = res.permute(*(q_idx + p_idx)).contiguous()
        W = res.view(self.out_features, self.in_features)

        if not self.training:
            self._cached_weight = W
            self._cached_version = self._current_version()
        return W

    @property
    def weight(self) -> torch.Tensor:
        return self.materialize_weight()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)

        # TT-FIX-09: generic path always works.
        W = self.materialize_weight()
        if W.dtype != x_flat.dtype or W.device != x_flat.device:
            W = W.to(device=x_flat.device, dtype=x_flat.dtype)
        out = F.linear(x_flat, W)

        if self.bias is not None:
            out = out + self.bias.to(out.device, out.dtype)
        return out.view(*orig_shape[:-1], self.out_features)

    # ---- Adaptive TT-SVD (TT-FIX-04) ----
    @classmethod
    def from_linear(
        cls,
        linear_layer: nn.Module,
        tt_rank: int = 16,
        num_factors: int = 2,
        max_error: float = 0.20,
        max_retries: int = 3,
        energy_threshold: float = 0.995,
        strict_rank: bool = False,
    ) -> TTLinearLayer | None:
        """
        Factorize a pretrained Linear/Conv1D into TT-cores with adaptive rank.

        Rank per core is chosen adaptively: the smallest r such that the
        retained energy ≥ energy_threshold, capped by tt_rank and by the
        max_rank_for_compression budget.

        Returns None if the layer cannot be compressed within the error
        threshold and parameter budget.
        """
        base_layer = getattr(linear_layer, "base_layer", linear_layer)
        raw_w = base_layer.weight.data
        orig_dtype = raw_w.dtype
        orig_device = raw_w.device

        cls_name = base_layer.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(base_layer, "nf")

        if hasattr(base_layer, "nx") and hasattr(base_layer, "nf"):
            in_f = int(base_layer.nx)
            out_f = int(base_layer.nf)
            W = raw_w.t().contiguous().float()
        elif hasattr(base_layer, "in_features") and hasattr(base_layer, "out_features"):
            in_f = int(base_layer.in_features)
            out_f = int(base_layer.out_features)
            W = (
                raw_w.t().contiguous().float()
                if (raw_w.shape == (in_f, out_f))
                else raw_w.float().contiguous()
            )
        elif is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float().contiguous()

        norm_orig = W.norm().item() + 1e-12

        p_dims, q_dims = factorize_dimension_pair(in_f, out_f, num_factors=num_factors)
        if int(np.prod(p_dims)) != in_f or int(np.prod(q_dims)) != out_f:
            return None

        d = len(p_dims)
        # Reshape to (q_0, q_1, ..., p_0, p_1, ...)
        W_tensor = W.view(*(list(q_dims) + list(p_dims)))
        perm = []
        for k in range(d):
            perm.append(d + k)  # p_k
            perm.append(k)  # q_k
        W_interleaved = W_tensor.permute(*perm).contiguous()

        max_rank = cls._max_rank_for_compression(out_f, in_f, p_dims, q_dims, safety_margin=0.9)
        rank_cap = min(int(tt_rank), max_rank) if not strict_rank else min(int(tt_rank), max_rank)

        bias = base_layer.bias.data if getattr(base_layer, "bias", None) is not None else None

        # ---- Build the layer with adaptive TT-SVD ----
        layer = cls(
            in_features=in_f,
            out_features=out_f,
            tt_rank=rank_cap,
            bias=bias is not None,
            shape=(p_dims, q_dims),
            dtype=orig_dtype,
            device=orig_device,
        )

        try:
            curr_matrix = W_interleaved.view(p_dims[0] * q_dims[0], -1)
            r_left = 1
            for i in range(d - 1):
                p_i, q_i = p_dims[i], q_dims[i]
                rows = r_left * p_i * q_i
                curr_matrix = curr_matrix.view(rows, -1)

                U, S, Vt = torch.linalg.svd(curr_matrix, full_matrices=False)
                # Adaptive rank: smallest r with retained energy >= threshold
                S_sq = S**2
                total_energy = S_sq.sum().clamp_min(1e-12)
                cum = torch.cumsum(S_sq, dim=0) / total_energy
                r_energy = int((cum >= energy_threshold).nonzero()[0].item()) + 1
                # Cap by: configured tt_rank, max_rank budget, and the shape bound
                r_right = min(r_energy, rank_cap, U.shape[1], S.numel())
                r_right = max(1, r_right)

                U_r = U[:, :r_right]
                S_r = S[:r_right]
                Vt_r = Vt[:r_right, :]

                core_data = U_r.view(r_left, p_i, q_i, r_right)
                with torch.no_grad():
                    layer.cores[i].data.zero_()
                    layer.cores[i].data.copy_(core_data.to(dtype=orig_dtype, device=orig_device))

                curr_matrix = torch.diag(S_r) @ Vt_r
                r_left = r_right

            p_last, q_last = p_dims[-1], q_dims[-1]
            final_core = curr_matrix.view(r_left, p_last, q_last, 1)
            with torch.no_grad():
                layer.cores[-1].data.zero_()
                layer.cores[-1].data.copy_(final_core.to(dtype=orig_dtype, device=orig_device))

            if bias is not None and layer.bias is not None:
                with torch.no_grad():
                    layer.bias.data.copy_(bias.to(dtype=orig_dtype, device=orig_device))

            # ---- Error check ----
            W_rec = layer.materialize_weight().float()
            rel_err = (torch.norm(W - W_rec) / norm_orig).item()

            # TT-FIX-05: fall back to SVD if TT error is too high.
            if rel_err > max_error:
                logger.info(
                    "TT: rel_err=%.4f > max_error=%.4f, trying SVD fallback",
                    rel_err,
                    max_error,
                )
                return None

            # Check parameter count
            tt_params = sum(p.numel() for p in layer.parameters())
            if tt_params >= out_f * in_f:
                logger.info(
                    "TT: no parameter reduction (%d >= %d), aborting",
                    tt_params,
                    out_f * in_f,
                )
                return None

            layer.reconstruction_error = float(rel_err)
            layer.tt_rank_actual = r_left
            return layer

        except Exception as e:
            logger.warning("TT-SVD failed: %s", e)
            return None

Tensor-Train (TT) linear layer.

Weight W ∈ R^{out × in} is reshaped to a multi-dim tensor of shape (q_0, q_1, …, p_0, p_1, …) with q_dims × p_dims = out × in, then factorized into TT cores G_1 … G_d of shape (r_{k-1}, p_k, q_k, r_k), with boundary ranks r_0 = r_d = 1.

Forward contracts TT cores with x directly, without materializing W for d == 2; for larger d, an optimized generic path is used.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Static methods

def from_linear(linear_layer: nn.Module, tt_rank: int = 16, num_factors: int = 2, max_error: float = 0.2, max_retries: int = 3, energy_threshold: float = 0.995, strict_rank: bool = False) ‑> TTLinearLayer | None: Factorize a pretrained Linear/Conv1D into TT-cores with adaptive rank.

Rank per core is chosen adaptively: the smallest r such that the retained energy ≥ energy_threshold, capped by tt_rank and by the max_rank_for_compression budget.

Returns None if the layer cannot be compressed within the error threshold and parameter budget.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    return self.materialize_weight()

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)

    # TT-FIX-09: generic path always works.
    W = self.materialize_weight()
    if W.dtype != x_flat.dtype or W.device != x_flat.device:
        W = W.to(device=x_flat.device, dtype=x_flat.dtype)
    out = F.linear(x_flat, W)

    if self.bias is not None:
        out = out + self.bias.to(out.device, out.dtype)
    return out.view(*orig_shape[:-1], self.out_features)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def materialize_weight(self) ‑> torch.Tensor

Expand source code

def materialize_weight(self) -> torch.Tensor:
    """Reconstruct W ∈ R^{out × in} by contracting TT-cores."""
    if (
        not self.training
        and self._cached_weight is not None
        and self._cached_version == self._current_version()
    ):
        return self._cached_weight

    res = self.cores[0][0]  # (p0, q0, r1)
    for i in range(1, self.d):
        core = self.cores[i]  # (r_i, p_i, q_i, r_{i+1})
        res = torch.tensordot(res, core, dims=([-1], [0]))
    res = res[..., 0]  # drop trailing r=1

    q_idx = [2 * k + 1 for k in range(self.d)]
    p_idx = [2 * k for k in range(self.d)]
    res = res.permute(*(q_idx + p_idx)).contiguous()
    W = res.view(self.out_features, self.in_features)

    if not self.training:
        self._cached_weight = W
        self._cached_version = self._current_version()
    return W

Reconstruct W ∈ R^{out × in} by contracting TT-cores.

class TTLinearLayer (in_features: int, out_features: int, tt_rank: int = 16, bias: bool = True, shape: tuple[tuple[int, ...], tuple[int, ...]] | None = None, dtype: torch.dtype | None = None, device: torch.device | None = None)

Expand source code

class TTLinearLayer(nn.Module):
    """
    Tensor-Train (TT) linear layer.

    Weight W ∈ R^{out × in} is reshaped to a multi-dim tensor of shape
    (q_0, q_1, ..., p_0, p_1, ...) with q_dims × p_dims = out × in, then
    factorized into TT cores G_1 ... G_d of shape (r_{k-1}, p_k, q_k, r_k),
    with boundary ranks r_0 = r_d = 1.

    Forward contracts TT cores with x directly, without materializing W
    for d == 2; for larger d, an optimized generic path is used.
    """

    # TT-FIX-01
    is_custom_attention: bool = False

    @staticmethod
    def _max_rank_for_compression(
        out_features: int,
        in_features: int,
        p_dims: tuple[int, ...],
        q_dims: tuple[int, ...],
        safety_margin: float = 0.9,
    ) -> int:
        """
        Max TT-rank so TT parameter count stays below dense count.
        Param count ≈ r · Σ_i (p_i q_i) for uniform rank.
        """
        if not p_dims or not q_dims or len(p_dims) != len(q_dims):
            return 1
        sum_pq = sum(p * q for p, q in zip(p_dims, q_dims))
        if sum_pq <= 0:
            return 1
        breakeven = (out_features * in_features) / sum_pq
        return max(1, int(breakeven * safety_margin))

    def __init__(
        self,
        in_features: int,
        out_features: int,
        tt_rank: int = 16,
        bias: bool = True,
        shape: tuple[tuple[int, ...], tuple[int, ...]] | None = None,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.nf = out_features
        self.nx = in_features
        self.tt_rank = int(tt_rank)

        if shape is None:
            self.p_dims, self.q_dims = factorize_dimension_pair(in_features, out_features, 2)
        else:
            self.p_dims, self.q_dims = shape
            if (
                int(np.prod(self.p_dims)) != in_features
                or int(np.prod(self.q_dims)) != out_features
            ):
                raise ValueError(
                    f"Shape mismatch: p={self.p_dims} prod={np.prod(self.p_dims)} in={in_features}, "
                    f"q={self.q_dims} prod={np.prod(self.q_dims)} out={out_features}"
                )

        self.d = len(self.p_dims)
        self.cores = nn.ParameterList()

        r_left = 1
        for i in range(self.d - 1):
            p_i, q_i = self.p_dims[i], self.q_dims[i]
            r_right = min(self.tt_rank, max(1, p_i * q_i * r_left))
            core = nn.Parameter(
                torch.empty(r_left, p_i, q_i, r_right, dtype=dtype, device=device),
                requires_grad=False,
            )
            bound = 1.0 / math.sqrt(max(1, r_left * p_i))
            nn.init.uniform_(core, -bound, bound)
            self.cores.append(core)
            r_left = r_right

        core_last = nn.Parameter(
            torch.empty(r_left, self.p_dims[-1], self.q_dims[-1], 1, dtype=dtype, device=device),
            requires_grad=False,
        )
        nn.init.uniform_(
            core_last,
            -1.0 / math.sqrt(max(1, r_left * self.p_dims[-1])),
            1.0 / math.sqrt(max(1, r_left * self.p_dims[-1])),
        )
        self.cores.append(core_last)

        if bias:
            self.bias = nn.Parameter(
                torch.zeros(out_features, dtype=dtype, device=device),
                requires_grad=False,
            )
        else:
            self.register_parameter("bias", None)

        self._cached_weight: torch.Tensor | None = None
        self._cached_version: int = -1

    # ---- cache management (TT-FIX-06) ----
    def _invalidate_cache(self) -> None:
        self._cached_weight = None
        self._cached_version = -1

    def _current_version(self) -> int:
        # Sum of _version counters for all cores + bias. Invalidates on any
        # in-place mutation of any parameter.
        v = sum(getattr(c, "_version", 0) for c in self.cores)
        if self.bias is not None:
            v += getattr(self.bias, "_version", 0)
        return v

    def materialize_weight(self) -> torch.Tensor:
        """Reconstruct W ∈ R^{out × in} by contracting TT-cores."""
        if (
            not self.training
            and self._cached_weight is not None
            and self._cached_version == self._current_version()
        ):
            return self._cached_weight

        res = self.cores[0][0]  # (p0, q0, r1)
        for i in range(1, self.d):
            core = self.cores[i]  # (r_i, p_i, q_i, r_{i+1})
            res = torch.tensordot(res, core, dims=([-1], [0]))
        res = res[..., 0]  # drop trailing r=1

        q_idx = [2 * k + 1 for k in range(self.d)]
        p_idx = [2 * k for k in range(self.d)]
        res = res.permute(*(q_idx + p_idx)).contiguous()
        W = res.view(self.out_features, self.in_features)

        if not self.training:
            self._cached_weight = W
            self._cached_version = self._current_version()
        return W

    @property
    def weight(self) -> torch.Tensor:
        return self.materialize_weight()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        x_flat = x.reshape(-1, self.in_features)

        # TT-FIX-09: generic path always works.
        W = self.materialize_weight()
        if W.dtype != x_flat.dtype or W.device != x_flat.device:
            W = W.to(device=x_flat.device, dtype=x_flat.dtype)
        out = F.linear(x_flat, W)

        if self.bias is not None:
            out = out + self.bias.to(out.device, out.dtype)
        return out.view(*orig_shape[:-1], self.out_features)

    # ---- Adaptive TT-SVD (TT-FIX-04) ----
    @classmethod
    def from_linear(
        cls,
        linear_layer: nn.Module,
        tt_rank: int = 16,
        num_factors: int = 2,
        max_error: float = 0.20,
        max_retries: int = 3,
        energy_threshold: float = 0.995,
        strict_rank: bool = False,
    ) -> TTLinearLayer | None:
        """
        Factorize a pretrained Linear/Conv1D into TT-cores with adaptive rank.

        Rank per core is chosen adaptively: the smallest r such that the
        retained energy ≥ energy_threshold, capped by tt_rank and by the
        max_rank_for_compression budget.

        Returns None if the layer cannot be compressed within the error
        threshold and parameter budget.
        """
        base_layer = getattr(linear_layer, "base_layer", linear_layer)
        raw_w = base_layer.weight.data
        orig_dtype = raw_w.dtype
        orig_device = raw_w.device

        cls_name = base_layer.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(base_layer, "nf")

        if hasattr(base_layer, "nx") and hasattr(base_layer, "nf"):
            in_f = int(base_layer.nx)
            out_f = int(base_layer.nf)
            W = raw_w.t().contiguous().float()
        elif hasattr(base_layer, "in_features") and hasattr(base_layer, "out_features"):
            in_f = int(base_layer.in_features)
            out_f = int(base_layer.out_features)
            W = (
                raw_w.t().contiguous().float()
                if (raw_w.shape == (in_f, out_f))
                else raw_w.float().contiguous()
            )
        elif is_conv1d:
            in_f, out_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.t().contiguous().float()
        else:
            out_f, in_f = raw_w.shape[0], raw_w.shape[1]
            W = raw_w.float().contiguous()

        norm_orig = W.norm().item() + 1e-12

        p_dims, q_dims = factorize_dimension_pair(in_f, out_f, num_factors=num_factors)
        if int(np.prod(p_dims)) != in_f or int(np.prod(q_dims)) != out_f:
            return None

        d = len(p_dims)
        # Reshape to (q_0, q_1, ..., p_0, p_1, ...)
        W_tensor = W.view(*(list(q_dims) + list(p_dims)))
        perm = []
        for k in range(d):
            perm.append(d + k)  # p_k
            perm.append(k)  # q_k
        W_interleaved = W_tensor.permute(*perm).contiguous()

        max_rank = cls._max_rank_for_compression(out_f, in_f, p_dims, q_dims, safety_margin=0.9)
        rank_cap = min(int(tt_rank), max_rank) if not strict_rank else min(int(tt_rank), max_rank)

        bias = base_layer.bias.data if getattr(base_layer, "bias", None) is not None else None

        # ---- Build the layer with adaptive TT-SVD ----
        layer = cls(
            in_features=in_f,
            out_features=out_f,
            tt_rank=rank_cap,
            bias=bias is not None,
            shape=(p_dims, q_dims),
            dtype=orig_dtype,
            device=orig_device,
        )

        try:
            curr_matrix = W_interleaved.view(p_dims[0] * q_dims[0], -1)
            r_left = 1
            for i in range(d - 1):
                p_i, q_i = p_dims[i], q_dims[i]
                rows = r_left * p_i * q_i
                curr_matrix = curr_matrix.view(rows, -1)

                U, S, Vt = torch.linalg.svd(curr_matrix, full_matrices=False)
                # Adaptive rank: smallest r with retained energy >= threshold
                S_sq = S**2
                total_energy = S_sq.sum().clamp_min(1e-12)
                cum = torch.cumsum(S_sq, dim=0) / total_energy
                r_energy = int((cum >= energy_threshold).nonzero()[0].item()) + 1
                # Cap by: configured tt_rank, max_rank budget, and the shape bound
                r_right = min(r_energy, rank_cap, U.shape[1], S.numel())
                r_right = max(1, r_right)

                U_r = U[:, :r_right]
                S_r = S[:r_right]
                Vt_r = Vt[:r_right, :]

                core_data = U_r.view(r_left, p_i, q_i, r_right)
                with torch.no_grad():
                    layer.cores[i].data.zero_()
                    layer.cores[i].data.copy_(core_data.to(dtype=orig_dtype, device=orig_device))

                curr_matrix = torch.diag(S_r) @ Vt_r
                r_left = r_right

            p_last, q_last = p_dims[-1], q_dims[-1]
            final_core = curr_matrix.view(r_left, p_last, q_last, 1)
            with torch.no_grad():
                layer.cores[-1].data.zero_()
                layer.cores[-1].data.copy_(final_core.to(dtype=orig_dtype, device=orig_device))

            if bias is not None and layer.bias is not None:
                with torch.no_grad():
                    layer.bias.data.copy_(bias.to(dtype=orig_dtype, device=orig_device))

            # ---- Error check ----
            W_rec = layer.materialize_weight().float()
            rel_err = (torch.norm(W - W_rec) / norm_orig).item()

            # TT-FIX-05: fall back to SVD if TT error is too high.
            if rel_err > max_error:
                logger.info(
                    "TT: rel_err=%.4f > max_error=%.4f, trying SVD fallback",
                    rel_err,
                    max_error,
                )
                return None

            # Check parameter count
            tt_params = sum(p.numel() for p in layer.parameters())
            if tt_params >= out_f * in_f:
                logger.info(
                    "TT: no parameter reduction (%d >= %d), aborting",
                    tt_params,
                    out_f * in_f,
                )
                return None

            layer.reconstruction_error = float(rel_err)
            layer.tt_rank_actual = r_left
            return layer

        except Exception as e:
            logger.warning("TT-SVD failed: %s", e)
            return None

Tensor-Train (TT) linear layer.

Forward contracts TT cores with x directly, without materializing W for d == 2; for larger d, an optimized generic path is used.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Static methods

def from_linear(linear_layer: nn.Module, tt_rank: int = 16, num_factors: int = 2, max_error: float = 0.2, max_retries: int = 3, energy_threshold: float = 0.995, strict_rank: bool = False) ‑> TTLinearLayer | None: Factorize a pretrained Linear/Conv1D into TT-cores with adaptive rank.

Rank per core is chosen adaptively: the smallest r such that the retained energy ≥ energy_threshold, capped by tt_rank and by the max_rank_for_compression budget.

Returns None if the layer cannot be compressed within the error threshold and parameter budget.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    return self.materialize_weight()

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_shape = x.shape
    x_flat = x.reshape(-1, self.in_features)

    # TT-FIX-09: generic path always works.
    W = self.materialize_weight()
    if W.dtype != x_flat.dtype or W.device != x_flat.device:
        W = W.to(device=x_flat.device, dtype=x_flat.dtype)
    out = F.linear(x_flat, W)

    if self.bias is not None:
        out = out + self.bias.to(out.device, out.dtype)
    return out.view(*orig_shape[:-1], self.out_features)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

def materialize_weight(self) ‑> torch.Tensor

Expand source code

def materialize_weight(self) -> torch.Tensor:
    """Reconstruct W ∈ R^{out × in} by contracting TT-cores."""
    if (
        not self.training
        and self._cached_weight is not None
        and self._cached_version == self._current_version()
    ):
        return self._cached_weight

    res = self.cores[0][0]  # (p0, q0, r1)
    for i in range(1, self.d):
        core = self.cores[i]  # (r_i, p_i, q_i, r_{i+1})
        res = torch.tensordot(res, core, dims=([-1], [0]))
    res = res[..., 0]  # drop trailing r=1

    q_idx = [2 * k + 1 for k in range(self.d)]
    p_idx = [2 * k for k in range(self.d)]
    res = res.permute(*(q_idx + p_idx)).contiguous()
    W = res.view(self.out_features, self.in_features)

    if not self.training:
        self._cached_weight = W
        self._cached_version = self._current_version()
    return W

Reconstruct W ∈ R^{out × in} by contracting TT-cores.

class ToMoECompressor (model: nn.Module, num_experts: int = 4, top_k: int = 2, max_error: float = 0.2, lora_r: int = 0, lora_alpha: float = 16.0, freeze_experts: bool = True, target_attention_only: bool = False, target_modules: list[str] | None = None, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("tomoe_class")
@COMPRESSOR_REGISTRY.register("tomoe_compressor_class")
class ToMoECompressor(BaseCompressor):
    """
    ToMoE Compressor: dynamically replaces dense feed-forward and projection layers
    with an MoE-fied equivalent.

    Args:
        model: PyTorch model.
        num_experts: Number of expert sub-layers to form (default: 4).
        top_k: Number of experts to route each token to (default: 2).
        max_error: Maximum allowed relative Frobenius error on reconstruction (default: 0.20).
        lora_r: LoRA adapter rank on experts and router (0 = disabled).
        lora_alpha: LoRA scaling factor (default: 16.0).
        freeze_experts: Whether to freeze expert base weights (default: True).
        target_attention_only: Whether to restrict to attention layers only.
        target_modules: Optional list of target layer name substrings.
    """

    # Fix 6: Protection of Critical Layers
    PROTECTED_KEYWORDS = [
        "classifier",
        "score",
        "lm_head",
        "pooler",
        "embed",
        "wte",
        "wpe",
        "shared",
        "head",
        "logits",
        "pre_classifier",
    ]

    def __init__(
        self,
        model: nn.Module,
        num_experts: int = 4,
        top_k: int = 2,
        max_error: float = 0.20,
        lora_r: int = 0,
        lora_alpha: float = 16.0,
        freeze_experts: bool = True,
        target_attention_only: bool = False,
        target_modules: list[str] | None = None,
        **kwargs: Any,
    ):
        super().__init__(model, name="tomoe", **kwargs)
        cfg = kwargs.pop("comp_config", kwargs.pop("cfg", None))
        if cfg is not None:
            num_experts = _cfg_get(
                cfg, "num_experts", _cfg_get(cfg, "tomoe_num_experts", num_experts)
            )
            top_k = _cfg_get(cfg, "top_k", _cfg_get(cfg, "tomoe_top_k", top_k))
            max_error = _cfg_get(
                cfg, "max_error", _cfg_get(cfg, "max_reconstruction_error", max_error)
            )
            lora_r = _cfg_get(cfg, "lora_r", _cfg_get(cfg, "tomoe_lora_r", lora_r))
            lora_alpha = _cfg_get(cfg, "lora_alpha", lora_alpha)
            freeze_experts = _cfg_get(cfg, "freeze_experts", freeze_experts)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)

        num_experts = int(kwargs.pop("num_experts", num_experts))
        top_k = int(kwargs.pop("top_k", top_k))
        max_error = float(kwargs.pop("max_error", max_error))
        lora_r = int(kwargs.pop("lora_r", lora_r))
        lora_alpha = float(kwargs.pop("lora_alpha", lora_alpha))
        freeze_experts = bool(kwargs.pop("freeze_experts", freeze_experts))
        target_attention_only = bool(kwargs.pop("target_attention_only", target_attention_only))

        self.num_experts = max(1, num_experts)
        self.top_k = max(1, min(top_k, self.num_experts))
        self.max_error = float(max_error)
        self.lora_r = max(0, lora_r)
        self.lora_alpha = float(lora_alpha)
        self.freeze_experts = bool(freeze_experts)
        self.target_attention_only = target_attention_only
        self.target_modules = target_modules
        self.stats: dict[str, Any] = {}

    def _is_protected(self, name: str) -> bool:
        low = name.lower()
        return any(kw in low for kw in self.PROTECTED_KEYWORDS)

    def _is_target(self, name: str) -> bool:
        if self._is_protected(name):
            return False
        if not should_compress_linear(self.model, name):
            return False
        # Exclude fused QKV (c_attn, in_proj_qkv) to prevent mixing semantic groups across heads in clustering
        if "c_attn" in name.lower() or "in_proj" in name.lower():
            return False
        if self.target_modules is not None:
            return any(t in name for t in self.target_modules)
        low = name.lower()
        if self.target_attention_only:
            attn_keywords = [
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_proj",
                "query",
                "key",
                "value",
            ]
            return any(k in low for k in attn_keywords)
        # Default: target FFN, MLP, and linear projections across all supported families
        return any(
            t in low
            for t in [
                "mlp",
                "ffn",
                "intermediate",
                "output.dense",
                "c_fc",
                "c_proj",
                "lin1",
                "lin2",
                "gate_proj",
                "up_proj",
                "down_proj",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "q_lin",
                "k_lin",
                "v_lin",
                "out_lin",
                "query",
                "key",
                "value",
            ]
        )

    def _get_parent(self, full_name: str) -> tuple[nn.Module | None, str]:
        parts = full_name.split(".")
        parent = self.model
        for part in parts[:-1]:
            parent = getattr(parent, part, None)
            if parent is None:
                return None, ""
        return parent, parts[-1]

    def compress(self, orig_model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
        """
        Compress target layers into ToMoE Mixture of Experts layers with full error validation.
        """
        if orig_model is not None:
            self.model = orig_model

        orig_total_params = sum(p.numel() for p in self.model.parameters())
        converted_layers = []
        skipped_layers = []
        layer_details = {}

        seen_weight_ptrs: set[int] = set()

        for name, module in list(self.model.named_modules()):
            if not self._is_target(name):
                continue

            # Fix 2: Check for supported layer type (nn.Linear or Conv1D) on base module
            base_mod = getattr(module, "base_layer", module)
            cls_name = base_mod.__class__.__name__.lower()
            is_linear = isinstance(base_mod, nn.Linear)
            is_conv1d = "conv1d" in cls_name or (
                hasattr(base_mod, "nx") and hasattr(base_mod, "nf")
            )

            if not (
                is_linear
                or is_conv1d
                or (hasattr(base_mod, "weight") and base_mod.weight.ndim == 2)
            ):
                logger.warning(
                    "ToMoE: Skipping unsupported layer type %s for %s",
                    type(base_mod).__name__,
                    name,
                )
                skipped_layers.append((name, f"unsupported_type_{type(base_mod).__name__}"))
                continue

            # H20: Detect and protect tied weights (e.g. wte and lm_head sharing same storage)
            if hasattr(base_mod, "weight") and base_mod.weight is not None:
                w_ptr = base_mod.weight.data_ptr()
                if w_ptr in seen_weight_ptrs:
                    logger.warning(
                        "ToMoE: Skipping tied module %s with shared weight storage (H20)",
                        name,
                    )
                    skipped_layers.append((name, "tied_weight_shared_storage"))
                    continue
                seen_weight_ptrs.add(w_ptr)

            # Determine weights and feature shapes ensuring orientation is (out_features, in_features)
            if is_conv1d:
                in_f = base_mod.weight.shape[0]
                out_f = base_mod.weight.shape[1]
                orig_w_2d = base_mod.weight.data.t().contiguous()
            else:
                in_f = getattr(base_mod, "in_features", base_mod.weight.shape[1])
                out_f = getattr(base_mod, "out_features", base_mod.weight.shape[0])
                orig_w_2d = base_mod.weight.data.contiguous()

            if torch.linalg.norm(orig_w_2d.float()) < 1e-4:
                skipped_layers.append((name, "frobenius_norm_too_small"))
                continue

            if out_f < 2:
                skipped_layers.append((name, f"out_features_too_small_{out_f}"))
                continue

            # Respect GQA dimensions (e.g. Qwen k_proj/v_proj with fewer heads)
            eff_num_experts = max(1, min(self.num_experts, out_f))
            eff_top_k = max(1, min(self.top_k, eff_num_experts))

            if eff_num_experts < 2:
                skipped_layers.append((name, "effective_experts_less_than_2"))
                continue

            parent, child_name = self._get_parent(name)
            if parent is None:
                skipped_layers.append((name, "parent_not_found"))
                continue

            # Construct ToMoELinear
            try:
                moe_layer = ToMoELinear(
                    original=module,
                    num_experts=eff_num_experts,
                    top_k=eff_top_k,
                    lora_r=self.lora_r,
                    lora_alpha=self.lora_alpha,
                    freeze_experts=self.freeze_experts,
                    max_error=self.max_error,
                )
            except Exception as e:
                logger.warning("ToMoE: Conversion failed for %s: %s", name, e)
                skipped_layers.append((name, f"construction_error_{e}"))
                continue

            # Meaningful fidelity check: activation reconstruction error on a calibration batch
            with torch.no_grad():
                device = orig_w_2d.device
                dtype = orig_w_2d.dtype
                calib_x = torch.randn(4, 16, in_f, device=device, dtype=dtype)
                orig_out = module(calib_x)
                moe_out = moe_layer(calib_x)
                err = (
                    torch.norm(moe_out.float() - orig_out.float())
                    / (torch.norm(orig_out.float()) + 1e-8)
                ).item()

            if err > self.max_error or np.isnan(err):
                logger.warning(
                    "ToMoE: %s activation reconstruction error (%.4f) exceeded max_error (%.4f). Keeping original layer.",
                    name,
                    err,
                    self.max_error,
                )
                skipped_layers.append((name, f"error_threshold_exceeded_{err:.4f}"))
                for p in module.parameters():
                    p.requires_grad = False
                continue

            # Replace layer in parent module preserving outer LoRALinear wrapper if present
            orig_p_count = sum(p.numel() for p in module.parameters())
            new_p_count = sum(p.numel() for p in moe_layer.parameters())
            if hasattr(module, "base_layer"):
                module.base_layer = moe_layer
                if hasattr(module, "lora_A") and module.lora_A is not None:
                    for p in (
                        module.lora_A.parameters()
                        if hasattr(module.lora_A, "parameters")
                        else [module.lora_A]
                    ):
                        p.requires_grad = True
                if hasattr(module, "lora_B") and module.lora_B is not None:
                    for p in (
                        module.lora_B.parameters()
                        if hasattr(module.lora_B, "parameters")
                        else [module.lora_B]
                    ):
                        p.requires_grad = True
            else:
                setattr(parent, child_name, moe_layer)

            converted_layers.append(name)
            layer_details[name] = {
                "num_experts": eff_num_experts,
                "top_k": eff_top_k,
                "reconstruction_error": err,
                "original_params": orig_p_count,
                "new_params": new_p_count,
                "equal_expert_size": moe_layer.equal_expert_size,
            }
            logger.info(
                "ToMoE: Converted %s -> ToMoELinear (experts=%d, top_k=%d, reconstruction_error=%.4e)",
                name,
                eff_num_experts,
                eff_top_k,
                err,
            )

        new_total_params = sum(p.numel() for p in self.model.parameters())
        self._compressed_params = new_total_params

        # Fix 5: Per-Layer Stats and Logging
        self.stats = {
            "converted_layers": converted_layers,
            "num_converted_layers": len(converted_layers),
            "skipped_layers": skipped_layers,
            "original_params": orig_total_params,
            "compressed_params": new_total_params,
            "compression_ratio": orig_total_params / max(1, new_total_params),
            "layer_details": layer_details,
            "num_experts": self.num_experts,
            "top_k": self.top_k,
            "max_error": self.max_error,
        }

        logger.info(
            "ToMoE Summary: Converted %d layers, Skipped %d layers. Params: %s -> %s (ratio: %.2fx)",
            len(converted_layers),
            len(skipped_layers),
            f"{orig_total_params:,}",
            f"{new_total_params:,}",
            self.stats["compression_ratio"],
        )

        return self.model, self.stats

ToMoE Compressor: dynamically replaces dense feed-forward and projection layers with an MoE-fied equivalent.

Args
-----=
model: PyTorch model.
num_experts: Number of expert sub-layers to form (default: 4).
top_k: Number of experts to route each token to (default: 2).
max_error: Maximum allowed relative Frobenius error on reconstruction (default: 0.20).
lora_r: LoRA adapter rank on experts and router (0 = disabled).
lora_alpha: LoRA scaling factor (default: 16.0).
freeze_experts: Whether to freeze expert base weights (default: True).
target_attention_only: Whether to restrict to attention layers only.
target_modules: Optional list of target layer name substrings.

Ancestors

BaseCompressor

Class variables

var PROTECTED_KEYWORDS: The type of the None singleton.

Methods

def compress(self, orig_model: nn.Module | None = None) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self, orig_model: nn.Module | None = None) -> tuple[nn.Module, dict[str, Any]]:
    """
    Compress target layers into ToMoE Mixture of Experts layers with full error validation.
    """
    if orig_model is not None:
        self.model = orig_model

    orig_total_params = sum(p.numel() for p in self.model.parameters())
    converted_layers = []
    skipped_layers = []
    layer_details = {}

    seen_weight_ptrs: set[int] = set()

    for name, module in list(self.model.named_modules()):
        if not self._is_target(name):
            continue

        # Fix 2: Check for supported layer type (nn.Linear or Conv1D) on base module
        base_mod = getattr(module, "base_layer", module)
        cls_name = base_mod.__class__.__name__.lower()
        is_linear = isinstance(base_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or (
            hasattr(base_mod, "nx") and hasattr(base_mod, "nf")
        )

        if not (
            is_linear
            or is_conv1d
            or (hasattr(base_mod, "weight") and base_mod.weight.ndim == 2)
        ):
            logger.warning(
                "ToMoE: Skipping unsupported layer type %s for %s",
                type(base_mod).__name__,
                name,
            )
            skipped_layers.append((name, f"unsupported_type_{type(base_mod).__name__}"))
            continue

        # H20: Detect and protect tied weights (e.g. wte and lm_head sharing same storage)
        if hasattr(base_mod, "weight") and base_mod.weight is not None:
            w_ptr = base_mod.weight.data_ptr()
            if w_ptr in seen_weight_ptrs:
                logger.warning(
                    "ToMoE: Skipping tied module %s with shared weight storage (H20)",
                    name,
                )
                skipped_layers.append((name, "tied_weight_shared_storage"))
                continue
            seen_weight_ptrs.add(w_ptr)

        # Determine weights and feature shapes ensuring orientation is (out_features, in_features)
        if is_conv1d:
            in_f = base_mod.weight.shape[0]
            out_f = base_mod.weight.shape[1]
            orig_w_2d = base_mod.weight.data.t().contiguous()
        else:
            in_f = getattr(base_mod, "in_features", base_mod.weight.shape[1])
            out_f = getattr(base_mod, "out_features", base_mod.weight.shape[0])
            orig_w_2d = base_mod.weight.data.contiguous()

        if torch.linalg.norm(orig_w_2d.float()) < 1e-4:
            skipped_layers.append((name, "frobenius_norm_too_small"))
            continue

        if out_f < 2:
            skipped_layers.append((name, f"out_features_too_small_{out_f}"))
            continue

        # Respect GQA dimensions (e.g. Qwen k_proj/v_proj with fewer heads)
        eff_num_experts = max(1, min(self.num_experts, out_f))
        eff_top_k = max(1, min(self.top_k, eff_num_experts))

        if eff_num_experts < 2:
            skipped_layers.append((name, "effective_experts_less_than_2"))
            continue

        parent, child_name = self._get_parent(name)
        if parent is None:
            skipped_layers.append((name, "parent_not_found"))
            continue

        # Construct ToMoELinear
        try:
            moe_layer = ToMoELinear(
                original=module,
                num_experts=eff_num_experts,
                top_k=eff_top_k,
                lora_r=self.lora_r,
                lora_alpha=self.lora_alpha,
                freeze_experts=self.freeze_experts,
                max_error=self.max_error,
            )
        except Exception as e:
            logger.warning("ToMoE: Conversion failed for %s: %s", name, e)
            skipped_layers.append((name, f"construction_error_{e}"))
            continue

        # Meaningful fidelity check: activation reconstruction error on a calibration batch
        with torch.no_grad():
            device = orig_w_2d.device
            dtype = orig_w_2d.dtype
            calib_x = torch.randn(4, 16, in_f, device=device, dtype=dtype)
            orig_out = module(calib_x)
            moe_out = moe_layer(calib_x)
            err = (
                torch.norm(moe_out.float() - orig_out.float())
                / (torch.norm(orig_out.float()) + 1e-8)
            ).item()

        if err > self.max_error or np.isnan(err):
            logger.warning(
                "ToMoE: %s activation reconstruction error (%.4f) exceeded max_error (%.4f). Keeping original layer.",
                name,
                err,
                self.max_error,
            )
            skipped_layers.append((name, f"error_threshold_exceeded_{err:.4f}"))
            for p in module.parameters():
                p.requires_grad = False
            continue

        # Replace layer in parent module preserving outer LoRALinear wrapper if present
        orig_p_count = sum(p.numel() for p in module.parameters())
        new_p_count = sum(p.numel() for p in moe_layer.parameters())
        if hasattr(module, "base_layer"):
            module.base_layer = moe_layer
            if hasattr(module, "lora_A") and module.lora_A is not None:
                for p in (
                    module.lora_A.parameters()
                    if hasattr(module.lora_A, "parameters")
                    else [module.lora_A]
                ):
                    p.requires_grad = True
            if hasattr(module, "lora_B") and module.lora_B is not None:
                for p in (
                    module.lora_B.parameters()
                    if hasattr(module.lora_B, "parameters")
                    else [module.lora_B]
                ):
                    p.requires_grad = True
        else:
            setattr(parent, child_name, moe_layer)

        converted_layers.append(name)
        layer_details[name] = {
            "num_experts": eff_num_experts,
            "top_k": eff_top_k,
            "reconstruction_error": err,
            "original_params": orig_p_count,
            "new_params": new_p_count,
            "equal_expert_size": moe_layer.equal_expert_size,
        }
        logger.info(
            "ToMoE: Converted %s -> ToMoELinear (experts=%d, top_k=%d, reconstruction_error=%.4e)",
            name,
            eff_num_experts,
            eff_top_k,
            err,
        )

    new_total_params = sum(p.numel() for p in self.model.parameters())
    self._compressed_params = new_total_params

    # Fix 5: Per-Layer Stats and Logging
    self.stats = {
        "converted_layers": converted_layers,
        "num_converted_layers": len(converted_layers),
        "skipped_layers": skipped_layers,
        "original_params": orig_total_params,
        "compressed_params": new_total_params,
        "compression_ratio": orig_total_params / max(1, new_total_params),
        "layer_details": layer_details,
        "num_experts": self.num_experts,
        "top_k": self.top_k,
        "max_error": self.max_error,
    }

    logger.info(
        "ToMoE Summary: Converted %d layers, Skipped %d layers. Params: %s -> %s (ratio: %.2fx)",
        len(converted_layers),
        len(skipped_layers),
        f"{orig_total_params:,}",
        f"{new_total_params:,}",
        self.stats["compression_ratio"],
    )

    return self.model, self.stats

Compress target layers into ToMoE Mixture of Experts layers with full error validation.

Inherited members

BaseCompressor:
- get_compression_stats
- summary

class ToMoELinear (original: nn.Module | None = None, in_features: int | None = None, out_features: int | None = None, weight: torch.Tensor | None = None, bias: torch.Tensor | None = None, num_experts: int = 4, top_k: int = 2, lora_r: int = 0, lora_alpha: float = 16.0, freeze_experts: bool = True, max_error: float = 0.2)

Expand source code

class ToMoELinear(nn.Module):
    """
    A single MoE-fied linear layer: learned routing + expert sub-layers.

    Features:
    - Clusters output neurons using FP32 KMeans (with equal-division fallback).
    - Vectorized forward pass using gather and scatter_add_ for equal-sized experts,
      or vectorized index-dispatch for variable-sized experts.
    - Numerical stability in FP16, BF16, and FP32.
    - Exposes reconstructed effective weight and bias properties for validation.
    - Supports optional trainable LoRA adapters on experts and router.
    """

    def __init__(
        self,
        original: nn.Module | None = None,
        in_features: int | None = None,
        out_features: int | None = None,
        weight: torch.Tensor | None = None,
        bias: torch.Tensor | None = None,
        num_experts: int = 4,
        top_k: int = 2,
        lora_r: int = 0,
        lora_alpha: float = 16.0,
        freeze_experts: bool = True,
        max_error: float = 0.20,
    ):
        super().__init__()

        # Fix 2: Handle both nn.Linear, Conv1D, and LoRALinear wrappers
        if original is not None:
            base_layer = getattr(original, "base_layer", original)
            cls_name = base_layer.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or (
                hasattr(base_layer, "nx") and hasattr(base_layer, "nf")
            )
            if is_conv1d:
                # Conv1D stores weight as (nx, nf) where nx=in_features, nf=out_features
                in_feat = base_layer.weight.shape[0]
                out_feat = base_layer.weight.shape[1]
                orig_w = (
                    base_layer.weight.data.t().contiguous().clone()
                )  # Transpose to (out_features, in_features)
                orig_b = base_layer.bias.data.clone() if base_layer.bias is not None else None
            elif isinstance(base_layer, nn.Linear):
                in_feat = base_layer.in_features
                out_feat = base_layer.out_features
                orig_w = base_layer.weight.data.contiguous().clone()
                orig_b = base_layer.bias.data.clone() if base_layer.bias is not None else None
            elif hasattr(base_layer, "weight") and base_layer.weight.ndim == 2:
                in_feat = getattr(base_layer, "in_features", base_layer.weight.shape[1])
                out_feat = getattr(base_layer, "out_features", base_layer.weight.shape[0])
                orig_w = base_layer.weight.data.contiguous().clone()
                orig_b = (
                    base_layer.bias.data.clone()
                    if getattr(base_layer, "bias", None) is not None
                    else None
                )
            else:
                raise TypeError(f"Unsupported module type: {type(base_layer)}")
        else:
            base_layer = None
            if in_features is None or out_features is None or weight is None:
                raise ValueError(
                    "Either 'original' module or ('in_features', 'out_features', 'weight') must be provided."
                )
            in_feat = in_features
            out_feat = out_features
            orig_w = weight.clone()
            orig_b = bias.clone() if bias is not None else None

        self.in_features = in_feat
        self.out_features = out_feat

        # Guard against num_experts > out_features
        self.num_experts = max(1, min(num_experts, out_feat))
        self.top_k = max(1, min(top_k, self.num_experts))

        # Check for existing LoRA in original
        has_orig_lora = (
            original is not None and hasattr(original, "lora_A") and hasattr(original, "lora_B")
        )
        if has_orig_lora and getattr(original, "r", 0) > 0:
            lora_r = getattr(original, "r", lora_r)

        self.lora_r = lora_r
        self.lora_alpha = lora_alpha
        self.lora_scale = (lora_alpha / lora_r) if lora_r > 0 else 1.0
        self.freeze_experts = freeze_experts
        self.max_error = max_error

        # Perform clustering with MiniBatchKMeans for large layers, standard KMeans otherwise
        W_fp32 = orig_w.detach().float().cpu().numpy()
        assignments = None

        if self.num_experts > 1:
            try:
                sig = inspect.signature(KMeans.__init__).parameters
                n_init_val = "auto" if "n_init" in sig else 10
                if self.out_features > 2048:
                    km = MiniBatchKMeans(
                        n_clusters=self.num_experts,
                        random_state=42,
                        batch_size=min(1024, self.out_features),
                        n_init=n_init_val,
                    )
                else:
                    km = KMeans(
                        n_clusters=self.num_experts,
                        random_state=42,
                        n_init=n_init_val,
                    )
                pred = km.fit_predict(W_fp32)
                counts = np.bincount(pred, minlength=self.num_experts)
                if (counts == 0).any():
                    logger.warning(
                        "ToMoE: KMeans produced empty cluster(s). Falling back to equal division."
                    )
                    assignments = None
                else:
                    # Balance KMeans clusters to equal size for fast vectorized path
                    target_size = self.out_features // self.num_experts
                    rem = self.out_features % self.num_experts
                    if rem == 0 and not (counts == target_size).all():
                        # Distances to cluster centers for balanced reassignment
                        dists = km.transform(W_fp32)  # [out_features, num_experts]
                        # Sort neurons by distance margin between assigned cluster and others
                        order = np.argsort(np.min(dists, axis=1))
                        balanced_assign = np.full(self.out_features, -1, dtype=int)
                        exp_occupancy = np.zeros(self.num_experts, dtype=int)
                        for idx in order:
                            pref_order = np.argsort(dists[idx])
                            for p in pref_order:
                                if exp_occupancy[p] < target_size:
                                    balanced_assign[idx] = p
                                    exp_occupancy[p] += 1
                                    break
                        assignments = balanced_assign
                    else:
                        assignments = pred
            except Exception as e:
                logger.warning(
                    "ToMoE: KMeans clustering failed (%s). Falling back to equal division.",
                    e,
                )
                assignments = None

        if assignments is None or (assignments < 0).any():
            # Deterministic equal contiguous division ensuring zero empty experts
            assignments = np.zeros(self.out_features, dtype=int)
            chunk = self.out_features // self.num_experts
            rem = self.out_features % self.num_experts
            start = 0
            for e in range(self.num_experts):
                end = start + chunk + (1 if e < rem else 0)
                assignments[start:end] = e
                start = end

        # Create expert sub-layers and record neuron assignment masks
        self.experts = nn.ModuleList()
        self.register_buffer(
            "_masks", torch.zeros(self.num_experts, self.out_features, dtype=torch.bool)
        )
        expert_indices_list = []
        cluster_sizes = []

        device = orig_w.device
        dtype = orig_w.dtype

        for e in range(self.num_experts):
            mask_bool = torch.tensor(assignments == e, dtype=torch.bool)
            self._masks[e] = mask_bool
            indices = torch.nonzero(mask_bool, as_tuple=True)[0]
            expert_indices_list.append(indices)
            n_neurons = len(indices)
            cluster_sizes.append(n_neurons)

            expert = nn.Linear(
                self.in_features,
                n_neurons,
                bias=orig_b is not None,
                device=device,
                dtype=dtype,
            )
            expert.weight.data.copy_(orig_w[mask_bool])
            if orig_b is not None:
                expert.bias.data.copy_(orig_b[mask_bool])

            if self.freeze_experts:
                expert.weight.requires_grad = False
                if expert.bias is not None:
                    expert.bias.requires_grad = False

            self.experts.append(expert)

        # Register padded expert indices as a proper buffer so it moves with .to(device)
        max_size = max(len(idx) for idx in expert_indices_list)
        padded_indices = torch.zeros(self.num_experts, max_size, dtype=torch.long)
        for e, idx in enumerate(expert_indices_list):
            padded_indices[e, : len(idx)] = idx
        self.register_buffer("expert_indices_padded", padded_indices)
        self.register_buffer("expert_lens", torch.tensor(cluster_sizes, dtype=torch.long))

        self.equal_expert_size = len(set(cluster_sizes)) == 1

        if self.equal_expert_size:
            self.expert_out_dim = cluster_sizes[0]
            col_map = torch.stack(expert_indices_list, dim=0).to(device)
            self.register_buffer("expert_col_map", col_map)
        else:
            self.expert_out_dim = -1

        # Router in FP32 for numerical stability & speed
        self.router = nn.Linear(
            self.in_features, self.num_experts, device=device, dtype=torch.float32
        )
        nn.init.normal_(self.router.weight, std=0.02)
        if self.router.bias is not None:
            nn.init.zeros_(self.router.bias)
        self.router.weight.requires_grad = True
        if self.router.bias is not None:
            self.router.bias.requires_grad = True

        # Trainable LoRA Adapters
        if self.lora_r > 0:
            self.expert_lora_A = nn.ParameterList()
            self.expert_lora_B = nn.ParameterList()
            for e in range(self.num_experts):
                n_neurons = cluster_sizes[e]
                l_A = nn.Parameter(
                    torch.empty(self.lora_r, self.in_features, device=device, dtype=dtype)
                )
                l_B = nn.Parameter(torch.zeros(n_neurons, self.lora_r, device=device, dtype=dtype))
                nn.init.kaiming_uniform_(l_A, a=math.sqrt(5))
                nn.init.zeros_(l_B)
                self.expert_lora_A.append(l_A)
                self.expert_lora_B.append(l_B)

            self.lora_A_router = nn.Parameter(
                torch.empty(self.lora_r, self.in_features, device=device, dtype=torch.float32)
            )
            self.lora_B_router = nn.Parameter(
                torch.zeros(self.num_experts, self.lora_r, device=device, dtype=torch.float32)
            )
            nn.init.kaiming_uniform_(self.lora_A_router, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B_router)

    @property
    def masks(self) -> list[torch.Tensor]:
        return [self._masks[e] for e in range(self.num_experts)]

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective full weight matrix (out_features, in_features)."""
        device = self.experts[0].weight.device
        dtype = self.experts[0].weight.dtype
        W = torch.zeros(self.out_features, self.in_features, device=device, dtype=dtype)
        for e, expert in enumerate(self.experts):
            mask = self._masks[e]
            W[mask] = expert.weight
        return W

    @property
    def bias(self) -> torch.Tensor | None:
        """Reconstructed effective full bias vector (out_features,)."""
        if self.experts[0].bias is None:
            return None
        device = self.experts[0].bias.device
        dtype = self.experts[0].bias.dtype
        b = torch.zeros(self.out_features, device=device, dtype=dtype)
        for e, expert in enumerate(self.experts):
            mask = self._masks[e]
            b[mask] = expert.bias
        return b

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        in_feat = orig_shape[-1]
        x_flat = x.reshape(-1, in_feat)
        N = x_flat.shape[0]

        # Compute router logits in FP32
        router_logits = self.router(x_flat.float())
        if self.lora_r > 0 and hasattr(self, "lora_A_router") and hasattr(self, "lora_B_router"):
            router_lora = (
                (x_flat.float() @ self.lora_A_router.t()) @ self.lora_B_router.t() * self.lora_scale
            )
            router_logits = router_logits + router_lora

        if torch.isnan(router_logits).any() or torch.isinf(router_logits).any():
            router_logits = torch.nan_to_num(router_logits, nan=0.0, posinf=1.0, neginf=-1.0)

        # Top-k routing
        topk_vals, topk_idx = torch.topk(router_logits, self.top_k, dim=-1)  # (N, top_k)
        topk_weights = torch.softmax(topk_vals, dim=-1).to(x.dtype)  # (N, top_k)

        out = torch.zeros(N, self.out_features, device=x.device, dtype=x.dtype)

        if self.equal_expert_size:
            # Fast vectorized gather & scatter_add_ for equal-sized expert partitions
            D_exp = self.expert_out_dim
            all_expert_outs = torch.stack(
                [exp(x_flat) for exp in self.experts], dim=1
            )  # (N, E, D_exp)

            if self.lora_r > 0 and hasattr(self, "expert_lora_A"):
                lora_outs = []
                for e in range(self.num_experts):
                    l_e = (
                        (x_flat.to(self.expert_lora_A[e].dtype) @ self.expert_lora_A[e].t())
                        @ self.expert_lora_B[e].t()
                        * self.lora_scale
                    )
                    lora_outs.append(l_e.to(x.dtype))
                all_expert_outs = all_expert_outs + torch.stack(lora_outs, dim=1)

            expanded_idx = topk_idx.unsqueeze(-1).expand(N, self.top_k, D_exp)
            selected_outs = torch.gather(
                all_expert_outs, dim=1, index=expanded_idx
            )  # (N, top_k, D_exp)
            weighted_selected = selected_outs * topk_weights.unsqueeze(-1)  # (N, top_k, D_exp)

            col_map = self.expert_col_map
            dest_cols = col_map[topk_idx]  # (N, top_k, D_exp)

            out.scatter_add_(
                dim=1,
                index=dest_cols.reshape(N, self.top_k * D_exp),
                src=weighted_selected.reshape(N, self.top_k * D_exp),
            )
        else:
            # Vectorized token dispatch per expert with scatter_add_
            for e in range(self.num_experts):
                expert_match = topk_idx == e  # (N, top_k)
                if not expert_match.any():
                    continue
                token_idx, rank_pos = torch.where(expert_match)
                selected_tokens = x_flat[token_idx]
                exp_out = self.experts[e](selected_tokens)
                if self.lora_r > 0 and hasattr(self, "expert_lora_A"):
                    l_e = (
                        (
                            selected_tokens.to(self.expert_lora_A[e].dtype)
                            @ self.expert_lora_A[e].t()
                        )
                        @ self.expert_lora_B[e].t()
                        * self.lora_scale
                    )
                    exp_out = exp_out + l_e.to(x.dtype)
                gate_w = topk_weights[token_idx, rank_pos].unsqueeze(-1)
                weighted_exp_out = exp_out * gate_w

                exp_cols = self.expert_indices_padded[e, : self.expert_lens[e]]
                flat_idx = token_idx[:, None] * self.out_features + exp_cols
                out.view(-1).scatter_add_(0, flat_idx.reshape(-1), weighted_exp_out.reshape(-1))

        return out.reshape(*orig_shape[:-1], self.out_features)

A single MoE-fied linear layer: learned routing + expert sub-layers.

Features: - Clusters output neurons using FP32 KMeans (with equal-division fallback). - Vectorized forward pass using gather and scatter_add_ for equal-sized experts, or vectorized index-dispatch for variable-sized experts. - Numerical stability in FP16, BF16, and FP32. - Exposes reconstructed effective weight and bias properties for validation. - Supports optional trainable LoRA adapters on experts and router.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Instance variables

prop bias : torch.Tensor | None

Expand source code

@property
def bias(self) -> torch.Tensor | None:
    """Reconstructed effective full bias vector (out_features,)."""
    if self.experts[0].bias is None:
        return None
    device = self.experts[0].bias.device
    dtype = self.experts[0].bias.dtype
    b = torch.zeros(self.out_features, device=device, dtype=dtype)
    for e, expert in enumerate(self.experts):
        mask = self._masks[e]
        b[mask] = expert.bias
    return b

Reconstructed effective full bias vector (out_features,).

prop masks : list[torch.Tensor]

Expand source code

@property
def masks(self) -> list[torch.Tensor]:
    return [self._masks[e] for e in range(self.num_experts)]

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective full weight matrix (out_features, in_features)."""
    device = self.experts[0].weight.device
    dtype = self.experts[0].weight.dtype
    W = torch.zeros(self.out_features, self.in_features, device=device, dtype=dtype)
    for e, expert in enumerate(self.experts):
        mask = self._masks[e]
        W[mask] = expert.weight
    return W

Reconstructed effective full weight matrix (out_features, in_features).

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_shape = x.shape
    in_feat = orig_shape[-1]
    x_flat = x.reshape(-1, in_feat)
    N = x_flat.shape[0]

    # Compute router logits in FP32
    router_logits = self.router(x_flat.float())
    if self.lora_r > 0 and hasattr(self, "lora_A_router") and hasattr(self, "lora_B_router"):
        router_lora = (
            (x_flat.float() @ self.lora_A_router.t()) @ self.lora_B_router.t() * self.lora_scale
        )
        router_logits = router_logits + router_lora

    if torch.isnan(router_logits).any() or torch.isinf(router_logits).any():
        router_logits = torch.nan_to_num(router_logits, nan=0.0, posinf=1.0, neginf=-1.0)

    # Top-k routing
    topk_vals, topk_idx = torch.topk(router_logits, self.top_k, dim=-1)  # (N, top_k)
    topk_weights = torch.softmax(topk_vals, dim=-1).to(x.dtype)  # (N, top_k)

    out = torch.zeros(N, self.out_features, device=x.device, dtype=x.dtype)

    if self.equal_expert_size:
        # Fast vectorized gather & scatter_add_ for equal-sized expert partitions
        D_exp = self.expert_out_dim
        all_expert_outs = torch.stack(
            [exp(x_flat) for exp in self.experts], dim=1
        )  # (N, E, D_exp)

        if self.lora_r > 0 and hasattr(self, "expert_lora_A"):
            lora_outs = []
            for e in range(self.num_experts):
                l_e = (
                    (x_flat.to(self.expert_lora_A[e].dtype) @ self.expert_lora_A[e].t())
                    @ self.expert_lora_B[e].t()
                    * self.lora_scale
                )
                lora_outs.append(l_e.to(x.dtype))
            all_expert_outs = all_expert_outs + torch.stack(lora_outs, dim=1)

        expanded_idx = topk_idx.unsqueeze(-1).expand(N, self.top_k, D_exp)
        selected_outs = torch.gather(
            all_expert_outs, dim=1, index=expanded_idx
        )  # (N, top_k, D_exp)
        weighted_selected = selected_outs * topk_weights.unsqueeze(-1)  # (N, top_k, D_exp)

        col_map = self.expert_col_map
        dest_cols = col_map[topk_idx]  # (N, top_k, D_exp)

        out.scatter_add_(
            dim=1,
            index=dest_cols.reshape(N, self.top_k * D_exp),
            src=weighted_selected.reshape(N, self.top_k * D_exp),
        )
    else:
        # Vectorized token dispatch per expert with scatter_add_
        for e in range(self.num_experts):
            expert_match = topk_idx == e  # (N, top_k)
            if not expert_match.any():
                continue
            token_idx, rank_pos = torch.where(expert_match)
            selected_tokens = x_flat[token_idx]
            exp_out = self.experts[e](selected_tokens)
            if self.lora_r > 0 and hasattr(self, "expert_lora_A"):
                l_e = (
                    (
                        selected_tokens.to(self.expert_lora_A[e].dtype)
                        @ self.expert_lora_A[e].t()
                    )
                    @ self.expert_lora_B[e].t()
                    * self.lora_scale
                )
                exp_out = exp_out + l_e.to(x.dtype)
            gate_w = topk_weights[token_idx, rank_pos].unsqueeze(-1)
            weighted_exp_out = exp_out * gate_w

            exp_cols = self.expert_indices_padded[e, : self.expert_lens[e]]
            flat_idx = token_idx[:, None] * self.out_features + exp_cols
            out.view(-1).scatter_add_(0, flat_idx.reshape(-1), weighted_exp_out.reshape(-1))

    return out.reshape(*orig_shape[:-1], self.out_features)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class UniQLCompressor (model: nn.Module, rank: int = 64, bits: int = 8, max_error: float = 0.2, max_retries: int = 3, fallback_to_original: bool = True, target_attention_only: bool = False, quantize_bias: bool = False, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("uniql_compressor")
@COMPRESSOR_REGISTRY.register("uniql")
class UniQLCompressor(BaseCompressor):
    """
    UniQL: Joint Truncated SVD + Uniform Quantization compressor.

    Unlike the naive "copy back" variant, this compressor installs a real
    UniQLinear module that stores int8 U, S, Vt buffers. Parameter count
    and memory footprint decrease by roughly the SVD factor
    (out + in) * r / (out * in).
    """

    # UQ-01: NOT an attention module.
    is_custom_attention: bool = False

    def __init__(
        self,
        model: nn.Module,
        rank: int = 64,
        bits: int = 8,
        max_error: float = 0.20,
        max_retries: int = 3,
        fallback_to_original: bool = True,
        target_attention_only: bool = False,
        quantize_bias: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="uniql", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            rank = _cfg_get(cfg, "uniql_rank", _cfg_get(cfg, "rank", rank))
            bits = _cfg_get(cfg, "uniql_bits", _cfg_get(cfg, "bits", bits))
            max_error = _cfg_get(
                cfg, "max_error", _cfg_get(cfg, "max_reconstruction_error", max_error)
            )
            max_retries = _cfg_get(cfg, "max_retries", max_retries)
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)
            quantize_bias = _cfg_get(cfg, "uniql_quantize_bias", quantize_bias)

        if kwargs.get("rank") is not None:
            rank = kwargs["rank"]
        if kwargs.get("bits") is not None:
            bits = kwargs["bits"]
        if kwargs.get("max_error") is not None:
            max_error = float(kwargs["max_error"])
        if kwargs.get("max_retries") is not None:
            max_retries = int(kwargs["max_retries"])
        if "fallback_to_original" in kwargs:
            fallback_to_original = bool(kwargs["fallback_to_original"])
        if "target_attention_only" in kwargs:
            target_attention_only = bool(kwargs["target_attention_only"])
        if "quantize_bias" in kwargs:
            quantize_bias = bool(kwargs["quantize_bias"])

        self.rank = max(1, int(rank))
        self.bits = max(2, min(16, int(bits)))  # cap at 16 bits (int8 storage assumption)
        self.max_error = float(max_error)
        self.max_retries = max(0, int(max_retries))
        self.fallback_to_original = bool(fallback_to_original)
        self.target_attention_only = target_attention_only
        self.quantize_bias = bool(quantize_bias)
        self.compressed_layers: list[dict[str, Any]] = []

    # ------------------------------------------------------------------
    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        # Protect classifier, pooler, lm_head, embeddings, norms.
        name_lower = name.lower()
        if any(
            k in name_lower
            for k in (
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "layernorm",
                "rmsnorm",
                "batchnorm",
            )
        ):
            return False

        if isinstance(
            module,
            (nn.LayerNorm, nn.RMSNorm, nn.Embedding, nn.BatchNorm1d, nn.BatchNorm2d),
        ):
            return False

        # Skip if already compressed.
        if isinstance(module, UniQLinear):
            return False

        if self.target_attention_only:
            attn_keywords = (
                "attn",
                "attention",
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "c_attn",
                "c_proj",
                "query",
                "key",
                "value",
            )
            if not any(k in name_lower for k in attn_keywords):
                return False

        base_layer = getattr(module, "base_layer", module)
        cls_name = base_layer.__class__.__name__.lower()
        # UQ-10: Conv1D detection across transformers versions.
        is_conv1d = "conv1d" in cls_name or (
            hasattr(base_layer, "nf") and hasattr(base_layer, "weight")
        )
        is_linear_like = (
            isinstance(base_layer, nn.Linear)
            or is_conv1d
            or (
                hasattr(base_layer, "weight")
                and hasattr(base_layer.weight, "ndim")
                and base_layer.weight.ndim == 2
            )
        )
        if not is_linear_like:
            return False

        # Respect framework policy (e.g. Qwen GQA asymmetric skips).
        try:
            if not should_compress_linear(self.model, name):
                return False
        except Exception:
            pass
        return True

    # ------------------------------------------------------------------
    def _normalize_weight(self, base_layer: nn.Module) -> tuple[torch.Tensor, bool]:
        """
        Return (W as (out, in) fp32, is_conv1d).
        UQ-04: explicit, not heuristic.
        """
        raw_w = base_layer.weight.data
        cls_name = base_layer.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or (
            hasattr(base_layer, "nf") and raw_w.shape[0] < raw_w.shape[1]
        )
        if is_conv1d:
            W = raw_w.t().contiguous().float()
        else:
            W = raw_w.float().contiguous()
        return W, is_conv1d

    # ------------------------------------------------------------------
    def _try_compress(self, W: torch.Tensor) -> dict[str, Any] | None:
        """
        Attempt SVD + quantization with adaptive bit/rank escalation.
        UQ-05: monotonic, bounded escalation.
        Returns a dict with quantized factors + reconstruction error, or None.
        """
        out_f, in_f = W.shape
        max_r = min(out_f, in_f)
        if max_r < 1:
            return None

        norm_orig = torch.linalg.norm(W).item() + 1e-12

        # ---- SVD once, in fp32 ----
        try:
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        except Exception as e:
            logger.warning("UniQL SVD failed: %s", e)
            return None

        # Candidate ranks: start at requested rank, expand up to max_r.
        base_r = max(1, min(self.rank, max_r))
        rank_candidates = [base_r]
        if base_r < max_r:
            rank_candidates.append(min(max_r, base_r * 2))
        if max_r not in rank_candidates:
            rank_candidates.append(max_r)

        # Candidate bits: bits, then 8, then 16 (int8 storage cap).
        bit_candidates = [self.bits]
        for b in (8, 16):
            if b > bit_candidates[-1]:
                bit_candidates.append(b)

        best = None
        best_err = float("inf")

        for cur_r in rank_candidates:
            U_r = U[:, :cur_r]
            S_r = S[:cur_r]
            Vt_r = Vt[:cur_r, :]
            for cur_bits in bit_candidates:
                # UQ-03: quantize the actual stored factors.
                U_q, scale_u, _ = _uniform_quantize(U_r, cur_bits)
                S_q, scale_s, _ = _uniform_quantize(S_r, cur_bits)
                Vt_q, scale_v, _ = _uniform_quantize(Vt_r, cur_bits)

                # Reconstruct in fp32 for error measurement.
                U_hat = U_q.float() * scale_u.unsqueeze(1)
                S_hat = S_q.float() * scale_s
                Vt_hat = Vt_q.float() * scale_v.unsqueeze(0)
                W_rec = (U_hat * S_hat.unsqueeze(0)) @ Vt_hat
                rel_err = (torch.linalg.norm(W - W_rec).item()) / norm_orig

                if rel_err < best_err:
                    best_err = rel_err
                    best = {
                        "U_q": U_q,
                        "S_q": S_q,
                        "Vt_q": Vt_q,
                        "scale_u": scale_u,
                        "scale_s": scale_s,
                        "scale_v": scale_v,
                        "rank": cur_r,
                        "bits": cur_bits,
                        "rel_err": rel_err,
                    }
                if rel_err <= self.max_error:
                    return best

        if best is None or best_err > self.max_error:
            return None
        return best

    # ------------------------------------------------------------------
    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        self.compressed_layers = []
        fallback_layers: list[str] = []
        total_err = 0.0
        orig_storage_bytes = 0
        comp_storage_bytes = 0

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            base_layer = getattr(module, "base_layer", module)
            raw_w = base_layer.weight.data
            orig_dtype = raw_w.dtype
            orig_device = raw_w.device

            W, is_conv1d = self._normalize_weight(base_layer)
            out_f, in_f = W.shape
            if out_f < 1 or in_f < 1:
                continue

            result = self._try_compress(W)
            if result is None:
                if self.fallback_to_original:
                    logger.info("UniQL: keeping original weights for %s", name)
                fallback_layers.append(name)
                continue

            # ---- Build the new module ----
            bias_tensor = None
            if getattr(base_layer, "bias", None) is not None and base_layer.bias is not None:
                b = base_layer.bias.data
                if self.quantize_bias:
                    bq, bs, _ = _uniform_quantize(b.float(), result["bits"])
                    b = bq.float() * bs
                bias_tensor = b.to(dtype=orig_dtype, device=orig_device)

            new_layer = UniQLinear(
                in_features=in_f,
                out_features=out_f,
                rank=result["rank"],
                U_q=result["U_q"],
                S_q=result["S_q"],
                Vt_q=result["Vt_q"],
                scale_u=result["scale_u"].to(orig_device),
                scale_s=result["scale_s"].to(orig_device),
                scale_v=result["scale_v"].to(orig_device),
                bias=bias_tensor,
                is_conv1d=is_conv1d,
                dtype=orig_dtype,
                device=orig_device,
            )
            new_layer.train(mode=module.training)

            # ---- Splice into the model ----
            if hasattr(module, "base_layer"):
                # Wrapped (e.g. LoRALinear). Keep wrapper; replace its base_layer.
                module.base_layer = new_layer
                # Freeze any remaining LoRA parameters so recovery can still
                # train them if the pipeline expects that.
                for p_name, p in module.named_parameters():
                    if "lora_" in p_name.lower() and not p_name.startswith("base_layer"):
                        p.requires_grad = False
            else:
                # Plain Linear / Conv1D: replace in the model.
                parent_name, _, child_name = name.rpartition(".")
                parent = self.model.get_submodule(parent_name) if parent_name else self.model
                setattr(parent, child_name, new_layer)

            # ---- Storage accounting (honest) ----
            # UQ-07: report real bytes.
            orig_bytes = raw_w.numel() * raw_w.element_size()
            comp_bytes = (
                result["U_q"].numel() + result["S_q"].numel() + result["Vt_q"].numel()
            )  # int8: 1 byte each
            comp_bytes += (
                result["scale_u"].numel() + result["scale_s"].numel() + result["scale_v"].numel()
            ) * 4  # fp32 scales
            if bias_tensor is not None:
                orig_bytes += base_layer.bias.numel() * base_layer.bias.element_size()
                comp_bytes += bias_tensor.numel() * bias_tensor.element_size()

            orig_storage_bytes += orig_bytes
            comp_storage_bytes += comp_bytes

            self.compressed_layers.append(
                {
                    "name": name,
                    "original_shape": (out_f, in_f),
                    "rank": result["rank"],
                    "bits": result["bits"],
                    "rel_error": float(result["rel_err"]),
                    "orig_bytes": orig_bytes,
                    "comp_bytes": comp_bytes,
                }
            )
            total_err += result["rel_err"]

        # ---- Stats ----
        try:
            self._compressed_params = self._count_params(self.model)
        except Exception:
            self._compressed_params = sum(p.numel() for p in self.model.parameters())

        stats = self.get_compression_stats()
        avg_err = (total_err / len(self.compressed_layers)) if self.compressed_layers else 0.0
        ratio = (orig_storage_bytes / comp_storage_bytes) if comp_storage_bytes > 0 else 1.0

        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "fallback_layers": fallback_layers,
                "num_compressed_layers": len(self.compressed_layers),
                "num_fallback": len(fallback_layers),
                "rank": self.rank,
                "bits": self.bits,
                "avg_reconstruction_error": avg_err,
                "max_error": self.max_error,
                # UQ-07: honest accounting
                "orig_storage_bytes": orig_storage_bytes,
                "comp_storage_bytes": comp_storage_bytes,
                "storage_compression_ratio": round(ratio, 3),
                "storage_reduction_pct": round(
                    100.0 * (1.0 - comp_storage_bytes / max(1, orig_storage_bytes)), 2
                ),
            }
        )

        logger.info(
            "UniQL: %d layers compressed, %d fallback, avg_err=%.4f, storage_ratio=%.2fx",
            len(self.compressed_layers),
            len(fallback_layers),
            avg_err,
            ratio,
        )
        return self.model, stats

UniQL: Joint Truncated SVD + Uniform Quantization compressor.

Unlike the naive "copy back" variant, this compressor installs a real UniQLinear module that stores int8 U, S, Vt buffers. Parameter count and memory footprint decrease by roughly the SVD factor (out + in) * r / (out * in).

Ancestors

BaseCompressor

Class variables

var is_custom_attention : bool: The type of the None singleton.

Inherited members

BaseCompressor:
- compress
- get_compression_stats
- summary

class UniQLinear (in_features: int, out_features: int, rank: int, U_q: torch.Tensor, S_q: torch.Tensor, Vt_q: torch.Tensor, scale_u: torch.Tensor, scale_s: torch.Tensor, scale_v: torch.Tensor, bias: torch.Tensor | None = None, is_conv1d: bool = False, dtype: torch.dtype = torch.float32, device: torch.device | None = None)

Expand source code

class UniQLinear(nn.Module):
    """
    Replaces a Linear (or Conv1D) with a stored factorization:

        W ≈ U @ diag(S) @ Vt
        where U  : (out, r) int8, scale_u  (out,)
              S  : (r,)    int8, scale_s  scalar-per-rank
              Vt : (r, in) int8, scale_v  (in,)

    Forward reconstructs W_hat on the fly and applies F.linear / F.linear with
    transposed Conv1D weight. Numerically this is exactly the same
    reconstruction as the original UniQL, but now the storage is genuinely
    reduced from out*in fp16 to (out + in) * r int8 plus small scale vectors.

    Original bias (if any) is stored as-is unless quantize_bias=True.
    """

    is_custom_attention: bool = False  # UQ-01: compressors are not attention.

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        U_q: torch.Tensor,  # (out, r) int8
        S_q: torch.Tensor,  # (r,)    int8
        Vt_q: torch.Tensor,  # (r, in) int8
        scale_u: torch.Tensor,  # (out,)
        scale_s: torch.Tensor,  # (r,)
        scale_v: torch.Tensor,  # (in,)
        bias: torch.Tensor | None = None,
        is_conv1d: bool = False,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = max(1, min(rank, in_features, out_features))
        self.is_conv1d = is_conv1d

        # Store quantized factors as non-trainable buffers.
        self.register_buffer("U_q", U_q.to(device=device), persistent=True)
        self.register_buffer("S_q", S_q.to(device=device), persistent=True)
        self.register_buffer("Vt_q", Vt_q.to(device=device), persistent=True)
        self.register_buffer(
            "scale_u", scale_u.to(device=device, dtype=torch.float32), persistent=True
        )
        self.register_buffer(
            "scale_s", scale_s.to(device=device, dtype=torch.float32), persistent=True
        )
        self.register_buffer(
            "scale_v", scale_v.to(device=device, dtype=torch.float32), persistent=True
        )

        if bias is not None:
            self.bias = nn.Parameter(bias.to(dtype=dtype, device=device), requires_grad=False)
        else:
            self.register_parameter("bias", None)

        self._weight_dtype = dtype

    # ---- Reconstruct in fp32 for numerical stability ----
    def _reconstruct_weight(self) -> torch.Tensor:
        U = self.U_q.float() * self.scale_u.unsqueeze(1)  # (out, r)
        S = self.S_q.float() * self.scale_s  # (r,)
        Vt = self.Vt_q.float() * self.scale_v.unsqueeze(0)  # (r, in)
        W = (U * S.unsqueeze(0)) @ Vt  # (out, in)
        return W.to(self._weight_dtype)

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed weight (out, in). Conv1D callers should transpose."""
        return self._reconstruct_weight()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        W = self._reconstruct_weight()
        if self.is_conv1d:
            # Conv1D stores weight as (in, out); our W is (out, in).
            out = F.linear(x, W)
        else:
            out = F.linear(x, W)
        if self.bias is not None:
            out = out + self.bias
        return out

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, out_features={self.out_features}, "
            f"rank={self.rank}, conv1d={self.is_conv1d}"
        )

Replaces a Linear (or Conv1D) with a stored factorization:

W ≈ U @ diag(S) @ Vt
where U  : (out, r) int8, scale_u  (out,)
      S  : (r,)    int8, scale_s  scalar-per-rank
      Vt : (r, in) int8, scale_v  (in,)

Forward reconstructs W_hat on the fly and applies F.linear / F.linear with transposed Conv1D weight. Numerically this is exactly the same reconstruction as the original UniQL, but now the storage is genuinely reduced from out*in fp16 to (out + in) * r int8 plus small scale vectors.

Original bias (if any) is stored as-is unless quantize_bias=True.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool: The type of the None singleton.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed weight (out, in). Conv1D callers should transpose."""
    return self._reconstruct_weight()

Reconstructed weight (out, in). Conv1D callers should transpose.

Methods

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    return (
        f"in_features={self.in_features}, out_features={self.out_features}, "
        f"rank={self.rank}, conv1d={self.is_conv1d}"
    )

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    W = self._reconstruct_weight()
    if self.is_conv1d:
        # Conv1D stores weight as (in, out); our W is (out, in).
        out = F.linear(x, W)
    else:
        out = F.linear(x, W)
    if self.bias is not None:
        out = out + self.bias
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class ZeroKVNeoCache (max_cache_len: int = 128, num_attention_sinks: int = 4, recent_window: int = 32)

Expand source code

class ZeroKVNeoCache:
    """
    Dynamic KV Cache Manager for ZeroKV-Neo.
    Maintains bounded O(1) cache memory by partitioning budget:
    - Pinned attention sinks (first N_sinks prompt tokens)
    - Pinned local sliding window (last N_window recent tokens)
    - Dynamic heavy hitters in the middle sequence retained via L2 norm / cumulative score.
    """

    def __init__(
        self,
        max_cache_len: int = 128,
        num_attention_sinks: int = 4,
        recent_window: int = 32,
    ) -> None:
        self.max_cache_len = int(max_cache_len)
        self.num_attention_sinks = int(num_attention_sinks)
        self.recent_window = int(recent_window)
        self.k_cache: torch.Tensor | None = None
        self.v_cache: torch.Tensor | None = None

    def reset(self) -> None:
        """Clear cache tensors to avoid memory leaks."""
        self.k_cache = None
        self.v_cache = None

    @torch.no_grad()
    def update_and_compress(
        self,
        new_k: torch.Tensor,
        new_v: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Append new KV tokens and evict non-essential tokens when exceeding capacity.
        Shapes: (B, H, N, D).
        """
        k_4d = new_k
        v_4d = new_v
        if new_k.ndim == 2:
            k_4d = new_k.unsqueeze(0).unsqueeze(0)
            v_4d = new_v.unsqueeze(0).unsqueeze(0)
        elif new_k.ndim == 3:
            k_4d = new_k.unsqueeze(1)
            v_4d = new_v.unsqueeze(1)

        B, H, N, D = k_4d.shape

        if self.k_cache is None:
            self.k_cache = k_4d
            self.v_cache = v_4d
        else:
            if self.k_cache.shape[0] != B:
                self.k_cache = self.k_cache[:1].expand(B, -1, -1, -1).clone()
                self.v_cache = self.v_cache[:1].expand(B, -1, -1, -1).clone()
            self.k_cache = torch.cat([self.k_cache, k_4d], dim=2)
            self.v_cache = torch.cat([self.v_cache, v_4d], dim=2)

        cur_len = self.k_cache.size(2)
        if cur_len <= self.max_cache_len:
            return self.k_cache, self.v_cache

        sinks = min(self.num_attention_sinks, cur_len)
        window = min(self.recent_window, cur_len - sinks)
        heavy_budget = max(0, self.max_cache_len - sinks - window)

        sink_k = self.k_cache[:, :, :sinks, :]
        sink_v = self.v_cache[:, :, :sinks, :]
        win_k = self.k_cache[:, :, -window:, :] if window > 0 else self.k_cache[:, :, :0, :]
        win_v = self.v_cache[:, :, -window:, :] if window > 0 else self.v_cache[:, :, :0, :]

        if heavy_budget > 0 and cur_len > (sinks + window):
            mid_k = self.k_cache[:, :, sinks : cur_len - window, :]
            mid_v = self.v_cache[:, :, sinks : cur_len - window, :]
            mid_norms = torch.norm(mid_k.float(), dim=-1)  # (B, H, mid_len)
            _, top_idx = torch.topk(mid_norms, min(heavy_budget, mid_norms.size(-1)), dim=-1)
            top_idx, _ = torch.sort(top_idx, dim=-1)

            idx_exp = top_idx.unsqueeze(-1).expand(-1, -1, -1, D)
            sel_mid_k = mid_k.gather(2, idx_exp)
            sel_mid_v = mid_v.gather(2, idx_exp)

            self.k_cache = torch.cat([sink_k, sel_mid_k, win_k], dim=2).contiguous()
            self.v_cache = torch.cat([sink_v, sel_mid_v, win_v], dim=2).contiguous()
        else:
            self.k_cache = torch.cat([sink_k, win_k], dim=2).contiguous() if window > 0 else sink_k
            self.v_cache = torch.cat([sink_v, win_v], dim=2).contiguous() if window > 0 else sink_v

        return self.k_cache, self.v_cache

Dynamic KV Cache Manager for ZeroKV-Neo. Maintains bounded O(1) cache memory by partitioning budget: - Pinned attention sinks (first N_sinks prompt tokens) - Pinned local sliding window (last N_window recent tokens) - Dynamic heavy hitters in the middle sequence retained via L2 norm / cumulative score.

Methods

def reset(self) ‑> None

Expand source code

def reset(self) -> None:
    """Clear cache tensors to avoid memory leaks."""
    self.k_cache = None
    self.v_cache = None

Clear cache tensors to avoid memory leaks.

def update_and_compress(self, new_k: torch.Tensor, new_v: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor]

Expand source code

@torch.no_grad()
def update_and_compress(
    self,
    new_k: torch.Tensor,
    new_v: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Append new KV tokens and evict non-essential tokens when exceeding capacity.
    Shapes: (B, H, N, D).
    """
    k_4d = new_k
    v_4d = new_v
    if new_k.ndim == 2:
        k_4d = new_k.unsqueeze(0).unsqueeze(0)
        v_4d = new_v.unsqueeze(0).unsqueeze(0)
    elif new_k.ndim == 3:
        k_4d = new_k.unsqueeze(1)
        v_4d = new_v.unsqueeze(1)

    B, H, N, D = k_4d.shape

    if self.k_cache is None:
        self.k_cache = k_4d
        self.v_cache = v_4d
    else:
        if self.k_cache.shape[0] != B:
            self.k_cache = self.k_cache[:1].expand(B, -1, -1, -1).clone()
            self.v_cache = self.v_cache[:1].expand(B, -1, -1, -1).clone()
        self.k_cache = torch.cat([self.k_cache, k_4d], dim=2)
        self.v_cache = torch.cat([self.v_cache, v_4d], dim=2)

    cur_len = self.k_cache.size(2)
    if cur_len <= self.max_cache_len:
        return self.k_cache, self.v_cache

    sinks = min(self.num_attention_sinks, cur_len)
    window = min(self.recent_window, cur_len - sinks)
    heavy_budget = max(0, self.max_cache_len - sinks - window)

    sink_k = self.k_cache[:, :, :sinks, :]
    sink_v = self.v_cache[:, :, :sinks, :]
    win_k = self.k_cache[:, :, -window:, :] if window > 0 else self.k_cache[:, :, :0, :]
    win_v = self.v_cache[:, :, -window:, :] if window > 0 else self.v_cache[:, :, :0, :]

    if heavy_budget > 0 and cur_len > (sinks + window):
        mid_k = self.k_cache[:, :, sinks : cur_len - window, :]
        mid_v = self.v_cache[:, :, sinks : cur_len - window, :]
        mid_norms = torch.norm(mid_k.float(), dim=-1)  # (B, H, mid_len)
        _, top_idx = torch.topk(mid_norms, min(heavy_budget, mid_norms.size(-1)), dim=-1)
        top_idx, _ = torch.sort(top_idx, dim=-1)

        idx_exp = top_idx.unsqueeze(-1).expand(-1, -1, -1, D)
        sel_mid_k = mid_k.gather(2, idx_exp)
        sel_mid_v = mid_v.gather(2, idx_exp)

        self.k_cache = torch.cat([sink_k, sel_mid_k, win_k], dim=2).contiguous()
        self.v_cache = torch.cat([sink_v, sel_mid_v, win_v], dim=2).contiguous()
    else:
        self.k_cache = torch.cat([sink_k, win_k], dim=2).contiguous() if window > 0 else sink_k
        self.v_cache = torch.cat([sink_v, win_v], dim=2).contiguous() if window > 0 else sink_v

    return self.k_cache, self.v_cache

Append new KV tokens and evict non-essential tokens when exceeding capacity. Shapes: (B, H, N, D).

class ZeroKVNeoCompressor (model: nn.Module, keep_ratio: float = 0.3, num_attention_sinks: int = 4, heavy_hitter_ratio: float = 0.1, max_error: float = 0.85, max_retries: int = 3, fallback_to_original: bool = True, target_attention_only: bool = True, **kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("zerokv_neo")
@COMPRESSOR_REGISTRY.register("zerokv")
@COMPRESSOR_REGISTRY.register("zerokv_neo_compressor")
class ZeroKVNeoCompressor(BaseCompressor):
    """
    ZeroKV-Neo Compressor: L2-norm weight pruning preserving attention sinks and heavy hitters.

    Args:
        model: PyTorch model.
        keep_ratio: Fraction of output channels to retain (default 0.3).
        num_attention_sinks: Number of initial sink channels pinned (default 4).
        heavy_hitter_ratio: Fraction allocated to heavy hitters (default 0.1).
        max_error: Maximum allowed relative Frobenius error (default 0.85).
        max_retries: Number of keep_ratio retry steps on error threshold excess (default 3).
        fallback_to_original: If True, keep original layer on high error (default True).
        target_attention_only: If True, only compress attention projections (default True for ZeroKV).
    """

    def __init__(
        self,
        model: nn.Module,
        keep_ratio: float = 0.3,
        num_attention_sinks: int = 4,
        heavy_hitter_ratio: float = 0.1,
        max_error: float = 0.85,
        max_retries: int = 3,
        fallback_to_original: bool = True,
        target_attention_only: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="zerokv_neo", **kwargs)
        cfg = kwargs.get("comp_config", kwargs.get("cfg"))
        if cfg is not None:
            keep_ratio = _cfg_get(cfg, "zerokv_keep_ratio", _cfg_get(cfg, "keep_ratio", keep_ratio))
            num_attention_sinks = _cfg_get(
                cfg,
                "zerokv_num_sinks",
                _cfg_get(cfg, "num_attention_sinks", num_attention_sinks),
            )
            heavy_hitter_ratio = _cfg_get(
                cfg,
                "zerokv_heavy_ratio",
                _cfg_get(cfg, "heavy_hitter_ratio", heavy_hitter_ratio),
            )
            max_error = _cfg_get(
                cfg, "max_error", _cfg_get(cfg, "max_reconstruction_error", max_error)
            )
            max_retries = _cfg_get(cfg, "max_retries", max_retries)
            fallback_to_original = _cfg_get(cfg, "fallback_to_original", fallback_to_original)
            target_attention_only = _cfg_get(cfg, "target_attention_only", target_attention_only)

        if "keep_ratio" in kwargs and kwargs["keep_ratio"] is not None:
            keep_ratio = kwargs["keep_ratio"]
        if "max_error" in kwargs and kwargs["max_error"] is not None:
            max_error = float(kwargs["max_error"])
        if "max_retries" in kwargs and kwargs["max_retries"] is not None:
            max_retries = int(kwargs["max_retries"])
        if "fallback_to_original" in kwargs:
            fallback_to_original = bool(kwargs["fallback_to_original"])
        if "target_attention_only" in kwargs:
            target_attention_only = bool(kwargs["target_attention_only"])

        self.keep_ratio = float(keep_ratio)
        self.num_attention_sinks = int(num_attention_sinks)
        self.heavy_hitter_ratio = float(heavy_hitter_ratio)
        self.max_error = float(max_error)
        self.max_retries = max(0, int(max_retries))
        self.fallback_to_original = bool(fallback_to_original)
        self.target_attention_only = target_attention_only
        self.compressed_layers: list[dict[str, Any]] = []
        self.fallback_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        # Protect critical classification heads, poolers, lm_heads, and embeddings
        if any(
            k in name.lower() for k in ["classifier", "score", "lm_head", "pooler", "embedding"]
        ):
            return False

        base_layer = getattr(module, "base_layer", module)
        cls_name = base_layer.__class__.__name__
        is_linear_like = (
            isinstance(base_layer, nn.Linear)
            or "conv1d" in cls_name.lower()
            or hasattr(base_layer, "nf")
            or (
                hasattr(base_layer, "weight")
                and hasattr(base_layer.weight, "ndim")
                and base_layer.weight.ndim == 2
                and not isinstance(base_layer, (nn.Embedding, nn.LayerNorm, nn.RMSNorm))
            )
        )
        if not is_linear_like:
            return False
        if not should_compress_linear(self.model, name):
            return False
        if self.target_attention_only:
            return self._is_attention_proj(name)
        return True

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Compress linear projection weights preserving attention sinks and heavy hitters."""
        self.compressed_layers = []
        self.fallback_layers = []
        self.layer_errors = {}
        total_error = 0.0

        for name, module in list(self.model.named_modules()):
            if not self._is_target_module(name, module):
                continue

            base_layer = getattr(module, "base_layer", module)
            raw_w = base_layer.weight.data
            orig_dtype = raw_w.dtype
            orig_device = raw_w.device
            cls_name = base_layer.__class__.__name__

            is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_layer, "nf")

            if is_conv1d:
                W = raw_w.t().float().contiguous()
                in_f = W.shape[1]
                out_f = W.shape[0]
            else:
                W = raw_w.float().contiguous()
                out_f, in_f = W.shape[0], W.shape[1]

            norm_orig = torch.linalg.norm(W).item() + 1e-8
            row_norms = torch.norm(W, dim=1)

            cur_keep = self.keep_ratio
            best_mask = None
            best_err = float("inf")
            best_keep = cur_keep
            best_sinks = 0
            best_heavy = 0

            for attempt in range(self.max_retries + 1):
                total_keep = max(1, min(int(round(out_f * cur_keep)), out_f))
                n_sinks = min(self.num_attention_sinks, total_keep)
                n_heavy = total_keep - n_sinks

                sink_idx = torch.arange(n_sinks, device=orig_device)
                remaining_mask = torch.ones(out_f, dtype=torch.bool, device=orig_device)
                remaining_mask[sink_idx] = False

                if n_heavy > 0 and remaining_mask.any():
                    rem_indices = torch.where(remaining_mask)[0]
                    rem_norms = row_norms[rem_indices]
                    _, top_rem = torch.topk(rem_norms, min(n_heavy, rem_norms.numel()))
                    heavy_idx = rem_indices[top_rem]
                else:
                    heavy_idx = torch.tensor([], device=orig_device, dtype=torch.long)

                keep_idx = torch.cat([sink_idx, heavy_idx])
                mask = torch.zeros(out_f, device=orig_device, dtype=torch.float32)
                mask[keep_idx] = 1.0

                W_pruned = W * mask.unsqueeze(1)
                rel_err = (torch.linalg.norm(W - W_pruned).item()) / norm_orig

                if rel_err < best_err:
                    best_err = rel_err
                    best_mask = mask
                    best_keep = cur_keep
                    best_sinks = n_sinks
                    best_heavy = len(heavy_idx)

                if rel_err <= self.max_error or cur_keep >= 1.0:
                    break

                step = max(0.1, (1.0 - cur_keep) / max(1, self.max_retries - attempt))
                cur_keep = min(1.0, cur_keep + step)

            if best_err > self.max_error:
                logger.warning(
                    "ZeroKV-Neo: %s error (%.4f) exceeded threshold (%.4f).",
                    name,
                    best_err,
                    self.max_error,
                )
                if self.fallback_to_original:
                    logger.info("ZeroKV-Neo: Keeping original module for %s.", name)
                    self.fallback_layers.append(name)
                    for p in module.parameters():
                        p.requires_grad = False
                    continue

            if best_mask is None:
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                continue

            # Assign pruned weights to base_layer
            if is_conv1d:
                W_assigned = (
                    (W * best_mask.unsqueeze(1))
                    .t()
                    .contiguous()
                    .to(dtype=orig_dtype, device=orig_device)
                )
            else:
                W_assigned = (W * best_mask.unsqueeze(1)).to(dtype=orig_dtype, device=orig_device)

            if W_assigned.shape != base_layer.weight.shape:
                logger.error(
                    "ZeroKV-Neo shape mismatch for %s: %s vs %s",
                    name,
                    W_assigned.shape,
                    base_layer.weight.shape,
                )
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                continue

            base_layer.weight.data.copy_(W_assigned)
            # ZeroKV-Neo compresses in-place: base_layer already holds the pruned weight
            compressed_layer = base_layer
            # Preserve existing LoRA adapter weights if the source module has them
            if (
                hasattr(module, "lora_A")
                and hasattr(module, "lora_B")
                and module.lora_A is not None
                and module.lora_B is not None
            ):
                src_lA = getattr(module.lora_A, "weight", module.lora_A)
                src_lB = getattr(module.lora_B, "weight", module.lora_B)
                if hasattr(compressed_layer, "lora_A") and compressed_layer.lora_A is not None:
                    if hasattr(src_lA, "data") and src_lA.shape == compressed_layer.lora_A.shape:
                        compressed_layer.lora_A.data.copy_(src_lA.data)
                if hasattr(compressed_layer, "lora_B") and compressed_layer.lora_B is not None:
                    if hasattr(src_lB, "data") and src_lB.shape == compressed_layer.lora_B.shape:
                        compressed_layer.lora_B.data.copy_(src_lB.data)

            if hasattr(module, "base_layer"):
                module.base_layer = compressed_layer  # put compressed layer inside wrapper
                # module itself (LoRALinear wrapper) remains the replacement
            else:
                pass
            if getattr(base_layer, "bias", None) is not None and base_layer.bias is not None:
                b_pruned = (base_layer.bias.data.float() * best_mask).to(
                    dtype=orig_dtype, device=orig_device
                )
                base_layer.bias.data.copy_(b_pruned)

            self.compressed_layers.append(
                {
                    "name": name,
                    "original_shape": (in_f, out_f),
                    "keep_ratio": float(best_keep),
                    "num_sinks": best_sinks,
                    "num_heavy": best_heavy,
                    "rel_error": float(best_err),
                }
            )
            self.layer_errors[name] = float(best_err)
            total_error += best_err

        self._compressed_params = self._count_params(self.model)
        stats = self.get_compression_stats()
        num_comp = len(self.compressed_layers)
        num_fb = len(self.fallback_layers)
        avg_err = (total_error / num_comp) if num_comp else 0.0
        stats.update(
            {
                "compressed_layers": self.compressed_layers,
                "num_compressed": num_comp,
                "num_compressed_layers": num_comp,
                "fallback_layers": self.fallback_layers,
                "num_fallback": num_fb,
                "num_fallback_layers": num_fb,
                "fallback_layer_count": num_fb,
                "layer_errors": self.layer_errors,
                "keep_ratio": self.keep_ratio,
                "num_attention_sinks": self.num_attention_sinks,
                "heavy_hitter_ratio": self.heavy_hitter_ratio,
                "avg_reconstruction_error": avg_err,
                "max_error": self.max_error,
                "fallback_to_original": self.fallback_to_original,
                "kv_cache_budget_saving": f"{(1.0 - self.keep_ratio) * 100:.1f}% Footprint Reduction",
            }
        )
        logger.info(
            "ZeroKV-Neo compression finished: %d layers compressed, avg error: %.4f",
            len(self.compressed_layers),
            avg_err,
        )
        return self.model, stats

    @staticmethod
    def _is_attention_proj(name: str) -> bool:
        """Check if module name corresponds to a QKV attention projection (not output projection).

        Explicitly excludes output projections first, then checks for QKV attention keywords.
        This prevents generic keywords like 'attn'/'attention' from matching out_proj, c_proj, etc.
        """
        name_lower = name.lower()

        # Exclusion list: output projections must never be treated as QKV targets
        output_proj_keywords = [
            "out_lin",
            "out_proj",
            "o_proj",
            "c_proj",
            "output.dense",
            "self.output",
            ".output",
        ]
        if any(kw in name_lower for kw in output_proj_keywords):
            return False

        # Inclusion list: only match true QKV projections
        qkv_keywords = [
            "q_lin",
            "k_lin",
            "v_lin",
            "q_proj",
            "k_proj",
            "v_proj",
            "query",
            "key",
            "value",
            "c_attn",
            "in_proj_qkv",
        ]
        return any(kw in name_lower for kw in qkv_keywords)

ZeroKV-Neo Compressor: L2-norm weight pruning preserving attention sinks and heavy hitters.

Args
-----=
model: PyTorch model.
keep_ratio: Fraction of output channels to retain (default 0.3).
num_attention_sinks: Number of initial sink channels pinned (default 4).
heavy_hitter_ratio: Fraction allocated to heavy hitters (default 0.1).
max_error: Maximum allowed relative Frobenius error (default 0.85).
max_retries: Number of keep_ratio retry steps on error threshold excess (default 3).
fallback_to_original: If True, keep original layer on high error (default True).
target_attention_only: If True, only compress attention projections (default True for ZeroKV).

Ancestors

BaseCompressor

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Compress linear projection weights preserving attention sinks and heavy hitters."""
    self.compressed_layers = []
    self.fallback_layers = []
    self.layer_errors = {}
    total_error = 0.0

    for name, module in list(self.model.named_modules()):
        if not self._is_target_module(name, module):
            continue

        base_layer = getattr(module, "base_layer", module)
        raw_w = base_layer.weight.data
        orig_dtype = raw_w.dtype
        orig_device = raw_w.device
        cls_name = base_layer.__class__.__name__

        is_conv1d = "conv1d" in cls_name.lower() or hasattr(base_layer, "nf")

        if is_conv1d:
            W = raw_w.t().float().contiguous()
            in_f = W.shape[1]
            out_f = W.shape[0]
        else:
            W = raw_w.float().contiguous()
            out_f, in_f = W.shape[0], W.shape[1]

        norm_orig = torch.linalg.norm(W).item() + 1e-8
        row_norms = torch.norm(W, dim=1)

        cur_keep = self.keep_ratio
        best_mask = None
        best_err = float("inf")
        best_keep = cur_keep
        best_sinks = 0
        best_heavy = 0

        for attempt in range(self.max_retries + 1):
            total_keep = max(1, min(int(round(out_f * cur_keep)), out_f))
            n_sinks = min(self.num_attention_sinks, total_keep)
            n_heavy = total_keep - n_sinks

            sink_idx = torch.arange(n_sinks, device=orig_device)
            remaining_mask = torch.ones(out_f, dtype=torch.bool, device=orig_device)
            remaining_mask[sink_idx] = False

            if n_heavy > 0 and remaining_mask.any():
                rem_indices = torch.where(remaining_mask)[0]
                rem_norms = row_norms[rem_indices]
                _, top_rem = torch.topk(rem_norms, min(n_heavy, rem_norms.numel()))
                heavy_idx = rem_indices[top_rem]
            else:
                heavy_idx = torch.tensor([], device=orig_device, dtype=torch.long)

            keep_idx = torch.cat([sink_idx, heavy_idx])
            mask = torch.zeros(out_f, device=orig_device, dtype=torch.float32)
            mask[keep_idx] = 1.0

            W_pruned = W * mask.unsqueeze(1)
            rel_err = (torch.linalg.norm(W - W_pruned).item()) / norm_orig

            if rel_err < best_err:
                best_err = rel_err
                best_mask = mask
                best_keep = cur_keep
                best_sinks = n_sinks
                best_heavy = len(heavy_idx)

            if rel_err <= self.max_error or cur_keep >= 1.0:
                break

            step = max(0.1, (1.0 - cur_keep) / max(1, self.max_retries - attempt))
            cur_keep = min(1.0, cur_keep + step)

        if best_err > self.max_error:
            logger.warning(
                "ZeroKV-Neo: %s error (%.4f) exceeded threshold (%.4f).",
                name,
                best_err,
                self.max_error,
            )
            if self.fallback_to_original:
                logger.info("ZeroKV-Neo: Keeping original module for %s.", name)
                self.fallback_layers.append(name)
                for p in module.parameters():
                    p.requires_grad = False
                continue

        if best_mask is None:
            self.fallback_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False
            continue

        # Assign pruned weights to base_layer
        if is_conv1d:
            W_assigned = (
                (W * best_mask.unsqueeze(1))
                .t()
                .contiguous()
                .to(dtype=orig_dtype, device=orig_device)
            )
        else:
            W_assigned = (W * best_mask.unsqueeze(1)).to(dtype=orig_dtype, device=orig_device)

        if W_assigned.shape != base_layer.weight.shape:
            logger.error(
                "ZeroKV-Neo shape mismatch for %s: %s vs %s",
                name,
                W_assigned.shape,
                base_layer.weight.shape,
            )
            self.fallback_layers.append(name)
            for p in module.parameters():
                p.requires_grad = False
            continue

        base_layer.weight.data.copy_(W_assigned)
        # ZeroKV-Neo compresses in-place: base_layer already holds the pruned weight
        compressed_layer = base_layer
        # Preserve existing LoRA adapter weights if the source module has them
        if (
            hasattr(module, "lora_A")
            and hasattr(module, "lora_B")
            and module.lora_A is not None
            and module.lora_B is not None
        ):
            src_lA = getattr(module.lora_A, "weight", module.lora_A)
            src_lB = getattr(module.lora_B, "weight", module.lora_B)
            if hasattr(compressed_layer, "lora_A") and compressed_layer.lora_A is not None:
                if hasattr(src_lA, "data") and src_lA.shape == compressed_layer.lora_A.shape:
                    compressed_layer.lora_A.data.copy_(src_lA.data)
            if hasattr(compressed_layer, "lora_B") and compressed_layer.lora_B is not None:
                if hasattr(src_lB, "data") and src_lB.shape == compressed_layer.lora_B.shape:
                    compressed_layer.lora_B.data.copy_(src_lB.data)

        if hasattr(module, "base_layer"):
            module.base_layer = compressed_layer  # put compressed layer inside wrapper
            # module itself (LoRALinear wrapper) remains the replacement
        else:
            pass
        if getattr(base_layer, "bias", None) is not None and base_layer.bias is not None:
            b_pruned = (base_layer.bias.data.float() * best_mask).to(
                dtype=orig_dtype, device=orig_device
            )
            base_layer.bias.data.copy_(b_pruned)

        self.compressed_layers.append(
            {
                "name": name,
                "original_shape": (in_f, out_f),
                "keep_ratio": float(best_keep),
                "num_sinks": best_sinks,
                "num_heavy": best_heavy,
                "rel_error": float(best_err),
            }
        )
        self.layer_errors[name] = float(best_err)
        total_error += best_err

    self._compressed_params = self._count_params(self.model)
    stats = self.get_compression_stats()
    num_comp = len(self.compressed_layers)
    num_fb = len(self.fallback_layers)
    avg_err = (total_error / num_comp) if num_comp else 0.0
    stats.update(
        {
            "compressed_layers": self.compressed_layers,
            "num_compressed": num_comp,
            "num_compressed_layers": num_comp,
            "fallback_layers": self.fallback_layers,
            "num_fallback": num_fb,
            "num_fallback_layers": num_fb,
            "fallback_layer_count": num_fb,
            "layer_errors": self.layer_errors,
            "keep_ratio": self.keep_ratio,
            "num_attention_sinks": self.num_attention_sinks,
            "heavy_hitter_ratio": self.heavy_hitter_ratio,
            "avg_reconstruction_error": avg_err,
            "max_error": self.max_error,
            "fallback_to_original": self.fallback_to_original,
            "kv_cache_budget_saving": f"{(1.0 - self.keep_ratio) * 100:.1f}% Footprint Reduction",
        }
    )
    logger.info(
        "ZeroKV-Neo compression finished: %d layers compressed, avg error: %.4f",
        len(self.compressed_layers),
        avg_err,
    )
    return self.model, stats

Compress linear projection weights preserving attention sinks and heavy hitters.

Inherited members

BaseCompressor:
- get_compression_stats
- summary