From 4388fa043da49db9ad250d8d5e33e2bffcd2179f Mon Sep 17 00:00:00 2001 From: Will Lin Date: Thu, 5 Jun 2025 22:50:54 -0700 Subject: [PATCH] exp --- fastvideo/v1/layers/layernorm.py | 29 ++++++---- fastvideo/v1/models/dits/wanvideo.py | 6 ++- fastvideo/v1/pipelines/stages/denoising.py | 19 ++++--- fastvideo/v1/training/training_pipeline.py | 2 + .../v1/training/wan_training_pipeline.py | 3 ++ finetune_v1.sh | 53 +++++++++++++++++++ finetune_v1_crush.sh | 53 +++++++++++++++++++ 7 files changed, 145 insertions(+), 20 deletions(-) create mode 100644 finetune_v1.sh create mode 100644 finetune_v1_crush.sh diff --git a/fastvideo/v1/layers/layernorm.py b/fastvideo/v1/layers/layernorm.py index ba98c46a2..f80cb0ab4 100644 --- a/fastvideo/v1/layers/layernorm.py +++ b/fastvideo/v1/layers/layernorm.py @@ -9,6 +9,17 @@ import torch.nn as nn from fastvideo.v1.layers.custom_op import CustomOp +class FP32LayerNorm(nn.LayerNorm): + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + origin_dtype = inputs.dtype + return torch.nn.functional.layer_norm( + inputs.float(), + self.normalized_shape, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ).to(origin_dtype) + @CustomOp.register("rms_norm") class RMSNorm(CustomOp): """Root mean square normalization. @@ -121,10 +132,9 @@ class ScaleResidualLayerNormScaleShift(nn.Module): eps=eps, dtype=dtype) elif norm_type == "layer": - self.norm = nn.LayerNorm(hidden_size, + self.norm = FP32LayerNorm(hidden_size, elementwise_affine=elementwise_affine, - eps=eps, - dtype=dtype) + eps=eps) else: raise NotImplementedError(f"Norm type {norm_type} not implemented") @@ -144,9 +154,9 @@ class ScaleResidualLayerNormScaleShift(nn.Module): # Apply residual connection with gating residual_output = residual + x * gate # Apply normalization - normalized = self.norm(residual_output) + normalized = self.norm(residual_output.float()).to(residual_output.dtype) # Apply scale and shift - modulated = normalized * (1.0 + scale) + shift + modulated = (normalized.float() * (1.0 + scale) + shift).to(residual_output.dtype) return modulated, residual_output @@ -171,15 +181,14 @@ class LayerNormScaleShift(nn.Module): has_weight=elementwise_affine, eps=eps) elif norm_type == "layer": - self.norm = nn.LayerNorm(hidden_size, + self.norm = FP32LayerNorm(hidden_size, elementwise_affine=elementwise_affine, - eps=eps, - dtype=dtype) + eps=eps) else: raise NotImplementedError(f"Norm type {norm_type} not implemented") def forward(self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: """Apply ln followed by scale and shift in a single fused operation.""" - normalized = self.norm(x) - return normalized * (1.0 + scale) + shift + normalized = self.norm(x.float()).to(x.dtype) + return (normalized.float() * (1.0 + scale) + shift).to(x.dtype) diff --git a/fastvideo/v1/models/dits/wanvideo.py b/fastvideo/v1/models/dits/wanvideo.py index 2470cd156..c93f96a2f 100644 --- a/fastvideo/v1/models/dits/wanvideo.py +++ b/fastvideo/v1/models/dits/wanvideo.py @@ -15,7 +15,8 @@ from fastvideo.v1.distributed.parallel_state import ( from fastvideo.v1.forward_context import get_forward_context from fastvideo.v1.layers.layernorm import (LayerNormScaleShift, RMSNorm, ScaleResidual, - ScaleResidualLayerNormScaleShift) + ScaleResidualLayerNormScaleShift, + FP32LayerNorm) from fastvideo.v1.layers.linear import ReplicatedLinear # from torch.nn import RMSNorm # TODO: RMSNorm .... @@ -229,7 +230,8 @@ class WanTransformerBlock(nn.Module): super().__init__() # 1. Self-attention - self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False) + # self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False) + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) self.to_q = ReplicatedLinear(dim, dim, bias=True) self.to_k = ReplicatedLinear(dim, dim, bias=True) self.to_v = ReplicatedLinear(dim, dim, bias=True) diff --git a/fastvideo/v1/pipelines/stages/denoising.py b/fastvideo/v1/pipelines/stages/denoising.py index a4afaa8a5..779356665 100644 --- a/fastvideo/v1/pipelines/stages/denoising.py +++ b/fastvideo/v1/pipelines/stages/denoising.py @@ -47,14 +47,15 @@ class DenoisingStage(PipelineStage): super().__init__() self.transformer = transformer self.scheduler = scheduler - attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads - self.attn_backend = get_attn_backend( - head_size=attn_head_size, - dtype=torch.float16, # TODO(will): hack - supported_attention_backends=(_Backend.SLIDING_TILE_ATTN, - _Backend.FLASH_ATTN, - _Backend.TORCH_SDPA) # hack - ) + if transformer is not None: + attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads + self.attn_backend = get_attn_backend( + head_size=attn_head_size, + dtype=torch.float16, # TODO(will): hack + supported_attention_backends=(_Backend.SLIDING_TILE_ATTN, + _Backend.FLASH_ATTN, + _Backend.TORCH_SDPA) # hack + ) def forward( self, @@ -71,6 +72,8 @@ class DenoisingStage(PipelineStage): Returns: The batch with denoised latents. """ + self.transformer.to(fastvideo_args.device) + # Prepare extra step kwargs for scheduler extra_step_kwargs = self.prepare_extra_func_kwargs( self.scheduler.step, diff --git a/fastvideo/v1/training/training_pipeline.py b/fastvideo/v1/training/training_pipeline.py index 37855464b..1a56ce271 100644 --- a/fastvideo/v1/training/training_pipeline.py +++ b/fastvideo/v1/training/training_pipeline.py @@ -208,6 +208,8 @@ class TrainingPipeline(ComposedPipelineBase, ABC): for _, embeddings, masks, infos in validation_dataloader: caption = infos['caption'] captions.extend(caption) + print(f"rank {self.rank} is running validation") + print(f"rank {self.rank} file_name: {infos['file_name']}") prompt_embeds = embeddings.to(training_args.device) prompt_attention_mask = masks.to(training_args.device) diff --git a/fastvideo/v1/training/wan_training_pipeline.py b/fastvideo/v1/training/wan_training_pipeline.py index 02d1d198d..b950c3a6a 100644 --- a/fastvideo/v1/training/wan_training_pipeline.py +++ b/fastvideo/v1/training/wan_training_pipeline.py @@ -272,6 +272,9 @@ class WanTrainingPipeline(TrainingPipeline): logger.info("GPU memory usage before train_one_step: %s MB", gpu_memory_usage) + # Do validation at the beginning of training + self._log_validation(self.transformer, self.training_args, 0) + for step in range(self.init_steps + 1, self.training_args.max_train_steps + 1): start_time = time.perf_counter() diff --git a/finetune_v1.sh b/finetune_v1.sh new file mode 100644 index 000000000..c8dfdfa41 --- /dev/null +++ b/finetune_v1.sh @@ -0,0 +1,53 @@ +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +DATA_DIR=data/cats_480_2_latents_parq_neg/combined_parquet_dataset +VALIDATION_DIR=data/cats_480_2_latents_parq_neg/validation_parquet_dataset +NUM_GPUS=4 +CUDA_VISIBLE_DEVICES=4,5,6,7 +# IP=[MASTER NODE IP] + +# 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/v1/training/wan_training_pipeline.py\ + --model_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ + --inference_mode False\ + --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ + --cache_dir "/home/ray/.cache"\ + --data_path "$DATA_DIR"\ + --validation_prompt_dir "$VALIDATION_DIR"\ + --train_batch_size=1 \ + --num_latent_t 20 \ + --num_gpus 4 \ + --sp_size 4 \ + --tp_size 4 \ + --dp_size 1 \ + --dp_shards 4 \ + --train_sp_batch_size 1\ + --dataloader_num_workers 1\ + --gradient_accumulation_steps=1 \ + --max_train_steps=5000 \ + --learning_rate=1e-6\ + --mixed_precision="bf16"\ + --checkpointing_steps=6000 \ + --validation_steps 10\ + --validation_sampling_steps "2,4,8" \ + --log_validation \ + --checkpoints_total_limit 3\ + --allow_tf32\ + --ema_start_step 0\ + --cfg 0.0\ + --output_dir="data/wan_finetune_crush"\ + --tracker_project_name wan_finetune \ + --num_height 480 \ + --num_width 832 \ + --num_frames 81 \ + --shift 3 \ + --validation_guidance_scale "1.0" \ + --num_euler_timesteps 50 \ + --multi_phased_distill_schedule "4000-1" \ + --weight_decay 0.01 \ + --not_apply_cfg_solver \ + --master_weight_type "fp32" \ + --max_grad_norm 1.0 diff --git a/finetune_v1_crush.sh b/finetune_v1_crush.sh new file mode 100644 index 000000000..9c262fcdf --- /dev/null +++ b/finetune_v1_crush.sh @@ -0,0 +1,53 @@ +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +DATA_DIR=data/crush-smol_parq/combined_parquet_dataset +VALIDATION_DIR=data/crush-smol_parq/validation_parquet_dataset +NUM_GPUS=1 +CUDA_VISIBLE_DEVICES=4,5,6,7 +# IP=[MASTER NODE IP] + +# 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/v1/training/wan_training_pipeline.py\ + --model_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ + --inference_mode False\ + --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ + --cache_dir "/home/ray/.cache"\ + --data_path "$DATA_DIR"\ + --validation_prompt_dir "$VALIDATION_DIR"\ + --train_batch_size=1 \ + --num_latent_t 20 \ + --num_gpus 1 \ + --sp_size 1 \ + --tp_size 1 \ + --dp_size 1 \ + --dp_shards 1 \ + --train_sp_batch_size 1\ + --dataloader_num_workers 1\ + --gradient_accumulation_steps=1 \ + --max_train_steps=5000 \ + --learning_rate=1e-5\ + --mixed_precision="bf16"\ + --checkpointing_steps=6000 \ + --validation_steps 100\ + --validation_sampling_steps "2,4,8" \ + --log_validation \ + --checkpoints_total_limit 3\ + --allow_tf32\ + --ema_start_step 0\ + --cfg 0.0\ + --output_dir="data/wan_finetune_crush"\ + --tracker_project_name finetrainers-wan \ + --num_height 480 \ + --num_width 832 \ + --num_frames 77 \ + --shift 3 \ + --validation_guidance_scale "1.0" \ + --num_euler_timesteps 50 \ + --multi_phased_distill_schedule "4000-1" \ + --weight_decay 0.01 \ + --not_apply_cfg_solver \ + --master_weight_type "fp32" \ + --max_grad_norm 1.0