diff --git a/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh b/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh index ee942605b..0efd6fe63 100644 --- a/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh +++ b/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh @@ -98,6 +98,7 @@ dmd_args=( torchrun \ --nnodes 1 \ --nproc_per_node $NUM_GPUS \ +--master_port $MASTER_PORT \ fastvideo/training/wan_distillation_pipeline.py \ "${parallel_args[@]}" \ "${model_args[@]}" \ diff --git a/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B_lora.sh b/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B_lora.sh new file mode 100644 index 000000000..0f0db0baf --- /dev/null +++ b/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B_lora.sh @@ -0,0 +1,112 @@ +#!/bin/bash + +# Basic Info +export WANDB_MODE="online" +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +export MASTER_PORT=29501 +export TOKENIZERS_PARALLELISM=false +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +# Configs +NUM_GPUS=1 +MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers" +DATA_DIR="data/crush-smol_processed_ti2v/combined_parquet_dataset/" +VALIDATION_DATASET_FILE="examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json" +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name wan_t2v_distill_dmd_VSA + --output_dir="checkpoints/wan_t2v_finetune" + --max_train_steps=4000 + --train_batch_size=1 + --train_sp_batch_size 1 + --gradient_accumulation_steps=1 + --num_latent_t 31 + --num_height 704 + --num_width 1280 + --num_frames 121 + --enable_gradient_checkpointing_type "full" + --training_state_checkpointing_steps=500 + --weight_only_checkpointing_steps=500 + --lora_rank 32 + --lora_training True +) + +# Parallel arguments +parallel_args=( + --num_gpus 1 + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 + --hsdp_shard_dim 1 +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 4 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 200 + --validation_sampling_steps "3" + --validation_guidance_scale "6.0" # not used for dmd inference +) + +# Optimizer arguments +optimizer_args=( + --learning_rate=1e-4 + --mixed_precision="bf16" + --weight_decay 0.01 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.0 + --dit_precision "fp32" + --ema_start_step 0 + --flow_shift 8 + --seed 1000 +) + +# DMD arguments +dmd_args=( + --dmd_denoising_steps '1000,757,522' + --min_timestep_ratio 0.02 + --max_timestep_ratio 0.98 + --generator_update_interval 5 + --real_score_guidance_scale 3.5 + --VSA_sparsity 0.8 +) + +torchrun \ +--nnodes 1 \ +--nproc_per_node $NUM_GPUS \ +--master_port $MASTER_PORT \ + fastvideo/training/wan_distillation_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ + "${dmd_args[@]}" \ No newline at end of file diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index 65d441169..ca53f69c1 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -79,7 +79,7 @@ class ComposedPipelineBase(ABC): for name, module in self.modules.items(): if not isinstance(module, torch.nn.Module): continue - if name == "transformer": + if "transformer" in name: module.requires_grad_(True) else: module.requires_grad_(False) diff --git a/fastvideo/pipelines/lora_pipeline.py b/fastvideo/pipelines/lora_pipeline.py index cfdc95a9d..17e16f377 100644 --- a/fastvideo/pipelines/lora_pipeline.py +++ b/fastvideo/pipelines/lora_pipeline.py @@ -7,7 +7,7 @@ import torch import torch.distributed as dist import torch.nn as nn from safetensors.torch import load_file -from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh from torch.distributed.tensor import DTensor from fastvideo.distributed import get_local_torch_device @@ -32,6 +32,7 @@ class LoRAPipeline(ComposedPipelineBase): cur_adapter_name: str = "" cur_adapter_path: str = "" lora_layers: dict[str, BaseLayerWithLoRA] = {} + lora_layers_critic: dict[str, BaseLayerWithLoRA] = {} fastvideo_args: FastVideoArgs | TrainingArgs exclude_lora_layers: list[str] = [] device: torch.device = get_local_torch_device() @@ -81,6 +82,17 @@ class LoRAPipeline(ComposedPipelineBase): def set_trainable(self) -> None: + def set_lora_grads(lora_layers: dict[str, BaseLayerWithLoRA], + device_mesh: DeviceMesh): + for name, layer in lora_layers.items(): + layer.lora_A.requires_grad_(True) + layer.lora_B.requires_grad_(True) + layer.base_layer.requires_grad_(False) + layer.lora_A = nn.Parameter( + DTensor.from_local(layer.lora_A, device_mesh=device_mesh)) + layer.lora_B = nn.Parameter( + DTensor.from_local(layer.lora_B, device_mesh=device_mesh)) + is_lora_training = self.training_mode and getattr( self.fastvideo_args, "lora_training", False) if not is_lora_training: @@ -88,18 +100,12 @@ class LoRAPipeline(ComposedPipelineBase): return self.modules["transformer"].requires_grad_(False) + if "fake_score_transformer" in self.modules: + self.modules["fake_score_transformer"].requires_grad_(False) device_mesh = init_device_mesh("cuda", (dist.get_world_size(), 1), mesh_dim_names=["fake", "replicate"]) - for name, layer in self.lora_layers.items(): - # Enable grads for lora weights only - # Must convert to DTensor for compatibility with other FSDP modules in grad calculation - layer.lora_A.requires_grad_(True) - layer.lora_B.requires_grad_(True) - layer.base_layer.requires_grad_(False) - layer.lora_A = nn.Parameter( - DTensor.from_local(layer.lora_A, device_mesh=device_mesh)) - layer.lora_B = nn.Parameter( - DTensor.from_local(layer.lora_B, device_mesh=device_mesh)) + set_lora_grads(self.lora_layers, device_mesh) + set_lora_grads(self.lora_layers_critic, device_mesh) def convert_to_lora_layers(self) -> None: """ @@ -131,6 +137,24 @@ class LoRAPipeline(ComposedPipelineBase): converted_count += 1 logger.info("Converted %d layers to LoRA layers", converted_count) + if "fake_score_transformer" in self.modules: + for name, layer in self.modules[ + "fake_score_transformer"].named_modules(): + if not self.is_target_layer(name): + continue + layer = get_lora_layer(layer, + lora_rank=self.lora_rank, + lora_alpha=self.lora_alpha, + training_mode=self.training_mode) + if layer is not None: + self.lora_layers_critic[name] = layer + replace_submodule(self.modules["fake_score_transformer"], + name, layer) + converted_count += 1 + logger.info( + "Converted %d layers to LoRA layers in the critic model", + converted_count) + def set_lora_adapter(self, lora_nickname: str, lora_path: str | None = None): # type: ignore @@ -224,4 +248,4 @@ class LoRAPipeline(ComposedPipelineBase): def unmerge_lora_weights(self) -> None: for name, layer in self.lora_layers.items(): - layer.unmerge_lora_weights() \ No newline at end of file + layer.unmerge_lora_weights() diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index b89489075..4c9ba481c 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -149,7 +149,8 @@ class DenoisingStage(PipelineStage): # Prepare image latents and embeddings for I2V generation image_embeds = batch.image_embeds if len(image_embeds) > 0: - assert torch.isnan(image_embeds[0]).sum() == 0 + assert not torch.isnan( + image_embeds[0]).any(), "image_embeds contains nan" image_embeds = [ image_embed.to(target_dtype) for image_embed in image_embeds ] @@ -186,11 +187,13 @@ class DenoisingStage(PipelineStage): # Get latents and embeddings latents = batch.latents prompt_embeds = batch.prompt_embeds - assert torch.isnan(prompt_embeds[0]).sum() == 0 + assert not torch.isnan( + prompt_embeds[0]).any(), "prompt_embeds contains nan" if batch.do_classifier_free_guidance: neg_prompt_embeds = batch.negative_prompt_embeds assert neg_prompt_embeds is not None - assert torch.isnan(neg_prompt_embeds[0]).sum() == 0 + assert not torch.isnan( + neg_prompt_embeds[0]).any(), "neg_prompt_embeds contains nan" # (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert if fastvideo_args.boundary_ratio is not None: @@ -268,6 +271,9 @@ class DenoisingStage(PipelineStage): latent_model_input = torch.cat( [latent_model_input, batch.image_latent], dim=1).to(target_dtype) + + assert not torch.isnan( + latent_model_input).any(), "latent_model_input contains nan" if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None: timestep = torch.stack([t]).to(get_local_torch_device()) temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten() @@ -280,7 +286,6 @@ class DenoisingStage(PipelineStage): else: t_expand = t.repeat(latent_model_input.shape[0]) - assert torch.isnan(latent_model_input).sum() == 0 latent_model_input = self.scheduler.scale_model_input( latent_model_input, t) @@ -737,7 +742,8 @@ class DmdDenoisingStage(DenoisingStage): video_raw_latent_shape = latents.shape prompt_embeds = batch.prompt_embeds - assert torch.isnan(prompt_embeds[0]).sum() == 0 + assert not torch.isnan( + prompt_embeds[0]).any(), "prompt_embeds contains nan" timesteps = torch.tensor( fastvideo_args.pipeline_config.dmd_denoising_steps, dtype=torch.long, @@ -776,7 +782,8 @@ class DmdDenoisingStage(DenoisingStage): batch.image_latent.permute(0, 2, 1, 3, 4) ], dim=2).to(target_dtype) - assert torch.isnan(latent_model_input).sum() == 0 + assert not torch.isnan( + latent_model_input).any(), "latent_model_input contains nan" # Prepare inputs for transformer t_expand = t.repeat(latent_model_input.shape[0]) diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 47381a91e..632f6a104 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -36,8 +36,9 @@ from fastvideo.training.activation_checkpoint import ( apply_activation_checkpointing) from fastvideo.training.training_pipeline import TrainingPipeline from fastvideo.training.training_utils import ( - clip_grad_norm_while_handling_failing_dtensor_cases, get_scheduler, - load_distillation_checkpoint, save_distillation_checkpoint, shift_timestep) + clip_grad_norm_while_handling_failing_dtensor_cases, count_trainable, + get_scheduler, load_distillation_checkpoint, save_distillation_checkpoint, + shift_timestep) from fastvideo.utils import is_vsa_available, set_random_seed import wandb # isort: skip @@ -68,11 +69,18 @@ class DistillationPipeline(TrainingPipeline): current_trainstep: int video_latent_shape: tuple[int, ...] video_latent_shape_sp: tuple[int, ...] + real_score_transformer: torch.nn.Module + fake_score_transformer: torch.nn.Module def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): raise RuntimeError( "create_pipeline_stages should not be called for training pipeline") + def set_trainable(self) -> None: + super().set_trainable() + self.modules["real_score_transformer"].requires_grad_(False) + self.modules["vae"].requires_grad_(False) + def initialize_training_pipeline(self, training_args: TrainingArgs): """Initialize the distillation training pipeline with multiple models.""" logger.info("Initializing distillation pipeline...") @@ -81,8 +89,6 @@ class DistillationPipeline(TrainingPipeline): self.noise_scheduler = self.get_module("scheduler") self.vae = self.get_module("vae") - self.vae.requires_grad_(False) - self.timestep_shift = self.training_args.pipeline_config.flow_shift self.noise_scheduler = FlowMatchEulerDiscreteScheduler( shift=self.timestep_shift) @@ -90,10 +96,7 @@ class DistillationPipeline(TrainingPipeline): # self.transformer is the generator model self.real_score_transformer = self.get_module("real_score_transformer") self.fake_score_transformer = self.get_module("fake_score_transformer") - - self.real_score_transformer.requires_grad_(False) self.real_score_transformer.eval() - self.fake_score_transformer.requires_grad_(True) self.fake_score_transformer.train() if training_args.enable_gradient_checkpointing_type is not None: @@ -167,9 +170,7 @@ class DistillationPipeline(TrainingPipeline): def _prepare_distillation(self, training_batch: TrainingBatch) -> TrainingBatch: """Prepare training environment for distillation.""" - self.transformer.requires_grad_(True) self.transformer.train() - self.fake_score_transformer.requires_grad_(True) self.fake_score_transformer.train() return training_batch @@ -895,6 +896,14 @@ class DistillationPipeline(TrainingPipeline): else: set_random_seed(seed + self.global_rank) + # Check trainable params + num_trainable_generator = round( + count_trainable(self.transformer) / 1e9, 3) + num_trainable_critic = round( + count_trainable(self.fake_score_transformer) / 1e9, 3) + logger.info( + "rank: %s: # of trainable params in generator: %sB, # of trainable params in critic: %sB", + self.global_rank, num_trainable_generator, num_trainable_critic) # Set random seeds for deterministic training self.noise_random_generator = torch.Generator(device="cpu").manual_seed( self.seed) diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 0615ed24f..dea70b277 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -38,8 +38,8 @@ from fastvideo.training.activation_checkpoint import ( apply_activation_checkpointing) from fastvideo.training.training_utils import ( clip_grad_norm_while_handling_failing_dtensor_cases, - compute_density_for_timestep_sampling, get_scheduler, get_sigmas, - load_checkpoint, normalize_dit_input, save_checkpoint, + compute_density_for_timestep_sampling, count_trainable, get_scheduler, + get_sigmas, load_checkpoint, normalize_dit_input, save_checkpoint, shard_latents_across_sp) from fastvideo.utils import is_vsa_available, set_random_seed, shallow_asdict @@ -50,10 +50,6 @@ vsa_available = is_vsa_available() logger = init_logger(__name__) -def _get_trainable_params(model: torch.nn.Module) -> int: - return sum(p.numel() for p in model.parameters() if p.requires_grad) - - class TrainingPipeline(LoRAPipeline, ABC): """ A pipeline for training a model. All training pipelines should inherit from this class. @@ -113,6 +109,7 @@ class TrainingPipeline(LoRAPipeline, ABC): enable_gradient_checkpointing_type) noise_scheduler = self.modules["scheduler"] + # Set grads for proper modules based on the training mode (Distill, LoRA, etc.) self.set_trainable() params_to_optimize = self.transformer.parameters() params_to_optimize = list( @@ -417,7 +414,7 @@ class TrainingPipeline(LoRAPipeline, ABC): local_main_process_only=False) if not self.post_init_called: self.post_init() - num_trainable_params = _get_trainable_params(self.transformer) + num_trainable_params = count_trainable(self.transformer) logger.info("Starting training with %s B trainable parameters", round(num_trainable_params / 1e9, 3)) @@ -506,7 +503,7 @@ class TrainingPipeline(LoRAPipeline, ABC): self._log_validation(self.transformer, self.training_args, step) gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2 trainable_params = round( - _get_trainable_params(self.transformer) / 1e9, 3) + count_trainable(self.transformer) / 1e9, 3) logger.info( "GPU memory usage after validation: %s MB, trainable params: %sB", gpu_memory_usage, trainable_params) @@ -542,7 +539,7 @@ class TrainingPipeline(LoRAPipeline, ABC): logger.info(" Total optimization steps = %s", self.training_args.max_train_steps) logger.info(" Total training parameters per FSDP shard = %s B", - round(_get_trainable_params(self.transformer) / 1e9, 3)) + round(count_trainable(self.transformer) / 1e9, 3)) # print dtype logger.info(" Master weight dtype: %s", self.transformer.parameters().__next__().dtype) @@ -610,7 +607,6 @@ class TrainingPipeline(LoRAPipeline, ABC): validation_dataloader = DataLoader(validation_dataset, batch_size=None, num_workers=0) - transformer.eval() validation_steps = training_args.validation_sampling_steps.split(",") diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index bdeb81b4f..aaef338af 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -1278,3 +1278,7 @@ def get_scheduler( num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, last_epoch=last_epoch) + + +def count_trainable(model: torch.nn.Module) -> int: + return sum(p.numel() for p in model.parameters() if p.requires_grad)