大模型网络结构：模型接口测试用例参考
大模型的网络文件是模型训练的核心参考。例如 MiniMind 项目中 `/model/model_minimind.py` 这份网络结构文件，它用来定义整套大模型的神经网络拓扑，包含配置参数、RMSNorm 归一化、GQA 注意力、RoPE 位置编码、Transformer 层、MoE 专家模块、因果语言模型输出头等全部网络组件。 在训练过程中，该文件会被训练脚本调用执行：完成输入 token 的前向计算、基础损失求取、MoE 路由辅助损失计算，输出 logits，结合标签计算交叉熵损失，为反向传播提供计算图，驱动模型权重参数更新，完成训练迭代。训练结束后，同一套网络代码可直接复用在推理阶段，加载训练完成的权重，实现文本续写与生成。 **注：代码完全是AI生成的，但经过了验证，可以应用到你自己的 自研大模型原型上测试。** 为快速验证网络实现是否存在代码缺陷、维度错误、模块兼容性问题，这里提供两个可独立运行的测试用例。执行测试用例能够对模型网络做基础完备性校验，自动输出模型运行状态报告，帮助开发者快速定位网络结构问题，可供用户参考与学习： ## 测试网络文件22项 此处我这边对Minimind模型进行了扩充功能，扩展增加了70%的代码，此处用户可自行测试自己的网络文件，代码仅供参考。 ``` import math import os import shutil import torch from transformers.cache_utils import DynamicCache # 此处导入网络模型文件 from model_minimind import MiniMindConfig, MiniMindForCausalLM def print_sep(name: str): print(f"\n{'='*70}") print(f"【{name}】") print(f"{'='*70}") def assert_with_info(cond: bool, msg: str, **kwargs): if not cond: print("\n断言失败信息：") for k, v in kwargs.items(): print(f" {k} = {v}") raise AssertionError(msg) def test_1_boundary_inputs(): """边界输入：seq_len=1、batch_size=1、单token推理""" print_sep("test_1_boundary_inputs") cfg = MiniMindConfig( hidden_size=192, num_hidden_layers=2, num_attention_heads=3, num_key_value_heads=1, vocab_size=512, use_moe=False ) model = MiniMindForCausalLM(cfg).eval() device = "cuda" if torch.cuda.is_available() else "cpu" model.to(device) print(f"model device: {device}") ids1 = torch.randint(0, cfg.vocab_size, (2, 1), device=device) print(f"case1 input_ids shape: {ids1.shape}") with torch.no_grad(): o1 = model(ids1) print(f"case1 output logits shape: {o1.logits.shape}") assert_with_info(o1.logits.shape == (2, 1, cfg.vocab_size), "seq_len=1 logits shape mismatch", expect=(2, 1, cfg.vocab_size), actual=o1.logits.shape) ids2 = torch.randint(0, cfg.vocab_size, (1, 32), device=device) print(f"case2 input_ids shape: {ids2.shape}") with torch.no_grad(): o2 = model(ids2) print(f"case2 output logits shape: {o2.logits.shape}") assert_with_info(o2.logits.shape == (1, 32, cfg.vocab_size), "batch=1 long seq shape mismatch", expect=(1, 32, cfg.vocab_size), actual=o2.logits.shape) prompt = torch.tensor([[42]], device=device) print(f"case3 prompt shape: {prompt.shape}, max_new_tokens=5") gen = model.generate(input_ids=prompt, max_new_tokens=5, do_sample=False) print(f"case3 generated output shape: {gen.shape}") assert_with_info(gen.shape[-1] == 1 + 5, "generate output length wrong", expect_len=1 + 5, actual_len=gen.shape[-1]) print("[√] test_1_boundary_inputs PASS") def test_2_rope_yarn_scaling(): """YaRN RoPE缩放，验证buffer维度与设备""" print_sep("test_2_rope_yarn_scaling") cfg = MiniMindConfig( hidden_size=256, num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, vocab_size=800, max_position_embeddings=32768, inference_rope_scaling=True, use_moe=False ) model = MiniMindForCausalLM(cfg).eval() device = "cuda" if torch.cuda.is_available() else "cpu" model.to(device) print(f"model device: {device}") print(f"freqs_cos buffer shape before forward: {model.model.freqs_cos.shape}, device={model.model.freqs_cos.device}") L = 4096 ids = torch.randint(0, cfg.vocab_size, (1, L), device=device) print(f"input long seq length = {L}, input shape {ids.shape}") with torch.no_grad(): out = model(ids) print(f"output logits shape {out.logits.shape}") print(f"freqs_cos buffer shape after forward: {model.model.freqs_cos.shape}, device={model.model.freqs_cos.device}") assert_with_info(out.logits.shape == (1, L, cfg.vocab_size), "yarn long seq logits shape error", expect=(1, L, cfg.vocab_size), actual=out.logits.shape) assert_with_info(model.model.freqs_cos.shape[0] == cfg.max_position_embeddings, "rope buffer size mismatch", expect=cfg.max_position_embeddings, actual=model.model.freqs_cos.shape[0]) print("[√] test_2_rope_yarn_scaling PASS") def test_3_attention_padding_mask(): """不等长padding mask训练，验证ignore_index生效""" print_sep("test_3_attention_padding_mask") cfg = MiniMindConfig( hidden_size=256, num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, vocab_size=1000, use_moe=False ) model = MiniMindForCausalLM(cfg).train() input_ids = torch.tensor([ [10, 11, 12, 13, 14, 15, 16, 17], [20, 21, 22, 23, 0, 0, 0, 0] ]) attention_mask = torch.tensor([ [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 0, 0, 0, 0] ]) labels = input_ids.clone() labels[1, 4:] = -100 print(f"input_ids:\n{input_ids}") print(f"attention_mask:\n{attention_mask}") print(f"labels:\n{labels}") out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) loss = out.loss print(f"computed training loss = {loss.item():.6f}") loss.backward() grad_emb = model.model.embed_tokens.weight.grad print(f"embed_tokens grad is None? {grad_emb is None}") assert_with_info(grad_emb is not None, "embedding grad should not be None") print("[√] test_3_attention_padding_mask PASS") def test_4_kv_cache_compatibility(): """DynamicCache / list-kv双向兼容""" print_sep("test_4_kv_cache_compatibility") cfg = MiniMindConfig( hidden_size=256, num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, vocab_size=600, use_moe=False ) model = MiniMindForCausalLM(cfg).eval() device = "cuda" if torch.cuda.is_available() else "cpu" model.to(device) prompt = torch.randint(0, cfg.vocab_size, (1, 12), device=device) print(f"prompt shape {prompt.shape}") cache_dyn = DynamicCache() with torch.no_grad(): o_dyn = model(prompt, past_key_values=cache_dyn, use_cache=True) logit_dyn_last = o_dyn.logits[0, -1].clone() print(f"DynamicCache output past_key_values type: {type(o_dyn.past_key_values)}") with torch.no_grad(): o_list = model(prompt, past_key_values=None, use_cache=True) logit_list_last = o_list.logits[0, -1].clone() print(f"list-kv output past_key_values type: {type(o_list.past_key_values)}") diff = torch.max(torch.abs(logit_dyn_last - logit_list_last)).item() print(f"logits max abs diff between two cache format: {diff:.2e}") assert_with_info(diff > skip fp16/bf16 test (no gpu or skip_gpu_case=True)") return for dtype in (torch.float16, torch.bfloat16): model = MiniMindForCausalLM(cfg).to(dev, dtype=dtype) model.train() ids = torch.randint(0, cfg.vocab_size, (2, 16), device=dev) out = model(ids, labels=ids) loss = out.loss print(f"dtype={dtype}, loss={loss.item():.6f}") loss.backward() grad_dtype = model.model.embed_tokens.weight.grad.dtype print(f" grad dtype = {grad_dtype}") assert_with_info(grad_dtype == dtype, f"grad dtype mismatch, expect {dtype}, got {grad_dtype}") print("[√] test_8_mixed_precision_fp16_bf16 PASS") def test_9_model_save_load_hf_style(tmp_dir="./tmp_minimind_test"): """HF格式保存加载验证""" print_sep("test_9_model_save_load_hf_style") cfg = MiniMindConfig( hidden_size=192, num_hidden_layers=2, num_attention_heads=3, num_key_value_heads=1, vocab_size=500, use_moe=True, num_experts=2 ) model1 = MiniMindForCausalLM(cfg).eval() device = "cuda" if torch.cuda.is_available() else "cpu" model1.to(device) inp = torch.randint(0, cfg.vocab_size, (1, 8), device=device) print(f"save-load test input shape {inp.shape}") with torch.no_grad(): logits_before = model1(inp, use_cache=False).logits.clone() if os.path.exists(tmp_dir): shutil.rmtree(tmp_dir) model1.save_pretrained(tmp_dir) print("[√] save_pretrained finished") model2 = MiniMindForCausalLM.from_pretrained(tmp_dir) model2.to(device) model2.eval() with torch.no_grad(): logits_after = model2(inp, use_cache=False).logits.clone() diff = torch.max(torch.abs(logits_before - logits_after)).item() print(f"logits max abs diff after reload = {diff:.2e}") assert_with_info(diff {losses[-1]}" assert losses[-1] 1e-4, f"Parameters did not update significantly (diff={param_diff})" # 3. 梯度应在 step 后被 optimizer 处理（虽然不一定为0，取决于实现，但通常检查是否还能反向传播） # 这里主要靠上面的 param_diff 证明更新发生了 print(f"[√] Loss decreased from {losses[0]:.4f} to {losses[-1]:.4f}") print(f"[√] Parameters updated (norm diff: {param_diff:.4f})") print("[√] test_21_minites PASS") def test_22_token_id_test(): """验证生成逻辑：词表范围、长度约束、返回结构、EOS 行为""" print_sep("test_22_token_id_test") torch.manual_seed(42) cfg = MiniMindConfig( hidden_size=256, num_hidden_layers=2, vocab_size=500, eos_token_id=2, pad_token_id=0, use_moe=False ) model = MiniMindForCausalLM(cfg).eval() prompt = torch.tensor([[10, 11, 12]]) max_new = 10 # ---------- Case 1: 基础生成 + 合法性 ---------- out = model.generate(input_ids=prompt, max_new_tokens=max_new, do_sample=False) seq = _pick_seq(out) assert isinstance(seq, torch.Tensor), "sequence must be a Tensor" assert seq.ndim == 2, f"expect [B, L], got shape {seq.shape}" assert seq.shape[0] == prompt.shape[0], "batch size changed" assert seq.max().item() = vocab_size" assert seq.min().item() >= 0, "token id 0, "nothing generated" # 生成段里应该出现 EOS，且 EOS 之后全是 pad if cfg.eos_token_id in new_part: first_eos = (new_part == cfg.eos_token_id).nonzero()[0, 0].item() after = new_part[first_eos + 1:] assert (after == cfg.pad_token_id).all() or after.numel() == 0, \ "tokens generated after EOS" # 恢复 lm_head bias if has_bias: lm_head.bias.data.copy_(bias_backup) else: lm_head.bias = original_bias # ---------- Case 4: 确定性（同 seed 同结果）---------- torch.manual_seed(7) m1 = MiniMindForCausalLM(cfg).eval() torch.manual_seed(7) m2 = MiniMindForCausalLM(cfg).eval() o1 = _pick_seq(m1.generate(input_ids=prompt, max_new_tokens=8, do_sample=False)) o2 = _pick_seq(m2.generate(input_ids=prompt, max_new_tokens=8, do_sample=False)) assert (o1 == o2).all(), "same seed -> different output, init not deterministic" print(f"generated: {seq.tolist()}") print(f"eos-truncated len: {seq_eos.shape[-1]} (max allowed {prompt.shape[-1]+20})") print("[√] test_22_token_id_test PASS") # ---- 两个小工具，避免测试被返回格式卡死 ---- def _pick_seq(out): if isinstance(out, dict): for k in ("sequences", "generated_ids", "output_ids"): if k in out: return out[k] raise KeyError(f"no sequence key in {out.keys()}") if hasattr(out, "sequences"): return out.sequences if isinstance(out, torch.Tensor): return out raise TypeError(f"unknown generate return type: {type(out)}") def _has_scores(out): if isinstance(out, dict): return "scores" in out return hasattr(out, "scores") if __name__ == "__main__": import sys skip_gpu = "--skip-gpu" in sys.argv print("Start MiniMind full test suite") test_1_boundary_inputs() test_2_rope_yarn_scaling() test_3_attention_padding_mask() test_4_kv_cache_compatibility() test_5_tie_word_embedding() test_6_moe_advanced() test_7_generate_decoding_modes() test_8_mixed_precision_fp16_bf16(skip_gpu_case=skip_gpu) test_9_model_save_load_hf_style() test_10_kv_cache_step_by_step_equivalence() test_11_logits_to_keep_slice() test_12_tie_weights_method() test_13_past_key_values_none_and_empty_list() test_14_moe_topk_1() test_15_attention_non_full_mask_no_flash() test_16_rope_no_yarn() test_17_generate_edge_cases() test_18_use_cache_false() test_19_logits_to_keep_zero() test_20_custom_head_dim_and_moe_intermediate() test_21_minites() test_22_token_id_test() print("\n" + "#" * 70) print("ALL TEST CASES PASSED!") print("#" * 70) ``` 运行输出模型测试结果： ``` Start MiniMind full test suite ====================================================================== 【test_1_boundary_inputs】 ====================================================================== model device: cpu case1 input_ids shape: torch.Size([2, 1]) case1 output logits shape: torch.Size([2, 1, 512]) case2 input_ids shape: torch.Size([1, 32]) case2 output logits shape: torch.Size([1, 32, 512]) case3 prompt shape: torch.Size([1, 1]), max_new_tokens=5 case3 generated output shape: torch.Size([1, 6]) [√] test_1_boundary_inputs PASS ====================================================================== 【test_2_rope_yarn_scaling】 ====================================================================== model device: cpu freqs_cos buffer shape before forward: torch.Size([32768, 64]), device=cpu input long seq length = 4096, input shape torch.Size([1, 4096]) output logits shape torch.Size([1, 4096, 800]) freqs_cos buffer shape after forward: torch.Size([32768, 64]), device=cpu [√] test_2_rope_yarn_scaling PASS ====================================================================== 【test_3_attention_padding_mask】 ====================================================================== input_ids: tensor([[10, 11, 12, 13, 14, 15, 16, 17], [20, 21, 22, 23, 0, 0, 0, 0]]) attention_mask: tensor([[1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 0, 0, 0, 0]]) labels: tensor([[ 10, 11, 12, 13, 14, 15, 16, 17], [ 20, 21, 22, 23, -100, -100, -100, -100]]) computed training loss = 6.919206 embed_tokens grad is None? False [√] test_3_attention_padding_mask PASS ====================================================================== 【test_4_kv_cache_compatibility】 ====================================================================== prompt shape torch.Size([1, 12]) DynamicCache output past_key_values type: list-kv output past_key_values type: logits max abs diff between two cache format: 0.00e+00 [√] test_4_kv_cache_compatibility PASS ====================================================================== 【test_5_tie_word_embedding】 ====================================================================== tie=True: lm_head.weight[0,:5] = tensor([ 0.0322, -0.0216, 0.0046, 0.0141, 0.0100], grad_fn=) after modify lm_head, embed_tokens[0,:5] = tensor([0.1322, 0.0784, 0.1046, 0.1141, 0.1100], grad_fn=) tie=False: lm_head and embed_tokens equal? False [√] test_5_tie_word_embedding PASS ====================================================================== 【test_6_moe_advanced】 ====================================================================== input shape B=4, S=16 train total loss(ce+aux+z) = 6.794805 [√] all experts & router have grad eval mode aux_loss layer0: 0.0 [√] test_6_moe_advanced PASS ====================================================================== 【test_7_generate_decoding_modes】 ====================================================================== prompt shape torch.Size([1, 1]) greedy output shape torch.Size([1, 13]) top-k sample output shape torch.Size([1, 9]) top-p sample output shape torch.Size([1, 9]) repetition_penalty output shape torch.Size([1, 9]) eos-prompt output shape torch.Size([1, 1]), expected > skip fp16/bf16 test (no gpu or skip_gpu_case=True) ====================================================================== 【test_9_model_save_load_hf_style】 ====================================================================== save-load test input shape torch.Size([1, 8]) Writing model shards: 100%|██████████████████████████████████████████████████████████████████████████████████ ████████████████████████████████████████████████████| 1/1 [00:00 past_key_values=[] empty-list ok [√] test_13_past_key_values_none_and_empty_list PASS ====================================================================== 【test_14_moe_topk_1】 ====================================================================== [√] MoE top-k=1 train grad ok eval aux_loss item: 0.0 [√] test_14_moe_topk_1 PASS ====================================================================== 【test_15_attention_non_full_mask_no_flash】 ====================================================================== no-flash non-full mask logits shape torch.Size([1, 5, 500]) flash vs no-flash logits max diff: 5.364e-07 [√] test_15_attention_non_full_mask_no_flash PASS ====================================================================== 【test_16_rope_no_yarn】 ====================================================================== no-yarn RoPE run ok, logits shape torch.Size([1, 1024, 600]) [√] test_16_rope_no_yarn PASS ====================================================================== 【test_17_generate_edge_cases】 ====================================================================== num_return_sequences=2 shape torch.Size([2, 5]) temperature=0, no sampling constraints shape torch.Size([1, 4]) return_kv is dict: True, keys=['generated_ids', 'past_kv'] eos_token_id=None output len=6 expect 6 [√] test_17_generate_edge_cases PASS ====================================================================== 【test_18_use_cache_false】 ====================================================================== use_cache=False past_key_values = None [√] test_18_use_cache_false PASS ====================================================================== 【test_19_logits_to_keep_zero】 ====================================================================== logits_to_keep=0 loss=6.5736 [√] test_19_logits_to_keep_zero PASS ====================================================================== 【test_20_custom_head_dim_and_moe_intermediate】 ====================================================================== custom head_dim=96, moe_intermediate_size=512, logits shape torch.Size([1, 8, 600]) [√] test_20_custom_head_dim_and_moe_intermediate PASS ====================================================================== 【test_21_minites】 ====================================================================== step 0, loss: 6.9440 step 1, loss: 4.8753 step 2, loss: 3.6879 step 3, loss: 2.9489 step 4, loss: 2.2793 step 5, loss: 1.8152 step 6, loss: 1.4206 step 7, loss: 1.1400 step 8, loss: 0.9259 step 9, loss: 0.7511 step 10, loss: 0.6075 step 11, loss: 0.4912 step 12, loss: 0.3980 step 13, loss: 0.3233 step 14, loss: 0.2633 step 15, loss: 0.2152 step 16, loss: 0.1765 step 17, loss: 0.1456 step 18, loss: 0.1209 step 19, loss: 0.1013 [√] Loss decreased from 6.9440 to 0.1013 [√] Parameters updated (norm diff: 2.7410) [√] test_21_minites PASS ====================================================================== 【test_22_token_id_test】 ====================================================================== generated: [[10, 11, 12, 2]] eos-truncated len: 4 (max allowed 23) [√] test_22_token_id_test PASS ###################################################################### ALL TEST CASES PASSED! ###################################################################### ``` ## 引擎多维测试套件 这里提供一个引擎功能测试套件代码，读者可自行理解测试不同的模型网络结构。 ``` # -*- coding: utf-8 -*- import argparse import gc import math import os import statistics import sys import tempfile import time import traceback from contextlib import contextmanager import torch sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import model_minimind as fixed try: import model_minimind_baseline as baseline HAS_BASELINE = True except Exception: HAS_BASELINE = False # torch 线程数固定，降低基准抖动 TORCH_THREADS = min(4, os.cpu_count() or 1) torch.set_num_threads(TORCH_THREADS) DEVICE = "cuda" if torch.cuda.is_available() else "cpu" DEFAULT_SEED = 20260919 # ============================================================ # 微型测试框架 # ============================================================ class CheckError(AssertionError): pass def expect(cond, msg="", **info): if not cond: extra = " | ".join(f"{k}={v}" for k, v in info.items()) raise CheckError(f"{msg} {extra}".strip()) def expect_raises(exc_type, fn, msg=""): try: fn() except exc_type: return except Exception as e: raise CheckError(f"{msg} 期望抛出 {exc_type.__name__}，实际抛出 {type(e).__name__}: {e}") raise CheckError(f"{msg} 期望抛出 {exc_type.__name__}，但未抛出任何异常") @contextmanager def seed_context(seed): cpu_state = torch.random.get_rng_state() cuda_state = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None torch.manual_seed(seed) try: yield finally: torch.random.set_rng_state(cpu_state) if cuda_state is not None: torch.cuda.set_rng_state_all(cuda_state) REGISTRY = [] # (group, name, fn) def testcase(group, name): def deco(fn): REGISTRY.append((group, name, fn)) return fn return deco def small_config(mod, **overrides): cfg = dict( hidden_size=96, num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, vocab_size=128, eos_token_id=2, flash_attn=False, use_moe=False, max_position_embeddings=4096, ) cfg.update(overrides) return mod.MiniMindConfig(**cfg) def build_model(mod=None, seed=DEFAULT_SEED, eval_mode=True, **cfg_overrides): mod = mod or fixed with seed_context(seed): model = mod.MiniMindForCausalLM(small_config(mod, **cfg_overrides)) model.to(DEVICE) if eval_mode: model.eval() return model def clone_weights(dst, src): dst.load_state_dict(src.state_dict()) def ids(*rows): return torch.tensor(list(rows), dtype=torch.long, device=DEVICE) # ============================================================ # 一、功能测试 # ============================================================ @testcase("functional", "F01 多batch/多序列长度前向形状") def f01(): m = build_model() for shape in [(1, 1), (1, 32), (3, 7), (2, 128)]: x = torch.randint(0, 128, shape, device=DEVICE) with torch.no_grad(): out = m(x) expect(tuple(out.logits.shape) == (*shape, 128), "logits 形状错误", expect=(*shape, 128), got=tuple(out.logits.shape)) expect(torch.isfinite(out.logits).all().item(), "logits 含 NaN/Inf") @testcase("functional", "F02 多语言 Unicode 字符级通路") def f02(): """模型本身只认 token id，多语言切词是 tokenizer 职责； 这里用确定性字符级映射验证 Unicode 文本→id→生成→映射回文本的整条通路。""" texts = { "zh": "大模型推理引擎正在测试中文字符序列", "en": "The quick brown fox jumps over the lazy dog", "ja": "日本語のテキストを生成エンジンで処理します", "ar": "مرحبا بالعالم اختبار محرك النموذج", "emoji": "大模型ABC", } specials = {"": 0, "": 1, "": 2, "": 3} vocab = dict(specials) for t in texts.values(): for ch in t: if ch not in vocab: vocab[ch] = len(vocab) expect(len(vocab) = self.schedule[r] else (10 + self.step) out[r, :, tok] = 1e9 self.step += 1 return out hook = m.lm_head.register_forward_hook(ForceEOS([2, 4], 90, 2)) prompt = ids([1, 5], [1, 6]) g = m.generate(input_ids=prompt, max_new_tokens=10, do_sample=False, top_k=0, eos_token_id=[2, 90]) hook.remove() # 两行都结束后停止 → 长度 = 2 + 4 = 6 expect(g.shape[1] == 6, "错峰 EOS 停止位置错误", got=g.shape[1]) expect(g[0, -1].item() == 90 and g[1, -1].item() == 90, "末位应为 EOS") expect(all(g[0, j].item() == 90 for j in range(4, 6)), "结束行应填充 EOS 占位", row0=g[0].tolist()) g0 = m.generate(input_ids=ids([2]), max_new_tokens=5, do_sample=False) expect(g0.shape[1] == 1, "末位 EOS 应立即终止", got=g0.shape[1]) hook2 = m.lm_head.register_forward_hook(ForceEOS([0, 0], 2, 1)) g1 = m.generate(input_ids=ids([1]), max_new_tokens=5, do_sample=False, top_k=0, eos_token_id=None) hook2.remove() expect(g1.shape[1] == 6, "禁用 EOS 后应生成满额", got=g1.shape[1]) @testcase("functional", "F05 流式生成（streamer）") def f05(): class ListStreamer: def __init__(self): self.chunks = [] self.ended = False def put(self, x): self.chunks.append(x.clone()) def end(self): self.ended = True m = build_model() st = ListStreamer() prompt = ids([1, 2, 3]) g = m.generate(input_ids=prompt, max_new_tokens=6, do_sample=False, top_k=0, eos_token_id=None, streamer=st) expect(st.ended, "streamer.end() 未被调用") expect(len(st.chunks) == 7, "streamer chunk 数应为 prompt+6", got=len(st.chunks)) rebuilt = torch.cat(st.chunks, dim=-1) expect(torch.equal(rebuilt, g.cpu()), "streamer 拼接结果与最终输出不一致") @testcase("functional", "F06 batch 生成与逐条生成一致（贪婪）") def f06(): m = build_model() prompts = ids([1, 7, 42], [8, 8, 9], [100, 50, 25]) g_batch = m.generate(input_ids=prompts, max_new_tokens=8, do_sample=False, top_k=0, eos_token_id=None) for r in range(3): g_single = m.generate(input_ids=prompts[r:r+1], max_new_tokens=8, do_sample=False, top_k=0, eos_token_id=None) expect(torch.equal(g_batch[r], g_single[0]), f"第{r}行 batch/单条结果不一致", batch=g_batch[r].tolist(), single=g_single[0].tolist()) @testcase("functional", "F07 右 padding 批量生成等价于无 padding 单条") def f07(): m = build_model() seqs = [[3, 11, 27, 44, 88], [60, 61]] pad_len = 5 padded = torch.full((2, pad_len), 0, dtype=torch.long, device=DEVICE) mask = torch.zeros(2, pad_len, dtype=torch.long, device=DEVICE) for r, s in enumerate(seqs): padded[r, :len(s)] = torch.tensor(s) mask[r, :len(s)] = 1 g_pad = m.generate(input_ids=padded, attention_mask=mask, max_new_tokens=8, do_sample=False, top_k=0, eos_token_id=None) for r, s in enumerate(seqs): g_single = m.generate(input_ids=ids(s), max_new_tokens=8, do_sample=False, top_k=0, eos_token_id=None) new_pad = g_pad[r, pad_len:].tolist() new_single = g_single[0, len(s):].tolist() expect(new_pad == new_single, f"行{r}右padding生成内容不一致", pad=new_pad, single=new_single) @testcase("functional", "F08 use_cache True/False 输出一致") def f08(): m = build_model() p = ids([1, 2, 3, 4, 5]) kw = dict(max_new_tokens=8, do_sample=False, top_k=0, eos_token_id=None) g1 = m.generate(input_ids=p, use_cache=True, **kw) g2 = m.generate(input_ids=p, use_cache=False, **kw) expect(torch.equal(g1, g2), "cache 开关输出不一致", on=g1.tolist(), off=g2.tolist()) @testcase("functional", "F09 return_kv / num_return_sequences / 1D / dict 输入") def f09(): m = build_model() ret = m.generate(input_ids=ids([1, 2]), max_new_tokens=3, do_sample=False, top_k=0, eos_token_id=None, return_kv=True) expect(isinstance(ret, dict) and "generated_ids" in ret and "past_kv" in ret, "return_kv 返回结构错误") expect(ret["past_kv"] is not None, "return_kv 中 past_kv 为空") g = m.generate(input_ids=ids([7], [8]), max_new_tokens=2, num_return_sequences=3, do_sample=False, top_k=0, eos_token_id=None) expect(g.shape[0] == 6, "num_return_sequences batch 错误", got=g.shape) expect(g[:, 0].tolist() == [7, 7, 7, 8, 8, 8], "repeat_interleave 顺序错误", got=g[:, 0].tolist()) g1d = m.generate(inputs=torch.tensor([1, 2, 3], device=DEVICE), max_new_tokens=2, do_sample=False, top_k=0, eos_token_id=None) expect(g1d.dim() == 2 and g1d.shape[0] == 1, "1D 输入未被正确归一化") gd = m.generate(inputs={"input_ids": ids([1, 2])}, max_new_tokens=2, do_sample=False, top_k=0, eos_token_id=None) expect(gd.shape[1] == 4, "dict 输入生成长度错误", got=gd.shape) @testcase("functional", "F10 训练前向反向（dense + MoE）") def f10(): for moe in (False, True): m = build_model(seed=11, eval_mode=False, use_moe=moe, num_experts=4, num_experts_per_tok=2) x = torch.randint(0, 128, (2, 16), device=DEVICE) out = m(x, labels=x) expect(torch.isfinite(out.loss).item(), f"MoE={moe} loss 非有限值", v=out.loss.item()) out.loss.backward() grads_ok = all(p.grad is not None and torch.isfinite(p.grad).all() for p in m.parameters() if p.requires_grad) expect(grads_ok, f"MoE={moe} 存在空/非有限梯度") m.eval() with torch.no_grad(): m(x) if moe: aux = m.model.layers[0].mlp.aux_loss.item() expect(abs(aux) 1") expect_raises(ValueError, lambda: m.generate(p, max_new_tokens=2, top_p=-0.1), "负 top_p") expect_raises(ValueError, lambda: m.generate(p, max_new_tokens=2, top_k=-3), "负 top_k") expect_raises(ValueError, lambda: m.generate(p, max_new_tokens=2, repetition_penalty=0.0), "penalty=0") expect_raises(ValueError, lambda: m.generate(p, max_new_tokens=2, repetition_penalty=-1.0), "负 penalty") expect_raises(ValueError, lambda: m.generate(p, max_new_tokens=-1), "负 max_new_tokens") expect_raises(ValueError, lambda: m.generate(p, max_new_tokens=2, num_return_sequences=0), "num_return_sequences=0") g = m.generate(p, max_new_tokens=2, top_k=99999, do_sample=False) expect(g.shape[1] == 3, "top_k clamp 后应正常生成", got=g.shape) @testcase("robustness", "R04 attention_mask 非法值/形状/左padding") def r04(): m = build_model() expect_raises(ValueError, lambda: m(ids([1, 2, 3, 4]), attention_mask=torch.tensor([[1, 1, 2, 0]], device=DEVICE)), "非二值 mask") expect_raises(ValueError, lambda: m(ids([1, 2, 3, 4]), attention_mask=torch.tensor([[1, 1, 1]], device=DEVICE)), "mask 长度不符") expect_raises(NotImplementedError, lambda: m.generate(input_ids=ids([0, 0, 1, 2]), attention_mask=torch.tensor([[0, 0, 1, 1]], device=DEVICE), max_new_tokens=2), "左 padding") @testcase("robustness", "R05 畸形 past_key_values") def r05(): m = build_model(num_hidden_layers=2) with torch.no_grad(): out = m(ids([1, 2, 3]), use_cache=True) expect_raises(ValueError, lambda: m(ids([4]), past_key_values=[out.past_key_values[0]], use_cache=True), "cache 层数不足") expect_raises(TypeError, lambda: m(ids([4]), past_key_values={"a": 1}, use_cache=True), "cache 类型错误") bad = [(torch.zeros(1), torch.zeros(1)), (None, None)] expect_raises(ValueError, lambda: m(ids([4]), past_key_values=bad, use_cache=True), "cache 层结构错误") @testcase("robustness", "R06 logits_to_keep / labels / position_ids 校验") def r06(): m = build_model(eval_mode=False) x = torch.randint(0, 128, (2, 8), device=DEVICE) expect_raises(ValueError, lambda: m(x, labels=x, logits_to_keep=1), "keep=1 + labels") expect_raises(ValueError, lambda: m(x, labels=x[:, :6]), "labels 形状不符") expect_raises(ValueError, lambda: m(x, logits_to_keep=-2), "负 keep") expect_raises(ValueError, lambda: m(x, logits_to_keep="all"), "非整数 keep") expect_raises(ValueError, lambda: m(x, position_ids=torch.zeros(2, 7, dtype=torch.long, device=DEVICE)), "position_ids 形状不符") @testcase("robustness", "R07 全 padding 行不产生 NaN") def r07(): x = ids([1, 2, 3, 4], [0, 0, 0, 0]) mask = torch.tensor([[1, 1, 1, 1], [0, 0, 0, 0]], device=DEVICE) for flash in (False, True): m2 = build_model(flash_attn=flash) with torch.no_grad(): out = m2(x, attention_mask=mask) expect(torch.isfinite(out.logits).all().item(), f"flash={flash} 全pad行出现 NaN/Inf") del m2 gc.collect() @testcase("robustness", "R08 对抗 prompt：超长重复/退化序列") def r08(): m = build_model() for tok in (0, 1, 127): x = torch.full((1, 512), tok, dtype=torch.long, device=DEVICE) with torch.no_grad(): out = m(x) expect(torch.isfinite(out.logits).all().item(), f"全 {tok} 退化序列 logits 异常") x = (torch.arange(1024, device=DEVICE) % 128).unsqueeze(0) g = m.generate(input_ids=x[:, :512], max_new_tokens=16, do_sample=False, top_k=0, eos_token_id=None) expect(torch.isfinite(m(g).logits).all().item(), "对抗序列生成后 logits 异常") @testcase("robustness", "R09 超 RoPE 长度自动扩展与告警") def r09(): import warnings m = build_model(max_position_embeddings=32) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") g = m.generate(input_ids=ids([1]), max_new_tokens=48, do_sample=False, top_k=0, eos_token_id=None) warned = any("max_position_embeddings" in str(rec.message) for rec in w) expect(g.shape[1] == 49, "超长生成失败", got=g.shape[1]) expect(warned, "越过 max_position_embeddings 应发出外推告警") @testcase("robustness", "R10 错误 dtype / 维度 / 类型输入") def r10(): m = build_model() expect_raises(TypeError, lambda: m(torch.randn(2, 4, device=DEVICE)), "浮点 input_ids") expect_raises(ValueError, lambda: m(torch.zeros(2, 2, 2, dtype=torch.long, device=DEVICE)), "3D input_ids") expect_raises(TypeError, lambda: m("not a tensor"), "字符串输入") # ============================================================ # 三、一致性回归测试 # ============================================================ GOLDEN_GREEDY = [1, 7, 42, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99] GOLDEN_SAMPLE = [1, 7, 42, 99, 99, 92, 50, 2, 107, 42, 20, 29, 121, 1] GOLDEN_PROMPT = [1, 7, 42, 99] def _golden_pair(seed=DEFAULT_SEED): mf = build_model(seed=seed, vocab_size=128, flash_attn=False) if HAS_BASELINE: mb = build_model(mod=baseline, seed=seed, vocab_size=128, flash_attn=False) clone_weights(mb, mf) else: mb = None return mf, mb @testcase("regression", "G01 golden 贪婪快照（防静默退化）") def g01(): m, _ = _golden_pair() g = m.generate(input_ids=ids(*GOLDEN_PROMPT), max_new_tokens=12, do_sample=False, top_k=0, eos_token_id=None) expect(g[0].tolist() == GOLDEN_GREEDY, "贪婪 golden 快照不匹配", got=g[0].tolist()) @testcase("regression", "G02 golden 采样快照（固定种子）") def g02(): m, _ = _golden_pair() with seed_context(777): g = m.generate(input_ids=ids(*GOLDEN_PROMPT), max_new_tokens=10, do_sample=True, temperature=0.8, top_k=50, top_p=0.9, eos_token_id=None) expect(g[0].tolist() == GOLDEN_SAMPLE, "采样 golden 快照不匹配", got=g[0].tolist()) @testcase("regression", "G03 修复版 vs 基线版：贪婪解码逐 token 一致") def g03(): expect(HAS_BASELINE, "基线模块 model_minimind_baseline 不可用，跳过") mf, mb = _golden_pair() p = ids(*GOLDEN_PROMPT) kw = dict(max_new_tokens=16, do_sample=False, top_k=0, eos_token_id=None) with torch.no_grad(): gf = mf.generate(input_ids=p, **kw) gb = mb.generate(input_ids=p, **kw) expect(torch.equal(gf, gb), "基线/修复版贪婪结果不一致", fixed=gf[0].tolist(), base=gb[0].tolist()) @testcase("regression", "G04 修复版 vs 基线版：同种子采样一致") def g04(): expect(HAS_BASELINE, "基线模块不可用，跳过") mf, mb = _golden_pair() p = ids(*GOLDEN_PROMPT) kw = dict(max_new_tokens=12, do_sample=True, temperature=0.9, top_k=40, top_p=0.8, eos_token_id=None) with seed_context(2024): gf = mf.generate(input_ids=p, **kw) with seed_context(2024): gb = mb.generate(input_ids=p, **kw) expect(torch.equal(gf, gb), "基线/修复版采样结果不一致", fixed=gf[0].tolist(), base=gb[0].tolist()) @testcase("regression", "G05 预填充 vs 逐 token 增量 logits 一致") def g05(): from transformers.cache_utils import DynamicCache m = build_model() seq = torch.randint(0, 128, (1, 24), device=DEVICE) cut = 10 with torch.no_grad(): cache_full = DynamicCache() full = m(seq, past_key_values=cache_full, use_cache=True).logits[0] cache_step = DynamicCache() prefix = m(seq[:, :cut], past_key_values=cache_step, use_cache=True).logits[0] outs = [prefix] for t in seq[0, cut:]: outs.append(m(t.reshape(1, 1), past_key_values=cache_step, use_cache=True).logits[0]) stepped = torch.cat(outs, dim=0) diff = (full - stepped).abs().max().item() expect(diff 1e-3, "基线版未复现位置错位（对照失效）", d=d_base) # ============================================================ # 四、性能基准测试 # ============================================================ def _median_time(fn, repeats=3, warmup=1): for _ in range(warmup): fn() if DEVICE == "cuda": torch.cuda.synchronize() ts = [] for _ in range(repeats): t0 = time.perf_counter() fn() if DEVICE == "cuda": torch.cuda.synchronize() ts.append(time.perf_counter() - t0) return statistics.median(ts) @testcase("performance", "P01 预填充吞吐 / 解码延迟 / cache 收益 / KV 占用") def p01(quick=False): m = build_model(seed=7, hidden_size=256, num_hidden_layers=4, num_attention_heads=8, num_key_value_heads=2, vocab_size=2048, max_position_embeddings=8192) n_layers, kv_heads, head_dim = 4, 2, 32 bytes_per = 2 if DEVICE == "cuda" else 4 rows = [] for L in ([256, 512] if quick else [256, 512, 1024]): x = torch.randint(0, 2048, (1, L), device=DEVICE) t = _median_time(lambda: m(x), repeats=2, warmup=1) rows.append(("prefill", L, f"{t*1000:.1f} ms", f"{L/t:,.0f} tok/s")) del x L, n_dec = 256, (16 if quick else 48) p = torch.randint(0, 2048, (1, L), device=DEVICE) def decode_with_cache(): from transformers.cache_utils import DynamicCache c = DynamicCache() with torch.no_grad(): m(p, past_key_values=c, use_cache=True) for _ in range(n_dec): m(torch.randint(0, 2048, (1, 1), device=DEVICE), past_key_values=c, use_cache=True) return c def decode_no_cache(): with torch.no_grad(): cur = p for _ in range(n_dec): cur = torch.cat([cur, torch.randint(0, 2048, (1, 1), device=DEVICE)], dim=-1) m(cur, use_cache=False) c = decode_with_cache() t_cache = _median_time(decode_with_cache, repeats=2, warmup=1) t_nocache = _median_time(decode_no_cache, repeats=2, warmup=1) rows.append(("decode cache", n_dec, f"{t_cache/n_dec*1000:.2f} ms/tok", f"{n_dec/t_cache:,.0f} tok/s")) rows.append(("decode no-cache", n_dec, f"{t_nocache/n_dec*1000:.2f} ms/tok", f"{n_dec/t_nocache:,.0f} tok/s")) expect(t_cache math.log(128) - 0.5, "高温应接近最大熵 ln(V)", e=es[-1]) @testcase("sampling", "S03 top-k 支撑集与定义逐元素一致") def s03(): torch.manual_seed(2) logits = torch.randn(4, 100) for k in (1, 5, 30, 100): p = fixed.warp_logits(logits.clone(), 1.0, k, 1.0) for r in range(4): kept = p[r] > 0 n_kept = int(kept.sum()) expect(n_kept = p_val - 1e-5, "保留质量低于 top_p", p=p_val, mass=kept_mass) nz = sp[r][sp[r] > 0] if nz.numel() > 1: expect(kept_mass - nz[-1].item() = stats[0.1], "高温生成 token 多样性未高于低温", stats=stats) # ============================================================ # 运行入口 # ============================================================ GROUPS = ["functional", "robustness", "regression", "performance", "sampling"] QUICK_AWARE = { "P01 预填充吞吐 / 解码延迟 / cache 收益 / KV 占用", "P02 预填充 logits_to_keep=1 加速 lm_head", "S05 大样本频率拟合理论分布（无偏性）", "S09 端到端分布：高温生成多样性高于低温", } def main(): ap = argparse.ArgumentParser(description="MiniMind 多维测试套件") ap.add_argument("groups", nargs="*", help="要执行的类别：" + "/".join(GROUPS)) ap.add_argument("--skip", nargs="*", default=[], help="跳过的类别") ap.add_argument("--list", action="store_true", help="列出全部用例后退出") ap.add_argument("--quick", action="store_true", help="快速模式（缩短基准/统计规模）") ap.add_argument("--seed", type=int, default=DEFAULT_SEED) args = ap.parse_args() if args.list: for g in GROUPS: print(f"[{g}]") for gg, name, _ in REGISTRY: if gg == g: print(f" - {name}") return 0 selected = args.groups if args.groups else GROUPS for g in list(selected) + list(args.skip): if g not in GROUPS: print(f"未知测试类别：{g}；可选 {GROUPS}") return 2 todo = [(g, n, f) for g, n, f in REGISTRY if g in selected and g not in args.skip] print("=" * 74) print(f"MiniMind 测试套件 | device={DEVICE} | threads={TORCH_THREADS} | " f"seed={args.seed} | quick={args.quick}") print(f"基线模块可用：{HAS_BASELINE} | 用例数：{len(todo)}") print("=" * 74) passed, failed = 0, 0 t_start = time.perf_counter() current_group = None for g, name, fn in todo: if g != current_group: current_group = g print(f"\n{'─'*30} {g} {'─'*30}") t0 = time.perf_counter() try: if name in QUICK_AWARE: fn(quick=args.quick) else: fn() dt = time.perf_counter() - t0 print(f" [√] {name} ({dt:.2f}s)") passed += 1 except CheckError as e: dt = time.perf_counter() - t0 print(f" [x] {name} ({dt:.2f}s)\n {e}") failed += 1 except Exception: dt = time.perf_counter() - t0 print(f" ! {name} ({dt:.2f}s) 抛出异常:") for line in traceback.format_exc().splitlines()[-4:]: print(f" {line}") failed += 1 gc.collect() total = time.perf_counter() - t_start print("\n" + "=" * 74) print(f"汇总：√ {passed} 通过 / X {failed} 失败 / 总耗时 {total:.1f}s") if failed == 0: print("全部用例通过") print("=" * 74) return 1 if failed else 0 if __name__ == "__main__": sys.exit(main()) ``` 运行命令行输出测试效果： ``` python .\main_baseline.py ========================================================================== 测试套件 | device=cpu | threads=4 | seed=20260919 | quick=False 基线模块可用：True | 用例数：41 ========================================================================== ────────────────────────────── functional ────────────────────────────── [√] F01 多batch/多序列长度前向形状 (0.05s) [√] F02 多语言 Unicode 字符级通路 (0.24s) [√] F03 极短/超长 prompt (0.31s) [x] F04 EOS 终止逻辑（单EOS/多EOS/禁用/错峰结束） (0.03s) 错峰 EOS 停止位置错误 got=7 [√] F05 流式生成（streamer） (0.03s) [√] F06 batch 生成与逐条生成一致（贪婪） (0.09s) [√] F07 右 padding 批量生成等价于无 padding 单条 (0.09s) [√] F08 use_cache True/False 输出一致 (0.06s) [√] F09 return_kv / num_return_sequences / 1D / dict 输入 (0.04s) [√] F10 训练前向反向（dense + MoE） (0.48s) [√] F11 YaRN 长上下文缩放 (0.30s) ────────────────────────────── robustness ────────────────────────────── [√] R01 空输入被拒绝 (0.02s) [√] R02 token id 越界（上界/下界） (0.02s) [√] R03 非法采样参数矩阵 (0.02s) [√] R04 attention_mask 非法值/形状/左padding (0.02s) [√] R05 畸形 past_key_values (0.02s) [√] R06 logits_to_keep / labels / position_ids 校验 (0.02s) [√] R07 全 padding 行不产生 NaN (0.41s) ! R08 对抗 prompt：超长重复/退化序列 (0.16s) 抛出异常: [√] R09 超 RoPE 长度自动扩展与告警 (0.13s) [√] R10 错误 dtype / 维度 / 类型输入 (0.02s) ────────────────────────────── regression ────────────────────────────── [√] G01 golden 贪婪快照（防静默退化） (0.06s) [√] G02 golden 采样快照（固定种子） (0.06s) [√] G03 修复版 vs 基线版：贪婪解码逐 token 一致 (0.11s) [√] G04 修复版 vs 基线版：同种子采样一致 (0.11s) [√] G05 预填充 vs 逐 token 增量 logits 一致 (0.07s) [√] G06 SDPA 与手写注意力一致 (0.04s) Writing model shards: 100%|██████████████████████████████████████████████████████████████████████████████████ ███████████████████████████████████████████████████| 1/1 [00:00<00:00, 112.85it/s] Loading weights: 100%|█████████████████████████████████████████████████████████████████████████████████████ ██████████████████████████████████████████████████| 32/32 [00:00<00:00, 5710.42it/s] ! G07 save_pretrained / from_pretrained 输出一致 (0.15s) 抛出异常: [x] G08 MoE eval 确定性与 checksum (0.04s) MoE logits checksum 偏离 golden s=26.449787139892578 | mx=1.3880141973495483 [√] G09 基线版右padding位置错位对照（证明回归测试有牙齿） (0.06s) ────────────────────────────── performance ────────────────────────────── +-----------------+----------+------------------+------------------+ | 阶段 | token 数 | 耗时 | 吞吐/占用 | +-----------------+----------+------------------+------------------+ | prefill | 256 | 30.6 ms | 8,374 tok/s | | prefill | 512 | 67.0 ms | 7,644 tok/s | | prefill | 1024 | 330.1 ms | 3,102 tok/s | | decode cache | 48 | 6.94 ms/tok | 144 tok/s | | decode no-cache | 48 | 23.05 ms/tok | 43 tok/s | | KV cache | 304 | theo=608.0 KiB | actual=608.0 KiB | | params | - | 3.74 M | 14.3 MiB(权重) | | 进程 RSS | - | 479.3 MiB | | +-----------------+----------+------------------+------------------+ [√] P01 预填充吞吐 / 解码延迟 / cache 收益 / KV 占用 (6.10s) 全量 logits 54.3 ms / keep=1 36.1 ms / 加速 1.51x / 末位 logits diff=7.15e-07 [√] P02 预填充 logits_to_keep=1 加速 lm_head (0.65s) ────────────────────────────── sampling ────────────────────────────── [√] S01 temperature 缩放严格等于 softmax(logits/T) (0.01s) [√] S02 熵随温度单调（低温尖锐/高温平坦） (0.00s) [√] S03 top-k 支撑集与定义逐元素一致 (0.01s) [√] S04 top-p(nucleus) 累积质量边界与最小集合 (0.01s) [√] S05 大样本频率拟合理论分布（无偏性） (0.10s) [√] S06 采样端到端：top-k 支撑集 / 贪婪=argmax (0.05s) [√] S07 temperature=0 强制贪婪；generator 可复现 (0.11s) [√] S08 重复惩罚公式方向与数值精确性 (0.00s) [√] S09 端到端分布：高温生成多样性高于低温 (3.51s) ========================================================================== 汇总：√ 37 通过 / X 4 失败 / 总耗时 21.6s ========================================================================== ```
