Re-export with public ExecuTorch 1.5.1 (recipe in export/)
Browse files- README.md +26 -17
- export/_common.py +330 -0
- export/_whisper.py +662 -0
- export/recipe.json +844 -0
- export/requirements.txt +18 -0
- export/whisper-base.py +784 -0
- manifest.json +10 -15
- whisper_cross_kv.pte +2 -2
- whisper_decode_stream.pte +2 -2
- whisper_decoder_split.pte +0 -3
- whisper_encoder.pte +2 -2
- whisper_lmhead.pte +2 -2
- whisper_mel.pte +2 -2
README.md
CHANGED
|
@@ -20,8 +20,8 @@ the visitor's GPU. It powers the corresponding demo in the ExecuTorch WebGPU Pla
|
|
| 20 |
|
| 21 |
- **Base model:** [openai/whisper-base](https://huggingface.co/openai/whisper-base)
|
| 22 |
- **Format:** ExecuTorch `.pte` programs delegated to the WebGPU backend
|
| 23 |
-
- **Quantization:** int8 token
|
| 24 |
-
- **Total size:**
|
| 25 |
|
| 26 |
> **Note:** OpenAI released Whisper's code under MIT; the weights are distributed under Apache-2.0, as on the upstream Hugging Face model card.
|
| 27 |
|
|
@@ -44,28 +44,34 @@ Accuracy varies by language, accent and recording conditions, and the model can
|
|
| 44 |
|
| 45 |
## What changed in this conversion
|
| 46 |
|
| 47 |
-
-
|
| 48 |
-
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 49 |
|
| 50 |
## Inputs and outputs
|
| 51 |
|
| 52 |
-
- 16 kHz mono audio in
|
| 53 |
-
- `whisper_mel.pte`:
|
| 54 |
-
- `
|
|
|
|
|
|
|
| 55 |
|
| 56 |
## Files
|
| 57 |
|
| 58 |
| File | Size | SHA-256 |
|
| 59 |
|---|---|---|
|
| 60 |
| `vocab.json` | 816.0 KB | `8f680bba319e01a6…` |
|
| 61 |
-
| `whisper_cross_kv.pte` |
|
| 62 |
-
| `whisper_decode_stream.pte` |
|
| 63 |
-
| `whisper_decoder_split.pte` | 13.1 MB | `84d9427851420bf7…` |
|
| 64 |
| `whisper_embed.i8.bin` | 25.3 MB | `bc7f8d8336ad935e…` |
|
| 65 |
| `whisper_embed.scale.f32.bin` | 202.6 KB | `1f58851911611b96…` |
|
| 66 |
-
| `whisper_encoder.pte` | 78.6 MB | `
|
| 67 |
-
| `whisper_lmhead.pte` |
|
| 68 |
-
| `whisper_mel.pte` |
|
| 69 |
| `whisper_pos_embed.f32.bin` | 896.0 KB | `88fa1cdbf2b06f86…` |
|
| 70 |
|
| 71 |
Full checksums are in `manifest.json`. The playground verifies them after download.
|
|
@@ -73,8 +79,9 @@ Full checksums are in `manifest.json`. The playground verifies them after downlo
|
|
| 73 |
## Running it
|
| 74 |
|
| 75 |
These programs were exported with the ExecuTorch WebGPU backend
|
| 76 |
-
([backends/webgpu](https://github.com/pytorch/executorch/tree/main/backends/webgpu)) and
|
| 77 |
-
WebAssembly runtime
|
|
|
|
| 78 |
the playground is the supported way to run them. It downloads these files once, caches them in the browser,
|
| 79 |
and runs everything locally.
|
| 80 |
|
|
@@ -82,8 +89,10 @@ Requirements: a browser with WebGPU enabled.
|
|
| 82 |
|
| 83 |
## Provenance
|
| 84 |
|
| 85 |
-
-
|
| 86 |
-
-
|
|
|
|
|
|
|
| 87 |
|
| 88 |
## Citation
|
| 89 |
|
|
|
|
| 20 |
|
| 21 |
- **Base model:** [openai/whisper-base](https://huggingface.co/openai/whisper-base)
|
| 22 |
- **Format:** ExecuTorch `.pte` programs delegated to the WebGPU backend
|
| 23 |
+
- **Quantization:** mixed: int8 output layer (int8 per-channel weights, int8 activations with a calibrated static scale) and int8 token-embedding table; mel front end, encoder and decoder in fp32
|
| 24 |
+
- **Total size:** 229.6 MB
|
| 25 |
|
| 26 |
> **Note:** OpenAI released Whisper's code under MIT; the weights are distributed under Apache-2.0, as on the upstream Hugging Face model card.
|
| 27 |
|
|
|
|
| 44 |
|
| 45 |
## What changed in this conversion
|
| 46 |
|
| 47 |
+
- Exported as six programs (mel front end, encoder, cross-attention K/V, two decoder forms, LM head); this repo ships the five used for streaming decoding. The full-recompute decoder (`whisper_decoder_split.pte`) is built by the recipe but not included
|
| 48 |
+
- Mel front end, encoder, cross-attention K/V projections and decoder run in fp32
|
| 49 |
+
- LM head: int8 weights (per output channel, symmetric) and int8 activations with a static per-tensor scale calibrated on 3 LibriSpeech clips that are not part of the evaluation set
|
| 50 |
+
- Token embedding table stored as int8 per row (row absmax / 127) with fp32 scales; the positional table is fp32
|
| 51 |
+
- Log-mel normalization runs on the host, between `whisper_mel.pte` (power mel spectrogram) and `whisper_encoder.pte`
|
| 52 |
+
- The LM head is padded from 51,865 to 51,872 outputs; the 7 padding logits are -1e9
|
| 53 |
+
- On 70 short clips from the LibriSpeech clean validation set, greedy transcripts have the same word error rate as transformers fp32 greedy decoding (8.42%), and 54 of 70 are token-identical to it. Only English transcription was tested
|
| 54 |
|
| 55 |
## Inputs and outputs
|
| 56 |
|
| 57 |
+
- 16 kHz mono audio in windows of up to 30 s: zero-pad to 480,000 samples, reflect-pad 200 samples on each side, and cut 3000 frames of 400 samples at hop 160
|
| 58 |
+
- `whisper_mel.pte`: frames `[1,3000,400]` -> power mel spectrogram `[1,3000,80]` (Hann window, DFT and mel filters inside the program). On the host: `x = log10(max(p, 1e-10))`, `x = max(x, max(x) - 8)`, `x = (x + 4) / 4`, then transpose to `[1,80,3000]`
|
| 59 |
+
- `whisper_encoder.pte`: log-mel `[1,80,3000]` -> hidden states `[1,1500,512]`. `whisper_cross_kv.pte`: hidden states -> 12 tensors `[1,8,1500,64]` (k0, v0, ..., k5, v5), computed once per window
|
| 60 |
+
- `whisper_decode_stream.pte`: `forward(embedding [1,1,512], input_pos int64 [1], 12 cross K/V tensors)` -> hidden `[1,1,512]`, one token per call; the self-attention cache (448 positions) lives inside the program. `whisper_lmhead.pte`: hidden -> logits `[1,1,51872]`; take the argmax over the first 51,865
|
| 61 |
+
- Embedding input: row `id` of `whisper_embed.i8.bin` (int8 [51865,512]) times `whisper_embed.scale.f32.bin[id]` (float32 [51865]), plus row `p` of `whisper_pos_embed.f32.bin` (float32 [448,512]). English prefix: 50258, 50259, 50359, 50363; stop at 50257. Token text from `vocab.json`
|
| 62 |
|
| 63 |
## Files
|
| 64 |
|
| 65 |
| File | Size | SHA-256 |
|
| 66 |
|---|---|---|
|
| 67 |
| `vocab.json` | 816.0 KB | `8f680bba319e01a6…` |
|
| 68 |
+
| `whisper_cross_kv.pte` | 12.0 MB | `6e8078924168c21a…` |
|
| 69 |
+
| `whisper_decode_stream.pte` | 85.1 MB | `bc2f33ad7aa16405…` |
|
|
|
|
| 70 |
| `whisper_embed.i8.bin` | 25.3 MB | `bc7f8d8336ad935e…` |
|
| 71 |
| `whisper_embed.scale.f32.bin` | 202.6 KB | `1f58851911611b96…` |
|
| 72 |
+
| `whisper_encoder.pte` | 78.6 MB | `841e7702ba8fa735…` |
|
| 73 |
+
| `whisper_lmhead.pte` | 25.9 MB | `68577fa240d29923…` |
|
| 74 |
+
| `whisper_mel.pte` | 694.3 KB | `7cae318294879584…` |
|
| 75 |
| `whisper_pos_embed.f32.bin` | 896.0 KB | `88fa1cdbf2b06f86…` |
|
| 76 |
|
| 77 |
Full checksums are in `manifest.json`. The playground verifies them after download.
|
|
|
|
| 79 |
## Running it
|
| 80 |
|
| 81 |
These programs were exported with the ExecuTorch WebGPU backend
|
| 82 |
+
([backends/webgpu](https://github.com/pytorch/executorch/tree/main/backends/webgpu)) and run in a browser (headless Chromium, WebGPU on SwiftShader) with a development build of the ExecuTorch
|
| 83 |
+
WebGPU WebAssembly runtime, where every program matched the CPU results (max abs error 1.6e-3 or less) and 3 test
|
| 84 |
+
transcripts were token-identical to the CPU pipeline. That runtime build is not published yet; until it is,
|
| 85 |
the playground is the supported way to run them. It downloads these files once, caches them in the browser,
|
| 86 |
and runs everything locally.
|
| 87 |
|
|
|
|
| 89 |
|
| 90 |
## Provenance
|
| 91 |
|
| 92 |
+
- **Upstream:** `openai/whisper-base` at revision `e37978b90ca9030d5170a5c07aadb050351a65bb` (official)
|
| 93 |
+
- **Exported with:** executorch 1.5.1+cpu, torch 2.14.1+cpu, torchao 0.18.0+cpu, transformers 5.19.0 (executorch git `3b6068392324`)
|
| 94 |
+
- **Reproduce:** `python recipes/whisper-base.py --out out/whisper-base/fp32dec-lmq8 --precision fp32dec-lmq8` (recipe in `export/`)
|
| 95 |
+
- Re-running the recipe produces byte-identical files.
|
| 96 |
|
| 97 |
## Citation
|
| 98 |
|
export/_common.py
ADDED
|
@@ -0,0 +1,330 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Shared helpers for the ExecuTorch WebGPU export recipes in this directory.
|
| 2 |
+
|
| 3 |
+
Each recipe builds an eager model from a pinned upstream revision and calls
|
| 4 |
+
`export_vision` (or its own LLM path), which writes:
|
| 5 |
+
|
| 6 |
+
<out>/<name>_webgpu.pte program lowered with WebGPUPartitioner (delegate id "VulkanBackend")
|
| 7 |
+
<out>/<name>_portable.pte same exported graph without a delegate, for CPU validation only
|
| 8 |
+
<out>/recipe.json upstream id + revision, versions, shapes, dtype, command line
|
| 9 |
+
<out>/validation.json eager vs exported-program vs portable .pte (and goldens, if given)
|
| 10 |
+
"""
|
| 11 |
+
|
| 12 |
+
import hashlib
|
| 13 |
+
import io
|
| 14 |
+
import json
|
| 15 |
+
import platform
|
| 16 |
+
import shlex
|
| 17 |
+
import sys
|
| 18 |
+
import urllib.request
|
| 19 |
+
import zipfile
|
| 20 |
+
from pathlib import Path
|
| 21 |
+
|
| 22 |
+
import numpy as np
|
| 23 |
+
import torch
|
| 24 |
+
|
| 25 |
+
PROXY_HINT = (
|
| 26 |
+
"set https_proxy if the machine needs a proxy for huggingface.co / pypi.org"
|
| 27 |
+
)
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def versions():
|
| 31 |
+
import executorch
|
| 32 |
+
import torchao
|
| 33 |
+
from executorch.version import git_version
|
| 34 |
+
|
| 35 |
+
out = {
|
| 36 |
+
"python": platform.python_version(),
|
| 37 |
+
"torch": torch.__version__,
|
| 38 |
+
"torchao": torchao.__version__,
|
| 39 |
+
"executorch": getattr(executorch, "__version__", None)
|
| 40 |
+
or __import__("executorch.version").version.__version__,
|
| 41 |
+
"executorch_git": git_version,
|
| 42 |
+
}
|
| 43 |
+
try:
|
| 44 |
+
import transformers
|
| 45 |
+
|
| 46 |
+
out["transformers"] = transformers.__version__
|
| 47 |
+
except ImportError:
|
| 48 |
+
pass
|
| 49 |
+
try:
|
| 50 |
+
import huggingface_hub
|
| 51 |
+
|
| 52 |
+
out["huggingface_hub"] = huggingface_hub.__version__
|
| 53 |
+
except ImportError:
|
| 54 |
+
pass
|
| 55 |
+
return out
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def sha256_file(path):
|
| 59 |
+
h = hashlib.sha256()
|
| 60 |
+
with open(path, "rb") as fh:
|
| 61 |
+
for chunk in iter(lambda: fh.read(1 << 20), b""):
|
| 62 |
+
h.update(chunk)
|
| 63 |
+
return h.hexdigest()
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def hf_file(repo_id, filename, revision, expected_sha256=None, cache_dir=None):
|
| 67 |
+
from huggingface_hub import hf_hub_download
|
| 68 |
+
|
| 69 |
+
path = hf_hub_download(
|
| 70 |
+
repo_id=repo_id, filename=filename, revision=revision, cache_dir=cache_dir
|
| 71 |
+
)
|
| 72 |
+
digest = sha256_file(path)
|
| 73 |
+
if expected_sha256 and digest != expected_sha256:
|
| 74 |
+
sys.exit(
|
| 75 |
+
f"{repo_id}@{revision}:{filename} sha256 {digest} != {expected_sha256}"
|
| 76 |
+
)
|
| 77 |
+
return Path(path), digest
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
def pypi_wheel_member(url, expected_sha256, member):
|
| 81 |
+
"""Read one file out of a PyPI wheel, verifying the wheel's sha256."""
|
| 82 |
+
with urllib.request.urlopen(url, timeout=120) as resp:
|
| 83 |
+
blob = resp.read()
|
| 84 |
+
digest = hashlib.sha256(blob).hexdigest()
|
| 85 |
+
if digest != expected_sha256:
|
| 86 |
+
sys.exit(f"{url}: sha256 {digest} != {expected_sha256} ({PROXY_HINT})")
|
| 87 |
+
with zipfile.ZipFile(io.BytesIO(blob)) as zf:
|
| 88 |
+
return zf.read(member).decode()
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
def scrub(text, out_dir):
|
| 92 |
+
"""Replace machine-specific paths so recipe.json can ship with the model."""
|
| 93 |
+
text = text.replace(str(Path(out_dir).resolve()), "<out>")
|
| 94 |
+
text = text.replace(str(Path(__file__).resolve().parent), "recipes")
|
| 95 |
+
return text.replace(str(Path.home()), "~")
|
| 96 |
+
|
| 97 |
+
|
| 98 |
+
def command_line(out_dir):
|
| 99 |
+
argv = [f"recipes/{Path(sys.argv[0]).name}", *sys.argv[1:]]
|
| 100 |
+
args = []
|
| 101 |
+
for a in argv:
|
| 102 |
+
if a.startswith("/") or a.startswith("."):
|
| 103 |
+
a = scrub(str(Path(a).resolve()), out_dir)
|
| 104 |
+
args.append(a)
|
| 105 |
+
return scrub(" ".join(["python", *(shlex.quote(a) for a in args)]), out_dir)
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
def lower_webgpu(ep):
|
| 109 |
+
from executorch.backends.webgpu.partitioner import WebGPUPartitioner
|
| 110 |
+
from executorch.exir import to_edge_transform_and_lower
|
| 111 |
+
|
| 112 |
+
return to_edge_transform_and_lower(
|
| 113 |
+
ep, partitioner=[WebGPUPartitioner()]
|
| 114 |
+
).to_executorch()
|
| 115 |
+
|
| 116 |
+
|
| 117 |
+
def lower_portable(ep):
|
| 118 |
+
from executorch.exir import to_edge_transform_and_lower
|
| 119 |
+
|
| 120 |
+
return to_edge_transform_and_lower(ep).to_executorch()
|
| 121 |
+
|
| 122 |
+
|
| 123 |
+
def delegate_summary(pte_bytes):
|
| 124 |
+
"""Non-delegated kernels and per-delegate op histograms of a serialized program."""
|
| 125 |
+
import collections
|
| 126 |
+
|
| 127 |
+
from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
|
| 128 |
+
extract_vk_flatbuffer,
|
| 129 |
+
flatbuffer_to_vk_graph,
|
| 130 |
+
)
|
| 131 |
+
from executorch.exir._serialize._program import deserialize_pte_binary
|
| 132 |
+
|
| 133 |
+
pte = deserialize_pte_binary(pte_bytes)
|
| 134 |
+
prog = pte.program if hasattr(pte, "program") else pte
|
| 135 |
+
seg_base = int.from_bytes(pte_bytes[24:32], "little")
|
| 136 |
+
plan = prog.execution_plan[0]
|
| 137 |
+
ops = [f"{o.name}.{o.overload}" for o in plan.operators]
|
| 138 |
+
kernels = collections.Counter()
|
| 139 |
+
for chain in plan.chains:
|
| 140 |
+
for ins in chain.instructions:
|
| 141 |
+
if type(ins.instr_args).__name__ == "KernelCall":
|
| 142 |
+
kernels[ops[ins.instr_args.op_index]] += 1
|
| 143 |
+
delegates = []
|
| 144 |
+
for d in plan.delegates:
|
| 145 |
+
proc = d.processed
|
| 146 |
+
if proc.location.name == "INLINE":
|
| 147 |
+
blob = prog.backend_delegate_data[proc.index].data
|
| 148 |
+
else:
|
| 149 |
+
seg = prog.segments[proc.index]
|
| 150 |
+
blob = pte_bytes[seg_base + seg.offset : seg_base + seg.offset + seg.size]
|
| 151 |
+
g = flatbuffer_to_vk_graph(extract_vk_flatbuffer(blob))
|
| 152 |
+
delegates.append(
|
| 153 |
+
{
|
| 154 |
+
"id": d.id,
|
| 155 |
+
"ops": dict(collections.Counter(c.name for c in g.chain)),
|
| 156 |
+
"inputs": len(g.input_ids),
|
| 157 |
+
"outputs": len(g.output_ids),
|
| 158 |
+
}
|
| 159 |
+
)
|
| 160 |
+
return {"non_delegated_kernels": dict(kernels), "delegates": delegates}
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
def compare(a, b):
|
| 164 |
+
a = np.asarray(a, dtype=np.float64).ravel()
|
| 165 |
+
b = np.asarray(b, dtype=np.float64).ravel()
|
| 166 |
+
n = min(a.size, b.size)
|
| 167 |
+
a, b = a[:n], b[:n]
|
| 168 |
+
cos = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30))
|
| 169 |
+
return {
|
| 170 |
+
"max_abs": float(np.max(np.abs(a - b))),
|
| 171 |
+
"mean_abs": float(np.mean(np.abs(a - b))),
|
| 172 |
+
"cosine": cos,
|
| 173 |
+
}
|
| 174 |
+
|
| 175 |
+
|
| 176 |
+
def run_pte(pte_bytes, inputs):
|
| 177 |
+
from executorch.extension.pybindings.portable_lib import (
|
| 178 |
+
_load_for_executorch_from_buffer,
|
| 179 |
+
)
|
| 180 |
+
|
| 181 |
+
m = _load_for_executorch_from_buffer(pte_bytes)
|
| 182 |
+
return [
|
| 183 |
+
o.detach().numpy() if isinstance(o, torch.Tensor) else o
|
| 184 |
+
for o in m.forward(inputs)
|
| 185 |
+
]
|
| 186 |
+
|
| 187 |
+
|
| 188 |
+
def try_run_webgpu_pte_on_cpu(pte_bytes, inputs):
|
| 189 |
+
try:
|
| 190 |
+
return run_pte(pte_bytes, inputs), None
|
| 191 |
+
except Exception as e: # noqa: BLE001 - the error text is the result
|
| 192 |
+
return None, f"{type(e).__name__}: {str(e).splitlines()[0][:300]}"
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
def export_vision(
|
| 196 |
+
*,
|
| 197 |
+
name,
|
| 198 |
+
model,
|
| 199 |
+
example_input,
|
| 200 |
+
out_dir,
|
| 201 |
+
upstream,
|
| 202 |
+
extra_recipe,
|
| 203 |
+
golden=None,
|
| 204 |
+
output_index=None,
|
| 205 |
+
):
|
| 206 |
+
"""Export, lower, write both .pte files, recipe.json, and validation.json.
|
| 207 |
+
|
| 208 |
+
`golden` is an optional (input ndarray, output ndarray) pair; when its input
|
| 209 |
+
shape matches `example_input`, all paths are also compared against it.
|
| 210 |
+
`output_index` picks one tensor when the eager model returns a tuple.
|
| 211 |
+
"""
|
| 212 |
+
out_dir = Path(out_dir)
|
| 213 |
+
out_dir.mkdir(parents=True, exist_ok=True)
|
| 214 |
+
torch.manual_seed(0)
|
| 215 |
+
model = model.eval()
|
| 216 |
+
|
| 217 |
+
class Wrapped(torch.nn.Module):
|
| 218 |
+
def __init__(self, m):
|
| 219 |
+
super().__init__()
|
| 220 |
+
self.m = m
|
| 221 |
+
|
| 222 |
+
def forward(self, x):
|
| 223 |
+
y = self.m(x)
|
| 224 |
+
return y[output_index] if output_index is not None else y
|
| 225 |
+
|
| 226 |
+
wrapped = Wrapped(model).eval()
|
| 227 |
+
with torch.no_grad():
|
| 228 |
+
ep = torch.export.export(wrapped, (example_input,), strict=False)
|
| 229 |
+
webgpu = lower_webgpu(ep)
|
| 230 |
+
portable = lower_portable(ep)
|
| 231 |
+
|
| 232 |
+
webgpu_path = out_dir / f"{name}_webgpu.pte"
|
| 233 |
+
portable_path = out_dir / f"{name}_portable.pte"
|
| 234 |
+
webgpu_path.write_bytes(webgpu.buffer)
|
| 235 |
+
portable_path.write_bytes(portable.buffer)
|
| 236 |
+
|
| 237 |
+
summary = delegate_summary(webgpu.buffer)
|
| 238 |
+
out_shape = list(ep.module()(example_input).shape)
|
| 239 |
+
recipe = {
|
| 240 |
+
"model": name,
|
| 241 |
+
"upstream": upstream,
|
| 242 |
+
"versions": versions(),
|
| 243 |
+
"export": {
|
| 244 |
+
"api": "torch.export.export(strict=False) -> to_edge_transform_and_lower(partitioner=[WebGPUPartitioner()]) -> to_executorch()",
|
| 245 |
+
"partitioner": "executorch.backends.webgpu.partitioner.WebGPUPartitioner() (default compile options; wraps VulkanPartitioner)",
|
| 246 |
+
"inputs": [
|
| 247 |
+
{
|
| 248 |
+
"shape": list(example_input.shape),
|
| 249 |
+
"dtype": "float32",
|
| 250 |
+
"layout": "NCHW",
|
| 251 |
+
}
|
| 252 |
+
],
|
| 253 |
+
"outputs": [{"shape": out_shape, "dtype": "float32"}],
|
| 254 |
+
"dtype": "float32",
|
| 255 |
+
"quantization": "none (fp32)",
|
| 256 |
+
**extra_recipe,
|
| 257 |
+
},
|
| 258 |
+
"files": {
|
| 259 |
+
webgpu_path.name: {
|
| 260 |
+
"bytes": webgpu_path.stat().st_size,
|
| 261 |
+
"sha256": sha256_file(webgpu_path),
|
| 262 |
+
},
|
| 263 |
+
portable_path.name: {
|
| 264 |
+
"bytes": portable_path.stat().st_size,
|
| 265 |
+
"sha256": sha256_file(portable_path),
|
| 266 |
+
"note": "CPU-only validation artifact (portable kernels, no delegate); not for release",
|
| 267 |
+
},
|
| 268 |
+
},
|
| 269 |
+
"graph": summary,
|
| 270 |
+
"command": command_line(out_dir),
|
| 271 |
+
}
|
| 272 |
+
(out_dir / "recipe.json").write_text(json.dumps(recipe, indent=2) + "\n")
|
| 273 |
+
|
| 274 |
+
# Validation on CPU.
|
| 275 |
+
gen = torch.Generator().manual_seed(1234)
|
| 276 |
+
rand_input = torch.rand(example_input.shape, generator=gen)
|
| 277 |
+
cases = {"random": rand_input}
|
| 278 |
+
if golden is not None and tuple(golden[0].shape) == tuple(example_input.shape):
|
| 279 |
+
cases["golden"] = torch.from_numpy(golden[0].copy())
|
| 280 |
+
report = {
|
| 281 |
+
"cpu_runtime": "executorch pybindings (portable_lib) from the pip wheel",
|
| 282 |
+
"cases": {},
|
| 283 |
+
}
|
| 284 |
+
webgpu_cpu_error = None
|
| 285 |
+
for label, x in cases.items():
|
| 286 |
+
with torch.no_grad():
|
| 287 |
+
eager = wrapped(x).numpy()
|
| 288 |
+
exported = ep.module()(x).numpy()
|
| 289 |
+
port = run_pte(portable.buffer, (x,))[0]
|
| 290 |
+
res = {
|
| 291 |
+
"exported_vs_eager": compare(exported, eager),
|
| 292 |
+
"portable_pte_vs_eager": compare(port, eager),
|
| 293 |
+
}
|
| 294 |
+
wg, err = try_run_webgpu_pte_on_cpu(webgpu.buffer, (x,))
|
| 295 |
+
if wg is not None:
|
| 296 |
+
res["webgpu_pte_vs_eager"] = compare(wg[0], eager)
|
| 297 |
+
else:
|
| 298 |
+
webgpu_cpu_error = err
|
| 299 |
+
if label == "golden":
|
| 300 |
+
ref = golden[1]
|
| 301 |
+
res["internal_golden_vs_eager"] = compare(eager, ref)
|
| 302 |
+
res["internal_golden_vs_portable_pte"] = compare(port, ref)
|
| 303 |
+
res["golden_note"] = (
|
| 304 |
+
"reference output from the internal goldens.manifest.json"
|
| 305 |
+
)
|
| 306 |
+
report["cases"][label] = res
|
| 307 |
+
report["webgpu_pte_on_cpu"] = (
|
| 308 |
+
"ran"
|
| 309 |
+
if webgpu_cpu_error is None
|
| 310 |
+
else f"not runnable on CPU pybindings: {webgpu_cpu_error}"
|
| 311 |
+
)
|
| 312 |
+
(out_dir / "validation.json").write_text(json.dumps(report, indent=2) + "\n")
|
| 313 |
+
print(
|
| 314 |
+
json.dumps(
|
| 315 |
+
{"recipe": recipe["files"], "graph": summary, "validation": report},
|
| 316 |
+
indent=2,
|
| 317 |
+
)
|
| 318 |
+
)
|
| 319 |
+
return ep, webgpu, portable
|
| 320 |
+
|
| 321 |
+
|
| 322 |
+
def load_golden(manifest_path, key):
|
| 323 |
+
man = json.loads(Path(manifest_path).read_text())[key]["tensors"]
|
| 324 |
+
base = Path(manifest_path).parent
|
| 325 |
+
|
| 326 |
+
def ld(t):
|
| 327 |
+
return np.fromfile(base / t["file"], dtype=np.float32).reshape(t["shape"])
|
| 328 |
+
|
| 329 |
+
inp = man.get("input")
|
| 330 |
+
return ld(inp), ld(man["output"])
|
export/_whisper.py
ADDED
|
@@ -0,0 +1,662 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Whisper base building blocks for recipes/whisper-base.py.
|
| 2 |
+
|
| 3 |
+
Everything here is built from the pinned openai/whisper-base checkpoint loaded through transformers
|
| 4 |
+
(WhisperForConditionalGeneration); the program wrappers reuse its nn.Modules (weights are shared, not
|
| 5 |
+
copied) and only re-express the forward pass so that it lowers to one WebGPU delegate per program.
|
| 6 |
+
|
| 7 |
+
Program contract (same split and I/O as the playground's previous whisper-base export):
|
| 8 |
+
|
| 9 |
+
mel forward(frames[1,3000,400] f32) -> power_mel[1,3000,80] f32
|
| 10 |
+
frames: 30 s of 16 kHz audio, zero-padded to 480000 samples, reflect-padded by 200 on
|
| 11 |
+
both ends, 3000 windows of 400 samples at hop 160 (= whisper.js buildFrames).
|
| 12 |
+
power_mel = |rfft(hann_periodic * frame)|^2 @ mel_filters (no log; see host_log_mel).
|
| 13 |
+
encoder forward(log_mel[1,80,3000] f32) -> encoder_hidden[1,1500,512] f32 (fp32 weights)
|
| 14 |
+
cross_kv forward(encoder_hidden[1,1500,512]) -> 12 x [1,8,1500,64] f32, order k0,v0,...,k5,v5
|
| 15 |
+
(per decoder layer: encoder_attn.k_proj / v_proj, heads-first [B,H,S,D])
|
| 16 |
+
decoder_split forward(tok_embeds[1,64,512], k0, v0, ..., k5, v5) -> hidden[1,64,512]
|
| 17 |
+
tok_embeds: token embeddings WITHOUT positions for positions 0..64 (zero rows beyond
|
| 18 |
+
the live prefix); the program adds positions 0..63, runs causal self-attention over
|
| 19 |
+
the 64 rows and returns final-LayerNorm hidden states (row p predicts token p+1).
|
| 20 |
+
decode_stream forward(embed[1,1,512], input_pos[1] i64, k0, v0, ..., k5, v5) -> hidden[1,1,512]
|
| 21 |
+
embed = token embedding + positional embedding of `input_pos` (added on the host);
|
| 22 |
+
self-attention K/V cache [1,448,8,64] x 6 layers x {k,v} lives inside the program
|
| 23 |
+
(mutable buffers owned by the delegate); feed positions 0,1,2,... in order.
|
| 24 |
+
lmhead forward(hidden[1,1,512]) -> logits[1,1,51872] (tied to decoder.embed_tokens; rows
|
| 25 |
+
51865..51871 are padding with logit -1e9, see LMHeadProgram)
|
| 26 |
+
|
| 27 |
+
Precision (whisper-base.py --precision):
|
| 28 |
+
q4 (previous-export parity) cross_kv / decoder_split / decode_stream / lmhead linears are 4-bit groupwise
|
| 29 |
+
weight-only (torchao IntxWeightOnlyConfig int4, group 64, symmetric, "affine" qparams ->
|
| 30 |
+
et_vk.linear_q4gsw); int8 token-embedding table. mel and encoder are fp32.
|
| 31 |
+
fp32 every program fp32 (aten.linear on 2-D views); fp32 token-embedding table.
|
| 32 |
+
|
| 33 |
+
Host-side files (reproduced from the checkpoint by the recipe, consumed by the playground JS):
|
| 34 |
+
whisper_embed.i8.bin int8 [51865,512] round(embed_tokens / scale), scale = rowwise absmax/127 (q4)
|
| 35 |
+
whisper_embed.scale.f32.bin f32 [51865] the per-row scales (q4)
|
| 36 |
+
whisper_embed.f32.bin f32 [51865,512] embed_tokens (fp32)
|
| 37 |
+
whisper_pos_embed.f32.bin f32 [448,512] decoder.embed_positions.weight
|
| 38 |
+
vocab.json upstream tokenizer vocab (byte-level BPE, detokenization only)
|
| 39 |
+
"""
|
| 40 |
+
|
| 41 |
+
import copy
|
| 42 |
+
import io
|
| 43 |
+
|
| 44 |
+
import numpy as np
|
| 45 |
+
import torch
|
| 46 |
+
import torch.nn.functional as F
|
| 47 |
+
|
| 48 |
+
SAMPLE_RATE = 16000
|
| 49 |
+
N_FFT = 400
|
| 50 |
+
HOP = 160
|
| 51 |
+
FRAMES = 3000
|
| 52 |
+
N_SAMPLES = FRAMES * HOP
|
| 53 |
+
PAD = N_FFT // 2
|
| 54 |
+
N_FREQ = N_FFT // 2 + 1
|
| 55 |
+
N_MELS = 80
|
| 56 |
+
ENC_LEN = 1500
|
| 57 |
+
D_MODEL = 512
|
| 58 |
+
N_HEADS = 8
|
| 59 |
+
HEAD_DIM = 64
|
| 60 |
+
N_LAYERS = 6
|
| 61 |
+
MAX_DEC = 64 # decoder_split rows; the playground's decode budget (prefix + generated)
|
| 62 |
+
KV_LEN = 448 # decode_stream self-attention cache length (= max_target_positions)
|
| 63 |
+
VOCAB = 51865
|
| 64 |
+
VOCAB_PADDED = 51872 # lmhead rows (multiple of 8, = the previous export)
|
| 65 |
+
PAD_LOGIT = -1e9
|
| 66 |
+
|
| 67 |
+
TOK = {
|
| 68 |
+
"sot": 50258,
|
| 69 |
+
"eot": 50257,
|
| 70 |
+
"en": 50259,
|
| 71 |
+
"transcribe": 50359,
|
| 72 |
+
"notimestamps": 50363,
|
| 73 |
+
}
|
| 74 |
+
PREFIX_EN = [TOK["sot"], TOK["en"], TOK["transcribe"], TOK["notimestamps"]]
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
# --------------------------------------------------------------------------------------------
|
| 78 |
+
# Host-side math that the playground does in JS (kept here so validation runs the same steps).
|
| 79 |
+
# --------------------------------------------------------------------------------------------
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
def build_frames(pcm):
|
| 83 |
+
"""Mirror of whisper.js buildFrames(): [3000, 400] float32 windows (no window function)."""
|
| 84 |
+
pcm = np.asarray(pcm, dtype=np.float32)
|
| 85 |
+
base = np.zeros(N_SAMPLES + 2 * PAD, dtype=np.float32)
|
| 86 |
+
n = min(pcm.size, N_SAMPLES)
|
| 87 |
+
base[PAD : PAD + n] = pcm[:n]
|
| 88 |
+
for i in range(PAD):
|
| 89 |
+
base[PAD - 1 - i] = base[PAD + 1 + i]
|
| 90 |
+
base[PAD + N_SAMPLES + i] = base[PAD + N_SAMPLES - 2 - i]
|
| 91 |
+
idx = np.arange(FRAMES)[:, None] * HOP + np.arange(N_FFT)[None, :]
|
| 92 |
+
return base[idx]
|
| 93 |
+
|
| 94 |
+
|
| 95 |
+
def host_log_mel(power_mel):
|
| 96 |
+
"""Host step between the mel and encoder programs (WhisperFeatureExtractor normalization).
|
| 97 |
+
|
| 98 |
+
power_mel [1,3000,80] (or [3000,80]) -> log_mel [1,80,3000]:
|
| 99 |
+
x = log10(max(p, 1e-10)); x = max(x, max(x) - 8); x = (x + 4) / 4; transpose.
|
| 100 |
+
float32 throughout (the JS equivalent uses Float32Array + Math.log10).
|
| 101 |
+
"""
|
| 102 |
+
p = np.asarray(power_mel, dtype=np.float32).reshape(FRAMES, N_MELS)
|
| 103 |
+
x = np.log10(np.maximum(p, np.float32(1e-10))).astype(np.float32)
|
| 104 |
+
x = np.maximum(x, x.max() - np.float32(8.0))
|
| 105 |
+
x = ((x + np.float32(4.0)) / np.float32(4.0)).astype(np.float32)
|
| 106 |
+
return np.ascontiguousarray(x.T)[None]
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
class HostTables:
|
| 110 |
+
"""Token-embedding lookup as the playground JS does it, plus the files that carry the tables.
|
| 111 |
+
|
| 112 |
+
kind="int8": whisper_embed.i8.bin (int8 [51865,512], round(w / s)) + whisper_embed.scale.f32.bin
|
| 113 |
+
(s = rowwise absmax / 127): byte-identical to the previous export's files.
|
| 114 |
+
kind="fp32": whisper_embed.f32.bin (fp32 [51865,512], the checkpoint's embed_tokens).
|
| 115 |
+
Both: whisper_pos_embed.f32.bin (fp32 [448,512], decoder.embed_positions)."""
|
| 116 |
+
|
| 117 |
+
def __init__(self, model, kind):
|
| 118 |
+
e = model.model.decoder.embed_tokens.weight.detach().float()
|
| 119 |
+
self.kind = kind
|
| 120 |
+
self.pos = model.model.decoder.embed_positions.weight.detach().float().numpy()
|
| 121 |
+
if kind == "int8":
|
| 122 |
+
scale = e.abs().amax(dim=1) / 127.0
|
| 123 |
+
self.table = torch.round(e / scale[:, None]).clamp(-127, 127).to(torch.int8).numpy()
|
| 124 |
+
self.scale = scale.numpy().astype(np.float32)
|
| 125 |
+
elif kind == "fp32":
|
| 126 |
+
self.table = e.numpy().astype(np.float32)
|
| 127 |
+
self.scale = None
|
| 128 |
+
else:
|
| 129 |
+
raise ValueError(kind)
|
| 130 |
+
|
| 131 |
+
def embed(self, ids, position=None):
|
| 132 |
+
e = self.table[ids].astype(np.float32)
|
| 133 |
+
if self.scale is not None:
|
| 134 |
+
e = e * self.scale[ids][..., None]
|
| 135 |
+
if position is not None:
|
| 136 |
+
e = e + self.pos[position]
|
| 137 |
+
return e.astype(np.float32)
|
| 138 |
+
|
| 139 |
+
def files(self):
|
| 140 |
+
out = {}
|
| 141 |
+
if self.kind == "int8":
|
| 142 |
+
out["whisper_embed.i8.bin"] = self.table.tobytes()
|
| 143 |
+
out["whisper_embed.scale.f32.bin"] = self.scale.astype("<f4").tobytes()
|
| 144 |
+
else:
|
| 145 |
+
out["whisper_embed.f32.bin"] = self.table.astype("<f4").tobytes()
|
| 146 |
+
out["whisper_pos_embed.f32.bin"] = self.pos.astype("<f4").tobytes()
|
| 147 |
+
return out
|
| 148 |
+
|
| 149 |
+
|
| 150 |
+
# --------------------------------------------------------------------------------------------
|
| 151 |
+
# Program wrappers.
|
| 152 |
+
# --------------------------------------------------------------------------------------------
|
| 153 |
+
|
| 154 |
+
|
| 155 |
+
def dft_mel_weights(feature_extractor):
|
| 156 |
+
"""cos / -sin DFT rows with the periodic Hann window folded in, plus the HF mel filters."""
|
| 157 |
+
n = np.arange(N_FFT)
|
| 158 |
+
k = np.arange(N_FREQ)
|
| 159 |
+
window = 0.5 - 0.5 * np.cos(
|
| 160 |
+
2.0 * np.pi * n / N_FFT
|
| 161 |
+
) # periodic Hann (torch / HF default)
|
| 162 |
+
ang = 2.0 * np.pi * ((k[:, None] * n[None, :]) % N_FFT) / N_FFT
|
| 163 |
+
wc = (np.cos(ang) * window[None, :]).astype(np.float32)
|
| 164 |
+
ws = (-np.sin(ang) * window[None, :]).astype(np.float32)
|
| 165 |
+
mel = np.asarray(feature_extractor.mel_filters, dtype=np.float64).T.astype(
|
| 166 |
+
np.float32
|
| 167 |
+
)
|
| 168 |
+
assert mel.shape == (N_MELS, N_FREQ), mel.shape
|
| 169 |
+
return torch.from_numpy(wc), torch.from_numpy(ws), torch.from_numpy(mel)
|
| 170 |
+
|
| 171 |
+
|
| 172 |
+
class MelProgram(torch.nn.Module):
|
| 173 |
+
def __init__(self, feature_extractor):
|
| 174 |
+
super().__init__()
|
| 175 |
+
wc, ws, mel = dft_mel_weights(feature_extractor)
|
| 176 |
+
self.register_buffer("wc", wc)
|
| 177 |
+
self.register_buffer("ws", ws)
|
| 178 |
+
self.register_buffer("mel", mel)
|
| 179 |
+
|
| 180 |
+
def forward(self, frames):
|
| 181 |
+
x = frames.reshape(
|
| 182 |
+
FRAMES, N_FFT
|
| 183 |
+
) # aten.linear on the WebGPU runtime is 2-D only
|
| 184 |
+
re = F.linear(x, self.wc)
|
| 185 |
+
im = F.linear(x, self.ws)
|
| 186 |
+
return F.linear(re * re + im * im, self.mel).reshape(1, FRAMES, N_MELS)
|
| 187 |
+
|
| 188 |
+
|
| 189 |
+
def _heads(x, s):
|
| 190 |
+
"""[s, 512] or [1, s, 512] -> [1, 8, s, 64]."""
|
| 191 |
+
return x.reshape(1, s, N_HEADS, HEAD_DIM).transpose(1, 2)
|
| 192 |
+
|
| 193 |
+
|
| 194 |
+
def _merge(o, s):
|
| 195 |
+
"""[1, 8, s, 64] -> [1, s, 512]."""
|
| 196 |
+
return o.transpose(1, 2).reshape(1, s, D_MODEL)
|
| 197 |
+
|
| 198 |
+
|
| 199 |
+
def attention(q, k, v, mask, impl):
|
| 200 |
+
"""q/k/v [1,H,S,D]; scale = 1/sqrt(64) = 0.125 (a power of two, so applying it to q as HF does
|
| 201 |
+
or to q.k as the kernels do is bit-identical)."""
|
| 202 |
+
if impl == "et_vk":
|
| 203 |
+
import executorch.backends.vulkan.custom_ops_lib # noqa: F401 (torch.ops.et_vk.*)
|
| 204 |
+
|
| 205 |
+
return torch.ops.et_vk.sdpa(q, k, v, mask, None)
|
| 206 |
+
s = torch.matmul(q * 0.125, k.transpose(-2, -1))
|
| 207 |
+
if mask is not None:
|
| 208 |
+
s = s + mask
|
| 209 |
+
return torch.matmul(torch.softmax(s, dim=-1), v)
|
| 210 |
+
|
| 211 |
+
|
| 212 |
+
class EncoderProgram(torch.nn.Module):
|
| 213 |
+
"""HF WhisperEncoder forward re-expressed with 2-D fp32 linears.
|
| 214 |
+
|
| 215 |
+
attn="eager": matmul + softmax (lowered to aten.bmm/_softmax; activation memory is planned).
|
| 216 |
+
attn="et_vk": the fused et_vk.sdpa op (what the previous export used); on the public WebGPU
|
| 217 |
+
runtime each call allocates its own two [8,1500,1500] f32 scratch buffers (72 MB each).
|
| 218 |
+
"""
|
| 219 |
+
|
| 220 |
+
def __init__(self, model, attn="eager"):
|
| 221 |
+
super().__init__()
|
| 222 |
+
self.enc = model.model.encoder
|
| 223 |
+
self.attn = attn
|
| 224 |
+
|
| 225 |
+
def forward(self, log_mel):
|
| 226 |
+
e = self.enc
|
| 227 |
+
x = F.gelu(e.conv1(log_mel))
|
| 228 |
+
x = F.gelu(e.conv2(x)) # [1, 512, 1500]
|
| 229 |
+
x = x.permute(0, 2, 1).reshape(ENC_LEN, D_MODEL) + e.embed_positions.weight
|
| 230 |
+
for layer in e.layers:
|
| 231 |
+
a = layer.self_attn
|
| 232 |
+
h = layer.self_attn_layer_norm(x)
|
| 233 |
+
q = _heads(a.q_proj(h), ENC_LEN)
|
| 234 |
+
k = _heads(a.k_proj(h), ENC_LEN)
|
| 235 |
+
v = _heads(a.v_proj(h), ENC_LEN)
|
| 236 |
+
o = attention(q, k, v, None, "et_vk" if self.attn == "et_vk" else "torch")
|
| 237 |
+
x = x + a.out_proj(o.transpose(1, 2).reshape(ENC_LEN, D_MODEL))
|
| 238 |
+
h = layer.final_layer_norm(x)
|
| 239 |
+
x = x + layer.fc2(F.gelu(layer.fc1(h)))
|
| 240 |
+
return e.layer_norm(x).reshape(1, ENC_LEN, D_MODEL)
|
| 241 |
+
|
| 242 |
+
|
| 243 |
+
def _lin(m, x):
|
| 244 |
+
"""nn.Linear that lowers on the public WebGPU runtime: fp32 aten.linear only accepts 2-D
|
| 245 |
+
inputs, so fp32 weights get a [-1, K] view; int4 (q4gsw) linears accept any rank and keep it."""
|
| 246 |
+
if type(m.weight.data) is torch.Tensor and x.dim() != 2:
|
| 247 |
+
return m(x.reshape(-1, x.shape[-1])).reshape(*x.shape[:-1], m.out_features)
|
| 248 |
+
return m(x)
|
| 249 |
+
|
| 250 |
+
|
| 251 |
+
class CrossKVProgram(torch.nn.Module):
|
| 252 |
+
def __init__(self, decoder):
|
| 253 |
+
super().__init__()
|
| 254 |
+
self.layers = decoder.layers
|
| 255 |
+
|
| 256 |
+
def forward(self, enc):
|
| 257 |
+
out = []
|
| 258 |
+
for layer in self.layers:
|
| 259 |
+
a = layer.encoder_attn
|
| 260 |
+
out.append(_heads(_lin(a.k_proj, enc), ENC_LEN))
|
| 261 |
+
out.append(_heads(_lin(a.v_proj, enc), ENC_LEN))
|
| 262 |
+
return tuple(out)
|
| 263 |
+
|
| 264 |
+
|
| 265 |
+
def _mlp(layer, x):
|
| 266 |
+
return x + _lin(layer.fc2, F.gelu(_lin(layer.fc1, layer.final_layer_norm(x))))
|
| 267 |
+
|
| 268 |
+
|
| 269 |
+
def _cross(layer, x, k, v, s, impl):
|
| 270 |
+
a = layer.encoder_attn
|
| 271 |
+
h = layer.encoder_attn_layer_norm(x)
|
| 272 |
+
o = attention(_heads(_lin(a.q_proj, h), s), k, v, None, impl)
|
| 273 |
+
return x + _lin(a.out_proj, _merge(o, s))
|
| 274 |
+
|
| 275 |
+
|
| 276 |
+
class DecoderSplitProgram(torch.nn.Module):
|
| 277 |
+
"""Full-recompute decoder over MAX_DEC rows with a causal mask (no KV cache)."""
|
| 278 |
+
|
| 279 |
+
def __init__(self, decoder, attn="et_vk"):
|
| 280 |
+
super().__init__()
|
| 281 |
+
self.dec = decoder
|
| 282 |
+
self.attn = attn
|
| 283 |
+
self.register_buffer(
|
| 284 |
+
"pos", decoder.embed_positions.weight.detach()[:MAX_DEC].clone()
|
| 285 |
+
)
|
| 286 |
+
causal = torch.full((MAX_DEC, MAX_DEC), -1e30).triu(1) # same fill as the previous export
|
| 287 |
+
# et_vk.sdpa on the WebGPU runtime requires the full [B, H, S_q, S_kv] mask.
|
| 288 |
+
self.register_buffer(
|
| 289 |
+
"mask", causal.expand(1, N_HEADS, MAX_DEC, MAX_DEC).contiguous()
|
| 290 |
+
)
|
| 291 |
+
|
| 292 |
+
def forward(self, tok_embeds, *cross_kv):
|
| 293 |
+
x = tok_embeds + self.pos
|
| 294 |
+
for i, layer in enumerate(self.dec.layers):
|
| 295 |
+
a = layer.self_attn
|
| 296 |
+
h = layer.self_attn_layer_norm(x)
|
| 297 |
+
q = _heads(_lin(a.q_proj, h), MAX_DEC)
|
| 298 |
+
k = _heads(_lin(a.k_proj, h), MAX_DEC)
|
| 299 |
+
v = _heads(_lin(a.v_proj, h), MAX_DEC)
|
| 300 |
+
o = attention(q, k, v, self.mask, self.attn)
|
| 301 |
+
x = x + _lin(a.out_proj, _merge(o, MAX_DEC))
|
| 302 |
+
x = _cross(
|
| 303 |
+
layer, x, cross_kv[2 * i], cross_kv[2 * i + 1], MAX_DEC, self.attn
|
| 304 |
+
)
|
| 305 |
+
x = _mlp(layer, x)
|
| 306 |
+
return self.dec.layer_norm(x)
|
| 307 |
+
|
| 308 |
+
|
| 309 |
+
class DecodeStreamProgram(torch.nn.Module):
|
| 310 |
+
"""One token per call; self-attention K/V cache [1,448,8,64] per layer kept in the program.
|
| 311 |
+
|
| 312 |
+
Self-attention uses llama.update_cache + llama.custom_sdpa (executorch.extension.llm.custom_ops),
|
| 313 |
+
which the Vulkan partitioner fuses into sdpa_with_kv_cache with an et_vk.select_as_symint
|
| 314 |
+
position. Cross-attention uses et_vk.sdpa (attn="et_vk") or matmul/softmax (attn="torch").
|
| 315 |
+
"""
|
| 316 |
+
|
| 317 |
+
def __init__(self, decoder, attn="et_vk"):
|
| 318 |
+
super().__init__()
|
| 319 |
+
from executorch.extension.llm.custom_ops import custom_ops # noqa: F401
|
| 320 |
+
|
| 321 |
+
self.dec = decoder
|
| 322 |
+
self.attn = attn
|
| 323 |
+
for i in range(N_LAYERS):
|
| 324 |
+
self.register_buffer(
|
| 325 |
+
f"k_cache_{i}",
|
| 326 |
+
torch.zeros(1, KV_LEN, N_HEADS, HEAD_DIM),
|
| 327 |
+
persistent=False,
|
| 328 |
+
)
|
| 329 |
+
self.register_buffer(
|
| 330 |
+
f"v_cache_{i}",
|
| 331 |
+
torch.zeros(1, KV_LEN, N_HEADS, HEAD_DIM),
|
| 332 |
+
persistent=False,
|
| 333 |
+
)
|
| 334 |
+
|
| 335 |
+
def reset(self):
|
| 336 |
+
for i in range(N_LAYERS):
|
| 337 |
+
getattr(self, f"k_cache_{i}").zero_()
|
| 338 |
+
getattr(self, f"v_cache_{i}").zero_()
|
| 339 |
+
|
| 340 |
+
def forward(self, embed, input_pos, *cross_kv):
|
| 341 |
+
pos = input_pos[0].item()
|
| 342 |
+
torch._check_is_size(pos)
|
| 343 |
+
torch._check(pos < KV_LEN)
|
| 344 |
+
x = embed
|
| 345 |
+
for i, layer in enumerate(self.dec.layers):
|
| 346 |
+
a = layer.self_attn
|
| 347 |
+
h = layer.self_attn_layer_norm(x)
|
| 348 |
+
q = _lin(a.q_proj, h).view(1, 1, N_HEADS, HEAD_DIM)
|
| 349 |
+
k = _lin(a.k_proj, h).view(1, 1, N_HEADS, HEAD_DIM)
|
| 350 |
+
v = _lin(a.v_proj, h).view(1, 1, N_HEADS, HEAD_DIM)
|
| 351 |
+
kc = getattr(self, f"k_cache_{i}")
|
| 352 |
+
vc = getattr(self, f"v_cache_{i}")
|
| 353 |
+
torch.ops.llama.update_cache(k, kc, pos)
|
| 354 |
+
torch.ops.llama.update_cache(v, vc, pos)
|
| 355 |
+
o = torch.ops.llama.custom_sdpa(q, kc, vc, pos, None, 0.0, True, None)
|
| 356 |
+
x = x + _lin(a.out_proj, o.reshape(1, 1, D_MODEL))
|
| 357 |
+
x = _cross(layer, x, cross_kv[2 * i], cross_kv[2 * i + 1], 1, self.attn)
|
| 358 |
+
x = _mlp(layer, x)
|
| 359 |
+
return self.dec.layer_norm(x)
|
| 360 |
+
|
| 361 |
+
|
| 362 |
+
class LMHeadProgram(torch.nn.Module):
|
| 363 |
+
"""Tied output projection. The vocab is padded to `pad_to` rows: executorch 1.5.1's q4gsw
|
| 364 |
+
lowering pads the scales to a multiple of 4 output channels and then fails its own shape
|
| 365 |
+
check for N = 51865, so N must already be a multiple of 4 (the previous export used 51872).
|
| 366 |
+
Padded rows have zero weights and a -1e9 bias, so they can never win an argmax."""
|
| 367 |
+
|
| 368 |
+
def __init__(self, model, pad_to=0):
|
| 369 |
+
super().__init__()
|
| 370 |
+
w = model.model.decoder.embed_tokens.weight.detach().clone()
|
| 371 |
+
n = w.shape[0]
|
| 372 |
+
b = torch.zeros(max(n, pad_to))
|
| 373 |
+
if pad_to and pad_to > n:
|
| 374 |
+
w = torch.cat([w, w.new_zeros(pad_to - n, w.shape[1])])
|
| 375 |
+
b[n:] = PAD_LOGIT
|
| 376 |
+
self.proj = torch.nn.Linear(D_MODEL, w.shape[0], bias=True)
|
| 377 |
+
self.proj.weight = torch.nn.Parameter(w, requires_grad=False)
|
| 378 |
+
self.proj.bias = torch.nn.Parameter(b, requires_grad=False)
|
| 379 |
+
|
| 380 |
+
def forward(self, hidden):
|
| 381 |
+
return _lin(self.proj, hidden)
|
| 382 |
+
|
| 383 |
+
|
| 384 |
+
# --------------------------------------------------------------------------------------------
|
| 385 |
+
# Quantization.
|
| 386 |
+
# --------------------------------------------------------------------------------------------
|
| 387 |
+
|
| 388 |
+
Q4_DESC = (
|
| 389 |
+
"torchao.quantization.quant_api.IntxWeightOnlyConfig(weight_dtype=torch.int4, "
|
| 390 |
+
"granularity=PerGroup(64), intx_choose_qparams_algorithm='affine') applied with quantize_ to "
|
| 391 |
+
"every nn.Linear of the program; symmetric, scale = group absmax / 7.5, q in [-8, 7]"
|
| 392 |
+
)
|
| 393 |
+
|
| 394 |
+
|
| 395 |
+
def quantize_q4(module, group_size=64):
|
| 396 |
+
from torchao.quantization.granularity import PerGroup
|
| 397 |
+
from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_
|
| 398 |
+
|
| 399 |
+
quantize_(
|
| 400 |
+
module,
|
| 401 |
+
IntxWeightOnlyConfig(
|
| 402 |
+
weight_dtype=torch.int4,
|
| 403 |
+
granularity=PerGroup(group_size),
|
| 404 |
+
intx_choose_qparams_algorithm="affine",
|
| 405 |
+
),
|
| 406 |
+
)
|
| 407 |
+
return module
|
| 408 |
+
|
| 409 |
+
|
| 410 |
+
Q8_HEAD_DESC = (
|
| 411 |
+
"PT2E with executorch.backends.vulkan.quantizer.VulkanQuantizer: weights int8 per-output-channel "
|
| 412 |
+
"symmetric (PerChannelMinMaxObserver, the quantizer's own weight spec), input activation int8 "
|
| 413 |
+
"per-tensor affine STATIC (HistogramObserver) calibrated on lmhead inputs; lowers to "
|
| 414 |
+
"quantized_decomposed.quantize_per_tensor + et_vk.linear_q8ta_q8csw"
|
| 415 |
+
)
|
| 416 |
+
|
| 417 |
+
|
| 418 |
+
def q8_lmhead(lm, calib_hidden):
|
| 419 |
+
"""Static-activation int8 x per-channel int8 lmhead (the only 8-bit linear executorch 1.5.1's
|
| 420 |
+
Vulkan lowering emits that the public WebGPU runtime registers). Returns the converted GraphModule
|
| 421 |
+
and the (input scale, zero point, dequantized weight) needed for an equivalent portable module."""
|
| 422 |
+
from executorch.backends.vulkan.quantizer.vulkan_quantizer import (
|
| 423 |
+
get_symmetric_quantization_config,
|
| 424 |
+
VulkanQuantizer,
|
| 425 |
+
)
|
| 426 |
+
from torchao.quantization.pt2e import HistogramObserver
|
| 427 |
+
from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
|
| 428 |
+
from torchao.quantization.pt2e.quantizer import QuantizationConfig, QuantizationSpec
|
| 429 |
+
|
| 430 |
+
base = get_symmetric_quantization_config(is_dynamic=False, weight_bits=8)
|
| 431 |
+
act = QuantizationSpec(
|
| 432 |
+
dtype=torch.int8,
|
| 433 |
+
quant_min=-128,
|
| 434 |
+
quant_max=127,
|
| 435 |
+
qscheme=torch.per_tensor_affine,
|
| 436 |
+
is_dynamic=False,
|
| 437 |
+
observer_or_fake_quant_ctr=HistogramObserver,
|
| 438 |
+
)
|
| 439 |
+
cfg = QuantizationConfig(
|
| 440 |
+
input_activation=act, output_activation=None, weight=base.weight, bias=None, is_qat=False
|
| 441 |
+
)
|
| 442 |
+
with torch.no_grad():
|
| 443 |
+
gm = torch.export.export(lm.eval(), (calib_hidden[:1],), strict=False).module()
|
| 444 |
+
gm = prepare_pt2e(gm, VulkanQuantizer().set_global(cfg))
|
| 445 |
+
for i in range(calib_hidden.shape[0]):
|
| 446 |
+
gm(calib_hidden[i : i + 1])
|
| 447 |
+
gm = convert_pt2e(gm)
|
| 448 |
+
qp, wdq = None, None
|
| 449 |
+
for n in gm.graph.nodes:
|
| 450 |
+
if n.op != "call_function":
|
| 451 |
+
continue
|
| 452 |
+
name = str(n.target)
|
| 453 |
+
if "quantize_per_tensor" in name and "dequantize" not in name and qp is None:
|
| 454 |
+
qp = (float(n.args[1]), int(n.args[2]))
|
| 455 |
+
if "dequantize_per_channel" in name:
|
| 456 |
+
w = getattr(gm, n.args[0].target)
|
| 457 |
+
sc = getattr(gm, n.args[1].target)
|
| 458 |
+
wdq = (w.float() * sc.float()[:, None]).contiguous()
|
| 459 |
+
if qp is None or wdq is None:
|
| 460 |
+
raise RuntimeError("q8 lmhead: could not find the quantize/dequantize nodes")
|
| 461 |
+
return gm, qp, wdq
|
| 462 |
+
|
| 463 |
+
|
| 464 |
+
class FakeQ8LMHead(torch.nn.Module):
|
| 465 |
+
"""Portable-runtime equivalent of the q8 lmhead: per-tensor int8 fake-quant of the input, then
|
| 466 |
+
an fp32 linear on the dequantized int8 weights (same bias)."""
|
| 467 |
+
|
| 468 |
+
def __init__(self, lm, qp, wdq):
|
| 469 |
+
super().__init__()
|
| 470 |
+
self.scale, self.zp = qp
|
| 471 |
+
self.register_buffer("w", wdq)
|
| 472 |
+
self.register_buffer("b", lm.proj.bias.detach().clone())
|
| 473 |
+
|
| 474 |
+
def forward(self, hidden):
|
| 475 |
+
x = hidden.reshape(-1, D_MODEL)
|
| 476 |
+
q = torch.clamp(torch.round(x / self.scale) + self.zp, -128, 127)
|
| 477 |
+
x = (q - self.zp) * self.scale
|
| 478 |
+
return F.linear(x, self.w, self.b).reshape(*hidden.shape[:-1], self.w.shape[0])
|
| 479 |
+
|
| 480 |
+
|
| 481 |
+
def word_errors(tokenizer, hyp_text, ref_text):
|
| 482 |
+
"""(edits, reference words) after WhisperTokenizer.normalize (English text normalizer)."""
|
| 483 |
+
h = tokenizer.normalize(hyp_text).split()
|
| 484 |
+
r = tokenizer.normalize(ref_text).split()
|
| 485 |
+
d = list(range(len(r) + 1))
|
| 486 |
+
for i in range(1, len(h) + 1):
|
| 487 |
+
prev, d[0] = d[0], i
|
| 488 |
+
for j in range(1, len(r) + 1):
|
| 489 |
+
cur = d[j]
|
| 490 |
+
d[j] = min(d[j] + 1, d[j - 1] + 1, prev + (h[i - 1] != r[j - 1]))
|
| 491 |
+
prev = cur
|
| 492 |
+
return d[len(r)], len(r)
|
| 493 |
+
|
| 494 |
+
|
| 495 |
+
def dequantized_copy(module):
|
| 496 |
+
"""Same module with every quantized linear weight replaced by its fp32 dequantization.
|
| 497 |
+
|
| 498 |
+
Used only for the CPU portable .pte (the portable runtime has no torchao kernels); its numerics
|
| 499 |
+
are the q4 numerics (dequantize-then-fp32-matmul, as the q4gsw kernels compute)."""
|
| 500 |
+
m = copy.deepcopy(module)
|
| 501 |
+
for sub in m.modules():
|
| 502 |
+
if isinstance(sub, torch.nn.Linear) and hasattr(sub.weight, "dequantize"):
|
| 503 |
+
sub.weight = torch.nn.Parameter(
|
| 504 |
+
sub.weight.dequantize().float(), requires_grad=False
|
| 505 |
+
)
|
| 506 |
+
return m
|
| 507 |
+
|
| 508 |
+
|
| 509 |
+
# --------------------------------------------------------------------------------------------
|
| 510 |
+
# Speech sample.
|
| 511 |
+
# --------------------------------------------------------------------------------------------
|
| 512 |
+
|
| 513 |
+
|
| 514 |
+
def load_librispeech_dummy(parquet_path, max_sec=15.0, count=8, longer=False):
|
| 515 |
+
"""First `count` clips (in file order) no longer than max_sec seconds, as 16 kHz float32.
|
| 516 |
+
longer=True selects the clips *longer* than max_sec instead (used only for calibration, so
|
| 517 |
+
calibration and evaluation clips never overlap)."""
|
| 518 |
+
import pyarrow.parquet as pq
|
| 519 |
+
import soundfile as sf
|
| 520 |
+
|
| 521 |
+
table = pq.read_table(parquet_path).to_pylist()
|
| 522 |
+
clips = []
|
| 523 |
+
for row in table:
|
| 524 |
+
audio, sr = sf.read(io.BytesIO(row["audio"]["bytes"]), dtype="float32")
|
| 525 |
+
if sr != SAMPLE_RATE:
|
| 526 |
+
raise ValueError(f"unexpected sample rate {sr}")
|
| 527 |
+
if audio.ndim != 1:
|
| 528 |
+
audio = audio.mean(axis=1)
|
| 529 |
+
if (audio.size / SAMPLE_RATE > max_sec) != longer:
|
| 530 |
+
continue
|
| 531 |
+
clips.append({"id": row["id"], "text": row["text"], "pcm": audio})
|
| 532 |
+
if len(clips) >= count:
|
| 533 |
+
break
|
| 534 |
+
return clips
|
| 535 |
+
|
| 536 |
+
|
| 537 |
+
def seconds(pcm):
|
| 538 |
+
return round(len(pcm) / SAMPLE_RATE, 3)
|
| 539 |
+
|
| 540 |
+
|
| 541 |
+
def isclose_report(a, b):
|
| 542 |
+
a = np.asarray(a, dtype=np.float64).ravel()
|
| 543 |
+
b = np.asarray(b, dtype=np.float64).ravel()
|
| 544 |
+
d = np.abs(a - b)
|
| 545 |
+
cos = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30))
|
| 546 |
+
return {
|
| 547 |
+
"max_abs": float(d.max()),
|
| 548 |
+
"mean_abs": float(d.mean()),
|
| 549 |
+
"max_rel_to_absmax": float(d.max() / (np.abs(b).max() + 1e-30)),
|
| 550 |
+
"cosine": cos,
|
| 551 |
+
}
|
| 552 |
+
|
| 553 |
+
|
| 554 |
+
|
| 555 |
+
# --------------------------------------------------------------------------------------------
|
| 556 |
+
# Lowered-graph inspection (what the WebGPU runtime will see).
|
| 557 |
+
# --------------------------------------------------------------------------------------------
|
| 558 |
+
|
| 559 |
+
|
| 560 |
+
def _vk_graphs(pte_bytes):
|
| 561 |
+
from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
|
| 562 |
+
extract_vk_flatbuffer,
|
| 563 |
+
flatbuffer_to_vk_graph,
|
| 564 |
+
)
|
| 565 |
+
from executorch.exir._serialize._program import deserialize_pte_binary
|
| 566 |
+
|
| 567 |
+
pte = deserialize_pte_binary(pte_bytes)
|
| 568 |
+
prog = pte.program if hasattr(pte, "program") else pte
|
| 569 |
+
seg_base = int.from_bytes(pte_bytes[24:32], "little")
|
| 570 |
+
out = []
|
| 571 |
+
for plan in prog.execution_plan:
|
| 572 |
+
for d in plan.delegates:
|
| 573 |
+
proc = d.processed
|
| 574 |
+
if proc.location.name == "INLINE":
|
| 575 |
+
blob = prog.backend_delegate_data[proc.index].data
|
| 576 |
+
else:
|
| 577 |
+
seg = prog.segments[proc.index]
|
| 578 |
+
blob = pte_bytes[seg_base + seg.offset : seg_base + seg.offset + seg.size]
|
| 579 |
+
out.append(flatbuffer_to_vk_graph(extract_vk_flatbuffer(blob)))
|
| 580 |
+
return prog, out
|
| 581 |
+
|
| 582 |
+
|
| 583 |
+
def program_io(pte_bytes):
|
| 584 |
+
"""Method names and their input/output tensor shapes + dtypes from the serialized program."""
|
| 585 |
+
from executorch.exir._serialize._program import deserialize_pte_binary
|
| 586 |
+
|
| 587 |
+
pte = deserialize_pte_binary(pte_bytes)
|
| 588 |
+
prog = pte.program if hasattr(pte, "program") else pte
|
| 589 |
+
res = {}
|
| 590 |
+
for plan in prog.execution_plan:
|
| 591 |
+
|
| 592 |
+
def desc(i):
|
| 593 |
+
v = plan.values[i].val
|
| 594 |
+
if hasattr(v, "sizes"):
|
| 595 |
+
return {"shape": list(v.sizes), "dtype": v.scalar_type.name.lower()}
|
| 596 |
+
return {"kind": type(v).__name__}
|
| 597 |
+
|
| 598 |
+
res[plan.name] = {
|
| 599 |
+
"inputs": [desc(i) for i in plan.inputs],
|
| 600 |
+
"outputs": [desc(i) for i in plan.outputs],
|
| 601 |
+
}
|
| 602 |
+
return res
|
| 603 |
+
|
| 604 |
+
|
| 605 |
+
def vk_checks(pte_bytes, runtime_ops=None):
|
| 606 |
+
"""Checks against constraints of the public WebGPU runtime's op implementations
|
| 607 |
+
(backends/webgpu/runtime/ops): op registered, aten.linear 2-D, et_vk.sdpa mask shape,
|
| 608 |
+
sdpa_with_kv_cache without mask / with causal, fp32 tensors only."""
|
| 609 |
+
_, graphs = _vk_graphs(pte_bytes)
|
| 610 |
+
problems, notes = [], {}
|
| 611 |
+
for g in graphs:
|
| 612 |
+
vals = [v.value for v in g.values]
|
| 613 |
+
|
| 614 |
+
def t(i):
|
| 615 |
+
v = vals[i]
|
| 616 |
+
return v if type(v).__name__ == "VkTensor" else None
|
| 617 |
+
|
| 618 |
+
dtypes = sorted({type(v).__name__ if type(v).__name__ != "VkTensor" else v.datatype.name for v in vals})
|
| 619 |
+
notes["value_kinds"] = dtypes
|
| 620 |
+
for op in g.chain:
|
| 621 |
+
name = op.name
|
| 622 |
+
if runtime_ops is not None and name not in runtime_ops:
|
| 623 |
+
problems.append(f"{name}: not registered in the WebGPU runtime op list")
|
| 624 |
+
a = op.args
|
| 625 |
+
if name == "aten.linear.default":
|
| 626 |
+
ri, rw = len(t(a[0]).dims), len(t(a[1]).dims)
|
| 627 |
+
notes.setdefault("aten.linear input ranks", set()).add(ri)
|
| 628 |
+
if ri != 2 or rw != 2:
|
| 629 |
+
problems.append(f"aten.linear with input rank {ri}, weight rank {rw} (runtime needs 2-D)")
|
| 630 |
+
if name == "et_vk.sdpa.default":
|
| 631 |
+
q, k = t(a[0]), t(a[1])
|
| 632 |
+
m = t(a[3])
|
| 633 |
+
if m is not None:
|
| 634 |
+
need = q.dims[0] * q.dims[1] * q.dims[2] * k.dims[2]
|
| 635 |
+
got = int(np.prod(m.dims))
|
| 636 |
+
if need != got:
|
| 637 |
+
problems.append(f"et_vk.sdpa mask numel {got} != B*H*Sq*Skv {need}")
|
| 638 |
+
notes.setdefault("et_vk.sdpa q/k/mask", set()).add(
|
| 639 |
+
(tuple(q.dims), tuple(k.dims), tuple(m.dims) if m is not None else None)
|
| 640 |
+
)
|
| 641 |
+
if name == "sdpa_with_kv_cache.default":
|
| 642 |
+
mask_kind = type(vals[a[7]]).__name__
|
| 643 |
+
causal = vals[a[9]]
|
| 644 |
+
if mask_kind != "Null":
|
| 645 |
+
problems.append("sdpa_with_kv_cache with attn_mask (runtime rejects)")
|
| 646 |
+
if not (type(causal).__name__ == "Bool" and causal.bool_val):
|
| 647 |
+
problems.append(f"sdpa_with_kv_cache is_causal={causal} (runtime needs true)")
|
| 648 |
+
notes.setdefault("sdpa_with_kv_cache q/k_cache", set()).add(
|
| 649 |
+
(tuple(t(a[0]).dims), tuple(t(a[3]).dims), type(vals[a[5]]).__name__)
|
| 650 |
+
)
|
| 651 |
+
if name == "aten.convolution.default":
|
| 652 |
+
notes.setdefault("aten.convolution input/weight", set()).add(
|
| 653 |
+
(tuple(t(a[0]).dims), tuple(t(a[1]).dims))
|
| 654 |
+
)
|
| 655 |
+
for v in vals:
|
| 656 |
+
if type(v).__name__ == "VkTensor" and v.datatype.name not in ("FLOAT32", "INT32", "INT64", "UINT8", "INT8"):
|
| 657 |
+
problems.append(f"tensor dtype {v.datatype.name}")
|
| 658 |
+
break
|
| 659 |
+
for k, v in list(notes.items()):
|
| 660 |
+
if isinstance(v, set):
|
| 661 |
+
notes[k] = sorted(str(x) for x in v)
|
| 662 |
+
return {"problems": problems, "ok": not problems, "notes": notes}
|
export/recipe.json
ADDED
|
@@ -0,0 +1,844 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model": "whisper-base",
|
| 3 |
+
"precision": "fp32dec-lmq8",
|
| 4 |
+
"upstream": {
|
| 5 |
+
"model": "openai/whisper-base",
|
| 6 |
+
"revision": "e37978b90ca9030d5170a5c07aadb050351a65bb",
|
| 7 |
+
"official": true,
|
| 8 |
+
"license": "apache-2.0 (model card)",
|
| 9 |
+
"files_sha256": {
|
| 10 |
+
"model.safetensors": "07cadb9f25677c8d50df603e66a98fbd842cce45047139baeb16e6219a1e807b",
|
| 11 |
+
"config.json": "a153c53883a6799b6f056b4a8d1a515c9926d03994682ba88a7616618d7da0c1",
|
| 12 |
+
"generation_config.json": "444b3f636d2fff89dd9ecf549e2a085b61f7ff0fa0246d4628bac6a3b8cc9ba4",
|
| 13 |
+
"preprocessor_config.json": "9b5cd03a36fbb8a627c64d98a5b5b126ead95a77720723944487311f0110b666",
|
| 14 |
+
"vocab.json": "8f680bba319e01a653d2e8a5dbc17a9157179e0576e6ce74ce0c06356c6e24f9",
|
| 15 |
+
"tokenizer.json": "27fc476bfe7f17299480be2273fc0608e4d5a99aba2ab5dec5374b4482d1a566",
|
| 16 |
+
"tokenizer_config.json": "2a4c4281cf9f51ac6ccc406fdc711a087afe6530f671fa7b80953edc498275ce",
|
| 17 |
+
"merges.txt": "2df2990a395e35e8dfbc7511e08c12d56018d8d04691e0133e5d63b21e154dc6",
|
| 18 |
+
"normalizer.json": "bf1c507dc8724ca9cf9903640dacfb69dae2f00edee4f21ceba106a7392f26dd",
|
| 19 |
+
"added_tokens.json": "9715fd2243b6f06a5858b5e32950d2853f73dd5bc201aafcf76f5082a2d8acd1",
|
| 20 |
+
"special_tokens_map.json": "e67ae3a0aaa99abcd9f187138e12db1f65c16a14761c50ef10eef2c174a7a691"
|
| 21 |
+
},
|
| 22 |
+
"loader": "transformers.WhisperForConditionalGeneration.from_pretrained(dtype=float32, attn_implementation='eager')"
|
| 23 |
+
},
|
| 24 |
+
"versions": {
|
| 25 |
+
"python": "3.12.15+meta",
|
| 26 |
+
"torch": "2.14.1+cpu",
|
| 27 |
+
"torchao": "0.18.0+cpu",
|
| 28 |
+
"executorch": "1.5.1+cpu",
|
| 29 |
+
"executorch_git": "3b60683923245cf472b7323426920e15623ba361",
|
| 30 |
+
"transformers": "5.19.0",
|
| 31 |
+
"huggingface_hub": "2.2.0"
|
| 32 |
+
},
|
| 33 |
+
"export": {
|
| 34 |
+
"api": "torch.export.export(strict=False) -> to_edge_transform_and_lower(partitioner=[WebGPUPartitioner()]) -> to_executorch()",
|
| 35 |
+
"partitioner": "executorch.backends.webgpu.partitioner.WebGPUPartitioner() (default compile options; wraps VulkanPartitioner; delegate id VulkanBackend)",
|
| 36 |
+
"model_code": "recipes/_whisper.py wrappers around the transformers modules (weights shared, forward re-expressed): MelProgram, EncoderProgram, CrossKVProgram, DecoderSplitProgram, DecodeStreamProgram, LMHeadProgram",
|
| 37 |
+
"quantization": {
|
| 38 |
+
"decoder_side_linears": {
|
| 39 |
+
"programs": [
|
| 40 |
+
"cross_kv",
|
| 41 |
+
"decoder_split",
|
| 42 |
+
"decode_stream"
|
| 43 |
+
],
|
| 44 |
+
"config": "none (fp32)",
|
| 45 |
+
"lowered_to": "aten.linear (2-D fp32)"
|
| 46 |
+
},
|
| 47 |
+
"lmhead": {
|
| 48 |
+
"config": "PT2E with executorch.backends.vulkan.quantizer.VulkanQuantizer: weights int8 per-output-channel symmetric (PerChannelMinMaxObserver, the quantizer's own weight spec), input activation int8 per-tensor affine STATIC (HistogramObserver) calibrated on lmhead inputs; lowers to quantized_decomposed.quantize_per_tensor + et_vk.linear_q8ta_q8csw",
|
| 49 |
+
"lowered_to": "quantized_decomposed.quantize_per_tensor + et_vk.linear_q8ta_q8csw",
|
| 50 |
+
"calibration_clips": [
|
| 51 |
+
"1272-128104-0004",
|
| 52 |
+
"1272-128104-0009",
|
| 53 |
+
"1272-128104-0011"
|
| 54 |
+
],
|
| 55 |
+
"calibration_vectors": 157,
|
| 56 |
+
"input_scale": 1.7133680582046509,
|
| 57 |
+
"input_zero_point": 84
|
| 58 |
+
},
|
| 59 |
+
"fp32_programs": [
|
| 60 |
+
"mel",
|
| 61 |
+
"encoder"
|
| 62 |
+
],
|
| 63 |
+
"embedding_table": "int8 rowwise (whisper_embed.i8.bin + whisper_embed.scale.f32.bin)"
|
| 64 |
+
},
|
| 65 |
+
"encoder_attention": "eager",
|
| 66 |
+
"decoder_attention": "et_vk.sdpa (decoder_split self-attention with a [1,8,64,64] additive causal mask of 0 / -1e30; cross-attention in decoder_split and decode_stream); decode_stream self-attention = llama.update_cache + llama.custom_sdpa, fused by the Vulkan partitioner into sdpa_with_kv_cache + et_vk.select_as_symint",
|
| 67 |
+
"lmhead_rows": "51872: vocab 51865 padded with zero weights and bias -1000000000.0 (executorch 1.5.1's q4gsw lowering fails for N % 4 != 0; kept for fp32 too so the contract is identical)",
|
| 68 |
+
"constants": {
|
| 69 |
+
"max_decoder_rows": 64,
|
| 70 |
+
"self_attention_cache": [
|
| 71 |
+
1,
|
| 72 |
+
448,
|
| 73 |
+
8,
|
| 74 |
+
64
|
| 75 |
+
],
|
| 76 |
+
"frames": [
|
| 77 |
+
3000,
|
| 78 |
+
400
|
| 79 |
+
],
|
| 80 |
+
"hop": 160,
|
| 81 |
+
"prefix_en": [
|
| 82 |
+
50258,
|
| 83 |
+
50259,
|
| 84 |
+
50359,
|
| 85 |
+
50363
|
| 86 |
+
]
|
| 87 |
+
},
|
| 88 |
+
"host_contract": {
|
| 89 |
+
"frames": "whisper.js buildFrames: zero-pad pcm to 480000 samples, reflect-pad 200 each side, 3000 x 400 windows at hop 160",
|
| 90 |
+
"log_mel": "between mel and encoder: x = log10(max(p, 1e-10)); x = max(x, max(x) - 8); x = (x + 4) / 4; transpose [1,3000,80] -> [1,80,3000]",
|
| 91 |
+
"embeddings": "i8[id] * scale[id]; decoder_split: rows 0..p = token embeddings (no positions), zero rows after; decode_stream: token embedding + pos_embed[p]",
|
| 92 |
+
"argmax": "over logits[..., :51865] (rows 51865..51871 are -1000000000.0)",
|
| 93 |
+
"cross_kv": "computed once per 30 s window, then passed as 12 inputs to every decoder call"
|
| 94 |
+
},
|
| 95 |
+
"io": {
|
| 96 |
+
"mel": {
|
| 97 |
+
"forward": {
|
| 98 |
+
"inputs": [
|
| 99 |
+
{
|
| 100 |
+
"shape": [
|
| 101 |
+
1,
|
| 102 |
+
3000,
|
| 103 |
+
400
|
| 104 |
+
],
|
| 105 |
+
"dtype": "float"
|
| 106 |
+
}
|
| 107 |
+
],
|
| 108 |
+
"outputs": [
|
| 109 |
+
{
|
| 110 |
+
"shape": [
|
| 111 |
+
1,
|
| 112 |
+
3000,
|
| 113 |
+
80
|
| 114 |
+
],
|
| 115 |
+
"dtype": "float"
|
| 116 |
+
}
|
| 117 |
+
]
|
| 118 |
+
}
|
| 119 |
+
},
|
| 120 |
+
"encoder": {
|
| 121 |
+
"forward": {
|
| 122 |
+
"inputs": [
|
| 123 |
+
{
|
| 124 |
+
"shape": [
|
| 125 |
+
1,
|
| 126 |
+
80,
|
| 127 |
+
3000
|
| 128 |
+
],
|
| 129 |
+
"dtype": "float"
|
| 130 |
+
}
|
| 131 |
+
],
|
| 132 |
+
"outputs": [
|
| 133 |
+
{
|
| 134 |
+
"shape": [
|
| 135 |
+
1,
|
| 136 |
+
1500,
|
| 137 |
+
512
|
| 138 |
+
],
|
| 139 |
+
"dtype": "float"
|
| 140 |
+
}
|
| 141 |
+
]
|
| 142 |
+
}
|
| 143 |
+
},
|
| 144 |
+
"cross_kv": {
|
| 145 |
+
"forward": {
|
| 146 |
+
"inputs": [
|
| 147 |
+
{
|
| 148 |
+
"shape": [
|
| 149 |
+
1,
|
| 150 |
+
1500,
|
| 151 |
+
512
|
| 152 |
+
],
|
| 153 |
+
"dtype": "float"
|
| 154 |
+
}
|
| 155 |
+
],
|
| 156 |
+
"outputs": [
|
| 157 |
+
{
|
| 158 |
+
"shape": [
|
| 159 |
+
1,
|
| 160 |
+
8,
|
| 161 |
+
1500,
|
| 162 |
+
64
|
| 163 |
+
],
|
| 164 |
+
"dtype": "float"
|
| 165 |
+
},
|
| 166 |
+
{
|
| 167 |
+
"shape": [
|
| 168 |
+
1,
|
| 169 |
+
8,
|
| 170 |
+
1500,
|
| 171 |
+
64
|
| 172 |
+
],
|
| 173 |
+
"dtype": "float"
|
| 174 |
+
},
|
| 175 |
+
{
|
| 176 |
+
"shape": [
|
| 177 |
+
1,
|
| 178 |
+
8,
|
| 179 |
+
1500,
|
| 180 |
+
64
|
| 181 |
+
],
|
| 182 |
+
"dtype": "float"
|
| 183 |
+
},
|
| 184 |
+
{
|
| 185 |
+
"shape": [
|
| 186 |
+
1,
|
| 187 |
+
8,
|
| 188 |
+
1500,
|
| 189 |
+
64
|
| 190 |
+
],
|
| 191 |
+
"dtype": "float"
|
| 192 |
+
},
|
| 193 |
+
{
|
| 194 |
+
"shape": [
|
| 195 |
+
1,
|
| 196 |
+
8,
|
| 197 |
+
1500,
|
| 198 |
+
64
|
| 199 |
+
],
|
| 200 |
+
"dtype": "float"
|
| 201 |
+
},
|
| 202 |
+
{
|
| 203 |
+
"shape": [
|
| 204 |
+
1,
|
| 205 |
+
8,
|
| 206 |
+
1500,
|
| 207 |
+
64
|
| 208 |
+
],
|
| 209 |
+
"dtype": "float"
|
| 210 |
+
},
|
| 211 |
+
{
|
| 212 |
+
"shape": [
|
| 213 |
+
1,
|
| 214 |
+
8,
|
| 215 |
+
1500,
|
| 216 |
+
64
|
| 217 |
+
],
|
| 218 |
+
"dtype": "float"
|
| 219 |
+
},
|
| 220 |
+
{
|
| 221 |
+
"shape": [
|
| 222 |
+
1,
|
| 223 |
+
8,
|
| 224 |
+
1500,
|
| 225 |
+
64
|
| 226 |
+
],
|
| 227 |
+
"dtype": "float"
|
| 228 |
+
},
|
| 229 |
+
{
|
| 230 |
+
"shape": [
|
| 231 |
+
1,
|
| 232 |
+
8,
|
| 233 |
+
1500,
|
| 234 |
+
64
|
| 235 |
+
],
|
| 236 |
+
"dtype": "float"
|
| 237 |
+
},
|
| 238 |
+
{
|
| 239 |
+
"shape": [
|
| 240 |
+
1,
|
| 241 |
+
8,
|
| 242 |
+
1500,
|
| 243 |
+
64
|
| 244 |
+
],
|
| 245 |
+
"dtype": "float"
|
| 246 |
+
},
|
| 247 |
+
{
|
| 248 |
+
"shape": [
|
| 249 |
+
1,
|
| 250 |
+
8,
|
| 251 |
+
1500,
|
| 252 |
+
64
|
| 253 |
+
],
|
| 254 |
+
"dtype": "float"
|
| 255 |
+
},
|
| 256 |
+
{
|
| 257 |
+
"shape": [
|
| 258 |
+
1,
|
| 259 |
+
8,
|
| 260 |
+
1500,
|
| 261 |
+
64
|
| 262 |
+
],
|
| 263 |
+
"dtype": "float"
|
| 264 |
+
}
|
| 265 |
+
]
|
| 266 |
+
}
|
| 267 |
+
},
|
| 268 |
+
"decoder_split": {
|
| 269 |
+
"forward": {
|
| 270 |
+
"inputs": [
|
| 271 |
+
{
|
| 272 |
+
"shape": [
|
| 273 |
+
1,
|
| 274 |
+
64,
|
| 275 |
+
512
|
| 276 |
+
],
|
| 277 |
+
"dtype": "float"
|
| 278 |
+
},
|
| 279 |
+
{
|
| 280 |
+
"shape": [
|
| 281 |
+
1,
|
| 282 |
+
8,
|
| 283 |
+
1500,
|
| 284 |
+
64
|
| 285 |
+
],
|
| 286 |
+
"dtype": "float"
|
| 287 |
+
},
|
| 288 |
+
{
|
| 289 |
+
"shape": [
|
| 290 |
+
1,
|
| 291 |
+
8,
|
| 292 |
+
1500,
|
| 293 |
+
64
|
| 294 |
+
],
|
| 295 |
+
"dtype": "float"
|
| 296 |
+
},
|
| 297 |
+
{
|
| 298 |
+
"shape": [
|
| 299 |
+
1,
|
| 300 |
+
8,
|
| 301 |
+
1500,
|
| 302 |
+
64
|
| 303 |
+
],
|
| 304 |
+
"dtype": "float"
|
| 305 |
+
},
|
| 306 |
+
{
|
| 307 |
+
"shape": [
|
| 308 |
+
1,
|
| 309 |
+
8,
|
| 310 |
+
1500,
|
| 311 |
+
64
|
| 312 |
+
],
|
| 313 |
+
"dtype": "float"
|
| 314 |
+
},
|
| 315 |
+
{
|
| 316 |
+
"shape": [
|
| 317 |
+
1,
|
| 318 |
+
8,
|
| 319 |
+
1500,
|
| 320 |
+
64
|
| 321 |
+
],
|
| 322 |
+
"dtype": "float"
|
| 323 |
+
},
|
| 324 |
+
{
|
| 325 |
+
"shape": [
|
| 326 |
+
1,
|
| 327 |
+
8,
|
| 328 |
+
1500,
|
| 329 |
+
64
|
| 330 |
+
],
|
| 331 |
+
"dtype": "float"
|
| 332 |
+
},
|
| 333 |
+
{
|
| 334 |
+
"shape": [
|
| 335 |
+
1,
|
| 336 |
+
8,
|
| 337 |
+
1500,
|
| 338 |
+
64
|
| 339 |
+
],
|
| 340 |
+
"dtype": "float"
|
| 341 |
+
},
|
| 342 |
+
{
|
| 343 |
+
"shape": [
|
| 344 |
+
1,
|
| 345 |
+
8,
|
| 346 |
+
1500,
|
| 347 |
+
64
|
| 348 |
+
],
|
| 349 |
+
"dtype": "float"
|
| 350 |
+
},
|
| 351 |
+
{
|
| 352 |
+
"shape": [
|
| 353 |
+
1,
|
| 354 |
+
8,
|
| 355 |
+
1500,
|
| 356 |
+
64
|
| 357 |
+
],
|
| 358 |
+
"dtype": "float"
|
| 359 |
+
},
|
| 360 |
+
{
|
| 361 |
+
"shape": [
|
| 362 |
+
1,
|
| 363 |
+
8,
|
| 364 |
+
1500,
|
| 365 |
+
64
|
| 366 |
+
],
|
| 367 |
+
"dtype": "float"
|
| 368 |
+
},
|
| 369 |
+
{
|
| 370 |
+
"shape": [
|
| 371 |
+
1,
|
| 372 |
+
8,
|
| 373 |
+
1500,
|
| 374 |
+
64
|
| 375 |
+
],
|
| 376 |
+
"dtype": "float"
|
| 377 |
+
},
|
| 378 |
+
{
|
| 379 |
+
"shape": [
|
| 380 |
+
1,
|
| 381 |
+
8,
|
| 382 |
+
1500,
|
| 383 |
+
64
|
| 384 |
+
],
|
| 385 |
+
"dtype": "float"
|
| 386 |
+
}
|
| 387 |
+
],
|
| 388 |
+
"outputs": [
|
| 389 |
+
{
|
| 390 |
+
"shape": [
|
| 391 |
+
1,
|
| 392 |
+
64,
|
| 393 |
+
512
|
| 394 |
+
],
|
| 395 |
+
"dtype": "float"
|
| 396 |
+
}
|
| 397 |
+
]
|
| 398 |
+
}
|
| 399 |
+
},
|
| 400 |
+
"decode_stream": {
|
| 401 |
+
"forward": {
|
| 402 |
+
"inputs": [
|
| 403 |
+
{
|
| 404 |
+
"shape": [
|
| 405 |
+
1,
|
| 406 |
+
1,
|
| 407 |
+
512
|
| 408 |
+
],
|
| 409 |
+
"dtype": "float"
|
| 410 |
+
},
|
| 411 |
+
{
|
| 412 |
+
"shape": [
|
| 413 |
+
1
|
| 414 |
+
],
|
| 415 |
+
"dtype": "long"
|
| 416 |
+
},
|
| 417 |
+
{
|
| 418 |
+
"shape": [
|
| 419 |
+
1,
|
| 420 |
+
8,
|
| 421 |
+
1500,
|
| 422 |
+
64
|
| 423 |
+
],
|
| 424 |
+
"dtype": "float"
|
| 425 |
+
},
|
| 426 |
+
{
|
| 427 |
+
"shape": [
|
| 428 |
+
1,
|
| 429 |
+
8,
|
| 430 |
+
1500,
|
| 431 |
+
64
|
| 432 |
+
],
|
| 433 |
+
"dtype": "float"
|
| 434 |
+
},
|
| 435 |
+
{
|
| 436 |
+
"shape": [
|
| 437 |
+
1,
|
| 438 |
+
8,
|
| 439 |
+
1500,
|
| 440 |
+
64
|
| 441 |
+
],
|
| 442 |
+
"dtype": "float"
|
| 443 |
+
},
|
| 444 |
+
{
|
| 445 |
+
"shape": [
|
| 446 |
+
1,
|
| 447 |
+
8,
|
| 448 |
+
1500,
|
| 449 |
+
64
|
| 450 |
+
],
|
| 451 |
+
"dtype": "float"
|
| 452 |
+
},
|
| 453 |
+
{
|
| 454 |
+
"shape": [
|
| 455 |
+
1,
|
| 456 |
+
8,
|
| 457 |
+
1500,
|
| 458 |
+
64
|
| 459 |
+
],
|
| 460 |
+
"dtype": "float"
|
| 461 |
+
},
|
| 462 |
+
{
|
| 463 |
+
"shape": [
|
| 464 |
+
1,
|
| 465 |
+
8,
|
| 466 |
+
1500,
|
| 467 |
+
64
|
| 468 |
+
],
|
| 469 |
+
"dtype": "float"
|
| 470 |
+
},
|
| 471 |
+
{
|
| 472 |
+
"shape": [
|
| 473 |
+
1,
|
| 474 |
+
8,
|
| 475 |
+
1500,
|
| 476 |
+
64
|
| 477 |
+
],
|
| 478 |
+
"dtype": "float"
|
| 479 |
+
},
|
| 480 |
+
{
|
| 481 |
+
"shape": [
|
| 482 |
+
1,
|
| 483 |
+
8,
|
| 484 |
+
1500,
|
| 485 |
+
64
|
| 486 |
+
],
|
| 487 |
+
"dtype": "float"
|
| 488 |
+
},
|
| 489 |
+
{
|
| 490 |
+
"shape": [
|
| 491 |
+
1,
|
| 492 |
+
8,
|
| 493 |
+
1500,
|
| 494 |
+
64
|
| 495 |
+
],
|
| 496 |
+
"dtype": "float"
|
| 497 |
+
},
|
| 498 |
+
{
|
| 499 |
+
"shape": [
|
| 500 |
+
1,
|
| 501 |
+
8,
|
| 502 |
+
1500,
|
| 503 |
+
64
|
| 504 |
+
],
|
| 505 |
+
"dtype": "float"
|
| 506 |
+
},
|
| 507 |
+
{
|
| 508 |
+
"shape": [
|
| 509 |
+
1,
|
| 510 |
+
8,
|
| 511 |
+
1500,
|
| 512 |
+
64
|
| 513 |
+
],
|
| 514 |
+
"dtype": "float"
|
| 515 |
+
},
|
| 516 |
+
{
|
| 517 |
+
"shape": [
|
| 518 |
+
1,
|
| 519 |
+
8,
|
| 520 |
+
1500,
|
| 521 |
+
64
|
| 522 |
+
],
|
| 523 |
+
"dtype": "float"
|
| 524 |
+
}
|
| 525 |
+
],
|
| 526 |
+
"outputs": [
|
| 527 |
+
{
|
| 528 |
+
"shape": [
|
| 529 |
+
1,
|
| 530 |
+
1,
|
| 531 |
+
512
|
| 532 |
+
],
|
| 533 |
+
"dtype": "float"
|
| 534 |
+
}
|
| 535 |
+
]
|
| 536 |
+
}
|
| 537 |
+
},
|
| 538 |
+
"lmhead": {
|
| 539 |
+
"forward": {
|
| 540 |
+
"inputs": [
|
| 541 |
+
{
|
| 542 |
+
"shape": [
|
| 543 |
+
1,
|
| 544 |
+
1,
|
| 545 |
+
512
|
| 546 |
+
],
|
| 547 |
+
"dtype": "float"
|
| 548 |
+
}
|
| 549 |
+
],
|
| 550 |
+
"outputs": [
|
| 551 |
+
{
|
| 552 |
+
"shape": [
|
| 553 |
+
1,
|
| 554 |
+
1,
|
| 555 |
+
51872
|
| 556 |
+
],
|
| 557 |
+
"dtype": "float"
|
| 558 |
+
}
|
| 559 |
+
]
|
| 560 |
+
}
|
| 561 |
+
}
|
| 562 |
+
}
|
| 563 |
+
},
|
| 564 |
+
"files": {
|
| 565 |
+
"whisper_mel.pte": {
|
| 566 |
+
"bytes": 710976,
|
| 567 |
+
"sha256": "7cae31829487958454c5a14aecfed0845f6320dd480c24883861fa39b07de497"
|
| 568 |
+
},
|
| 569 |
+
"whisper_encoder.pte": {
|
| 570 |
+
"bytes": 82420996,
|
| 571 |
+
"sha256": "841e7702ba8fa73594f0b4e1400230d6b686af958628ead6b22919e9a3e27c69"
|
| 572 |
+
},
|
| 573 |
+
"whisper_cross_kv.pte": {
|
| 574 |
+
"bytes": 12610560,
|
| 575 |
+
"sha256": "6e8078924168c21a863ff662e9ccba55ed86bb3826dabd28fea77010e3151012"
|
| 576 |
+
},
|
| 577 |
+
"whisper_decoder_split.pte": {
|
| 578 |
+
"bytes": 88619264,
|
| 579 |
+
"sha256": "ade1aedc9e5795ac8c224347b88f1b826f5598e68d1df0961695fd843fa0e9b5"
|
| 580 |
+
},
|
| 581 |
+
"whisper_decode_stream.pte": {
|
| 582 |
+
"bytes": 89270912,
|
| 583 |
+
"sha256": "bc2f33ad7aa164059e819247365e2e0a75a3f1a2973024e4fdc67a254d7faa8d"
|
| 584 |
+
},
|
| 585 |
+
"whisper_lmhead.pte": {
|
| 586 |
+
"bytes": 27184000,
|
| 587 |
+
"sha256": "68577fa240d29923058fd72b94db583f3ff171302ad12834c2b610e04c559f8e"
|
| 588 |
+
},
|
| 589 |
+
"whisper_embed.i8.bin": {
|
| 590 |
+
"bytes": 26554880,
|
| 591 |
+
"sha256": "bc7f8d8336ad935eb8b6374036478d92d679eae156a1d4188d5cc016b5f2dda2"
|
| 592 |
+
},
|
| 593 |
+
"whisper_embed.scale.f32.bin": {
|
| 594 |
+
"bytes": 207460,
|
| 595 |
+
"sha256": "1f58851911611b9686851c586365df2fcaaa0727957c05bd163d7c3ccc231f4d"
|
| 596 |
+
},
|
| 597 |
+
"whisper_pos_embed.f32.bin": {
|
| 598 |
+
"bytes": 917504,
|
| 599 |
+
"sha256": "88fa1cdbf2b06f86b0ecb7be0fccfc39e906502986572b8cf5319c250e857169"
|
| 600 |
+
},
|
| 601 |
+
"vocab.json": {
|
| 602 |
+
"bytes": 835550,
|
| 603 |
+
"sha256": "8f680bba319e01a653d2e8a5dbc17a9157179e0576e6ce74ce0c06356c6e24f9"
|
| 604 |
+
}
|
| 605 |
+
},
|
| 606 |
+
"graph": {
|
| 607 |
+
"mel": {
|
| 608 |
+
"non_delegated_kernels": {},
|
| 609 |
+
"delegates": [
|
| 610 |
+
{
|
| 611 |
+
"id": "VulkanBackend",
|
| 612 |
+
"ops": {
|
| 613 |
+
"aten.squeeze_copy.dims": 1,
|
| 614 |
+
"aten.linear.default": 3,
|
| 615 |
+
"aten.mul.Tensor": 2,
|
| 616 |
+
"aten.add.Tensor": 1,
|
| 617 |
+
"aten.view_copy.default": 1
|
| 618 |
+
},
|
| 619 |
+
"inputs": 1,
|
| 620 |
+
"outputs": 1
|
| 621 |
+
}
|
| 622 |
+
]
|
| 623 |
+
},
|
| 624 |
+
"encoder": {
|
| 625 |
+
"non_delegated_kernels": {},
|
| 626 |
+
"delegates": [
|
| 627 |
+
{
|
| 628 |
+
"id": "VulkanBackend",
|
| 629 |
+
"ops": {
|
| 630 |
+
"et_vk.prepack.default": 7,
|
| 631 |
+
"aten.convolution.default": 2,
|
| 632 |
+
"aten.gelu.default": 8,
|
| 633 |
+
"aten.clone.default": 1,
|
| 634 |
+
"aten.permute_copy.default": 31,
|
| 635 |
+
"aten.squeeze_copy.dims": 25,
|
| 636 |
+
"aten.add.Tensor": 13,
|
| 637 |
+
"aten.native_layer_norm.default": 13,
|
| 638 |
+
"aten.linear.default": 36,
|
| 639 |
+
"aten.view_copy.default": 37,
|
| 640 |
+
"aten.mul.Tensor": 6,
|
| 641 |
+
"aten.bmm.default": 12,
|
| 642 |
+
"aten._softmax.default": 6
|
| 643 |
+
},
|
| 644 |
+
"inputs": 1,
|
| 645 |
+
"outputs": 1
|
| 646 |
+
}
|
| 647 |
+
]
|
| 648 |
+
},
|
| 649 |
+
"cross_kv": {
|
| 650 |
+
"non_delegated_kernels": {},
|
| 651 |
+
"delegates": [
|
| 652 |
+
{
|
| 653 |
+
"id": "VulkanBackend",
|
| 654 |
+
"ops": {
|
| 655 |
+
"aten.squeeze_copy.dims": 12,
|
| 656 |
+
"aten.linear.default": 12,
|
| 657 |
+
"aten.view_copy.default": 12,
|
| 658 |
+
"aten.permute_copy.default": 12
|
| 659 |
+
},
|
| 660 |
+
"inputs": 1,
|
| 661 |
+
"outputs": 12
|
| 662 |
+
}
|
| 663 |
+
]
|
| 664 |
+
},
|
| 665 |
+
"decoder_split": {
|
| 666 |
+
"non_delegated_kernels": {},
|
| 667 |
+
"delegates": [
|
| 668 |
+
{
|
| 669 |
+
"id": "VulkanBackend",
|
| 670 |
+
"ops": {
|
| 671 |
+
"et_vk.prepack.default": 2,
|
| 672 |
+
"aten.add.Tensor": 19,
|
| 673 |
+
"aten.native_layer_norm.default": 19,
|
| 674 |
+
"aten.squeeze_copy.dims": 30,
|
| 675 |
+
"aten.linear.default": 48,
|
| 676 |
+
"aten.view_copy.default": 60,
|
| 677 |
+
"aten.permute_copy.default": 36,
|
| 678 |
+
"et_vk.sdpa.default": 12,
|
| 679 |
+
"aten.gelu.default": 6
|
| 680 |
+
},
|
| 681 |
+
"inputs": 13,
|
| 682 |
+
"outputs": 1
|
| 683 |
+
}
|
| 684 |
+
]
|
| 685 |
+
},
|
| 686 |
+
"decode_stream": {
|
| 687 |
+
"non_delegated_kernels": {},
|
| 688 |
+
"delegates": [
|
| 689 |
+
{
|
| 690 |
+
"id": "VulkanBackend",
|
| 691 |
+
"ops": {
|
| 692 |
+
"et_vk.select_as_symint.default": 1,
|
| 693 |
+
"aten.native_layer_norm.default": 19,
|
| 694 |
+
"aten.squeeze_copy.dims": 30,
|
| 695 |
+
"aten.linear.default": 48,
|
| 696 |
+
"aten.view_copy.default": 42,
|
| 697 |
+
"sdpa_with_kv_cache.default": 6,
|
| 698 |
+
"aten.unsqueeze_copy.default": 18,
|
| 699 |
+
"aten.add.Tensor": 18,
|
| 700 |
+
"aten.permute_copy.default": 12,
|
| 701 |
+
"et_vk.sdpa.default": 6,
|
| 702 |
+
"aten.gelu.default": 6
|
| 703 |
+
},
|
| 704 |
+
"inputs": 14,
|
| 705 |
+
"outputs": 1
|
| 706 |
+
}
|
| 707 |
+
]
|
| 708 |
+
},
|
| 709 |
+
"lmhead": {
|
| 710 |
+
"non_delegated_kernels": {},
|
| 711 |
+
"delegates": [
|
| 712 |
+
{
|
| 713 |
+
"id": "VulkanBackend",
|
| 714 |
+
"ops": {
|
| 715 |
+
"quantized_decomposed.quantize_per_tensor.default": 1,
|
| 716 |
+
"aten.clone.default": 1,
|
| 717 |
+
"et_vk.linear_q8ta_q8csw.default": 1
|
| 718 |
+
},
|
| 719 |
+
"inputs": 1,
|
| 720 |
+
"outputs": 1
|
| 721 |
+
}
|
| 722 |
+
]
|
| 723 |
+
}
|
| 724 |
+
},
|
| 725 |
+
"webgpu_runtime_checks": {
|
| 726 |
+
"op_list": "checked against --runtime-ops",
|
| 727 |
+
"mel": {
|
| 728 |
+
"problems": [],
|
| 729 |
+
"ok": true,
|
| 730 |
+
"notes": {
|
| 731 |
+
"value_kinds": [
|
| 732 |
+
"FLOAT32",
|
| 733 |
+
"Int",
|
| 734 |
+
"IntList",
|
| 735 |
+
"Null"
|
| 736 |
+
],
|
| 737 |
+
"aten.linear input ranks": [
|
| 738 |
+
"2"
|
| 739 |
+
]
|
| 740 |
+
}
|
| 741 |
+
},
|
| 742 |
+
"encoder": {
|
| 743 |
+
"problems": [],
|
| 744 |
+
"ok": true,
|
| 745 |
+
"notes": {
|
| 746 |
+
"value_kinds": [
|
| 747 |
+
"Bool",
|
| 748 |
+
"Double",
|
| 749 |
+
"FLOAT32",
|
| 750 |
+
"Int",
|
| 751 |
+
"IntList",
|
| 752 |
+
"Null",
|
| 753 |
+
"String",
|
| 754 |
+
"ValueList"
|
| 755 |
+
],
|
| 756 |
+
"aten.convolution input/weight": [
|
| 757 |
+
"((1, 512, 3000), (512, 512, 3))",
|
| 758 |
+
"((1, 80, 3000), (512, 80, 3))"
|
| 759 |
+
],
|
| 760 |
+
"aten.linear input ranks": [
|
| 761 |
+
"2"
|
| 762 |
+
]
|
| 763 |
+
}
|
| 764 |
+
},
|
| 765 |
+
"cross_kv": {
|
| 766 |
+
"problems": [],
|
| 767 |
+
"ok": true,
|
| 768 |
+
"notes": {
|
| 769 |
+
"value_kinds": [
|
| 770 |
+
"FLOAT32",
|
| 771 |
+
"IntList",
|
| 772 |
+
"Null"
|
| 773 |
+
],
|
| 774 |
+
"aten.linear input ranks": [
|
| 775 |
+
"2"
|
| 776 |
+
]
|
| 777 |
+
}
|
| 778 |
+
},
|
| 779 |
+
"decoder_split": {
|
| 780 |
+
"problems": [],
|
| 781 |
+
"ok": true,
|
| 782 |
+
"notes": {
|
| 783 |
+
"value_kinds": [
|
| 784 |
+
"Double",
|
| 785 |
+
"FLOAT32",
|
| 786 |
+
"Int",
|
| 787 |
+
"IntList",
|
| 788 |
+
"Null",
|
| 789 |
+
"String",
|
| 790 |
+
"ValueList"
|
| 791 |
+
],
|
| 792 |
+
"aten.linear input ranks": [
|
| 793 |
+
"2"
|
| 794 |
+
],
|
| 795 |
+
"et_vk.sdpa q/k/mask": [
|
| 796 |
+
"((1, 8, 64, 64), (1, 8, 1500, 64), None)",
|
| 797 |
+
"((1, 8, 64, 64), (1, 8, 64, 64), (1, 8, 64, 64))"
|
| 798 |
+
]
|
| 799 |
+
}
|
| 800 |
+
},
|
| 801 |
+
"decode_stream": {
|
| 802 |
+
"problems": [],
|
| 803 |
+
"ok": true,
|
| 804 |
+
"notes": {
|
| 805 |
+
"value_kinds": [
|
| 806 |
+
"Bool",
|
| 807 |
+
"Double",
|
| 808 |
+
"FLOAT32",
|
| 809 |
+
"INT32",
|
| 810 |
+
"Int",
|
| 811 |
+
"IntList",
|
| 812 |
+
"Null",
|
| 813 |
+
"String",
|
| 814 |
+
"SymInt",
|
| 815 |
+
"ValueList"
|
| 816 |
+
],
|
| 817 |
+
"aten.linear input ranks": [
|
| 818 |
+
"2"
|
| 819 |
+
],
|
| 820 |
+
"sdpa_with_kv_cache q/k_cache": [
|
| 821 |
+
"((1, 1, 8, 64), (1, 448, 8, 64), 'SymInt')"
|
| 822 |
+
],
|
| 823 |
+
"et_vk.sdpa q/k/mask": [
|
| 824 |
+
"((1, 8, 1, 64), (1, 8, 1500, 64), None)"
|
| 825 |
+
]
|
| 826 |
+
}
|
| 827 |
+
},
|
| 828 |
+
"lmhead": {
|
| 829 |
+
"problems": [],
|
| 830 |
+
"ok": true,
|
| 831 |
+
"notes": {
|
| 832 |
+
"value_kinds": [
|
| 833 |
+
"Double",
|
| 834 |
+
"FLOAT32",
|
| 835 |
+
"INT32",
|
| 836 |
+
"INT8",
|
| 837 |
+
"Int",
|
| 838 |
+
"Null"
|
| 839 |
+
]
|
| 840 |
+
}
|
| 841 |
+
}
|
| 842 |
+
},
|
| 843 |
+
"command": "python recipes/whisper-base.py --out out/whisper-base/fp32dec-lmq8 --precision fp32dec-lmq8"
|
| 844 |
+
}
|
export/requirements.txt
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Environment the re-exports in REEXPORT.md were produced with (Python 3.12, Linux x86_64, CPU only).
|
| 2 |
+
# python -m venv venv
|
| 3 |
+
# venv/bin/pip install -r recipes/requirements.txt --extra-index-url https://download.pytorch.org/whl/cpu
|
| 4 |
+
# executorch 1.5.1 = pytorch/executorch git 3b60683923245cf472b7323426920e15623ba361; its wheel ships
|
| 5 |
+
# backends/webgpu (WebGPUPartitioner) and examples/models/llama (export_llm).
|
| 6 |
+
executorch==1.5.1
|
| 7 |
+
torch==2.14.1+cpu
|
| 8 |
+
torchvision==0.29.1+cpu
|
| 9 |
+
torchao==0.18.0+cpu
|
| 10 |
+
transformers==5.19.0
|
| 11 |
+
huggingface_hub==2.2.0
|
| 12 |
+
safetensors==0.8.0
|
| 13 |
+
tokenizers==0.23.3
|
| 14 |
+
numpy==2.5.4
|
| 15 |
+
pillow==12.3.0
|
| 16 |
+
# whisper-base: decode the LibriSpeech clips used for the output-layer calibration and validation
|
| 17 |
+
soundfile==0.14.0
|
| 18 |
+
pyarrow==26.0.0
|
export/whisper-base.py
ADDED
|
@@ -0,0 +1,784 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Whisper base (openai/whisper-base) -> six ExecuTorch WebGPU programs + host tables, public ExecuTorch.
|
| 3 |
+
|
| 4 |
+
Same split and program I/O contract as the playground's previous export (per-program contract in
|
| 5 |
+
recipes/_whisper.py):
|
| 6 |
+
|
| 7 |
+
whisper_mel.pte frames[1,3000,400] -> power mel[1,3000,80] fp32
|
| 8 |
+
whisper_encoder.pte log-mel[1,80,3000] -> encoder hidden[1,1500,512] fp32
|
| 9 |
+
whisper_cross_kv.pte encoder hidden -> 12 x [1,8,1500,64] (k0,v0,...,k5,v5)
|
| 10 |
+
whisper_decoder_split.pte (tok embeds[1,64,512], 12 cross K/V) -> hidden[1,64,512]
|
| 11 |
+
whisper_decode_stream.pte (embed+pos[1,1,512], pos[1] i64, 12 cross K/V) -> hidden[1,1,512]
|
| 12 |
+
whisper_lmhead.pte hidden[1,1,512] -> logits[1,1,51872] (rows >= 51865 are -1e9)
|
| 13 |
+
whisper_pos_embed.f32.bin, vocab.json, and the token-embedding table:
|
| 14 |
+
|
| 15 |
+
--precision q4 (default; the previous export's numerics) cross_kv / decoder_split / decode_stream
|
| 16 |
+
/ lmhead linears int4 groupwise (et_vk.linear_q4gsw, group 64; quantized weights
|
| 17 |
+
bit-identical to the previous files) + whisper_embed.i8.bin / .scale.f32.bin.
|
| 18 |
+
--precision fp32 every program fp32 + whisper_embed.f32.bin. Larger, but its greedy transcripts
|
| 19 |
+
are token-identical to transformers (int4 weights are not; see validation.json).
|
| 20 |
+
--precision hybrid-lmfp32 as q4, but lmhead fp32 (aten.linear).
|
| 21 |
+
--precision hybrid-lmq8 as q4, but lmhead int8 x int8 (static per-tensor input scale calibrated on
|
| 22 |
+
the 3 LibriSpeech-dummy clips longer than 15 s; et_vk.linear_q8ta_q8csw).
|
| 23 |
+
--precision fp32dec-lmq8 fp32 cross_kv / decoder_split / decode_stream, the int8 x int8 lmhead above
|
| 24 |
+
(calibrated on the fp32 decoder's hidden states) and the int8 embedding table.
|
| 25 |
+
|
| 26 |
+
Every program is lowered with executorch.backends.webgpu.partitioner.WebGPUPartitioner() into one
|
| 27 |
+
VulkanBackend delegate with no CPU kernels. The host does: framing, the log-mel normalization between
|
| 28 |
+
mel and encoder (log10 / max-8 / (x+4)/4 / transpose; the public runtime has no log op), token
|
| 29 |
+
embedding lookup (+ positional row for decode_stream), and argmax.
|
| 30 |
+
|
| 31 |
+
python recipes/whisper-base.py --out out/whisper-base/q4 # export + validation
|
| 32 |
+
python recipes/whisper-base.py --out out/whisper-base/fp32 --precision fp32
|
| 33 |
+
python recipes/whisper-base.py --out /tmp/w --no-validate # export only
|
| 34 |
+
|
| 35 |
+
Validation (CPU only; the WebGPU kernels cannot run here): eager vs torch.export program vs portable
|
| 36 |
+
.pte (ExecuTorch pybindings) for every program, and greedy transcripts of LibriSpeech clips
|
| 37 |
+
(hf-internal-testing/librispeech_asr_dummy) through the full host + program pipeline, compared token
|
| 38 |
+
by token with transformers WhisperForConditionalGeneration.
|
| 39 |
+
"""
|
| 40 |
+
|
| 41 |
+
import argparse
|
| 42 |
+
import copy
|
| 43 |
+
import json
|
| 44 |
+
import shutil
|
| 45 |
+
import sys
|
| 46 |
+
import time
|
| 47 |
+
from pathlib import Path
|
| 48 |
+
|
| 49 |
+
import numpy as np
|
| 50 |
+
import torch
|
| 51 |
+
|
| 52 |
+
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
| 53 |
+
import _common # noqa: E402
|
| 54 |
+
import _whisper as W # noqa: E402
|
| 55 |
+
|
| 56 |
+
UPSTREAM = {
|
| 57 |
+
"repo": "openai/whisper-base",
|
| 58 |
+
"revision": "e37978b90ca9030d5170a5c07aadb050351a65bb",
|
| 59 |
+
"files": {
|
| 60 |
+
"model.safetensors": "07cadb9f25677c8d50df603e66a98fbd842cce45047139baeb16e6219a1e807b",
|
| 61 |
+
"config.json": "a153c53883a6799b6f056b4a8d1a515c9926d03994682ba88a7616618d7da0c1",
|
| 62 |
+
"generation_config.json": "444b3f636d2fff89dd9ecf549e2a085b61f7ff0fa0246d4628bac6a3b8cc9ba4",
|
| 63 |
+
"preprocessor_config.json": "9b5cd03a36fbb8a627c64d98a5b5b126ead95a77720723944487311f0110b666",
|
| 64 |
+
"vocab.json": "8f680bba319e01a653d2e8a5dbc17a9157179e0576e6ce74ce0c06356c6e24f9",
|
| 65 |
+
"tokenizer.json": "27fc476bfe7f17299480be2273fc0608e4d5a99aba2ab5dec5374b4482d1a566",
|
| 66 |
+
"tokenizer_config.json": "2a4c4281cf9f51ac6ccc406fdc711a087afe6530f671fa7b80953edc498275ce",
|
| 67 |
+
"merges.txt": "2df2990a395e35e8dfbc7511e08c12d56018d8d04691e0133e5d63b21e154dc6",
|
| 68 |
+
"normalizer.json": "bf1c507dc8724ca9cf9903640dacfb69dae2f00edee4f21ceba106a7392f26dd",
|
| 69 |
+
"added_tokens.json": "9715fd2243b6f06a5858b5e32950d2853f73dd5bc201aafcf76f5082a2d8acd1",
|
| 70 |
+
"special_tokens_map.json": "e67ae3a0aaa99abcd9f187138e12db1f65c16a14761c50ef10eef2c174a7a691",
|
| 71 |
+
},
|
| 72 |
+
}
|
| 73 |
+
SAMPLE = {
|
| 74 |
+
"repo": "hf-internal-testing/librispeech_asr_dummy",
|
| 75 |
+
"repo_type": "dataset",
|
| 76 |
+
"revision": "5be91486e11a2d616f4ec5db8d3fd248585ac07a",
|
| 77 |
+
"file": "clean/validation-00000-of-00001.parquet",
|
| 78 |
+
"sha256": "4e69a06fa5edc90921e5e7e39a7084881f8b3ed9c805c574f4f39c6fde27c603",
|
| 79 |
+
}
|
| 80 |
+
NAMES = ["mel", "encoder", "cross_kv", "decoder_split", "decode_stream", "lmhead"]
|
| 81 |
+
DECODER_SIDE = ["cross_kv", "decoder_split", "decode_stream", "lmhead"]
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def fetch_upstream():
|
| 85 |
+
snap, digests = None, {}
|
| 86 |
+
for f, sha in UPSTREAM["files"].items():
|
| 87 |
+
p, d = _common.hf_file(UPSTREAM["repo"], f, UPSTREAM["revision"], sha)
|
| 88 |
+
digests[f] = d
|
| 89 |
+
snap = p.parent
|
| 90 |
+
return snap, digests
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
def fetch_sample():
|
| 94 |
+
from huggingface_hub import hf_hub_download
|
| 95 |
+
|
| 96 |
+
p = hf_hub_download(
|
| 97 |
+
SAMPLE["repo"],
|
| 98 |
+
SAMPLE["file"],
|
| 99 |
+
repo_type=SAMPLE["repo_type"],
|
| 100 |
+
revision=SAMPLE["revision"],
|
| 101 |
+
)
|
| 102 |
+
d = _common.sha256_file(p)
|
| 103 |
+
if d != SAMPLE["sha256"]:
|
| 104 |
+
sys.exit(f"{SAMPLE['repo']}:{SAMPLE['file']} sha256 {d} != {SAMPLE['sha256']}")
|
| 105 |
+
return p
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
def example_inputs(name):
|
| 109 |
+
g = torch.Generator().manual_seed(0)
|
| 110 |
+
ckv = tuple(torch.randn(1, 8, W.ENC_LEN, 64, generator=g) for _ in range(12))
|
| 111 |
+
return {
|
| 112 |
+
"mel": (torch.randn(1, W.FRAMES, W.N_FFT, generator=g),),
|
| 113 |
+
"encoder": (torch.randn(1, W.N_MELS, W.FRAMES, generator=g),),
|
| 114 |
+
"cross_kv": (torch.randn(1, W.ENC_LEN, W.D_MODEL, generator=g),),
|
| 115 |
+
"decoder_split": (torch.randn(1, W.MAX_DEC, W.D_MODEL, generator=g),) + ckv,
|
| 116 |
+
"decode_stream": (torch.randn(1, 1, W.D_MODEL, generator=g), torch.tensor([0]))
|
| 117 |
+
+ ckv,
|
| 118 |
+
"lmhead": (torch.randn(1, 1, W.D_MODEL, generator=g),),
|
| 119 |
+
}[name]
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
PRECISIONS = {
|
| 123 |
+
# name: (decoder-side linears, lmhead, host embedding table)
|
| 124 |
+
"q4": ("int4", "int4", "int8"),
|
| 125 |
+
"hybrid-lmfp32": ("int4", "fp32", "int8"),
|
| 126 |
+
"hybrid-lmq8": ("int4", "int8 static-act", "int8"),
|
| 127 |
+
"fp32dec-lmq8": ("fp32", "int8 static-act", "int8"),
|
| 128 |
+
"fp32": ("fp32", "fp32", "fp32"),
|
| 129 |
+
}
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
def calibration_hidden(model, fe, dec, tables, max_sec):
|
| 133 |
+
"""lmhead inputs for the q8 head's static activation scale: the variant's decoder's hidden states
|
| 134 |
+
for every decoded position of the LibriSpeech dummy clips LONGER than max_sec (3 clips, first
|
| 135 |
+
30 s window), which are disjoint from the evaluation clips. Tokens are chosen with the fp32
|
| 136 |
+
head (greedy, playground prefix)."""
|
| 137 |
+
clips = W.load_librispeech_dummy(fetch_sample(), max_sec, 100, longer=True)
|
| 138 |
+
stream = W.DecodeStreamProgram(dec, "torch")
|
| 139 |
+
ckvp = W.CrossKVProgram(dec)
|
| 140 |
+
mel, enc = W.MelProgram(fe), W.EncoderProgram(model, "eager")
|
| 141 |
+
lm32 = W.LMHeadProgram(model, W.VOCAB_PADDED)
|
| 142 |
+
hs = []
|
| 143 |
+
with torch.no_grad():
|
| 144 |
+
for c in clips:
|
| 145 |
+
lmel = W.host_log_mel(mel(torch.from_numpy(W.build_frames(c["pcm"]))[None]).numpy())
|
| 146 |
+
ckv = ckvp(enc(torch.from_numpy(lmel)))
|
| 147 |
+
stream.reset()
|
| 148 |
+
ids = list(W.PREFIX_EN)
|
| 149 |
+
for p in range(W.MAX_DEC):
|
| 150 |
+
h = stream(torch.from_numpy(tables.embed(ids[p], p)[None, None]), torch.tensor([p]), *ckv)
|
| 151 |
+
if p < len(W.PREFIX_EN) - 1:
|
| 152 |
+
continue
|
| 153 |
+
hs.append(h)
|
| 154 |
+
nxt = int(lm32(h).reshape(-1)[: W.VOCAB].argmax())
|
| 155 |
+
if nxt == W.TOK["eot"] or len(ids) + 1 >= W.MAX_DEC:
|
| 156 |
+
break
|
| 157 |
+
ids.append(nxt)
|
| 158 |
+
return torch.cat(hs, 0), [c["id"] for c in clips]
|
| 159 |
+
|
| 160 |
+
|
| 161 |
+
def build_programs(model, fe, precision, encoder_attn, tables, max_sec):
|
| 162 |
+
"""name -> (module lowered for WebGPU, module for the CPU portable .pte, fp32 reference or None).
|
| 163 |
+
|
| 164 |
+
The portable module computes the same function with ops the portable runtime has: attention as
|
| 165 |
+
matmul+softmax instead of et_vk.sdpa, the fp32 dequantization of the same int4 weights, and for
|
| 166 |
+
the q8 head an explicit per-tensor fake-quant + dequantized int8 weights."""
|
| 167 |
+
dec_kind, lm_kind, _ = PRECISIONS[precision]
|
| 168 |
+
fdec = model.model.decoder
|
| 169 |
+
lm32 = W.LMHeadProgram(model, W.VOCAB_PADDED)
|
| 170 |
+
info = {}
|
| 171 |
+
if dec_kind == "int4":
|
| 172 |
+
wdec = W.quantize_q4(copy.deepcopy(fdec))
|
| 173 |
+
pdec = W.dequantized_copy(wdec)
|
| 174 |
+
dref = lambda m: m # noqa: E731
|
| 175 |
+
else:
|
| 176 |
+
wdec = pdec = fdec
|
| 177 |
+
dref = lambda m: None # noqa: E731
|
| 178 |
+
if lm_kind == "int4":
|
| 179 |
+
wlm = W.quantize_q4(copy.deepcopy(lm32))
|
| 180 |
+
plm, lref = W.dequantized_copy(wlm), lm32
|
| 181 |
+
elif lm_kind == "fp32":
|
| 182 |
+
wlm = plm = lm32
|
| 183 |
+
lref = None
|
| 184 |
+
else:
|
| 185 |
+
calib, calib_ids = calibration_hidden(model, fe, wdec, tables, max_sec)
|
| 186 |
+
wlm, qp, wdq = W.q8_lmhead(lm32, calib)
|
| 187 |
+
plm, lref = W.FakeQ8LMHead(lm32, qp, wdq), lm32
|
| 188 |
+
info = {
|
| 189 |
+
"calibration_clips": calib_ids,
|
| 190 |
+
"calibration_vectors": int(calib.shape[0]),
|
| 191 |
+
"input_scale": qp[0],
|
| 192 |
+
"input_zero_point": qp[1],
|
| 193 |
+
}
|
| 194 |
+
progs = {
|
| 195 |
+
"mel": (W.MelProgram(fe), W.MelProgram(fe), None),
|
| 196 |
+
"encoder": (W.EncoderProgram(model, encoder_attn), W.EncoderProgram(model, "eager"), None),
|
| 197 |
+
"cross_kv": (W.CrossKVProgram(wdec), W.CrossKVProgram(pdec), dref(W.CrossKVProgram(fdec))),
|
| 198 |
+
"decoder_split": (
|
| 199 |
+
W.DecoderSplitProgram(wdec, "et_vk"),
|
| 200 |
+
W.DecoderSplitProgram(pdec, "torch"),
|
| 201 |
+
dref(W.DecoderSplitProgram(fdec, "torch")),
|
| 202 |
+
),
|
| 203 |
+
"decode_stream": (
|
| 204 |
+
W.DecodeStreamProgram(wdec, "et_vk"),
|
| 205 |
+
W.DecodeStreamProgram(pdec, "torch"),
|
| 206 |
+
dref(W.DecodeStreamProgram(fdec, "torch")),
|
| 207 |
+
),
|
| 208 |
+
"lmhead": (wlm, plm, lref),
|
| 209 |
+
}
|
| 210 |
+
return progs, info
|
| 211 |
+
|
| 212 |
+
|
| 213 |
+
def export_one(module, inputs):
|
| 214 |
+
if isinstance(module, W.DecodeStreamProgram):
|
| 215 |
+
module.reset() # the serialized K/V cache must start as zeros
|
| 216 |
+
with torch.no_grad():
|
| 217 |
+
if not isinstance(module, torch.fx.GraphModule): # pt2e-converted modules have no eval()
|
| 218 |
+
module = module.eval()
|
| 219 |
+
return torch.export.export(module, inputs, strict=False)
|
| 220 |
+
|
| 221 |
+
|
| 222 |
+
# ---------------------------------------------------------------------------------------------
|
| 223 |
+
# Runners: run(name, [arrays]) -> [np arrays]; reset() clears decode_stream's K/V cache.
|
| 224 |
+
# ---------------------------------------------------------------------------------------------
|
| 225 |
+
|
| 226 |
+
|
| 227 |
+
class ModuleRunner:
|
| 228 |
+
def __init__(self, modules):
|
| 229 |
+
self.m = modules
|
| 230 |
+
|
| 231 |
+
def reset(self):
|
| 232 |
+
for m in self.m.values():
|
| 233 |
+
if isinstance(m, W.DecodeStreamProgram):
|
| 234 |
+
m.reset()
|
| 235 |
+
elif hasattr(m, "named_buffers"):
|
| 236 |
+
for n, b in m.named_buffers():
|
| 237 |
+
if "cache" in n:
|
| 238 |
+
b.zero_()
|
| 239 |
+
|
| 240 |
+
def __call__(self, name, inputs):
|
| 241 |
+
with torch.no_grad():
|
| 242 |
+
out = self.m[name](
|
| 243 |
+
*[torch.from_numpy(np.ascontiguousarray(x)) for x in inputs]
|
| 244 |
+
)
|
| 245 |
+
out = out if isinstance(out, (tuple, list)) else (out,)
|
| 246 |
+
return [o.detach().numpy() for o in out]
|
| 247 |
+
|
| 248 |
+
|
| 249 |
+
class PteRunner:
|
| 250 |
+
def __init__(self, ptes):
|
| 251 |
+
from executorch.extension.pybindings.portable_lib import (
|
| 252 |
+
_load_for_executorch_from_buffer,
|
| 253 |
+
)
|
| 254 |
+
|
| 255 |
+
self._load = _load_for_executorch_from_buffer
|
| 256 |
+
self.ptes = ptes
|
| 257 |
+
self.m = {k: self._load(v) for k, v in ptes.items()}
|
| 258 |
+
|
| 259 |
+
def reset(self):
|
| 260 |
+
if "decode_stream" in self.ptes: # fresh method = zeroed mutable K/V buffers
|
| 261 |
+
self.m["decode_stream"] = self._load(self.ptes["decode_stream"])
|
| 262 |
+
|
| 263 |
+
def __call__(self, name, inputs):
|
| 264 |
+
out = self.m[name].forward(
|
| 265 |
+
[torch.from_numpy(np.ascontiguousarray(x)) for x in inputs]
|
| 266 |
+
)
|
| 267 |
+
return [
|
| 268 |
+
o.detach().numpy() if isinstance(o, torch.Tensor) else np.asarray(o)
|
| 269 |
+
for o in out
|
| 270 |
+
]
|
| 271 |
+
|
| 272 |
+
|
| 273 |
+
def encode_window(run, pcm):
|
| 274 |
+
power = run("mel", [W.build_frames(pcm)[None]])[0]
|
| 275 |
+
log_mel = W.host_log_mel(power)
|
| 276 |
+
enc = run("encoder", [log_mel])[0]
|
| 277 |
+
return log_mel, enc, run("cross_kv", [enc])
|
| 278 |
+
|
| 279 |
+
|
| 280 |
+
def greedy_stream(run, ckv, tables, max_dec=W.MAX_DEC):
|
| 281 |
+
"""whisper.js #window() with the streaming decoder: one token per call, argmax over 51865."""
|
| 282 |
+
run.reset()
|
| 283 |
+
ids = list(W.PREFIX_EN)
|
| 284 |
+
for p in range(max_dec):
|
| 285 |
+
e = tables.embed(ids[p], p)[None, None]
|
| 286 |
+
h = run("decode_stream", [e, np.array([p], dtype=np.int64), *ckv])[0]
|
| 287 |
+
if p < len(W.PREFIX_EN) - 1:
|
| 288 |
+
continue
|
| 289 |
+
nxt = int(np.argmax(run("lmhead", [h])[0].reshape(-1)[: W.VOCAB]))
|
| 290 |
+
if nxt == W.TOK["eot"]:
|
| 291 |
+
break
|
| 292 |
+
ids.append(nxt)
|
| 293 |
+
if len(ids) >= max_dec:
|
| 294 |
+
break
|
| 295 |
+
return ids[len(W.PREFIX_EN) :]
|
| 296 |
+
|
| 297 |
+
|
| 298 |
+
def greedy_split(run, ckv, tables, max_dec=W.MAX_DEC):
|
| 299 |
+
"""whisper.js #window() with the full-recompute decoder_split (no positions on the host)."""
|
| 300 |
+
ids = list(W.PREFIX_EN)
|
| 301 |
+
for p in range(len(W.PREFIX_EN) - 1, max_dec):
|
| 302 |
+
te = np.zeros((1, W.MAX_DEC, W.D_MODEL), dtype=np.float32)
|
| 303 |
+
te[0, : p + 1] = tables.embed(np.array(ids[: p + 1]))
|
| 304 |
+
h = run("decoder_split", [te, *ckv])[0]
|
| 305 |
+
nxt = int(np.argmax(run("lmhead", [h[:, p : p + 1]])[0].reshape(-1)[: W.VOCAB]))
|
| 306 |
+
if nxt == W.TOK["eot"]:
|
| 307 |
+
break
|
| 308 |
+
ids.append(nxt)
|
| 309 |
+
if len(ids) >= max_dec:
|
| 310 |
+
break
|
| 311 |
+
return ids[len(W.PREFIX_EN) :]
|
| 312 |
+
|
| 313 |
+
|
| 314 |
+
def hf_generate(model, feats, max_new):
|
| 315 |
+
with torch.no_grad():
|
| 316 |
+
seq = model.generate(
|
| 317 |
+
feats,
|
| 318 |
+
language="en",
|
| 319 |
+
task="transcribe",
|
| 320 |
+
return_timestamps=False,
|
| 321 |
+
do_sample=False,
|
| 322 |
+
num_beams=1,
|
| 323 |
+
max_new_tokens=max_new,
|
| 324 |
+
)[0].tolist()
|
| 325 |
+
# transformers 5 returns only the generated tokens; older versions also return the prompt.
|
| 326 |
+
out = seq[len(W.PREFIX_EN) :] if seq[: len(W.PREFIX_EN)] == W.PREFIX_EN else seq
|
| 327 |
+
return out[: out.index(W.TOK["eot"])] if W.TOK["eot"] in out else out
|
| 328 |
+
|
| 329 |
+
|
| 330 |
+
def hf_greedy(model, feats, max_dec=W.MAX_DEC):
|
| 331 |
+
"""Pure argmax with the forced prefix and no logit processors (the playground's rule)."""
|
| 332 |
+
with torch.no_grad():
|
| 333 |
+
enc = model.model.encoder(feats)
|
| 334 |
+
ids = list(W.PREFIX_EN)
|
| 335 |
+
while len(ids) < max_dec:
|
| 336 |
+
logits = model(
|
| 337 |
+
encoder_outputs=enc, decoder_input_ids=torch.tensor([ids])
|
| 338 |
+
).logits
|
| 339 |
+
nxt = int(logits[0, -1].argmax())
|
| 340 |
+
if nxt == W.TOK["eot"]:
|
| 341 |
+
break
|
| 342 |
+
ids.append(nxt)
|
| 343 |
+
return ids[len(W.PREFIX_EN) :]
|
| 344 |
+
|
| 345 |
+
|
| 346 |
+
def first_diff(a, b):
|
| 347 |
+
for i, (x, y) in enumerate(zip(a, b)):
|
| 348 |
+
if x != y:
|
| 349 |
+
return i
|
| 350 |
+
return None if len(a) == len(b) else min(len(a), len(b))
|
| 351 |
+
|
| 352 |
+
|
| 353 |
+
def summarize(rs):
|
| 354 |
+
if len(rs) == 1:
|
| 355 |
+
return rs[0]
|
| 356 |
+
return {
|
| 357 |
+
"outputs": len(rs),
|
| 358 |
+
"worst_max_abs": max(d["max_abs"] for d in rs),
|
| 359 |
+
"worst_max_rel_to_absmax": max(d["max_rel_to_absmax"] for d in rs),
|
| 360 |
+
"min_cosine": min(d["cosine"] for d in rs),
|
| 361 |
+
}
|
| 362 |
+
|
| 363 |
+
|
| 364 |
+
def cmp(a, b):
|
| 365 |
+
return summarize([W.isclose_report(x, y) for x, y in zip(a, b)])
|
| 366 |
+
|
| 367 |
+
|
| 368 |
+
# ---------------------------------------------------------------------------------------------
|
| 369 |
+
|
| 370 |
+
|
| 371 |
+
def validate(args, out, model, fe, tok, progs, eps, ptes, tables):
|
| 372 |
+
t0 = time.time()
|
| 373 |
+
clips = W.load_librispeech_dummy(fetch_sample(), args.max_sec, args.clips)
|
| 374 |
+
report = {
|
| 375 |
+
"precision": args.precision,
|
| 376 |
+
"cpu_runtime": "executorch pybindings (portable_lib) from the pip wheel, with "
|
| 377 |
+
"executorch.extension.llm.custom_ops loaded (llama::update_cache / custom_sdpa kernels)",
|
| 378 |
+
"webgpu_pte_on_cpu": "not runnable: the pip wheel's pybindings have no WebGPU/Vulkan backend, "
|
| 379 |
+
"so the lowered GPU kernels themselves are not exercised here",
|
| 380 |
+
"portable_pte_note": "portable .pte = the same function exported without a delegate, with "
|
| 381 |
+
"attention as matmul+softmax (et_vk.sdpa has no portable kernel) and, for q4, the fp32 "
|
| 382 |
+
"dequantization of the same int4 weights (torchao ops have no portable kernels). It checks graph "
|
| 383 |
+
"structure, the host contract and the quantized numerics on the ExecuTorch runtime",
|
| 384 |
+
"sample": {
|
| 385 |
+
**SAMPLE,
|
| 386 |
+
"clips_used": [c["id"] for c in clips],
|
| 387 |
+
"max_sec": args.max_sec,
|
| 388 |
+
},
|
| 389 |
+
"programs": {},
|
| 390 |
+
}
|
| 391 |
+
pcm0 = clips[0]["pcm"]
|
| 392 |
+
eager = ModuleRunner({k: v[0] for k, v in progs.items()})
|
| 393 |
+
exported = ModuleRunner({k: e.module() for k, e in eps.items()})
|
| 394 |
+
pte = PteRunner(ptes)
|
| 395 |
+
log_mel0, enc0, ckv0 = encode_window(eager, pcm0)
|
| 396 |
+
dec_in = np.zeros((1, W.MAX_DEC, W.D_MODEL), np.float32)
|
| 397 |
+
dec_in[0, :4] = tables.embed(np.array(W.PREFIX_EN))
|
| 398 |
+
speech = {
|
| 399 |
+
"mel": [W.build_frames(pcm0)[None]],
|
| 400 |
+
"encoder": [log_mel0],
|
| 401 |
+
"cross_kv": [enc0],
|
| 402 |
+
"decoder_split": [dec_in, *ckv0],
|
| 403 |
+
}
|
| 404 |
+
for name in NAMES:
|
| 405 |
+
mod, _, ref32 = progs[name]
|
| 406 |
+
res = {}
|
| 407 |
+
if name != "decode_stream":
|
| 408 |
+
cases = {"random": [x.numpy() for x in example_inputs(name)]}
|
| 409 |
+
if name in speech:
|
| 410 |
+
cases["speech"] = speech[name]
|
| 411 |
+
if name == "lmhead":
|
| 412 |
+
cases["speech"] = [
|
| 413 |
+
eager("decoder_split", speech["decoder_split"])[0][:, 3:4]
|
| 414 |
+
]
|
| 415 |
+
for label, inp in cases.items():
|
| 416 |
+
e = eager(name, inp)
|
| 417 |
+
r = {
|
| 418 |
+
"exported_vs_eager": cmp(exported(name, inp), e),
|
| 419 |
+
"portable_pte_vs_eager": cmp(pte(name, inp), e),
|
| 420 |
+
}
|
| 421 |
+
if ref32 is not None:
|
| 422 |
+
r["quantized_eager_vs_fp32_eager"] = cmp(
|
| 423 |
+
e, ModuleRunner({name: ref32})(name, inp)
|
| 424 |
+
)
|
| 425 |
+
res[label] = r
|
| 426 |
+
else:
|
| 427 |
+
# 12 steps from a fresh cache: the 4 prefix tokens, then a fixed continuation.
|
| 428 |
+
seq = W.PREFIX_EN + [440, 1029, 291, 4532, 13, 50, 2034, 293]
|
| 429 |
+
runners = {
|
| 430 |
+
"eager": ModuleRunner({name: mod}),
|
| 431 |
+
"exported": exported,
|
| 432 |
+
"portable_pte": PteRunner({name: ptes[name]}),
|
| 433 |
+
}
|
| 434 |
+
if ref32 is not None:
|
| 435 |
+
runners["fp32_eager"] = ModuleRunner({name: ref32})
|
| 436 |
+
outs = {}
|
| 437 |
+
for k, rr in runners.items():
|
| 438 |
+
rr.reset()
|
| 439 |
+
hs = [
|
| 440 |
+
rr(
|
| 441 |
+
name,
|
| 442 |
+
[
|
| 443 |
+
tables.embed(t_, p_)[None, None],
|
| 444 |
+
np.array([p_], dtype=np.int64),
|
| 445 |
+
*ckv0,
|
| 446 |
+
],
|
| 447 |
+
)[0]
|
| 448 |
+
for p_, t_ in enumerate(seq)
|
| 449 |
+
]
|
| 450 |
+
outs[k] = np.concatenate(hs, 1)
|
| 451 |
+
r = {
|
| 452 |
+
"exported_vs_eager": W.isclose_report(outs["exported"], outs["eager"]),
|
| 453 |
+
"portable_pte_vs_eager": W.isclose_report(
|
| 454 |
+
outs["portable_pte"], outs["eager"]
|
| 455 |
+
),
|
| 456 |
+
}
|
| 457 |
+
if ref32 is not None:
|
| 458 |
+
r["quantized_eager_vs_fp32_eager"] = W.isclose_report(
|
| 459 |
+
outs["eager"], outs["fp32_eager"]
|
| 460 |
+
)
|
| 461 |
+
te = np.zeros((1, W.MAX_DEC, W.D_MODEL), np.float32)
|
| 462 |
+
te[0, : len(seq)] = tables.embed(np.array(seq))
|
| 463 |
+
hsplit = eager("decoder_split", [te, *ckv0])[0]
|
| 464 |
+
r["stream_vs_decoder_split_eager"] = W.isclose_report(
|
| 465 |
+
outs["eager"], hsplit[:, : len(seq)]
|
| 466 |
+
)
|
| 467 |
+
res["speech_12_steps"] = r
|
| 468 |
+
report["programs"][name] = res
|
| 469 |
+
print(f"[validate] {name}: {json.dumps(res)[:300]}", flush=True)
|
| 470 |
+
|
| 471 |
+
e32 = model.model.decoder.embed_tokens.weight.detach().numpy()
|
| 472 |
+
report["host_tables"] = {
|
| 473 |
+
"kind": tables.kind,
|
| 474 |
+
"embed_vs_checkpoint": W.isclose_report(tables.embed(np.arange(W.VOCAB)), e32),
|
| 475 |
+
"pos_embed_equals_checkpoint": bool(
|
| 476 |
+
np.array_equal(
|
| 477 |
+
tables.pos, model.model.decoder.embed_positions.weight.detach().numpy()
|
| 478 |
+
)
|
| 479 |
+
),
|
| 480 |
+
}
|
| 481 |
+
|
| 482 |
+
# End-to-end greedy transcripts.
|
| 483 |
+
keys = ["exported_stream", "exported_split"]
|
| 484 |
+
e2e = []
|
| 485 |
+
for ci, c in enumerate(clips):
|
| 486 |
+
t1 = time.time()
|
| 487 |
+
pcm = c["pcm"]
|
| 488 |
+
feats = fe(pcm, sampling_rate=W.SAMPLE_RATE, return_tensors="pt").input_features
|
| 489 |
+
hf_gen = hf_generate(model, feats, W.MAX_DEC - len(W.PREFIX_EN))
|
| 490 |
+
hf_arg = hf_greedy(model, feats)
|
| 491 |
+
log_mel_x, _, ckv_x = encode_window(exported, pcm)
|
| 492 |
+
row = {
|
| 493 |
+
"id": c["id"],
|
| 494 |
+
"seconds": W.seconds(pcm),
|
| 495 |
+
"reference_text": c["text"],
|
| 496 |
+
"log_mel_vs_hf_features": W.isclose_report(log_mel_x, feats.numpy()),
|
| 497 |
+
"hf_generate": hf_gen,
|
| 498 |
+
"hf_greedy_argmax": hf_arg,
|
| 499 |
+
"exported_stream": greedy_stream(exported, ckv_x, tables),
|
| 500 |
+
"exported_split": greedy_split(exported, ckv_x, tables),
|
| 501 |
+
}
|
| 502 |
+
if ci < args.portable_clips:
|
| 503 |
+
_, _, ckv_p = encode_window(pte, pcm)
|
| 504 |
+
row["portable_pte_stream"] = greedy_stream(pte, ckv_p, tables)
|
| 505 |
+
row["portable_pte_split"] = greedy_split(pte, ckv_p, tables)
|
| 506 |
+
for k in [k for k in row if k.startswith(("exported_", "portable_pte_"))]:
|
| 507 |
+
row[f"{k}_matches_hf_generate"] = row[k] == hf_gen
|
| 508 |
+
row[f"{k}_matches_hf_greedy_argmax"] = row[k] == hf_arg
|
| 509 |
+
if row[k] != hf_gen:
|
| 510 |
+
row[f"{k}_first_diff_vs_hf_generate"] = first_diff(row[k], hf_gen)
|
| 511 |
+
row["text_exported_stream"] = tok.decode(
|
| 512 |
+
row["exported_stream"], skip_special_tokens=True
|
| 513 |
+
).strip()
|
| 514 |
+
row["text_hf_generate"] = tok.decode(hf_gen, skip_special_tokens=True).strip()
|
| 515 |
+
row["seconds_to_validate"] = round(time.time() - t1, 1)
|
| 516 |
+
e2e.append(row)
|
| 517 |
+
flags = {
|
| 518 |
+
k: row[f"{k}_matches_hf_generate"]
|
| 519 |
+
for k in row
|
| 520 |
+
if f"{k}_matches_hf_generate" in row
|
| 521 |
+
}
|
| 522 |
+
print(
|
| 523 |
+
f"[validate] {c['id']} ({row['seconds']} s) {flags} | {row['text_exported_stream']}",
|
| 524 |
+
flush=True,
|
| 525 |
+
)
|
| 526 |
+
|
| 527 |
+
port = [r for r in e2e if "portable_pte_stream" in r]
|
| 528 |
+
summary = {
|
| 529 |
+
"clips": len(e2e),
|
| 530 |
+
"audio_seconds": round(sum(r["seconds"] for r in e2e), 2),
|
| 531 |
+
"hf_tokens": sum(len(r["hf_generate"]) for r in e2e),
|
| 532 |
+
"hf_generate_equals_hf_greedy_argmax": sum(
|
| 533 |
+
r["hf_generate"] == r["hf_greedy_argmax"] for r in e2e
|
| 534 |
+
),
|
| 535 |
+
}
|
| 536 |
+
for k in keys:
|
| 537 |
+
summary[f"{k}_token_exact_vs_hf_generate"] = sum(
|
| 538 |
+
r[f"{k}_matches_hf_generate"] for r in e2e
|
| 539 |
+
)
|
| 540 |
+
summary[f"{k}_token_exact_vs_hf_greedy_argmax"] = sum(
|
| 541 |
+
r[f"{k}_matches_hf_greedy_argmax"] for r in e2e
|
| 542 |
+
)
|
| 543 |
+
summary["portable_pte_clips"] = len(port)
|
| 544 |
+
for k in ["portable_pte_stream", "portable_pte_split"]:
|
| 545 |
+
summary[f"{k}_token_exact_vs_hf_generate"] = sum(
|
| 546 |
+
r[f"{k}_matches_hf_generate"] for r in port
|
| 547 |
+
)
|
| 548 |
+
summary[f"{k}_equals_exported_stream"] = sum(
|
| 549 |
+
r[k] == r["exported_stream"] for r in port
|
| 550 |
+
)
|
| 551 |
+
summary["exported_stream_equals_exported_split"] = sum(
|
| 552 |
+
r["exported_stream"] == r["exported_split"] for r in e2e
|
| 553 |
+
)
|
| 554 |
+
wer = {"hf_greedy_argmax": [0, 0], "exported_stream": [0, 0]}
|
| 555 |
+
for r in e2e:
|
| 556 |
+
for k in wer:
|
| 557 |
+
e_, n_ = W.word_errors(tok, tok.decode(r[k], skip_special_tokens=True), r["reference_text"])
|
| 558 |
+
wer[k][0] += e_
|
| 559 |
+
wer[k][1] += n_
|
| 560 |
+
summary["wer_vs_librispeech_text"] = {k: round(v[0] / max(1, v[1]), 4) for k, v in wer.items()}
|
| 561 |
+
summary["word_errors"] = {k: v for k, v in wer.items()}
|
| 562 |
+
summary["wer_note"] = "WhisperTokenizer.normalize (English normalizer) on both sides, word-level edit distance"
|
| 563 |
+
summary["worst_log_mel_vs_hf_features_max_abs"] = max(
|
| 564 |
+
r["log_mel_vs_hf_features"]["max_abs"] for r in e2e
|
| 565 |
+
)
|
| 566 |
+
summary["decode_budget"] = (
|
| 567 |
+
f"{W.MAX_DEC} positions including the 4-token prefix (playground maxDec)"
|
| 568 |
+
)
|
| 569 |
+
summary["pipeline"] = (
|
| 570 |
+
"build_frames -> mel -> host_log_mel -> encoder -> cross_kv -> [decode_stream | decoder_split] -> "
|
| 571 |
+
f"lmhead -> argmax over the first {W.VOCAB} logits; token embeddings from the host table "
|
| 572 |
+
"(+ whisper_pos_embed.f32.bin rows for decode_stream)"
|
| 573 |
+
)
|
| 574 |
+
summary["keys"] = {
|
| 575 |
+
"hf_generate": "WhisperForConditionalGeneration.generate(language='en', task='transcribe', "
|
| 576 |
+
"return_timestamps=False, greedy, max_new_tokens=60) with the checkpoint's generation_config "
|
| 577 |
+
"(applies suppress_tokens and begin_suppress_tokens), input_features from WhisperFeatureExtractor",
|
| 578 |
+
"hf_greedy_argmax": "same model and features, forced prefix <|startoftranscript|><|en|>"
|
| 579 |
+
"<|transcribe|><|notimestamps|>, plain argmax, no logit processors (the playground's rule)",
|
| 580 |
+
"exported_*": "the torch.export programs that were lowered for WebGPU (ExportedProgram.module(), "
|
| 581 |
+
"PyTorch CPU kernels), driven by the host pipeline exactly as whisper.js does",
|
| 582 |
+
"portable_pte_*": "the portable .pte programs on the ExecuTorch CPU runtime (first "
|
| 583 |
+
f"{args.portable_clips} clips; slow single-threaded kernels)",
|
| 584 |
+
}
|
| 585 |
+
report["end_to_end_summary"] = summary
|
| 586 |
+
report["end_to_end"] = e2e
|
| 587 |
+
report["validation_seconds"] = round(time.time() - t0, 1)
|
| 588 |
+
(out / "validation.json").write_text(json.dumps(report, indent=2) + "\n")
|
| 589 |
+
print(json.dumps(summary, indent=2))
|
| 590 |
+
|
| 591 |
+
|
| 592 |
+
def main():
|
| 593 |
+
ap = argparse.ArgumentParser(
|
| 594 |
+
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
| 595 |
+
)
|
| 596 |
+
ap.add_argument("--out", required=True, help="output directory")
|
| 597 |
+
ap.add_argument("--precision", choices=list(PRECISIONS), default="q4")
|
| 598 |
+
ap.add_argument(
|
| 599 |
+
"--no-validate", action="store_true", help="export only (no CPU validation)"
|
| 600 |
+
)
|
| 601 |
+
ap.add_argument(
|
| 602 |
+
"--clips",
|
| 603 |
+
type=int,
|
| 604 |
+
default=70,
|
| 605 |
+
help="LibriSpeech clips for the transcript check",
|
| 606 |
+
)
|
| 607 |
+
ap.add_argument(
|
| 608 |
+
"--portable-clips",
|
| 609 |
+
type=int,
|
| 610 |
+
default=2,
|
| 611 |
+
help="of those, how many also run as portable .pte",
|
| 612 |
+
)
|
| 613 |
+
ap.add_argument(
|
| 614 |
+
"--max-sec",
|
| 615 |
+
type=float,
|
| 616 |
+
default=15.0,
|
| 617 |
+
help="longest clip (the playground's window is <= 15 s)",
|
| 618 |
+
)
|
| 619 |
+
ap.add_argument(
|
| 620 |
+
"--encoder-attn",
|
| 621 |
+
choices=["eager", "et_vk"],
|
| 622 |
+
default="eager",
|
| 623 |
+
help="encoder attention: matmul+softmax (default; planned activation memory) or the fused "
|
| 624 |
+
"et_vk.sdpa op as in the previous export (public runtime: 2 x 72 MB scratch per layer)",
|
| 625 |
+
)
|
| 626 |
+
ap.add_argument(
|
| 627 |
+
"--keep-portable",
|
| 628 |
+
action="store_true",
|
| 629 |
+
help="also write the CPU-only portable .pte files",
|
| 630 |
+
)
|
| 631 |
+
ap.add_argument(
|
| 632 |
+
"--runtime-ops",
|
| 633 |
+
help="optional text file, one registered WebGPU runtime op per line; every delegated op is checked against it",
|
| 634 |
+
)
|
| 635 |
+
args = ap.parse_args()
|
| 636 |
+
out = Path(args.out)
|
| 637 |
+
out.mkdir(parents=True, exist_ok=True)
|
| 638 |
+
torch.manual_seed(0)
|
| 639 |
+
|
| 640 |
+
from transformers import (
|
| 641 |
+
WhisperFeatureExtractor,
|
| 642 |
+
WhisperForConditionalGeneration,
|
| 643 |
+
WhisperTokenizer,
|
| 644 |
+
)
|
| 645 |
+
|
| 646 |
+
snap, digests = fetch_upstream()
|
| 647 |
+
model = WhisperForConditionalGeneration.from_pretrained(
|
| 648 |
+
snap, dtype=torch.float32, attn_implementation="eager"
|
| 649 |
+
).eval()
|
| 650 |
+
fe = WhisperFeatureExtractor.from_pretrained(snap)
|
| 651 |
+
tables = W.HostTables(model, PRECISIONS[args.precision][2])
|
| 652 |
+
progs, quant_info = build_programs(model, fe, args.precision, args.encoder_attn, tables, args.max_sec)
|
| 653 |
+
runtime_ops = None
|
| 654 |
+
if args.runtime_ops:
|
| 655 |
+
runtime_ops = {
|
| 656 |
+
l.strip()
|
| 657 |
+
for l in Path(args.runtime_ops).read_text().splitlines()
|
| 658 |
+
if l.strip()
|
| 659 |
+
}
|
| 660 |
+
|
| 661 |
+
eps, ptes, files, graphs, checks, io_contract = {}, {}, {}, {}, {}, {}
|
| 662 |
+
for name in NAMES:
|
| 663 |
+
t = time.time()
|
| 664 |
+
wg_mod, port_mod, _ = progs[name]
|
| 665 |
+
inputs = example_inputs(name)
|
| 666 |
+
ep = export_one(wg_mod, inputs)
|
| 667 |
+
prog = _common.lower_webgpu(ep)
|
| 668 |
+
path = out / f"whisper_{name}.pte"
|
| 669 |
+
path.unlink(missing_ok=True) # never write through a hard link shared with another variant
|
| 670 |
+
path.write_bytes(prog.buffer)
|
| 671 |
+
eps[name] = ep
|
| 672 |
+
graphs[name] = _common.delegate_summary(prog.buffer)
|
| 673 |
+
checks[name] = W.vk_checks(prog.buffer, runtime_ops)
|
| 674 |
+
io_contract[name] = W.program_io(prog.buffer)
|
| 675 |
+
files[path.name] = {
|
| 676 |
+
"bytes": path.stat().st_size,
|
| 677 |
+
"sha256": _common.sha256_file(path),
|
| 678 |
+
}
|
| 679 |
+
if not args.no_validate or args.keep_portable:
|
| 680 |
+
ptes[name] = _common.lower_portable(export_one(port_mod, inputs)).buffer
|
| 681 |
+
if args.keep_portable:
|
| 682 |
+
(out / "portable").mkdir(exist_ok=True)
|
| 683 |
+
(out / "portable" / f"whisper_{name}_portable.pte").write_bytes(
|
| 684 |
+
ptes[name]
|
| 685 |
+
)
|
| 686 |
+
del prog
|
| 687 |
+
print(
|
| 688 |
+
f"[export] {path.name}: {files[path.name]['bytes']} B, {time.time() - t:.1f} s, "
|
| 689 |
+
f"checks ok={checks[name]['ok']}, {json.dumps(graphs[name])[:240]}",
|
| 690 |
+
flush=True,
|
| 691 |
+
)
|
| 692 |
+
|
| 693 |
+
for fname, blob in tables.files().items():
|
| 694 |
+
(out / fname).unlink(missing_ok=True)
|
| 695 |
+
(out / fname).write_bytes(blob)
|
| 696 |
+
(out / "vocab.json").unlink(missing_ok=True)
|
| 697 |
+
shutil.copyfile(snap / "vocab.json", out / "vocab.json")
|
| 698 |
+
for fname in [*tables.files(), "vocab.json"]:
|
| 699 |
+
files[fname] = {
|
| 700 |
+
"bytes": (out / fname).stat().st_size,
|
| 701 |
+
"sha256": _common.sha256_file(out / fname),
|
| 702 |
+
}
|
| 703 |
+
|
| 704 |
+
dec_kind, lm_kind, emb_kind = PRECISIONS[args.precision]
|
| 705 |
+
lm_desc = {
|
| 706 |
+
"int4": {"config": W.Q4_DESC, "lowered_to": "et_vk.linear_q4gsw (group 64)"},
|
| 707 |
+
"fp32": {"config": "none", "lowered_to": "aten.linear (2-D fp32)"},
|
| 708 |
+
"int8 static-act": {"config": W.Q8_HEAD_DESC, "lowered_to": "quantized_decomposed.quantize_per_tensor + et_vk.linear_q8ta_q8csw", **quant_info},
|
| 709 |
+
}[lm_kind]
|
| 710 |
+
recipe = {
|
| 711 |
+
"model": "whisper-base",
|
| 712 |
+
"precision": args.precision,
|
| 713 |
+
"upstream": {
|
| 714 |
+
"model": UPSTREAM["repo"],
|
| 715 |
+
"revision": UPSTREAM["revision"],
|
| 716 |
+
"official": True,
|
| 717 |
+
"license": "apache-2.0 (model card)",
|
| 718 |
+
"files_sha256": digests,
|
| 719 |
+
"loader": "transformers.WhisperForConditionalGeneration.from_pretrained(dtype=float32, attn_implementation='eager')",
|
| 720 |
+
},
|
| 721 |
+
"versions": _common.versions(),
|
| 722 |
+
"export": {
|
| 723 |
+
"api": "torch.export.export(strict=False) -> to_edge_transform_and_lower(partitioner=[WebGPUPartitioner()]) -> to_executorch()",
|
| 724 |
+
"partitioner": "executorch.backends.webgpu.partitioner.WebGPUPartitioner() (default compile options; wraps VulkanPartitioner; delegate id VulkanBackend)",
|
| 725 |
+
"model_code": "recipes/_whisper.py wrappers around the transformers modules (weights shared, forward re-expressed): "
|
| 726 |
+
"MelProgram, EncoderProgram, CrossKVProgram, DecoderSplitProgram, DecodeStreamProgram, LMHeadProgram",
|
| 727 |
+
"quantization": {
|
| 728 |
+
"decoder_side_linears": {
|
| 729 |
+
"programs": ["cross_kv", "decoder_split", "decode_stream"],
|
| 730 |
+
"config": W.Q4_DESC if dec_kind == "int4" else "none (fp32)",
|
| 731 |
+
"lowered_to": "et_vk.linear_q4gsw (group 64)" if dec_kind == "int4" else "aten.linear (2-D fp32)",
|
| 732 |
+
},
|
| 733 |
+
"lmhead": lm_desc,
|
| 734 |
+
"fp32_programs": ["mel", "encoder"],
|
| 735 |
+
"embedding_table": "int8 rowwise (whisper_embed.i8.bin + whisper_embed.scale.f32.bin)"
|
| 736 |
+
if emb_kind == "int8"
|
| 737 |
+
else "fp32 (whisper_embed.f32.bin)",
|
| 738 |
+
},
|
| 739 |
+
"encoder_attention": args.encoder_attn,
|
| 740 |
+
"decoder_attention": "et_vk.sdpa (decoder_split self-attention with a [1,8,64,64] additive causal mask "
|
| 741 |
+
"of 0 / -1e30; cross-attention in decoder_split and decode_stream); decode_stream self-attention = "
|
| 742 |
+
"llama.update_cache + llama.custom_sdpa, fused by the Vulkan partitioner into sdpa_with_kv_cache + "
|
| 743 |
+
"et_vk.select_as_symint",
|
| 744 |
+
"lmhead_rows": f"{W.VOCAB_PADDED}: vocab {W.VOCAB} padded with zero weights and bias {W.PAD_LOGIT} "
|
| 745 |
+
"(executorch 1.5.1's q4gsw lowering fails for N % 4 != 0; kept for fp32 too so the contract is identical)",
|
| 746 |
+
"constants": {
|
| 747 |
+
"max_decoder_rows": W.MAX_DEC,
|
| 748 |
+
"self_attention_cache": [1, W.KV_LEN, W.N_HEADS, W.HEAD_DIM],
|
| 749 |
+
"frames": [W.FRAMES, W.N_FFT],
|
| 750 |
+
"hop": W.HOP,
|
| 751 |
+
"prefix_en": W.PREFIX_EN,
|
| 752 |
+
},
|
| 753 |
+
"host_contract": {
|
| 754 |
+
"frames": "whisper.js buildFrames: zero-pad pcm to 480000 samples, reflect-pad 200 each side, 3000 x 400 windows at hop 160",
|
| 755 |
+
"log_mel": "between mel and encoder: x = log10(max(p, 1e-10)); x = max(x, max(x) - 8); x = (x + 4) / 4; "
|
| 756 |
+
"transpose [1,3000,80] -> [1,80,3000]",
|
| 757 |
+
"embeddings": ("i8[id] * scale[id]" if emb_kind == "int8" else "f32[id]")
|
| 758 |
+
+ "; decoder_split: rows 0..p = token embeddings (no positions), zero rows after; "
|
| 759 |
+
"decode_stream: token embedding + pos_embed[p]",
|
| 760 |
+
"argmax": f"over logits[..., :{W.VOCAB}] (rows {W.VOCAB}..{W.VOCAB_PADDED - 1} are {W.PAD_LOGIT})",
|
| 761 |
+
"cross_kv": "computed once per 30 s window, then passed as 12 inputs to every decoder call",
|
| 762 |
+
},
|
| 763 |
+
"io": io_contract,
|
| 764 |
+
},
|
| 765 |
+
"files": files,
|
| 766 |
+
"graph": graphs,
|
| 767 |
+
"webgpu_runtime_checks": {
|
| 768 |
+
"op_list": "checked against --runtime-ops"
|
| 769 |
+
if runtime_ops
|
| 770 |
+
else "not checked (no --runtime-ops given)",
|
| 771 |
+
**checks,
|
| 772 |
+
},
|
| 773 |
+
"command": _common.command_line(out),
|
| 774 |
+
}
|
| 775 |
+
(out / "recipe.json").write_text(json.dumps(recipe, indent=2) + "\n")
|
| 776 |
+
print(json.dumps({"files": files}, indent=2), flush=True)
|
| 777 |
+
|
| 778 |
+
if not args.no_validate:
|
| 779 |
+
tok = WhisperTokenizer.from_pretrained(snap)
|
| 780 |
+
validate(args, out, model, fe, tok, progs, eps, ptes, tables)
|
| 781 |
+
|
| 782 |
+
|
| 783 |
+
if __name__ == "__main__":
|
| 784 |
+
main()
|
manifest.json
CHANGED
|
@@ -9,18 +9,13 @@
|
|
| 9 |
},
|
| 10 |
{
|
| 11 |
"path": "whisper_cross_kv.pte",
|
| 12 |
-
"bytes":
|
| 13 |
-
"sha256": "
|
| 14 |
},
|
| 15 |
{
|
| 16 |
"path": "whisper_decode_stream.pte",
|
| 17 |
-
"bytes":
|
| 18 |
-
"sha256": "
|
| 19 |
-
},
|
| 20 |
-
{
|
| 21 |
-
"path": "whisper_decoder_split.pte",
|
| 22 |
-
"bytes": 13719936,
|
| 23 |
-
"sha256": "84d9427851420bf73d521649e1bb6af52458b36906d9a9a064d6057239e0dd1e"
|
| 24 |
},
|
| 25 |
{
|
| 26 |
"path": "whisper_embed.i8.bin",
|
|
@@ -34,18 +29,18 @@
|
|
| 34 |
},
|
| 35 |
{
|
| 36 |
"path": "whisper_encoder.pte",
|
| 37 |
-
"bytes":
|
| 38 |
-
"sha256": "
|
| 39 |
},
|
| 40 |
{
|
| 41 |
"path": "whisper_lmhead.pte",
|
| 42 |
-
"bytes":
|
| 43 |
-
"sha256": "
|
| 44 |
},
|
| 45 |
{
|
| 46 |
"path": "whisper_mel.pte",
|
| 47 |
-
"bytes":
|
| 48 |
-
"sha256": "
|
| 49 |
},
|
| 50 |
{
|
| 51 |
"path": "whisper_pos_embed.f32.bin",
|
|
|
|
| 9 |
},
|
| 10 |
{
|
| 11 |
"path": "whisper_cross_kv.pte",
|
| 12 |
+
"bytes": 12610560,
|
| 13 |
+
"sha256": "6e8078924168c21a863ff662e9ccba55ed86bb3826dabd28fea77010e3151012"
|
| 14 |
},
|
| 15 |
{
|
| 16 |
"path": "whisper_decode_stream.pte",
|
| 17 |
+
"bytes": 89270912,
|
| 18 |
+
"sha256": "bc2f33ad7aa164059e819247365e2e0a75a3f1a2973024e4fdc67a254d7faa8d"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 19 |
},
|
| 20 |
{
|
| 21 |
"path": "whisper_embed.i8.bin",
|
|
|
|
| 29 |
},
|
| 30 |
{
|
| 31 |
"path": "whisper_encoder.pte",
|
| 32 |
+
"bytes": 82420996,
|
| 33 |
+
"sha256": "841e7702ba8fa73594f0b4e1400230d6b686af958628ead6b22919e9a3e27c69"
|
| 34 |
},
|
| 35 |
{
|
| 36 |
"path": "whisper_lmhead.pte",
|
| 37 |
+
"bytes": 27184000,
|
| 38 |
+
"sha256": "68577fa240d29923058fd72b94db583f3ff171302ad12834c2b610e04c559f8e"
|
| 39 |
},
|
| 40 |
{
|
| 41 |
"path": "whisper_mel.pte",
|
| 42 |
+
"bytes": 710976,
|
| 43 |
+
"sha256": "7cae31829487958454c5a14aecfed0845f6320dd480c24883861fa39b07de497"
|
| 44 |
},
|
| 45 |
{
|
| 46 |
"path": "whisper_pos_embed.f32.bin",
|
whisper_cross_kv.pte
CHANGED
|
@@ -1,3 +1,3 @@
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:
|
| 3 |
-
size
|
|
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:6e8078924168c21a863ff662e9ccba55ed86bb3826dabd28fea77010e3151012
|
| 3 |
+
size 12610560
|
whisper_decode_stream.pte
CHANGED
|
@@ -1,3 +1,3 @@
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:
|
| 3 |
-
size
|
|
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:bc2f33ad7aa164059e819247365e2e0a75a3f1a2973024e4fdc67a254d7faa8d
|
| 3 |
+
size 89270912
|
whisper_decoder_split.pte
DELETED
|
@@ -1,3 +0,0 @@
|
|
| 1 |
-
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:84d9427851420bf73d521649e1bb6af52458b36906d9a9a064d6057239e0dd1e
|
| 3 |
-
size 13719936
|
|
|
|
|
|
|
|
|
|
|
|
whisper_encoder.pte
CHANGED
|
@@ -1,3 +1,3 @@
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:
|
| 3 |
-
size
|
|
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:841e7702ba8fa73594f0b4e1400230d6b686af958628ead6b22919e9a3e27c69
|
| 3 |
+
size 82420996
|
whisper_lmhead.pte
CHANGED
|
@@ -1,3 +1,3 @@
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:
|
| 3 |
-
size
|
|
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:68577fa240d29923058fd72b94db583f3ff171302ad12834c2b610e04c559f8e
|
| 3 |
+
size 27184000
|
whisper_mel.pte
CHANGED
|
@@ -1,3 +1,3 @@
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:
|
| 3 |
-
size
|
|
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:7cae31829487958454c5a14aecfed0845f6320dd480c24883861fa39b07de497
|
| 3 |
+
size 710976
|