From a48e08150caaade9d5770a85a430910b37a949e3 Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Thu, 17 Sep 2026 18:51:37 +0000 Subject: [PATCH] feat(wan): fast serving with persistent AOT caching and tuned inference recipe Delivers production-grade fast serving for Wan 2.2 T2V-A14B on Cloud TPU: - Persistent AOT Compilation Caching with source hash and static graphdef metadata - Tuned Serving Recipe: Production launcher with platform-specific v6e and v7 profiles - Clean Key-Centering & 4-Operand Kernel (mk, q, k, v) without VMEM register pressure - Merged Pre-A2A Norm Reduction (pmax) removing ring all_gather relayout barrier - Hoisted scalar-prefetch metadata in _lse_scan eliminating collective-permute stalls - Lossless 6D unpatchify transpose optimization (collapsing p_t=1 to avoid 8D stride copies) --- end_to_end/tpu/run_wan_fast_inference.sh | 276 +++++++++++++-- src/maxdiffusion/aot_cache.py | 122 ++++--- src/maxdiffusion/configs/base_wan_14b.yml | 3 +- src/maxdiffusion/configs/base_wan_animate.yml | 2 +- src/maxdiffusion/configs/base_wan_i2v_14b.yml | 2 +- src/maxdiffusion/configs/base_wan_i2v_27b.yml | 2 +- src/maxdiffusion/generate_wan.py | 185 ++++++++-- src/maxdiffusion/max_utils.py | 30 +- src/maxdiffusion/models/attention_flax.py | 4 +- .../wan/transformers/transformer_wan.py | 52 ++- .../pipelines/wan/wan_pipeline.py | 1 + src/maxdiffusion/pyconfig.py | 12 +- src/maxdiffusion/tests/aot_cache_test.py | 70 ++++ .../tests/wan/wan_transformer_test.py | 322 +++++++++++++++++- src/maxdiffusion/utils/export_utils.py | 40 ++- 15 files changed, 969 insertions(+), 154 deletions(-) diff --git a/end_to_end/tpu/run_wan_fast_inference.sh b/end_to_end/tpu/run_wan_fast_inference.sh index 89680760b..fe62b6341 100755 --- a/end_to_end/tpu/run_wan_fast_inference.sh +++ b/end_to_end/tpu/run_wan_fast_inference.sh @@ -14,7 +14,13 @@ # limitations under the License. # WAN T2V fast-serving example: AOT executable cache + converted-weights -# cache + zero-exec warmup, with a tuned v7 2D-ring attention recipe. +# cache + zero-exec warmup, with a tuned 2D-ring attention recipe. +# +# The XLA flag set, attention tile and text-encoder options differ per TPU +# generation, so the platform is auto-detected from the GCE metadata server +# and the matching recipe is selected (see "TPU platform detection" below). +# Supported profiles: v6e (Trillium) and v7. Anything else falls back to a +# conservative generic profile. # # First run per (model, shape) pays one-time conversion + compile and # populates the caches; every later process start is ~25s to ready. @@ -26,15 +32,18 @@ # OUTPUT_DIR video/metrics output (default /tmp/wan_out) # COMPILE_TE=true torch.compile the text encoder (adds ~30s to load, # saves ~10s/encode; worth it for long-lived processes) -# FIXEDM=0 plain online softmax instead of fixed-m # EXTRA_LIBTPU extra libtpu flags, appended to the tuned set -# -# 720p 81f / 40 steps denoise: fixed-m 105.3s, plain 109.6s. Each mode gets its -# own optimal tile below; plain degrades badly on fixed-m's. +# TPU_PROFILE force a platform profile (v6e|v7|generic), skipping +# autodetection +# ACCEL_TYPE force the raw accelerator type (e.g. v6e-8) +# ATTENTION / ULYSSES_SHARDS / BQ / BKV / BQ_DKV / VMEM_LIMIT_BYTES +# override the per-platform attention recipe +# DP / CP override the derived ici_data/context_parallelism set -u MODEL=${1:-22} STEPS=${2:-40} PROMPT=${3:-""} +shift $(($# > 3 ? 3 : $#)) PROJECT_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &> /dev/null && pwd)" cd "$PROJECT_ROOT" || exit 1 @@ -42,38 +51,232 @@ export PYTHONPATH="$PROJECT_ROOT/src:${PYTHONPATH:-}" export HF_HUB_ENABLE_HF_TRANSFER=1 export JAX_DEFAULT_MATMUL_PRECISION=bfloat16 export TORCHINDUCTOR_FX_GRAPH_CACHE=1 +export XLA_PYTHON_CLIENT_MEM_FRACTION=0.95 +# Without these the JAX persistent cache silently skips most entries, so the +# "warm" start still recompiles a large part of the graph. +export JAX_PERSISTENT_CACHE_MIN_ENTRY_SIZE_BYTES=-1 +export JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS=0 CACHE_ROOT=${WAN_CACHE_ROOT:-$HOME/.cache/maxdiffusion_wan} -OUTPUT_DIR=${OUTPUT_DIR:-/tmp/wan_out} -mkdir -p "$CACHE_ROOT/jax" "$CACHE_ROOT/aot_wan$MODEL" "$CACHE_ROOT/converted" "$OUTPUT_DIR" - -# Tuned collective/scheduler flag set for v7 (from the PR #430 2D-ring -# baseline). One line: libtpu stops parsing at a literal backslash. -export LIBTPU_INIT_ARGS="--xla_tpu_spmd_rng_bit_generator_unsafe=true --xla_tpu_enable_dot_strength_reduction=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_enable_async_collective_permute=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true --xla_tpu_scoped_vmem_limit_kib=65536 --xla_tpu_enable_async_all_to_all=true --xla_tpu_enable_all_experimental_scheduler_features=true --xla_tpu_enable_scheduler_memory_pressure_tracking=true --xla_tpu_host_transfer_overlap_limit=24 --xla_tpu_aggressive_opt_barrier_removal=ENABLED --xla_lhs_prioritize_async_depth_over_stall=ENABLED --xla_should_allow_loop_variant_parameter_in_chain=ENABLED --xla_should_add_loop_invariant_op_in_chain=ENABLED --xla_tpu_enable_ici_ag_pipelining=true --xla_max_concurrent_host_send_recv=100 --xla_tpu_scheduler_percent_shared_memory_limit=100 --xla_latency_hiding_scheduler_rerun=2 --xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 --xla_tpu_enable_latency_hiding_scheduler=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_megacore_fusion=true --xla_tpu_megacore_fusion_allow_ags=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true --xla_tpu_sparse_core_all_gather_latency_multiplier=1 --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 --xla_tpu_enable_sparse_core_collective_aggregator=true --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true --xla_tpu_enable_sparse_core_reduce_scatter_v2=true --xla_tpu_enable_sparse_core_collective_offload_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true --xla_tpu_enable_concurrent_sparse_core_offloading=true --xla_tpu_assign_all_reduce_scatter_layout=true" -# Timings only compare across runs passing the same extra flags. -export LIBTPU_INIT_ARGS="${LIBTPU_INIT_ARGS} ${EXTRA_LIBTPU:-}" - -# fixed-m on by default: faster, and covered by tests/ring_fixed_m_test.py. -if [ "${FIXEDM:-1}" = "1" ]; then - ATTENTION=ulysses_ring_custom_fixed_m - BQ=6400; BKV=2048 +OUTPUT_DIR=${OUTPUT_DIR:-$HOME/maxdiffusion_wan_output} +export TMPDIR=${TMPDIR:-$CACHE_ROOT/tmp} +export TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR:-$CACHE_ROOT/torch_compile} +mkdir -p "$CACHE_ROOT/jax" "$CACHE_ROOT/aot_wan$MODEL" "$CACHE_ROOT/converted" \ + "$OUTPUT_DIR" "$TMPDIR" "$TORCHINDUCTOR_CACHE_DIR" + +# --------------------------------------------------------------------------- +# TPU platform detection +# --------------------------------------------------------------------------- +# Preference order: explicit override -> TPU_ACCELERATOR_TYPE env (set by some +# runtimes) -> GCE metadata "accelerator-type" (e.g. "v6e-8") -> the +# ACCELERATOR_TYPE line inside the "tpu-env" metadata blob. Detection is pure +# metadata/env: it must not initialise the TPU, or it would take the device +# before the real process starts. +_tpu_metadata() { + curl -s -m 2 -H 'Metadata-Flavor: Google' \ + "http://metadata.google.internal/computeMetadata/v1/instance/attributes/$1" 2> /dev/null +} + +_detect_accel_type() { + local t="${TPU_ACCELERATOR_TYPE:-}" + [ -z "$t" ] && t="$(_tpu_metadata accelerator-type)" + [ -z "$t" ] && t="$(_tpu_metadata tpu-env | sed -n "s/^ACCELERATOR_TYPE: *'\([^']*\)'.*/\1/p")" + # Guard against metadata returning an HTML error page. Real names seen in the + # wild: "v6e-8", "v5litepod-8", "tpu7x-8" (v7 reports as tpu7x, not v7x). + case "$t" in + v[0-9]* | tpu[0-9]*) printf '%s' "$t" ;; + *) printf '' ;; + esac +} + +ACCEL_TYPE=${ACCEL_TYPE:-$(_detect_accel_type)} +if [ -n "$ACCEL_TYPE" ]; then + TPU_GEN="${ACCEL_TYPE%%-*}" # v6e-8 -> v6e + TPU_CHIPS="${ACCEL_TYPE##*-}" # v6e-8 -> 8 else - ATTENTION=ulysses_ring_custom - BQ=9472; BKV=1024 + TPU_GEN="" + TPU_CHIPS="" fi +case "$TPU_CHIPS" in + '' | *[!0-9]*) TPU_CHIPS="" ;; +esac + +if [ -z "${TPU_PROFILE:-}" ]; then + case "$TPU_GEN" in + v6e | v5litepod | v5e) TPU_PROFILE=v6e ;; + v7 | v7x | v7p | v7e | tpu7 | tpu7x | tpu7p | tpu7e) TPU_PROFILE=v7 ;; + *) TPU_PROFILE=generic ;; + esac +fi + +# --------------------------------------------------------------------------- +# Per-platform XLA flags +# --------------------------------------------------------------------------- +# One line per variable: libtpu tokenises LIBTPU_INIT_ARGS on whitespace and +# stops parsing at the first literal backslash, so a multi-line continuation +# inside a *single-quoted* string silently drops every flag after the first +# line break. Keep these as single-line strings. +COMMON_LIBTPU="--xla_tpu_spmd_rng_bit_generator_unsafe=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_memory_bound_loop_optimizer_options=enabled:true --xla_tpu_enable_dot_strength_reduction=true --xla_enable_async_collective_permute=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true --xla_tpu_scoped_vmem_limit_kib=65536 --xla_tpu_enable_async_all_to_all=true --xla_tpu_enable_all_experimental_scheduler_features=true --xla_tpu_enable_scheduler_memory_pressure_tracking=true --xla_tpu_host_transfer_overlap_limit=24 --xla_tpu_aggressive_opt_barrier_removal=ENABLED --xla_lhs_prioritize_async_depth_over_stall=ENABLED --xla_should_allow_loop_variant_parameter_in_chain=ENABLED --xla_should_add_loop_invariant_op_in_chain=ENABLED --xla_tpu_enable_ici_ag_pipelining=true --xla_max_concurrent_host_send_recv=100 --xla_tpu_scheduler_percent_shared_memory_limit=100 --xla_latency_hiding_scheduler_rerun=2 --xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 --xla_tpu_enable_latency_hiding_scheduler=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true --xla_tpu_sparse_core_all_gather_latency_multiplier=1 --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 --xla_tpu_enable_sparse_core_collective_aggregator=true --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true --xla_tpu_enable_sparse_core_reduce_scatter_v2=true --xla_tpu_enable_sparse_core_collective_offload_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true --xla_tpu_enable_concurrent_sparse_core_offloading=true --xla_tpu_assign_all_reduce_scatter_layout=true" + +# v6e (Trillium). xla_tpu_enable_async_collective_fusion MUST stay false here: +# continuation fusion for AllGather is a Viperlite-only path and aborts backend +# init on v6e with +# FAILED_PRECONDITION: Continuation fusion for AllGather is enabled on +# platform ghostlite ... Please disable it with +# xla_tpu_enable_async_collective_fusion=false +# The two sub-flags are inert while the parent is false; they are kept so the +# set stays a one-line edit away from a platform that does support it. +# Note xla_tpu_memory_bound_loop_optimizer_options is proto-valued, so +# "enabled:true" is its correct syntax -- unlike the bare enum flags above, +# which only accept true/enabled/ENABLED. +V6E_LIBTPU="--xla_tpu_enable_async_collective_fusion=false" + +# v7 has two TensorCores per chip, so megacore fusion applies; v6e has one and +# ignores these. +V7_LIBTPU="--xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_enable_megacore_fusion=true --xla_tpu_megacore_fusion_allow_ags=true" + +# --------------------------------------------------------------------------- +# Attention recipe: measurements behind the per-platform defaults +# --------------------------------------------------------------------------- +# The kernel, Ulysses degree and tile are set per platform in the profile case +# below: +# - v6e profile: ulysses_custom_fixed_m_per_q_block, U=4 (R=1 on CP=4), BQ=9472, BKV=1024 +# - v7 profile: ulysses_ring_custom_fixed_m_per_q_block, U=2 (R=2 on CP=4), BQ=6400, BKV=2048 +# +# Cloud TPU v6e-8 (Wan 2.2 T2V-A14B, 720p, 81f, 40 steps, CP=4, DP=2): +# this script, single prompt, end to end: 138.7s (denoise 136.0s, VAE 0.8s) +# origin/main @ a4ff3aec, single prompt: 156.4s (denoise 153.3s, VAE 2.9s) +# Denoise improvement: -17.3s (-11.3%); VAE decode: 3.6x faster (0.8s vs 2.9s). +# Outputs are not bit-identical to the baseline (fixed-m reorders float +# accumulation): mean pairwise PSNR +3.25 dB vs exact online softmax, all VBench dimensions within 0.25%. +# +# Cloud TPU v7 (tpu7x-8, Wan 2.2 T2V-A14B, 720p, 81f, 40 steps, CP=4, DP=2): +# this script (profile v7, U=2, R=2, BQ=6400, BKV=2048): 115.2s (denoise 114.4s, VAE 0.8s) +# origin/main @ 1bc54811, single prompt: 119.8s (denoise 119.0s, VAE 0.8s) +# Total inference beats origin/main by -4.6s (-3.9%); xprof collective-permute-done: 2405ms -> 47ms. +# persistent AOT: compile 94.7s cold -> 15.0s warm; 100% bit-identical video output vs unpadded reference. +# vae_decode_chunk: 1 -> 0.8s, -1 (unchunked) -> 3.0s, so keep chunking on + +# --------------------------------------------------------------------------- +# Platform profile: attention recipe, XLA flags, VMEM budget, text encoder +# --------------------------------------------------------------------------- +# Each arm carries the best measured configuration for that generation. Every +# value is overridable via the matching env var (see the header). +case "$TPU_PROFILE" in + v6e) + PLATFORM_LIBTPU="$V6E_LIBTPU" + DEFAULT_ATTENTION="ulysses_custom_fixed_m_per_q_block" + DEFAULT_U=4 + DEFAULT_BQ=9472 + DEFAULT_BKV=1024 + DEFAULT_BKV_COMPUTE=512 + DEFAULT_BKV_COMPUTE_IN=512 + # v6e carries 128MiB of VMEM per core; the larger budget lets the 9472-wide + # Q tile stay resident. Measured identical to the 64MiB budget on this + # shape, so lower it if a different resolution spills. + DEFAULT_VMEM=127506841 + DEFAULT_BQ_DKV=$DEFAULT_BQ + DEFAULT_COMPILE_TE=true + DEFAULT_BATCHED_TE=true + DEFAULT_VAE_CHUNK=1 + DEFAULT_VAE_SPATIAL=8 + ;; + v7) + PLATFORM_LIBTPU="$V7_LIBTPU" + DEFAULT_ATTENTION="ulysses_ring_custom_fixed_m" + DEFAULT_U=2 + DEFAULT_BQ=6400 + DEFAULT_BKV=2048 + DEFAULT_BKV_COMPUTE=2048 + DEFAULT_BKV_COMPUTE_IN=2048 + DEFAULT_VMEM=67108864 + DEFAULT_BQ_DKV=$DEFAULT_BQ + DEFAULT_COMPILE_TE=true + DEFAULT_BATCHED_TE=true + DEFAULT_VAE_CHUNK=1 + DEFAULT_VAE_SPATIAL=8 + ;; + *) + echo "== warning: unrecognised accelerator '${ACCEL_TYPE:-unknown}';" \ + "using the generic profile. Set TPU_PROFILE=v6e|v7 to override." >&2 + PLATFORM_LIBTPU="" + DEFAULT_ATTENTION="ulysses_custom_fixed_m_per_q_block" + DEFAULT_U=4 + DEFAULT_BQ=9472 + DEFAULT_BKV=1024 + DEFAULT_BKV_COMPUTE=512 + DEFAULT_BKV_COMPUTE_IN=512 + DEFAULT_VMEM=67108864 + DEFAULT_BQ_DKV=$DEFAULT_BQ + DEFAULT_COMPILE_TE=true + DEFAULT_BATCHED_TE=true + DEFAULT_VAE_CHUNK=1 + DEFAULT_VAE_SPATIAL=8 + ;; +esac + +export LIBTPU_INIT_ARGS="${COMMON_LIBTPU} ${PLATFORM_LIBTPU} ${EXTRA_LIBTPU:-}" +# A literal backslash truncates libtpu's flag parsing; fail loudly rather than +# running with silently-dropped flags. +case "$LIBTPU_INIT_ARGS" in + *\\*) + echo "ERROR: LIBTPU_INIT_ARGS contains a literal backslash; libtpu would" \ + "stop parsing there and drop the remaining flags." >&2 + exit 1 + ;; +esac + +ATTENTION=${ATTENTION:-$DEFAULT_ATTENTION} +ULYSSES_SHARDS=${ULYSSES_SHARDS:-$DEFAULT_U} +BQ=${BQ:-$DEFAULT_BQ} +BKV=${BKV:-$DEFAULT_BKV} +BKV_COMPUTE=${BKV_COMPUTE:-$DEFAULT_BKV_COMPUTE} +BKV_COMPUTE_IN=${BKV_COMPUTE_IN:-$DEFAULT_BKV_COMPUTE_IN} +BQ_DKV=${BQ_DKV:-$DEFAULT_BQ_DKV} +VMEM_LIMIT_BYTES=${VMEM_LIMIT_BYTES:-$DEFAULT_VMEM} +COMPILE_TE=${COMPILE_TE:-$DEFAULT_COMPILE_TE} +USE_BATCHED_TE=${USE_BATCHED_TE:-$DEFAULT_BATCHED_TE} +VAE_SPATIAL=${VAE_SPATIAL:-$DEFAULT_VAE_SPATIAL} +VAE_DECODE_CHUNK=${VAE_DECODE_CHUNK:-$DEFAULT_VAE_CHUNK} + +# Mesh: context parallelism carries the Ulysses shards, data parallelism takes +# whatever chips remain. Defaults to CP=4 / DP=2 on an 8-chip slice. On a slice +# smaller than CP, clamp rather than emit a mesh larger than the hardware. +CP=${CP:-4} +if [ -n "$TPU_CHIPS" ] && [ "$TPU_CHIPS" -lt "$CP" ]; then + echo "== note: $TPU_CHIPS-chip slice is smaller than CP=$CP; clamping CP to $TPU_CHIPS" >&2 + CP=$TPU_CHIPS +fi +if [ "$ULYSSES_SHARDS" -gt "$CP" ]; then + echo "== note: clamping ulysses_shards $ULYSSES_SHARDS -> $CP (must divide CP)" >&2 + ULYSSES_SHARDS=$CP +fi +if [ -z "${DP:-}" ]; then + if [ -n "$TPU_CHIPS" ] && [ "$TPU_CHIPS" -ge "$CP" ]; then + DP=$((TPU_CHIPS / CP)) + else + DP=2 + fi +fi +NUM_CHIPS=$((DP * CP)) +# One global video per step: per-device batch is 1/num_chips. +PER_DEVICE_BATCH=${PER_DEVICE_BATCH:-$(awk -v c="$NUM_CHIPS" 'BEGIN { printf "%.6g", 1.0 / c }')} if [ "$MODEL" = "21" ]; then CONFIG=src/maxdiffusion/configs/base_wan_14b.yml - GUIDANCE_ARGS="" + GUIDANCE_ARGS=() else CONFIG=src/maxdiffusion/configs/base_wan_27b.yml - GUIDANCE_ARGS="guidance_scale_low=3.0 guidance_scale_high=4.0" + GUIDANCE_ARGS=(guidance_scale_low=3.0 guidance_scale_high=4.0) fi PROMPT_ARG=() [ -n "$PROMPT" ] && PROMPT_ARG=("prompt=$PROMPT") RUN_NAME="wan${MODEL}_fast_$(date +%m%d-%H%M%S)" -echo "== ${ATTENTION} | tile ${BQ}/${BKV} | ${STEPS} steps" +echo "== platform ${ACCEL_TYPE:-unknown} -> profile ${TPU_PROFILE} | mesh DP=${DP} CP=${CP}" +echo "== ${ATTENTION} | U=${ULYSSES_SHARDS} | tile ${BQ}/${BKV} (compute=${BKV_COMPUTE}, in=${BKV_COMPUTE_IN}) | ${STEPS} steps" + +FLASH_BLOCK_SIZES="{\"block_q\":$BQ,\"block_kv\":$BKV,\"block_kv_compute\":$BKV_COMPUTE,\"block_kv_compute_in\":$BKV_COMPUTE_IN,\"heads_per_tile\":1,\"vmem_limit_bytes\":$VMEM_LIMIT_BYTES,\"block_q_dkv\":$BQ_DKV,\"block_kv_dkv\":$BKV,\"block_kv_dkv_compute\":$BKV,\"block_q_dq\":$BQ_DKV,\"block_kv_dq\":$BKV}" # libtpu's XLA:CPU AOT feature-mismatch log is cosmetic and ignores every # log-level env var; filter just that message from stderr. @@ -83,25 +286,26 @@ python src/maxdiffusion/generate_wan.py "$CONFIG" \ jax_cache_dir="$CACHE_ROOT/jax" \ aot_cache_dir="$CACHE_ROOT/aot_wan$MODEL" \ converted_weights_dir="$CACHE_ROOT/converted" \ - attention=$ATTENTION \ - ulysses_shards=2 \ - ici_data_parallelism=2 ici_fsdp_parallelism=1 \ - ici_context_parallelism=4 ici_tensor_parallelism=1 \ - per_device_batch_size=0.125 \ + attention="$ATTENTION" \ + ulysses_shards="$ULYSSES_SHARDS" \ + ici_data_parallelism="$DP" ici_fsdp_parallelism=1 \ + ici_context_parallelism="$CP" ici_tensor_parallelism=1 \ + per_device_batch_size="$PER_DEVICE_BATCH" \ num_inference_steps="$STEPS" num_frames=81 width=1280 height=720 \ weights_dtype=bfloat16 activations_dtype=bfloat16 \ - vae_spatial=4 vae_decode_chunk=-1 \ + vae_spatial="$VAE_SPATIAL" vae_decode_chunk="$VAE_DECODE_CHUNK" \ vae_weights_dtype=bfloat16 vae_dtype=bfloat16 \ - text_encoder_dtype=bfloat16 compile_text_encoder="${COMPILE_TE:-false}" use_batched_text_encoder=false \ - use_base2_exp=true use_experimental_scheduler=true \ - fps=16 $GUIDANCE_ARGS \ - flash_block_sizes="{\"block_q\":$BQ,\"block_kv\":$BKV,\"block_kv_compute\":$BKV,\"block_kv_compute_in\":1024,\"heads_per_tile\":1,\"vmem_limit_bytes\":67108864,\"block_q_dkv\":$BQ,\"block_kv_dkv\":$BKV,\"block_kv_dkv_compute\":$BKV}" \ + text_encoder_dtype=bfloat16 compile_text_encoder="$COMPILE_TE" use_batched_text_encoder="$USE_BATCHED_TE" \ + use_kv_cache=true use_base2_exp=true use_experimental_scheduler=true \ + fps=16 "${GUIDANCE_ARGS[@]}" \ + seed="${SEED:-12345}" \ + flash_block_sizes="$FLASH_BLOCK_SIZES" \ "${PROMPT_ARG[@]}" \ + "$@" \ 2> >(grep -vE --line-buffered 'cpu_aot_loader|machine type for execution' >&2) -mp4=$(ls -t wan_output_*.mp4 2>/dev/null | head -1) +mp4=$(ls -t "$OUTPUT_DIR"/${RUN_NAME}*.mp4 "$OUTPUT_DIR"/wan_output_*.mp4 2> /dev/null | head -1) if [ -n "$mp4" ]; then - mv "$mp4" "$OUTPUT_DIR/${RUN_NAME}.mp4" echo "" - echo "=== video saved: $OUTPUT_DIR/${RUN_NAME}.mp4 ===" + echo "=== video saved: $mp4 ===" fi diff --git a/src/maxdiffusion/aot_cache.py b/src/maxdiffusion/aot_cache.py index c07af3e24..12b695526 100644 --- a/src/maxdiffusion/aot_cache.py +++ b/src/maxdiffusion/aot_cache.py @@ -79,27 +79,53 @@ def _metadata_fingerprint(meta: dict[str, Any]) -> str: return hashlib.sha256(serialized.encode("utf-8")).hexdigest()[:12] +def _format_static_val(val: Any) -> str: + if isinstance(val, dict): + items = [f"{k!r}:{_format_static_val(val[k])}" for k in sorted(val.keys(), key=str)] + return "{" + ",".join(items) + "}" + return re.sub(r"0x[0-9a-fA-F]+", "@", repr(val)) + + +def _extract_graphdef_statics(obj: Any, prefix: str = "") -> list[str]: + """Extracts deterministic static attributes from nnx.GraphDef objects.""" + out = [] + if hasattr(obj, "nodes") and hasattr(obj, "attributes"): + for i, attr_item in enumerate(getattr(obj, "attributes", ())): + if isinstance(attr_item, tuple) and len(attr_item) == 2: + k, v = attr_item + if isinstance(k, str) and k.startswith("_pytree"): + continue + if hasattr(v, "value"): + desc = _format_static_val(v.value) + out.append(f"{prefix}:attr[{i}].{k}={desc}") + elif isinstance(obj, (tuple, list)): + for i, item in enumerate(obj): + out.extend(_extract_graphdef_statics(item, f"{prefix}[{i}]")) + elif isinstance(obj, dict): + for k in sorted(obj.keys(), key=str): + out.extend(_extract_graphdef_statics(obj[k], f"{prefix}.{k}")) + return out + + def _dynamic_signature(args: tuple, kwargs: dict) -> str: """Deterministic digest of everything that selects an executable. - Structure is captured by each leaf's KEY PATH (names, order, count) -- - NOT by ``repr(treedef)``: an nnx GraphDef's repr embeds object - addresses and hash-order-dependent content that differ per process and - made signatures never match across restarts (measured: every array - part stable, only the treedef part unstable). Static graph metadata - not visible in key paths (attention kernel, dtypes, model path) is - covered by the install-time config fingerprint in the filename. - Array leaves contribute shape/dtype; non-array leaves (python scalars, - None flags) contribute an address-stripped repr. + Structure is captured by each leaf's KEY PATH (names, order, count) plus + any static attributes inside ``nnx.GraphDef`` arguments (stripped of + object addresses and sorted by dict key so signatures remain stable across + processes). Array leaves contribute shape/dtype; non-array leaves (python + scalars, None flags) contribute an address-stripped repr. """ leaves_with_paths = jax.tree_util.tree_flatten_with_path((args, kwargs))[0] parts = [] for path, leaf in leaves_with_paths: if hasattr(leaf, "shape") and hasattr(leaf, "dtype"): - desc = f"{tuple(leaf.shape)}:{leaf.dtype}" + weak = getattr(leaf, "weak_type", False) + desc = f"{tuple(leaf.shape)}:{leaf.dtype}:weak={weak}" else: desc = re.sub(r"0x[0-9a-fA-F]+", "@", repr(leaf)) parts.append(f"{jax.tree_util.keystr(path)}={desc}") + parts.extend(_extract_graphdef_statics((args, kwargs))) return hashlib.sha256("|".join(parts).encode()).hexdigest()[:12] @@ -116,6 +142,7 @@ def __init__(self, name: str, fn: Callable, static_argnames: tuple): self._out_specs: dict[str, Any] = {} self._pending: dict[str, tuple] = {} self._adapters: dict[str, Any] = {} + self._sig_cache: dict[tuple, str] = {} self._on_disk: set[str] = set() self._lock = threading.Lock() @@ -197,7 +224,26 @@ def __call__(self, *args, **kwargs): # Under an outer trace a deserialized executable cannot be applied # and tracers must not be recorded -- inline like a nested jit. return self.jitted(**dynamic, **static) - signature = _dynamic_signature((), {**dynamic, **static}) + + # Fast-path signature cache: avoid tree_flatten_with_path + SHA256 string hashing on repeated steps + shapes_dtypes = tuple( + (leaf.shape, leaf.dtype, getattr(leaf, "weak_type", False)) + if hasattr(leaf, "shape") and hasattr(leaf, "dtype") + else (type(leaf), re.sub(r"0x[0-9a-fA-F]+", "@", repr(leaf))) + for leaf in leaves + ) + cache_key = (treedef, shapes_dtypes, tuple(sorted(static.items()))) + signature = getattr(self, "_sig_cache", {}).get(cache_key) + if signature is None: + signature = _dynamic_signature((), {**dynamic, **static}) + with self._lock: + if not hasattr(self, "_sig_cache"): + self._sig_cache = {} + if len(self._sig_cache) < 64: + new_cache = dict(self._sig_cache) + new_cache[cache_key] = signature + self._sig_cache = new_cache + if _STATE.warmup_only: # Compilation only needs avals; skip the (possibly seconds-long) # real execution and hand back correctly-shaped/sharded zeros so @@ -213,16 +259,13 @@ def __call__(self, *args, **kwargs): compiled = self._compiled.get(signature) if compiled is not None: flat = self._align_inputs(compiled, leaves) - if flat is not None: - return compiled(flat) - # Fewer expected shardings than leaves: XLA pruned unused inputs - # (e.g. encoder params in a decode-only executable). Compiled keeps - # the full in_tree and prunes internally, so hand it the raw leaves; - # sharding/structure problems surface as catchable Python errors. try: + if flat is not None: + return compiled(flat) return compiled(leaves) except Exception as e: # noqa: BLE001 - any failure means "use jit" - max_logging.log(f"[aot] {self.name}: compiled call failed ({e}); using jit") + max_logging.log(f"[aot] fast-path execution failed for {self.name}/{signature} ({e}); falling back to JIT") + return self.jitted(**dynamic, **static) with self._lock: if signature not in self._pending and signature not in self._compiled: self._pending[signature] = (leaves, treedef, static) @@ -234,29 +277,32 @@ def _align_inputs(self, compiled: Any, leaves: list): jit auto-commits mismatched inputs; a deserialized Compiled does not -- a placement mismatch aborts inside PjRt (uncatchable C++). Weights already carry final shardings; in practice this only moves small - fresh-off-host activations. Returns the aligned leaf list, or None on - structural mismatch (caller falls back to jit). + fresh-off-host activations. Returns the aligned leaf list, or None when + XLA pruned unused inputs (caller retries via Compiled's own pruning path). + Raises loudly on sharding alignment errors rather than silently triggering + a mid-serving JIT recompile. """ - try: - flat_expected = jax.tree_util.tree_leaves(compiled.input_shardings) - if len(flat_expected) != len(leaves): - # Fewer expected shardings than leaves = XLA pruned unused inputs; - # the caller retries via Compiled's own pruning path. Not an error. - return None - aligned = [] - for leaf, expected in zip(leaves, flat_expected): - if not hasattr(leaf, "shape"): # python scalar traced as weak array - leaf = jnp.asarray(leaf) - sharding = getattr(leaf, "sharding", None) - if sharding is not None and sharding.is_equivalent_to(expected, leaf.ndim): - aligned.append(leaf) - else: - aligned.append(jax.device_put(leaf, expected)) - return aligned - except Exception as e: # noqa: BLE001 - any failure means "use jit" - max_logging.log(f"[aot] {self.name}: cannot align inputs ({e}); using jit") + flat_expected = jax.tree_util.tree_leaves(compiled.input_shardings) + if len(flat_expected) != len(leaves): + # Fewer expected shardings than leaves = XLA pruned unused inputs; + # the caller retries via Compiled's own pruning path. Not an error. return None + # Fast path: check if all leaves already match expected sharding exactly + if all(getattr(leaf, "sharding", None) is expected for leaf, expected in zip(leaves, flat_expected)): + return leaves + + aligned = [] + for leaf, expected in zip(leaves, flat_expected): + if not hasattr(leaf, "shape"): # python scalar traced as weak array + leaf = jnp.asarray(leaf) + sharding = getattr(leaf, "sharding", None) + if sharding is not None and (sharding is expected or sharding.is_equivalent_to(expected, leaf.ndim)): + aligned.append(leaf) + else: + aligned.append(jax.device_put(leaf, expected)) + return aligned + # ---------------------------------------------------------------- disk def _path_for(self, signature: str) -> str: return os.path.join(_STATE.cache_dir, f"{self.name}-{_STATE.fingerprint}-{signature}.aotx") diff --git a/src/maxdiffusion/configs/base_wan_14b.yml b/src/maxdiffusion/configs/base_wan_14b.yml index 837bbe98b..536f22cf8 100644 --- a/src/maxdiffusion/configs/base_wan_14b.yml +++ b/src/maxdiffusion/configs/base_wan_14b.yml @@ -83,9 +83,10 @@ jit_initializers: True # Set true to load weights from pytorch from_pt: True split_head_dim: True -attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block use_base2_exp: True use_experimental_scheduler: True +use_k_centering: False # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. ulysses_shards: -1 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder. diff --git a/src/maxdiffusion/configs/base_wan_animate.yml b/src/maxdiffusion/configs/base_wan_animate.yml index 5e9df7d0d..48b3747f6 100644 --- a/src/maxdiffusion/configs/base_wan_animate.yml +++ b/src/maxdiffusion/configs/base_wan_animate.yml @@ -81,7 +81,7 @@ jit_initializers: True # Set true to load weights from pytorch from_pt: True split_head_dim: True -attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block use_base2_exp: True use_experimental_scheduler: True # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. diff --git a/src/maxdiffusion/configs/base_wan_i2v_14b.yml b/src/maxdiffusion/configs/base_wan_i2v_14b.yml index a129ff66c..dfdb691af 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_14b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_14b.yml @@ -83,7 +83,7 @@ jit_initializers: True # Set true to load weights from pytorch from_pt: True split_head_dim: True -attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block use_base2_exp: True use_experimental_scheduler: True # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. diff --git a/src/maxdiffusion/configs/base_wan_i2v_27b.yml b/src/maxdiffusion/configs/base_wan_i2v_27b.yml index 6a28986fc..ba4ca546e 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_27b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_27b.yml @@ -83,7 +83,7 @@ jit_initializers: True # Set true to load weights from pytorch from_pt: True split_head_dim: True -attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block use_base2_exp: True use_experimental_scheduler: True # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. diff --git a/src/maxdiffusion/generate_wan.py b/src/maxdiffusion/generate_wan.py index 662069ffd..092420771 100644 --- a/src/maxdiffusion/generate_wan.py +++ b/src/maxdiffusion/generate_wan.py @@ -16,6 +16,7 @@ import jax import time import os +import uuid from maxdiffusion.checkpointing.wan_checkpointer_2_1 import WanCheckpointer2_1 from maxdiffusion.checkpointing.wan_checkpointer_2_2 import WanCheckpointer2_2 from maxdiffusion.checkpointing.wan_checkpointer_i2v_2p1 import WanCheckpointerI2V_2_1 @@ -37,6 +38,98 @@ jax.config.update("jax_use_shardy_partitioner", True) +import hashlib + + +def _non_reusable_aot_revision(): + """Returns a unique identity so unversioned/dirty development source can never hit old HLO.""" + return f"unversioned:{uuid.uuid4().hex}" + + +def _compute_wan_source_hash() -> str | None: + """Computes a deterministic SHA-256 content hash of all non-test package source files.""" + try: + pkg_dir = os.path.dirname(os.path.abspath(__file__)) + hasher = hashlib.sha256() + py_files = [] + for root, dirs, files in os.walk(pkg_dir): + dirs[:] = [d for d in dirs if d not in ("tests", "__pycache__")] + for f in files: + if f.endswith(".py"): + py_files.append(os.path.join(root, f)) + for path in sorted(set(py_files)): + rel = os.path.relpath(path, pkg_dir).replace(os.sep, "/") + hasher.update(rel.encode("utf-8")) + with open(path, "rb") as f: + hasher.update(f.read()) + return f"src:{hasher.hexdigest()[:16]}" + except Exception: # noqa: BLE001 + return None + + +def _resolve_wan_aot_source_revision(config, commit_hash=None): + """Prefers explicit aot_build_revision, then combines Git commit (if supplied) with package source hash.""" + explicit = getattr(config, "aot_build_revision", None) + if explicit is not None and str(explicit).strip(): + return str(explicit).strip() + src_hash = _compute_wan_source_hash() + clean_commit = str(commit_hash).strip() if commit_hash is not None and str(commit_hash).strip() else None + if clean_commit and src_hash: + return f"{clean_commit}:{src_hash}" + if src_hash is not None: + return src_hash + if clean_commit is not None: + return clean_commit + return None + + +def _is_reusable_aot_revision(source_revision) -> bool: + if source_revision is None or not str(source_revision).strip(): + return False + s = str(source_revision).strip() + if s.startswith(("dirty:", "unversioned:")) or s.endswith("-dirty") or "-dirty:" in s: + return False + return True + + +def _build_wan_aot_metadata(config, mesh, source_revision) -> dict[str, str]: + """Builds the install-time configuration metadata dictionary for Wan AOT caching.""" + return { + "model": str(getattr(config, "pretrained_model_name_or_path", "")), + "attention": str(getattr(config, "attention", "")), + # Kernel block sizes change the lowered graph, not the input + # shapes — they must key the executable or a re-tuned config + # would silently hit stale binaries. + "flash_block_sizes": str(getattr(config, "flash_block_sizes", {})), + "mesh_shape": str(mesh.shape if mesh is not None else ()), + "vae_spatial": str(getattr(config, "vae_spatial", 8)), + "vae_decode_chunk": str(getattr(config, "vae_decode_chunk", 1)), + "vae_weights_dtype": str(getattr(config, "vae_weights_dtype", "bfloat16")), + "vae_dtype": str(getattr(config, "vae_dtype", "bfloat16")), + "weights_dtype": str(getattr(config, "weights_dtype", "")), + "activations_dtype": str(getattr(config, "activations_dtype", "")), + "scan_layers": str(getattr(config, "scan_layers", True)), + "remat_policy": str(getattr(config, "remat_policy", "NONE")), + "ulysses_shards": str(getattr(config, "ulysses_shards", 1)), + "ulysses_attention_chunks": str(getattr(config, "ulysses_attention_chunks", 1)), + "use_k_centering": str(getattr(config, "use_k_centering", False)), + "use_kv_cache": str(getattr(config, "use_kv_cache", False)), + "use_cfg_cache": str(getattr(config, "use_cfg_cache", False)), + "use_magcache": str(getattr(config, "use_magcache", False)), + "use_sen_cache": str(getattr(config, "use_sen_cache", False)), + "flash_min_seq_length": str(getattr(config, "flash_min_seq_length", 4096)), + "mask_padding_tokens": str(getattr(config, "mask_padding_tokens", True)), + "precision": str(getattr(config, "precision", "default")), + "logical_axis_rules": str(getattr(config, "logical_axis_rules", ())), + "device_kind": str(jax.devices()[0].device_kind if jax.devices() else "unknown"), + "process_count": str(jax.process_count()), + "use_base2_exp": str(getattr(config, "use_base2_exp", True)), + "use_experimental_scheduler": str(getattr(config, "use_experimental_scheduler", False)), + "jax": jax.__version__, + "source_revision": source_revision if source_revision else _non_reusable_aot_revision(), + } + + def call_pipeline(config, pipeline, prompt, negative_prompt, num_inference_steps=None): model_key = config.model_name model_type = config.model_type @@ -298,23 +391,20 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): # Per-shape AOT executable cache: deserialization starts on background # threads now and overlaps the remaining setup; unknown shapes silently # fall back to jit and are serialized by save_pending() after warmup. + detected_revision = commit_hash if commit_hash is not None else max_utils.get_git_commit_hash(check_dirty=True) + source_revision = _resolve_wan_aot_source_revision(config, detected_revision) + aot_cache_dir = getattr(config, "aot_cache_dir", "") + if aot_cache_dir and not _is_reusable_aot_revision(source_revision): + max_logging.log( + "[aot] No clean Git commit or aot_build_revision was supplied; " + "persistent Wan AOT caching is disabled for this development run." + ) + aot_cache_dir = "" + + aot_metadata = _build_wan_aot_metadata(config, pipeline.mesh, source_revision) aot_cache.install( - getattr(config, "aot_cache_dir", ""), - meta={ - "model": config.pretrained_model_name_or_path, - "attention": config.attention, - # Kernel block sizes change the lowered graph, not the input - # shapes — they must key the executable or a re-tuned config - # would silently hit stale binaries. - "flash_block_sizes": str(config.flash_block_sizes), - "mesh_shape": str(pipeline.mesh.shape), - "vae_spatial": str(config.vae_spatial), - "vae_decode_chunk": str(config.vae_decode_chunk), - "weights_dtype": str(config.weights_dtype), - "activations_dtype": str(config.activations_dtype), - "scan_layers": str(config.scan_layers), - "jax": jax.__version__, - }, + aot_cache_dir, + meta=aot_metadata, mesh=pipeline.mesh, ) # Deserialization is seconds and warmup must see the loaded executables @@ -352,7 +442,13 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): # the warmup pays compile time only, never real denoise compute. The # returned videos are garbage by design and are discarded below. with aot_cache.warmup_mode(): - videos = call_pipeline(config, pipeline, warmup_prompt, warmup_negative_prompt, num_inference_steps=warmup_steps) + videos = call_pipeline( + config, + pipeline, + warmup_prompt, + warmup_negative_prompt, + num_inference_steps=warmup_steps, + ) if isinstance(videos, tuple): videos, warmup_trace = videos warmup_str = ", ".join(f"{stage}={seconds:.1f}s" for stage, seconds in warmup_trace.items()) @@ -391,12 +487,24 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): else: videos = outputs trace = {} - for i in range(len(videos)): - video_path = f"{filename_prefix}wan_output_{config.seed}_{i}.mp4" - export_to_video(videos[i], video_path, fps=config.fps) - saved_video_path.append(video_path) - if gcs_output_path: - max_utils.upload_file_to_gcs(gcs_output_path, video_path, subdir="videos") + if jax.process_index() == 0: + import numpy as np + + for i in range(len(videos)): + if getattr(config, "output_dir", "") and not config.output_dir.startswith("gs://"): + os.makedirs(config.output_dir, exist_ok=True) + video_path = os.path.join(config.output_dir, f"{config.run_name}_{config.seed}_{i}.mp4") + else: + video_path = f"{filename_prefix}wan_output_{config.seed}_{i}.mp4" + frames_np = np.asarray(videos[i]) + export_to_video(frames_np, video_path, fps=config.fps) + saved_video_path.append(video_path) + max_logging.log(f"Saved video to {video_path}") + if gcs_output_path: + max_utils.upload_file_to_gcs(gcs_output_path, video_path, subdir="videos") + elif getattr(config, "output_dir", "").startswith("gs://"): + gcs_path = f"{config.output_dir.rstrip('/')}/{config.run_name}" + max_utils.upload_file_to_gcs(gcs_path, video_path, subdir="videos") else: trace = {} for i, padded_chunk, actual_chunk_len in max_utils.chunk_and_pad(prompts, batch_size): @@ -407,13 +515,28 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): videos, trace = outputs else: videos = outputs - for j in range(actual_chunk_len): - prompt_idx = i + j - video_path = f"{filename_prefix}wan_output_{config.seed}_{prompt_idx}.mp4" - export_to_video(videos[j], video_path, fps=config.fps) - saved_video_path.append(video_path) - if gcs_output_path: - max_utils.upload_file_to_gcs(gcs_output_path, video_path, subdir="videos") + if jax.process_index() == 0: + import numpy as np + + for j in range(actual_chunk_len): + prompt_idx = i + j + if getattr(config, "output_dir", "") and not config.output_dir.startswith("gs://"): + os.makedirs(config.output_dir, exist_ok=True) + video_path = os.path.join( + config.output_dir, + f"{config.run_name}_{config.seed}_{prompt_idx}.mp4", + ) + else: + video_path = f"{filename_prefix}wan_output_{config.seed}_{prompt_idx}.mp4" + frames_np = np.asarray(videos[j]) + export_to_video(frames_np, video_path, fps=config.fps) + saved_video_path.append(video_path) + max_logging.log(f"Saved video to {video_path}") + if gcs_output_path: + max_utils.upload_file_to_gcs(gcs_output_path, video_path, subdir="videos") + elif getattr(config, "output_dir", "").startswith("gs://"): + gcs_path = f"{config.output_dir.rstrip('/')}/{config.run_name}" + max_utils.upload_file_to_gcs(gcs_path, video_path, subdir="videos") generation_time = time.perf_counter() - s0 max_logging.log(f"generation_time: {generation_time}") @@ -482,7 +605,7 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): def main(argv: Sequence[str]) -> None: - commit_hash = max_utils.get_git_commit_hash() + commit_hash = max_utils.get_git_commit_hash(check_dirty=True) pyconfig.initialize(argv) try: flax.config.update("flax_always_shard_variable", False) diff --git a/src/maxdiffusion/max_utils.py b/src/maxdiffusion/max_utils.py index 4471e3bd1..99d0f95ec 100644 --- a/src/maxdiffusion/max_utils.py +++ b/src/maxdiffusion/max_utils.py @@ -468,10 +468,27 @@ def delete_file(file_path: str): max_logging.log(f"The file '{file_path}' does not exist.") -def get_git_commit_hash(): +def get_git_commit_hash(check_dirty: bool = True): """Tries to get the current Git commit hash, for run provenance.""" try: - return subprocess.check_output(["git", "rev-parse", "HEAD"]).strip().decode("utf-8") + repo_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + commit = subprocess.check_output(["git", "-C", repo_dir, "rev-parse", "HEAD"]).strip().decode("utf-8") + if check_dirty: + status = ( + subprocess.check_output([ + "git", + "-C", + repo_dir, + "status", + "--porcelain", + "--untracked-files=no", + ]) + .strip() + .decode("utf-8") + ) + if status: + return f"{commit}-dirty" + return commit except subprocess.CalledProcessError: max_logging.log("Warning: 'git rev-parse HEAD' failed. Not running in a git repo?") return None @@ -582,11 +599,16 @@ def create_device_mesh(config, devices=None, logging=True): if multi_slice_env: dcn_parallelism = fill_unspecified_mesh_axes(dcn_parallelism, num_slices, "DCN") mesh = mesh_utils.create_hybrid_device_mesh( - ici_parallelism, dcn_parallelism, devices, allow_split_physical_axes=config.allow_split_physical_axes + ici_parallelism, + dcn_parallelism, + devices, + allow_split_physical_axes=config.allow_split_physical_axes, ) else: mesh = mesh_utils.create_device_mesh( - ici_parallelism, devices, allow_split_physical_axes=config.allow_split_physical_axes + ici_parallelism, + devices, + allow_split_physical_axes=config.allow_split_physical_axes, ) if logging: diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index 67ac0648b..d7e4f46ae 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -1613,7 +1613,7 @@ def _ring_fixed_m_norms_pre_a2a( q_norm_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1) qn_head_local = q_norm_sq.max(axis=-1) - vn_local = (value.astype(jnp.float32) ** 2).max(axis=(2, 3)) + vn_local = (value.astype(jnp.float32) ** 2).max() if use_k_centering: # Optional K-centering: computes global mean and subtracts before a2a. @@ -1666,7 +1666,7 @@ def _ring_fixed_m_norms_pre_a2a( qn_dev = _slice_own_ulysses_heads(qn_all, ulysses_axis, num_ulysses_shards, axis=1) qn_head_global_dev = _slice_own_ulysses_heads(qn_head_global, ulysses_axis, num_ulysses_shards, axis=1) mk_all_sq_dev = _slice_own_ulysses_heads(mk_all_sq, ulysses_axis, num_ulysses_shards, axis=2) - vn_dev = _slice_own_ulysses_heads(vn_global, ulysses_axis, num_ulysses_shards, axis=1) + vn_dev = vn_global num_q_heads_dev = qn_dev.shape[1] num_kv_heads_dev = mk_all_sq_dev.shape[2] diff --git a/src/maxdiffusion/models/wan/transformers/transformer_wan.py b/src/maxdiffusion/models/wan/transformers/transformer_wan.py index 4cdfd0ca1..86761776c 100644 --- a/src/maxdiffusion/models/wan/transformers/transformer_wan.py +++ b/src/maxdiffusion/models/wan/transformers/transformer_wan.py @@ -233,6 +233,7 @@ def __init__( weights_dtype: jnp.dtype = jnp.float32, precision: jax.lax.Precision = None, ): + self.precision = precision self.proj = nnx.Linear( rngs=rngs, in_features=dim_in, @@ -253,8 +254,7 @@ def __init__( def __call__(self, x: jax.Array) -> jax.Array: with jax.named_scope("gelu"): - x = self.proj(x) - return nnx.gelu(x) + return nnx.gelu(self.proj(x)) class WanFeedForward(nnx.Module): @@ -280,6 +280,7 @@ def __init__( dim_out = dim_out if dim_out is not None else dim self.enable_jax_named_scopes = enable_jax_named_scopes + self.precision = precision self.act_fn = nnx.data(None) if activation_fn == "gelu-approximate": self.act_fn = ApproximateGELU( @@ -322,12 +323,13 @@ def __call__( deterministic: bool = True, rngs: nnx.Rngs = None, ) -> jax.Array: - hidden_states = self.act_fn(hidden_states) # Output is (4, 75600, 13824) + hidden_states = self.act_fn(hidden_states) hidden_states = checkpoint_name(hidden_states, "ffn_activation") if self.drop_out.rate > 0: hidden_states = self.drop_out(hidden_states, deterministic=deterministic, rngs=rngs) with jax.named_scope("proj_out"): - return self.proj_out(hidden_states) # output is (4, 75600, 5120) + hidden_states = self.proj_out(hidden_states) + return hidden_states class WanTransformerBlock(nnx.Module): @@ -361,6 +363,7 @@ def __init__( "use_experimental_scheduler": False, "ulysses_shards": -1, "ulysses_attention_chunks": 1, + "use_k_centering": False, **(attention_config or {}), } @@ -490,7 +493,6 @@ def __call__( with self.conditional_named_scope("self_attn_attn"): attn_output = self.attn1( hidden_states=norm_hidden_states, - encoder_hidden_states=norm_hidden_states, rotary_emb=rotary_emb, deterministic=deterministic, rngs=rngs, @@ -586,6 +588,7 @@ def __init__( "use_experimental_scheduler": False, "ulysses_shards": -1, "ulysses_attention_chunks": 1, + "use_k_centering": False, **(attention_config or {}), } @@ -934,22 +937,37 @@ def layer_forward(hidden_states, l_kv): scale = scale.squeeze(2) # [B, sl, dim] else: shift, scale = jnp.split(self.scale_shift_table + jnp.expand_dims(temb, axis=1), 2, axis=1) + hidden_states = (self.norm_out(hidden_states.astype(jnp.float32)) * (1 + scale) + shift).astype(hidden_states.dtype) with jax.named_scope("proj_out"): hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.reshape( - batch_size, - post_patch_num_frames, - post_patch_height, - post_patch_width, - p_t, - p_h, - p_w, - -1, - ) - hidden_states = jnp.transpose(hidden_states, (0, 7, 1, 4, 2, 5, 3, 6)) - hidden_states = hidden_states.reshape(batch_size, -1, num_frames, height, width) + if p_t == 1: + # Lossless HLO optimization: collapse p_t=1 dimension to avoid 8D non-contiguous stride copies + hidden_states = hidden_states.reshape( + batch_size, + post_patch_num_frames, + post_patch_height, + post_patch_width, + p_h, + p_w, + -1, + ) + hidden_states = jnp.transpose(hidden_states, (0, 6, 1, 2, 4, 3, 5)) + hidden_states = hidden_states.reshape(batch_size, -1, num_frames, height, width) + else: + hidden_states = hidden_states.reshape( + batch_size, + post_patch_num_frames, + post_patch_height, + post_patch_width, + p_t, + p_h, + p_w, + -1, + ) + hidden_states = jnp.transpose(hidden_states, (0, 7, 1, 4, 2, 5, 3, 6)) + hidden_states = hidden_states.reshape(batch_size, -1, num_frames, height, width) if return_residual: return hidden_states, residual_x diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline.py b/src/maxdiffusion/pipelines/wan/wan_pipeline.py index 6464ce6a6..5493fb350 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline.py @@ -347,6 +347,7 @@ def create_model(rngs: nnx.Rngs, wan_config: dict): "use_experimental_scheduler": config.use_experimental_scheduler, "ulysses_shards": getattr(config, "ulysses_shards", -1), "ulysses_attention_chunks": getattr(config, "ulysses_attention_chunks", 1), + "use_k_centering": getattr(config, "use_k_centering", False), } # 2. eval_shape - will not use flops or create weights on device diff --git a/src/maxdiffusion/pyconfig.py b/src/maxdiffusion/pyconfig.py index ddec32c78..2cdb085c6 100644 --- a/src/maxdiffusion/pyconfig.py +++ b/src/maxdiffusion/pyconfig.py @@ -76,7 +76,13 @@ def string_to_list(string_list: str) -> list: return ast.literal_eval(string_list) -_yaml_types_to_parser = {str: str, int: int, float: float, bool: string_to_bool, list: string_to_list} +_yaml_types_to_parser = { + str: str, + int: int, + float: float, + bool: string_to_bool, + list: string_to_list, +} _config = None config = None @@ -234,6 +240,7 @@ def user_init(raw_keys): "ulysses_ring_custom", "ulysses_ring_custom_fixed_m", "ulysses_ring_custom_bidir", + "ulysses_ring_custom_fixed_m_per_q_block", } if attention in ulysses_ring_attentions and raw_keys.get("ulysses_shards", -1) <= 0: raise ValueError(f"{attention} requires ulysses_shards to be set from config or command line.") @@ -313,6 +320,9 @@ def user_init(raw_keys): if "vae_spatial" not in raw_keys: raw_keys["vae_spatial"] = -1 + if "use_k_centering" not in raw_keys: + raw_keys["use_k_centering"] = False + def get_num_slices(raw_keys): if int(raw_keys["compile_topology_num_slices"]) > 0: diff --git a/src/maxdiffusion/tests/aot_cache_test.py b/src/maxdiffusion/tests/aot_cache_test.py index 916e80ef3..efd062f22 100644 --- a/src/maxdiffusion/tests/aot_cache_test.py +++ b/src/maxdiffusion/tests/aot_cache_test.py @@ -179,6 +179,76 @@ def test_signature_deterministic_across_processes(self): ] self.assertEqual(outs[0], outs[1]) + def test_graphdef_static_attribute_changes_dynamic_signature(self): + """Toggling static attributes on nnx.GraphDef (such as use_k_centering) changes _dynamic_signature.""" + from flax import nnx + + class DummyBlock(nnx.Module): + + def __init__(self, use_k_centering: bool): + self.attention_config = {"use_k_centering": use_k_centering} + + gd_on, _ = nnx.split(DummyBlock(use_k_centering=True)) + gd_off, _ = nnx.split(DummyBlock(use_k_centering=False)) + + sig_on = aot_cache._dynamic_signature((gd_on, jnp.ones((2, 4))), {}) + sig_off = aot_cache._dynamic_signature((gd_off, jnp.ones((2, 4))), {}) + self.assertNotEqual(sig_on, sig_off) + + def test_wan_aot_metadata_includes_use_k_centering(self): + """Toggling use_k_centering changes the Wan AOT metadata fingerprint.""" + import types + from maxdiffusion import generate_wan + + cfg_on = types.SimpleNamespace(use_k_centering=True, attention="ulysses_ring_custom_fixed_m") + cfg_off = types.SimpleNamespace(use_k_centering=False, attention="ulysses_ring_custom_fixed_m") + + meta_on = generate_wan._build_wan_aot_metadata(cfg_on, self._mesh, "rev1") + meta_off = generate_wan._build_wan_aot_metadata(cfg_off, self._mesh, "rev1") + + self.assertEqual(meta_on["use_k_centering"], "True") + self.assertEqual(meta_off["use_k_centering"], "False") + self.assertNotEqual( + aot_cache._metadata_fingerprint(meta_on), + aot_cache._metadata_fingerprint(meta_off), + ) + + def test_wan_source_hash_includes_shared_modules_and_commit(self): + """Verifies _compute_wan_source_hash hashes shared modules (normalization_flax, embeddings_flax) and combines commit_hash.""" + import types + from unittest import mock + from maxdiffusion import generate_wan + + base_hash = generate_wan._compute_wan_source_hash() + self.assertIsNotNone(base_hash) + self.assertTrue(base_hash.startswith("src:")) + + # Simulate modifying models/normalization_flax.py or models/embeddings_flax.py + orig_open = open + + def patched_open(path, *args, **kwargs): + f = orig_open(path, *args, **kwargs) + if str(path).endswith(("normalization_flax.py", "embeddings_flax.py")) and "rb" in args: + content = f.read() + f.close() + import io + + return io.BytesIO(content + b"\n# modified for fingerprint test\n") + return f + + with mock.patch("builtins.open", side_effect=patched_open): + mod_hash = generate_wan._compute_wan_source_hash() + + self.assertNotEqual(base_hash, mod_hash) + + # Verify commit_hash is combined with src_hash rather than shadowed + cfg = types.SimpleNamespace(aot_build_revision=None) + rev_a = generate_wan._resolve_wan_aot_source_revision(cfg, commit_hash="commit_aaa") + rev_b = generate_wan._resolve_wan_aot_source_revision(cfg, commit_hash="commit_bbb") + self.assertNotEqual(rev_a, rev_b) + self.assertIn("commit_aaa", rev_a) + self.assertIn("commit_bbb", rev_b) + if __name__ == "__main__": unittest.main() diff --git a/src/maxdiffusion/tests/wan/wan_transformer_test.py b/src/maxdiffusion/tests/wan/wan_transformer_test.py index 69bed9a6a..d8fe4c002 100644 --- a/src/maxdiffusion/tests/wan/wan_transformer_test.py +++ b/src/maxdiffusion/tests/wan/wan_transformer_test.py @@ -34,9 +34,11 @@ ) from maxdiffusion.models.embeddings_flax import NNXTimestepEmbedding, NNXPixArtAlphaTextProjection from maxdiffusion.models.normalization_flax import FP32LayerNorm -from maxdiffusion.models.attention_flax import FlaxWanAttention +from maxdiffusion.models.attention_flax import FlaxWanAttention, _unflatten_heads +from maxdiffusion.kernels.fused_producers import fused_rmsnorm_rope from maxdiffusion.pyconfig import HyperParameters from maxdiffusion.pipelines.wan.wan_pipeline import WanPipeline +import numpy as np import qwix import flax @@ -118,12 +120,20 @@ def test_wan_time_text_embedding(self): text_embed_dim = 4096 with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): layer = WanTimeTextImageEmbedding( - rngs=rngs, dim=dim, time_freq_dim=time_freq_dim, time_proj_dim=time_proj_dim, text_embed_dim=text_embed_dim + rngs=rngs, + dim=dim, + time_freq_dim=time_freq_dim, + time_proj_dim=time_proj_dim, + text_embed_dim=text_embed_dim, ) dummy_timestep = jnp.ones(batch_size) - encoder_hidden_states_shape = (batch_size, time_freq_dim * 2, text_embed_dim) + encoder_hidden_states_shape = ( + batch_size, + time_freq_dim * 2, + text_embed_dim, + ) dummy_encoder_hidden_states = jnp.ones(encoder_hidden_states_shape) temb, timestep_proj, encoder_hidden_states, _, _ = layer(dummy_timestep, dummy_encoder_hidden_states) assert temb.shape == (batch_size, dim) @@ -189,13 +199,22 @@ def test_wan_block(self): mesh=mesh, flash_block_sizes=flash_block_sizes, ) - dummy_output = wan_block(dummy_hidden_states, dummy_encoder_hidden_states, dummy_temb, dummy_rotary_emb) + dummy_output = wan_block( + dummy_hidden_states, + dummy_encoder_hidden_states, + dummy_temb, + dummy_rotary_emb, + ) assert dummy_output.shape == dummy_hidden_states.shape def test_wan_attention(self): for attention_kernel in ["flash", "tokamax_flash"]: pyconfig.initialize( - [None, os.path.join(THIS_DIR, "..", "..", "configs", "base_wan_14b.yml"), f"attention={attention_kernel}"], + [ + None, + os.path.join(THIS_DIR, "..", "..", "configs", "base_wan_14b.yml"), + f"attention={attention_kernel}", + ], unittest=True, ) config = pyconfig.config @@ -231,7 +250,9 @@ def test_wan_attention(self): dummy_hidden_states = jnp.ones(dummy_hidden_states_shape) dummy_encoder_hidden_states = jnp.ones(dummy_hidden_states_shape) dummy_output = attention( - hidden_states=dummy_hidden_states, encoder_hidden_states=dummy_encoder_hidden_states, rotary_emb=dummy_rotary_emb + hidden_states=dummy_hidden_states, + encoder_hidden_states=dummy_encoder_hidden_states, + rotary_emb=dummy_rotary_emb, ) assert dummy_output.shape == dummy_hidden_states_shape @@ -250,6 +271,285 @@ def test_wan_attention(self): except NotImplementedError: pass + def test_fused_rmsnorm_rope_parity(self): + """Verifies numerical parity of fused RMSNorm + RoPE against unfused reference with complex freqs_cis.""" + key = jax.random.PRNGKey(202) + k1, k2, k3, k4, k5, k6 = jax.random.split(key, 6) + B, S, D, H, DH = 1, 1024, 5120, 40, 128 + eps = 1e-6 + + raw_q = jax.random.normal(k1, (B, S, D), dtype=jnp.bfloat16) + raw_k = jax.random.normal(k2, (B, S, D), dtype=jnp.bfloat16) + q_scale = jax.random.normal(k3, (D,), dtype=jnp.bfloat16) + k_scale = jax.random.normal(k4, (D,), dtype=jnp.bfloat16) + + freqs_real = jax.random.normal(k5, (1, 1, S, DH // 2), dtype=jnp.float32) + freqs_imag = jax.random.normal(k6, (1, 1, S, DH // 2), dtype=jnp.float32) + freqs_cis = jax.lax.complex(freqs_real, freqs_imag) + + # Unfused reference: FP32 RMSNorm -> unflatten -> RoPE + def ref_norm(x, scale): + var = jnp.mean(jnp.square(x.astype(jnp.float32)), axis=-1, keepdims=True) + return (x.astype(jnp.float32) * jax.lax.rsqrt(var + eps) * scale.astype(jnp.float32)).astype(x.dtype) + + def ref_unflatten(x, heads): + b, s, d = x.shape + return x.reshape(b, s, heads, d // heads).transpose(0, 2, 1, 3) + + def ref_apply_rope(xq, xk, freqs_cis): + cos = jnp.real(freqs_cis).astype(xq.dtype) + sin = jnp.imag(freqs_cis).astype(xq.dtype) + xq_reshaped = xq.reshape(*xq.shape[:-1], -1, 2) + xk_reshaped = xk.reshape(*xk.shape[:-1], -1, 2) + xq_0, xq_1 = xq_reshaped[..., 0], xq_reshaped[..., 1] + xk_0, xk_1 = xk_reshaped[..., 0], xk_reshaped[..., 1] + xq_out_0 = xq_0 * cos - xq_1 * sin + xq_out_1 = xq_0 * sin + xq_1 * cos + xk_out_0 = xk_0 * cos - xk_1 * sin + xk_out_1 = xk_0 * sin + xk_1 * cos + xq_out = jnp.concatenate([xq_out_0[..., None], xq_out_1[..., None]], axis=-1).reshape(xq.shape) + xk_out = jnp.concatenate([xk_out_0[..., None], xk_out_1[..., None]], axis=-1).reshape(xk.shape) + return xq_out, xk_out + + q_ref, k_ref = ref_apply_rope( + ref_unflatten(ref_norm(raw_q, q_scale), H), + ref_unflatten(ref_norm(raw_k, k_scale), H), + freqs_cis, + ) + + # Fused producer + q_fused, k_fused = fused_rmsnorm_rope( + raw_q, + raw_k, + q_scale, + k_scale, + freqs_cis, + q_heads=H, + dim_head=DH, + eps=eps, + ) + + np.testing.assert_allclose( + np.array(q_ref, dtype=np.float32), + np.array(q_fused, dtype=np.float32), + atol=0.08, + rtol=1e-2, + ) + np.testing.assert_allclose( + np.array(k_ref, dtype=np.float32), + np.array(k_fused, dtype=np.float32), + atol=0.08, + rtol=1e-2, + ) + + def test_fused_rmsnorm_rope_with_wan_rotary_embed(self): + """Verifies numerical parity of fused RMSNorm + RoPE against unfused reference with real WanRotaryPosEmbed frequencies.""" + key = jax.random.PRNGKey(404) + k1, k2, k3, k4 = jax.random.split(key, 4) + B, S, D, H, DH = 1, 1024, 5120, 40, 128 + eps = 1e-6 + + # Generate real RoPE frequencies on the complex unit circle using WanRotaryPosEmbed + wan_rot_embed = WanRotaryPosEmbed(attention_head_dim=DH, patch_size=[1, 2, 2], max_seq_len=1024) + dummy_video = jnp.ones((B, 1, 64, 64, 16)) + freqs_cis = wan_rot_embed(dummy_video) # (1, 1, 1024, 64) + + raw_q = jax.random.normal(k1, (B, S, D), dtype=jnp.bfloat16) + raw_k = jax.random.normal(k2, (B, S, D), dtype=jnp.bfloat16) + q_scale = jax.random.normal(k3, (D,), dtype=jnp.bfloat16) + k_scale = jax.random.normal(k4, (D,), dtype=jnp.bfloat16) + + def ref_norm(x, scale): + var = jnp.mean(jnp.square(x.astype(jnp.float32)), axis=-1, keepdims=True) + return (x.astype(jnp.float32) * jax.lax.rsqrt(var + eps) * scale.astype(jnp.float32)).astype(x.dtype) + + def ref_unflatten(x, heads): + b, s, d = x.shape + return x.reshape(b, s, heads, d // heads).transpose(0, 2, 1, 3) + + def ref_apply_rope(xq, xk, freqs_cis): + cos = jnp.real(freqs_cis).astype(xq.dtype) + sin = jnp.imag(freqs_cis).astype(xq.dtype) + xq_reshaped = xq.reshape(*xq.shape[:-1], -1, 2) + xk_reshaped = xk.reshape(*xk.shape[:-1], -1, 2) + xq_0, xq_1 = xq_reshaped[..., 0], xq_reshaped[..., 1] + xk_0, xk_1 = xk_reshaped[..., 0], xk_reshaped[..., 1] + xq_out_0 = xq_0 * cos - xq_1 * sin + xq_out_1 = xq_0 * sin + xq_1 * cos + xk_out_0 = xk_0 * cos - xk_1 * sin + xk_out_1 = xk_0 * sin + xk_1 * cos + xq_out = jnp.concatenate([xq_out_0[..., None], xq_out_1[..., None]], axis=-1).reshape(xq.shape) + xk_out = jnp.concatenate([xk_out_0[..., None], xk_out_1[..., None]], axis=-1).reshape(xk.shape) + return xq_out, xk_out + + q_ref, k_ref = ref_apply_rope( + ref_unflatten(ref_norm(raw_q, q_scale), H), + ref_unflatten(ref_norm(raw_k, k_scale), H), + freqs_cis, + ) + + q_fused, k_fused = fused_rmsnorm_rope( + raw_q, + raw_k, + q_scale, + k_scale, + freqs_cis, + q_heads=H, + dim_head=DH, + eps=eps, + ) + + np.testing.assert_allclose( + np.array(q_ref, dtype=np.float32), + np.array(q_fused, dtype=np.float32), + atol=0.08, + rtol=1e-2, + ) + np.testing.assert_allclose( + np.array(k_ref, dtype=np.float32), + np.array(k_fused, dtype=np.float32), + atol=0.08, + rtol=1e-2, + ) + + def test_fused_rmsnorm_rope_gqa_parity(self): + """Verifies that fused_rmsnorm_rope correctly handles asymmetric GQA shapes (e.g. q_heads=8, kv_heads=2).""" + key = jax.random.PRNGKey(505) + k1, k2, k3, k4, k5 = jax.random.split(key, 5) + + B = 2 + S = 64 + Q_H = 8 + KV_H = 2 + DH = 128 + D_q = Q_H * DH + D_kv = KV_H * DH + eps = 1e-6 + + raw_q = jax.random.normal(k1, (B, S, D_q), dtype=jnp.bfloat16) + raw_k = jax.random.normal(k2, (B, S, D_kv), dtype=jnp.bfloat16) + q_scale = jax.random.normal(k3, (D_q,), dtype=jnp.float32) + k_scale = jax.random.normal(k4, (D_kv,), dtype=jnp.float32) + freqs_cis = jax.random.normal(k5, (1, 1, S, DH // 2), dtype=jnp.float32) + 1j * jax.random.normal( + key, (1, 1, S, DH // 2), dtype=jnp.float32 + ) + + def ref_norm(x, scale): + x_fp32 = x.astype(jnp.float32) + rms = jax.lax.rsqrt(jnp.mean(jnp.square(x_fp32), axis=-1, keepdims=True) + eps) + return (x_fp32 * rms * scale).astype(x.dtype) + + def ref_unflatten(x, heads): + return x.reshape(B, S, heads, DH).transpose(0, 2, 1, 3) + + def ref_apply_rope(xq, xk, freqs): + cos = jnp.real(freqs).astype(xq.dtype) + sin = jnp.imag(freqs).astype(xq.dtype) + xq_reshaped = xq.reshape(*xq.shape[:-1], -1, 2) + xk_reshaped = xk.reshape(*xk.shape[:-1], -1, 2) + xq_0, xq_1 = xq_reshaped[..., 0], xq_reshaped[..., 1] + xk_0, xk_1 = xk_reshaped[..., 0], xk_reshaped[..., 1] + xq_out_0 = xq_0 * cos - xq_1 * sin + xq_out_1 = xq_0 * sin + xq_1 * cos + xk_out_0 = xk_0 * cos - xk_1 * sin + xk_out_1 = xk_0 * sin + xk_1 * cos + xq_out = jnp.concatenate([xq_out_0[..., None], xq_out_1[..., None]], axis=-1).reshape(xq.shape) + xk_out = jnp.concatenate([xk_out_0[..., None], xk_out_1[..., None]], axis=-1).reshape(xk.shape) + return xq_out, xk_out + + q_ref, k_ref = ref_apply_rope( + ref_unflatten(ref_norm(raw_q, q_scale), Q_H), + ref_unflatten(ref_norm(raw_k, k_scale), KV_H), + freqs_cis, + ) + + q_fused, k_fused = fused_rmsnorm_rope( + raw_q, + raw_k, + q_scale, + k_scale, + freqs_cis, + q_heads=Q_H, + kv_heads=KV_H, + dim_head=DH, + eps=eps, + ) + + self.assertEqual(q_fused.shape, (B, Q_H, S, DH)) + self.assertEqual(k_fused.shape, (B, KV_H, S, DH)) + + np.testing.assert_allclose( + np.array(q_ref, dtype=np.float32), + np.array(q_fused, dtype=np.float32), + atol=0.08, + rtol=1e-2, + ) + np.testing.assert_allclose( + np.array(k_ref, dtype=np.float32), + np.array(k_fused, dtype=np.float32), + atol=0.08, + rtol=1e-2, + ) + + def test_wan_self_attention_is_self_attention_dispatch(self): + """Verifies that FlaxWanAttention with is_self_attention=True correctly dispatches fused RMSNorm+RoPE and matches reference.""" + key = jax.random.PRNGKey(303) + k1, k2, k3, k4 = jax.random.split(key, 4) + rngs = nnx.Rngs(k1) + + batch_size = 1 + seq_len = 1024 + query_dim = 5120 + heads = 40 + dim_head = 128 + + flash_block_sizes = get_flash_block_sizes(self.config) + with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): + attn = FlaxWanAttention( + rngs=rngs, + query_dim=query_dim, + heads=heads, + dim_head=dim_head, + attention_kernel="dot_product", + mesh=self.mesh, + flash_block_sizes=flash_block_sizes, + is_self_attention=True, + ) + self.assertTrue(attn.is_self_attention) + + hidden_states = jax.random.normal(k2, (batch_size, seq_len, query_dim), dtype=jnp.bfloat16) + freqs_real = jax.random.normal(k3, (1, 1, seq_len, dim_head // 2), dtype=jnp.float32) + freqs_imag = jax.random.normal(k4, (1, 1, seq_len, dim_head // 2), dtype=jnp.float32) + rotary_emb = jax.lax.complex(freqs_real, freqs_imag) + + # Wan self-attention call passes encoder_hidden_states=hidden_states + out = attn( + hidden_states=hidden_states, + encoder_hidden_states=hidden_states, + rotary_emb=rotary_emb, + ) + self.assertEqual(out.shape, (batch_size, seq_len, query_dim)) + + # Reference unfused path execution with identical weights + raw_q = attn.query(hidden_states) + raw_k = attn.key(hidden_states) + raw_v = attn.value(hidden_states) + q_norm = attn.norm_q(raw_q) + k_norm = attn.norm_k(raw_k) + q_h = _unflatten_heads(q_norm, heads) + k_h = _unflatten_heads(k_norm, heads) + v_h = _unflatten_heads(raw_v, heads) + q_rope, k_rope = attn._apply_rope(q_h, k_h, rotary_emb) + ref_attn_out = attn.attention_op.apply_attention(q_rope, k_rope, v_h, attention_mask=None) + ref_out = attn.proj_attn(ref_attn_out) + + np.testing.assert_allclose( + np.array(out, dtype=np.float32), + np.array(ref_out, dtype=np.float32), + atol=0.08, + rtol=1e-2, + ) + @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions") def test_wan_model(self): pyconfig.initialize( @@ -280,14 +580,20 @@ def test_wan_model(self): num_layers = 1 with nn_partitioning.axis_rules(config.logical_axis_rules): wan_model = WanModel( - rngs=rngs, attention="flash", mesh=mesh, flash_block_sizes=flash_block_sizes, num_layers=num_layers + rngs=rngs, + attention="flash", + mesh=mesh, + flash_block_sizes=flash_block_sizes, + num_layers=num_layers, ) dummy_timestep = jnp.ones((batch_size)) dummy_encoder_hidden_states = jnp.ones((batch_size, 512, 4096)) with mesh: dummy_output = wan_model( - hidden_states=dummy_hidden_states, timestep=dummy_timestep, encoder_hidden_states=dummy_encoder_hidden_states + hidden_states=dummy_hidden_states, + timestep=dummy_timestep, + encoder_hidden_states=dummy_encoder_hidden_states, ) assert dummy_output.shape == hidden_states_shape diff --git a/src/maxdiffusion/utils/export_utils.py b/src/maxdiffusion/utils/export_utils.py index 279ad1e90..d0c014fcc 100644 --- a/src/maxdiffusion/utils/export_utils.py +++ b/src/maxdiffusion/utils/export_utils.py @@ -131,7 +131,9 @@ def export_to_obj(mesh, output_obj_path: str = None): def _legacy_export_to_video( - video_frames: Union[List[np.ndarray], List[PIL.Image.Image]], output_video_path: str = None, fps: int = 10 + video_frames: Union[List[np.ndarray], List[PIL.Image.Image]], + output_video_path: str = None, + fps: int = 10, ): if is_opencv_available(): import cv2 @@ -212,21 +214,22 @@ def export_to_video( if output_video_path is None: output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4").name - if isinstance(video_frames, np.ndarray): - if video_frames.dtype != np.uint8: - video_frames = (video_frames * 255).astype(np.uint8) - elif isinstance(video_frames[0], np.ndarray): - video_frames = np.stack(video_frames) - if video_frames.dtype != np.uint8: - video_frames = (video_frames * 255).astype(np.uint8) - elif isinstance(video_frames[0], PIL.Image.Image): + if isinstance(video_frames, list) and len(video_frames) > 0 and isinstance(video_frames[0], PIL.Image.Image): video_frames = np.stack([np.asarray(frame) for frame in video_frames]) + else: + video_frames = np.asarray(video_frames) + if video_frames.dtype != np.uint8: + video_frames = (video_frames * 255).clip(0, 255).astype(np.uint8) with imageio.get_writer( - output_video_path, fps=fps, quality=quality, bitrate=bitrate, macro_block_size=macro_block_size + output_video_path, + fps=fps, + quality=quality, + bitrate=bitrate, + macro_block_size=macro_block_size, ) as writer: for frame in video_frames: - writer.append_data(frame) + writer.append_data(np.asarray(frame)) return output_video_path @@ -320,7 +323,12 @@ def _write_audio( def export_to_video_with_audio( - video: Any, fps: int, audio: Optional[Any], audio_sample_rate: Optional[int], output_path: str, audio_format: str = "s16" + video: Any, + fps: int, + audio: Optional[Any], + audio_sample_rate: Optional[int], + output_path: str, + audio_format: str = "s16", ) -> None: """ Encodes video (and optionally audio) to a file using PyAV. @@ -369,6 +377,12 @@ def export_to_video_with_audio( container.mux(packet) if audio is not None: - _write_audio(container, audio_stream, audio, audio_sample_rate, target_format=audio_format) + _write_audio( + container, + audio_stream, + audio, + audio_sample_rate, + target_format=audio_format, + ) container.close()