PetraAI/Nashmi

1import os2import sys3from pathlib import Path4from setuptools import setup, find_packages5 6 7common_setup_kwargs = {8    "version": "0.4.1",9    "name": "auto_gptq",10    "author": "PanQiWei",11    "description": "An easy-to-use LLMs quantization package with user-friendly apis, based on GPTQ algorithm.",12    "long_description": (Path(__file__).parent / "README.md").read_text(encoding="UTF-8"),13    "long_description_content_type": "text/markdown",14    "url": "https://github.com/PanQiWei/AutoGPTQ",15    "keywords": ["gptq", "quantization", "large-language-models", "transformers"],16    "platforms": ["windows", "linux"],17    "classifiers": [18        "Environment :: GPU :: NVIDIA CUDA :: 11.7",19        "Environment :: GPU :: NVIDIA CUDA :: 11.8",20        "Environment :: GPU :: NVIDIA CUDA :: 12.0",21        "License :: OSI Approved :: MIT License",22        "Natural Language :: Chinese (Simplified)",23        "Natural Language :: English",24        "Programming Language :: Python :: 3.8",25        "Programming Language :: Python :: 3.9",26        "Programming Language :: Python :: 3.10",27        "Programming Language :: Python :: 3.11",28        "Programming Language :: C++",29    ]30}31 32 33BUILD_CUDA_EXT = int(os.environ.get('BUILD_CUDA_EXT', '1')) == 134if BUILD_CUDA_EXT:35    try:36        import torch37    except:38        print("Building cuda extension requires PyTorch(>=1.13.0) been installed, please install PyTorch first!")39        sys.exit(-1)40 41    CUDA_VERSION = None42    ROCM_VERSION = os.environ.get('ROCM_VERSION', None)43    if ROCM_VERSION and not torch.version.hip:44        print(45            f"Trying to compile auto-gptq for RoCm, but PyTorch {torch.__version__} "46            "is installed without RoCm support."47        )48        sys.exit(-1)49 50    if not ROCM_VERSION:51        default_cuda_version = torch.version.cuda52        CUDA_VERSION = "".join(os.environ.get("CUDA_VERSION", default_cuda_version).split("."))53 54    if ROCM_VERSION:55        common_setup_kwargs['version'] += f"+rocm{ROCM_VERSION}"56    else:57        if not CUDA_VERSION:58            print(59                f"Trying to compile auto-gptq for CUDA, byt Pytorch {torch.__version__} "60                "is installed without CUDA support."61            )62            sys.exit(-1)63        common_setup_kwargs['version'] += f"+cu{CUDA_VERSION}"64 65 66requirements = [67    "accelerate>=0.19.0",68    "datasets",69    "numpy",70    "rouge",71    "torch>=1.13.0",72    "safetensors",73    "transformers>=4.31.0",74    "peft"75]76 77extras_require = {78    "triton": ["triton==2.0.0"],79    "test": ["parameterized"]80}81 82include_dirs = ["autogptq_cuda"]83 84additional_setup_kwargs = dict()85if BUILD_CUDA_EXT:86    from torch.utils import cpp_extension87 88    if not ROCM_VERSION:89        from distutils.sysconfig import get_python_lib90        conda_cuda_include_dir = os.path.join(get_python_lib(), "nvidia/cuda_runtime/include")91 92        print("conda_cuda_include_dir", conda_cuda_include_dir)93        if os.path.isdir(conda_cuda_include_dir):94            include_dirs.append(conda_cuda_include_dir)95            print(f"appending conda cuda include dir {conda_cuda_include_dir}")96    extensions = [97        cpp_extension.CUDAExtension(98            "autogptq_cuda_64",99            [100                "autogptq_cuda/autogptq_cuda_64.cpp",101                "autogptq_cuda/autogptq_cuda_kernel_64.cu"102            ]103        ),104        cpp_extension.CUDAExtension(105            "autogptq_cuda_256",106            [107                "autogptq_cuda/autogptq_cuda_256.cpp",108                "autogptq_cuda/autogptq_cuda_kernel_256.cu"109            ]110        )111    ]112 113    if os.environ.get("INCLUDE_EXLLAMA_KERNELS", "1") == "1":  # TODO: improve github action to always compile exllama_kernels114        extensions.append(115            cpp_extension.CUDAExtension(116                "exllama_kernels",117                [118                    "autogptq_cuda/exllama/exllama_ext.cpp",119                    "autogptq_cuda/exllama/cuda_buffers.cu",120                    "autogptq_cuda/exllama/cuda_func/column_remap.cu",121                    "autogptq_cuda/exllama/cuda_func/q4_matmul.cu",122                    "autogptq_cuda/exllama/cuda_func/q4_matrix.cu"123                ]124            )125        )126 127    additional_setup_kwargs = {128        "ext_modules": extensions,129        "cmdclass": {'build_ext': cpp_extension.BuildExtension}130    }131common_setup_kwargs.update(additional_setup_kwargs)132setup(133    packages=find_packages(),134    install_requires=requirements,135    extras_require=extras_require,136    include_dirs=include_dirs,137    python_requires=">=3.8.0",138    **common_setup_kwargs139)140