apple / apple/coreai-torch

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

Open
#51 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.