[Quantization] int8 compute with PostTraining Quantization (W8A8)
- Dominant language
- Python
- Stars
- 5.4k
- Forks
- 850
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 10
Description
## ❓Question
Im trying to understand the runtime behavior of W8A8 quantized networks on CoreML. I have written up a very simple model as follows :
```python
class SimpleModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(1, 3, kernel_size=3, stride=1, padding=1)
self.relu = nn.ReLU()
def forward(self, x):
x = self.conv(x)
x = self.relu(x)
return x
```
I convert this model to W8A8 format using post-training quantisation and export it to an `.mlpackage` via the following code :
```python
def quantise(
model: torch.nn.Module,
) -> torch.nn.Module:
import tqdm
from coremltools.optimize.torch.quantization import (
LinearQuantizer,
LinearQuantizerConfig,
ModuleLinearQuantizerConfig,
)
# make dataloader for calibration
dataloader = ...
config = LinearQuantizerConfig(
global_config=ModuleLinearQuantizerConfig(
weight_dtype="qint8",
activation_dtype="quint8",
quantization_scheme="symmetric",
milestones=[0, 1000, 1000, 0],
)
)
quantizer = LinearQuantizer(model, config)
example_inputs = next(iter(dataloader))
quantizer.prepare(example_inputs=example_inputs, inplace=True)
quantizer.step()
# Do a forward pass through the model with calibration data
for data in tqdm.tqdm(dataloader, desc="calibrating"):
with torch.no_grad():
model(data)
quantized_model = quantizer.finalize()
return quantized_model
def export_super_simple():
import coremltools as ct
import torch
torch_model = SimpleModel()
torch_model.eval()
torch_model.cpu()
# quantise model
torch_model = quantise(torch_model)
# trace model for export
inputs = torch.rand(1, 1, 480, 640)
input_shapes = [ct.TensorType(shape=inputs.shape)]
outputs = [
ct.TensorType(name="scores"),
]
export_model = torch.jit.trace(torch_model, inputs, strict=True)
_ = export_model(inputs)
# coreml export
model = ct.convert(
export_model,
inputs=input_shapes,
minimum_deployment_target=ct.target.iOS18,
debug=True,
outputs=outputs,
)
model.save(f"super_simple_model.mlpackage")
```
From the graph of the exported coreml-model :
* the inputs and weights of the convolution operator are dequantised to fp-16
* conv + relu run in fp-16
* relu outputs are quantized again to int8
My questions are :
* Is this the expected behavior to have dequant / quant applied at operator boundaries while actual compute runs in f16? From the documentation of coremltools [(link)](https://apple.github.io/coremltools/docs-guides/source/opt-quantization-perf.html#results) it looks like there's quite a bit of speed up that one can expect with W8A8 models. Do those w8a8 models run compute at fp-16? If so is the speedup coming from having lesser activation data between layers?
* How do I get the exported model to run integer computation? that is instead of having these dequant -> quant operators before and after each op, I would like to see the network inputs quantized to int8 right at the beginning -> whole pipeline runs only on int8 tensors -> final network outputs dequantised back to f16. My goal is to see if a pure integer only compute gives me any inference speedup.
Contributor guide
Research direction
Review the quantise() and export_super_simple() entry points, along with the linked quantization-performance documentation, to determine what graph behavior is expected. A useful resolution should explain the observed dequantize/quantize boundaries and state whether and how integer-only computation can be requested, including any supported constraints.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, tooling
- Issue type
- Documentation
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100