Skip to content

GitLab

Commit e5d5fc7e authored by gaoqiong's avatar gaoqiong
Browse files

Update setup.py

parent 224ab8d9
Hide whitespace changes
Inline Side-by-side
Showing with 225 additions and 224 deletions
setup.py
View file @ e5d5fc7e
import os import os
import torch import torch
import platform import platform
import requests import requests
import subprocess import subprocess
from pathlib import Path from pathlib import Path
from setuptools import setup, find_packages from setuptools import setup, find_packages
from torch.utils.cpp_extension import CUDAExtension from torch.utils.cpp_extension import CUDAExtension
from typing import Optional, Union from typing import Optional, Union
add_git_version = False pwd = os.path.dirname(__file__)
if int(os.environ.get('ADD_GIT_VERSION', '0')) == 1: add_git_version = False
add_git_version = True if int(os.environ.get('ADD_GIT_VERSION', '0')) == 1:
add_git_version = True
def get_latest_kernels_version(repo):
""" def get_latest_kernels_version(repo):
Get the latest version of the kernels from the github repo. """
""" Get the latest version of the kernels from the github repo.
response = requests.get(f"https://api.github.com/repos/{repo}/releases/latest") """
data = response.json() response = requests.get(f"https://api.github.com/repos/{repo}/releases/latest")
tag_name = data["tag_name"] data = response.json()
version = tag_name.replace("v", "") tag_name = data["tag_name"]
return version version = tag_name.replace("v", "")
return version
def get_kernels_whl_url(
gpu_system_version, def get_kernels_whl_url(
release_version, gpu_system_version,
python_version, release_version,
platform, python_version,
architecture, platform,
): architecture,
""" ):
Get the url for the kernels wheel file. """
""" Get the url for the kernels wheel file.
return f"https://github.com/casper-hansen/AutoAWQ_kernels/releases/download/v{release_version}/autoawq_kernels-{release_version}+{gpu_system_version}-cp{python_version}-cp{python_version}-{platform}_{architecture}.whl" """
return f"https://github.com/casper-hansen/AutoAWQ_kernels/releases/download/v{release_version}/autoawq_kernels-{release_version}+{gpu_system_version}-cp{python_version}-cp{python_version}-{platform}_{architecture}.whl"
def get_sha(pytorch_root: Union[str, Path]) -> str:
try: def get_sha(pytorch_root: Union[str, Path]) -> str:
return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=pytorch_root).decode('ascii').strip() try:
except Exception: return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=pytorch_root).decode('ascii').strip()
return 'Unknown' except Exception:
return 'Unknown'
def get_abi():
try: def get_abi():
command = "echo '#include <string>' | gcc -x c++ -E -dM - | fgrep _GLIBCXX_USE_CXX11_ABI" try:
result = subprocess.run(command, shell=True, capture_output=True, text=True) command = "echo '#include <string>' | gcc -x c++ -E -dM - | fgrep _GLIBCXX_USE_CXX11_ABI"
output = result.stdout.strip() result = subprocess.run(command, shell=True, capture_output=True, text=True)
abi = "abi" + output.split(" ")[-1] output = result.stdout.strip()
return abi abi = "abi" + output.split(" ")[-1]
except Exception: return abi
return 'abiUnknown' except Exception:
return 'abiUnknown'
def get_version_add(sha: Optional[str] = None) -> str:
command = "git config --global --add safe.directory "+pwd def get_version_add(sha: Optional[str] = None) -> str:
result = subprocess.run(command, shell=True, capture_output=False, text=True) command = "git config --global --add safe.directory "+pwd
result = subprocess.run(command, shell=True, capture_output=False, text=True)
version=''
autoawq_root = os.path.dirname(os.path.abspath(__file__)) version=''
add_version_path = os.path.join(os.path.join(autoawq_root, "awq"), "__init__.py") autoawq_root = os.path.dirname(os.path.abspath(__file__))
if add_git_version: add_version_path = os.path.join(os.path.join(autoawq_root, "awq"), "__init__.py")
if sha != 'Unknown': if add_git_version:
if sha is None: if sha != 'Unknown':
sha = get_sha(autoawq_root) if sha is None:
version = 'das.opt1.' + sha[:7] sha = get_sha(autoawq_root)
else: version = 'das.opt1.' + sha[:7]
version = 'das.opt1' else:
# abi version = 'das.opt1'
#version += "." + get_abi() # abi
#version += "." + get_abi()
# dtk version
if os.getenv("ROCM_PATH"): # dtk version
rocm_path = os.getenv('ROCM_PATH', "") if os.getenv("ROCM_PATH"):
rocm_version_path = os.path.join(rocm_path, '.info', "rocm_version") rocm_path = os.getenv('ROCM_PATH', "")
with open(rocm_version_path, 'r',encoding='utf-8') as file: rocm_version_path = os.path.join(rocm_path, '.info', "rocm_version")
lines = file.readlines() with open(rocm_version_path, 'r',encoding='utf-8') as file:
rocm_version=lines[0][:-2].replace(".", "") lines = file.readlines()
version += ".dtk" + rocm_version rocm_version=lines[0][:-2].replace(".", "")
version += ".dtk" + rocm_version
# torch version
version += ".torch" + torch.__version__[:5] # torch version
version += ".torch" + torch.__version__[:5]
lines=[]
with open(add_version_path, 'r',encoding='utf-8') as file: lines=[]
lines = file.readlines() with open(add_version_path, 'r',encoding='utf-8') as file:
lines = file.readlines()
lines[1] = "__dcu_version__ = '0.2.5+{}'\n".format(version)
with open(add_version_path, encoding="utf-8",mode="w") as file: lines[1] = "__dcu_version__ = '0.2.5+{}'\n".format(version)
file.writelines(lines) with open(add_version_path, encoding="utf-8",mode="w") as file:
file.close() file.writelines(lines)
file.close()
def get_version():
get_version_add() def get_version():
version_file = 'awq/__init__.py' get_version_add()
with open(version_file, encoding='utf-8') as f: version_file = 'awq/__init__.py'
exec(compile(f.read(), version_file, 'exec')) with open(version_file, encoding='utf-8') as f:
return locals()['__dcu_version__'] exec(compile(f.read(), version_file, 'exec'))
return locals()['__dcu_version__']
AUTOAWQ_VERSION = ""
PYPI_BUILD = os.getenv("PYPI_BUILD", "0") == "1" AUTOAWQ_VERSION = ""
PYPI_BUILD = os.getenv("PYPI_BUILD", "0") == "1"
CUDA_VERSION = os.getenv("CUDA_VERSION", None) or torch.version.cuda
if CUDA_VERSION: CUDA_VERSION = os.getenv("CUDA_VERSION", None) or torch.version.cuda
CUDA_VERSION = "".join(CUDA_VERSION.split("."))[:3] if CUDA_VERSION:
CUDA_VERSION = "".join(CUDA_VERSION.split("."))[:3]
ROCM_VERSION = os.getenv("ROCM_VERSION", None) or torch.version.hip
if ROCM_VERSION: ROCM_VERSION = os.getenv("ROCM_VERSION", None) or torch.version.hip
if ROCM_VERSION.startswith("5.6"): if ROCM_VERSION:
ROCM_VERSION = "5.6.1" if ROCM_VERSION.startswith("5.6"):
elif ROCM_VERSION.startswith("5.7"): ROCM_VERSION = "5.6.1"
ROCM_VERSION = "5.7.1" elif ROCM_VERSION.startswith("5.7"):
ROCM_VERSION = "5.7.1"
ROCM_VERSION = "".join(ROCM_VERSION.split("."))[:3]
ROCM_VERSION = "".join(ROCM_VERSION.split("."))[:3]
if not PYPI_BUILD:
if CUDA_VERSION: if not PYPI_BUILD:
AUTOAWQ_VERSION += f"+cu{CUDA_VERSION}" if CUDA_VERSION:
elif ROCM_VERSION: AUTOAWQ_VERSION += f"+cu{CUDA_VERSION}"
#version_info = get_version() elif ROCM_VERSION:
AUTOAWQ_VERSION += get_version()#f"+rocm{ROCM_VERSION}" #version_info = get_version()
else: AUTOAWQ_VERSION += get_version()#f"+rocm{ROCM_VERSION}"
raise RuntimeError( else:
"Your system must have either Nvidia or AMD GPU to build this package." raise RuntimeError(
) "Your system must have either Nvidia or AMD GPU to build this package."
)
common_setup_kwargs = {
"version": AUTOAWQ_VERSION, common_setup_kwargs = {
"name": "autoawq", "version": AUTOAWQ_VERSION,
"author": "Casper Hansen", "name": "autoawq",
"license": "MIT", "author": "Casper Hansen",
"python_requires": ">=3.8.0", "license": "MIT",
"description": "AutoAWQ implements the AWQ algorithm for 4-bit quantization with a 2x speedup during inference.", "python_requires": ">=3.8.0",
"long_description": (Path(__file__).parent / "README.md").read_text( "description": "AutoAWQ implements the AWQ algorithm for 4-bit quantization with a 2x speedup during inference.",
encoding="UTF-8" "long_description": (Path(__file__).parent / "README.md").read_text(
), encoding="UTF-8"
"long_description_content_type": "text/markdown", ),
"url": "https://github.com/casper-hansen/AutoAWQ", "long_description_content_type": "text/markdown",
"keywords": ["awq", "autoawq", "quantization", "transformers"], "url": "https://github.com/casper-hansen/AutoAWQ",
"platforms": ["linux", "windows"], "keywords": ["awq", "autoawq", "quantization", "transformers"],
"classifiers": [ "platforms": ["linux", "windows"],
"Environment :: GPU :: NVIDIA CUDA :: 11.8", "classifiers": [
"Environment :: GPU :: NVIDIA CUDA :: 12", "Environment :: GPU :: NVIDIA CUDA :: 11.8",
"License :: OSI Approved :: MIT License", "Environment :: GPU :: NVIDIA CUDA :: 12",
"Natural Language :: English", "License :: OSI Approved :: MIT License",
"Programming Language :: Python :: 3.8", "Natural Language :: English",
"Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.8",
"Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.10",
"Programming Language :: C++", "Programming Language :: Python :: 3.11",
], "Programming Language :: C++",
} ],
}
requirements = [
"torch>=2.0.1", requirements = [
"transformers>=4.35.0", "torch>=2.0.1",
"tokenizers>=0.12.1", "transformers>=4.35.0",
"typing_extensions>=4.8.0", "tokenizers>=0.12.1",
"accelerate", "typing_extensions>=4.8.0",
"datasets", "accelerate",
"zstandard", "datasets",
] "zstandard",
]
try:
if ROCM_VERSION: try:
import exlv2_ext if ROCM_VERSION:
else: import exlv2_ext
import awq_ext else:
import awq_ext
KERNELS_INSTALLED = True
except ImportError: KERNELS_INSTALLED = True
KERNELS_INSTALLED = False except ImportError:
KERNELS_INSTALLED = False
# kernels can be downloaded from pypi for cuda+121 only
# for everything else, we need to download the wheels from github # kernels can be downloaded from pypi for cuda+121 only
if not KERNELS_INSTALLED and (CUDA_VERSION or ROCM_VERSION): # for everything else, we need to download the wheels from github
if CUDA_VERSION and CUDA_VERSION.startswith("12"): if not KERNELS_INSTALLED and (CUDA_VERSION or ROCM_VERSION):
requirements.append("autoawq-kernels") if CUDA_VERSION and CUDA_VERSION.startswith("12"):
elif CUDA_VERSION and CUDA_VERSION.startswith("11") or ROCM_VERSION in ["561", "571"]: requirements.append("autoawq-kernels")
gpu_system_version = ( elif CUDA_VERSION and CUDA_VERSION.startswith("11") or ROCM_VERSION in ["561", "571"]:
f"cu{CUDA_VERSION}" if CUDA_VERSION else f"rocm{ROCM_VERSION}" gpu_system_version = (
) f"cu{CUDA_VERSION}" if CUDA_VERSION else f"rocm{ROCM_VERSION}"
kernels_version = get_latest_kernels_version("casper-hansen/AutoAWQ_kernels") )
python_version = "".join(platform.python_version_tuple()[:2]) kernels_version = get_latest_kernels_version("casper-hansen/AutoAWQ_kernels")
platform_name = platform.system().lower() python_version = "".join(platform.python_version_tuple()[:2])
architecture = platform.machine().lower() platform_name = platform.system().lower()
latest_rocm_kernels_wheels = get_kernels_whl_url( architecture = platform.machine().lower()
gpu_system_version, latest_rocm_kernels_wheels = get_kernels_whl_url(
kernels_version, gpu_system_version,
python_version, kernels_version,
platform_name, python_version,
architecture, platform_name,
) architecture,
requirements.append(f"autoawq-kernels@{latest_rocm_kernels_wheels}") )
else: requirements.append(f"autoawq-kernels@{latest_rocm_kernels_wheels}")
raise RuntimeError( else:
"Your system have a GPU with an unsupported CUDA or ROCm version. " raise RuntimeError(
"Please install the kernels manually from https://github.com/casper-hansen/AutoAWQ_kernels" "Your system have a GPU with an unsupported CUDA or ROCm version. "
) "Please install the kernels manually from https://github.com/casper-hansen/AutoAWQ_kernels"
)
force_extension = os.getenv("PYPI_FORCE_TAGS", "0")
if force_extension == "1": force_extension = os.getenv("PYPI_FORCE_TAGS", "0")
# NOTE: We create an empty CUDAExtension because torch helps us with if force_extension == "1":
# creating the right boilerplate to enable correct targeting of # NOTE: We create an empty CUDAExtension because torch helps us with
# the autoawq-kernels package # creating the right boilerplate to enable correct targeting of
common_setup_kwargs["ext_modules"] = [ # the autoawq-kernels package
CUDAExtension( common_setup_kwargs["ext_modules"] = [
name="test_kernel", CUDAExtension(
sources=[], name="test_kernel",
) sources=[],
] )
]
setup(
packages=find_packages(), setup(
install_requires=requirements, packages=find_packages(),
extras_require={ install_requires=requirements,
"eval": ["lm_eval==0.4.1", "tabulate", "protobuf", "evaluate", "scipy"], extras_require={
"dev": ["black", "mkdocstrings-python", "mkdocs-material", "griffe-typingdoc"] "eval": ["lm_eval==0.4.1", "tabulate", "protobuf", "evaluate", "scipy"],
}, "dev": ["black", "mkdocstrings-python", "mkdocs-material", "griffe-typingdoc"]
**common_setup_kwargs, },
) **common_setup_kwargs,
)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment