Enable timestep warping & using SelfForcing scheduler

2025-09-09 23:30:55 +00:00
parent 918180048e
commit 80baffd540
5 changed files with 15 additions and 16 deletions
@@ -16,7 +16,7 @@ export NCCL_P2P_DISABLE=1
 export TORCH_NCCL_ENABLE_MONITORING=0
 # different cache dir for different processes
 export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29500
+export MASTER_PORT=29502
 export TOKENIZERS_PARALLELISM=false
 export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
 export WANDB_BASE_URL="https://api.wandb.ai"
@@ -40,7 +40,7 @@ VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
 training_args=(
  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd  # Updated for self-forcing DMD
  --output_dir "/mnt/sharefs/users/hao.zhang/SFwan_t2v_finetune"
-  --max_train_steps 500
+  --max_train_steps 4000
  --train_batch_size 1
  --train_sp_batch_size 1
  --gradient_accumulation_steps 1
@@ -84,7 +84,7 @@ dataset_args=(
 validation_args=(
  --log_validation
  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 10
+  --validation_steps 20
  --validation_sampling_steps "4"
  --validation_guidance_scale "6.0" # not used for dmd inference
 )
@@ -93,8 +93,8 @@ validation_args=(
 optimizer_args=(
  --learning_rate 1e-5
  --mixed_precision "bf16"
-  --training_state_checkpointing_steps 50
-  --weight_only_checkpointing_steps 50
+  --training_state_checkpointing_steps 200
+  --weight_only_checkpointing_steps 200
  --weight_decay 0.01
  --betas '0.0,0.999'
  --max_grad_norm 1.0
@@ -123,6 +123,7 @@ dmd_args=(
  --real_score_guidance_scale 3.0
  --fake_score_learning_rate 8e-6
  --fake_score_betas '0.0,0.999'
+  --warp_denoising_step
 )

 # Self-forcing specific arguments
@@ -136,6 +137,7 @@ self_forcing_args=(

 torchrun \
 --nnodes 1 \
+--master_port $MASTER_PORT \
 --nproc_per_node $NUM_GPUS \
    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
    "${parallel_args[@]}" \
@@ -22,8 +22,10 @@ class SelfForcingFlowMatchSchedulerOutput(BaseOutput):
    prev_sample: torch.FloatTensor

 class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
-
+    
+    config_name = "scheduler_config.json"
    order = 1
+    @register_to_config
    def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False, training=False):
        self.num_train_timesteps = num_train_timesteps
        self.shift = shift
@@ -28,10 +28,6 @@ class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
    ]

-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
-            shift=fastvideo_args.pipeline_config.flow_shift)
-
    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
        """Set up pipeline stages with proper dependency injection."""

@@ -22,6 +22,7 @@ from fastvideo.training.training_utils import (
    save_distillation_checkpoint,
 )
 from fastvideo.models.utils import pred_noise_to_pred_video
+from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import SelfForcingFlowMatchScheduler
 from fastvideo.distributed import get_world_group
 import torch.distributed as dist
 import numpy as np
@@ -50,6 +51,9 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
        logger.info("Initializing self-forcing distillation pipeline...")

        super().initialize_training_pipeline(training_args)
+        self.noise_scheduler = SelfForcingFlowMatchScheduler(
+            num_inference_steps=1000, shift=5.0, sigma_min=0.0, extra_one_step=True, training=True
+        )
        self.dfake_gen_update_ratio = getattr(training_args, 'dfake_gen_update_ratio', 5)

        # Self-forcing specific properties
@@ -158,7 +162,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
            noise = noise[:, self.rank_in_sp_group, :, :, :, :]
        noisy_latent = self.noise_scheduler.add_noise(latents.flatten(0, 1),
                                                      noise.flatten(0, 1),
-                                                      torch.tensor([timestep], device=noise.device))
+                                                      timestep * torch.ones([latents.shape[0] * latents.shape[1]], device=noise.device, dtype=torch.long))

        # Step 4: Build input kwargs with KV cache support
        training_batch = self._build_distill_input_kwargs(
@@ -25,11 +25,6 @@ class WanSelfForcingDistillationPipeline(SelfForcingDistillationPipeline):
        "fake_score_transformer"
    ]

-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        """Initialize Wan-specific scheduler."""
-        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
-            shift=fastvideo_args.pipeline_config.flow_shift)
-
    def create_training_stages(self, training_args: TrainingArgs):
        """
        May be used in future refactors.