366 lines
14 KiB
Python
366 lines
14 KiB
Python
"""Benchmark MLX FastWan quantization modes with one shared prompt encode."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
from typing import cast
|
|
|
|
import numpy as np
|
|
|
|
from examples.inference.basic.mlx_wan_prompt_to_video import (
|
|
DEFAULT_MODEL_ROOT,
|
|
decode_latents_to_video,
|
|
encode_prompt,
|
|
make_rotary_embeddings,
|
|
)
|
|
from fastvideo.mlx_runtime.memory import cleanup_mlx
|
|
|
|
|
|
def _parse_modes(raw: str) -> list[str]:
|
|
"""
|
|
Parse and validate a comma-separated list of quantization modes.
|
|
|
|
Parameters:
|
|
raw (str): Comma-separated mode names.
|
|
|
|
Returns:
|
|
list[str]: Normalized, whitespace-trimmed mode names.
|
|
|
|
Raises:
|
|
ValueError: If any mode is unsupported.
|
|
"""
|
|
modes = [mode.strip() for mode in raw.split(",") if mode.strip()]
|
|
allowed = {"none", "int8", "int4", "mxfp8", "mxfp4", "nvfp4"}
|
|
unknown = sorted(set(modes) - allowed)
|
|
if unknown:
|
|
raise ValueError(f"Unsupported modes: {unknown}")
|
|
return modes
|
|
|
|
|
|
def _latent_delta_metrics(candidate: np.ndarray, baseline: np.ndarray) -> dict[str, float]:
|
|
"""
|
|
Compare candidate and baseline latent arrays using error and signal-quality metrics.
|
|
|
|
Parameters:
|
|
candidate (np.ndarray): Latent array to evaluate.
|
|
baseline (np.ndarray): Reference latent array for comparison.
|
|
|
|
Returns:
|
|
dict[str, float]: Mean squared error, mean absolute error, maximum absolute
|
|
error, and signal-to-noise ratio in decibels between the arrays.
|
|
"""
|
|
diff = candidate.astype(np.float32) - baseline.astype(np.float32)
|
|
mse = float(np.mean(np.square(diff)))
|
|
mae = float(np.mean(np.abs(diff)))
|
|
max_abs = float(np.max(np.abs(diff)))
|
|
signal = float(np.mean(np.square(baseline.astype(np.float32))))
|
|
return {
|
|
"latent_mse_vs_fp16": mse,
|
|
"latent_mae_vs_fp16": mae,
|
|
"latent_max_abs_vs_fp16": max_abs,
|
|
"latent_snr_db_vs_fp16": float(10.0 * np.log10(signal / mse)) if mse > 0 else float("inf"),
|
|
}
|
|
|
|
|
|
def _torch_mps_memory() -> dict[str, int | None]:
|
|
"""
|
|
Report PyTorch MPS memory statistics when PyTorch MPS is available.
|
|
|
|
Returns:
|
|
dict[str, int | None]: A mapping of MPS memory metric names to byte counts, or `None` values when PyTorch or MPS is unavailable.
|
|
"""
|
|
try:
|
|
import torch
|
|
except ImportError:
|
|
return {
|
|
"torch_mps_current_allocated_bytes": None,
|
|
"torch_mps_driver_allocated_bytes": None,
|
|
"torch_mps_recommended_max_bytes": None,
|
|
}
|
|
if not torch.backends.mps.is_available():
|
|
return {
|
|
"torch_mps_current_allocated_bytes": None,
|
|
"torch_mps_driver_allocated_bytes": None,
|
|
"torch_mps_recommended_max_bytes": None,
|
|
}
|
|
return {
|
|
"torch_mps_current_allocated_bytes": int(torch.mps.current_allocated_memory()),
|
|
"torch_mps_driver_allocated_bytes": int(torch.mps.driver_allocated_memory()),
|
|
"torch_mps_recommended_max_bytes": int(torch.mps.recommended_max_memory()),
|
|
}
|
|
|
|
|
|
def _decode_with_metrics(*, args, latents: np.ndarray, output_path: Path) -> dict[str, float | int | None | str]:
|
|
"""
|
|
Decode latents to a video and collect export timing and PyTorch MPS memory metrics.
|
|
|
|
Parameters:
|
|
args: Configuration values for decoding and video export.
|
|
latents (np.ndarray): Latent representation to decode.
|
|
output_path (Path): Destination path for the exported video.
|
|
|
|
Returns:
|
|
dict[str, float | int | None | str]: Video export duration and PyTorch MPS memory measurements.
|
|
"""
|
|
before = _torch_mps_memory()
|
|
decode_start = time.perf_counter()
|
|
decode_latents_to_video(
|
|
model_root=args.model_root,
|
|
latents_np=latents,
|
|
output_path=output_path,
|
|
fps=args.fps,
|
|
device_arg=args.torch_device,
|
|
dtype_arg=args.torch_dtype,
|
|
backend=args.decode_backend,
|
|
taehv_source_path=args.taehv_source_path,
|
|
taehv_checkpoint_path=args.taehv_checkpoint_path,
|
|
taehv_parallel=args.taehv_parallel,
|
|
)
|
|
decode_time = time.perf_counter() - decode_start
|
|
after = _torch_mps_memory()
|
|
return {
|
|
"decode_export_s": decode_time,
|
|
"decode_torch_mps_current_before_bytes": before["torch_mps_current_allocated_bytes"],
|
|
"decode_torch_mps_current_after_bytes": after["torch_mps_current_allocated_bytes"],
|
|
"decode_torch_mps_driver_before_bytes": before["torch_mps_driver_allocated_bytes"],
|
|
"decode_torch_mps_driver_after_bytes": after["torch_mps_driver_allocated_bytes"],
|
|
"decode_torch_mps_recommended_max_bytes": after["torch_mps_recommended_max_bytes"],
|
|
}
|
|
|
|
|
|
def _run_one_mode(
|
|
*,
|
|
mode: str,
|
|
args,
|
|
config: dict,
|
|
checkpoint_path: Path,
|
|
config_path: Path,
|
|
prompt_embeds,
|
|
freqs_cis,
|
|
):
|
|
"""
|
|
Run denoising for one quantization mode and collect performance and memory metrics.
|
|
|
|
Parameters:
|
|
mode (str): Quantization mode to benchmark.
|
|
args: Benchmark configuration, including dtype, dimensions, seed, scheduler, and denoising settings.
|
|
config (dict): Model configuration containing the input channel count.
|
|
checkpoint_path (Path): Path to the transformer checkpoint.
|
|
config_path (Path): Path to the transformer configuration.
|
|
prompt_embeds: Encoded prompt embeddings shared across benchmark modes.
|
|
freqs_cis: Rotary positional embeddings used during denoising.
|
|
|
|
Returns:
|
|
dict: The mode name, generated latent array, and metrics for model loading,
|
|
denoising, step timing, and MLX memory usage.
|
|
"""
|
|
import mlx.core as mx
|
|
import torch
|
|
|
|
from fastvideo.benchmarks.mlx_fastwan_bench import denoise_dmd_on_device
|
|
from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
|
|
from fastvideo.mlx_runtime.fastwan import mlx_dit_from_diffusers_safetensors
|
|
from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, dmd_step
|
|
|
|
mx_dtype = mx.float16 if args.mlx_dtype == "fp16" else mx.float32
|
|
quantization = None if mode == "none" else mode
|
|
latent_frames = (args.num_frames - 1) // 4 + 1
|
|
latent_height = args.height // 8
|
|
latent_width = args.width // 8
|
|
|
|
load_start = time.perf_counter()
|
|
mx.clear_cache()
|
|
mx.reset_peak_memory()
|
|
dit = mlx_dit_from_diffusers_safetensors(
|
|
checkpoint_path,
|
|
config_path,
|
|
dtype=args.mlx_dtype,
|
|
quantization=quantization,
|
|
)
|
|
load_time = time.perf_counter() - load_start
|
|
load_peak_memory = mx.get_peak_memory()
|
|
|
|
scheduler = FlowMatchEulerDiscreteScheduler(shift=args.flow_shift)
|
|
schedule = MLXDMDSchedule.from_torch_scheduler(scheduler)
|
|
timesteps = [int(step.strip()) for step in args.dmd_denoising_steps.split(",") if step.strip()]
|
|
# Same torch generator sequence as the original host-round-trip loop
|
|
# (initial latents first, then one re-noise draw per intermediate step),
|
|
# so every mode still shares identical stochasticity.
|
|
generator = torch.Generator(device="cpu").manual_seed(args.seed)
|
|
latents_seed = torch.randn(
|
|
(1, int(config["in_channels"]), latent_frames, latent_height, latent_width),
|
|
generator=generator,
|
|
dtype=torch.float32,
|
|
).numpy()
|
|
renoise_by_step = [
|
|
torch.randn(latents_seed.shape, generator=generator, dtype=torch.float32).numpy()
|
|
for _ in range(max(0, len(timesteps) - 1))
|
|
]
|
|
latents = mx.array(latents_seed).astype(mx_dtype)
|
|
encoder_hidden_states = mx.array(prompt_embeds.numpy()).astype(mx_dtype)
|
|
|
|
denoise_start = time.perf_counter()
|
|
mx.reset_peak_memory()
|
|
latents_np, step_times = denoise_dmd_on_device(
|
|
mx=mx,
|
|
dit=dit,
|
|
latents=latents,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
freqs_cis=freqs_cis,
|
|
timesteps=timesteps,
|
|
renoise_by_step=renoise_by_step,
|
|
schedule=schedule,
|
|
dmd_step=dmd_step,
|
|
mx_dtype=mx_dtype,
|
|
)
|
|
denoise_time = time.perf_counter() - denoise_start
|
|
denoise_peak_memory = mx.get_peak_memory()
|
|
active_memory = mx.get_active_memory()
|
|
return {
|
|
"mode": mode,
|
|
"latents": latents_np,
|
|
"metrics": {
|
|
"mlx_dit_load_s": load_time,
|
|
"mlx_denoise_s": denoise_time,
|
|
"mlx_denoise_first_step_s": step_times[0] if step_times else None,
|
|
"mlx_load_peak_bytes": int(load_peak_memory),
|
|
"mlx_denoise_peak_bytes": int(denoise_peak_memory),
|
|
"mlx_active_after_denoise_bytes": int(active_memory),
|
|
},
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
"""
|
|
Run the MLX FastWan quantization benchmark for the selected modes and write latency, memory, output, and latent-difference metrics to the output directory.
|
|
"""
|
|
parser = argparse.ArgumentParser(description="Benchmark MLX FastWan quantization modes.")
|
|
parser.add_argument("--model-root", type=Path, default=DEFAULT_MODEL_ROOT)
|
|
parser.add_argument("--prompt", default="A snow leopard walks across a windy mountain ridge.")
|
|
parser.add_argument("--height", type=int, default=192)
|
|
parser.add_argument("--width", type=int, default=320)
|
|
parser.add_argument("--num-frames", type=int, default=17)
|
|
parser.add_argument("--dmd-denoising-steps", default="1000,757,522")
|
|
parser.add_argument("--flow-shift", type=float, default=8.0)
|
|
parser.add_argument("--max-sequence-length", type=int, default=256)
|
|
parser.add_argument("--seed", type=int, default=1024)
|
|
parser.add_argument("--fps", type=int, default=16)
|
|
parser.add_argument("--torch-device", default="auto")
|
|
parser.add_argument("--torch-dtype", choices=("fp16", "fp32"), default="fp16")
|
|
parser.add_argument("--mlx-dtype", choices=("fp16", "fp32"), default="fp16")
|
|
parser.add_argument("--modes", default="none,int8,int4,mxfp8,mxfp4,nvfp4")
|
|
parser.add_argument("--output-dir", type=Path, default=Path("video_samples/mlx_quant_benchmark"))
|
|
parser.add_argument("--decode-backend", choices=("none", "wan-vae", "taehv"), default="taehv")
|
|
parser.add_argument("--taehv-source-path", type=Path, default=None)
|
|
parser.add_argument("--taehv-checkpoint-path", type=Path, default=None)
|
|
parser.add_argument("--taehv-parallel", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
import mlx.core as mx
|
|
import torch
|
|
|
|
mx.random.seed(args.seed)
|
|
torch.manual_seed(args.seed)
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
config_path = args.model_root / "transformer/config.json"
|
|
checkpoint_path = args.model_root / "transformer/diffusion_pytorch_model.safetensors"
|
|
config = json.loads(config_path.read_text())
|
|
latent_frames = (args.num_frames - 1) // 4 + 1
|
|
latent_height = args.height // 8
|
|
latent_width = args.width // 8
|
|
|
|
prompt_start = time.perf_counter()
|
|
prompt_embeds = encode_prompt(
|
|
model_root=args.model_root,
|
|
prompt=args.prompt,
|
|
max_sequence_length=args.max_sequence_length,
|
|
device_arg=args.torch_device,
|
|
dtype_arg=args.torch_dtype,
|
|
)
|
|
prompt_time = time.perf_counter() - prompt_start
|
|
freqs_cis = make_rotary_embeddings(
|
|
config,
|
|
latent_frames=latent_frames,
|
|
latent_height=latent_height,
|
|
latent_width=latent_width,
|
|
)
|
|
|
|
from fastvideo.mlx_runtime.fastwan import UnsupportedMLXQuantizationError
|
|
|
|
baseline_latents = None
|
|
rows = []
|
|
for mode in _parse_modes(args.modes):
|
|
print(f"=== MLX quant mode: {mode} ===")
|
|
mode_start = time.perf_counter()
|
|
try:
|
|
result = _run_one_mode(
|
|
mode=mode,
|
|
args=args,
|
|
config=config,
|
|
checkpoint_path=checkpoint_path,
|
|
config_path=config_path,
|
|
prompt_embeds=prompt_embeds,
|
|
freqs_cis=freqs_cis,
|
|
)
|
|
except UnsupportedMLXQuantizationError as exc:
|
|
print(f"skipping mode (unsupported by this MLX build): {exc}")
|
|
rows.append({"mode": mode, "status": "unsupported_by_mlx", "error": str(exc)})
|
|
continue
|
|
cleanup_mlx(mx)
|
|
latents = result["latents"]
|
|
if baseline_latents is None:
|
|
baseline_latents = latents
|
|
latent_path = args.output_dir / f"latents_{mode}.npy"
|
|
np.save(latent_path, latents)
|
|
|
|
decode_time = 0.0
|
|
decode_metrics = {}
|
|
output_path = None
|
|
if args.decode_backend != "none":
|
|
output_path = args.output_dir / f"video_{mode}_{args.decode_backend}_{args.height}x{args.width}x{args.num_frames}.mp4"
|
|
decode_metrics = _decode_with_metrics(args=args, latents=latents, output_path=output_path)
|
|
decode_time = cast(float, decode_metrics["decode_export_s"])
|
|
|
|
mode_total = time.perf_counter() - mode_start
|
|
mlx_denoise_peak_bytes = int(result["metrics"]["mlx_denoise_peak_bytes"])
|
|
mlx_active_bytes = int(result["metrics"]["mlx_active_after_denoise_bytes"])
|
|
metrics = {
|
|
"mode": mode,
|
|
"status": "ok",
|
|
"prompt_encode_shared_s": prompt_time,
|
|
"height": args.height,
|
|
"width": args.width,
|
|
"num_frames": args.num_frames,
|
|
"decode_backend": args.decode_backend,
|
|
"decode_export_s": decode_time,
|
|
"mode_total_excluding_shared_prompt_s": mode_total,
|
|
"mode_total_including_shared_prompt_s": mode_total + prompt_time,
|
|
"latents_path": str(latent_path),
|
|
"output_path": str(output_path) if output_path else None,
|
|
"mlx_denoise_peak_gib": mlx_denoise_peak_bytes / (1024**3),
|
|
"mlx_active_after_denoise_gib": mlx_active_bytes / (1024**3),
|
|
"mlx_dit_peak_under_16gb": mlx_denoise_peak_bytes < 16 * 1024**3,
|
|
"mlx_dit_active_under_16gb": mlx_active_bytes < 16 * 1024**3,
|
|
"mac_16gb_status": (
|
|
"dit_memory_fits_16gb_measured_decode_separately"
|
|
if mlx_denoise_peak_bytes < 16 * 1024**3 else "dit_memory_exceeds_16gb"
|
|
),
|
|
**result["metrics"],
|
|
**decode_metrics,
|
|
**_latent_delta_metrics(latents, baseline_latents),
|
|
}
|
|
rows.append(metrics)
|
|
print(json.dumps(metrics, indent=2))
|
|
|
|
metrics_path = args.output_dir / "metrics.json"
|
|
metrics_path.write_text(json.dumps(rows, indent=2))
|
|
print(f"Wrote benchmark metrics to: {metrics_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|