Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 “BrianChen1129”
|
63c00ced0c |
@@ -187,6 +187,3 @@ class TrainingBatch:
|
||||
# Distillation losses
|
||||
generator_loss: float = 0.0
|
||||
fake_score_loss: float = 0.0
|
||||
|
||||
dmd_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
|
||||
fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
|
||||
@@ -12,7 +12,7 @@ from typing import Any
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
import time
|
||||
from fastvideo.attention import get_attn_backend
|
||||
from fastvideo.configs.pipelines.base import STA_Mode
|
||||
from fastvideo.distributed import (get_local_torch_device, get_sp_parallel_rank,
|
||||
@@ -88,6 +88,11 @@ class DenoisingStage(PipelineStage):
|
||||
Returns:
|
||||
The batch with denoised latents.
|
||||
"""
|
||||
# Start timing for the entire DenoisingStage
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
stage_start_time = time.time()
|
||||
|
||||
pipeline = self.pipeline() if self.pipeline else None
|
||||
if not fastvideo_args.model_loaded["transformer"]:
|
||||
loader = TransformerLoader()
|
||||
@@ -601,6 +606,11 @@ class DmdDenoisingStage(DenoisingStage):
|
||||
Returns:
|
||||
The batch with denoised latents.
|
||||
"""
|
||||
# Start timing for the entire DmdDenoisingStage
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
dmd_stage_start_time = time.time()
|
||||
|
||||
# Setup precision and autocast settings
|
||||
# TODO(will): make the precision configurable for inference
|
||||
# target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision]
|
||||
@@ -681,7 +691,13 @@ class DmdDenoisingStage(DenoisingStage):
|
||||
image_latent = image_latent[:, :, rank_in_sp_group, :, :, :]
|
||||
batch.image_latent = image_latent
|
||||
|
||||
# Run denoising loop
|
||||
|
||||
|
||||
# Synchronize CUDA before starting timer
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
denoise_start_time = time.time()
|
||||
|
||||
with self.progress_bar(total=len(timesteps)) as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
# Skip if interrupted
|
||||
@@ -715,9 +731,7 @@ class DmdDenoisingStage(DenoisingStage):
|
||||
with torch.autocast(device_type="cuda",
|
||||
dtype=target_dtype,
|
||||
enabled=autocast_enabled):
|
||||
if (st_attn_available
|
||||
and self.attn_backend == SlidingTileAttentionBackend
|
||||
) or (vsa_available and self.attn_backend
|
||||
if (vsa_available and self.attn_backend
|
||||
== VideoSparseAttentionBackend):
|
||||
self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
|
||||
)
|
||||
@@ -798,6 +812,20 @@ class DmdDenoisingStage(DenoisingStage):
|
||||
if sp_group:
|
||||
latents = sequence_model_parallel_all_gather(latents, dim=1)
|
||||
latents = latents.permute(0, 2, 1, 3, 4)
|
||||
|
||||
# Calculate total denoising time
|
||||
# Synchronize CUDA before ending timer
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
denoise_end_time = time.time()
|
||||
total_denoise_time = denoise_end_time - denoise_start_time
|
||||
print(f"Total denoising time: {total_denoise_time:.2f} seconds")
|
||||
|
||||
# Calculate total stage time
|
||||
dmd_stage_end_time = time.time()
|
||||
total_dmd_stage_time = dmd_stage_end_time - dmd_stage_start_time
|
||||
print(f"Total DmdDenoisingStage time: {total_dmd_stage_time:.2f} seconds")
|
||||
|
||||
# Update batch with final latents
|
||||
batch.latents = latents
|
||||
|
||||
|
||||
@@ -192,7 +192,6 @@ class DistillationPipeline(TrainingPipeline):
|
||||
device=self.device,
|
||||
dtype=torch.long)
|
||||
timestep = self.denoising_step_list[index]
|
||||
training_batch.dmd_latent_vis_dict["generator_timestep"] = timestep
|
||||
|
||||
noise = torch.randn(self.video_latent_shape,
|
||||
device=self.device,
|
||||
@@ -259,7 +258,7 @@ class DistillationPipeline(TrainingPipeline):
|
||||
fake_score_pred_noise = self.fake_score_transformer(
|
||||
**training_batch.input_kwargs).permute(0, 2, 1, 3, 4)
|
||||
|
||||
faker_score_pred_video = pred_noise_to_pred_video(
|
||||
pred_fake_video = pred_noise_to_pred_video(
|
||||
pred_noise=fake_score_pred_noise.flatten(0, 1),
|
||||
noise_input_latent=training_batch.noise_latents.flatten(0, 1),
|
||||
timestep=timestep,
|
||||
@@ -298,7 +297,7 @@ class DistillationPipeline(TrainingPipeline):
|
||||
pred_real_video_cond -
|
||||
pred_real_video_uncond) * self.real_score_guidance_scale
|
||||
|
||||
grad = (faker_score_pred_video - real_score_pred_video) / torch.abs(
|
||||
grad = (pred_fake_video - real_score_pred_video) / torch.abs(
|
||||
generator_pred_video - real_score_pred_video).mean()
|
||||
grad = torch.nan_to_num(grad)
|
||||
|
||||
@@ -306,19 +305,6 @@ class DistillationPipeline(TrainingPipeline):
|
||||
generator_pred_video.float(),
|
||||
(generator_pred_video.float() - grad.float()).detach())
|
||||
|
||||
training_batch.dmd_latent_vis_dict.update({
|
||||
"training_batch_dmd_fwd_clean_latent":
|
||||
training_batch.latents,
|
||||
"generator_pred_video":
|
||||
generator_pred_video,
|
||||
"real_score_pred_video":
|
||||
real_score_pred_video,
|
||||
"faker_score_pred_video":
|
||||
faker_score_pred_video,
|
||||
"dmd_timestep":
|
||||
timestep,
|
||||
})
|
||||
|
||||
return dmd_loss
|
||||
|
||||
def faker_score_forward(
|
||||
@@ -367,16 +353,9 @@ class DistillationPipeline(TrainingPipeline):
|
||||
**training_batch.input_kwargs).permute(0, 2, 1, 3, 4)
|
||||
|
||||
target = fake_score_noise - generator_pred_video
|
||||
flow_matching_loss = torch.mean((fake_score_pred_noise - target)**2)
|
||||
denoising_loss = torch.mean((fake_score_pred_noise - target)**2)
|
||||
|
||||
training_batch.fake_score_latent_vis_dict = {
|
||||
"training_batch_fakerscore_fwd_clean_latent":
|
||||
training_batch.latents,
|
||||
"generator_pred_video": generator_pred_video,
|
||||
"fake_score_timestep": fake_score_timestep,
|
||||
}
|
||||
|
||||
return training_batch, flow_matching_loss
|
||||
return training_batch, denoising_loss
|
||||
|
||||
def _clip_model_grad_norm_(self, training_batch: TrainingBatch,
|
||||
transformer) -> TrainingBatch:
|
||||
@@ -410,9 +389,6 @@ class DistillationPipeline(TrainingPipeline):
|
||||
"encoder_attention_mask": self.negative_prompt_attention_mask,
|
||||
}
|
||||
|
||||
training_batch.dmd_latent_vis_dict = {}
|
||||
training_batch.fake_score_latent_vis_dict = {}
|
||||
|
||||
training_batch.conditional_dict = conditional_dict
|
||||
training_batch.unconditional_dict = unconditional_dict
|
||||
training_batch.raw_latent_shape = training_batch.latents.shape
|
||||
@@ -449,29 +425,27 @@ class DistillationPipeline(TrainingPipeline):
|
||||
|
||||
self.optimizer.zero_grad()
|
||||
total_dmd_loss = 0.0
|
||||
dmd_latent_vis_dict = {}
|
||||
fake_score_latent_vis_dict = {}
|
||||
if (self.current_trainstep % self.generator_update_interval == 0):
|
||||
for batch in batches:
|
||||
batch_gen = copy.deepcopy(batch)
|
||||
batch_stu = copy.deepcopy(batch)
|
||||
|
||||
with set_forward_context(
|
||||
current_timestep=batch_gen.timesteps,
|
||||
attn_metadata=batch_gen.attn_metadata_vsa):
|
||||
generator_pred_video = self._generator_forward(batch_gen)
|
||||
current_timestep=batch_stu.timesteps,
|
||||
attn_metadata=batch_stu.attn_metadata_vsa):
|
||||
generator_pred_video = self._generator_forward(batch_stu)
|
||||
|
||||
with set_forward_context(current_timestep=batch_gen.timesteps,
|
||||
attn_metadata=batch_gen.attn_metadata):
|
||||
with set_forward_context(current_timestep=batch_stu.timesteps,
|
||||
attn_metadata=batch_stu.attn_metadata):
|
||||
dmd_loss = self._dmd_forward(
|
||||
generator_pred_video=generator_pred_video,
|
||||
training_batch=batch_gen)
|
||||
training_batch=batch_stu)
|
||||
|
||||
with set_forward_context(
|
||||
current_timestep=batch_gen.timesteps,
|
||||
attn_metadata=batch_gen.attn_metadata_vsa):
|
||||
current_timestep=batch_stu.timesteps,
|
||||
attn_metadata=batch_stu.attn_metadata_vsa):
|
||||
(dmd_loss / gradient_accumulation_steps).backward()
|
||||
total_dmd_loss += dmd_loss.detach().item()
|
||||
self._clip_model_grad_norm_(batch_gen, self.transformer)
|
||||
self._clip_model_grad_norm_(batch_stu, self.transformer)
|
||||
self.optimizer.step()
|
||||
self.lr_scheduler.step()
|
||||
self.optimizer.zero_grad(set_to_none=True)
|
||||
@@ -482,22 +456,20 @@ class DistillationPipeline(TrainingPipeline):
|
||||
world_group.all_reduce(avg_dmd_loss,
|
||||
op=torch.distributed.ReduceOp.AVG)
|
||||
training_batch.generator_loss = avg_dmd_loss.item()
|
||||
dmd_latent_vis_dict = batch_gen.dmd_latent_vis_dict
|
||||
else:
|
||||
training_batch.generator_loss = 0.0
|
||||
|
||||
self.fake_score_optimizer.zero_grad()
|
||||
total_fake_score_loss = 0.0
|
||||
for batch in batches:
|
||||
batch_fake = copy.deepcopy(batch)
|
||||
batch_fake, fake_score_loss = self.faker_score_forward(batch_fake)
|
||||
with set_forward_context(current_timestep=batch_fake.timesteps,
|
||||
attn_metadata=batch_fake.attn_metadata):
|
||||
batch_critic = copy.deepcopy(batch)
|
||||
batch_critic, fake_score_loss = self.faker_score_forward(
|
||||
batch_critic)
|
||||
with set_forward_context(current_timestep=batch_critic.timesteps,
|
||||
attn_metadata=batch_critic.attn_metadata):
|
||||
(fake_score_loss / gradient_accumulation_steps).backward()
|
||||
total_fake_score_loss += fake_score_loss.detach().item()
|
||||
fake_score_latent_vis_dict.update(
|
||||
batch_fake.fake_score_latent_vis_dict)
|
||||
self._clip_model_grad_norm_(batch_fake, self.fake_score_transformer)
|
||||
self._clip_model_grad_norm_(batch_critic, self.fake_score_transformer)
|
||||
self.fake_score_optimizer.step()
|
||||
self.fake_score_lr_scheduler.step()
|
||||
self.fake_score_optimizer.zero_grad(set_to_none=True)
|
||||
@@ -508,8 +480,6 @@ class DistillationPipeline(TrainingPipeline):
|
||||
world_group.all_reduce(avg_fake_score_loss,
|
||||
op=torch.distributed.ReduceOp.AVG)
|
||||
training_batch.fake_score_loss = avg_fake_score_loss.item()
|
||||
training_batch.dmd_latent_vis_dict = dmd_latent_vis_dict
|
||||
training_batch.fake_score_latent_vis_dict = batch_fake.fake_score_latent_vis_dict
|
||||
|
||||
training_batch.total_loss = training_batch.generator_loss + training_batch.fake_score_loss
|
||||
return training_batch
|
||||
@@ -695,88 +665,6 @@ class DistillationPipeline(TrainingPipeline):
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def visualize_intermediate_latents(self, training_batch: TrainingBatch,
|
||||
training_args: TrainingArgs, step: int):
|
||||
"""Add visualization data to wandb logging and save frames to disk."""
|
||||
wandb_loss_dict = {}
|
||||
torch.cuda.empty_cache()
|
||||
dmd_latents_vis_dict = training_batch.dmd_latent_vis_dict
|
||||
fake_score_latents_vis_dict = training_batch.fake_score_latent_vis_dict
|
||||
fake_score_log_keys = [
|
||||
'training_batch_fakerscore_fwd_clean_latent', 'generator_pred_video'
|
||||
]
|
||||
dmd_log_keys = [
|
||||
'training_batch_dmd_fwd_clean_latent', 'generator_pred_video',
|
||||
'real_score_pred_video', 'faker_score_pred_video'
|
||||
]
|
||||
|
||||
for latent_key in fake_score_log_keys:
|
||||
latents = fake_score_latents_vis_dict[latent_key] # bs, t,c, h, w
|
||||
latents = latents.permute(0, 2, 1, 3, 4)
|
||||
# decoded_latent = decode_stage(ForwardBatch(data_type="video", latents=latents), training_args)
|
||||
|
||||
if isinstance(self.vae.scaling_factor, torch.Tensor):
|
||||
latents = latents / self.vae.scaling_factor.to(
|
||||
latents.device, latents.dtype)
|
||||
else:
|
||||
latents = latents / self.vae.scaling_factor
|
||||
|
||||
# Apply shifting if needed
|
||||
if (hasattr(self.vae, "shift_factor")
|
||||
and self.vae.shift_factor is not None):
|
||||
if isinstance(self.vae.shift_factor, torch.Tensor):
|
||||
latents += self.vae.shift_factor.to(latents.device,
|
||||
latents.dtype)
|
||||
else:
|
||||
latents += self.vae.shift_factor
|
||||
with torch.autocast("cuda", dtype=torch.bfloat16):
|
||||
video = self.vae.decode(latents)
|
||||
video = (video / 2 + 0.5).clamp(0, 1)
|
||||
video = video.cpu().float()
|
||||
video = video.permute(0, 2, 1, 3, 4)
|
||||
video = (video * 255).numpy().astype(np.uint8)
|
||||
wandb_loss_dict[latent_key] = wandb.Video(
|
||||
video, fps=24, format="mp4") # change to 16 for Wan2.1
|
||||
# Clean up references
|
||||
del video, latents
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Process DMD training data if available - use decode_stage instead of self.vae.decode
|
||||
if 'generator_pred_video' in dmd_latents_vis_dict:
|
||||
for latent_key in dmd_log_keys:
|
||||
latents = dmd_latents_vis_dict[latent_key]
|
||||
latents = latents.permute(0, 2, 1, 3, 4)
|
||||
# decoded_latent = decode_stage(ForwardBatch(data_type="video", latents=latents), training_args)
|
||||
if isinstance(self.vae.scaling_factor, torch.Tensor):
|
||||
latents = latents / self.vae.scaling_factor.to(
|
||||
latents.device, latents.dtype)
|
||||
else:
|
||||
latents = latents / self.vae.scaling_factor
|
||||
|
||||
# Apply shifting if needed
|
||||
if (hasattr(self.vae, "shift_factor")
|
||||
and self.vae.shift_factor is not None):
|
||||
if isinstance(self.vae.shift_factor, torch.Tensor):
|
||||
latents += self.vae.shift_factor.to(
|
||||
latents.device, latents.dtype)
|
||||
else:
|
||||
latents += self.vae.shift_factor
|
||||
with torch.autocast("cuda", dtype=torch.bfloat16):
|
||||
video = self.vae.decode(latents)
|
||||
video = (video / 2 + 0.5).clamp(0, 1)
|
||||
video = video.cpu().float()
|
||||
video = video.permute(0, 2, 1, 3, 4)
|
||||
video = (video * 255).numpy().astype(np.uint8)
|
||||
wandb_loss_dict[latent_key] = wandb.Video(
|
||||
video, fps=24, format="mp4") # change to 16 for Wan2.1
|
||||
# Clean up references
|
||||
del video, latents
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Log to wandb
|
||||
if self.global_rank == 0:
|
||||
wandb.log(wandb_loss_dict, step=step)
|
||||
|
||||
def train(self) -> None:
|
||||
"""Main training loop with distillation-specific logging."""
|
||||
assert self.training_args.seed is not None, "seed must be set"
|
||||
@@ -880,25 +768,6 @@ class DistillationPipeline(TrainingPipeline):
|
||||
log_data["train_generator_loss"] = generator_loss
|
||||
if use_vsa:
|
||||
log_data["VSA_train_sparsity"] = current_vsa_sparsity
|
||||
|
||||
if training_batch.dmd_latent_vis_dict:
|
||||
dmd_additional_logs = {
|
||||
"generator_timestep":
|
||||
training_batch.
|
||||
dmd_latent_vis_dict["generator_timestep"].item(),
|
||||
"dmd_timestep":
|
||||
training_batch.dmd_latent_vis_dict["dmd_timestep"].item(
|
||||
),
|
||||
}
|
||||
log_data.update(dmd_additional_logs)
|
||||
|
||||
faker_score_additional_logs = {
|
||||
"fake_score_timestep":
|
||||
training_batch.
|
||||
fake_score_latent_vis_dict["fake_score_timestep"].item(),
|
||||
}
|
||||
log_data.update(faker_score_additional_logs)
|
||||
|
||||
wandb.log(log_data, step=step)
|
||||
|
||||
# Save training state checkpoint (for resuming training)
|
||||
@@ -932,10 +801,6 @@ class DistillationPipeline(TrainingPipeline):
|
||||
only_save_generator_weight=True)
|
||||
|
||||
if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
|
||||
if self.training_args.log_visualization:
|
||||
self.visualize_intermediate_latents(training_batch,
|
||||
self.training_args,
|
||||
step)
|
||||
self._log_validation(self.transformer, self.training_args, step)
|
||||
|
||||
wandb.finish()
|
||||
|
||||
Reference in New Issue
Block a user