lora
Training-side LoRA utilities for fastvideo.train model plugins.
fastvideo.train.utils.lora.get_local_torch_device
get_local_torch_device() -> torch.deviceReturn the torch device for the current rank.
fastvideo.train.utils.lora.init_logger
init_logger(name: str) -> _FastvideoLoggerThe 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"),
)</code></pre></details>
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"),
)</code></pre></details>
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</code></pre></details>
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)</code></pre></details>