diff --git a/optimized/tensorRT/README.md b/optimized/tensorRT/README.md index bb95beaa..0a84dcfd 100644 --- a/optimized/tensorRT/README.md +++ b/optimized/tensorRT/README.md @@ -137,21 +137,28 @@ Omit `--dit` / `--decoder` for an interactive arrow-key picker. Relative before QK^T fixed that (**4.3× faster at L=4096**), which retired the reason to prefer `bf16`. Engines built before 2026-07 are the slow variant; rebuild with `build_from_onnx.py sa3-m`. -- **`fp8`** — *medium only; the max-speed clean tier.* fp8 E4M3 on the 176 linear - GEMMs + bf16 fused FMHA (96 nodes) + a **baked fp32 RoPE constant table** (position +- **`fp8`** — *medium only; the max-speed clean tier, calibrated.* fp8 E4M3 on the 176 + linear GEMMs + bf16 fused FMHA (96 nodes) + a **baked fp32 RoPE constant table** (position cos/sin computed host-side at build and frozen as a graph constant — no in-graph trig, so the island is precision-policy-robust and cross-runtime-stable). **~1.3× faster than `fp16mixed` at every length** (H200, same-run round-robin: 1.40× @L129 / 1.32× @L1292 / 1.34× @L4092 — also ahead of `bf16`) and **clean at long sequence** (latent std 0.86 vs eager 0.95, **0.000% clip** at 2-min and 6-min), so it stays clean exactly where `bf16` - clips. It is a **speed tier over an already-good default, not a fidelity upgrade**: - single-step velocity cos vs the FP32 engine is ~0.92–0.97 (below fp16mixed's ~1.0) but - the 8-step render stays coherent. **Capped at L≤4096** — the baked table is sized to the - profile max, which is *also* the SAME-L decoder's own cap, so this is not a new - end-to-end limit; longer renders are rejected (a re-bake would be needed). **Not - seed-reproducible vs fp16-mixed.** Built weakly-typed `EXPLICIT_BATCH` + `BF16` + `FP8` + - `OBEY_PRECISION_CONSTRAINTS` (rebuild: `build_from_onnx.py sa3-m-fp8`; producer: - `build/make_rope_baked_onnx.py`, identity check: `scripts/verify_fp8_rope.py`). + clips. The fp8 scales are **calibrated on real conditioning** (per-tensor activation amax + + per-channel weight scales from @ryanontheinside's #47, captured via `make_calib.py`) — a + **speed-free** change (31.2 vs the earlier uncalibrated 30.8 ms/fwd) that lifts worst-step + velocity-cos vs fp32 on adversarial seeds **0.52/0.57/0.64 → 0.92/0.94/0.92** (steps 1–7 + match the fully-calibrated reference within ~0.001). It remains a **speed tier over an + already-good default, not a fidelity upgrade over it**: single-step velocity cos ~0.92–0.94 + (below fp16mixed's ~1.0) but the 8-step render stays coherent — it just no longer collapses + at the highest-noise first step the way the uncalibrated engine did. **Capped at L≤4096** — + the baked table is sized to the profile max, which is *also* the SAME-L decoder's own cap, + so this is not a new end-to-end limit; longer renders are rejected (a re-bake would be + needed). **Not seed-reproducible vs fp16-mixed.** Built weakly-typed `EXPLICIT_BATCH` + + `BF16` + `FP8` + `OBEY_PRECISION_CONSTRAINTS` (rebuild: `build_from_onnx.py sa3-m-fp8`; + producer: `build/build_dit_bf16.py` RoPE-baker + `build/transplant_scales.py` calibrated-scale + transplant; identity check: `scripts/verify_fp8_rope.py`; calibration by @ryanontheinside, + [#47](https://github.com/Stability-AI/stable-audio-3/pull/47)). - **`bf16`** — *medium only.* Same `dit.onnx` as fp32, built with `BuilderFlag.BF16`; a uniform bf16 trunk also lets the FMHA fuser fire, and it is ~3% faster than `fp16mixed`. **But it drifts at long sequence**: weakly-typed BF16 lets TRT diff --git a/optimized/tensorRT/build/README.md b/optimized/tensorRT/build/README.md index f78d632d..4c23938f 100644 --- a/optimized/tensorRT/build/README.md +++ b/optimized/tensorRT/build/README.md @@ -251,39 +251,53 @@ git commit -m "Refresh canonical ONNX" git push ``` -## Medium `fp8` — the max-speed RoPE-baked engine +## Medium `fp8` — the max-speed RoPE-baked engine, calibrated On top of the `fp16mixed` default and the selectable `bf16`, the medium DiT ships an **`fp8`** engine (`dit_fp8.trt`, `--precision fp8`): fp8 E4M3 on the 176 linear GEMMs + 96 bf16 fused FMHA + the **same baked fp32 RoPE constant table** as bf16. Measured on H200 it is ~1.3× faster than `fp16mixed` at every length (1.40× / 1.32× / 1.34× @L129/1292/4092, also ahead of bf16) and clean at long sequence (latent std 0.86, 0.000% clip at 6-min), so it stays clean -exactly where bf16 clips. It is a speed tier over the default, **not** a fidelity upgrade -(single-step velocity cos vs fp32 ~0.92–0.97 < fp16mixed's ~1.0; the 8-step render stays -coherent). See the runtime README's precision section for positioning. +exactly where bf16 clips. + +The fp8 scales are **calibrated on real conditioning** (updated 2026-07-31). The scale VALUES — +per-tensor activation amax + per-channel weight scales — come from @ryanontheinside's fp8 work +([#47](https://github.com/Stability-AI/stable-audio-3/pull/47)): captured with his `make_calib.py` +and grafted onto this baked / bf16-fused structure by `transplant_scales.py`. Calibration is +**speed-free** (31.2 vs the earlier uncalibrated 30.8 ms/fwd, inside run noise) and lifts +worst-case per-step fidelity — worst-step velocity-cos vs fp32 on adversarial seeds goes +**0.52/0.57/0.64 → 0.92/0.94/0.92**, with sampling steps 1–7 tracking the fully-calibrated #47 +reference within ~0.001. It is still a **speed tier over the `fp16mixed` default, not a fidelity +upgrade over it** (single-step velocity cos ~0.92–0.94 < fp16mixed's ~1.0; the 8-step render +stays coherent), but it no longer collapses at the highest-noise first step the way the +uncalibrated engine did. See the runtime README's precision section for positioning. **Consumer (per-arch rebuild).** The `.trt` is `sm_90`-specific; `sm_89` / `sm_120` / `sm_100` are a rebuild away (run on the target GPU): ```bash -python build_from_onnx.py sa3-m-fp8 # pulls dit_fp8.onnx (+ dit_fp8lin.onnx.data) from HF +python build_from_onnx.py sa3-m-fp8 # pulls the calibrated dit_fp8.onnx (+ dit_fp8lin.onnx.data) from HF ``` `sa3-m-fp8` is the `sa3-m-bf16` recipe **plus `BuilderFlag.FP8`**: weakly-typed `EXPLICIT_BATCH` + `BF16` + `FP8` + `OBEY_PRECISION_CONSTRAINTS`, reusing the same `_pin_fourier_fp32` island (on the RoPE-baked ONNX the only `Cos`/`Sin` left are the two -runtime Fourier chains). The fp8 E4M3 Quantize/Dequantize pairs ride in the ONNX, so TRT fires -fp8 tensor-core GEMMs on the Linears while attention stays bf16 fused-MHA. Identity of the -shipped engine — **176 fp8 GEMMs + 96 bf16 fused MHA + fp32 RoPE constant** — verify with +runtime Fourier chains). The fp8 E4M3 Quantize/Dequantize pairs — carrying the calibrated +scales — ride in the ONNX, so TRT fires fp8 tensor-core GEMMs on the Linears while attention +stays bf16 fused-MHA. The build recipe is **unchanged by calibration**: calibration lives in the +ONNX scale values, not the builder flags. Identity of the shipped engine — **176 fp8 GEMMs + +96 bf16 fused MHA + baked fp32 RoPE constant** — verify with `python ../scripts/verify_fp8_rope.py .trt` (needs a DETAILED-verbosity build, as the shipped engine is; a plain consumer rebuild renders identically but is not introspectable). -**Producer (refresh the ONNX).** RoPE-baking is the SAME step as bf16 — `build_dit_bf16.py` is -the shared baker (it handles both an external `inv_freq`, as in the fp32 `dit.onnx`, and an -inline one, as in the fp8-linear ONNX): +**Producer (refresh the ONNX).** Two independent pieces: RoPE-baking (structure) and +calibration (scale values). + +RoPE-baking is the SAME step as bf16 — `build_dit_bf16.py` is the shared baker (it handles both +an external `inv_freq`, as in the fp32 `dit.onnx`, and an inline one, as in the fp8-linear ONNX): ```bash -python build_dit_bf16.py --input dit_fp8lin.onnx --output dit_fp8.onnx --max-t 4160 +python build_dit_bf16.py --input dit_fp8lin.onnx --output dit_fp8lin_ropebaked.onnx --max-t 4160 ``` `--max-t` (4160 = profile max 4096 + 64 global tokens) sizes the baked table; rendering past @@ -291,6 +305,25 @@ L=4096 would need a larger `--max-t` **and** a matching TRT profile — the runt L>4096 up front, coinciding with the SAME-L decoder's own cap, so it is not a new end-to-end limit. +Then graft the calibrated fp8 scales onto the baked graph with `transplant_scales.py`. It matches +every quantized Linear by weight-initializer name and swaps only the scale VALUES (per-tensor +activation amax + per-channel weight scales); the 5.8 GB fp32 weights are untouched and TRT +re-quantizes them at build: + +```bash +python transplant_scales.py \ + --ours-onnx dit_fp8lin_ropebaked.onnx \ + --calib-onnx /dit_fp8_calib.onnx \ + --out dit_fp8.onnx # publish this as onnx/sa3-m/dit_fp8.onnx +``` + +The calibrated `--calib-onnx` comes from @ryanontheinside's fp8 pipeline +([#47](https://github.com/Stability-AI/stable-audio-3/pull/47)). A **from-scratch recalibration** +(only needed on a model retrain) reruns `make_calib.py` (real-conditioning capture, in this +repo) → `build_dit_fp8.py` (max-PTQ + per-channel weight scales; that builder lives in #47, not +merged here). Everyday consumers never recalibrate — they pull the published calibrated +`dit_fp8.onnx`. + ## File map | File | Role | Flow | @@ -300,7 +333,9 @@ limit. | `build_dit_profile.py` | Build a DiT with custom `(min, opt, max)` profile shapes (experimental — short-form / fixed-shape variants). Operates on either ONNX flavor. | consumer | | `build_dit_fp16mixed.py` | **Producer-side** ONNX surgery: takes the canonical FP32 `dit.onnx`, finds RMSNorm chains + attention `Softmax` + RoPE region, wraps each in `Cast(FP32) ↔ Cast(FP16)` islands, converts non-island weights to FP16, then bounds the RoPE island before QK^T (`bound_attention_core()`, `--no-bound-attn` to skip) so the attention core runs FP16 and TRT's FMHA fuser fires — 96/96 attentions on the medium DiT, 4.3× at L=4096. Writes both the modified `dit_fp16mixed.onnx` AND the TRT engine, which **must** be `STRONGLY_TYPED` (weakly-typed + `BuilderFlag.FP16` re-casts the FP32 islands and silently degrades to naive FP16). Only re-run when the model retrains or the island recipe changes. Requires `onnx` + `onnx-graphsurgeon`. | producer | | `build_dit_bf16.py` | **Producer-side** shared RoPE-baker for the medium `bf16` AND `fp8` engines: precomputes RoPE's cos/sin in fp64 on the host, freezes them as fp32 constant tables (`--max-t`), rewires the 96 trig sites and lets DCE delete the runtime angle chain — so the trunk runs bf16/fp8 without the long-angle drift. Weights are never loaded (keeps the input's `.data` sidecar). Handles both external `inv_freq` (fp32 `dit.onnx`) and inline (fp8-linear ONNX). Consumer compile: `build_from_onnx.py sa3-m-bf16` / `sa3-m-fp8`. Requires `onnx`. | producer | -| `../scripts/verify_fp8_rope.py` | EngineInspector identity check for the fp8 engine (176 fp8 GEMMs + 96 bf16 fused MHA + fp32 RoPE constant). Needs a DETAILED-verbosity build. | consumer | +| `make_calib.py` | **Producer-side** FP8 calibration capture: drives the model's own pingpong `generate()` and records the six DiT engine inputs at every sampling step into a `.npz` (real-conditioning, deployment-matched prompts). Feeds `build_dit_fp8.py` (#47). Only re-run for a from-scratch recalibration on a model retrain. Calibration tooling by @ryanontheinside (#47). Requires `torch` + `stable_audio_3`. | producer | +| `transplant_scales.py` | **Producer-side** fp8 scale transplant: grafts #47's calibrated activation + per-channel weight scales onto the RoPE-baked bakedmin ONNX (matches Linears by weight-initializer name, swaps only the scale VALUES; weights untouched). This is what makes the shipped `dit_fp8.onnx` calibrated while keeping bakedmin's bf16 fused-MHA speed. Scale values / calibration by @ryanontheinside (#47). Requires `onnx`. | producer | +| `../scripts/verify_fp8_rope.py` | EngineInspector identity check for the fp8 engine (176 fp8 GEMMs + 96 bf16 fused MHA + baked fp32 RoPE constant). Needs a DETAILED-verbosity build. | consumer | | `build_t5gemma.py` | Trace + export T5Gemma encoder ONNX + build TRT | producer | | `build_same_s_decoder.py` | Trace + export SAME-S decoder ONNX + build TRT | producer | | `build_same_s_encoder.py` | Trace + export SAME-S encoder ONNX + build TRT | producer | diff --git a/optimized/tensorRT/build/build_from_onnx.py b/optimized/tensorRT/build/build_from_onnx.py index 03872848..e59c6fb8 100755 --- a/optimized/tensorRT/build/build_from_onnx.py +++ b/optimized/tensorRT/build/build_from_onnx.py @@ -210,29 +210,43 @@ # still need a (tiny, ~0.1%) fp32 island. See _pin_fourier_fp32. "pin_fourier_fp32": True, }, - # SA3 medium DiT in fp8 — SELECTABLE, medium-only, the MAX-SPEED clean tier. - # Identical recipe to sa3-m-bf16 (same RoPE-baked ONNX flow from build_dit_bf16.py, - # same EXPLICIT_BATCH + OBEY + Fourier fp32 island) PLUS BuilderFlag.FP8: the ONNX - # carries fp8 E4M3 Quantize/Dequantize pairs on the 176 Linear GEMMs, so TRT fires - # fp8 tensor-core GEMMs there while the 96 attention blocks stay bf16 fused-MHA - # (_gemm_mha_v2) and the RoPE cos/sin ride as baked fp32 constants. Measured on - # H200 (same-run round-robin): ~1.3x faster than fp16-mixed at every length - # (1.40x@L129 / 1.32x@L1292 / 1.34x@L4092, also ahead of bf16) and clean at long - # sequence (latent std 0.86 vs eager 0.95, 0.000% clip @6-min) — the baked RoPE - # dodges bf16's long-angle drift, so it stays clean where bf16 clips. It is a SPEED - # tier over the fp16-mixed default, NOT a fidelity upgrade (single-step velocity cos - # vs the fp32 engine ~0.92-0.97 < fp16mixed's ~1.0; the 8-step render stays - # coherent). Identity of the shipped engine: 176 fp8 GEMMs + 96 bf16 fused MHA + - # fp32 RoPE constant — check with scripts/verify_fp8_rope.py. Same _DIT_PROFILE + # SA3 medium DiT in fp8 — SELECTABLE, medium-only, the MAX-SPEED clean tier, + # now CALIBRATED. Identical recipe to sa3-m-bf16 (same RoPE-baked ONNX flow from + # build_dit_bf16.py, same EXPLICIT_BATCH + OBEY + Fourier fp32 island) PLUS + # BuilderFlag.FP8: the ONNX carries fp8 E4M3 Quantize/Dequantize pairs on the 176 + # Linear GEMMs, so TRT fires fp8 tensor-core GEMMs there while the 96 attention + # blocks stay bf16 fused-MHA (_gemm_mha_v2) and the RoPE cos/sin ride as baked fp32 + # constants. Measured on H200 (same-run round-robin): ~1.3x faster than fp16-mixed + # at every length (1.40x@L129 / 1.32x@L1292 / 1.34x@L4092, also ahead of bf16) and + # clean at long sequence (latent std 0.86 vs eager 0.95, 0.000% clip @6-min) — the + # baked RoPE dodges bf16's long-angle drift, so it stays clean where bf16 clips. + # + # CALIBRATED SCALES (2026-07-31): the published dit_fp8.onnx now carries the + # real-conditioning calibrated activation + per-channel weight scales from + # @ryanontheinside's fp8 work (#47), grafted onto this bakedmin structure by + # build/transplant_scales.py. Calibration is speed-free (31.2 vs the uncalibrated + # 30.8 ms/fwd, within run noise) and lifts worst-step velocity-cos vs fp32 from + # 0.52/0.57/0.64 to 0.92/0.94/0.92 on adversarial seeds; sampling steps 1-7 match + # the fully-calibrated #47 reference within ~0.001. It is still a SPEED tier over the + # fp16-mixed default rather than a fidelity upgrade over it (single-step velocity cos + # ~0.92-0.94 < fp16mixed's ~1.0), but it no longer collapses at the highest-noise + # first step the way the uncalibrated engine did. #47's fp16-attention variant buys + # the last ~0.03 of step-0 fidelity for +2 ms/fwd; this tier keeps bf16 attention for + # the speed. Identity of the shipped engine: 176 fp8 GEMMs + 96 bf16 fused MHA + + # baked fp32 RoPE constant — check with scripts/verify_fp8_rope.py. Same _DIT_PROFILE # (batch=1, L in [1,4096]) → identical CLI/feature surface. medium-only, not # seed-reproducible vs fp16mixed. # - # Producer: the fp8-linear ONNX (fp8 QDQ inserted; inv_freq inline) is RoPE-baked - # by build_dit_bf16.py, the SAME baker as the bf16 tier (it handles inline OR - # external inv_freq). The baked cos/sin tables are sized to profile max L=4096; - # renders past 4096 (the SAME-L decoder's own cap → rejected by sa3_trt) would need - # a re-bake (--max-t). The .onnx keeps its external-data reference to - # `dit_fp8lin.onnx.data`, so that exact sidecar name is what ships alongside. + # Producer: the fp8-linear ONNX (fp8 QDQ inserted; inv_freq inline) is RoPE-baked by + # build_dit_bf16.py, the SAME baker as the bf16 tier (it handles inline OR external + # inv_freq), then build/transplant_scales.py swaps in #47's calibrated scale VALUES + # (the 5.8 GB fp32 weights are untouched — TRT re-quantizes them at build). A + # from-scratch recalibration (model retrain) regenerates those scales with #47's full + # pipeline: make_calib.py (real-conditioning capture, in this repo) -> build_dit_fp8.py + # (max-PTQ + per-channel scales, lives in #47). The baked cos/sin tables are sized to + # profile max L=4096; renders past 4096 (the SAME-L decoder's own cap → rejected by + # sa3_trt) would need a re-bake (--max-t). The .onnx keeps its external-data reference + # to `dit_fp8lin.onnx.data`, so that exact sidecar name is what ships alongside. "sa3-m-fp8": { "onnx_hf": ["sa3-m/dit_fp8.onnx", "sa3-m/dit_fp8lin.onnx.data"], "trt_local": "sa3-m/dit_fp8.trt", diff --git a/optimized/tensorRT/build/make_calib.py b/optimized/tensorRT/build/make_calib.py new file mode 100644 index 00000000..a5a84b20 --- /dev/null +++ b/optimized/tensorRT/build/make_calib.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python3 +"""Capture FP8 calibration data for the SA3 DiT (producer-side). + +``build_dit_fp8.py`` needs a calibration ``.npz``: real DiT inputs sampled +across the pingpong sampling schedule. This script produces one from nothing +but the model checkpoint, by driving the model's own ``generate()`` and +recording the six DiT engine inputs at every sampling step. + +FP8 is the only quantization path in this repo that needs calibration data +(``fp16mixed`` / ``fp32`` need none, MLX is weight-only), so this also +establishes the convention: a ``.calib.npz`` whose keys are the six ONNX +input names, each a leading-axis batch of samples. The npz is a reproducible +producer artifact: gitignored, never committed, regenerated on demand. + +Native by construction: + * Loads with ``loading_utils.load_diffusion_cond`` (the repo's own loader: + ``create_diffusion_cond_from_config`` + ``copy_state_dict``), wrapped in a + ``StableAudioModel``, so it needs only the checkpoint dir + (``model_config.json`` + ``model.safetensors`` + the bundled + ``t5gemma-b-b-ul2/`` encoder, loaded locally; falls back to the config's HF + model_name if absent). No ``StableAudioModel.from_pretrained`` (its + hard-coded checkpoint registry), no ``stable-audio-tools``. + * Prompts come from the repo's own ``interface/reprompt.py``: the Music + few-shot examples (``SYSTEM_PROMPTS["Music"]``), pulled at runtime via the + module's ``_extract_examples`` so there is one source of truth, not a + copy. These are the exact post-reprompt format the model is driven with at + deployment (genre + instruments + rhythm + mood + BPM + Length). + * Captures from the real ``StableAudioModel.generate()`` with + ``sampler_type="pingpong"``, ``steps=8``, ``cfg_scale=1.0`` (the settings + the model is sampled with): the conditioner, the dist-shifted sigma + schedule, and the pingpong renoise are all the real inference path, not a + re-implementation. + * Hooks the ``DiTWrapper`` forward, where conditioning arrives as the + 257-token ``cross_attn_cond`` (``[t5(256) | seconds(1)]``) and the + 257-channel ``local_add_cond`` (``[inpaint_masked_input(256) | + inpaint_mask(1)]``, all zeros for plain text-to-music). The trailing + seconds token is dropped to ``t5_hidden`` (256); ``seconds_total`` is the + raw duration scalar (the engine recomputes the seconds embedding itself). + +The six keys and the verified mapping: + x (N, 256, L) pre-conv latent at each step + t (N,) sigma at each step + t5_hidden (N, 256, 768) raw T5Gemma hidden states (cross-attn token 257 dropped) + t5_mask (N, 256) T5 attention mask + seconds_total (N,) raw duration scalar + local_add_cond (N, 257, L) inpaint local conditioning (zeros for t2m) + +Why these prompts, 1 seed, one duration: + Calibration is amax (``calibration_method="max"``), so the question is "have + we seen the tail that sets each tensor's scale", not "have we covered the + distribution". The quantized prompt-driven GEMMs (``to_cond_embed`` and the + cross-attn projections; attention BMMs are excluded by ``disable_mha_qdq``) + have activation maxima dominated by T5Gemma's systematic outlier channels, + which fire for essentially every prompt, so amax saturates after a handful of + prompts. The reprompt example set (47 prompts, a few hundred samples + over the 8 sigmas) sits comfortably inside the standard PTQ calibration band + (128-512 samples), past the knee. Broader or out-of-distribution prompts can + make calibration *worse*: under max calibration a single OOD outlier inflates + the scale and coarsens fp8 resolution for the in-distribution inputs the + engine actually serves, so the deployment-matched in-repo set is the right + default. A second seed re-rolls noise the outlier channels already maxed, so + seeds are the least useful axis. Duration is irrelevant to fp8 scales (the + seconds-embedding front-end is inside the re-applied FP32 islands, never + quantized), and a single duration keeps the npz rectangular for the build + contract. + +Usage: + python make_calib.py \ + --model-config /model_config.json \ + --checkpoint /model.safetensors \ + --out sa3-m.calib.npz + [--duration 54.0] [--steps 8] [--seed 1528] [--device cuda] + +Requires (producer): torch, numpy, and this repo's ``stable_audio_3`` package +with the model checkpoint. +""" +import argparse +import json +import sys +from pathlib import Path + +import numpy as np + +T5_TOKENS = 256 +T5_HIDDEN_DIM = 768 +IO_CHANNELS = 256 +LOCAL_ADD_CHANNELS = 257 # [inpaint_masked_input(256) | inpaint_mask(1)] + +# Default duration: with the default 6 s duration_padding_sec this adapts to a +# latent length L=646, a representative shape inside the published sa3-m profile's +# dynamic range (min 1, opt 1292, max 4096). Calibration scales are per-tensor, +# so the capture length need not equal the profile's opt point. +DEFAULT_DURATION_S = 54.0 +DEFAULT_STEPS = 8 +DEFAULT_SEED = 1528 + + +def _load_prompts() -> list: + """The Music few-shot examples from the repo's own ``reprompt.py`` (the + exact post-reprompt format the model is driven with at deployment), via the + module's own ``_extract_examples`` parser. Imported lazily so ``--help`` + doesn't pull in transformers (which the model load needs anyway).""" + from stable_audio_3.interface.reprompt import ( + SYSTEM_PROMPTS, _extract_examples) + prompts = _extract_examples(SYSTEM_PROMPTS["Music"]) + if not prompts: + raise RuntimeError( + "no Music examples found in reprompt.SYSTEM_PROMPTS['Music']") + return prompts + + +# --------------------------------------------------------------------------- +# capture +# --------------------------------------------------------------------------- + + +def _load_model(model_config_path: Path, checkpoint_path: Path, device: str): + """Load via the repo's own ``load_diffusion_cond`` (config dict + safetensors) + and wrap in ``StableAudioModel`` so ``generate()`` is the real entry point. + fp32 (model_half=False) for a clean, deterministic capture; amax from fp32 is + a safe upper bound for the fp16/fp8 engine.""" + from stable_audio_3.loading_utils import load_diffusion_cond + from stable_audio_3.model import StableAudioModel + + with open(model_config_path) as f: + model_config = json.load(f) + + # Point the T5Gemma conditioner at the encoder bundled in the checkpoint dir + # (snapshot_download lays it out as //), so capture needs + # only the checkpoint and runs offline. Drop repo_id so a stale network repo + # can never win. Falls back to the config's HF model_name if not bundled. + dest = model_config_path.parent + for c in model_config.get("model", {}).get("conditioning", {}).get("configs", []): + if c.get("type") == "t5gemma": + cc = c.setdefault("config", {}) + subfolder = cc.get("subfolder", "t5gemma-b-b-ul2") + if (dest / subfolder).is_dir(): + cc["model_path"] = str(dest) + cc.pop("repo_id", None) + + print(f" loading model: {checkpoint_path}", flush=True) + model = load_diffusion_cond( + model_config, str(checkpoint_path), device=device, model_half=False) + return StableAudioModel(model, model_config, device, model_half=False) + + +def capture(sa3, prompts, duration, steps, base_seed): + """Hook the DiTWrapper forward and run the pingpong generate() once per + prompt. Each generate() emits exactly ``steps`` DiT calls, one per sigma, so + sigma coverage is balanced without per-sigma capping.""" + rows = {"x": [], "t": [], "t5_hidden": [], "t5_mask": [], "local_add_cond": []} + dit = sa3.dit # DiTWrapper + orig_forward = dit.forward + + def hook(x, t, **kw): + ca = kw["cross_attn_cond"] + mask = kw["cross_attn_mask"] + if ca.shape[1] == T5_TOKENS + 1: # drop the trailing seconds token + ca = ca[:, :T5_TOKENS] + mask = mask[:, :T5_TOKENS] + lac = kw.get("local_add_cond", None) + for i in range(x.shape[0]): # batch is 1 at cfg_scale=1.0 + rows["x"].append(x[i].float().cpu().numpy()) + rows["t"].append(np.array([float(t[i])], dtype=np.float32)) + rows["t5_hidden"].append(ca[i].float().cpu().numpy()) + rows["t5_mask"].append(mask[i].float().cpu().numpy()) + if lac is not None: + rows["local_add_cond"].append(lac[i].float().cpu().numpy()) + else: + rows["local_add_cond"].append( + np.zeros((LOCAL_ADD_CHANNELS, x.shape[-1]), dtype=np.float32)) + return orig_forward(x, t, **kw) + + try: + dit.forward = hook + for p_idx, prompt in enumerate(prompts): + sa3.generate( + prompt=prompt, duration=duration, steps=steps, + cfg_scale=1.0, sampler_type="pingpong", + seed=base_seed + p_idx, + return_latents=True, # skip decode; we only need DiT inputs + ) + print(f" [{p_idx + 1}/{len(prompts)}] {len(rows['x'])} samples", + flush=True) + finally: + dit.forward = orig_forward + + return rows + + +def _save(rows, duration, out_path: Path): + n = len(rows["x"]) + if n == 0: + raise RuntimeError("captured 0 samples") + lengths = {a.shape[-1] for a in rows["x"]} + if len(lengths) != 1: + raise RuntimeError( + f"captured ragged latent lengths {sorted(lengths)} - the npz must be " + f"rectangular; use a single --duration") + L = lengths.pop() + + out = { + "x": np.stack(rows["x"]).astype(np.float32), + "t": np.concatenate(rows["t"]).astype(np.float32), + "t5_hidden": np.stack(rows["t5_hidden"]).astype(np.float32), + "t5_mask": np.stack(rows["t5_mask"]).astype(np.float32), + "seconds_total": np.full(n, float(duration), dtype=np.float32), + "local_add_cond": np.stack(rows["local_add_cond"]).astype(np.float32), + } + out_path.parent.mkdir(parents=True, exist_ok=True) + np.savez(out_path, **out) + + sched = sorted({round(float(t), 6) for t in out["t"]}, reverse=True) + print(f"\n wrote {out_path}: {n} samples, L={L}") + print(f" sigma schedule ({len(sched)}): {[round(s, 4) for s in sched]}") + lo, hi = float(out['t5_hidden'].min()), float(out['t5_hidden'].max()) + print(f" ranges: t5_hidden [{lo:.2f}, {hi:.2f}] " + f"x [{float(out['x'].min()):.2f}, {float(out['x'].max()):.2f}] " + f"local_add_cond absmax {float(np.abs(out['local_add_cond']).max()):.3g}") + return out_path + + +# --------------------------------------------------------------------------- +# main +# --------------------------------------------------------------------------- + + +def main() -> int: + ap = argparse.ArgumentParser( + description=__doc__.split("\n\n")[0], + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--model-config", required=True, + help="Path to model_config.json") + ap.add_argument("--checkpoint", required=True, + help="Path to model.safetensors") + ap.add_argument("--out", default="sa3-m.calib.npz", + help="Output calibration npz") + ap.add_argument("--duration", type=float, default=DEFAULT_DURATION_S, + help="Generation duration in seconds (one value: keeps the " + "npz rectangular; default adapts to L=646)") + ap.add_argument("--steps", type=int, default=DEFAULT_STEPS, + help="Pingpong steps (= sigmas captured per generate)") + ap.add_argument("--seed", type=int, default=DEFAULT_SEED, + help="Base seed (a per-prompt offset is derived)") + ap.add_argument("--device", default="cuda") + args = ap.parse_args() + + prompts = _load_prompts() + n_expected = len(prompts) * args.steps + print(f"capturing {len(prompts)} prompts x {args.steps} sigmas = " + f"{n_expected} samples @ {args.duration}s") + + sa3 = _load_model(Path(args.model_config), Path(args.checkpoint), args.device) + rows = capture(sa3, prompts, args.duration, args.steps, args.seed) + _save(rows, args.duration, Path(args.out)) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/optimized/tensorRT/build/transplant_scales.py b/optimized/tensorRT/build/transplant_scales.py new file mode 100644 index 00000000..7d9e7891 --- /dev/null +++ b/optimized/tensorRT/build/transplant_scales.py @@ -0,0 +1,265 @@ +#!/usr/bin/env python3 +"""Graft #47's calibrated fp8 scales onto OUR RoPE-baked bakedmin ONNX. + +This is the transplant step that produces the *calibrated* medium-DiT fp8 tier +(`dit_fp8.onnx` / `dit_fp8.trt`). It keeps our fast "bakedmin" structure — baked +fp32 RoPE constant + bf16 fused MHA, built weakly-typed — while swapping in the +calibrated fp8 scale VALUES from @ryanontheinside's fp8 work (#47, feat/dit-fp8). + +Pipeline (medium-DiT fp8 "calibrated bakedmin" tier) +---------------------------------------------------- + 1. make_calib.py capture a real-conditioning calibration .npz + (@ryanontheinside, #47) + 2. build_dit_fp8.py max-PTQ + per-channel weight scales -> a + *calibrated* fp8 ONNX (@ryanontheinside, #47; + full pipeline lives in that PR, not this repo) + 3. build_dit_bf16.py RoPE-bake OUR fp8-linear ONNX -> the bakedmin + structure (baked fp32 RoPE constant; this is the + SHARED baker for the bf16 AND fp8 tiers) + 4. transplant_scales.py <-- THIS SCRIPT: copy step 2's calibrated scale + VALUES onto step 3's graph, so the shipped fp8 + engine keeps bakedmin's bf16 fused-MHA speed AND + gains #47's worst-step fidelity. + 5. build_from_onnx.py sa3-m-fp8 compile the calibrated ONNX to dit_fp8.trt + +Why transplant instead of shipping #47's engine directly: #47's calibrated ONNX +is fp16mixed-derived and its engine uses fp16 attention + a runtime fp32 RoPE +island (+~2 ms/fwd). Our bakedmin structure keeps the faster bf16 fused MHA and +a baked fp32 RoPE constant. Calibration is a property of the SCALE VALUES, not of +the structure, so grafting #47's scales onto the bakedmin graph gives #47's +worst-step fidelity at bakedmin's speed (measured H200: 31.2 vs 30.8 ms/fwd, i.e. +speed-free; worst-step velocity-cos vs fp32 0.52->0.92 on adversarial seeds). A +from-scratch recalibration (model retrain) reruns steps 1-2 with #47's full +pipeline; everyday consumers just pull the published `dit_fp8.onnx`. + +What it does +------------ +OUR ropebaked ONNX has fp32 trunk, baked RoPE constant, and per-MatMul dedicated +fp8 Q/DQ with PER-TENSOR scales. Match every quantized Linear by weight-initializer +name (identical across both graphs) and swap ONLY the fp8 scale VALUES: + * activation : per-tensor scalar <- #47's real-conditioning amax/448 + * weight : per-tensor -> PER-CHANNEL axis=0 (pre-transpose N/out axis) + <- #47's per-channel vector. New per-channel fp8-zero vectors are + added (the graph's global scalar `fp8_zero` is shared, so weight + Q/DQ are repointed to fresh vectors; activations keep the scalar). +The 5.8 GB fp32 weights (`dit_fp8lin.onnx.data`) are UNTOUCHED — only scales change; +TRT re-quantizes the weights at build with the new scales. The 2 `to_global_embed` +linears #47 leaves unquantized are kept as-is (bakedmin's original fp8), preserving +bakedmin's 176-fp8-GEMM structure. Verify the result with +`../scripts/verify_fp8_rope.py` (expect 176 fp8 GEMMs + 96 bf16 fused MHA). + +Calibration approach + scale VALUES by @ryanontheinside (#47). + +usage: + python transplant_scales.py \ + --ours-onnx onnx/sa3-m/dit_fp8lin_ropebaked.onnx \ + --calib-onnx /dit_fp8_calib.onnx \ + --out onnx/sa3-m/dit_fp8.onnx +""" +import os, sys, argparse +import numpy as np +import onnx +from onnx import TensorProto, numpy_helper, helper +from collections import defaultdict + +E4M3_MAX = 448.0 + + +def read_ext(tensor, base): + meta = {d.key: d.value for d in tensor.external_data} + path = os.path.join(base, meta["location"]) + with open(path, "rb") as f: + f.seek(int(meta.get("offset", 0))) + raw = f.read(int(meta["length"])) + dt = onnx.helper.tensor_dtype_to_np_dtype(tensor.data_type) + return np.frombuffer(raw, dtype=dt).reshape([d for d in tensor.dims]).copy() + + +class G: + def __init__(self, path, load_data=False): + self.path = path; self.base = os.path.dirname(os.path.abspath(path)) + self.m = onnx.load(path, load_external_data=False) + self.g = self.m.graph + self.inits = {i.name: i for i in self.g.initializer} + self.producer = {}; self.consumers = defaultdict(list) + for n in self.g.node: + for o in n.output: self.producer[o] = n + for i in n.input: + if i: self.consumers[i].append(n) + + def arr(self, name): + if name in self.inits: + t = self.inits[name] + return read_ext(t, self.base) if t.data_location == TensorProto.EXTERNAL \ + else numpy_helper.to_array(t) + p = self.producer.get(name) + if p is not None and p.op_type == "Constant": + return numpy_helper.to_array(p.attribute[0].t) + return None + + def dims(self, name): + return list(self.inits[name].dims) if name in self.inits else None + + def axis(self, node): + for a in node.attribute: + if a.name == "axis": return a.i + return None + + +def scale_dtype(g, name): + if name in g.inits: return g.inits[name].data_type + return None + + +def trace_weight_init(g, tensor, hops=6): + cur = tensor + for _ in range(hops): + if cur in g.inits: return cur + p = g.producer.get(cur) + if p is None: return None + if p.op_type in ("Transpose", "DequantizeLinear", "QuantizeLinear", "Cast"): + cur = p.input[0] + else: return None + return None + + +def build_map(g): + out = {} + for mm in g.g.node: + if mm.op_type not in ("MatMul", "Gemm"): continue + w_side = a_side = None + for i in mm.input: + wi = trace_weight_init(g, i) + if wi is not None: w_side = (i, wi) + else: a_side = i + if w_side is None or a_side is None: continue + w_in, w_init = w_side + + def find_dq(t, hops=4): + cur = t + for _ in range(hops): + p = g.producer.get(cur) + if p is None: return None + if p.op_type == "DequantizeLinear": return p + if p.op_type in ("Transpose", "Cast"): cur = p.input[0] + else: return None + return None + + wdq = find_dq(w_in) + pa = g.producer.get(a_side); adq = pa if (pa and pa.op_type == "DequantizeLinear") else None + if wdq is None or adq is None: continue + wq = g.producer.get(wdq.input[0]); aq = g.producer.get(adq.input[0]) + out[w_init] = dict(matmul=mm.name, wq=wq, wdq=wdq, aq=aq, adq=adq) + return out + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__.split("\n\n")[0], + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--ours-onnx", required=True, + help="OUR RoPE-baked fp8-linear ONNX (build_dit_bf16.py output): " + "the bakedmin graph whose per-tensor scales are overwritten") + ap.add_argument("--calib-onnx", required=True, + help="#47's calibrated fp8 ONNX (build_dit_fp8.py output): the " + "SOURCE of the calibrated activation + per-channel weight scales") + ap.add_argument("--out", required=True, help="output calibrated bakedmin ONNX") + a = ap.parse_args() + + print(f"loading OURS ({a.ours_onnx}) + #47 calib ({a.calib_onnx}) metadata ...", flush=True) + go = G(a.ours_onnx); gc = G(a.calib_onnx) + mo = build_map(go); mc = build_map(gc) + both = sorted(set(mo) & set(mc)); only_ours = sorted(set(mo) - set(mc)) + print(f"matched linears: {len(both)} | ours-only (kept as-is): {len(only_ours)} " + f"-> {only_ours}", flush=True) + + # index #47 scale arrays by weight_init + c47scale = {} + for w in both: + cwdq = mc[w]['wdq']; cadq = mc[w]['adq'] + wscale = gc.arr(cwdq.input[1]); waxis = gc.axis(cwdq) + ascale = gc.arr(cadq.input[1]) + c47scale[w] = dict(wscale=np.asarray(wscale), waxis=waxis, ascale=np.asarray(ascale)) + + # sanity: #47 per-channel weight scale ~= amax(our fp32 weight channel)/448 + checked = 0; maxrel = 0.0 + for w in both[:6]: + W = go.arr(w) # our fp32 weight [N, K] + if W is None: continue + amax = np.abs(W.astype(np.float64)).max(axis=1) # per output channel N + want = amax / E4M3_MAX + got = c47scale[w]['wscale'].astype(np.float64).ravel() + if got.size == want.size: + rel = np.abs(got - want) / (want + 1e-12) + maxrel = max(maxrel, float(np.median(rel))) + checked += 1 + print(f"weight-scale amax cross-check on {checked} layers: median rel-diff ~ {maxrel:.3f} " + f"(~0 means #47 scales fit OUR fp32 weights and axis=0 is right)", flush=True) + + # ---- edit OUR graph in place ---- + inits_by_name = {i.name: i for i in go.g.initializer} + + def set_init_array(name, arr): + """Replace an existing (inline) initializer's data with arr, same name.""" + new = numpy_helper.from_array(np.ascontiguousarray(arr), name=name) + old = inits_by_name.get(name) + if old is not None: old.CopyFrom(new) + else: go.g.initializer.append(new); inits_by_name[name] = new + + def add_fp8_zeros(name, n): + """fp8 E4M3 zero = byte 0x00; build via raw bytes (robust vs ml_dtypes).""" + t = helper.make_tensor(name, TensorProto.FLOAT8E4M3FN, [n], vals=b"\x00" * n, raw=True) + old = inits_by_name.get(name) + if old is not None: old.CopyFrom(t) + else: go.g.initializer.append(t); inits_by_name[name] = t + + def set_axis(node, ax): + found = False + for at in node.attribute: + if at.name == "axis": at.i = ax; found = True + if not found: + node.attribute.append(helper.make_attribute("axis", ax)) + + n_w = n_a = 0 + for w in both: + o = mo[w]; c = c47scale[w] + N = go.dims(w)[0] + # weight scale is fp32 (matches bakedmin's fp32 scale -> DQ output dtype unchanged) + wscale = c['wscale'].ravel().astype(np.float32) + assert wscale.size == N, f"{w}: wscale {wscale.size} != N {N}" + ascale = np.asarray(c['ascale']).ravel().astype(np.float32)[0] # scalar fp32 + # --- weight Q + DQ: per-channel axis=0, new per-channel fp8 zero-point vector --- + wzp_name = o['wdq'].input[1].replace("wscale", "wzpvec") + if wzp_name == o['wdq'].input[1]: + wzp_name = o['wdq'].input[1] + "_zpvec" + add_fp8_zeros(wzp_name, N) + for nd in (o['wq'], o['wdq']): + set_axis(nd, 0) + set_init_array(nd.input[1], wscale) # scalar->vector, same name, fp32 + if len(nd.input) > 2 and nd.input[2]: + nd.input[2] = wzp_name # repoint off shared scalar fp8_zero + # --- activation Q + DQ: scalar value swap, keep shared scalar fp8_zero zp --- + for nd in (o['aq'], o['adq']): + set_init_array(nd.input[1], np.float32(ascale)) + n_w += 1; n_a += 1 + + os.makedirs(os.path.dirname(os.path.abspath(a.out)) or ".", exist_ok=True) + onnx.save(go.m, a.out) + print(f"wrote {a.out} ({os.path.getsize(a.out)/1e6:.1f} MB) " + f"edited {n_w} weight + {n_a} act quantizers", flush=True) + + # link the external-data sidecar(s) next to the output (weights are shared, untouched) + locs = {d.value for i in go.g.initializer if i.data_location == TensorProto.EXTERNAL + for d in i.external_data if d.key == "location"} + for loc in sorted(locs): + link = os.path.join(os.path.dirname(os.path.abspath(a.out)), loc) + tgt = os.path.realpath(os.path.join(go.base, loc)) + if os.path.exists(link) or os.path.islink(link): + print(f"external data {loc}: already present -> {os.path.realpath(link)}", flush=True) + else: + os.symlink(tgt, link) + print(f"external data {loc}: symlinked -> {tgt}", flush=True) + + +if __name__ == "__main__": + main()