mirror of
https://github.com/deepspeedai/DeepSpeed.git
synced 2025-10-20 23:53:48 +08:00
66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
import torch
|
|
import warnings
|
|
import importlib
|
|
import warnings
|
|
|
|
GREEN = '\033[92m'
|
|
RED = '\033[91m'
|
|
YELLOW = '\033[93m'
|
|
END = '\033[0m'
|
|
SUCCESS = f"{GREEN} [SUCCESS] {END}"
|
|
WARNING = f"{YELLOW} [WARNING] {END}"
|
|
FAIL = f'{RED} [FAIL] {END}'
|
|
INFO = ' [INFO]'
|
|
|
|
try:
|
|
import deepspeed
|
|
print(f"{SUCCESS} deepspeed successfully imported.")
|
|
except ImportError as err:
|
|
raise err
|
|
|
|
print(f"{INFO} torch install path: {torch.__path__}")
|
|
print(f"{INFO} torch version: {torch.__version__}, torch.cuda: {torch.version.cuda}")
|
|
print(f"{INFO} deepspeed install path: {deepspeed.__path__}")
|
|
print(
|
|
f"{INFO} deepspeed info: {deepspeed.__version__}, {deepspeed.__git_hash__}, {deepspeed.__git_branch__}"
|
|
)
|
|
|
|
try:
|
|
apex_C = importlib.import_module('apex_C')
|
|
print(f"{SUCCESS} apex extensions successfully installed")
|
|
except Exception as err:
|
|
print(f'{WARNING} apex extensions are not installed')
|
|
|
|
try:
|
|
from apex.optimizers import FP16_Optimizer
|
|
print(f"{INFO} using old-style apex")
|
|
except ImportError:
|
|
print(f"{INFO} using new-style apex")
|
|
|
|
try:
|
|
importlib.import_module('deepspeed.ops.lamb.fused_lamb_cuda')
|
|
print(f'{SUCCESS} fused lamb successfully installed.')
|
|
except Exception as err:
|
|
print(f"{WARNING} fused lamb is NOT installed.")
|
|
|
|
try:
|
|
importlib.import_module('deepspeed.ops.transformer.transformer_cuda')
|
|
print(f'{SUCCESS} transformer kernels successfully installed.')
|
|
except Exception as err:
|
|
print(f'{WARNING} transformer kernels are NOT installed.')
|
|
|
|
try:
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter('ignore')
|
|
importlib.import_module('deepspeed.ops.sparse_attention.cpp_utils')
|
|
import triton
|
|
print(f'{SUCCESS} sparse attention successfully installed.')
|
|
except ImportError:
|
|
print(f'{WARNING} sparse attention is NOT installed.')
|
|
|
|
try:
|
|
importlib.import_module('deepspeed.ops.adam.cpu_adam_op')
|
|
print(f'{SUCCESS} cpu-adam (used by ZeRO-offload) successfully installed.')
|
|
except ImportError:
|
|
print(f'{WARNING} cpu-adam (used by ZeRO-offload) is NOT installed.')
|