load different models for dmd distill

fastvideo/v1/fastvideo_args.py

@@ -396,7 +396,8 @@ class TrainingArgs(FastVideoArgs):
 
     # text encoder & vae & diffusion model
     pretrained_model_name_or_path: str = ""
-    dit_model_name_or_path: str = ""
+    override_dit_model_name_or_path: str = ""
+
 
     # diffusion setting
     ema_decay: float = 0.0
@@ -464,11 +465,17 @@ class TrainingArgs(FastVideoArgs):
     VSA_decay_rate: float = 0.01  # decay rate -> 0.02
     VSA_decay_interval_steps: int = 1  # decay interval steps -> 50
     
-    # distillation args
+    # DMD distillation args
     student_critic_update_ratio: int = 5
     min_step_ratio: float = 0.2
     max_step_ratio: float = 0.98
     teacher_guidance_scale: float = 3.5
+    override_teacher_dit_model_name_or_path: str = ""
+    override_critic_dit_model_name_or_path: str = ""
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.check_training_args()
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
@@ -534,7 +541,7 @@ class TrainingArgs(FastVideoArgs):
                             type=str,
                             required=True,
                             help="Path to pretrained model or model name")
-        parser.add_argument("--dit-model-name-or-path",
+        parser.add_argument("--override-dit-model-name-or-path",
                             type=str,
                             required=False,
                             help="Path to DiT model or model name")
@@ -736,7 +743,7 @@ class TrainingArgs(FastVideoArgs):
             default=TrainingArgs.VSA_decay_interval_steps,
             help="VSA decay interval steps")
         
-        # Distillation arguments
+        # DMD Distillation arguments
         parser.add_argument("--student-critic-update-ratio",
             type=int,
             default=TrainingArgs.student_critic_update_ratio,
@@ -753,11 +760,21 @@ class TrainingArgs(FastVideoArgs):
             type=float,
             default=TrainingArgs.teacher_guidance_scale,
             help="Teacher guidance scale")
+        parser.add_argument("--override-teacher-dit-model-name-or-path",
+            type=str,
+            help="Override teacher DiT model name or path")
+        parser.add_argument("--override-critic-dit-model-name-or-path",
+            type=str,
+            help="Override critic DiT model name or path")
 
         return parser
 
+    def check_training_args(self) -> None:
+        pass
+
+
 def parse_int_list(value: str) -> List[int]:
     """Parse a comma-separated string of integers into a list."""
     if not value:
         return []
-    return [int(x.strip()) for x in value.split(",")]
+    return [int(x.strip()) for x in value.split(",")]

fastvideo/v1/models/loader/component_loader.py

@@ -556,6 +556,5 @@ class PipelineComponentLoader:
         # Get the appropriate loader for this module type
         loader = ComponentLoader.for_module_type(module_name,
                                                  transformers_or_diffusers)
-
         # Load the module
         return loader.load(component_model_path, fastvideo_args)

fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py

@@ -587,301 +587,3 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
 
     def __len__(self):
         return self.config.num_train_timesteps
-
-class FlowMatchScheduler():
-    order = 1
-    def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False):
-        self.num_train_timesteps = num_train_timesteps
-        self.shift = shift
-        self.sigma_max = sigma_max
-        self.sigma_min = sigma_min
-        self.inverse_timesteps = inverse_timesteps
-        self.extra_one_step = extra_one_step
-        self.reverse_sigmas = reverse_sigmas
-        self.set_timesteps(num_inference_steps)
-
-    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False):
-        sigma_start = self.sigma_min + \
-            (self.sigma_max - self.sigma_min) * denoising_strength
-        if self.extra_one_step:
-            self.sigmas = torch.linspace(
-                sigma_start, self.sigma_min, num_inference_steps + 1)[:-1]
-        else:
-            self.sigmas = torch.linspace(
-                sigma_start, self.sigma_min, num_inference_steps)
-        if self.inverse_timesteps:
-            self.sigmas = torch.flip(self.sigmas, dims=[0])
-        self.sigmas = self.shift * self.sigmas / \
-            (1 + (self.shift - 1) * self.sigmas)
-        if self.reverse_sigmas:
-            self.sigmas = 1 - self.sigmas
-        self.timesteps = self.sigmas * self.num_train_timesteps
-        if training:
-            x = self.timesteps
-            y = torch.exp(-2 * ((x - num_inference_steps / 2) /
-                          num_inference_steps) ** 2)
-            y_shifted = y - y.min()
-            bsmntw_weighing = y_shifted * \
-                (num_inference_steps / y_shifted.sum())
-            self.linear_timesteps_weights = bsmntw_weighing
-
-    def step(self, model_output, timestep, sample, to_final=False):
-        self.sigmas = self.sigmas.to(model_output.device)
-        self.timesteps = self.timesteps.to(model_output.device)
-        timestep_id = torch.argmin(
-            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
-        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
-        if to_final or (timestep_id + 1 >= len(self.timesteps)).any():
-            sigma_ = 1 if (
-                self.inverse_timesteps or self.reverse_sigmas) else 0
-        else:
-            sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
-        prev_sample = sample + model_output * (sigma_ - sigma)
-        return prev_sample
-
-    def add_noise(self, original_samples, noise, timestep):
-        """
-        Diffusion forward corruption process.
-        Input:
-            - clean_latent: the clean latent with shape [B, C, H, W]
-            - noise: the noise with shape [B, C, H, W]
-            - timestep: the timestep with shape [B]
-        Output: the corrupted latent with shape [B, C, H, W]
-        """
-        self.sigmas = self.sigmas.to(noise.device)
-        self.timesteps = self.timesteps.to(noise.device)
-        timestep_id = torch.argmin(
-            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
-        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) # [21, 1, 1, 1]
-        sample = (1 - sigma) * original_samples + sigma * noise
-        return sample.type_as(noise)
-
-    def training_target(self, sample, noise, timestep):
-        target = noise - sample
-        return target
-
-    def training_weight(self, timestep):
-        timestep_id = torch.argmin(
-            (self.timesteps - timestep.to(self.timesteps.device)).abs())
-        weights = self.linear_timesteps_weights[timestep_id]
-        return weights
-    
-# class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
-#     """
-#     Euler scheduler.
-
-#     This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
-#     methods the library implements for all schedulers such as loading and saving.
-
-#     Args:
-#         num_train_timesteps (`int`, defaults to 1000):
-#             The number of diffusion steps to train the model.
-#         timestep_spacing (`str`, defaults to `"linspace"`):
-#             The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
-#             Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
-#         shift (`float`, defaults to 1.0):
-#             The shift value for the timestep schedule.
-#         reverse (`bool`, defaults to `True`):
-#             Whether to reverse the timestep schedule.
-#     """
-
-#     _compatibles: list[Any] = []
-#     order = 1
-
-#     @register_to_config
-#     def __init__(
-#         self,
-#         num_train_timesteps: int = 1000,
-#         shift: float = 1.0,
-#         reverse: bool = True,
-#         solver: str = "euler",
-#         n_tokens: Optional[int] = None,
-#         **kwargs,
-#     ):
-#         sigmas = torch.linspace(1, 0, num_train_timesteps + 1)
-
-#         if not reverse:
-#             sigmas = sigmas.flip(0)
-
-#         self.sigmas = sigmas
-#         # the value fed to model
-#         self.timesteps = (sigmas[:-1] *
-#                           num_train_timesteps).to(dtype=torch.float32)
-
-#         self._step_index: int | None = None
-#         self._begin_index = 0
-
-#         self.supported_solver = ["euler"]
-#         if solver not in self.supported_solver:
-#             raise ValueError(
-#                 f"Solver {solver} not supported. Supported solvers: {self.supported_solver}"
-#             )
-
-#         BaseScheduler.__init__(self)
-
-#     @property
-#     def step_index(self):
-#         """
-#         The index counter for current timestep. It will increase 1 after each scheduler step.
-#         """
-#         return self._step_index
-
-#     @property
-#     def begin_index(self):
-#         """
-#         The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
-#         """
-#         return self._begin_index
-
-#     # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
-#     def set_begin_index(self, begin_index: int = 0):
-#         """
-#         Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
-
-#         Args:
-#             begin_index (`int`):
-#                 The begin index for the scheduler.
-#         """
-#         self._begin_index = begin_index
-
-#     def _sigma_to_t(self, sigma):
-#         return sigma * self.config.num_train_timesteps
-
-#     def set_timesteps(
-#         self,
-#         num_inference_steps: int,
-#         device: Union[str, torch.device] = None,
-#         n_tokens: int = 0,
-#     ):
-#         """
-#         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
-
-#         Args:
-#             num_inference_steps (`int`):
-#                 The number of diffusion steps used when generating samples with a pre-trained model.
-#             device (`str` or `torch.device`, *optional*):
-#                 The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-#             n_tokens (`int`, *optional*):
-#                 Number of tokens in the input sequence.
-#         """
-#         self.num_inference_steps = num_inference_steps
-
-#         sigmas = torch.linspace(1, 0, num_inference_steps + 1)
-#         sigmas = self.sd3_time_shift(sigmas)
-
-#         if not self.config.reverse:
-#             sigmas = 1 - sigmas
-
-#         self.sigmas = sigmas
-#         if not getattr(self.config, "timesteps_scale", True):
-#             self.timesteps = sigmas[:-1]  # for stepvideo
-#         else:
-#             self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(
-#                 dtype=torch.float32, device=device)
-#         # Reset step index
-#         self._step_index = None
-
-#     def index_for_timestep(self, timestep, schedule_timesteps=None) -> int:
-#         if schedule_timesteps is None:
-#             schedule_timesteps = self.timesteps
-
-#         indices = (schedule_timesteps == timestep).nonzero()
-
-#         # The sigma index that is taken for the **very** first `step`
-#         # is always the second index (or the last index if there is only 1)
-#         # This way we can ensure we don't accidentally skip a sigma in
-#         # case we start in the middle of the denoising schedule (e.g. for image-to-image)
-#         pos = 1 if len(indices) > 1 else 0
-
-#         idx: int = indices[pos].item()
-
-#         return idx
-
-#     def set_shift(self, shift: float) -> None:
-#         self.config.shift = shift
-
-#     def set_timesteps_scale(self, timesteps_scale: bool) -> None:
-#         self.config.timesteps_scale = timesteps_scale
-
-#     def _init_step_index(self, timestep) -> None:
-#         if self.begin_index is None:
-#             if isinstance(timestep, torch.Tensor):
-#                 timestep = timestep.to(self.timesteps.device)
-#             self._step_index = self.index_for_timestep(timestep)
-#         else:
-#             self._step_index = self._begin_index
-
-#     def scale_model_input(self,
-#                           sample: torch.Tensor,
-#                           timestep: Optional[int] = None) -> torch.Tensor:
-#         return sample
-
-#     def sd3_time_shift(self, t: torch.Tensor):
-#         return (self.config.shift * t) / (1 + (self.config.shift - 1) * t)
-
-#     def step(
-#         self,
-#         model_output: torch.FloatTensor,
-#         timestep: Union[float, torch.FloatTensor],
-#         sample: torch.FloatTensor,
-#         return_dict: bool = True,
-#         **kwargs,
-#     ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]:
-#         """
-#         Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
-#         process from the learned model outputs (most often the predicted noise).
-
-#         Args:
-#             model_output (`torch.FloatTensor`):
-#                 The direct output from learned diffusion model.
-#             timestep (`float`):
-#                 The current discrete timestep in the diffusion chain.
-#             sample (`torch.FloatTensor`):
-#                 A current instance of a sample created by the diffusion process.
-#             generator (`torch.Generator`, *optional*):
-#                 A random number generator.
-#             n_tokens (`int`, *optional*):
-#                 Number of tokens in the input sequence.
-#             return_dict (`bool`):
-#                 Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
-#                 tuple.
-
-#         Returns:
-#             [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
-#                 If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
-#                 returned, otherwise a tuple is returned where the first element is the sample tensor.
-#         """
-
-#         if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
-#             raise ValueError((
-#                 "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
-#                 " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
-#                 " one of the `scheduler.timesteps` as a timestep."), )
-
-#         if self.step_index is None:
-#             self._init_step_index(timestep)
-
-#         # Upcast to avoid precision issues when computing prev_sample
-#         sample = sample.to(torch.float32)
-
-#         assert self.step_index is not None
-#         dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index]
-
-#         if self.config.solver == "euler":
-#             prev_sample = sample + model_output.to(torch.float32) * dt
-#         else:
-#             raise ValueError(
-#                 f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}"
-#             )
-
-#         # upon completion increase step index by one
-#         assert self._step_index is not None
-#         self._step_index += 1
-
-#         if not return_dict:
-#             return (prev_sample, )
-
-#         return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample)
-    
-#     def __len__(self):
-#         return self.config.num_train_timesteps

fastvideo/v1/pipelines/composed_pipeline_base.py

@@ -51,6 +51,7 @@ class ComposedPipelineBase(ABC):
         Initialize the pipeline. After __init__, the pipeline should be ready to
         use. The pipeline should be stateless and not hold any batch state.
         """
+
         self.fastvideo_args = fastvideo_args
 
         self.model_path: str = model_path
@@ -122,7 +123,7 @@ class ComposedPipelineBase(ABC):
             for key, value in kwargs.items():
                 setattr(fastvideo_args, key, value)
 
-            fastvideo_args.use_cpu_offload = False
+            # fastvideo_args.use_cpu_offload = False
             # make sure we are in training mode
             fastvideo_args.inference_mode = False
             # we hijack the precision to be the master weight type so that the
@@ -130,7 +131,7 @@ class ComposedPipelineBase(ABC):
             # use FSDP2's MixedPrecisionPolicy to set the precision for the
             # fwd, bwd, and other operations' precision.
             # fastvideo_args.precision = fastvideo_args.master_weight_type
-            assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training'
+            # assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training'
             # assert fastvideo_args.precision == 'fp32', 'only fp32 is supported for training'
 
         logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args)
@@ -200,6 +201,33 @@ class ComposedPipelineBase(ABC):
         """
         return
 
+    def _prepare_component_model_metadata_for_loading(self, model_index: dict[str, Any], fastvideo_args: FastVideoArgs) -> dict[str, Any]:
+        """
+        Prepare the model index for loading. Checks the following:
+            - If the model index does not contain a module, raise an error.
+            - If the module is in the required_config_modules, check if the module is in the model_index.
+        Can be overridden by subclasses to add additional checks or modify model index.
+        """
+        # some sanity checks
+        assert len(
+            model_index
+        ) > 1, "model_index.json must contain at least one pipeline module"
+
+        for module_name in self.required_config_modules:
+            if module_name not in model_index:
+                raise ValueError(
+                    f"model_index.json must contain a {module_name} module")
+
+        component_model_metadata= {}
+        for module_name, (transformers_or_diffusers, architecture) in model_index.items():
+            component_model_path = os.path.join(self.model_path, module_name)
+            if not os.path.exists(component_model_path):
+                raise ValueError(
+                    f"Component model path {component_model_path} does not exist")
+            component_model_metadata[module_name] = (transformers_or_diffusers, architecture, component_model_path)
+
+        return component_model_metadata
+
     def load_modules(
         self,
         fastvideo_args: FastVideoArgs,
@@ -218,24 +246,15 @@ class ComposedPipelineBase(ABC):
         model_index.pop("_class_name")
         model_index.pop("_diffusers_version")
 
-        # some sanity checks
-        assert len(
-            model_index
-        ) > 1, "model_index.json must contain at least one pipeline module"
+        component_model_metadata = self._prepare_component_model_metadata_for_loading(model_index, fastvideo_args)
         
-        for module_name in self.required_config_modules:
-            if module_name not in model_index:
-                logger.warning(
-                    f"model_index.json does not contain a {module_name} module, adding {module_name} to model_index")
-                if 'transformer' in module_name:
-                    model_index[module_name] = model_index['transformer']
         # all the component models used by the pipeline
         required_modules = self.required_config_modules
         logger.info("Loading required modules: %s", required_modules)
 
         modules = {}
         for module_name, (transformers_or_diffusers,
-                          architecture) in model_index.items():
+                          architecture, component_model_path) in component_model_metadata.items():
             if module_name not in required_modules:
                 logger.info("Skipping module %s", module_name)
                 continue
@@ -243,11 +262,6 @@ class ComposedPipelineBase(ABC):
                 logger.info("Using module %s already provided", module_name)
                 modules[module_name] = loaded_modules[module_name]
                 continue
-            if 'transformer' in module_name:
-                loading_module_name = module_name.split("_")[-1]
-            else:
-                loading_module_name = module_name
-            component_model_path = os.path.join(self.model_path, loading_module_name)
             module = PipelineComponentLoader.load_module(
                 module_name=module_name,
                 component_model_path=component_model_path,

fastvideo/v1/pipelines/pipeline_batch_info.py

@@ -156,6 +156,7 @@ class TrainingBatch:
 
     # Dataloader batch outputs
     latents: torch.Tensor | None = None
+    raw_latent_shape: torch.Tensor | None = None
     noise_latents: torch.Tensor | None = None
     encoder_hidden_states: torch.Tensor | None = None
     encoder_attention_mask: torch.Tensor | None = None
@@ -194,6 +195,7 @@ class TrainingBatch:
     # Distillation losses
     student_loss: float = 0.0
     critic_loss: float = 0.0
+    regression_loss: float = 0.0
     
     # Training control
     dmd_log_dict: dict[str, Any] = field(default_factory=dict)

fastvideo/v1/pipelines/stages/denoising.py

@@ -106,13 +106,13 @@ class DenoisingStage(PipelineStage):
         sp_group = sp_world_size > 1
         if sp_group:
             latents = rearrange(batch.latents,
-                                "b t (n s) h w -> b t n s h w",
+                                "b c (n t) h w -> b c n t h w",
                                 n=sp_world_size).contiguous()
             latents = latents[:, :, rank_in_sp_group, :, :, :]
             batch.latents = latents
             if batch.image_latent is not None:
                 image_latent = rearrange(batch.image_latent,
-                                         "b t (n s) h w -> b t n s h w",
+                                         "b c (n t) h w -> b c n t h w",
                                          n=sp_world_size).contiguous()
                 image_latent = image_latent[:, :, rank_in_sp_group, :, :, :]
                 batch.image_latent = image_latent
@@ -585,22 +585,6 @@ class DmdDenoisingStage(DenoisingStage):
         autocast_enabled = (target_dtype != torch.float32
                             ) and not fastvideo_args.disable_autocast
 
-        # Handle sequence parallelism if enabled
-        sp_world_size, rank_in_sp_group = get_sp_world_size(
-        ), get_sp_parallel_rank()
-        sp_group = sp_world_size > 1
-        if sp_group:
-            latents = rearrange(batch.latents,
-                                "b t (n s) h w -> b t n s h w",
-                                n=sp_world_size).contiguous()
-            latents = latents[:, :, rank_in_sp_group, :, :, :]
-            batch.latents = latents
-            if batch.image_latent is not None:
-                image_latent = rearrange(batch.image_latent,
-                                         "b t (n s) h w -> b t n s h w",
-                                         n=sp_world_size).contiguous()
-                image_latent = image_latent[:, :, rank_in_sp_group, :, :, :]
-                batch.image_latent = image_latent
         # Get timesteps and calculate warmup steps
         timesteps = batch.timesteps
 
@@ -652,12 +636,32 @@ class DmdDenoisingStage(DenoisingStage):
         latents = batch.latents
         # TODO(yongqi) hard code prepare latents
         latents = torch.randn(latents.permute(0, 2, 1, 3, 4).shape, dtype=torch.bfloat16, device="cuda", generator=torch.Generator(device="cuda").manual_seed(42))
-
+        video_raw_latent_shape = latents.shape
         prompt_embeds = batch.prompt_embeds
         assert torch.isnan(prompt_embeds[0]).sum() == 0
         timesteps = torch.tensor(
             fastvideo_args.denoising_step_list, dtype=torch.long, device=get_local_torch_device())
 
+        # Handle sequence parallelism if enabled
+        sp_world_size, rank_in_sp_group = get_sp_world_size(
+        ), get_sp_parallel_rank()
+        sp_group = sp_world_size > 1
+        if sp_group:
+            latents = rearrange(latents,
+                                "b (n t) c h w -> b n t c h w",
+                                n=sp_world_size).contiguous()
+            latents = latents[:, rank_in_sp_group, :, :, :, :]
+            if batch.image_latent is not None:
+                image_latent = rearrange(batch.image_latent,
+                                         "b c (n t) h w -> b c n t h w",
+                                         n=sp_world_size).contiguous()
+
+                image_latent = image_latent[:, :, rank_in_sp_group, :, :, :]
+                batch.image_latent = image_latent
+
+        # timesteps = batch.timesteps
+        # num_inference_steps = batch.num_inference_steps
+        
         # Run denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
@@ -746,13 +750,20 @@ class DmdDenoisingStage(DenoisingStage):
                     if i < len(timesteps) - 1:
                         next_timestep = timesteps[i + 1] * torch.ones(
                             pred_video.shape[:2], dtype=torch.long, device=pred_video.device)
+                        noise = torch.randn(
+                            video_raw_latent_shape, device=self.device, dtype=pred_video.dtype)
+                        if sp_group:
+                            noise = rearrange(noise,
+                                                "b (n t) c h w -> b n t c h w",
+                                                n=sp_world_size).contiguous()
+                            noise = noise[:, rank_in_sp_group, :, :, :, :]
                         latents = self.scheduler.add_noise(
                             pred_video.flatten(0, 1),
-                            torch.randn_like(pred_video.flatten(0, 1)),
+                            noise.flatten(0, 1),
                             next_timestep.flatten(0, 1)
                         ).unflatten(0, pred_video.shape[:2])
                     else:
-                        latents = pred_video.permute(0, 2, 1, 3, 4)
+                        latents = pred_video
 
                     # Update progress bar
                     if i == len(timesteps) - 1 or (
@@ -763,8 +774,8 @@ class DmdDenoisingStage(DenoisingStage):
 
         # Gather results if using sequence parallelism
         if sp_group:
-            latents = sequence_model_parallel_all_gather(latents, dim=2)
-
+            latents = sequence_model_parallel_all_gather(latents, dim=1)
+        latents = latents.permute(0, 2, 1, 3, 4)
         # Update batch with final latents
         batch.latents = latents
 

fastvideo/v1/training/distillation_pipeline.py

@@ -22,7 +22,7 @@ import fastvideo.v1.envs as envs
 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_local_torch_device, get_world_group)
+                                      get_local_torch_device, get_world_group, get_sp_parallel_rank, get_sp_world_size)
 from fastvideo.v1.fastvideo_args import FastVideoArgs,TrainingArgs
 from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.logger import init_logger
@@ -32,15 +32,14 @@ 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, prepare_for_saving)
-from fastvideo.v1.utils import set_random_seed, is_vsa_available
-from fastvideo.v1.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler, FlowMatchScheduler
+    normalize_dit_input, save_checkpoint, prepare_for_saving)
+from fastvideo.v1.utils import set_random_seed, is_vsa_available, maybe_download_model
+from fastvideo.v1.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
 from fastvideo.v1.training.activation_checkpoint import (
     apply_activation_checkpointing)
-from fastvideo.v1.attention.backends.video_sparse_attn import (
-    VideoSparseAttentionMetadata)
 from fastvideo.v1.dataset.validation_dataset import ValidationDataset
 from fastvideo.v1.training.training_utils import DiffusionWrapper
+
 import wandb  # isort: skip
 
 vsa_available = is_vsa_available()
@@ -59,6 +58,93 @@ class DistillationPipeline(TrainingPipeline):
                                       Dict[str, Any]]]
     current_epoch: int = 0
 
+    def _prepare_component_model_metadata_for_loading(self, model_index: dict[str, Any], fastvideo_args: FastVideoArgs) -> dict[str, Any]:
+        """
+        For DMD distillation, we need to override model index to include the
+        teacher and critic transformer model names.  By default, the teacher and
+        critic transformer will be the same as the student transformer
+        (transformer). However, if the user provides
+        training_args.override_teacher_dit_model_name_or_path and
+        training_args.override_critic_dit_model_name_or_path, we will override
+        the teacher and critic transformer model names.
+        """
+        # some sanity checks
+        assert len(
+            model_index
+        ) > 1, "model_index.json must contain at least one pipeline module"
+
+        for module_name in self.required_config_modules:
+            if module_name not in model_index:
+                logger.warning(
+                    f"model_index.json does not contain a {module_name} module, adding {module_name} to model_index")
+                if 'transformer' in module_name:
+                    model_index[module_name] = model_index['transformer']
+
+        component_model_metadata= {}
+        for module_name, (transformers_or_diffusers, architecture) in model_index.items():
+            if 'transformer' in module_name:
+                loading_module_name = module_name.split("_")[-1]
+            else:
+                loading_module_name = module_name
+            component_model_path = os.path.join(self.model_path, loading_module_name)
+            if not os.path.exists(component_model_path):
+                raise ValueError(
+                    f"Component model path {component_model_path} does not exist")
+            component_model_metadata[module_name] = (transformers_or_diffusers, architecture, component_model_path)
+
+        if fastvideo_args.override_teacher_dit_model_name_or_path:
+            # we accept the following formats:
+            # 1. a HF string to a diffusers model
+            # 5. a local path to a directory for DiT model
+            # 6. a local path to a diffusers model
+
+            try:
+                local_path = maybe_download_model(fastvideo_args.override_teacher_dit_model_name_or_path)
+            except ValueError as e:
+                raise ValueError(f"Failed to download teacher model {fastvideo_args.override_teacher_dit_model_name_or_path}: {e}")
+            
+            config_json_path = os.path.join(local_path, "config.json")
+            model_index_json_path = os.path.join(local_path, "model_index.json")
+
+            if os.path.exists(model_index_json_path):
+                # This is a diffusers model directory
+                logger.info(f"Teacher model path {local_path} contains model_index.json - treating as diffusers model")
+                teacher_dit_path = os.path.join(local_path, "transformer")
+            elif os.path.exists(config_json_path):
+                # This is a DiT model directory
+                logger.info(f"Teacher model path {local_path} contains config.json - treating as DiT model")
+                teacher_dit_path = local_path
+            else:
+                raise ValueError(f"Teacher model directory {local_path} must contain either config.json or model_index.json")
+            
+            logger.info(f"Overriding teacher model path to {teacher_dit_path}")
+            component_model_metadata['teacher_transformer'] = (component_model_metadata['transformer'][0], component_model_metadata['transformer'][1], teacher_dit_path)
+        
+        if fastvideo_args.override_critic_dit_model_name_or_path:
+            try:
+                local_path = maybe_download_model(fastvideo_args.override_critic_dit_model_name_or_path)
+            except ValueError as e:
+                raise ValueError(f"Failed to download critic model {fastvideo_args.override_critic_dit_model_name_or_path}: {e}")
+            
+            config_json_path = os.path.join(local_path, "config.json")
+            model_index_json_path = os.path.join(local_path, "model_index.json")
+
+            if os.path.exists(model_index_json_path):
+                # This is a diffusers model directory
+                logger.info(f"Critic model path {local_path} contains model_index.json - treating as diffusers model")
+                critic_dit_path = os.path.join(local_path, "transformer")
+            elif os.path.exists(config_json_path):
+                # This is a DiT model directory
+                logger.info(f"Critic model path {local_path} contains config.json - treating as DiT model")
+                critic_dit_path = local_path
+            else:
+                raise ValueError(f"Critic model directory {local_path} must contain either config.json or model_index.json")
+            
+            logger.info(f"Overriding critic model path to {critic_dit_path}")
+            component_model_metadata['critic_transformer'] = (component_model_metadata['transformer'][0], component_model_metadata['transformer'][1], critic_dit_path)
+
+        return component_model_metadata
+
     def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
         raise RuntimeError(
             "create_pipeline_stages should not be called for training pipeline")
@@ -76,18 +162,17 @@ class DistillationPipeline(TrainingPipeline):
         self.vae.requires_grad_(False)
         
         self.timestep_shift = self.training_args.pipeline_config.flow_shift
-        # self.noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=self.timestep_shift)
-        self.noise_scheduler = FlowMatchScheduler(
-            shift=8.0, sigma_min=0.0, extra_one_step=True
-        )
-        self.noise_scheduler.set_timesteps(1000, training=True)
+        self.noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=self.timestep_shift)
         
         # 2. Distillation-specific initialization
         # The parent class already sets self.transformer as the student model
         self.student_transformer = DiffusionWrapper(self.transformer, self.noise_scheduler)
         self.teacher_transformer = DiffusionWrapper(self.get_module("teacher_transformer"), self.noise_scheduler)
         self.critic_transformer = DiffusionWrapper(self.get_module("critic_transformer"), self.noise_scheduler)
-        
+        # self.student_transformer.to(torch.bfloat16)
+        # self.teacher_transformer.to(torch.bfloat16)
+        # self.critic_transformer.to(torch.bfloat16)
+        # torch.distributed.breakpoint()
         self.teacher_transformer.requires_grad_(False)
         self.teacher_transformer.eval()
         self.critic_transformer.requires_grad_(True)
@@ -98,6 +183,10 @@ class DistillationPipeline(TrainingPipeline):
                 self.critic_transformer,
                 checkpointing_type=training_args.
                 enable_gradient_checkpointing_type)
+            self.teacher_transformer = apply_activation_checkpointing(
+                self.teacher_transformer,
+                checkpointing_type=training_args.
+                enable_gradient_checkpointing_type)
 
         # Initialize optimizers
         critic_params = list(filter(lambda p: p.requires_grad, self.critic_transformer.parameters()))
@@ -172,9 +261,11 @@ class DistillationPipeline(TrainingPipeline):
             assert timestep.shape[1] == 1
             return timestep
         elif type == "bidirectional_video":
+            # Create a new tensor to avoid in-place operations
+            new_timestep = timestep.clone()
             for index in range(timestep.shape[0]):
-                timestep[index] = timestep[index, 0]
-            return timestep
+                new_timestep[index] = timestep[index, 0]
+            return new_timestep
         elif type == "causal_video":
             # make the noise level the same within every motion block
             timestep = timestep.reshape(
@@ -194,16 +285,21 @@ class DistillationPipeline(TrainingPipeline):
             # Use cross-codebase generator for reproducible noise generation
             noise = torch.randn(
                 self.video_latent_shape, device=self.device, dtype=dtype)
+            if self.sp_world_size > 1:
+                noise = rearrange(noise,
+                                    "b (n t) c h w -> b n t c h w",
+                                    n=self.sp_world_size).contiguous()
+                noise = noise[:, self.rank_in_sp_group, :, :, :, :]
 
             noisy_timestep = timestep * torch.ones(
-                self.video_latent_shape[:2], device=self.device, dtype=torch.long)
+                self.video_latent_shape_sp[:2], device=self.device, dtype=torch.long)
 
             if timestep != 0:
                 noisy_video = self.noise_scheduler.add_noise(
                     latents.flatten(0, 1),
                     noise.flatten(0, 1),
                     noisy_timestep.flatten(0, 1)
-                ).unflatten(0, self.video_latent_shape[:2])
+                ).unflatten(0, self.video_latent_shape_sp[:2])
             else:
                 noisy_video = latents
 
@@ -214,7 +310,7 @@ class DistillationPipeline(TrainingPipeline):
         # Step 2: Randomly sample a timestep and pick the corresponding input
         # Use cross-codebase generator for reproducible index generation
         index = torch.randint(0, len(self.denoising_step_list), [
-                              self.video_latent_shape[0], self.video_latent_shape[1]], device=self.device, dtype=torch.long)
+                              self.video_latent_shape_sp[0], self.video_latent_shape_sp[1]], device=self.device, dtype=torch.long)
 
         index = self._process_timestep(index, type=self.distill_task_type)
 
@@ -223,16 +319,16 @@ class DistillationPipeline(TrainingPipeline):
         noisy_input = torch.gather(
             simulated_noisy_input, dim=1,
             index=index.reshape(index.shape[0], 1, index.shape[1], 1, 1, 1).expand(
-                -1, -1, -1, *self.video_latent_shape[2:])
+                -1, -1, -1, *self.video_latent_shape_sp[2:])
         ).squeeze(1)
 
         timestep = self.denoising_step_list[index]
 
         training_batch = self._build_input_kwargs(noisy_input, timestep, training_batch.conditional_dict, training_batch)
+
         pred_video = self.student_transformer(training_batch, timestep)
 
         pred_video = pred_video.type_as(noisy_input)
-
         return pred_video, timestep.float().detach()
 
     def _compute_kl_grad(
@@ -279,13 +375,13 @@ class DistillationPipeline(TrainingPipeline):
         """Compute DMD (Diffusion Model Distillation) loss."""
         
         original_latent = pred_video
-        batch_size, num_frame = self.video_latent_shape[:2]
+        batch_size, latent_t = self.video_latent_shape_sp[:2]
         with torch.no_grad():
             # Use cross-codebase generator for reproducible timestep generation
             timestep = torch.randint(
                 0,
                 self.num_train_timestep,
-                [batch_size, num_frame],
+                [batch_size, latent_t],
                 device=self.device,
                 dtype=torch.long
             )
@@ -301,12 +397,19 @@ class DistillationPipeline(TrainingPipeline):
             timestep = timestep.clamp(self.min_step, self.max_step)
 
             # Use cross-codebase generator for reproducible noise generation
-            noise = torch.randn_like(pred_video)
+            noise = torch.randn(
+                self.video_latent_shape, device=self.device, dtype=pred_video.dtype)
+            if self.sp_world_size > 1:
+                noise = rearrange(noise,
+                                    "b (n t) c h w -> b n t c h w",
+                                    n=self.sp_world_size).contiguous()
+                noise = noise[:, self.rank_in_sp_group, :, :, :, :]
+                
             noisy_latent = self.noise_scheduler.add_noise(
                 pred_video.flatten(0, 1),
                 noise.flatten(0, 1),
                 timestep.flatten(0, 1)
-            ).detach().unflatten(0, (batch_size, num_frame))
+            ).detach().unflatten(0, (batch_size, latent_t))
 
             grad, dmd_log_dict = self._compute_kl_grad(
                 noisy_video=noisy_latent,
@@ -322,9 +425,6 @@ class DistillationPipeline(TrainingPipeline):
 
     def _student_forward_and_compute_dmd_loss(self, training_batch: TrainingBatch) -> Tuple[TrainingBatch, torch.Tensor, dict]:
         """Forward pass through student transformer and compute student losses."""
-        assert training_batch.conditional_dict is not None
-        assert training_batch.unconditional_dict is not None
-        assert training_batch.latents is not None
         with set_forward_context(
                 current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata_vsa):
             pred_video, timestep_dmd = self._student_forward(training_batch)
@@ -335,25 +435,21 @@ class DistillationPipeline(TrainingPipeline):
                 pred_video=pred_video,
                 training_batch=training_batch
             )
-        
+
         dmd_log_dict['dmd_timestep_stu'] = timestep_dmd
-        
+
         return training_batch, dmd_loss, dmd_log_dict
 
     def _critic_forward_and_compute_loss(self, training_batch: TrainingBatch) -> Tuple[TrainingBatch, torch.Tensor, dict]:
-        assert training_batch.conditional_dict is not None
-        assert training_batch.unconditional_dict is not None
-        assert training_batch.latents is not None
-
         with torch.no_grad():
             with set_forward_context(
                 current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata_vsa):
                 generated_video, timestep_gen = self._student_forward(training_batch)
-
+        
         critic_timestep = torch.randint(
             0,
             self.num_train_timestep,
-            self.video_latent_shape[:2],
+            self.video_latent_shape_sp[:2],
             device=self.device,
             dtype=torch.long
         )
@@ -368,17 +464,24 @@ class DistillationPipeline(TrainingPipeline):
         critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
 
         # Use cross-codebase generator for reproducible noise generation
-        critic_noise = torch.randn_like(generated_video)
+        critic_noise = torch.randn(
+            self.video_latent_shape, device=self.device, dtype=generated_video.dtype)
+        if self.sp_world_size > 1:
+            critic_noise = rearrange(critic_noise,
+                                    "b (n t) c h w -> b n t c h w",
+                                    n=self.sp_world_size).contiguous()
+            critic_noise = critic_noise[:, self.rank_in_sp_group, :, :, :, :]
         
         noisy_generated_video = self.noise_scheduler.add_noise(
             generated_video.flatten(0, 1),
             critic_noise.flatten(0, 1),
             critic_timestep.flatten(0, 1)
-        ).unflatten(0, self.video_latent_shape[:2])
+        ).unflatten(0, self.video_latent_shape_sp[:2])
 
         with set_forward_context(
                 current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata):
             training_batch = self._build_input_kwargs(noisy_generated_video, critic_timestep, training_batch.conditional_dict, training_batch)
+            
             pred_fake_video = self.critic_transformer(training_batch, critic_timestep)
 
         # # Step 3: Compute the denoising loss for the fake critic
@@ -400,7 +503,7 @@ class DistillationPipeline(TrainingPipeline):
             "critictrain_noisy_latent": noisy_generated_video.detach(),
             "critictrain_pred_video": pred_fake_video.detach(),
             "critic_timestep": critic_timestep.float().detach(),
-            "critic_timestep_stu": timestep_gen.float().detach()    
+            "critic_timestep_stu": timestep_gen.float().detach(),
         }
 
         return training_batch, denoising_loss, critic_log_dict
@@ -445,66 +548,89 @@ class DistillationPipeline(TrainingPipeline):
         assert training_batch.latents is not None
         training_batch.latents = training_batch.latents.permute(0, 2, 1, 3, 4)
         self.video_latent_shape = training_batch.latents.shape # [B, C, T, H, W]
+        training_batch.raw_latent_shape = training_batch.latents.shape
 
+        if self.sp_world_size > 1:
+            training_batch.latents = rearrange(training_batch.latents,
+                                "b (n t) c h w -> b n t c h w",
+                                n=self.sp_world_size).contiguous()
+            training_batch.latents = training_batch.latents[:, self.rank_in_sp_group, :, :, :, :]
+
+        self.video_latent_shape_sp = training_batch.latents.shape
 
         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_distillation(training_batch)
-        TRAIN_STUDENT = self.current_trainstep % self.student_critic_update_ratio == 0
-        # for _ in range(self.training_args.gradient_accumulation_steps):
-        training_batch = self._get_next_batch(training_batch)
-
-        training_batch = self._normalize_dit_input(training_batch)
-        training_batch = self._prepare_dit_inputs(training_batch)
-        
-        training_batch = self._build_attention_metadata(training_batch)
-
+        """Train one step with alternating student and critic updates, supporting gradient accumulation."""
         import copy
-        training_batch.attn_metadata_vsa = copy.deepcopy(training_batch.attn_metadata)
-        if training_batch.attn_metadata is not None:
-            training_batch.attn_metadata.VSA_sparsity = 0.0
-        
-        if TRAIN_STUDENT:
-            training_batch, dmd_loss, dmd_log_dict = self._student_forward_and_compute_dmd_loss(training_batch)
-            training_batch.dmd_log_dict = dmd_log_dict
-            self.optimizer.zero_grad()
-            with set_forward_context(
-                current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata_vsa): 
-                dmd_loss.backward()
-            training_batch = self._clip_grad_norm(training_batch, self.student_transformer)
+        gradient_accumulation_steps = getattr(self.training_args, 'gradient_accumulation_steps', 1)
+        batches = []
+        # Collect N batches for gradient accumulation
+        for _ in range(gradient_accumulation_steps):
+            batch = self._prepare_distillation(training_batch)
+            batch = self._get_next_batch(batch)
+            batch = self._normalize_dit_input(batch)
+            batch = self._prepare_dit_inputs(batch)
+            batch = self._build_attention_metadata(batch)
+            batch.attn_metadata_vsa = copy.deepcopy(batch.attn_metadata)
+            if batch.attn_metadata is not None:
+                batch.attn_metadata.VSA_sparsity = 0.0
+            batches.append(batch)
+
+        # Student accumulation
+        self.optimizer.zero_grad()
+        total_dmd_loss = 0.0
+        total_dmd_log_dict = None
+        if (self.current_trainstep % self.student_critic_update_ratio == 0):
+            for batch in batches:
+                batch_stu = copy.deepcopy(batch)
+                batch_stu, dmd_loss, dmd_log_dict = self._student_forward_and_compute_dmd_loss(batch_stu)
+                # Ensure backward is under the correct forward context
+                with set_forward_context(
+                    current_timestep=batch_stu.timesteps, attn_metadata=batch_stu.attn_metadata):
+                    (dmd_loss / gradient_accumulation_steps).backward()
+                total_dmd_loss += dmd_loss.detach().item()
+                if total_dmd_log_dict is None:
+                    total_dmd_log_dict = dmd_log_dict
+                # Only keep the first log dict, ignore subsequent ones
+            self._clip_grad_norm(batch_stu, self.student_transformer)
             self.optimizer.step()
             self.lr_scheduler.step()
-            
-            avg_dmd_loss = dmd_loss.detach().clone()
+            self.optimizer.zero_grad(set_to_none=True)
+            avg_dmd_loss = torch.tensor(total_dmd_loss / gradient_accumulation_steps, device=self.device)
             world_group = get_world_group()
             world_group.all_reduce(avg_dmd_loss, op=torch.distributed.ReduceOp.AVG)
-            
-            training_batch.student_loss = avg_dmd_loss.item() 
-                            
-        training_batch, critic_loss, critic_log_dict = self._critic_forward_and_compute_loss(training_batch)
-        training_batch.critic_log_dict = critic_log_dict
+            training_batch.student_loss = avg_dmd_loss.item()
+            training_batch.dmd_log_dict = total_dmd_log_dict if total_dmd_log_dict is not None else {}
+        else:
+            training_batch.student_loss = 0.0
+            training_batch.dmd_log_dict = {}
+
+        # Critic accumulation
         self.critic_transformer_optimizer.zero_grad()
-        with set_forward_context(
-            current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata): 
-            critic_loss.backward()
-        training_batch = self._clip_grad_norm(training_batch, self.critic_transformer)
+        total_critic_loss = 0.0
+        total_critic_log_dict = None
+        for batch in batches:
+            batch_critic = copy.deepcopy(batch)
+            batch_critic, critic_loss, critic_log_dict = self._critic_forward_and_compute_loss(batch_critic)
+            with set_forward_context(
+                current_timestep=batch_critic.timesteps, attn_metadata=batch_critic.attn_metadata):
+                (critic_loss / gradient_accumulation_steps).backward()
+            total_critic_loss += critic_loss.detach().item()
+            if total_critic_log_dict is None:
+                total_critic_log_dict = critic_log_dict
+            # Only keep the first log dict, ignore subsequent ones
+        self._clip_grad_norm(batch_critic, self.critic_transformer)
         self.critic_transformer_optimizer.step()
         self.critic_lr_scheduler.step()
-        
-        avg_critic_loss = critic_loss.detach().clone()
+        self.critic_transformer_optimizer.zero_grad(set_to_none=True)
+        avg_critic_loss = torch.tensor(total_critic_loss / gradient_accumulation_steps, device=self.device)
         world_group = get_world_group()
         world_group.all_reduce(avg_critic_loss, op=torch.distributed.ReduceOp.AVG)
-
-        # Record loss values for logging
-                
         training_batch.critic_loss = avg_critic_loss.item()
-        
+        training_batch.critic_log_dict = total_critic_log_dict if total_critic_log_dict is not None else {}
+
         training_batch.total_loss = training_batch.student_loss + training_batch.critic_loss
-        
         return training_batch
 
     def _resume_from_checkpoint(self) -> None: #TODO(yongqi)
@@ -546,7 +672,7 @@ class DistillationPipeline(TrainingPipeline):
                     sum(p.numel() for p in self.critic_transformer.parameters()) / 1e9)
         
     def add_visualization(self, generator_log_dict: Dict[str, Any], critic_log_dict: Dict[str, Any], training_args: TrainingArgs, step: int):
-        """Add visualization data to wandb logging."""
+        """Add visualization data to wandb logging and save frames to disk."""
         wandb_loss_dict = {}
         
         # Clear GPU cache before VAE decoding to prevent OOM
@@ -558,11 +684,12 @@ class DistillationPipeline(TrainingPipeline):
         # Process critic training data
         critic_latents_name = ['critictrain_latent', 'critictrain_noisy_latent', 'critictrain_pred_video']
         # critic_latents_name = ['critictrain_pred_video']
+
         for latent_key in critic_latents_name:
-            latents = critic_log_dict[latent_key]
+            latents = critic_log_dict[latent_key] # bs, t,c, h, w
             latents = latents.permute(0, 2, 1, 3, 4)
             # decoded_latent = decode_stage(ForwardBatch(data_type="video", latents=latents), training_args)
-
+            
             if isinstance(self.vae.scaling_factor, torch.Tensor):
                 latents = latents / self.vae.scaling_factor.to(
                     latents.device, latents.dtype)
@@ -577,44 +704,47 @@ class DistillationPipeline(TrainingPipeline):
                                                         latents.dtype)
                 else:
                     latents += self.vae.shift_factor
+
             video = self.vae.decode(latents)
             video = (video / 2 + 0.5).clamp(0, 1)
             video = video.cpu().float()
             video = video.permute(0, 2, 1, 3, 4)
+
             wandb_loss_dict[latent_key] = prepare_for_saving(video)
             # Clean up references
             del video, latents
             torch.cuda.empty_cache()
 
         # Process DMD training data if available - use decode_stage instead of self.vae.decode
-
-        dmd_latents_name = ['dmdtrain_pred_fake_video', 'dmdtrain_pred_real_video', 'dmdtrain_latents', 'dmdtrain_noisy_latent']
-        for latent_key in dmd_latents_name:
-            latents = generator_log_dict[latent_key]
-            latents = latents.permute(0, 2, 1, 3, 4)
-            # decoded_latent = decode_stage(ForwardBatch(data_type="video", latents=latents), training_args)
-            if isinstance(self.vae.scaling_factor, torch.Tensor):
-                latents = latents / self.vae.scaling_factor.to(
-                    latents.device, latents.dtype)
-            else:
-                latents = latents / self.vae.scaling_factor
-
-            # Apply shifting if needed
-            if (hasattr(self.vae, "shift_factor")
-                    and self.vae.shift_factor is not None):
-                if isinstance(self.vae.shift_factor, torch.Tensor):
-                    latents += self.vae.shift_factor.to(latents.device,
-                                                        latents.dtype)
+        if 'dmdtrain_pred_fake_video' in generator_log_dict:
+            dmd_latents_name = ['dmdtrain_pred_fake_video', 'dmdtrain_pred_real_video', 'dmdtrain_latents', 'dmdtrain_noisy_latent']
+            for latent_key in dmd_latents_name:
+                latents = generator_log_dict[latent_key]
+                latents = latents.permute(0, 2, 1, 3, 4)
+                # decoded_latent = decode_stage(ForwardBatch(data_type="video", latents=latents), training_args)
+                if isinstance(self.vae.scaling_factor, torch.Tensor):
+                    latents = latents / self.vae.scaling_factor.to(
+                        latents.device, latents.dtype)
                 else:
-                    latents += self.vae.shift_factor
-            video = self.vae.decode(latents)
-            video = (video / 2 + 0.5).clamp(0, 1)
-            video = video.cpu().float()
-            video = video.permute(0, 2, 1, 3, 4)
-            wandb_loss_dict[latent_key] = prepare_for_saving(video)
-            # Clean up references
-            del video, latents
-            torch.cuda.empty_cache()
+                    latents = latents / self.vae.scaling_factor
+
+                # Apply shifting if needed
+                if (hasattr(self.vae, "shift_factor")
+                        and self.vae.shift_factor is not None):
+                    if isinstance(self.vae.shift_factor, torch.Tensor):
+                        latents += self.vae.shift_factor.to(latents.device,
+                                                            latents.dtype)
+                    else:
+                        latents += self.vae.shift_factor
+                video = self.vae.decode(latents)
+                video = (video / 2 + 0.5).clamp(0, 1)
+                video = video.cpu().float()
+                video = video.permute(0, 2, 1, 3, 4)
+                
+                wandb_loss_dict[latent_key] = prepare_for_saving(video)
+                # Clean up references
+                del video, latents
+                torch.cuda.empty_cache()
         
         # Log to wandb
         if self.global_rank == 0:
@@ -624,7 +754,7 @@ class DistillationPipeline(TrainingPipeline):
     def _log_validation(self, transformer, training_args, global_step) -> None:
         assert training_args is not None
         training_args.inference_mode = True
-        training_args.use_cpu_offload = False
+        training_args.use_cpu_offload = True
         if not training_args.log_validation:
             return
         if self.validation_pipeline is None:
@@ -684,19 +814,20 @@ class DistillationPipeline(TrainingPipeline):
                 self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[
                     0], result_batch.prompt_attention_mask[0]
 
-                logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
-                            self.global_rank,
-                            self.rank_in_sp_group,
-                            batch.prompt,
-                            local_main_process_only=False)
+                # logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
+                #             self.global_rank,
+                #             self.rank_in_sp_group,
+                #             batch.prompt,
+                #             local_main_process_only=False)
 
                 assert batch.prompt is not None and isinstance(
                     batch.prompt, str)
                 step_captions.append(batch.prompt)
 
-                # Run validation inference
-                output_batch = self.validation_pipeline.forward(
-                    batch, training_args)
+            # #     # Run validation inference
+                with torch.no_grad():
+                    output_batch = self.validation_pipeline.forward(
+                        batch, training_args)
                 samples = output_batch.output
                 if self.rank_in_sp_group != 0:
                     continue
@@ -751,14 +882,15 @@ class DistillationPipeline(TrainingPipeline):
                     wandb.log(logs, step=global_step)
                     
                     # Save all prompts from all cards to txt file
-                    prompt_filename = os.path.join(
-                        training_args.output_dir,
-                        f"validation_step_{global_step}_inference_steps_{num_inference_steps}_prompts.txt"
-                    )
-                    with open(prompt_filename, 'w', encoding='utf-8') as f:
-                        for i, caption in enumerate(all_captions):
-                            f.write(f"Video_{i}: {caption}\n")
-                    logger.info(f"Saved {len(all_captions)} prompts to {prompt_filename}")
+                    # prompt_filename = os.path.join(
+                    #     training_args.output_dir,
+                    #     f"validation_step_{global_step}_inference_steps_{num_inference_steps}_prompts.txt"
+                    # )
+                    # with open(prompt_filename, 'w', encoding='utf-8') as f:
+                    #     for i, caption in enumerate(all_captions):
+                    #         f.write(f"Video_{i}: {caption}\n")
+                    # logger.info(f"Saved {len(all_captions)} prompts to {prompt_filename}")
+                    
                 else:
                     # Other sp_group leaders send their results to global rank 0
                     world_group.send_object(step_videos, dst=0)
@@ -775,7 +907,15 @@ class DistillationPipeline(TrainingPipeline):
 
         assert self.training_args.seed is not None, "seed must be set"
         seed = self.training_args.seed
-        set_random_seed(seed + self.global_rank)
+        
+        # Set the same seed within each SP group to ensure reproducibility
+        if self.sp_world_size > 1:
+            # Use the same seed for all processes within the same SP group
+            sp_group_seed = seed + (self.global_rank // self.sp_world_size)
+            set_random_seed(sp_group_seed)
+            logger.info(f"Rank {self.global_rank}: Using SP group seed {sp_group_seed}")
+        else:
+            set_random_seed(seed + self.global_rank)
 
         self.noise_random_generator = torch.Generator(
             device="cpu").manual_seed(seed)
@@ -801,8 +941,8 @@ class DistillationPipeline(TrainingPipeline):
             disable=self.local_rank > 0,
         )
         
-        for step in range(self.init_steps,
-                          self.training_args.max_train_steps + 1):
+        for step in range(self.init_steps + 1,
+                        self.training_args.max_train_steps + 1):
             start_time = time.perf_counter()
             current_vsa_sparsity = self.training_args.VSA_sparsity if vsa_available else 0.0
             
@@ -861,32 +1001,30 @@ class DistillationPipeline(TrainingPipeline):
                     log_data.update(critic_metrics)
                 wandb.log(log_data, step=step)
                 
-            # if step % self.training_args.checkpointing_steps == 0:
-            #     print("rank", self.global_rank, "save checkpoint at step", step)
-            #     save_checkpoint(self.transformer, self.global_rank, #TODO(yongqi)
-            #                     self.training_args.output_dir, step,
-            #                     self.optimizer, self.train_dataloader,
-            #                     self.lr_scheduler, self.noise_random_generator)
-            #     if self.transformer:
-            #         self.transformer.train()
-            #     self.sp_group.barrier()
+            if step % self.training_args.checkpointing_steps == 0:
+                print("rank", self.global_rank, "save checkpoint at step", step)
+                save_checkpoint(self.transformer, self.global_rank, #TODO(yongqi)
+                                self.training_args.output_dir, step,
+                                self.optimizer, self.train_dataloader,
+                                self.lr_scheduler, self.noise_random_generator)
+                if self.transformer:
+                    self.transformer.train()
+                self.sp_group.barrier()
                 
             if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
-                gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
-                logger.info("GPU memory usage before validation: %s MB",
-                            gpu_memory_usage)
+
+
                 self.add_visualization(training_batch.dmd_log_dict, training_batch.critic_log_dict, self.training_args, step)
                 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.model, 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)
+        # save_checkpoint(self.student_transformer.model, 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()

fastvideo/v1/training/training_pipeline.py

@@ -5,7 +5,7 @@ import time
 from abc import ABC, abstractmethod
 from collections import deque
 from collections.abc import Iterator
-from typing import Any
+from typing import Any, Dict
 
 import imageio
 import numpy as np
@@ -39,7 +39,7 @@ from fastvideo.v1.training.activation_checkpoint import (
 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)
+    normalize_dit_input, save_checkpoint, shard_latents_across_sp, prepare_for_saving)
 from fastvideo.v1.utils import is_vsa_available, set_random_seed, shallow_asdict
 
 import wandb  # isort: skip
@@ -243,23 +243,22 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         training_batch.timesteps = timesteps
         training_batch.sigmas = sigmas
         training_batch.noise = noise
+        training_batch.raw_latent_shape = training_batch.latents.shape
 
         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
+        latents_shape = training_batch.raw_latent_shape
         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] * self.sp_world_size // patch_size[0],
-                latents.shape[3] // patch_size[1],
-                latents.shape[4] // patch_size[2]
+                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,
@@ -327,6 +326,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
 
             # make sure no implicit broadcasting happens
             assert model_pred.shape == target.shape, f"model_pred.shape: {model_pred.shape}, target.shape: {target.shape}"
+
             loss = (torch.mean((model_pred.float() - target.float())**2) /
                     self.training_args.gradient_accumulation_steps)
 
@@ -441,7 +441,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         step_times: deque[float] = deque(maxlen=100)
 
         self._log_training_info()
-        self._log_validation(self.transformer, self.training_args, 1)
+        self._log_validation(self.transformer, self.training_args, 0)
 
         # Train!
         progress_bar = tqdm(
@@ -451,7 +451,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             # Only show the progress bar once on each machine.
             disable=self.local_rank > 0,
         )
-        for step in range(self.init_steps + 1,
+        for step in range(self.init_steps,
                           self.training_args.max_train_steps + 1):
             start_time = time.perf_counter()
             if vsa_available:
@@ -469,6 +469,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             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
 

fastvideo/v1/training/training_utils.py

@@ -575,7 +575,6 @@ class DiffusionWrapper(torch.nn.Module, ABC):
         
     def forward(self, training_batch: TrainingBatch, timestep: torch.Tensor):
         pred_noise = self.model(**training_batch.input_kwargs).permute(0, 2, 1, 3, 4)
-        torch.distributed.breakpoint()
         pred_video = self._convert_flow_pred_to_x0(
             flow_pred=pred_noise.flatten(0, 1),
             xt=training_batch.noise_latents.flatten(0, 1),

fastvideo/v1/training/wan_distillation_pipeline.py

@@ -6,8 +6,8 @@ import torch
 from fastvideo.v1.distributed import get_local_torch_device
 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.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler)
 from fastvideo.v1.pipelines.wan.wan_dmd_pipeline import WanDmdPipeline
 from fastvideo.v1.training.distillation_pipeline import DistillationPipeline
 from fastvideo.v1.pipelines.pipeline_batch_info import (ForwardBatch,
@@ -29,7 +29,7 @@ class WanDistillationPipeline(DistillationPipeline):
     
     def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
         """Initialize Wan-specific scheduler."""
-        self.modules["scheduler"] = FlowUniPCMultistepScheduler(
+        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
             shift=fastvideo_args.pipeline_config.flow_shift)
 
     def create_training_stages(self, training_args: TrainingArgs):
@@ -43,11 +43,11 @@ class WanDistillationPipeline(DistillationPipeline):
         args_copy = deepcopy(training_args)
 
         args_copy.inference_mode = True
-        args_copy.use_cpu_offload = False
+        args_copy.use_cpu_offload = True
         args_copy.pipeline_config.vae_config.load_encoder = False
         validation_pipeline = WanDmdPipeline.from_pretrained(
             training_args.model_path,
-            args=None,
+            args=args_copy,
             inference_mode=True,
             loaded_modules={"transformer": self.get_module("transformer")},
             tp_size=training_args.tp_size,
@@ -77,7 +77,6 @@ def main(args) -> None:
         args.pretrained_model_name_or_path, args=args)
     
     args = pipeline.training_args
-    
     # Start training
     pipeline.train()
     logger.info("Wan distillation pipeline completed")
@@ -91,5 +90,4 @@ if __name__ == "__main__":
     parser = TrainingArgs.add_cli_args(parser)
     parser = FastVideoArgs.add_cli_args(parser)   
     args = parser.parse_args()
-    args.use_cpu_offload = False
     main(args) 

fastvideo/v1/training/wan_i2v_distillation_pipeline.py

@@ -4,8 +4,9 @@ from copy import deepcopy
 from typing import Any
 
 import torch
+from einops import rearrange
 from fastvideo.v1.configs.sample import SamplingParam
-from fastvideo.v1.distributed import get_local_torch_device
+from fastvideo.v1.distributed import get_local_torch_device, get_sp_parallel_rank, get_sp_world_size
 from fastvideo.v1.dataset.dataloader.schema import (
     pyarrow_schema_i2v, pyarrow_schema_i2v_validation)
 from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
@@ -173,9 +174,15 @@ class WanI2VDistillationPipeline(DistillationPipeline):
         image_latents = torch.cat(
             [mask_lat_size, image_latents],
             dim=1)
-
         training_batch.image_latents = image_latents
-        
+
+        if self.sp_world_size > 1:
+            image_latents = rearrange(image_latents,
+                                        "b c (n t) h w -> b c n t h w",
+                                        n=self.sp_world_size).contiguous()
+            image_latents = image_latents[:, :, self.rank_in_sp_group, :, :, :]
+            training_batch.image_latents = image_latents
+    
         return training_batch
 
     def _build_input_kwargs(self, noise_input: torch.Tensor, timestep: torch.Tensor, text_dict: dict[str, torch.Tensor],
@@ -193,7 +200,7 @@ class WanI2VDistillationPipeline(DistillationPipeline):
             [noise_input, training_batch.image_latents.permute(0, 2, 1, 3, 4)], dim=2)
 
         training_batch.input_kwargs = {
-            "hidden_states": noisy_model_input.permute(0, 2, 1, 3, 4),
+            "hidden_states": noisy_model_input.permute(0, 2, 1, 3, 4), # bs, c, t, h, w
             "encoder_hidden_states": text_dict["encoder_hidden_states"],
             "encoder_attention_mask": text_dict["encoder_attention_mask"],
             "timestep": timestep[0][:1],

fastvideo/v1/training/wan_i2v_training_pipeline.py

@@ -28,7 +28,7 @@ class WanI2VTrainingPipeline(TrainingPipeline):
     """
     A training pipeline for Wan.
     """
-    _required_config_modules = ["scheduler", "transformer"]
+    _required_config_modules = ["scheduler", "transformer", "vae"]
 
     def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
         self.modules["scheduler"] = FlowUniPCMultistepScheduler(

fastvideo/v1/training/wan_training_pipeline.py

@@ -40,7 +40,7 @@ class WanTrainingPipeline(TrainingPipeline):
         args_copy.pipeline_config.vae_config.load_encoder = False
         validation_pipeline = WanPipeline.from_pretrained(
             training_args.model_path,
-            args=None,
+            args=args_copy,
             inference_mode=True,
             loaded_modules={"transformer": self.get_module("transformer")},
             tp_size=training_args.tp_size,

scripts/finetune_wan_distillation_ci.sh

@@ -0,0 +1,67 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export WANDB_API_KEY='73190d8c0de18a14eb3444e222f9432d247d1e30'
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+export TRITON_CACHE_DIR=/tmp/triton_cache
+DATA_DIR=mini_i2v_dataset/crush-smol_preprocessed/combined_parquet_dataset
+# DATA_DIR=/mnt/sharefs/users/hao.zhang/Vchitect-2M/Vchitect-2M-laten-93x512x512/val/
+VALIDATION_DIR=mini_i2v_dataset/crush-smol_raw/validation.json
+NUM_GPUS=2
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+# export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+export TOKENIZERS_PARALLELISM=false
+CHECKPOINT_PATH="outputs_train_test/wan_finetune/checkpoint-10"
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
+    fastvideo/v1/training/wan_distillation_pipeline.py \
+    --model_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+    --inference_mode False\
+    --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+    --cache_dir "/home/ray/.cache" \
+    --data_path "$DATA_DIR" \
+    --validation_dataset_file  "$VALIDATION_DIR" \
+    --train_batch_size 1 \
+    --num_latent_t 8 \
+    --sp_size 1 \
+    --tp_size 1 \
+    --num_gpus $NUM_GPUS \
+    --hsdp_replicate_dim $NUM_GPUS  \
+    --hsdp-shard-dim 1 \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 0 \
+    --gradient_accumulation_steps 1 \
+    --max_train_steps 5 \
+    --learning_rate 1e-5 \
+    --mixed_precision "bf16" \
+    --checkpointing_steps 10 \
+    --validation_steps 3 \
+    --validation_sampling_steps "3" \
+    --log_validation \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --training_cfg_rate 0.0 \
+    --output_dir "outputs_dmd/wan_finetune" \
+    --tracker_project_name Wan_distillation \
+    --num_height 448 \
+    --num_width 832 \
+    --num_frames  29 \
+    --flow_shift 8 \
+    --validation_guidance_scale "1.0" \
+    --master_weight_type "fp32" \
+    --dit_precision "fp32" \
+    --vae_precision "bf16" \
+    --weight_decay 0.01 \
+    --max_grad_norm 1.0 \
+    --student_critic_update_ratio 3 \
+    --denoising_step_list '1000,757,522' \
+    --min_step_ratio 0.02 \
+    --max_step_ratio 0.98 \
+    --teacher_guidance_scale 3.5 \
+    --enable_gradient_checkpointing_type "full" \
+    --seed 1024 \
+
+# validation_preprocessed_path