2024-03-05 13:52:30 +00:00
|
|
|
from .checkpoint import MoECheckpointIO
|
2023-11-02 02:21:24 +00:00
|
|
|
from .experts import MLPExperts
|
2023-12-14 09:52:05 +00:00
|
|
|
from .layers import SparseMLP, apply_load_balance
|
|
|
|
|
from .manager import MOE_MANAGER
|
2023-11-02 02:21:24 +00:00
|
|
|
from .routers import MoeRouter, Top1Router, Top2Router, TopKRouter
|
|
|
|
|
from .utils import NormalNoiseGenerator, UniformNoiseGenerator
|
|
|
|
|
|
|
|
|
|
__all__ = [
|
|
|
|
|
"MLPExperts",
|
|
|
|
|
"MoeRouter",
|
|
|
|
|
"Top1Router",
|
|
|
|
|
"Top2Router",
|
|
|
|
|
"TopKRouter",
|
|
|
|
|
"NormalNoiseGenerator",
|
|
|
|
|
"UniformNoiseGenerator",
|
|
|
|
|
"SparseMLP",
|
2024-03-05 13:52:30 +00:00
|
|
|
"MoECheckpointIO",
|
2023-12-14 09:52:05 +00:00
|
|
|
"MOE_MANAGER",
|
|
|
|
|
"apply_load_balance",
|
2023-11-02 02:21:24 +00:00
|
|
|
]
|
