From 01dbcfa16366eb64741c53da60949f3c01d6728f Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Fri, 5 Sep 2025 06:27:23 +0000 Subject: [PATCH] Add training --- fastvideo/training/training_pipeline.py | 28 +++++++---- fastvideo/training/wan_training_pipeline.py | 1 + wan2.2_train.sh | 56 +++++++++++++++++++++ 3 files changed, 76 insertions(+), 9 deletions(-) create mode 100644 wan2.2_train.sh diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 450f8aafc..50977fa80 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +import gc import dataclasses import math import os @@ -286,7 +287,7 @@ class TrainingPipeline(LoRAPipeline, ABC): def _sample_timesteps(self, batch_size, device): # Determine which model to train based on the boundary timestep if (self.transformer_2 is not None and self.boundary_timestep is not None and - torch.rand(1, generator=self.noise_random_generator).item() > self.training_args.boundary_ratio): + torch.rand(1, generator=self.noise_random_generator).item() <= self.training_args.boundary_ratio): self.train_transformer_2 = True else: self.train_transformer_2 = False @@ -308,9 +309,9 @@ class TrainingPipeline(LoRAPipeline, ABC): boundary_ratio = self.training_args.boundary_ratio if self.train_transformer_2: - u = boundary_ratio + u * (1.0 - boundary_ratio) + u = (1 - boundary_ratio) + u * boundary_ratio # min: 1 - boundary_ratio, max: 1 else: - u = u * boundary_ratio + u = u * (1 - boundary_ratio) # min: 0, max: 1 - boundary_ratio indices = (u * self.noise_scheduler.config.num_train_timesteps).long() return self.noise_scheduler.timesteps[indices].to(device=device) @@ -490,6 +491,11 @@ class TrainingPipeline(LoRAPipeline, ABC): logger.info("Starting training with %s B trainable parameters", round(num_trainable_params / 1e9, 3)) + if getattr(self, "transformer_2", None) is not None: + num_trainable_params = _get_trainable_params(self.transformer_2) + logger.info("Transformer 2: Starting training with %s B trainable parameters", + round(num_trainable_params / 1e9, 3)) + # Set random seeds for deterministic training self.noise_random_generator = torch.Generator(device="cpu").manual_seed( self.seed) @@ -510,7 +516,7 @@ class TrainingPipeline(LoRAPipeline, ABC): self._log_training_info() - self._log_validation(self.transformer, self.training_args, + self._log_validation(self.training_args, self.init_steps) # Train! @@ -572,7 +578,7 @@ class TrainingPipeline(LoRAPipeline, ABC): self.transformer.train() self.sp_group.barrier() if self.training_args.log_validation and step % self.training_args.validation_steps == 0: - self._log_validation(self.transformer, self.training_args, step) + self._log_validation(self.training_args, step) gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2 trainable_params = round( _get_trainable_params(self.transformer) / 1e9, 3) @@ -653,12 +659,12 @@ class TrainingPipeline(LoRAPipeline, ABC): return batch @torch.no_grad() - def _log_validation(self, transformer, training_args, global_step) -> None: + def _log_validation(self, training_args, global_step) -> None: """ Generate a validation video and log it to wandb to check the quality during training. """ training_args.inference_mode = True - training_args.dit_cpu_offload = True + training_args.dit_cpu_offload = False if not training_args.log_validation: return if self.validation_pipeline is None: @@ -680,7 +686,9 @@ class TrainingPipeline(LoRAPipeline, ABC): batch_size=None, num_workers=0) - transformer.eval() + self.transformer.eval() + if getattr(self, "transformer_2", None) is not None: + self.transformer_2.eval() validation_steps = training_args.validation_sampling_steps.split(",") validation_steps = [int(step) for step in validation_steps] @@ -772,4 +780,6 @@ class TrainingPipeline(LoRAPipeline, ABC): # Re-enable gradients for training training_args.inference_mode = False - transformer.train() + self.transformer.train() + if getattr(self, "transformer_2", None) is not None: + self.transformer_2.train() diff --git a/fastvideo/training/wan_training_pipeline.py b/fastvideo/training/wan_training_pipeline.py index a87e3de47..b1fbc94ee 100644 --- a/fastvideo/training/wan_training_pipeline.py +++ b/fastvideo/training/wan_training_pipeline.py @@ -42,6 +42,7 @@ class WanTrainingPipeline(TrainingPipeline): inference_mode=True, loaded_modules={ "transformer": self.get_module("transformer"), + "transformer_2": self.get_module("transformer_2") }, tp_size=training_args.tp_size, sp_size=training_args.sp_size, diff --git a/wan2.2_train.sh b/wan2.2_train.sh new file mode 100644 index 000000000..64b7e1ebd --- /dev/null +++ b/wan2.2_train.sh @@ -0,0 +1,56 @@ +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=offline +export WANDB_API_KEY='50632ebd88ffd970521cec9ab4a1a2d7e85bfc45' +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA +export TRITON_CACHE_DIR=/tmp/triton_cache +# DATA_DIR=/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/ +DATA_DIR=/mnt/weka/home/hao.zhang/wei/FastVideo/data/crush-smol_processed_t2v/combined_parquet_dataset +# VALIDATION_DIR=/mnt/sharefs/users/hao.zhang/Vchitect-2M/mixkit/validation_8.json +VALIDATION_DIR=/mnt/weka/home/hao.zhang/wei/FastVideo/data/crush-smol-single_processed_t2v/validation.json +NUM_GPUS=8 +export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN +# export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] +CHECKPOINT_PATH="outputs_train_test/wan_finetune/checkpoint-10" +# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t +torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \ + fastvideo/training/wan_training_pipeline.py \ + --model_path Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --inference_mode False\ + --pretrained_model_name_or_path Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --cache_dir "/home/ray/.cache" \ + --data_path "$DATA_DIR" \ + --validation_dataset_file "$VALIDATION_DIR" \ + --train_batch_size 1 \ + --num_latent_t 16\ + --sp_size 4 \ + --tp_size 1 \ + --num_gpus $NUM_GPUS \ + --hsdp_replicate_dim 1 \ + --hsdp-shard-dim 8 \ + --train_sp_batch_size 1 \ + --dataloader_num_workers 4 \ + --gradient_accumulation_steps 1 \ + --max_train_steps 30000 \ + --learning_rate 1e-5 \ + --mixed_precision "bf16" \ + --checkpointing_steps 1000 \ + --validation_steps 30 \ + --validation_sampling_steps "40" \ + --checkpoints_total_limit 3 \ + --ema_start_step 0 \ + --training_cfg_rate 0.1 \ + --seed 1024 \ + --output_dir "outputs_train_test/wan_finetune_v1" \ + --tracker_project_name VSA_finetune \ + --num_height 448 \ + --num_width 832 \ + --num_frames 61 \ + --flow_shift 5 \ + --validation_guidance_scale "5.0" \ + --num_euler_timesteps 50 \ + --master_weight_type "fp32" \ + --dit_precision "fp32" \ + --weight_decay 0.01 \ + --max_grad_norm 1.0 \ No newline at end of file