diff --git a/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh b/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh index 17e75d4bd..a50741ec7 100644 --- a/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh +++ b/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh @@ -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[@]}" \ diff --git a/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py b/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py index 1d76c14ae..651518292 100644 --- a/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py +++ b/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py @@ -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 diff --git a/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py b/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py index 42a8baa9b..8e812228e 100644 --- a/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py +++ b/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py @@ -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.""" diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index 745177cb2..2f39b57b6 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -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( diff --git a/fastvideo/training/wan_self_forcing_distillation_pipeline.py b/fastvideo/training/wan_self_forcing_distillation_pipeline.py index 6b3adf9b6..c1edf60fb 100644 --- a/fastvideo/training/wan_self_forcing_distillation_pipeline.py +++ b/fastvideo/training/wan_self_forcing_distillation_pipeline.py @@ -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.