mirror of https://github.com/hpcaitech/ColossalAI
[builder] reconfig op_builder for pypi install (#2314)
Jiarui Fang
2 years ago
committed by
GitHub
parent
a9b27b9265
commit
db6eea3583
@ -1,3 +1,4 @@
|
||||
include *.txt README.md
|
||||
recursive-include requirements *.txt
|
||||
recursive-include colossalai *.cpp *.h *.cu *.tr *.cuh *.cc *.pyi
|
||||
recursive-include op_builder *.py
|
||||
|
||||
@ -0,0 +1 @@
|
||||
../../op_builder
|
||||
@ -1,7 +0,0 @@
|
||||
from .cpu_adam import CPUAdamBuilder
|
||||
from .fused_optim import FusedOptimBuilder
|
||||
from .moe import MOEBuilder
|
||||
from .multi_head_attn import MultiHeadAttnBuilder
|
||||
from .scaled_upper_triang_masked_softmax import ScaledSoftmaxBuilder
|
||||
|
||||
__all__ = ['CPUAdamBuilder', 'FusedOptimBuilder', 'MultiHeadAttnBuilder', 'ScaledSoftmaxBuilder', 'MOEBuilder']
|
||||
@ -1,104 +0,0 @@
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def get_cuda_cc_flag() -> List:
|
||||
"""get_cuda_cc_flag
|
||||
|
||||
cc flag for your GPU arch
|
||||
"""
|
||||
cc_flag = []
|
||||
for arch in torch.cuda.get_arch_list():
|
||||
res = re.search(r'sm_(\d+)', arch)
|
||||
if res:
|
||||
arch_cap = res[1]
|
||||
if int(arch_cap) >= 60:
|
||||
cc_flag.extend(['-gencode', f'arch=compute_{arch_cap},code={arch}'])
|
||||
|
||||
return cc_flag
|
||||
|
||||
|
||||
class Builder(object):
|
||||
|
||||
def colossalai_src_path(self, code_path):
|
||||
if os.path.isabs(code_path):
|
||||
return code_path
|
||||
else:
|
||||
return os.path.join(Path(__file__).parent.parent.absolute(), code_path)
|
||||
|
||||
def get_cuda_home_include(self):
|
||||
"""
|
||||
return include path inside the cuda home.
|
||||
"""
|
||||
from torch.utils.cpp_extension import CUDA_HOME
|
||||
if CUDA_HOME is None:
|
||||
raise RuntimeError("CUDA_HOME is None, please set CUDA_HOME to compile C++/CUDA kernels in ColossalAI.")
|
||||
cuda_include = os.path.join(CUDA_HOME, "include")
|
||||
return cuda_include
|
||||
|
||||
# functions must be overrided begin
|
||||
def sources_files(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def include_dirs(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def cxx_flags(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def nvcc_flags(self):
|
||||
raise NotImplementedError
|
||||
|
||||
# functions must be overrided over
|
||||
|
||||
def strip_empty_entries(self, args):
|
||||
'''
|
||||
Drop any empty strings from the list of compile and link flags
|
||||
'''
|
||||
return [x for x in args if len(x) > 0]
|
||||
|
||||
def load(self, verbose=True):
|
||||
"""
|
||||
|
||||
load and compile cpu_adam lib at runtime
|
||||
|
||||
Args:
|
||||
verbose (bool, optional): show detailed info. Defaults to True.
|
||||
"""
|
||||
import time
|
||||
|
||||
from torch.utils.cpp_extension import load
|
||||
start_build = time.time()
|
||||
|
||||
op_module = load(name=self.name,
|
||||
sources=self.strip_empty_entries(self.sources_files()),
|
||||
extra_include_paths=self.strip_empty_entries(self.include_dirs()),
|
||||
extra_cflags=self.cxx_flags(),
|
||||
extra_cuda_cflags=self.nvcc_flags(),
|
||||
extra_ldflags=[],
|
||||
verbose=verbose)
|
||||
|
||||
build_duration = time.time() - start_build
|
||||
if verbose:
|
||||
print(f"Time to load {self.name} op: {build_duration} seconds")
|
||||
|
||||
return op_module
|
||||
|
||||
def builder(self, name) -> 'CUDAExtension':
|
||||
"""
|
||||
get a CUDAExtension instance used for setup.py
|
||||
"""
|
||||
from torch.utils.cpp_extension import CUDAExtension
|
||||
|
||||
return CUDAExtension(
|
||||
name=name,
|
||||
sources=[os.path.join('colossalai/kernel/cuda_native/csrc', path) for path in self.sources_files()],
|
||||
include_dirs=self.include_dirs(),
|
||||
extra_compile_args={
|
||||
'cxx': self.cxx_flags(),
|
||||
'nvcc': self.nvcc_flags()
|
||||
})
|
||||
@ -1,42 +0,0 @@
|
||||
import os
|
||||
|
||||
from .builder import Builder
|
||||
from .utils import append_nvcc_threads
|
||||
|
||||
|
||||
class CPUAdamBuilder(Builder):
|
||||
NAME = "cpu_adam"
|
||||
BASE_DIR = "cuda_native"
|
||||
|
||||
def __init__(self):
|
||||
self.name = CPUAdamBuilder.NAME
|
||||
super().__init__()
|
||||
|
||||
self.version_dependent_macros = ['-DVERSION_GE_1_1', '-DVERSION_GE_1_3', '-DVERSION_GE_1_5']
|
||||
|
||||
# necessary 4 functions
|
||||
def sources_files(self):
|
||||
ret = [
|
||||
os.path.join(CPUAdamBuilder.BASE_DIR, "csrc/cpu_adam.cpp"),
|
||||
]
|
||||
return [self.colossalai_src_path(path) for path in ret]
|
||||
|
||||
def include_dirs(self):
|
||||
return [
|
||||
self.colossalai_src_path(os.path.join(CPUAdamBuilder.BASE_DIR, "includes")),
|
||||
self.get_cuda_home_include()
|
||||
]
|
||||
|
||||
def cxx_flags(self):
|
||||
extra_cxx_flags = ['-std=c++14', '-lcudart', '-lcublas', '-g', '-Wno-reorder', '-fopenmp', '-march=native']
|
||||
return ['-O3'] + self.version_dependent_macros + extra_cxx_flags
|
||||
|
||||
def nvcc_flags(self):
|
||||
extra_cuda_flags = [
|
||||
'-std=c++14', '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__',
|
||||
'-U__CUDA_NO_HALF2_OPERATORS__', '-DTHRUST_IGNORE_CUB_VERSION_CHECK'
|
||||
]
|
||||
|
||||
return append_nvcc_threads(['-O3', '--use_fast_math'] + self.version_dependent_macros + extra_cuda_flags)
|
||||
|
||||
# necessary 4 functions
|
||||
@ -1,35 +0,0 @@
|
||||
import os
|
||||
|
||||
from .builder import Builder, get_cuda_cc_flag
|
||||
|
||||
|
||||
class FusedOptimBuilder(Builder):
|
||||
NAME = 'fused_optim'
|
||||
BASE_DIR = "cuda_native/csrc"
|
||||
|
||||
def __init__(self):
|
||||
self.name = FusedOptimBuilder.NAME
|
||||
super().__init__()
|
||||
self.version_dependent_macros = ['-DVERSION_GE_1_1', '-DVERSION_GE_1_3', '-DVERSION_GE_1_5']
|
||||
|
||||
def sources_files(self):
|
||||
ret = [
|
||||
self.colossalai_src_path(os.path.join(FusedOptimBuilder.BASE_DIR, fname)) for fname in [
|
||||
'colossal_C_frontend.cpp', 'multi_tensor_sgd_kernel.cu', 'multi_tensor_scale_kernel.cu',
|
||||
'multi_tensor_adam.cu', 'multi_tensor_l2norm_kernel.cu', 'multi_tensor_lamb.cu'
|
||||
]
|
||||
]
|
||||
return ret
|
||||
|
||||
def include_dirs(self):
|
||||
ret = [os.path.join(FusedOptimBuilder.BASE_DIR, "includes"), self.get_cuda_home_include()]
|
||||
return [self.colossalai_src_path(path) for path in ret]
|
||||
|
||||
def cxx_flags(self):
|
||||
extra_cxx_flags = []
|
||||
return ['-O3'] + self.version_dependent_macros + extra_cxx_flags
|
||||
|
||||
def nvcc_flags(self):
|
||||
extra_cuda_flags = ['-lineinfo']
|
||||
extra_cuda_flags.extend(get_cuda_cc_flag())
|
||||
return ['-O3', '--use_fast_math'] + extra_cuda_flags
|
||||
@ -1,33 +0,0 @@
|
||||
import os
|
||||
|
||||
from .builder import Builder, get_cuda_cc_flag
|
||||
|
||||
|
||||
class MOEBuilder(Builder):
|
||||
|
||||
def __init__(self):
|
||||
self.base_dir = "cuda_native/csrc"
|
||||
self.name = 'moe'
|
||||
super().__init__()
|
||||
|
||||
def include_dirs(self):
|
||||
ret = []
|
||||
ret = [os.path.join(self.base_dir, "includes"), self.get_cuda_home_include()]
|
||||
ret.append(os.path.join(self.base_dir, "kernels", "include"))
|
||||
return [self.colossalai_src_path(path) for path in ret]
|
||||
|
||||
def sources_files(self):
|
||||
ret = [os.path.join(self.base_dir, fname) for fname in ['moe_cuda.cpp', 'moe_cuda_kernel.cu']]
|
||||
return [self.colossalai_src_path(path) for path in ret]
|
||||
|
||||
def cxx_flags(self):
|
||||
return ['-O3', '-DVERSION_GE_1_1', '-DVERSION_GE_1_3', '-DVERSION_GE_1_5']
|
||||
|
||||
def nvcc_flags(self):
|
||||
extra_cuda_flags = [
|
||||
'-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '--expt-relaxed-constexpr',
|
||||
'--expt-extended-lambda'
|
||||
]
|
||||
extra_cuda_flags.extend(get_cuda_cc_flag())
|
||||
ret = ['-O3', '--use_fast_math'] + extra_cuda_flags
|
||||
return ret
|
||||
@ -1,41 +0,0 @@
|
||||
import os
|
||||
|
||||
from .builder import Builder, get_cuda_cc_flag
|
||||
|
||||
|
||||
class MultiHeadAttnBuilder(Builder):
|
||||
|
||||
def __init__(self):
|
||||
self.base_dir = "cuda_native/csrc"
|
||||
self.name = 'multihead_attention'
|
||||
super().__init__()
|
||||
|
||||
self.version_dependent_macros = ['-DVERSION_GE_1_1', '-DVERSION_GE_1_3', '-DVERSION_GE_1_5']
|
||||
|
||||
def include_dirs(self):
|
||||
ret = []
|
||||
ret = [os.path.join(self.base_dir, "includes"), self.get_cuda_home_include()]
|
||||
ret.append(os.path.join(self.base_dir, "kernels", "include"))
|
||||
return [self.colossalai_src_path(path) for path in ret]
|
||||
|
||||
def sources_files(self):
|
||||
ret = [
|
||||
os.path.join(self.base_dir, fname) for fname in [
|
||||
'multihead_attention_1d.cpp', 'kernels/cublas_wrappers.cu', 'kernels/transform_kernels.cu',
|
||||
'kernels/dropout_kernels.cu', 'kernels/normalize_kernels.cu', 'kernels/softmax_kernels.cu',
|
||||
'kernels/general_kernels.cu', 'kernels/cuda_util.cu'
|
||||
]
|
||||
]
|
||||
return [self.colossalai_src_path(path) for path in ret]
|
||||
|
||||
def cxx_flags(self):
|
||||
return ['-O3'] + self.version_dependent_macros
|
||||
|
||||
def nvcc_flags(self):
|
||||
extra_cuda_flags = [
|
||||
'-std=c++14', '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__',
|
||||
'-U__CUDA_NO_HALF2_OPERATORS__', '-DTHRUST_IGNORE_CUB_VERSION_CHECK'
|
||||
]
|
||||
extra_cuda_flags.extend(get_cuda_cc_flag())
|
||||
ret = ['-O3', '--use_fast_math'] + extra_cuda_flags
|
||||
return ret
|
||||
@ -1,36 +0,0 @@
|
||||
import os
|
||||
|
||||
from .builder import Builder, get_cuda_cc_flag
|
||||
|
||||
|
||||
class ScaledSoftmaxBuilder(Builder):
|
||||
|
||||
def __init__(self):
|
||||
self.base_dir = "cuda_native/csrc"
|
||||
self.name = 'scaled_upper_triang_masked_softmax'
|
||||
super().__init__()
|
||||
|
||||
def include_dirs(self):
|
||||
ret = []
|
||||
ret = [os.path.join(self.base_dir, "includes"), self.get_cuda_home_include()]
|
||||
ret.append(os.path.join(self.base_dir, "kernels", "include"))
|
||||
return [self.colossalai_src_path(path) for path in ret]
|
||||
|
||||
def sources_files(self):
|
||||
ret = [
|
||||
os.path.join(self.base_dir, fname)
|
||||
for fname in ['scaled_upper_triang_masked_softmax.cpp', 'scaled_upper_triang_masked_softmax_cuda.cu']
|
||||
]
|
||||
return [self.colossalai_src_path(path) for path in ret]
|
||||
|
||||
def cxx_flags(self):
|
||||
return ['-O3']
|
||||
|
||||
def nvcc_flags(self):
|
||||
extra_cuda_flags = [
|
||||
'-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '--expt-relaxed-constexpr',
|
||||
'--expt-extended-lambda'
|
||||
]
|
||||
extra_cuda_flags.extend(get_cuda_cc_flag())
|
||||
ret = ['-O3', '--use_fast_math'] + extra_cuda_flags
|
||||
return ret
|
||||
@ -1,20 +0,0 @@
|
||||
import subprocess
|
||||
|
||||
|
||||
def get_cuda_bare_metal_version(cuda_dir):
|
||||
raw_output = subprocess.check_output([cuda_dir + "/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]
|
||||
|
||||
return raw_output, bare_metal_major, bare_metal_minor
|
||||
|
||||
|
||||
def append_nvcc_threads(nvcc_extra_args):
|
||||
from torch.utils.cpp_extension import CUDA_HOME
|
||||
_, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
|
||||
if int(bare_metal_major) >= 11 and int(bare_metal_minor) >= 2:
|
||||
return nvcc_extra_args + ["--threads", "4"]
|
||||
return nvcc_extra_args
|
||||
Loading…
Reference in new issue