大模型网络结构：模型接口测试用例参考
大模型的网络文件是模型训练的核心参考。例如 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 < 1e-4, "cache format output diverge too large", max_diff=diff)
print("[√] test_4_kv_cache_compatibility PASS")
def test_5_tie_word_embedding():
"""权重绑定验证"""
print_sep("test_5_tie_word_embedding")
cfg_tie = MiniMindConfig(hidden_size=192, num_hidden_layers=2, vocab_size=512, tie_word_embeddings=True, use_moe=False)
m1 = MiniMindForCausalLM(cfg_tie)
print(f"tie=True: lm_head.weight[0,:5] = {m1.lm_head.weight[0,:5]}")
m1.lm_head.weight.data[0] += 0.1
print(f"after modify lm_head, embed_tokens[0,:5] = {m1.model.embed_tokens.weight[0,:5]}")
assert_with_info(torch.allclose(m1.lm_head.weight[0], m1.model.embed_tokens.weight[0]),
"tie weight not sync")
cfg_no_tie = MiniMindConfig(hidden_size=192, num_hidden_layers=2, vocab_size=512, tie_word_embeddings=False, use_moe=False)
m2 = MiniMindForCausalLM(cfg_no_tie)
eq = torch.equal(m2.lm_head.weight, m2.model.embed_tokens.weight)
print(f"tie=False: lm_head and embed_tokens equal? {eq}")
assert_with_info(not eq, "tie=False but weight still shared")
print("[√] test_5_tie_word_embedding PASS")
def test_6_moe_advanced():
"""MoE：loss计算、梯度检查、eval模式aux_loss"""
print_sep("test_6_moe_advanced")
cfg = MiniMindConfig(
hidden_size=256, num_hidden_layers=3,
num_attention_heads=4, num_key_value_heads=2,
vocab_size=800,
use_moe=True,
num_experts=4,
num_experts_per_tok=2,
norm_topk_prob=True,
router_aux_loss_coef=1e-3,
router_z_loss_coef=1e-3
)
model = MiniMindForCausalLM(cfg).train()
B, S = 4, 16
input_ids = torch.randint(0, cfg.vocab_size, (B, S))
labels = input_ids.clone()
print(f"input shape B={B}, S={S}")
out = model(input_ids=input_ids, labels=labels)
total_loss = out.loss
print(f"train total loss(ce+aux+z) = {total_loss.item():.6f}")
total_loss.backward()
grad_check_ok = True
for li, layer in enumerate(model.model.layers):
moeff = layer.mlp
for ei, exp in enumerate(moeff.experts):
for p in exp.parameters():
if p.grad is None:
print(f"[!] layer{li} expert{ei} param grad is None!")
grad_check_ok = False
if moeff.gate.weight.grad is None:
print(f"[!] layer{li} router gate grad None!")
grad_check_ok = False
assert_with_info(grad_check_ok, "MoE grad check failed")
print("[√] all experts & router have grad")
model.eval()
with torch.no_grad():
out_eval = model(input_ids)
aux_loss_layer0 = model.model.layers[0].mlp.aux_loss
print(f"eval mode aux_loss layer0: {aux_loss_layer0.item()}")
assert abs(aux_loss_layer0.item()) < 1e-8
print("[√] test_6_moe_advanced PASS")
def test_7_generate_decoding_modes():
"""多种解码策略验证"""
print_sep("test_7_generate_decoding_modes")
cfg = MiniMindConfig(
hidden_size=192, num_hidden_layers=2,
num_attention_heads=3, num_key_value_heads=1,
vocab_size=400, eos_token_id=2, use_moe=False
)
model = MiniMindForCausalLM(cfg).eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
prompt = torch.tensor([[100]], device=device)
print(f"prompt shape {prompt.shape}")
g1 = model.generate(input_ids=prompt, max_new_tokens=12, do_sample=False)
print(f"greedy output shape {g1.shape}")
assert_with_info(g1.shape[-1] == 1 + 12, "greedy length error")
g2 = model.generate(input_ids=prompt, max_new_tokens=8, do_sample=True, temperature=0.7, top_k=20)
print(f"top-k sample output shape {g2.shape}")
assert_with_info(g2.shape[-1] <= 1 + 8, "top-k length error")
g3 = model.generate(input_ids=prompt, max_new_tokens=8, do_sample=True, temperature=0.7, top_p=0.6)
print(f"top-p sample output shape {g3.shape}")
assert_with_info(g3.shape[-1] <= 1 + 8, "top-p length error")
g4 = model.generate(input_ids=prompt, max_new_tokens=8, repetition_penalty=1.2, repetition_window=32)
print(f"repetition_penalty output shape {g4.shape}")
assert_with_info(g4.shape[-1] <= 1 + 8, "rep-penalty length error")
prompt_eos = torch.tensor([[cfg.eos_token_id]], device=device)
g5 = model.generate(input_ids=prompt_eos, max_new_tokens=20)
print(f"eos-prompt output shape {g5.shape}, expected < {1+20}")
assert_with_info(g5.shape[-1] < 1 + 20, "eos should early-stop")
print("[√] test_7_generate_decoding_modes PASS")
def test_8_mixed_precision_fp16_bf16(skip_gpu_case: bool = False):
"""混合精度验证"""
print_sep("test_8_mixed_precision_fp16_bf16")
cfg = MiniMindConfig(
hidden_size=256, num_hidden_layers=2,
num_attention_heads=4, num_key_value_heads=2,
vocab_size=600, use_moe=False
)
dev = "cuda" if torch.cuda.is_available() else "cpu"
if skip_gpu_case or dev == "cpu":
print(">> 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 < 1.0, "save-load weight changed too much", diff=diff)
del model1, model2
import gc
gc.collect()
if os.path.exists(tmp_dir):
shutil.rmtree(tmp_dir)
print("[√] test_9_model_save_load_hf_style PASS")
def test_10_kv_cache_step_by_step_equivalence():
"""核心回归：全量前向 vs 逐token增量解码一致性"""
print_sep("test_10_kv_cache_step_by_step_equivalence")
cfg = MiniMindConfig(
hidden_size=256, num_hidden_layers=2,
num_attention_heads=4, num_key_value_heads=2,
vocab_size=800, use_moe=False
)
model = MiniMindForCausalLM(cfg).eval()
dev = "cuda" if torch.cuda.is_available() else "cpu"
model.to(dev)
full_seq = torch.randint(0, cfg.vocab_size, (1, 16), device=dev)
print(f"full test sequence shape {full_seq.shape}")
cache_full = DynamicCache()
with torch.no_grad():
out_full = model(full_seq, past_key_values=cache_full, use_cache=True)
logits_full = out_full.logits[0]
prefix = full_seq[:, :8]
remain_tokens = full_seq[:, 8:].squeeze(0)
cache_step = DynamicCache()
with torch.no_grad():
out_prefix = model(prefix, past_key_values=cache_step, use_cache=True)
collected = [out_prefix.logits[0]]
for idx, tok in enumerate(remain_tokens):
inp_tok = tok.reshape(1, 1)
with torch.no_grad():
o = model(inp_tok, past_key_values=cache_step, use_cache=True)
collected.append(o.logits[0])
print(f" incremental step {idx}, input token id={tok.item()}, logits shape {o.logits.shape}")
logits_step = torch.cat(collected, dim=0)
max_diff = torch.max(torch.abs(logits_full - logits_step)).item()
print(f"max abs diff full-prefill vs step-by-step = {max_diff:.2e}")
assert_with_info(max_diff < 1e-3, "KV-cache incremental output mismatch", max_diff=max_diff)
print("[√] test_10_kv_cache_step_by_step_equivalence PASS")
def test_11_logits_to_keep_slice():
"""logits_to_keep 尾部切片与loss计算"""
print_sep("test_11_logits_to_keep_slice")
cfg = MiniMindConfig(hidden_size=256, num_hidden_layers=2, vocab_size=600, use_moe=False)
model = MiniMindForCausalLM(cfg).train()
B, S = 2, 20
ids = torch.randint(0, cfg.vocab_size, (B, S))
labels = ids.clone()
keep = 4
print(f"seq_len={S}, logits_to_keep={keep}")
out = model(input_ids=ids, labels=labels, logits_to_keep=keep)
loss = out.loss
print(f"loss = {loss.item():.6f}")
loss.backward()
assert_with_info(model.model.embed_tokens.weight.grad is not None, "embed grad None")
print(f"[√] logits_to_keep={keep} loss & grad ok")
print("[√] test_11_logits_to_keep_slice PASS")
def test_12_tie_weights_method():
"""tie_weights() 方法兼容性"""
print_sep("test_12_tie_weights_method")
cfg = MiniMindConfig(hidden_size=192, num_hidden_layers=2, vocab_size=512, tie_word_embeddings=True, use_moe=False)
model = MiniMindForCausalLM(cfg)
model.tie_weights()
eq = torch.equal(model.lm_head.weight, model.model.embed_tokens.weight)
print(f"call tie_weights(), weight shared = {eq}")
assert_with_info(eq, "tie_weights() broke weight sharing")
cfg_no_tie = MiniMindConfig(hidden_size=192, num_hidden_layers=2, vocab_size=512, tie_word_embeddings=False, use_moe=False)
m2 = MiniMindForCausalLM(cfg_no_tie)
m2.tie_weights()
eq2 = torch.equal(m2.lm_head.weight, m2.model.embed_tokens.weight)
print(f"tie_word_embeddings=False, shared={eq2}")
assert_with_info(not eq2, "tie=False should not share weights")
print("[√] test_12_tie_weights_method PASS")
def test_13_past_key_values_none_and_empty_list():
"""past_key_values=None / 空列表边界输入"""
print_sep("test_13_past_key_values_none_and_empty_list")
cfg = MiniMindConfig(hidden_size=256, num_hidden_layers=2, vocab_size=600, use_moe=False)
model = MiniMindForCausalLM(cfg).eval()
dev = "cuda" if torch.cuda.is_available() else "cpu"
model.to(dev)
ids = torch.randint(0, cfg.vocab_size, (1, 8), device=dev)
with torch.no_grad():
out_none = model(ids, past_key_values=None, use_cache=True)
print(f"past_key_values=None ok, past type {type(out_none.past_key_values)}")
with torch.no_grad():
out_empty = model(ids, past_key_values=[], use_cache=True)
print(f"past_key_values=[] empty-list ok")
assert out_none.logits.shape == (1, 8, cfg.vocab_size)
assert out_empty.logits.shape == (1, 8, cfg.vocab_size)
print("[√] test_13_past_key_values_none_and_empty_list PASS")
def test_14_moe_topk_1():
"""MoE top-k=1 路径验证"""
print_sep("test_14_moe_topk_1")
# 固定随机种子：top-1 路由下小批量可能让某个专家完全未被选中（无梯度），
# 加种子保证测试可复现、稳定通过
torch.manual_seed(0)
cfg = MiniMindConfig(
hidden_size=256, num_hidden_layers=2,
num_attention_heads=4, num_key_value_heads=2,
vocab_size=800, use_moe=True,
num_experts=4, num_experts_per_tok=1
)
model = MiniMindForCausalLM(cfg)
B, S = 2, 10
x = torch.randint(0, cfg.vocab_size, (B, S))
model.train()
out_train = model(x, labels=x)
loss = out_train.loss
loss.backward()
for layer in model.model.layers:
for exp in layer.mlp.experts:
for p in exp.parameters():
assert p.grad is not None
print("[√] MoE top-k=1 train grad ok")
model.eval()
with torch.no_grad():
out_eval = model(x)
aux = model.model.layers[0].mlp.aux_loss
print(f"eval aux_loss item: {aux.item()}")
assert abs(aux.item()) < 1e-8
print("[√] test_14_moe_topk_1 PASS")
def test_15_attention_non_full_mask_no_flash():
"""关闭flash与开启flash结果一致性验证"""
print_sep("test_15_attention_non_full_mask_no_flash")
cfg_flash_off = MiniMindConfig(
hidden_size=256, num_hidden_layers=1,
num_attention_heads=4, num_key_value_heads=2,
vocab_size=500, use_moe=False,
flash_attn=False
)
model = MiniMindForCausalLM(cfg_flash_off).eval()
dev = "cuda" if torch.cuda.is_available() else "cpu"
model.to(dev)
input_ids = torch.tensor([[10, 11, 12, 0, 0]], device=dev)
attn_mask = torch.tensor([[1, 1, 1, 0, 0]], device=dev)
with torch.no_grad():
out_no_flash = model(input_ids, attention_mask=attn_mask)
print(f"no-flash non-full mask logits shape {out_no_flash.logits.shape}")
cfg_flash_on_dict = cfg_flash_off.to_dict()
cfg_flash_on_dict["flash_attn"] = True # to_dict() 已含 flash_attn 键，需原地修改避免重复传参
cfg_flash_on = MiniMindConfig(**cfg_flash_on_dict)
model2 = MiniMindForCausalLM(cfg_flash_on).eval().to(dev)
# 复用同一份权重：两个模型仅 flash_attn 开关不同，否则随机初始化不同会导致输出不可比
model2.load_state_dict(model.state_dict())
with torch.no_grad():
out_flash = model2(input_ids, attention_mask=attn_mask)
diff = torch.max(torch.abs(out_no_flash.logits - out_flash.logits)).item()
print(f"flash vs no-flash logits max diff: {diff:.3e}")
assert diff < 1e-3
print("[√] test_15_attention_non_full_mask_no_flash PASS")
def test_16_rope_no_yarn():
"""原始RoPE分支（无YaRN缩放）"""
print_sep("test_16_rope_no_yarn")
cfg = MiniMindConfig(
hidden_size=256, num_hidden_layers=2,
num_attention_heads=4, num_key_value_heads=2,
vocab_size=600, use_moe=False,
inference_rope_scaling=False,
rope_scaling=None
)
model = MiniMindForCausalLM(cfg).eval()
dev = "cuda" if torch.cuda.is_available() else "cpu"
model.to(dev)
L = 1024
ids = torch.randint(0, cfg.vocab_size, (1, L), device=dev)
with torch.no_grad():
out = model(ids)
print(f"no-yarn RoPE run ok, logits shape {out.logits.shape}")
assert out.logits.shape == (1, L, cfg.vocab_size)
print("[√] test_16_rope_no_yarn PASS")
def test_17_generate_edge_cases():
"""generate 边界参数验证"""
print_sep("test_17_generate_edge_cases")
cfg = MiniMindConfig(hidden_size=192, num_hidden_layers=1, vocab_size=300, eos_token_id=2, use_moe=False)
model = MiniMindForCausalLM(cfg).eval()
dev = "cuda" if torch.cuda.is_available() else "cpu"
model.to(dev)
prompt = torch.tensor([[10]], device=dev)
g1 = model.generate(input_ids=prompt, max_new_tokens=4, num_return_sequences=2, do_sample=True)
print(f"num_return_sequences=2 shape {g1.shape}")
assert g1.shape[0] == 2
g2 = model.generate(input_ids=prompt, max_new_tokens=3, temperature=0.0, top_k=0, top_p=1.0, do_sample=False)
print(f"temperature=0, no sampling constraints shape {g2.shape}")
ret_dict = model.generate(input_ids=prompt, max_new_tokens=2, return_kv=True, do_sample=False)
print(f"return_kv is dict: {isinstance(ret_dict, dict)}, keys={list(ret_dict.keys())}")
assert "generated_ids" in ret_dict and "past_kv" in ret_dict
g4 = model.generate(input_ids=prompt, max_new_tokens=5, eos_token_id=None, do_sample=False)
print(f"eos_token_id=None output len={g4.shape[-1]} expect {1+5}")
assert g4.shape[-1] == 1 + 5
print("[√] test_17_generate_edge_cases PASS")
def test_18_use_cache_false():
"""use_cache=False 关闭KV缓存"""
print_sep("test_18_use_cache_false")
cfg = MiniMindConfig(hidden_size=256, num_hidden_layers=2, vocab_size=600, use_moe=False)
model = MiniMindForCausalLM(cfg).eval()
dev = "cuda" if torch.cuda.is_available() else "cpu"
model.to(dev)
ids = torch.randint(0, cfg.vocab_size, (1, 16), device=dev)
with torch.no_grad():
out_no_cache = model(ids, use_cache=False)
print(f"use_cache=False past_key_values = {out_no_cache.past_key_values}")
assert out_no_cache.past_key_values is None
assert out_no_cache.logits.shape == (1, 16, cfg.vocab_size)
print("[√] test_18_use_cache_false PASS")
def test_19_logits_to_keep_zero():
"""logits_to_keep=0 默认取全部序列"""
print_sep("test_19_logits_to_keep_zero")
cfg = MiniMindConfig(hidden_size=256, num_hidden_layers=2, vocab_size=600, use_moe=False)
model = MiniMindForCausalLM(cfg).train()
B, S = 2, 12
ids = torch.randint(0, cfg.vocab_size, (B, S))
labels = ids.clone()
out = model(input_ids=ids, labels=labels, logits_to_keep=0)
loss = out.loss
print(f"logits_to_keep=0 loss={loss.item():.4f}")
loss.backward()
assert model.model.embed_tokens.weight.grad is not None
print("[√] test_19_logits_to_keep_zero PASS")
def test_20_custom_head_dim_and_moe_intermediate():
"""自定义head_dim与moe_intermediate_size"""
print_sep("test_20_custom_head_dim_and_moe_intermediate")
cfg = MiniMindConfig(
hidden_size=256, num_hidden_layers=2,
num_attention_heads=4, num_key_value_heads=2,
head_dim=96,
vocab_size=600, use_moe=True,
num_experts=2,
moe_intermediate_size=512
)
assert cfg.head_dim == 96
assert cfg.moe_intermediate_size == 512
model = MiniMindForCausalLM(cfg).eval()
dev = "cuda" if torch.cuda.is_available() else "cpu"
model.to(dev)
ids = torch.randint(0, cfg.vocab_size, (1, 8), device=dev)
with torch.no_grad():
out = model(ids)
print(f"custom head_dim={cfg.head_dim}, moe_intermediate_size={cfg.moe_intermediate_size}, logits shape {out.logits.shape}")
assert out.logits.shape == (1, 8, cfg.vocab_size)
print("[√] test_20_custom_head_dim_and_moe_intermediate PASS")
def test_21_minites():
"""验证模型训练循环：Loss下降、参数更新、梯度清零"""
print_sep("test_21_minites")
cfg = MiniMindConfig(hidden_size=256, num_hidden_layers=2, vocab_size=1000, use_moe=False)
model = MiniMindForCausalLM(cfg)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3) # 适当调大学习率以便快速收敛
# 构造一个有规律的序列，模型应该很容易学会预测
seq_len = 16
batch_size = 4
# 规律: 1,2,3,4 重复
template = torch.tensor([1,2,3,4,1,2,3,4,1,2,3,4,1,2,3,4])
x = template.unsqueeze(0).repeat(batch_size, 1)
input_ids = x[:, :-1]
labels = x[:, 1:]
# 记录初始参数快照
first_param_snapshot = model.model.layers[0].mlp.gate_proj.weight.data.clone()
losses = []
for step in range(20):
out = model(input_ids, labels=labels)
loss = out.loss
losses.append(loss.item())
opt.zero_grad()
loss.backward()
# 检查梯度是否存在
assert model.model.layers[0].mlp.gate_proj.weight.grad is not None, "Grad is None"
opt.step()
print(f"step {step}, loss: {loss.item():.4f}")
# --- 核心断言 ---
# 1. Loss 必须下降 (允许一定的波动，但总体趋势要降)
assert losses[-1] < losses[0], f"Loss did not decrease: {losses[0]} -> {losses[-1]}"
assert losses[-1] < 2.0, f"Final loss too high: {losses[-1]}" # 对于简单规律，loss应很低
# 2. 参数必须更新
last_param_snapshot = model.model.layers[0].mlp.gate_proj.weight.data
param_diff = torch.norm(first_param_snapshot - last_param_snapshot).item()
assert param_diff > 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() < cfg.vocab_size, "token id >= vocab_size"
assert seq.min().item() >= 0, "token id < 0"
assert seq.shape[-1] <= prompt.shape[-1] + max_new, \
f"length {seq.shape[-1]} exceeds prompt({prompt.shape[-1]}) + max_new({max_new})"
# ---------- Case 2: 返回结构一致性 ----------
assert _has_scores(out) or True, \
"scores 未实现时请显式返回 None，不要抛 KeyError"
# ---------- Case 3: EOS 截断 ----------
# 构造"必出 EOS"的场景：把 lm_head 偏置压向 eos_token_id
# 先检查 lm_head 是否有 bias 参数
lm_head = model.lm_head
has_bias = lm_head.bias is not None
if has_bias:
bias_backup = lm_head.bias.data.clone()
lm_head.bias.data.fill_(-20.0)
lm_head.bias.data[2] = 20.0 # 强行让 token 2 (eos) 概率最大
else:
# 没有 bias 时，临时给 lm_head 加上 bias
original_bias = lm_head.bias
lm_head.bias = torch.nn.Parameter(torch.zeros(lm_head.out_features, device=lm_head.weight.device))
lm_head.bias.data.fill_(-20.0)
lm_head.bias.data[2] = 20.0
out_eos = model.generate(input_ids=prompt, max_new_tokens=20, do_sample=False)
seq_eos = _pick_seq(out_eos)
new_part = seq_eos[0, prompt.shape[-1]:] # 只看新生成的部分
assert new_part.numel() > 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: <class 'transformers.cache_utils.DynamicCache'>
list-kv output past_key_values type: <class 'list'>
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=<SliceBackward0>)
after modify lm_head, embed_tokens[0,:5] = tensor([0.1322, 0.0784, 0.1046, 0.1141, 0.1100], grad_fn=<SliceBackward0>)
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 < 21
[√] test_7_generate_decoding_modes PASS
======================================================================
【test_8_mixed_precision_fp16_bf16】
======================================================================
>> 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<00:00, 75.01it/s]
[√] save_pretrained finished
Loading weights: 100%|█████████████████████████████████████████████████████████████████████████████████████
██████████████████████████████████████████████████| 32/32 [00:00<00:00, 3534.00it/s]
logits max abs diff after reload = 0.00e+00
[√] test_9_model_save_load_hf_style PASS
======================================================================
【test_10_kv_cache_step_by_step_equivalence】
======================================================================
full test sequence shape torch.Size([1, 16])
incremental step 0, input token id=333, logits shape torch.Size([1, 1, 800])
incremental step 1, input token id=230, logits shape torch.Size([1, 1, 800])
incremental step 2, input token id=649, logits shape torch.Size([1, 1, 800])
incremental step 3, input token id=306, logits shape torch.Size([1, 1, 800])
incremental step 4, input token id=459, logits shape torch.Size([1, 1, 800])
incremental step 5, input token id=485, logits shape torch.Size([1, 1, 800])
incremental step 6, input token id=196, logits shape torch.Size([1, 1, 800])
incremental step 7, input token id=759, logits shape torch.Size([1, 1, 800])
max abs diff full-prefill vs step-by-step = 1.01e-06
[√] test_10_kv_cache_step_by_step_equivalence PASS
======================================================================
【test_11_logits_to_keep_slice】
======================================================================
seq_len=20, logits_to_keep=4
loss = 6.404029
[√] logits_to_keep=4 loss & grad ok
[√] test_11_logits_to_keep_slice PASS
======================================================================
【test_12_tie_weights_method】
======================================================================
call tie_weights(), weight shared = True
tie_word_embeddings=False, shared=False
[√] test_12_tie_weights_method PASS
======================================================================
【test_13_past_key_values_none_and_empty_list】
======================================================================
past_key_values=None ok, past type <class 'list'>
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 = {"<pad>": 0, "<bos>": 1, "<eos>": 2, "<unk>": 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) < 128, "字符表超过测试模型 vocab_size", v=len(vocab))
m = build_model(vocab_size=128)
inv = {v: k for k, v in vocab.items()}
for lang, text in texts.items():
enc = torch.tensor([[vocab[c] for c in text]], device=DEVICE)
with torch.no_grad():
g1 = m.generate(input_ids=enc, max_new_tokens=8, do_sample=False,
top_k=0, eos_token_id=None)
g2 = m.generate(input_ids=enc, max_new_tokens=8, do_sample=False,
top_k=0, eos_token_id=None)
expect(torch.equal(g1, g2), f"{lang} 贪婪生成不可复现")
expect(g1.shape[1] == len(text) + 8, f"{lang} 生成长度错误", got=g1.shape[1])
roundtrip = "".join(inv[i] for i in g1[0, :len(text)].tolist())
expect(roundtrip == text, f"{lang} 字符映射往返不一致", got=roundtrip)
expect(torch.isfinite(m(enc).logits).all().item(), f"{lang} logits 非有限值")
@testcase("functional", "F03 极短/超长 prompt")
def f03():
m = build_model()
g = m.generate(input_ids=ids([5]), max_new_tokens=4, do_sample=False,
top_k=0, eos_token_id=None)
expect(g.shape == (1, 5), "单 token prompt 生成形状错误", got=tuple(g.shape))
x = torch.randint(0, 128, (1, 2048), device=DEVICE)
with torch.no_grad():
out = m(x)
expect(tuple(out.logits.shape) == (1, 2048, 128), "长序列形状错误")
expect(torch.isfinite(out.logits).all().item(), "长序列 logits 非有限值")
@testcase("functional", "F04 EOS 终止逻辑（单EOS/多EOS/禁用/错峰结束）")
def f04():
m = build_model(vocab_size=128, eos_token_id=2)
class ForceEOS:
"""按行、按步强制 argmax token：第 r 行在 schedule[r] 步输出 eos。"""
def __init__(self, schedule, eos_tok, batch):
self.schedule, self.eos_tok, self.batch = schedule, eos_tok, batch
self.step = 0
def __call__(self, module, module_in, output):
out = output.clone()
out.fill_(-1e9)
for r in range(self.batch):
tok = self.eos_tok if self.step >= 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) < 1e-8, "eval 模式 aux_loss 应为 0", v=aux)
del m
gc.collect()
@testcase("functional", "F11 YaRN 长上下文缩放")
def f11():
m = build_model(max_position_embeddings=4096, inference_rope_scaling=True)
expect(m.config.rope_scaling is not None and m.config.rope_scaling["type"] == "yarn",
"YaRN 配置未生效")
x = torch.randint(0, 128, (1, 2048), device=DEVICE)
with torch.no_grad():
out = m(x)
expect(tuple(out.logits.shape) == (1, 2048, 128), "YaRN 长序列形状错误")
expect(torch.isfinite(out.logits).all().item(), "YaRN logits 非有限值")
expect(m.model.freqs_cos.shape[0] == 4096, "RoPE buffer 长度错误")
# ============================================================
# 二、鲁棒性测试
# ============================================================
@testcase("robustness", "R01 空输入被拒绝")
def r01():
m = build_model()
expect_raises(ValueError,
lambda: m(torch.zeros(1, 0, dtype=torch.long, device=DEVICE)),
"空序列前向")
expect_raises(ValueError,
lambda: m.generate(inputs=torch.zeros(1, 0, dtype=torch.long, device=DEVICE),
max_new_tokens=2),
"空 prompt 生成")
@testcase("robustness", "R02 token id 越界（上界/下界）")
def r02():
m = build_model(vocab_size=128)
expect_raises(IndexError, lambda: m(ids([0, 128])), "id==vocab_size")
expect_raises(IndexError, lambda: m(ids([0, 10000])), "id 远超上界")
expect_raises(IndexError, lambda: m(ids([-1, 5])), "负 id")
expect_raises(IndexError,
lambda: m.generate(input_ids=ids([200]), max_new_tokens=1),
"generate 越界 id")
@testcase("robustness", "R03 非法采样参数矩阵")
def r03():
m = build_model()
p = ids([1])
expect_raises(ValueError, lambda: m.generate(p, max_new_tokens=2, temperature=-0.5), "负温度")
expect_raises(ValueError,
lambda: m.generate(p, max_new_tokens=2, temperature=float("nan")), "NaN 温度")
expect_raises(ValueError, lambda: m.generate(p, max_new_tokens=2, top_p=0.0), "top_p=0")
expect_raises(ValueError, lambda: m.generate(p, max_new_tokens=2, top_p=1.5), "top_p>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-4, "增量解码与全量前向 logits 偏差过大", diff=diff)
@testcase("regression", "G06 SDPA 与手写注意力一致")
def g06():
m_off = build_model(seed=321, flash_attn=False)
m_on = build_model(seed=999, flash_attn=True)
clone_weights(m_on, m_off)
x = ids([1, 2, 3, 0, 0])
mask = torch.tensor([[1, 1, 1, 0, 0]], device=DEVICE)
with torch.no_grad():
a = m_off(x, attention_mask=mask).logits
b = m_on(x, attention_mask=mask).logits
diff = (a - b).abs().max().item()
expect(diff < 1e-3, "flash/手写分支 logits 偏差过大", diff=diff)
@testcase("regression", "G07 save_pretrained / from_pretrained 输出一致")
def g07():
m = build_model(seed=555, use_moe=True, num_experts=2, num_experts_per_tok=1)
x = torch.randint(0, 128, (1, 10), device=DEVICE)
with torch.no_grad():
before = m(x).logits.clone()
with tempfile.TemporaryDirectory() as d:
m.save_pretrained(d)
m2 = fixed.MiniMindForCausalLM.from_pretrained(d).to(DEVICE).eval()
with torch.no_grad():
after = m2(x).logits
diff = (before - after).abs().max().item()
expect(diff < 1e-5, "存取后 logits 偏差过大", diff=diff)
expect(torch.equal(m2.lm_head.weight, m2.model.embed_tokens.weight),
"存取后权重绑定失效")
@testcase("regression", "G08 MoE eval 确定性与 checksum")
def g08():
m = build_model(seed=4242, use_moe=True, num_experts=4, num_experts_per_tok=2)
x = torch.randint(0, 128, (1, 16), device=DEVICE)
with torch.no_grad():
a = m(x).logits
b = m(x).logits
expect(torch.equal(a, b), "MoE eval 前向不确定")
expect(abs(a.sum().item() - (-2.204725)) < 2e-3
and abs(a.abs().max().item() - 1.595536) < 2e-3,
"MoE logits checksum 偏离 golden", s=a.sum().item(), mx=a.abs().max().item())
@testcase("regression", "G09 基线版右padding位置错位对照（证明回归测试有牙齿）")
def g09():
"""对照实验：基线版右 padding+cache 解码位置错位，修复版对齐。"""
expect(HAS_BASELINE, "基线模块不可用，跳过")
from transformers.cache_utils import DynamicCache
mf, mb = _golden_pair()
def padded_decode(model):
c = DynamicCache()
with torch.no_grad():
model(torch.tensor([[5, 6, 0, 0]], device=DEVICE),
attention_mask=torch.tensor([[1, 1, 0, 0]], device=DEVICE),
past_key_values=c, use_cache=True)
kw = {}
if model is mf:
kw["position_ids"] = torch.tensor([[2]], device=DEVICE)
o = model(torch.tensor([[9]], device=DEVICE),
attention_mask=torch.tensor([[1, 1, 0, 0, 1]], device=DEVICE),
past_key_values=c, use_cache=True, **kw)
c2 = DynamicCache()
with torch.no_grad():
model(torch.tensor([[5, 6]], device=DEVICE),
past_key_values=c2, use_cache=True)
ref = model(torch.tensor([[9]], device=DEVICE),
past_key_values=c2, use_cache=True)
return (o.logits[0, -1] - ref.logits[0, -1]).abs().max().item()
d_fixed = padded_decode(mf)
d_base = padded_decode(mb)
expect(d_fixed < 1e-5, "修复版右padding仍错位", d=d_fixed)
expect(d_base > 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 < t_nocache, "KV cache 未带来加速（异常）",
cache=round(t_cache, 3), no_cache=round(t_nocache, 3))
total_len = L + n_dec
theo = 2 * total_len * n_layers * kv_heads * head_dim * bytes_per
actual = sum(layer.keys.numel() + layer.values.numel()
for layer in c.layers) * c.layers[0].keys.element_size()
rows.append(("KV cache", total_len, f"theo={theo/1024:.1f} KiB",
f"actual={actual/1024:.1f} KiB"))
expect(abs(theo - actual) / theo < 0.02, "KV cache 实测字节与公式不符",
theo=theo, actual=actual)
n_params = sum(p.numel() for p in m.parameters())
rows.append(("params", "-", f"{n_params/1e6:.2f} M",
f"{n_params*bytes_per/1024/1024:.1f} MiB(权重)"))
if DEVICE == "cuda":
rows.append(("peak GPU mem", "-",
f"{torch.cuda.max_memory_allocated()/1024/1024:.1f} MiB", ""))
else:
try:
import psutil
rss = psutil.Process().memory_info().rss / 1024 / 1024
rows.append(("进程 RSS", "-", f"{rss:.1f} MiB", ""))
except ImportError:
rows.append(("进程 RSS", "-", "需 pip install psutil", ""))
print("\n +-----------------+----------+------------------+------------------+")
print(" | 阶段 | token 数 | 耗时 | 吞吐/占用 |")
print(" +-----------------+----------+------------------+------------------+")
for a, b, c1, d1 in rows:
print(f" | {a:<15} | {str(b):<8} | {c1:<16} | {d1:<16} |")
print(" +-----------------+----------+------------------+------------------+")
@testcase("performance", "P02 预填充 logits_to_keep=1 加速 lm_head")
def p02(quick=False):
m = build_model(seed=3, hidden_size=256, num_hidden_layers=2, num_attention_heads=8,
num_key_value_heads=2, vocab_size=16384)
x = torch.randint(0, 16384, (1, 512), device=DEVICE)
t_full = _median_time(lambda: m(x, logits_to_keep=0), repeats=3, warmup=1)
t_keep = _median_time(lambda: m(x, logits_to_keep=1), repeats=3, warmup=1)
with torch.no_grad():
a = m(x, logits_to_keep=0).logits[:, -1]
b = m(x, logits_to_keep=1).logits[:, -1]
diff = (a - b).abs().max().item()
print(f" 全量 logits {t_full*1000:.1f} ms / keep=1 {t_keep*1000:.1f} ms "
f"/ 加速 {t_full/t_keep:.2f}x / 末位 logits diff={diff:.2e}")
expect(diff < 1e-4, "logits_to_keep=1 末位结果不一致", diff=diff)
expect(t_keep <= t_full * 1.05, "keep=1 未体现加速（环境抖动？）")
# ============================================================
# 五、采样分布测试
# ============================================================
@testcase("sampling", "S01 temperature 缩放严格等于 softmax(logits/T)")
def s01():
torch.manual_seed(0)
logits = torch.randn(1, 64)
for T in (0.3, 0.7, 1.0, 1.5):
p = fixed.warp_logits(logits.clone(), T, 0, 1.0)
ref = torch.softmax(logits / T, dim=-1)
diff = (p - ref).abs().max().item()
expect(diff < 1e-6, "温度缩放与定义不符", T=T, diff=diff)
@testcase("sampling", "S02 熵随温度单调（低温尖锐/高温平坦）")
def s02():
torch.manual_seed(1)
logits = torch.randn(1, 128)
ent = lambda p: float((-(p.clamp_min(1e-12) * p.clamp_min(1e-12).log()).sum(-1)))
temps = [0.1, 0.4, 1.0, 2.0, 5.0]
es = [ent(fixed.warp_logits(logits.clone(), T, 0, 1.0)) for T in temps]
for a, b in zip(es, es[1:]):
expect(a < b, "熵未随温度升高而增大", es=[round(x, 4) for x in es])
expect(es[-1] > 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 <= k, "保留 token 数超过 k", k=k, n=n_kept)
topk_idx = set(torch.topk(logits[r], k).indices.tolist())
kept_idx = set(torch.where(kept)[0].tolist())
expect(kept_idx.issubset(topk_idx), "top-k 支撑集错误")
expect(abs(p[r].sum().item() - 1.0) < 1e-5, "top-k 后质量不为 1",
s=p[r].sum().item())
@testcase("sampling", "S04 top-p(nucleus) 累积质量边界与最小集合")
def s04():
torch.manual_seed(3)
logits = torch.randn(8, 200)
for p_val in (0.1, 0.5, 0.9):
p = fixed.warp_logits(logits.clone(), 1.0, 0, p_val)
sp, _ = torch.sort(p, descending=True, dim=-1)
for r in range(8):
kept_mass = p[r].sum().item()
expect(kept_mass >= 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() < p_val + 1e-5,
"存在可移除的多余 token（集合非最小）", p=p_val)
@testcase("sampling", "S05 大样本频率拟合理论分布（无偏性）")
def s05(quick=False):
torch.manual_seed(4)
V = 20
probs = torch.softmax(torch.randn(1, V), dim=-1)
n = 20000 if quick else 120000
gen = torch.Generator().manual_seed(2026)
draws = torch.multinomial(probs.expand(n, V), num_samples=1,
generator=gen if DEVICE == "cpu" else None).squeeze(-1)
freq = torch.bincount(draws, minlength=V).float() / n
max_err = (freq - probs[0]).abs().max().item()
bound = 3.04 / math.sqrt(2 * n) # Hoeffding 99.9% 保守界
expect(max_err < bound, "经验频率偏离理论概率（采样可能有偏）",
err=round(max_err, 5), bound=round(bound, 5))
@testcase("sampling", "S06 采样端到端：top-k 支撑集 / 贪婪=argmax")
def s06():
m = build_model(seed=9, eos_token_id=None)
p = ids([1, 2, 3])
with torch.no_grad():
greedy_first = m(p).logits[0, -1].argmax().item()
g_greedy = m.generate(input_ids=p, max_new_tokens=1, do_sample=False,
top_k=0, eos_token_id=None)
expect(g_greedy[0, -1].item() == greedy_first, "贪婪首 token 与 argmax 不符")
g_logits = torch.Generator().manual_seed(7)
logits = torch.randn(1000, 128, generator=g_logits)
nxt = fixed.sample_next_token(logits, do_sample=True, temperature=1.0,
top_k=5, top_p=1.0,
generator=torch.Generator().manual_seed(8))
for r in range(1000):
support = set(torch.topk(logits[r], 5).indices.tolist())
expect(nxt[r].item() in support, "采样 token 落在 top-k 支撑集之外")
@testcase("sampling", "S07 temperature=0 强制贪婪；generator 可复现")
def s07():
m = build_model(seed=12, eos_token_id=None)
p = ids([1, 2])
g_t0 = m.generate(input_ids=p, max_new_tokens=8, do_sample=True,
temperature=0.0, top_k=0, top_p=1.0, eos_token_id=None)
g_greedy = m.generate(input_ids=p, max_new_tokens=8, do_sample=False,
top_k=0, eos_token_id=None)
expect(torch.equal(g_t0, g_greedy), "temperature=0 未退化为贪婪")
kw = dict(input_ids=p, max_new_tokens=10, do_sample=True, temperature=0.9,
top_k=40, top_p=0.9, eos_token_id=None)
ga = m.generate(generator=torch.Generator().manual_seed(31337), **kw)
gb = m.generate(generator=torch.Generator().manual_seed(31337), **kw)
expect(torch.equal(ga, gb), "同 generator 种子结果不一致")
@testcase("sampling", "S08 重复惩罚公式方向与数值精确性")
def s08():
logits = torch.tensor([[2.0, -2.0, 1.0, -1.0]])
history = torch.tensor([[0, 1]])
out = fixed.apply_repetition_penalty(logits.clone(), history, 1.5, 1024)
expect(abs(out[0, 0].item() - 2.0 / 1.5) < 1e-6, "正 logit 惩罚分支错误",
v=out[0, 0].item())
expect(abs(out[0, 1].item() - (-2.0 * 1.5)) < 1e-6, "负 logit 惩罚分支错误",
v=out[0, 1].item())
expect(abs(out[0, 2].item() - 1.0) < 1e-6 and abs(out[0, 3].item() + 1.0) < 1e-6,
"未见 token 不应被惩罚")
p_before = torch.softmax(logits, -1)
p_after = torch.softmax(out, -1)
expect(bool(p_after[0, 0] < p_before[0, 0] and p_after[0, 1] < p_before[0, 1]),
"惩罚后已见 token 概率未下降")
same = fixed.apply_repetition_penalty(logits.clone(), history, 1.0, 1024)
expect(torch.allclose(same, logits), "penalty=1 应恒等")
@testcase("sampling", "S09 端到端分布：高温生成多样性高于低温")
def s09(quick=False):
m = build_model(seed=21, eos_token_id=None)
p = ids([1, 2, 3, 4])
n_seq = 16 if quick else 40
stats = {}
for T in (0.1, 1.0):
toks = []
for i in range(n_seq):
gen = torch.Generator().manual_seed(1000 + i) if DEVICE == "cpu" else None
with seed_context(1000 + i):
out = m.generate(input_ids=p, max_new_tokens=16, do_sample=True,
temperature=T, top_k=0, top_p=1.0,
eos_token_id=None, generator=gen)
toks.append(out[0, p.shape[1]:])
stats[T] = len(torch.unique(torch.stack(toks)))
expect(stats[1.0] >= 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
==========================================================================
文章出处：https://www.cnblogs.com/LyShark/p/23039691
本博客所有文章除特别声明外，均采用 BY-NC-SA 许可协议。转载请注明出处！
本博客所有文章除特别声明外，均采用 BY-NC-SA 许可协议。转载请注明出处！
浙公网安备 33010602011771号