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

[builder] reconfig op_builder for pypi install (#2314)

pull/2317/head
Jiarui Fang 2 years ago committed by GitHub
parent
a9b27b9265
commit
db6eea3583
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
17 changed files with 13 additions and 332 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. 1
      MANIFEST.in
  2. 1
      colossalai/kernel/op_builder
  3. 7
      colossalai/kernel/op_builder/__init__.py
  4. 104
      colossalai/kernel/op_builder/builder.py
  5. 42
      colossalai/kernel/op_builder/cpu_adam.py
  6. 35
      colossalai/kernel/op_builder/fused_optim.py
  7. 33
      colossalai/kernel/op_builder/moe.py
  8. 41
      colossalai/kernel/op_builder/multi_head_attn.py
  9. 36
      colossalai/kernel/op_builder/scaled_upper_triang_masked_softmax.py
  10. 20
      colossalai/kernel/op_builder/utils.py
  11. 8
      op_builder/builder.py
  12. 2
      op_builder/cpu_adam.py
  13. 2
      op_builder/fused_optim.py
  14. 2
      op_builder/moe.py
  15. 2
      op_builder/multi_head_attn.py
  16. 2
      op_builder/scaled_upper_triang_masked_softmax.py
  17. 7
      tests/test_optimizer/test_cpu_adam.py

1
MANIFEST.in
Unescape View File

@ -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

1
colossalai/kernel/op_builder
Unescape View File

@ -0,0 +1 @@
../../op_builder

7
colossalai/kernel/op_builder/__init__.py
Unescape View File

@ -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']

104
colossalai/kernel/op_builder/builder.py
Unescape View File

@ -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()
})

42
colossalai/kernel/op_builder/cpu_adam.py
Unescape View File

@ -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

35
colossalai/kernel/op_builder/fused_optim.py
Unescape View File

@ -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

33
colossalai/kernel/op_builder/moe.py
Unescape View File

@ -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

41
colossalai/kernel/op_builder/multi_head_attn.py
Unescape View File

@ -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

36
colossalai/kernel/op_builder/scaled_upper_triang_masked_softmax.py
Unescape View File

@ -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

20
colossalai/kernel/op_builder/utils.py
Unescape View File

@ -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

8
op_builder/builder.py
Unescape View File

@ -25,10 +25,12 @@ def get_cuda_cc_flag() -> List:
class Builder(object):
def colossalai_src_path(self, code_path):
if os.path.isabs(code_path):
return code_path
current_file_path = Path(__file__)
if os.path.islink(current_file_path.parent):
# symbolic link
return os.path.join(current_file_path.parent.parent.absolute(), code_path)
else:
return os.path.join(Path(__file__).parent.parent.absolute(), code_path)
return os.path.join(current_file_path.parent.parent.absolute(), "colossalai", "kernel", code_path)
def get_cuda_home_include(self):
"""

2
op_builder/cpu_adam.py
Unescape View File

@ -6,7 +6,7 @@ from .utils import append_nvcc_threads
class CPUAdamBuilder(Builder):
NAME = "cpu_adam"
BASE_DIR = "colossalai/kernel/cuda_native"
BASE_DIR = "cuda_native"
def __init__(self):
self.name = CPUAdamBuilder.NAME

2
op_builder/fused_optim.py
Unescape View File

@ -5,7 +5,7 @@ from .builder import Builder, get_cuda_cc_flag
class FusedOptimBuilder(Builder):
NAME = 'fused_optim'
BASE_DIR = "colossalai/kernel/cuda_native/csrc"
BASE_DIR = "cuda_native/csrc"
def __init__(self):
self.name = FusedOptimBuilder.NAME

2
op_builder/moe.py
Unescape View File

@ -6,7 +6,7 @@ from .builder import Builder, get_cuda_cc_flag
class MOEBuilder(Builder):
def __init__(self):
self.base_dir = "colossalai/kernel/cuda_native/csrc"
self.base_dir = "cuda_native/csrc"
self.name = 'moe'
super().__init__()

2
op_builder/multi_head_attn.py
Unescape View File

@ -6,7 +6,7 @@ from .builder import Builder, get_cuda_cc_flag
class MultiHeadAttnBuilder(Builder):
def __init__(self):
self.base_dir = "colossalai/kernel/cuda_native/csrc"
self.base_dir = "cuda_native/csrc"
self.name = 'multihead_attention'
super().__init__()

2
op_builder/scaled_upper_triang_masked_softmax.py
Unescape View File

@ -6,7 +6,7 @@ from .builder import Builder, get_cuda_cc_flag
class ScaledSoftmaxBuilder(Builder):
def __init__(self):
self.base_dir = "colossalai/kernel/cuda_native/csrc"
self.base_dir = "cuda_native/csrc"
self.name = 'scaled_upper_triang_masked_softmax'
super().__init__()

7
tests/test_optimizer/test_cpu_adam.py
Unescape View File

@ -66,12 +66,7 @@ def test_cpu_adam(adamw, step, p_dtype, g_dtype):
exp_avg_sq = torch.rand(p_data.shape)
exp_avg_sq_copy = exp_avg_sq.clone()
try:
from colossalai._C import cpu_optim
except:
from colossalai.kernel.op_builder import CPUAdamBuilder
cpu_optim = CPUAdamBuilder().load()
print("build CPUAdamOptimizer at runtime")
from colossalai.kernel import cpu_optim
cpu_adam_op = cpu_optim.CPUAdamOptimizer(lr, beta1, beta2, eps, weight_decay, adamw)

Write Preview
Loading…
Cancel
Save