From 55d8c1e5fbe516886a0516c557ee8f60e2f5a184 Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Tue, 7 Oct 2025 05:03:36 +0000 Subject: [PATCH] Fix small runtime issues --- .../distill/SFWan2.2-A14B/distill_dmd.slurm | 2 +- fastvideo/training/distillation_pipeline.py | 23 +++++++++++-------- .../self_forcing_distillation_pipeline.py | 17 ++++++++------ 3 files changed, 24 insertions(+), 18 deletions(-) diff --git a/examples/distill/SFWan2.2-A14B/distill_dmd.slurm b/examples/distill/SFWan2.2-A14B/distill_dmd.slurm index bd4a6f7f9..c2fd4edff 100644 --- a/examples/distill/SFWan2.2-A14B/distill_dmd.slurm +++ b/examples/distill/SFWan2.2-A14B/distill_dmd.slurm @@ -62,7 +62,7 @@ training_args=( --max_train_steps 4000 --train_batch_size 1 --train_sp_batch_size 1 - --gradient_accumulation_steps 2 + --gradient_accumulation_steps 1 --num_latent_t 21 --num_height 448 # Updated to match Wan2.2 config --num_width 832 # Updated to match Wan2.2 config diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 87a7c16ee..0a852e4d4 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -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, Tuple import imageio import numpy as np @@ -640,7 +640,9 @@ class DistillationPipeline(TrainingPipeline): return pred_video def _generator_multi_step_simulation_forward( - self, training_batch: TrainingBatch) -> torch.Tensor: + self, + training_batch: TrainingBatch, + exit_flags: list[int] | None = None) -> Tuple[torch.Tensor, float]: """Forward pass through student transformer matching inference procedure.""" latents = training_batch.latents dtype = latents.dtype @@ -742,11 +744,11 @@ 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, target_timestep + return pred_video, target_timestep.item() def _dmd_forward(self, generator_pred_video: torch.Tensor, training_batch: TrainingBatch, - exit_timestep: int) -> torch.Tensor: + exit_timestep: float) -> torch.Tensor: """Compute DMD (Diffusion Model Distillation) loss.""" original_latent = generator_pred_video with torch.no_grad(): @@ -754,9 +756,9 @@ class DistillationPipeline(TrainingPipeline): if self.boundary_timestep is not None: if exit_timestep < self.boundary_timestep: start_timestep = 0 - end_timestep = self.boundary_timestep + end_timestep = int(self.boundary_timestep) else: - start_timestep = self.boundary_timestep + start_timestep = int(self.boundary_timestep) end_timestep = self.num_train_timestep else: start_timestep = 0 @@ -883,9 +885,9 @@ class DistillationPipeline(TrainingPipeline): if self.boundary_timestep is not None: if exit_timestep < self.boundary_timestep: start_timestep = 0 - end_timestep = self.boundary_timestep + end_timestep = int(self.boundary_timestep) else: - start_timestep = self.boundary_timestep + start_timestep = int(self.boundary_timestep) end_timestep = self.num_train_timestep else: start_timestep = 0 @@ -1037,7 +1039,7 @@ class DistillationPipeline(TrainingPipeline): current_timestep=batch_gen.timesteps, attn_metadata=batch_gen.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( batch_gen) else: generator_pred_video = self._generator_forward( @@ -1047,7 +1049,8 @@ class DistillationPipeline(TrainingPipeline): attn_metadata=batch_gen.attn_metadata): dmd_loss = self._dmd_forward( generator_pred_video=generator_pred_video, - training_batch=batch_gen) + training_batch=batch_gen, + exit_timestep=exit_timestep) with set_forward_context( current_timestep=batch_gen.timesteps, diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index 5d1043d9e..e806bbe97 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -2,7 +2,7 @@ import copy import time from collections import deque -from typing import Any +from typing import Any, Tuple import gc import numpy as np @@ -133,8 +133,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): if self.boundary_timestep is not None and self.transformer_2 is not None: assert self.same_step_across_blocks, "same_step_across_blocks must be True for MoE generator. Otherwise we might need to train both transformers which will cause OOM" exit_flags = self.generate_and_sync_list(1, len(self.denoising_step_list), training_batch.latents.device) - exit_timestep = self.denoising_step_list[exit_flags[0]] - logger.info("Exit timestep in generator_loss(): %s", exit_timestep) + exit_timestep = self.denoising_step_list[exit_flags[0]].item() + # logger.info("Exit timestep in generator_loss(): %s", exit_timestep) if exit_timestep < self.boundary_timestep: self.train_transformer_2 = True self._enable_training(self.transformer_2, self.optimizer_2) @@ -198,8 +198,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): def _generator_multi_step_simulation_forward( self, training_batch: TrainingBatch, - return_sim_steps: bool = False, - exit_flags: list[int] | None = None) -> torch.Tensor: + exit_flags: list[int] | None = None) -> Tuple[torch.Tensor, float]: """Forward pass through student transformer matching inference procedure with KV cache management. This function is adapted from the reference self-forcing implementation's inference_with_trajectory @@ -321,7 +320,7 @@ 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]] + exit_timestep = self.denoising_step_list[exit_flags[0]].item() start_gradient_frame_index = max(0, num_output_frames - 21) @@ -564,7 +563,11 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # assert crossattn_cache is not None # self._reset_simulation_caches(self.kv_cache1, self.crossattn_cache) - return final_output, exit_timestep if gradient_mask is not None else pred_image_or_video, exit_timestep + if gradient_mask is not None: + return final_output, exit_timestep + else: + return pred_image_or_video, exit_timestep + def _initialize_simulation_caches( self, batch_size: int,