Add NeMo backend probe and architecture validator

tools/validate_model_architecture.py

@@ -0,0 +1,76 @@
+#!/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 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)
+
+    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
+    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
+
+    print("architecture_ok")
+    print(f"tokenizer_sha256={tok_sha}")
+    print(f"vocab_size={vocab_size}")
+    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_total_params={dense_total/1e9:.3f}B")
+    print(f"rough_active_params={active_total/1e9:.3f}B")
+
+
+if __name__ == "__main__":
+    main()