[ANE] FP16 numerical discrepancy in MobileNetV3 (2D MatMul + Hardswish)

[zli96]

Summary

When running FP16 models exported with coreai-torch on the Apple Neural Engine (NPU), MobileNetV3 exhibits significant output numerical discrepancies vs GPU/CPU (Max Abs Diff: ~0.199, Rel L2 Diff: ~3.92%). In contrast, MobileNetV2 outputs match GPU/CPU within standard FP16 tolerance (Max Abs Diff: < 0.003).

Key Empirical Findings

Test Case (FP16 NPU vs GPU)	Max Abs Diff	Rel L2 Diff
MobileNet V2 (Linear + ReLU/Identity in classifier)	0.002686	0.001025
MobileNet V3 Small (Linear + Hardswish in classifier)	0.199219	0.039235

Root Cause & Isolation

In my own experiment I was able to narrow down this problem to the following part at the end of mobileNet V3 compute graph produced by coreai-torch:

A 2D Linear / MatMul operation feeds directly into a Hardswish activation.

Transforming the 2D matrix into a 4D matrix (1 x 1 x m x n) avoids the issue on the NPU.

Reproduction Script

The reproduction script is available at reproduce_mobilenet_v3_fp16_npu_bug.py:

import asyncio
import numpy as np
import pathlib
import torch
import torchvision
from coreai_torch import TorchConverter, get_decomp_table
import coreai.runtime as r

async def run_reproduction():
    out_dir = pathlib.Path.cwd()
    spec_gpu = r.SpecializationOptions.from_preferred_compute_unit_kind(r.ComputeUnitKind.gpu())
    spec_npu = r.SpecializationOptions.from_preferred_compute_unit_kind(r.ComputeUnitKind.neural_engine())

    torch.manual_seed(42)
    in_np = torch.randn(1, 3, 224, 224).numpy().astype(np.float16)
    nd_in = r.NDArray(in_np)

    # MobileNet V3 Small (Linear + Hardswish classifier)
    model_v3 = torchvision.models.mobilenet_v3_small(weights=torchvision.models.MobileNet_V3_Small_Weights.DEFAULT).half().eval()
    ep_v3 = torch.export.export(model_v3, (torch.randn(1, 3, 224, 224).half(),))
    ep_v3 = ep_v3.run_decompositions(get_decomp_table())
    prog_v3 = TorchConverter().add_exported_program(ep_v3, input_names=['image'], output_names=['logits']).to_coreai()
    prog_v3.optimize()
    path_v3 = out_dir / "mobilenet_v3_fp16.aimodel"
    prog_v3.save_asset(path_v3)

    m_v3_gpu = await r.AIModel.load(path_v3, specialization_options=spec_gpu)
    m_v3_npu = await r.AIModel.load(path_v3, specialization_options=spec_npu)
    v3_gpu = (await m_v3_gpu.load_function("main")({"image": nd_in}))["logits"].numpy().astype(np.float32)
    v3_npu = (await m_v3_npu.load_function("main")({"image": nd_in}))["logits"].numpy().astype(np.float32)

    diff = float(np.max(np.abs(v3_npu - v3_gpu)))
    l2 = float(np.linalg.norm(v3_npu - v3_gpu) / np.linalg.norm(v3_gpu))
    print(f"MobileNet V3 Small FP16 (NPU vs GPU) -> Max Abs Diff: {diff:.6f}, Rel L2: {l2:.6f}")

if __name__ == "__main__":
    asyncio.run(run_reproduction())

Environment

• OS: macOS 27 beta 3
• version: coreai-torch (v0.4.1)