Support directly quantizing huggingface transformers into 4bit format (#8371)

* Support directly quantizing huggingface transformers into 4bit format

* refine example

* license

* fix bias

* address comments

* move to ggml transformers

* fix example

* fix style

* fix style

* address comments

* rename

* change API

* fix style

* add lm head to conversion

* address comments

python/llm/example/transformers_int4.py

@@ -0,0 +1,36 @@
+#
+# Copyright 2016 The BigDL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import torch
+import os
+from bigdl.llm.transformers import AutoModelForCausalLM
+from transformers import LlamaTokenizer
+
+if __name__ == '__main__':
+    model_path = 'decapoda-research/llama-7b-hf'
+
+    # load_in_4bit=True in bigdl.llm.transformers will convert
+    # the relevant layers in the model into int4 format
+    model = AutoModelForCausalLM.from_pretrained(model_path, load_in_4bit=True)
+    tokenizer = LlamaTokenizer.from_pretrained(model_path)
+
+    input_str = "Once upon a time, there existed a little girl who liked to have adventures. She wanted to go to places and meet new people, and have fun"
+
+    with torch.inference_mode():
+        input_ids = tokenizer.encode(input_str, return_tensors="pt")
+        output = model.generate(input_ids, do_sample=False, max_new_tokens=32)
+        output_str = tokenizer.decode(output[0], skip_special_tokens=True)
+    print(output_str)

python/llm/src/bigdl/llm/ggml/model/llama/llama_cpp.py

@@ -953,6 +953,57 @@ def llama_print_system_info() -> bytes:
 _lib.llama_print_system_info.argtypes = []
 _lib.llama_print_system_info.restype = c_char_p
 
+
+# GGML API
+def ggml_quantize_q4_0(
+    src,  # type: ctypes.Array[ctypes.c_float] # type: ignore
+    dst: ctypes.c_void_p,
+    n: ctypes.c_int,
+    k: ctypes.c_int,
+    hist,  # type: ctypes.Array[ctypes.c_int64] # type: ignore
+) -> int:
+    return _lib.ggml_quantize_q4_0(src, dst, n, k, hist)
+
+
+_lib.ggml_quantize_q4_0.argtypes = [
+    ctypes.POINTER(ctypes.c_float),
+    ctypes.c_void_p,
+    ctypes.c_int,
+    ctypes.c_int,
+    ctypes.POINTER(ctypes.c_int64),
+]
+_lib.ggml_quantize_q4_0.restype = ctypes.c_size_t
+
+
+def ggml_compute_forward_mul_mat_q_fp32(src_0_ne,  # type: ctypes.Array[ctypes.c_int64]
+                                        src_0_data,  # type: ctypes.c_void_p
+                                        src_1_ne,  # type: ctypes.Array[ctypes.c_int64]
+                                        src_1_data,  # type: ctypes.c_void_p
+                                        result,  # type: ctypes.c_void_p
+                                        ) -> None:
+
+    # ctx = ctx.context
+    # src_0_ne = (ctypes.c_int64 * 2)(*src_0_ne)
+    # src_0_data = ctypes.c_void_p(src_0_data)
+    # src_1_ne = (ctypes.c_int64 * 2)(*src_1_ne)
+    # src_1_data = ctypes.c_void_p(src_1_data)
+    # result = ctypes.c_void_p(result)
+
+    return _lib.ggml_compute_forward_mul_mat_q_fp32(src_0_ne,
+                                                    src_0_data,
+                                                    src_1_ne,
+                                                    src_1_data,
+                                                    result)
+
+
+_lib.ggml_compute_forward_mul_mat_q_fp32.argtypes = [
+    ctypes.POINTER(ctypes.c_int64),
+    ctypes.c_void_p,
+    ctypes.POINTER(ctypes.c_int64),
+    ctypes.c_void_p,
+    ctypes.c_void_p
+]
+_lib.ggml_compute_forward_mul_mat_q_fp32.restype = None
 ###################################################################################################
 
 

python/llm/src/bigdl/llm/transformers/__init__.py

@@ -0,0 +1,18 @@
+#
+# Copyright 2016 The BigDL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from .convert import ggml_convert_int4
+from .model import AutoModelForCausalLM, AutoModel

python/llm/src/bigdl/llm/transformers/convert.py

@@ -0,0 +1,96 @@
+#
+# Copyright 2016 The BigDL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+# Some parts of this file is adapted from
+# https://github.com/huggingface/transformers/blob/v4.30.2/src/transformers/utils/bitsandbytes.py
+# which is licensed under Apache License 2.0:
+#
+# Copyright 2021 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+import torch.nn as nn
+from accelerate import init_empty_weights
+from bigdl.llm.transformers.linear_int4 import LinearInt4, ParamsInt4
+import warnings
+
+
+def _replace_with_int4_linear(model, modules_to_not_convert=None, current_key_name=None):
+    has_been_replaced = False
+    for name, module in model.named_children():
+        if current_key_name is None:
+            current_key_name = []
+
+        if isinstance(module, nn.Linear) and name not in modules_to_not_convert:
+            # Check if the current key is not in the `modules_to_not_convert`
+            if not any(key in ".".join(current_key_name) for key in modules_to_not_convert):
+                with init_empty_weights():
+
+                    new_linear = LinearInt4(
+                        module.in_features,
+                        module.out_features,
+                        module.bias is not None,
+                    )
+
+                    # Copy the weights
+                    new_linear._parameters['weight'] = ParamsInt4(data=module.weight.data,
+                                                                  requires_grad=False,
+                                                                  quantized=False,
+                                                                  _shape=None).to("cpu")
+                    if module.bias is not None:
+                        new_linear._parameters['bias'] = nn.Parameter(module.bias.data).to("cpu")
+
+                    model._modules[name] = new_linear
+                    has_been_replaced = True
+                    # Force requires grad to False to avoid unexpected errors
+                    model._modules[name].requires_grad_(False)
+
+        # Remove the last key for recursion
+        if len(list(module.children())) > 0:
+            _, has_been_replaced = _replace_with_int4_linear(
+                module,
+                modules_to_not_convert,
+                current_key_name,
+            )
+    return model, has_been_replaced
+
+
+def ggml_convert_int4(model):
+    modules_to_not_convert = []  # ["lm_head"]
+    model, has_been_replaced = _replace_with_int4_linear(
+        model, modules_to_not_convert, None
+    )
+    if not has_been_replaced:
+        warnings.warn(
+            "No linear modules were found in "
+            "your model. This can happen for some architectures such as gpt2 that uses Conv1D "
+            "instead of Linear layers. Please double check your model architecture, or submit "
+            "an issue on github if you think this is a bug."
+        )
+    return model

python/llm/src/bigdl/llm/transformers/linear_int4.py

@@ -0,0 +1,200 @@
+#
+# Copyright 2016 The BigDL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# Some parts of this file is adapted from
+# https://github.com/TimDettmers/bitsandbytes/blob/0.39.1/bitsandbytes/nn/modules.py
+# which is licensed under the MIT license:
+#
+# MIT License
+#
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+
+from typing import Optional, TypeVar, Union, overload
+from bigdl.llm.utils.common import invalidInputError
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, device, dtype, nn
+
+T = TypeVar("T", bound="torch.nn.Module")
+
+import bigdl.llm.ggml.model.llama.llama_cpp as ggml
+
+import torch
+import ctypes
+
+QK = 64  # todo read this value from libllama.so
+scale_size_in_bytes = 4
+block_size_in_bytes = QK // 2 + scale_size_in_bytes
+
+
+def ggml_convert_int4(tensor: torch.Tensor):
+
+    invalidInputError(tensor.dtype == torch.float,
+                      "Input tensor must be float32")
+    src = tensor.data.data_ptr()
+    src = ctypes.cast(src, ctypes.POINTER(ctypes.c_float))
+    n = tensor.numel()
+    invalidInputError(n % QK == 0,
+                      "Input tensor size must be multiple of 64")
+    k = tensor.shape[-1]
+    invalidInputError(k % QK == 0,
+                      "Last dim of input tensor must be multiple of 64")
+
+    dst_size = (n // QK) * block_size_in_bytes
+    dst_tensor = torch.empty(dst_size, dtype=torch.uint8)
+    dst = ctypes.c_void_p(dst_tensor.data.data_ptr())
+
+    hist = (ctypes.c_int64 * 16)()
+
+    ggml.ggml_quantize_q4_0(src, dst, n, k, hist)
+    return dst_tensor
+
+
+class ParamsInt4(torch.nn.Parameter):
+    def __new__(cls, data=None, requires_grad=True, old_data=None, quantized=False, _shape=None):
+        if data is None:
+            data = torch.empty(0)
+
+        self = torch.Tensor._make_subclass(cls, data, requires_grad)
+        self.data = data
+        self.quantized = quantized
+        self._shape = _shape
+        return self
+
+    def quantize(self, device):
+        if not self.quantized:
+            w = self.data.contiguous().float()
+            # self.old_data = self.data
+            w_4bit = ggml_convert_int4(w)
+            self.data = w_4bit
+            self.quantized = True
+            self._shape = w.shape
+        return self
+
+    def get_shape(self):
+        return self._shape
+
+    @overload
+    def to(self: T, device: Optional[Union[int, device]]=...,
+           dtype: Optional[Union[dtype, str]]=..., non_blocking: bool=...,) -> T:
+        ...
+
+    @overload
+    def to(self: T, dtype: Union[dtype, str], non_blocking: bool=...) -> T:
+        ...
+
+    @overload
+    def to(self: T, tensor: Tensor, non_blocking: bool=...) -> T:
+        ...
+
+    def to(self, *args, **kwargs):
+        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
+
+        if (device is not None and device.type == "cpu" and self.data.device.type == "cpu"):
+            return self.quantize(device)
+        else:
+            new_param = ParamsInt4(super().to(device=device,
+                                              dtype=dtype,
+                                              non_blocking=non_blocking),
+                                   requires_grad=self.requires_grad,
+                                   quantized=self.quantized,
+                                   _shape=self._shape)
+
+            return new_param
+
+
+def ggml_matmul_src1_x_src0_t(src0: torch.Tensor, src1: torch.Tensor, src0_shape: torch.Size):
+    if src1.dtype != torch.float32:
+        src1 = src1.float()
+
+    src0_ptr = src0.data_ptr() + (src0.storage_offset() * src0.element_size())
+    src1_ptr = src1.data_ptr() + (src1.storage_offset() * src1.element_size())
+
+    result_shape = (src1.shape[0], src0_shape[0])
+
+    result_t = torch.empty(result_shape, dtype=torch.float32)
+    result_ptr = result_t.data_ptr() + (result_t.storage_offset() * result_t.element_size())
+
+    src0_shape = tuple(reversed(src0_shape))
+    src1_shape = tuple(reversed(src1.shape))
+
+    # ctx_p = ctx.context
+    src_0_ne = (ctypes.c_int64 * 2)(*src0_shape)
+    src_0_data = ctypes.c_void_p(src0_ptr)
+    src_1_ne = (ctypes.c_int64 * 2)(*src1_shape)
+    src_1_data = ctypes.c_void_p(src1_ptr)
+    result_ptr = ctypes.c_void_p(result_ptr)
+
+    ggml.ggml_compute_forward_mul_mat_q_fp32(
+        # ctx=ctx_p,
+        src_0_ne=src_0_ne,
+        src_0_data=src_0_data,
+        src_1_ne=src_1_ne,
+        src_1_data=src_1_data,
+        result=result_ptr,
+    )
+
+    return result_t
+
+
+class LinearInt4(nn.Linear):
+    def __init__(self, input_features, output_features, bias=True):
+        super().__init__(input_features, output_features, bias)
+        self.weight = ParamsInt4(self.weight.data, requires_grad=False,
+                                 old_data=self.weight.data,
+                                 quantized=False, _shape=None)
+        self.in_len = input_features
+        self.out_len = output_features
+        self.weight_shape = (self.out_len, self.in_len)
+
+    def forward(self, x: torch.Tensor):
+        # weights are cast automatically as Int8Params, but the bias has to be cast manually
+        if self.bias is not None and self.bias.dtype != x.dtype:
+            self.bias.data = self.bias.data.to(x.dtype)
+
+        x_shape = x.shape
+        x = x.view(-1, x_shape[-1])
+
+        x0 = self.weight.data
+
+        result = ggml_matmul_src1_x_src0_t(x0, x, self.weight_shape)
+        new_shape = x_shape[:-1] + (self.out_len,)
+        result = result.view(new_shape)
+
+        if self.bias is not None:
+            result += self.bias
+
+        return result

python/llm/src/bigdl/llm/transformers/model.py

@@ -0,0 +1,45 @@
+#
+# Copyright 2016 The BigDL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import transformers
+import torch
+
+
+class _BaseAutoModelClass:
+
+    HF_MODEL = None
+
+    @classmethod
+    def from_pretrained(cls,
+                        *args,
+                        **kwargs):
+        load_in_4bit = kwargs.pop("load_in_4bit", False)
+        model = cls.HF_Model.from_pretrained(*args, **kwargs)
+
+        if load_in_4bit:
+            from .convert import ggml_convert_int4
+            model = model.to("cpu", torch.float32)
+            model = ggml_convert_int4(model)
+
+        return model
+
+
+class AutoModelForCausalLM(_BaseAutoModelClass):
+    HF_Model = transformers.AutoModelForCausalLM
+
+
+class AutoModel(_BaseAutoModelClass):
+    HF_Model = transformers.AutoModel