From 80ce083ecb9d79bdc85ade7f1c237f0c0124a43f Mon Sep 17 00:00:00 2001
From: RandNMR73 <notomatthew31@gmail.com>
Date: Mon, 6 Oct 2025 03:13:13 +0000
Subject: [PATCH] fix gradient handling

---
 fastvideo/training/distillation_pipeline.py   | 27 ++++++++++++++-----
 .../self_forcing_distillation_pipeline.py     |  3 ---
 fastvideo/training/training_pipeline.py       |  2 +-
 3 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py
index f860cddec..97215fc92 100644
--- a/fastvideo/training/distillation_pipeline.py
+++ b/fastvideo/training/distillation_pipeline.py
@@ -741,8 +741,22 @@ class DistillationPipeline(TrainingPipeline):
         """Compute DMD (Diffusion Model Distillation) loss."""
         original_latent = generator_pred_video
         with torch.no_grad():
-            timestep = torch.randint(0,
-                                     self.num_train_timestep, [1],
+            # MoE logic
+            if exit_timestep is not None and self.boundary_timestep is not None:
+                exit_timestep_type = self._get_moe_timestep_type(exit_timestep)
+                if exit_timestep_type == "low":
+                    start_timestep = 0
+                    end_timestep = self.boundary_timestep
+                else:
+                    start_timestep = self.boundary_timestep
+                    end_timestep = self.num_train_timestep
+            # non-MoE logic
+            else:
+                start_timestep = 0
+                end_timestep = self.num_train_timestep
+
+            timestep = torch.randint(start_timestep,
+                                     end_timestep, [1],
                                      device=self.device,
                                      dtype=torch.long)
             world_group = get_world_group()
@@ -754,12 +768,11 @@ class DistillationPipeline(TrainingPipeline):
                 self.timestep_shift,  # type: ignore
                 self.num_train_timestep)
 
-            timestep = timestep.clamp(self.min_timestep, self.max_timestep)
+            # Ensure shifted timestep stays within the selected MoE noise range
+            if exit_timestep is not None and self.boundary_timestep is not None:
+                timestep = timestep.clamp(start_timestep, end_timestep - 1)
 
-            timestep_type = self._get_moe_timestep_type(timestep)
-            exit_timestep_type = self._get_moe_timestep_type(exit_timestep)
-            if timestep_type != exit_timestep_type:
-                return 0.0
+            timestep = timestep.clamp(self.min_timestep, self.max_timestep)
 
             noise = torch.randn(self.video_latent_shape,
                                 device=self.device,
diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py
index 21e7a6530..66fcd0611 100644
--- a/fastvideo/training/self_forcing_distillation_pipeline.py
+++ b/fastvideo/training/self_forcing_distillation_pipeline.py
@@ -842,9 +842,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
             world_group.all_reduce(avg_generator_loss,
                                    op=torch.distributed.ReduceOp.AVG)
             training_batch.generator_loss = avg_generator_loss.item()
-            # training_batch.fake_score_loss = 0
-            # training_batch.total_loss = training_batch.generator_loss
-            # return training_batch
         else:
             training_batch.generator_loss = 0.0
 
diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py
index 989227f5d..c481269f8 100644
--- a/fastvideo/training/training_pipeline.py
+++ b/fastvideo/training/training_pipeline.py
@@ -308,7 +308,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 if self.transformer_2 is not None:
                     self._disable_training(self.transformer_2, self.optimizer_2)
         else:
-            timesteps = 0 # Fill in a dummy value
+            timesteps = 0  # Fill in a dummy value
 
         if self.training_args.sp_size > 1:
             # Make sure that the timesteps are the same across all sp processes.