Add context forcing, ode_init to sf training

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 \

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

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
 

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,

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,

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":

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()