Tie embeddings for small MoE config

model/README.md

@@ -19,4 +19,4 @@
 
 ## Tokenizer 兼容策略
 
-为了兼容后续 GLM-5.2 风格服务、SFT 数据和推理工具链，本模型不再使用 Jiayi 旧工程的 65K `laoyaomodel.bytes_bpe` tokenizer。模型词表直接设计为 GLM-5.2 同款 tokenizer JSON 和 token-id vocab。这样做的代价是 embedding/output head 参数量上升，但可以避免后续 tokenizer swap、数据重切和推理侧 token id 不一致的问题。
+为了兼容后续 GLM-5.2 风格服务、SFT 数据和推理工具链，本模型不再使用 Jiayi 旧工程的 65K `laoyaomodel.bytes_bpe` tokenizer。模型词表直接设计为 GLM-5.2 同款 tokenizer JSON 和 token-id vocab。由于 GLM-5.2 vocab 较大，当前小模型采用 tied embedding，输入 embedding 与输出 lm head 共享权重，从而减少大词表带来的参数开销。

model/nemo_megatron/laoyao_2b_moe_nemo_megatron.yaml

@@ -20,7 +20,7 @@ model:
   normalization: rmsnorm
   activation: swiglu
   position_embedding_type: rope
-  share_embeddings_and_output_weights: false
+  share_embeddings_and_output_weights: true
 
   moe:
     enabled: true

tools/validate_model_architecture.py

@@ -27,6 +27,13 @@ def read_list(name: str, text: str) -> list[int]:
     return [int(x.strip()) for x in match.group(1).split(",") if x.strip()]
 
 
+def read_bool(name: str, text: str) -> bool:
+    match = re.search(rf"^\s*{re.escape(name)}:\s*(true|false)\s*$", text, re.M)
+    if not match:
+        raise SystemExit(f"missing boolean field: {name}")
+    return match.group(1) == "true"
+
+
 def main() -> None:
     text = CONFIG.read_text(encoding="utf-8")
     tok_bytes = TOKENIZER.read_bytes()
@@ -45,6 +52,7 @@ def main() -> None:
     topk = read_scalar("moe_router_topk", text)
     expert_ffn = read_scalar("moe_ffn_hidden_size", text)
     moe_layers = read_list("layer_indices", text)
+    share_embeddings = read_bool("share_embeddings_and_output_weights", text)
 
     assert hidden % heads == 0, "hidden_size must divide num_attention_heads"
     assert heads % groups == 0, "num_attention_heads must divide num_query_groups for GQA"
@@ -56,18 +64,27 @@ def main() -> None:
     dense_layers = layers - len(moe_layers)
     # Rough parameter accounting for audit only. SwiGLU FFN uses 3 matrices.
     embed = cfg_vocab * hidden
+    embedding_total = embed if share_embeddings else embed * 2
     attn_per_layer = hidden * hidden * (2 + 2 * groups / heads)  # q,o full; k,v grouped
     dense_ffn_per_layer = 3 * hidden * dense_ffn
     moe_ffn_per_layer_total = num_experts * 3 * hidden * expert_ffn
     moe_ffn_per_layer_active = topk * 3 * hidden * expert_ffn
-    dense_total = embed * 2 + layers * attn_per_layer + dense_layers * dense_ffn_per_layer + len(moe_layers) * moe_ffn_per_layer_total
-    active_total = embed * 2 + layers * attn_per_layer + dense_layers * dense_ffn_per_layer + len(moe_layers) * moe_ffn_per_layer_active
+    shared_total = embedding_total + layers * attn_per_layer + dense_layers * dense_ffn_per_layer
+    active_expert_total = len(moe_layers) * moe_ffn_per_layer_active
+    total_expert_total = len(moe_layers) * moe_ffn_per_layer_total
+    dense_total = shared_total + total_expert_total
+    active_total = shared_total + active_expert_total
 
     print("architecture_ok")
     print(f"tokenizer_sha256={tok_sha}")
     print(f"vocab_size={vocab_size}")
+    print(f"share_embeddings_and_output_weights={str(share_embeddings).lower()}")
     print(f"layers={layers} hidden={hidden} heads={heads} query_groups={groups} seq_len={seq_len}")
     print(f"moe_layers={moe_layers} experts={num_experts} topk={topk} expert_ffn={expert_ffn}")
+    print(f"rough_embedding_params={embedding_total/1e9:.3f}B")
+    print(f"rough_shared_params={shared_total/1e9:.3f}B")
+    print(f"rough_active_expert_params={active_expert_total/1e9:.3f}B")
+    print(f"rough_total_expert_params={total_expert_total/1e9:.3f}B")
     print(f"rough_total_params={dense_total/1e9:.3f}B")
     print(f"rough_active_params={active_total/1e9:.3f}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,