diff --git a/docker/nemo/README.md b/docker/nemo/README.md new file mode 100644 index 0000000..b6b85a3 --- /dev/null +++ b/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。 diff --git a/scripts/probe_nemo_backend.sh b/scripts/probe_nemo_backend.sh new file mode 100755 index 0000000..842bc23 --- /dev/null +++ b/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 - < 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()