Hello,
I am trying to apply different kinds of quantization (static, dynamic and quantization-aware trainings) to a BERT model taken from the transformers library. My model works as expected with no quantization or with dynamic quantization, but both static quantization and quantization-aware training return a strange error which makes it seem as though QuantStub and DeQuantStub are changing tensor shapes.
Here is a simplified version, the comments indicate where the errors are and the reasoning for adding the various tweaks. I included the two error stacks described in the comments (error A and B) at the end this post.
def __init__(self, ...):
    ...
    self.model = AutoModelForTokenClassification.from_pretrained('bert-base-multilingual-cased', config=...)
    ...

def train(self, trainset, devset, num_epochs, **kwargs):
    self.trainer = Trainer(model=self.model, args=..., train_dataset=trainset, eval_dataset=devset)
    
    # Without any kind of quantization, the model works as expected
    
    # Case 1: dynamic quantization for Linear layers, works as expected
    if self.quantization == "dynamic":
        self.trainer.train()
        self.model = quant.quantize_dynamic(
            self.model, {nn.Linear}, dtype=getattr(torch, self.quant_type)
        )
        self.trainer.model = self.model

    # Case 2: quantization-aware training. The errors are exactly the same as for static quantization.
    elif self.quantization == "qat":
        self.model.quant = quant.QuantStub()
        self.model.dequant = quant.DeQuantStub()

            # This snippet is necessary in the first place because of https://discuss.pytorch.org/t/89154, otherwise I get a "AssertionError: The only supported dtype for nnq.Embedding is torch.quint8"
            if isinstance(module, nn.Embedding):
                # Both alternatives here cause different errors
                module.qconfig = None # Causes Error A (presumably because of operating on both quantized and non-quantized tensors)
                module.qconfig = quant.float_qparams_dynamic_qconfig # Causes Error B (why?)

        self.model.qconfig = quant.get_default_qat_qconfig('fbgemm')
        quant.prepare_qat(self.model, inplace=True)

        Trainer.compute_loss = quant_compute_loss
        Trainer.prediction_step = quant_prediction_step
        self.model.train()
        self.trainer.train()
        self.model.eval()

        self.model = quant.convert(self.model, inplace=True)
        self.trainer.model = self.model

    # Case 3: static quantization, not included here. The code is very similar and I get the exact same errors in the same place.
    elif self.quantization == "static":
        pass

    self.trainer.save_model('model')

Error A (quantization disabled for Embedding layers):
Traceback (most recent calls WITHOUT Sacred internals):
  File "/home/username/project/filename.py", line 402, in predict
    batch_pred_ids, label_ids, _ = self.trainer.predict(dataset)
  File "/home/username/.local/lib/python3.6/site-packages/transformers/trainer.py", line 1355, in predict
    return self.prediction_loop(test_dataloader, description="Prediction")
  File "/home/username/.local/lib/python3.6/site-packages/transformers/trainer.py", line 1417, in prediction_loop
    loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only)
  File "/home/username/project/filename.py", line 79, in quant_prediction_step
    outputs = model(**inputs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/username/.local/lib/python3.6/site-packages/transformers/modeling_bert.py", line 1539, in forward
    return_dict=return_dict,
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/username/.local/lib/python3.6/site-packages/transformers/modeling_bert.py", line 838, in forward
    input_ids=input_ids, position_ids=position_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/username/.local/lib/python3.6/site-packages/transformers/modeling_bert.py", line 202, in forward
    embeddings = self.LayerNorm(embeddings)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/quantized/modules/normalization.py", line 25, in forward
    eps=self.eps, output_scale=self.scale, output_zero_point=self.zero_point)
RuntimeError: Could not run 'quantized::layer_norm' with arguments from the 'CPU' backend. 'quantized::layer_norm' is only available for these backends: [QuantizedCPU, BackendSelect, Named, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, Tracer, Autocast, Batched,
 VmapMode].

Error B (modified qconfig for embedding layers):
Traceback (most recent calls WITHOUT Sacred internals):
  File "/home/username/project/filename.py", line 401, in predict
    batch_pred_ids, label_ids, _ = self.trainer.predict(dataset)
  File "/home/username/.local/lib/python3.6/site-packages/transformers/trainer.py", line 1355, in predict
    return self.prediction_loop(test_dataloader, description="Prediction")
  File "/home/username/.local/lib/python3.6/site-packages/transformers/trainer.py", line 1417, in prediction_loop
    loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only)
  File "/home/username/project/filename.py", line 79, in quant_prediction_step
    outputs = model(**inputs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/username/.local/lib/python3.6/site-packages/transformers/modeling_bert.py", line 1539, in forward
    return_dict=return_dict,
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/username/.local/lib/python3.6/site-packages/transformers/modeling_bert.py", line 838, in forward
    input_ids=input_ids, position_ids=position_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/username/.local/lib/python3.6/site-packages/transformers/modeling_bert.py", line 201, in forward
    embeddings = inputs_embeds + position_embeddings + token_type_embeddings
RuntimeError: The size of tensor a (1024) must match the size of tensor b (128) at non-singleton dimension 0

Errors A seems understandable to me, but I’m not sure what is going on with that second one. If I check manually, without quantization, inputs_embeds, position_embeddings and token_type_embeddings have shape (8, 128, 768), or sometimes (1, 128, 768) at the point where the error occurs. But with QAT (or static quantization), they have shape (1024, 768) or (128, 768) instead, as if the first two dimensions had been concatenated. Is there a way to change them back to their correct shape, maybe by changing the qconfig? I’d rather not edit the library, especially considering that my model type is determined at runtime.