mirror of https://github.com/hpcaitech/ColossalAI
54 lines
2.0 KiB
Python
54 lines
2.0 KiB
Python
|
import os
|
||
|
|
import re
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from .builder import Builder
|
||
|
|
|
||
|
|
|
||
|
|
class FusedOptimBuilder(Builder):
|
||
|
|
NAME = "fused_optim"
|
||
|
|
BASE_DIR = "cuda_native/csrc"
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self.name = FusedOptimBuilder.NAME
|
||
|
|
super().__init__()
|
||
|
|
|
||
|
|
self.extra_cxx_flags = []
|
||
|
|
self.extra_cuda_flags = ['-lineinfo']
|
||
|
|
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:
|
||
|
|
self.extra_cuda_flags.extend(['-gencode', f'arch=compute_{arch_cap},code={arch}'])
|
||
|
|
|
||
|
|
self.sources = [self.colossalai_src_path(path) for path in self.sources_files()]
|
||
|
|
self.extra_include_paths = [self.colossalai_src_path(path) for path in self.include_paths()]
|
||
|
|
self.version_dependent_macros = ['-DVERSION_GE_1_1', '-DVERSION_GE_1_3', '-DVERSION_GE_1_5']
|
||
|
|
|
||
|
|
def sources_files(self):
|
||
|
|
return [
|
||
|
|
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'
|
||
|
|
]
|
||
|
|
]
|
||
|
|
|
||
|
|
def include_paths(self):
|
||
|
|
import torch
|
||
|
|
from torch.utils.cpp_extension import CUDA_HOME
|
||
|
|
cuda_include = os.path.join(CUDA_HOME, "include")
|
||
|
|
return [os.path.join(FusedOptimBuilder.BASE_DIR, "includes"), cuda_include]
|
||
|
|
|
||
|
|
def builder(self):
|
||
|
|
from torch.utils.cpp_extension import CUDAExtension
|
||
|
|
return CUDAExtension(
|
||
|
|
name=self.name,
|
||
|
|
sources=[os.path.join('colossalai/kernel/cuda_native/csrc', path) for path in self.sources],
|
||
|
|
include_dirs=self.extra_include_paths,
|
||
|
|
extra_compile_args={
|
||
|
|
'cxx': ['-O3'] + self.version_dependent_macros + self.extra_cxx_flags,
|
||
|
|
'nvcc': ['-O3', '--use_fast_math'] + self.extra_cuda_flags
|
||
|
|
})
|