FP16 casting pass does not guard against activation-level overflow (softplus, exp, logsumexp)

Open
#7 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
35/100
Issue type
Bug
Clarity
Mostly clear
Activity status
Quiet
Tech stack
python

Research direction

Start in casting.py by reading check_tensor_overflow_fp16(), handle_overflow_op(), and handle_non_overflow_op() to understand the current classification flow. Trace how softplus, logsumexp, and logcumsumexp are represented in the graph, then determine how an activation_overflow_audit would identify risky intermediates and connect to stable decomposition or fp32 accumulation. Done means the affected activation paths are audited and protected, with coverage for the listed overflow cases.

Written by the indexing model from the issue text.

Description

Problem

The FP16 casting pass in casting.py correctly guards against static tensor overflow (weights/constants > FP16_MAX = 65504), but does not account for activation-level overflow — where operations like exp(x) produce intermediate values > 65504 at runtime even though the inputs are within fp16 range.

Example: softplus

A model with a Softplus() activation layer:

  1. cast_fp32_to_fp16() successfully converts all weights and inputs to fp16 ✅
  2. The converter decomposes softplus(x) = log(1 + exp(x))
  3. At runtime on Apple Neural Engine, when x > 10.4:
    • exp(10.4) ≈ 32,900 (fits fp16)
    • exp(11.0) ≈ 59,874 (barely fits fp16)
    • exp(11.1) ≈ 66,686 → OVERFLOW → output collapses to 0

The current check_tensor_overflow_fp16() and handle_overflow_op() logic only checks scalar/tensor values, not whether the computation graph will produce intermediate overflows.

Affected Operations
Operation Naive Form fp16 Overflow Threshold
softplus exp(x) x ≈ 10.4
logsumexp sum(exp(x_i)) x ≈ 7.63
logcumsumexp cumsum(exp(x_i)) x ≈ 11.09
The Compound Effect

When coreai-optimization applies weight compression (palettization, quantization) AND fp16 casting together:

  1. Quantization introduces rounding errors in weights
  2. These errors can shift activation distributions
  3. Values that were safely below the overflow threshold may now exceed it
  4. The casting pass has no mechanism to detect or prevent this
Proposed Fix

Add an activation_overflow_audit pass that:

  1. Identifies ops in the graph whose intermediates can overflow fp16 (exp, log(1+exp(...)))
  2. Flags them for stable decomposition or fp32 accumulation
  3. Integrates with the existing handle_overflow_op / handle_non_overflow_op classification
Prior Art
  • apple/coremltools PRs #2725, #2726, #2727 fix the converter-level decomposition
  • apple/coreai-torch PR #22 adds stable converters for softplus/mish/logsumexp
  • This issue addresses the optimization layer, where the interaction between quantization and fp16 creates compound failures
Environment
  • coreai-optimization: latest (cloned June 21, 2026)
  • Related: apple/coreai-torch#21, apple/coremltools#2687
Dominant language
Python
Stars
131
Forks
36
Avg merge
2d 13h
Merged PRs (30d)
24

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from apple/coreai-optimization

All issues in apple/coreai-optimization

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.