diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py index 355e60c5b..0858ad658 100644 --- a/fastvideo/v1/fastvideo_args.py +++ b/fastvideo/v1/fastvideo_args.py @@ -712,3 +712,55 @@ class TrainingArgs(FastVideoArgs): help="VSA decay interval steps") return parser + + +@dataclasses.dataclass +class DistillationArgs(TrainingArgs): + """ + Distillation arguments. Inherits from TrainingArgs and adds distillation-specific + arguments. + """ + student_critic_update_ratio: int = 1 + distillation_weight: float = 1.0 + dmd_weight: float = 0.1 + adversarial_weight: float = 0.01 + + @staticmethod + def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: + parser.add_argument("--student-critic-update-ratio", + type=int, + default=DistillationArgs.student_critic_update_ratio, + help="Ratio of student updates to critic updates.") + parser.add_argument("--distillation-weight", + type=float, + default=DistillationArgs.distillation_weight, + help="Weight for distillation loss.") + parser.add_argument("--dmd-weight", + type=float, + default=DistillationArgs.dmd_weight, + help="Weight for DMD loss.") + parser.add_argument("--adversarial-weight", + type=float, + default=DistillationArgs.adversarial_weight, + help="Weight for student's adversarial loss.") + return parser + + @classmethod + def from_cli_args(cls, args: argparse.Namespace) -> "DistillationArgs": + provided_args = clean_cli_args(args) + # Get all fields from the dataclass + attrs = [attr.name for attr in dataclasses.fields(cls)] + logger.info(provided_args) + # Create a dictionary of attribute values, with defaults for missing attributes + kwargs = {} + for attr in attrs: + if attr == 'pipeline_config': + pipeline_config = PipelineConfig.from_kwargs(provided_args) + kwargs[attr] = pipeline_config + # Use getattr with default value from the dataclass for potentially missing attributes + else: + default_value = getattr(cls, attr, None) + value = getattr(args, attr, default_value) + kwargs[attr] = value # type: ignore + + return cls(**kwargs) # type: ignore diff --git a/fastvideo/v1/training/distillation_pipeline.py b/fastvideo/v1/training/distillation_pipeline.py new file mode 100644 index 000000000..25176ce63 --- /dev/null +++ b/fastvideo/v1/training/distillation_pipeline.py @@ -0,0 +1,665 @@ +# SPDX-License-Identifier: Apache-2.0 +import gc +import math +import os +import time +from collections import deque +from typing import Any, Dict, Iterator, List, Optional + +import imageio +import numpy as np +import torch +import torchvision +from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler +from diffusers.optimization import get_scheduler +from einops import rearrange +from torchdata.stateful_dataloader import StatefulDataLoader +from tqdm.auto import tqdm + +import fastvideo.v1.envs as envs +from fastvideo.v1.attention.backends.video_sparse_attn import ( + VideoSparseAttentionMetadata) +from fastvideo.v1.configs.sample import SamplingParam +from fastvideo.v1.dataset import build_parquet_map_style_dataloader +from fastvideo.v1.distributed import (cleanup_dist_env_and_memory, get_sp_group, + get_torch_device, get_world_group) +from fastvideo.v1.fastvideo_args import (DistillationArgs, FastVideoArgs, + TrainingArgs) +from fastvideo.v1.forward_context import set_forward_context +from fastvideo.v1.logger import init_logger +from fastvideo.v1.pipelines import (ComposedPipelineBase, ForwardBatch, + TrainingBatch) +from fastvideo.v1.training.training_pipeline import TrainingPipeline +from fastvideo.v1.training.training_utils import ( + clip_grad_norm_while_handling_failing_dtensor_cases, + compute_density_for_timestep_sampling, get_sigmas, load_checkpoint, + normalize_dit_input, save_checkpoint, shard_latents_across_sp) +from fastvideo.v1.utils import is_vsa_available, set_random_seed + +import wandb # isort: skip + +vsa_available = is_vsa_available() + +logger = init_logger(__name__) + + +class DistillationPipeline(TrainingPipeline): + """ + A distillation pipeline for training a student model using teacher model guidance. + Inherits from TrainingPipeline to reuse training infrastructure. + """ + _required_config_modules = ["scheduler", "transformer"] + + def __init__(self): + super().__init__() + self.student_transformer = None + self.teacher_transformer = None + self.critic_transformer = None + self.student_transformer_optimizer = None + self.critic_transformer_optimizer = None + self.student_critic_update_ratio = 1 # How often to update student vs critic + self.max_grad_norm = 10.0 + self.unconditional_dict = None # Cache for unconditional embeddings + + def initialize_training_pipeline(self, training_args: TrainingArgs): + """Initialize the distillation training pipeline with multiple models.""" + logger.info("Initializing distillation training pipeline...") + + # Call parent initialization first + super().initialize_training_pipeline(training_args) + assert isinstance(self.training_args, DistillationArgs) + + # Rename transformer to student_transformer for clarity + self.student_transformer = self.transformer + + # Initialize distillation-specific models + self._initialize_distillation_models(training_args) + + # Initialize distillation-specific optimizers + self._initialize_distillation_optimizers(training_args) + + # Set distillation-specific parameters + self.student_critic_update_ratio = self.training_args.student_critic_update_ratio + self.max_grad_norm = self.training_args.max_grad_norm + + logger.info(f"Distillation pipeline initialized with student_critic_update_ratio={self.student_critic_update_ratio}") + + def _initialize_distillation_models(self, training_args: TrainingArgs): + """Initialize teacher and critic models by copying the student model.""" + assert self.student_transformer is not None, "Student model must be loaded first" + + # Create teacher model by copying student model + self.teacher_transformer = self._copy_model(self.student_transformer) + assert self.teacher_transformer is not None, "Failed to create teacher model" + + # Create critic model by copying student model + self.critic_transformer = self._copy_model(self.student_transformer) + assert self.critic_transformer is not None, "Failed to create critic model" + + # Set teacher model to eval mode and freeze parameters + self.teacher_transformer.requires_grad_(False) + self.teacher_transformer.eval() + + # Set critic model to train mode + self.critic_transformer.requires_grad_(True) + self.critic_transformer.train() + + logger.info("Distillation models initialized:") + logger.info(" - Student transformer: %s B parameters (trainable)", + sum(p.numel() for p in self.student_transformer.parameters()) / 1e9) + logger.info(" - Teacher transformer: %s B parameters (frozen)", + sum(p.numel() for p in self.teacher_transformer.parameters()) / 1e9) + logger.info(" - Critic transformer: %s B parameters (trainable)", + sum(p.numel() for p in self.critic_transformer.parameters()) / 1e9) + + def _copy_model(self, original_model): + """Create a deep copy of a model.""" + import copy + copied_model = copy.deepcopy(original_model) + device = next(original_model.parameters()).device + copied_model = copied_model.to(device) + if hasattr(copied_model, 'reset_parameters'): + copied_model.reset_parameters() + return copied_model + + def _initialize_distillation_optimizers(self, training_args: TrainingArgs): + """Initialize separate optimizers for student and critic models.""" + assert self.student_transformer is not None + student_params = list(filter(lambda p: p.requires_grad, self.student_transformer.parameters())) + self.student_transformer_optimizer = torch.optim.AdamW( + student_params, + lr=training_args.learning_rate, + betas=(0.9, 0.999), + weight_decay=training_args.weight_decay, + eps=1e-8, + ) + + assert self.critic_transformer is not None + 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, + betas=(0.9, 0.999), + weight_decay=training_args.weight_decay, + eps=1e-8, + ) + + logger.info("Distillation optimizers initialized: student and critic") + + def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch: + """Prepare training environment for distillation.""" + assert self.student_transformer is not None + self.student_transformer.requires_grad_(True) + self.student_transformer.train() + assert self.critic_transformer is not None + self.critic_transformer.requires_grad_(True) + self.critic_transformer.train() + + assert self.student_transformer_optimizer is not None + self.student_transformer_optimizer.zero_grad() + assert self.critic_transformer_optimizer is not None + self.critic_transformer_optimizer.zero_grad() + + training_batch.total_loss = 0.0 + setattr(training_batch, 'student_transformer_loss', 0.0) + setattr(training_batch, 'critic_transformer_loss', 0.0) + + return training_batch + + def _compute_teacher_outputs(self, training_batch: TrainingBatch) -> TrainingBatch: + """Compute teacher model outputs for distillation.""" + assert self.teacher_transformer is not None + assert training_batch.noisy_model_input is not None + assert training_batch.encoder_hidden_states is not None + assert training_batch.encoder_attention_mask is not None + assert training_batch.timesteps is not None + + with torch.no_grad(): + teacher_input_kwargs = { + "hidden_states": training_batch.noisy_model_input, + "encoder_hidden_states": training_batch.encoder_hidden_states, + "timestep": training_batch.timesteps.to(get_torch_device(), dtype=torch.bfloat16), + "encoder_attention_mask": training_batch.encoder_attention_mask, + "return_dict": False, + } + + teacher_output = self.teacher_transformer(**teacher_input_kwargs) + setattr(training_batch, 'teacher_output', teacher_output) + + return training_batch + + def _compute_distillation_loss(self, training_batch: TrainingBatch) -> torch.Tensor: + """Compute distillation loss between teacher and student outputs.""" + teacher_output = getattr(training_batch, 'teacher_output', None) + assert teacher_output is not None + model_pred = getattr(training_batch, 'model_pred', None) + assert model_pred is not None + + distillation_loss = torch.mean((teacher_output.float() - model_pred.float())**2) + return distillation_loss + + def _compute_dmd_loss(self, training_batch) -> torch.Tensor: + """Compute DMD (Diffusion Model Distillation) loss.""" + assert training_batch.latents is not None + model_pred = getattr(training_batch, 'model_pred', None) + assert model_pred is not None + assert training_batch.encoder_hidden_states is not None + assert training_batch.encoder_attention_mask is not None + assert training_batch.timesteps is not None + + with torch.no_grad(): + assert self.teacher_transformer is not None + teacher_clean_pred = self.teacher_transformer( + hidden_states=training_batch.latents, + encoder_hidden_states=training_batch.encoder_hidden_states, + timestep=training_batch.timesteps, + encoder_attention_mask=training_batch.encoder_attention_mask, + return_dict=False + ) + + assert self.student_transformer is not None + student_clean_pred = self.student_transformer( + hidden_states=training_batch.latents, + encoder_hidden_states=training_batch.encoder_hidden_states, + timestep=training_batch.timesteps, + encoder_attention_mask=training_batch.encoder_attention_mask, + return_dict=False + ) + + dmd_loss = torch.mean((teacher_clean_pred.float() - student_clean_pred.float())**2) + return dmd_loss + + def _compute_critic_transformer_loss(self, training_batch: TrainingBatch) -> torch.Tensor: + """Compute critic loss for adversarial training.""" + assert self.critic_transformer is not None + assert training_batch.latents is not None + model_pred = getattr(training_batch, 'model_pred', None) + assert model_pred is not None + + real_output = self.critic_transformer(training_batch.latents) + real_loss = torch.mean((real_output - 1.0) ** 2) + + fake_output = self.critic_transformer(model_pred.detach()) + fake_loss = torch.mean(fake_output ** 2) + + critic_transformer_loss = real_loss + fake_loss + return critic_transformer_loss + + def _compute_student_adversarial_loss(self, training_batch: TrainingBatch) -> torch.Tensor: + """Compute student's adversarial loss.""" + assert self.critic_transformer is not None + model_pred = getattr(training_batch, 'model_pred', None) + assert model_pred is not None + + fake_output = self.critic_transformer(model_pred) + adversarial_loss = torch.mean((fake_output - 1.0) ** 2) + return adversarial_loss + + def _student_transformer_forward_and_compute_loss(self, training_batch: TrainingBatch) -> TrainingBatch: + """Forward pass through student and compute all losses.""" + assert self.student_transformer is not None + assert training_batch.noisy_model_input is not None + assert training_batch.encoder_hidden_states is not None + assert training_batch.encoder_attention_mask is not None + assert training_batch.timesteps is not None + + model_output = self.student_transformer( + hidden_states=training_batch.noisy_model_input, + encoder_hidden_states=training_batch.encoder_hidden_states, + timestep=training_batch.timesteps.to(get_torch_device(), dtype=torch.bfloat16), + encoder_attention_mask=training_batch.encoder_attention_mask, + return_dict=False, + ) + setattr(training_batch, 'model_pred', model_output) + + training_batch = self._compute_teacher_outputs(training_batch) + + task_loss = self._compute_task_loss(training_batch) + distillation_loss = self._compute_distillation_loss(training_batch) + dmd_loss = self._compute_dmd_loss(training_batch) + adversarial_loss = self._compute_student_adversarial_loss(training_batch) + critic_transformer_loss = self._compute_critic_transformer_loss(training_batch) + + assert isinstance(self.training_args, DistillationArgs) + distillation_weight = self.training_args.distillation_weight + dmd_weight = self.training_args.dmd_weight + adversarial_weight = self.training_args.adversarial_weight + + student_transformer_loss = (task_loss + + distillation_weight * distillation_loss + + dmd_weight * dmd_loss + + adversarial_weight * adversarial_loss) + + setattr(training_batch, 'task_loss', task_loss) + setattr(training_batch, 'distillation_loss', distillation_loss) + setattr(training_batch, 'dmd_loss', dmd_loss) + setattr(training_batch, 'adversarial_loss', adversarial_loss) + setattr(training_batch, 'critic_transformer_loss', critic_transformer_loss) + setattr(training_batch, 'student_transformer_loss', student_transformer_loss) + training_batch.total_loss = float(student_transformer_loss + critic_transformer_loss) + + student_transformer_loss.backward() + critic_transformer_loss.backward() + + return training_batch + + def _compute_task_loss(self, training_batch: TrainingBatch) -> torch.Tensor: + """Compute the main task loss (e.g., noise prediction loss).""" + model_pred = getattr(training_batch, 'model_pred', None) + assert model_pred is not None + assert training_batch.noise is not None + + task_loss = torch.mean((model_pred - training_batch.noise) ** 2) + return task_loss + + def _clip_grad_norm(self, training_batch: TrainingBatch) -> TrainingBatch: + """Clip gradients for both student and critic.""" + assert self.student_transformer is not None + student_transformer_grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases( + self.student_transformer, self.max_grad_norm) + + assert self.critic_transformer is not None + critic_transformer_grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases( + self.critic_transformer, self.max_grad_norm) + + setattr(training_batch, 'student_transformer_grad_norm', student_transformer_grad_norm) + setattr(training_batch, 'critic_transformer_grad_norm', critic_transformer_grad_norm) + training_batch.grad_norm = float(student_transformer_grad_norm) if student_transformer_grad_norm is not None else 0.0 + + return training_batch + + def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch: + """Train one step with alternating student and critic updates.""" + assert self.training_args is not None + + training_batch = self._prepare_training(training_batch) + + TRAIN_STUDENT_TRANSFORMER = training_batch.current_timestep % self.student_critic_update_ratio == 0 + + for _ in range(self.training_args.gradient_accumulation_steps): + training_batch = self._get_next_batch(training_batch) + + assert training_batch.latents is not None + training_batch.latents = shard_latents_across_sp( + training_batch.latents, + num_latent_t=self.training_args.num_latent_t) + + training_batch = self._normalize_dit_input(training_batch) + training_batch = self._prepare_dit_inputs(training_batch) + training_batch = self._build_attention_metadata(training_batch) + + training_batch = self._student_transformer_forward_and_compute_loss(training_batch) + + training_batch = self._clip_grad_norm(training_batch) + + if TRAIN_STUDENT_TRANSFORMER and self.student_transformer_optimizer is not None: + self.student_transformer_optimizer.step() + if self.critic_transformer_optimizer is not None: + self.critic_transformer_optimizer.step() + + if TRAIN_STUDENT_TRANSFORMER and hasattr(self, 'lr_scheduler') and self.lr_scheduler is not None: + self.lr_scheduler.step() + + return training_batch + + def _resume_from_checkpoint(self) -> None: + """Resume training from checkpoint with distillation models.""" + assert self.training_args is not None + logger.info("Loading distillation checkpoint from %s", + self.training_args.resume_from_checkpoint) + + resumed_step = load_checkpoint( + self.student_transformer, self.global_rank, + self.training_args.resume_from_checkpoint, self.student_transformer_optimizer, + self.train_dataloader, self.lr_scheduler, + self.noise_random_generator) + + # TODO: Add checkpoint loading for critic and teacher models + + if resumed_step > 0: + self.init_steps = resumed_step + logger.info("Successfully resumed from step %s", resumed_step) + else: + logger.warning("Failed to load checkpoint, starting from step 0") + self.init_steps = 0 + + def _log_training_info(self) -> None: + """Log distillation-specific training information.""" + super()._log_training_info() + + logger.info("Distillation-specific settings:") + logger.info(" Student/Critic update ratio: %s", self.student_critic_update_ratio) + logger.info(" Max gradient norm: %s", self.max_grad_norm) + assert self.teacher_transformer is not None + logger.info(" Teacher transformer parameters: %s B", + sum(p.numel() for p in self.teacher_transformer.parameters()) / 1e9) + assert self.critic_transformer is not None + logger.info(" Critic transformer parameters: %s B", + sum(p.numel() for p in self.critic_transformer.parameters()) / 1e9) + + @torch.no_grad() + def _log_validation(self, student_transformer, training_args, global_step) -> None: + """Log validation results for distillation training.""" + assert training_args is not None + training_args.inference_mode = True + training_args.use_cpu_offload = False + if not training_args.log_validation: + return + if self.validation_pipeline is None: + raise ValueError("Validation pipeline is not set") + + logger.info("Starting validation") + + sampling_param = SamplingParam.from_pretrained(training_args.model_path) + + validation_seed = training_args.seed if training_args.seed is not None else 42 + torch.manual_seed(validation_seed) + torch.cuda.manual_seed_all(validation_seed) + + logger.info("Using validation seed: %s", validation_seed) + + logger.info('fastvideo_args.validation_preprocessed_path: %s', + training_args.validation_preprocessed_path) + validation_dataset, validation_dataloader = build_parquet_map_style_dataloader( + training_args.validation_preprocessed_path, + batch_size=1, + num_data_workers=0, + drop_last=False, + drop_first_row=sampling_param.negative_prompt is not None, + ) + if sampling_param.negative_prompt: + _, negative_prompt_embeds, negative_prompt_attention_mask, _ = validation_dataset.get_validation_negative_prompt( + ) + + if self.student_transformer: + self.student_transformer.eval() + + validation_steps = training_args.validation_sampling_steps.split(",") + validation_steps = [int(step) for step in validation_steps] + validation_steps = [step for step in validation_steps if step > 0] + + for num_inference_steps in validation_steps: + step_videos: List[np.ndarray] = [] + step_captions: List[str | None] = [] + + for _, embeddings, masks, infos in validation_dataloader: + step_captions.extend([None]) + prompt_embeds = embeddings.to(get_torch_device()) + prompt_attention_mask = masks.to(get_torch_device()) + + latents_size = [(sampling_param.num_frames - 1) // 4 + 1, + sampling_param.height // 8, + sampling_param.width // 8] + n_tokens = latents_size[0] * latents_size[1] * latents_size[2] + + temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio + num_frames = (training_args.num_latent_t - + 1) * temporal_compression_factor + 1 + + batch = ForwardBatch( + data_type="video", + latents=None, + seed=validation_seed, + generator=torch.Generator( + device="cpu").manual_seed(validation_seed), + prompt_embeds=[prompt_embeds], + prompt_attention_mask=[prompt_attention_mask], + negative_prompt_embeds=[negative_prompt_embeds], + negative_attention_mask=[negative_prompt_attention_mask], + height=training_args.num_height, + width=training_args.num_width, + num_frames=num_frames, + num_inference_steps= + num_inference_steps, + guidance_scale=sampling_param.guidance_scale, + n_tokens=n_tokens, + eta=0.0, + VSA_sparsity=training_args.VSA_sparsity, + ) + + with torch.no_grad(), torch.autocast("cuda", + dtype=torch.bfloat16): + output_batch = self.validation_pipeline.forward( + batch, training_args) + samples = output_batch.output + + if self.rank_in_sp_group != 0: + continue + + video = rearrange(samples, "b c t h w -> t b c h w") + frames: List[np.ndarray] = [] + for x in video: + x = torchvision.utils.make_grid(x, nrow=6) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + frames.append((x * 255).numpy().astype(np.uint8)) + step_videos.append(np.array(frames)) + + world_group = get_world_group() + num_sp_groups = world_group.world_size // self.sp_group.world_size + + if self.rank_in_sp_group == 0: + if self.global_rank == 0: + all_videos: List[np.ndarray] = [] + all_videos.extend(step_videos) + all_captions = step_captions + + for sp_group_idx in range(1, num_sp_groups): + src_rank = sp_group_idx * self.sp_world_size + recv_videos = world_group.recv_object(src=src_rank) + recv_captions = world_group.recv_object(src=src_rank) + all_videos.extend(recv_videos) + all_captions.extend(recv_captions) + + video_filenames = [] + for i, (video_frames, + caption) in enumerate(zip(all_videos, + all_captions)): + os.makedirs(training_args.output_dir, exist_ok=True) + filename = os.path.join( + training_args.output_dir, + f"validation_step_{global_step}_inference_steps_{num_inference_steps}_video_{i}.mp4" + ) + imageio.mimsave(filename, list(video_frames), fps=sampling_param.fps) + video_filenames.append(filename) + + logs = { + f"validation_videos_{num_inference_steps}_steps": [ + wandb.Video(filename, caption=caption) + for filename, caption in zip( + video_filenames, all_captions) + ] + } + wandb.log(logs, step=global_step) + else: + world_group.send_object(step_videos, dst=0) + world_group.send_object(step_captions, dst=0) + + if self.student_transformer: + self.student_transformer.train() + gc.collect() + torch.cuda.empty_cache() + + def train(self) -> None: + """Main training loop with distillation-specific logging.""" + assert self.training_args is not None + + assert self.training_args.seed is not None, "seed must be set" + seed = self.training_args.seed + set_random_seed(seed) + + self.noise_random_generator = torch.Generator( + device="cpu").manual_seed(seed) + + logger.info("Initialized random seeds with seed: %s", seed) + + self.noise_scheduler = FlowMatchEulerDiscreteScheduler() + + if self.training_args.resume_from_checkpoint: + self._resume_from_checkpoint() + + self.train_loader_iter = iter(self.train_dataloader) + + step_times: deque[float] = deque(maxlen=100) + + self._log_training_info() + self._log_validation(self.student_transformer, self.training_args, 1) + + progress_bar = tqdm( + range(0, self.training_args.max_train_steps), + initial=self.init_steps, + desc="Steps", + disable=self.local_rank > 0, + ) + + for step in range(self.init_steps + 1, + self.training_args.max_train_steps + 1): + start_time = time.perf_counter() + + if vsa_available: + vsa_sparsity = self.training_args.VSA_sparsity + vsa_decay_rate = self.training_args.VSA_decay_rate + vsa_decay_interval_steps = self.training_args.VSA_decay_interval_steps + current_decay_times = min(step // vsa_decay_interval_steps, + vsa_sparsity // vsa_decay_rate) + current_vsa_sparsity = current_decay_times * vsa_decay_rate + else: + current_vsa_sparsity = 0.0 + + training_batch = TrainingBatch() + training_batch.current_timestep = step + training_batch.current_vsa_sparsity = current_vsa_sparsity + training_batch = self.train_one_step(training_batch) + + total_loss = training_batch.total_loss + student_transformer_loss = getattr(training_batch, 'student_transformer_loss', 0.0) + student_transformer_loss = student_transformer_loss + critic_transformer_loss = getattr(training_batch, 'critic_transformer_loss', 0.0) + critic_transformer_loss = critic_transformer_loss + distillation_loss = getattr(training_batch, 'distillation_loss', 0.0) + distillation_loss = distillation_loss + task_loss = getattr(training_batch, 'task_loss', 0.0) + task_loss = task_loss + adversarial_loss = getattr(training_batch, 'adversarial_loss', 0.0) + adversarial_loss = adversarial_loss + grad_norm = training_batch.grad_norm + + step_time = time.perf_counter() - start_time + step_times.append(step_time) + avg_step_time = sum(step_times) / len(step_times) + + progress_bar.set_postfix({ + "total_loss": f"{total_loss:.4f}", + "student_loss": f"{student_transformer_loss:.4f}", + "critic_loss": f"{critic_transformer_loss:.4f}", + "step_time": f"{step_time:.2f}s", + "grad_norm": grad_norm, + }) + progress_bar.update(1) + + if self.global_rank == 0: + wandb.log( + { + "train_total_loss": total_loss, + "train_student_transformer_loss": student_transformer_loss, + "train_critic_transformer_loss": critic_transformer_loss, + "train_distillation_loss": distillation_loss, + "train_task_loss": task_loss, + "train_adversarial_loss": adversarial_loss, + "learning_rate": self.lr_scheduler.get_last_lr()[0], + "step_time": step_time, + "avg_step_time": avg_step_time, + "grad_norm": grad_norm, + "vsa_sparsity": current_vsa_sparsity, + }, + step=step, + ) + + if step % self.training_args.checkpointing_steps == 0: + save_checkpoint(self.student_transformer, self.global_rank, + self.training_args.output_dir, step, + self.student_transformer_optimizer, self.train_dataloader, + self.lr_scheduler, self.noise_random_generator) + if self.student_transformer: + self.student_transformer.train() + self.sp_group.barrier() + + if self.training_args.log_validation and step % self.training_args.validation_steps == 0: + self._log_validation(self.student_transformer, self.training_args, step) + gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2 + logger.info("GPU memory usage after validation: %s MB", + gpu_memory_usage) + + wandb.finish() + save_checkpoint(self.student_transformer, self.global_rank, + self.training_args.output_dir, + self.training_args.max_train_steps, self.student_transformer_optimizer, + self.train_dataloader, self.lr_scheduler, + self.noise_random_generator) + + if get_sp_group(): + cleanup_dist_env_and_memory() + + def initialize_validation_pipeline(self, training_args: TrainingArgs): + """Initialize validation pipeline - must be implemented by subclasses.""" + raise NotImplementedError( + "Distillation pipelines must implement this method") \ No newline at end of file diff --git a/fastvideo/v1/training/wan_distillation_pipeline.py b/fastvideo/v1/training/wan_distillation_pipeline.py new file mode 100644 index 000000000..c6aebbec2 --- /dev/null +++ b/fastvideo/v1/training/wan_distillation_pipeline.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +import sys +from copy import deepcopy + +from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler) +from fastvideo.v1.pipelines.wan.wan_pipeline import WanValidationPipeline +from fastvideo.v1.training.distillation_pipeline import DistillationPipeline +from fastvideo.v1.utils import is_vsa_available + +vsa_available = is_vsa_available() + +logger = init_logger(__name__) + + +class WanDistillationPipeline(DistillationPipeline): + """ + A distillation pipeline for Wan that uses a single transformer model. + The main transformer serves as the student model, and copies are made for teacher and critic. + """ + _required_config_modules = ["scheduler", "transformer"] # Only need scheduler and transformer + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + """Initialize Wan-specific scheduler.""" + self.modules["scheduler"] = FlowUniPCMultistepScheduler( + shift=fastvideo_args.pipeline_config.flow_shift) + + def create_training_stages(self, training_args: TrainingArgs): + """ + May be used in future refactors. + """ + pass + + def initialize_validation_pipeline(self, training_args: TrainingArgs): + """Initialize Wan validation pipeline.""" + logger.info("Initializing Wan validation pipeline...") + args_copy = deepcopy(training_args) + + args_copy.inference_mode = True + args_copy.pipeline_config.vae_config.load_encoder = False + validation_pipeline = WanValidationPipeline.from_pretrained( + training_args.model_path, + args=None, + inference_mode=True, + loaded_modules={"transformer": self.get_module("transformer")}, + tp_size=training_args.tp_size, + sp_size=training_args.sp_size, + num_gpus=training_args.num_gpus) + + self.validation_pipeline = validation_pipeline + + +def main(args) -> None: + logger.info("Starting distillation pipeline...") + + pipeline = WanDistillationPipeline.from_pretrained( + args.pretrained_model_name_or_path, args=args) + args = pipeline.training_args + pipeline.train() + logger.info("Wan distillation pipeline completed") + + +if __name__ == "__main__": + argv = sys.argv + from fastvideo.v1.fastvideo_args import TrainingArgs + from fastvideo.v1.utils import FlexibleArgumentParser + parser = FlexibleArgumentParser() + parser = TrainingArgs.add_cli_args(parser) + from fastvideo.v1.fastvideo_args import FastVideoArgs + parser = FastVideoArgs.add_cli_args(parser) + + # Remove the need for separate teacher and discriminator model paths + # since we now copy from the main transformer + args = parser.parse_args() + args.use_cpu_offload = False + main(args) \ No newline at end of file