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

Add matthew's timestep change & add real score guidance scale 2

Browse Source
This commit is contained in:
JerryZhou54
2025-10-06 22:24:44 +00:00
parent f755dd1ad5
commit aeb8f2e5ac
8 changed files with 120 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/distill/SFWan2.2-A14B/distill_dmd.slurm
+24 -23
View File
@@ -1,14 +1,14 @@
#!/bin/bash
#SBATCH --job-name=moe_sf_distill
#SBATCH --job-name=moe_4n_sf_distill
#SBATCH --partition=main
#SBATCH --nodes=2
#SBATCH --ntasks=2
#SBATCH --nodes=4
#SBATCH --ntasks=4
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:8
#SBATCH --cpus-per-task=128
#SBATCH --mem=1440G
#SBATCH --output=moe_sf_distill_output/moe_sf_distill_%j.out
#SBATCH --error=moe_sf_distill_output/moe_sf_distill_%j.err
#SBATCH --output=moe_4n_sf_distill_output/moe_sf_distill_%j.out
#SBATCH --error=moe_4n_sf_distill_output/moe_sf_distill_%j.err
#SBATCH --exclusive
# Basic Info
@@ -35,16 +35,16 @@ export HOME="/mnt/weka/home/hao.zhang/wei"
# Configs
NUM_GPUS=8
NUM_TOTAL_GPUS=16
NUM_TOTAL_GPUS=32
# Model paths for Self-Forcing DMD distillation with Wan2.2:
# GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"
GENERATOR_MODEL_PATH="rand0nmr/SFWan2.2-T2V-A14B-Diffusers" # Updated to Wan2.2
# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model
# FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model
REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model
FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model
# GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" # Updated to Wan2.2
REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Teacher model
FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model
# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Teacher model
# FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model
# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
@@ -62,7 +62,7 @@ training_args=(
--max_train_steps 4000
--train_batch_size 1
--train_sp_batch_size 1
--gradient_accumulation_steps 1
--gradient_accumulation_steps 2
--num_latent_t 21
--num_height 448 # Updated to match Wan2.2 config
--num_width 832 # Updated to match Wan2.2 config
@@ -73,8 +73,8 @@ training_args=(
--num_frame_per_block 3 # Frame generation block size for self-forcing
--enable_gradient_masking
--gradient_mask_last_n_frames 21
# --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/high/3k/
# --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/low/3k/
--init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/high/3k/
--init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/low/3k/
# --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors
)
@@ -99,21 +99,21 @@ dataset_args=(
)
validation_args=(
# --log_validation
--log_validation
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 20
--validation_sampling_steps "4"
--validation_steps 40
--validation_sampling_steps "8"
--validation_guidance_scale "6.0" # not used for dmd inference
)
optimizer_args=(
--learning_rate 1e-5
--learning_rate 2e-6
--mixed_precision "bf16"
--training_state_checkpointing_steps 500
--weight_only_checkpointing_steps 500
--weight_decay 0.01
--betas '0.0,0.999'
--max_grad_norm 1.0
--max_grad_norm 10.0
)
miscellaneous_args=(
@@ -121,20 +121,21 @@ miscellaneous_args=(
--checkpoints_total_limit 3
--training_cfg_rate 0.0
--dit_precision "fp32"
--flow_shift 5
--flow_shift 12
--seed 1000
--use_ema True
--ema_decay 0.99
--ema_start_step 20
--ema_start_step 200
)
dmd_args=(
--dmd_denoising_steps '1000,750,500,250'
--dmd_denoising_steps '1000,850,700,550,350,275,200,125'
--min_timestep_ratio 0.02
--max_timestep_ratio 0.98
--dfake_gen_update_ratio 5
--real_score_guidance_scale 3.0
--fake_score_learning_rate 8e-6
--real_score_guidance_scale 4.0
--real_score_guidance_scale_2 3.0
--fake_score_learning_rate 4e-7
--fake_score_betas '0.0,0.999'
--warp_denoising_step
)
fastvideo/fastvideo_args.py
+5
View File
@@ -717,6 +717,7 @@ class TrainingArgs(FastVideoArgs):
min_timestep_ratio: float = 0.2
max_timestep_ratio: float = 0.98
real_score_guidance_scale: float = 3.5
real_score_guidance_scale_2: float = 3.5
fake_score_learning_rate: float = 0.0 # separate learning rate for fake_score_transformer, if 0.0, use learning_rate
fake_score_lr_scheduler: str = "constant" # separate lr scheduler for fake_score_transformer, if not set, use lr_scheduler
fake_score_betas: str = "0.9,0.999" # betas for fake score optimizer, format: "beta1,beta2"
@@ -1102,6 +1103,10 @@ class TrainingArgs(FastVideoArgs):
type=float,
default=TrainingArgs.real_score_guidance_scale,
help="Teacher guidance scale")
parser.add_argument("--real-score-guidance-scale-2",
type=float,
default=TrainingArgs.real_score_guidance_scale_2,
help="Teacher guidance scale")
parser.add_argument("--fake-score-learning-rate",
type=float,
default=TrainingArgs.fake_score_learning_rate,
fastvideo/pipelines/composed_pipeline_base.py
+1 -1
View File
@@ -71,7 +71,7 @@ class ComposedPipelineBase(ABC):
# Load modules directly in initialization
logger.info("Loading pipeline modules...")
# fastvideo_args.dit_cpu_offload = True
fastvideo_args.dit_cpu_offload = True
self.modules = self.load_modules(fastvideo_args, loaded_modules)
def set_trainable(self) -> None:
fastvideo/pipelines/stages/causal_denoising.py
+6
View File
@@ -1,4 +1,5 @@
import torch # type: ignore
import gc
from fastvideo.distributed import get_local_torch_device
from fastvideo.fastvideo_args import FastVideoArgs
@@ -364,6 +365,11 @@ class CausalDMDDenosingStage(DenoisingStage):
)
start_index += current_num_frames
del kv_cache1
del crossattn_cache
gc.collect()
torch.cuda.empty_cache()
batch.latents = latents
return batch
fastvideo/training/distillation_pipeline.py
+48 -10
View File
@@ -260,6 +260,7 @@ class DistillationPipeline(TrainingPipeline):
self.num_train_timestep)
self.real_score_guidance_scale = self.training_args.real_score_guidance_scale
self.real_score_guidance_scale_2 = self.training_args.real_score_guidance_scale_2
self.generator_ema: EMA_FSDP | None = None
self.generator_ema_2: EMA_FSDP | None = None
@@ -541,6 +542,15 @@ class DistillationPipeline(TrainingPipeline):
self.generator_ema_2.decay = original_decay_2
logger.info("EMA_2 reset completed")
def _get_real_score_guidance_scale(self, timestep: torch.Tensor):
"""
Get the appropriate real score guidance scale based on timestep and boundary logic.
"""
if self.boundary_timestep is not None:
if timestep.item() < self.boundary_timestep:
return self.real_score_guidance_scale_2
return self.real_score_guidance_scale
def _get_real_score_transformer(self, timestep: torch.Tensor):
"""
Get the appropriate real score transformer based on timestep and boundary logic.
@@ -732,15 +742,28 @@ class DistillationPipeline(TrainingPipeline):
scheduler=self.noise_scheduler).unflatten(0, pred_noise.shape[:2])
training_batch.dmd_latent_vis_dict[
"generator_timestep"] = target_timestep.float().detach()
return pred_video
return pred_video, target_timestep
def _dmd_forward(self, generator_pred_video: torch.Tensor,
training_batch: TrainingBatch) -> torch.Tensor:
training_batch: TrainingBatch,
exit_timestep: int) -> torch.Tensor:
"""Compute DMD (Diffusion Model Distillation) loss."""
original_latent = generator_pred_video
with torch.no_grad():
timestep = torch.randint(0,
self.num_train_timestep, [1],
# Detect whether we're doing MoE DMD
if self.boundary_timestep is not None:
if exit_timestep < self.boundary_timestep:
start_timestep = 0
end_timestep = self.boundary_timestep
else:
start_timestep = self.boundary_timestep
end_timestep = self.num_train_timestep
else:
start_timestep = 0
end_timestep = self.num_train_timestep
timestep = torch.randint(start_timestep,
end_timestep, [1],
device=self.device,
dtype=torch.long)
world_group = get_world_group()
@@ -750,7 +773,8 @@ class DistillationPipeline(TrainingPipeline):
timestep = shift_timestep(
timestep,
self.timestep_shift, # type: ignore
self.num_train_timestep)
min_timestep=start_timestep,
max_timestep=end_timestep)
timestep = timestep.clamp(self.min_timestep, self.max_timestep)
@@ -817,7 +841,7 @@ class DistillationPipeline(TrainingPipeline):
real_score_pred_video = pred_real_video_cond + (
pred_real_video_cond -
pred_real_video_uncond) * self.real_score_guidance_scale
pred_real_video_uncond) * self._get_real_score_guidance_scale(timestep)
grad = (faker_score_pred_video - real_score_pred_video) / torch.abs(
original_latent - real_score_pred_video).mean()
@@ -849,13 +873,26 @@ class DistillationPipeline(TrainingPipeline):
current_timestep=training_batch.timesteps,
attn_metadata=training_batch.attn_metadata_vsa):
if self.training_args.simulate_generator_forward:
generator_pred_video = self._generator_multi_step_simulation_forward(
generator_pred_video, exit_timestep = self._generator_multi_step_simulation_forward(
training_batch)
else:
generator_pred_video = self._generator_forward(training_batch)
fake_score_timestep = torch.randint(0,
self.num_train_timestep, [1],
# Detect whether we're doing MoE DMD
# If yes, then we need to use the appropriate timestep range
if self.boundary_timestep is not None:
if exit_timestep < self.boundary_timestep:
start_timestep = 0
end_timestep = self.boundary_timestep
else:
start_timestep = self.boundary_timestep
end_timestep = self.num_train_timestep
else:
start_timestep = 0
end_timestep = self.num_train_timestep
fake_score_timestep = torch.randint(start_timestep,
end_timestep, [1],
device=self.device,
dtype=torch.long)
world_group = get_world_group()
@@ -865,7 +902,8 @@ class DistillationPipeline(TrainingPipeline):
fake_score_timestep = shift_timestep(
fake_score_timestep,
self.timestep_shift, # type: ignore
self.num_train_timestep)
min_timestep=start_timestep,
max_timestep=end_timestep)
fake_score_timestep = fake_score_timestep.clamp(self.min_timestep,
self.max_timestep)
fastvideo/training/self_forcing_distillation_pipeline.py
+30 -5
View File
@@ -155,14 +155,15 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
with set_forward_context(
current_timestep=training_batch.timesteps,
attn_metadata=training_batch.attn_metadata_vsa):
generator_pred_video = self._generator_multi_step_simulation_forward(
generator_pred_video, exit_timestep = self._generator_multi_step_simulation_forward(
training_batch, exit_flags=exit_flags)
with set_forward_context(current_timestep=training_batch.timesteps,
attn_metadata=training_batch.attn_metadata):
dmd_loss = self._dmd_forward(
generator_pred_video=generator_pred_video,
training_batch=training_batch)
training_batch=training_batch,
exit_timestep=exit_timestep)
log_dict = {
"dmdtrain_gradient_norm": torch.tensor(0.0, device=self.device)
@@ -320,6 +321,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
exit_flags = self.generate_and_sync_list(len(all_num_frames),
num_denoising_steps,
device=noise.device)
exit_timestep = self.denoising_step_list[exit_flags[0]]
start_gradient_frame_index = max(0, num_output_frames - 21)
for block_index, current_num_frames in enumerate(all_num_frames):
@@ -561,7 +564,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
# assert crossattn_cache is not None
# self._reset_simulation_caches(self.kv_cache1, self.crossattn_cache)
return final_output if gradient_mask is not None else pred_image_or_video
return final_output, exit_timestep if gradient_mask is not None else pred_image_or_video, exit_timestep
def _initialize_simulation_caches(
self,
batch_size: int,
@@ -784,7 +787,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
grad_sum += p.grad.sum().item()
else:
raise ValueError("Transformer 2 param %s has no gradient", n)
logger.info("Transformer 2 gradient sum for all params %s", grad_sum)
assert grad_sum != 0, "Transformer 2 did not receive gradients"
else:
assert all(p.grad is None for p in self.transformer_2.parameters())
grad_sum = 0
@@ -793,7 +796,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
grad_sum += p.grad.sum().item()
else:
raise ValueError("Transformer param %s has no gradient", n)
logger.info("Transformer gradient sum for all params %s", grad_sum)
assert grad_sum != 0, "Transformer did not receive gradients"
total_generator_loss += generator_loss.detach().item()
generator_log_dict.update(gen_log_dict)
@@ -839,6 +842,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
logger.debug("Training critic at step %s", self.current_trainstep)
self.fake_score_optimizer.zero_grad()
if self.fake_score_transformer_2 is not None:
self.fake_score_optimizer_2.zero_grad()
total_critic_loss = 0.0
critic_log_dict = {}
@@ -863,6 +868,26 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
with set_forward_context(current_timestep=batch_critic.timesteps,
attn_metadata=batch_critic.attn_metadata):
(critic_loss / gradient_accumulation_steps).backward()
if self.train_fake_score_transformer_2:
assert all(p.grad is None for p in self.fake_score_transformer.parameters())
grad_sum = 0
for n, p in self.fake_score_transformer_2.named_parameters():
if p.grad is not None:
grad_sum += p.grad.sum().item()
else:
raise ValueError("Fake score transformer 2 param %s has no gradient", n)
assert grad_sum != 0, "Fake score transformer 2 did not receive gradients"
else:
assert all(p.grad is None for p in self.fake_score_transformer_2.parameters())
grad_sum = 0
for n, p in self.fake_score_transformer.named_parameters():
if p.grad is not None:
grad_sum += p.grad.sum().item()
else:
raise ValueError("Fake score transformer param %s has no gradient", n)
assert grad_sum != 0, "Fake score transformer did not receive gradients"
total_critic_loss += critic_loss.detach().item()
critic_log_dict.update(crit_log_dict)
# Store visualization data from critic training
fastvideo/training/training_pipeline.py
+2 -2
View File
@@ -236,7 +236,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
for param in model.parameters():
param.requires_grad = True
model.train()
optimizer.zero_grad()
# optimizer.zero_grad()
def _disable_training(self, model: torch.nn.Module,
optimizer: torch.optim.Optimizer) -> None:
@@ -244,7 +244,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
for param in model.parameters():
param.requires_grad = False
model.eval()
optimizer.zero_grad(set_to_none=True)
# optimizer.zero_grad(set_to_none=True)
def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
batch = next(self.train_loader_iter, None) # type: ignore
fastvideo/training/training_utils.py
+4 -4
View File
@@ -1152,14 +1152,14 @@ def custom_to_hf_state_dict(
return new_state_dict
# More generalized version of shift_timestep
def shift_timestep(timestep: torch.Tensor, shift: float,
num_train_timestep: float) -> torch.Tensor:
min_timestep: int = 0, max_timestep: int = 1000) -> torch.Tensor:
if shift == 1:
return timestep
t = timestep / num_train_timestep
t = (timestep - min_timestep) / (max_timestep - min_timestep)
denominator = 1 + (shift - 1) * t
return num_train_timestep * (shift * t / denominator)
return min_timestep + (max_timestep - min_timestep) * (shift * t / denominator)
# coding=utf-8