| """Adapt only the selected decoder projections with rank-32 LoRA.""" |
| from __future__ import annotations |
| import math |
| from collections.abc import Iterator, Mapping |
| from contextlib import contextmanager |
| from typing import Final |
| import torch |
| from torch import Tensor, nn |
| from torch.nn import functional |
| from .backbone import Generator |
| PROJECTIONS: Final = frozenset(('q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj')) |
| RANK: Final = 32 |
| SCALE: Final = 2.0 |
|
|
| class AdaptedLinear(nn.Module): |
|
|
| def __init__(self, original: nn.Linear) -> None: |
| super().__init__() |
| self.original = original |
| self.original.requires_grad_(False) |
| self.factor_a = nn.Parameter(torch.empty(RANK, original.in_features, dtype=torch.float32, device=original.weight.device)) |
| self.factor_b = nn.Parameter(torch.zeros(original.out_features, RANK, dtype=torch.float32, device=original.weight.device)) |
| nn.init.kaiming_uniform_(self.factor_a, a=math.sqrt(5)) |
| self.is_enabled = True |
| self.inference_scale = 1.0 |
|
|
| def forward(self, inputs: Tensor) -> Tensor: |
| output: Tensor = self.original(inputs) |
| if not self.is_enabled: |
| return output |
| adapted = functional.linear(functional.linear(inputs.to(self.factor_a.dtype), self.factor_a), self.factor_b) |
| return output + (adapted * (SCALE * self.inference_scale)).to(output.dtype) |
|
|
| def merged_weight(self) -> Tensor: |
| delta = self.factor_b.float() @ self.factor_a.float() * (SCALE * self.inference_scale) |
| return (self.original.weight.float() + delta).to(self.original.weight.dtype) |
|
|
| def inject_adapter(model: Generator) -> tuple[str, ...]: |
| model.requires_grad_(False) |
| selected: list[str] = [] |
| for name, layer in tuple(model.named_modules()): |
| if isinstance(layer, nn.Linear) and name.rsplit('.', 1)[-1] in PROJECTIONS: |
| parent_name, attribute = name.rsplit('.', 1) |
| model.get_submodule(parent_name).set_submodule(attribute, AdaptedLinear(layer)) |
| selected.append(name) |
| expected = model.geometry.layers * len(PROJECTIONS) |
| trainable = sum((parameter.numel() for parameter in model.parameters() if parameter.requires_grad)) |
| if len(selected) != expected or trainable != model.geometry.adapter_parameters: |
| message = f'adapter contract failed: projections={len(selected)}/{expected}' |
| message += f', parameters={trainable}/{model.geometry.adapter_parameters}' |
| raise ValueError(message) |
| return tuple(selected) |
|
|
| def adapter_state(model: Generator) -> dict[str, Tensor]: |
| return {name: parameter.detach().cpu().clone() for name, parameter in model.named_parameters() if parameter.requires_grad} |
|
|
| def restore_adapter(model: Generator, state: Mapping[str, Tensor]) -> None: |
| expected = {name: parameter for name, parameter in model.named_parameters() if parameter.requires_grad} |
| if set(state) != set(expected): |
| message = f'adapter keys mismatch: missing={set(expected) - set(state)}' |
| message += f', extra={set(state) - set(expected)}' |
| raise ValueError(message) |
| with torch.no_grad(): |
| for name, parameter in expected.items(): |
| tensor = state[name] |
| if tensor.shape != parameter.shape or not bool(torch.isfinite(tensor).all().item()): |
| message = f'invalid adapter tensor {name}: shape={tensor.shape}' |
| message += f', expected={parameter.shape}' |
| raise ValueError(message) |
| parameter.copy_(tensor) |
|
|
| def merge_adapter(model: Generator) -> None: |
| for name, layer in tuple(model.named_modules()): |
| if isinstance(layer, AdaptedLinear): |
| with torch.no_grad(): |
| layer.original.weight.copy_(layer.merged_weight()) |
| parent, attribute = name.rsplit('.', 1) |
| model.get_submodule(parent).set_submodule(attribute, layer.original) |
|
|
| @contextmanager |
| def disabled_adapter(model: Generator) -> Iterator[None]: |
| layers = [(layer, layer.is_enabled) for layer in model.modules() if isinstance(layer, AdaptedLinear)] |
| try: |
| for layer, _ in layers: |
| layer.is_enabled = False |
| yield |
| finally: |
| for layer, previous in layers: |
| layer.is_enabled = previous |
|
|
| @contextmanager |
| def scaled_adapter(model: Generator, multiplier: float) -> Iterator[None]: |
| if not math.isfinite(multiplier) or multiplier < 0: |
| message = f'adapter multiplier must be finite and nonnegative, got {multiplier!r}' |
| raise ValueError(message) |
| layers = [(layer, layer.inference_scale) for layer in model.modules() if isinstance(layer, AdaptedLinear)] |
| if not layers and multiplier != 1: |
| message = 'adapter scaling requires unmerged adapter weights' |
| raise ValueError(message) |
| try: |
| for layer, _ in layers: |
| layer.inference_scale = multiplier |
| yield |
| finally: |
| for layer, previous in layers: |
| layer.inference_scale = previous |
|
|
