logcumsumexp overflows on Apple Neural Engine in fp16 for inputs > ~11.09 (output collapses to inf then NaN)
- Dominant language
- Python
- Stars
- 5.4k
- Forks
- 850
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 10
Description
## Problem
The PyTorch `torch.logcumsumexp` converter produces `inf` / `NaN` on Apple Neural Engine (ANE) in fp16 for input values above ~11.09. This is because the converter computes `exp(x)` on raw input without any stabilization, and `exp(11.09) ~ 65,504` which is the fp16 maximum.
## Root Cause
The current converter at `converters/mil/frontend/torch/ops.py` line 2230 computes:
```python
exp = mb.exp(x=x) # raw exp, no max-shift
cumsumexp = mb.cumsum(x=exp, axis=dim)
res = mb.log(x=cumsumexp)
```
`exp(x)` is computed on raw input with zero stabilization. For any `x > ~11.09`, this overflows to `inf` in fp16, and the cumulative sum then propagates infinity.
## Reproduction
```python
import torch
import coremltools as ct
class LogCumSumExpModel(torch.nn.Module):
def forward(self, x):
return torch.logcumsumexp(x, dim=-1)
model = LogCumSumExpModel().eval()
x = torch.tensor([[1.0, 5.0, 10.0, 12.0, 15.0, 20.0, 50.0]])
traced = torch.jit.trace(model, x)
mlmodel = ct.convert(traced,
inputs=[ct.TensorType(shape=x.shape)],
compute_precision=ct.precision.FLOAT16)
pytorch_out = model(x).detach().numpy()
coreml_out = list(mlmodel.predict({'x_1': x.numpy()}).values())[0]
print('PyTorch:', pytorch_out)
print('CoreML:', coreml_out) # inf/NaN for positions with x > 11.09
```
## Fix
Use the standard max-shift stabilization:
`
logcumsumexp(x) = max(x) + log(cumsum(exp(x - max(x))))
`
By subtracting the global `max(x)` first, all `exp()` arguments are <= 0, so values are in (0, 1]. This is the same pattern used in the `logsumexp` stable decomposition (PR #2726).
Note: The global max is used rather than a running (cumulative) max because MIL does not provide a `cummax` op. The global max is always >= the running max at every position, so `exp(x_i - global_max) <= 1` for all i, guaranteeing no overflow. The trade-off is slightly more underflow for early positions when a much larger value appears later, but this does not affect correctness.
## Impact
Models using `torch.logcumsumexp` -- CTC decoders, autoregressive attention mechanisms, sequential probability models.
## Environment
- coremltools version: 9.0 (main branch)
- Affected compute unit: Neural Engine (fp16)
- Unaffected: CPU, GPU
## Related Issues
- #2690 -- reduce_log_sum_exp fp16 overflow (same pattern, fixed in PR #2726)
- #2687 -- Softplus fp16 overflow (same class, fixed in PR #2725)
- #2728 -- log_softmax fp16 underflow (same class, fix in PR #2727)
Contributor guide
Research direction
Start in converters/mil/frontend/torch/ops.py around line 2230, where the logcumsumexp converter builds exp, cumsum, and log. Reproduce the issue with the provided traced model and FLOAT16 conversion, then verify that the converted output remains finite and matches PyTorch for values above 11.09, including the supplied input sequence.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, tooling
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 76/100