linear_quantize_weights: divide-by-zero and NaN-to-int cast when a weight block is all zeros
- Dominant language
- Python
- Stars
- 5.4k
- Forks
- 850
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 10
Description
## Description
When a quantization block is entirely zero, `linear_quantize_weights` computes `scale = 0` and then evaluates `0 / 0`, producing `NaN` values that are subsequently cast to the integer weight dtype. Four `RuntimeWarning`s are emitted and a `scale` of exactly `0` is written into the resulting model.
An all-zero output channel is exactly what pruning produces, and joint pruning + quantization is a supported workflow (`prune_weights` followed by `linear_quantize_weights`).
## Reproduction
Self-contained, public APIs only. coremltools 9.0.
```python
import warnings
import numpy as np, torch, coremltools as ct
from coremltools.optimize.coreml import (
OpLinearQuantizerConfig, OptimizationConfig,
linear_quantize_weights, decompress_weights, get_weights_metadata,
)
torch.manual_seed(0)
linear = torch.nn.Linear(64, 32, bias=False).eval()
with torch.no_grad():
linear.weight[0].zero_() # output channel 0 pruned -> all zeros
mlmodel = ct.convert(
torch.jit.trace(linear, torch.randn(1, 64)),
inputs=[ct.TensorType(shape=(1, 64))],
minimum_deployment_target=ct.target.iOS18,
)
config = OptimizationConfig(global_config=OpLinearQuantizerConfig(
mode="LINEAR", dtype="int8", granularity="per_channel", weight_threshold=0))
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always")
quantized = linear_quantize_weights(mlmodel, config=config)
for w in caught:
if issubclass(w.category, RuntimeWarning):
print(f"{w.category.__name__}: {w.message} "
f"[coremltools/{w.filename.split('coremltools/')[-1]}:{w.lineno}]")
weight = next(md.val for md in get_weights_metadata(decompress_weights(quantized)).values()
if md.val.shape == (32, 64))
print("dequantized channel 0:", weight[0][:8])
```
### Output I get
```
RuntimeWarning: invalid value encountered in divide [coremltools/optimize/_utils.py:113]
RuntimeWarning: invalid value encountered in divide [coremltools/optimize/_utils.py:120]
RuntimeWarning: invalid value encountered in cast [coremltools/optimize/_utils.py:294]
RuntimeWarning: invalid value encountered in cast [coremltools/optimize/_utils.py:297]
dequantized channel 0: [0. 0. 0. 0. 0. 0. 0. 0.]
```
Inspecting the emitted `constexpr_blockwise_shift_scale` op shows `scale[0] == 0.0`, with `data[0]` and `offset[0]` being the result of casting `NaN` to `int8`.
### Output I expected
No warnings, and a valid (non-zero) scale for the all-zero channel. An all-zero block has nothing to quantize: it should deterministically produce `data = 0`, `offset = 0`, and a non-degenerate `scale`, dequantizing back to `0.0`.
## Root cause
`coremltools/optimize/_utils.py`, `quantize_weight_by_dtype`:
```python
scale = (val_max - val_min) / (q_val_max - q_val_min)
quantized_data = np.round(weight / scale) # line 113
...
zero_point = (q_val_min * val_max - q_val_max * val_min) / (val_max - val_min) # line 120
```
For an all-zero block, `val_min == val_max == 0`, so `scale == 0`. Both `weight / scale` and the `zero_point` expression evaluate `0 / 0` and yield `NaN`. Those `NaN`s reach the integer casts at lines 294 and 297.
## Precedent in this repo
The palettization path already guards precisely this condition, in `coremltools/optimize/coreml/_quantization_passes.py`:
```python
per_channel_scale[per_channel_scale == 0] = 1
```
The linear quantization path has no equivalent guard.
## Note on observed impact
I want to be accurate about severity rather than overstate it. On the platform I tested (x86-64, NumPy 2.5.3), the `NaN -> int8` casts land on `0`, and because `scale == 0` the dequantized channel reads back as `0.0` — the correct value, but arrived at by accident rather than by computation. The concrete problems are:
1. Four `RuntimeWarning`s surfaced to users running a supported prune + quantize workflow.
2. A `scale` of exactly `0` serialized into `constexpr_blockwise_shift_scale`, which is a degenerate quantization parameter.
3. Model contents that depend on `NaN -> int` cast behaviour, which NumPy documents as undefined.
## Related: `scale` can also underflow to zero
The same divide-by-zero is reachable without an all-zero block. Core ML weights are fp16 by default, and `scale` is computed in the weight dtype, so a channel whose entire range is below roughly `7.6e-6` underflows to `scale == 0`:
```
channel magnitude 1e-5 -> scale = 1.788e-07, dequantizes correctly
channel magnitude 1e-6 -> scale = 0.0, 'divide by zero encountered in divide', channel dequantizes to all zeros
```
This suggests the guard belongs on `scale == 0` (mirroring the palettization line above) rather than only on `val_max == val_min`.
## Environment
- coremltools 9.0 (pip)
- Python 3.12, NumPy 2.5.3, PyTorch 2.14.0+cpu
- Linux x86-64 (conversion and compression only; no on-device prediction)
A proposed fix is in #2851.
Contributor guide
Research direction
Start in coremltools/optimize/_utils.py at quantize_weight_by_dtype and trace the scale, quantized_data, and zero_point calculations shown in the report. Compare the guard in coremltools/optimize/coreml/_quantization_passes.py, then verify the all-zero and tiny-range cases produce no warnings, a nonzero scale, and zero data and offset while dequantizing to zero.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- numpy, python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100