94 lines
4.1 KiB
Python
Executable File
94 lines
4.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import re
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
CONFIG = ROOT / "model/nemo_megatron/laoyao_2b_moe_nemo_megatron.yaml"
|
|
TOKENIZER = ROOT / "tokenizer/glm5.2/tokenizer.json"
|
|
EXPECTED_TOKENIZER_SHA256 = "19e773648cb4e65de8660ea6365e10acca112d42a854923df93db4a6f333a82d"
|
|
|
|
|
|
def read_scalar(name: str, text: str) -> int:
|
|
match = re.search(rf"^\s*{re.escape(name)}:\s*([0-9]+)\s*$", text, re.M)
|
|
if not match:
|
|
raise SystemExit(f"missing integer field: {name}")
|
|
return int(match.group(1))
|
|
|
|
|
|
def read_list(name: str, text: str) -> list[int]:
|
|
match = re.search(rf"^\s*{re.escape(name)}:\s*\[([^\]]*)\]", text, re.M)
|
|
if not match:
|
|
raise SystemExit(f"missing list field: {name}")
|
|
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()
|
|
tok_sha = hashlib.sha256(tok_bytes).hexdigest()
|
|
tok = json.loads(tok_bytes)
|
|
vocab_size = len(tok.get("model", {}).get("vocab", {}))
|
|
|
|
hidden = read_scalar("hidden_size", text)
|
|
layers = read_scalar("num_layers", text)
|
|
heads = read_scalar("num_attention_heads", text)
|
|
groups = read_scalar("num_query_groups", text)
|
|
dense_ffn = read_scalar("ffn_hidden_size", text)
|
|
seq_len = read_scalar("seq_length", text)
|
|
cfg_vocab = read_scalar("vocab_size", text)
|
|
num_experts = read_scalar("num_moe_experts", text)
|
|
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"
|
|
assert cfg_vocab == vocab_size, f"config vocab_size={cfg_vocab} != tokenizer vocab_size={vocab_size}"
|
|
assert tok_sha == EXPECTED_TOKENIZER_SHA256, f"unexpected tokenizer sha256: {tok_sha}"
|
|
assert all(0 <= i < layers for i in moe_layers), "MoE layer index out of range"
|
|
assert 0 < topk <= num_experts, "moe_router_topk must be in [1, num_experts]"
|
|
|
|
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
|
|
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")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|