separate critic lr

examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh

@@ -15,10 +15,9 @@ NUM_GPUS=8
 # Training arguments
 training_args=(
   --tracker_project_name "wan_i2v_finetune"
-  --wandb_run_name "wan_i2v_finetune"
-  --output_dir "outputs/wan_i2v_finetune"
+  --output_dir "$DATA_DIR/outputs/wan_i2v_finetune"
   --max_train_steps 2000
-  --train_batch_size 1
+  --train_batch_size 4
   --train_sp_batch_size 1
   --gradient_accumulation_steps 1
   --num_latent_t 8
@@ -30,10 +29,10 @@ training_args=(
 # Parallel arguments
 parallel_args=(
   --num_gpus $NUM_GPUS
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 8
-  --hsdp_shard_dim 1
+  --sp_size 4
+  --tp_size 4
+  --hsdp_replicate_dim 2
+  --hsdp_shard_dim 4
 )
 
 # Model arguments
@@ -53,7 +52,7 @@ validation_args=(
   --log_validation
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
   --validation_steps 100
-  --validation_sampling_steps "50"
+  --validation_sampling_steps "40"
   --validation_guidance_scale "1.0"
 )
 
@@ -62,7 +61,7 @@ optimizer_args=(
   --learning_rate 2e-5
   --mixed_precision "bf16"
   --checkpointing_steps 2000
-  --weight_decay 0.1
+  --weight_decay 1e-4
   --max_grad_norm 1.0
 )
 
@@ -71,14 +70,14 @@ miscellaneous_args=(
   --inference_mode False
   --allow_tf32
   --checkpoints_total_limit 3
-  --training_cfg_rate 0.0
+  --training_cfg_rate 0.1
   --multi_phased_distill_schedule "4000-1"
   --not_apply_cfg_solver
   --dit_precision "fp32"
   --num_euler_timesteps 50
   --ema_start_step 0
+  --enable_gradient_checkpointing_type "full"
 )
-  # --enable_gradient_checkpointing_type "full"
 
 # If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
 torchrun \

fastvideo/v1/configs/sample/registry.py

@@ -6,11 +6,11 @@ from typing import Any
 from fastvideo.v1.configs.sample.hunyuan import (FastHunyuanSamplingParam,
                                                  HunyuanSamplingParam)
 from fastvideo.v1.configs.sample.stepvideo import StepVideoT2VSamplingParam
-from fastvideo.v1.configs.sample.wan import (WanI2V_14B_480P_SamplingParam,
+from fastvideo.v1.configs.sample.wan import (Wan2_1_Fun_1_3B_InP_SamplingParam,
+                                             WanI2V_14B_480P_SamplingParam,
                                              WanI2V_14B_720P_SamplingParam,
                                              WanT2V_1_3B_SamplingParam,
-                                             WanT2V_14B_SamplingParam,
-                                             Wan2_1_Fun_1_3B_InP_SamplingParam)
+                                             WanT2V_14B_SamplingParam)
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import (maybe_download_model_index,
                                 verify_model_config_and_directory)
@@ -24,7 +24,8 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam,
     "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam,
     "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam,
-    "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": Wan2_1_Fun_1_3B_InP_SamplingParam,
+    "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers":
+    Wan2_1_Fun_1_3B_InP_SamplingParam,
     "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam,
     # Add other specific weight variants
 }

fastvideo/v1/fastvideo_args.py

@@ -428,6 +428,7 @@ class TrainingArgs(FastVideoArgs):
     learning_rate: float = 0.0
     scale_lr: bool = False
     lr_scheduler: str = "constant"
+    lr_step_rules: str | None = None
     lr_warmup_steps: int = 0
     max_grad_norm: float = 0.0
     enable_gradient_checkpointing_type: str | None = None
@@ -466,15 +467,14 @@ class TrainingArgs(FastVideoArgs):
     
     # distillation args
     student_critic_update_ratio: int = 5
+    critic_learning_rate: float = 1e-5
+    critic_lr_scheduler: str = "constant"
+    critic_lr_step_rules: str | None = None
     min_step_ratio: float = 0.2
     max_step_ratio: float = 0.98
     teacher_guidance_scale: float = 3.5
-    
-    # I2V frame weighting for DMD loss
-    i2v_frame_weighting: bool = False  # Enable frame-wise weighting for I2V
-    i2v_first_frame_weight: float = 0.1  # Weight for first frame (0.1 = 10% weight)
-    i2v_weighting_scheme: str = "temporal_consistency"  # Weighting scheme: "first_frame_only", "linear_increase", "exponential_increase", "motion_aware", "temporal_consistency", "first_frame_masking"
-    i2v_temporal_scale_factor: float = 2.0  # Scale factor for temporal weighting schemes
+    simulate_student_forward: bool = False
+    num_teacher_noisy_ground_truth_steps: int = 0
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
@@ -634,6 +634,9 @@ class TrainingArgs(FastVideoArgs):
                             type=str,
                             default="constant",
                             help="Learning rate scheduler type")
+        parser.add_argument("--lr-step-rules",
+                            type=str,
+                            help="Learning rate step rules")
         parser.add_argument("--lr-warmup-steps",
                             type=int,
                             default=10,
@@ -747,6 +750,17 @@ class TrainingArgs(FastVideoArgs):
             type=int,
             default=TrainingArgs.student_critic_update_ratio,
             help="Ratio of student updates to critic updates.")
+        parser.add_argument("--critic-learning-rate",
+            type=float,
+            default=TrainingArgs.critic_learning_rate,
+            help="Learning rate for critic")
+        parser.add_argument("--critic-lr-scheduler",
+            type=str,
+            default=TrainingArgs.critic_lr_scheduler,
+            help="Learning rate scheduler type for critic")
+        parser.add_argument("--critic-lr-step-rules",
+            type=str,
+            help="Learning rate step rules for critic")
         parser.add_argument("--min-step-ratio",
             type=float,
             default=TrainingArgs.min_step_ratio,
@@ -759,25 +773,14 @@ class TrainingArgs(FastVideoArgs):
             type=float,
             default=TrainingArgs.teacher_guidance_scale,
             help="Teacher guidance scale")
-        
-        # I2V frame weighting for DMD loss
-        parser.add_argument("--i2v-frame-weighting",
+        parser.add_argument("--simulate-student-forward",
             action=StoreBoolean,
-            default=TrainingArgs.i2v_frame_weighting,
-            help="Enable frame-wise weighting for I2V distillation to discount first frame")
-        parser.add_argument("--i2v-first-frame-weight",
-            type=float,
-            default=TrainingArgs.i2v_first_frame_weight,
-            help="Weight for first frame in I2V distillation (0.1 = 10% weight)")
-        parser.add_argument("--i2v-weighting-scheme",
-            type=str,
-            default=TrainingArgs.i2v_weighting_scheme,
-            choices=["first_frame_only", "linear_increase", "exponential_increase", "motion_aware", "temporal_consistency", "first_frame_masking"],
-            help="Weighting scheme for I2V frame weighting")
-        parser.add_argument("--i2v-temporal-scale-factor",
-            type=float,
-            default=TrainingArgs.i2v_temporal_scale_factor,
-            help="Scale factor for temporal weighting schemes in I2V")
+            default=TrainingArgs.simulate_student_forward,
+            help="Whether to simulate student forward")
+        parser.add_argument("--num-teacher-noisy-ground-truth-steps",
+            type=int,
+            default=TrainingArgs.num_teacher_noisy_ground_truth_steps,
+            help="Number of steps to use noisy ground truth for teacher")
 
         return parser
 

fastvideo/v1/training/distillation_pipeline.py

@@ -77,11 +77,8 @@ class DistillationPipeline(TrainingPipeline):
         
         self.timestep_shift = self.training_args.pipeline_config.flow_shift
         # self.noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=self.timestep_shift)
-        # self.noise_scheduler = FlowMatchScheduler(
-        #     shift=8.0, sigma_min=0.0, extra_one_step=True
-        # )
         self.noise_scheduler = FlowMatchScheduler(
-            shift=3.0, sigma_min=0.0, extra_one_step=True
+            shift=8.0, sigma_min=0.0, extra_one_step=True
         )
         self.noise_scheduler.set_timesteps(1000, training=True)
         
@@ -106,15 +103,18 @@ class DistillationPipeline(TrainingPipeline):
         critic_params = list(filter(lambda p: p.requires_grad, self.critic_transformer.parameters()))
         self.critic_transformer_optimizer = torch.optim.AdamW(
             critic_params,
-            lr=training_args.learning_rate,
-            # lr=5e-4,
+            lr=training_args.critic_learning_rate,
             betas=(0.9, 0.999),
             weight_decay=training_args.weight_decay,
             eps=1e-8,
         )
         
+        if training_args.critic_lr_scheduler == "piecewise_constant":
+            assert training_args.critic_lr_step_rules is not None, "critic lr step rules is required when using piecewise_constant lr scheduler"
+        
         self.critic_lr_scheduler = get_scheduler(
-            training_args.lr_scheduler,
+            training_args.critic_lr_scheduler,
+            step_rules=training_args.critic_lr_step_rules,
             optimizer=self.critic_transformer_optimizer,
             num_warmup_steps=training_args.lr_warmup_steps * self.world_size,
             num_training_steps=training_args.max_train_steps * self.world_size,
@@ -190,6 +190,56 @@ class DistillationPipeline(TrainingPipeline):
             raise NotImplementedError("Unsupported model type {}".format(type))
 
     def _student_forward(self, training_batch: TrainingBatch) -> torch.Tensor:
+        """Forward pass through student transformer and compute student losses."""
+        latents = training_batch.latents
+        dtype = latents.dtype
+        simulated_noisy_input = []
+        for timestep in self.denoising_step_list:
+            # Use cross-codebase generator for reproducible noise generation
+            noise = torch.randn(
+                self.video_latent_shape, device=self.device, dtype=dtype)
+
+            noisy_timestep = timestep * torch.ones(
+                self.video_latent_shape[:2], device=self.device, dtype=torch.long)
+
+            if timestep != 0:
+                noisy_video = self.noise_scheduler.add_noise(
+                    latents.flatten(0, 1),
+                    noise.flatten(0, 1),
+                    noisy_timestep.flatten(0, 1)
+                ).unflatten(0, self.video_latent_shape[:2])
+            else:
+                noisy_video = latents
+
+            simulated_noisy_input.append(noisy_video)
+
+        simulated_noisy_input = torch.stack(simulated_noisy_input, dim=1)
+
+        # Step 2: Randomly sample a timestep and pick the corresponding input
+        # Use cross-codebase generator for reproducible index generation
+        index = torch.randint(0, len(self.denoising_step_list), [
+                              self.video_latent_shape[0], self.video_latent_shape[1]], device=self.device, dtype=torch.long)
+
+        index = self._process_timestep(index, type=self.distill_task_type)
+
+        # select the corresponding timestep's noisy input from the stacked tensor [B, T, F, C, H, W]
+
+        noisy_input = torch.gather(
+            simulated_noisy_input, dim=1,
+            index=index.reshape(index.shape[0], 1, index.shape[1], 1, 1, 1).expand(
+                -1, -1, -1, *self.video_latent_shape[2:])
+        ).squeeze(1)
+
+        timestep = self.denoising_step_list[index]
+
+        training_batch = self._build_input_kwargs(noisy_input, timestep, training_batch.conditional_dict, training_batch)
+        pred_video = self.student_transformer(training_batch, timestep)
+
+        pred_video = pred_video.type_as(noisy_input)
+
+        return pred_video, timestep.float().detach()
+    
+    def _multi_step_simulation_student_forward(self, training_batch: TrainingBatch) -> torch.Tensor:
         """Forward pass through student transformer matching inference procedure."""
         from fastvideo.v1.training.training_utils import DiffusionWrapper
         
@@ -254,31 +304,44 @@ class DistillationPipeline(TrainingPipeline):
         
         return pred_video, target_timestep.float().detach()
 
+
     def _compute_kl_grad(
-        self, noisy_ground_truth: torch.Tensor,
+        self, 
         noisy_video: torch.Tensor,
         estimated_clean_video: torch.Tensor,
+        noise: torch.Tensor,
         timestep: torch.Tensor,
         training_batch: TrainingBatch,
         normalization: bool = True
     ) -> Tuple[torch.Tensor, dict]:
+        assert self.training_args is not None
+
         # critic_transformer forward
         training_batch = self._build_input_kwargs(noisy_video, timestep, training_batch.conditional_dict, training_batch)
         pred_fake_video = self.critic_transformer(training_batch, timestep)
         
-        # teacher_transformer cond forward
-
-        if self.current_trainstep > 600: 
+        if self.current_trainstep > self.training_args.num_teacher_noisy_ground_truth_steps: 
             teacher_noisy_input = noisy_video
+            teacher_timestep = timestep
         else:
+            teacher_timestep = timestep
+            logger.info(f"Using noisy ground truth for teacher with timestep {teacher_timestep}")
+            batch_size, num_frame = self.video_latent_shape[:2]
+            noisy_ground_truth = self.noise_scheduler.add_noise(
+                training_batch.latents.flatten(0, 1),
+                noise.flatten(0, 1),
+                teacher_timestep.flatten(0, 1)
+            ).detach().unflatten(0, (batch_size, num_frame))
+
             teacher_noisy_input = noisy_ground_truth
 
-        training_batch = self._build_input_kwargs(teacher_noisy_input, timestep, training_batch.conditional_dict, training_batch)
-        pred_real_video_cond = self.teacher_transformer(training_batch, timestep)
+        # teacher_transformer cond forward
+        training_batch = self._build_input_kwargs(teacher_noisy_input, teacher_timestep, training_batch.conditional_dict, training_batch)
+        pred_real_video_cond = self.teacher_transformer(training_batch, teacher_timestep)
         
         # teacher_transformer uncond forward
-        training_batch = self._build_input_kwargs(teacher_noisy_input, timestep, training_batch.unconditional_dict, training_batch)
-        pred_real_video_uncond = self.teacher_transformer(training_batch, timestep)
+        training_batch = self._build_input_kwargs(teacher_noisy_input, teacher_timestep, training_batch.unconditional_dict, training_batch)
+        pred_real_video_uncond = self.teacher_transformer(training_batch, teacher_timestep)
         
         pred_real_video = pred_real_video_cond + (
             pred_real_video_cond - pred_real_video_uncond
@@ -333,145 +396,32 @@ class DistillationPipeline(TrainingPipeline):
                 noise.flatten(0, 1),
                 timestep.flatten(0, 1)
             ).detach().unflatten(0, (batch_size, num_frame))
-        
-            noisy_ground_truth = self.noise_scheduler.add_noise(
-                training_batch.latents.flatten(0, 1),
-                noise.flatten(0, 1),
-                timestep.flatten(0, 1)
-            ).detach().unflatten(0, (batch_size, num_frame))
 
             grad, dmd_log_dict = self._compute_kl_grad(
-                noisy_ground_truth=noisy_ground_truth,
                 noisy_video=noisy_latent,
                 estimated_clean_video=original_latent,
+                noise=noise,
                 timestep=timestep,
                 training_batch=training_batch
             )
 
-        # Apply frame-wise weighting for I2V to discount the first frame
-        if hasattr(self.training_args, 'i2v_frame_weighting') and self.training_args.i2v_frame_weighting:
-            weighting_scheme = getattr(self.training_args, 'i2v_weighting_scheme', 'linear_increase')
-            first_frame_weight = getattr(self.training_args, 'i2v_first_frame_weight', 0.1)
-            temporal_scale_factor = getattr(self.training_args, 'i2v_temporal_scale_factor', 2.0)
-            
-            logger.info(f"Applying I2V frame weighting - scheme: {weighting_scheme}, first_frame_weight: {first_frame_weight}, scale_factor: {temporal_scale_factor}")
-            
-            # Create frame weights: [B, C, T, H, W]
-            frame_weights = torch.ones_like(original_latent)
-            batch_size, channels, num_frames, height, width = original_latent.shape
-            
-            if weighting_scheme == "first_frame_only":
-                # Simple: just discount first frame
-                frame_weights[:, :, 0] = first_frame_weight
-                
-            elif weighting_scheme == "linear_increase":
-                # Linear increase: first frame gets first_frame_weight, last frame gets temporal_scale_factor
-                for t in range(num_frames):
-                    if t == 0:
-                        weight = first_frame_weight
-                    else:
-                        # Linear interpolation from first_frame_weight to temporal_scale_factor
-                        weight = first_frame_weight + (temporal_scale_factor - first_frame_weight) * t / (num_frames - 1)
-                    frame_weights[:, :, t] = weight
-                    
-            elif weighting_scheme == "exponential_increase":
-                # Exponential increase: emphasizes later frames more aggressively
-                for t in range(num_frames):
-                    if t == 0:
-                        weight = first_frame_weight
-                    else:
-                        # Exponential scaling: later frames get exponentially higher weights
-                        progress = t / (num_frames - 1)  # 0 to 1
-                        weight = first_frame_weight + (temporal_scale_factor - first_frame_weight) * (progress ** 2)
-                    frame_weights[:, :, t] = weight
-                    
-            elif weighting_scheme == "motion_aware":
-                # Motion-aware: weight based on difference from first frame
-                with torch.no_grad():
-                    first_frame = original_latent[:, :, 0:1]  # [B, C, 1, H, W]
-                    differences = torch.abs(original_latent - first_frame)  # [B, C, T, H, W]
-                    # Normalize differences to [0, 1] range per sample
-                    max_diff = differences.view(batch_size, -1).max(dim=1)[0].view(batch_size, 1, 1, 1, 1)
-                    normalized_diff = differences / (max_diff + 1e-8)
-                    # Scale weights: first frame gets first_frame_weight, high-motion frames get up to temporal_scale_factor
-                    frame_weights = first_frame_weight + (temporal_scale_factor - first_frame_weight) * normalized_diff
-                    # Ensure first frame always gets the specified weight
-                    frame_weights[:, :, 0] = first_frame_weight
-                    
-            elif weighting_scheme == "temporal_consistency":
-                # New approach: Add temporal consistency loss instead of just reweighting
-                # Standard DMD loss first
-                dmd_loss = 0.5 * F.mse_loss(original_latent.double(), 
-                                           (original_latent.double() - grad.double()).detach(), 
-                                           reduction="mean")
-                
-                # Add temporal consistency loss that encourages smooth transitions
-                if num_frames > 1:
-                    # Temporal differences in prediction
-                    pred_diff = original_latent[:, :, 1:] - original_latent[:, :, :-1]
-                    target_diff = (original_latent.double() - grad.double()).detach()[:, :, 1:] - (original_latent.double() - grad.double()).detach()[:, :, :-1]
-                    
-                    temporal_consistency_loss = F.mse_loss(pred_diff.double(), target_diff, reduction="mean")
-                    dmd_loss = dmd_loss + temporal_scale_factor * temporal_consistency_loss
-                    
-                # Log temporal consistency for monitoring
-                logger.info(f"DMD loss: {dmd_loss:.6f}, temporal_consistency_scale: {temporal_scale_factor}")
-                return dmd_loss, dmd_log_dict
-                
-            elif weighting_scheme == "first_frame_masking":
-                # Alternative approach: Apply DMD loss only to non-first frames
-                if num_frames > 1:
-                    # Exclude first frame completely from DMD loss computation
-                    original_latent_masked = original_latent[:, :, 1:]  # Remove first frame
-                    grad_masked = grad[:, :, 1:]  # Remove first frame
-                    
-                    dmd_loss = 0.5 * F.mse_loss(original_latent_masked.double(),
-                                               (original_latent_masked.double() - grad_masked.double()).detach(),
-                                               reduction="mean")
-                    
-                    # Optional: Add small first-frame loss with very low weight
-                    if first_frame_weight > 0:
-                        first_frame_loss = 0.5 * F.mse_loss(original_latent[:, :, 0:1].double(),
-                                                           (original_latent[:, :, 0:1].double() - grad[:, :, 0:1].double()).detach(),
-                                                           reduction="mean")
-                        dmd_loss = dmd_loss + first_frame_weight * first_frame_loss
-                        
-                    logger.info(f"First-frame masked DMD loss: {dmd_loss:.6f}")
-                    return dmd_loss, dmd_log_dict
-                else:
-                    # Fallback for single frame
-                    dmd_loss = 0.5 * F.mse_loss(original_latent.double(),
-                                               (original_latent.double() - grad.double()).detach(),
-                                               reduction="mean")
-                    return dmd_loss, dmd_log_dict
-            
-            # For traditional weighting schemes (not temporal_consistency or first_frame_masking)
-            # Compute weighted MSE loss
-            mse_loss = (original_latent.double() - (original_latent.double() - grad.double()).detach())**2
-            weighted_loss = mse_loss * frame_weights.double()
-            
-            # Normalize by the sum of weights to maintain scale
-            total_weight = frame_weights.sum()
-            dmd_loss = 0.5 * weighted_loss.sum() / total_weight
-            
-            # Log weighting statistics for debugging
-            avg_weights_per_frame = frame_weights.mean(dim=[0, 1, 3, 4])  # Average weight per frame
-            logger.info(f"Frame weights (avg per frame): {avg_weights_per_frame.tolist()}")
-        else:
-            # Original unweighted loss
-            dmd_loss = 0.5 * F.mse_loss(original_latent.double(
-            ), (original_latent.double() - grad.double()).detach(), reduction="mean")
+        dmd_loss = 0.5 * F.mse_loss(original_latent.double(
+        ), (original_latent.double() - grad.double()).detach(), reduction="mean")
         
         return dmd_loss, dmd_log_dict
 
     def _student_forward_and_compute_dmd_loss(self, training_batch: TrainingBatch) -> Tuple[TrainingBatch, torch.Tensor, dict]:
         """Forward pass through student transformer and compute student losses."""
+        assert self.training_args is not None
         assert training_batch.conditional_dict is not None
         assert training_batch.unconditional_dict is not None
         assert training_batch.latents is not None
         with set_forward_context(
                 current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata_vsa):
-            pred_video, timestep_dmd = self._student_forward(training_batch)
+            if self.training_args.simulate_student_forward:
+                pred_video, timestep_dmd = self._multi_step_simulation_student_forward(training_batch)
+            else:
+                pred_video, timestep_dmd = self._student_forward(training_batch)
 
         with set_forward_context(
                 current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata):
@@ -485,6 +435,7 @@ class DistillationPipeline(TrainingPipeline):
         return training_batch, dmd_loss, dmd_log_dict
 
     def _critic_forward_and_compute_loss(self, training_batch: TrainingBatch) -> Tuple[TrainingBatch, torch.Tensor, dict]:
+        assert self.training_args is not None
         assert training_batch.conditional_dict is not None
         assert training_batch.unconditional_dict is not None
         assert training_batch.latents is not None
@@ -492,7 +443,10 @@ class DistillationPipeline(TrainingPipeline):
         with torch.no_grad():
             with set_forward_context(
                 current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata_vsa):
-                generated_video, timestep_gen = self._student_forward(training_batch)
+                if self.training_args.simulate_student_forward:
+                    generated_video, timestep_gen = self._multi_step_simulation_student_forward(training_batch)
+                else:
+                    generated_video, timestep_gen = self._student_forward(training_batch)
 
         critic_timestep = torch.randint(
             0,
@@ -945,7 +899,7 @@ class DistillationPipeline(TrainingPipeline):
             disable=self.local_rank > 0,
         )
         
-        for step in range(self.init_steps,
+        for step in range(self.init_steps+1,
                           self.training_args.max_train_steps + 1):
             start_time = time.perf_counter()
             current_vsa_sparsity = self.training_args.VSA_sparsity if vsa_available else 0.0

fastvideo/v1/training/training_pipeline.py

@@ -109,8 +109,13 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         self.init_steps = 0
         logger.info("optimizer: %s", self.optimizer)
 
+
+        if training_args.lr_scheduler == "piecewise_constant":
+            assert training_args.lr_step_rules is not None, "lr step rules is required when using piecewise_constant lr scheduler"
+
         self.lr_scheduler = get_scheduler(
             training_args.lr_scheduler,
+            step_rules=training_args.lr_step_rules,
             optimizer=self.optimizer,
             num_warmup_steps=training_args.lr_warmup_steps * self.world_size,
             num_training_steps=training_args.max_train_steps * self.world_size,

fastvideo/v1/training/training_utils.py

@@ -3,26 +3,26 @@ import json
 import math
 import os
 import time
-from abc import ABC
-from typing import Any
+from typing import Any, Dict
 
-import numpy as np
 import torch
 import torch.distributed as dist
 import torch.distributed.checkpoint as dcp
+from torchvision.utils import make_grid
 from einops import rearrange
 from safetensors.torch import save_file
-from torchvision.utils import make_grid
-
 import wandb
+import numpy as np
+
 from fastvideo.v1.distributed.parallel_state import (get_sp_parallel_rank,
                                                      get_sp_world_size)
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines.pipeline_batch_info import TrainingBatch
 from fastvideo.v1.training.checkpointing_utils import (ModelWrapper,
                                                        OptimizerWrapper,
                                                        RandomStateWrapper,
                                                        SchedulerWrapper)
+from fastvideo.v1.pipelines.pipeline_batch_info import TrainingBatch
+from abc import ABC
 
 logger = init_logger(__name__)
 
@@ -556,49 +556,36 @@ def convert_custom_format_to_diffusers_format(state_dict: dict[str, Any],
 
     return new_state_dict
 
-
-def prepare_for_saving(tensor: torch.Tensor,
-                       fps: int = 16,
-                       caption: str | None = None) -> wandb.Image | wandb.Video:
+def prepare_for_saving(tensor: torch.Tensor, fps: int = 16, caption: str | None = None) -> wandb.Image | wandb.Video:
     if tensor.ndim == 4:
         # Assuming it's an image and has shape [batch_size, 3, height, width]
         tensor = make_grid(tensor, 4, padding=0, normalize=False)
-        return wandb.Image((tensor * 255).numpy().astype(np.uint8),
-                           caption=caption)
+        return wandb.Image((tensor * 255).numpy().astype(np.uint8), caption=caption)
     elif tensor.ndim == 5:
         # Assuming it's a video and has shape [batch_size, num_frames, 3, height, width]
-        return wandb.Video((tensor * 255).numpy().astype(np.uint8),
-                           fps=fps,
-                           format="webm",
-                           caption=caption)
+        return wandb.Video((tensor * 255).numpy().astype(np.uint8), fps=fps, format="webm", caption=caption)
     else:
-        raise ValueError(
-            "Unsupported tensor shape for saving. Expected 4D (image) or 5D (video) tensor."
-        )
-
-
-class DiffusionWrapper(torch.nn.Module, ABC):
+        raise ValueError("Unsupported tensor shape for saving. Expected 4D (image) or 5D (video) tensor.")
 
+class DiffusionWrapper(torch.nn.Module, ABC): 
     def __init__(self, transformer, scheduler):
         super().__init__()
         self.model = transformer
         self.scheduler = scheduler
-
+        
     def forward(self, training_batch: TrainingBatch, timestep: torch.Tensor):
-        pred_noise = self.model(**training_batch.input_kwargs).permute(
-            0, 2, 1, 3, 4)
+        pred_noise = self.model(**training_batch.input_kwargs).permute(0, 2, 1, 3, 4)
         pred_video = self._convert_flow_pred_to_x0(
             flow_pred=pred_noise.flatten(0, 1),
             xt=training_batch.noise_latents.flatten(0, 1),
             timestep=timestep.flatten(0, 1),
-            scheduler=self.scheduler).unflatten(0, pred_noise.shape[:2])
+            scheduler=self.scheduler
+        ).unflatten(0, pred_noise.shape[:2])
 
         return pred_video
-
+    
     @staticmethod
-    def _convert_x0_to_flow_pred(x0_pred: torch.Tensor, xt: torch.Tensor,
-                                 timestep: torch.Tensor,
-                                 scheduler: Any) -> torch.Tensor:
+    def _convert_x0_to_flow_pred(x0_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor, scheduler: Any) -> torch.Tensor:
         """
         Convert x0 prediction to flow matching's prediction.
         x0_pred: the x0 prediction with shape [B, C, H, W]
@@ -610,8 +597,10 @@ class DiffusionWrapper(torch.nn.Module, ABC):
         # use higher precision for calculations
         original_dtype = x0_pred.dtype
         x0_pred, xt, sigmas, timesteps = map(
-            lambda x: x.double().to(x0_pred.device),
-            [x0_pred, xt, scheduler.sigmas, scheduler.timesteps])
+            lambda x: x.double().to(x0_pred.device), [x0_pred, xt,
+                                                      scheduler.sigmas,
+                                                      scheduler.timesteps] 
+        )
         timestep_id = torch.argmin(
             (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
         sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
@@ -619,9 +608,7 @@ class DiffusionWrapper(torch.nn.Module, ABC):
         return flow_pred.to(original_dtype)
 
     @staticmethod
-    def _convert_flow_pred_to_x0(flow_pred: torch.Tensor, xt: torch.Tensor,
-                                 timestep: torch.Tensor,
-                                 scheduler: Any) -> torch.Tensor:
+    def _convert_flow_pred_to_x0(flow_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor, scheduler: Any) -> torch.Tensor:
         """
         Convert flow matching's prediction to x0 prediction.
         flow_pred: the prediction with shape [B, C, H, W]
@@ -636,10 +623,12 @@ class DiffusionWrapper(torch.nn.Module, ABC):
         # use higher precision for calculations
         original_dtype = flow_pred.dtype
         flow_pred, xt, sigmas, timesteps = map(
-            lambda x: x.double().to(flow_pred.device),
-            [flow_pred, xt, scheduler.sigmas, scheduler.timesteps])
+            lambda x: x.double().to(flow_pred.device), [flow_pred, xt,
+                                                        scheduler.sigmas,
+                                                        scheduler.timesteps]
+        )
         timestep_id = torch.argmin(
             (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
         sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
         x0_pred = xt - sigma_t * flow_pred
-        return x0_pred.to(original_dtype)
+        return x0_pred.to(original_dtype)

finetune_dmd.sh

@@ -1,78 +0,0 @@
-#!/bin/bash
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export WANDB_API_KEY='73190d8c0de18a14eb3444e222f9432d247d1e30'
-export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-DATA_DIR=data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/
-# DATA_DIR=/mnt/sharefs/users/hao.zhang/Vchitect-2M/Vchitect-2M-laten-93x512x512/val/
-VALIDATION_DIR=examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json
-NUM_GPUS=8
-export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
-# export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-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 NODE_RANK=$SLURM_PROCID
-# nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
-# export MASTER_ADDR=${nodes[0]}
-# export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
-export TOKENIZERS_PARALLELISM=false
-# echo "MASTER_ADDR: $MASTER_ADDR"
-# echo "NODE_RANK: $NODE_RANK"
-
-# 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_distillation_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_dataset_file  "$VALIDATION_DIR" \
-    --train_batch_size 1 \
-    --num_latent_t 8 \
-    --sp_size 1 \
-    --tp_size 1 \
-    --num_gpus 8 \
-    --hsdp_replicate_dim 8 \
-    --hsdp-shard-dim 1 \
-    --train_sp_batch_size 1 \
-    --dataloader_num_workers 0 \
-    --gradient_accumulation_steps 1 \
-    --max_train_steps 30000 \
-    --learning_rate 1e-5 \
-    --mixed_precision "bf16" \
-    --checkpointing_steps 500 \
-    --validation_steps 100 \
-    --validation_sampling_steps "3" \
-    --log_validation \
-    --checkpoints_total_limit 3 \
-    --allow_tf32 \
-    --ema_start_step 0 \
-    --training_cfg_rate 0.0 \
-    --output_dir "outputs_dmd_train/wan_finetune_1e5" \
-    --tracker_project_name Wan_distillation \
-    --wandb_run_name "crush_smol_dmd_test" \
-    --num_height 448 \
-    --num_width 832 \
-    --num_frames 61 \
-    --flow_shift 8 \
-    --validation_guidance_scale "1.0" \
-    --master_weight_type "fp32" \
-    --dit_precision "fp32" \
-    --vae_precision "bf16" \
-    --weight_decay 0.01 \
-    --max_grad_norm 1.0 \
-    --student_critic_update_ratio 5 \
-    --denoising_step_list '1000,757,522' \
-    --min_step_ratio 0.02 \
-    --max_step_ratio 0.98 \
-    --seed 1000 \
-    --teacher_guidance_scale 3.5 \
-    --enable_gradient_checkpointing_type "full" 

finetune_dmd_i2v.sh

@@ -1,79 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export WANDB_API_KEY='73190d8c0de18a14eb3444e222f9432d247d1e30'
-export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-# DATA_DIR=data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/
-DATA_DIR=/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn/latents_i2v/train/
-# VALIDATION_DIR=examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json
-# VALIDATION_DIR=/mnt/weka/home/hao.zhang/wl/FastVideo/data/mixkit/validation.json
-VALIDATION_DIR=/mnt/sharefs/users/hao.zhang/wl/mixkit/validation.json
-NUM_GPUS=8
-export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
-# export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
-# export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6
-# IP=[MASTER NODE IP]
-
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-# different cache dir for different processes
-export MASTER_PORT=29501
-export TOKENIZERS_PARALLELISM=false
-
-# 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_i2v_distillation_pipeline.py \
-    --model_path weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers  \
-    --inference_mode False\
-    --pretrained_model_name_or_path weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers  \
-    --cache_dir "/home/ray/.cache" \
-    --data_path "$DATA_DIR" \
-    --validation_dataset_file  "$VALIDATION_DIR" \
-    --train_batch_size 1 \
-    --num_latent_t 8 \
-    --sp_size 1 \
-    --tp_size 1 \
-    --num_gpus 8 \
-    --hsdp_replicate_dim 8 \
-    --hsdp-shard-dim 1 \
-    --train_sp_batch_size 1 \
-    --dataloader_num_workers 0 \
-    --gradient_accumulation_steps 1 \
-    --max_train_steps 6000 \
-    --learning_rate 1e-5 \
-    --mixed_precision "bf16" \
-    --checkpointing_steps 1000 \
-    --validation_steps 50 \
-    --validation_sampling_steps "3" \
-    --log_validation \
-    --checkpoints_total_limit 3 \
-    --allow_tf32 \
-    --ema_start_step 0 \
-    --training_cfg_rate 0.0 \
-    --output_dir "outputs_dmd_train_i2v/wan_i2v_finetune_9e6" \
-    --tracker_project_name Wan_distillation \
-    --wandb_run_name "backward_sim_synth" \
-    --num_height 480 \
-    --num_width 832 \
-    --num_frames 61 \
-    --flow_shift 3 \
-    --validation_guidance_scale "1.0" \
-    --master_weight_type "fp32" \
-    --dit_precision "fp32" \
-    --vae_precision "bf16" \
-    --weight_decay 0.01 \
-    --max_grad_norm 1.0 \
-    --student_critic_update_ratio 5 \
-    --denoising_step_list '1000,757,522' \
-    --min_step_ratio 0.02 \
-    --max_step_ratio 0.98 \
-    --seed 1000 \
-    --teacher_guidance_scale 3.5
-
-    # --i2v_frame_weighting \
-    #--denoising_step_list '1000,757,522' \
-    # --i2v_weighting_scheme "first_frame_only" \
-    # --i2v_temporal_scale_factor 1.0 \
-    # --i2v_first_frame_weight 0.01 \

scripts/finetune_wan_distillation.sh

@@ -1,13 +1,11 @@
 export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
+export WANDB_MODE=offline
 export WANDB_API_KEY='73190d8c0de18a14eb3444e222f9432d247d1e30'
-export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
 export TRITON_CACHE_DIR=/tmp/triton_cache
-# DATA_DIR=/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn/latents_i2v/val/
-DATA_DIR=data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/
+DATA_DIR=/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn/latents_i2v/val/
 # DATA_DIR=/mnt/sharefs/users/hao.zhang/Vchitect-2M/Vchitect-2M-laten-93x512x512/val/
-VALIDATION_DIR=examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json
+VALIDATION_DIR=/mnt/sharefs/users/hao.zhang/Vchitect-2M/mixkit/validation_8.json
 NUM_GPUS=8
 export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
 # export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
@@ -48,7 +46,6 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
     --training_cfg_rate 0.0 \
     --output_dir "outputs_dmd/wan_finetune" \
     --tracker_project_name Wan_distillation \
-    --wandb_run_name "crush_smol_dmd_test" \
     --num_height 448 \
     --num_width 832 \
     --num_frames  29 \