From be27c17a7a9c1c8ec623308e76fde82bb2fdaa45 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CBrianChen1129=E2=80=9D?= Date: Thu, 14 Aug 2025 03:51:06 +0000 Subject: [PATCH] update --- fastvideo/fastvideo_args.py | 2 +- fastvideo/training/training_pipeline.py | 129 +++++++++++++++++--- fastvideo/training/wan_training_pipeline.py | 4 +- 3 files changed, 113 insertions(+), 22 deletions(-) diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index cce97ce2b..ebf33f2a5 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -119,7 +119,7 @@ class FastVideoArgs: output_type: str = "pil" # CPU offload parameters - dit_cpu_offload: bool = True + dit_cpu_offload: bool = False use_fsdp_inference: bool = True text_encoder_cpu_offload: bool = True image_encoder_cpu_offload: bool = True diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 0615ed24f..16cf7ae58 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -64,6 +64,7 @@ class TrainingPipeline(LoRAPipeline, ABC): train_dataloader: StatefulDataLoader train_loader_iter: Iterator[dict[str, Any]] current_epoch: int = 0 + train_transformer_2: bool = False def __init__( self, @@ -99,6 +100,7 @@ class TrainingPipeline(LoRAPipeline, ABC): self.sp_world_size = self.sp_group.world_size self.local_rank = world_group.local_rank self.transformer = self.get_module("transformer") + self.transformer_2 = self.get_module("transformer_2", None) self.seed = training_args.seed self.set_schemas() @@ -111,10 +113,16 @@ class TrainingPipeline(LoRAPipeline, ABC): self.transformer, checkpointing_type=training_args. enable_gradient_checkpointing_type) + if self.transformer_2 is not None: + self.transformer_2 = apply_activation_checkpointing( + self.transformer_2, + checkpointing_type=training_args. + enable_gradient_checkpointing_type) noise_scheduler = self.modules["scheduler"] self.set_trainable() params_to_optimize = self.transformer.parameters() + params_to_optimize = list( filter(lambda p: p.requires_grad, params_to_optimize)) self.optimizer = torch.optim.AdamW( @@ -138,6 +146,30 @@ class TrainingPipeline(LoRAPipeline, ABC): min_lr_ratio=training_args.min_lr_ratio, last_epoch=self.init_steps - 1, ) + if self.transformer_2 is not None: + # Ensure transformer_2 has trainable parameters before creating optimizer + self.transformer_2.train() + self.transformer_2.requires_grad_(True) + params_to_optimize_2 = self.transformer_2.parameters() + params_to_optimize_2 = list( + filter(lambda p: p.requires_grad, params_to_optimize_2)) + self.optimizer_2 = torch.optim.AdamW( + params_to_optimize_2, + lr=training_args.learning_rate, + betas=(0.9, 0.999), + weight_decay=training_args.weight_decay, + eps=1e-8, + ) + self.lr_scheduler_2 = get_scheduler( + training_args.lr_scheduler, + optimizer=self.optimizer_2, + num_warmup_steps=training_args.lr_warmup_steps, + num_training_steps=training_args.max_train_steps, + num_cycles=training_args.lr_num_cycles, + power=training_args.lr_power, + min_lr_ratio=training_args.min_lr_ratio, + last_epoch=self.init_steps - 1, + ) self.train_dataset, self.train_dataloader = build_parquet_map_style_dataloader( training_args.data_path, @@ -152,7 +184,7 @@ class TrainingPipeline(LoRAPipeline, ABC): seed=self.seed) self.noise_scheduler = noise_scheduler - + self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps self.num_update_steps_per_epoch = math.ceil( len(self.train_dataloader) / training_args.gradient_accumulation_steps * training_args.sp_size / @@ -176,11 +208,27 @@ class TrainingPipeline(LoRAPipeline, ABC): "Training pipelines must implement this method") def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch: - self.transformer.train() - self.optimizer.zero_grad() + # At the beginning, disable training for all models + self._disable_training(self.transformer, self.optimizer) + if self.transformer_2 is not None: + self._disable_training(self.transformer_2, self.optimizer_2) + training_batch.total_loss = 0.0 return training_batch + def _enable_training(self, model: torch.nn.Module, optimizer: torch.optim.Optimizer) -> None: + """Enable training mode and gradients for the specified model.""" + for param in model.parameters(): + param.requires_grad = True + model.train() + optimizer.zero_grad() + + def _disable_training(self, model: torch.nn.Module, optimizer: torch.optim.Optimizer) -> None: + """Disable training mode and gradients for the specified model.""" + for param in model.parameters(): + param.requires_grad = False + optimizer.zero_grad(set_to_none=True) + def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: batch = next(self.train_loader_iter, None) # type: ignore if batch is None: @@ -224,17 +272,17 @@ class TrainingPipeline(LoRAPipeline, ABC): generator=self.noise_gen_cuda, device=latents.device, dtype=latents.dtype) - u = compute_density_for_timestep_sampling( - weighting_scheme=self.training_args.weighting_scheme, - batch_size=batch_size, - generator=self.noise_random_generator, - logit_mean=self.training_args.logit_mean, - logit_std=self.training_args.logit_std, - mode_scale=self.training_args.mode_scale, - ) - indices = (u * self.noise_scheduler.config.num_train_timesteps).long() - timesteps = self.noise_scheduler.timesteps[indices].to( - device=latents.device) + timesteps = self._sample_timesteps(batch_size, latents.device) + + # Enable training for the model that will be trained next and disable the other + if self.train_transformer_2: + self._enable_training(self.transformer_2, self.optimizer_2) + self._disable_training(self.transformer, self.optimizer) + else: + self._enable_training(self.transformer, self.optimizer) + if self.transformer_2 is not None: + self._disable_training(self.transformer_2, self.optimizer_2) + if self.training_args.sp_size > 1: # Make sure that the timesteps are the same across all sp processes. sp_group = get_sp_group() @@ -257,6 +305,38 @@ class TrainingPipeline(LoRAPipeline, ABC): return training_batch + def _sample_timesteps(self, batch_size, device): + # Determine which model to train based on the boundary timestep + if (self.transformer_2 is not None and self.boundary_timestep is not None and + torch.rand(1, generator=self.noise_random_generator).item() > self.training_args.boundary_ratio): + self.train_transformer_2 = True + else: + self.train_transformer_2 = False + + # Broadcast the decision to all processes + decision = torch.tensor(1.0 if self.train_transformer_2 else 0.0, device=self.device) + dist.broadcast(decision, src=0) + self.train_transformer_2 = decision.item() == 1.0 + + # Sample u from the appropriate range + u = compute_density_for_timestep_sampling( + weighting_scheme=self.training_args.weighting_scheme, + batch_size=batch_size, + generator=self.noise_random_generator, + logit_mean=self.training_args.logit_mean, + logit_std=self.training_args.logit_std, + mode_scale=self.training_args.mode_scale, + ) + + boundary_ratio = self.training_args.boundary_ratio + if self.train_transformer_2: + u = boundary_ratio + u * (1.0 - boundary_ratio) + else: + u = u * boundary_ratio + + indices = (u * self.noise_scheduler.config.num_train_timesteps).long() + return self.noise_scheduler.timesteps[indices].to(device=device) + def _build_attention_metadata( self, training_batch: TrainingBatch) -> TrainingBatch: latents_shape = training_batch.raw_latent_shape @@ -307,11 +387,12 @@ class TrainingPipeline(LoRAPipeline, ABC): # [1000.0], # device=training_batch.noisy_model_input.device, # dtype=torch.bfloat16) + current_model = self.transformer_2 if self.train_transformer_2 else self.transformer with set_forward_context( current_timestep=training_batch.current_timestep, attn_metadata=training_batch.attn_metadata): - model_pred = self.transformer(**input_kwargs) + model_pred = current_model(**input_kwargs) if self.training_args.precondition_outputs: assert training_batch.sigmas is not None model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas @@ -342,7 +423,12 @@ class TrainingPipeline(LoRAPipeline, ABC): # the following: # grad_norm = transformer.clip_grad_norm_(max_grad_norm) if max_grad_norm is not None: - model_parts = [self.transformer] + # Only clip gradients for the model that is currently training + if self.train_transformer_2 and self.transformer_2 is not None: + model_parts = [self.transformer_2] + else: + model_parts = [self.transformer] + grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases( [p for m in model_parts for p in m.parameters()], max_grad_norm, @@ -387,9 +473,14 @@ class TrainingPipeline(LoRAPipeline, ABC): training_batch = self._clip_grad_norm(training_batch) - self.optimizer.step() - self.lr_scheduler.step() - + # Only step the optimizer and scheduler for the model that is currently training + if self.train_transformer_2 and self.transformer_2 is not None: + self.optimizer_2.step() + self.lr_scheduler_2.step() + else: + self.optimizer.step() + self.lr_scheduler.step() + training_batch.total_loss = training_batch.total_loss training_batch.grad_norm = training_batch.grad_norm return training_batch diff --git a/fastvideo/training/wan_training_pipeline.py b/fastvideo/training/wan_training_pipeline.py index 1b3785197..a87e3de47 100644 --- a/fastvideo/training/wan_training_pipeline.py +++ b/fastvideo/training/wan_training_pipeline.py @@ -34,7 +34,7 @@ class WanTrainingPipeline(TrainingPipeline): def initialize_validation_pipeline(self, training_args: TrainingArgs): logger.info("Initializing validation pipeline...") args_copy = deepcopy(training_args) - + assert training_args.dit_cpu_offload == False args_copy.inference_mode = True validation_pipeline = WanPipeline.from_pretrained( training_args.model_path, @@ -47,7 +47,7 @@ class WanTrainingPipeline(TrainingPipeline): sp_size=training_args.sp_size, num_gpus=training_args.num_gpus, pin_cpu_memory=training_args.pin_cpu_memory, - dit_cpu_offload=True) + dit_cpu_offload=training_args.dit_cpu_offload) self.validation_pipeline = validation_pipeline