--- title: lora ---

Training-side LoRA utilities for fastvideo.train model plugins.

fastvideo.train.utils.lora.get_local_torch_device

get_local_torch_device() -> torch.device

Return the torch device for the current rank.

fastvideo.train.utils.lora.init_logger

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

fastvideo.train.utils.lora.LoraConfig

LoraConfig(enable: bool = False, rank: int | None = None, alpha: int | None = None, target_modules: list[str] | None = None)

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.

Source
@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"),
        )

fastvideo.train.utils.lora.LoraConfig.coerce

coerce(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.

Source
@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"),
    )

fastvideo.train.utils.lora.sync_replicated_lora_gradients

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.

Source
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

fastvideo.train.utils.lora.finalize_lora_training

finalize_lora_training(transformer: torch.nn.Module) -> int

Freeze base weights and enable only adapter parameters after loading.

Source
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

fastvideo.train.utils.lora.enable_lora_training

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.

Source
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)