DeepSpeed

1
# Copyright (c) Microsoft Corporation.
2
# SPDX-License-Identifier: Apache-2.0
3

4
# DeepSpeed Team
5
"""
6
DeepSpeed library
7

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

14

15
Create a new wheel via the following command:
16
build_win.bat
17

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

21
import pathlib
22
import os
23
import shutil
24
import sys
25
import subprocess
26
from setuptools import setup, find_packages
27
from setuptools.command import egg_info
28
import time
29
import typing
30
import shlex
31

32
torch_available = True
33
try:
34
    import torch
35
except ImportError:
36
    torch_available = False
37
    print('[WARNING] Unable to import torch, pre-compiling ops will be disabled. ' \
38
        'Please visit https://pytorch.org/ to see how to properly install torch on your system.')
39

40
from op_builder import get_default_compute_capabilities, OpBuilder
41
from op_builder.all_ops import ALL_OPS, accelerator_name
42
from op_builder.builder import installed_cuda_version
43

44
from accelerator import get_accelerator
45

46
# Fetch rocm state.
47
is_rocm_pytorch = OpBuilder.is_rocm_pytorch()
48
rocm_version = OpBuilder.installed_rocm_version()
49

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

54

55
def abort(msg):
56
    print(f"{ERROR} {msg}")
57
    assert False, msg
58

59

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

64

65
def is_env_set(key):
66
    """
67
    Checks if an environment variable is set and not "".
68
    """
69
    return bool(os.environ.get(key, None))
70

71

72
def get_env_if_set(key, default: typing.Any = ""):
73
    """
74
    Returns an environment variable if it is set and not "",
75
    otherwise returns a default value. In contrast, the fallback
76
    parameter of os.environ.get() is skipped if the variable is set to "".
77
    """
78
    return os.environ.get(key, None) or default
79

80

81
install_requires = fetch_requirements('requirements/requirements.txt')
82
extras_require = {
83
    '1bit': [],  # add cupy based on cuda/rocm version
84
    '1bit_mpi': fetch_requirements('requirements/requirements-1bit-mpi.txt'),
85
    'readthedocs': fetch_requirements('requirements/requirements-readthedocs.txt'),
86
    'dev': fetch_requirements('requirements/requirements-dev.txt'),
87
    'autotuning': fetch_requirements('requirements/requirements-autotuning.txt'),
88
    'autotuning_ml': fetch_requirements('requirements/requirements-autotuning-ml.txt'),
89
    'sparse_attn': fetch_requirements('requirements/requirements-sparse_attn.txt'),
90
    'sparse': fetch_requirements('requirements/requirements-sparse_pruning.txt'),
91
    'inf': fetch_requirements('requirements/requirements-inf.txt'),
92
    'sd': fetch_requirements('requirements/requirements-sd.txt'),
93
    'triton': fetch_requirements('requirements/requirements-triton.txt'),
94
}
95

96
# Only install pynvml on nvidia gpus.
97
if torch_available and get_accelerator().device_name() == 'cuda' and not is_rocm_pytorch:
98
    install_requires.append('nvidia-ml-py')
99

100
# Add specific cupy version to both onebit extension variants.
101
if torch_available and get_accelerator().device_name() == 'cuda':
102
    cupy = None
103
    if is_rocm_pytorch:
104
        rocm_major, rocm_minor = rocm_version
105
        # XXX cupy support for rocm 5 is not available yet.
106
        if rocm_major <= 4:
107
            cupy = f"cupy-rocm-{rocm_major}-{rocm_minor}"
108
    else:
109
        cuda_major_ver, cuda_minor_ver = installed_cuda_version()
110
        if (cuda_major_ver < 11) or ((cuda_major_ver == 11) and (cuda_minor_ver < 3)):
111
            cupy = f"cupy-cuda{cuda_major_ver}{cuda_minor_ver}"
112
        else:
113
            cupy = f"cupy-cuda{cuda_major_ver}x"
114

115
    if cupy:
116
        extras_require['1bit'].append(cupy)
117
        extras_require['1bit_mpi'].append(cupy)
118

119
# Make an [all] extra that installs all needed dependencies.
120
all_extras = set()
121
for extra in extras_require.items():
122
    for req in extra[1]:
123
        all_extras.add(req)
124
extras_require['all'] = list(all_extras)
125

126
cmdclass = {}
127

128
# For any pre-installed ops force disable ninja.
129
if torch_available:
130
    use_ninja = is_env_set("DS_ENABLE_NINJA")
131
    cmdclass['build_ext'] = get_accelerator().build_extension().with_options(use_ninja=use_ninja)
132

133
if torch_available:
134
    TORCH_MAJOR = torch.__version__.split('.')[0]
135
    TORCH_MINOR = torch.__version__.split('.')[1]
136
else:
137
    TORCH_MAJOR = "0"
138
    TORCH_MINOR = "0"
139

140
if torch_available and not get_accelerator().device_name() == 'cuda':
141
    # Fix to allow docker builds, similar to https://github.com/NVIDIA/apex/issues/486.
142
    print("[WARNING] Torch did not find cuda available, if cross-compiling or running with cpu only "
143
          "you can ignore this message. Adding compute capability for Pascal, Volta, and Turing "
144
          "(compute capabilities 6.0, 6.1, 6.2)")
145
    if not is_env_set("TORCH_CUDA_ARCH_LIST"):
146
        os.environ["TORCH_CUDA_ARCH_LIST"] = get_default_compute_capabilities()
147

148
ext_modules = []
149

150
# Default to pre-install kernels to false so we rely on JIT on Linux, opposite on Windows.
151
BUILD_OP_PLATFORM = 1 if sys.platform == "win32" else 0
152
BUILD_OP_DEFAULT = int(get_env_if_set('DS_BUILD_OPS', BUILD_OP_PLATFORM))
153
print(f"DS_BUILD_OPS={BUILD_OP_DEFAULT}")
154

155
if BUILD_OP_DEFAULT:
156
    assert torch_available, "Unable to pre-compile ops without torch installed. Please install torch before attempting to pre-compile ops."
157

158

159
def command_exists(cmd):
160
    if sys.platform == "win32":
161
        safe_cmd = shlex.split(f'{cmd}')
162
        result = subprocess.Popen(safe_cmd, stdout=subprocess.PIPE)
163
        return result.wait() == 1
164
    else:
165
        safe_cmd = shlex.split(f"bash -c type {cmd}")
166
        result = subprocess.Popen(safe_cmd, stdout=subprocess.PIPE)
167
        return result.wait() == 0
168

169

170
def op_envvar(op_name):
171
    assert hasattr(ALL_OPS[op_name], 'BUILD_VAR'), \
172
        f"{op_name} is missing BUILD_VAR field"
173
    return ALL_OPS[op_name].BUILD_VAR
174

175

176
def op_enabled(op_name):
177
    env_var = op_envvar(op_name)
178
    return int(get_env_if_set(env_var, BUILD_OP_DEFAULT))
179

180

181
install_ops = dict.fromkeys(ALL_OPS.keys(), False)
182
for op_name, builder in ALL_OPS.items():
183
    op_compatible = builder.is_compatible()
184

185
    # If op is requested but not available, throw an error.
186
    if op_enabled(op_name) and not op_compatible:
187
        env_var = op_envvar(op_name)
188
        if not is_env_set(env_var):
189
            builder.warning(f"Skip pre-compile of incompatible {op_name}; One can disable {op_name} with {env_var}=0")
190
        continue
191

192
    # If op is compatible but install is not enabled (JIT mode).
193
    if is_rocm_pytorch and op_compatible and not op_enabled(op_name):
194
        builder.hipify_extension()
195

196
    # If op install enabled, add builder to extensions.
197
    if op_enabled(op_name) and op_compatible:
198
        assert torch_available, f"Unable to pre-compile {op_name}, please first install torch"
199
        install_ops[op_name] = op_enabled(op_name)
200
        ext_modules.append(builder.builder())
201

202
print(f'Install Ops={install_ops}')
203

204
# Write out version/git info.
205
git_hash_cmd = shlex.split("bash -c git rev-parse --short HEAD")
206
git_branch_cmd = shlex.split("bash -c git rev-parse --abbrev-ref HEAD")
207
if command_exists('git') and not is_env_set('DS_BUILD_STRING'):
208
    try:
209
        result = subprocess.check_output(git_hash_cmd)
210
        git_hash = result.decode('utf-8').strip()
211
        result = subprocess.check_output(git_branch_cmd)
212
        git_branch = result.decode('utf-8').strip()
213
    except subprocess.CalledProcessError:
214
        git_hash = "unknown"
215
        git_branch = "unknown"
216
else:
217
    git_hash = "unknown"
218
    git_branch = "unknown"
219

220
if sys.platform == "win32":
221
    shutil.rmtree('.\\deepspeed\\ops\\csrc', ignore_errors=True)
222
    pathlib.Path('.\\deepspeed\\ops\\csrc').unlink(missing_ok=True)
223
    shutil.copytree('.\\csrc', '.\\deepspeed\\ops\\csrc', dirs_exist_ok=True)
224
    shutil.rmtree('.\\deepspeed\\ops\\op_builder', ignore_errors=True)
225
    pathlib.Path('.\\deepspeed\\ops\\op_builder').unlink(missing_ok=True)
226
    shutil.copytree('.\\op_builder', '.\\deepspeed\\ops\\op_builder', dirs_exist_ok=True)
227
    shutil.rmtree('.\\deepspeed\\accelerator', ignore_errors=True)
228
    pathlib.Path('.\\deepspeed\\accelerator').unlink(missing_ok=True)
229
    shutil.copytree('.\\accelerator', '.\\deepspeed\\accelerator', dirs_exist_ok=True)
230
    egg_info.manifest_maker.template = 'MANIFEST_win.in'
231

232
# Parse the DeepSpeed version string from version.txt.
233
version_str = open('version.txt', 'r').read().strip()
234

235
# Build specifiers like .devX can be added at install time. Otherwise, add the git hash.
236
# Example: DS_BUILD_STRING=".dev20201022" python setup.py sdist bdist_wheel.
237

238
# Building wheel for distribution, update version file.
239
if is_env_set('DS_BUILD_STRING'):
240
    # Build string env specified, probably building for distribution.
241
    with open('build.txt', 'w') as fd:
242
        fd.write(os.environ['DS_BUILD_STRING'])
243
    version_str += os.environ['DS_BUILD_STRING']
244
elif os.path.isfile('build.txt'):
245
    # build.txt exists, probably installing from distribution.
246
    with open('build.txt', 'r') as fd:
247
        version_str += fd.read().strip()
248
else:
249
    # None of the above, probably installing from source.
250
    version_str += f'+{git_hash}'
251

252
torch_version = ".".join([TORCH_MAJOR, TORCH_MINOR])
253
bf16_support = False
254
# Set cuda_version to 0.0 if cpu-only.
255
cuda_version = "0.0"
256
nccl_version = "0.0"
257
# Set hip_version to 0.0 if cpu-only.
258
hip_version = "0.0"
259
if torch_available and torch.version.cuda is not None:
260
    cuda_version = ".".join(torch.version.cuda.split('.')[:2])
261
    if sys.platform != "win32":
262
        if isinstance(torch.cuda.nccl.version(), int):
263
            # This will break if minor version > 9.
264
            nccl_version = ".".join(str(torch.cuda.nccl.version())[:2])
265
        else:
266
            nccl_version = ".".join(map(str, torch.cuda.nccl.version()[:2]))
267
    if hasattr(torch.cuda, 'is_bf16_supported') and torch.cuda.is_available():
268
        bf16_support = torch.cuda.is_bf16_supported()
269
if torch_available and hasattr(torch.version, 'hip') and torch.version.hip is not None:
270
    hip_version = ".".join(torch.version.hip.split('.')[:2])
271
torch_info = {
272
    "version": torch_version,
273
    "bf16_support": bf16_support,
274
    "cuda_version": cuda_version,
275
    "nccl_version": nccl_version,
276
    "hip_version": hip_version
277
}
278

279
print(f"version={version_str}, git_hash={git_hash}, git_branch={git_branch}")
280
with open('deepspeed/git_version_info_installed.py', 'w') as fd:
281
    fd.write(f"version='{version_str}'\n")
282
    fd.write(f"git_hash='{git_hash}'\n")
283
    fd.write(f"git_branch='{git_branch}'\n")
284
    fd.write(f"installed_ops={install_ops}\n")
285
    fd.write(f"accelerator_name='{accelerator_name}'\n")
286
    fd.write(f"torch_info={torch_info}\n")
287

288
print(f'install_requires={install_requires}')
289
print(f'ext_modules={ext_modules}')
290

291
# Parse README.md to make long_description for PyPI page.
292
thisdir = os.path.abspath(os.path.dirname(__file__))
293
with open(os.path.join(thisdir, 'README.md'), encoding='utf-8') as fin:
294
    readme_text = fin.read()
295

296
if sys.platform == "win32":
297
    scripts = ['bin/deepspeed.bat', 'bin/ds', 'bin/ds_report.bat', 'bin/ds_report']
298
else:
299
    scripts = [
300
        'bin/deepspeed', 'bin/deepspeed.pt', 'bin/ds', 'bin/ds_ssh', 'bin/ds_report', 'bin/ds_bench', 'bin/dsr',
301
        'bin/ds_elastic', 'bin/ds_nvme_tune', 'bin/ds_io'
302
    ]
303

304
start_time = time.time()
305

306
setup(name='deepspeed',
307
      version=version_str,
308
      description='DeepSpeed library',
309
      long_description=readme_text,
310
      long_description_content_type='text/markdown',
311
      author='DeepSpeed Team',
312
      author_email='deepspeed-info@microsoft.com',
313
      url='http://deepspeed.ai',
314
      project_urls={
315
          'Documentation': 'https://deepspeed.readthedocs.io',
316
          'Source': 'https://github.com/microsoft/DeepSpeed',
317
      },
318
      install_requires=install_requires,
319
      extras_require=extras_require,
320
      packages=find_packages(include=['deepspeed', 'deepspeed.*']),
321
      include_package_data=True,
322
      scripts=scripts,
323
      classifiers=[
324
          'Programming Language :: Python :: 3.6', 'Programming Language :: Python :: 3.7',
325
          'Programming Language :: Python :: 3.8', 'Programming Language :: Python :: 3.9',
326
          'Programming Language :: Python :: 3.10'
327
      ],
328
      license='Apache Software License 2.0',
329
      ext_modules=ext_modules,
330
      cmdclass=cmdclass)
331

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