diff --git a/model/README.md b/model/README.md index a76725b..8594756 100644 --- a/model/README.md +++ b/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 共享权重,从而减少大词表带来的参数开销。 diff --git a/model/nemo_megatron/laoyao_2b_moe_nemo_megatron.yaml b/model/nemo_megatron/laoyao_2b_moe_nemo_megatron.yaml index 8b3ebff..b49efe3 100644 --- a/model/nemo_megatron/laoyao_2b_moe_nemo_megatron.yaml +++ b/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 diff --git a/tools/validate_model_architecture.py b/tools/validate_model_architecture.py index 65d0f02..9af3a46 100755 --- a/tools/validate_model_architecture.py +++ b/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") diff --git a/training/megatron_bridge/laoyao_2b_moe_pretrain.py b/training/megatron_bridge/laoyao_2b_moe_pretrain.py index 7c25889..3c3c9ab 100755 --- 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,