nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Refactor ode trajectory generation

Browse Source 
Add env variable to enforce running flash attn
This commit is contained in:
JerryZhou54
2025-09-26 01:26:03 +00:00
parent d31695b736
commit 92dfbfe4fc
9 changed files with 207 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/attention/layer.py
+3 -2
View File
@@ -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,
fastvideo/configs/pipelines/base.py
+1
View File
@@ -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)
fastvideo/envs.py
+5
View File
@@ -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
fastvideo/models/dits/wanvideo.py
+13 -3
View File
@@ -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)
fastvideo/models/loader/component_loader.py
+4 -1
View File
@@ -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)
fastvideo/models/loader/fsdp_load.py
+3 -2
View File
@@ -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,
fastvideo/models/utils.py
+22
View File
@@ -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)
fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py
+124 -12
View File
@@ -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 = {
fastvideo/tests/transformers/test_ori_wanvideo.py
+32 -21
View File
@@ -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()}"