# /// script
# requires-python = ">=3.11"
# dependencies = [
# "coreai-core==1.0.0b2",
# "coreai-torch==0.4.1",
# "diffusers",
# "timm",
# "einops",
# "pyyaml",
# "numpy",
# "huggingface_hub",
# ]
#
# [tool.uv]
# index-url = "https://pypi.org/simple"
# prerelease = "allow"
# index-strategy = "unsafe-best-match"
# ///
"""ANECCompile internal error: two transformer instances at 64x64 spatial extent,
value-dependent on LayerNorm affine weights.
"Compiler internal error: failed create split by input channel, graph is changed"
"- From NEFUSED_CONV Layer ... SCALE_BIAS: PerCoutScale: Y ... TRANSPOSE: [W-C] [C-W]"
Environment: macOS 27.0 (26A5388g), CoreAICompiler 3600.79.1, M5 Max, coreai-torch 0.4.1,
coreai-core 1.0.0b2 (USE_OS runtime), torch 2.x.
Setup (public code + public MIT weights):
git clone https://github.com/hustvl/Moebius reference
uv run repro_upstream.py # trained weights -> ANECCompile FAILED
uv run repro_upstream.py --random # same graph, random weights -> compiles + runs
Findings from bisection (each line one measured compile attempt):
* ONE such transformer at 64x64: compiles for neuralEngine.
* TWO chained at 64x64, trained weights: FAILS as above.
* The same two at 16x16 / 32x32 (inside their real down-blocks): compile.
* Same two modules, same graph, weights re-randomized: compiles (--random below).
* Random weights + ONLY the 4 trained `transformer_blocks.0.norm2.{weight,bias}` vectors
(the LayerNorm feeding the cross-attention, [320] each): FAILS.
* Synthetic values in the same range (uniform 0.42-1.17) at the same site: compiles —
the specific trained values matter, not their range.
The graph is otherwise fully ANE-clean: with the patches below every submodule AND every
single-transformer composition compiles and runs on the Neural Engine.
"""
import argparse
import importlib
import json
import os
import subprocess
import sys
import tempfile
import types
from pathlib import Path
import numpy as np
import torch
import yaml
HERE = Path(__file__).resolve().parent
# MOEBIUS_REF / MOEBIUS_CKPT are optional overrides for an existing clone / local checkpoint.
REF = Path(os.environ.get("MOEBIUS_REF", HERE / "reference"))
NUM_EMBEDDINGS = 20
# --------------------------------------------------------------------------- model loading
def load_unet():
"""Load the reference UNet without executing model_lib/__init__.py (it imports a
CUDA-only variant)."""
sys.path.insert(0, str(REF))
for name, path in [("model_lib", REF / "model_lib"),
("model_lib.nets", REF / "model_lib/nets"),
("model_lib.nets.layers", REF / "model_lib/nets/layers")]:
m = types.ModuleType(name)
m.__path__ = [str(path)]
sys.modules[name] = m
mod = importlib.import_module("model_lib.nets.unet_lambda_prune_lite")
cfg = yaml.safe_load((REF / "config/model_cfg/moebius.yaml").read_text())
model_cfg = dict(cfg["model"])
model_type = model_cfg.pop("model_type")
model_cfg["sample_size"] = cfg["data"]["image_size"] // cfg["vae"]["downsample_ratio"]
model_cfg["num_embeddings"] = NUM_EMBEDDINGS
net = getattr(mod, model_type)(**model_cfg)
ckpt = os.environ.get("MOEBIUS_CKPT")
if not ckpt:
from huggingface_hub import hf_hub_download
ckpt = hf_hub_download("hustvl/Moebius", "ft_places2/diffusion_pytorch_model.bin")
sd = torch.load(ckpt, map_location="cpu", weights_only=True)
unet_sd = {k[len("diff_model."):]: v for k, v in sd.items() if k.startswith("diff_model.")}
net.load_state_dict(unet_sd, strict=True)
return net.eval()
# ----------------------------------------------------- ANE-eligibility patches (all exact)
# These make the rest of the graph ANE-clean so the reported bug is reachable; each rewrite
# is numerically identical to the original (verified at fp32 to ~1e-7 relative).
def patch_lambda_einsums():
vλ = importlib.import_module("model_lib.nets.layers.λ.vanillaλ")
original = vλ._einsum
def _patched(eq, *ops):
if eq == 'n m k u, b u v m -> b n k v': # rank-6 decomposition otherwise
rel, V = ops
N, M, K, U = rel.shape
B, _, Vd, _ = V.shape
A = rel.permute(0, 2, 1, 3).reshape(1, N, K, M * U)
Bm = V.permute(0, 3, 1, 2).reshape(B, 1, M * U, Vd)
return (A @ Bm).contiguous()
if eq == 'b h k n, b n k v -> b h v n':
Q, lam = ops
return (Q.permute(0, 3, 1, 2) @ lam).permute(0, 2, 3, 1).contiguous()
return original(eq, *ops)
vλ._einsum = _patched
def patch_self_lambda_forward():
"""Rank-5-free self-attention: fold the Conv3d (depth kernel 1) into a batched Conv2d."""
import torch.nn.functional as F
vλ = importlib.import_module("model_lib.nets.layers.λ.vanillaλ")
def forward(self, x): # [b, hh, ww, c]
b, hh, ww, _ = x.shape
n = hh * ww
xc = x.permute(0, 3, 1, 2)
q, k, v = self.to_q(xc), self.to_k(xc), self.to_v(xc)
Q, V = self.norm_q(q), self.norm_v(v)
h, u = self.heads, self.u
dk, dv = q.shape[1] // h, V.shape[1] // u
Q = Q.reshape(b, h, dk, n)
k = k.reshape(b, u, dk, n).softmax(dim=-1)
V = V.reshape(b, u, dv, n)
lam_c = torch.einsum('b u k m, b u v m -> b k v', k, V)
Yc = torch.einsum('b h k n, b k v -> b h v n', Q, lam_c)
w2d = self.pos_conv.weight.squeeze(2)
Vb = V.reshape(b * dv, u, hh, ww)
lam = F.conv2d(Vb, w2d, self.pos_conv.bias, padding=self.pos_conv.padding[1])
lam = lam.reshape(b, dv, dk, n).permute(0, 3, 2, 1)
Yp = (Q.permute(0, 3, 1, 2) @ lam).permute(0, 2, 3, 1)
Y = Yc + Yp
return Y.reshape(b, h * dv, n).permute(0, 2, 1).reshape(b, hh, ww, h * dv)
vλ.MultiQuerySelfLambda.forward = forward
# ------------------------------------------------------------------------------- harness
def child_load(asset: str) -> None:
import asyncio
from coreai.runtime import AIModel, ComputeUnitKind, NDArray, SpecializationOptions
async def go():
options = SpecializationOptions.from_preferred_compute_unit_kind(
ComputeUnitKind.neural_engine())
model = await AIModel.load(asset, specialization_options=options)
fn = model.load_function("main")
meta = json.loads(Path(asset + ".inputs.json").read_text())
out = await fn({k: NDArray(np.zeros(v, dtype=np.float16)) for k, v in meta.items()})
print(f"[child] ran ok, outputs={list(out)}")
asyncio.run(go())
class TwoTx(torch.nn.Module):
def __init__(self, t0, t1):
super().__init__()
self.t0, self.t1 = t0, t1
def forward(self, x, ctx):
h = self.t0(x, encoder_hidden_states=ctx)[0]
return self.t1(h, encoder_hidden_states=ctx)[0]
def randomized(module):
"""Same architecture, weights re-randomized — the passing control."""
import copy
m = copy.deepcopy(module)
for sub in m.modules():
if isinstance(sub, (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d)):
sub.running_mean.data.zero_(); sub.running_var.data.fill_(1.0)
sub.weight.data.fill_(1.0); sub.bias.data.zero_()
elif isinstance(sub, (torch.nn.LayerNorm, torch.nn.GroupNorm)):
sub.weight.data.fill_(1.0); sub.bias.data.zero_()
elif isinstance(sub, (torch.nn.Conv2d, torch.nn.Conv3d, torch.nn.Linear)):
torch.nn.init.normal_(sub.weight, std=0.02)
if sub.bias is not None:
sub.bias.data.zero_()
for name, p in m.named_parameters():
if "rel_pos_emb" in name:
torch.nn.init.normal_(p, std=0.02)
return m
def main() -> None:
if len(sys.argv) > 2 and sys.argv[1] == "--load":
child_load(sys.argv[2])
return
ap = argparse.ArgumentParser()
ap.add_argument("--random", action="store_true",
help="re-randomize weights (same graph) — compiles fine")
args = ap.parse_args()
from coreai_torch import TorchConverter, get_decomp_table
net = load_unet()
patch_lambda_einsums()
patch_self_lambda_forward()
t0 = net.down_blocks[0].attentions[0]
t1 = net.down_blocks[0].attentions[1]
if args.random:
torch.manual_seed(0)
t0, t1 = randomized(t0), randomized(t1)
model = TwoTx(t0, t1).half().eval()
x = torch.randn(2, 320, 64, 64, dtype=torch.float16)
ctx = torch.randn(2, 10, 768, dtype=torch.float16)
with torch.no_grad():
model(x, ctx)
ep = torch.export.export(model, args=(x, ctx))
ep = ep.run_decompositions(get_decomp_table())
program = (TorchConverter()
.add_exported_program(ep, input_names=["x", "ctx"], output_names=["out"])
.to_coreai())
program.optimize()
out = Path(tempfile.mkdtemp()) / "repro.aimodel"
program.save_asset(out)
Path(str(out) + ".inputs.json").write_text(
json.dumps({"x": list(x.shape), "ctx": list(ctx.shape)}))
proc = subprocess.run([sys.executable, __file__, "--load", str(out)],
capture_output=True, text=True, timeout=900)
err = proc.stderr + proc.stdout
label = "random weights" if args.random else "trained weights"
if "ANECCompile" in err and "FAILED" in err:
detail = next((line for line in err.splitlines() if "split by input channel" in line), "")
print(f"[repro] {label}: ANE COMPILE FAILED {detail[:160]}")
elif proc.returncode != 0:
print(f"[repro] {label}: LOAD FAILED\n{err[-400:]}")
else:
print(f"[repro] {label}: compiled + ran on neuralEngine OK")
if __name__ == "__main__":
main()
Summary
ANECCompile()fails with a compiler internal error when a graph contains two instances of a particular transformer block at 64×64 spatial extent, while one instance compiles and runs on the Neural Engine. The failure is value-dependent: re-randomizing the weights of the same two-instance graph makes it compile, and bisection isolates the trigger to four small LayerNorm affine vectors.After the failure the runtime falls back and executes the graph on the GPU (a completed run is not evidence of ANE execution — worth noting for anyone benchmarking).
Environment
/System/Library/SubFrameworks/CoreAICompiler.framework)Reproduction
(
repro_upstream.pyattached below; it downloads the public checkpoint viahuggingface_hub, exports two chainedMixTransformer2DModelinstances withcoreai-torch, and requestsneuralEngineviaSpecializationOptions. The script includes three exact graph rewrites — einsum→matmul folds and a Conv3d→batched-Conv2d fold — that make the architecture otherwise fully ANE-clean, so the reported failure is the only remaining blocker.)Bisection evidence (each row one measured compile attempt)
failed create split by input channelnorm2.{weight,bias}vectors ([320] each; the LayerNorm feeding the cross-attention)So the failing pass appears to be the input-channel splitter (engaged only at the larger spatial extent), on a fused conv layer carrying a per-channel scale/bias derived from a LayerNorm affine — and only for the specific trained values, not their range. Every individual submodule of the model compiles for ANE in isolation.
Also possibly relevant: "graph is changed" suggests the split pass mutates shared state and cannot be applied to a second identical region in the same compile.
Two hand-built structural replicas did NOT reproduce
We attempted twice to reproduce with synthetic modules (same op inventory: GroupNorm + 1×1 conv projections, token↔spatial transposes, conv-projected multi-query linear attention with BatchNorm, gated depthwise-conv FFN, and the trained trigger vectors transplanted in) — both compile fine. The full real module structure is load-bearing in a way we could not distill further, which is why the repro leans on the public model.
Happy to provide the intermediate bisection scripts or run further experiments on request.
repro_upstream.py (validated on the environment above: trained → FAILED, --random → OK)