mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-05-15 11:53:00 +00:00
40d376c566
* [setup] support pre-build and jit-build of cuda kernels * polish code * polish code * polish code * polish code * polish code * polish code
24 lines
999 B
Python
24 lines
999 B
Python
from .cpu_adam import CPUAdamBuilder
|
|
from .fused_optim import FusedOptimBuilder
|
|
from .layernorm import LayerNormBuilder
|
|
from .moe import MOEBuilder
|
|
from .multi_head_attn import MultiHeadAttnBuilder
|
|
from .scaled_masked_softmax import ScaledMaskedSoftmaxBuilder
|
|
from .scaled_upper_triangle_masked_softmax import ScaledUpperTrainglemaskedSoftmaxBuilder
|
|
|
|
ALL_OPS = {
|
|
'cpu_adam': CPUAdamBuilder,
|
|
'fused_optim': FusedOptimBuilder,
|
|
'moe': MOEBuilder,
|
|
'multi_head_attn': MultiHeadAttnBuilder,
|
|
'scaled_masked_softmax': ScaledMaskedSoftmaxBuilder,
|
|
'scaled_upper_triangle_masked_softmax': ScaledUpperTrainglemaskedSoftmaxBuilder,
|
|
'layernorm': LayerNormBuilder,
|
|
}
|
|
|
|
__all__ = [
|
|
'ALL_OPS', 'CPUAdamBuilder', 'FusedOptimBuilder', 'MultiHeadAttnBuilder', 'ScaledMaskedSoftmaxBuilder',
|
|
'ScaledUpperTrainglemaskedSoftmaxBuilder', 'MOEBuilder', 'MultiTensorSGDBuilder', 'MultiTensorAdamBuilder',
|
|
'MultiTensorLambBuilder', 'MultiTensorScaleBuilder', 'MultiTensorL2NormBuilder'
|
|
]
|