diff --git a/fastvideo/layers/lora/linear.py b/fastvideo/layers/lora/linear.py index c6dcd965c..014cc9c5b 100644 --- a/fastvideo/layers/lora/linear.py +++ b/fastvideo/layers/lora/linear.py @@ -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 diff --git a/fastvideo/train/models/wan/wan_genrl.py b/fastvideo/train/models/wan/wan_genrl.py index f01d97695..9531b41d9 100644 --- a/fastvideo/train/models/wan/wan_genrl.py +++ b/fastvideo/train/models/wan/wan_genrl.py @@ -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 # ------------------------------------------------------------------