Module neuralops.utils.compat.param_recovery

Parameter recovery selection and budget pruning for compressed models.

Functions

def get_recovery_params(model: nn.Module) ‑> list[torch.nn.parameter.Parameter]

Expand source code

def get_recovery_params(model: nn.Module) -> list[nn.Parameter]:
    """Select recovery parameters for LoRA / SVD / injected attention modules.

    Fix CP-01: SVDLinear U/V factor matrices are now preserved during the
    parameter budget pruning step. Previously, their numel (e.g. 64 * 768 =
    49152) exceeded the 32768 threshold and they were pruned, reducing recovery
    to bias-only training.

    Fix CP-02: Pruning now uses priority-based allocation:
        Priority 0 (highest): protected params (LoRA / gates / classifier / alpha / beta)
        Priority 1: SVDLinear U/V factors and factorized parameters
        Priority 2: small params (1D, or numel <= 32768)
        Priority 3 (lowest): other large 2D params

    Args:
        model: The PyTorch neural network module.

    Returns:
        List of parameters chosen for gradient optimization.
    """
    svd_modules = [
        m
        for m in model.modules()
        if m.__class__.__name__ in ("SVDLinear", "SVDLinearLayer", "FactorisedLinear")
    ]
    svd_param_ids = {id(p) for m in svd_modules for p in m.parameters()}
    has_svd = len(svd_modules) > 0

    has_lora = any(any(t in n.lower() for t in LORA_KEYWORDS) for n, _ in model.named_parameters())

    trainable_params = _collect_trainable_params(
        model, svd_param_ids=svd_param_ids, has_svd=has_svd, has_lora=has_lora
    )

    # Set requires_grad across the model.
    trainable_id_set = {id(p) for p in trainable_params}
    for _, param in model.named_parameters():
        param.requires_grad = id(param) in trainable_id_set

    total_trainable = sum(p.numel() for p in trainable_params)
    logger.info(
        "Recovery fine-tuning selected %d trainable parameters across %d tensors.",
        total_trainable,
        len(trainable_params),
    )

    is_nested_or_deep = any(
        "encoder" in n.lower() or "roberta" in n.lower() or "bert" in n.lower()
        for n, _ in model.named_parameters()
    )
    prune_threshold = 10_000_000 if (has_lora or is_nested_or_deep) else 2_000_000

    # Fix CP-03: Shrink LoRA rank properly, preserving the r/alpha scaling.
    if total_trainable > 5_000_000:
        for m in model.modules():
            if hasattr(m, "lora_rank") and getattr(m, "lora_rank", 0) > 8:
                old_r = m.lora_rank
                m.lora_rank = 8
                if hasattr(m, "lora_alpha"):
                    m.lora_alpha = m.lora_alpha * (8.0 / max(1, old_r))
                if hasattr(m, "lora_scaling") and hasattr(m, "lora_alpha"):
                    m.lora_scaling = m.lora_alpha / 8

    # Fix CP-01/CP-02: Priority-based pruning.
    if total_trainable > prune_threshold:
        logger.warning(
            "Total trainable parameters (%d) exceeds %d threshold for recovery. "
            "Pruning with priority-based allocation.",
            total_trainable,
            prune_threshold,
        )
        trainable_params = _prune_to_budget(
            model=model,
            trainable_id_set=trainable_id_set,
            svd_param_ids=svd_param_ids,
            prune_threshold=prune_threshold,
        )
        total_trainable = sum(p.numel() for p in trainable_params)
        logger.info(
            "Reduced recovery parameters to %d trainable parameters across %d tensors.",
            total_trainable,
            len(trainable_params),
        )

    logger.info(
        "Discovered %d trainable tensors (%d params)",
        len(trainable_params),
        total_trainable,
    )
    for name, param in model.named_parameters():
        if param.requires_grad:
            logger.debug("  %s: %s (%d params)", name, tuple(param.shape), param.numel())

    return trainable_params

Select recovery parameters for LoRA / SVD / injected attention modules.

Fix CP-01: SVDLinear U/V factor matrices are now preserved during the parameter budget pruning step. Previously, their numel (e.g. 64 * 768 = 49152) exceeded the 32768 threshold and they were pruned, reducing recovery to bias-only training.

Fix CP-02: Pruning now uses priority-based allocation: Priority 0 (highest): protected params (LoRA / gates / classifier / alpha / beta) Priority 1: SVDLinear U/V factors and factorized parameters Priority 2: small params (1D, or numel <= 32768) Priority 3 (lowest): other large 2D params

Args

-----=

model

The PyTorch neural network module.

Returns -----= List of parameters chosen for gradient optimization.