20260801-moe model impl

need to add aux loss for load balancing
2026-08-01 22:48:58 +08:00
parent 925cbedc93
commit 6d98bb4f9f
6 changed files with 182 additions and 6 deletions
@@ -3,7 +3,7 @@ from astrai.model.components.attention import GQA, MLA
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.embedding import Embedding
 from astrai.model.components.linear import Linear
-from astrai.model.components.mlp import MLP
+from astrai.model.components.mlp import MLP, DeepSeekMoE
 from astrai.model.components.norm import RMSNorm
 from astrai.model.components.rope import (
    RotaryEmbedding,
@@ -14,6 +14,7 @@ __all__ = [
    "Linear",
    "RMSNorm",
    "MLP",
+    "DeepSeekMoE",
    "Embedding",
    "GQA",
    "MLA",