[ANE] FP16 numerical discrepancy in MobileNetV3 (2D MatMul + Hardswish)
- Dominant language
- Python
- Stars
- 156
- Forks
- 45
- Avg merge
- 1d 7m
- Merged PRs (30d)
- 12
Description
## 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](file:///Users/provision/LiteRT/temp_mobilenet_experiment/reproduce_mobilenet_v3_fp16_npu_bug.py):
```python
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)
Contributor guide
Research direction
Start by running reproduce_mobilenet_v3_fp16_npu_bug.py with coreai-torch v0.4.1 and compare the MobileNetV3 FP16 GPU/NPU outputs. Trace TorchConverter and get_decomp_table around the 2D Linear/MatMul feeding Hardswish, then verify the NPU discrepancy is resolved without changing the MobileNetV2 result.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100