fix sp

examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh

@@ -16,7 +16,7 @@ training_args=(
   --max_train_steps 5000
   --train_batch_size 1
   --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
+  --gradient_accumulation_steps 8
   --num_latent_t 8 
   --num_height 480
   --num_width 832
@@ -48,17 +48,17 @@ dataset_args=(
 validation_args=(
   --log_validation 
   --validation_preprocessed_path $VALIDATION_DIR
-  --validation_steps 100
+  --validation_steps 50
   --validation_sampling_steps "50" 
   --validation_guidance_scale "1.0"
 )
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
+  --learning_rate 5e-5
   --mixed_precision "bf16"
   --checkpointing_steps 6000
-  --weight_decay 0.01
+  --weight_decay 1e-4
   --max_grad_norm 1.0
 )
 

fastvideo/v1/pipelines/pipeline_batch_info.py

@@ -149,6 +149,7 @@ class TrainingBatch:
 
     # Dataloader batch outputs
     latents: Optional[torch.Tensor] = None
+    original_latents: Optional[torch.Tensor] = None
     encoder_hidden_states: Optional[torch.Tensor] = None
     encoder_attention_mask: Optional[torch.Tensor] = None
     # i2v

fastvideo/v1/training/training_pipeline.py

@@ -25,7 +25,8 @@ from fastvideo.v1.dataset import build_parquet_map_style_dataloader
 from fastvideo.v1.dataset.dataloader.schema import (
     pyarrow_schema_t2v, pyarrow_schema_t2v_validation)
 from fastvideo.v1.distributed import (cleanup_dist_env_and_memory, get_sp_group,
-                                      get_torch_device, get_world_group)
+                                      get_torch_device, get_world_group,
+                                      sequence_model_parallel_all_gather)
 from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.logger import init_logger
@@ -318,10 +319,15 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             model_pred = self.transformer(**input_kwargs)
         if self.training_args.precondition_outputs:
             model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas
-        target = training_batch.latents if self.training_args.precondition_outputs else training_batch.noise - training_batch.latents
+        target = training_batch.original_latents if self.training_args.precondition_outputs else training_batch.noise - training_batch.latents
 
+        if self.training_args.sp_size > 1:
+            model_pred = sequence_model_parallel_all_gather(model_pred, dim=2)
+
+        assert model_pred.shape == target.shape, f"model_pred.shape: {model_pred.shape}, target.shape: {target.shape}"
         loss = (torch.mean((model_pred.float() - target.float())**2) /
                 self.training_args.gradient_accumulation_steps)
+        # target = sequence_model_parallel_all_gather(target, dim=2)
 
         loss.backward()
         avg_loss = loss.detach().clone()
@@ -363,14 +369,40 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         for _ in range(self.training_args.gradient_accumulation_steps):
             training_batch = self._get_next_batch(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)
+            logger.info(
+                f"training_batch.latents.shape: {training_batch.latents.shape}")
+            training_batch.latents = training_batch.latents[:, :, :self.
+                                                            training_args.
+                                                            num_latent_t]
+            logger.info(
+                f"after slice training_batch.latents.shape: {training_batch.latents.shape}"
+            )
 
             # Normalize DIT input
             training_batch = self._normalize_dit_input(training_batch)
+            logger.info(
+                f"after normalize training_batch.latents.shape: {training_batch.latents.shape}"
+            )
+
+            # Save original latents
+            training_batch.original_latents = training_batch.latents.clone()
+
+            # # 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._prepare_dit_inputs(training_batch)
+            logger.info(
+                f"after prepare dit inputs training_batch.noisy_model_input.shape: {training_batch.noisy_model_input.shape}"
+            )
+            # shard noisy_model_input to match
+            training_batch.noisy_model_input = shard_latents_across_sp(
+                training_batch.noisy_model_input,
+                num_latent_t=self.training_args.num_latent_t)
+            logger.info(
+                f"after shard noisy_model_input training_batch.noisy_model_input.shape: {training_batch.noisy_model_input.shape}"
+            )
             training_batch = self._build_attention_metadata(training_batch)
             training_batch = self._build_input_kwargs(training_batch)
             training_batch = self._transformer_forward_and_compute_loss(