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 os
22
import sys
23
import subprocess
24
from setuptools import setup, find_packages
25
from setuptools.command import egg_info
26
import time
27
import typing
28

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

37
from op_builder import get_default_compute_capabilities, OpBuilder
38
from op_builder.all_ops import ALL_OPS
39
from op_builder.builder import installed_cuda_version
40

41
# Fetch rocm state.
42
is_rocm_pytorch = OpBuilder.is_rocm_pytorch()
43
rocm_version = OpBuilder.installed_rocm_version()
44

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

49

50
def abort(msg):
51
    print(f"{ERROR} {msg}")
52
    assert False, msg
53

54

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

59

60
def is_env_set(key):
61
    """
62
    Checks if an environment variable is set and not "".
63
    """
64
    return bool(os.environ.get(key, None))
65

66

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

75

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

91
# Add specific cupy version to both onebit extension variants.
92
if torch_available and torch.cuda.is_available():
93
    cupy = None
94
    if is_rocm_pytorch:
95
        rocm_major, rocm_minor = rocm_version
96
        # XXX cupy support for rocm 5 is not available yet.
97
        if rocm_major <= 4:
98
            cupy = f"cupy-rocm-{rocm_major}-{rocm_minor}"
99
    else:
100
        cuda_major_ver, cuda_minor_ver = installed_cuda_version()
101
        if (cuda_major_ver < 11) or ((cuda_major_ver == 11) and (cuda_minor_ver < 3)):
102
            cupy = f"cupy-cuda{cuda_major_ver}{cuda_minor_ver}"
103
        else:
104
            cupy = f"cupy-cuda{cuda_major_ver}x"
105

106
    if cupy:
107
        extras_require['1bit'].append(cupy)
108
        extras_require['1bit_mpi'].append(cupy)
109

110
# Make an [all] extra that installs all needed dependencies.
111
all_extras = set()
112
for extra in extras_require.items():
113
    for req in extra[1]:
114
        all_extras.add(req)
115
extras_require['all'] = list(all_extras)
116

117
cmdclass = {}
118

119
# For any pre-installed ops force disable ninja.
120
if torch_available:
121
    from accelerator import get_accelerator
122
    use_ninja = is_env_set("DS_ENABLE_NINJA")
123
    cmdclass['build_ext'] = get_accelerator().build_extension().with_options(use_ninja=use_ninja)
124

125
if torch_available:
126
    TORCH_MAJOR = torch.__version__.split('.')[0]
127
    TORCH_MINOR = torch.__version__.split('.')[1]
128
else:
129
    TORCH_MAJOR = "0"
130
    TORCH_MINOR = "0"
131

132
if torch_available and not torch.cuda.is_available():
133
    # Fix to allow docker builds, similar to https://github.com/NVIDIA/apex/issues/486.
134
    print("[WARNING] Torch did not find cuda available, if cross-compiling or running with cpu only "
135
          "you can ignore this message. Adding compute capability for Pascal, Volta, and Turing "
136
          "(compute capabilities 6.0, 6.1, 6.2)")
137
    if not is_env_set("TORCH_CUDA_ARCH_LIST"):
138
        os.environ["TORCH_CUDA_ARCH_LIST"] = get_default_compute_capabilities()
139

140
ext_modules = []
141

142
# Default to pre-install kernels to false so we rely on JIT on Linux, opposite on Windows.
143
BUILD_OP_PLATFORM = 1 if sys.platform == "win32" else 0
144
BUILD_OP_DEFAULT = int(get_env_if_set('DS_BUILD_OPS', BUILD_OP_PLATFORM))
145
print(f"DS_BUILD_OPS={BUILD_OP_DEFAULT}")
146

147
if BUILD_OP_DEFAULT:
148
    assert torch_available, "Unable to pre-compile ops without torch installed. Please install torch before attempting to pre-compile ops."
149

150

151
def command_exists(cmd):
152
    if sys.platform == "win32":
153
        result = subprocess.Popen(f'{cmd}', stdout=subprocess.PIPE, shell=True)
154
        return result.wait() == 1
155
    else:
156
        result = subprocess.Popen(f'type {cmd}', stdout=subprocess.PIPE, shell=True)
157
        return result.wait() == 0
158

159

160
def op_envvar(op_name):
161
    assert hasattr(ALL_OPS[op_name], 'BUILD_VAR'), \
162
        f"{op_name} is missing BUILD_VAR field"
163
    return ALL_OPS[op_name].BUILD_VAR
164

165

166
def op_enabled(op_name):
167
    env_var = op_envvar(op_name)
168
    return int(get_env_if_set(env_var, BUILD_OP_DEFAULT))
169

170

171
compatible_ops = dict.fromkeys(ALL_OPS.keys(), False)
172
install_ops = dict.fromkeys(ALL_OPS.keys(), False)
173
for op_name, builder in ALL_OPS.items():
174
    op_compatible = builder.is_compatible()
175
    compatible_ops[op_name] = op_compatible
176
    compatible_ops["deepspeed_not_implemented"] = False
177

178
    # If op is requested but not available, throw an error.
179
    if op_enabled(op_name) and not op_compatible:
180
        env_var = op_envvar(op_name)
181
        if not is_env_set(env_var):
182
            builder.warning(f"One can disable {op_name} with {env_var}=0")
183
        abort(f"Unable to pre-compile {op_name}")
184

185
    # If op is compatible but install is not enabled (JIT mode).
186
    if is_rocm_pytorch and op_compatible and not op_enabled(op_name):
187
        builder.hipify_extension()
188

189
    # If op install enabled, add builder to extensions.
190
    if op_enabled(op_name) and op_compatible:
191
        assert torch_available, f"Unable to pre-compile {op_name}, please first install torch"
192
        install_ops[op_name] = op_enabled(op_name)
193
        ext_modules.append(builder.builder())
194

195
print(f'Install Ops={install_ops}')
196

197
# Write out version/git info.
198
git_hash_cmd = "git rev-parse --short HEAD"
199
git_branch_cmd = "git rev-parse --abbrev-ref HEAD"
200
if command_exists('git') and not is_env_set('DS_BUILD_STRING'):
201
    try:
202
        result = subprocess.check_output(git_hash_cmd, shell=True)
203
        git_hash = result.decode('utf-8').strip()
204
        result = subprocess.check_output(git_branch_cmd, shell=True)
205
        git_branch = result.decode('utf-8').strip()
206
    except subprocess.CalledProcessError:
207
        git_hash = "unknown"
208
        git_branch = "unknown"
209
else:
210
    git_hash = "unknown"
211
    git_branch = "unknown"
212

213

214
def create_dir_symlink(src, dest):
215
    if not os.path.islink(dest):
216
        if os.path.exists(dest):
217
            os.remove(dest)
218
        assert not os.path.exists(dest)
219
        os.symlink(src, dest)
220

221

222
if sys.platform == "win32":
223
    # This creates a symbolic links on Windows.
224
    # It needs Administrator privilege to create symlinks on Windows.
225
    create_dir_symlink('..\\..\\csrc', '.\\deepspeed\\ops\\csrc')
226
    create_dir_symlink('..\\..\\op_builder', '.\\deepspeed\\ops\\op_builder')
227
    create_dir_symlink('..\\accelerator', '.\\deepspeed\\accelerator')
228
    egg_info.manifest_maker.template = 'MANIFEST_win.in'
229

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

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

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

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

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

286
print(f'install_requires={install_requires}')
287
print(f'compatible_ops={compatible_ops}')
288
print(f'ext_modules={ext_modules}')
289

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

295
start_time = time.time()
296

297
setup(name='deepspeed',
298
      version=version_str,
299
      description='DeepSpeed library',
300
      long_description=readme_text,
301
      long_description_content_type='text/markdown',
302
      author='DeepSpeed Team',
303
      author_email='deepspeed-info@microsoft.com',
304
      url='http://deepspeed.ai',
305
      project_urls={
306
          'Documentation': 'https://deepspeed.readthedocs.io',
307
          'Source': 'https://github.com/microsoft/DeepSpeed',
308
      },
309
      install_requires=install_requires,
310
      extras_require=extras_require,
311
      packages=find_packages(include=['deepspeed', 'deepspeed.*']),
312
      include_package_data=True,
313
      scripts=[
314
          'bin/deepspeed', 'bin/deepspeed.pt', 'bin/ds', 'bin/ds_ssh', 'bin/ds_report', 'bin/ds_bench', 'bin/dsr',
315
          'bin/ds_elastic'
316
      ],
317
      classifiers=[
318
          'Programming Language :: Python :: 3.6', 'Programming Language :: Python :: 3.7',
319
          'Programming Language :: Python :: 3.8', 'Programming Language :: Python :: 3.9',
320
          'Programming Language :: Python :: 3.10'
321
      ],
322
      license='Apache Software License 2.0',
323
      ext_modules=ext_modules,
324
      cmdclass=cmdclass)
325

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