nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

exp

Browse Source
This commit is contained in:
Will Lin
2025-06-05 22:50:54 -07:00
parent d6ef6c6ae4
commit 4388fa043d
7 changed files with 145 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/v1/layers/layernorm.py
+19 -10
View File
@@ -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)
fastvideo/v1/models/dits/wanvideo.py
+4 -2
View File
@@ -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)
fastvideo/v1/pipelines/stages/denoising.py
+11 -8
View File
@@ -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,
fastvideo/v1/training/training_pipeline.py
+2
View File
@@ -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)
fastvideo/v1/training/wan_training_pipeline.py
+3
View File
@@ -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()
finetune_v1.sh
+53
View File
@@ -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
finetune_v1_crush.sh
+53
View File
@@ -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