Tie embeddings for small MoE config

2026-07-01 23:27:56 +08:00
parent 2c2b7ccc24
commit 816eccb5b5
4 changed files with 22 additions and 5 deletions
--- a/training/megatron_bridge/laoyao_2b_moe_pretrain.py
+++ b/training/megatron_bridge/laoyao_2b_moe_pretrain.py
@@ -30,7 +30,7 @@ def build_config(args: argparse.Namespace) -> ConfigContainer:
        seq_length=args.seq_length,
        vocab_size=154820,
        should_pad_vocab=True,
-        share_embeddings_and_output_weights=False,
+        share_embeddings_and_output_weights=True,
        position_embedding_type="rope",
        normalization="RMSNorm",
        gated_linear_unit=True,