apple / apple/coremltools

linear_quantize_weights: divide-by-zero and NaN-to-int cast when a weight block is all zeros

Open
#2,853 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

## 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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.