2020-02-01 02:57:11 +03:00
|
|
|
"""
|
|
|
|
Copyright 2020 The Microsoft DeepSpeed Team
|
|
|
|
|
|
|
|
DeepSpeed library
|
|
|
|
|
|
|
|
Create a new wheel via the following command: python setup.py bdist_wheel
|
|
|
|
|
|
|
|
The wheel will be located at: dist/*.whl
|
|
|
|
"""
|
|
|
|
|
2020-03-12 23:08:58 +03:00
|
|
|
import os
|
2020-02-01 02:57:11 +03:00
|
|
|
import torch
|
|
|
|
from setuptools import setup, find_packages
|
|
|
|
from torch.utils.cpp_extension import CUDAExtension, BuildExtension
|
|
|
|
|
|
|
|
cmdclass = {}
|
|
|
|
cmdclass['build_ext'] = BuildExtension
|
|
|
|
|
2020-03-12 23:08:58 +03:00
|
|
|
TORCH_MAJOR = int(torch.__version__.split('.')[0])
|
|
|
|
TORCH_MINOR = int(torch.__version__.split('.')[1])
|
|
|
|
|
|
|
|
if not torch.cuda.is_available():
|
|
|
|
# Fix to allow docker buils, similar to https://github.com/NVIDIA/apex/issues/486
|
|
|
|
print(
|
|
|
|
"[WARNING] Torch did not find cuda available, if cross-compling or running with cpu only "
|
|
|
|
"you can ignore this message. Adding compute capability for Pascal, Volta, and Turing "
|
|
|
|
"(compute capabilities 6.0, 6.1, 6.2)")
|
|
|
|
if os.environ.get("TORCH_CUDA_ARCH_LIST", None) is None:
|
|
|
|
os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5"
|
|
|
|
|
|
|
|
# Fix from apex that might be relevant for us as well, related to https://github.com/NVIDIA/apex/issues/456
|
|
|
|
version_ge_1_1 = []
|
|
|
|
if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 0):
|
|
|
|
version_ge_1_1 = ['-DVERSION_GE_1_1']
|
|
|
|
version_ge_1_3 = []
|
|
|
|
if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 2):
|
|
|
|
version_ge_1_3 = ['-DVERSION_GE_1_3']
|
|
|
|
version_ge_1_5 = []
|
|
|
|
if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 4):
|
|
|
|
version_ge_1_5 = ['-DVERSION_GE_1_5']
|
|
|
|
version_dependent_macros = version_ge_1_1 + version_ge_1_3 + version_ge_1_5
|
|
|
|
|
2020-05-29 23:15:36 +03:00
|
|
|
ext_modules = [
|
|
|
|
CUDAExtension(
|
|
|
|
name='deepspeed_lamb_cuda',
|
|
|
|
sources=['csrc/lamb/fused_lamb_cuda.cpp',
|
|
|
|
'csrc/lamb/fused_lamb_cuda_kernel.cu'],
|
|
|
|
include_dirs=['csrc/includes'],
|
|
|
|
extra_compile_args={
|
|
|
|
'cxx': [
|
|
|
|
'-O3',
|
|
|
|
] + version_dependent_macros,
|
|
|
|
'nvcc': ['-O3',
|
|
|
|
'--use_fast_math'] + version_dependent_macros
|
|
|
|
}),
|
|
|
|
CUDAExtension(name='deepspeed_transformer_cuda',
|
|
|
|
sources=[
|
|
|
|
'csrc/transformer/ds_transformer_cuda.cpp',
|
|
|
|
'csrc/transformer/cublas_wrappers.cu',
|
|
|
|
'csrc/transformer/transform_kernels.cu',
|
|
|
|
'csrc/transformer/gelu_kernels.cu',
|
|
|
|
'csrc/transformer/dropout_kernels.cu',
|
|
|
|
'csrc/transformer/normalize_kernels.cu',
|
|
|
|
'csrc/transformer/softmax_kernels.cu',
|
|
|
|
'csrc/transformer/general_kernels.cu'
|
|
|
|
],
|
|
|
|
include_dirs=['csrc/includes'],
|
2020-02-01 02:57:11 +03:00
|
|
|
extra_compile_args={
|
2020-05-29 23:15:36 +03:00
|
|
|
'cxx': ['-O3',
|
|
|
|
'-std=c++14',
|
|
|
|
'-g',
|
|
|
|
'-Wno-reorder'],
|
|
|
|
'nvcc': [
|
2020-02-01 02:57:11 +03:00
|
|
|
'-O3',
|
2020-05-29 23:15:36 +03:00
|
|
|
'--use_fast_math',
|
|
|
|
'-gencode',
|
|
|
|
'arch=compute_61,code=compute_61',
|
|
|
|
'-gencode',
|
|
|
|
'arch=compute_70,code=compute_70',
|
|
|
|
'-std=c++14',
|
|
|
|
'-U__CUDA_NO_HALF_OPERATORS__',
|
|
|
|
'-U__CUDA_NO_HALF_CONVERSIONS__',
|
|
|
|
'-U__CUDA_NO_HALF2_OPERATORS__'
|
|
|
|
]
|
|
|
|
}),
|
|
|
|
CUDAExtension(name='deepspeed_stochastic_transformer_cuda',
|
|
|
|
sources=[
|
|
|
|
'csrc/transformer/ds_transformer_cuda.cpp',
|
|
|
|
'csrc/transformer/cublas_wrappers.cu',
|
|
|
|
'csrc/transformer/transform_kernels.cu',
|
|
|
|
'csrc/transformer/gelu_kernels.cu',
|
|
|
|
'csrc/transformer/dropout_kernels.cu',
|
|
|
|
'csrc/transformer/normalize_kernels.cu',
|
|
|
|
'csrc/transformer/softmax_kernels.cu',
|
|
|
|
'csrc/transformer/general_kernels.cu'
|
|
|
|
],
|
|
|
|
include_dirs=['csrc/includes'],
|
|
|
|
extra_compile_args={
|
|
|
|
'cxx': ['-O3',
|
|
|
|
'-std=c++14',
|
|
|
|
'-g',
|
|
|
|
'-Wno-reorder'],
|
|
|
|
'nvcc': [
|
|
|
|
'-O3',
|
|
|
|
'--use_fast_math',
|
|
|
|
'-gencode',
|
|
|
|
'arch=compute_61,code=compute_61',
|
|
|
|
'-gencode',
|
|
|
|
'arch=compute_70,code=compute_70',
|
|
|
|
'-std=c++14',
|
|
|
|
'-U__CUDA_NO_HALF_OPERATORS__',
|
|
|
|
'-U__CUDA_NO_HALF_CONVERSIONS__',
|
|
|
|
'-U__CUDA_NO_HALF2_OPERATORS__',
|
|
|
|
'-D__STOCHASTIC_MODE__'
|
|
|
|
]
|
|
|
|
}),
|
|
|
|
]
|
2020-02-01 02:57:11 +03:00
|
|
|
|
|
|
|
setup(name='deepspeed',
|
2020-05-29 23:15:36 +03:00
|
|
|
version='0.2.0',
|
2020-02-01 02:57:11 +03:00
|
|
|
description='DeepSpeed library',
|
|
|
|
author='DeepSpeed Team',
|
|
|
|
author_email='deepspeed@microsoft.com',
|
|
|
|
url='http://aka.ms/deepspeed',
|
2020-02-03 22:12:45 +03:00
|
|
|
packages=find_packages(exclude=["docker",
|
|
|
|
"third_party",
|
|
|
|
"csrc"]),
|
|
|
|
scripts=['bin/deepspeed',
|
2020-02-09 09:00:24 +03:00
|
|
|
'bin/deepspeed.pt',
|
|
|
|
'bin/ds',
|
|
|
|
'bin/ds_ssh'],
|
2020-02-01 02:57:11 +03:00
|
|
|
classifiers=['Programming Language :: Python :: 3.6'],
|
|
|
|
ext_modules=ext_modules,
|
|
|
|
cmdclass=cmdclass)
|