Module neuralops.compression_engine.liftr

LIFTR – Linearized Function Trick Compressor. Registered as 'liftr', 'liftr_compressor', 'liftr_class'.

Mathematical Formulation: Activation: ReLU(x) → Identity(x), GELU(x) / SiLU(x) ≈ α·x Attention: Softmax(QK^T/√d)V → (1/N) Σ V_i

Functions

def liftr_compressor_entry(model: nn.Module, comp_config: Any = None, **kwargs: Any) ‑> Any

Expand source code

@COMPRESSOR_REGISTRY.register("liftr")
@COMPRESSOR_REGISTRY.register("liftr_compressor")
@COMPRESSOR_REGISTRY.register("liftr_factory")
def liftr_compressor_entry(
    model: nn.Module,
    comp_config: Any = None,
    **kwargs: Any,
) -> Any:
    """Factory function for LIFTR compressor."""
    replace_attention = _cfg_get(
        comp_config,
        "liftr_replace_attention",
        _cfg_get(comp_config, "replace_attention", kwargs.pop("replace_attention", False)),
    )
    gelu_coefficient = _cfg_get(
        comp_config,
        "liftr_gelu_coefficient",
        _cfg_get(comp_config, "gelu_coefficient", kwargs.pop("gelu_coefficient", 0.8)),
    )
    target_attention_only = _cfg_get(
        comp_config,
        "liftr_target_attention_only",
        _cfg_get(
            comp_config,
            "target_attention_only",
            kwargs.pop("target_attention_only", False),
        ),
    )
    skip_filter = _cfg_get(
        comp_config,
        "skip_should_compress_filter",
        kwargs.pop("skip_should_compress_filter", False),
    )
    return_stats = kwargs.pop("return_stats", kwargs.pop("return_tuple", False))

    compressor = LIFTRCompressor(
        model,
        replace_attention=replace_attention,
        gelu_coefficient=gelu_coefficient,
        target_attention_only=target_attention_only,
        skip_should_compress_filter=skip_filter,
        comp_config=comp_config,
        **kwargs,
    )
    compressed_model, stats = compressor.compress()
    if return_stats:
        return compressed_model, stats
    return compressed_model

Factory function for LIFTR compressor.

Classes

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