github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
Browse Source

[cli] updated installation check cli for aot/jit build (#2395)

pull/2397/head
Frank Lee 2 years ago committed by GitHub
parent
2add870138
commit
ce08661eb1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 150 additions and 45 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 195
      colossalai/cli/check/check_installation.py

195
colossalai/cli/check/check_installation.py
Unescape View File

@ -7,30 +7,100 @@ from torch.utils.cpp_extension import CUDA_HOME
import colossalai
def to_click_output(val):
# installation check output to understandable symbols for readability
VAL_TO_SYMBOL = {True: u'\u2713', False: 'x', None: 'N/A'}
if val in VAL_TO_SYMBOL:
return VAL_TO_SYMBOL[val]
else:
return val
def check_installation():
cuda_ext_installed = _check_cuda_extension_installed()
cuda_version, torch_version, torch_cuda_version = _check_cuda_torch()
"""
This function will check the installation of colossalai, specifically, the version compatibility of
colossalai, pytorch and cuda.
Example:
```text
```
Returns: A table of installation information.
"""
found_aot_cuda_ext = _check_aot_built_cuda_extension_installed()
cuda_version = _check_cuda_version()
torch_version, torch_cuda_version = _check_torch_version()
colossalai_verison, torch_version_required, cuda_version_required = _parse_colossalai_version()
cuda_compatibility = _get_compatibility_string([cuda_version, torch_cuda_version, cuda_version_required])
torch_compatibility = _get_compatibility_string([torch_version, torch_version_required])
click.echo(f'#### Installation Report ####\n')
click.echo(f"Colossal-AI version: {colossalai_verison}")
click.echo(f'----------------------------')
click.echo(f"PyTorch Version: {torch_version}")
click.echo(f"PyTorch Version required by Colossal-AI: {torch_version_required}")
click.echo(f'PyTorch version match: {torch_compatibility}')
click.echo(f'----------------------------')
click.echo(f"System CUDA Version: {cuda_version}")
click.echo(f"CUDA Version required by PyTorch: {torch_cuda_version}")
click.echo(f"CUDA Version required by Colossal-AI: {cuda_version_required}")
click.echo(f"CUDA Version Match: {cuda_compatibility}")
click.echo(f'----------------------------')
click.echo(f"CUDA Extension: {cuda_ext_installed}")
# if cuda_version is None, that means either
# CUDA_HOME is not found, thus cannot compare the version compatibility
if not cuda_version:
sys_torch_cuda_compatibility = None
else:
sys_torch_cuda_compatibility = _is_compatible([cuda_version, torch_cuda_version])
# if cuda_version or cuda_version_required is None, that means either
# CUDA_HOME is not found or AOT compilation is not enabled
# thus, there is no need to compare the version compatibility at all
if not cuda_version or not cuda_version_required:
sys_colossalai_cuda_compatibility = None
else:
sys_colossalai_cuda_compatibility = _is_compatible([cuda_version, cuda_version_required])
def _get_compatibility_string(versions):
# if torch_version_required is None, that means AOT compilation is not enabled
# thus there is no need to compare the versions
if torch_version_required is None:
torch_compatibility = None
else:
torch_compatibility = _is_compatible([torch_version, torch_version_required])
click.echo(f'#### Installation Report ####')
click.echo(f'\n------------ Environment ------------')
click.echo(f"Colossal-AI version: {to_click_output(colossalai_verison)}")
click.echo(f"PyTorch version: {to_click_output(torch_version)}")
click.echo(f"CUDA version: {to_click_output(cuda_version)}")
click.echo(f"CUDA version required by PyTorch: {to_click_output(torch_cuda_version)}")
click.echo("")
click.echo(f"Note:")
click.echo(f"1. The table above checks the versions of the libraries/tools in the current environment")
click.echo(f"2. If the CUDA version is N/A, you can set the CUDA_HOME environment variable to locate it")
click.echo(f'\n------------ CUDA Extensions AOT Compilation ------------')
click.echo(f"Found AOT CUDA Extension: {to_click_output(found_aot_cuda_ext)}")
click.echo(f"PyTorch version used for AOT compilation: {to_click_output(torch_version_required)}")
click.echo(f"CUDA version used for AOT compilation: {to_click_output(cuda_version_required)}")
click.echo("")
click.echo(f"Note:")
click.echo(
f"1. AOT (ahead-of-time) compilation of the CUDA kernels occurs during installation when the environment varialbe CUDA_EXT=1 is set"
)
click.echo(f"2. If AOT compilation is not enabled, stay calm as the CUDA kernels can still be built during runtime")
click.echo(f"\n------------ Compatibility ------------")
click.echo(f'PyTorch version match: {to_click_output(torch_compatibility)}')
click.echo(f"System and PyTorch CUDA version match: {to_click_output(sys_torch_cuda_compatibility)}")
click.echo(f"System and Colossal-AI CUDA version match: {to_click_output(sys_colossalai_cuda_compatibility)}")
click.echo(f"")
click.echo(f"Note:")
click.echo(f"1. The table above checks the version compatibility of the libraries/tools in the current environment")
click.echo(
f" - PyTorch version mistach: whether the PyTorch version in the current environment is compatible with the PyTorch version used for AOT compilation"
)
click.echo(
f" - System and PyTorch CUDA version match: whether the CUDA version in the current environment is compatible with the CUDA version required by PyTorch"
)
click.echo(
f" - System and Colossal-AI CUDA version match: whether the CUDA version in the current environment is compatible with the CUDA version used for AOT compilation"
)
def _is_compatible(versions):
"""
Compare the list of versions and return whether they are compatible.
"""
if None in versions:
return False
# split version into [major, minor, patch]
versions = [version.split('.') for version in versions]
@ -44,37 +114,81 @@ def _get_compatibility_string(versions):
equal = len(set(version_values)) == 1
if idx in [0, 1] and not equal:
# if the major/minor versions do not match
# return a cross
return 'x'
return False
elif idx == 1:
# if the minor versions match
# return a tick
return u'\u2713'
return True
else:
continue
def _parse_colossalai_version():
"""
Get the Colossal-AI version information.
Returns:
colossalai_version: Colossal-AI version.
torch_version_for_aot_build: PyTorch version used for AOT compilation of CUDA kernels.
cuda_version_for_aot_build: CUDA version used for AOT compilation of CUDA kernels.
"""
# colossalai version can be in two formats
# 1. X.X.X+torchX.XXcuXX.X (when colossalai is installed with CUDA extensions)
# 2. X.X.X (when colossalai is not installed with CUDA extensions)
# where X represents an integer.
colossalai_verison = colossalai.__version__.split('+')[0]
torch_version_required = colossalai.__version__.split('torch')[1].split('cu')[0]
cuda_version_required = colossalai.__version__.split('cu')[1]
return colossalai_verison, torch_version_required, cuda_version_required
def _check_cuda_extension_installed():
try:
torch_version_for_aot_build = colossalai.__version__.split('torch')[1].split('cu')[0]
cuda_version_for_aot_build = colossalai.__version__.split('cu')[1]
except:
torch_version_for_aot_build = None
cuda_version_for_aot_build = None
return colossalai_verison, torch_version_for_aot_build, cuda_version_for_aot_build
def _check_aot_built_cuda_extension_installed():
"""
According to `op_builder/README.md`, the CUDA extension can be built with either
AOT (ahead-of-time) or JIT (just-in-time) compilation.
AOT compilation will build CUDA extensions to `colossalai._C` during installation.
JIT (just-in-time) compilation will build CUDA extensions to `~/.cache/colossalai/torch_extensions` during runtime.
"""
try:
import colossalai._C.fused_optim
is_cuda_extension_installed = u'\u2713'
found_aot_cuda_ext = True
except ImportError:
is_cuda_extension_installed = 'x'
return is_cuda_extension_installed
found_aot_cuda_ext = False
return found_aot_cuda_ext
def _check_torch_version():
"""
Get the PyTorch version information.
Returns:
torch_version: PyTorch version.
torch_cuda_version: CUDA version required by PyTorch.
"""
# get torch version
torch_version = torch.__version__.split('+')[0]
def _check_cuda_torch():
# get cuda version in pytorch build
torch_cuda_major = torch.version.cuda.split(".")[0]
torch_cuda_minor = torch.version.cuda.split(".")[1]
torch_cuda_version = f'{torch_cuda_major}.{torch_cuda_minor}'
return torch_version, torch_cuda_version
def _check_cuda_version():
"""
Get the CUDA version information.
Returns:
cuda_version: CUDA version found on the system.
"""
# get cuda version
if CUDA_HOME is None:
cuda_version = 'N/A (CUDA_HOME is not set)'
cuda_version = CUDA_HOME
else:
raw_output = subprocess.check_output([CUDA_HOME + "/bin/nvcc", "-V"], universal_newlines=True)
output = raw_output.split()
@ -83,13 +197,4 @@ def _check_cuda_torch():
bare_metal_major = release[0]
bare_metal_minor = release[1][0]
cuda_version = f'{bare_metal_major}.{bare_metal_minor}'
# get torch version
torch_version = torch.__version__.split('+')[0]
# get cuda version in pytorch build
torch_cuda_major = torch.version.cuda.split(".")[0]
torch_cuda_minor = torch.version.cuda.split(".")[1]
torch_cuda_version = f'{torch_cuda_major}.{torch_cuda_minor}'
return cuda_version, torch_version, torch_cuda_version
return cuda_version

Write Preview
Loading…
Cancel
Save