diff --git a/fastvideo/distributed/communication_op.py b/fastvideo/distributed/communication_op.py index f29484c3d..c17b06e88 100644 --- a/fastvideo/distributed/communication_op.py +++ b/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) diff --git a/fastvideo/models/dits/hunyuanvideo15.py b/fastvideo/models/dits/hunyuanvideo15.py index 508ea9648..1abf8611c 100644 --- a/fastvideo/models/dits/hunyuanvideo15.py +++ b/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 \ No newline at end of file + return x diff --git a/fastvideo/tests/nightly/test_e2e_hunyuan15_overfit_single_sample.py b/fastvideo/tests/nightly/test_e2e_hunyuan15_overfit_single_sample.py index e1feb2d22..81e673aef 100644 --- a/fastvideo/tests/nightly/test_e2e_hunyuan15_overfit_single_sample.py +++ b/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), diff --git a/fastvideo/training/hunyuan15_training_pipeline.py b/fastvideo/training/hunyuan15_training_pipeline.py index a6c12e01f..f798f8760 100644 --- a/fastvideo/training/hunyuan15_training_pipeline.py +++ b/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 diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 01fe83490..ecb917d0d 100644 --- a/fastvideo/training/training_pipeline.py +++ b/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: diff --git a/kill.sh b/kill.sh new file mode 100755 index 000000000..196697fe4 --- /dev/null +++ b/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