mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-01 01:06:00 +00:00
f71e63b0f3
* Refactor MoE Manager setup method * unshard optim ckpt * optim io * update transformer version * update requirements * update ckpt * update ckpt * update ckpt * fix engine * fix engine
18 lines
431 B
Python
18 lines
431 B
Python
from .checkpoint import MoECheckpintIO
|
|
from .experts import MLPExperts
|
|
from .layers import SparseMLP
|
|
from .routers import MoeRouter, Top1Router, Top2Router, TopKRouter
|
|
from .utils import NormalNoiseGenerator, UniformNoiseGenerator
|
|
|
|
__all__ = [
|
|
"MLPExperts",
|
|
"MoeRouter",
|
|
"Top1Router",
|
|
"Top2Router",
|
|
"TopKRouter",
|
|
"NormalNoiseGenerator",
|
|
"UniformNoiseGenerator",
|
|
"SparseMLP",
|
|
"MoECheckpintIO",
|
|
]
|