[feat] GenRL: add Wan LoRA adapter support

fastvideo/layers/lora/linear.py

@@ -336,8 +336,28 @@ def forward(self, input_: torch.Tensor):
         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_B, DTensor):
+                lora_B = lora_B.to_local()
+                lora_A = lora_A.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)

fastvideo/train/models/wan/wan_genrl.py

@@ -10,8 +10,12 @@
 
 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,105 @@
 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}")
+        or target in module_name
+        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 +160,12 @@ def __init__(
         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 +188,41 @@ def __init__(
                 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
     # ------------------------------------------------------------------