update

fastvideo/distributed/communication_op.py

@@ -53,6 +53,9 @@ def sequence_model_parallel_all_gather_with_unpad(
         Tensor: Gathered and unpadded tensor
     """
 
+    # NCCL all_gather expects contiguous inputs.
+    if not input_.is_contiguous():
+        input_ = input_.contiguous()
     # First gather across all ranks
     gathered = get_sp_group().all_gather(input_, dim)
 

fastvideo/models/dits/hunyuanvideo15.py

@@ -690,12 +690,15 @@ class SingleTokenRefiner(nn.Module):
         timestep_aware_representations = self.t_embedder(t)
 
         # Get context-aware representations
-        original_dtype = x.dtype
+        target_dtype = self.c_embedder.fc_in.weight.dtype
+        if x.dtype != target_dtype:
+            x = x.to(dtype=target_dtype)
         if mask is None:
             context_aware_representations = x.mean(dim=1)
         else:
-            mask_float = mask.float().unsqueeze(-1)  # [B, L, 1]
-            context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1)
+            mask_float = mask.to(dtype=target_dtype).unsqueeze(-1)  # [B, L, 1]
+            context_aware_representations = (
+                x * mask_float).sum(dim=1) / mask_float.sum(dim=1)
 
         context_aware_representations = self.c_embedder(
             context_aware_representations)
@@ -850,4 +853,4 @@ class FinalLayer(nn.Module):
         scale, shift = self.adaLN_modulation(c).chunk(2, dim=-1)
         x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
         x, _ = self.linear(x)
-        return x
+        return x

fastvideo/tests/nightly/test_e2e_hunyuan15_overfit_single_sample.py

@@ -18,7 +18,7 @@ LOCAL_PREPROCESSED_DATA_DIR = Path(DATA_DIR) / "cats_processed_t2v_hunyuan15"
 LOCAL_OUTPUT_DIR = Path(DATA_DIR) / "outputs_hunyuan15"
 
 NUM_GPUS_PER_NODE_PREPROCESSING = "1"
-NUM_GPUS_PER_NODE_TRAINING = "4"
+NUM_GPUS_PER_NODE_TRAINING = "8"
 
 # entrypoints (adjust to what hunyuan15 scripts use)
 PREPROCESS_ENTRY = ["-m", "fastvideo.pipelines.preprocess.v1_preprocessing_new"]
@@ -145,6 +145,7 @@ def run_training():
         "--ema_start_step", "0",
         "--training_cfg_rate", "0.0",
         "--dit_precision", "fp32",
+        "--enable_gradient_checkpointing_type", "full",
 
         # Output
         "--output_dir", str(LOCAL_OUTPUT_DIR),

fastvideo/training/hunyuan15_training_pipeline.py

@@ -5,8 +5,7 @@ from copy import deepcopy
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.logger import init_logger
 from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler,
-)
+    FlowMatchEulerDiscreteScheduler, )
 from fastvideo.training.training_pipeline import TrainingPipeline
 
 from fastvideo.pipelines.basic.hunyuan15.hunyuan15_pipeline import HunyuanVideo15Pipeline
@@ -20,10 +19,14 @@ class Hunyuan15TrainingPipeline(TrainingPipeline):
     """
     _required_config_modules = ["scheduler", "transformer", "vae"]
 
+    def initialize_training_pipeline(self, training_args: TrainingArgs):
+        if training_args.enable_gradient_checkpointing_type is None:
+            training_args.enable_gradient_checkpointing_type = "full"
+        super().initialize_training_pipeline(training_args)
+
     def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
         self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
-            shift=fastvideo_args.pipeline_config.flow_shift
-        )
+            shift=fastvideo_args.pipeline_config.flow_shift)
 
     def create_training_stages(self, training_args: TrainingArgs):
         # reserved for future refactors
@@ -47,6 +50,8 @@ class Hunyuan15TrainingPipeline(TrainingPipeline):
             num_gpus=training_args.num_gpus,
             pin_cpu_memory=training_args.pin_cpu_memory,
             dit_cpu_offload=True,
+            # dit_layerwise_offload=True,
+            use_fsdp_inference=True,
         )
 
         self.validation_pipeline = validation_pipeline

fastvideo/training/training_pipeline.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import asdict
+import inspect
 import math
 import os
 import time
@@ -174,6 +175,16 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 last_epoch=self.init_steps - 1,
             )
 
+        text_padding_length = training_args.pipeline_config.text_encoder_configs[
+            0].arch_config.text_len  # type: ignore[attr-defined]
+        if not text_padding_length:
+            text_max_lengths = getattr(training_args.pipeline_config,
+                                       "text_encoder_max_lengths", None)
+            if text_max_lengths:
+                text_padding_length = text_max_lengths[0]
+            else:
+                text_padding_length = 512
+
         self.train_dataset, self.train_dataloader = build_parquet_map_style_dataloader(
             training_args.data_path,
             training_args.train_batch_size,
@@ -181,9 +192,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
             num_data_workers=training_args.dataloader_num_workers,
             cfg_rate=training_args.training_cfg_rate,
             drop_last=True,
-            text_padding_length=training_args.pipeline_config.
-            text_encoder_configs[0].arch_config.
-            text_len,  # type: ignore[attr-defined]
+            text_padding_length=text_padding_length,
             seed=self.seed)
 
         self.noise_scheduler = noise_scheduler
@@ -275,7 +284,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
         # Automatically detect model type based on VAE class name
         vae = self.get_module("vae")
         vae_class_name = vae.__class__.__name__
-        
+
         # Map VAE class names to model types
         if 'Hunyuan' in vae_class_name:
             model_type = 'hunyuan'
@@ -287,7 +296,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 model_type = 'wan'
             else:
                 model_type = 'hunyuan'
-        
+
         with self.tracker.timed("timing/normalize_input"):
             training_batch.latents = normalize_dit_input(
                 model_type,
@@ -444,6 +453,87 @@ class TrainingPipeline(LoRAPipeline, ABC):
         #         device=training_batch.noisy_model_input.device,
         #         dtype=torch.bfloat16)
         current_model = self.transformer_2 if self.train_transformer_2 else self.transformer
+        if "return_dict" in input_kwargs:
+            try:
+                sig = inspect.signature(current_model.forward)
+            except (TypeError, ValueError):
+                sig = None
+            if sig is not None:
+                params = sig.parameters.values()
+                accepts_kwargs = any(param.kind == inspect.Parameter.VAR_KEYWORD
+                                     for param in params)
+                if ("return_dict" not in sig.parameters and not accepts_kwargs):
+                    input_kwargs = dict(input_kwargs)
+                    input_kwargs.pop("return_dict", None)
+        base_model = current_model
+        for attr in ("module", "_orig_mod"):
+            if hasattr(base_model, attr):
+                base_model = getattr(base_model, attr)
+        if hasattr(base_model, "_fsdp_wrapped_module"):
+            base_model = base_model._fsdp_wrapped_module
+            if hasattr(base_model, "module"):
+                base_model = base_model.module
+        model_config = getattr(base_model, "config", None)
+        if model_config is None:
+            model_config = getattr(current_model, "config", None)
+        if model_config is not None and hasattr(model_config,
+                                                "text_embed_2_dim"):
+            input_kwargs = dict(input_kwargs)
+            text_states = training_batch.encoder_hidden_states
+            text_mask = training_batch.encoder_attention_mask
+            if text_states is None or text_mask is None:
+                raise ValueError(
+                    "HunyuanVideo15 training requires text embeddings and masks"
+                )
+            if not isinstance(input_kwargs.get("encoder_hidden_states"),
+                              (list, tuple)):
+                batch_size = text_states.shape[0]
+                text_2_len = 256
+                if hasattr(self.training_args.pipeline_config,
+                           "text_encoder_max_lengths"):
+                    text_2_len = self.training_args.pipeline_config.text_encoder_max_lengths[
+                        1]
+                text_2_dim = model_config.text_embed_2_dim
+                text_2 = torch.zeros((batch_size, text_2_len, text_2_dim),
+                                     device=text_states.device,
+                                     dtype=text_states.dtype)
+                input_kwargs["encoder_hidden_states"] = [text_states, text_2]
+            if not isinstance(input_kwargs.get("encoder_attention_mask"),
+                              (list, tuple)):
+                batch_size = text_mask.shape[0]
+                text_2_len = 256
+                if hasattr(self.training_args.pipeline_config,
+                           "text_encoder_max_lengths"):
+                    text_2_len = self.training_args.pipeline_config.text_encoder_max_lengths[
+                        1]
+                text_mask_2 = torch.zeros((batch_size, text_2_len),
+                                          device=text_mask.device,
+                                          dtype=text_mask.dtype)
+                input_kwargs["encoder_attention_mask"] = [
+                    text_mask, text_mask_2
+                ]
+            if "encoder_hidden_states_image" not in input_kwargs:
+                image_embeds = training_batch.image_embeds
+                if image_embeds is None or image_embeds.numel() == 0:
+                    batch_size = text_states.shape[0]
+                    image_len = 729
+                    image_dim = model_config.image_embed_dim
+                    image_embeds = torch.zeros(
+                        (batch_size, image_len, image_dim),
+                        device=text_states.device,
+                        dtype=text_states.dtype,
+                    )
+                input_kwargs["encoder_hidden_states_image"] = [image_embeds]
+            hidden_states = input_kwargs.get("hidden_states")
+            if (isinstance(hidden_states, torch.Tensor)
+                    and hidden_states.shape[1] != model_config.in_channels):
+                batch_size, _, t, h, w = hidden_states.shape
+                video_latent = torch.zeros((batch_size, 1, t, h, w),
+                                           device=hidden_states.device,
+                                           dtype=hidden_states.dtype)
+                zeros_latent = torch.zeros_like(hidden_states)
+                input_kwargs["hidden_states"] = torch.cat(
+                    [hidden_states, video_latent, zeros_latent], dim=1)
 
         with self.tracker.timed("timing/forward_backward"), set_forward_context(
                 current_timestep=training_batch.current_timestep,
@@ -610,8 +700,8 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
         self._log_training_info()
 
-        self._log_validation(self.transformer, self.training_args,
-                             self.init_steps)
+        # self._log_validation(self.transformer, self.training_args,
+        #                      self.init_steps)
 
         # Train!
         progress_bar = tqdm(
@@ -666,7 +756,13 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 }
                 metrics["batch_size"] = int(training_batch.raw_latent_shape[0])
 
-                patch_t, patch_h, patch_w = self.training_args.pipeline_config.dit_config.patch_size
+                patch_size = self.training_args.pipeline_config.dit_config.patch_size
+                if isinstance(patch_size, tuple):
+                    patch_t, patch_h, patch_w = patch_size
+                else:
+                    patch_t = patch_size
+                    patch_h = patch_size
+                    patch_w = patch_size
                 seq_len = (training_batch.raw_latent_shape[2] // patch_t) * (
                     training_batch.raw_latent_shape[3] //
                     patch_h) * (training_batch.raw_latent_shape[4] // patch_w)
@@ -679,7 +775,13 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
                 metrics["hidden_dim"] = arch_config.hidden_size
                 metrics["num_layers"] = arch_config.num_layers
-                metrics["ffn_dim"] = arch_config.ffn_dim
+                ffn_dim = getattr(arch_config, "ffn_dim", None)
+                if ffn_dim is None:
+                    mlp_ratio = getattr(arch_config, "mlp_ratio", None)
+                    if mlp_ratio and arch_config.hidden_size:
+                        ffn_dim = int(arch_config.hidden_size * mlp_ratio)
+                if ffn_dim is not None:
+                    metrics["ffn_dim"] = ffn_dim
 
                 self.tracker.log(metrics, step)
             if step % self.training_args.training_state_checkpointing_steps == 0:
@@ -854,8 +956,12 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 step_captions.append(batch.prompt)
 
                 # Run validation inference
+                # Use the validation pipeline args to match its loaded modules.
+                validation_args = self.validation_pipeline.fastvideo_args
+                if validation_args is None:
+                    validation_args = training_args
                 output_batch = self.validation_pipeline.forward(
-                    batch, training_args)
+                    batch, validation_args)
                 samples = output_batch.output
 
                 if self.rank_in_sp_group != 0:

kill.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+# Kill all GPU processes shown by nvidia-smi
+
+pids=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null | tr -d ' ')
+
+if [ -z "$pids" ]; then
+    echo "No GPU processes found."
+    exit 0
+fi
+
+echo "Killing GPU processes: $pids"
+for pid in $pids; do
+    kill -9 "$pid" 2>/dev/null && echo "Killed PID $pid" || echo "Failed to kill PID $pid"
+done