mirror of https://github.com/hpcaitech/ColossalAI
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'
|
|
]
|