Fine-tuning

import torch
import torch.nn as nn

class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, r: int = 16, alpha: int = 32, dropout: float = 0.05):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad = False          # freeze the pretrained weight entirely

        self.A = nn.Parameter(torch.randn(r, base.in_features) * 0.01)   # small random init
        self.B = nn.Parameter(torch.zeros(base.out_features, r))         # zero init -> no-op at start
        self.scale = alpha / r
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        base_out = self.base(x)                              # frozen path, no grad
        lora_out = self.drop(x) @ self.A.T @ self.B.T         # low-rank path, trainable
        return base_out + lora_out * self.scale

    @torch.no_grad()
    def merge_(self):
        """Fold the adapter into the base weight for zero-overhead inference."""
        self.base.weight.data += (self.B @ self.A) * self.scale

    def trainable_parameters(self):
        return [self.A, self.B]


def inject_lora(model: nn.Module, target_names=("q_proj", "k_proj", "v_proj", "o_proj"), r=16, alpha=32):
    """Walk the model and replace target nn.Linear modules with LoRALinear in place."""
    for name, module in model.named_modules():
        for child_name, child in module.named_children():
            if any(t in child_name for t in target_names) and isinstance(child, nn.Linear):
                setattr(module, child_name, LoRALinear(child, r=r, alpha=alpha))
    return model