From 657ea0ee50e1544bb9ece19f58c31ccd822a8b32 Mon Sep 17 00:00:00 2001 From: Yina Chen <33650826+cyita@users.noreply.github.com> Date: Fri, 2 Jun 2023 17:03:17 +0800 Subject: [PATCH] [LLM] Fix linux load libs for NeoX and llama (#8257) * init * add lisence * fix style --- .../llm/ggml/model/gptneox/gptneox_cpp.py | 11 ++++--- .../bigdl/llm/ggml/model/llama/llama_cpp.py | 11 ++++--- python/llm/src/bigdl/llm/utils/__init__.py | 2 ++ .../src/bigdl/llm/utils/common/__init__.py | 2 +- python/llm/src/bigdl/llm/utils/utils.py | 32 +++++++++++++++++++ 5 files changed, 47 insertions(+), 11 deletions(-) create mode 100644 python/llm/src/bigdl/llm/utils/utils.py diff --git a/python/llm/src/bigdl/llm/ggml/model/gptneox/gptneox_cpp.py b/python/llm/src/bigdl/llm/ggml/model/gptneox/gptneox_cpp.py index 08ca9b16..b6ab13d8 100644 --- a/python/llm/src/bigdl/llm/ggml/model/gptneox/gptneox_cpp.py +++ b/python/llm/src/bigdl/llm/ggml/model/gptneox/gptneox_cpp.py @@ -63,28 +63,29 @@ from ctypes import ( ) import pathlib from bigdl.llm.utils.common import invalidInputError +from bigdl.llm.utils import get_avx_flags # Load the library def _load_shared_library(lib_base_name: str): # Determine the file extension based on the platform - if sys.platform.startswith("linux"): - lib_ext = ".so" - elif sys.platform == "darwin": + if sys.platform.startswith("linux") or sys.platform == "darwin": lib_ext = ".so" elif sys.platform == "win32": lib_ext = ".dll" else: invalidInputError(False, "Unsupported platform.") + avx = get_avx_flags() + # Construct the paths to the possible shared library names (python/llm/src/bigdl/llm/libs) _base_path = pathlib.Path(__file__).parent.parent.parent.parent.resolve() _base_path = _base_path / 'libs' # Searching for the library in the current directory under the name "libgptneox" (default name # for gptneoxcpp) and "gptneox" (default name for this repo) _lib_paths = [ - _base_path / f"lib{lib_base_name}{lib_ext}", - _base_path / f"{lib_base_name}{lib_ext}", + _base_path / f"lib{lib_base_name}{avx}{lib_ext}", + _base_path / f"{lib_base_name}{avx}{lib_ext}", ] if "GPTNEOX_CPP_LIB" in os.environ: 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 7f9e50ff..60f982c0 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 @@ -63,28 +63,29 @@ from ctypes import ( ) import pathlib from bigdl.llm.utils.common import invalidInputError +from bigdl.llm.utils import get_avx_flags # Load the library def _load_shared_library(lib_base_name: str): # Determine the file extension based on the platform - if sys.platform.startswith("linux"): - lib_ext = ".so" - elif sys.platform == "darwin": + if sys.platform.startswith("linux") or sys.platform == "darwin": lib_ext = ".so" elif sys.platform == "win32": lib_ext = ".dll" else: invalidInputError(False, "Unsupported platform.") + avx = get_avx_flags() + # Construct the paths to the possible shared library names _base_path = pathlib.Path(__file__).parent.parent.parent.parent.resolve() _base_path = _base_path / 'libs' # Searching for the library in the current directory under the name "libllama" (default name # for llamacpp) and "llama" (default name for this repo) _lib_paths = [ - _base_path / f"lib{lib_base_name}{lib_ext}", - _base_path / f"{lib_base_name}{lib_ext}", + _base_path / f"lib{lib_base_name}{avx}{lib_ext}", + _base_path / f"{lib_base_name}{avx}{lib_ext}", ] if "LLAMA_CPP_LIB" in os.environ: diff --git a/python/llm/src/bigdl/llm/utils/__init__.py b/python/llm/src/bigdl/llm/utils/__init__.py index dbdafd2a..0656c8b5 100644 --- a/python/llm/src/bigdl/llm/utils/__init__.py +++ b/python/llm/src/bigdl/llm/utils/__init__.py @@ -18,3 +18,5 @@ # physically located elsewhere. # Otherwise there would be module not found error in non-pip's setting as Python would # only search the first bigdl package and end up finding only one sub-package. + +from .utils import * diff --git a/python/llm/src/bigdl/llm/utils/common/__init__.py b/python/llm/src/bigdl/llm/utils/common/__init__.py index 7d318395..5b36cc3e 100644 --- a/python/llm/src/bigdl/llm/utils/common/__init__.py +++ b/python/llm/src/bigdl/llm/utils/common/__init__.py @@ -19,4 +19,4 @@ # Otherwise there would be module not found error in non-pip's setting as Python would # only search the first bigdl package and end up finding only one sub-package. -from .log4Error import invalidInputError +from .log4Error import invalidInputError, invalidOperationError diff --git a/python/llm/src/bigdl/llm/utils/utils.py b/python/llm/src/bigdl/llm/utils/utils.py new file mode 100644 index 00000000..f6064bca --- /dev/null +++ b/python/llm/src/bigdl/llm/utils/utils.py @@ -0,0 +1,32 @@ +# +# 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 sys +from bigdl.llm.utils.common import invalidInputError, invalidOperationError + + +def get_avx_flags(): + avx = "" + if sys.platform != "win32": + import subprocess + msg = subprocess.check_output(["lscpu"]).decode("utf-8") + if "avx512_vnni" in msg: + avx = "_avx512" + elif "avx2" in msg: + avx = "_avx2" + else: + invalidOperationError(False, "Unsupported CPUFLAGS.") + return avx