diff --git a/fastvideo/v1/forward_context.py b/fastvideo/v1/forward_context.py index 6fc196883..44eb533bc 100644 --- a/fastvideo/v1/forward_context.py +++ b/fastvideo/v1/forward_context.py @@ -11,10 +11,10 @@ import torch from fastvideo.v1.fastvideo_args import FastVideoArgs from fastvideo.v1.logger import init_logger -from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch if TYPE_CHECKING: from fastvideo.v1.attention import AttentionMetadata + from fastvideo.v1.pipelines import ForwardBatch logger = init_logger(__name__) @@ -37,13 +37,13 @@ class ForwardContext: # attn_layers: Dict[str, Any] # TODO: extend to support per-layer dynamic forward context attn_metadata: "AttentionMetadata" # set dynamically for each forward pass - forward_batch: Optional[ForwardBatch] = None + forward_batch: Optional["ForwardBatch"] = None -_forward_context: Optional[ForwardContext] = None +_forward_context: Optional["ForwardContext"] = None -def get_forward_context() -> ForwardContext: +def get_forward_context() -> "ForwardContext": """Get the current forward context.""" assert _forward_context is not None, ( "Forward context is not set. " @@ -55,7 +55,7 @@ def get_forward_context() -> ForwardContext: @contextmanager def set_forward_context(current_timestep, attn_metadata, - forward_batch: Optional[ForwardBatch] = None, + forward_batch: Optional["ForwardBatch"] = None, fastvideo_args: Optional[FastVideoArgs] = None): """A context manager that stores the current forward context, can be attention metadata, etc. diff --git a/fastvideo/v1/models/dits/wanvideo.py b/fastvideo/v1/models/dits/wanvideo.py index a09891e01..4a4cb6f4c 100644 --- a/fastvideo/v1/models/dits/wanvideo.py +++ b/fastvideo/v1/models/dits/wanvideo.py @@ -674,7 +674,7 @@ class WanTransformer3DModel(CachableDiT): # 5. Output norm, projection & unpatchify shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) - hidden_states = self.norm_out(hidden_states.float(), shift, scale) + hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, diff --git a/fastvideo/v1/models/utils.py b/fastvideo/v1/models/utils.py index 36642fa23..81c83cc91 100644 --- a/fastvideo/v1/models/utils.py +++ b/fastvideo/v1/models/utils.py @@ -51,11 +51,6 @@ def auto_attributes(init_func): return wrapper -def set_random_seed(seed: int) -> None: - from fastvideo.v1.platforms import current_platform - current_platform.seed_everything(seed) - - def set_weight_attrs( weight: torch.Tensor, weight_attrs: Optional[Dict[str, Any]], diff --git a/fastvideo/v1/pipelines/__init__.py b/fastvideo/v1/pipelines/__init__.py index a60a6c6db..4a78dec0f 100644 --- a/fastvideo/v1/pipelines/__init__.py +++ b/fastvideo/v1/pipelines/__init__.py @@ -11,7 +11,8 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs from fastvideo.v1.logger import init_logger from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.v1.pipelines.lora_pipeline import LoRAPipeline -from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.pipelines.pipeline_batch_info import (ForwardBatch, + TrainingBatch) from fastvideo.v1.pipelines.pipeline_registry import PipelineRegistry from fastvideo.v1.utils import (maybe_download_model, verify_model_config_and_directory) @@ -63,4 +64,5 @@ __all__ = [ "PipelineRegistry", "ForwardBatch", "LoRAPipeline", + "TrainingBatch", ] diff --git a/fastvideo/v1/pipelines/composed_pipeline_base.py b/fastvideo/v1/pipelines/composed_pipeline_base.py index e9997b1b8..55dae4ccf 100644 --- a/fastvideo/v1/pipelines/composed_pipeline_base.py +++ b/fastvideo/v1/pipelines/composed_pipeline_base.py @@ -298,3 +298,7 @@ class ComposedPipelineBase(ABC): # Return the output return batch + + def train(self) -> None: + raise NotImplementedError( + "if training_mode is True, the pipeline must implement this method") diff --git a/fastvideo/v1/pipelines/pipeline_batch_info.py b/fastvideo/v1/pipelines/pipeline_batch_info.py index 12d24dd77..4e7bd2993 100644 --- a/fastvideo/v1/pipelines/pipeline_batch_info.py +++ b/fastvideo/v1/pipelines/pipeline_batch_info.py @@ -13,6 +13,7 @@ from typing import Any, Dict, List, Optional, Union import torch +from fastvideo.v1.attention import AttentionMetadata from fastvideo.v1.configs.sample.teacache import (TeaCacheParams, WanTeaCacheParams) @@ -136,3 +137,30 @@ class ForwardBatch: def __str__(self): return pprint.pformat(asdict(self), indent=2, width=120) + + +@dataclass +class TrainingBatch: + current_timestep: int = 0 + current_vsa_sparsity: float = 0.0 + + # Dataloader batch outputs + latents: Optional[torch.Tensor] = None + encoder_hidden_states: Optional[torch.Tensor] = None + encoder_attention_mask: Optional[torch.Tensor] = None + info: Optional[Dict[str, Any]] = None + + # Transformer inputs + noisy_model_input: Optional[torch.Tensor] = None + timesteps: Optional[torch.Tensor] = None + sigmas: Optional[torch.Tensor] = None + noise: Optional[torch.Tensor] = None + + attn_metadata: Optional[AttentionMetadata] = None + + # Training loss + loss: torch.Tensor | None = None + + # Training outputs + total_loss: float | None = None + grad_norm: float | None = None diff --git a/fastvideo/v1/platforms/interface.py b/fastvideo/v1/platforms/interface.py index 8dc5297fb..a1c5c7ef0 100644 --- a/fastvideo/v1/platforms/interface.py +++ b/fastvideo/v1/platforms/interface.py @@ -169,6 +169,7 @@ class Platform: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) @classmethod def verify_model_arch(cls, model_arch: str) -> None: diff --git a/fastvideo/v1/tests/training/Vanilla/test_training_loss.py b/fastvideo/v1/tests/training/Vanilla/test_training_loss.py index 1b75725af..94d7b66b4 100644 --- a/fastvideo/v1/tests/training/Vanilla/test_training_loss.py +++ b/fastvideo/v1/tests/training/Vanilla/test_training_loss.py @@ -114,7 +114,7 @@ def test_distributed_training(): fields_and_thresholds = { 'avg_step_time': 1.0, - 'grad_norm': 0.1, + 'grad_norm': 0.2, 'step_time': 0.5, 'train_loss': 0.001 } diff --git a/fastvideo/v1/training/training_pipeline.py b/fastvideo/v1/training/training_pipeline.py index 0311254d4..0ae3f3d04 100644 --- a/fastvideo/v1/training/training_pipeline.py +++ b/fastvideo/v1/training/training_pipeline.py @@ -2,28 +2,43 @@ import gc import math import os +import time from abc import ABC, abstractmethod +from collections import deque from typing import Any, Dict, Iterator, List import imageio import numpy as np import torch import torchvision +from diffusers 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 (get_sp_group, get_torch_device, - get_world_group) +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 FastVideoArgs, TrainingArgs +from fastvideo.v1.forward_context import set_forward_context from fastvideo.v1.logger import init_logger -from fastvideo.v1.pipelines import ComposedPipelineBase -from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.pipelines import (ComposedPipelineBase, ForwardBatch, + TrainingBatch) +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__) @@ -45,6 +60,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC): def initialize_training_pipeline(self, training_args: TrainingArgs): logger.info("Initializing training pipeline...") + self.training_args = training_args self.device = get_torch_device() world_group = get_world_group() self.world_size = world_group.world_size @@ -122,17 +138,372 @@ class TrainingPipeline(ComposedPipelineBase, ABC): raise NotImplementedError( "Training pipelines must implement this method") - @abstractmethod - def train_one_step(self, transformer, model_type, optimizer, lr_scheduler, - loader, noise_scheduler, noise_random_generator, - gradient_accumulation_steps, sp_size, - precondition_outputs, max_grad_norm, weighting_scheme, - logit_mean, logit_std, mode_scale): - """ - Train one step of the model. - """ - raise NotImplementedError( - "Training pipeline must implement this method") + def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch: + self.transformer.requires_grad_(True) + self.transformer.train() + self.optimizer.zero_grad() + training_batch.total_loss = 0.0 + return training_batch + + def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: + assert self.train_loader_iter is not None + assert self.train_dataloader is not None + + 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 + training_batch.latents = latents.to(get_torch_device(), + dtype=torch.bfloat16) + training_batch.encoder_hidden_states = encoder_hidden_states.to( + get_torch_device(), dtype=torch.bfloat16) + training_batch.encoder_attention_mask = encoder_attention_mask.to( + get_torch_device(), dtype=torch.bfloat16) + training_batch.info = infos + + return training_batch + + def _normalize_dit_input(self, + training_batch: TrainingBatch) -> TrainingBatch: + # TODO(will): support other models + training_batch.latents = normalize_dit_input('wan', + training_batch.latents) + return training_batch + + def _prepare_dit_inputs(self, + training_batch: TrainingBatch) -> TrainingBatch: + assert self.training_args is not None + assert training_batch.latents is not None + assert training_batch.encoder_hidden_states is not None + assert training_batch.encoder_attention_mask is not None + assert self.noise_random_generator is not None + + batch_size = training_batch.latents.shape[0] + noise = torch.randn_like(training_batch.latents) + 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=training_batch.latents.device) + if self.training_args.sp_size > 1: + # Make sure that the timesteps are the same across all sp processes. + sp_group = get_sp_group() + sp_group.broadcast(timesteps, src=0) + sigmas = get_sigmas( + self.noise_scheduler, + training_batch.latents.device, + timesteps, + n_dim=training_batch.latents.ndim, + dtype=training_batch.latents.dtype, + ) + noisy_model_input = (1.0 - + sigmas) * training_batch.latents + sigmas * noise + + training_batch.noisy_model_input = noisy_model_input + training_batch.timesteps = timesteps + training_batch.sigmas = sigmas + training_batch.noise = noise + + return training_batch + + def _build_attention_metadata( + self, training_batch: TrainingBatch) -> TrainingBatch: + assert self.training_args is not None + latents = training_batch.latents + assert latents is not None + assert training_batch.timesteps is not None + patch_size = self.training_args.pipeline_config.dit_config.patch_size + current_vsa_sparsity = training_batch.current_vsa_sparsity + + if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": + + dit_seq_shape = [ + latents.shape[2] // patch_size[0], + latents.shape[3] // patch_size[1], + latents.shape[4] // patch_size[2] + ] + training_batch.attn_metadata = VideoSparseAttentionMetadata( + current_timestep=training_batch.timesteps, + dit_seq_shape=dit_seq_shape, + VSA_sparsity=current_vsa_sparsity) + else: + training_batch.attn_metadata = None + + return training_batch + + def _transformer_forward_and_compute_loss( + self, training_batch: TrainingBatch) -> TrainingBatch: + assert self.transformer is not None + assert self.training_args 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 + # assert training_batch.attn_metadata is not None + assert training_batch.latents is not None + assert training_batch.noise is not None + assert training_batch.sigmas is not None + + 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 + + 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, + } + # if 'hunyuan' in self.training_args.model_type: + # input_kwargs["guidance"] = torch.tensor( + # [1000.0], + # device=training_batch.noisy_model_input.device, + # dtype=torch.bfloat16) + + with set_forward_context( + current_timestep=training_batch.current_timestep, + attn_metadata=training_batch.attn_metadata): + model_pred = self.transformer(**input_kwargs) + if self.training_args.precondition_outputs: + model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas + target = training_batch.latents if self.training_args.precondition_outputs else training_batch.noise - training_batch.latents + + loss = (torch.mean((model_pred.float() - target.float())**2) / + self.training_args.gradient_accumulation_steps) + + loss.backward() + avg_loss = loss.detach().clone() + # logger.info(f"rank: {self.rank}, avg_loss: {avg_loss.item()}", + # local_main_process_only=False) + world_group = get_world_group() + world_group.all_reduce(avg_loss, op=torch.distributed.ReduceOp.AVG) + training_batch.total_loss += avg_loss.item() + + return training_batch + + def _clip_grad_norm(self, training_batch: TrainingBatch) -> TrainingBatch: + assert self.training_args is not None + max_grad_norm = self.training_args.max_grad_norm + + # TODO(will): perhaps move this into transformer api so that we can do + # the following: + # grad_norm = transformer.clip_grad_norm_(max_grad_norm) + if max_grad_norm is not None: + 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, + foreach=None, + ) + assert grad_norm is not float('nan') or grad_norm is not float( + 'inf') + grad_norm = grad_norm.item() if grad_norm is not None else 0.0 + else: + grad_norm = 0.0 + training_batch.grad_norm = grad_norm + return training_batch + + def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch: + assert self.training_args is not None + + training_batch = self._prepare_training(training_batch) + + for _ in range(self.training_args.gradient_accumulation_steps): + training_batch = self._get_next_batch(training_batch) + + # Shard latents across sp groups + training_batch.latents = shard_latents_across_sp( + training_batch.latents, + num_latent_t=self.training_args.num_latent_t) + + # Normalize DIT input + 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._transformer_forward_and_compute_loss( + training_batch) + + training_batch = self._clip_grad_norm(training_batch) + + 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 + + def _resume_from_checkpoint(self) -> None: + assert self.training_args is not None + logger.info("Loading checkpoint from %s", + self.training_args.resume_from_checkpoint) + resumed_step = load_checkpoint( + self.transformer, self.global_rank, + self.training_args.resume_from_checkpoint, self.optimizer, + self.train_dataloader, self.lr_scheduler, + self.noise_random_generator) + 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 train(self) -> None: + assert self.training_args is not None + + # Set random seeds for deterministic training + 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.transformer, self.training_args, 1) + + # Train! + progress_bar = tqdm( + range(0, self.training_args.max_train_steps), + initial=self.init_steps, + desc="Steps", + # Only show the progress bar once on each machine. + 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) + + loss = training_batch.total_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({ + "loss": f"{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_loss": 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.transformer, self.global_rank, + self.training_args.output_dir, step, + self.optimizer, self.train_dataloader, + self.lr_scheduler, self.noise_random_generator) + self.transformer.train() + self.sp_group.barrier() + if self.training_args.log_validation and step % self.training_args.validation_steps == 0: + self._log_validation(self.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.transformer, self.global_rank, + self.training_args.output_dir, + self.training_args.max_train_steps, self.optimizer, + self.train_dataloader, self.lr_scheduler, + self.noise_random_generator) + + if get_sp_group(): + cleanup_dist_env_and_memory() + + def _log_training_info(self) -> None: + assert self.training_args is not None + assert self.training_args.sp_size is not None + assert self.training_args.gradient_accumulation_steps is not None + total_batch_size = (self.world_size * + self.training_args.gradient_accumulation_steps / + self.training_args.sp_size * + self.training_args.train_sp_batch_size) + logger.info("***** Running training *****") + logger.info(" Num examples = %s", len(self.train_dataset)) + logger.info(" Dataloader size = %s", len(self.train_dataloader)) + logger.info(" Num Epochs = %s", self.num_train_epochs) + logger.info(" Resume training from step %s", + self.init_steps) # type: ignore + logger.info(" Instantaneous batch size per device = %s", + self.training_args.train_batch_size) + logger.info( + " Total train batch size (w. data & sequence parallel, accumulation) = %s", + total_batch_size) + logger.info(" Gradient Accumulation steps = %s", + self.training_args.gradient_accumulation_steps) + logger.info(" Total optimization steps = %s", + self.training_args.max_train_steps) + logger.info( + " Total training parameters per FSDP shard = %s B", + sum(p.numel() + for p in self.transformer.parameters() if p.requires_grad) / + 1e9) + # print dtype + logger.info(" Master weight dtype: %s", + self.transformer.parameters().__next__().dtype) + + gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2 + logger.info("GPU memory usage before train_one_step: %s MB", + gpu_memory_usage) + logger.info("VSA validation sparsity: %s", + self.training_args.VSA_sparsity) @torch.no_grad() def _log_validation(self, transformer, training_args, global_step) -> None: diff --git a/fastvideo/v1/training/wan_training_pipeline.py b/fastvideo/v1/training/wan_training_pipeline.py index 7f4d2a17a..6657a1d80 100644 --- a/fastvideo/v1/training/wan_training_pipeline.py +++ b/fastvideo/v1/training/wan_training_pipeline.py @@ -1,39 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -import importlib.util -import random import sys -import time -from collections import deque from copy import deepcopy -import numpy as np -import torch -from diffusers import FlowMatchEulerDiscreteScheduler -from tqdm.auto import tqdm - -import fastvideo.v1.envs as envs -from fastvideo.v1.attention.backends.video_sparse_attn import ( - VideoSparseAttentionMetadata) -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 FastVideoArgs, TrainingArgs -from fastvideo.v1.forward_context import set_forward_context from fastvideo.v1.logger import init_logger from fastvideo.v1.models.schedulers.scheduling_flow_unipc_multistep import ( FlowUniPCMultistepScheduler) -from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.v1.pipelines.wan.wan_pipeline import WanValidationPipeline 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 -import wandb # isort: skip - -vsa_available = False -if importlib.util.find_spec("vsa") is not None: - vsa_available = True +vsa_available = is_vsa_available() logger = init_logger(__name__) @@ -71,316 +48,6 @@ class WanTrainingPipeline(TrainingPipeline): self.validation_pipeline = validation_pipeline - def train_one_step( # type: ignore[override] - self, - transformer, - model_type, - optimizer, - lr_scheduler, - loader_iter, - noise_scheduler, - noise_random_generator, - gradient_accumulation_steps, - sp_size, - precondition_outputs, - max_grad_norm, - weighting_scheme, - logit_mean, - logit_std, - mode_scale, - patch_size, - current_vsa_sparsity, - ) -> tuple[float, float]: - assert self.training_args is not None - self.modules["transformer"].requires_grad_(True) - self.modules["transformer"].train() - - total_loss = 0.0 - optimizer.zero_grad() - - for _ in range(gradient_accumulation_steps): - # Get next batch, handling epoch boundaries gracefully - 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, _ = batch - - latents = latents.to(get_torch_device(), dtype=torch.bfloat16) - encoder_hidden_states = encoder_hidden_states.to( - get_torch_device(), dtype=torch.bfloat16) - latents = shard_latents_across_sp( - latents, num_latent_t=self.training_args.num_latent_t) - - dit_seq_shape = [ - latents.shape[2] // patch_size[0], - latents.shape[3] // patch_size[1], - latents.shape[4] // patch_size[2] - ] - - latents = normalize_dit_input(model_type, latents) - batch_size = latents.shape[0] - noise = torch.randn_like(latents) - u = compute_density_for_timestep_sampling( - weighting_scheme=weighting_scheme, - batch_size=batch_size, - generator=noise_random_generator, - logit_mean=logit_mean, - logit_std=logit_std, - mode_scale=mode_scale, - ) - indices = (u * noise_scheduler.config.num_train_timesteps).long() - timesteps = noise_scheduler.timesteps[indices].to( - device=latents.device) - if sp_size > 1: - # Make sure that the timesteps are the same across all sp processes. - sp_group = get_sp_group() - sp_group.broadcast(timesteps, src=0) - sigmas = get_sigmas( - noise_scheduler, - latents.device, - timesteps, - n_dim=latents.ndim, - dtype=latents.dtype, - ) - noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise - with torch.autocast("cuda", dtype=torch.bfloat16): - input_kwargs = { - "hidden_states": noisy_model_input, - "encoder_hidden_states": encoder_hidden_states, - "timestep": timesteps, - "encoder_attention_mask": encoder_attention_mask, # B, L - "return_dict": False, - } - if 'hunyuan' in model_type: - input_kwargs["guidance"] = torch.tensor( - [1000.0], - device=noisy_model_input.device, - dtype=torch.bfloat16) - - if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": - attn_metadata = VideoSparseAttentionMetadata( - current_timestep=timesteps, - dit_seq_shape=dit_seq_shape, - VSA_sparsity=current_vsa_sparsity) - else: - attn_metadata = None - - with set_forward_context(current_timestep=timesteps, - attn_metadata=attn_metadata): - model_pred = transformer(**input_kwargs) - - if precondition_outputs: - model_pred = noisy_model_input - model_pred * sigmas - target = latents if precondition_outputs else noise - latents - - loss = (torch.mean((model_pred.float() - target.float())**2) / - gradient_accumulation_steps) - - loss.backward() - - avg_loss = loss.detach().clone() - # logger.info(f"rank: {self.rank}, avg_loss: {avg_loss.item()}", - # local_main_process_only=False) - world_group = get_world_group() - world_group.all_reduce(avg_loss, op=torch.distributed.ReduceOp.AVG) - total_loss += avg_loss.item() - - # TODO(will): perhaps move this into transformer api so that we can do - # the following: - # grad_norm = transformer.clip_grad_norm_(max_grad_norm) - if max_grad_norm is not None: - 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, - foreach=None, - ) - assert grad_norm is not float('nan') or grad_norm is not float( - 'inf') - grad_norm = grad_norm.item() if grad_norm is not None else 0.0 - else: - grad_norm = 0.0 - - optimizer.step() - lr_scheduler.step() - return total_loss, grad_norm - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ): - assert self.training_args is not None - - # Set random seeds for deterministic training - seed = self.training_args.seed if self.training_args.seed is not None else 42 - - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - noise_random_generator = torch.Generator(device="cpu").manual_seed(seed) - - logger.info("Initialized random seeds with seed: %s", seed) - - noise_scheduler = FlowMatchEulerDiscreteScheduler() - - # Train! - assert self.training_args.sp_size is not None - assert self.training_args.gradient_accumulation_steps is not None - total_batch_size = (self.world_size * - self.training_args.gradient_accumulation_steps / - self.training_args.sp_size * - self.training_args.train_sp_batch_size) - logger.info("***** Running training *****") - logger.info(" Num examples = %s", len(self.train_dataset)) - logger.info(" Dataloader size = %s", len(self.train_dataloader)) - logger.info(" Num Epochs = %s", self.num_train_epochs) - logger.info(" Resume training from step %s", - self.init_steps) # type: ignore - logger.info(" Instantaneous batch size per device = %s", - self.training_args.train_batch_size) - logger.info( - " Total train batch size (w. data & sequence parallel, accumulation) = %s", - total_batch_size) - logger.info(" Gradient Accumulation steps = %s", - self.training_args.gradient_accumulation_steps) - logger.info(" Total optimization steps = %s", - self.training_args.max_train_steps) - logger.info( - " Total training parameters per FSDP shard = %s B", - sum(p.numel() - for p in self.transformer.parameters() if p.requires_grad) / - 1e9) - # print dtype - logger.info(" Master weight dtype: %s", - self.transformer.parameters().__next__().dtype) - - if self.training_args.resume_from_checkpoint: - logger.info("Loading checkpoint from %s", - self.training_args.resume_from_checkpoint) - resumed_step = load_checkpoint( - self.transformer, self.global_rank, - self.training_args.resume_from_checkpoint, self.optimizer, - self.train_dataloader, self.lr_scheduler, - noise_random_generator) - 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 - - progress_bar = tqdm( - range(0, self.training_args.max_train_steps), - initial=self.init_steps, - desc="Steps", - # Only show the progress bar once on each machine. - disable=self.local_rank > 0, - ) - - self.train_loader_iter = iter(self.train_dataloader) - - step_times: deque[float] = deque(maxlen=100) - - # TODO(will): fix this - # for i in range(self.init_steps): - # next(loader_iter) - # get gpu memory usage - gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2 - logger.info("GPU memory usage before train_one_step: %s MB", - gpu_memory_usage) - logger.info("VSA validation sparsity: %s", - self.training_args.VSA_sparsity) - self._log_validation(self.transformer, self.training_args, 1) - 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 - - for step in range(self.init_steps + 1, - self.training_args.max_train_steps + 1): - start_time = time.perf_counter() - if vsa_available: - 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 - loss, grad_norm = self.train_one_step( - self.transformer, - # args.model_type, - "wan", - self.optimizer, - self.lr_scheduler, - self.train_loader_iter, - noise_scheduler, - noise_random_generator, - self.training_args.gradient_accumulation_steps, - self.training_args.sp_size, - self.training_args.precondition_outputs, - self.training_args.max_grad_norm, - self.training_args.weighting_scheme, - self.training_args.logit_mean, - self.training_args.logit_std, - self.training_args.mode_scale, - self.training_args.pipeline_config.dit_config.patch_size, - current_vsa_sparsity, - ) - - 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({ - "loss": f"{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_loss": 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.transformer, self.global_rank, - self.training_args.output_dir, step, - self.optimizer, self.train_dataloader, - self.lr_scheduler, noise_random_generator) - self.transformer.train() - self.sp_group.barrier() - if self.training_args.log_validation and step % self.training_args.validation_steps == 0: - self._log_validation(self.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.transformer, self.global_rank, - self.training_args.output_dir, - self.training_args.max_train_steps, self.optimizer, - self.train_dataloader, self.lr_scheduler, - noise_random_generator) - - if get_sp_group(): - cleanup_dist_env_and_memory() - def main(args) -> None: logger.info("Starting training pipeline...") @@ -388,7 +55,7 @@ def main(args) -> None: pipeline = WanTrainingPipeline.from_pretrained( args.pretrained_model_name_or_path, args=args) args = pipeline.training_args - pipeline.forward(None, args) + pipeline.train() logger.info("Training pipeline done") diff --git a/fastvideo/v1/utils.py b/fastvideo/v1/utils.py index db102b28a..2abd2ab7b 100644 --- a/fastvideo/v1/utils.py +++ b/fastvideo/v1/utils.py @@ -5,6 +5,7 @@ import argparse import ctypes import hashlib import importlib +import importlib.util import inspect import json import math @@ -16,7 +17,7 @@ import tempfile import threading import traceback from dataclasses import dataclass, fields, is_dataclass -from functools import partial, wraps +from functools import lru_cache, partial, wraps from typing import (Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union, cast) @@ -785,3 +786,13 @@ def dict_to_3d_list( # else: silently ignore any key that doesn’t fit return result + + +def set_random_seed(seed: int) -> None: + from fastvideo.v1.platforms import current_platform + current_platform.seed_everything(seed) + + +@lru_cache(maxsize=1) +def is_vsa_available() -> bool: + return importlib.util.find_spec("vsa") is not None