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

Fix small runtime issues

Browse Source
This commit is contained in:
JerryZhou54
2025-10-07 05:03:36 +00:00
parent aeb8f2e5ac
commit 55d8c1e5fb
3 changed files with 24 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/distill/SFWan2.2-A14B/distill_dmd.slurm
+1 -1
View File
@@ -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
fastvideo/training/distillation_pipeline.py
+13 -10
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, 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,
fastvideo/training/self_forcing_distillation_pipeline.py
+10 -7
View File
@@ -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,