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 f0d3d1c43..61d8a01f4 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 @@ -18,20 +18,20 @@ export TORCH_NCCL_ENABLE_MONITORING=0 export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} export MASTER_PORT=29503 export TOKENIZERS_PARALLELISM=false -export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45" +export WANDB_API_KEY="2f25ad37933894dbf0966c838c0b8494987f9f2f" export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN # Configs -NUM_GPUS=4 +NUM_GPUS=1 # Model paths for Self-Forcing DMD distillation: GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers" # Teacher model FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model -DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/" +DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/" VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json" # export CUDA_VISIBLE_DEVICES=4,5 # IP=[MASTER NODE IP] diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index 2f39b57b6..e4d563d9b 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -23,7 +23,7 @@ from fastvideo.training.training_utils import ( ) 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 +from fastvideo.distributed import get_world_group, get_local_torch_device import torch.distributed as dist import numpy as np from fastvideo.utils import set_random_seed, is_vsa_available @@ -579,6 +579,42 @@ class SelfForcingDistillationPipeline(DistillationPipeline): assert cache_dict["k"].shape[2] == num_attention_heads, f"Attention heads mismatch in crossattn_cache block {i}" assert cache_dict["k"].shape[3] == attention_head_dim, f"Attention head dim mismatch in crossattn_cache block {i}" + def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: + batch = next(self.train_loader_iter, None) # type: ignore + if batch is None: + self.current_epoch += 1 + logger.info("Starting epoch %s", self.current_epoch) + # Reset iterator for next epoch + self.train_loader_iter = iter(self.train_dataloader) + # Get first batch of new epoch + batch = next(self.train_loader_iter) + + # latents, encoder_hidden_states, encoder_attention_mask, infos = batch + encoder_hidden_states = batch['text_embedding'] + encoder_attention_mask = batch['text_attention_mask'] + infos = batch['info_list'] + + batch_size = encoder_hidden_states.shape[0] + vae_config = self.training_args.pipeline_config.vae_config.arch_config + num_channels = vae_config.z_dim + spatial_compression_ratio = vae_config.spatial_compression_ratio + + latent_height = self.training_args.num_height // spatial_compression_ratio + latent_width = self.training_args.num_width // spatial_compression_ratio + + latents = torch.randn(batch_size, num_channels, self.training_args.num_latent_t, + latent_height, latent_width).to(get_local_torch_device(), dtype=torch.bfloat16) + + training_batch.latents = latents.to(get_local_torch_device(), + dtype=torch.bfloat16) + training_batch.encoder_hidden_states = encoder_hidden_states.to( + get_local_torch_device(), dtype=torch.bfloat16) + training_batch.encoder_attention_mask = encoder_attention_mask.to( + get_local_torch_device(), dtype=torch.bfloat16) + training_batch.infos = infos + + return training_batch + def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch: """ Self-forcing training step that alternates between generator and critic training. diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 1adcee37d..8ac2469cd 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -22,7 +22,7 @@ from tqdm.auto import tqdm import fastvideo.envs as envs from fastvideo.attention.backends.video_sparse_attn import ( VideoSparseAttentionMetadataBuilder) -from fastvideo.attention.backends.vmoba import VideoMobaAttentionMetadataBuilder +# from fastvideo.attention.backends.vmoba import VideoMobaAttentionMetadataBuilder from fastvideo.configs.sample import SamplingParam from fastvideo.dataset import build_parquet_map_style_dataloader from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v @@ -42,13 +42,15 @@ from fastvideo.training.training_utils import ( compute_density_for_timestep_sampling, get_scheduler, get_sigmas, load_checkpoint, normalize_dit_input, save_checkpoint, shard_latents_across_sp) -from fastvideo.utils import (is_vmoba_available, is_vsa_available, +# from fastvideo.utils import (is_vmoba_available, is_vsa_available, +# set_random_seed, shallow_asdict) +from fastvideo.utils import (is_vsa_available, set_random_seed, shallow_asdict) import wandb # isort: skip vsa_available = is_vsa_available() -vmoba_available = is_vmoba_available() +# vmoba_available = is_vmoba_available() logger = init_logger(__name__) @@ -279,20 +281,20 @@ class TrainingPipeline(LoRAPipeline, ABC): patch_size=patch_size, VSA_sparsity=current_vsa_sparsity, device=get_local_torch_device()) - elif vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN": - moba_params = self.training_args.moba_config.copy() - moba_params.update({ - "current_timestep": - training_batch.timesteps, - "raw_latent_shape": - training_batch.raw_latent_shape[2:5], - "patch_size": - self.training_args.pipeline_config.dit_config.patch_size, - "device": - get_local_torch_device(), - }) - training_batch.attn_metadata = VideoMobaAttentionMetadataBuilder( - ).build(**moba_params) + # elif vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN": + # moba_params = self.training_args.moba_config.copy() + # moba_params.update({ + # "current_timestep": + # training_batch.timesteps, + # "raw_latent_shape": + # training_batch.raw_latent_shape[2:5], + # "patch_size": + # self.training_args.pipeline_config.dit_config.patch_size, + # "device": + # get_local_torch_device(), + # }) + # training_batch.attn_metadata = VideoMobaAttentionMetadataBuilder( + # ).build(**moba_params) else: training_batch.attn_metadata = None @@ -317,7 +319,8 @@ class TrainingPipeline(LoRAPipeline, ABC): def _transformer_forward_and_compute_loss( self, training_batch: TrainingBatch) -> TrainingBatch: - if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN" or vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN": + # if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN" or vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN": + if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": assert training_batch.attn_metadata is not None else: assert training_batch.attn_metadata is None @@ -483,9 +486,9 @@ class TrainingPipeline(LoRAPipeline, ABC): current_decay_times = min(step // vsa_decay_interval_steps, vsa_sparsity // vsa_decay_rate) current_vsa_sparsity = current_decay_times * vsa_decay_rate - elif vmoba_available: - # TODO: add vmoba sparsity scheduling here - current_vsa_sparsity = 0.0 + # elif vmoba_available: + # # TODO: add vmoba sparsity scheduling here + # pass else: current_vsa_sparsity = 0.0