FastH3-Omni-MLX / bench /README.md
aryan5v's picture
Build + benchmark kit for FastH3 Omni on MLX
0ac7396 verified
|
Raw History Blame Contribute Delete
7.35 kB

FastH3 Omni (Ref2VA) MLX benchmark matrix

Runs the FastH3 OmniRef pipeline on Apple Silicon over the same workload matrix as the DGX Spark benchmark, and writes results in the Spark schema so the two can be compared. You only install and run it; nothing needs editing.

Files here: run_matrix.py (the harness), aggregate.py (summary writer, imported by the harness, also runnable alone), this README.

What it measures

  • Compositions (all 5 s, 124 frames, 24 fps): image, image-voice, images3, clip, video (also accepted: mixed, clips3, images9).
  • Decoders: 36l (quality 36-layer) and 8l (fast 8-layer).
  • Canvases: 480p = 832x480, 768p = 1344x768.
  • Per cell: one excluded warmup at seed 42, then measured seeds 0, 1, 2, one request at a time. Latency is wall time of pipeline.generate(...) through the finished MP4, with model load excluded. Mean and sample SD over seeds.
  • Default order: all 480p cells first, 36L before 8L, compositions in the order above (20 cells). Each cell runs in a fresh subprocess so memory never accumulates between cells.

Naming note: columns such as nvfp4_36l_480p_mean_s keep the Spark spelling. On a Mac the NVFP4 VAE decoders run as W4A16 (4-bit packed weights, bf16 activations) and the DiT is the INT6 conversion, not native FP4 hardware math.

Requirements

  • Apple Silicon Mac, macOS, Python 3.10+ environment that can import mlx.core and import fastvideo (FastVideo branch aryan/omniref-mlx-speed). If FastVideo is a checkout that is not pip installed, pass --fastvideo-root /path/to/FastVideo (or set FASTVIDEO_ROOT).

  • ffmpeg and ffprobe on PATH (brew install ffmpeg).

  • The model directory with this layout (this is the layout of the Hugging Face repo):

    <model-dir>/root/                      composed model root (root/text_encoder -> ../text_encoder-nvfp4)
    <model-dir>/dit/int6/                  mlx_h3_dit.json + mlx_h3_dit.safetensors
    <model-dir>/text_encoder-nvfp4/
    <model-dir>/vae/fasth3-quality-36L-decoder-nvfp4.safetensors
    <model-dir>/vae/fasth3-fast-8L-decoder-nvfp4.safetensors
    <model-dir>/vae/lightvae-lite-h3-encoder.safetensors
    

    If bench/ sits inside that repo, --model-dir defaults to the repo root. If the root/text_encoder symlink is missing, the harness automatically uses <model-dir>/text_encoder-nvfp4.

  • The data bundle omniref-handoff/ (workload manifests, reference media, SHA256SUMS). Pass --data-root or set OMNIREF_DATA (default ~/omniref-handoff). Every file is verified against SHA256SUMS at start.

The harness itself is pure standard library; run it with the Python that has mlx and fastvideo installed (that interpreter is also used for the cell subprocesses; override with --python).

Commands

  1. Check everything without a GPU (manifests parse, CSV schema, paths):

    python3 run_matrix.py --self-test --data-root ~/omniref-handoff
    python3 run_matrix.py --dry-run --model-dir /path/to/model --data-root ~/omniref-handoff
    
  2. Smoke test (a few minutes): one image request, seed 42, 36L, 448x256, with reference canvases capped to the smoke canvas (--reference-max-pixels match; pass your own value to override).

    python3 run_matrix.py --smoke --model-dir /path/to/model --data-root ~/omniref-handoff
    

    Success: the last lines say COMPLETE, a SEND_BACK.tar.gz is written, and results-<host>-<date>-smoke/media/*/*-smoke.mp4 plays with audio. Other smoke compositions: --smoke --compositions image-voice,clip.

  3. Full matrix (long; safe to interrupt and re-run):

    python3 run_matrix.py --model-dir /path/to/model --data-root ~/omniref-handoff --out ./results-mystudio
    
  4. Subsets:

    python3 run_matrix.py ... --canvases 480p --decoders 36l              # one lane
    python3 run_matrix.py ... --compositions image,clip --seeds 0         # fewer cells / seeds
    python3 run_matrix.py ... --no-warmup                                 # skip the seed-42 warmup
    python3 run_matrix.py ... --dit int8                                  # other DiT conversion under dit/
    python3 run_matrix.py ... --reference-max-pixels match                # cap reference canvases (default: trained canvas)
    

When done, send back SEND_BACK.tar.gz from the --out directory.

Resuming, failures and Ctrl-C

  • Re-run the exact same command: finished seeds and cells are skipped; an interrupted cell resumes with a fresh process (it re-runs the warmup first so the measured seeds still start warm).
  • A cell that crashes or runs out of memory is recorded as failed (error tail and signal in per-request-records.jsonl, full output in logs/<cell>.log) and the matrix continues. Failed cells are not retried on re-run unless you pass --retry-failed.
  • Ctrl-C (or SIGTERM) stops the running cell, releases the GPU lock, rebuilds the summary and bundle, and exits with code 130.
  • A cell is killed after --cell-timeout-s (default 4 h).

GPU lock

Before any model work the harness takes /tmp/fh3mlx/gpu.lock (atomic mkdir, retry every 30 s, owner PID written to owner) and removes it when done, including on errors and Ctrl-C. If a crash left a lock whose owner PID is dead, pass --steal-stale-lock or rm -rf /tmp/fh3mlx/gpu.lock.

Memory settings

--vae-dtype bf16; --metal-wired-limit-gib auto (0.66 x RAM, capped at 64 GiB; off disables); --mlx-cache-limit-gib 8 (applied through the pipeline's mx.set_cache_limit; negative leaves it unset). A background sampler records MLX active/cache/peak memory and sysctl vm.swapusage every 2 s per request (media/<cell>/*.mem.jsonl, summary in each record). Close other heavy apps for clean numbers.

Output directory

per-request-records.jsonl   every warmup/measured request: wall_s, pipeline timings, peak_memory_gib,
                            memory-sampler summary, ffprobe of the MP4, output/prompt/condition/weight hashes
latency.csv                 Spark column schema (nvfp4_{36l,8l}_{480p,768p}_{mean_s,sample_std_s,status});
                            custom canvases are appended as extra columns
benchmark.json              profiles + performance (Spark benchmark.json layout)
environment_receipt.json    chip, memory, macOS, MLX/Python versions, mx.device_info(), FastVideo git sha,
                            weight-manifest hashes (first 64 MB + size), data-bundle identity, settings
README.md                   generated results table (mean ± sample SD) and peak memory table
media/  logs/  status/      MP4s + memory traces; one log per cell; per-cell state
SEND_BACK.tar.gz            everything above (MP4s trimmed to one seed-0 per cell if the total exceeds ~2 GB)

Rebuild the summary files at any time with python3 aggregate.py --out <results-dir>.

M5 fastest-config sweep

--m5-sweep runs the image composition at 480p with the 36L decoder, seed 0 only (plus the excluded warmup), for DiT int6 and int8 x --quant-gemm {dequant,qmm}, and for bf16 when <model-dir>/dit/bf16 exists (bf16 ignores the GEMM path). Cell ids carry -dit-<fmt>[-gemm-<mode>]; results are ranked in m5-sweep.md / m5-sweep.csv. Resume and the GPU lock behave as in the full matrix. --dit bf16 and --quant-gemm also work on ordinary runs.