setup.py 11.4 KB
Newer Older
aiss's avatar
aiss committed
1
2
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
Jeff Rasley's avatar
Jeff Rasley committed
3

aiss's avatar
aiss committed
4
5
# DeepSpeed Team
"""
Jeff Rasley's avatar
Jeff Rasley committed
6
7
DeepSpeed library

aiss's avatar
aiss committed
8
To build wheel on Windows:
aiss's avatar
aiss committed
9
10
11
12
13
1. Install pytorch, such as pytorch 1.12 + cuda 11.6.
2. Install visual cpp build tool.
3. Include cuda toolkit.
4. Launch cmd console with Administrator privilege for creating required symlink folders.

aiss's avatar
aiss committed
14
15

Create a new wheel via the following command:
aiss's avatar
aiss committed
16
build_win.bat
Jeff Rasley's avatar
Jeff Rasley committed
17
18
19
20

The wheel will be located at: dist/*.whl
"""

21
import os
aiss's avatar
aiss committed
22
import sys
23
import subprocess
aiss's avatar
aiss committed
24
from setuptools import setup, find_packages, find_namespace_packages
aiss's avatar
aiss committed
25
from setuptools.command import egg_info
Samyam Rajbhandari's avatar
Samyam Rajbhandari committed
26
import time
27

aiss's avatar
aiss committed
28
torch_available = True
29
30
31
try:
    import torch
except ImportError:
aiss's avatar
aiss committed
32
33
34
35
    torch_available = False
    print('[WARNING] Unable to import torch, pre-compiling ops will be disabled. ' \
        'Please visit https://pytorch.org/ to see how to properly install torch on your system.')

aiss's avatar
aiss committed
36
37
38
from op_builder import get_default_compute_capabilities, OpBuilder
from op_builder.all_ops import ALL_OPS
from op_builder.builder import installed_cuda_version
aiss's avatar
aiss committed
39

aiss's avatar
aiss committed
40
# Fetch rocm state.
aiss's avatar
aiss committed
41
42
43
44
45
46
47
is_rocm_pytorch = OpBuilder.is_rocm_pytorch()
rocm_version = OpBuilder.installed_rocm_version()

RED_START = '\033[31m'
RED_END = '\033[0m'
ERROR = f"{RED_START} [ERROR] {RED_END}"

48

aiss's avatar
aiss committed
49
50
51
def abort(msg):
    print(f"{ERROR} {msg}")
    assert False, msg
52
53
54
55
56
57
58
59


def fetch_requirements(path):
    with open(path, 'r') as fd:
        return [r.strip() for r in fd.readlines()]


install_requires = fetch_requirements('requirements/requirements.txt')
60
extras_require = {
aiss's avatar
aiss committed
61
    '1bit': [],  # add cupy based on cuda/rocm version
aiss's avatar
aiss committed
62
    '1bit_mpi': fetch_requirements('requirements/requirements-1bit-mpi.txt'),
63
64
    'readthedocs': fetch_requirements('requirements/requirements-readthedocs.txt'),
    'dev': fetch_requirements('requirements/requirements-dev.txt'),
aiss's avatar
aiss committed
65
66
    'autotuning': fetch_requirements('requirements/requirements-autotuning.txt'),
    'autotuning_ml': fetch_requirements('requirements/requirements-autotuning-ml.txt'),
aiss's avatar
aiss committed
67
68
69
    'sparse_attn': fetch_requirements('requirements/requirements-sparse_attn.txt'),
    'inf': fetch_requirements('requirements/requirements-inf.txt'),
    'sd': fetch_requirements('requirements/requirements-sd.txt')
70
}
71

aiss's avatar
aiss committed
72
# Add specific cupy version to both onebit extension variants.
aiss's avatar
aiss committed
73
74
75
76
if torch_available and torch.cuda.is_available():
    cupy = None
    if is_rocm_pytorch:
        rocm_major, rocm_minor = rocm_version
aiss's avatar
aiss committed
77
        # XXX cupy support for rocm 5 is not available yet.
aiss's avatar
aiss committed
78
79
80
        if rocm_major <= 4:
            cupy = f"cupy-rocm-{rocm_major}-{rocm_minor}"
    else:
aiss's avatar
aiss committed
81
82
83
84
85
86
        cuda_major_ver, cuda_minor_ver = installed_cuda_version()
        if (cuda_major_ver < 11) or ((cuda_major_ver == 11) and (cuda_minor_ver < 3)):
            cupy = f"cupy-cuda{cuda_major_ver}{cuda_minor_ver}"
        else:
            cupy = f"cupy-cuda{cuda_major_ver}x"

aiss's avatar
aiss committed
87
88
89
    if cupy:
        extras_require['1bit'].append(cupy)
        extras_require['1bit_mpi'].append(cupy)
90

aiss's avatar
aiss committed
91
# Make an [all] extra that installs all needed dependencies.
92
93
94
95
96
all_extras = set()
for extra in extras_require.items():
    for req in extra[1]:
        all_extras.add(req)
extras_require['all'] = list(all_extras)
Jeff Rasley's avatar
Jeff Rasley committed
97
98

cmdclass = {}
99

aiss's avatar
aiss committed
100
# For any pre-installed ops force disable ninja.
aiss's avatar
aiss committed
101
if torch_available:
aiss's avatar
aiss committed
102
    from accelerator import get_accelerator
aiss's avatar
aiss committed
103
    cmdclass['build_ext'] = get_accelerator().build_extension().with_options(use_ninja=False)
Jeff Rasley's avatar
Jeff Rasley committed
104

aiss's avatar
aiss committed
105
106
107
108
109
110
if torch_available:
    TORCH_MAJOR = torch.__version__.split('.')[0]
    TORCH_MINOR = torch.__version__.split('.')[1]
else:
    TORCH_MAJOR = "0"
    TORCH_MINOR = "0"
111

aiss's avatar
aiss committed
112
if torch_available and not torch.cuda.is_available():
aiss's avatar
aiss committed
113
114
115
116
    # Fix to allow docker builds, similar to https://github.com/NVIDIA/apex/issues/486.
    print("[WARNING] Torch did not find cuda available, if cross-compiling 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)")
aiss's avatar
aiss committed
117
          
118
    if os.environ.get("TORCH_CUDA_ARCH_LIST", None) is None:
aiss's avatar
aiss committed
119
        os.environ["TORCH_CUDA_ARCH_LIST"] = get_default_compute_capabilities()
120

121
122
ext_modules = []

aiss's avatar
aiss committed
123
124
125
# Default to pre-install kernels to false so we rely on JIT on Linux, opposite on Windows.
BUILD_OP_PLATFORM = 1 if sys.platform == "win32" else 0
BUILD_OP_DEFAULT = int(os.environ.get('DS_BUILD_OPS', BUILD_OP_PLATFORM))
126
print(f"DS_BUILD_OPS={BUILD_OP_DEFAULT}")
127

aiss's avatar
aiss committed
128
129
130
if BUILD_OP_DEFAULT:
    assert torch_available, "Unable to pre-compile ops without torch installed. Please install torch before attempting to pre-compile ops."

131
132

def command_exists(cmd):
aiss's avatar
aiss committed
133
134
135
136
137
138
    if sys.platform == "win32":
        result = subprocess.Popen(f'{cmd}', stdout=subprocess.PIPE, shell=True)
        return result.wait() == 1
    else:
        result = subprocess.Popen(f'type {cmd}', stdout=subprocess.PIPE, shell=True)
        return result.wait() == 0
139
140


aiss's avatar
aiss committed
141
def op_envvar(op_name):
142
143
    assert hasattr(ALL_OPS[op_name], 'BUILD_VAR'), \
        f"{op_name} is missing BUILD_VAR field"
aiss's avatar
aiss committed
144
145
    return ALL_OPS[op_name].BUILD_VAR

aiss's avatar
aiss committed
146

aiss's avatar
aiss committed
147
148
def op_enabled(op_name):
    env_var = op_envvar(op_name)
149
150
151
    return int(os.environ.get(env_var, BUILD_OP_DEFAULT))


aiss's avatar
aiss committed
152
compatible_ops = dict.fromkeys(ALL_OPS.keys(), False)
153
154
155
install_ops = dict.fromkeys(ALL_OPS.keys(), False)
for op_name, builder in ALL_OPS.items():
    op_compatible = builder.is_compatible()
aiss's avatar
aiss committed
156
    compatible_ops[op_name] = op_compatible
aiss's avatar
aiss committed
157
158
159
    print("op_name: ", op_name)
    print("op_enabled: ", op_enabled(op_name))
    print("op_compatible: ", op_compatible)
aiss's avatar
aiss committed
160
161

    # If op is requested but not available, throw an error.
aiss's avatar
aiss committed
162
163
164
165
166
    if op_enabled(op_name) and not op_compatible:
        env_var = op_envvar(op_name)
        if env_var not in os.environ:
            builder.warning(f"One can disable {op_name} with {env_var}=0")
        abort(f"Unable to pre-compile {op_name}")
aiss's avatar
aiss committed
167

aiss's avatar
aiss committed
168
    # If op is compatible but install is not enabled (JIT mode).
aiss's avatar
aiss committed
169
170
    if is_rocm_pytorch and op_compatible and not op_enabled(op_name):
        builder.hipify_extension()
aiss's avatar
aiss committed
171

aiss's avatar
aiss committed
172
    # If op install enabled, add builder to extensions.
173
    if op_enabled(op_name) and op_compatible:
aiss's avatar
aiss committed
174
        assert torch_available, f"Unable to pre-compile {op_name}, please first install torch"
175
176
177
178
        install_ops[op_name] = op_enabled(op_name)
        ext_modules.append(builder.builder())

print(f'Install Ops={install_ops}')
179

aiss's avatar
aiss committed
180
# Write out version/git info.
181
182
git_hash_cmd = "git rev-parse --short HEAD"
git_branch_cmd = "git rev-parse --abbrev-ref HEAD"
aiss's avatar
aiss committed
183
184
#if command_exists('git') and 'DS_BUILD_STRING' not in os.environ:
if command_exists('git'):
185
186
187
188
189
    try:
        result = subprocess.check_output(git_hash_cmd, shell=True)
        git_hash = result.decode('utf-8').strip()
        result = subprocess.check_output(git_branch_cmd, shell=True)
        git_branch = result.decode('utf-8').strip()
aiss's avatar
aiss committed
190

191
192
193
    except subprocess.CalledProcessError:
        git_hash = "unknown"
        git_branch = "unknown"
194
195
196
else:
    git_hash = "unknown"
    git_branch = "unknown"
197

aiss's avatar
aiss committed
198
199
200
201
202
203
204
205
206
207
208
209
210
211

def create_dir_symlink(src, dest):
    if not os.path.islink(dest):
        if os.path.exists(dest):
            os.remove(dest)
        assert not os.path.exists(dest)
        os.symlink(src, dest)


if sys.platform == "win32":
    # This creates a symbolic links on Windows.
    # It needs Administrator privilege to create symlinks on Windows.
    create_dir_symlink('..\\..\\csrc', '.\\deepspeed\\ops\\csrc')
    create_dir_symlink('..\\..\\op_builder', '.\\deepspeed\\ops\\op_builder')
aiss's avatar
aiss committed
212
    create_dir_symlink('..\\accelerator', '.\\deepspeed\\accelerator')
aiss's avatar
aiss committed
213
214
    egg_info.manifest_maker.template = 'MANIFEST_win.in'

aiss's avatar
aiss committed
215
# Parse the DeepSpeed version string from version.txt.
216
version_str = open('version.txt', 'r').read().strip()
aiss's avatar
aiss committed
217
dcu_version_str = version_str
aiss's avatar
aiss committed
218
219
220
221
222
223
224
225
226
227
228
229
abi_version = subprocess.getoutput("echo '#include <string>' | gcc -x c++ -E -dM - | fgrep _GLIBCXX_USE_CXX11_ABI >tmp.txt")
abi_version = 'abi' + subprocess.getoutput("awk '{print $3}' tmp.txt")
# print(abi_version)

dtk_version_path = subprocess.getoutput("echo $DTKROOT")

with open("{0}/.info/rocm_version".format(dtk_version_path), 'r') as fd:
    dtk_version_str = fd.read().split('.')[:2]
    dtk_version_str ='dtk' +  "".join(dtk_version_str)

torch_version = "torch"+ subprocess.getoutput("python3 -c \"import torch;print(torch.__version__)\"")
# print(torch_version)
230

aiss's avatar
aiss committed
231
232
dcu_version_str += f'+git{git_hash}' + '.' + abi_version + '.'+ dtk_version_str + '.' + torch_version
print("dcu_version_str=", dcu_version_str)
233
torch_version = ".".join([TORCH_MAJOR, TORCH_MINOR])
aiss's avatar
aiss committed
234
bf16_support = False
aiss's avatar
aiss committed
235
# Set cuda_version to 0.0 if cpu-only.
236
cuda_version = "0.0"
aiss's avatar
aiss committed
237
nccl_version = "0.0"
aiss's avatar
aiss committed
238
# Set hip_version to 0.0 if cpu-only.
aiss's avatar
aiss committed
239
240
hip_version = "0.0"
if torch_available and torch.version.cuda is not None:
241
    cuda_version = ".".join(torch.version.cuda.split('.')[:2])
aiss's avatar
aiss committed
242
243
    if sys.platform != "win32":
        if isinstance(torch.cuda.nccl.version(), int):
aiss's avatar
aiss committed
244
            # This will break if minor version > 9.
aiss's avatar
aiss committed
245
246
247
248
249
            nccl_version = ".".join(str(torch.cuda.nccl.version())[:2])
        else:
            nccl_version = ".".join(map(str, torch.cuda.nccl.version()[:2]))
    if hasattr(torch.cuda, 'is_bf16_supported') and torch.cuda.is_available():
        bf16_support = torch.cuda.is_bf16_supported()
aiss's avatar
aiss committed
250
251
252
253
if torch_available and hasattr(torch.version, 'hip') and torch.version.hip is not None:
    hip_version = ".".join(torch.version.hip.split('.')[:2])
torch_info = {
    "version": torch_version,
aiss's avatar
aiss committed
254
    "bf16_support": bf16_support,
aiss's avatar
aiss committed
255
    "cuda_version": cuda_version,
aiss's avatar
aiss committed
256
    "nccl_version": nccl_version,
aiss's avatar
aiss committed
257
258
    "hip_version": hip_version
}
259

aiss's avatar
aiss committed
260
print(f"version={version_str}, dcu_version={dcu_version_str}, git_hash={git_hash}, git_branch={git_branch}")
261
with open('deepspeed/git_version_info_installed.py', 'w') as fd:
262
    fd.write(f"version='{version_str}'\n")
aiss's avatar
aiss committed
263
    fd.write(f"dcu_version='{dcu_version_str}'\n")
264
265
    fd.write(f"git_hash='{git_hash}'\n")
    fd.write(f"git_branch='{git_branch}'\n")
Jeff Rasley's avatar
Jeff Rasley committed
266
    fd.write(f"installed_ops={install_ops}\n")
267
268
    fd.write(f"compatible_ops={compatible_ops}\n")
    fd.write(f"torch_info={torch_info}\n")
269
270

print(f'install_requires={install_requires}')
271
272
print(f'compatible_ops={compatible_ops}')
print(f'ext_modules={ext_modules}')
Jeff Rasley's avatar
Jeff Rasley committed
273

274
275
276
277
278
# Parse README.md to make long_description for PyPI page.
thisdir = os.path.abspath(os.path.dirname(__file__))
with open(os.path.join(thisdir, 'README.md'), encoding='utf-8') as fin:
    readme_text = fin.read()

Samyam Rajbhandari's avatar
Samyam Rajbhandari committed
279
280
start_time = time.time()

Jeff Rasley's avatar
Jeff Rasley committed
281
setup(name='deepspeed',
aiss's avatar
aiss committed
282
      version=dcu_version_str,
Jeff Rasley's avatar
Jeff Rasley committed
283
      description='DeepSpeed library',
284
285
      long_description=readme_text,
      long_description_content_type='text/markdown',
Jeff Rasley's avatar
Jeff Rasley committed
286
      author='DeepSpeed Team',
aiss's avatar
aiss committed
287
      author_email='deepspeed-info@microsoft.com',
288
      url='http://deepspeed.ai',
aiss's avatar
aiss committed
289
290
291
292
      project_urls={
          'Documentation': 'https://deepspeed.readthedocs.io',
          'Source': 'https://github.com/microsoft/DeepSpeed',
      },
293
      install_requires=install_requires,
294
      extras_require=extras_require,
aiss's avatar
aiss committed
295
      #packages=find_packages(include=['deepspeed', 'deepspeed.*']),
aiss's avatar
aiss committed
296
297
      packages=find_namespace_packages(include=['deepspeed',
                                      'deepspeed.*']),
298
299
      include_package_data=True,
      scripts=[
aiss's avatar
aiss committed
300
          'bin/deepspeed', 'bin/deepspeed.pt', 'bin/ds', 'bin/ds_ssh', 'bin/ds_report', 'bin/ds_bench', 'bin/dsr',
Jeff Rasley's avatar
Jeff Rasley committed
301
          'bin/ds_elastic'
302
      ],
303
      classifiers=[
aiss's avatar
aiss committed
304
305
          'Programming Language :: Python :: 3.6', 'Programming Language :: Python :: 3.7',
          'Programming Language :: Python :: 3.8', 'Programming Language :: Python :: 3.9',
aiss's avatar
aiss committed
306
          'Programming Language :: Python :: 3.10'
307
308
      ],
      license='MIT',
Jeff Rasley's avatar
Jeff Rasley committed
309
310
      ext_modules=ext_modules,
      cmdclass=cmdclass)
Samyam Rajbhandari's avatar
Samyam Rajbhandari committed
311
312
313

end_time = time.time()
print(f'deepspeed build time = {end_time - start_time} secs')