mirror of https://github.com/hpcaitech/ColossalAI
96 lines
3.4 KiB
Python
96 lines
3.4 KiB
Python
import subprocess
|
|
|
|
import click
|
|
import torch
|
|
from torch.utils.cpp_extension import CUDA_HOME
|
|
|
|
import colossalai
|
|
|
|
|
|
def check_installation():
|
|
cuda_ext_installed = _check_cuda_extension_installed()
|
|
cuda_version, torch_version, torch_cuda_version = _check_cuda_torch()
|
|
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}")
|
|
|
|
|
|
def _get_compatibility_string(versions):
|
|
|
|
# split version into [major, minor, patch]
|
|
versions = [version.split('.') for version in versions]
|
|
|
|
for version in versions:
|
|
if len(version) == 2:
|
|
# x means unknown
|
|
version.append('x')
|
|
|
|
for idx, version_values in enumerate(zip(*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'
|
|
elif idx == 1:
|
|
# if the minor versions match
|
|
# return a tick
|
|
return u'\u2713'
|
|
else:
|
|
continue
|
|
|
|
|
|
def _parse_colossalai_version():
|
|
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:
|
|
import colossalai._C.fused_optim
|
|
is_cuda_extension_installed = u'\u2713'
|
|
except ImportError:
|
|
is_cuda_extension_installed = 'x'
|
|
return is_cuda_extension_installed
|
|
|
|
|
|
def _check_cuda_torch():
|
|
# get cuda version
|
|
if CUDA_HOME is None:
|
|
cuda_version = 'N/A (CUDA_HOME is not set)'
|
|
else:
|
|
raw_output = subprocess.check_output([CUDA_HOME + "/bin/nvcc", "-V"], universal_newlines=True)
|
|
output = raw_output.split()
|
|
release_idx = output.index("release") + 1
|
|
release = output[release_idx].split(".")
|
|
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
|