exp

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")
 

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 ....

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(

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",

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):

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 \