diff --git a/end_to_end/tpu/run_wan_fast_inference.sh b/end_to_end/tpu/run_wan_fast_inference.sh new file mode 100755 index 000000000..e712d6868 --- /dev/null +++ b/end_to_end/tpu/run_wan_fast_inference.sh @@ -0,0 +1,76 @@ +#!/bin/bash +# WAN T2V fast-serving example: AOT executable cache + converted-weights +# cache + zero-exec warmup, with a tuned v7 2D-ring attention recipe. +# +# First run per (model, shape) pays one-time conversion + compile and +# populates the caches; every later process start is ~25s to ready. +# +# Usage: +# ./run_wan_fast_inference.sh [21|22] [steps] ["prompt..."] +# Env overrides: +# WAN_CACHE_ROOT cache root (default ~/.cache/maxdiffusion_wan) +# 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) +set -u +MODEL=${1:-22} +STEPS=${2:-40} +PROMPT=${3:-""} + +PROJECT_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &> /dev/null && pwd)" +cd "$PROJECT_ROOT" || exit 1 +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 + +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_dvfs_p_state=7 --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" + +if [ "$MODEL" = "21" ]; then + CONFIG=src/maxdiffusion/configs/base_wan_14b.yml + GUIDANCE_ARGS="" +else + CONFIG=src/maxdiffusion/configs/base_wan_27b.yml + 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)" + +# libtpu's XLA:CPU AOT feature-mismatch log is cosmetic and ignores every +# log-level env var; filter just that message from stderr. +python src/maxdiffusion/generate_wan.py "$CONFIG" \ + run_name="$RUN_NAME" \ + output_dir="$OUTPUT_DIR" \ + jax_cache_dir="$CACHE_ROOT/jax" \ + aot_cache_dir="$CACHE_ROOT/aot_wan$MODEL" \ + converted_weights_dir="$CACHE_ROOT/converted" \ + attention=ulysses_ring_custom \ + 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 \ + 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_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":9472,"block_kv":1024,"block_kv_compute":1024,"block_kv_compute_in":1024,"heads_per_tile":1,"vmem_limit_bytes":67108864,"block_q_dkv":9472,"block_kv_dkv":1024,"block_kv_dkv_compute":1024}' \ + "${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) +if [ -n "$mp4" ]; then + mv "$mp4" "$OUTPUT_DIR/${RUN_NAME}.mp4" + echo "" + echo "=== video saved: $OUTPUT_DIR/${RUN_NAME}.mp4 ===" +fi diff --git a/src/maxdiffusion/aot_cache.py b/src/maxdiffusion/aot_cache.py new file mode 100644 index 000000000..5c0392a4b --- /dev/null +++ b/src/maxdiffusion/aot_cache.py @@ -0,0 +1,413 @@ +"""Per-shape AOT executable cache for jitted inference entry points. + +A JAX persistent-compilation-cache hit still pays trace + lowering + +cache-key hashing on every process start (~seconds per big executable). +This module serializes the compiled executable itself +(``jax.experimental.serialize_executable``) after the first warmup, so +subsequent processes deserialize on a background thread and call it +directly: no trace, no lowering, no cache lookup. + +Design (ported from DiffusionServing runners/torchax_aot.py, PR #38/#39, +minus the torch interop): + * ``cached_jit`` replaces ``jax.jit`` at the definition site. Until + ``install()`` is called it delegates to plain ``jax.jit`` with zero + behavioral difference, so tests and trainers are unaffected. + * One executable is kept PER dynamic input signature (shapes/dtypes of + array leaves + treedef + non-array leaves). Different resolutions or + frame counts never collide on disk. + * Unknown signature -> silent jit fallback; the first call's args are + recorded so ``save_pending()`` (call it after warmup, synchronously) + can lower + serialize that shape without touching other shapes. + * ``deserialize_and_load`` must receive ``execution_devices`` in the + mesh's topology-pinned order; the default ``jax.devices()`` order can + bind logical slots to the wrong physical chips and abort in C++. + * A deserialized ``Compiled`` does not auto-reshard inputs like jit + does; inputs are aligned to ``compiled.input_shardings`` in Python + before the call. + +Usage:: + + @partial(aot_cache.cached_jit, static_argnames=("guidance_scale",)) + def transformer_forward_pass(...): + ... + + # at pipeline construction (e.g. generate_wan.run): + aot_cache.install(cache_dir, meta={...config fingerprint...}, mesh=mesh) + # ... run warmup ... + aot_cache.save_pending() +""" + +from __future__ import annotations + +import contextlib +import glob +import hashlib +import inspect +import os +import pickle +import re +import threading +from typing import Any, Callable + +import jax +import jax.numpy as jnp +from jax.experimental import serialize_executable + +from maxdiffusion import max_logging + +_FORMAT_VERSION = 1 + + +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. + """ + 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}" + else: + desc = re.sub(r"0x[0-9a-fA-F]+", "@", repr(leaf)) + parts.append(f"{jax.tree_util.keystr(path)}={desc}") + return hashlib.sha256("|".join(parts).encode()).hexdigest()[:12] + + +class _AotEntry: + """Executables for one wrapped fn, keyed by dynamic input signature.""" + + def __init__(self, name: str, fn: Callable, static_argnames: tuple): + self.name = name + self.fn = fn + self.static_argnames = tuple(static_argnames) + self.py_signature = inspect.signature(fn) + self.jitted = jax.jit(fn, static_argnames=static_argnames or None) + self._compiled: dict[str, Any] = {} + self._out_specs: dict[str, Any] = {} + self._pending: dict[str, tuple] = {} + self._adapters: dict[str, Any] = {} + self._on_disk: set[str] = set() + self._lock = threading.Lock() + + def _zeros_output(self, signature: str): + """Builds all-zero outputs matching a compiled signature's out specs. + + Compilation only needs avals — executing a 14B transformer to warm the + pipeline wastes the full step compute. In warmup mode the caller gets + zeros with the executable's exact shapes/dtypes/shardings, so every + DOWNSTREAM executable (scheduler step, VAE) still compiles against + faithful inputs. Returns None if out specs are unknown for this + signature (caller must execute normally). + """ + spec = self._out_specs.get(signature) + if spec is None: + return None + out_treedef, shapes_dtypes, out_shardings = spec + zeros = [ + jax.device_put(jnp.zeros(shape, dtype), sharding) for (shape, dtype), sharding in zip(shapes_dtypes, out_shardings) + ] + return jax.tree_util.tree_unflatten(out_treedef, zeros) + + def _adapter_for(self, signature: str, treedef: Any, static: dict): + """Returns the per-signature flat-leaf-list jit of fn. + + The input treedef (which may embed unpicklable statics, e.g. an nnx + GraphDef holding initializer closures) stays inside this process's + closure and is never serialized — the adapter's own in/out trees are + plain lists/tuples of arrays. Warmup compiles this adapter, so + ``save_pending``'s lower().compile() hits the in-memory pjit cache + instead of recompiling. + """ + adapter_jit = self._adapters.get(signature) + if adapter_jit is None: + + def adapter(flat, _treedef=treedef, _static=static): + return self.fn(**jax.tree_util.tree_unflatten(_treedef, flat), **_static) + + adapter_jit = jax.jit(adapter) + with self._lock: + self._adapters.setdefault(signature, adapter_jit) + adapter_jit = self._adapters[signature] + return adapter_jit + + def _canonicalize(self, args: tuple, kwargs: dict) -> tuple[dict, dict]: + """Splits a call into (dynamic kwargs, static kwargs) by param name. + + A deserialized ``Compiled`` must be called with the static args + STRIPPED (they are baked into the graph and absent from its input + pytree), and positional/keyword form must match how it was lowered. + Canonicalizing every call to keyword form on both the lower and call + paths makes the pytrees agree by construction. + """ + bound = self.py_signature.bind(*args, **kwargs) + dynamic, static = {}, {} + for name, val in bound.arguments.items(): + (static if name in self.static_argnames else dynamic)[name] = val + return dynamic, static + + def _compile_and_record(self, signature: str, leaves: list, treedef: Any, static: dict): + """Lower+compile one signature (no execution) and capture out specs.""" + lowered = self._adapter_for(signature, treedef, static).lower(leaves) + compiled = lowered.compile() + info_leaves, out_treedef = jax.tree_util.tree_flatten(lowered.out_info) + shapes_dtypes = [(tuple(x.shape), jnp.dtype(x.dtype)) for x in info_leaves] + out_shardings = jax.tree_util.tree_leaves(compiled.output_shardings) + with self._lock: + self._compiled[signature] = compiled + self._out_specs[signature] = (out_treedef, shapes_dtypes, out_shardings) + return compiled + + # ---------------------------------------------------------------- call + def __call__(self, *args, **kwargs): + if not _STATE.enabled: + return self.jitted(*args, **kwargs) + dynamic, static = self._canonicalize(args, kwargs) + leaves, treedef = jax.tree_util.tree_flatten(dynamic) + if any(isinstance(leaf, jax.core.Tracer) for leaf in leaves): + # 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}) + if _STATE.warmup_only: + # Compilation only needs avals; skip the (possibly seconds-long) + # real execution and hand back correctly-shaped/sharded zeros so + # downstream executables still warm against faithful inputs. + if signature not in self._compiled: + self._compile_and_record(signature, leaves, treedef, static) + with self._lock: + if signature not in self._on_disk: + self._pending[signature] = (leaves, treedef, static) + zeros = self._zeros_output(signature) + if zeros is not None: + return zeros + 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: + 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") + with self._lock: + if signature not in self._pending and signature not in self._compiled: + self._pending[signature] = (leaves, treedef, static) + return self._adapter_for(signature, treedef, static)(leaves) + + def _align_inputs(self, compiled: Any, leaves: list): + """Reshards the flat input leaves onto the executable's shardings. + + 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). + """ + 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") + return None + + # ---------------------------------------------------------------- disk + def _path_for(self, signature: str) -> str: + return os.path.join(_STATE.cache_dir, f"{self.name}-{_STATE.fingerprint}-{signature}.aotx") + + def load_from_disk(self) -> None: + """Deserializes every on-disk executable for this fn. Never raises.""" + pattern = os.path.join(_STATE.cache_dir, f"{self.name}-{_STATE.fingerprint}-*.aotx") + for path in glob.glob(pattern): + try: + with open(path, "rb") as f: + blob = pickle.load(f) + if blob["format_version"] != _FORMAT_VERSION: + continue + # Topology-pinned device order; the default reconstruction binds + # logical slots to the wrong physical chips and aborts in C++. + execution_devices = list(_STATE.mesh.devices.flatten()) if _STATE.mesh is not None else None + compiled = serialize_executable.deserialize_and_load( + blob["payload"], + blob["in_tree"], + blob["out_tree"], + execution_devices=execution_devices, + ) + signature = blob["dynamic_signature"] + with self._lock: + self._compiled[signature] = compiled + self._on_disk.add(signature) + if "out_shapes_dtypes" in blob: + # Out specs let warmup mode return zeros instead of executing. + self._out_specs[signature] = ( + blob["out_tree"], + [(tuple(shape), jnp.dtype(dtype)) for shape, dtype in blob["out_shapes_dtypes"]], + jax.tree_util.tree_leaves(compiled.output_shardings), + ) + max_logging.log(f"[aot] {self.name}: loaded {os.path.basename(path)} ({len(blob['payload']) / 1e6:.1f}MB)") + except Exception as e: # noqa: BLE001 - fall back to jit for this shape + max_logging.log(f"[aot] {self.name}: load failed for {os.path.basename(path)} ({e}); will re-jit") + + def save_pending(self) -> int: + """Lowers + serializes every recorded signature. Returns count saved.""" + saved = 0 + with self._lock: + pending, self._pending = self._pending, {} + for signature, (leaves, treedef, static) in pending.items(): + if signature in self._on_disk: + # Background deserialization landed after this shape was recorded. + continue + try: + compiled = self._compiled.get(signature) + if compiled is None: + # Re-lowering retraces on this thread (the compile itself hits + # the in-memory cache from warmup); sharding constraints inside + # the model need the mesh context that warmup provided. + mesh_ctx = _STATE.mesh if _STATE.mesh is not None else contextlib.nullcontext() + with mesh_ctx: + compiled = self._compile_and_record(signature, leaves, treedef, static) + payload, in_tree, out_tree = serialize_executable.serialize(compiled) + blob = { + "format_version": _FORMAT_VERSION, + "payload": payload, + "in_tree": in_tree, + "out_tree": out_tree, + "dynamic_signature": signature, + } + out_spec = self._out_specs.get(signature) + if out_spec is not None: + blob["out_shapes_dtypes"] = [(list(shape), str(dtype)) for shape, dtype in out_spec[1]] + path = self._path_for(signature) + tmp_path = f"{path}.tmp.{os.getpid()}" + with open(tmp_path, "wb") as f: + pickle.dump(blob, f) + os.replace(tmp_path, path) + with self._lock: + self._compiled[signature] = compiled + self._on_disk.add(signature) + saved += 1 + max_logging.log(f"[aot] {self.name}: serialized {os.path.basename(path)} ({len(payload) / 1e6:.1f}MB)") + except Exception as e: # noqa: BLE001 - saving is best-effort + max_logging.log(f"[aot] {self.name}: serialize failed ({e}); shape stays on jit") + return saved + + +class _State: + """Process-global install state (null until install() is called).""" + + def __init__(self): + self.enabled = False + self.cache_dir = "" + self.fingerprint = "" + self.mesh = None + self.warmup_only = False + + +_STATE = _State() +_REGISTRY: list[_AotEntry] = [] +_LOAD_THREADS: list[threading.Thread] = [] + + +def cached_jit(fn: Callable, static_argnames: tuple = ()) -> Callable: + """Drop-in replacement for ``jax.jit`` with an optional AOT layer. + + Behaves exactly like ``jax.jit(fn, static_argnames=...)`` until + ``install()`` enables the executable cache. + """ + # Qualify by module: same-named fns (e.g. the VACE and base + # transformer_forward_pass) must not glob each other's files. + name = f"{fn.__module__.rsplit('.', 1)[-1]}.{fn.__name__}" + entry = _AotEntry(name, fn, static_argnames) + _REGISTRY.append(entry) + return entry + + +def install(cache_dir: str, meta: dict[str, Any], mesh: Any) -> None: + """Enables the AOT cache and starts background deserialization. + + Args: + cache_dir: Directory for .aotx files (created if missing). + meta: Everything the executables depend on beyond input shapes: + model path, mesh shape, sharding/attention config, jax version. + Hashed into the filename so incompatible executables never load. + mesh: The pipeline mesh; pins device order for deserialization and + provides the context for re-lowering at save time. + """ + if not cache_dir: + return + os.makedirs(cache_dir, exist_ok=True) + _STATE.cache_dir = cache_dir + _STATE.fingerprint = hashlib.sha256(repr(sorted(meta.items())).encode()).hexdigest()[:12] + _STATE.mesh = mesh + _STATE.enabled = True + for entry in _REGISTRY: + with entry._lock: + # Cached state belongs to the previous install's dir/fingerprint. + entry._compiled.clear() + entry._out_specs.clear() + entry._pending.clear() + entry._adapters.clear() + entry._on_disk.clear() + thread = threading.Thread(target=entry.load_from_disk, name=f"aot-load-{entry.name}", daemon=True) + thread.start() + _LOAD_THREADS.append(thread) + + +def wait_for_loads() -> None: + """Blocks until background deserialization finishes (call before timing).""" + for thread in _LOAD_THREADS: + thread.join() + _LOAD_THREADS.clear() + + +def save_pending() -> int: + """Serializes all recorded shapes across wrapped fns. Call after warmup, + synchronously — a background save competes with the first real request + (DiffusionServing PR#39 first-generation-stall lesson).""" + if not _STATE.enabled: + return 0 + return sum(entry.save_pending() for entry in _REGISTRY) + + +@contextlib.contextmanager +def warmup_mode(): + """Zero-execution warmup: wrapped fns lower+compile but never execute. + + Compilation only needs avals, not values — a 14B transformer step costs + seconds to run and nothing to skip. Inside this context a wrapped call + compiles its signature (or reuses the deserialized executable) and + returns all-zero outputs with the executable's exact shapes, dtypes and + shardings, so downstream executables (scheduler step, VAE decode) still + compile against faithful inputs. Outputs of a warmup pass are garbage by + design; callers must discard them. No-op when the cache is disabled. + """ + _STATE.warmup_only = _STATE.enabled + try: + yield + finally: + _STATE.warmup_only = False diff --git a/src/maxdiffusion/configs/base_wan_14b.yml b/src/maxdiffusion/configs/base_wan_14b.yml index 493dfd6eb..18a2b4eee 100644 --- a/src/maxdiffusion/configs/base_wan_14b.yml +++ b/src/maxdiffusion/configs/base_wan_14b.yml @@ -274,6 +274,10 @@ load_tfrecord_cached: True train_data_dir: '' dataset_config_name: '' jax_cache_dir: '' +# Directory for per-shape AOT serialized executables ('' = disabled). +aot_cache_dir: '' +# Directory for memoized torch->flax converted weights ('' = disabled). +converted_weights_dir: '' hf_data_dir: '' hf_train_files: '' hf_access_token: '' diff --git a/src/maxdiffusion/configs/base_wan_1_3b.yml b/src/maxdiffusion/configs/base_wan_1_3b.yml index 994b1da36..c0de05c9e 100644 --- a/src/maxdiffusion/configs/base_wan_1_3b.yml +++ b/src/maxdiffusion/configs/base_wan_1_3b.yml @@ -227,6 +227,10 @@ load_tfrecord_cached: True train_data_dir: '' dataset_config_name: '' jax_cache_dir: '' +# Directory for per-shape AOT serialized executables ('' = disabled). +aot_cache_dir: '' +# Directory for memoized torch->flax converted weights ('' = disabled). +converted_weights_dir: '' hf_data_dir: '' hf_train_files: '' hf_access_token: '' diff --git a/src/maxdiffusion/configs/base_wan_27b.yml b/src/maxdiffusion/configs/base_wan_27b.yml index 0e919a479..4e5f7642f 100644 --- a/src/maxdiffusion/configs/base_wan_27b.yml +++ b/src/maxdiffusion/configs/base_wan_27b.yml @@ -248,6 +248,10 @@ load_tfrecord_cached: True train_data_dir: '' dataset_config_name: '' jax_cache_dir: '' +# Directory for per-shape AOT serialized executables ('' = disabled). +aot_cache_dir: '' +# Directory for memoized torch->flax converted weights ('' = disabled). +converted_weights_dir: '' hf_data_dir: '' hf_train_files: '' hf_access_token: '' diff --git a/src/maxdiffusion/configs/base_wan_animate.yml b/src/maxdiffusion/configs/base_wan_animate.yml index 35a063203..e0abf4515 100644 --- a/src/maxdiffusion/configs/base_wan_animate.yml +++ b/src/maxdiffusion/configs/base_wan_animate.yml @@ -237,6 +237,10 @@ load_tfrecord_cached: True train_data_dir: '' dataset_config_name: '' jax_cache_dir: '.jax_cache' +# Directory for per-shape AOT serialized executables ('' = disabled). +aot_cache_dir: '' +# Directory for memoized torch->flax converted weights ('' = disabled). +converted_weights_dir: '' hf_data_dir: '' hf_train_files: '' hf_access_token: '' diff --git a/src/maxdiffusion/configs/base_wan_i2v_14b.yml b/src/maxdiffusion/configs/base_wan_i2v_14b.yml index 3188e7849..5c59ddbdc 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_14b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_14b.yml @@ -240,6 +240,10 @@ load_tfrecord_cached: True train_data_dir: '' dataset_config_name: '' jax_cache_dir: '' +# Directory for per-shape AOT serialized executables ('' = disabled). +aot_cache_dir: '' +# Directory for memoized torch->flax converted weights ('' = disabled). +converted_weights_dir: '' hf_data_dir: '' hf_train_files: '' hf_access_token: '' diff --git a/src/maxdiffusion/configs/base_wan_i2v_27b.yml b/src/maxdiffusion/configs/base_wan_i2v_27b.yml index 5f2e8c884..8c4bf8853 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_27b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_27b.yml @@ -241,6 +241,10 @@ load_tfrecord_cached: True train_data_dir: '' dataset_config_name: '' jax_cache_dir: '' +# Directory for per-shape AOT serialized executables ('' = disabled). +aot_cache_dir: '' +# Directory for memoized torch->flax converted weights ('' = disabled). +converted_weights_dir: '' hf_data_dir: '' hf_train_files: '' hf_access_token: '' diff --git a/src/maxdiffusion/generate_wan.py b/src/maxdiffusion/generate_wan.py index a07fc5632..49f35f490 100644 --- a/src/maxdiffusion/generate_wan.py +++ b/src/maxdiffusion/generate_wan.py @@ -21,7 +21,7 @@ from maxdiffusion.checkpointing.wan_checkpointer_2_2 import WanCheckpointer2_2 from maxdiffusion.checkpointing.wan_checkpointer_i2v_2p1 import WanCheckpointerI2V_2_1 from maxdiffusion.checkpointing.wan_checkpointer_i2v_2p2 import WanCheckpointerI2V_2_2 -from maxdiffusion import pyconfig, max_logging, max_utils +from maxdiffusion import aot_cache, pyconfig, max_logging, max_utils from absl import app from maxdiffusion.train_utils import transformer_engine_context from maxdiffusion.utils import export_to_video @@ -279,6 +279,30 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): dtype=config.weights_dtype, ) + # 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. + 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), + "weights_dtype": str(config.weights_dtype), + "activations_dtype": str(config.activations_dtype), + "scan_layers": str(config.scan_layers), + "jax": jax.__version__, + }, + mesh=pipeline.mesh, + ) + # Deserialization is seconds and warmup must see the loaded executables + # to hit them; without this the first call races the loader threads. + aot_cache.wait_for_loads() + s0 = time.perf_counter() # Disable profiler for the first two runs to avoid duplicate uploads @@ -299,11 +323,21 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): # step count only changes the Python loop trip count, not traced shapes. warmup_steps = min(2, config.num_inference_steps) max_logging.log(f"Compile warmup: {warmup_steps} denoising steps") - videos = call_pipeline(config, pipeline, prompt, negative_prompt, num_inference_steps=warmup_steps) + # Zero-execution warmup: wrapped transformer passes lower+compile (or + # reuse the deserialized AOT executable) and return sharded zeros, so + # 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, prompt, negative_prompt, num_inference_steps=warmup_steps) if isinstance(videos, tuple): videos, warmup_trace = videos max_logging.log("Warmup breakdown: " + ", ".join(f"{stage}={seconds:.1f}s" for stage, seconds in warmup_trace.items())) + # Serialize any newly-compiled shapes synchronously while still inside + # warmup-accounted time; a background save would compete with the first + # real generation (DiffusionServing PR#39 first-generation-stall lesson). + aot_cache.save_pending() + max_logging.log("===================== Model details =======================") max_logging.log(f"model name: {config.model_name}") max_logging.log(f"model path: {config.pretrained_model_name_or_path}") @@ -356,15 +390,17 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): vae_decode_total = trace.get("vae_decode", 0.0) vae_decode_tpu = trace.get("vae_decode_tpu", 0.0) vae_decode_post = vae_decode_total - vae_decode_tpu - summary.extend([ - f" {'─' * 40}", - f" Conditioning: {trace.get('conditioning', 0.0):>7.1f}s", - f" - VAE Encode: {trace.get('vae_encode', 0.0):>7.1f}s", - f" Denoise Total: {trace.get('denoise_total', 0.0):>7.1f}s", - f" VAE Decode: {vae_decode_total:>7.1f}s", - f" - TPU Compute: {vae_decode_tpu:>7.1f}s", - f" - Host Formatting: {vae_decode_post:>7.1f}s", - ]) + summary.extend( + [ + f" {'─' * 40}", + f" Conditioning: {trace.get('conditioning', 0.0):>7.1f}s", + f" - VAE Encode: {trace.get('vae_encode', 0.0):>7.1f}s", + f" Denoise Total: {trace.get('denoise_total', 0.0):>7.1f}s", + f" VAE Decode: {vae_decode_total:>7.1f}s", + f" - TPU Compute: {vae_decode_tpu:>7.1f}s", + f" - Host Formatting: {vae_decode_post:>7.1f}s", + ] + ) summary.append(f"{'=' * 50}") max_logging.log("\n".join(summary)) diff --git a/src/maxdiffusion/models/wan/wan_utils.py b/src/maxdiffusion/models/wan/wan_utils.py index 883ef7d23..e60ffd31c 100644 --- a/src/maxdiffusion/models/wan/wan_utils.py +++ b/src/maxdiffusion/models/wan/wan_utils.py @@ -17,6 +17,7 @@ import concurrent.futures import json import os +import shutil import threading import time from typing import Callable, Optional @@ -286,6 +287,7 @@ def load_wan_transformer( scan_layers: bool = True, subfolder: str = "", cast_dtype_fn: Optional[Callable] = None, + converted_cache_dir: str = "", ): if pretrained_model_name_or_path == CAUSVID_TRANSFORMER_MODEL_NAME_OR_PATH: return load_causvid_transformer(pretrained_model_name_or_path, eval_shapes, device, hf_download, num_layers, scan_layers) @@ -293,7 +295,15 @@ def load_wan_transformer( return load_fusionx_transformer(pretrained_model_name_or_path, eval_shapes, device, hf_download, num_layers, scan_layers) else: return load_base_wan_transformer( - pretrained_model_name_or_path, eval_shapes, device, hf_download, num_layers, scan_layers, subfolder, cast_dtype_fn + pretrained_model_name_or_path, + eval_shapes, + device, + hf_download, + num_layers, + scan_layers, + subfolder, + cast_dtype_fn, + converted_cache_dir, ) @@ -308,6 +318,76 @@ def _torch_tensor_to_numpy(tensor: torch.Tensor) -> np.ndarray: return tensor.numpy() +def _converted_key_to_filename(flax_key: tuple) -> str: + return ".".join(str(k) for k in flax_key) + ".npy" + + +def try_load_converted_weights(cache_dir: str, eval_shapes: dict, cast_dtype_fn: Optional[Callable]) -> Optional[dict]: + """Loads a converted-weights cache as mmapped arrays, or None on mismatch. + + The torch->flax conversion (transpose + scan-stack + cast) is a pure + function of the checkpoint, so it is paid once and memoized on disk. + Keys/shapes are validated against eval_shapes and dtypes against + cast_dtype_fn, so a policy or model change falls back to a fresh + conversion (which re-saves). + """ + manifest_path = os.path.join(cache_dir, "manifest.json") + if not os.path.isfile(manifest_path): + return None + try: + with open(manifest_path, "r") as f: + manifest = json.load(f) + expected_keys = set(flatten_dict(eval_shapes).keys()) + + def load_one(key_str, meta): + flax_key = _tuple_str_to_int(tuple(key_str.split("."))) + logical_dtype = np.dtype(meta["dtype"]) + if cast_dtype_fn is not None and logical_dtype != np.dtype(cast_dtype_fn(flax_key)): + raise ValueError(f"dtype policy changed for {key_str}") + # Eager parallel read (page-cache/RAM speed): an mmap would defer the + # read into the device_put as serial page faults, halving put speed. + value = np.load(os.path.join(cache_dir, meta["file"])) + if meta.get("bitview"): + # Non-native dtypes (bf16/fp8) are stored as same-width uints: + # npy cannot resolve ml_dtypes descriptors on all paths. + value = value.view(logical_dtype) + if tuple(value.shape) != tuple(meta["shape"]): + raise ValueError(f"shape changed for {key_str}") + return flax_key, value + + flax_state_dict = {} + with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor: + for flax_key, value in executor.map(lambda kv: load_one(*kv), manifest.items()): + flax_state_dict[flax_key] = value + if set(flax_state_dict.keys()) != expected_keys: + return None + return unflatten_dict(flax_state_dict) + except (OSError, ValueError, KeyError, TypeError) as e: + max_logging.log(f"Converted-weights cache unusable ({e}); reconverting") + return None + + +def save_converted_weights(cache_dir: str, flat_state_dict: dict) -> None: + """Writes the converted tree as per-tensor .npy + manifest, atomically.""" + tmp_dir = f"{cache_dir}.tmp.{os.getpid()}" + os.makedirs(tmp_dir, exist_ok=True) + manifest = {} + uint_by_width = {1: np.uint8, 2: np.uint16, 4: np.uint32} + for flax_key, value in flat_state_dict.items(): + filename = _converted_key_to_filename(flax_key) + bitview = value.dtype.kind not in "fiub" # ml_dtypes (bf16/fp8) etc. + stored = value.view(uint_by_width[value.dtype.itemsize]) if bitview else value + np.save(os.path.join(tmp_dir, filename), stored) + key_str = ".".join(str(k) for k in flax_key) + manifest[key_str] = {"file": filename, "shape": list(value.shape), "dtype": str(value.dtype), "bitview": bitview} + with open(os.path.join(tmp_dir, "manifest.json"), "w") as f: + json.dump(manifest, f) + try: + os.rename(tmp_dir, cache_dir) + except OSError: + shutil.rmtree(tmp_dir, ignore_errors=True) # another process won the race + + def load_base_wan_transformer( pretrained_model_name_or_path: str, eval_shapes: dict, @@ -317,6 +397,7 @@ def load_base_wan_transformer( scan_layers: bool = True, subfolder: str = "", cast_dtype_fn: Optional[Callable] = None, + converted_cache_dir: str = "", ): """Loads WAN transformer weights from diffusers safetensors shards. @@ -335,6 +416,14 @@ def load_base_wan_transformer( Returns a nested dict of numpy arrays (host memory). """ del device # weights stay in plain host numpy until device_put by the caller + if converted_cache_dir: + t_start = time.perf_counter() + cached = try_load_converted_weights(converted_cache_dir, eval_shapes, cast_dtype_fn) + if cached is not None: + max_logging.log( + f"Loaded converted {subfolder or 'transformer'} weights (mmap) in {time.perf_counter() - t_start:.1f}s" + ) + return cached filename = "diffusion_pytorch_model.safetensors.index.json" local_files = False if os.path.isdir(pretrained_model_name_or_path): @@ -425,6 +514,10 @@ def convert_chunk(ckpt_shard_path, chunk_keys): future.result() # re-raise conversion errors validate_flax_state_dict(eval_shapes, flax_state_dict) + if converted_cache_dir and not os.path.isdir(converted_cache_dir): + t_save = time.perf_counter() + save_converted_weights(converted_cache_dir, flax_state_dict) + max_logging.log(f"Saved converted-weights cache to {converted_cache_dir} in {time.perf_counter() - t_save:.1f}s") flax_state_dict = unflatten_dict(flax_state_dict) max_logging.log(f"Converted {subfolder or 'transformer'} weights to host arrays in {time.perf_counter() - t_start:.1f}s") return flax_state_dict diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline.py b/src/maxdiffusion/pipelines/wan/wan_pipeline.py index 730d67fc7..9ebfa68e9 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline.py @@ -18,6 +18,7 @@ from maxdiffusion.image_processor import PipelineImageInput import numpy as np import math +import os import jax import jax.numpy as jnp import threading @@ -28,6 +29,7 @@ from flax import nnx from flax.linen import partitioning as nn_partitioning from ...pyconfig import HyperParameters +from ... import aot_cache from ... import max_logging from ... import max_utils from ...max_utils import get_flash_block_sizes, get_precision, device_put_replicated @@ -157,6 +159,149 @@ def _select_restored_transformer_state(restored_checkpoint, subfolder: str): raise ValueError(f"Unsupported WAN checkpoint transformer subfolder `{subfolder}`.") +# Concurrent transformer loads (WAN 2.2's two experts) must not interleave +# their device transfers: shared PCIe lanes degrade ~50% under contention. +_DEVICE_PUT_LOCK = threading.Lock() + + +def converted_weights_cache_dir(config, subfolder: str) -> str: + """Per-(model, subfolder, dtype, scan) dir for memoized converted weights.""" + base = getattr(config, "converted_weights_dir", "") + if not base: + return "" + model_tag = (config.wan_transformer_pretrained_model_name_or_path or config.pretrained_model_name_or_path).replace( + "/", "--" + ) + return os.path.join(base, f"{model_tag}--{subfolder or 'transformer'}--{config.weights_dtype}--scan{config.scan_layers}") + + +def put_params_into_state( + state: dict, + params: dict, + logical_state_sharding: dict, + mesh: Mesh, + restored_checkpoint=None, + subfolder: str = "", +) -> dict: + """Moves host params into the flat nnx state on device. + + Shared by the WAN 2.x and VACE pipelines. Single-process: replicated + params are the bulk of the bytes; a direct device_put broadcasts the same + bytes over every device's PCIe stream (~2GB/s each). Instead, stage them + sharded along dim0 (each device receives only 1/n of the bytes over PCIe) + and replicate on-device through ICI, which is an order of magnitude + faster than host links. Multi-process: per-param device_put_replicated + with a process_allgather fallback. + + Args: + state: Flat nnx state dict (path tuple -> VariableState) to fill. + params: Host-side param tree with final dtypes. + logical_state_sharding: Flat dict of target NamedShardings per path. + mesh: Device mesh the shardings refer to. + restored_checkpoint: If set, params came from an orbax restore and + paths need 'value' suffix / block-index normalization. + subfolder: Label used only for logging. + + Returns: + The same `state` dict with `.value` set to on-device arrays. + """ + t_put_start = time.perf_counter() + put_specs = [] + for path, val in flax.traverse_util.flatten_dict(params).items(): + if restored_checkpoint: + if path[-1] == "value": + path = path[:-1] # remove 'value' + + try: + # Convert block indices to integers, as they might have been loaded as strings from the checkpoint. + path = path[:1] + (int(path[1]),) + path[2:] + except Exception: + pass + + put_specs.append((path, val, logical_state_sharding[path].value)) + + if jax.process_count() == 1: + n_devices = mesh.devices.size + dim0_sharding = NamedSharding(mesh, P(mesh.axis_names)) + + def stages_via_ici(val, sharding) -> bool: + return ( + sharding.is_fully_replicated + and val.ndim > 0 + and val.shape[0] % n_devices == 0 + and val.nbytes >= 1 << 26 # 64MB: below this, staging overhead wins + ) + + staged_ids = [i for i, (_, val, sharding) in enumerate(put_specs) if stages_via_ici(val, sharding)] + direct_ids = [i for i in range(len(put_specs)) if i not in set(staged_ids)] + + put_arrays = [None] * len(put_specs) + # Lock through block_until_ready: puts are async, and concurrent expert + # transfers on shared PCIe lanes degrade ~50%. + with _DEVICE_PUT_LOCK: + if staged_ids: + # Per-device slice puts run each device's PCIe lane in parallel + # (a sharded device_put of the whole list serializes near single- + # lane speed). Gathers go in ~6GB chunks: chunk N replicates over + # ICI while chunk N+1's host transfers stream, and the transient + # HBM reservation stays bounded. + chunk_limit_bytes = 6 << 30 + chunks, current, current_bytes = [], [], 0 + for i in staged_ids: + current.append(i) + current_bytes += put_specs[i][1].nbytes + if current_bytes >= chunk_limit_bytes: + chunks.append(current) + current, current_bytes = [], 0 + if current: + chunks.append(current) + for chunk in chunks: + # One batched put per device: per-tensor-per-device calls pay + # dispatch overhead 8x per tensor and defeat lane pipelining. + slices_by_device = {} + index_maps = [] + for i in chunk: + val = put_specs[i][1] + indices_map = dim0_sharding.addressable_devices_indices_map(val.shape) + index_maps.append(list(indices_map.items())) + for d, index in indices_map.items(): + slices_by_device.setdefault(d, []).append(val[index]) + shards_by_device = {d: iter(jax.device_put(slices, d)) for d, slices in slices_by_device.items()} + sharded_arrays = [] + for i, device_indices in zip(chunk, index_maps): + val = put_specs[i][1] + shards = [next(shards_by_device[d]) for d, _ in device_indices] + sharded_arrays.append(jax.make_array_from_single_device_arrays(val.shape, dim0_sharding, shards)) + # out_shardings must be the exact target sharding objects (not an + # equivalent P()): downstream jit cache keys include arg shardings, + # so a different-but-equivalent spec would force a full recompile. + replicate_fn = jax.jit(lambda xs: xs, out_shardings=[put_specs[i][2] for i in chunk]) + for i, replicated in zip(chunk, replicate_fn(sharded_arrays)): + put_arrays[i] = replicated + if direct_ids: + for i, put_array in zip( + direct_ids, + jax.device_put([put_specs[i][1] for i in direct_ids], [put_specs[i][2] for i in direct_ids]), + ): + put_arrays[i] = put_array + jax.block_until_ready([a for a in put_arrays if a is not None]) + for (path, _, _), put_array in zip(put_specs, put_arrays): + state[path].value = put_array + else: + for path, val, sharding in put_specs: + try: + state[path].value = device_put_replicated(val, sharding) + except Exception as e: + max_logging.log(f"Failed to device_put_replicated {path}: {e}") + max_logging.log(f"Trying to use process_allgather for {path}") + val_on_host = jax.experimental.multihost_utils.process_allgather(val, tiled=True) + state[path].value = device_put_replicated(val_on_host, sharding) + del val_on_host + jax.block_until_ready([state[path].value for path, _, _ in put_specs]) + max_logging.log(f"Transformer {subfolder or 'transformer'} weights on device in {time.perf_counter() - t_put_start:.1f}s") + return state + + # For some reason, jitting this function increases the memory significantly, so instead manually move weights to device. def create_sharded_logical_transformer( devices_array: np.array, @@ -234,6 +379,7 @@ def create_model(rngs: nnx.Rngs, wan_config: dict): scan_layers=config.scan_layers, subfolder=subfolder, cast_dtype_fn=partial(_final_param_dtype, dtype_to_cast=config.weights_dtype), + converted_cache_dir=converted_weights_cache_dir(config, subfolder), ) # No-op (returns leaves unchanged) when the loader already cast to the @@ -242,70 +388,14 @@ def create_model(rngs: nnx.Rngs, wan_config: dict): lambda path, x: cast_with_exclusion(path, x, dtype_to_cast=config.weights_dtype), params, ) - t_put_start = time.perf_counter() - put_specs = [] - for path, val in flax.traverse_util.flatten_dict(params).items(): - if restored_checkpoint: - if path[-1] == "value": - path = path[:-1] # remove 'value' - - try: - # Convert block indices to integers, as they might have been loaded as strings from the checkpoint. - path = path[:1] + (int(path[1]),) + path[2:] - except Exception: - pass - - put_specs.append((path, val, logical_state_sharding[path].value)) - - if jax.process_count() == 1: - # Replicated params are the bulk of the bytes; a direct device_put - # broadcasts the same bytes over every device's PCIe stream (~2GB/s - # each). Instead, stage them sharded along dim0 (each device receives - # only 1/n of the bytes over PCIe) and replicate on-device through ICI, - # which is an order of magnitude faster than host links. - n_devices = mesh.devices.size - dim0_sharding = NamedSharding(mesh, P(mesh.axis_names)) - - def stages_via_ici(val, sharding) -> bool: - return ( - sharding.is_fully_replicated - and val.ndim > 0 - and val.shape[0] % n_devices == 0 - and val.nbytes >= 1 << 26 # 64MB: below this, staging overhead wins - ) - - staged_ids = [i for i, (_, val, sharding) in enumerate(put_specs) if stages_via_ici(val, sharding)] - direct_ids = [i for i in range(len(put_specs)) if i not in set(staged_ids)] - - put_arrays = [None] * len(put_specs) - if staged_ids: - staged = jax.device_put([put_specs[i][1] for i in staged_ids], [dim0_sharding] * len(staged_ids)) - # out_shardings must be the exact target sharding objects (not an - # equivalent P()): downstream jit cache keys include arg shardings, so - # a different-but-equivalent spec would force a full recompile. - replicate_fn = jax.jit(lambda xs: xs, out_shardings=[put_specs[i][2] for i in staged_ids]) - for i, replicated in zip(staged_ids, replicate_fn(staged)): - put_arrays[i] = replicated - if direct_ids: - for i, put_array in zip( - direct_ids, - jax.device_put([put_specs[i][1] for i in direct_ids], [put_specs[i][2] for i in direct_ids]), - ): - put_arrays[i] = put_array - for (path, _, _), put_array in zip(put_specs, put_arrays): - state[path].value = put_array - else: - for path, val, sharding in put_specs: - try: - state[path].value = device_put_replicated(val, sharding) - except Exception as e: - max_logging.log(f"Failed to device_put_replicated {path}: {e}") - max_logging.log(f"Trying to use process_allgather for {path}") - val_on_host = jax.experimental.multihost_utils.process_allgather(val, tiled=True) - state[path].value = device_put_replicated(val_on_host, sharding) - del val_on_host - jax.block_until_ready([state[path].value for path, _, _ in put_specs]) - max_logging.log(f"Transformer {subfolder or 'transformer'} weights on device in {time.perf_counter() - t_put_start:.1f}s") + state = put_params_into_state( + state, + params, + logical_state_sharding, + mesh, + restored_checkpoint=restored_checkpoint, + subfolder=subfolder, + ) state = nnx.from_flat_state(state) wan_transformer = nnx.merge(graphdef, state, rest_of_state) @@ -379,6 +469,9 @@ def __init__( self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) self.p_run_inference = None + # encode_prompt result cache: same-prompt calls (warmup + real run, + # repeated serving requests) skip the ~10s/call CPU text encoder. + self._prompt_embeds_cache = {} @classmethod def load_text_encoder(cls, config: HyperParameters): @@ -659,6 +752,15 @@ def encode_prompt( elif isinstance(negative_prompt, str): negative_prompt = [negative_prompt] * batch_size + # Same-prompt calls (warmup then the real generation, or repeated + # serving requests) should not re-run the ~10s/call CPU text encoder. + cache_key = None + if prompt is not None and prompt_embeds is None and negative_prompt_embeds is None: + cache_key = (tuple(prompt), tuple(negative_prompt), num_videos_per_prompt, max_sequence_length) + cached = self._prompt_embeds_cache.get(cache_key) + if cached is not None: + return cached + use_batched_text_encoder = self.config.use_batched_text_encoder if use_batched_text_encoder and prompt_embeds is None and negative_prompt_embeds is None: # Batch both together @@ -692,6 +794,11 @@ def encode_prompt( ) negative_prompt_embeds = jnp.array(negative_prompt_embeds.detach().float().numpy(), dtype=jnp.float32) + if cache_key is not None: + if len(self._prompt_embeds_cache) >= 16: # bound long-serving growth + self._prompt_embeds_cache.pop(next(iter(self._prompt_embeds_cache))) + self._prompt_embeds_cache[cache_key] = (prompt_embeds, negative_prompt_embeds) + return prompt_embeds, negative_prompt_embeds def prepare_latents( @@ -763,7 +870,8 @@ def prepare_latents_i2v_base( video_condition = video_condition.astype(vae_dtype) t_vae_encode_start = time.perf_counter() with self.vae_mesh, nn_partitioning.axis_rules(self.vae_logical_axis_rules): - encoded_output = self.vae.encode(video_condition, self.vae_cache)[0].mode() + graphdef, state, rest_of_state = nnx.split(self.vae, nnx.Param, ...) + encoded_output = vae_encode_pass(graphdef, state, rest_of_state, video_condition) if hasattr(encoded_output, "block_until_ready"): encoded_output.block_until_ready() @@ -791,10 +899,8 @@ def _decode_latents_to_video(self, latents: jax.Array, trace: Optional[dict] = N """Decodes latents to video frames and postprocesses.""" t_vae_tpu_start = time.perf_counter() with self.vae_mesh, nn_partitioning.axis_rules(self.vae_logical_axis_rules): - video = self.vae.decode(latents, self.vae_cache)[0] - video = (video / 2.0) + 0.5 - video = jnp.clip(video, 0.0, 1.0) - video = (video * 255.0).astype(jnp.uint8) + graphdef, state, rest_of_state = nnx.split(self.vae, nnx.Param, ...) + video = vae_decode_pass(graphdef, state, rest_of_state, latents) video.block_until_ready() if trace is not None: trace["vae_decode_tpu"] = time.perf_counter() - t_vae_tpu_start @@ -1138,6 +1244,8 @@ def _prepare_model_inputs( batch_size = len(prompt) if prompt is not None else prompt_embeds.shape[0] // num_videos_per_prompt + debug_timers = bool(os.environ.get("WAN_DEBUG_COND_TIMERS")) + t_probe = time.perf_counter() with jax.named_scope("Encode-Prompt"): prompt_embeds, negative_prompt_embeds = self.encode_prompt( prompt=prompt, @@ -1146,6 +1254,9 @@ def _prepare_model_inputs( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, ) + if debug_timers: + max_logging.log(f"[cond] encode_prompt {time.perf_counter() - t_probe:.1f}s") + t_probe = time.perf_counter() num_channel_latents = self._get_num_channel_latents() if latents is None: @@ -1158,6 +1269,9 @@ def _prepare_model_inputs( num_frames=num_frames, num_channels_latents=num_channel_latents, ) + if debug_timers: + max_logging.log(f"[cond] prepare_latents {time.perf_counter() - t_probe:.1f}s") + t_probe = time.perf_counter() data_sharding = NamedSharding(self.mesh, P()) # Using global_batch_size_to_train_on so not to create more config variables @@ -1167,12 +1281,18 @@ def _prepare_model_inputs( latents = jax.device_put(latents, data_sharding) prompt_embeds = jax.device_put(prompt_embeds, data_sharding) negative_prompt_embeds = jax.device_put(negative_prompt_embeds, data_sharding) + if debug_timers: + jax.block_until_ready([latents, prompt_embeds, negative_prompt_embeds]) + max_logging.log(f"[cond] device_put {time.perf_counter() - t_probe:.1f}s") + t_probe = time.perf_counter() scheduler_state = self.scheduler.set_timesteps( self.scheduler_state, num_inference_steps=num_inference_steps, shape=latents.shape, ) + if debug_timers: + max_logging.log(f"[cond] set_timesteps {time.perf_counter() - t_probe:.1f}s") return ( latents, @@ -1189,7 +1309,7 @@ def __call__(self, **kwargs): @partial( - jax.jit, + aot_cache.cached_jit, static_argnames=( "do_classifier_free_guidance", "guidance_scale", @@ -1246,7 +1366,30 @@ def transformer_forward_pass( return noise_pred, latents -@partial(jax.jit, static_argnames=("guidance_scale",)) +@aot_cache.cached_jit +def vae_encode_pass(graphdef, state, rest_of_state, video): + """Encodes conditioning video to its deterministic latent (I2V path).""" + wan_vae = nnx.merge(graphdef, state, rest_of_state) + return wan_vae.encode(video, AutoencoderKLWanCache(wan_vae), return_dict=False)[0].mode() + + +@aot_cache.cached_jit +def vae_decode_pass(graphdef, state, rest_of_state, latents): + """Decodes latents and postprocesses to uint8 frames as ONE executable. + + Wrapped in the AOT cache so warm processes deserialize instead of paying + the deep conv-stack trace + compile-cache lookup, and zero-exec warmup + skips the decode compute entirely. The feat cache is rebuilt from the + merged module — it is pure per-call temporary state. + """ + wan_vae = nnx.merge(graphdef, state, rest_of_state) + video = wan_vae.decode(latents, AutoencoderKLWanCache(wan_vae), return_dict=False)[0] + video = (video / 2.0) + 0.5 + video = jnp.clip(video, 0.0, 1.0) + return (video * 255.0).astype(jnp.uint8) + + +@partial(aot_cache.cached_jit, static_argnames=("guidance_scale",)) def transformer_forward_pass_full_cfg( graphdef, sharded_state, @@ -1287,7 +1430,7 @@ def transformer_forward_pass_full_cfg( return noise_pred_merged, noise_cond, noise_uncond -@partial(jax.jit, static_argnames=("guidance_scale",)) +@partial(aot_cache.cached_jit, static_argnames=("guidance_scale",)) def transformer_forward_pass_cfg_cache( graphdef, sharded_state, diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py index 85daec33a..54629d181 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py @@ -16,9 +16,11 @@ from ...models.wan.transformers.transformer_wan import WanModel from typing import List, Union, Optional from ...pyconfig import HyperParameters +import concurrent.futures from functools import partial from flax import nnx from flax.linen import partitioning as nn_partitioning +from jax.sharding import Mesh import jax import jax.numpy as jnp from ...schedulers.scheduling_unipc_multistep_flax import FlaxUniPCMultistepScheduler @@ -44,22 +46,37 @@ def _load_and_init( load_transformer=True, load_scheduler=True, ): - common_components = cls._create_common_components( + # Load VAE/tokenizer/text-encoder/scheduler in a background thread while + # the main thread converts the 14B transformer: the small components and + # the text-encoder torch.compile warmup (compile_text_encoder) are hidden + # behind the transformer conversion time. The mesh/rngs built here are + # deterministic duplicates of the ones _create_common_components builds + # (same devices, same seed). + common_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) + common_future = common_executor.submit( + cls._create_common_components, config, load_vae=load_vae, load_text_encoder=load_text_encoder, load_scheduler=load_scheduler, ) transformer = None - if load_transformer: - transformer = super().load_transformer( - devices_array=common_components["devices_array"], - mesh=common_components["mesh"], - rngs=common_components["rngs"], - config=config, - restored_checkpoint=restored_checkpoint, - subfolder="transformer", - ) + try: + if load_transformer: + devices_array = max_utils.create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + rngs = nnx.Rngs(jax.random.key(config.seed)) + transformer = super().load_transformer( + devices_array=devices_array, + mesh=mesh, + rngs=rngs, + config=config, + restored_checkpoint=restored_checkpoint, + subfolder="transformer", + ) + common_components = common_future.result() + finally: + common_executor.shutdown(wait=True) pipeline = cls( tokenizer=common_components["tokenizer"], @@ -273,13 +290,15 @@ def run_inference_2_1( transformer_obj = nnx.merge(graphdef, sharded_state, rest_of_state) # Compute RoPE once as it only depends on shape - dummy_hidden_states = jnp.zeros(( - latents.shape[0], - latents.shape[2], - latents.shape[3], - latents.shape[4], - latents.shape[1], - )) + dummy_hidden_states = jnp.zeros( + ( + latents.shape[0], + latents.shape[2], + latents.shape[3], + latents.shape[4], + latents.shape[1], + ) + ) rotary_emb = transformer_obj.rope(dummy_hidden_states) kv_cache = None diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_animate.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_animate.py index ffd87ddfd..8a63e8c3e 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline_animate.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_animate.py @@ -31,7 +31,6 @@ from functools import partial from typing import List, Optional, Tuple, Union -import flax import flax.linen as nn import jax import jax.numpy as jnp @@ -43,13 +42,13 @@ from jax.sharding import NamedSharding, PartitionSpec as P from maxdiffusion import max_logging from maxdiffusion.image_processor import PipelineImageInput, VaeImageProcessor -from maxdiffusion.max_utils import device_put_replicated, get_flash_block_sizes, get_precision +from maxdiffusion.max_utils import get_flash_block_sizes, get_precision from maxdiffusion.video_processor import VideoProcessor from ...models.wan.transformers.transformer_wan_animate import WanAnimateTransformer3DModel from ...models.wan.wan_utils import load_wan_animate_transformer from ...pyconfig import HyperParameters -from .wan_pipeline import WanPipeline, cast_with_exclusion +from .wan_pipeline import WanPipeline, cast_with_exclusion, put_params_into_state def create_sharded_animate_transformer( @@ -124,11 +123,14 @@ def _create_model(rngs: nnx.Rngs, wan_config: dict): params = jax.tree_util.tree_map_with_path( lambda path, x: cast_with_exclusion(path, x, dtype_to_cast=config.weights_dtype), params ) - for path, val in flax.traverse_util.flatten_dict(params).items(): - if restored_checkpoint: - path = path[:-1] - sharding = logical_state_sharding[path].value - state[path].value = device_put_replicated(val, sharding) + state = put_params_into_state( + state, + params, + logical_state_sharding, + mesh, + restored_checkpoint=restored_checkpoint, + subfolder=subfolder, + ) state = nnx.from_flat_state(state) wan_transformer = nnx.merge(graphdef, state, rest_of_state) diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py index 3863b05c4..f156fedd7 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py @@ -967,11 +967,11 @@ def scan_body(carry, t): latent_model_input = jnp.concatenate([latents_input, condition], axis=-1) timestep = jnp.broadcast_to(t, latents_input.shape[0]) - use_high_noise = jnp.greater_equal(t, boundary) - noise_pred, _ = jax.lax.cond( - use_high_noise, - high_noise_branch, - low_noise_branch, + # Timesteps are host-known: Python dispatch (like the T2V loop) avoids + # tracing both 14B branches per step and keeps the AOT cache usable. + use_high_noise = bool(np.asarray(scheduler_state.timesteps)[step] >= np.asarray(boundary)) + branch = high_noise_branch if use_high_noise else low_noise_branch + noise_pred, _ = branch( ( latent_model_input, timestep, @@ -982,7 +982,7 @@ def scan_body(carry, t): rotary_emb, encoder_attention_mask_high, encoder_attention_mask_low, - ), + ) ) noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1)) latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() diff --git a/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py b/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py index dcdf9396d..2f0adbf0f 100644 --- a/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py +++ b/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py @@ -18,19 +18,19 @@ import jax import jax.numpy as jnp from jax.sharding import Mesh, NamedSharding, PartitionSpec as P -import flax import flax.linen as nn from flax import nnx from flax.linen import partitioning as nn_partitioning from ...pyconfig import HyperParameters +from ... import aot_cache from ... import max_logging from ...image_processor import PipelineImageInput -from ...max_utils import get_flash_block_sizes, get_precision, device_put_replicated +from ...max_utils import get_flash_block_sizes, get_precision from ...models.wan.wan_utils import load_wan_transformer from ...models.wan.transformers.transformer_wan_vace import WanVACEModel from ...schedulers.scheduling_unipc_multistep_flax import FlaxUniPCMultistepScheduler from ...models.modeling_flax_pytorch_utils import torch2jax -from .wan_pipeline import cast_with_exclusion +from .wan_pipeline import _final_param_dtype, cast_with_exclusion, converted_weights_cache_dir, put_params_into_state from .wan_pipeline_2_1 import WanPipeline2_1 import torch import PIL @@ -123,31 +123,23 @@ def create_model(rngs: nnx.Rngs, wan_config: dict): num_layers=wan_config["num_layers"], scan_layers=config.scan_layers, subfolder=subfolder, + cast_dtype_fn=partial(_final_param_dtype, dtype_to_cast=config.weights_dtype), + converted_cache_dir=converted_weights_cache_dir(config, subfolder), ) + # No-op (returns leaves unchanged) when the loader already cast to the + # final dtypes; still needed for restored orbax checkpoints. params = jax.tree_util.tree_map_with_path( lambda path, x: cast_with_exclusion(path, x, dtype_to_cast=config.weights_dtype), params ) - for path, val in flax.traverse_util.flatten_dict(params).items(): - if restored_checkpoint: - if path[-1] == "value": - path = path[:-1] # remove 'value' - - try: - # Convert block indices to integers, as they might have been loaded as strings from the checkpoint. - path = path[:1] + (int(path[1]),) + path[2:] - except Exception: - pass - - sharding = logical_state_sharding[path].value - try: - state[path].value = device_put_replicated(val, sharding) - except Exception as e: - max_logging.log(f"Failed to device_put_replicated {path}: {e}") - max_logging.log(f"Trying to use process_allgather for {path}") - val_on_host = jax.experimental.multihost_utils.process_allgather(val, tiled=True) - state[path].value = device_put_replicated(val_on_host, sharding) - del val_on_host + state = put_params_into_state( + state, + params, + logical_state_sharding, + mesh, + restored_checkpoint=restored_checkpoint, + subfolder=subfolder, + ) state = nnx.from_flat_state(state) wan_transformer = nnx.merge(graphdef, state, rest_of_state) @@ -716,7 +708,7 @@ def prepare_video_latents( return jnp.stack(latent_list) -@partial(jax.jit, static_argnames=("do_classifier_free_guidance", "guidance_scale")) +@partial(aot_cache.cached_jit, static_argnames=("do_classifier_free_guidance", "guidance_scale")) def transformer_forward_pass( graphdef: nnx.graph.GraphDef, sharded_state: nnx.State, diff --git a/src/maxdiffusion/tests/aot_cache_test.py b/src/maxdiffusion/tests/aot_cache_test.py new file mode 100644 index 000000000..d0f5e2d9f --- /dev/null +++ b/src/maxdiffusion/tests/aot_cache_test.py @@ -0,0 +1,159 @@ +"""Tests for the per-shape AOT executable cache (CPU backend).""" + +import functools +import os +import tempfile +import unittest + +os.environ.setdefault("JAX_PLATFORMS", "cpu") + +import jax +import jax.numpy as jnp +import numpy as np +from jax.sharding import Mesh + +from maxdiffusion import aot_cache + + +@functools.partial(aot_cache.cached_jit, static_argnames=("flag",)) +def _toy_fn(x, y, flag=False): + return x @ y + (1.0 if flag else 0.0) + + +class AotCacheTest(unittest.TestCase): + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self._mesh = Mesh(np.array(jax.devices()[:1]), ("d",)) + self._a = jnp.ones((8, 8)) + self._b = jnp.eye(8) + # Reset process-global install state between tests. + aot_cache._STATE.enabled = False + for entry in aot_cache._REGISTRY: + entry._compiled.clear() + entry._pending.clear() + + def tearDown(self): + aot_cache._STATE.enabled = False + self._tmp.cleanup() + + def _install(self): + aot_cache.install(self._tmp.name, meta={"m": 1}, mesh=self._mesh) + aot_cache.wait_for_loads() + + def test_disabled_is_plain_jit(self): + result = _toy_fn(self._a, self._b, True) + np.testing.assert_allclose(result, self._a @ self._b + 1.0) + self.assertFalse(aot_cache._STATE.enabled) + + def test_record_save_hit_roundtrip(self): + self._install() + first = _toy_fn(self._a, self._b, True) # miss -> jit + record + self.assertEqual(aot_cache.save_pending(), 1) + hit = _toy_fn(self._a, self._b, True) # compiled hit + np.testing.assert_allclose(np.asarray(first), np.asarray(hit)) + + def test_reload_from_disk(self): + self._install() + first = _toy_fn(self._a, self._b, True) + aot_cache.save_pending() + entry = next(e for e in aot_cache._REGISTRY if e.name.endswith("._toy_fn")) + entry._compiled.clear() # simulate a fresh process + entry.load_from_disk() + self.assertTrue(entry._compiled) + reloaded = _toy_fn(self._a, self._b, True) + np.testing.assert_allclose(np.asarray(first), np.asarray(reloaded)) + + def test_static_value_gets_own_executable(self): + self._install() + _toy_fn(self._a, self._b, True) + aot_cache.save_pending() + # Different static value -> different signature -> jit fallback, correct. + off = _toy_fn(self._a, self._b, False) + np.testing.assert_allclose(off, self._a @ self._b) + + def test_new_shape_falls_back_and_saves(self): + self._install() + _toy_fn(self._a, self._b, True) + self.assertEqual(aot_cache.save_pending(), 1) + small = _toy_fn(jnp.ones((4, 8)), jnp.ones((8, 4)), True) + self.assertEqual(small.shape, (4, 4)) + self.assertEqual(aot_cache.save_pending(), 1) # only the new shape + + def test_warmup_mode_compiles_without_executing(self): + self._install() + with aot_cache.warmup_mode(): + warm = _toy_fn(self._a, self._b, True) + # Zeros prove the fn body never ran (real output would be a@b+1). + self.assertEqual(warm.shape, (8, 8)) + np.testing.assert_allclose(np.asarray(warm), np.zeros((8, 8))) + # The signature was compiled during warmup and is serializable. + self.assertEqual(aot_cache.save_pending(), 1) + # Outside warmup mode the compiled executable returns real values. + real = _toy_fn(self._a, self._b, True) + np.testing.assert_allclose(np.asarray(real), self._a @ self._b + 1.0) + + def test_warmup_mode_disabled_cache_executes_normally(self): + with aot_cache.warmup_mode(): # cache not installed -> no-op + result = _toy_fn(self._a, self._b, True) + np.testing.assert_allclose(np.asarray(result), self._a @ self._b + 1.0) + + def test_traced_call_inlines_like_nested_jit(self): + # I2V's denoise loop invokes wrapped fns inside lax.cond branches; a + # deserialized executable cannot be applied to tracers. The wrapper + # must inline (plain nested-jit behavior) and never crash or record. + self._install() + _toy_fn(self._a, self._b, True) + aot_cache.save_pending() # compiled entry exists for this signature + + def branch_true(x): + return _toy_fn(x, self._b, True) + + def branch_false(x): + return x + + result = jax.lax.cond(True, branch_true, branch_false, self._a) + np.testing.assert_allclose(np.asarray(result), self._a @ self._b + 1.0) + self.assertEqual(aot_cache.save_pending(), 0) # nothing recorded + + def test_signature_deterministic_across_processes(self): + # Signatures live in filenames; a process-dependent component (e.g. + # object addresses inside a GraphDef repr) would make every restart + # miss its own cache. Compute the same signature in two interpreters. + import subprocess + import sys + + snippet = "\n".join( + ( + "import os", + "os.environ['JAX_PLATFORMS'] = 'cpu'", + "import jax", + "import jax.numpy as jnp", + "from flax import nnx", + "from maxdiffusion import aot_cache", + "", + "class T(nnx.Module):", + " def __init__(self, rngs):", + " self.lin = nnx.Linear(4, 4, rngs=rngs)", + "", + "graphdef, state = nnx.split(T(nnx.Rngs(0)))", + "sig = aot_cache._dynamic_signature(", + " (graphdef, state.to_pure_dict(), jnp.ones((2, 4))), {})", + "print(sig)", + ) + ) + outs = [ + subprocess.run( + [sys.executable, "-c", snippet], + capture_output=True, + text=True, + check=True, + env={**os.environ, "JAX_PLATFORMS": "cpu"}, + ).stdout.strip() + for _ in range(2) + ] + self.assertEqual(outs[0], outs[1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/converted_weights_cache_test.py b/src/maxdiffusion/tests/converted_weights_cache_test.py new file mode 100644 index 000000000..a3d0e4937 --- /dev/null +++ b/src/maxdiffusion/tests/converted_weights_cache_test.py @@ -0,0 +1,80 @@ +"""Tests for the memoized torch->flax converted-weights cache.""" + +import os +import tempfile +import unittest + +import ml_dtypes +import numpy as np + +from maxdiffusion.models.wan.wan_utils import save_converted_weights, try_load_converted_weights + + +def _flat_tree(): + return { + ("blocks", "attn1", "kernel"): np.arange(24, dtype=np.float32).reshape(2, 3, 4), + ("proj_out", 0, "bias"): np.ones(5, dtype=np.float16), + # bf16 exercises the uint bit-view path (npy mmap cannot resolve + # ml_dtypes descriptors directly). + ("blocks", "ffn", "kernel"): np.arange(12, dtype=np.float32).astype(ml_dtypes.bfloat16).reshape(3, 4), + } + + +def _eval_shapes(flat): + shapes = {} + for key, value in flat.items(): + node = shapes + for part in key[:-1]: + node = node.setdefault(part, {}) + node[key[-1]] = value # only keys/structure are validated + return shapes + + +class ConvertedWeightsCacheTest(unittest.TestCase): + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.cache_dir = os.path.join(self._tmp.name, "cache") + self.flat = _flat_tree() + self.eval_shapes = _eval_shapes(self.flat) + + def tearDown(self): + self._tmp.cleanup() + + def test_round_trip(self): + save_converted_weights(self.cache_dir, self.flat) + loaded = try_load_converted_weights(self.cache_dir, self.eval_shapes, None) + self.assertIsNotNone(loaded) + np.testing.assert_array_equal(loaded["blocks"]["attn1"]["kernel"], self.flat[("blocks", "attn1", "kernel")]) + np.testing.assert_array_equal(loaded["proj_out"][0]["bias"], self.flat[("proj_out", 0, "bias")]) + bf16 = loaded["blocks"]["ffn"]["kernel"] + self.assertEqual(bf16.dtype, np.dtype(ml_dtypes.bfloat16)) + np.testing.assert_array_equal(bf16.view(np.uint16), self.flat[("blocks", "ffn", "kernel")].view(np.uint16)) + + def test_missing_cache_returns_none(self): + self.assertIsNone(try_load_converted_weights(self.cache_dir, self.eval_shapes, None)) + + def test_dtype_policy_change_invalidates(self): + save_converted_weights(self.cache_dir, self.flat) + loaded = try_load_converted_weights(self.cache_dir, self.eval_shapes, lambda key: np.dtype(np.float64)) + self.assertIsNone(loaded) + + def test_key_set_change_invalidates(self): + save_converted_weights(self.cache_dir, self.flat) + bigger = dict(self.flat) + bigger[("new_param", "kernel")] = np.zeros(2, dtype=np.float32) + loaded = try_load_converted_weights(self.cache_dir, _eval_shapes(bigger), None) + self.assertIsNone(loaded) + + def test_dtypes_preserved(self): + save_converted_weights(self.cache_dir, self.flat) + loaded = try_load_converted_weights(self.cache_dir, self.eval_shapes, None) + for key, value in self.flat.items(): + node = loaded + for part in key: + node = node[part] + self.assertEqual(node.dtype, value.dtype) + + +if __name__ == "__main__": + unittest.main()