| """Minimal LoRA for the Stage-2 fake-score critic (no external dependency). |
| |
| ``inject_lora`` replaces selected ``nn.Linear`` layers inside ``model.blocks`` |
| with ``LoRALinear`` (frozen base + trainable low-rank A/B, B zero-initialised so |
| the wrapped model starts identical to the base). Everything outside the blocks |
| (patch/time/text embeddings, output head) stays frozen. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from typing import Iterable |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| TARGET_GROUPS = { |
| "self_attn": ("self_attn.q", "self_attn.k", "self_attn.v", "self_attn.o"), |
| "cross_attn": ("cross_attn.q", "cross_attn.k", "cross_attn.v", "cross_attn.o"), |
| "ffn": ("ffn.0", "ffn.2"), |
| } |
|
|
|
|
| class LoRALinear(nn.Module): |
| def __init__( |
| self, base: nn.Linear, rank: int, alpha: float, dropout: float = 0.0 |
| ) -> None: |
| super().__init__() |
| if rank <= 0: |
| raise ValueError(f"LoRA rank must be positive, got {rank}") |
| self.base = base |
| self.base.requires_grad_(False) |
| self.rank = int(rank) |
| self.alpha = float(alpha) |
| self.scale = self.alpha / self.rank |
| self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() |
| |
| |
| self.lora_A = nn.Parameter(torch.empty(self.rank, base.in_features, device=base.weight.device)) |
| self.lora_B = nn.Parameter(torch.zeros(base.out_features, self.rank, device=base.weight.device)) |
| nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| out = self.base(x) |
| update = nn.functional.linear( |
| nn.functional.linear(self.dropout(x), self.lora_A.to(x.dtype)), |
| self.lora_B.to(x.dtype), |
| ) |
| return out + self.scale * update |
|
|
| def extra_repr(self) -> str: |
| return f"rank={self.rank}, alpha={self.alpha}, base={self.base.in_features}->{self.base.out_features}" |
|
|
|
|
| def resolve_targets(spec: str | Iterable[str]) -> tuple[str, ...]: |
| """'self_attn,ffn' -> module-name suffixes; individual suffixes also accepted.""" |
| names = spec.split(",") if isinstance(spec, str) else list(spec) |
| out: list[str] = [] |
| for name in (n.strip() for n in names if n.strip()): |
| if name == "all": |
| for group in TARGET_GROUPS.values(): |
| out.extend(group) |
| elif name in TARGET_GROUPS: |
| out.extend(TARGET_GROUPS[name]) |
| else: |
| out.append(name) |
| if not out: |
| raise ValueError("No LoRA target modules selected") |
| return tuple(dict.fromkeys(out)) |
|
|
|
|
| def inject_lora( |
| model: nn.Module, |
| *, |
| rank: int, |
| alpha: float | None = None, |
| targets: str | Iterable[str] = "all", |
| dropout: float = 0.0, |
| ) -> list[nn.Parameter]: |
| """Freeze ``model`` and wrap the target Linears of every ``model.blocks[i]``. |
| |
| Returns the trainable LoRA parameters. |
| """ |
| if not hasattr(model, "blocks"): |
| raise ValueError("inject_lora expects a Wan model with a .blocks list") |
| alpha = float(rank if alpha is None else alpha) |
| suffixes = resolve_targets(targets) |
| model.requires_grad_(False) |
| trainable: list[nn.Parameter] = [] |
| replaced = 0 |
| for block in model.blocks: |
| for name in suffixes: |
| parent_name, _, attr = name.rpartition(".") |
| parent = block.get_submodule(parent_name) if parent_name else block |
| child = getattr(parent, attr) if not attr.isdigit() else parent[int(attr)] |
| if not isinstance(child, nn.Linear): |
| raise TypeError(f"{name} is {type(child).__name__}, not nn.Linear") |
| wrapped = LoRALinear(child, rank, alpha, dropout) |
| if attr.isdigit(): |
| parent[int(attr)] = wrapped |
| else: |
| setattr(parent, attr, wrapped) |
| trainable.extend([wrapped.lora_A, wrapped.lora_B]) |
| replaced += 1 |
| if replaced == 0: |
| raise RuntimeError("inject_lora replaced no layers") |
| return trainable |
|
|
|
|
| def lora_state_dict(model: nn.Module) -> dict[str, torch.Tensor]: |
| return { |
| key: value.detach().cpu().contiguous() |
| for key, value in model.state_dict().items() |
| if ".lora_A" in key or ".lora_B" in key |
| } |
|
|
|
|
| def load_lora_state_dict(model: nn.Module, state: dict[str, torch.Tensor]) -> None: |
| expected = set(lora_state_dict(model)) |
| if set(state) != expected: |
| missing = sorted(expected - set(state)) |
| unexpected = sorted(set(state) - expected) |
| raise RuntimeError(f"LoRA state mismatch: missing={missing[:3]}..., unexpected={unexpected[:3]}...") |
| incompatible = model.load_state_dict(state, strict=False) |
| if incompatible.unexpected_keys: |
| raise RuntimeError(f"Unexpected LoRA keys: {incompatible.unexpected_keys[:3]}") |
|
|
