From 92dfbfe4fc4451a998f29585dc508f6356712336 Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Fri, 26 Sep 2025 01:20:37 +0000 Subject: [PATCH] Refactor ode trajectory generation Add env variable to enforce running flash attn --- fastvideo/attention/layer.py | 5 +- fastvideo/configs/pipelines/base.py | 1 + fastvideo/envs.py | 5 + fastvideo/models/dits/wanvideo.py | 16 ++- fastvideo/models/loader/component_loader.py | 5 +- fastvideo/models/loader/fsdp_load.py | 5 +- fastvideo/models/utils.py | 22 +++ .../preprocess_pipeline_ode_trajectory.py | 136 ++++++++++++++++-- .../tests/transformers/test_ori_wanvideo.py | 53 ++++--- 9 files changed, 207 insertions(+), 41 deletions(-) diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py index 7b086a3c6..2ba1c9f7f 100644 --- a/fastvideo/attention/layer.py +++ b/fastvideo/attention/layer.py @@ -3,6 +3,7 @@ import torch import torch.nn as nn +import fastvideo.envs as envs from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend from fastvideo.distributed.communication_op import ( sequence_model_parallel_all_gather, sequence_model_parallel_all_to_all_4D) @@ -36,7 +37,7 @@ class DistributedAttention(nn.Module): if num_kv_heads is None: num_kv_heads = num_heads - dtype = get_compute_dtype() + dtype = torch.bfloat16 if envs.FASTVIDEO_FORCE_ATTN_BF16 else get_compute_dtype() attn_backend = get_attn_backend( head_size, dtype, @@ -221,7 +222,7 @@ class LocalAttention(nn.Module): if num_kv_heads is None: num_kv_heads = num_heads - dtype = get_compute_dtype() + dtype = torch.bfloat16 if envs.FASTVIDEO_FORCE_ATTN_BF16 else get_compute_dtype() attn_backend = get_attn_backend( head_size, dtype, diff --git a/fastvideo/configs/pipelines/base.py b/fastvideo/configs/pipelines/base.py index 4c3803dd9..54ef28c2d 100644 --- a/fastvideo/configs/pipelines/base.py +++ b/fastvideo/configs/pipelines/base.py @@ -51,6 +51,7 @@ class PipelineConfig: # Model configuration dit_config: DiTConfig = field(default_factory=DiTConfig) dit_precision: str = "bf16" + dit_forward_precision: str = "bf16" # VAE configuration vae_config: VAEConfig = field(default_factory=VAEConfig) diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 25444854f..d8634b1ed 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -18,6 +18,7 @@ if TYPE_CHECKING: FASTVIDEO_LOGGING_PREFIX: str = "" FASTVIDEO_LOGGING_CONFIG_PATH: str | None = None FASTVIDEO_TRACE_FUNCTION: int = 0 + FASTVIDEO_FORCE_ATTN_BF16: bool = False FASTVIDEO_ATTENTION_BACKEND: str | None = None FASTVIDEO_ATTENTION_CONFIG: str | None = None FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "fork" @@ -168,6 +169,10 @@ environment_variables: dict[str, Callable[[], Any]] = { "FASTVIDEO_TRACE_FUNCTION": lambda: int(os.getenv("FASTVIDEO_TRACE_FUNCTION", "0")), + # if set, fastvideo will force attention to be computed in bfloat16 + "FASTVIDEO_FORCE_ATTN_BF16": + lambda: bool(int(os.getenv("FASTVIDEO_FORCE_ATTN_BF16", "0"))), + # Backend for attention computation # Available options: # - "TORCH_SDPA": use torch.nn.MultiheadAttention diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index 3262a9a1e..9c013f946 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -172,8 +172,13 @@ class WanT2VCrossAttention(WanSelfAttention): k = self.norm_k.forward_native(self.to_k(context)[0]).view(b, -1, n, d) v = self.to_v(context)[0].view(b, -1, n, d) - # compute attention - x = self.attn(q, k, v) + if envs.FASTVIDEO_FORCE_ATTN_BF16: + out_dtype = v.dtype + # compute attention + x = self.attn(q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)).to(out_dtype) + else: + # compute attention + x = self.attn(q, k, v) # output x = x.flatten(2) @@ -359,7 +364,12 @@ class WanTransformerBlock(nn.Module): is_neox_style=False), _apply_rotary_emb( key, cos, sin, is_neox_style=False) - attn_output, _ = self.attn1(query, key, value) + if envs.FASTVIDEO_FORCE_ATTN_BF16: + out_dtype = value.dtype + attn_output, _ = self.attn1(query.to(torch.bfloat16), key.to(torch.bfloat16), value.to(torch.bfloat16)) + attn_output = attn_output.to(out_dtype) + else: + attn_output, _ = self.attn1(query, key, value) attn_output = attn_output.flatten(2) attn_output, _ = self.to_out(attn_output) attn_output = attn_output.squeeze(1) diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 0a3d8e097..989341966 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -450,6 +450,8 @@ class TransformerLoader(ComponentLoader): default_dtype = PRECISION_TO_TYPE[ fastvideo_args.pipeline_config.dit_precision] + param_dtype = PRECISION_TO_TYPE[ + fastvideo_args.pipeline_config.dit_forward_precision] # Load the model using FSDP loader logger.info("Loading model from %s, default_dtype: %s", cls_name, @@ -469,7 +471,8 @@ class TransformerLoader(ComponentLoader): pin_cpu_memory=fastvideo_args.pin_cpu_memory, fsdp_inference=fastvideo_args.use_fsdp_inference, # TODO(will): make these configurable - param_dtype=torch.bfloat16, + default_dtype=default_dtype, + param_dtype=param_dtype, reduce_dtype=torch.float32, output_dtype=None, training_mode=fastvideo_args.training_mode) diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 4a7bbb4a1..8a9dd4067 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -62,6 +62,7 @@ def maybe_load_fsdp_model( device: torch.device, hsdp_replicate_dim: int, hsdp_shard_dim: int, + default_dtype: torch.dtype, param_dtype: torch.dtype, reduce_dtype: torch.dtype, cpu_offload: bool = False, @@ -87,7 +88,7 @@ def maybe_load_fsdp_model( mp_policy=mp_policy, ) - with set_default_dtype(torch.float32), torch.device("meta"): + with set_default_dtype(default_dtype), torch.device("meta"): model = model_cls(**init_params) # Check if we should use FSDP @@ -125,7 +126,7 @@ def maybe_load_fsdp_model( model, weight_iterator, device, - torch.float32, + default_dtype, strict=True, cpu_offload=cpu_offload, param_names_mapping=param_names_mapping_fn, diff --git a/fastvideo/models/utils.py b/fastvideo/models/utils.py index 6da2f4481..adc55f491 100644 --- a/fastvideo/models/utils.py +++ b/fastvideo/models/utils.py @@ -180,3 +180,25 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor, sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) pred_video = noise_input_latent - sigma_t * pred_noise return pred_video.to(dtype) + +def pred_video_to_pred_noise(x0_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor, scheduler: Any) -> torch.Tensor: + """ + Convert x0 prediction to flow matching's prediction. + x0_pred: the x0 prediction with shape [B, C, H, W] + xt: the input noisy data with shape [B, C, H, W] + timestep: the timestep with shape [B] + + pred = (x_t - x_0) / sigma_t + """ + # use higher precision for calculations + original_dtype = x0_pred.dtype + x0_pred, xt, sigmas, timesteps = map( + lambda x: x.double().to(x0_pred.device), [x0_pred, xt, + scheduler.sigmas, + scheduler.timesteps] + ) + timestep_id = torch.argmin( + (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) + flow_pred = (xt - x0_pred) / sigma_t + return flow_pred.to(original_dtype) diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py index 0e23f8886..c81abd474 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py @@ -42,6 +42,11 @@ from fastvideo.pipelines.stages import (DecodingStage, DenoisingStage, TimestepPreparationStage) from fastvideo.utils import save_decoded_latents_as_video, shallow_asdict +from fastvideo.forward_context import set_forward_context +from fastvideo.models.utils import pred_noise_to_pred_video, pred_video_to_pred_noise +from fastvideo.models.loader.component_loader import TransformerLoader +from fastvideo.distributed import get_local_torch_device + logger = init_logger(__name__) @@ -70,6 +75,15 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): self.modules["scheduler"].set_timesteps(num_inference_steps=48, denoising_strength=1.0) + fastvideo_args.model_loaded["transformer"] = False + loader = TransformerLoader() + fastvideo_args.pipeline_config.dit_precision = "fp32" # Overwrite the precision to fp32 for transformer + fastvideo_args.pipeline_config.dit_forward_precision = "fp32" + self.transformer = loader.load( + fastvideo_args.model_paths["transformer"], fastvideo_args) + self.add_module("transformer", self.transformer) + fastvideo_args.model_loaded["transformer"] = True + self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage()) self.add_stage(stage_name="prompt_encoding_stage", @@ -98,6 +112,7 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): """Preprocess text-only data and generate trajectory information.""" for batch_idx, data in enumerate(self.pbar): + # logger.info("transformer weight sum: %s", sum(p.float().sum().item() for p in self.transformer.parameters())) if data is None: continue @@ -143,10 +158,12 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): sampling_params = SamplingParam.from_pretrained(args.model_path) + negative_prompt = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走' + # encode negative prompt for trajectory collection if sampling_params.guidance_scale > 1 and sampling_params.negative_prompt is not None: negative_prompt_embeds_list, negative_prompt_masks_list = self.prompt_encoding_stage.encode_text( - sampling_params.negative_prompt, + negative_prompt, fastvideo_args, encoder_index=[0], return_attention_mask=True, @@ -180,7 +197,7 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): batch.return_trajectory_latents = True # Enabling this will save the decoded trajectory videos. # Used for debugging. - batch.return_trajectory_decoded = False + batch.return_trajectory_decoded = True batch.height = args.max_height batch.width = args.max_width batch.fps = args.train_fps @@ -191,22 +208,117 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): batch, fastvideo_args) result_batch = self.timestep_preparation_stage( batch, fastvideo_args) - result_batch = self.latent_preparation_stage( - result_batch, fastvideo_args) - result_batch = self.denoising_stage(result_batch, - fastvideo_args) - # result_batch = self.decoding_stage(result_batch, - # fastvideo_args) + # result_batch = self.latent_preparation_stage( + # result_batch, fastvideo_args) + # result_batch = self.denoising_stage(result_batch, + # fastvideo_args) + noisy_input = [] + # latents = result_batch.latents.permute(0, 2, 1, 3, 4) + latents = torch.randn( + [1, 21, 16, 60, 104], dtype=torch.float32, device=get_local_torch_device() + ) + # logger.info("transformer weight sum: %s", sum(p.float().sum().item() for p in self.transformer.parameters())) + logger.info("latents sum: %s, latents shape: %s, latents dtype: %s", latents.float().sum(), latents.shape, latents.dtype) + + logger.info("scheduler timesteps: %s", self.get_module("scheduler").timesteps) + for progress_id, t in enumerate(tqdm(self.get_module("scheduler").timesteps)): + timestep = t * \ + torch.ones([1, 21], device=latents.device, dtype=torch.float32) + + noisy_input.append(latents) + + with set_forward_context( + current_timestep=0, + attn_metadata=None, + forward_batch=None, + ): + # logger.info("prompt_embed sum: %s, prompt_embed shape: %s, prompt_embed dtype: %s", prompt_embed.float().sum(), prompt_embed.shape, prompt_embed.dtype) + # logger.info("timestep: %s", timestep[:, 0]) + # Run transformer + cond_pred_noise_btchw = self.transformer( + hidden_states=latents.permute(0, 2, 1, 3, 4), + encoder_hidden_states=prompt_embed, + timestep=timestep[:, 0] + ).permute(0, 2, 1, 3, 4) + + # logger.info("cond_pred_noise_btchw sum: %s, cond_pred_noise_btchw shape: %s, cond_pred_noise_btchw dtype: %s", cond_pred_noise_btchw.float().sum(), cond_pred_noise_btchw.shape, cond_pred_noise_btchw.dtype) + + cond_pred_video_btchw = pred_noise_to_pred_video( + pred_noise=cond_pred_noise_btchw.flatten(0, 1), + noise_input_latent=latents.flatten(0, 1), + timestep=timestep.flatten(0, 1), + scheduler=self.get_module("scheduler")).unflatten( + 0, cond_pred_noise_btchw.shape[:2]) + + # logger.info("cond_pred_video_btchw sum: %s, cond_pred_video_btchw shape: %s, cond_pred_video_btchw dtype: %s", cond_pred_video_btchw.float().sum(), cond_pred_video_btchw.shape, cond_pred_video_btchw.dtype) + + with set_forward_context( + current_timestep=t, + attn_metadata=None, + forward_batch=result_batch, + ): + # logger.info("latents sum: %s, latents shape: %s, latents dtype: %s", latents.float().sum(), latents.shape, latents.dtype) + # Run transformer + uncond_pred_noise_btchw = self.transformer( + latents.permute(0, 2, 1, 3, 4), + negative_prompt_embed, + timestep[:, 0] + ).permute(0, 2, 1, 3, 4) + + # logger.info("uncond_pred_noise_btchw sum: %s, uncond_pred_noise_btchw shape: %s, uncond_pred_noise_btchw dtype: %s", uncond_pred_noise_btchw.float().sum(), uncond_pred_noise_btchw.shape, uncond_pred_noise_btchw.dtype) + + uncond_pred_video_btchw = pred_noise_to_pred_video( + pred_noise=uncond_pred_noise_btchw.flatten(0, 1), + noise_input_latent=latents.flatten(0, 1), + timestep=timestep.flatten(0, 1), + scheduler=self.get_module("scheduler")).unflatten( + 0, uncond_pred_noise_btchw.shape[:2]) + + pred_video_btchw = uncond_pred_video_btchw + batch.guidance_scale * ( + cond_pred_video_btchw - uncond_pred_video_btchw + ) + + # logger.info("pred_video_btchw sum: %s, pred_video_btchw shape: %s, pred_video_btchw dtype: %s", pred_video_btchw.float().sum(), pred_video_btchw.shape, pred_video_btchw.dtype) + + pred_noise_btchw = pred_video_to_pred_noise( + x0_pred=pred_video_btchw.flatten(0, 1), + xt=latents.flatten(0, 1), + timestep=timestep.flatten(0, 1), + scheduler=self.get_module("scheduler")).unflatten( + 0, pred_video_btchw.shape[:2]) + + # logger.info("pred_noise_btchw sum: %s, pred_noise_btchw shape: %s, pred_noise_btchw dtype: %s", pred_noise_btchw.float().sum(), pred_noise_btchw.shape, pred_noise_btchw.dtype) + + latents = self.get_module("scheduler").step( + pred_noise_btchw.flatten(0, 1), + self.get_module("scheduler").timesteps[progress_id] * torch.ones( + [1, 21], device=latents.device, dtype=torch.long).flatten(0, 1), + latents.flatten(0, 1) + )[0].unflatten(dim=0, sizes=pred_noise_btchw.shape[:2]) + + # logger.info("latents sum: %s, latents shape: %s, latents dtype: %s", latents.float().sum(), latents.shape, latents.dtype) + + noisy_input.append(latents) + + noisy_inputs = torch.stack(noisy_input, dim=1) + + noisy_inputs = noisy_inputs[:, [0, 12, 24, 36, -1]].half() + + logger.info("noisy inputs sum: %s, noisy inputs shape: %s, noisy inputs dtype: %s", noisy_inputs.float().sum(), noisy_inputs.shape, noisy_inputs.dtype) + + result_batch.trajectory_latents = noisy_inputs.permute(0, 1, 3, 2, 4, 5) + result_batch.trajectory_timesteps = torch.tensor([self.get_module("scheduler").timesteps[i] for i in [0, 12, 24, 36, -1]]) + result_batch.latents = latents.permute(0, 2, 1, 3, 4) + result_batch = self.decoding_stage(result_batch, + fastvideo_args) trajectory_latents.append( result_batch.trajectory_latents.cpu()) trajectory_timesteps.append( result_batch.trajectory_timesteps.cpu()) - # trajectory_decoded.append(result_batch.trajectory_decoded) + trajectory_decoded.append(result_batch.trajectory_decoded) trajectory_latents = torch.stack(trajectory_latents, dim=0).squeeze(0) - trajecotry_latents = trajectory_latents[:, [0, 12, 24, 36, -1]] - logger.info("trajecotry_latents sum: %s, trajecotry_latents shape: %s, trajecotry_latents dtype: %s", trajecotry_latents.float().sum(), trajecotry_latents.shape, trajecotry_latents.dtype) - assert False + # trajecotry_latents = trajectory_latents[:, [0, 12, 24, 36, -1]] # Prepare extra features for text-only processing extra_features = { diff --git a/fastvideo/tests/transformers/test_ori_wanvideo.py b/fastvideo/tests/transformers/test_ori_wanvideo.py index 11b5a3460..f15fe4619 100644 --- a/fastvideo/tests/transformers/test_ori_wanvideo.py +++ b/fastvideo/tests/transformers/test_ori_wanvideo.py @@ -22,6 +22,7 @@ logger = init_logger(__name__) os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29503" +os.environ["FASTVIDEO_FORCE_ATTN_BF16"] = "1" BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join( @@ -32,19 +33,21 @@ TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer") @pytest.mark.usefixtures("distributed_setup") def test_ori_wan_transformer(): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - precision = torch.bfloat16 - precision_str = "bf16" + precision = torch.float32 + precision_str = "fp32" args = FastVideoArgs(model_path=TRANSFORMER_PATH, dit_cpu_offload=True, - pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str)) + pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str, dit_forward_precision=precision_str)) args.device = device loader = TransformerLoader() - model2 = loader.load(TRANSFORMER_PATH, args).to(dtype=precision) + model2 = loader.load(TRANSFORMER_PATH, args) model1 = WanModel.from_pretrained( - "/mnt/weka/home/hao.zhang/wei/Self-Forcing/wan_models/Wan2.1-T2V-1.3B", device=device, - torch_dtype=precision).to(device, dtype=precision).requires_grad_(False) + "/mnt/weka/home/hao.zhang/wei/Self-Forcing-clean/wan_models/Wan2.1-T2V-1.3B") + model1.eval() + model1 = model1.to(device).to(precision) + model1.requires_grad_(False) total_params = sum(p.numel() for p in model1.parameters()) # Calculate weight sum for model1 (converting to float64 to avoid overflow) @@ -75,29 +78,35 @@ def test_ori_wan_transformer(): # Create identical inputs for both models batch_size = 1 - text_seq_len = 30 - seq_len = math.ceil((160 * 90) / + text_seq_len = 120 + seq_len = math.ceil((104 * 60) / (2 * 2) * 21) # Video latents [B, C, T, H, W] hidden_states = torch.randn(batch_size, - 16, 21, - 160, - 90, + 16, + 60, + 104, device=device, + generator=torch.Generator("cuda").manual_seed(1024), dtype=precision) + logger.info("Hidden states sum: %s, Hidden states shape: %s, Hidden states dtype: %s", hidden_states.float().sum(), hidden_states.shape, hidden_states.dtype) + # Text embeddings [B, L, D] (including global token) - encoder_hidden_states = torch.randn(batch_size, - text_seq_len + 1, - 4096, - device=device, - dtype=precision) + # encoder_hidden_states = torch.randn(batch_size, + # text_seq_len + 1, + # 4096, + # device=device, + # dtype=precision) + encoder_hidden_states = torch.load("../sf_cond_prompt_embeds.pt").to(device, dtype=precision) + + logger.info("Encoder hidden states sum: %s, Encoder hidden states shape: %s, Encoder hidden states dtype: %s", encoder_hidden_states.float().sum(), encoder_hidden_states.shape, encoder_hidden_states.dtype) # Timestep - timestep = torch.tensor([500], device=device, dtype=precision) + timestep = torch.tensor([995.7627], device=device, dtype=precision) forward_batch = ForwardBatch( data_type="dummy", @@ -105,19 +114,19 @@ def test_ori_wan_transformer(): # with torch.amp.autocast('cuda', dtype=precision): output1 = model1( - x=hidden_states, + x=hidden_states.permute(0, 2, 1, 3, 4), context=encoder_hidden_states, t=timestep, seq_len=seq_len, - ) + ).permute(0, 2, 1, 3, 4) with set_forward_context( current_timestep=0, attn_metadata=None, forward_batch=forward_batch, ): - output2 = model2(hidden_states=hidden_states, + output2 = model2(hidden_states=hidden_states.permute(0, 2, 1, 3, 4), encoder_hidden_states=encoder_hidden_states, - timestep=timestep) + timestep=timestep).permute(0, 2, 1, 3, 4) # Check if outputs have the same shape assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" @@ -128,6 +137,8 @@ def test_ori_wan_transformer(): mean_diff = torch.mean(torch.abs(output1 - output2)) logger.info("Max Diff: %s", max_diff.item()) logger.info("Mean Diff: %s", mean_diff.item()) + logger.info("Output 1 sum: %s", output1.float().sum()) + logger.info("Output 2 sum: %s", output2.float().sum()) assert max_diff < 1e-4, f"Maximum difference between outputs: {max_diff.item()}" # mean diff assert mean_diff < 1e-4, f"Mean difference between outputs: {mean_diff.item()}" \ No newline at end of file