apple / apple/coremltools

logcumsumexp overflows on Apple Neural Engine in fp16 for inputs > ~11.09 (output collapses to inf then NaN)

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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.