From 38f9c41e46d058d93667cf1ea94fdf38b5ea99ec Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Mon, 26 Jan 2026 20:40:11 +0000 Subject: [PATCH] refactor hy15 sf distill --- fastvideo/fastvideo_args.py | 5 + .../pipelines/stages/causal_denoising.py | 12 +- fastvideo/training/distillation_pipeline.py | 114 +- ...hy15_self_forcing_distillation_pipeline.py | 1048 ++++++++++++++++- .../self_forcing_distillation_pipeline.py | 687 ++--------- fastvideo/training/training_pipeline.py | 98 +- 6 files changed, 1260 insertions(+), 704 deletions(-) diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index fef1a2cd7..104b364b0 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -876,6 +876,7 @@ class TrainingArgs(FastVideoArgs): num_train_epochs: int = 0 max_train_steps: int = 0 gradient_accumulation_steps: int = 0 + optimizer_type: str = "adamw" learning_rate: float = 0.0 scale_lr: bool = False lr_scheduler: str = "constant" @@ -1150,6 +1151,10 @@ class TrainingArgs(FastVideoArgs): parser.add_argument("--gradient-accumulation-steps", type=int, help="Number of steps to accumulate gradients") + parser.add_argument("--optimizer-type", + type=str, + choices=["adamw", "muon"], + help="Optimizer type") parser.add_argument("--learning-rate", type=float, required=True, diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 777314547..1f5341506 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -45,12 +45,12 @@ class CausalDMDDenosingStage(DenoisingStage): self.transformer_2 = transformer_2 self.vae = vae # Model-dependent constants (aligned with causal_inference.py assumptions) - self.num_transformer_blocks = len(self.transformer.blocks) - self.num_frames_per_block = self.transformer.config.arch_config.num_frames_per_block - self.sliding_window_num_frames = self.transformer.config.arch_config.sliding_window_num_frames + self.num_transformer_blocks = self.transformer.config.num_layers + self.num_frames_per_block = self.transformer.config.num_frames_per_block + self.sliding_window_num_frames = self.transformer.config.sliding_window_num_frames try: - self.local_attn_size = getattr(self.transformer.model, + self.local_attn_size = getattr(self.transformer.config, "local_attn_size", -1) # type: ignore except Exception: @@ -412,8 +412,8 @@ class CausalDMDDenosingStage(DenoisingStage): Initialize a Per-GPU KV cache aligned with the Wan model assumptions. """ kv_cache1 = [] - num_attention_heads = self.transformer.num_attention_heads - attention_head_dim = self.transformer.attention_head_dim + num_attention_heads = self.transformer.config.num_attention_heads + attention_head_dim = self.transformer.config.attention_head_dim if self.local_attn_size != -1: kv_cache_size = self.local_attn_size * self.frame_seq_length else: diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 0694b6c1c..ba5358bd6 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -84,7 +84,6 @@ class DistillationPipeline(TrainingPipeline): if training_args.real_score_model_path: logger.info("Loading real score transformer from: %s", training_args.real_score_model_path) - # training_args.override_transformer_cls_name = "WanTransformer3DModel" # TODO(will): can use deepcopy instead if the model is the same self.real_score_transformer = self.load_module_from_path( training_args.real_score_model_path, "transformer", @@ -110,7 +109,6 @@ class DistillationPipeline(TrainingPipeline): if training_args.fake_score_model_path: logger.info("Loading fake score transformer from: %s", training_args.fake_score_model_path) - # training_args.override_transformer_cls_name = "WanTransformer3DModel" self.fake_score_transformer = self.load_module_from_path( training_args.fake_score_model_path, "transformer", training_args) @@ -145,10 +143,8 @@ class DistillationPipeline(TrainingPipeline): self.vae.requires_grad_(False) self.timestep_shift = self.training_args.pipeline_config.flow_shift - # self.noise_scheduler.set_timesteps(sigmas=list(np.linspace(1.0, 0.0, 50 + 1)[:-1]), device=get_local_torch_device()) - self.noise_scheduler.set_timesteps(num_inference_steps=1000, - extra_one_step=True, - device=get_local_torch_device()) + # self.noise_scheduler = FlowMatchEulerDiscreteScheduler( + # shift=self.timestep_shift) if self.training_args.boundary_ratio is not None: self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps @@ -183,19 +179,34 @@ class DistillationPipeline(TrainingPipeline): checkpointing_type=training_args. enable_gradient_checkpointing_type) + # Initialize optimizers + fake_score_params = list( + filter(lambda p: p.requires_grad, + self.fake_score_transformer.parameters())) + # Use separate learning rate for fake_score_transformer if specified fake_score_lr = training_args.fake_score_learning_rate if fake_score_lr == 0.0: fake_score_lr = training_args.learning_rate - betas_str = training_args.fake_score_betas - betas = tuple(float(x.strip()) for x in betas_str.split(",")) - - self.fake_score_optimizer = get_muon_optimizer( - self.fake_score_transformer, - lr=fake_score_lr, - weight_decay=training_args.weight_decay, - ) + if training_args.optimizer_type == "adamw": + betas_str = training_args.fake_score_betas + betas = tuple(float(x.strip()) for x in betas_str.split(",")) + self.fake_score_optimizer = torch.optim.AdamW( + fake_score_params, + lr=fake_score_lr, + betas=betas, + weight_decay=training_args.weight_decay, + eps=1e-8, + ) + elif training_args.optimizer_type == "muon": + self.fake_score_optimizer = get_muon_optimizer( + self.fake_score_transformer, + lr=fake_score_lr, + weight_decay=training_args.weight_decay, + ) + else: + raise ValueError(f"Unsupported optimizer type: {training_args.optimizer_type}") self.fake_score_lr_scheduler = get_scheduler( training_args.fake_score_lr_scheduler, @@ -212,13 +223,23 @@ class DistillationPipeline(TrainingPipeline): fake_score_params_2 = list( filter(lambda p: p.requires_grad, self.fake_score_transformer_2.parameters())) - self.fake_score_optimizer_2 = torch.optim.AdamW( - fake_score_params_2, - lr=fake_score_lr, - betas=betas, - weight_decay=training_args.weight_decay, - eps=1e-8, - ) + if training_args.optimizer_type == "adamw": + self.fake_score_optimizer_2 = torch.optim.AdamW( + fake_score_params_2, + lr=fake_score_lr, + betas=betas, + weight_decay=training_args.weight_decay, + eps=1e-8, + ) + elif training_args.optimizer_type == "muon": + self.fake_score_optimizer_2 = get_muon_optimizer( + self.fake_score_transformer_2, + lr=fake_score_lr, + weight_decay=training_args.weight_decay, + ) + else: + raise ValueError(f"Unsupported optimizer type: {training_args.optimizer_type}") + self.fake_score_lr_scheduler_2 = get_scheduler( training_args.fake_score_lr_scheduler, optimizer=self.fake_score_optimizer_2, @@ -263,25 +284,9 @@ class DistillationPipeline(TrainingPipeline): self.num_train_timestep) self.real_score_guidance_scale = self.training_args.real_score_guidance_scale - # assert self.real_score_guidance_scale == 3.5, "Real score guidance scale must be 3.5" self.generator_ema: EMA_FSDP | None = None self.generator_ema_2: EMA_FSDP | None = None - # if (self.training_args.ema_decay - # is not None) and (self.training_args.ema_decay > 0.0): - # self.generator_ema = EMA_FSDP(self.transformer, - # decay=self.training_args.ema_decay) - # logger.info("Initialized generator EMA with decay=%s", - # self.training_args.ema_decay) - - # # Initialize EMA for transformer_2 if it exists - # if self.transformer_2 is not None: - # self.generator_ema_2 = EMA_FSDP( - # self.transformer_2, decay=self.training_args.ema_decay) - # logger.info("Initialized generator EMA_2 with decay=%s", - # self.training_args.ema_decay) - # else: - # logger.info("Generator EMA disabled (ema_decay <= 0.0)") def load_module_from_path(self, model_path: str, module_type: str, training_args: "TrainingArgs"): @@ -708,11 +713,6 @@ class DistillationPipeline(TrainingPipeline): timestep, self.timestep_shift, # type: ignore self.num_train_timestep) - # timestep_idx = torch.randint(1, - # len(self.denoising_step_list), [1], - # device=self.device, - # dtype=torch.long) - # timestep = self.denoising_step_list[timestep_idx] timestep = timestep.clamp(self.min_timestep, self.max_timestep) @@ -783,7 +783,6 @@ class DistillationPipeline(TrainingPipeline): real_score_pred_video = pred_real_video_cond + ( pred_real_video_cond - pred_real_video_uncond) * self.real_score_guidance_scale - # real_score_pred_video = pred_real_video_cond grad = (faker_score_pred_video - real_score_pred_video) / torch.abs( original_latent - real_score_pred_video).mean() @@ -836,11 +835,6 @@ class DistillationPipeline(TrainingPipeline): fake_score_timestep, self.timestep_shift, # type: ignore self.num_train_timestep) - # fake_score_timestep_idx = torch.randint(1, - # len(self.denoising_step_list), [1], - # device=self.device, - # dtype=torch.long) - # fake_score_timestep = self.denoising_step_list[fake_score_timestep_idx] fake_score_timestep = fake_score_timestep.clamp(self.min_timestep, self.max_timestep) @@ -921,26 +915,13 @@ class DistillationPipeline(TrainingPipeline): training_batch.dmd_latent_vis_dict = {} training_batch.fake_score_latent_vis_dict = {} - training_batch.timesteps = 0 training_batch.conditional_dict = conditional_dict training_batch.raw_latent_shape = training_batch.latents.shape training_batch.latents = training_batch.latents.permute(0, 2, 1, 3, 4) self.video_latent_shape = training_batch.latents.shape self.video_latent_shape_sp = training_batch.latents.shape - training_batch.image_embeds = [ - torch.zeros(1, - 729, - 1152, - device=get_local_torch_device(), - dtype=torch.bfloat16) - ] - raw_latent_shape = list(self.video_latent_shape) - raw_latent_shape[2] = 1 - training_batch.video_latent = torch.zeros( - tuple(raw_latent_shape), - device=get_local_torch_device(), - dtype=torch.bfloat16) + return training_batch def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: @@ -1284,7 +1265,6 @@ class DistillationPipeline(TrainingPipeline): batch.prompt_attention_mask = [] negative_prompt = batch.negative_prompt - assert negative_prompt == "" batch_negative = ForwardBatch( data_type="video", prompt=negative_prompt, @@ -1293,9 +1273,11 @@ class DistillationPipeline(TrainingPipeline): ) result_batch = self.validation_pipeline.prompt_encoding_stage( # type: ignore batch_negative, training_args) - # self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[ - # 0], result_batch.prompt_attention_mask[0] - self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds, result_batch.prompt_attention_mask + if len(result_batch.prompt_embeds) == 1: + self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[ + 0], result_batch.prompt_attention_mask[0] + else: + self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds, result_batch.prompt_attention_mask logger.info( "rank: %s: rank_in_sp_group: %s, batch.prompt: %s", diff --git a/fastvideo/training/hy15_self_forcing_distillation_pipeline.py b/fastvideo/training/hy15_self_forcing_distillation_pipeline.py index b486005c8..877836025 100644 --- a/fastvideo/training/hy15_self_forcing_distillation_pipeline.py +++ b/fastvideo/training/hy15_self_forcing_distillation_pipeline.py @@ -1,19 +1,36 @@ # SPDX-License-Identifier: Apache-2.0 +import copy +import os +import gc import sys +from typing import Any from copy import deepcopy -from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs +import torch +import torch.distributed as dist +import torch.nn.functional as F +from einops import rearrange + +import fastvideo.envs as envs +from fastvideo.distributed import (get_local_torch_device, + get_world_group) +from fastvideo.fastvideo_args import TrainingArgs, FastVideoArgs +from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger + +from fastvideo.models.utils import pred_noise_to_pred_video +from fastvideo.pipelines import TrainingBatch +from fastvideo.training.self_forcing_distillation_pipeline import SelfForcingDistillationPipeline +from fastvideo.training.distillation_pipeline import DistillationPipeline from fastvideo.pipelines.basic.hunyuan15.hunyuan15_causal_dmd_pipeline import ( Hy15CausalDMDPipeline) -from fastvideo.training.self_forcing_distillation_pipeline import ( - SelfForcingDistillationPipeline) +from fastvideo.training.training_utils import (EMA_FSDP) from fastvideo.utils import is_vsa_available -vsa_available = is_vsa_available() - logger = init_logger(__name__) +vsa_available = is_vsa_available() + class Hy15SelfForcingDistillationPipeline(SelfForcingDistillationPipeline): """ @@ -31,6 +48,38 @@ class Hy15SelfForcingDistillationPipeline(SelfForcingDistillationPipeline): "vae", ] + def initialize_training_pipeline(self, training_args: TrainingArgs): + """Initialize the self-forcing training pipeline.""" + # Check if FSDP2 auto wrap is enabled - not supported for self-forcing distillation + if os.environ.get("FASTVIDEO_FSDP2_AUTOWRAP", "0") == "1": + raise NotImplementedError( + "FASTVIDEO_FSDP2_AUTOWRAP is not implemented for self-forcing distillation. " + "Please set FASTVIDEO_FSDP2_AUTOWRAP=0 or unset the environment variable." + ) + + logger.info("Initializing self-forcing distillation pipeline...") + + DistillationPipeline.initialize_training_pipeline(self, training_args) + self.noise_scheduler.set_timesteps(num_inference_steps=1000, + extra_one_step=True, + device=get_local_torch_device()) + self.generator_ema: EMA_FSDP | None = None + + self.dfake_gen_update_ratio = getattr(training_args, + 'dfake_gen_update_ratio', 5) + + self.num_frame_per_block = getattr(training_args, 'num_frame_per_block', + 3) + self.independent_first_frame = getattr(training_args, + 'independent_first_frame', False) + self.same_step_across_blocks = getattr(training_args, + 'same_step_across_blocks', False) + self.last_step_only = getattr(training_args, 'last_step_only', False) + self.context_noise = getattr(training_args, 'context_noise', 0) + + logger.info("Self-forcing generator update ratio: %s", + self.dfake_gen_update_ratio) + def create_training_stages(self, training_args: TrainingArgs): """ May be used in future refactors. @@ -57,6 +106,995 @@ class Hy15SelfForcingDistillationPipeline(SelfForcingDistillationPipeline): self.validation_pipeline = validation_pipeline + def _prepare_dit_inputs(self, training_batch: TrainingBatch) -> TrainingBatch: + training_batch = super()._prepare_dit_inputs(training_batch) + training_batch.timesteps = 0 + training_batch.image_embeds = [ + torch.zeros(1, + 729, + 1152, + device=get_local_torch_device(), + dtype=torch.bfloat16) + ] + raw_latent_shape = list(self.video_latent_shape) + raw_latent_shape[2] = 1 + training_batch.video_latent = torch.zeros( + tuple(raw_latent_shape), + device=get_local_torch_device(), + dtype=torch.bfloat16) + return training_batch + + def _prepare_gt_inputs(self, + training_batch: TrainingBatch) -> TrainingBatch: + trajectory_latents = training_batch.trajectory_latents + + start_timestep_index = torch.randint(1, + len(self.denoising_step_list), + (1, ), + device=self.device) + if dist.is_initialized(): + dist.broadcast(start_timestep_index, src=0) + start_timestep_index = start_timestep_index.cpu().item() + + trajectory_latents = trajectory_latents[:, -1, :21] + training_batch.trajectory_latents = self.noise_scheduler.add_noise( + trajectory_latents.flatten(0, 1), + torch.randn_like(trajectory_latents.flatten(0, 1)), + self.denoising_step_list[start_timestep_index] * torch.ones( + [trajectory_latents.shape[0] * trajectory_latents.shape[1]], + device=self.device, + dtype=torch.long)).unflatten( + 0, + (trajectory_latents.shape[0], trajectory_latents.shape[1])) + del trajectory_latents + training_batch.start_timestep_index = start_timestep_index + return training_batch + + def _prepare_ode_init_inputs( + self, training_batch: TrainingBatch) -> TrainingBatch: + trajectory_latents = training_batch.trajectory_latents + + trajectory_latents = trajectory_latents[:, :, :21] + _cached_closest_idx_per_dmd = torch.tensor([0, 12, 24, 36, 50], + dtype=torch.long).cpu() + training_batch.trajectory_latents = torch.index_select( + trajectory_latents, + dim=1, + index=_cached_closest_idx_per_dmd.to(trajectory_latents.device)) + del trajectory_latents + return training_batch + + def _prepare_context_forcing_inputs( + self, training_batch: TrainingBatch) -> TrainingBatch: + trajectory_latents = training_batch.trajectory_latents + + total_num_chunks = self.training_args.num_latent_t // self.num_frame_per_block + # Gradually increase maximum number of chunks + # Increase max_num_chunks by 1 every 100 training steps until it reaches total_num_chunks + max_num_chunks = min(total_num_chunks, + self.current_trainstep // 100 + 2) + context_forcing_length = torch.randint( + 0, max_num_chunks, + (1, ), device=get_local_torch_device()) * self.num_frame_per_block + # broadcast context_forcing_length to all ranks + if dist.is_initialized(): + dist.broadcast(context_forcing_length, src=0) + context_forcing_length = context_forcing_length.cpu().item() + assert context_forcing_length < self.training_args.num_latent_t, "Context forcing length is greater than the number of latent frames" + if context_forcing_length > 0: + training_batch.trajectory_latents = trajectory_latents[:, -1, : + context_forcing_length] + else: + training_batch.trajectory_latents = None + del trajectory_latents + return training_batch + + def _generator_multi_step_simulation_forward( + self, + training_batch: TrainingBatch, + return_sim_steps: bool = False) -> torch.Tensor: + """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 + and includes gradient masking logic for dynamic frame generation. + """ + latents = training_batch.latents + dtype = latents.dtype + batch_size = latents.shape[0] + initial_latent = getattr(training_batch, 'image_latent', None) + + # Dynamic frame generation logic (adapted from _run_generator) + num_training_frames = getattr(self.training_args, 'num_latent_t', 21) + + min_num_frames = 20 if self.independent_first_frame else 21 + + num_generated_frames = num_training_frames + if self.independent_first_frame and initial_latent is None: + num_generated_frames += 1 + min_num_frames += 1 + + # Create noise with dynamic shape + if initial_latent is not None: + noise_shape = [ + batch_size, num_generated_frames - 1, + *self.video_latent_shape[2:] + ] + else: + noise_shape = [ + batch_size, num_generated_frames, *self.video_latent_shape[2:] + ] + + if training_batch.use_gt_trajectory and training_batch.trajectory_latents is not None: + noise = training_batch.trajectory_latents.to(self.device, + dtype=dtype) + else: + noise = torch.randn(noise_shape, device=self.device, dtype=dtype) + if self.sp_world_size > 1: + noise = rearrange(noise, + "b (n t) c h w -> b n t c h w", + n=self.sp_world_size).contiguous() + noise = noise[:, self.rank_in_sp_group, :, :, :, :] + + batch_size, num_frames, num_channels, height, width = noise.shape + + import math + num_blocks = math.ceil(num_frames / self.num_frame_per_block) + + num_input_frames = initial_latent.shape[ + 1] if initial_latent is not None else 0 + num_output_frames = num_frames + num_input_frames + output = torch.zeros( + [batch_size, num_output_frames, num_channels, height, width], + device=noise.device, + dtype=noise.dtype) + + # Step 1: Initialize KV cache to all zeros + cache_frames = num_generated_frames + num_input_frames + kv_cache1, crossattn_cache = self._initialize_simulation_caches( + batch_size, dtype, self.device, max_num_frames=cache_frames) + + current_start_frame = 0 + + # Step 2: Temporal denoising loop + all_num_frames = [self.num_frame_per_block] * num_blocks + # all_num_frames[0] = 1 + if self.independent_first_frame and initial_latent is None: + all_num_frames = [1] + all_num_frames + + if training_batch.use_gt_trajectory and training_batch.trajectory_latents is not None: + start_timestep_index = training_batch.start_timestep_index + end_timestep_index = len(self.denoising_step_list) + else: + start_timestep_index = 0 + end_timestep_index = len(self.denoising_step_list) + exit_flags = self.generate_and_sync_list(len(all_num_frames), + start_timestep_index, + end_timestep_index, + device=noise.device) + start_gradient_frame_index = max(0, num_output_frames - 31) + + logger.info("distillation step list: %s", self.denoising_step_list) + + # Initialize txt kv cache + txt_kv_cache = self.transformer( + txt_inference=True, + vision_inference=False, + encoder_hidden_states=training_batch. + conditional_dict["encoder_hidden_states"], + encoder_hidden_states_image=training_batch.image_embeds, + encoder_attention_mask=training_batch. + conditional_dict["encoder_attention_mask"], + timestep=torch.zeros([batch_size], device=noise.device), + cache_txt=True, + ) + + # Context forcing logic + if self.training_args.use_context_forcing and training_batch.trajectory_latents is not None: + gt_latents = training_batch.trajectory_latents + context_forcing_length = gt_latents.shape[1] + output[:, :context_forcing_length] = gt_latents.detach() + num_residual_blocks = ( + self.training_args.num_latent_t - + context_forcing_length) // self.num_frame_per_block + assert num_residual_blocks > 0, "Number of residual blocks is less than 0" + all_num_frames = [self.num_frame_per_block] * num_residual_blocks + + if training_batch.video_latent is not None: + gt_latents = torch.cat([ + gt_latents, + training_batch. + video_latent[:, current_start_frame:current_start_frame + + context_forcing_length], + torch.zeros_like(gt_latents), + ], + dim=2) + + with torch.no_grad(): + training_batch_temp = self._build_distill_input_kwargs( + gt_latents, + torch.zeros([batch_size, context_forcing_length], + device=noise.device, + dtype=torch.int64), + training_batch.conditional_dict, training_batch) + + self.transformer( + txt_inference=False, + vision_inference=True, + hidden_states=training_batch_temp. + input_kwargs['hidden_states'], + timestep=training_batch_temp.input_kwargs['timestep'], + kv_cache=kv_cache1, + txt_kv_cache=txt_kv_cache, + current_start=current_start_frame * self.frame_seq_length, + rope_start_idx=current_start_frame) + + current_start_frame += context_forcing_length + + for block_index, current_num_frames in enumerate(all_num_frames): + noisy_input = noise[:, current_start_frame - + num_input_frames:current_start_frame + + current_num_frames - num_input_frames] + + index = start_timestep_index + current_timestep = self.denoising_step_list[index] + + assert index < len( + self.denoising_step_list + ), "Index is greater than the number of denoising steps" + # Step 3.1: Spatial denoising loop + while index < len(self.denoising_step_list): + noisy_input_latent = noisy_input.clone() + if training_batch.video_latent is not None: + noisy_input_latent = torch.cat([ + noisy_input_latent, + training_batch. + video_latent[:, + current_start_frame:current_start_frame + + current_num_frames], + torch.zeros_like(noisy_input_latent), + ], + dim=2) + + if self.same_step_across_blocks: + exit_flag = (index == exit_flags[0]) + else: + exit_flag = (index == exit_flags[block_index]) + + timestep = torch.ones([batch_size, current_num_frames], + device=noise.device, + dtype=torch.int64) * current_timestep + + if not exit_flag: + with torch.no_grad(): + # Build input kwargs + training_batch_temp = self._build_distill_input_kwargs( + noisy_input_latent, timestep, + training_batch.conditional_dict, training_batch) + + pred_flow, kv_cache1 = self.transformer( + txt_inference=False, + vision_inference=True, + hidden_states=training_batch_temp. + input_kwargs['hidden_states'], + timestep=training_batch_temp. + input_kwargs['timestep'], + kv_cache=kv_cache1, + txt_kv_cache=txt_kv_cache, + current_start=current_start_frame * + self.frame_seq_length, + rope_start_idx=current_start_frame) + pred_flow = pred_flow.permute(0, 2, 1, 3, 4) + + denoised_pred = pred_noise_to_pred_video( + pred_noise=pred_flow.flatten(0, 1), + noise_input_latent=noisy_input.flatten(0, 1), + timestep=timestep, + scheduler=self.noise_scheduler).unflatten( + 0, pred_flow.shape[:2]) + + next_timestep = self.denoising_step_list[index + 1] + noisy_input = self.noise_scheduler.add_noise( + denoised_pred.flatten(0, 1), + torch.randn_like(denoised_pred.flatten(0, 1)), + next_timestep * + torch.ones([batch_size * current_num_frames], + device=noise.device, + dtype=torch.long)).unflatten( + 0, denoised_pred.shape[:2]) + else: + # Final prediction with gradient control + if current_start_frame < start_gradient_frame_index: + with torch.no_grad(): + training_batch_temp = self._build_distill_input_kwargs( + noisy_input_latent, timestep, + training_batch.conditional_dict, training_batch) + + pred_flow, kv_cache1 = self.transformer( + txt_inference=False, + vision_inference=True, + hidden_states=training_batch_temp. + input_kwargs['hidden_states'], + timestep=training_batch_temp. + input_kwargs['timestep'], + kv_cache=kv_cache1, + txt_kv_cache=txt_kv_cache, + current_start=current_start_frame * + self.frame_seq_length, + rope_start_idx=current_start_frame) + pred_flow = pred_flow.permute(0, 2, 1, 3, 4) + else: + training_batch_temp = self._build_distill_input_kwargs( + noisy_input_latent, timestep, + training_batch.conditional_dict, training_batch) + + pred_flow, kv_cache1 = self.transformer( + txt_inference=False, + vision_inference=True, + hidden_states=training_batch_temp. + input_kwargs['hidden_states'], + timestep=training_batch_temp. + input_kwargs['timestep'], + kv_cache=kv_cache1, + txt_kv_cache=txt_kv_cache, + current_start=current_start_frame * + self.frame_seq_length, + rope_start_idx=current_start_frame) + pred_flow = pred_flow.permute(0, 2, 1, 3, 4) + + denoised_pred = pred_noise_to_pred_video( + pred_noise=pred_flow.flatten(0, 1), + noise_input_latent=noisy_input.flatten(0, 1), + timestep=timestep, + scheduler=self.noise_scheduler).unflatten( + 0, pred_flow.shape[:2]) + break + + index += 1 + current_timestep = self.denoising_step_list[index] + + # Step 3.2: record the model's output + output[:, current_start_frame:current_start_frame + + current_num_frames] = denoised_pred + + # Step 3.3: rerun with timestep zero to update the cache + context_timestep = torch.ones_like(timestep) * self.context_noise + denoised_pred = self.noise_scheduler.add_noise( + denoised_pred.flatten(0, 1), + torch.randn_like(denoised_pred.flatten(0, 1)), + context_timestep).unflatten(0, denoised_pred.shape[:2]) + + with torch.no_grad(): + if training_batch.video_latent is not None: + denoised_pred = torch.cat([ + denoised_pred, + training_batch. + video_latent[:, + current_start_frame:current_start_frame + + current_num_frames], + torch.zeros_like(denoised_pred), + ], + dim=2) + + training_batch_temp = self._build_distill_input_kwargs( + denoised_pred, context_timestep, + training_batch.conditional_dict, training_batch) + + _, kv_cache1 = self.transformer( + txt_inference=False, + vision_inference=True, + hidden_states=training_batch_temp. + input_kwargs['hidden_states'], + timestep=training_batch_temp.input_kwargs['timestep'], + kv_cache=kv_cache1, + txt_kv_cache=txt_kv_cache, + current_start=current_start_frame * self.frame_seq_length, + rope_start_idx=current_start_frame) + + # Step 3.4: update the start and end frame indices + current_start_frame += current_num_frames + + # Handle last 21 frames logic + pred_image_or_video = output + if num_input_frames > 0: + pred_image_or_video = output[:, num_input_frames:] + + # Slice last 21 frames if we generated more + gradient_mask = None + if pred_image_or_video.shape[1] > 31: + with torch.no_grad(): + # Re-encode to get image latent + latent_to_decode = pred_image_or_video[:, :-20, ...] + # Decode to video + latent_to_decode = latent_to_decode.permute( + 0, 2, 1, 3, 4) # [B, C, F, H, W] + + # Apply VAE scaling and shift factors + if isinstance(self.vae.scaling_factor, torch.Tensor): + latent_to_decode = latent_to_decode / self.vae.scaling_factor.to( + latent_to_decode.device, latent_to_decode.dtype) + else: + latent_to_decode = latent_to_decode / self.vae.scaling_factor + + if hasattr( + self.vae, + "shift_factor") and self.vae.shift_factor is not None: + if isinstance(self.vae.shift_factor, torch.Tensor): + latent_to_decode += self.vae.shift_factor.to( + latent_to_decode.device, latent_to_decode.dtype) + else: + latent_to_decode += self.vae.shift_factor + + # Decode to pixels + pixels = self.vae.decode(latent_to_decode) + frame = pixels[:, :, -1:, :, :].to( + dtype) # Last frame [B, C, 1, H, W] + + # Encode frame back to get image latent + image_latent = self.vae.encode(frame).to(dtype) + image_latent = image_latent.permute(0, 2, 1, 3, + 4) # [B, F, C, H, W] + + pred_image_or_video_last_21 = torch.cat( + [image_latent, pred_image_or_video[:, -20:, ...]], dim=1) + else: + pred_image_or_video_last_21 = pred_image_or_video + + # Apply gradient masking if needed + final_output = pred_image_or_video_last_21.to(dtype) + if gradient_mask is not None: + # Apply gradient masking: detach frames that shouldn't contribute gradients + final_output = torch.where( + gradient_mask, + pred_image_or_video_last_21, # Keep original values where gradient_mask is True + pred_image_or_video_last_21.detach( + ) # Detach where gradient_mask is False + ) + + # Store visualization data + training_batch.dmd_latent_vis_dict["generator_timestep"] = torch.tensor( + self.denoising_step_list[exit_flags[0]], + dtype=torch.float32, + device=self.device) + + # Store gradient mask information for debugging + if gradient_mask is not None: + training_batch.dmd_latent_vis_dict[ + "gradient_mask"] = gradient_mask.float() + training_batch.dmd_latent_vis_dict[ + "num_generated_frames"] = torch.tensor(num_generated_frames, + dtype=torch.float32, + device=self.device) + training_batch.dmd_latent_vis_dict["min_num_frames"] = torch.tensor( + min_num_frames, dtype=torch.float32, device=self.device) + + # Clean up caches + assert kv_cache1 is not None + self._reset_simulation_caches(kv_cache1, crossattn_cache) + + assert pred_image_or_video.shape[ + 1] == num_training_frames, "pred_image_or_video must have %d frames" % num_training_frames + + return final_output if gradient_mask is not None else pred_image_or_video + + def _initialize_simulation_caches( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + *, + max_num_frames: int | None = None, + ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """Initialize KV cache and cross-attention cache for multi-step simulation.""" + num_transformer_blocks = self.transformer.config.num_layers + latent_shape = self.video_latent_shape_sp + _, num_frames, _, height, width = latent_shape + + if isinstance(self.transformer.config.patch_size, tuple): + ph, pw = self.transformer.config.patch_size[ + 1], self.transformer.config.patch_size[2] + elif isinstance(self.transformer.config.patch_size, int): + ph, pw = self.transformer.config.patch_size, self.transformer.config.patch_size + else: + raise ValueError( + f"Unsupported patch size type: {type(self.transformer.config.patch_size)}" + ) + post_patch_height = height // ph + post_patch_width = width // pw + + frame_seq_length = post_patch_height * post_patch_width + self.frame_seq_length = frame_seq_length + + # Get model configuration parameters - handle FSDP wrapping + num_attention_heads = self.transformer.config.num_attention_heads + attention_head_dim = self.transformer.config.attention_head_dim + + if max_num_frames is None: + max_num_frames = num_frames + num_max_frames = max(max_num_frames, num_frames) + local_attn_size = getattr(self.transformer.config, 'local_attn_size', + -1) + kv_cache_size = num_max_frames * frame_seq_length if local_attn_size == -1 else local_attn_size * frame_seq_length + + kv_cache = [] + for _ in range(num_transformer_blocks): + kv_cache.append({ + "k": + torch.zeros([ + batch_size, kv_cache_size, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "v": + torch.zeros([ + batch_size, kv_cache_size, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "global_end_index": + torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": + torch.tensor([0], dtype=torch.long, device=device) + }) + + # Initialize cross-attention cache + crossattn_cache = None + + return kv_cache, crossattn_cache + + def _step_predict_next_latent( + self, training_batch: TrainingBatch + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, + torch.Tensor]]: + latent_vis_dict: dict[str, torch.Tensor] = {} + device = get_local_torch_device() + target_latent = training_batch.trajectory_latents[:, -1] + + # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S] + B, S, num_frames, num_channels, height, width = training_batch.trajectory_latents.shape + + dmd_denoising_steps = torch.cat([ + self.denoising_step_list, + torch.tensor([0], device=self.device, dtype=torch.long) + ]) + indexes = self._get_timestep( # [B, num_frames] + 0, + len(dmd_denoising_steps), + B, + num_frames, + 3, + uniform_timestep=False) + # noisy_input = relevant_traj_latents[indexes] + latents = torch.gather( + training_batch.trajectory_latents, + dim=1, + index=indexes.reshape(B, 1, num_frames, 1, 1, + 1).expand(-1, -1, -1, num_channels, height, + width).to(self.device)).squeeze(1) + noisy_input = torch.cat([ + latents, + torch.zeros_like(latents), + torch.zeros_like(latents[:, :, 0:1]) + ], + dim=2) + timestep = dmd_denoising_steps[indexes] + + logger.info("timestep: %s", timestep) + txt_input_kwargs = { + "txt_inference": + True, + "vision_inference": + False, + "encoder_hidden_states": + training_batch.encoder_hidden_states, + "encoder_hidden_states_image": + training_batch.image_embeds, + "encoder_attention_mask": + training_batch.encoder_attention_mask, + "timestep": + torch.zeros([latents.shape[0]], + device=latents.device, + dtype=torch.bfloat16), + "cache_txt": + True, + } + with set_forward_context(current_timestep=timestep, + attn_metadata=None, + forward_batch=None): + txt_kv_cache = self.transformer(**txt_input_kwargs) + + vision_input_kwargs = { + "txt_inference": False, + "vision_inference": True, + "hidden_states": noisy_input.permute(0, 2, 1, 3, 4), + "timestep": timestep.to(device, dtype=torch.bfloat16), + "txt_kv_cache": txt_kv_cache, + } + # Predict noise and step the scheduler to obtain next latent + with set_forward_context(current_timestep=timestep, + attn_metadata=None, + forward_batch=None): + noise_pred = self.transformer(**vision_input_kwargs).permute( + 0, 2, 1, 3, 4) + + from fastvideo.models.utils import pred_noise_to_pred_video + pred_video = pred_noise_to_pred_video( + pred_noise=noise_pred.flatten(0, 1), + noise_input_latent=latents.flatten(0, 1), + timestep=timestep.to(dtype=torch.bfloat16).flatten(0, 1), + scheduler=self.noise_scheduler).unflatten(0, noise_pred.shape[:2]) + latent_vis_dict["pred_video"] = pred_video.permute( + 0, 2, 1, 3, 4).detach().clone().cpu() + + return pred_video, target_latent, timestep + + def _get_timestep(self, + min_timestep: int, + max_timestep: int, + batch_size: int, + num_frame: int, + num_frame_per_block: int, + uniform_timestep: bool = False) -> torch.Tensor: + if uniform_timestep: + timestep = torch.randint(min_timestep, + max_timestep, [batch_size, 1], + device=self.device, + dtype=torch.long).repeat(1, num_frame) + return timestep + else: + num_pad_frames = 0 + if num_frame % num_frame_per_block != 0: + # Pad num_frame to be divisible by num_frame_per_block + num_pad_frames = num_frame_per_block - (num_frame % + num_frame_per_block) + num_frame += num_pad_frames + timestep = torch.randint(min_timestep, + max_timestep, [batch_size, num_frame], + device=self.device, + dtype=torch.long) + # logger.info(f"individual timestep: {timestep}") + # make the noise level the same within every block + timestep = timestep.reshape(timestep.shape[0], -1, + num_frame_per_block) + timestep[:, :, 1:] = timestep[:, :, 0:1] + timestep = timestep.reshape(timestep.shape[0], -1) + if num_pad_frames > 0: + timestep = timestep[:, num_pad_frames:] + return timestep + + def _get_next_batch_2(self, training_batch: TrainingBatch) -> TrainingBatch: + batch = next(self.train_loader_iter_2, None) # type: ignore + if batch is None: + self.current_epoch += 1 + self.train_loader_iter_2 = iter(self.train_dataloader_2) + batch = next(self.train_loader_iter_2) + + # Required fields from parquet (ODE trajectory schema) + device = get_local_torch_device() + encoder_hidden_states = batch['text_embedding'].to( + device, dtype=torch.bfloat16).squeeze(0) + encoder_hidden_states_2 = batch['text_embedding_2'].to( + device, dtype=torch.bfloat16).squeeze(0) + encoder_attention_mask = batch['text_mask'].to( + device, dtype=torch.bfloat16).squeeze(0) + encoder_attention_mask_2 = batch['text_mask_2'].to( + device, dtype=torch.bfloat16).squeeze(0) + infos = batch['info_list'] + + if encoder_hidden_states.dim() < 3: + prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text( + infos[0]["caption"], + self.training_args, + encoder_index=[0], + return_attention_mask=True, + ) + encoder_hidden_states = prompt_embeds_list[0].to( + device, dtype=torch.bfloat16) + encoder_attention_mask = prompt_masks_list[0].to( + device, dtype=torch.bfloat16) + + # Trajectory tensors may include a leading singleton batch dim per row + trajectory_latents = batch['trajectory_latents'] + if trajectory_latents.dim() == 7: + # [B, 1, S, C, T, H, W] -> [B, S, C, T, H, W] + trajectory_latents = trajectory_latents[:, 0] + elif trajectory_latents.dim() == 6: + # already [B, S, C, T, H, W] + pass + else: + raise ValueError( + f"Unexpected trajectory_latents dim: {trajectory_latents.dim()}" + ) + + trajectory_timesteps = batch['trajectory_timesteps'] + if trajectory_timesteps.dim() == 3: + # [B, 1, S] -> [B, S] + trajectory_timesteps = trajectory_timesteps[:, 0] + elif trajectory_timesteps.dim() == 2: + # [B, S] + pass + else: + raise ValueError( + f"Unexpected trajectory_timesteps dim: {trajectory_timesteps.dim()}" + ) + # [B, S, C, T, H, W] -> [B, S, T, C, H, W] to match self-forcing + trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5) + + batch_size = trajectory_latents.shape[0] + vae_config = self.training_args.pipeline_config.vae_config.arch_config + num_channels = vae_config.latent_channels + spatial_compression_ratio = vae_config.spatial_compression_ratio + + latent_height = self.training_args.num_height // spatial_compression_ratio + latent_width = self.training_args.num_width // spatial_compression_ratio + training_batch.latents = torch.randn( + batch_size, num_channels, self.training_args.num_latent_t, + latent_height, latent_width).to(get_local_torch_device(), + dtype=torch.bfloat16) + training_batch.trajectory_latents = trajectory_latents.to( + device, dtype=torch.bfloat16) + training_batch.trajectory_timesteps = trajectory_timesteps.to(device) + training_batch.encoder_hidden_states = [ + encoder_hidden_states, encoder_hidden_states_2 + ] + training_batch.encoder_attention_mask = [ + encoder_attention_mask, encoder_attention_mask_2 + ] + training_batch.infos = infos + return training_batch + + def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: + batch = next(self.train_loader_iter, None) # type: ignore + if batch is None: + self.current_epoch += 1 + # Reset iterator for next epoch + self.train_loader_iter = iter(self.train_dataloader) + # Get first batch of new epoch + batch = next(self.train_loader_iter) + + # latents, encoder_hidden_states, encoder_attention_mask, infos = batch + # encoder_hidden_states = batch['text_embedding'] + # encoder_attention_mask = batch['text_attention_mask'] + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + infos = batch['info_list'] + batch_prompt = ForwardBatch( + data_type="video", + prompt=infos[0]["caption"], + prompt_embeds=[], + prompt_attention_mask=[], + ) + result_batch = self.validation_pipeline.prompt_encoding_stage( # type: ignore + batch_prompt, self.training_args) + encoder_hidden_states, encoder_attention_mask = result_batch.prompt_embeds, result_batch.prompt_attention_mask + for i in range(len(encoder_hidden_states)): + encoder_hidden_states[i] = encoder_hidden_states[i].to( + get_local_torch_device(), dtype=torch.bfloat16) + encoder_attention_mask[i] = encoder_attention_mask[i].to( + get_local_torch_device(), dtype=torch.bfloat16) + + batch_size = encoder_hidden_states[0].shape[0] + vae_config = self.training_args.pipeline_config.vae_config.arch_config + num_channels = vae_config.latent_channels + spatial_compression_ratio = vae_config.spatial_compression_ratio + + latent_height = self.training_args.num_height // spatial_compression_ratio + latent_width = self.training_args.num_width // spatial_compression_ratio + + latents = torch.randn(batch_size, num_channels, + self.training_args.num_latent_t, latent_height, + latent_width).to(get_local_torch_device(), + dtype=torch.bfloat16) + + training_batch.latents = latents.to(get_local_torch_device(), + dtype=torch.bfloat16) + training_batch.encoder_hidden_states = encoder_hidden_states + training_batch.encoder_attention_mask = encoder_attention_mask + training_batch.infos = infos + return training_batch + + def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch: + """ + Self-forcing training step that alternates between generator and critic training. + """ + gradient_accumulation_steps = getattr(self.training_args, + 'gradient_accumulation_steps', 1) + train_generator = (self.current_trainstep % + self.dfake_gen_update_ratio == 0) + + if getattr( + self, "train_dataloader_2", None + ) is not None and not self.training_args.use_ode_init and not self.training_args.use_context_forcing: + use_gt_trajectory = torch.rand(1, device=get_local_torch_device()) + if dist.is_initialized(): + dist.broadcast(use_gt_trajectory, src=0) + training_batch.use_gt_trajectory = (use_gt_trajectory < 0.5).item() + else: + training_batch.use_gt_trajectory = False + + batches = [] + for _ in range(gradient_accumulation_steps): + if training_batch.use_gt_trajectory: + batch = self._get_next_batch_2(training_batch) + else: + if self.training_args.use_context_forcing: + batch = self._get_next_batch_2(training_batch) + else: + batch = self._get_next_batch(training_batch) + batch = self._prepare_dit_inputs(batch) + if self.training_args.use_context_forcing: + batch = self._prepare_context_forcing_inputs(batch) + if training_batch.use_gt_trajectory: + batch = self._prepare_gt_inputs(batch) + batch = self._build_attention_metadata(batch) + batch.attn_metadata_vsa = copy.deepcopy(batch.attn_metadata) + if batch.attn_metadata is not None: + batch.attn_metadata.VSA_sparsity = 0.0 + batches.append(batch) + + training_batch.dmd_latent_vis_dict = {} + training_batch.fake_score_latent_vis_dict = {} + + if train_generator: + logger.debug("Training generator at step %s", + self.current_trainstep) + self.optimizer.zero_grad() + if self.transformer_2 is not None: + self.optimizer_2.zero_grad() + total_generator_loss = 0.0 + generator_log_dict = {} + + for batch in batches: + # Create a new batch with detached tensors + batch_gen = TrainingBatch() + for key, value in batch.__dict__.items(): + if isinstance(value, torch.Tensor): + setattr(batch_gen, key, value.detach().clone()) + elif isinstance(value, dict): + setattr( + batch_gen, key, { + k: + v.detach().clone() if isinstance( + v, torch.Tensor) else copy.deepcopy(v) + for k, v in value.items() + }) + else: + setattr(batch_gen, key, copy.deepcopy(value)) + + generator_loss, gen_log_dict = self.generator_loss(batch_gen) + if self.training_args.use_ode_init: + batch_gt = self._get_next_batch_2(training_batch) + batch_gt = self._prepare_dit_inputs(batch_gt) + batch_gt = self._prepare_ode_init_inputs(batch_gt) + batch_gt = self._build_attention_metadata(batch_gt) + batch_gt.attn_metadata_vsa = copy.deepcopy( + batch_gt.attn_metadata) + batch_gt_gen = TrainingBatch() + for key, value in batch_gt.__dict__.items(): + if isinstance(value, torch.Tensor): + setattr(batch_gt_gen, key, value.detach().clone()) + elif isinstance(value, dict): + setattr( + batch_gt_gen, key, { + k: + v.detach().clone() if isinstance( + v, torch.Tensor) else copy.deepcopy(v) + for k, v in value.items() + }) + else: + setattr(batch_gt_gen, key, copy.deepcopy(value)) + + gt_pred_video, gt_target_latent, gt_timestep = self._step_predict_next_latent( + batch_gt_gen) + mask = gt_timestep != 0 + # Compute loss + gt_loss = F.mse_loss(gt_pred_video[mask], + gt_target_latent[mask], + reduction="mean") + generator_loss += 0.5 * gt_loss + + # Store visualization data from generator training + with set_forward_context(current_timestep=batch_gen.timesteps, + attn_metadata=batch_gen.attn_metadata): + (generator_loss / gradient_accumulation_steps).backward() + total_generator_loss += generator_loss.detach().item() + generator_log_dict.update(gen_log_dict) + # Store visualization data from generator training + if hasattr(batch_gen, 'dmd_latent_vis_dict'): + training_batch.dmd_latent_vis_dict.update( + batch_gen.dmd_latent_vis_dict) + + # Only clip gradients and step optimizer for the model that is currently training + if hasattr( + self, 'train_transformer_2' + ) and self.train_transformer_2 and self.transformer_2 is not None: + self._clip_model_grad_norm_(batch_gen, self.transformer_2) + self.optimizer_2.step() + self.lr_scheduler_2.step() + else: + self._clip_model_grad_norm_(batch_gen, self.transformer) + self.optimizer.step() + self.lr_scheduler.step() + + if self.generator_ema is not None: + if hasattr( + self, 'train_transformer_2' + ) and self.train_transformer_2 and self.transformer_2 is not None: + # Update EMA for transformer_2 when training it + if self.generator_ema_2 is not None: + self.generator_ema_2.update(self.transformer_2) + else: + self.generator_ema.update(self.transformer) + + avg_generator_loss = torch.tensor(total_generator_loss / + gradient_accumulation_steps, + device=self.device) + world_group = get_world_group() + world_group.all_reduce(avg_generator_loss, + op=torch.distributed.ReduceOp.AVG) + training_batch.generator_loss = avg_generator_loss.item() + else: + training_batch.generator_loss = 0.0 + + logger.debug("Training critic at step %s", self.current_trainstep) + self.fake_score_optimizer.zero_grad() + total_critic_loss = 0.0 + critic_log_dict = {} + + for batch in batches: + # Create a new batch with detached tensors + batch_critic = TrainingBatch() + for key, value in batch.__dict__.items(): + if isinstance(value, torch.Tensor): + setattr(batch_critic, key, value.detach().clone()) + elif isinstance(value, dict): + setattr( + batch_critic, key, { + k: + v.detach().clone() + if isinstance(v, torch.Tensor) else copy.deepcopy(v) + for k, v in value.items() + }) + else: + setattr(batch_critic, key, copy.deepcopy(value)) + + critic_loss, crit_log_dict = self.critic_loss(batch_critic) + with set_forward_context(current_timestep=batch_critic.timesteps, + attn_metadata=batch_critic.attn_metadata): + (critic_loss / gradient_accumulation_steps).backward() + total_critic_loss += critic_loss.detach().item() + critic_log_dict.update(crit_log_dict) + # Store visualization data from critic training + if hasattr(batch_critic, 'fake_score_latent_vis_dict'): + training_batch.fake_score_latent_vis_dict.update( + batch_critic.fake_score_latent_vis_dict) + + if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None: + self._clip_model_grad_norm_(batch_critic, + self.fake_score_transformer_2) + self.fake_score_optimizer_2.step() + self.fake_score_lr_scheduler_2.step() + else: + self._clip_model_grad_norm_(batch_critic, + self.fake_score_transformer) + self.fake_score_optimizer.step() + self.fake_score_lr_scheduler.step() + + avg_critic_loss = torch.tensor(total_critic_loss / + gradient_accumulation_steps, + device=self.device) + world_group = get_world_group() + world_group.all_reduce(avg_critic_loss, + op=torch.distributed.ReduceOp.AVG) + training_batch.fake_score_loss = avg_critic_loss.item() + + training_batch.total_loss = training_batch.generator_loss + training_batch.fake_score_loss + + if training_batch.current_timestep % 10 == 0: + gc.collect() + torch.cuda.empty_cache() + return training_batch + def main(args) -> None: logger.info("Starting Hy15 self-forcing distillation pipeline...") diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index 584ccac71..95e587d7e 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -22,6 +22,7 @@ from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger from fastvideo.models.utils import pred_noise_to_pred_video +from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import SelfForcingFlowMatchScheduler from fastvideo.pipelines import TrainingBatch from fastvideo.training.distillation_pipeline import DistillationPipeline from fastvideo.training.training_utils import (EMA_FSDP, @@ -67,12 +68,12 @@ class SelfForcingDistillationPipeline(DistillationPipeline): except Exception: logger.info("Entered initialize_training_pipeline (rank unknown)") - # self.noise_scheduler = SelfForcingFlowMatchScheduler( - # num_inference_steps=1000, - # shift=5.0, - # sigma_min=0.0, - # extra_one_step=True, - # training=True) + self.noise_scheduler = SelfForcingFlowMatchScheduler( + num_inference_steps=1000, + shift=5.0, + sigma_min=0.0, + extra_one_step=True, + training=True) self.dfake_gen_update_ratio = getattr(training_args, 'dfake_gen_update_ratio', 5) @@ -150,71 +151,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): return flow_matching_loss, log_dict - def _prepare_context_forcing_inputs( - self, training_batch: TrainingBatch) -> TrainingBatch: - trajectory_latents = training_batch.trajectory_latents - - total_num_chunks = self.training_args.num_latent_t // self.num_frame_per_block - # Gradually increase maximum number of chunks - # Increase max_num_chunks by 1 every 100 training steps until it reaches total_num_chunks - max_num_chunks = min(total_num_chunks, - self.current_trainstep // 100 + 2) - context_forcing_length = torch.randint( - 0, max_num_chunks, - (1, ), device=get_local_torch_device()) * self.num_frame_per_block - # broadcast context_forcing_length to all ranks - if dist.is_initialized(): - dist.broadcast(context_forcing_length, src=0) - context_forcing_length = context_forcing_length.cpu().item() - assert context_forcing_length < self.training_args.num_latent_t, "Context forcing length is greater than the number of latent frames" - if context_forcing_length > 0: - training_batch.trajectory_latents = trajectory_latents[:, -1, : - context_forcing_length] - else: - training_batch.trajectory_latents = None - del trajectory_latents - return training_batch - - def _prepare_ode_init_inputs( - self, training_batch: TrainingBatch) -> TrainingBatch: - trajectory_latents = training_batch.trajectory_latents - - trajectory_latents = trajectory_latents[:, :, :21] - _cached_closest_idx_per_dmd = torch.tensor([0, 12, 24, 36, 50], - dtype=torch.long).cpu() - training_batch.trajectory_latents = torch.index_select( - trajectory_latents, - dim=1, - index=_cached_closest_idx_per_dmd.to(trajectory_latents.device)) - del trajectory_latents - return training_batch - - def _prepare_gt_inputs(self, - training_batch: TrainingBatch) -> TrainingBatch: - trajectory_latents = training_batch.trajectory_latents - - start_timestep_index = torch.randint(1, - len(self.denoising_step_list), - (1, ), - device=self.device) - if dist.is_initialized(): - dist.broadcast(start_timestep_index, src=0) - start_timestep_index = start_timestep_index.cpu().item() - - trajectory_latents = trajectory_latents[:, -1, :21] - training_batch.trajectory_latents = self.noise_scheduler.add_noise( - trajectory_latents.flatten(0, 1), - torch.randn_like(trajectory_latents.flatten(0, 1)), - self.denoising_step_list[start_timestep_index] * torch.ones( - [trajectory_latents.shape[0] * trajectory_latents.shape[1]], - device=self.device, - dtype=torch.long)).unflatten( - 0, - (trajectory_latents.shape[0], trajectory_latents.shape[1])) - del trajectory_latents - training_batch.start_timestep_index = start_timestep_index - return training_batch - def _generator_multi_step_simulation_forward( self, training_batch: TrainingBatch, @@ -235,20 +171,20 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # During training, the number of generated frames should be uniformly sampled from # [21, self.num_training_frames], but still being a multiple of self.num_frame_per_block min_num_frames = 20 if self.independent_first_frame else 21 - # max_num_frames = num_training_frames - 1 if self.independent_first_frame else num_training_frames - # assert max_num_frames % self.num_frame_per_block == 0 - # assert min_num_frames % self.num_frame_per_block == 0 - # max_num_blocks = max_num_frames // self.num_frame_per_block - # min_num_blocks = min_num_frames // self.num_frame_per_block + max_num_frames = num_training_frames - 1 if self.independent_first_frame else num_training_frames + assert max_num_frames % self.num_frame_per_block == 0 + assert min_num_frames % self.num_frame_per_block == 0 + max_num_blocks = max_num_frames // self.num_frame_per_block + min_num_blocks = min_num_frames // self.num_frame_per_block # Sample number of blocks and sync across processes - # num_generated_blocks = torch.randint(min_num_blocks, - # max_num_blocks + 1, (1, ), - # device=self.device) - # if dist.is_initialized(): - # dist.broadcast(num_generated_blocks, src=0) - # num_generated_blocks = num_generated_blocks.item() - num_generated_frames = num_training_frames + num_generated_blocks = torch.randint(min_num_blocks, + max_num_blocks + 1, (1, ), + device=self.device) + if dist.is_initialized(): + dist.broadcast(num_generated_blocks, src=0) + num_generated_blocks = num_generated_blocks.item() + num_generated_frames = num_generated_blocks * self.num_frame_per_block if self.independent_first_frame and initial_latent is None: num_generated_frames += 1 min_num_frames += 1 @@ -264,7 +200,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): batch_size, num_generated_frames, *self.video_latent_shape[2:] ] - if training_batch.use_gt_trajectory: + if training_batch.use_gt_trajectory and training_batch.trajectory_latents is not None: noise = training_batch.trajectory_latents.to(self.device, dtype=dtype) else: @@ -278,15 +214,13 @@ class SelfForcingDistillationPipeline(DistillationPipeline): batch_size, num_frames, num_channels, height, width = noise.shape # Block size calculation - # if not self.independent_first_frame or (self.independent_first_frame - # and initial_latent is not None): - # assert num_frames % self.num_frame_per_block == 0 - # num_blocks = num_frames // self.num_frame_per_block - # else: - # assert (num_frames - 1) % self.num_frame_per_block == 0 - # num_blocks = (num_frames - 1) // self.num_frame_per_block - import math - num_blocks = math.ceil(num_frames / self.num_frame_per_block) + if not self.independent_first_frame or (self.independent_first_frame + and initial_latent is not None): + assert num_frames % self.num_frame_per_block == 0 + num_blocks = num_frames // self.num_frame_per_block + else: + assert (num_frames - 1) % self.num_frame_per_block == 0 + num_blocks = (num_frames - 1) // self.num_frame_per_block num_input_frames = initial_latent.shape[ 1] if initial_latent is not None else 0 @@ -338,11 +272,10 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # Step 3: Temporal denoising loop all_num_frames = [self.num_frame_per_block] * num_blocks - # all_num_frames[0] = 1 if self.independent_first_frame and initial_latent is None: all_num_frames = [1] + all_num_frames - if training_batch.use_gt_trajectory: + if training_batch.use_gt_trajectory and training_batch.trajectory_latents is not None: start_timestep_index = training_batch.start_timestep_index end_timestep_index = len(self.denoising_step_list) else: @@ -352,64 +285,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): start_timestep_index, end_timestep_index, device=noise.device) - start_gradient_frame_index = max(0, num_output_frames - 31) - - logger.info("distillation step list: %s", self.denoising_step_list) - - # Initialize txt kv cache - txt_kv_cache = self.transformer( - txt_inference=True, - vision_inference=False, - encoder_hidden_states=training_batch. - conditional_dict["encoder_hidden_states"], - encoder_hidden_states_image=training_batch.image_embeds, - encoder_attention_mask=training_batch. - conditional_dict["encoder_attention_mask"], - timestep=torch.zeros([batch_size], device=noise.device), - cache_txt=True, - ) - - # Context forcing logic - if self.training_args.use_context_forcing and training_batch.trajectory_latents is not None: - gt_latents = training_batch.trajectory_latents - context_forcing_length = gt_latents.shape[1] - output[:, :context_forcing_length] = gt_latents.detach() - num_residual_blocks = ( - self.training_args.num_latent_t - - context_forcing_length) // self.num_frame_per_block - assert num_residual_blocks > 0, "Number of residual blocks is less than 0" - all_num_frames = [self.num_frame_per_block] * num_residual_blocks - - if training_batch.video_latent is not None: - gt_latents = torch.cat([ - gt_latents, - training_batch. - video_latent[:, current_start_frame:current_start_frame + - context_forcing_length], - torch.zeros_like(gt_latents), - ], - dim=2) - - with torch.no_grad(): - training_batch_temp = self._build_distill_input_kwargs( - gt_latents, - torch.zeros([batch_size, context_forcing_length], - device=noise.device, - dtype=torch.int64), - training_batch.conditional_dict, training_batch) - - self.transformer( - txt_inference=False, - vision_inference=True, - hidden_states=training_batch_temp. - input_kwargs['hidden_states'], - timestep=training_batch_temp.input_kwargs['timestep'], - kv_cache=kv_cache1, - txt_kv_cache=txt_kv_cache, - current_start=current_start_frame * self.frame_seq_length, - rope_start_idx=current_start_frame) - - current_start_frame += context_forcing_length + start_gradient_frame_index = max(0, num_output_frames - 21) for block_index, current_num_frames in enumerate(all_num_frames): noisy_input = noise[:, current_start_frame - @@ -424,18 +300,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): ), "Index is greater than the number of denoising steps" # Step 3.1: Spatial denoising loop while index < len(self.denoising_step_list): - noisy_input_latent = noisy_input.clone() - if training_batch.video_latent is not None: - noisy_input_latent = torch.cat([ - noisy_input_latent, - training_batch. - video_latent[:, - current_start_frame:current_start_frame + - current_num_frames], - torch.zeros_like(noisy_input_latent), - ], - dim=2) - if self.same_step_across_blocks: exit_flag = (index == exit_flags[0]) else: @@ -454,33 +318,24 @@ class SelfForcingDistillationPipeline(DistillationPipeline): with torch.no_grad(): # Build input kwargs training_batch_temp = self._build_distill_input_kwargs( - noisy_input_latent, timestep, + noisy_input, timestep, training_batch.conditional_dict, training_batch) - model_output = current_model( - txt_inference=False, - vision_inference=True, + pred_flow = current_model( hidden_states=training_batch_temp. input_kwargs['hidden_states'], - # encoder_hidden_states=training_batch_temp. - # input_kwargs['encoder_hidden_states'], - # encoder_attention_mask=training_batch_temp. - # input_kwargs['encoder_attention_mask'], + encoder_hidden_states=training_batch_temp. + input_kwargs['encoder_hidden_states'], timestep=training_batch_temp. input_kwargs['timestep'], - # encoder_hidden_states_image=training_batch_temp. - # input_kwargs.get('encoder_hidden_states_image'), + encoder_hidden_states_image=training_batch_temp. + input_kwargs.get('encoder_hidden_states_image'), kv_cache=kv_cache1, - txt_kv_cache=txt_kv_cache, - # crossattn_cache=crossattn_cache, + crossattn_cache=crossattn_cache, current_start=current_start_frame * self.frame_seq_length, - rope_start_idx=current_start_frame) - if isinstance(model_output, tuple): - pred_flow, kv_cache1 = model_output - else: - pred_flow = model_output - pred_flow = pred_flow.permute(0, 2, 1, 3, 4) + start_frame=current_start_frame).permute( + 0, 2, 1, 3, 4) denoised_pred = pred_noise_to_pred_video( pred_noise=pred_flow.flatten(0, 1), @@ -490,13 +345,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): 0, pred_flow.shape[:2]) next_timestep = self.denoising_step_list[index + 1] - # noisy_input = pred_noise_to_x_bound( - # pred_noise=pred_flow.flatten(0, 1), - # noise_input_latent=noisy_input.flatten(0, 1), - # timestep=timestep, - # boundary_timestep=next_timestep * torch.ones_like(timestep).flatten(0, 1), - # scheduler=self.noise_scheduler).unflatten( - # 0, pred_flow.shape[:2]) noisy_input = self.noise_scheduler.add_noise( denoised_pred.flatten(0, 1), torch.randn_like(denoised_pred.flatten(0, 1)), @@ -510,62 +358,44 @@ class SelfForcingDistillationPipeline(DistillationPipeline): if current_start_frame < start_gradient_frame_index: with torch.no_grad(): training_batch_temp = self._build_distill_input_kwargs( - noisy_input_latent, timestep, + noisy_input, timestep, training_batch.conditional_dict, training_batch) - model_output = current_model( - txt_inference=False, - vision_inference=True, + pred_flow = current_model( hidden_states=training_batch_temp. input_kwargs['hidden_states'], - # encoder_hidden_states=training_batch_temp. - # input_kwargs['encoder_hidden_states'], - # encoder_attention_mask=training_batch_temp. - # input_kwargs['encoder_attention_mask'], + encoder_hidden_states=training_batch_temp. + input_kwargs['encoder_hidden_states'], timestep=training_batch_temp. input_kwargs['timestep'], - # encoder_hidden_states_image=training_batch_temp. - # input_kwargs.get('encoder_hidden_states_image'), + encoder_hidden_states_image=training_batch_temp. + input_kwargs.get('encoder_hidden_states_image'), kv_cache=kv_cache1, - txt_kv_cache=txt_kv_cache, - # crossattn_cache=crossattn_cache, + crossattn_cache=crossattn_cache, current_start=current_start_frame * self.frame_seq_length, - rope_start_idx=current_start_frame) - if isinstance(model_output, tuple): - pred_flow, kv_cache1 = model_output - else: - pred_flow = model_output - pred_flow = pred_flow.permute(0, 2, 1, 3, 4) + start_frame=current_start_frame).permute( + 0, 2, 1, 3, 4) else: training_batch_temp = self._build_distill_input_kwargs( - noisy_input_latent, timestep, + noisy_input, timestep, training_batch.conditional_dict, training_batch) - model_output = current_model( - txt_inference=False, - vision_inference=True, + pred_flow = current_model( hidden_states=training_batch_temp. input_kwargs['hidden_states'], - # encoder_hidden_states=training_batch_temp. - # input_kwargs['encoder_hidden_states'], - # encoder_attention_mask=training_batch_temp. - # input_kwargs['encoder_attention_mask'], + encoder_hidden_states=training_batch_temp. + input_kwargs['encoder_hidden_states'], timestep=training_batch_temp. input_kwargs['timestep'], - # encoder_hidden_states_image=training_batch_temp. - # input_kwargs.get('encoder_hidden_states_image'), + encoder_hidden_states_image=training_batch_temp. + input_kwargs.get('encoder_hidden_states_image'), kv_cache=kv_cache1, - txt_kv_cache=txt_kv_cache, - # crossattn_cache=crossattn_cache, + crossattn_cache=crossattn_cache, current_start=current_start_frame * self.frame_seq_length, - rope_start_idx=current_start_frame) - if isinstance(model_output, tuple): - pred_flow, kv_cache1 = model_output - else: - pred_flow = model_output - pred_flow = pred_flow.permute(0, 2, 1, 3, 4) + start_frame=current_start_frame).permute( + 0, 2, 1, 3, 4) denoised_pred = pred_noise_to_pred_video( pred_noise=pred_flow.flatten(0, 1), @@ -607,25 +437,18 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # context_timestep is 0 so we use transformer_2 current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer - model_output = current_model( - txt_inference=False, - vision_inference=True, + current_model( hidden_states=training_batch_temp. input_kwargs['hidden_states'], - # encoder_hidden_states=training_batch_temp. - # input_kwargs['encoder_hidden_states'], - # encoder_attention_mask=training_batch_temp. - # input_kwargs['encoder_attention_mask'], + encoder_hidden_states=training_batch_temp. + input_kwargs['encoder_hidden_states'], timestep=training_batch_temp.input_kwargs['timestep'], - # encoder_hidden_states_image=training_batch_temp. - # input_kwargs.get('encoder_hidden_states_image'), + encoder_hidden_states_image=training_batch_temp. + input_kwargs.get('encoder_hidden_states_image'), kv_cache=kv_cache1, - txt_kv_cache=txt_kv_cache, - # crossattn_cache=crossattn_cache, + crossattn_cache=crossattn_cache, current_start=current_start_frame * self.frame_seq_length, - rope_start_idx=current_start_frame) - if isinstance(model_output, tuple): - _, kv_cache1 = model_output + start_frame=current_start_frame) # Step 3.4: update the start and end frame indices current_start_frame += current_num_frames @@ -637,7 +460,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # Slice last 21 frames if we generated more gradient_mask = None - if pred_image_or_video.shape[1] > 31: + if pred_image_or_video.shape[1] > 21: with torch.no_grad(): # Re-encode to get image latent latent_to_decode = pred_image_or_video[:, :-20, ...] @@ -677,14 +500,14 @@ class SelfForcingDistillationPipeline(DistillationPipeline): pred_image_or_video_last_21 = pred_image_or_video # Set up gradient mask if we generated more than minimum frames - # if num_generated_frames != min_num_frames: - # # Currently, we do not use gradient for the first chunk, since it contains image latents - # gradient_mask = torch.ones_like(pred_image_or_video_last_21, - # dtype=torch.bool) - # if self.independent_first_frame: - # gradient_mask[:, :1] = False - # else: - # gradient_mask[:, :self.num_frame_per_block] = False + if num_generated_frames != min_num_frames: + # Currently, we do not use gradient for the first chunk, since it contains image latents + gradient_mask = torch.ones_like(pred_image_or_video_last_21, + dtype=torch.bool) + if self.independent_first_frame: + gradient_mask[:, :1] = False + else: + gradient_mask[:, :self.num_frame_per_block] = False # Apply gradient masking if needed final_output = pred_image_or_video_last_21.to(dtype) @@ -716,11 +539,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # Clean up caches assert kv_cache1 is not None - # assert self.crossattn_cache is not None - self._cleanup_simulation_caches(kv_cache1, crossattn_cache) - - assert pred_image_or_video.shape[ - 1] == 21, "pred_image_or_video must have 21 frames" + assert crossattn_cache is not None + self._reset_simulation_caches(kv_cache1, crossattn_cache) return final_output if gradient_mask is not None else pred_image_or_video @@ -755,7 +575,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # Get model configuration parameters - handle FSDP wrapping num_attention_heads = self.transformer.config.num_attention_heads attention_head_dim = self.transformer.config.attention_head_dim - # text_len = getattr(self.transformer, 'text_len', None) + text_len = getattr(self.transformer.config, 'text_len', None) if max_num_frames is None: max_num_frames = num_frames @@ -788,57 +608,29 @@ class SelfForcingDistillationPipeline(DistillationPipeline): }) # Initialize cross-attention cache - crossattn_cache = None - # crossattn_cache = [] - # for _ in range(num_transformer_blocks): - # crossattn_cache.append({ - # "k": - # torch.zeros([ - # batch_size, text_len, num_attention_heads, - # attention_head_dim - # ], - # dtype=dtype, - # device=device), - # "v": - # torch.zeros([ - # batch_size, text_len, num_attention_heads, - # attention_head_dim - # ], - # dtype=dtype, - # device=device), - # "is_init": - # False - # }) + crossattn_cache = [] + for _ in range(num_transformer_blocks): + crossattn_cache.append({ + "k": + torch.zeros([ + batch_size, text_len, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "v": + torch.zeros([ + batch_size, text_len, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "is_init": + False + }) return kv_cache, crossattn_cache - def _cleanup_simulation_caches( - self, kv_cache: list[dict[str, Any]], - crossattn_cache: list[dict[str, Any]]) -> None: - """Properly clean up KV cache and cross-attention cache GPU memory.""" - if kv_cache is not None: - for cache_dict in kv_cache: - # Clear tensor references to free GPU memory - if "k" in cache_dict and cache_dict["k"] is not None: - cache_dict["k"] = None - if "v" in cache_dict and cache_dict["v"] is not None: - cache_dict["v"] = None - if "global_end_index" in cache_dict and cache_dict[ - "global_end_index"] is not None: - cache_dict["global_end_index"] = None - if "local_end_index" in cache_dict and cache_dict[ - "local_end_index"] is not None: - cache_dict["local_end_index"] = None - - if crossattn_cache is not None: - for cache_dict in crossattn_cache: - # Clear tensor references to free GPU memory - if "k" in cache_dict and cache_dict["k"] is not None: - cache_dict["k"] = None - if "v" in cache_dict and cache_dict["v"] is not None: - cache_dict["v"] = None - cache_dict["is_init"] = False - def _reset_simulation_caches(self, kv_cache: list[dict[str, Any]], crossattn_cache: list[dict[str, Any]]) -> None: """Reset KV cache and cross-attention cache to clean state.""" @@ -855,207 +647,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): cache_dict["k"].zero_() cache_dict["v"].zero_() - def _step_predict_next_latent( - self, training_batch: TrainingBatch - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, - torch.Tensor]]: - latent_vis_dict: dict[str, torch.Tensor] = {} - device = get_local_torch_device() - target_latent = training_batch.trajectory_latents[:, -1] - - # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S] - B, S, num_frames, num_channels, height, width = training_batch.trajectory_latents.shape - - dmd_denoising_steps = torch.cat([ - self.denoising_step_list, - torch.tensor([0], device=self.device, dtype=torch.long) - ]) - indexes = self._get_timestep( # [B, num_frames] - 0, - len(dmd_denoising_steps), - B, - num_frames, - 3, - uniform_timestep=False) - # noisy_input = relevant_traj_latents[indexes] - latents = torch.gather( - training_batch.trajectory_latents, - dim=1, - index=indexes.reshape(B, 1, num_frames, 1, 1, - 1).expand(-1, -1, -1, num_channels, height, - width).to(self.device)).squeeze(1) - noisy_input = torch.cat([ - latents, - torch.zeros_like(latents), - torch.zeros_like(latents[:, :, 0:1]) - ], - dim=2) - timestep = dmd_denoising_steps[indexes] - - logger.info("timestep: %s", timestep) - txt_input_kwargs = { - "txt_inference": - True, - "vision_inference": - False, - "encoder_hidden_states": - training_batch.encoder_hidden_states, - "encoder_hidden_states_image": - training_batch.image_embeds, - "encoder_attention_mask": - training_batch.encoder_attention_mask, - "timestep": - torch.zeros([latents.shape[0]], - device=latents.device, - dtype=torch.bfloat16), - "cache_txt": - True, - } - with set_forward_context(current_timestep=timestep, - attn_metadata=None, - forward_batch=None): - txt_kv_cache = self.transformer(**txt_input_kwargs) - - vision_input_kwargs = { - "txt_inference": False, - "vision_inference": True, - "hidden_states": noisy_input.permute(0, 2, 1, 3, 4), - "timestep": timestep.to(device, dtype=torch.bfloat16), - "txt_kv_cache": txt_kv_cache, - } - # Predict noise and step the scheduler to obtain next latent - with set_forward_context(current_timestep=timestep, - attn_metadata=None, - forward_batch=None): - noise_pred = self.transformer(**vision_input_kwargs).permute( - 0, 2, 1, 3, 4) - - from fastvideo.models.utils import pred_noise_to_pred_video - pred_video = pred_noise_to_pred_video( - pred_noise=noise_pred.flatten(0, 1), - noise_input_latent=latents.flatten(0, 1), - timestep=timestep.to(dtype=torch.bfloat16).flatten(0, 1), - scheduler=self.noise_scheduler).unflatten(0, noise_pred.shape[:2]) - latent_vis_dict["pred_video"] = pred_video.permute( - 0, 2, 1, 3, 4).detach().clone().cpu() - - return pred_video, target_latent, timestep - - def _get_timestep(self, - min_timestep: int, - max_timestep: int, - batch_size: int, - num_frame: int, - num_frame_per_block: int, - uniform_timestep: bool = False) -> torch.Tensor: - if uniform_timestep: - timestep = torch.randint(min_timestep, - max_timestep, [batch_size, 1], - device=self.device, - dtype=torch.long).repeat(1, num_frame) - return timestep - else: - num_pad_frames = 0 - if num_frame % num_frame_per_block != 0: - # Pad num_frame to be divisible by num_frame_per_block - num_pad_frames = num_frame_per_block - (num_frame % - num_frame_per_block) - num_frame += num_pad_frames - timestep = torch.randint(min_timestep, - max_timestep, [batch_size, num_frame], - device=self.device, - dtype=torch.long) - # logger.info(f"individual timestep: {timestep}") - # make the noise level the same within every block - timestep = timestep.reshape(timestep.shape[0], -1, - num_frame_per_block) - timestep[:, :, 1:] = timestep[:, :, 0:1] - timestep = timestep.reshape(timestep.shape[0], -1) - if num_pad_frames > 0: - timestep = timestep[:, num_pad_frames:] - return timestep - - def _get_next_batch_2(self, training_batch: TrainingBatch) -> TrainingBatch: - batch = next(self.train_loader_iter_2, None) # type: ignore - if batch is None: - self.current_epoch += 1 - self.train_loader_iter_2 = iter(self.train_dataloader_2) - batch = next(self.train_loader_iter_2) - - # Required fields from parquet (ODE trajectory schema) - device = get_local_torch_device() - encoder_hidden_states = batch['text_embedding'].to( - device, dtype=torch.bfloat16).squeeze(0) - encoder_hidden_states_2 = batch['text_embedding_2'].to( - device, dtype=torch.bfloat16).squeeze(0) - encoder_attention_mask = batch['text_mask'].to( - device, dtype=torch.bfloat16).squeeze(0) - encoder_attention_mask_2 = batch['text_mask_2'].to( - device, dtype=torch.bfloat16).squeeze(0) - infos = batch['info_list'] - - if encoder_hidden_states.dim() < 3: - prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text( - infos[0]["caption"], - self.training_args, - encoder_index=[0], - return_attention_mask=True, - ) - encoder_hidden_states = prompt_embeds_list[0].to( - device, dtype=torch.bfloat16) - encoder_attention_mask = prompt_masks_list[0].to( - device, dtype=torch.bfloat16) - - # Trajectory tensors may include a leading singleton batch dim per row - trajectory_latents = batch['trajectory_latents'] - if trajectory_latents.dim() == 7: - # [B, 1, S, C, T, H, W] -> [B, S, C, T, H, W] - trajectory_latents = trajectory_latents[:, 0] - elif trajectory_latents.dim() == 6: - # already [B, S, C, T, H, W] - pass - else: - raise ValueError( - f"Unexpected trajectory_latents dim: {trajectory_latents.dim()}" - ) - - trajectory_timesteps = batch['trajectory_timesteps'] - if trajectory_timesteps.dim() == 3: - # [B, 1, S] -> [B, S] - trajectory_timesteps = trajectory_timesteps[:, 0] - elif trajectory_timesteps.dim() == 2: - # [B, S] - pass - else: - raise ValueError( - f"Unexpected trajectory_timesteps dim: {trajectory_timesteps.dim()}" - ) - # [B, S, C, T, H, W] -> [B, S, T, C, H, W] to match self-forcing - trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5) - - batch_size = trajectory_latents.shape[0] - vae_config = self.training_args.pipeline_config.vae_config.arch_config - num_channels = vae_config.latent_channels - spatial_compression_ratio = vae_config.spatial_compression_ratio - - latent_height = self.training_args.num_height // spatial_compression_ratio - latent_width = self.training_args.num_width // spatial_compression_ratio - training_batch.latents = torch.randn( - batch_size, num_channels, self.training_args.num_latent_t, - latent_height, latent_width).to(get_local_torch_device(), - dtype=torch.bfloat16) - training_batch.trajectory_latents = trajectory_latents.to( - device, dtype=torch.bfloat16) - training_batch.trajectory_timesteps = trajectory_timesteps.to(device) - training_batch.encoder_hidden_states = [ - encoder_hidden_states, encoder_hidden_states_2 - ] - training_batch.encoder_attention_mask = [ - encoder_attention_mask, encoder_attention_mask_2 - ] - training_batch.infos = infos - return training_batch - def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: batch = next(self.train_loader_iter, None) # type: ignore if batch is None: @@ -1066,28 +657,13 @@ class SelfForcingDistillationPipeline(DistillationPipeline): batch = next(self.train_loader_iter) # latents, encoder_hidden_states, encoder_attention_mask, infos = batch - # encoder_hidden_states = batch['text_embedding'] - # encoder_attention_mask = batch['text_attention_mask'] - from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + encoder_hidden_states = batch['text_embedding'] + encoder_attention_mask = batch['text_attention_mask'] infos = batch['info_list'] - batch_prompt = ForwardBatch( - data_type="video", - prompt=infos[0]["caption"], - prompt_embeds=[], - prompt_attention_mask=[], - ) - result_batch = self.validation_pipeline.prompt_encoding_stage( # type: ignore - batch_prompt, self.training_args) - encoder_hidden_states, encoder_attention_mask = result_batch.prompt_embeds, result_batch.prompt_attention_mask - for i in range(len(encoder_hidden_states)): - encoder_hidden_states[i] = encoder_hidden_states[i].to( - get_local_torch_device(), dtype=torch.bfloat16) - encoder_attention_mask[i] = encoder_attention_mask[i].to( - get_local_torch_device(), dtype=torch.bfloat16) - batch_size = encoder_hidden_states[0].shape[0] + batch_size = encoder_hidden_states.shape[0] vae_config = self.training_args.pipeline_config.vae_config.arch_config - num_channels = vae_config.latent_channels + num_channels = vae_config.z_dim spatial_compression_ratio = vae_config.spatial_compression_ratio latent_height = self.training_args.num_height // spatial_compression_ratio @@ -1100,8 +676,10 @@ class SelfForcingDistillationPipeline(DistillationPipeline): training_batch.latents = latents.to(get_local_torch_device(), dtype=torch.bfloat16) - training_batch.encoder_hidden_states = encoder_hidden_states - training_batch.encoder_attention_mask = encoder_attention_mask + training_batch.encoder_hidden_states = encoder_hidden_states.to( + get_local_torch_device(), dtype=torch.bfloat16) + training_batch.encoder_attention_mask = encoder_attention_mask.to( + get_local_torch_device(), dtype=torch.bfloat16) training_batch.infos = infos return training_batch @@ -1114,32 +692,11 @@ class SelfForcingDistillationPipeline(DistillationPipeline): train_generator = (self.current_trainstep % self.dfake_gen_update_ratio == 0) - if getattr( - self, "train_dataloader_2", None - ) is not None and not self.training_args.use_ode_init and not self.training_args.use_context_forcing: - use_gt_trajectory = torch.rand(1, device=get_local_torch_device()) - if dist.is_initialized(): - dist.broadcast(use_gt_trajectory, src=0) - training_batch.use_gt_trajectory = (use_gt_trajectory < 0.5).item() - # training_batch.use_gt_trajectory = True - else: - training_batch.use_gt_trajectory = False - batches = [] for _ in range(gradient_accumulation_steps): - if training_batch.use_gt_trajectory: - batch = self._get_next_batch_2(training_batch) - else: - if self.training_args.use_context_forcing: - batch = self._get_next_batch_2(training_batch) - else: - batch = self._get_next_batch(training_batch) + batch = self._get_next_batch(training_batch) batch = self._normalize_dit_input(batch) batch = self._prepare_dit_inputs(batch) - if self.training_args.use_context_forcing: - batch = self._prepare_context_forcing_inputs(batch) - if training_batch.use_gt_trajectory: - batch = self._prepare_gt_inputs(batch) batch = self._build_attention_metadata(batch) batch.attn_metadata_vsa = copy.deepcopy(batch.attn_metadata) if batch.attn_metadata is not None: @@ -1176,38 +733,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): setattr(batch_gen, key, copy.deepcopy(value)) generator_loss, gen_log_dict = self.generator_loss(batch_gen) - if self.training_args.use_ode_init: - batch_gt = self._get_next_batch_2(training_batch) - batch_gt = self._prepare_dit_inputs(batch_gt) - batch_gt = self._prepare_ode_init_inputs(batch_gt) - batch_gt = self._build_attention_metadata(batch_gt) - batch_gt.attn_metadata_vsa = copy.deepcopy( - batch_gt.attn_metadata) - batch_gt_gen = TrainingBatch() - for key, value in batch_gt.__dict__.items(): - if isinstance(value, torch.Tensor): - setattr(batch_gt_gen, key, value.detach().clone()) - elif isinstance(value, dict): - setattr( - batch_gt_gen, key, { - k: - v.detach().clone() if isinstance( - v, torch.Tensor) else copy.deepcopy(v) - for k, v in value.items() - }) - else: - setattr(batch_gt_gen, key, copy.deepcopy(value)) - - gt_pred_video, gt_target_latent, gt_timestep = self._step_predict_next_latent( - batch_gt_gen) - mask = gt_timestep != 0 - # Compute loss - gt_loss = F.mse_loss(gt_pred_video[mask], - gt_target_latent[mask], - reduction="mean") - generator_loss += 0.5 * gt_loss - - # Store visualization data from generator training with set_forward_context(current_timestep=batch_gen.timesteps, attn_metadata=batch_gen.attn_metadata): (generator_loss / gradient_accumulation_steps).backward() @@ -1303,10 +828,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): training_batch.fake_score_loss = avg_critic_loss.item() training_batch.total_loss = training_batch.generator_loss + training_batch.fake_score_loss - - if training_batch.current_timestep % 10 == 0: - gc.collect() - torch.cuda.empty_cache() return training_batch def _log_training_info(self) -> None: diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 1a8087d49..c908d44fd 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -14,6 +14,7 @@ import numpy as np import torch import torch.distributed as dist import torchvision +from diffusers import FlowMatchEulerDiscreteScheduler from einops import rearrange from torch.utils.data import DataLoader from torchdata.stateful_dataloader import StatefulDataLoader @@ -110,7 +111,7 @@ class TrainingPipeline(LoRAPipeline, ABC): # Set random seeds for deterministic training assert self.seed is not None, "seed must be set" - set_random_seed(self.seed) + set_random_seed(self.seed + self.global_rank) self.transformer.train() if training_args.enable_gradient_checkpointing_type is not None: self.transformer = apply_activation_checkpointing( @@ -129,22 +130,24 @@ class TrainingPipeline(LoRAPipeline, ABC): params_to_optimize = list( filter(lambda p: p.requires_grad, params_to_optimize)) # Parse betas from string format "beta1,beta2" - # betas_str = training_args.betas - # betas = tuple(float(x.strip()) for x in betas_str.split(",")) - - # self.optimizer = torch.optim.AdamW( - # params_to_optimize, - # lr=training_args.learning_rate, - # betas=betas, - # weight_decay=training_args.weight_decay, - # eps=1e-8, - # ) - - self.optimizer = get_muon_optimizer( - self.transformer, - lr=training_args.learning_rate, - weight_decay=training_args.weight_decay, - ) + if training_args.optimizer_type == "adamw": + betas_str = training_args.betas + betas = tuple(float(x.strip()) for x in betas_str.split(",")) + self.optimizer = torch.optim.AdamW( + params_to_optimize, + lr=training_args.learning_rate, + betas=betas, + weight_decay=training_args.weight_decay, + eps=1e-8, + ) + elif training_args.optimizer_type == "muon": + self.optimizer = get_muon_optimizer( + self.transformer, + lr=training_args.learning_rate, + weight_decay=training_args.weight_decay, + ) + else: + raise ValueError(f"Unsupported optimizer type: {training_args.optimizer_type}") self.init_steps = 0 logger.info("optimizer: %s", self.optimizer) @@ -164,13 +167,23 @@ class TrainingPipeline(LoRAPipeline, ABC): params_to_optimize_2 = self.transformer_2.parameters() params_to_optimize_2 = list( filter(lambda p: p.requires_grad, params_to_optimize_2)) - self.optimizer_2 = torch.optim.AdamW( - params_to_optimize_2, - lr=training_args.learning_rate, - betas=(0.9, 0.999), - weight_decay=training_args.weight_decay, - eps=1e-8, - ) + if training_args.optimizer_type == "adamw": + self.optimizer_2 = torch.optim.AdamW( + params_to_optimize_2, + lr=training_args.learning_rate, + betas=(0.9, 0.999), + weight_decay=training_args.weight_decay, + eps=1e-8, + ) + elif training_args.optimizer_type == "muon": + self.optimizer_2 = get_muon_optimizer( + self.transformer_2, + lr=training_args.learning_rate, + weight_decay=training_args.weight_decay, + ) + else: + raise ValueError(f"Unsupported optimizer type: {training_args.optimizer_type}") + self.lr_scheduler_2 = get_scheduler( training_args.lr_scheduler, optimizer=self.optimizer_2, @@ -295,7 +308,7 @@ class TrainingPipeline(LoRAPipeline, ABC): # TODO(will): support other models with self.tracker.timed("timing/normalize_input"): training_batch.latents = normalize_dit_input( - 'hunyuan', + 'wan', training_batch.latents, self.get_module("vae"), ) @@ -606,7 +619,7 @@ class TrainingPipeline(LoRAPipeline, ABC): logger.info("Initialized random seeds with seed: %s", self.seed + self.global_rank) - # self.noise_scheduler = FlowMatchEulerDiscreteScheduler() + self.noise_scheduler = FlowMatchEulerDiscreteScheduler() if self.training_args.resume_from_checkpoint: self._resume_from_checkpoint() @@ -671,22 +684,25 @@ class TrainingPipeline(LoRAPipeline, ABC): "grad_norm": grad_norm, "vsa_sparsity": current_vsa_sparsity, } - # metrics["batch_size"] = int(training_batch.raw_latent_shape[0]) + try: + metrics["batch_size"] = int(training_batch.raw_latent_shape[0]) - # patch_t, patch_h, patch_w = self.training_args.pipeline_config.dit_config.patch_size - # seq_len = (training_batch.raw_latent_shape[2] // patch_t) * ( - # training_batch.raw_latent_shape[3] // - # patch_h) * (training_batch.raw_latent_shape[4] // patch_w) - # context_len = int(training_batch.encoder_hidden_states.shape[1]) + patch_t, patch_h, patch_w = self.training_args.pipeline_config.dit_config.patch_size + seq_len = (training_batch.raw_latent_shape[2] // patch_t) * ( + training_batch.raw_latent_shape[3] // + patch_h) * (training_batch.raw_latent_shape[4] // patch_w) + context_len = int(training_batch.encoder_hidden_states.shape[1]) - # metrics["dit_seq_len"] = int(seq_len) - # metrics["context_len"] = context_len + metrics["dit_seq_len"] = int(seq_len) + metrics["context_len"] = context_len - # arch_config = self.training_args.pipeline_config.dit_config.arch_config + arch_config = self.training_args.pipeline_config.dit_config.arch_config - # metrics["hidden_dim"] = arch_config.hidden_size - # metrics["num_layers"] = arch_config.num_layers - # metrics["ffn_dim"] = arch_config.ffn_dim + metrics["hidden_dim"] = arch_config.hidden_size + metrics["num_layers"] = arch_config.num_layers + metrics["ffn_dim"] = arch_config.ffn_dim + except: + pass self.tracker.log(metrics, step) if step % self.training_args.training_state_checkpointing_steps == 0: @@ -709,12 +725,6 @@ class TrainingPipeline(LoRAPipeline, ABC): "profiler_region_training_validation"): self._log_validation(self.transformer, self.training_args, step) - for module in self.transformer.modules(): - for name, buf in module._buffers.items(): - if "k_cache" in name or "v_cache" in name: - module._buffers[name] = None - gc.collect() - torch.cuda.empty_cache() gpu_memory_usage = current_platform.get_torch_device( ).memory_allocated() / 1024**2 trainable_params = round(