[feat] GenRL: add Wan LoRA adapter support (#1404)
This commit is contained in:
Adam
@@ -317,8 +317,29 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
|
||||
input_parallel = splitted_input[tp_rank].contiguous()
|
||||
output_parallel = self.base_layer.quant_method.apply(self.base_layer, input_parallel)
|
||||
|
||||
if self.set_lora:
|
||||
output_parallel = self.apply_lora(output_parallel, input_parallel)
|
||||
if not self.merged and not self.disable_lora:
|
||||
lora_A = self.lora_A
|
||||
lora_B = self.lora_B
|
||||
if lora_A is None or lora_B is None:
|
||||
raise RuntimeError(
|
||||
"LoRA weights (lora_A, lora_B) must be initialized "
|
||||
"before forward pass when LoRA is enabled."
|
||||
)
|
||||
if isinstance(lora_A, DTensor):
|
||||
lora_A = lora_A.to_local()
|
||||
if isinstance(lora_B, DTensor):
|
||||
lora_B = lora_B.to_local()
|
||||
|
||||
lora_A_sliced = self.slice_lora_a_weights(
|
||||
lora_A.to(input_parallel, non_blocking=True))
|
||||
lora_B_sliced = self.slice_lora_b_weights(
|
||||
lora_B.to(output_parallel, non_blocking=True))
|
||||
delta = input_parallel @ lora_A_sliced.T @ lora_B_sliced.T
|
||||
if self.lora_alpha != self.lora_rank:
|
||||
delta = delta * (
|
||||
self.lora_alpha / self.lora_rank # type: ignore
|
||||
)
|
||||
output_parallel = output_parallel + delta
|
||||
|
||||
if self.base_layer.reduce_results and self.base_layer.tp_size > 1:
|
||||
output_ = tensor_model_parallel_all_reduce(output_parallel)
|
||||
@@ -334,12 +355,17 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
|
||||
return output, output_bias
|
||||
|
||||
def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
|
||||
tp_rank = get_tp_rank()
|
||||
shard_size = self.base_layer.input_size_per_partition
|
||||
# LoRA A gets its input size from base_layer.weight.shape[1].
|
||||
# If that size is already input_size_per_partition, A is already
|
||||
# sharded; otherwise it is a global tensor and needs TP slicing.
|
||||
if A.shape[1] == shard_size:
|
||||
return A.contiguous()
|
||||
|
||||
tp_rank = get_tp_rank()
|
||||
start_idx = tp_rank * shard_size
|
||||
end_idx = (tp_rank + 1) * shard_size
|
||||
A = A[:, start_idx:end_idx].contiguous()
|
||||
return A
|
||||
return A[:, start_idx:end_idx].contiguous()
|
||||
|
||||
def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor:
|
||||
return B
|
||||
|
||||
@@ -10,8 +10,12 @@ real prompt dataloaders are created by the GenRLMethod.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from types import MethodType
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.distributed import (
|
||||
get_sp_group,
|
||||
get_world_group,
|
||||
@@ -30,6 +34,104 @@ if TYPE_CHECKING:
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _is_lora_target(
|
||||
module_name: str,
|
||||
target_modules: list[str],
|
||||
) -> bool:
|
||||
return any(
|
||||
module_name == target
|
||||
or module_name.endswith(f".{target}")
|
||||
for target in target_modules
|
||||
)
|
||||
|
||||
|
||||
def _apply_fastvideo_lora(
|
||||
transformer: Any,
|
||||
*,
|
||||
lora_rank: int,
|
||||
lora_alpha: int,
|
||||
target_modules: list[str],
|
||||
init_weights: str,
|
||||
) -> int:
|
||||
from fastvideo.layers.lora.linear import (
|
||||
get_lora_layer,
|
||||
replace_submodule,
|
||||
)
|
||||
|
||||
transformer.requires_grad_(False)
|
||||
converted_count = 0
|
||||
for name, layer in list(transformer.named_modules()):
|
||||
if not _is_lora_target(name, target_modules):
|
||||
continue
|
||||
lora_layer = get_lora_layer(
|
||||
layer,
|
||||
lora_rank=lora_rank,
|
||||
lora_alpha=lora_alpha,
|
||||
training_mode=True,
|
||||
)
|
||||
if lora_layer is None:
|
||||
continue
|
||||
_init_lora_weights(lora_layer, init_weights, lora_rank)
|
||||
replace_submodule(transformer, name, lora_layer)
|
||||
converted_count += 1
|
||||
return converted_count
|
||||
|
||||
|
||||
def _init_lora_weights(
|
||||
lora_layer: Any,
|
||||
init_weights: str,
|
||||
lora_rank: int,
|
||||
) -> None:
|
||||
"""Match PEFT's useful LoRA initialization modes."""
|
||||
init = init_weights.lower()
|
||||
if init == "default":
|
||||
return
|
||||
|
||||
lora_A = getattr(lora_layer, "lora_A", None)
|
||||
lora_B = getattr(lora_layer, "lora_B", None)
|
||||
if lora_A is None or lora_B is None:
|
||||
return
|
||||
|
||||
if init == "gaussian":
|
||||
torch.nn.init.normal_(lora_A, std=1 / max(1, lora_rank))
|
||||
torch.nn.init.zeros_(lora_B)
|
||||
return
|
||||
|
||||
raise ValueError(
|
||||
"Unsupported GenRLWanModel LoRA init_weights="
|
||||
f"{init_weights!r}. Use 'gaussian' or 'default'."
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _disable_lora_adapters(transformer: Any):
|
||||
"""Temporarily run a LoRA-wrapped transformer as its frozen base model."""
|
||||
lora_layers = [
|
||||
module for module in transformer.modules()
|
||||
if hasattr(module, "disable_lora")
|
||||
]
|
||||
previous = [bool(module.disable_lora) for module in lora_layers]
|
||||
try:
|
||||
for module in lora_layers:
|
||||
module.disable_lora = True
|
||||
yield
|
||||
finally:
|
||||
for module, was_disabled in zip(lora_layers, previous, strict=True):
|
||||
module.disable_lora = was_disabled
|
||||
|
||||
|
||||
def _attach_disable_adapter(transformer: Any) -> None:
|
||||
"""Expose a PEFT-compatible disable_adapter context manager."""
|
||||
|
||||
def disable_adapter(self):
|
||||
return _disable_lora_adapters(self)
|
||||
|
||||
transformer.disable_adapter = MethodType( # type: ignore[attr-defined]
|
||||
disable_adapter,
|
||||
transformer,
|
||||
)
|
||||
|
||||
|
||||
class _InfiniteDummyLoader:
|
||||
"""Trivial iterable that yields empty dicts forever."""
|
||||
|
||||
@@ -57,6 +159,12 @@ class GenRLWanModel(WanModel):
|
||||
init_from: str,
|
||||
training_config: TrainingConfig,
|
||||
trainable: bool = True,
|
||||
use_lora: bool = False,
|
||||
lora_r: int = 32,
|
||||
lora_alpha: int = 64,
|
||||
lora_target_modules: list[str] | None = None,
|
||||
lora_path: str | None = None,
|
||||
lora_init_weights: str = "gaussian",
|
||||
disable_custom_init_weights: bool = False,
|
||||
flow_shift: float = 3.0,
|
||||
enable_gradient_checkpointing_type: str
|
||||
@@ -79,9 +187,41 @@ class GenRLWanModel(WanModel):
|
||||
transformer_override_safetensor
|
||||
),
|
||||
)
|
||||
if use_lora:
|
||||
if lora_target_modules is None:
|
||||
raise ValueError(
|
||||
"GenRLWanModel use_lora=True requires "
|
||||
"lora_target_modules."
|
||||
)
|
||||
if lora_path:
|
||||
raise ValueError(
|
||||
"GenRLWanModel lora_path is not supported for "
|
||||
"FastVideo LoRA training yet."
|
||||
)
|
||||
converted_count = _apply_fastvideo_lora(
|
||||
self.transformer,
|
||||
lora_rank=int(lora_r),
|
||||
lora_alpha=int(lora_alpha),
|
||||
target_modules=lora_target_modules,
|
||||
init_weights=lora_init_weights,
|
||||
)
|
||||
if converted_count == 0:
|
||||
raise ValueError(
|
||||
"GenRLWanModel use_lora=True did not match any "
|
||||
f"FastVideo linear layers: {lora_target_modules}"
|
||||
)
|
||||
logger.info(
|
||||
"Converted %d GenRL Wan transformer layers to LoRA",
|
||||
converted_count,
|
||||
)
|
||||
_attach_disable_adapter(self.transformer)
|
||||
self.text_encoder: Any = None
|
||||
self.tokenizer: Any = None
|
||||
|
||||
def disable_adapter(self):
|
||||
"""PEFT-compatible context manager for reference KL with LoRA."""
|
||||
return _disable_lora_adapters(self.transformer)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Lifecycle
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user