search

zama-ai / concrete-ml

Concrete ML: Privacy Preserving ML framework using Fully Homomorphic Encryption (FHE), built on top of Concrete, with bindings to traditional ML frameworks.
Other
954 stars 139 forks source link

Two consecutive Unsqueeze operations in QAT model throws error at compilation time #682

Open raeell opened 5 months ago

raeell commented 5 months ago

Summary

When using two consecutive reshape operations in a QAT model, such as two torch.unsqueeze operations, the compilation throws the ValueError "Could not determine a unique scale for the quantization! Please check the ONNX graph of this model.", even though the two torch.unsqueeze are between two QuantIdentity layers. It seems that this especially happens when passing a parameter n_bits to compile_brevitas_qat_model that is smaller than the bit width used for the QAT model. For instance, using a bit width of 8 for the QuantIdentity layers in the QAT model, but choosing n_bits=6.

Description

Minimal code to reproduce the bug:

```python import brevitas.nn as qnn import torch import torch.nn as nn from concrete.ml.torch.compile import compile_brevitas_qat_model class Unsqueeze(nn.Module): def __init__(self, bit_width): super().__init__() self.id1 = qnn.QuantIdentity(bit_width=bit_width) self.conv1 = qnn.QuantConv2d(1, 1, 1, bit_width=bit_width, bias=False) def forward(self, x): """Forward pass of the model.""" x = self.id1(x) x = x.unsqueeze(1) x = x.unsqueeze(1) x = self.id1(x) x = self.conv1(x) return x model = Unsqueeze(bit_width=8) tensor_ = torch.randn(1, 200) compile_brevitas_qat_model(model, tensor_, verbose=False, n_bits=8) print("Compilation with 8 bits successful") compile_brevitas_qat_model(model, tensor_, verbose=False, n_bits=7) print("Compilation with 7 bits successful") try: compile_brevitas_qat_model(model, tensor_, verbose=False, n_bits=6) except Exception as e: print(e) print("Compilation with 6 bits failed") ```

andrei-stoian-zama commented 5 months ago

Thank you for raising this issue, we'll look into it!

jfrery commented 5 months ago

Hi @raeell,

You are correct the error message isn't really clear here. Some assertion seem to sometimes fail when setting less bits in the compilation than in brevitas layers.

Brevitas already has all the necessary n_bits information so the n_bits parameter at compilation time should be hidden from the user. We will improve this. Thanks for the issue!