From 79b5d82d9865a1bed82bc4e64896b444a01e78b2 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Wed, 25 Jun 2025 14:55:49 -0700 Subject: [PATCH] fix sp --- .../wan_t2v_1_3b/crush_smol/finetune_t2v.sh | 8 ++-- fastvideo/v1/pipelines/pipeline_batch_info.py | 1 + fastvideo/v1/training/training_pipeline.py | 44 ++++++++++++++++--- 3 files changed, 43 insertions(+), 10 deletions(-) diff --git a/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh b/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh index fcdbf1aa8..b79694d87 100644 --- a/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh +++ b/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 ) diff --git a/fastvideo/v1/pipelines/pipeline_batch_info.py b/fastvideo/v1/pipelines/pipeline_batch_info.py index 4a11227da..a81868cd1 100644 --- a/fastvideo/v1/pipelines/pipeline_batch_info.py +++ b/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 diff --git a/fastvideo/v1/training/training_pipeline.py b/fastvideo/v1/training/training_pipeline.py index a4d8f8115..a6fb397d1 100644 --- a/fastvideo/v1/training/training_pipeline.py +++ b/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(