simcop2387/AutoGPTQ
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages 1 Wiki Activity

update setup.py

Browse source
This commit is contained in:
PanQiWei 2023-04-25 18:50:21 +08:00
parent 419160b733
commit 0bb84053a2
4 changed files with 39 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
auto_gptq/quantization/setup_cuda.py
View file

@ -1,10 +0,0 @@
from setuptools import setup, Extension
from torch.utils import cpp_extension
setup(
name='quant_cuda',
ext_modules=[cpp_extension.CUDAExtension(
'quant_cuda', ['quant_cuda.cpp', 'quant_cuda_kernel.cu']
)],
cmdclass={'build_ext': cpp_extension.BuildExtension}
)

0
auto_gptq/quantization/quant_cuda.cpp → quant_cuda/quant_cuda.cpp
View file

0
auto_gptq/quantization/quant_cuda_kernel.cu → quant_cuda/quant_cuda_kernel.cu
View file

62
setup.py
View file

@ -1,9 +1,11 @@
from os.path import abspath, dirname, join from setuptools import setup, find_packages
from setuptools import setup, find_packages, Extension
from torch.utils import cpp_extension try:
import torch
TORCH_AVAILABLE = True
except ImportError:
TORCH_AVAILABLE = False
project_root = dirname(abspath(__file__))
requirements = [ requirements = [
"datasets", "datasets",
@ -11,29 +13,43 @@ requirements = [
"rouge", "rouge",
"torch>=1.13.0", "torch>=1.13.0",
"safetensors", "safetensors",
"transformers>=4.26.1" "transformers>=4.26.1",
"triton>=2.0.0"
] ]
extras_require = { extras_require = {
"llama": ["transformers>=4.28.0"] "llama": ["transformers>=4.28.0"]
} }
extensions = [
cpp_extension.CUDAExtension(
"quant_cuda",
[
join(project_root, "auto_gptq/quantization/quant_cuda.cpp"),
join(project_root, "auto_gptq/quantization/quant_cuda_kernel.cu")
]
)
]
setup( if TORCH_AVAILABLE:
name="auto_gptq", from torch.utils import cpp_extension
packages=find_packages(),
version="v0.0.4-dev", extensions = [
install_requires=requirements, cpp_extension.CUDAExtension(
extras_require=extras_require, "quant_cuda",
ext_modules=extensions, [
cmdclass={'build_ext': cpp_extension.BuildExtension} "quant_cuda/quant_cuda.cpp",
) "quant_cuda/quant_cuda_kernel.cu"
]
)
]
setup(
name="auto_gptq",
packages=find_packages(),
version="v0.0.4-dev",
install_requires=requirements,
extras_require=extras_require,
include_dirs=["quant_cuda"],
ext_modules=extensions,
cmdclass={'build_ext': cpp_extension.BuildExtension}
)
else:
setup(
name="auto_gptq",
packages=find_packages(),
version="v0.0.4-dev",
install_requires=requirements,
extras_require=extras_require,
include_dirs=["quant_cuda"]
)