github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-15 22:19:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

[misc] add verbose arg for zero and op builder (#3552)

Browse Source 
* [misc] add print verbose

* [gemini] add print verbose

* [zero] add print verbose for low level

* [misc] add print verbose for op builder
This commit is contained in:
Hongxin Liu
2023-04-17 11:25:35 +08:00
committed by GitHub
parent 4341f5e8e6
commit 173dad0562
8 changed files with 55 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
op_builder/builder.py
View File

@@ -7,7 +7,7 @@ import os
import time
from abc import ABC, abstractmethod
from pathlib import Path
from typing import List
from typing import List, Optional
from .utils import check_cuda_availability, check_system_pytorch_cuda_match, print_rank_0
@@ -138,7 +138,7 @@ class Builder(ABC):
# make sure system CUDA and pytorch CUDA match, an error will raised inside the function if not
check_system_pytorch_cuda_match(CUDA_HOME)
def load(self, verbose=True):
def load(self, verbose: Optional[bool] = None):
"""
load the kernel during runtime. If the kernel is not built during pip install, it will build the kernel.
If the kernel is built during runtime, it will be stored in `~/.cache/colossalai/torch_extensions/`. If the
@@ -149,6 +149,8 @@ class Builder(ABC):
Args:
verbose (bool, optional): show detailed info. Defaults to True.
"""
if verbose is None:
verbose = os.environ.get('CAI_KERNEL_VERBOSE', '0') == '1'
# if the kernel has be compiled and cached, we directly use it
if self.cached_op_module is not None:
return self.cached_op_module

20
op_builder/utils.py
View File

@@ -90,7 +90,6 @@ def check_system_pytorch_cuda_match(cuda_dir):
'Please make sure you have set the CUDA_HOME correctly and installed the correct PyTorch in https://pytorch.org/get-started/locally/ .'
)
print(bare_metal_minor != torch_cuda_minor)
if bare_metal_minor != torch_cuda_minor:
warnings.warn(
f"[extension] The CUDA version on the system ({bare_metal_major}.{bare_metal_minor}) does not match with the version ({torch_cuda_major}.{torch_cuda_minor}) torch was compiled with. "
@@ -156,16 +155,15 @@ def set_cuda_arch_list(cuda_dir):
# we only need to set this when CUDA is not available for cross-compilation
if not cuda_available:
warnings.warn(
'\n[extension] PyTorch did not find available GPUs on this system.\n'
'If your intention is to cross-compile, this is not an error.\n'
'By default, Colossal-AI will cross-compile for \n'
'1. Pascal (compute capabilities 6.0, 6.1, 6.2),\n'
'2. Volta (compute capability 7.0)\n'
'3. Turing (compute capability 7.5),\n'
'4. Ampere (compute capability 8.0, 8.6)if the CUDA version is >= 11.0\n'
'\nIf you wish to cross-compile for a single specific architecture,\n'
'export TORCH_CUDA_ARCH_LIST="compute capability" before running setup.py.\n')
warnings.warn('\n[extension] PyTorch did not find available GPUs on this system.\n'
'If your intention is to cross-compile, this is not an error.\n'
'By default, Colossal-AI will cross-compile for \n'
'1. Pascal (compute capabilities 6.0, 6.1, 6.2),\n'
'2. Volta (compute capability 7.0)\n'
'3. Turing (compute capability 7.5),\n'
'4. Ampere (compute capability 8.0, 8.6)if the CUDA version is >= 11.0\n'
'\nIf you wish to cross-compile for a single specific architecture,\n'
'export TORCH_CUDA_ARCH_LIST="compute capability" before running setup.py.\n')
if os.environ.get("TORCH_CUDA_ARCH_LIST", None) is None:
bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(cuda_dir)