初始提交：3 款纯音频游戏 + 共享引擎 + 完整文档与资产管线
- 共享引擎（app/src/main）：GestureEngine 6 手势识别、AudioEngine 三层音频 （ExoPlayer BGM 交叉淡入 / 对白 ducking / SoundPool 低延迟）、Game/GameEngine 状态机接口、MainActivity 全屏封面 + TalkBack - 《声命线》storyline：末世互动广播剧，3 幕 8 抉择点 4 结局 - 《声东击西》reflex：听声辨位反应，12 波战斗 + BGM 三档 BPM 变速 - 《暗室》escape：3x3 网格密室逃脱，钥匙→铁柜→保险丝→电箱→大门谜题链 - 文档：setup（从零搭建）/ architecture（架构）/ engine-api（引擎 API）/ asset-pipeline（资产流水线）/ games/*（三款游戏设计文档） - 资产管线：Qwen3-TTS 对白 + 程序合成音效/BGM + Mureka BGM 的 gen 脚本与台账， 全部 ogg 资产入库可复现 - env.d 环境脚本：android/tts 入库；sfx/mureka 含密钥不入库，仅提交 .example.sh 模板
This commit is contained in:
SoundGame Dev
@@ -0,0 +1,182 @@
|
||||
# SA3 Encoder + Decoder — Quantized Tiers (SAME-S / SAME-L)
|
||||
|
||||
Train-free quantized variants of the Stable Audio 3 **encoders and decoders**, each a single
|
||||
**self-contained ONNX** derived from the published bf16 export. TensorRT builds the engine; quality is
|
||||
preserved by GPTQ weight-error compensation and per-op precision placement (details below).
|
||||
|
||||
Two tiers, both derived from the published bf16 export:
|
||||
|
||||
- **`fp8`** — fp8 **compute** on the FFN GEMMs (fp8 tensor cores) → real speedup **and** a smaller engine, near-transparent.
|
||||
- **`fp8_fast`** — fp8 pushed into the attention projections too (SAME-S only) — faster still, lower quality.
|
||||
|
||||
For max fidelity, use the bf16 baseline (or `--precision fp32`).
|
||||
|
||||
> **Retired — `w8_bf16` (int8 weight-only).** Its `DequantizeLinear` constant-folds to bf16 at build, so
|
||||
> the *engine* came out byte-for-byte the size and speed of the bf16 baseline (only the *ONNX* was
|
||||
> smaller) with slightly lossier weights — strictly dominated by bf16. The lesson worth keeping: **ONNX
|
||||
> size ≠ engine size.** fp8 stays 1-byte all the way through (small ONNX *and* small engine *and* fp8
|
||||
> math); an int8-weight-only tier only shrinks the download, then expands back to bf16 in the engine.
|
||||
|
||||
---
|
||||
|
||||
## Tiers at a glance
|
||||
|
||||
The **bf16 baseline is the first row** of each table. Decoder dB is PSNR vs that bf16 baseline; encoder
|
||||
dB/cos is vs **eager** encode (the latent). dB is a heuristic — confirm near-transparent tiers by ear;
|
||||
`_fast` has audible artifacts. Speed is whole-model median at seq 1292 (~2 min audio).
|
||||
|
||||
### SAME-S decoder
|
||||
|
||||
| file | quantized | compute | onnx | speed | dB vs bf16 (music / sfx) |
|
||||
|---|---|---|---|---|---|
|
||||
| `dec_dynamic_bf16.onnx` | — *(bf16 baseline)* | bf16 | 219 MB | 1.00× | ref (7.25 ms @1292 · 23.2 @4096) |
|
||||
| `dec_fp8.onnx` | fp8 weights + activations (FFN) | fp8 FFN GEMM | 58 MB | **1.14×** | 30.9 / 41.0 · near-transparent |
|
||||
| `dec_fp8_fast.onnx` | fp8 weights + activations (all linears) | fp8 all GEMMs | 58 MB | **1.22×** | 26.1 / 36.5 · lossy |
|
||||
|
||||
### SAME-L decoder
|
||||
|
||||
| file | quantized | compute | onnx | speed | dB vs bf16 |
|
||||
|---|---|---|---|---|---|
|
||||
| `dec_dynamic_triton_swa.onnx` | — *(bf16 baseline)* | bf16 | 1193 MB | 1.00× | ref (54.0 ms @1292 · 174.4 @4096) |
|
||||
| `dec_fp8.onnx` | fp8 weights + activations (FFN) | fp8 FFN GEMM | 937 MB | **1.15×** | 43.4 · near-transparent |
|
||||
|
||||
SAME-L has **no `_fast` tier**: its attention projections are fp32 islands (fp8 there collapses quality — see islands below).
|
||||
|
||||
### SAME-S encoder
|
||||
|
||||
The **encoder is cheap** (~1 ms SAME-S, ~7 ms SAME-L), so quant mostly buys **size** — `fp8` shrinks both the onnx and the engine.
|
||||
|
||||
| file | quantized | compute | onnx | speed | dB / cos vs eager |
|
||||
|---|---|---|---|---|---|
|
||||
| `enc_dynamic_bf16.onnx` | — *(bf16 baseline)* | bf16 | 216 MB | 1.00× | 36.4 / 0.997 (~1 ms) |
|
||||
| `enc_fp8.onnx` | fp8 (FFN) | fp8 FFN GEMM | 54 MB | 1.08× | 29.9 / 0.986 · near-transparent |
|
||||
| `enc_fp8_fast.onnx` | fp8 (all linears) | fp8 all GEMMs | 54 MB | 1.14× | 22.3 / 0.924 · lossy |
|
||||
|
||||
### SAME-L encoder
|
||||
|
||||
| file | quantized | compute | onnx | speed | dB / cos vs eager |
|
||||
|---|---|---|---|---|---|
|
||||
| `enc_dynamic_triton_swa.onnx` | — *(bf16 baseline)* | bf16 | 1192 MB | 1.00× | 46.0 / 0.9997 (~7 ms) |
|
||||
| `enc_fp8.onnx` | fp8 (FFN) | fp8 FFN GEMM | 895 MB | 1.03× | 30.8 / 0.989 · near-transparent |
|
||||
|
||||
Encoder input is audio `(1,2,N)`; output latent `(1,256,N/4096)`. **Every encoder onnx carries a
|
||||
silence-pad node** (see below), so any `N` is accepted.
|
||||
|
||||
### `fp8` vs `fp8_fast`
|
||||
|
||||
Both run the **FFN** GEMMs (`ff.0` / `ff.2`) in fp8. The difference is the **attention projections**:
|
||||
- **`fp8`** leaves `to_qkv` / `to_out` in bf16 compute (weight-only fp8 for storage) — attention inputs
|
||||
stay full-precision → near-transparent.
|
||||
- **`fp8_fast`** runs `to_qkv` / `to_out` in fp8 too — the extra **+0.08×** (1.14 → 1.22×), but it
|
||||
fp8-rounds Q/K/V *before* the still-bf16 attention core, and the softmax path is fp8-sensitive → the
|
||||
quality drop (31 → 26 dB). SAME-S only.
|
||||
|
||||
---
|
||||
|
||||
## Sequence length, padding & building
|
||||
|
||||
The onnx graphs are **fully dynamic** — every model is correct at **any** length (verified L=1→8192
|
||||
vs eager, odd/even/prime, no size-dependent error; quant tiers track bf16 identically at all lengths).
|
||||
The only bound is the **TensorRT optimization profile** you build with. Two rules:
|
||||
|
||||
- **Build wide.** Use `min=1` and a `max` covering your longest track. Reference build uses
|
||||
**latent `[1, 1292, 8192]`** (decoders, ≈12:40) and **audio `[1, 2097152, 33554432]`** (encoders).
|
||||
`L=1–31` then work and match eager (these can't be chunked); long files run **natively, no chunking**
|
||||
(a 12:40 decode / 6:20 encode is a single call). The old `[32,4096]` was just a conservative default.
|
||||
- **Encoders auto-pad.** The SAME-L encoder needs the audio length to be a multiple of 4096 (the
|
||||
downsample ratio); the built-in pad node silence-pads the tail up to the next multiple, matching what
|
||||
eager does internally. Callers feed any `N`. (SAME-S never needed it; padded anyway for a uniform API.)
|
||||
|
||||
The `encode_chunked` path in the runtime is a **stale workaround** — it predated wide profiles and was
|
||||
chasing what turned out to be out-of-profile garbage, not a real divergence. Single-shot is accurate at
|
||||
any length (SAME-L enc cos 0.9995 on a real 285 s track). Chunking is unnecessary.
|
||||
|
||||
---
|
||||
|
||||
## Where the quantization goes, per tier
|
||||
|
||||
Every decoder block is `latent_proj → [ attn: to_qkv → (QKᵀ·softmax·V) → to_out ] → [ ff: ff.0 → GELU → ff.2 ]`,
|
||||
with RMSNorm + RoPE around attention. Each tier touches a different subset:
|
||||
|
||||
### `dec_fp8` — fp8 on the FFN
|
||||
- **FFN** (`ff.0` up-proj, `ff.2` down-proj) `+ latent_proj`: **weight fp8 + activation fp8** (per-tensor clipping scale) → **fp8 GEMM**. Weights GPTQ-compensated.
|
||||
- **Attention projections** (`to_qkv`, `to_out`): SAME-S → **weight-only fp8** (bf16 compute — kept out of fp8 math because they feed softmax); SAME-L → **fp32 islands** (untouched).
|
||||
- **Attention core** (QKᵀ·softmax·V): bf16 flash (unchanged).
|
||||
- **RMSNorm / RoPE:** fp32.
|
||||
|
||||
### `dec_fp8_fast` — fp8 everywhere it pays (SAME-S only)
|
||||
- **All linears** (FFN + `to_qkv` + `to_out` + `latent_proj`): **weight fp8 + activation fp8** → fp8 GEMM.
|
||||
- **Attention core** still bf16 — but now it's fed **fp8-rounded Q/K/V** from the projections. *That rounding is the quality cost* (the softmax path is fp8-sensitive), buying +0.08× speed over `dec_fp8`.
|
||||
- **RMSNorm / RoPE:** fp32.
|
||||
|
||||
---
|
||||
|
||||
## The levers (reference)
|
||||
|
||||
### What you quantize — the two halves of every GEMM
|
||||
|
||||
| lever | format | compute or storage | speed | quality | notes |
|
||||
|---|---|---|---|---|---|
|
||||
| **Weights** | fp8 (per-tensor*) or int8 (per-channel) | **storage** (dequant→bf16 at build) | none | near-lossless; **GPTQ** compensates rounding | 1 byte either format — fp8 ≠ smaller than int8 |
|
||||
| **Activations** | fp8 only (dynamic, per-tensor **clip** scale) | turns the GEMM into fp8 tensor-core compute | **unlocks the speedup** | the **lossy half** (~31 dB floor; GPTQ can't fix it) | calibration-critical — amax scale → collapse; use a clipping scale |
|
||||
|
||||
\* fp8 weight scale **must be per-tensor in a strongly-typed build** (per-channel silently collapses the GEMM). int8 is happily per-channel. int8 *activations* are a non-starter — softmax breaks.
|
||||
|
||||
### Where you apply it — which op
|
||||
|
||||
| lever | touches | fp8? | speed | quality | verdict |
|
||||
|---|---|---|---|---|---|
|
||||
| **FFN GEMM** | `ff.0`, `ff.2` | ✅ weight+act | 1.14× (both) / 1.07× (ff.0 only) | ff.0 transparent; **ff.2 sensitive** (±outliers) | the main speed win → `dec_fp8` |
|
||||
| **Attn projections** | `to_qkv`, `to_out` | ✅ weight+act | +0.08× on top | lossy — fp8-rounds Q/K/V *before* attention | `dec_fp8_fast` (SAME-S) |
|
||||
| **Attn core** | QKᵀ·softmax·V | ❌ no fp8 MHA kernel @ hd=64 | ~none | — | **dead at hd=64** (Sage 0.46–0.97×, fully-fp8 softmax slower); the real lever is algorithmic **windowing** (~1.6× on SAME-L long-seq), not quant |
|
||||
|
||||
### Leave-alone islands (constraints, not levers)
|
||||
|
||||
| op | rule | why |
|
||||
|---|---|---|
|
||||
| **RMSNorm** | keep fp32/bf16 | cheap; quantizing buys nothing |
|
||||
| **RoPE** | **must stay fp32** | bf16 RoPE angle → long-sequence clip bug (≥2 min renders) |
|
||||
| **Attn projections (SAME-L)** | fp32 islands | fp8 compute there collapses quality in the strongly-typed graph |
|
||||
|
||||
The mental model: **activations = speed (and the quality floor); weights = size (near-free with GPTQ); attn core = neither (use windowing).**
|
||||
|
||||
---
|
||||
|
||||
## Which to pick
|
||||
|
||||
- **Max fidelity** → the bf16 baseline (or `--precision fp32` for bit-exact eager)
|
||||
- **Faster + smaller, near-transparent (default quant)** → `dec_fp8`
|
||||
- **Fastest, quality-tolerant** → `dec_fp8_fast` (SAME-S)
|
||||
|
||||
## Reducing SAME-L size (optional)
|
||||
Both SAME-L tiers are 937 MB because the attention projections stay fp32 (680 MB). Weight-only-quantizing
|
||||
those islands (int8/fp8 weights, bf16 compute — no speed or quality cost) pulls both to ~500 MB.
|
||||
|
||||
## Building the engines
|
||||
|
||||
The tier **ONNX files ship on HuggingFace** (`stabilityai/stable-audio-3-optimized`, alongside the
|
||||
existing `dec_dynamic_bf16.onnx` / `enc_dynamic_*`). To build an engine with the wide profile:
|
||||
|
||||
```bash
|
||||
python build_tiers.py <tier.onnx> <out.trt> --arch {same-s|same-l} --kind {enc|dec} [--fp8]
|
||||
# SAME-S fp8 / fp8_fast tiers need --fp8; SAME-L carries fp8 in-graph (strongly-typed) and needs no flag.
|
||||
```
|
||||
|
||||
`build_tiers.py` bakes the **wide** profile (decoder `latent [1,1292,8192]`, encoder `audio [1,·,33.5M]`),
|
||||
so L=1..31 work and long files run natively. SAME-L uses the `diff_attn_swa` plugin (PREFER_JIT).
|
||||
|
||||
## Reproducing the tier ONNX (grafters)
|
||||
|
||||
The quantized ONNX are produced by grafting QDQ / int8 / fp8-stored weights onto the published bf16 export:
|
||||
|
||||
| grafter | produces |
|
||||
|---|---|
|
||||
| `fp8_gptq.py` | SAME-S `dec_fp8` / `dec_fp8_fast` (fp8-stored + GPTQ, ORT Hessian capture) |
|
||||
| `fp8_gptq_samel.py` | SAME-L `dec_fp8` (fp8-stored + GPTQ, eager Hessian capture — SWA blocks ORT) |
|
||||
| `gptq_w8.py` / `gptq_samel.py` | int8 weight-only ONNX — the *retired* `w8_bf16` tier (kept for reference) |
|
||||
| `enc_quant_sames.py` / `enc_quant_samel.py` | the encoder tiers |
|
||||
| `pad_encoder.py` | adds the silence-pad node so any audio length is accepted |
|
||||
|
||||
> **Note:** the grafters reference campaign-local calibration data (real-audio latents) and model
|
||||
> checkpoints via absolute paths at the top of each file — adjust those to reproduce. Calibration +
|
||||
> Hessians are computed from real-audio latents. The ready-made ONNX on HF need none of this.
|
||||
@@ -0,0 +1,68 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Build a SA3 encoder/decoder tier engine with a WIDE dynamic profile.
|
||||
|
||||
python build_tiers.py <tier.onnx> <out.trt> --arch {same-s,same-l} --kind {enc,dec} [--fp8]
|
||||
|
||||
Why "wide": the historical [32,4096]-frame profile was a conservative default, not a model limit.
|
||||
With min=1 the models are correct at L=1..31 (verified vs eager — these can't be chunked), and long
|
||||
files run natively (a ~12:40 decode / 6:20 encode is a single call, no chunking). See README.
|
||||
|
||||
Profiles built here:
|
||||
decoder : latent (1,256,1) .. (1,256,1292) .. (1,256,8192) # 1 frame .. ~12:40
|
||||
encoder : audio (1,2,1) .. (1,2,2097152) .. (1,2,33554432) # 1 samp .. ~12:40
|
||||
|
||||
SAME-S -> weakly-typed EXPLICIT_BATCH + BF16 (+ FP8 builder flag for fp8/fp8_fast tiers).
|
||||
SAME-L -> strongly-typed + diff_attn_swa plugin + NO FP8 flag (the in-graph fp8 QDQ nodes
|
||||
carry the precision). Plugin impl is chosen by GPU arch: AOT on sm_120 (JIT isn't
|
||||
stream-capturable there), JIT on sm_90 (override with SA3_SWA_PLUGIN=aot|jit).
|
||||
`--fp8` is required for fp8 / fp8_fast SAME-S tiers; harmless/ignored for SAME-L and w8/bf16.
|
||||
"""
|
||||
import argparse, os, sys
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "scripts"))
|
||||
import tensorrt as trt
|
||||
|
||||
DEC_PROFILE = ((1, 256, 1), (1, 256, 1292), (1, 256, 8192))
|
||||
ENC_PROFILE = ((1, 2, 1), (1, 2, 2097152), (1, 2, 33554432))
|
||||
|
||||
def build(onnx_path, out_path, arch, kind, fp8):
|
||||
lg = trt.Logger(trt.Logger.ERROR)
|
||||
trt.init_libnvinfer_plugins(lg, "")
|
||||
strong = arch == "same-l"
|
||||
flags = 0
|
||||
if strong:
|
||||
# SWA plugin: sm_120 (Blackwell) needs AOT — the JIT impl isn't stream-capturable there
|
||||
# and silently drops the decode inside the runtime's mega-graph. sm_90 uses JIT. Choose by
|
||||
# the GPU arch (override with SA3_SWA_PLUGIN=aot|jit); set it before importing the plugin.
|
||||
import torch
|
||||
gpu = "sm_%d%d" % torch.cuda.get_device_capability()
|
||||
want_aot = gpu == "sm_120" or os.environ.get("SA3_SWA_PLUGIN") == "aot"
|
||||
os.environ.setdefault("SA3_SWA_PLUGIN", "aot" if want_aot else "jit")
|
||||
import diff_attn_nocast_plugin # registers samel::diff_attn_swa (AOT + JIT impls) # noqa: F401
|
||||
flags |= 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
|
||||
flags |= 1 << int(trt.NetworkDefinitionCreationFlag.PREFER_AOT_PYTHON_PLUGINS if want_aot
|
||||
else trt.NetworkDefinitionCreationFlag.PREFER_JIT_PYTHON_PLUGINS)
|
||||
else:
|
||||
flags |= 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
|
||||
b = trt.Builder(lg); net = b.create_network(flags); p = trt.OnnxParser(net, lg)
|
||||
if not p.parse_from_file(onnx_path):
|
||||
for i in range(p.num_errors): print(p.get_error(i))
|
||||
raise SystemExit("parse failed")
|
||||
cfg = b.create_builder_config(); cfg.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 20 << 30)
|
||||
if not strong:
|
||||
cfg.set_flag(trt.BuilderFlag.BF16)
|
||||
if fp8: cfg.set_flag(trt.BuilderFlag.FP8) # SAME-L carries fp8 in-graph; only SAME-S needs the flag
|
||||
name, prof = ("audio", ENC_PROFILE) if kind == "enc" else ("latent", DEC_PROFILE)
|
||||
pr = b.create_optimization_profile(); pr.set_shape(name, *prof); cfg.add_optimization_profile(pr)
|
||||
ser = b.build_serialized_network(net, cfg)
|
||||
if ser is None: raise SystemExit("build failed")
|
||||
with open(out_path, "wb") as f: f.write(ser)
|
||||
print(f"built {out_path} ({ser.nbytes/1e6:.0f} MB) — {arch} {kind} {'fp8 ' if fp8 else ''}wide", flush=True)
|
||||
|
||||
if __name__ == "__main__":
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("onnx"); ap.add_argument("out")
|
||||
ap.add_argument("--arch", required=True, choices=["same-s", "same-l"])
|
||||
ap.add_argument("--kind", required=True, choices=["enc", "dec"])
|
||||
ap.add_argument("--fp8", action="store_true", help="required for SAME-S fp8/fp8_fast tiers")
|
||||
a = ap.parse_args()
|
||||
build(a.onnx, a.out, a.arch, a.kind, a.fp8)
|
||||
@@ -0,0 +1,124 @@
|
||||
#!/usr/bin/env python3
|
||||
"""SAME-L ENCODER quant tiers (w8 + fp8). Eager-capture the encoder FFN activations (Hessian + p99.9
|
||||
clip scale) via hooks (SWA plugin blocks ORT), map eager<->onnx FFN linears by order+weight-value,
|
||||
graft int8 weight-only (w8) and fp8 FFN GEMM (fp8, per-tensor weight, per-linear clip scale). Attention
|
||||
proj stay fp32 islands. Self-contained."""
|
||||
import sys, os, re, json, glob
|
||||
import numpy as np, torch, onnx
|
||||
from onnx import TensorProto, helper, numpy_helper
|
||||
sys.path.insert(0, "/weka2/cj/clod/sa3s/stable-audio-3/optimized/tensorRT/scripts"); sys.path.insert(0, "/weka2/cj/clod/fp8_calib/build")
|
||||
from stable_audio_3.factory import create_autoencoder_from_config
|
||||
from stable_audio_3.loading_utils import copy_state_dict
|
||||
torch.set_grad_enabled(False)
|
||||
WD = "/weka2/cj/clod/sames_fp8"; E4M3 = 448.0
|
||||
cfg = json.load(open("/weka2/cj/clod/sa3s/models/SAME-L/SAME-L.json"))
|
||||
ae = create_autoencoder_from_config(cfg["model"], cfg["sample_rate"])
|
||||
ck = torch.load("/weka2/cj/clod/sa3s/models/SAME-L/SAME-L.ckpt", map_location="cpu", weights_only=False)
|
||||
copy_state_dict(ae, ck.get("state_dict", ck) if isinstance(ck, dict) else ck); ae = ae.cuda().eval()
|
||||
# eager encoder FFN linears, in module order
|
||||
eproj = [m for nm, m in ae.named_modules() if re.search(r"encoder\..*\.ff\.ff\.0\.proj$", nm)]
|
||||
eout = [m for nm, m in ae.named_modules() if re.search(r"encoder\..*\.ff\.ff\.2$", nm)]
|
||||
print(f"eager encoder FFN: {len(eproj)} proj + {len(eout)} out", flush=True)
|
||||
HP = {}; PP = {} # keyed by (role, idx)
|
||||
def mk(role, i):
|
||||
def hook(m, inp):
|
||||
A = inp[0].reshape(-1, inp[0].shape[-1]).double(); h = (A.T @ A).cpu().numpy()
|
||||
HP[(role, i)] = h if (role, i) not in HP else HP[(role, i)] + h
|
||||
p = float(np.percentile(np.abs(inp[0].detach().float().cpu().numpy()), 99.9)); PP.setdefault((role, i), []).append(p)
|
||||
return hook
|
||||
H = [m.register_forward_pre_hook(mk("proj", i)) for i, m in enumerate(eproj)] + [m.register_forward_pre_hook(mk("out", i)) for i, m in enumerate(eout)]
|
||||
AUD = np.load(f"{WD}/calib_audio_samel.npz")["audio"].astype(np.float32)
|
||||
for ci in range(AUD.shape[0]):
|
||||
ae.encode(torch.tensor(AUD[ci:ci + 1], device="cuda")); print(f" captured clip {ci}", flush=True)
|
||||
for h in H: h.remove()
|
||||
ASC = {k: max(float(np.percentile(v, 90)) / E4M3, 1e-4) for k, v in PP.items()} # per-linear clip scale (robust over clips)
|
||||
def fp8_gptq(Wnp, Hnp, sc, bs=128, pd=0.05):
|
||||
dev = "cuda"; W = torch.tensor(np.asarray(Wnp, np.float32), device=dev); Hh = torch.tensor(np.asarray(Hnp, np.float32), device=dev)
|
||||
cin = W.shape[1]; idx = torch.arange(cin, device=dev); Hh[~torch.isfinite(Hh)] = 0
|
||||
d = torch.diag(Hh).clone(); dead = d <= 0; d[dead] = 1; Hh[idx, idx] = d; W[:, dead] = 0; Hh = (Hh + Hh.T) / 2; damp = pd * d.mean(); Hinv = None
|
||||
for k in range(6):
|
||||
try: L = torch.linalg.cholesky(Hh + torch.eye(cin, device=dev) * (damp * 4.0 ** k)); Hinv = torch.linalg.cholesky(torch.cholesky_inverse(L), upper=True); break
|
||||
except Exception: continue
|
||||
if Hinv is None: Hinv = torch.diag(1.0 / torch.sqrt(d + damp))
|
||||
Q = torch.zeros_like(W)
|
||||
for i1 in range(0, cin, bs):
|
||||
i2 = min(i1 + bs, cin); W1 = W[:, i1:i2].clone(); Q1 = torch.zeros_like(W1); Err = torch.zeros_like(W1); Hi = Hinv[i1:i2, i1:i2]
|
||||
for j in range(i2 - i1):
|
||||
w = W1[:, j]; dd = Hi[j, j].clamp_min(1e-12); qf = (w / sc).to(torch.float8_e4m3fn).to(torch.float32)
|
||||
q = qf * sc; Q1[:, j] = qf; err = (w - q) / dd; Err[:, j] = err; W1[:, j:] -= err[:, None] * Hi[j, j:][None, :]
|
||||
Q[:, i1:i2] = Q1
|
||||
if i2 < cin: W[:, i2:] -= Err @ Hinv[i1:i2, i2:]
|
||||
return Q
|
||||
def warr(inits, prod, n):
|
||||
w = n.input[1]
|
||||
if w in inits: return w, numpy_helper.to_array(inits[w]), False, None
|
||||
p = prod.get(w)
|
||||
if p and p.op_type == "Transpose" and p.input and p.input[0] in inits: return p.input[0], numpy_helper.to_array(inits[p.input[0]]), True, p
|
||||
return None, None, None, None
|
||||
SRC = os.environ["ENC_ONNX"]
|
||||
def onnx_ffn(g, inits, prod):
|
||||
pr, ou = {}, {}
|
||||
for n in g.node:
|
||||
if n.op_type != "MatMul": continue
|
||||
mp = re.search(r"blocks\.(\d+)/ff/proj", n.name); mo = re.search(r"blocks\.(\d+)/ff/out", n.name)
|
||||
if mp: pr[int(mp.group(1))] = n
|
||||
elif mo: ou[int(mo.group(1))] = n
|
||||
return [pr[i] for i in sorted(pr)], [ou[i] for i in sorted(ou)]
|
||||
def build_tier(mode, out):
|
||||
m = onnx.load(SRC, load_external_data=True); g = m.graph
|
||||
inits = {i.name: i for i in g.initializer}; prod = {o: n for n in g.node for o in n.output}
|
||||
oproj, oout = onnx_ffn(g, inits, prod)
|
||||
assert len(oproj) == len(eproj) and len(oout) == len(eout), f"{len(oproj)}/{len(eproj)} {len(oout)}/{len(eout)}"
|
||||
fp8 = mode == "fp8"
|
||||
if fp8: g.initializer.append(helper.make_tensor("fp8_zero", TensorProto.FLOAT8E4M3FN, [], [0.0]))
|
||||
new_nodes, new_inits, drop = [], [], set()
|
||||
for role, elist, olist in [("proj", eproj, oproj), ("out", eout, oout)]:
|
||||
for i, (em, n) in enumerate(zip(elist, olist)):
|
||||
ws, W, vt, tn = warr(inits, prod, n)
|
||||
Wg = W if vt else W.T # [out,in]
|
||||
Weag = em.weight.detach().double().cpu().numpy()
|
||||
assert np.allclose(Wg, Weag, atol=1e-2), f"map mismatch {role}{i} (L2 {np.abs(Wg-Weag).mean():.4f})"
|
||||
pfx = n.name.strip("/").replace("/", "_"); qax = 0 if vt else 1; wdq = f"{pfx}_wdq"
|
||||
if fp8:
|
||||
scale = np.float32(max(np.abs(Wg).max() / E4M3, 1e-8)) # per-tensor
|
||||
Qi = fp8_gptq(Wg, HP[(role, i)], scale); Qs = Qi if vt else Qi.T.contiguous()
|
||||
raw = Qs.to(torch.float8_e4m3fn).view(torch.uint8).cpu().numpy().tobytes()
|
||||
new_inits.append(helper.make_tensor(f"{pfx}_w8", TensorProto.FLOAT8E4M3FN, list(Qs.shape), raw, raw=True))
|
||||
new_inits.append(numpy_helper.from_array(np.array(scale, np.float16), f"{pfx}_wsc"))
|
||||
new_nodes.append(helper.make_node("DequantizeLinear", [f"{pfx}_w8", f"{pfx}_wsc", "fp8_zero"], [wdq], name=f"{pfx}_DQw"))
|
||||
a = ASC[(role, i)]; new_inits.append(numpy_helper.from_array(np.array(a, np.float16), f"{pfx}_asc"))
|
||||
aq, adq = f"{pfx}_aq", f"{pfx}_adq"
|
||||
new_nodes.append(helper.make_node("QuantizeLinear", [n.input[0], f"{pfx}_asc", "fp8_zero"], [aq], name=f"{pfx}_Qa"))
|
||||
new_nodes.append(helper.make_node("DequantizeLinear", [aq, f"{pfx}_asc", "fp8_zero"], [adq], name=f"{pfx}_DQa"))
|
||||
n.input[0] = adq
|
||||
else:
|
||||
scale = np.maximum(np.abs(Wg).max(axis=1) / 127.0, 1e-4).astype(np.float32) # per-channel int8
|
||||
sc_b = scale.reshape(-1, 1) if not vt else scale.reshape(-1, 1)
|
||||
Wq = np.round((Wg / scale.reshape(-1, 1))).clip(-127, 127).astype(np.int8); Qs = Wq if vt else Wq.T
|
||||
new_inits.append(numpy_helper.from_array(Qs, f"{pfx}_w8"))
|
||||
new_inits.append(numpy_helper.from_array(scale.astype(np.float16), f"{pfx}_wsc"))
|
||||
new_inits.append(helper.make_tensor(f"{pfx}_wzp", TensorProto.INT8, [scale.size], [0] * scale.size))
|
||||
new_nodes.append(helper.make_node("DequantizeLinear", [f"{pfx}_w8", f"{pfx}_wsc", f"{pfx}_wzp"], [wdq], name=f"{pfx}_DQw", axis=qax))
|
||||
if vt:
|
||||
for k, inp in enumerate(tn.input):
|
||||
if inp == ws: tn.input[k] = wdq
|
||||
else: n.input[1] = wdq
|
||||
drop.add(ws)
|
||||
g.initializer.extend(new_inits); g.node.extend(new_nodes)
|
||||
keep = [i for i in g.initializer if i.name not in drop]; del g.initializer[:]; g.initializer.extend(keep)
|
||||
avail = {i.name for i in g.initializer} | {i.name for i in g.input} | {""}; order, rem = [], list(g.node)
|
||||
while rem:
|
||||
nx, prog = [], False
|
||||
for nd in rem:
|
||||
if all(i in avail for i in nd.input): order.append(nd); [avail.add(o) for o in nd.output]; prog = True
|
||||
else: nx.append(nd)
|
||||
rem = nx
|
||||
if not prog: raise RuntimeError("topo stuck")
|
||||
del g.node[:]; g.node.extend(order)
|
||||
for f in (out, out + ".data"):
|
||||
if os.path.exists(f): os.remove(f)
|
||||
onnx.save(m, out, save_as_external_data=True, all_tensors_to_one_file=True, location=os.path.basename(out) + ".data", size_threshold=1024)
|
||||
print(f" {mode} -> {out}", flush=True)
|
||||
build_tier("w8", f"{WD}/enc_samel_w8.onnx")
|
||||
build_tier("fp8", f"{WD}/enc_samel_fp8.onnx")
|
||||
print("done", flush=True)
|
||||
@@ -0,0 +1,122 @@
|
||||
#!/usr/bin/env python3
|
||||
"""SAME-S ENCODER quant tiers. ORT-capture FFN/attn-proj activations from calibration AUDIO (Hessian +
|
||||
p99.9 clip scale), then emit 3 self-contained onnx: w8 (int8 weight-only), fp8 (FFN fp8-compute +
|
||||
attn weight-only fp8), fast (all-linears fp8-compute). Per-channel weight scales (weakly-typed OK)."""
|
||||
import os, re
|
||||
import numpy as np, onnx, torch
|
||||
from onnx import TensorProto, helper, numpy_helper
|
||||
import onnxruntime as ort
|
||||
torch.set_grad_enabled(False)
|
||||
WD = "/weka2/cj/clod/sames_fp8"; E4M3 = 448.0
|
||||
SRC = os.environ.get("ENC_ONNX")
|
||||
if not SRC:
|
||||
from huggingface_hub import hf_hub_download
|
||||
SRC = hf_hub_download("stabilityai/stable-audio-3-optimized", "onnx/same-s/enc_dynamic_bf16.onnx")
|
||||
AUD = np.load(f"{WD}/calib_audio_sames.npz")["audio"].astype(np.float32) # [6,2,N]
|
||||
def warr(inits, prod, n):
|
||||
w = n.input[1]
|
||||
if w in inits: return w, numpy_helper.to_array(inits[w]), False, None
|
||||
p = prod.get(w)
|
||||
if p and p.op_type == "Transpose" and p.input and p.input[0] in inits: return p.input[0], numpy_helper.to_array(inits[p.input[0]]), True, p
|
||||
return None, None, None, None
|
||||
# ---- capture (ORT): Hessian H=AtA + p99.9 abs per linear ----
|
||||
m = onnx.load(SRC, load_external_data=True); g = m.graph
|
||||
inits = {i.name: i for i in g.initializer}; prod = {o: n for n in g.node for o in n.output}
|
||||
tin = {}
|
||||
for n in g.node:
|
||||
if n.op_type == "MatMul":
|
||||
ws, W, vt, tn = warr(inits, prod, n)
|
||||
if ws is not None and W.ndim == 2: tin[n.name] = n.input[0]
|
||||
n0 = len(g.output); ex = {o.name for o in g.output}
|
||||
for t in set(tin.values()):
|
||||
if t not in ex: g.output.append(helper.make_tensor_value_info(t, TensorProto.FLOAT, None))
|
||||
for f in (f"{WD}/_enc.onnx", f"{WD}/_enc.onnx.data"):
|
||||
if os.path.exists(f): os.remove(f)
|
||||
onnx.save(m, f"{WD}/_enc.onnx", save_as_external_data=True, all_tensors_to_one_file=True, location="_enc.onnx.data", size_threshold=1024)
|
||||
so = ort.SessionOptions(); so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
|
||||
sess = ort.InferenceSession(f"{WD}/_enc.onnx", so, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
|
||||
iname = sess.get_inputs()[0].name; ks = list(tin.keys()); ts = [tin[k] for k in ks]
|
||||
Hs = {k: None for k in ks}; P99 = {k: [] for k in ks}
|
||||
for ci in range(AUD.shape[0]):
|
||||
acts = sess.run(ts, {iname: AUD[ci:ci + 1]})
|
||||
for k, a in zip(ks, acts):
|
||||
A = a.reshape(-1, a.shape[-1]).astype(np.float64); Hs[k] = A.T @ A if Hs[k] is None else Hs[k] + A.T @ A
|
||||
av = np.abs(a).ravel(); P99[k].append(av if av.size < 400000 else np.random.default_rng(ci).choice(av, 400000, replace=False))
|
||||
print(f" captured clip {ci}", flush=True)
|
||||
del sess
|
||||
ASCALE = {k: max(float(np.percentile(np.concatenate(v), 99.9)) / E4M3, 1e-4) for k, v in P99.items()} # per-linear clip scale
|
||||
print(f" clip scales: min {min(ASCALE.values()):.4f} max {max(ASCALE.values()):.4f}", flush=True)
|
||||
# ---- GPTQ (int8 or fp8), blocked, per-output-channel ----
|
||||
def gptq(Wnp, Hnp, sc_np, fmt, bs=128, pd=0.05):
|
||||
dev = "cuda"; W = torch.tensor(np.asarray(Wnp, np.float32), device=dev); H = torch.tensor(np.asarray(Hnp, np.float32), device=dev)
|
||||
sc = torch.tensor(sc_np.astype(np.float32), device=dev); cin = W.shape[1]; idx = torch.arange(cin, device=dev)
|
||||
H[~torch.isfinite(H)] = 0; d = torch.diag(H).clone(); dead = d <= 0; d[dead] = 1; H[idx, idx] = d; W[:, dead] = 0
|
||||
H = (H + H.T) / 2; damp = pd * d.mean(); Hinv = None
|
||||
for k in range(6):
|
||||
try: L = torch.linalg.cholesky(H + torch.eye(cin, device=dev) * (damp * 4.0 ** k)); Hinv = torch.linalg.cholesky(torch.cholesky_inverse(L), upper=True); break
|
||||
except Exception: continue
|
||||
if Hinv is None: Hinv = torch.diag(1.0 / torch.sqrt(d + damp))
|
||||
Q = torch.zeros_like(W)
|
||||
for i1 in range(0, cin, bs):
|
||||
i2 = min(i1 + bs, cin); W1 = W[:, i1:i2].clone(); Q1 = torch.zeros_like(W1); Err = torch.zeros_like(W1); Hi = Hinv[i1:i2, i1:i2]
|
||||
for j in range(i2 - i1):
|
||||
w = W1[:, j]; dd = Hi[j, j].clamp_min(1e-12)
|
||||
qf = torch.clamp(torch.round(w / sc), -127, 127) if fmt == "int8" else (w / sc).to(torch.float8_e4m3fn).to(torch.float32)
|
||||
q = qf * sc; Q1[:, j] = qf; err = (w - q) / dd; Err[:, j] = err; W1[:, j:] -= err[:, None] * Hi[j, j:][None, :]
|
||||
Q[:, i1:i2] = Q1
|
||||
if i2 < cin: W[:, i2:] -= Err @ Hinv[i1:i2, i2:]
|
||||
return Q
|
||||
# ---- graft one tier (fresh load -> self-contained) ----
|
||||
def build_tier(mode, out):
|
||||
m = onnx.load(SRC, load_external_data=True); g = m.graph
|
||||
inits = {i.name: i for i in g.initializer}; prod = {o: n for n in g.node for o in n.output}
|
||||
fp8 = mode in ("fp8", "fast"); noact = re.compile(r"to_qkv|to_out") if mode == "fp8" else re.compile(r"$^")
|
||||
if fp8: g.initializer.append(helper.make_tensor("fp8_zero", TensorProto.FLOAT8E4M3FN, [], [0.0]))
|
||||
new_nodes, new_inits, drop = [], [], set()
|
||||
for n in [x for x in g.node if x.op_type == "MatMul" and x.name in Hs]:
|
||||
ws, W, vt, tn = warr(inits, prod, n)
|
||||
if ws is None: continue
|
||||
Wg = W if vt else W.T; div = E4M3 if fp8 else 127.0
|
||||
scale = np.maximum(np.abs(Wg).max(axis=1) / div, 1e-8).astype(np.float32)
|
||||
Qi = gptq(Wg, Hs[n.name], scale, "fp8" if fp8 else "int8"); Qs = Qi if vt else Qi.T.contiguous()
|
||||
pfx = n.name.strip("/").replace("/", "_"); qax = 0 if vt else 1; wdq = f"{pfx}_wdq"
|
||||
do_act = fp8 and not noact.search(n.name)
|
||||
if fp8:
|
||||
raw = Qs.to(torch.float8_e4m3fn).view(torch.uint8).cpu().numpy().tobytes()
|
||||
new_inits.append(helper.make_tensor(f"{pfx}_w8", TensorProto.FLOAT8E4M3FN, list(Qs.shape), raw, raw=True))
|
||||
new_inits.append(numpy_helper.from_array(scale, f"{pfx}_wsc"))
|
||||
new_nodes.append(helper.make_node("DequantizeLinear", [f"{pfx}_w8", f"{pfx}_wsc", "fp8_zero"], [wdq], name=f"{pfx}_DQw", axis=qax))
|
||||
if do_act:
|
||||
a = ASCALE[n.name]; new_inits.append(helper.make_tensor(f"{pfx}_asc", TensorProto.FLOAT, [], [a]))
|
||||
aq, adq = f"{pfx}_aq", f"{pfx}_adq"
|
||||
new_nodes.append(helper.make_node("QuantizeLinear", [n.input[0], f"{pfx}_asc", "fp8_zero"], [aq], name=f"{pfx}_Qa"))
|
||||
new_nodes.append(helper.make_node("DequantizeLinear", [aq, f"{pfx}_asc", "fp8_zero"], [adq], name=f"{pfx}_DQa"))
|
||||
n.input[0] = adq
|
||||
else:
|
||||
new_inits.append(numpy_helper.from_array(Qs.to(torch.int8).cpu().numpy(), f"{pfx}_w8"))
|
||||
new_inits.append(numpy_helper.from_array(scale, f"{pfx}_wsc"))
|
||||
new_inits.append(helper.make_tensor(f"{pfx}_wzp", TensorProto.INT8, [scale.size], [0] * scale.size))
|
||||
new_nodes.append(helper.make_node("DequantizeLinear", [f"{pfx}_w8", f"{pfx}_wsc", f"{pfx}_wzp"], [wdq], name=f"{pfx}_DQw", axis=qax))
|
||||
if vt:
|
||||
for i, inp in enumerate(tn.input):
|
||||
if inp == ws: tn.input[i] = wdq
|
||||
else: n.input[1] = wdq
|
||||
drop.add(ws)
|
||||
g.initializer.extend(new_inits); g.node.extend(new_nodes)
|
||||
keep = [i for i in g.initializer if i.name not in drop]; del g.initializer[:]; g.initializer.extend(keep)
|
||||
avail = {i.name for i in g.initializer} | {i.name for i in g.input} | {""}; order, rem = [], list(g.node)
|
||||
while rem:
|
||||
nx, prog = [], False
|
||||
for nd in rem:
|
||||
if all(i in avail for i in nd.input): order.append(nd); [avail.add(o) for o in nd.output]; prog = True
|
||||
else: nx.append(nd)
|
||||
rem = nx
|
||||
if not prog: raise RuntimeError("topo stuck")
|
||||
del g.node[:]; g.node.extend(order)
|
||||
for f in (out, out + ".data"):
|
||||
if os.path.exists(f): os.remove(f)
|
||||
onnx.save(m, out, save_as_external_data=True, all_tensors_to_one_file=True, location=os.path.basename(out) + ".data", size_threshold=1024)
|
||||
print(f" {mode:5s} -> {out}", flush=True)
|
||||
for mode, out in [("w8", f"{WD}/enc_sames_w8.onnx"), ("fp8", f"{WD}/enc_sames_fp8.onnx"), ("fast", f"{WD}/enc_sames_fast.onnx")]:
|
||||
build_tier(mode, out)
|
||||
print("done", flush=True)
|
||||
@@ -0,0 +1,104 @@
|
||||
#!/usr/bin/env python3
|
||||
"""fp8-STORED + GPTQ decoder tier — self-contained, every linear 8-bit.
|
||||
ALL linears: per-output-channel fp8 weight, GPTQ-compensated, STORED as fp8 (small onnx).
|
||||
- FFN / non-attention: ALSO per-tensor fp8 activation QDQ -> fp8 GEMM (speed, 1.14x)
|
||||
- attention proj (to_qkv/to_out): weight-only fp8, DQ to bf16 -> bf16 compute (quality-safe)
|
||||
Fresh reload before grafting so the output .data is self-contained. argv: <fp32_onnx> <out_onnx>."""
|
||||
import sys, os, json, re
|
||||
import numpy as np, onnx, torch
|
||||
from onnx import TensorProto, helper, numpy_helper
|
||||
import onnxruntime as ort
|
||||
torch.set_grad_enabled(False)
|
||||
DEC, OUT = sys.argv[1], sys.argv[2]; WD = "/weka2/cj/clod/sames_fp8"
|
||||
AS = json.load(open(f"{WD}/dec_fp8_act_scales_real.json")); nsc = AS["node_scale"]; gs = AS["global_scale"]
|
||||
LATS = np.load(f"{WD}/calib_latents.npz")["latents"]
|
||||
NOACT = re.compile(os.environ.get("NOACT_RE", r"to_qkv|to_out")) # weight-only fp8 (bf16); set "$^" for all-linears fp8
|
||||
E4M3 = 448.0
|
||||
def warr(inits, prod, n):
|
||||
w = n.input[1]
|
||||
if w in inits: return w, numpy_helper.to_array(inits[w]), False, None
|
||||
p = prod.get(w)
|
||||
if p and p.op_type == "Transpose" and p.input and p.input[0] in inits: return p.input[0], numpy_helper.to_array(inits[p.input[0]]), True, p
|
||||
return None, None, None, None
|
||||
# ---- capture Hessians for ALL linears (ORT), on a temp copy ----
|
||||
m = onnx.load(DEC, load_external_data=True); g = m.graph
|
||||
inits = {i.name: i for i in g.initializer}; prod = {o: n for n in g.node for o in n.output}
|
||||
tin = {n.name: n.input[0] for n in g.node if n.op_type == "MatMul" and warr(inits, prod, n)[0] is not None and warr(inits, prod, n)[1].ndim == 2}
|
||||
n0 = len(g.output); ex = {o.name for o in g.output}
|
||||
for t in set(tin.values()):
|
||||
if t not in ex: g.output.append(helper.make_tensor_value_info(t, TensorProto.FLOAT, None))
|
||||
for _f in (f"{WD}/_fpg.onnx", f"{WD}/_fpg.onnx.data"):
|
||||
if os.path.exists(_f): os.remove(_f)
|
||||
onnx.save(m, f"{WD}/_fpg.onnx", save_as_external_data=True, all_tensors_to_one_file=True, location="_fpg.onnx.data", size_threshold=1024)
|
||||
so = ort.SessionOptions(); so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
|
||||
sess = ort.InferenceSession(f"{WD}/_fpg.onnx", so, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
|
||||
iname = sess.get_inputs()[0].name; ks = list(tin.keys()); ts = [tin[k] for k in ks]; Hs = {k: None for k in ks}
|
||||
for li in range(LATS.shape[0]):
|
||||
acts = sess.run(ts, {iname: LATS[li:li + 1, :, :768].astype(np.float32)})
|
||||
for k, a in zip(ks, acts):
|
||||
A = a.reshape(-1, a.shape[-1]).astype(np.float64); Hs[k] = A.T @ A if Hs[k] is None else Hs[k] + A.T @ A
|
||||
print(f" captured {li}", flush=True)
|
||||
del sess
|
||||
# ---- GPTQ (torch, blocked, per-output-channel fp8 weight) ----
|
||||
def fp8_gptq(Wnp, Hnp, sc_np, bs=128, pd=0.05):
|
||||
dev = "cuda"; W = torch.tensor(np.asarray(Wnp, np.float32), device=dev); H = torch.tensor(np.asarray(Hnp, np.float32), device=dev)
|
||||
sc = torch.tensor(sc_np.astype(np.float32), device=dev); cin = W.shape[1]; idx = torch.arange(cin, device=dev)
|
||||
H[~torch.isfinite(H)] = 0; d = torch.diag(H).clone(); dead = d <= 0; d[dead] = 1; H[idx, idx] = d; W[:, dead] = 0
|
||||
H = (H + H.T) / 2; damp = pd * d.mean(); Hinv = None
|
||||
for k in range(6):
|
||||
try: L = torch.linalg.cholesky(H + torch.eye(cin, device=dev) * (damp * 4.0 ** k)); Hinv = torch.linalg.cholesky(torch.cholesky_inverse(L), upper=True); break
|
||||
except Exception: continue
|
||||
if Hinv is None: Hinv = torch.diag(1.0 / torch.sqrt(d + damp))
|
||||
Q = torch.zeros_like(W)
|
||||
for i1 in range(0, cin, bs):
|
||||
i2 = min(i1 + bs, cin); W1 = W[:, i1:i2].clone(); Q1 = torch.zeros_like(W1); Err = torch.zeros_like(W1); Hi = Hinv[i1:i2, i1:i2]
|
||||
for j in range(i2 - i1):
|
||||
w = W1[:, j]; dd = Hi[j, j].clamp_min(1e-12)
|
||||
qf = (w / sc).to(torch.float8_e4m3fn).to(torch.float32); q = qf * sc
|
||||
Q1[:, j] = qf; err = (w - q) / dd; Err[:, j] = err; W1[:, j:] -= err[:, None] * Hi[j, j:][None, :]
|
||||
Q[:, i1:i2] = Q1
|
||||
if i2 < cin: W[:, i2:] -= Err @ Hinv[i1:i2, i2:]
|
||||
return Q
|
||||
# ---- fresh reload (all raw_data present -> self-contained output), then graft ----
|
||||
m = onnx.load(DEC, load_external_data=True); g = m.graph
|
||||
inits = {i.name: i for i in g.initializer}; prod = {o: n for n in g.node for o in n.output}
|
||||
g.initializer.append(helper.make_tensor("fp8_zero", TensorProto.FLOAT8E4M3FN, [], [0.0]))
|
||||
new_nodes, new_inits, drop, nact = [], [], set(), 0
|
||||
for n in [x for x in g.node if x.op_type == "MatMul" and x.name in Hs]:
|
||||
ws, W, vt, tn = warr(inits, prod, n)
|
||||
if ws is None: continue
|
||||
Wg = W if vt else W.T
|
||||
scale = np.maximum(np.abs(Wg).max(axis=1) / E4M3, 1e-8).astype(np.float32)
|
||||
Qs = fp8_gptq(Wg, Hs[n.name], scale); Qs = Qs if vt else Qs.T.contiguous()
|
||||
raw = Qs.to(torch.float8_e4m3fn).view(torch.uint8).cpu().numpy().tobytes()
|
||||
pfx = n.name.strip("/").replace("/", "_"); qax = 0 if vt else 1
|
||||
new_inits.append(helper.make_tensor(f"{pfx}_w8", TensorProto.FLOAT8E4M3FN, list(Qs.shape), raw, raw=True))
|
||||
new_inits.append(numpy_helper.from_array(scale, f"{pfx}_wsc"))
|
||||
wdq = f"{pfx}_wdq"
|
||||
new_nodes.append(helper.make_node("DequantizeLinear", [f"{pfx}_w8", f"{pfx}_wsc", "fp8_zero"], [wdq], name=f"{pfx}_DQw", axis=qax))
|
||||
if not NOACT.search(n.name): # FFN -> fp8 activation too (fp8 GEMM)
|
||||
a = float(max(nsc.get(n.name, gs), 1e-4)); new_inits.append(helper.make_tensor(f"{pfx}_asc", TensorProto.FLOAT, [], [a]))
|
||||
aq, adq = f"{pfx}_aq", f"{pfx}_adq"
|
||||
new_nodes.append(helper.make_node("QuantizeLinear", [n.input[0], f"{pfx}_asc", "fp8_zero"], [aq], name=f"{pfx}_Qa"))
|
||||
new_nodes.append(helper.make_node("DequantizeLinear", [aq, f"{pfx}_asc", "fp8_zero"], [adq], name=f"{pfx}_DQa"))
|
||||
n.input[0] = adq; nact += 1
|
||||
if vt:
|
||||
for i, inp in enumerate(tn.input):
|
||||
if inp == ws: tn.input[i] = wdq
|
||||
else: n.input[1] = wdq
|
||||
drop.add(ws)
|
||||
g.initializer.extend(new_inits); g.node.extend(new_nodes)
|
||||
_keep = [i for i in g.initializer if i.name not in drop]; del g.initializer[:]; g.initializer.extend(_keep)
|
||||
avail = {i.name for i in g.initializer} | {i.name for i in g.input} | {""}; order, rem = [], list(g.node)
|
||||
while rem:
|
||||
nx, prog = [], False
|
||||
for nd in rem:
|
||||
if all(i in avail for i in nd.input): order.append(nd); [avail.add(o) for o in nd.output]; prog = True
|
||||
else: nx.append(nd)
|
||||
rem = nx
|
||||
if not prog: raise RuntimeError("topo stuck")
|
||||
del g.node[:]; g.node.extend(order)
|
||||
for f in (OUT, OUT + ".data"):
|
||||
if os.path.exists(f): os.remove(f)
|
||||
onnx.save(m, OUT, save_as_external_data=True, all_tensors_to_one_file=True, location=os.path.basename(OUT) + ".data", size_threshold=1024)
|
||||
print(f"fp8-stored: {len(drop)} linears 8-bit ({nact} with fp8 activation/GEMM, {len(drop)-nact} weight-only) -> {OUT}", flush=True)
|
||||
@@ -0,0 +1,119 @@
|
||||
#!/usr/bin/env python3
|
||||
"""SAME-L fp8-STORED + GPTQ. Eager Hessian+amax capture (ORT can't run the SWA plugin), fp8-GPTQ the
|
||||
FFN weights, graft fp8-stored weight + per-tensor fp8 activation QDQ onto the SAME-L onnx. Attention
|
||||
proj stay fp32 islands (fp8 compute breaks them). Self-contained. Writes samel_fp8_gptq.onnx."""
|
||||
import sys, os, re, json
|
||||
from pathlib import Path
|
||||
import numpy as np, torch, onnx
|
||||
from onnx import TensorProto, helper, numpy_helper
|
||||
sys.path.insert(0, "/weka2/cj/clod/sa3s/stable-audio-3/optimized/tensorRT/scripts")
|
||||
sys.path.insert(0, "/weka2/cj/clod/fp8_calib/build")
|
||||
from stable_audio_3.factory import create_autoencoder_from_config
|
||||
from stable_audio_3.loading_utils import copy_state_dict
|
||||
torch.set_grad_enabled(False)
|
||||
ML = "/weka2/cj/clod/sa3s/models/SAME-L"; AB = "/weka2/cj/clod/sames_fp8/decoder_ab"; E4M3 = 448.0
|
||||
cfg = json.load(open(f"{ML}/SAME-L.json"))
|
||||
ae = create_autoencoder_from_config(cfg["model"], cfg["sample_rate"])
|
||||
ck = torch.load(f"{ML}/SAME-L.ckpt", map_location="cpu", weights_only=False)
|
||||
copy_state_dict(ae, ck.get("state_dict", ck) if isinstance(ck, dict) else ck); ae = ae.to("cuda").eval()
|
||||
def ekey(nm):
|
||||
m = re.search(r"transformers\.(\d+)\.ff\.ff\.0\.proj$", nm); x = f"b{m.group(1)}.proj" if m else None
|
||||
if x: return x
|
||||
m = re.search(r"transformers\.(\d+)\.ff\.ff\.2$", nm)
|
||||
if m: return f"b{m.group(1)}.out"
|
||||
return "latent_proj" if re.search(r"latent_proj$", nm) else None
|
||||
def okey(nm):
|
||||
m = re.search(r"blocks\.(\d+)/ff/proj", nm); x = f"b{m.group(1)}.proj" if m else None
|
||||
if x: return x
|
||||
m = re.search(r"blocks\.(\d+)/ff/out", nm)
|
||||
if m: return f"b{m.group(1)}.out"
|
||||
return "latent_proj" if re.search(r"latent_proj", nm) else None
|
||||
mods = {}
|
||||
for nm, mod in ae.named_modules():
|
||||
if hasattr(mod, "weight") and getattr(mod.weight, "ndim", 0) == 2:
|
||||
k = ekey(nm)
|
||||
if k: mods[k] = mod
|
||||
print(f"hooked {len(mods)} FFN linears", flush=True)
|
||||
H = {k: None for k in mods}; AMAX = {k: 0.0 for k in mods}
|
||||
def mk_hook(k):
|
||||
def hook(m, inp):
|
||||
A = inp[0].reshape(-1, inp[0].shape[-1]).double(); h = (A.T @ A).cpu().numpy()
|
||||
H[k] = h if H[k] is None else H[k] + h; AMAX[k] = max(AMAX[k], float(inp[0].abs().max()))
|
||||
return hook
|
||||
handles = [m.register_forward_pre_hook(mk_hook(k)) for k, m in mods.items()]
|
||||
lats = [np.load("/weka2/cj/clod/fp8_listening/latents_2min.npz")["bf16"]]
|
||||
lats += [np.load(f) for f in sorted(Path(AB).glob("samel_lat_*.npy"))]
|
||||
for i, L in enumerate(lats):
|
||||
ae.decode(torch.tensor(L, device="cuda", dtype=torch.float32)); print(f" captured {i} ({L.shape[-1]})", flush=True)
|
||||
for h in handles: h.remove()
|
||||
def fp8_gptq(Wnp, Hnp, sc_np, bs=128, pd=0.05):
|
||||
dev = "cuda"; W = torch.tensor(np.asarray(Wnp, np.float32), device=dev); Hh = torch.tensor(np.asarray(Hnp, np.float32), device=dev)
|
||||
sc = torch.tensor(sc_np.astype(np.float32), device=dev); cin = W.shape[1]; idx = torch.arange(cin, device=dev)
|
||||
Hh[~torch.isfinite(Hh)] = 0; d = torch.diag(Hh).clone(); dead = d <= 0; d[dead] = 1; Hh[idx, idx] = d; W[:, dead] = 0
|
||||
Hh = (Hh + Hh.T) / 2; damp = pd * d.mean(); Hinv = None
|
||||
for k in range(6):
|
||||
try: L = torch.linalg.cholesky(Hh + torch.eye(cin, device=dev) * (damp * 4.0 ** k)); Hinv = torch.linalg.cholesky(torch.cholesky_inverse(L), upper=True); break
|
||||
except Exception: continue
|
||||
if Hinv is None: Hinv = torch.diag(1.0 / torch.sqrt(d + damp))
|
||||
Q = torch.zeros_like(W)
|
||||
for i1 in range(0, cin, bs):
|
||||
i2 = min(i1 + bs, cin); W1 = W[:, i1:i2].clone(); Q1 = torch.zeros_like(W1); Err = torch.zeros_like(W1); Hi = Hinv[i1:i2, i1:i2]
|
||||
for j in range(i2 - i1):
|
||||
w = W1[:, j]; dd = Hi[j, j].clamp_min(1e-12)
|
||||
qf = (w / sc).to(torch.float8_e4m3fn).to(torch.float32); q = qf * sc
|
||||
Q1[:, j] = qf; err = (w - q) / dd; Err[:, j] = err; W1[:, j:] -= err[:, None] * Hi[j, j:][None, :]
|
||||
Q[:, i1:i2] = Q1
|
||||
if i2 < cin: W[:, i2:] -= Err @ Hinv[i1:i2, i2:]
|
||||
return Q
|
||||
# ---- graft onto the SAME-L onnx (fresh load -> self-contained) ----
|
||||
from huggingface_hub import hf_hub_download
|
||||
DEC = hf_hub_download("stabilityai/stable-audio-3-optimized", "onnx/same-l/dec_dynamic_triton_swa.onnx")
|
||||
m = onnx.load(DEC, load_external_data=True); g = m.graph
|
||||
inits = {i.name: i for i in g.initializer}; prod = {o: n for n in g.node for o in n.output}
|
||||
def warr(n):
|
||||
w = n.input[1]
|
||||
if w in inits: return w, numpy_helper.to_array(inits[w]), False, None
|
||||
p = prod.get(w)
|
||||
if p and p.op_type == "Transpose" and p.input and p.input[0] in inits: return p.input[0], numpy_helper.to_array(inits[p.input[0]]), True, p
|
||||
return None, None, None, None
|
||||
g.initializer.append(helper.make_tensor("fp8_zero", TensorProto.FLOAT8E4M3FN, [], [0.0]))
|
||||
new_nodes, new_inits, drop, done = [], [], set(), 0
|
||||
for n in [x for x in g.node if x.op_type == "MatMul"]:
|
||||
key = okey(n.name)
|
||||
if key not in mods: continue # skip attn proj (fp32 islands) / non-FFN
|
||||
ws, W, vt, tn = warr(n)
|
||||
if ws is None: continue
|
||||
Wg = mods[key].weight.detach().double().cpu().numpy() # eager [out,in]
|
||||
scale = np.float32(max(np.abs(Wg).max() / E4M3, 1e-8)) # per-TENSOR scalar (strongly-typed fp8 GEMM needs per-tensor weight scale)
|
||||
Qi = fp8_gptq(Wg, H[key], scale); Qstore = Qi if vt else Qi.T.contiguous() # onnx orientation
|
||||
raw = Qstore.to(torch.float8_e4m3fn).view(torch.uint8).cpu().numpy().tobytes()
|
||||
pfx = n.name.strip("/").replace("/", "_")
|
||||
new_inits.append(helper.make_tensor(f"{pfx}_w8", TensorProto.FLOAT8E4M3FN, list(Qstore.shape), raw, raw=True))
|
||||
new_inits.append(numpy_helper.from_array(np.array(scale, np.float16), f"{pfx}_wsc")) # scalar fp16 (fp16 trunk)
|
||||
a = 0.05; new_inits.append(numpy_helper.from_array(np.array(a, np.float16), f"{pfx}_asc")) # repo-calibrated FFN activation CLIPPING scale
|
||||
wdq, aq, adq = f"{pfx}_wdq", f"{pfx}_aq", f"{pfx}_adq"
|
||||
new_nodes.append(helper.make_node("DequantizeLinear", [f"{pfx}_w8", f"{pfx}_wsc", "fp8_zero"], [wdq], name=f"{pfx}_DQw"))
|
||||
new_nodes.append(helper.make_node("QuantizeLinear", [n.input[0], f"{pfx}_asc", "fp8_zero"], [aq], name=f"{pfx}_Qa"))
|
||||
new_nodes.append(helper.make_node("DequantizeLinear", [aq, f"{pfx}_asc", "fp8_zero"], [adq], name=f"{pfx}_DQa"))
|
||||
n.input[0] = adq
|
||||
if vt:
|
||||
for i, inp in enumerate(tn.input):
|
||||
if inp == ws: tn.input[i] = wdq
|
||||
else: n.input[1] = wdq
|
||||
drop.add(ws); done += 1
|
||||
g.initializer.extend(new_inits); g.node.extend(new_nodes)
|
||||
_keep = [i for i in g.initializer if i.name not in drop]; del g.initializer[:]; g.initializer.extend(_keep)
|
||||
avail = {i.name for i in g.initializer} | {i.name for i in g.input} | {""}; order, rem = [], list(g.node)
|
||||
while rem:
|
||||
nx, prog = [], False
|
||||
for nd in rem:
|
||||
if all(i in avail for i in nd.input): order.append(nd); [avail.add(o) for o in nd.output]; prog = True
|
||||
else: nx.append(nd)
|
||||
rem = nx
|
||||
if not prog: raise RuntimeError("topo stuck")
|
||||
del g.node[:]; g.node.extend(order)
|
||||
OUT = "/weka2/cj/clod/sames_fp8/samel_fp8_gptq.onnx"
|
||||
for f in (OUT, OUT + ".data"):
|
||||
if os.path.exists(f): os.remove(f)
|
||||
onnx.save(m, OUT, save_as_external_data=True, all_tensors_to_one_file=True, location=os.path.basename(OUT) + ".data", size_threshold=1024)
|
||||
print(f"SAME-L fp8-stored+GPTQ on {done} FFN linears -> {OUT}", flush=True)
|
||||
@@ -0,0 +1,130 @@
|
||||
#!/usr/bin/env python3
|
||||
"""GPTQ int8 for the SAME-L decoder FFN linears. ORT can't run the SWA plugin, so capture the
|
||||
per-linear input activations from the EAGER SAME-L autoencoder (hooks), build Hessians, GPTQ, and
|
||||
graft int8 (fp16 scale) onto the SAME-L onnx (attention proj kept fp32). Writes samel_w8_gptq.onnx."""
|
||||
import sys, os, re
|
||||
from pathlib import Path
|
||||
import numpy as np, torch, onnx
|
||||
from onnx import TensorProto, helper, numpy_helper
|
||||
sys.path.insert(0, "/weka2/cj/clod/sa3s/stable-audio-3/optimized/tensorRT/scripts")
|
||||
sys.path.insert(0, "/weka2/cj/clod/fp8_calib/build")
|
||||
from stable_audio_3.factory import create_autoencoder_from_config
|
||||
from stable_audio_3.loading_utils import copy_state_dict
|
||||
import json
|
||||
torch.set_grad_enabled(False)
|
||||
ML = "/weka2/cj/clod/sa3s/models/SAME-L"; AB = "/weka2/cj/clod/sames_fp8/decoder_ab"
|
||||
cfg = json.load(open(f"{ML}/SAME-L.json"))
|
||||
ae = create_autoencoder_from_config(cfg["model"], cfg["sample_rate"])
|
||||
ck = torch.load(f"{ML}/SAME-L.ckpt", map_location="cpu", weights_only=False)
|
||||
copy_state_dict(ae, ck.get("state_dict", ck) if isinstance(ck, dict) else ck)
|
||||
ae = ae.to("cuda").eval()
|
||||
# hook the decoder FFN linears (exclude attention proj). Normalize eager+onnx names to (block, role).
|
||||
def ekey(nm):
|
||||
m = re.search(r"transformers\.(\d+)\.ff\.ff\.0\.proj$", nm)
|
||||
if m: return f"b{m.group(1)}.proj"
|
||||
m = re.search(r"transformers\.(\d+)\.ff\.ff\.2$", nm)
|
||||
if m: return f"b{m.group(1)}.out"
|
||||
if re.search(r"latent_proj$", nm): return "latent_proj"
|
||||
return None
|
||||
def okey(nm):
|
||||
m = re.search(r"blocks\.(\d+)/ff/proj", nm)
|
||||
if m: return f"b{m.group(1)}.proj"
|
||||
m = re.search(r"blocks\.(\d+)/ff/out", nm)
|
||||
if m: return f"b{m.group(1)}.out"
|
||||
if re.search(r"latent_proj", nm): return "latent_proj"
|
||||
return None
|
||||
mods = {}
|
||||
for nm, mod in ae.named_modules():
|
||||
if hasattr(mod, "weight") and getattr(mod.weight, "ndim", 0) == 2:
|
||||
k = ekey(nm)
|
||||
if k: mods[k] = mod
|
||||
print(f"hooked {len(mods)} eager FFN linears: {sorted(mods)[:3]} ...", flush=True)
|
||||
H = {k: None for k in mods}
|
||||
def mk_hook(k):
|
||||
def hook(m, inp):
|
||||
A = inp[0].reshape(-1, inp[0].shape[-1]).double()
|
||||
h = (A.T @ A).cpu().numpy()
|
||||
H[k] = h if H[k] is None else H[k] + h
|
||||
return hook
|
||||
handles = [m.register_forward_pre_hook(mk_hook(k)) for k, m in mods.items()]
|
||||
# calibration latents: medium 2-min + 380s + the 10 real-song SAME-L latents
|
||||
lats = [np.load("/weka2/cj/clod/fp8_listening/latents_2min.npz")["bf16"]] # 1292 samples ...
|
||||
lats += [np.load(f) for f in sorted(Path(AB).glob("samel_lat_*.npy"))] # + 10×376 = ample for the Hessian
|
||||
for i, L in enumerate(lats):
|
||||
ae.decode(torch.tensor(L, device="cuda", dtype=torch.float32)); print(f" captured latent {i} ({L.shape[-1]})", flush=True)
|
||||
for h in handles: h.remove()
|
||||
def gptq(Wnp, Hnp, scale_np, pd=0.05, bs=128):
|
||||
dev = "cuda"
|
||||
W = torch.tensor(np.asarray(Wnp, np.float32), device=dev) # [out, in]
|
||||
H = torch.tensor(np.asarray(Hnp, np.float32), device=dev)
|
||||
sc = torch.tensor(scale_np.astype(np.float32), device=dev) # [out]
|
||||
cin = W.shape[1]; idx = torch.arange(cin, device=dev)
|
||||
H[~torch.isfinite(H)] = 0.0
|
||||
d = torch.diag(H).clone(); dead = d <= 0; d[dead] = 1.0; H[idx, idx] = d; W[:, dead] = 0.0
|
||||
H = (H + H.T) / 2; damp = pd * d.mean(); Hinv = None
|
||||
for k in range(6): # damped Cholesky of H⁻¹ (upper)
|
||||
try:
|
||||
L = torch.linalg.cholesky(H + torch.eye(cin, device=dev) * (damp * 4.0 ** k))
|
||||
Hinv = torch.linalg.cholesky(torch.cholesky_inverse(L), upper=True); break
|
||||
except Exception:
|
||||
continue
|
||||
if Hinv is None: Hinv = torch.diag(1.0 / torch.sqrt(d + damp))
|
||||
Q = torch.zeros_like(W)
|
||||
for i1 in range(0, cin, bs): # BLOCKED: small updates in-block, 1 matmul for the rest
|
||||
i2 = min(i1 + bs, cin); W1 = W[:, i1:i2].clone(); Q1 = torch.zeros_like(W1); Err = torch.zeros_like(W1)
|
||||
Hi = Hinv[i1:i2, i1:i2]
|
||||
for j in range(i2 - i1):
|
||||
w = W1[:, j]; dd = Hi[j, j].clamp_min(1e-12)
|
||||
qi = torch.clamp(torch.round(w / sc), -127, 127); err = (w - qi * sc) / dd
|
||||
Q1[:, j] = qi; Err[:, j] = err; W1[:, j:] -= err[:, None] * Hi[j, j:][None, :]
|
||||
Q[:, i1:i2] = Q1
|
||||
if i2 < cin: W[:, i2:] -= Err @ Hinv[i1:i2, i2:]
|
||||
return torch.clamp(torch.round(Q), -127, 127).to(torch.int8).cpu().numpy()
|
||||
# graft onto the onnx
|
||||
from huggingface_hub import hf_hub_download
|
||||
DEC = hf_hub_download("stabilityai/stable-audio-3-optimized", "onnx/same-l/dec_dynamic_triton_swa.onnx")
|
||||
m = onnx.load(DEC, load_external_data=True); g = m.graph
|
||||
inits = {i.name: i for i in g.initializer}; prod = {o: n for n in g.node for o in n.output}
|
||||
def warr(n):
|
||||
w = n.input[1]
|
||||
if w in inits: return w, numpy_helper.to_array(inits[w]), False, None
|
||||
p = prod.get(w)
|
||||
if p and p.op_type == "Transpose" and p.input and p.input[0] in inits: return p.input[0], numpy_helper.to_array(inits[p.input[0]]), True, p
|
||||
return None, None, None, None
|
||||
new_nodes, new_inits, done = [], [], 0
|
||||
for n in [x for x in g.node if x.op_type == "MatMul"]:
|
||||
key = okey(n.name)
|
||||
if key not in mods: continue # skip attn proj / non-FFN
|
||||
ws, W, vt, tn = warr(n)
|
||||
if ws is None: continue
|
||||
Wg = mods[key].weight.detach().double().cpu().numpy() # eager [out,in]
|
||||
scale = np.maximum(np.abs(Wg).max(axis=1) / 127.0, 1e-4).astype(np.float32)
|
||||
Qi = gptq(Wg, H[key], scale) # [out,in]
|
||||
qax = 0 if vt else 1
|
||||
inits[ws].CopyFrom(numpy_helper.from_array(Qi if vt else Qi.T, ws))
|
||||
pfx = n.name.strip("/").replace("/", "_")
|
||||
new_inits += [numpy_helper.from_array(scale.astype(np.float16), f"{pfx}_wsc"),
|
||||
helper.make_tensor(f"{pfx}_wzp", TensorProto.INT8, [scale.size], [0] * scale.size)]
|
||||
dq = f"{pfx}_wdq"
|
||||
new_nodes.append(helper.make_node("DequantizeLinear", [ws, f"{pfx}_wsc", f"{pfx}_wzp"], [dq], name=f"{pfx}_DQw", axis=qax))
|
||||
if vt:
|
||||
for i, inp in enumerate(tn.input):
|
||||
if inp == ws: tn.input[i] = dq
|
||||
else: n.input[1] = dq
|
||||
done += 1
|
||||
g.initializer.extend(new_inits); g.node.extend(new_nodes)
|
||||
avail = {i.name for i in g.initializer} | {i.name for i in g.input} | {""}
|
||||
order, rem = [], list(g.node)
|
||||
while rem:
|
||||
nx, prog = [], False
|
||||
for nd in rem:
|
||||
if all(i in avail for i in nd.input): order.append(nd); [avail.add(o) for o in nd.output]; prog = True
|
||||
else: nx.append(nd)
|
||||
rem = nx
|
||||
if not prog: raise RuntimeError("topo stuck")
|
||||
del g.node[:]; g.node.extend(order)
|
||||
OUT = "/weka2/cj/clod/sames_fp8/samel_w8_gptq.onnx"
|
||||
for f in (OUT, OUT + ".data"):
|
||||
if os.path.exists(f): os.remove(f)
|
||||
onnx.save(m, OUT, save_as_external_data=True, all_tensors_to_one_file=True, location=os.path.basename(OUT) + ".data", size_threshold=1024)
|
||||
print(f"GPTQ int8 grafted on {done} SAME-L FFN linears -> {OUT}", flush=True)
|
||||
@@ -0,0 +1,87 @@
|
||||
#!/usr/bin/env python3
|
||||
"""GPTQ int8 per-channel weight-only quantization of a decoder onnx. Captures each linear's input
|
||||
activations (ORT), builds the Hessian H = AᵀA, and quantizes weights column-by-column with
|
||||
inverse-Hessian error compensation (the OBQ/GPTQ update). Writes a W8 onnx. argv: <fp32_onnx> <out>
|
||||
[Lslice]."""
|
||||
import os, sys
|
||||
import numpy as np, onnx
|
||||
from onnx import TensorProto, helper, numpy_helper
|
||||
import onnxruntime as ort
|
||||
DEC, OUT = sys.argv[1], sys.argv[2]
|
||||
LSL = int(sys.argv[3]) if len(sys.argv) > 3 else 768
|
||||
LATS = np.load("/weka2/cj/clod/sames_fp8/calib_latents.npz")["latents"]
|
||||
m = onnx.load(DEC, load_external_data=True); g = m.graph
|
||||
inits = {i.name: i for i in g.initializer}; prod = {o: n for n in g.node for o in n.output}
|
||||
def warr(n):
|
||||
w = n.input[1]
|
||||
if w in inits: return w, numpy_helper.to_array(inits[w]), False, None
|
||||
p = prod.get(w)
|
||||
if p and p.op_type == "Transpose" and p.input and p.input[0] in inits: return p.input[0], numpy_helper.to_array(inits[p.input[0]]), True, p
|
||||
return None, None, None, None
|
||||
linears = []
|
||||
for n in g.node:
|
||||
if n.op_type == "MatMul":
|
||||
ws, W, vt, tn = warr(n)
|
||||
if ws is not None and W.ndim == 2: linears.append((n, ws, W, vt, tn))
|
||||
print(f"{len(linears)} linears; capturing activations (L={LSL}) ...", flush=True)
|
||||
# --- capture inputs via ORT (add as temp outputs, then restore) ---
|
||||
in_names = [n.input[0] for n, *_ in linears]
|
||||
n0 = len(g.output)
|
||||
existing = {o.name for o in g.output}
|
||||
for t in in_names:
|
||||
if t not in existing: g.output.append(helper.make_tensor_value_info(t, TensorProto.FLOAT, None))
|
||||
onnx.save(m, "/weka2/cj/clod/sames_fp8/_gptq_cap.onnx", save_as_external_data=True, all_tensors_to_one_file=True, location="_gptq_cap.onnx.data", size_threshold=1024)
|
||||
del g.output[n0:]
|
||||
so = ort.SessionOptions(); so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
|
||||
sess = ort.InferenceSession("/weka2/cj/clod/sames_fp8/_gptq_cap.onnx", so, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
|
||||
iname = sess.get_inputs()[0].name
|
||||
Hs = {n.name: None for n, *_ in linears}
|
||||
for li in range(LATS.shape[0]):
|
||||
acts = sess.run(in_names, {iname: LATS[li:li + 1, :, :LSL].astype(np.float32)})
|
||||
for (n, *_), a in zip(linears, acts):
|
||||
A = a.reshape(-1, a.shape[-1]).astype(np.float64)
|
||||
Hs[n.name] = A.T @ A if Hs[n.name] is None else Hs[n.name] + A.T @ A
|
||||
print(f" latent {li} done", flush=True)
|
||||
# --- GPTQ ---
|
||||
def gptq(W, H, scale, pd=0.01):
|
||||
W = W.astype(np.float64).copy(); cin = W.shape[1]; H = H.copy()
|
||||
dead = np.diag(H) == 0; H[dead, dead] = 1.0; W[:, dead] = 0.0
|
||||
H[np.diag_indices(cin)] += pd * np.mean(np.diag(H))
|
||||
Hinv = np.linalg.cholesky(np.linalg.inv(H)).T # upper-tri Cholesky of H⁻¹
|
||||
Qi = np.zeros(W.shape, dtype=np.int8); sc = scale
|
||||
for i in range(cin):
|
||||
w = W[:, i]; d = Hinv[i, i]
|
||||
qi = np.clip(np.round(w / sc), -127, 127); q = qi * sc
|
||||
Qi[:, i] = qi.astype(np.int8)
|
||||
W[:, i:] -= np.outer((w - q) / d, Hinv[i, i:])
|
||||
return Qi
|
||||
new_nodes, new_inits = [], []
|
||||
for (n, ws, W, vt, tn) in linears:
|
||||
Wg = W if vt else W.T # -> [out, in]
|
||||
scale = np.maximum(np.abs(Wg).max(axis=1) / 127.0, 1e-9).astype(np.float32) # per-output [out]
|
||||
Qi = gptq(Wg, Hs[n.name], scale) # [out, in] int8
|
||||
inits[ws].CopyFrom(numpy_helper.from_array(Qi if vt else Qi.T, ws))
|
||||
pfx = n.name.strip("/").replace("/", "_"); qax = 0 if vt else 1
|
||||
new_inits += [helper.make_tensor(f"{pfx}_wsc", TensorProto.FLOAT, [scale.size], scale.tolist()),
|
||||
helper.make_tensor(f"{pfx}_wzp", TensorProto.INT8, [scale.size], [0] * scale.size)]
|
||||
dq = f"{pfx}_wdq"
|
||||
new_nodes.append(helper.make_node("DequantizeLinear", [ws, f"{pfx}_wsc", f"{pfx}_wzp"], [dq], name=f"{pfx}_DQw", axis=qax))
|
||||
if vt:
|
||||
for i, inp in enumerate(tn.input):
|
||||
if inp == ws: tn.input[i] = dq
|
||||
else: n.input[1] = dq
|
||||
g.initializer.extend(new_inits); g.node.extend(new_nodes)
|
||||
avail = {i.name for i in g.initializer} | {i.name for i in g.input} | {""}
|
||||
order, rem = [], list(g.node)
|
||||
while rem:
|
||||
nx, prog = [], False
|
||||
for nd in rem:
|
||||
if all(i in avail for i in nd.input): order.append(nd); [avail.add(o) for o in nd.output]; prog = True
|
||||
else: nx.append(nd)
|
||||
rem = nx
|
||||
if not prog: raise RuntimeError("topo stuck")
|
||||
del g.node[:]; g.node.extend(order)
|
||||
for f in (OUT, OUT + ".data"):
|
||||
if os.path.exists(f): os.remove(f)
|
||||
onnx.save(m, OUT, save_as_external_data=True, all_tensors_to_one_file=True, location=os.path.basename(OUT) + ".data", size_threshold=1024)
|
||||
print(f"GPTQ int8 on {len(linears)} linears -> {OUT}", flush=True)
|
||||
@@ -0,0 +1,43 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Add a dynamic silence-pad node to an encoder onnx so it accepts ANY audio length: pads the sample
|
||||
axis up to the next multiple of 4096 with zeros (matches what eager does internally). argv: <in> <out>."""
|
||||
import sys, os
|
||||
import onnx
|
||||
from onnx import helper, TensorProto
|
||||
IN, OUT = sys.argv[1], sys.argv[2]
|
||||
m = onnx.load(IN, load_external_data=True); g = m.graph
|
||||
audio = g.input[0].name; adt = g.input[0].type.tensor_type.elem_type
|
||||
opset = {o.domain: o.version for o in m.opset_import}.get("", 0)
|
||||
print(f"input '{audio}' dtype={adt} opset={opset}", flush=True)
|
||||
g.initializer.extend([
|
||||
helper.make_tensor("_pc4096", TensorProto.INT64, [1], [4096]),
|
||||
helper.make_tensor("_pzeros5", TensorProto.INT64, [5], [0, 0, 0, 0, 0]),
|
||||
helper.make_tensor("_pidx2", TensorProto.INT64, [1], [2]),
|
||||
helper.make_tensor("_pc0", adt, [], [0]),
|
||||
])
|
||||
for n in g.node: # rewire consumers to the padded tensor
|
||||
for i, inp in enumerate(n.input):
|
||||
if inp == audio: n.input[i] = "_audio_pad"
|
||||
g.node.extend([ # pad = (4096 - N%4096) % 4096, pad end of axis 2
|
||||
helper.make_node("Shape", [audio], ["_pshape"], name="_pad_shape"),
|
||||
helper.make_node("Gather", ["_pshape", "_pidx2"], ["_pN"], name="_pad_gather", axis=0),
|
||||
helper.make_node("Mod", ["_pN", "_pc4096"], ["_pr"], name="_pad_mod1"),
|
||||
helper.make_node("Sub", ["_pc4096", "_pr"], ["_pt"], name="_pad_sub"),
|
||||
helper.make_node("Mod", ["_pt", "_pc4096"], ["_ppad"], name="_pad_mod2"),
|
||||
helper.make_node("Concat", ["_pzeros5", "_ppad"], ["_ppads"], name="_pad_concat", axis=0),
|
||||
helper.make_node("Pad", [audio, "_ppads", "_pc0"], ["_audio_pad"], name="_pad_pad", mode="constant"),
|
||||
])
|
||||
avail = {i.name for i in g.initializer} | {i.name for i in g.input} | {""}
|
||||
order, rem = [], list(g.node)
|
||||
while rem:
|
||||
nx, prog = [], False
|
||||
for nd in rem:
|
||||
if all(i in avail for i in nd.input): order.append(nd); [avail.add(o) for o in nd.output]; prog = True
|
||||
else: nx.append(nd)
|
||||
rem = nx
|
||||
if not prog: raise RuntimeError("topo stuck")
|
||||
del g.node[:]; g.node.extend(order)
|
||||
for f in (OUT, OUT + ".data"):
|
||||
if os.path.exists(f): os.remove(f)
|
||||
onnx.save(m, OUT, save_as_external_data=True, all_tensors_to_one_file=True, location=os.path.basename(OUT) + ".data", size_threshold=1024)
|
||||
print(f"padded encoder -> {OUT}", flush=True)
|
||||
@@ -0,0 +1,93 @@
|
||||
"""Retighten the SAME-S fp8 DECODER activation scales.
|
||||
|
||||
The shipped scales are amax*1.5 over 6 calibration latents, and on real music that leaves
|
||||
1.19-1.80x of headroom (median 1.50) with nothing clipping. Headroom above 1.0 is wasted fp8
|
||||
range, so this recalibrates from a WIDER real-music set with a smaller margin.
|
||||
|
||||
⚠ Calibrated on tracks DISJOINT from the evaluation set -- calibrating and evaluating on the same
|
||||
audio would report an overfit number.
|
||||
"""
|
||||
import argparse, json, os, random, sys
|
||||
sys.path.append("/weka2/cj/clod/sa3s/stable-audio-3/optimized/tensorRT/scripts")
|
||||
sys.path.insert(0, "gradio"); sys.path.insert(0, "lora")
|
||||
os.environ.setdefault("SA3_SWA_PLUGIN", "aot")
|
||||
import numpy as np, onnx, onnxruntime as ort, torch, soundfile as sf
|
||||
import sa3_trt_core as canon; canon._import_heavy()
|
||||
from onnx import numpy_helper as nh
|
||||
E4M3 = 448.0; SR, SPL, L = 44100, 4096, 256; N = L * SPL
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--fp8", required=True); ap.add_argument("--out", required=True)
|
||||
ap.add_argument("--margin", type=float, default=1.1)
|
||||
ap.add_argument("--ntracks", type=int, default=16)
|
||||
ap.add_argument("--seed", type=int, default=99) # != the eval seed (5)
|
||||
a = ap.parse_args()
|
||||
def load_excerpt(path, start_s=35.0):
|
||||
from math import gcd
|
||||
from scipy.signal import resample_poly
|
||||
try:
|
||||
i = sf.info(path); need = int(N * i.samplerate / SR) + 256
|
||||
if i.frames < int(start_s * i.samplerate) + need: return None
|
||||
w, sr = sf.read(path, start=int(start_s*i.samplerate), frames=need,
|
||||
dtype="float32", always_2d=True)
|
||||
except Exception: return None
|
||||
if w.shape[1] == 1: w = np.repeat(w, 2, 1)
|
||||
w = w[:, :2]
|
||||
if sr != SR:
|
||||
g = gcd(SR, sr); w = resample_poly(w, SR//g, sr//g, axis=0).astype(np.float32)
|
||||
if w.shape[0] < N: return None
|
||||
w = w[:N]
|
||||
if not np.isfinite(w).all() or np.abs(w).max() < 1e-3: return None
|
||||
return np.ascontiguousarray(w.T)
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
cat = json.load(open("/weka2/cj/knn/promptlists/asx_comma_audio.json"))
|
||||
mus = [e for e in cat if "tracktype: sfx" not in e.get("p", "").lower()]
|
||||
random.Random(a.seed).shuffle(mus)
|
||||
picked = []
|
||||
with ThreadPoolExecutor(32) as ex:
|
||||
for e, w in zip(mus[:900], ex.map(lambda x: load_excerpt(x["audio"]), mus[:900])):
|
||||
if w is not None:
|
||||
picked.append(w)
|
||||
if len(picked) >= a.ntracks: break
|
||||
print(f" calibrating on {len(picked)} real tracks (seed {a.seed}, disjoint from eval)", flush=True)
|
||||
E = canon.TRTRunner("engines/ship/same-s_enc_2prof_v3.trt", None, False, 0)
|
||||
lats = [canon.encode_chunked(E, torch.from_numpy(w).unsqueeze(0).cuda(),
|
||||
warmup_passes=0, balance=True).float().cpu().numpy() for w in picked]
|
||||
E.free(); torch.cuda.empty_cache()
|
||||
from huggingface_hub import hf_hub_download
|
||||
base = hf_hub_download("stabilityai/stable-audio-3-optimized", "onnx/same-s/dec_dynamic_bf16.onnx")
|
||||
m = onnx.load(base, load_external_data=True); g = m.graph
|
||||
inits = {i.name for i in g.initializer}; prod = {o: n for n in g.node for o in n.output}
|
||||
targets = {}
|
||||
for n in g.node:
|
||||
if n.op_type == "MatMul":
|
||||
w = n.input[1]
|
||||
if w in inits or (prod.get(w) and prod[w].op_type == "Transpose" and prod[w].input[0] in inits):
|
||||
targets[n.name] = n.input[0]
|
||||
have = {o.name for o in g.output}
|
||||
for t in set(targets.values()):
|
||||
if t not in have:
|
||||
g.output.append(onnx.helper.make_tensor_value_info(t, onnx.TensorProto.FLOAT, None))
|
||||
WK = "/weka2/cj/tmp/_dec_recal.onnx"
|
||||
onnx.save(m, WK, save_as_external_data=True, all_tensors_to_one_file=True,
|
||||
location=os.path.basename(WK)+".data", size_threshold=1024)
|
||||
so = ort.SessionOptions(); so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
|
||||
sess = ort.InferenceSession(WK, so, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
|
||||
iname = sess.get_inputs()[0].name; ks = list(targets); ts = [targets[k] for k in ks]
|
||||
amax = {k: 0.0 for k in ks}
|
||||
for i, z in enumerate(lats):
|
||||
for k, o in zip(ks, sess.run(ts, {iname: z.astype(np.float32)})):
|
||||
amax[k] = max(amax[k], float(np.abs(o).max()))
|
||||
del sess
|
||||
new = {k.strip("/").replace("/", "_") + "_asc": amax[k] * a.margin / E4M3 for k in ks}
|
||||
mq = onnx.load(a.fp8, load_external_data=True)
|
||||
old_v, new_v, hits = [], [], 0
|
||||
for init in mq.graph.initializer:
|
||||
if init.name in new:
|
||||
arr = nh.to_array(init); o = float(np.asarray(arr).ravel()[0])
|
||||
old_v.append(o); new_v.append(new[init.name]); hits += 1
|
||||
init.CopyFrom(nh.from_array(np.full(arr.shape, new[init.name], dtype=arr.dtype), init.name))
|
||||
r = np.array(new_v) / np.array(old_v)
|
||||
print(f" {hits} activation scales rewritten (margin {a.margin})")
|
||||
print(f" new/old ratio: min {r.min():.2f}x median {np.median(r):.2f}x max {r.max():.2f}x")
|
||||
onnx.save(mq, a.out, save_as_external_data=False)
|
||||
print(f" wrote {a.out}")
|
||||
@@ -0,0 +1,115 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Recompute the SAME-L fp8 ENCODER activation scales with AMAX.
|
||||
|
||||
WHY. enc_quant_samel.py set each activation quantiser's scale to
|
||||
|
||||
max(percentile(p99.9_per_calibration_clip, 90) / 448, 1e-4)
|
||||
|
||||
— the 99.9th percentile within each clip, then the 90th percentile of *that* across clips, then a
|
||||
floor. Three tail-discarding steps on a heavy-tailed distribution. The decoder's own quantiser
|
||||
(fp8_gptq_samel.py) used plain `amax`, and the two disagree wildly:
|
||||
|
||||
decoder clip point = scale*448 = 22.4 on all 24 quantisers (one global amax)
|
||||
encoder clip point = 0.04 .. 1.22, per layer, with every ff_out pinned to the 1e-4 floor
|
||||
|
||||
Measured effect: the encoder's latents landed 18.3 dB from the fp16 reference where the decoder's
|
||||
output landed 28. End to end that cost 0.15-2.16 dB of round trip on real music — worst on clean
|
||||
acoustic material (jazz), where quantisation noise has nothing to hide behind, and least on loud
|
||||
clipped sources where the limiter masks it.
|
||||
|
||||
WHAT THIS DOES. Re-runs the SAME capture — same 6 calibration clips, same hook points — with
|
||||
`scale = amax / 448`, and rewrites only the 24 activation-scale initializers. The GPTQ'd fp8
|
||||
WEIGHTS are untouched, so this is a scale correction rather than a requantisation. Measured: the
|
||||
old scales were 1.3-28.4x too small (median 3.0x); after the fix latent SNR is 30.9 dB and the
|
||||
round trip equals fp16 on every source tested.
|
||||
|
||||
⚠ The scale initializers are fp16. Writing them back as fp32 makes the DequantizeLinear output
|
||||
FLOAT while the matching weight stays HALF, and TRT rejects the MatMul with "must have same input
|
||||
types". Preserve the dtype.
|
||||
|
||||
usage: recalib_enc_fp8.py --src enc_samel_fp8.onnx --out enc_fp8_amax.onnx \
|
||||
--ckpt-dir /path/to/SAME-L --calib calib_audio_samel.npz
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
|
||||
import numpy as np
|
||||
import onnx
|
||||
import torch
|
||||
from onnx import numpy_helper as nh
|
||||
|
||||
E4M3_MAX = 448.0
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--src", required=True, help="the fp8 encoder ONNX to correct")
|
||||
ap.add_argument("--out", required=True)
|
||||
ap.add_argument("--ckpt-dir", required=True, help="dir holding SAME-L.json + SAME-L.ckpt")
|
||||
ap.add_argument("--calib", required=True, help="calib_audio_samel.npz (key 'audio')")
|
||||
a = ap.parse_args()
|
||||
|
||||
from stable_audio_3.factory import create_autoencoder_from_config
|
||||
from stable_audio_3.loading_utils import copy_state_dict
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
cfg = json.load(open(f"{a.ckpt_dir}/SAME-L.json"))
|
||||
ae = create_autoencoder_from_config(cfg["model"], cfg["sample_rate"])
|
||||
ck = torch.load(f"{a.ckpt_dir}/SAME-L.ckpt", map_location="cpu", weights_only=False)
|
||||
copy_state_dict(ae, ck.get("state_dict", ck) if isinstance(ck, dict) else ck)
|
||||
ae = ae.cuda().eval()
|
||||
|
||||
eproj = [m for nm, m in ae.named_modules() if re.search(r"encoder\..*\.ff\.ff\.0\.proj$", nm)]
|
||||
eout = [m for nm, m in ae.named_modules() if re.search(r"encoder\..*\.ff\.ff\.2$", nm)]
|
||||
print(f"eager encoder FFN: {len(eproj)} proj + {len(eout)} out", flush=True)
|
||||
|
||||
amax = {}
|
||||
|
||||
def mk(role, i):
|
||||
def hook(_m, inp):
|
||||
amax[(role, i)] = max(amax.get((role, i), 0.0), float(inp[0].detach().abs().max()))
|
||||
return hook
|
||||
|
||||
handles = ([m.register_forward_pre_hook(mk("proj", i)) for i, m in enumerate(eproj)] +
|
||||
[m.register_forward_pre_hook(mk("out", i)) for i, m in enumerate(eout)])
|
||||
aud = np.load(a.calib)["audio"].astype(np.float32)
|
||||
for ci in range(aud.shape[0]):
|
||||
ae.encode(torch.tensor(aud[ci:ci + 1], device="cuda"))
|
||||
print(f" captured clip {ci}", flush=True)
|
||||
for h in handles:
|
||||
h.remove()
|
||||
del ae
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
new = {k: float(v) / E4M3_MAX for k, v in amax.items()}
|
||||
model = onnx.load(a.src, load_external_data=True)
|
||||
inits = {i.name: i for i in model.graph.initializer}
|
||||
qa = {}
|
||||
for n in model.graph.node:
|
||||
if n.op_type == "QuantizeLinear" and n.input[0] not in inits and n.input[1] in inits:
|
||||
m = re.match(r"blocks\.(\d+)_ff_(proj|out)_MatMul_Qa", n.name or "")
|
||||
if not m:
|
||||
raise SystemExit(f"unrecognised activation quantiser name: {n.name!r}")
|
||||
qa[(m.group(2), int(m.group(1)))] = n.input[1]
|
||||
if set(qa) != set(new):
|
||||
raise SystemExit(f"mapping mismatch: onnx has {len(qa)}, eager captured {len(new)}")
|
||||
|
||||
ratios = []
|
||||
for key, init_name in qa.items():
|
||||
init = inits[init_name]
|
||||
arr = nh.to_array(init)
|
||||
old = float(np.asarray(arr).ravel()[0])
|
||||
ratios.append(new[key] / old)
|
||||
# dtype preserved on purpose — see the module docstring.
|
||||
init.CopyFrom(nh.from_array(np.full(arr.shape, new[key], dtype=arr.dtype), init_name))
|
||||
r = np.array(ratios)
|
||||
print(f"\n {len(qa)} activation scales rewritten with amax/448")
|
||||
print(f" amax/old ratio: min {r.min():.1f}x median {np.median(r):.1f}x max {r.max():.1f}x")
|
||||
onnx.save(model, a.out, save_as_external_data=False)
|
||||
print(f" wrote {a.out}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,92 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Recompute the SAME-S fp8 ENCODER activation scales with AMAX.
|
||||
|
||||
quantize/enc_quant_sames.py sets each scale to
|
||||
|
||||
max(percentile(|activations|, 99.9) / 448, 1e-4)
|
||||
|
||||
where the percentile is taken over a 400,000-element RANDOM SUBSAMPLE per calibration clip — so
|
||||
the tail is discarded twice: once by the percentile and once by the subsampling. The SAME-S
|
||||
DECODER's own quantiser uses amax, and the two disagree by 26x across layers.
|
||||
|
||||
Same ORT capture as the original (same base ONNX, same 6 calibration clips, same MatMul selection
|
||||
rule) with `amax` in place of the percentile, computed exactly over the full tensor. Only the 13
|
||||
activation-scale initializers are rewritten; the quantised WEIGHTS are untouched.
|
||||
|
||||
Mapping is exact, not positional: the quantiser writes its scale as `{pfx}_asc` where
|
||||
pfx = matmul_node_name.strip("/").replace("/", "_").
|
||||
|
||||
usage: recalib_enc_sames_fp8.py --fp8 enc_sames_fp8_pad.onnx --out enc_sames_fp8_amax.onnx
|
||||
"""
|
||||
import argparse, os
|
||||
import numpy as np, onnx, onnxruntime as ort
|
||||
from onnx import TensorProto, helper, numpy_helper as nh
|
||||
E4M3 = 448.0
|
||||
WD = "/weka2/cj/clod/sames_fp8"
|
||||
|
||||
def warr(inits, prod, n):
|
||||
w = n.input[1]
|
||||
if w in inits: return w, nh.to_array(inits[w])
|
||||
p = prod.get(w)
|
||||
if p and p.op_type == "Transpose" and p.input and p.input[0] in inits:
|
||||
return p.input[0], nh.to_array(inits[p.input[0]])
|
||||
return None, None
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--fp8", required=True, help="the fp8 encoder ONNX to correct")
|
||||
ap.add_argument("--out", required=True)
|
||||
ap.add_argument("--base", default="", help="base ONNX to capture on (default: HF bf16 encoder)")
|
||||
ap.add_argument("--calib", default=f"{WD}/calib_audio_sames.npz")
|
||||
ap.add_argument("--work", default="/weka2/cj/tmp/_enc_amax.onnx")
|
||||
a = ap.parse_args()
|
||||
base = a.base
|
||||
if not base:
|
||||
from huggingface_hub import hf_hub_download
|
||||
base = hf_hub_download("stabilityai/stable-audio-3-optimized",
|
||||
"onnx/same-s/enc_dynamic_bf16.onnx")
|
||||
m = onnx.load(base, load_external_data=True); g = m.graph
|
||||
inits = {i.name: i for i in g.initializer}
|
||||
prod = {o: n for n in g.node for o in n.output}
|
||||
tin = {}
|
||||
for n in g.node:
|
||||
if n.op_type == "MatMul":
|
||||
ws, W = warr(inits, prod, n)
|
||||
if ws is not None and W.ndim == 2: tin[n.name] = n.input[0]
|
||||
have = {o.name for o in g.output}
|
||||
for t in set(tin.values()):
|
||||
if t not in have: g.output.append(helper.make_tensor_value_info(t, TensorProto.FLOAT, None))
|
||||
onnx.save(m, a.work, save_as_external_data=True, all_tensors_to_one_file=True,
|
||||
location=os.path.basename(a.work) + ".data", size_threshold=1024)
|
||||
so = ort.SessionOptions(); so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
|
||||
sess = ort.InferenceSession(a.work, so, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
|
||||
iname = sess.get_inputs()[0].name
|
||||
ks = list(tin); ts = [tin[k] for k in ks]
|
||||
aud = np.load(a.calib)["audio"].astype(np.float32)
|
||||
amax = {k: 0.0 for k in ks}
|
||||
for ci in range(aud.shape[0]):
|
||||
for k, act in zip(ks, sess.run(ts, {iname: aud[ci:ci + 1]})):
|
||||
amax[k] = max(amax[k], float(np.abs(act).max())) # exact, whole tensor
|
||||
print(f" captured clip {ci}", flush=True)
|
||||
del sess
|
||||
mq = onnx.load(a.fp8, load_external_data=True)
|
||||
qi = {i.name: i for i in mq.graph.initializer}
|
||||
new = {k.strip("/").replace("/", "_") + "_asc": v / E4M3 for k, v in amax.items()}
|
||||
hits, ratios = 0, []
|
||||
for name, init in qi.items():
|
||||
if name in new:
|
||||
old = float(np.asarray(nh.to_array(init)).ravel()[0])
|
||||
ratios.append(new[name] / old); hits += 1
|
||||
arr = nh.to_array(init)
|
||||
init.CopyFrom(nh.from_array(np.full(arr.shape, new[name], dtype=arr.dtype), name))
|
||||
if hits == 0:
|
||||
raise SystemExit("no activation scales matched — check the pfx mapping")
|
||||
r = np.array(ratios)
|
||||
print(f"\n {hits} activation scales rewritten with amax/448")
|
||||
print(f" amax/old ratio: min {r.min():.1f}x median {np.median(r):.1f}x max {r.max():.1f}x")
|
||||
print(f" clip points now: {min(new.values())*E4M3:.3f} .. {max(new.values())*E4M3:.3f}")
|
||||
onnx.save(mq, a.out, save_as_external_data=False)
|
||||
print(f" wrote {a.out}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user