You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
FP16 casting pass does not guard against activation-level overflow (softplus, exp, logsumexp) #7
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:
cast_fp32_to_fp16() successfully converts all weights and inputs to fp16 ✅
The converter decomposes softplus(x) = log(1 + exp(x))
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:
Quantization introduces rounding errors in weights
These errors can shift activation distributions
Values that were safely below the overflow threshold may now exceed it
The casting pass has no mechanism to detect or prevent this
Proposed Fix
Add an activation_overflow_audit pass that:
Identifies ops in the graph whose intermediates can overflow fp16 (exp, log(1+exp(...)))
Flags them for stable decomposition or fp32 accumulation
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
Thanks for raising this issue @Ashutosh0x . You're right that the casting utility currently only considers statically available tensors when choosing whether or not to cast parts of the model to lower precision.
Irrespective of any specific decompositions, it would be useful to enhance the casting utility to handle one or both of the following:
using calibration logic to analyze activation tensor ranges when determining casting locations
allowing the user to specify specific ops or op types to ignore being casted
Marking this as a feature request issue.
Before such handling is in place, a workaround could be to manually edit the original Pytorch model definition to substitute stable versions of ops like Softplus, avoiding the need for changes in either coreai-opt or coreai-torch.
I have opened PR #112 implementing the user-specified op skipping approach (ignored_ops: Sequence[Callable | str] | None = None) for cast_fp32_to_fp16 and cast_to_16_bit_precision.
It supports passing:
OpOverload (e.g. torch.ops.aten.exp.default)
OpOverloadPacket (e.g. torch.ops.aten.exp)
Callable (e.g. torch.exp)
String op names (e.g. "exp", "aten.exp")
Specific FX node names (e.g. "exp_1")
Ignored ops remain in FP32 with automatic input/output boundary casts, and consecutive ignored ops (like exp followed by log1p) have their intermediate round-trip casts automatically eliminated by _cleanup_casts. Unit and end-to-end tests and documentation are included in the PR.
I have opened PR #117 implementing the dynamic activation range calibration logic for cast_fp32_to_fp16 and cast_to_16_bit_precision via an optional calibration_data parameter.
Key details:
Automatic Overflow Observation: Simulates execution over representative calibration data (single tensors, tuples, kwargs dicts, or a DataLoader) using a lightweight torch.fx.Interpreter. Operations producing values exceeding the IEEE 754 float16 limit (> 65,504.0 or non-finite values) are automatically flagged to stay in FP32 precision with boundary casts inserted.
Consumer Cascade Propagation: Propagates overflow status to direct consumers and through operator.getitem (for multi-output ops like torch.chunk), preventing immediate saturation upon downcasting.
Minimal Footprint & Full Compatibility: All internal calibration logic is modularized in _utils/casting_utils.py without expanding the public API surface. Callers passing ignored_ops positionally remain 100% backward compatible.
Comprehensive Verification: 130 tests passing (including attention logits, residual blocks, kwargs, DataLoader, and validation error cases) with all 26 make check quality gates passing.
Please see #117 for the full implementation and diff.
Activity
crowbat commentedon Jun 23, 2026
Thanks for raising this issue @Ashutosh0x . You're right that the casting utility currently only considers statically available tensors when choosing whether or not to cast parts of the model to lower precision.
Irrespective of any specific decompositions, it would be useful to enhance the casting utility to handle one or both of the following:
Marking this as a feature request issue.
Before such handling is in place, a workaround could be to manually edit the original Pytorch model definition to substitute stable versions of ops like
Softplus, avoiding the need for changes in eithercoreai-optorcoreai-torch.rohith500 commentedon Sep 19, 2026
I have opened PR #112 implementing the user-specified op skipping approach (
ignored_ops: Sequence[Callable | str] | None = None) forcast_fp32_to_fp16andcast_to_16_bit_precision.It supports passing:
OpOverload(e.g.torch.ops.aten.exp.default)OpOverloadPacket(e.g.torch.ops.aten.exp)Callable(e.g.torch.exp)"exp","aten.exp")"exp_1")Ignored ops remain in FP32 with automatic input/output boundary casts, and consecutive ignored ops (like
expfollowed bylog1p) have their intermediate round-trip casts automatically eliminated by_cleanup_casts. Unit and end-to-end tests and documentation are included in the PR.feat(casting): support ignored_ops in cast_fp32_to_fp16 and cast_to_1…
crowbat commentedon Sep 23, 2026
Reopening this issue as the core idea of observing dynamic activation ranges to automatically exclude ops is still to be implemented.
rohith500 commentedon Sep 24, 2026
Hi @crowbat,
I have opened PR #117 implementing the dynamic activation range calibration logic for
cast_fp32_to_fp16andcast_to_16_bit_precisionvia an optionalcalibration_dataparameter.Key details:
DataLoader) using a lightweighttorch.fx.Interpreter. Operations producing values exceeding the IEEE 754 float16 limit (> 65,504.0 or non-finite values) are automatically flagged to stay in FP32 precision with boundary casts inserted.operator.getitem(for multi-output ops liketorch.chunk), preventing immediate saturation upon downcasting._utils/casting_utils.pywithout expanding the public API surface. Callers passingignored_opspositionally remain 100% backward compatible.make checkquality gates passing.Please see #117 for the full implementation and diff.
feat(casting): add dynamic activation range calibration in FP16 casti…
feat(casting): add dynamic activation range calibration in FP16 casti…
feat(casting): add dynamic activation range calibration in FP16 casti…