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

[Feature] Support Lora for DMD (#755)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Wenxuan Tan
2025-09-08 14:18:21 -07:00
committed by GitHub
co-authored by gemini-code-assist[bot]
parent 85eb7265db
commit 076f45c1ee
8 changed files with 191 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh
+1
View File
@@ -98,6 +98,7 @@ dmd_args=(
torchrun \
--nnodes 1 \
--nproc_per_node $NUM_GPUS \
--master_port $MASTER_PORT \
fastvideo/training/wan_distillation_pipeline.py \
"${parallel_args[@]}" \
"${model_args[@]}" \
examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B_lora.sh
+112
View File
@@ -0,0 +1,112 @@
#!/bin/bash
# Basic Info
export WANDB_MODE="online"
export NCCL_P2P_DISABLE=1
export TORCH_NCCL_ENABLE_MONITORING=0
export MASTER_PORT=29501
export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
# Configs
NUM_GPUS=1
MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
DATA_DIR="data/crush-smol_processed_ti2v/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
--tracker_project_name wan_t2v_distill_dmd_VSA
--output_dir="checkpoints/wan_t2v_finetune"
--max_train_steps=4000
--train_batch_size=1
--train_sp_batch_size 1
--gradient_accumulation_steps=1
--num_latent_t 31
--num_height 704
--num_width 1280
--num_frames 121
--enable_gradient_checkpointing_type "full"
--training_state_checkpointing_steps=500
--weight_only_checkpointing_steps=500
--lora_rank 32
--lora_training True
)
# Parallel arguments
parallel_args=(
--num_gpus 1
--sp_size 1
--tp_size 1
--hsdp_replicate_dim 1
--hsdp_shard_dim 1
)
# Model arguments
model_args=(
--model_path $MODEL_PATH
--pretrained_model_name_or_path $MODEL_PATH
)
# Dataset arguments
dataset_args=(
--data_path "$DATA_DIR"
--dataloader_num_workers 4
)
# Validation arguments
validation_args=(
--log_validation
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 200
--validation_sampling_steps "3"
--validation_guidance_scale "6.0" # not used for dmd inference
)
# Optimizer arguments
optimizer_args=(
--learning_rate=1e-4
--mixed_precision="bf16"
--weight_decay 0.01
--max_grad_norm 1.0
)
# Miscellaneous arguments
miscellaneous_args=(
--inference_mode False
--checkpoints_total_limit 3
--training_cfg_rate 0.0
--dit_precision "fp32"
--ema_start_step 0
--flow_shift 8
--seed 1000
)
# DMD arguments
dmd_args=(
--dmd_denoising_steps '1000,757,522'
--min_timestep_ratio 0.02
--max_timestep_ratio 0.98
--generator_update_interval 5
--real_score_guidance_scale 3.5
--VSA_sparsity 0.8
)
torchrun \
--nnodes 1 \
--nproc_per_node $NUM_GPUS \
--master_port $MASTER_PORT \
fastvideo/training/wan_distillation_pipeline.py \
"${parallel_args[@]}" \
"${model_args[@]}" \
"${dataset_args[@]}" \
"${training_args[@]}" \
"${optimizer_args[@]}" \
"${validation_args[@]}" \
"${miscellaneous_args[@]}" \
"${dmd_args[@]}"
fastvideo/pipelines/composed_pipeline_base.py
+1 -1
View File
@@ -79,7 +79,7 @@ class ComposedPipelineBase(ABC):
for name, module in self.modules.items():
if not isinstance(module, torch.nn.Module):
continue
if name == "transformer":
if "transformer" in name:
module.requires_grad_(True)
else:
module.requires_grad_(False)
fastvideo/pipelines/lora_pipeline.py
+36 -12
View File
@@ -7,7 +7,7 @@ import torch
import torch.distributed as dist
import torch.nn as nn
from safetensors.torch import load_file
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
from torch.distributed.tensor import DTensor
from fastvideo.distributed import get_local_torch_device
@@ -32,6 +32,7 @@ class LoRAPipeline(ComposedPipelineBase):
cur_adapter_name: str = ""
cur_adapter_path: str = ""
lora_layers: dict[str, BaseLayerWithLoRA] = {}
lora_layers_critic: dict[str, BaseLayerWithLoRA] = {}
fastvideo_args: FastVideoArgs | TrainingArgs
exclude_lora_layers: list[str] = []
device: torch.device = get_local_torch_device()
@@ -81,6 +82,17 @@ class LoRAPipeline(ComposedPipelineBase):
def set_trainable(self) -> None:
def set_lora_grads(lora_layers: dict[str, BaseLayerWithLoRA],
device_mesh: DeviceMesh):
for name, layer in lora_layers.items():
layer.lora_A.requires_grad_(True)
layer.lora_B.requires_grad_(True)
layer.base_layer.requires_grad_(False)
layer.lora_A = nn.Parameter(
DTensor.from_local(layer.lora_A, device_mesh=device_mesh))
layer.lora_B = nn.Parameter(
DTensor.from_local(layer.lora_B, device_mesh=device_mesh))
is_lora_training = self.training_mode and getattr(
self.fastvideo_args, "lora_training", False)
if not is_lora_training:
@@ -88,18 +100,12 @@ class LoRAPipeline(ComposedPipelineBase):
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"])
for name, layer in self.lora_layers.items():
# Enable grads for lora weights only
# Must convert to DTensor for compatibility with other FSDP modules in grad calculation
layer.lora_A.requires_grad_(True)
layer.lora_B.requires_grad_(True)
layer.base_layer.requires_grad_(False)
layer.lora_A = nn.Parameter(
DTensor.from_local(layer.lora_A, device_mesh=device_mesh))
layer.lora_B = nn.Parameter(
DTensor.from_local(layer.lora_B, device_mesh=device_mesh))
set_lora_grads(self.lora_layers, device_mesh)
set_lora_grads(self.lora_layers_critic, device_mesh)
def convert_to_lora_layers(self) -> None:
"""
@@ -131,6 +137,24 @@ class LoRAPipeline(ComposedPipelineBase):
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():
if not self.is_target_layer(name):
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)
converted_count += 1
logger.info(
"Converted %d layers to LoRA layers in the critic model",
converted_count)
def set_lora_adapter(self,
lora_nickname: str,
lora_path: str | None = None): # type: ignore
@@ -224,4 +248,4 @@ class LoRAPipeline(ComposedPipelineBase):
def unmerge_lora_weights(self) -> None:
for name, layer in self.lora_layers.items():
layer.unmerge_lora_weights()
layer.unmerge_lora_weights()
fastvideo/pipelines/stages/denoising.py
+13 -6
View File
@@ -149,7 +149,8 @@ class DenoisingStage(PipelineStage):
# Prepare image latents and embeddings for I2V generation
image_embeds = batch.image_embeds
if len(image_embeds) > 0:
assert torch.isnan(image_embeds[0]).sum() == 0
assert not torch.isnan(
image_embeds[0]).any(), "image_embeds contains nan"
image_embeds = [
image_embed.to(target_dtype) for image_embed in image_embeds
]
@@ -186,11 +187,13 @@ class DenoisingStage(PipelineStage):
# Get latents and embeddings
latents = batch.latents
prompt_embeds = batch.prompt_embeds
assert torch.isnan(prompt_embeds[0]).sum() == 0
assert not torch.isnan(
prompt_embeds[0]).any(), "prompt_embeds contains nan"
if batch.do_classifier_free_guidance:
neg_prompt_embeds = batch.negative_prompt_embeds
assert neg_prompt_embeds is not None
assert torch.isnan(neg_prompt_embeds[0]).sum() == 0
assert not torch.isnan(
neg_prompt_embeds[0]).any(), "neg_prompt_embeds contains nan"
# (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
if fastvideo_args.boundary_ratio is not None:
@@ -268,6 +271,9 @@ class DenoisingStage(PipelineStage):
latent_model_input = torch.cat(
[latent_model_input, batch.image_latent],
dim=1).to(target_dtype)
assert not torch.isnan(
latent_model_input).any(), "latent_model_input contains nan"
if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
timestep = torch.stack([t]).to(get_local_torch_device())
temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten()
@@ -280,7 +286,6 @@ class DenoisingStage(PipelineStage):
else:
t_expand = t.repeat(latent_model_input.shape[0])
assert torch.isnan(latent_model_input).sum() == 0
latent_model_input = self.scheduler.scale_model_input(
latent_model_input, t)
@@ -737,7 +742,8 @@ class DmdDenoisingStage(DenoisingStage):
video_raw_latent_shape = latents.shape
prompt_embeds = batch.prompt_embeds
assert torch.isnan(prompt_embeds[0]).sum() == 0
assert not torch.isnan(
prompt_embeds[0]).any(), "prompt_embeds contains nan"
timesteps = torch.tensor(
fastvideo_args.pipeline_config.dmd_denoising_steps,
dtype=torch.long,
@@ -776,7 +782,8 @@ class DmdDenoisingStage(DenoisingStage):
batch.image_latent.permute(0, 2, 1, 3, 4)
],
dim=2).to(target_dtype)
assert torch.isnan(latent_model_input).sum() == 0
assert not torch.isnan(
latent_model_input).any(), "latent_model_input contains nan"
# Prepare inputs for transformer
t_expand = t.repeat(latent_model_input.shape[0])
fastvideo/training/distillation_pipeline.py
+18 -9
View File
@@ -36,8 +36,9 @@ from fastvideo.training.activation_checkpoint import (
apply_activation_checkpointing)
from fastvideo.training.training_pipeline import TrainingPipeline
from fastvideo.training.training_utils import (
clip_grad_norm_while_handling_failing_dtensor_cases, get_scheduler,
load_distillation_checkpoint, save_distillation_checkpoint, shift_timestep)
clip_grad_norm_while_handling_failing_dtensor_cases, count_trainable,
get_scheduler, load_distillation_checkpoint, save_distillation_checkpoint,
shift_timestep)
from fastvideo.utils import is_vsa_available, set_random_seed
import wandb # isort: skip
@@ -68,11 +69,18 @@ class DistillationPipeline(TrainingPipeline):
current_trainstep: int
video_latent_shape: tuple[int, ...]
video_latent_shape_sp: tuple[int, ...]
real_score_transformer: torch.nn.Module
fake_score_transformer: torch.nn.Module
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
raise RuntimeError(
"create_pipeline_stages should not be called for training pipeline")
def set_trainable(self) -> None:
super().set_trainable()
self.modules["real_score_transformer"].requires_grad_(False)
self.modules["vae"].requires_grad_(False)
def initialize_training_pipeline(self, training_args: TrainingArgs):
"""Initialize the distillation training pipeline with multiple models."""
logger.info("Initializing distillation pipeline...")
@@ -81,8 +89,6 @@ class DistillationPipeline(TrainingPipeline):
self.noise_scheduler = self.get_module("scheduler")
self.vae = self.get_module("vae")
self.vae.requires_grad_(False)
self.timestep_shift = self.training_args.pipeline_config.flow_shift
self.noise_scheduler = FlowMatchEulerDiscreteScheduler(
shift=self.timestep_shift)
@@ -90,10 +96,7 @@ class DistillationPipeline(TrainingPipeline):
# self.transformer is the generator model
self.real_score_transformer = self.get_module("real_score_transformer")
self.fake_score_transformer = self.get_module("fake_score_transformer")
self.real_score_transformer.requires_grad_(False)
self.real_score_transformer.eval()
self.fake_score_transformer.requires_grad_(True)
self.fake_score_transformer.train()
if training_args.enable_gradient_checkpointing_type is not None:
@@ -167,9 +170,7 @@ class DistillationPipeline(TrainingPipeline):
def _prepare_distillation(self,
training_batch: TrainingBatch) -> TrainingBatch:
"""Prepare training environment for distillation."""
self.transformer.requires_grad_(True)
self.transformer.train()
self.fake_score_transformer.requires_grad_(True)
self.fake_score_transformer.train()
return training_batch
@@ -895,6 +896,14 @@ class DistillationPipeline(TrainingPipeline):
else:
set_random_seed(seed + self.global_rank)
# Check trainable params
num_trainable_generator = round(
count_trainable(self.transformer) / 1e9, 3)
num_trainable_critic = round(
count_trainable(self.fake_score_transformer) / 1e9, 3)
logger.info(
"rank: %s: # of trainable params in generator: %sB, # of trainable params in critic: %sB",
self.global_rank, num_trainable_generator, num_trainable_critic)
# Set random seeds for deterministic training
self.noise_random_generator = torch.Generator(device="cpu").manual_seed(
self.seed)
fastvideo/training/training_pipeline.py
+6 -10
View File
@@ -38,8 +38,8 @@ from fastvideo.training.activation_checkpoint import (
apply_activation_checkpointing)
from fastvideo.training.training_utils import (
clip_grad_norm_while_handling_failing_dtensor_cases,
compute_density_for_timestep_sampling, get_scheduler, get_sigmas,
load_checkpoint, normalize_dit_input, save_checkpoint,
compute_density_for_timestep_sampling, count_trainable, get_scheduler,
get_sigmas, load_checkpoint, normalize_dit_input, save_checkpoint,
shard_latents_across_sp)
from fastvideo.utils import is_vsa_available, set_random_seed, shallow_asdict
@@ -50,10 +50,6 @@ vsa_available = is_vsa_available()
logger = init_logger(__name__)
def _get_trainable_params(model: torch.nn.Module) -> int:
return sum(p.numel() for p in model.parameters() if p.requires_grad)
class TrainingPipeline(LoRAPipeline, ABC):
"""
A pipeline for training a model. All training pipelines should inherit from this class.
@@ -113,6 +109,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
enable_gradient_checkpointing_type)
noise_scheduler = self.modules["scheduler"]
# Set grads for proper modules based on the training mode (Distill, LoRA, etc.)
self.set_trainable()
params_to_optimize = self.transformer.parameters()
params_to_optimize = list(
@@ -417,7 +414,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
local_main_process_only=False)
if not self.post_init_called:
self.post_init()
num_trainable_params = _get_trainable_params(self.transformer)
num_trainable_params = count_trainable(self.transformer)
logger.info("Starting training with %s B trainable parameters",
round(num_trainable_params / 1e9, 3))
@@ -506,7 +503,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
self._log_validation(self.transformer, self.training_args, step)
gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
trainable_params = round(
_get_trainable_params(self.transformer) / 1e9, 3)
count_trainable(self.transformer) / 1e9, 3)
logger.info(
"GPU memory usage after validation: %s MB, trainable params: %sB",
gpu_memory_usage, trainable_params)
@@ -542,7 +539,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
logger.info(" Total optimization steps = %s",
self.training_args.max_train_steps)
logger.info(" Total training parameters per FSDP shard = %s B",
round(_get_trainable_params(self.transformer) / 1e9, 3))
round(count_trainable(self.transformer) / 1e9, 3))
# print dtype
logger.info(" Master weight dtype: %s",
self.transformer.parameters().__next__().dtype)
@@ -610,7 +607,6 @@ class TrainingPipeline(LoRAPipeline, ABC):
validation_dataloader = DataLoader(validation_dataset,
batch_size=None,
num_workers=0)
transformer.eval()
validation_steps = training_args.validation_sampling_steps.split(",")
fastvideo/training/training_utils.py
+4
View File
@@ -1278,3 +1278,7 @@ def get_scheduler(
num_warmup_steps=num_warmup_steps,
num_training_steps=num_training_steps,
last_epoch=last_epoch)
def count_trainable(model: torch.nn.Module) -> int:
return sum(p.numel() for p in model.parameters() if p.requires_grad)