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
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
import torch
|
|
|
|
|
from setuptools import setup, find_packages
|
|
|
|
|
from torch.utils.cpp_extension import CUDAExtension, BuildExtension
|
|
|
|
|
|
|
|
|
|
cmdclass = {}
|
|
|
|
|
ext_modules = []
|
|
|
|
|
cmdclass['build_ext'] = BuildExtension
|
|
|
|
|
|
|
|
|
|
ext_modules.append(
|
|
|
|
|
CUDAExtension(name='fused_lamb_cuda',
|
|
|
|
|
sources=['csrc/fused_lamb_cuda.cpp',
|
|
|
|
|
'csrc/fused_lamb_cuda_kernel.cu'],
|
|
|
|
|
extra_compile_args={
|
|
|
|
|
'cxx': [
|
|
|
|
|
'-O3',
|
|
|
|
|
],
|
|
|
|
|
'nvcc': ['-O3',
|
|
|
|
|
'--use_fast_math']
|
|
|
|
|
}))
|
|
|
|
|
|
|
|
|
|
setup(name='deepspeed',
|
|
|
|
|
version='0.1',
|
|
|
|
|
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',
|
|
|
|
|
'bin/ds'],
|
2020-02-01 02:57:11 +03:00
|
|
|
classifiers=['Programming Language :: Python :: 3.6'],
|
|
|
|
|
ext_modules=ext_modules,
|
|
|
|
|
cmdclass=cmdclass)
|
