Add NeMo backend probe and architecture validator

docker/nemo/README.md

@@ -0,0 +1,11 @@
+# NVIDIA NeMo Backend
+
+本项目的训练后端目标是 NVIDIA 官方 NeMo/Megatron 栈，而不是旧的手写 PyTorch trainer。
+
+默认镜像：
+
+```text
+nvcr.io/nvidia/nemo:26.06
+```
+
+如果 g0033 拉取 NGC 镜像受限，可以先在可联网机器拉取后导入，或在本机配置 Docker daemon 代理。训练脚本只依赖 `/mnt/beegfs` 挂载，不把模型权重、数据 shard 或日志提交到 git。

scripts/probe_nemo_backend.sh

@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe}"
+IMAGE="${IMAGE:-nvcr.io/nvidia/nemo:26.06}"
+DOCKER_PULL="${DOCKER_PULL:-0}"
+USE_GPUS="${USE_GPUS:-0}"
+
+cd "$REPO_ROOT"
+python3 tools/validate_model_architecture.py
+
+if [ "$DOCKER_PULL" = "1" ]; then
+  docker pull "$IMAGE"
+fi
+
+GPU_ARGS=()
+if [ "$USE_GPUS" = "1" ]; then
+  GPU_ARGS=(--gpus all)
+fi
+
+docker run --rm "${GPU_ARGS[@]}" --ipc=host --network=host \
+  -v /mnt/beegfs:/mnt/beegfs \
+  -w "$REPO_ROOT" \
+  "$IMAGE" \
+  bash -lc 'set -euo pipefail
+python - <<PY
+import importlib
+mods = ["torch", "nemo", "megatron.core"]
+for name in mods:
+    mod = importlib.import_module(name)
+    version = getattr(mod, "__version__", "")
+    print("IMPORT_OK", name, version)
+PY
+python tools/validate_model_architecture.py
+'

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()