Bloom float16 quantization fail

Annie-Sihan-Chen · September 26, 2022, 8:07pm

I try to use quantize_fx to quantize huggingface bloom model to float16.
I use exactly the same way as the torch quantization test/tutorial file. The code runs well except that when I try to print the result type, it still float32 type which means the quantization never happens.
Here is the patch I use:

gist.github.com

https://gist.github.com/AmosLewis/20811123b15bb5d818a3e8da545075c8

bloom_fp16.py

import torch
from transformers import AutoModelForSequenceClassification

class HuggingFaceLanguage(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = AutoModelForSequenceClassification.from_pretrained(
            "bigscience/bloom-560m",  # The pretrained model.
            num_labels=2,  # The number of output labels--2 for binary classification.
            output_attentions=False,  # Whether the model returns attentions weights.

This file has been truncated. show original

jerryzh168 · September 27, 2022, 9:17pm

the result will stay in float32 by default, you will need to check the ops used in the model to see if the model is actually quantized or not, by printing the model itself.