| |
| """Low-rank adapters for DiT linear layers (BaseLinear, NativeMoELinear, nn.Linear). |
| |
| Scope: the middle DiT only (adapter embedders + TransformerBlock + final heads). |
| Video VAE / audio VAE / text encoder stay outside training (precomputed latents). |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import asdict, dataclass |
| from typing import Any, Dict, List, Optional, Sequence, Tuple |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from .dit_module import BaseLinear, NativeMoELinear |
|
|
|
|
| |
| ATTN_MLP_SUFFIXES: Tuple[str, ...] = ( |
| "linear_qkv", |
| "linear_proj", |
| "up_gate_proj", |
| "down_proj", |
| ) |
| PLAIN_LINEAR_NAMES: Tuple[str, ...] = ( |
| "video_embedder", |
| "text_embedder", |
| "audio_embedder", |
| "final_linear_video", |
| "final_linear_audio", |
| ) |
|
|
|
|
| @dataclass |
| class LoraConfig: |
| r: int = 16 |
| alpha: float = 16.0 |
| dropout: float = 0.0 |
| targets: str = "all" |
| """Which DiT linears to wrap. See TRAIN.md in v8/LoRA.""" |
|
|
| def scaling(self) -> float: |
| return float(self.alpha) / float(max(self.r, 1)) |
|
|
| def to_dict(self) -> Dict[str, Any]: |
| return asdict(self) |
|
|
|
|
| def _match_target(name: str, targets: str) -> bool: |
| leaf = name.rsplit(".", 1)[-1] |
| t = (targets or "all").strip().lower() |
| if t == "all": |
| return leaf in ATTN_MLP_SUFFIXES or leaf in PLAIN_LINEAR_NAMES |
| if t == "attn": |
| return leaf in ("linear_qkv", "linear_proj") |
| if t == "attn_mlp": |
| return leaf in ATTN_MLP_SUFFIXES |
| if t in ("embed_final", "heads"): |
| return leaf in PLAIN_LINEAR_NAMES |
| raise ValueError(f"Unknown lora targets={targets!r}; use all|attn|attn_mlp|embed_final") |
|
|
|
|
| def _init_lora_ab( |
| in_features: int, out_features: int, r: int, device: torch.device, dtype: torch.dtype |
| ) -> Tuple[nn.Parameter, nn.Parameter]: |
| |
| a = nn.Parameter(torch.empty(r, in_features, device=device, dtype=dtype)) |
| b = nn.Parameter(torch.zeros(out_features, r, device=device, dtype=dtype)) |
| nn.init.kaiming_uniform_(a, a=5**0.5) |
| return a, b |
|
|
|
|
| class LoRAPlainLinear(nn.Module): |
| """LoRA around a standard ``nn.Linear`` (adapter embedders / final heads).""" |
|
|
| def __init__(self, base: nn.Linear, *, r: int, alpha: float, dropout: float = 0.0): |
| super().__init__() |
| if not isinstance(base, nn.Linear): |
| raise TypeError(f"LoRAPlainLinear expects nn.Linear, got {type(base)}") |
| self.base = base |
| self.r = int(r) |
| self.alpha = float(alpha) |
| self.scaling = float(alpha) / float(max(self.r, 1)) |
| self.dropout_p = float(dropout) |
| for p in self.base.parameters(): |
| p.requires_grad = False |
| dtype = self.base.weight.dtype |
| device = self.base.weight.device |
| self.lora_A, self.lora_B = _init_lora_ab( |
| self.base.in_features, self.base.out_features, self.r, device, dtype |
| ) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| y = self.base(x) |
| x_d = F.dropout(x, p=self.dropout_p, training=self.training) if self.dropout_p > 0 else x |
| |
| a = self.lora_A.to(dtype=x_d.dtype) |
| b = self.lora_B.to(dtype=x_d.dtype) |
| delta = (x_d @ a.t() @ b.t()) * self.scaling |
| return y + delta.to(dtype=y.dtype) |
|
|
|
|
| class LoRABaseLinear(nn.Module): |
| """LoRA around ``BaseLinear`` / ``NativeMoELinear`` (DiT attn + MLP). |
| |
| For ``num_experts==1``: one (A, B) pair. |
| For ``NativeMoELinear`` (mm_layers): one (A, B) per modality expert, applied |
| after ``modality_dispatcher.dispatch``, matching the base forward split. |
| """ |
|
|
| def __init__(self, base: BaseLinear, *, r: int, alpha: float, dropout: float = 0.0): |
| super().__init__() |
| if not isinstance(base, BaseLinear): |
| raise TypeError(f"LoRABaseLinear expects BaseLinear, got {type(base)}") |
| self.base = base |
| self.r = int(r) |
| self.alpha = float(alpha) |
| self.scaling = float(alpha) / float(max(self.r, 1)) |
| self.dropout_p = float(dropout) |
| self.num_experts = int(base.num_experts) |
| for p in self.base.parameters(): |
| p.requires_grad = False |
|
|
| dtype = self.base.weight.dtype |
| device = self.base.weight.device |
| in_f = int(base.in_features) |
| out_f = int(base.out_features) |
| if self.num_experts <= 1: |
| self.lora_A, self.lora_B = _init_lora_ab(in_f, out_f, self.r, device, dtype) |
| self.lora_A_list = None |
| self.lora_B_list = None |
| else: |
| self.lora_A = None |
| self.lora_B = None |
| self.lora_A_list = nn.ParameterList() |
| self.lora_B_list = nn.ParameterList() |
| for _ in range(self.num_experts): |
| a, b = _init_lora_ab(in_f, out_f, self.r, device, dtype) |
| self.lora_A_list.append(a) |
| self.lora_B_list.append(b) |
|
|
| def forward( |
| self, |
| input: torch.Tensor, |
| output_dtype: Optional[torch.dtype] = None, |
| modality_dispatcher: Optional[Any] = None, |
| ) -> torch.Tensor: |
| y = self.base(input, output_dtype=output_dtype, modality_dispatcher=modality_dispatcher) |
| if self.num_experts <= 1: |
| x_d = F.dropout(input, p=self.dropout_p, training=self.training) if self.dropout_p > 0 else input |
| a = self.lora_A.to(dtype=x_d.dtype) |
| b = self.lora_B.to(dtype=x_d.dtype) |
| delta = (x_d @ a.t() @ b.t()) * self.scaling |
| return y + delta.to(dtype=y.dtype) |
|
|
| if modality_dispatcher is None: |
| raise RuntimeError("NativeMoELinear LoRA requires modality_dispatcher") |
| input_list = modality_dispatcher.dispatch(input) |
| deltas: List[torch.Tensor] = [] |
| for i in range(self.num_experts): |
| x_i = input_list[i] |
| if x_i.numel() == 0: |
| deltas.append(x_i.new_zeros(x_i.shape[0], self.base.out_features)) |
| continue |
| x_d = F.dropout(x_i, p=self.dropout_p, training=self.training) if self.dropout_p > 0 else x_i |
| a = self.lora_A_list[i].to(dtype=x_d.dtype) |
| b = self.lora_B_list[i].to(dtype=x_d.dtype) |
| deltas.append(((x_d @ a.t() @ b.t()) * self.scaling).to(dtype=y.dtype)) |
| delta = modality_dispatcher.undispatch(*deltas) |
| return y + delta |
|
|
|
|
| def _wrap_module(module: nn.Module, cfg: LoraConfig) -> Optional[nn.Module]: |
| if isinstance(module, (LoRAPlainLinear, LoRABaseLinear)): |
| return None |
| if isinstance(module, BaseLinear): |
| return LoRABaseLinear(module, r=cfg.r, alpha=cfg.alpha, dropout=cfg.dropout) |
| if isinstance(module, nn.Linear): |
| return LoRAPlainLinear(module, r=cfg.r, alpha=cfg.alpha, dropout=cfg.dropout) |
| return None |
|
|
|
|
| def inject_lora(model: nn.Module, cfg: LoraConfig) -> List[str]: |
| """Replace matching linear modules in-place. Returns list of wrapped paths.""" |
| if int(cfg.r) <= 0: |
| raise ValueError(f"lora r must be > 0, got {cfg.r}") |
| to_wrap: List[Tuple[str, nn.Module]] = [] |
| for name, module in model.named_modules(): |
| if not name: |
| continue |
| if not _match_target(name, cfg.targets): |
| continue |
| if isinstance(module, (LoRAPlainLinear, LoRABaseLinear)): |
| continue |
| if isinstance(module, (BaseLinear, nn.Linear)): |
| to_wrap.append((name, module)) |
|
|
| wrapped: List[str] = [] |
| for name, module in to_wrap: |
| new_mod = _wrap_module(module, cfg) |
| if new_mod is None: |
| continue |
| parent = model |
| parts = name.split(".") |
| for p in parts[:-1]: |
| parent = getattr(parent, p) |
| setattr(parent, parts[-1], new_mod) |
| wrapped.append(name) |
| if not wrapped: |
| raise RuntimeError(f"inject_lora: no modules matched targets={cfg.targets!r}") |
| return wrapped |
|
|
|
|
| def freeze_non_lora_parameters(model: nn.Module) -> Tuple[int, int]: |
| """Freeze everything except LoRA A/B. Returns (trainable_numel, frozen_numel).""" |
| n_train, n_freeze = 0, 0 |
| for name, p in model.named_parameters(): |
| is_lora = ("lora_A" in name) or ("lora_B" in name) or ("lora_A_list" in name) or ("lora_B_list" in name) |
| if is_lora: |
| p.requires_grad = True |
| n_train += p.numel() |
| else: |
| p.requires_grad = False |
| n_freeze += p.numel() |
| return n_train, n_freeze |
|
|
|
|
| def lora_state_dict(model: nn.Module) -> Dict[str, torch.Tensor]: |
| out: Dict[str, torch.Tensor] = {} |
| for name, p in model.named_parameters(): |
| if ("lora_A" in name) or ("lora_B" in name): |
| out[name] = p.detach().cpu() |
| return out |
|
|
|
|
| def load_lora_state_dict(model: nn.Module, state: Dict[str, torch.Tensor], strict: bool = True) -> None: |
| missing, unexpected = [], [] |
| own = {n: p for n, p in model.named_parameters() if ("lora_A" in n) or ("lora_B" in n)} |
| for k, v in state.items(): |
| if k not in own: |
| unexpected.append(k) |
| continue |
| if tuple(own[k].shape) != tuple(v.shape): |
| raise RuntimeError(f"LoRA shape mismatch for {k}: ckpt {tuple(v.shape)} vs model {tuple(own[k].shape)}") |
| own[k].data.copy_(v.to(device=own[k].device, dtype=own[k].dtype)) |
| for k in own: |
| if k not in state: |
| missing.append(k) |
| if strict and (missing or unexpected): |
| raise RuntimeError(f"load_lora_state_dict strict fail missing={missing[:8]} unexpected={unexpected[:8]}") |
|
|
|
|
| def mark_lora_text_param_ids(model: nn.Module) -> set: |
| """Parameter ids under adapter.text_embedder LoRA (for text_lr_mult group).""" |
| ids: set = set() |
| for name, p in model.named_parameters(): |
| if "adapter.text_embedder" in name and (("lora_A" in name) or ("lora_B" in name)): |
| ids.add(id(p)) |
| return ids |
|
|
|
|
| def summarize_lora(model: nn.Module, wrapped: Sequence[str]) -> str: |
| n_train, n_freeze = 0, 0 |
| for name, p in model.named_parameters(): |
| if p.requires_grad: |
| n_train += p.numel() |
| else: |
| n_freeze += p.numel() |
| return ( |
| f"lora_wrapped={len(wrapped)} trainable={n_train / 1e6:.3f}M " |
| f"frozen={n_freeze / 1e6:.1f}M sample={list(wrapped)[:6]}" |
| ) |
|
|
