From 36dab19ff95d55d72025ed7b899e99ebcc14c5cb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CBrianChen1129=E2=80=9D?= Date: Sun, 29 Jun 2025 08:45:56 +0000 Subject: [PATCH] update --- .../v1/training/distillation_pipeline.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/fastvideo/v1/training/distillation_pipeline.py b/fastvideo/v1/training/distillation_pipeline.py index d99f4bcaf..323f5334f 100644 --- a/fastvideo/v1/training/distillation_pipeline.py +++ b/fastvideo/v1/training/distillation_pipeline.py @@ -527,13 +527,14 @@ class DistillationPipeline(TrainingPipeline): for _ in range(self.training_args.gradient_accumulation_steps): training_batch = self._get_next_batch(training_batch) - # assert training_batch.latents is not None - # training_batch.latents = shard_latents_across_sp( - # training_batch.latents, - # num_latent_t=self.training_args.num_latent_t) - training_batch = self._normalize_dit_input(training_batch) training_batch = self._prepare_dit_inputs(training_batch) + + # Shard latents across sp groups + training_batch.latents = shard_latents_across_sp( + training_batch.latents, + num_latent_t=self.training_args.num_latent_t) + training_batch = self._build_attention_metadata(training_batch) if TRAIN_STUDENT: @@ -557,8 +558,6 @@ class DistillationPipeline(TrainingPipeline): self.lr_scheduler.step() self.critic_lr_scheduler.step() - - # Record loss values for logging if TRAIN_STUDENT: training_batch.student_loss = dmd_loss.item() @@ -734,7 +733,7 @@ class DistillationPipeline(TrainingPipeline): self.lr_scheduler, self.noise_random_generator) if self.transformer: self.transformer.train() - # self.sp_group.barrier() + self.sp_group.barrier() if self.training_args.log_validation and step % self.training_args.validation_steps == 0: gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2 @@ -754,8 +753,8 @@ class DistillationPipeline(TrainingPipeline): self.train_dataloader, self.lr_scheduler, self.noise_random_generator) - # if get_sp_group(): - # cleanup_dist_env_and_memory() + if get_sp_group(): + cleanup_dist_env_and_memory() class FlowPredLoss(): def __call__(