diff --git a/examples/distill/SFWan2.2-A14B/distill_dmd.slurm b/examples/distill/SFWan2.2-A14B/distill_dmd.slurm index 1dd592d70..bd4a6f7f9 100644 --- a/examples/distill/SFWan2.2-A14B/distill_dmd.slurm +++ b/examples/distill/SFWan2.2-A14B/distill_dmd.slurm @@ -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 ) diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 1f43b232a..7fa6d6e1b 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -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, diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index 87b9e03ff..823b23423 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -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: diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 64a528e23..80e8bbfae 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -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 diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 1df8255fd..87a7c16ee 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -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) diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index 8ed9f4692..5d1043d9e 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -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 diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index a058f5b38..e8c70b08e 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -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 diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index 01eeb3e5e..87c2ae715 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -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