[misc] [training] Reorganize training pipeline (#533)

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.

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,

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]],

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",
 ]

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")

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

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:

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
     }

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:

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")
 
 

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