[Bug][Metal] Metal codegen produces silently-wrong output on a HiFi-GAN-style Relax graph (cosine 0.756 vs LLVM 1.0000)
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
## Environment
- TVM: `apache-tvm==0.25.0.post1` (pip wheel)
- OS: macOS 15.5, arm64 (Apple Silicon, 12-core Apple Silicon Mac)
- Python: 3.13
- Backends compiled into the wheel: LLVM (host), Metal. Vulkan and
CUDA codegen absent (verified — `tvm.compile(mod, Target("vulkan"))`
raises `Cannot find global function target.build.vulkan`).
## Expected behaviour
Compiling the same Relax IR to `Target("llvm", ...)` and
`Target("metal", host="llvm")` and executing under the Relax VM with
identical inputs should produce numerically equivalent outputs within
`float32` rounding tolerance.
## Actual behaviour
LLVM produces the expected output. Metal produces output that is
uncorrelated with LLVM:
- `cosine(LLVM, Metal) = 0.755631`
- `max_abs_delta = 4.0493e+00`
- `RMS(Metal - LLVM) = 76% of RMS(LLVM)`
The bug is **deterministic** across fresh Python subprocesses — a
single `python3 repro.py` invocation reproduces byte-exact-identical
wrong Metal output every time (verified across 8 consecutive fresh
subprocesses, all producing identical SHA-256 of the output tensor).
## Reproducer
Repo: **https://github.com/zacharywhitley/tvm-metal-codegen-bug-repro**
Contents (all with published SHA-256 in the repo README):
- `minimum_repro_synthetic.onnx` (~3.4 MB, 174 nodes, opset 17) —
a 174-node HiFi-GAN-style Relax subgraph. All float initializers
are i.i.d. `N(0, 0.1)` samples with a fixed seed. Graph topology,
node names, initializer names/dtypes/shapes, and value_info are
exactly those of a HiFi-GAN-style ONNX decoder subgraph extracted
with `onnx.utils.extract_model`; only float weight values are
synthetic.
- `expected_output_synthetic.f32` (589 824 bytes) — LLVM oracle output
for the deterministic input, raw fp32 little-endian, shape
`[1, 144, 1024]`.
- `expected_output_synthetic.npy` — same, NumPy format.
- `repro.py` — self-contained driver.
To reproduce:
```
git clone https://github.com/zacharywhitley/tvm-metal-codegen-bug-repro
cd tvm-metal-codegen-bug-repro
pip install "apache-tvm==0.25.0.post1" "onnx>=1.16"
python3 repro.py
```
Op inventory (post-simplification, from `onnx.load`):
- 49× Slice
- 36× Mul
- 34× Conv (1-D)
- 19× Add
- 13× Tanh
- 13× Sigmoid
- 4× Split
- 3× Sub
- 3× Concat
- (174 total)
No `ConvTranspose` in this minimum subgraph — an important isolation
point since we previously suspected the Metal ConvTranspose lowering
(see "Isolation performed" below).
## Reproduction
```python
import numpy as np, onnx, tvm
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx
model = onnx.load("minimum_repro_synthetic.onnx")
mod = from_onnx(model, keep_params_in_input=False)
# Deterministic inputs
ch = np.arange(128, dtype=np.float32).reshape(-1, 1)
f = np.arange(1024, dtype=np.float32).reshape(1, -1)
z_p = (np.sin(2 * np.pi * (ch * 3 + f) / 128) * 0.5).reshape(1, 128, 1024).astype("float32")
y_mask = np.ones((1, 1, 1024), dtype=np.float32)
def run(target, dev):
ex = tvm.compile(mod, target)
vm = relax.VirtualMachine(ex, dev)
out = vm["main"](tvm.runtime.tensor(z_p, dev),
tvm.runtime.tensor(y_mask, dev))
dev.sync()
return (out.numpy() if hasattr(out, "numpy") else out[0].numpy()).ravel()
wf_llvm = run(tvm.target.Target({"kind":"llvm","mtriple":"arm64-apple-darwin"}),
tvm.runtime.cpu())
wf_metal = run(tvm.target.Target({"kind":"metal"},
host={"kind":"llvm","mtriple":"arm64-apple-darwin"}),
tvm.metal(0))
def cos(a, b):
return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))
print("LLVM vs METAL cosine:", cos(wf_llvm, wf_metal))
# Expected: ~1.0. Observed: 0.756.
```
Recommend running LLVM and Metal in separate Python subprocesses —
see "Additional observations" for why. The attached `repro.py`
handles this automatically (default `--fresh-subprocess` mode) and
prints a correctness table.
## Isolation performed
- **Not a single-op ConvTranspose1D bug.** A hand-written Relax
module with a single `nn.conv1d_transpose` at the same shape the
parent decoder uses (`[1,128,1024] × [128,64,16] stride=8 pad=8`)
produces byte-exact LLVM ↔ Metal agreement (`cosine=1.0000000`,
`max_abs_delta=0.0e+00`).
- **Not the first three flow blocks.** The exact same op pattern
(WaveNet-like `Conv → gated (tanh × sigmoid) → Conv`) is stacked
three times in the prefix of this graph without triggering the
divergence. Only when the fourth flow block's first `in_layers.1`
Conv is added does Metal diverge from LLVM. Bisection table:
| Anchor node index | Nodes | Fresh-proc cosine | In-process cosine |
|---|---:|---:|---:|
| flow.6/post/Conv (1 flow block) | 49 | 1.000000 | 1.000000 |
| flow.4/post/Conv (2 blocks) | 102 | – | 1.000000 |
| flow.2/post/Conv (3 blocks) | 155 | – | 1.000000 |
| flow.0/pre/Conv (4 blocks, pre only) | 162 | – | 1.000000 |
| flow.0/enc/in_layers.0/Conv | 164 | – | 1.000000 |
| flow.0/enc/res_skip_layers.0/Conv | 170 | **1.000000** | **0.905** *(in-proc, run ≥ 2)* |
| flow.0/enc/Slice_1_output_0 | 171 | 1.000000 | 0.914 |
| flow.0/enc/Add_output_0 | 172 | 1.000000 | 1.000000 |
| flow.0/enc/Mul_output_0 (mask × add) | 173 | 1.000000 | 0.832 |
| **flow.0/enc/in_layers.1/Conv** *(this repro)* | **174** | **~0.80** | **~0.80** |
| flow.0/enc/res_skip_layers.1/Conv | 180 | 0.794 | 0.706 |
| flow.0/post/Conv (4 blocks) | 208 | 0.645 | 0.645 |
| /dec/conv_pre/Conv | 215 | 0.648 | 0.648 |
| /dec/ups.0/ConvTranspose | 217 | 0.046 | 0.046 |
| waveform (full 410-node decoder) | 410 | 0.05 – 0.11 | 0.01 – 0.11 |
Cosine numbers in the bisection table are from the original trained
weights. With synthetic `N(0, 0.1)` weights the divergence pattern
is qualitatively the same (correct through node 170; wrong at node
174) with slightly different magnitudes.
## Additional observations (in-process state carryover)
A related pattern surfaced during bisection that may narrow the
search space:
- On a strictly smaller subgraph (170 nodes; extracted at
`/flow/flows.0/enc/res_skip_layers.0/Conv_output_0`), a fresh
Python subprocess produces `cosine=1.0000000` on the first Metal
compile — correct.
- Running Metal a **second time** on the same graph within the same
process produces `cosine=0.905` — deterministically wrong, and
byte-exact-identical across further in-process runs.
- The 174-node graph in this reproducer produces `cosine=0.756`
(synthetic weights) or `cosine=0.802` (trained weights) on the
**first** Metal compile in a fresh process.
Reading: at least one component of the failure appears to be state
that persists across `tvm.compile(mod, Target("metal"))` invocations
within a single process (compilation cache, device kernel cache, or a
static-scoped IR mutation). A large-enough graph triggers the bug
cold; a smaller graph only triggers it once that state has been
perturbed by a prior Metal compile.
The attached `repro.py --in-process` toggles this mode for
comparison.
## Ask
- Is this a known-fixed issue in `main` / 0.26?
- Which Relax/TIR pass is the most likely culprit given the minimum
reproducer's op profile (Conv + Slice + gated activation chain)
and the in-process state-carryover observation? A pointer to
which pass to bisect against would let me narrow further.
- Is the compilation-cache / device-kernel-cache lifecycle expected
to be process-scoped and safely mutable across `tvm.compile(...)`
calls?
Happy to reduce further if pointed at a specific pass or op class,
share intermediate outputs at any bisection anchor, or run against a
`tvm` build with candidate patches.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with repro.py and the 174-node synthetic ONNX graph, running LLVM and Metal in separate processes as described. Bisect the Relax/TIR compilation path around flow.0/enc/in_layers.1/Conv at node 174, comparing the relevant passes and compilation-cache behavior against the 170-node case. Done means Metal output matches LLVM within float32 rounding tolerance across fresh and repeated compilations.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100