diff --git a/python/llm/example/transformers_int4.py b/python/llm/example/transformers_int4.py new file mode 100644 index 00000000..6d128af1 --- /dev/null +++ b/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) diff --git a/python/llm/src/bigdl/llm/ggml/model/llama/llama_cpp.py b/python/llm/src/bigdl/llm/ggml/model/llama/llama_cpp.py index 60f982c0..956c3bfa 100644 --- a/python/llm/src/bigdl/llm/ggml/model/llama/llama_cpp.py +++ b/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 ################################################################################################### diff --git a/python/llm/src/bigdl/llm/transformers/__init__.py b/python/llm/src/bigdl/llm/transformers/__init__.py new file mode 100644 index 00000000..c6713d32 --- /dev/null +++ b/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 diff --git a/python/llm/src/bigdl/llm/transformers/convert.py b/python/llm/src/bigdl/llm/transformers/convert.py new file mode 100644 index 00000000..d95e0a64 --- /dev/null +++ b/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 diff --git a/python/llm/src/bigdl/llm/transformers/linear_int4.py b/python/llm/src/bigdl/llm/transformers/linear_int4.py new file mode 100644 index 00000000..63e5978a --- /dev/null +++ b/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 diff --git a/python/llm/src/bigdl/llm/transformers/model.py b/python/llm/src/bigdl/llm/transformers/model.py new file mode 100644 index 00000000..6a31fa89 --- /dev/null +++ b/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