nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[bugfix] [lora] [distillation] Fix lora distillation bug (#933)

Browse Source
This commit is contained in:
William Lin
2025-12-15 04:12:02 -08:00
committed by GitHub
parent e6066223e6
commit 424fc2b4ae
12 changed files with 260 additions and 230 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/pipelines/composed_pipeline_base.py
+25 -7
View File
@@ -43,6 +43,9 @@ class ComposedPipelineBase(ABC):
training_args: TrainingArgs | None = None
fastvideo_args: FastVideoArgs | TrainingArgs | None = None
modules: dict[str, Any] = {}
# do not need to include moe related transformers
trainable_transformer_names: list[str] = ["transformer"]
trainable_transformer_modules: dict[str, torch.nn.Module] = {}
post_init_called: bool = False
# TODO(will): args should support both inference args and training args
@@ -87,13 +90,14 @@ class ComposedPipelineBase(ABC):
def set_trainable(self) -> None:
# Only train DiT
if getattr(self.fastvideo_args, "training_mode", False):
for name, module in self.modules.items():
for name, module in self.trainable_transformer_modules.items():
logger.info("Setting %s to requires_grad=True", name)
if not isinstance(module, torch.nn.Module):
logger.info(
"Skipping %s because it is not a torch.nn.Module", name)
continue
if "transformer" in name:
module.requires_grad_(True)
else:
module.requires_grad_(False)
module.requires_grad_(True)
module.train()
def post_init(self) -> None:
assert self.fastvideo_args is not None, "fastvideo_args must be set"
@@ -122,18 +126,32 @@ class ComposedPipelineBase(ABC):
fsdp_module_cls = FSDPModule
except Exception: # pragma: no cover - FSDP not always available
fsdp_module_cls = None
compile_kwargs = self.fastvideo_args.torch_compile_kwargs or {}
if fsdp_module_cls is not None and isinstance(
transformer_module, fsdp_module_cls):
logger.info(
"Transformer is already FSDP-wrapped; skipping torch.compile in pipeline"
)
else:
compile_kwargs = self.fastvideo_args.torch_compile_kwargs or {}
logger.info("Enabling torch.compile for DiT with kwargs=%s",
compile_kwargs)
self.modules["transformer"] = torch.compile(
transformer_module, **compile_kwargs)
logger.info("Torch Compile enabled for DiT")
if "transformer_2" in self.modules:
transformer_module_2 = self.modules["transformer_2"]
if fsdp_module_cls is not None and isinstance(
transformer_module_2, fsdp_module_cls):
logger.info(
"Transformer_2 is already FSDP-wrapped; skipping torch.compile in pipeline"
)
else:
logger.info(
"Enabling torch.compile for Transformer_2 with kwargs=%s",
compile_kwargs)
self.modules["transformer_2"] = torch.compile(
transformer_module_2, **compile_kwargs)
logger.info("Torch Compile enabled for DiT")
if not self.fastvideo_args.training_mode:
logger.info("Creating pipeline stages...")
fastvideo/pipelines/lora_pipeline.py
+101 -72
View File
@@ -32,10 +32,9 @@ class LoRAPipeline(ComposedPipelineBase):
) # state dicts of loaded lora adapters (includes lora_A, lora_B, and lora_alpha)
cur_adapter_name: str = ""
cur_adapter_path: str = ""
lora_layers: dict[str, BaseLayerWithLoRA] = {}
lora_layers_critic: dict[str, BaseLayerWithLoRA] = {}
lora_layers: dict[str, dict[str, BaseLayerWithLoRA]] = {}
fastvideo_args: FastVideoArgs | TrainingArgs
exclude_lora_layers: list[str] = []
exclude_lora_layers: dict[str, list[str]] = {}
device: torch.device = get_local_torch_device()
lora_target_modules: list[str] | None = None
lora_path: str | None = None
@@ -47,8 +46,32 @@ class LoRAPipeline(ComposedPipelineBase):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.device = get_local_torch_device()
self.exclude_lora_layers = self.modules[
"transformer"].config.arch_config.exclude_lora_layers
# build list of trainable transformers
for transformer_name in self.trainable_transformer_names:
if transformer_name in self.modules and self.modules[
transformer_name] is not None:
self.trainable_transformer_modules[
transformer_name] = self.modules[transformer_name]
# check for transformer_2 in case of Wan2.2 MoE or fake_score_transformer_2
if transformer_name.endswith("_2"):
raise ValueError(
f"trainable_transformer_name override in pipelines should not include _2 suffix: {transformer_name}"
)
secondary_transformer_name = transformer_name + "_2"
if secondary_transformer_name in self.modules and self.modules[
secondary_transformer_name] is not None:
self.trainable_transformer_modules[
secondary_transformer_name] = self.modules[
secondary_transformer_name]
logger.info("trainable_transformer_modules: %s",
self.trainable_transformer_modules.keys())
for transformer_name, transformer_module in self.trainable_transformer_modules.items(
):
self.exclude_lora_layers[
transformer_name] = transformer_module.config.arch_config.exclude_lora_layers
self.lora_target_modules = self.fastvideo_args.lora_target_modules
self.lora_path = self.fastvideo_args.lora_path
self.lora_nickname = self.fastvideo_args.lora_nickname
@@ -65,8 +88,16 @@ class LoRAPipeline(ComposedPipelineBase):
if self.lora_target_modules is None:
self.lora_target_modules = [
"q_proj", "k_proj", "v_proj", "o_proj", "to_q", "to_k",
"to_v", "to_out", "to_qkv"
"to_v", "to_out", "to_qkv", "to_gate_compress"
]
logger.info(
"Using default lora_target_modules for all transformers: %s",
self.lora_target_modules)
else:
logger.warning(
"Using custom lora_target_modules for all transformers, which may not be intended: %s",
self.lora_target_modules)
self.convert_to_lora_layers()
# Inference
elif not self.training_mode and self.lora_path is not None:
@@ -100,13 +131,18 @@ class LoRAPipeline(ComposedPipelineBase):
super().set_trainable()
return
self.modules["transformer"].requires_grad_(False)
if "fake_score_transformer" in self.modules:
self.modules["fake_score_transformer"].requires_grad_(False)
device_mesh = init_device_mesh("cuda", (dist.get_world_size(), 1),
mesh_dim_names=["fake", "replicate"])
set_lora_grads(self.lora_layers, device_mesh)
set_lora_grads(self.lora_layers_critic, device_mesh)
for transformer_name, transformer_module in self.trainable_transformer_modules.items(
):
transformer_module.train()
transformer_module.requires_grad_(False)
if transformer_name in self.lora_layers:
set_lora_grads(self.lora_layers[transformer_name], device_mesh)
else:
raise ValueError(
f"Transformer {transformer_name} should be trainable but not found in lora_layers"
)
def convert_to_lora_layers(self) -> None:
"""
@@ -115,46 +151,33 @@ class LoRAPipeline(ComposedPipelineBase):
if self.lora_initialized:
return
self.lora_initialized = True
converted_count = 0
for name, layer in self.modules["transformer"].named_modules():
if not self.is_target_layer(name):
continue
excluded = False
for exclude_layer in self.exclude_lora_layers:
if exclude_layer in name:
excluded = True
break
if excluded:
continue
layer = get_lora_layer(layer,
lora_rank=self.lora_rank,
lora_alpha=self.lora_alpha,
training_mode=self.training_mode)
if layer is not None:
self.lora_layers[name] = layer
replace_submodule(self.modules["transformer"], name, layer)
converted_count += 1
logger.info("Converted %d layers to LoRA layers", converted_count)
if "fake_score_transformer" in self.modules:
for name, layer in self.modules[
"fake_score_transformer"].named_modules():
for transformer_name, transformer_module in self.trainable_transformer_modules.items(
):
converted_count = 0
if transformer_name not in self.lora_layers:
self.lora_layers[transformer_name] = {}
logger.info("Converting %s to LoRA Transformer", transformer_name)
for name, layer in transformer_module.named_modules():
if not self.is_target_layer(name):
continue
excluded = False
for exclude_layer in self.exclude_lora_layers[transformer_name]:
if exclude_layer in name:
excluded = True
break
if excluded:
continue
layer = get_lora_layer(layer,
lora_rank=self.lora_rank,
lora_alpha=self.lora_alpha,
training_mode=self.training_mode)
if layer is not None:
self.lora_layers_critic[name] = layer
replace_submodule(self.modules["fake_score_transformer"],
name, layer)
self.lora_layers[transformer_name][name] = layer
replace_submodule(transformer_module, name, layer)
converted_count += 1
logger.info(
"Converted %d layers to LoRA layers in the critic model",
converted_count)
logger.info("Converted %d layers to LoRA layers", converted_count)
def set_lora_adapter(self,
lora_nickname: str,
@@ -238,39 +261,45 @@ class LoRAPipeline(ComposedPipelineBase):
# Merge the new adapter
adapted_count = 0
for name, layer in self.lora_layers.items():
lora_A_name = name + ".lora_A"
lora_B_name = name + ".lora_B"
lora_alpha_name = name + ".lora_alpha"
if lora_A_name in self.lora_adapters[lora_nickname]\
and lora_B_name in self.lora_adapters[lora_nickname]:
# Get alpha value for this layer (defaults to None if not present)
lora_A = self.lora_adapters[lora_nickname][lora_A_name]
lora_B = self.lora_adapters[lora_nickname][lora_B_name]
# Simple lookup - alpha stored with same naming scheme as lora_A/lora_B
alpha = self.lora_adapters[lora_nickname].get(
lora_alpha_name) if adapter_updated else None
for transformer_name, transformer_lora_layers in self.lora_layers.items(
):
for name, layer in transformer_lora_layers.items():
lora_A_name = name + ".lora_A"
lora_B_name = name + ".lora_B"
lora_alpha_name = name + ".lora_alpha"
if lora_A_name in self.lora_adapters[lora_nickname]\
and lora_B_name in self.lora_adapters[lora_nickname]:
# Get alpha value for this layer (defaults to None if not present)
lora_A = self.lora_adapters[lora_nickname][lora_A_name]
lora_B = self.lora_adapters[lora_nickname][lora_B_name]
# Simple lookup - alpha stored with same naming scheme as lora_A/lora_B
alpha = self.lora_adapters[lora_nickname].get(
lora_alpha_name) if adapter_updated else None
layer.set_lora_weights(
lora_A,
lora_B,
lora_alpha=alpha,
training_mode=self.fastvideo_args.training_mode,
lora_path=lora_path)
adapted_count += 1
else:
if rank == 0:
logger.warning(
"LoRA adapter %s does not contain the weights for layer %s. LoRA will not be applied to it.",
lora_path, name)
layer.disable_lora = True
layer.set_lora_weights(
lora_A,
lora_B,
lora_alpha=alpha,
training_mode=self.fastvideo_args.training_mode,
lora_path=lora_path)
adapted_count += 1
else:
if rank == 0:
logger.warning(
"LoRA adapter %s does not contain the weights for layer %s. LoRA will not be applied to it.",
lora_path, name)
layer.disable_lora = True
logger.info("Rank %d: LoRA adapter %s applied to %d layers", rank,
lora_path, adapted_count)
def merge_lora_weights(self) -> None:
for name, layer in self.lora_layers.items():
layer.merge_lora_weights()
for transformer_name, transformer_lora_layers in self.lora_layers.items(
):
for name, layer in transformer_lora_layers.items():
layer.merge_lora_weights()
def unmerge_lora_weights(self) -> None:
for name, layer in self.lora_layers.items():
layer.unmerge_lora_weights()
for transformer_name, transformer_lora_layers in self.lora_layers.items(
):
for name, layer in transformer_lora_layers.items():
layer.unmerge_lora_weights()
fastvideo/tests/training/VSA/h200_reference_wandb_summary_VSA.json
+18
View File
@@ -0,0 +1,18 @@
{
"step_time": 0.8274769484996796,
"avg_step_time": 1.4095150619745254,
"timing/optimizer_step": 0.02941332757472992,
"grad_norm": 0.24477368593215942,
"timing/forward_backward": 0.693624597042799,
"_timestamp": 1765590605.595408,
"timing/prepare_dit_inputs": 0.009604107588529587,
"learning_rate": 1e-05,
"train_loss": 0.30774877965450287,
"_step": 5,
"timing/normalize_input": 0.0003825277090072632,
"vsa_sparsity": 0.05,
"timing/clip_grad_norm": 0.025074530392885208,
"_runtime": 90.44591733,
"timing/get_next_batch": 0.0055067278444767,
"timing/reduce_loss": 0.00021916255354881287
}
fastvideo/tests/training/VSA/reference_wandb_summary_VSA.json → fastvideo/tests/training/VSA/l40s_reference_wandb_summary_VSA.json
View File
fastvideo/tests/training/VSA/test_training_loss_VSA.py
+11 -1
View File
@@ -4,6 +4,7 @@ import subprocess
from pathlib import Path
import json
from huggingface_hub import snapshot_download
import torch
# Import the training pipeline
sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
@@ -12,7 +13,8 @@ from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
from fastvideo.utils import FlexibleArgumentParser
wandb_name = "test_training_loss_VSA"
reference_wandb_summary_file = "fastvideo/tests/training/VSA/reference_wandb_summary_VSA.json"
l40s_reference_wandb_summary_file = "fastvideo/tests/training/VSA/reference_wandb_summary_VSA.json"
h200_reference_wandb_summary_file = "fastvideo/tests/training/VSA/h200_reference_wandb_summary_VSA.json"
NUM_NODES = "1"
NUM_GPUS_PER_NODE = "2"
@@ -104,6 +106,14 @@ def test_distributed_training():
summary_file = 'data/wan_finetune_test_VSA/tracker/wandb/latest-run/files/wandb-summary.json'
device_name = torch.cuda.get_device_name()
if "L40S" in device_name:
reference_wandb_summary_file = l40s_reference_wandb_summary_file
elif "H200" in device_name:
reference_wandb_summary_file = h200_reference_wandb_summary_file
else:
raise ValueError(f"Unknown device: {device_name}")
reference_wandb_summary = json.load(open(reference_wandb_summary_file))
wandb_summary = json.load(open(summary_file))
fastvideo/tests/training/Vanilla/h200_reference_wandb_summary.json
+23
View File
@@ -0,0 +1,23 @@
{
"step_time": 1.171190332621336,
"vsa_sparsity": 0,
"_step": 5,
"avg_step_time": 1.2594494499266147,
"timing/forward_backward": 0.6414474211633205,
"timing/reduce_loss": 0.00025020167231559753,
"_runtime": 58.873311423,
"batch_size": 2,
"timing/prepare_dit_inputs": 0.011771254241466522,
"timing/normalize_input": 0.0005782917141914368,
"learning_rate": 1e-06,
"hidden_dim": 1536,
"train_loss": 0.13831664621829987,
"_timestamp": 1765589891.5887935,
"context_len": 512,
"grad_norm": 0.0643853023648262,
"ffn_dim": 8960,
"timing/optimizer_step": 0.02862849459052086,
"dit_seq_len": 6240,
"timing/clip_grad_norm": 0.02274264767765999,
"timing/get_next_batch": 0.40074965357780457
}
fastvideo/tests/training/Vanilla/l40s_reference_wandb_summary.json
+11 -1
View File
@@ -1 +1,11 @@
{"train_loss":0.1021774671971798,"_step":5,"_wandb":{"runtime":35},"step_time":2.2575189135968685,"grad_norm":0.11582941561937332,"_runtime":35.701958502,"avg_step_time":2.4754387199878694,"_timestamp":1.7525528420185745e+09,"learning_rate":1e-06,"vsa_sparsity":0,"validation_videos_8_steps":{"videos":[{"caption":"A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.","_type":"video-file","sha256":"ee83e83df073a648f89dcd288cccaed9af765e11fc35c77a6e0cb2ebaa1be5b0","size":475248,"path":"media/videos/validation_videos_8_steps_0_ee83e83df073a648f89d.mp4"},{"size":341490,"path":"media/videos/validation_videos_8_steps_0_d0b2758549d5c82845ca.mp4","caption":"A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.","_type":"video-file","sha256":"d0b2758549d5c82845ca3c1ac0db6812877631495e77ac041acdc8ea31f0a5ee"},{"caption":"A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.","_type":"video-file","sha256":"08638381f1607d6ab10684772be38a58b18628fd402a208b25f6af96e765d454","size":436814,"path":"media/videos/validation_videos_8_steps_0_08638381f1607d6ab106.mp4"}],"captions":["A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.","A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.","A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table."],"_type":"videos","count":3}}
{
"train_loss": 0.1021774671971798,
"_step": 5,
"step_time": 2.2575189135968685,
"grad_norm": 0.11582941561937332,
"_runtime": 35.701958502,
"avg_step_time": 2.4754387199878694,
"_timestamp": 1752552842.0185745,
"learning_rate": 1e-06,
"vsa_sparsity": 0
}
fastvideo/tests/training/Vanilla/test_training_loss.py
+3
View File
@@ -18,6 +18,7 @@ from fastvideo.training.wan_training_pipeline import WanTrainingPipeline
wandb_name = "test_training_loss"
a40_reference_wandb_summary_file = "fastvideo/tests/training/Vanilla/a40_reference_wandb_summary.json"
l40s_reference_wandb_summary_file = "fastvideo/tests/training/Vanilla/l40s_reference_wandb_summary.json"
h200_reference_wandb_summary_file = "fastvideo/tests/training/Vanilla/h200_reference_wandb_summary.json"
NUM_NODES = "1"
NUM_GPUS_PER_NODE = "2"
@@ -126,6 +127,8 @@ def test_distributed_training():
reference_wandb_summary_file = a40_reference_wandb_summary_file
elif "L40S" in device_name:
reference_wandb_summary_file = l40s_reference_wandb_summary_file
elif "H200" in device_name:
reference_wandb_summary_file = h200_reference_wandb_summary_file
else:
raise ValueError(f"Unknown device: {device_name}")
fastvideo/tests/training/lora/l40s_reference_lora_wandb_summary.json
+1 -33
View File
@@ -8,37 +8,5 @@
"learning_rate": 0.00005,
"step_time": 9.380061847000036,
"train_loss": 0.1361890360713005,
"validation_videos_50_steps": {
"_type": "videos",
"captions": [
"A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
"A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
"A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table."
],
"count": 3,
"videos": [
{
"_type": "video-file",
"caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
"path": "media/videos/validation_videos_50_steps_0_54dafffa2c2dea9fcb21.mp4",
"sha256": "54dafffa2c2dea9fcb21b734a6b4abb072b99a32e57603899d12b8eaf4324434",
"size": 611595
},
{
"_type": "video-file",
"caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
"path": "media/videos/validation_videos_50_steps_0_6d82abfb6ab979f9df94.mp4",
"sha256": "6d82abfb6ab979f9df94e96fcd1ac344f40f03cb3f106a6b1d9c510dc17e8f91",
"size": 460476
},
{
"_type": "video-file",
"caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.",
"path": "media/videos/validation_videos_50_steps_0_cedf3954f184052149c8.mp4",
"sha256": "cedf3954f184052149c83ac23776ad83f75cf911351fcb6f8e05854ba4d4c614",
"size": 738913
}
]
},
"vsa_sparsity": 0
}
}
fastvideo/training/distillation_pipeline.py
+66 -86
View File
@@ -7,7 +7,7 @@ import time
from abc import abstractmethod
from collections import deque
from collections.abc import Iterator
from typing import Any
from typing import Any, cast
import imageio
import numpy as np
@@ -58,6 +58,10 @@ class DistillationPipeline(TrainingPipeline):
"transformer",
"vae",
]
# used by lora training
trainable_transformer_names: list[str] = [
"transformer", "fake_score_transformer"
]
validation_pipeline: ComposedPipelineBase
train_dataloader: StatefulDataLoader
train_loader_iter: Iterator[dict[str, Any]]
@@ -72,6 +76,65 @@ class DistillationPipeline(TrainingPipeline):
raise RuntimeError(
"create_pipeline_stages should not be called for training pipeline")
def load_modules(self,
fastvideo_args: FastVideoArgs,
loaded_modules: dict[str, torch.nn.Module] | None = None):
modules = super().load_modules(fastvideo_args, loaded_modules)
training_args = cast(TrainingArgs, fastvideo_args)
if training_args.real_score_model_path:
logger.info("Loading real score transformer from: %s",
training_args.real_score_model_path)
training_args.override_transformer_cls_name = "WanTransformer3DModel"
# TODO(will): can use deepcopy instead if the model is the same
self.real_score_transformer = self.load_module_from_path(
training_args.real_score_model_path, "transformer",
training_args)
modules["real_score_transformer"] = self.real_score_transformer
try:
self.real_score_transformer_2 = self.load_module_from_path(
training_args.real_score_model_path, "transformer_2",
training_args)
logger.info("Loaded real score transformer_2 for MoE support")
modules[
"real_score_transformer_2"] = self.real_score_transformer_2
except Exception:
logger.info(
"real score transformer_2 not found, using single transformer"
)
self.real_score_transformer_2 = None
else:
raise ValueError(
"real_score_model_path is required for DMD distillation pipeline"
)
if training_args.fake_score_model_path:
logger.info("Loading fake score transformer from: %s",
training_args.fake_score_model_path)
training_args.override_transformer_cls_name = "WanTransformer3DModel"
self.fake_score_transformer = self.load_module_from_path(
training_args.fake_score_model_path, "transformer",
training_args)
modules["fake_score_transformer"] = self.fake_score_transformer
try:
self.fake_score_transformer_2 = self.load_module_from_path(
training_args.fake_score_model_path, "transformer_2",
training_args)
logger.info("Loaded fake score transformer_2 for MoE support")
modules[
"fake_score_transformer_2"] = self.fake_score_transformer_2
except Exception:
logger.info(
"fake score transformer_2 not found, using single transformer"
)
self.fake_score_transformer_2 = None
else:
raise ValueError(
"fake_score_model_path is required for DMD distillation pipeline"
)
return modules
def initialize_training_pipeline(self, training_args: TrainingArgs):
"""Initialize the distillation training pipeline with multiple models."""
logger.info("Initializing distillation pipeline...")
@@ -91,65 +154,13 @@ class DistillationPipeline(TrainingPipeline):
else:
self.boundary_timestep = None
if training_args.real_score_model_path:
logger.info("Loading real score transformer from: %s",
training_args.real_score_model_path)
training_args.override_transformer_cls_name = "WanTransformer3DModel"
self.real_score_transformer = self.load_module_from_path(
training_args.real_score_model_path, "transformer",
training_args)
try:
self.real_score_transformer_2 = self.load_module_from_path(
training_args.real_score_model_path, "transformer_2",
training_args)
logger.info("Loaded real score transformer_2 for MoE support")
except Exception:
logger.info(
"real score transformer_2 not found, using single transformer"
)
self.real_score_transformer_2 = None
else:
self.real_score_transformer = self.get_module(
"real_score_transformer")
self.real_score_transformer_2 = self.get_module(
"real_score_transformer_2")
if training_args.fake_score_model_path:
logger.info("Loading fake score transformer from: %s",
training_args.fake_score_model_path)
training_args.override_transformer_cls_name = "WanTransformer3DModel"
self.fake_score_transformer = self.load_module_from_path(
training_args.fake_score_model_path, "transformer",
training_args)
try:
self.fake_score_transformer_2 = self.load_module_from_path(
training_args.fake_score_model_path, "transformer_2",
training_args)
logger.info("Loaded fake score transformer_2 for MoE support")
except Exception:
logger.info(
"fake score transformer_2 not found, using single transformer"
)
self.fake_score_transformer_2 = None
else:
self.fake_score_transformer = self.get_module(
"fake_score_transformer")
self.fake_score_transformer_2 = self.get_module(
"fake_score_transformer_2")
# make sure the real score transformer is not trainable
self.real_score_transformer.requires_grad_(False)
self.real_score_transformer.eval()
if self.real_score_transformer_2 is not None:
self.real_score_transformer_2.requires_grad_(False)
self.real_score_transformer_2.eval()
# Set training modes for fake score transformers (trainable)
self.fake_score_transformer.requires_grad_(True)
self.fake_score_transformer.train()
if self.fake_score_transformer_2 is not None:
self.fake_score_transformer_2.requires_grad_(True)
self.fake_score_transformer_2.train()
if training_args.enable_gradient_checkpointing_type is not None:
self.fake_score_transformer = apply_activation_checkpointing(
self.fake_score_transformer,
@@ -349,22 +360,6 @@ class DistillationPipeline(TrainingPipeline):
raise NotImplementedError(
"Distillation pipelines must implement this method")
def _prepare_distillation(self,
training_batch: TrainingBatch) -> TrainingBatch:
"""Prepare training environment for distillation."""
self.transformer.requires_grad_(True)
self.transformer.train()
if self.transformer_2 is not None:
self.transformer_2.requires_grad_(True)
self.transformer_2.train()
self.fake_score_transformer.requires_grad_(True)
self.fake_score_transformer.train()
if self.fake_score_transformer_2 is not None:
self.fake_score_transformer_2.requires_grad_(True)
self.fake_score_transformer_2.train()
return training_batch
def apply_ema_to_model(self, model):
"""Apply EMA weights to the model for validation or inference."""
if model is self.transformer and self.generator_ema is not None:
@@ -948,8 +943,7 @@ class DistillationPipeline(TrainingPipeline):
batches = []
# Collect N batches for gradient accumulation
for _ in range(gradient_accumulation_steps):
batch = self._prepare_distillation(training_batch)
batch = self._get_next_batch(batch)
batch = self._get_next_batch(training_batch)
batch = self._normalize_dit_input(batch)
batch = self._prepare_dit_inputs(batch)
batch = self._build_attention_metadata(batch)
@@ -990,9 +984,6 @@ class DistillationPipeline(TrainingPipeline):
# Only clip gradients for the model that is currently training
self._clip_model_grad_norm_(batch_gen, self.transformer)
for param in self.transformer.parameters():
# check if the gradient is not None and not zero
assert param.grad is not None and param.grad.abs().sum() > 0
self.optimizer.step()
self.optimizer.zero_grad(set_to_none=True)
@@ -1031,17 +1022,6 @@ class DistillationPipeline(TrainingPipeline):
else:
self._clip_model_grad_norm_(batch_fake, self.fake_score_transformer)
# Check gradients for fake score transformer
for param in self.fake_score_transformer.parameters():
if param.requires_grad:
assert param.grad is not None and param.grad.abs().sum() > 0
# Check gradients for fake score transformer_2 if available
if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None:
for param in self.fake_score_transformer_2.parameters():
if param.requires_grad:
assert param.grad is not None and param.grad.abs().sum() > 0
if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None:
self.fake_score_optimizer_2.step()
self.fake_score_lr_scheduler_2.step()
fastvideo/training/self_forcing_distillation_pipeline.py
+1 -2
View File
@@ -677,8 +677,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
batches = []
for _ in range(gradient_accumulation_steps):
batch = self._prepare_distillation(training_batch)
batch = self._get_next_batch(batch)
batch = self._get_next_batch(training_batch)
batch = self._normalize_dit_input(batch)
batch = self._prepare_dit_inputs(batch)
batch = self._build_attention_metadata(batch)
fastvideo/training/training_pipeline.py
-28
View File
@@ -153,8 +153,6 @@ class TrainingPipeline(LoRAPipeline, ABC):
)
if self.transformer_2 is not None:
# Ensure transformer_2 has trainable parameters before creating optimizer
self.transformer_2.train()
self.transformer_2.requires_grad_(True)
params_to_optimize_2 = self.transformer_2.parameters()
params_to_optimize_2 = list(
filter(lambda p: p.requires_grad, params_to_optimize_2))
@@ -238,29 +236,12 @@ class TrainingPipeline(LoRAPipeline, ABC):
"Training pipelines must implement this method")
def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch:
self.transformer.train()
self.optimizer.zero_grad()
if self.transformer_2 is not None:
self.transformer_2.train()
self.optimizer_2.zero_grad()
training_batch.total_loss = 0.0
return training_batch
def _enable_training(self, model: torch.nn.Module,
optimizer: torch.optim.Optimizer) -> None:
"""Enable training mode and gradients for the specified model."""
for param in model.parameters():
param.requires_grad = True
model.train()
optimizer.zero_grad()
def _disable_training(self, model: torch.nn.Module,
optimizer: torch.optim.Optimizer) -> None:
"""Disable training mode and gradients for the specified model."""
for param in model.parameters():
param.requires_grad = False
optimizer.zero_grad(set_to_none=True)
def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
with self.tracker.timed("timing/get_next_batch"):
batch = next(self.train_loader_iter, None) # type: ignore
@@ -311,15 +292,6 @@ class TrainingPipeline(LoRAPipeline, ABC):
dtype=latents.dtype)
timesteps = self._sample_timesteps(batch_size, latents.device)
# Enable training for the model that will be trained next and disable the other
if self.train_transformer_2:
self._enable_training(self.transformer_2, self.optimizer_2)
self._disable_training(self.transformer, self.optimizer)
else:
self._enable_training(self.transformer, self.optimizer)
if self.transformer_2 is not None:
self._disable_training(self.transformer_2, self.optimizer_2)
if self.training_args.sp_size > 1:
# Make sure that the timesteps are the same across all sp processes.
sp_group = get_sp_group()