KappaTune-LoRA feature

[oswaldoludwig]

Feature request

Is your feature request related to a problem? Please describe.
Manually choosing target_modules for LoRA (e.g. linear_qkv, linear_proj, etc.) becomes impractical on very large models (70B+). The recent KappaTune paper introduces a simple, principled way to automatically select the best target modules: compute the matrix condition number κ (σ_max / σ_min) for every nn.Linear layer and keep the most isotropic ones (lowest κ). This gives better or comparable performance to full LoRA while using fewer modules.

A clean implementation of KappaTuneSelector (with full 4-bit bitsandbytes support) has already been contributed to Hugging Face PEFT:
→ huggingface/peft#3106

Describe the solution you'd like
Add a lightweight utility similar to KappaTuneSelector (or integrate it directly into NeMo's LoRA/PEFT recipe system) so users can do:

from nemo.collections.llm import LoRA 

# After loading your NeMo model
from nemo.collections import llm
from kappatune_selector_nemo import find_kappa_target_modules   

selector = KappaTuneSelector(model)          # or use find_kappa_target_modules
stable_modules = find_kappa_target_modules(model, top_p=0.2)

lora_config = llm.peft.LoRA(
    target_modules=stable_modules,   # ← NeMo accepts the full names automatically
    dim=64,                          # rank
    alpha=128,
    dropout=0.05,
)

[oswaldoludwig]

# kappatune_selector_nemo.py
# =============================================================================
# KappaTuneSelector for NVIDIA NeMo (Megatron-Core)
# Automatic LoRA target selection based on matrix condition number (κ)
# Works with NeMo's linear_qkv, linear_proj, linear_fc1, linear_fc2, etc.
# Fully compatible with 4-bit quantization and TP/PP models.
# See: https://arxiv.org/abs/2506.16289
# =============================================================================

import torch
import torch.nn as nn
from typing import List, Optional, Dict

try:
    import bitsandbytes as bnb
except ImportError:
    bnb = None


class KappaTuneSelector:
    """
    Lightweight utility to compute per-module condition numbers (κ = σ_max / σ_min)
    and return the best LoRA target modules. Now supports bnb 4-bit models.
    Lower κ = more isotropic weight matrix = better for LoRA.

    Works out-of-the-box with NeMo/Megatron models (linear_qkv, linear_fc1, etc.).
    """

    def __init__(self, model: nn.Module, max_dim_size_to_analyze: int = 16384):
        """
        Args:
            model: The base model (Megatron model, any PEFT-compatible model, including 4-bit).
            max_dim_size_to_analyze: Skip any matrix larger than this to avoid OOM
                (default 16384 is safe on consumer GPUs).
        """
        self.model = model
        self.max_dim_size_to_analyze = max_dim_size_to_analyze
        self._condition_numbers: Optional[Dict[str, float]] = None

    def _compute_kappas(self) -> None:
        if self._condition_numbers is not None:
            return

        condition_numbers: Dict[str, float] = {}

        for module_name, module in self.model.named_modules():
            if not isinstance(module, nn.Linear):
                continue

            # Handle bnb 4-bit quantization (for QLoRA / paper example)
            weight = module.weight
            if bnb is not None and hasattr(weight, "quant_state"):
                try:
                    w = bnb.functional.dequantize_4bit(weight.data, weight.quant_state).float()
                except Exception:
                    w = weight.data.detach().float()
            else:
                w = weight.data.detach().float()

            # Skip huge matrices
            if any(dim > self.max_dim_size_to_analyze for dim in w.shape):
                continue

            try:
                S = torch.linalg.svdvals(w.view(w.size(0), -1))
                kappa = (S[0] / (S[-1] + 1e-8)).item()
                condition_numbers[module_name] = kappa
            except (torch.linalg.LinAlgError, RuntimeError):
                condition_numbers[module_name] = float("inf")

        self._condition_numbers = condition_numbers

    def get_best_targets(
        self,
        top_p: Optional[float] = None,
        num_modules: Optional[int] = None,
        threshold: Optional[float] = None,
    ) -> List[str]:
        """
        Return the best target modules according to one of three mutually-exclusive strategies.

        Args:
            top_p: Return the top X% of best modules (e.g. 0.2 = paper default).
            num_modules: Return exactly this many best modules (fixed budget).
            threshold: Return every module with κ ≤ threshold (quality cutoff).

        Returns:
            List of module names ready for NeMo's LoRA (e.g. ["decoder.layers.0.self_attention.linear_qkv", ...])

        Notes:
            - Precedence: `num_modules` → `top_p` → `threshold` → all modules.
            - Modules are always sorted by ascending κ (lowest = best).
            - Recommended: `top_p=0.2` for most models.
        """
        self._compute_kappas()

        if not self._condition_numbers:
            return []

        sorted_modules = sorted(self._condition_numbers.items(), key=lambda x: x[1])

        if num_modules is not None:
            k = min(num_modules, len(sorted_modules))
            return [name for name, _ in sorted_modules[:k]]

        if top_p is not None:
            k = max(1, int(len(sorted_modules) * top_p))
            return [name for name, _ in sorted_modules[:k]]

        if threshold is not None:
            return [name for name, kappa in sorted_modules if kappa <= threshold]

        return [name for name, _ in sorted_modules]


def find_kappa_target_modules(
    model: nn.Module, top_p: float = 0.2, max_dim_size_to_analyze: int = 16384
) -> List[str]:
    """
    One-liner convenience function (recommended for most users).
    """
    selector = KappaTuneSelector(model, max_dim_size_to_analyze)
    return selector.get_best_targets(top_p=top_p)