From 7d373cd2c4212eace03cd0bbad093d16b9fc0e7d Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Mon, 26 Jan 2026 00:51:08 +0000 Subject: [PATCH] Add context forcing, ode_init to sf training --- .../SFHunyuan1.5-T2V/distill_dmd.slurm | 17 +- fastvideo/configs/pipelines/hunyuan15.py | 4 +- fastvideo/fastvideo_args.py | 10 + .../models/dits/causal_hunyuanvideo15.py | 107 ++++++----- .../pipelines/stages/causal_denoising.py | 9 +- fastvideo/training/distillation_pipeline.py | 13 +- .../self_forcing_distillation_pipeline.py | 179 +++++++++++++----- 7 files changed, 222 insertions(+), 117 deletions(-) diff --git a/examples/distill/SFHunyuan1.5-T2V/distill_dmd.slurm b/examples/distill/SFHunyuan1.5-T2V/distill_dmd.slurm index 8780f00ec..a3bfbbb0c 100644 --- a/examples/distill/SFHunyuan1.5-T2V/distill_dmd.slurm +++ b/examples/distill/SFHunyuan1.5-T2V/distill_dmd.slurm @@ -47,7 +47,7 @@ FAKE_SCORE_MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2 # DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/" # DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing" DATA_DIR="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/vidprom_16k_text_embed" -DATA_DIR_2="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/fv-ode-preprocessing-16k-hy15-distilled-121" +DATA_DIR_2="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/fv-ode-preprocessing-16k-hy15-121" # DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset # DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/" # VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json" @@ -59,7 +59,7 @@ VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/ training_args=( --tracker_project_name SFhy1.5_t2v_distill_self_forcing_dmd # Updated for Wan2.2 --output_dir "/mnt/weka/home/hao.zhang/wei/SFhy1.5_distill_3333_1e-5_1e-5_cfg6_corrected_scheduler" - --wandb_run_name "self_forcing_3333_1e-5_1e-5_cfg3.5_corrected_scheduler" + --wandb_run_name "self_forcing_3333_context_forcing" --max_train_steps 4000 --train_batch_size 1 --train_sp_batch_size 1 @@ -73,7 +73,8 @@ training_args=( --num_frames 81 --num_frame_per_block 3 # Frame generation block size for self-forcing # --resume-from-checkpoint "/mnt/weka/home/hao.zhang/wei/SFhy1.5_distill_1333/checkpoint-300" - --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wei/hy15_distilled_ode_init_3333/checkpoint-2400/transformer/diffusion_pytorch_model.safetensors + --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wei/hy15_worldplay_df_init_3333/checkpoint-2000/transformer/diffusion_pytorch_model.safetensors + # --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wei/hy15_distilled_ode_init_3333/checkpoint-2400/transformer/diffusion_pytorch_model.safetensors ) parallel_args=( @@ -93,7 +94,7 @@ model_args=( dataset_args=( --data_path "$DATA_DIR" - --data_path_2 "$DATA_DIR_2" + # --data_path_2 "$DATA_DIR_2" --dataloader_num_workers 4 ) @@ -108,7 +109,7 @@ validation_args=( ) optimizer_args=( - --learning_rate 1e-6 + --learning_rate 1e-5 --mixed_precision "bf16" --training_state_checkpointing_steps 100 --weight_only_checkpointing_steps 500 @@ -122,7 +123,7 @@ miscellaneous_args=( --checkpoints_total_limit 3 --training_cfg_rate 0.0 --dit_precision "fp32" - --flow_shift 7 + --flow_shift 5 --seed 1000 --use_ema True --ema_decay 0.99 @@ -133,11 +134,12 @@ dmd_args=( --warp_denoising_step --dmd_denoising_steps '1000,875,750,625,500,375,250,125' # --dmd_denoising_steps '1000,760,520,280' + # --dmd_denoising_steps '1000,750,500,250' --min_timestep_ratio 0.02 --max_timestep_ratio 0.98 --dfake_gen_update_ratio 5 --real_score_guidance_scale 3.5 - --fake_score_learning_rate 2e-7 + --fake_score_learning_rate 5e-6 --fake_score_betas '0.0,0.999' ) @@ -146,6 +148,7 @@ self_forcing_args=( --same_step_across_blocks True # Whether to use same denoising step across all blocks --last_step_only False # Whether to only use the last denoising step --context_noise 0 # Amount of noise to add during context caching (0 = no noise) + # --use-context-forcing True ) srun torchrun \ diff --git a/fastvideo/configs/pipelines/hunyuan15.py b/fastvideo/configs/pipelines/hunyuan15.py index 2d26057d0..02c268c51 100644 --- a/fastvideo/configs/pipelines/hunyuan15.py +++ b/fastvideo/configs/pipelines/hunyuan15.py @@ -143,10 +143,10 @@ class Hunyuan15T2V720PConfig(Hunyuan15T2V480PConfig): @dataclass class SelfForcingHunyuan15T2V480PConfig(Hunyuan15T2V480PConfig): - flow_shift: int = 7 + flow_shift: int = 5 is_causal: bool = True # dmd_denoising_steps: list[int] | None = field( - # default_factory=lambda: [1000, 760, 520, 280]) + # default_factory=lambda: [1000, 750, 500, 250]) dmd_denoising_steps: list[int] | None = field( default_factory=lambda: [1000, 875, 750, 625, 500, 375, 250, 125]) warp_denoising_step: bool = True \ No newline at end of file diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 5564568d2..cf2a000f6 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -945,6 +945,8 @@ class TrainingArgs(FastVideoArgs): same_step_across_blocks: bool = False # Use same exit timestep for all blocks last_step_only: bool = False # Only use the last timestep for training context_noise: int = 0 # Context noise level for cache updates + use_context_forcing: bool = False + use_ode_init: bool = False @classmethod def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs": @@ -1374,6 +1376,14 @@ class TrainingArgs(FastVideoArgs): type=int, default=TrainingArgs.context_noise, help="Context noise level for cache updates") + parser.add_argument( + "--use-context-forcing", + action=StoreBoolean, + help="Whether to use context forcing") + parser.add_argument( + "--use-ode-init", + action=StoreBoolean, + help="Whether to use ODE init") return parser diff --git a/fastvideo/models/dits/causal_hunyuanvideo15.py b/fastvideo/models/dits/causal_hunyuanvideo15.py index 50eb1dcb7..3d7627397 100644 --- a/fastvideo/models/dits/causal_hunyuanvideo15.py +++ b/fastvideo/models/dits/causal_hunyuanvideo15.py @@ -56,25 +56,6 @@ from fastvideo.models.dits.hunyuanvideo15 import ( logger = init_logger(__name__) -class CacheAppend(torch.autograd.Function): - """ - KV cache with shape [batch, seq_len, heads, head_dim]. - """ - @staticmethod - def forward(ctx, storage, active_cache, x, start, end): - # Ensure storage has the same dtype as x - storage.data[:, start:end] = x - ctx.save_for_backward(storage.to(x.dtype)) - ctx.start = start - ctx.end = end - return storage[:, :end].to(x.dtype) # Ensure returned value has same dtype as input - - @staticmethod - def backward(ctx, grad_output): - start = ctx.start - end = ctx.end - return None, grad_output[:, :start], grad_output[:, start:end], None, None - class MMDoubleStreamBlock(nn.Module): """ A multimodal DiT block with separate modulation for text and image/video, @@ -86,6 +67,8 @@ class MMDoubleStreamBlock(nn.Module): hidden_size: int, num_attention_heads: int, mlp_ratio: float, + local_attn_size: int = -1, + sink_size: int = 0, dtype: torch.dtype | None = None, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, @@ -96,6 +79,8 @@ class MMDoubleStreamBlock(nn.Module): self.deterministic = False self.num_attention_heads = num_attention_heads self.head_dim = hidden_size // num_attention_heads + self.local_attn_size = local_attn_size + self.sink_size = sink_size mlp_hidden_dim = int(hidden_size * mlp_ratio) # Image modulation components @@ -179,9 +164,7 @@ class MMDoubleStreamBlock(nn.Module): self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype) - self.max_attention_size = 31 * 1590 - self.k_cache = None - self.v_cache = None + self.max_attention_size = 21 * 1590 if local_attn_size == -1 else local_attn_size * 1590 self.attn = LocalAttention( num_heads=self.num_attention_heads, @@ -252,17 +235,6 @@ class MMDoubleStreamBlock(nn.Module): txt_kv_cache: list | None = None, current_start: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: - if kv_cache is not None: - # Then we are in autoregressive mode without block-wise attention mask - if self.k_cache is None: - del self.k_cache - self.register_buffer("k_cache", torch.empty(1, self.max_attention_size, self.num_attention_heads, self.head_dim, device=img.device, dtype=img.dtype), persistent=False) - self.k_cache.requires_grad_(True) - if self.v_cache is None: - del self.v_cache - self.register_buffer("v_cache", torch.empty(1, self.max_attention_size, self.num_attention_heads, self.head_dim, device=img.device, dtype=img.dtype), persistent=False) - self.v_cache.requires_grad_(True) - # Process modulation vectors if vec.dim() == 3: img_mod_outputs = self.img_mod(vec).unflatten(dim=-1, sizes=(6, -1)) @@ -342,25 +314,54 @@ class MMDoubleStreamBlock(nn.Module): )[:, :, :-padded_length].transpose(2, 1) assert img_attn.shape[1] == image_seq_len + updated_kv_cache = None else: current_end = current_start + img_q.shape[1] num_new_tokens = img_q.shape[1] - assert kv_cache["local_end_index"].item() == kv_cache["global_end_index"].item(), "local_end_index and global_end_index must be the same" - # Assign new keys/values directly up to current_end - local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() - local_start_index = local_end_index - num_new_tokens - img_local_k = CacheAppend.apply(self.k_cache, kv_cache["k"], img_k, local_start_index, local_end_index) - img_local_v = CacheAppend.apply(self.v_cache, kv_cache["v"], img_v, local_start_index, local_end_index) + sink_tokens = self.sink_size * 1590 + # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache + kv_cache_size = self.max_attention_size + + # Clone cache to avoid in-place modification during gradient checkpointing + k_cache = kv_cache["k"].clone() + v_cache = kv_cache["v"].clone() + + if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and ( + num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size): + # Calculate the number of new tokens added in this step + # Shift existing cache content left to discard oldest tokens + # Clone the source slice to avoid overlapping memory error + num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size + num_rolled_tokens = kv_cache_size - num_new_tokens - sink_tokens + k_cache[:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + k_cache[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + v_cache[:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + v_cache[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + # Insert the new keys/values at the end + local_end_index = kv_cache["local_end_index"].item() + current_end - \ + kv_cache["global_end_index"].item() - num_evicted_tokens + local_start_index = local_end_index - num_new_tokens + assert local_end_index == self.max_attention_size + else: + # Assign new keys/values directly up to current_end + local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() + local_start_index = local_end_index - num_new_tokens + assert local_start_index >= 0 q = img_q - k = torch.cat([img_local_k, txt_kv_cache["k_txt"]], dim=1) - v = torch.cat([img_local_v, txt_kv_cache["v_txt"]], dim=1) + k = torch.cat([k_cache[:, :local_start_index], img_k, txt_kv_cache["k_txt"]], dim=1) + v = torch.cat([v_cache[:, :local_start_index], img_v, txt_kv_cache["v_txt"]], dim=1) img_attn = self.attn(q, k, v) - kv_cache["k"] = img_local_k - kv_cache["v"] = img_local_v - kv_cache["global_end_index"].fill_(current_end) - kv_cache["local_end_index"].fill_(local_end_index) + k_cache[:, local_start_index:local_end_index] = img_k + v_cache[:, local_start_index:local_end_index] = img_v + + updated_kv_cache = { + "k": k_cache, + "v": v_cache, + "global_end_index": torch.tensor([current_end], dtype=torch.long, device=k_cache.device), + "local_end_index": torch.tensor([local_end_index], dtype=torch.long, device=k_cache.device) + } img_attn_out, _ = self.img_attn_proj( img_attn.view(batch_size, image_seq_len, -1)) @@ -372,7 +373,7 @@ class MMDoubleStreamBlock(nn.Module): img_mlp_out = self.img_mlp(img_mlp_input) img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate) - return img + return img, updated_kv_cache def forward( self, @@ -444,6 +445,8 @@ class CausalHunyuanVideo15Transformer3DModel(CachableDiT): hidden_size=self.hidden_size, num_attention_heads=config.num_attention_heads, mlp_ratio=config.mlp_ratio, + local_attn_size=config.local_attn_size, + sink_size=config.sink_size, dtype=None, supported_attention_backends=self._supported_attention_backends, prefix=f"{config.prefix}.double_blocks.{i}" @@ -751,7 +754,7 @@ class CausalHunyuanVideo15Transformer3DModel(CachableDiT): # 4. Transformer blocks if torch.is_grad_enabled() and self.gradient_checkpointing: for block_index, block in enumerate(self.double_blocks): - hidden_states = self._gradient_checkpointing_func( + hidden_states, new_cache = self._gradient_checkpointing_func( block, txt_inference=False, vision_inference=True, @@ -763,10 +766,13 @@ class CausalHunyuanVideo15Transformer3DModel(CachableDiT): txt_kv_cache=txt_kv_cache[block_index], current_start=current_start ) + if new_cache is not None and kv_cache is not None: + for k in new_cache.keys(): + kv_cache[block_index][k] = new_cache[k].clone() else: for block_index, block in enumerate(self.double_blocks): - hidden_states = block( + hidden_states, new_cache = block( txt_inference=False, vision_inference=True, img=hidden_states, @@ -777,13 +783,16 @@ class CausalHunyuanVideo15Transformer3DModel(CachableDiT): txt_kv_cache=txt_kv_cache[block_index], current_start=current_start ) + if new_cache is not None and kv_cache is not None: + for k in new_cache.keys(): + kv_cache[block_index][k] = new_cache[k].clone() # Final layer processing hidden_states = self.final_layer(hidden_states, temb) # Unpatchify to get original shape hidden_states = unpatchify(hidden_states, post_patch_num_frames, post_patch_height, post_patch_width, self.patch_size, self.out_channels) - return hidden_states + return hidden_states, kv_cache def forward( self, diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index d00a5a374..86555cacb 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -52,7 +52,7 @@ class CausalDMDDenosingStage(DenoisingStage): 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: @@ -364,7 +364,7 @@ class CausalDMDDenosingStage(DenoisingStage): (latent_model_input.shape[0], current_num_frames), device=latent_model_input.device, dtype=torch.long) - pred_noise_btchw = current_model( + pred_noise_btchw, kv_cache1 = current_model( hidden_states=latent_model_input, timestep=t_expanded_noise, kv_cache=_get_kv_cache(t_cur), @@ -374,7 +374,8 @@ class CausalDMDDenosingStage(DenoisingStage): self.frame_seq_length, rope_start_idx=start_index, **vision_input_kwargs, - ).permute(0, 2, 1, 3, 4) + ) + pred_noise_btchw = pred_noise_btchw.permute(0, 2, 1, 3, 4) # Convert pred noise to pred video with FM Euler scheduler utilities # if i < len(timesteps) - 1: @@ -489,7 +490,7 @@ class CausalDMDDenosingStage(DenoisingStage): **vision_input_kwargs, ) - self.transformer( + _, kv_cache1 = self.transformer( hidden_states=context_bcthw, timestep=t_context, kv_cache=kv_cache1, diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 01192d66d..51adad023 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -147,7 +147,6 @@ class DistillationPipeline(TrainingPipeline): self.vae.requires_grad_(False) self.timestep_shift = self.training_args.pipeline_config.flow_shift - assert self.timestep_shift == 7.0, "flow_shift must be 7.0" # 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()) @@ -789,9 +788,15 @@ class DistillationPipeline(TrainingPipeline): original_latent - real_score_pred_video).mean() grad = torch.nan_to_num(grad) - dmd_loss = 0.5 * F.mse_loss( - original_latent.float(), - (original_latent.float() - grad.float()).detach()) + if self.training_args.use_context_forcing and training_batch.trajectory_latents is not None: + context_forcing_length = training_batch.trajectory_latents.shape[1] + dmd_loss = 0.5 * F.mse_loss( + original_latent.float()[:, context_forcing_length:], + (original_latent.float()[:, context_forcing_length:] - grad.float()[:, context_forcing_length:]).detach()) + else: + dmd_loss = 0.5 * F.mse_loss( + original_latent.float(), + (original_latent.float() - grad.float()).detach()) training_batch.dmd_latent_vis_dict.update({ "training_batch_dmd_fwd_clean_latent": diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index 87d238f7f..7a27b47b7 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -150,24 +150,54 @@ class SelfForcingDistillationPipeline(DistillationPipeline): return flow_matching_loss, log_dict - def _prepare_gt_inputs(self, training_batch: TrainingBatch) -> TrainingBatch: + def _prepare_context_forcing_inputs(self, training_batch: TrainingBatch) -> TrainingBatch: trajectory_latents = training_batch.trajectory_latents - assert trajectory_latents.shape[1] == len(self.denoising_step_list) + 1 - - # start_timestep_index = torch.randint(0, 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() - training_batch.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)) + 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 - # 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_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( @@ -219,10 +249,10 @@ class SelfForcingDistillationPipeline(DistillationPipeline): batch_size, num_generated_frames, *self.video_latent_shape[2:] ] - # if training_batch.use_gt_trajectory: - # noise = training_batch.trajectory_latents.to(self.device, dtype=dtype) - # else: - noise = torch.randn(noise_shape, device=self.device, dtype=dtype) + if training_batch.use_gt_trajectory: + 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", @@ -292,16 +322,16 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # Step 3: Temporal denoising loop all_num_frames = [self.num_frame_per_block] * num_blocks - all_num_frames[0] = 1 + # 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: - # 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) + if training_batch.use_gt_trajectory: + 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, @@ -321,6 +351,45 @@ class SelfForcingDistillationPipeline(DistillationPipeline): 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 + @@ -328,6 +397,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): 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() @@ -359,7 +430,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): noisy_input_latent, timestep, training_batch.conditional_dict, training_batch) - pred_flow = current_model( + pred_flow, kv_cache1 = current_model( txt_inference=False, vision_inference=True, hidden_states=training_batch_temp. @@ -377,8 +448,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # crossattn_cache=crossattn_cache, current_start=current_start_frame * self.frame_seq_length, - rope_start_idx=current_start_frame).permute( - 0, 2, 1, 3, 4) + 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), @@ -411,7 +482,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): noisy_input_latent, timestep, training_batch.conditional_dict, training_batch) - pred_flow = current_model( + pred_flow, kv_cache1 = current_model( txt_inference=False, vision_inference=True, hidden_states=training_batch_temp. @@ -429,14 +500,14 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # crossattn_cache=crossattn_cache, current_start=current_start_frame * self.frame_seq_length, - rope_start_idx=current_start_frame).permute( - 0, 2, 1, 3, 4) + 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 = current_model( + pred_flow, kv_cache1 = current_model( txt_inference=False, vision_inference=True, hidden_states=training_batch_temp. @@ -454,8 +525,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # crossattn_cache=crossattn_cache, current_start=current_start_frame * self.frame_seq_length, - rope_start_idx=current_start_frame).permute( - 0, 2, 1, 3, 4) + 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), @@ -493,7 +564,7 @@ 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 - current_model( + _, kv_cache1 = current_model( txt_inference=False, vision_inference=True, hidden_states=training_batch_temp. @@ -640,7 +711,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): if max_num_frames is None: max_num_frames = num_frames num_max_frames = max(max_num_frames, num_frames) - kv_cache_size = num_max_frames * frame_seq_length + 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): @@ -959,25 +1031,30 @@ 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 train_generator: - # 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.8).item() - training_batch.use_gt_trajectory = True + 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: - batch = self._get_next_batch(training_batch) + 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._normalize_dit_input(batch) batch = self._prepare_dit_inputs(batch) - # if training_batch.use_gt_trajectory: - # batch = self._prepare_gt_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: @@ -1014,10 +1091,10 @@ class SelfForcingDistillationPipeline(DistillationPipeline): setattr(batch_gen, key, copy.deepcopy(value)) generator_loss, gen_log_dict = self.generator_loss(batch_gen) - if training_batch.use_gt_trajectory: + 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_gt_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()