hao-ai-lab-FastVideo/fastvideo/train/utils/lora.py

# 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 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")


@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 enable_lora_training(
    transformer: torch.nn.Module,
    *,
    lora_rank: int,
    lora_alpha: int | None = None,
    lora_target_modules: Sequence[str] | None = None,
) -> 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]] = []
    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))

    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)

    _replicate_lora_parameters(transformer)
    transformer.train()

    logger.info(
        "Enabled LoRA training with rank=%d alpha=%d on %d layers",
        rank,
        alpha,
        len(replacements),
    )
    return len(replacements)