diff --git a/fastvideo/v1/layers/layernorm.py b/fastvideo/v1/layers/layernorm.py index f80cb0ab4..1c4f321ea 100644 --- a/fastvideo/v1/layers/layernorm.py +++ b/fastvideo/v1/layers/layernorm.py @@ -10,6 +10,7 @@ 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( @@ -20,6 +21,7 @@ class FP32LayerNorm(nn.LayerNorm): self.eps, ).to(origin_dtype) + @CustomOp.register("rms_norm") class RMSNorm(CustomOp): """Root mean square normalization. @@ -133,8 +135,8 @@ class ScaleResidualLayerNormScaleShift(nn.Module): dtype=dtype) elif norm_type == "layer": self.norm = FP32LayerNorm(hidden_size, - elementwise_affine=elementwise_affine, - eps=eps) + elementwise_affine=elementwise_affine, + eps=eps) else: raise NotImplementedError(f"Norm type {norm_type} not implemented") @@ -154,9 +156,11 @@ class ScaleResidualLayerNormScaleShift(nn.Module): # Apply residual connection with gating residual_output = residual + x * gate # Apply normalization - normalized = self.norm(residual_output.float()).to(residual_output.dtype) + normalized = self.norm(residual_output.float()).to( + residual_output.dtype) # Apply scale and shift - modulated = (normalized.float() * (1.0 + scale) + shift).to(residual_output.dtype) + modulated = (normalized.float() * (1.0 + scale) + shift).to( + residual_output.dtype) return modulated, residual_output @@ -182,8 +186,8 @@ class LayerNormScaleShift(nn.Module): eps=eps) elif norm_type == "layer": self.norm = FP32LayerNorm(hidden_size, - elementwise_affine=elementwise_affine, - eps=eps) + elementwise_affine=elementwise_affine, + eps=eps) else: raise NotImplementedError(f"Norm type {norm_type} not implemented") diff --git a/fastvideo/v1/models/dits/wanvideo.py b/fastvideo/v1/models/dits/wanvideo.py index c93f96a2f..a364ea050 100644 --- a/fastvideo/v1/models/dits/wanvideo.py +++ b/fastvideo/v1/models/dits/wanvideo.py @@ -13,10 +13,9 @@ from fastvideo.v1.configs.sample.wan import WanTeaCacheParams from fastvideo.v1.distributed.parallel_state import ( get_sequence_model_parallel_world_size) from fastvideo.v1.forward_context import get_forward_context -from fastvideo.v1.layers.layernorm import (LayerNormScaleShift, RMSNorm, - ScaleResidual, - ScaleResidualLayerNormScaleShift, - FP32LayerNorm) +from fastvideo.v1.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, + RMSNorm, ScaleResidual, + ScaleResidualLayerNormScaleShift) from fastvideo.v1.layers.linear import ReplicatedLinear # from torch.nn import RMSNorm # TODO: RMSNorm .... diff --git a/fastvideo/v1/pipelines/stages/denoising.py b/fastvideo/v1/pipelines/stages/denoising.py index 779356665..f092e2041 100644 --- a/fastvideo/v1/pipelines/stages/denoising.py +++ b/fastvideo/v1/pipelines/stages/denoising.py @@ -53,8 +53,8 @@ class DenoisingStage(PipelineStage): 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 + _Backend.FLASH_ATTN, + _Backend.TORCH_SDPA) # hack ) def forward( diff --git a/fastvideo/v1/training/training_pipeline.py b/fastvideo/v1/training/training_pipeline.py index 1a56ce271..6217c55b6 100644 --- a/fastvideo/v1/training/training_pipeline.py +++ b/fastvideo/v1/training/training_pipeline.py @@ -223,6 +223,8 @@ class TrainingPipeline(ComposedPipelineBase, ABC): num_frames = (training_args.num_latent_t - 1) * temporal_compression_factor + 1 + logger.info(f"rank {self.rank} num_frames: {num_frames}") + # Prepare batch for validation batch = ForwardBatch( data_type="video", diff --git a/fastvideo/v1/training/wan_training_pipeline.py b/fastvideo/v1/training/wan_training_pipeline.py index b950c3a6a..69d9fd19d 100644 --- a/fastvideo/v1/training/wan_training_pipeline.py +++ b/fastvideo/v1/training/wan_training_pipeline.py @@ -273,7 +273,7 @@ class WanTrainingPipeline(TrainingPipeline): gpu_memory_usage) # Do validation at the beginning of training - self._log_validation(self.transformer, self.training_args, 0) + # self._log_validation(self.transformer, self.training_args, 0) for step in range(self.init_steps + 1, self.training_args.max_train_steps + 1): diff --git a/finetune_v1_crush.sh b/finetune_v1_crush.sh index 9c262fcdf..139682331 100644 --- a/finetune_v1_crush.sh +++ b/finetune_v1_crush.sh @@ -18,7 +18,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\ --data_path "$DATA_DIR"\ --validation_prompt_dir "$VALIDATION_DIR"\ --train_batch_size=1 \ - --num_latent_t 20 \ + --num_latent_t 14 \ --num_gpus 1 \ --sp_size 1 \ --tp_size 1 \