| """ |
| lora.py — LoRA adapter for KaizenLM Phase 5 online learning. |
| |
| Architecture: rank-4 LoRA on Q and V projections of all 8 attention layers. |
| Q slice: qkv_out[..., :D] ← corrected by lora_q |
| V slice: qkv_out[..., 2*D:] ← corrected by lora_v |
| |
| Adapter params per layer: 2 × (rank×D + D×rank) = 2 × (4×512 + 512×4) = 8,192 |
| Total (8 layers): 65,536 params |
| Storage per adapter: 65,536 × 4 bytes ≈ 256 KB |
| |
| Usage: |
| adapter = LoRAAdapter(n_layers=8, d=512, rank=4, alpha=8) |
| with torch.no_grad(): |
| out = model_with_lora_forward(base_model, x, adapter) |
| loss.backward() # gradients only on adapter params |
| optimizer.step() |
| """ |
| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class LoRALinear(nn.Module): |
| """Additive low-rank correction: output += B(A(x)) * scale.""" |
|
|
| def __init__(self, in_features: int, out_features: int, |
| rank: int = 4, alpha: float = 8.0): |
| super().__init__() |
| self.rank = rank |
| self.scale = alpha / rank |
| self.A = nn.Linear(in_features, rank, bias=False) |
| self.B = nn.Linear(rank, out_features, bias=False) |
| nn.init.kaiming_uniform_(self.A.weight, a=math.sqrt(5)) |
| nn.init.zeros_(self.B.weight) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.B(self.A(x)) * self.scale |
|
|
|
|
| class LoRAAdapter(nn.Module): |
| """ |
| Per-layer LoRA corrections for KaizenLM. |
| Applies to Q and V slices of the fused QKV projection in each attention block. |
| """ |
|
|
| def __init__(self, n_layers: int = 8, d: int = 512, |
| rank: int = 4, alpha: float = 8.0): |
| super().__init__() |
| self.d = d |
| self.lora_q = nn.ModuleList( |
| [LoRALinear(d, d, rank, alpha) for _ in range(n_layers)]) |
| self.lora_v = nn.ModuleList( |
| [LoRALinear(d, d, rank, alpha) for _ in range(n_layers)]) |
|
|
| def correct_qkv(self, layer_idx: int, x: torch.Tensor, |
| qkv: torch.Tensor) -> torch.Tensor: |
| """ |
| Add LoRA corrections to Q and V slices. |
| qkv shape: [B, T, 3*D] |
| Uses torch.cat (not in-place) to preserve autograd graph for backward. |
| """ |
| D = self.d |
| q_corrected = qkv[..., :D] + self.lora_q[layer_idx](x) |
| k_unchanged = qkv[..., D:2*D] |
| v_corrected = qkv[..., 2*D:] + self.lora_v[layer_idx](x) |
| return torch.cat([q_corrected, k_unchanged, v_corrected], dim=-1) |
|
|
| def param_count(self) -> int: |
| return sum(p.numel() for p in self.parameters()) |
|
|
| def size_bytes(self) -> int: |
| return sum(p.numel() * p.element_size() for p in self.parameters()) |
|
|
| @classmethod |
| def merged(cls, adapters: list, weights: list, |
| n_layers: int = 8, d: int = 512, |
| rank: int = 4, alpha: float = 8.0) -> 'LoRAAdapter': |
| """ |
| Weighted average merge of multiple adapters. |
| weights need not sum to 1 (normalised internally). |
| """ |
| assert len(adapters) == len(weights) and len(adapters) > 0 |
| total = sum(weights) |
| merged = cls(n_layers=n_layers, d=d, rank=rank, alpha=alpha) |
| |
| with torch.no_grad(): |
| for p in merged.parameters(): |
| p.zero_() |
| norm_w = [w / total for w in weights] |
| for adapter, w in zip(adapters, norm_w): |
| for p_m, p_a in zip(merged.parameters(), adapter.parameters()): |
| p_m.add_(p_a * w) |
| return merged |
|
|
|
|
| |
| |
|
|
| def _rotate_half(x: torch.Tensor) -> torch.Tensor: |
| h = x.shape[-1] // 2 |
| return torch.cat([-x[..., h:], x[..., :h]], dim=-1) |
|
|
|
|
| def _apply_rope(q, k, cos, sin): |
| return q * cos + _rotate_half(q) * sin, k * cos + _rotate_half(k) * sin |
|
|
|
|
| def _precompute_rope(head_dim, max_seq, theta=10000.0): |
| inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim)) |
| t = torch.arange(max_seq).float() |
| f = torch.outer(t, inv) |
| return torch.cat([f, f], dim=-1).cos(), torch.cat([f, f], dim=-1).sin() |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, d: int, eps: float = 1e-6): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(d)) |
| self.eps = eps |
|
|
| def forward(self, x): |
| n = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt() |
| return (x.float() * n * self.weight).type_as(x) |
|
|
|
|
| class CausalSelfAttentionWithLoRA(nn.Module): |
| def __init__(self, d: int, n_heads: int, drop: float = 0.0): |
| super().__init__() |
| self.nh = n_heads |
| self.hd = d // n_heads |
| self.drop = drop |
| self.qkv = nn.Linear(d, 3 * d, bias=False) |
| self.out_proj = nn.Linear(d, d, bias=False) |
|
|
| def forward(self, x, cos, sin, lora_adapter=None, layer_idx=None): |
| B, T, C = x.shape |
| qkv = self.qkv(x) |
| if lora_adapter is not None: |
| qkv = lora_adapter.correct_qkv(layer_idx, x, qkv) |
| q, k, v = qkv.split(C, dim=-1) |
| q = q.view(B, T, self.nh, self.hd) |
| k = k.view(B, T, self.nh, self.hd) |
| v = v.view(B, T, self.nh, self.hd) |
| q, k = _apply_rope(q, k, cos[:T].unsqueeze(1), sin[:T].unsqueeze(1)) |
| q = q.transpose(1, 2); k = k.transpose(1, 2); v = v.transpose(1, 2) |
| o = F.scaled_dot_product_attention(q, k, v, is_causal=True, |
| dropout_p=self.drop if self.training else 0.0) |
| return self.out_proj(o.transpose(1, 2).contiguous().view(B, T, C)) |
|
|
|
|
| class FFN(nn.Module): |
| def __init__(self, d: int, d_ff: int): |
| super().__init__() |
| self.w1 = nn.Linear(d, d_ff, bias=False) |
| self.w2 = nn.Linear(d_ff, d, bias=False) |
|
|
| def forward(self, x): |
| return self.w2(F.gelu(self.w1(x), approximate='tanh')) |
|
|
|
|
| class BlockWithLoRA(nn.Module): |
| def __init__(self, d: int, n_heads: int, d_ff: int): |
| super().__init__() |
| self.ln1 = RMSNorm(d) |
| self.attn = CausalSelfAttentionWithLoRA(d, n_heads) |
| self.ln2 = RMSNorm(d) |
| self.ff = FFN(d, d_ff) |
|
|
| def forward(self, x, cos, sin, lora_adapter=None, layer_idx=None): |
| x = x + self.attn(self.ln1(x), cos, sin, lora_adapter, layer_idx) |
| return x + self.ff(self.ln2(x)) |
|
|
|
|
| class KaizenWithLoRA(nn.Module): |
| """ |
| KaizenLM base (frozen) + optional LoRAAdapter (trainable). |
| Base weights are never updated. LoRA adapter is swappable at inference time. |
| """ |
|
|
| VOCAB_SIZE = 32768 |
| D_MODEL = 512 |
| N_HEADS = 8 |
| N_LAYERS = 8 |
| D_FF = 2048 |
| BLOCK_SIZE = 1024 |
|
|
| def __init__(self): |
| super().__init__() |
| d, h, l, ff, b = (self.D_MODEL, self.N_HEADS, self.N_LAYERS, |
| self.D_FF, self.BLOCK_SIZE) |
| self.embed = nn.Embedding(self.VOCAB_SIZE, d) |
| self.blocks = nn.ModuleList( |
| [BlockWithLoRA(d, h, ff) for _ in range(l)]) |
| self.ln_f = RMSNorm(d) |
| self.head = nn.Linear(d, self.VOCAB_SIZE, bias=False) |
| self.head.weight = self.embed.weight |
|
|
| cos, sin = _precompute_rope(d // h, b * 2) |
| self.register_buffer('rope_cos', cos) |
| self.register_buffer('rope_sin', sin) |
|
|
| def load_base(self, ckpt_path: str): |
| """Load base weights, freeze all base params.""" |
| ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=True) |
| sd = ckpt.get('model_state', ckpt.get('model', ckpt)) |
| self.load_state_dict(sd, strict=True) |
| for p in self.parameters(): |
| p.requires_grad_(False) |
|
|
| def forward(self, idx: torch.Tensor, targets=None, |
| adapter: 'LoRAAdapter' = None) -> torch.Tensor: |
| x = self.embed(idx) |
| cos = self.rope_cos[:idx.shape[1]] |
| sin = self.rope_sin[:idx.shape[1]] |
| for i, blk in enumerate(self.blocks): |
| x = blk(x, cos, sin, lora_adapter=adapter, layer_idx=i) |
| logits = self.head(self.ln_f(x)) |
| if targets is None: |
| return logits |
| loss = F.cross_entropy( |
| logits.view(-1, logits.size(-1)), |
| targets.view(-1), |
| ignore_index=-100, |
| ) |
| return logits, loss |
|
|
| def embed_task(self, idx: torch.Tensor, |
| adapter: 'LoRAAdapter' = None) -> torch.Tensor: |
| """ |
| Returns mean-pooled last hidden state (before lm_head) as task embedding. |
| Shape: [D_MODEL] |
| """ |
| cos = self.rope_cos[:idx.shape[1]] |
| sin = self.rope_sin[:idx.shape[1]] |
| with torch.no_grad(): |
| x = self.embed(idx) |
| for i, blk in enumerate(self.blocks): |
| x = blk(x, cos, sin, lora_adapter=adapter, layer_idx=i) |
| x = self.ln_f(x) |
| return x[0].mean(0) |
|
|
