d5287f8352
Co-authored-by: William Lin <8941107+SolitaryThinker@users.noreply.github.com> Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
414 lines
15 KiB
Python
414 lines
15 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Training-side LoRA utilities for ``fastvideo.train`` model plugins."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
from functools import partial
|
|
import hashlib
|
|
import math
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
from torch.distributed.device_mesh import DeviceMesh
|
|
from torch.distributed.tensor import DTensor, Replicate
|
|
|
|
from fastvideo.distributed import get_local_torch_device
|
|
from fastvideo.layers.lora.linear import (
|
|
BaseLayerWithLoRA,
|
|
get_lora_layer,
|
|
replace_submodule,
|
|
)
|
|
from fastvideo.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
DEFAULT_LORA_TARGET_MODULES = [
|
|
"q_proj",
|
|
"k_proj",
|
|
"v_proj",
|
|
"o_proj",
|
|
"to_q",
|
|
"to_k",
|
|
"to_v",
|
|
"to_out",
|
|
"to_qkv",
|
|
"to_gate_compress",
|
|
]
|
|
|
|
_LORA_CONFIG_KEYS = ("enable", "rank", "alpha", "target_modules")
|
|
|
|
# Private loader attributes intentionally live on the transformed model. The
|
|
# component loader constructs the model on ``meta`` before FSDP sees it, so the
|
|
# structural transform and the checkpoint-loading plan have to travel together.
|
|
LORA_CHECKPOINT_KEY_ALIASES_ATTR = "_fastvideo_checkpoint_key_aliases"
|
|
LORA_MISSING_PARAMETER_INITIALIZER_ATTR = "_fastvideo_missing_parameter_initializer"
|
|
|
|
|
|
@dataclass
|
|
class LoraConfig:
|
|
"""Structured LoRA settings for one ``fastvideo.train`` model role.
|
|
|
|
Parsed from the nested ``models.<role>.lora`` YAML block::
|
|
|
|
lora:
|
|
enable: true # default false
|
|
rank: 16
|
|
alpha: 32 # defaults to rank when omitted
|
|
target_modules: [to_q, to_k, to_v, to_out]
|
|
|
|
``enable`` is an explicit on/off switch so a config states its intent
|
|
plainly: the presence of ``rank`` alone never silently flips a run into
|
|
LoRA-only training. When ``enable`` is false a still-present ``rank`` is
|
|
ignored (with an INFO log), so a configured-but-off block is valid.
|
|
"""
|
|
|
|
enable: bool = False
|
|
rank: int | None = None
|
|
alpha: int | None = None
|
|
target_modules: list[str] | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.rank is not None:
|
|
self.rank = int(self.rank)
|
|
if self.alpha is not None:
|
|
self.alpha = int(self.alpha)
|
|
if self.target_modules is not None:
|
|
self.target_modules = list(self.target_modules)
|
|
|
|
if self.enable:
|
|
if self.rank is None:
|
|
raise ValueError("models.<role>.lora.enable is true but lora.rank is unset "
|
|
"— an explicit positive rank is required to enable LoRA")
|
|
if self.rank <= 0:
|
|
raise ValueError(f"models.<role>.lora.rank must be > 0, got {self.rank!r}")
|
|
elif self.rank is not None:
|
|
logger.info(
|
|
"models.<role>.lora.rank=%s is set but lora.enable is false — "
|
|
"LoRA will NOT be applied (model trains on its normal "
|
|
"trainable path).", self.rank)
|
|
|
|
@classmethod
|
|
def coerce(
|
|
cls,
|
|
obj: LoraConfig | dict[str, Any] | None,
|
|
) -> LoraConfig | None:
|
|
"""Normalize a raw YAML mapping (or existing config) into a LoraConfig.
|
|
|
|
Returns ``None`` when no ``lora`` block was given, which callers treat
|
|
as "LoRA not configured" — identical in effect to ``enable: false``.
|
|
"""
|
|
if obj is None:
|
|
return None
|
|
if isinstance(obj, LoraConfig):
|
|
return obj
|
|
if not isinstance(obj, dict):
|
|
raise TypeError("models.<role>.lora must be a mapping or LoraConfig, got "
|
|
f"{type(obj).__name__}")
|
|
unknown = set(obj) - set(_LORA_CONFIG_KEYS)
|
|
if unknown:
|
|
logger.warning("LoraConfig: ignoring unrecognized lora keys %s "
|
|
"(valid keys: %s)", sorted(unknown), list(_LORA_CONFIG_KEYS))
|
|
return cls(
|
|
enable=bool(obj.get("enable", False)),
|
|
rank=obj.get("rank"),
|
|
alpha=obj.get("alpha"),
|
|
target_modules=obj.get("target_modules"),
|
|
)
|
|
|
|
|
|
def _is_target_layer(
|
|
module_name: str,
|
|
target_modules: Sequence[str],
|
|
) -> bool:
|
|
return any(target_name in module_name for target_name in target_modules)
|
|
|
|
|
|
def _is_excluded_layer(
|
|
module_name: str,
|
|
excluded_modules: Sequence[str],
|
|
) -> bool:
|
|
return any(excluded in module_name for excluded in excluded_modules)
|
|
|
|
|
|
def _replicated_dims(parameter: DTensor) -> list[int]:
|
|
return [
|
|
mesh_dim for mesh_dim, placement in enumerate(parameter.placements)
|
|
if isinstance(placement, Replicate) and parameter.device_mesh.size(mesh_dim) > 1
|
|
]
|
|
|
|
|
|
def sync_replicated_lora_gradients(transformer: torch.nn.Module) -> int:
|
|
"""Average every replicated LoRA gradient over its mesh, on every rank.
|
|
|
|
A per-parameter ``register_post_accumulate_grad_hook`` is not enough: the
|
|
hook only fires on ranks where that parameter actually received a
|
|
gradient, so a rank whose loss skipped a layer issues fewer collectives
|
|
than its peers -- the ranks then pair the wrong tensors together, or hang
|
|
outright when the totals differ. This pass runs after backward and before
|
|
clipping/the optimizer step, walks the LoRA parameters in a stable order,
|
|
and joins the collective on every rank, zero-filling the gradient where the
|
|
rank contributed none (a rank that skipped the layer contributes exactly
|
|
zero to the average).
|
|
|
|
Returns the number of parameters synchronized.
|
|
"""
|
|
|
|
if not dist.is_available() or not dist.is_initialized():
|
|
return 0
|
|
|
|
synced = 0
|
|
for module in transformer.modules():
|
|
for attr_name in ("lora_A", "lora_B"):
|
|
parameter = getattr(module, attr_name, None)
|
|
if not isinstance(parameter, DTensor) or not parameter.requires_grad:
|
|
continue
|
|
replicated_dims = _replicated_dims(parameter)
|
|
if not replicated_dims:
|
|
continue
|
|
grad = parameter.grad
|
|
if grad is None:
|
|
local_grad = torch.zeros_like(parameter.to_local())
|
|
parameter.grad = DTensor.from_local(
|
|
local_grad,
|
|
device_mesh=parameter.device_mesh,
|
|
placements=parameter.placements,
|
|
run_check=False,
|
|
)
|
|
else:
|
|
local_grad = grad.to_local() if isinstance(grad, DTensor) else grad
|
|
for mesh_dim in replicated_dims:
|
|
dist.all_reduce(
|
|
local_grad,
|
|
group=parameter.device_mesh.get_group(mesh_dim),
|
|
)
|
|
local_grad.div_(parameter.device_mesh.size(mesh_dim))
|
|
synced += 1
|
|
return synced
|
|
|
|
|
|
def _make_replicated_lora_parameter(
|
|
parameter: nn.Parameter,
|
|
mesh: DeviceMesh,
|
|
) -> nn.Parameter:
|
|
"""Create a synchronized replicated DTensor for a late-added LoRA weight."""
|
|
|
|
placements = [Replicate()] * mesh.ndim
|
|
replicated = DTensor.from_local(
|
|
parameter.detach(),
|
|
device_mesh=mesh,
|
|
placements=placements,
|
|
run_check=True,
|
|
)
|
|
replicated_parameter = nn.Parameter(
|
|
replicated,
|
|
requires_grad=parameter.requires_grad,
|
|
)
|
|
return replicated_parameter
|
|
|
|
|
|
def _replicate_lora_parameters(transformer: torch.nn.Module, ) -> None:
|
|
"""Wrap LoRA params in replicated DTensors when distributed is active.
|
|
|
|
The training loaders shard the base transformer with FSDP/HSDP before the
|
|
model plugin sees it. Newly-added LoRA parameters therefore need to be
|
|
explicit replicated DTensors so optimizers/checkpointing can treat them the
|
|
same way across ranks. Replicated values are broadcast during creation,
|
|
and their rank-local gradients are averaged before the optimizer step.
|
|
|
|
The mesh is reused from the FSDP-wrapped base_layer parameters rather than
|
|
rebuilt via ``init_device_mesh`` — building a parallel mesh with a different
|
|
topology than the one FSDP already registered can conflict with the
|
|
existing mesh init. ``placements=[Replicate()] * mesh.ndim`` is passed
|
|
explicitly so the local tensor is treated as a replicated copy across all
|
|
mesh dimensions (instead of falling back to a default Shard layout).
|
|
"""
|
|
|
|
if not dist.is_available() or not dist.is_initialized():
|
|
return
|
|
|
|
device = get_local_torch_device()
|
|
if device.type != "cuda":
|
|
return
|
|
|
|
# Look up the mesh that FSDP/HSDP already attached to a base_layer
|
|
# parameter. Non-FSDP runs (e.g. single-GPU / non-distributed) won't have
|
|
# any DTensor params here; in that case we leave LoRA params as plain
|
|
# tensors, which is the correct local-only behavior.
|
|
mesh: DeviceMesh | None = None
|
|
for module in transformer.modules():
|
|
if not isinstance(module, BaseLayerWithLoRA):
|
|
continue
|
|
for p in module.base_layer.parameters():
|
|
if isinstance(p, DTensor):
|
|
mesh = p.device_mesh
|
|
break
|
|
if mesh is not None:
|
|
break
|
|
|
|
if mesh is None:
|
|
return
|
|
|
|
for module in transformer.modules():
|
|
if not isinstance(module, BaseLayerWithLoRA):
|
|
continue
|
|
|
|
module.base_layer.requires_grad_(False)
|
|
|
|
for attr_name in ("lora_A", "lora_B"):
|
|
param = getattr(module, attr_name, None)
|
|
if param is None:
|
|
continue
|
|
param.requires_grad_(True)
|
|
if isinstance(param, DTensor):
|
|
continue
|
|
setattr(
|
|
module,
|
|
attr_name,
|
|
_make_replicated_lora_parameter(param, mesh),
|
|
)
|
|
|
|
|
|
def _initialize_lora_parameter(
|
|
name: str,
|
|
shape: torch.Size,
|
|
dtype: torch.dtype,
|
|
*,
|
|
parameter_names: frozenset[str],
|
|
seed: int,
|
|
) -> torch.Tensor | None:
|
|
"""Build one deterministic full LoRA tensor for the FSDP loader.
|
|
|
|
Every rank constructs the same full tensor and lets the loader distribute
|
|
it according to the parameter's FSDP placement. Deriving a seed from the
|
|
fully-qualified name keeps initialization independent of state-dict/set
|
|
iteration order.
|
|
"""
|
|
if name not in parameter_names:
|
|
return None
|
|
|
|
if name.endswith(".lora_B"):
|
|
return torch.zeros(tuple(shape), dtype=dtype, device="cpu")
|
|
if not name.endswith(".lora_A") or len(shape) != 2:
|
|
raise ValueError(f"Unsupported LoRA parameter initialization request: {name} {tuple(shape)}")
|
|
|
|
name_seed = int.from_bytes(hashlib.sha256(name.encode("utf-8")).digest()[:8], "little")
|
|
generator = torch.Generator(device="cpu").manual_seed((int(seed) + name_seed) % (2**63 - 1))
|
|
# nn.init.kaiming_uniform_(a=sqrt(5)) simplifies to +/- 1/sqrt(fan_in).
|
|
bound = 1.0 / math.sqrt(int(shape[1]))
|
|
value = torch.empty(tuple(shape), dtype=torch.float32, device="cpu")
|
|
value.uniform_(-bound, bound, generator=generator)
|
|
return value.to(dtype=dtype)
|
|
|
|
|
|
def finalize_lora_training(transformer: torch.nn.Module) -> int:
|
|
"""Freeze base weights and enable only adapter parameters after loading."""
|
|
transformer.requires_grad_(False)
|
|
count = 0
|
|
for module in transformer.modules():
|
|
if not isinstance(module, BaseLayerWithLoRA):
|
|
continue
|
|
if module.lora_A is None or module.lora_B is None:
|
|
raise RuntimeError("Training LoRA wrapper is missing adapter parameters")
|
|
module.base_layer.requires_grad_(False)
|
|
module.lora_A.requires_grad_(True)
|
|
module.lora_B.requires_grad_(True)
|
|
count += 1
|
|
if count == 0:
|
|
raise ValueError("No training LoRA wrappers were found after checkpoint loading")
|
|
transformer.train()
|
|
return count
|
|
|
|
|
|
def enable_lora_training(
|
|
transformer: torch.nn.Module,
|
|
*,
|
|
lora_rank: int,
|
|
lora_alpha: int | None = None,
|
|
lora_target_modules: Sequence[str] | None = None,
|
|
prepare_for_fsdp: bool = False,
|
|
initialization_seed: int = 0,
|
|
) -> int:
|
|
"""Replace supported linear layers with trainable LoRA wrappers.
|
|
|
|
Returns the number of layers converted to LoRA.
|
|
"""
|
|
|
|
rank = int(lora_rank)
|
|
if rank <= 0:
|
|
raise ValueError(f"lora_rank must be > 0, got {lora_rank!r}")
|
|
|
|
alpha = int(lora_alpha) if lora_alpha is not None else rank
|
|
target_modules = list(lora_target_modules or DEFAULT_LORA_TARGET_MODULES)
|
|
arch_config = getattr(
|
|
getattr(transformer, "config", None),
|
|
"arch_config",
|
|
None,
|
|
)
|
|
excluded_modules = list(getattr(arch_config, "exclude_lora_layers", []), )
|
|
|
|
transformer.requires_grad_(False)
|
|
|
|
replacements: list[tuple[str, BaseLayerWithLoRA]] = []
|
|
checkpoint_key_aliases: dict[str, str] = {}
|
|
lora_parameter_names: set[str] = set()
|
|
for module_name, module in transformer.named_modules():
|
|
if not module_name:
|
|
continue
|
|
if not _is_target_layer(module_name, target_modules):
|
|
continue
|
|
if _is_excluded_layer(module_name, excluded_modules):
|
|
continue
|
|
|
|
lora_layer = get_lora_layer(
|
|
module,
|
|
lora_rank=rank,
|
|
lora_alpha=alpha,
|
|
training_mode=True,
|
|
)
|
|
if lora_layer is None:
|
|
continue
|
|
replacements.append((module_name, lora_layer))
|
|
for state_name in module.state_dict():
|
|
checkpoint_key_aliases[f"{module_name}.{state_name}"] = f"{module_name}.base_layer.{state_name}"
|
|
lora_parameter_names.update({f"{module_name}.lora_A", f"{module_name}.lora_B"})
|
|
|
|
if not replacements:
|
|
raise ValueError("No LoRA-compatible layers were found for the requested "
|
|
f"target modules: {target_modules}")
|
|
|
|
for module_name, lora_layer in replacements:
|
|
replace_submodule(transformer, module_name, lora_layer)
|
|
|
|
if prepare_for_fsdp:
|
|
# The FSDP loader consumes these aliases after the normal HF-to-native
|
|
# mapping, then initializes adapter tensors that are intentionally absent
|
|
# from the base checkpoint. FSDP therefore owns base and adapter params
|
|
# in the same parameter groups from the outset.
|
|
setattr(transformer, LORA_CHECKPOINT_KEY_ALIASES_ATTR, checkpoint_key_aliases)
|
|
setattr(
|
|
transformer,
|
|
LORA_MISSING_PARAMETER_INITIALIZER_ATTR,
|
|
partial(
|
|
_initialize_lora_parameter,
|
|
parameter_names=frozenset(lora_parameter_names),
|
|
seed=int(initialization_seed),
|
|
),
|
|
)
|
|
else:
|
|
# Retain the legacy post-load path for existing model plugins. New
|
|
# distributed integrations must use ``prepare_for_fsdp=True``.
|
|
_replicate_lora_parameters(transformer)
|
|
finalize_lora_training(transformer)
|
|
|
|
logger.info(
|
|
"Enabled LoRA training with rank=%d alpha=%d on %d layers",
|
|
rank,
|
|
alpha,
|
|
len(replacements),
|
|
)
|
|
return len(replacements)
|