forked from deepspeedai/DeepSpeed
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathsetup.py
More file actions
47 lines (40 loc) · 1.41 KB
/
setup.py
File metadata and controls
47 lines (40 loc) · 1.41 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
"""
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 deepspeed import __version__ as ds_version
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=ds_version,
description='DeepSpeed library',
author='DeepSpeed Team',
url='http://aka.ms/deepspeed',
packages=find_packages(exclude=["docker",
"third_party",
"csrc"]),
scripts=['bin/deepspeed',
'bin/deepspeed.pt',
'bin/ds',
'bin/ds_ssh'],
classifiers=['Programming Language :: Python :: 3.6'],
ext_modules=ext_modules,
cmdclass=cmdclass)