98 lines
2.8 KiB
Python
98 lines
2.8 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
测试实数版本的位置编码
|
||
"""
|
||
|
||
import torch
|
||
from model.model import precompute_pos_cis, precompute_pos_cis_real, apply_rotary_emb, apply_rotary_emb_real
|
||
from model.LMConfig import LMConfig
|
||
from model.model import MiniMindLM
|
||
|
||
def test_pos_encoding_equivalence():
|
||
"""测试复数版本和实数版本的位置编码是否等价"""
|
||
print("测试位置编码等价性...")
|
||
|
||
# 参数设置
|
||
dim = 64
|
||
seq_len = 10
|
||
|
||
# 生成复数版本的位置编码
|
||
pos_cis = precompute_pos_cis(dim=dim, end=seq_len)
|
||
|
||
# 生成实数版本的位置编码
|
||
pos_cis_real = precompute_pos_cis_real(dim=dim, end=seq_len)
|
||
|
||
# 创建随机查询和键
|
||
batch_size = 2
|
||
n_heads = 4
|
||
head_dim = dim
|
||
|
||
xq = torch.randn(batch_size, seq_len, n_heads, head_dim)
|
||
xk = torch.randn(batch_size, seq_len, n_heads, head_dim)
|
||
|
||
# 应用复数版本的旋转位置编码
|
||
xq_complex, xk_complex = apply_rotary_emb(xq, xk, pos_cis)
|
||
|
||
# 应用实数版本的旋转位置编码
|
||
xq_real, xk_real = apply_rotary_emb_real(xq, xk, pos_cis_real)
|
||
|
||
# 计算差异
|
||
q_diff = torch.abs(xq_complex - xq_real).mean().item()
|
||
k_diff = torch.abs(xk_complex - xk_real).mean().item()
|
||
|
||
print(f"查询差异: {q_diff:.6f}")
|
||
print(f"键差异: {k_diff:.6f}")
|
||
|
||
# 检查差异是否在可接受范围内
|
||
tolerance = 1e-5
|
||
if q_diff < tolerance and k_diff < tolerance:
|
||
print("✅ 测试通过: 复数版本和实数版本的位置编码在数值上等价")
|
||
else:
|
||
print("❌ 测试失败: 复数版本和实数版本的位置编码存在显著差异")
|
||
|
||
def test_model_forward():
|
||
"""测试模型前向传播"""
|
||
print("\n测试模型前向传播...")
|
||
|
||
# 创建模型配置
|
||
config = LMConfig(
|
||
dim=128,
|
||
n_layers=2,
|
||
n_heads=4,
|
||
n_kv_heads=4, # 确保n_kv_heads被设置，且n_heads能被n_kv_heads整除
|
||
vocab_size=1000,
|
||
max_seq_len=128,
|
||
disable_db=True # 禁用数据库功能，避免额外的复杂性
|
||
)
|
||
|
||
# 创建模型
|
||
try:
|
||
model = MiniMindLM(config)
|
||
print(f"✅ 模型初始化成功")
|
||
except Exception as e:
|
||
print(f"❌ 模型初始化失败: {str(e)}")
|
||
return
|
||
|
||
# 创建输入
|
||
batch_size = 2
|
||
seq_len = 10
|
||
input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len))
|
||
|
||
# 前向传播
|
||
try:
|
||
with torch.no_grad():
|
||
outputs = model(input_ids)
|
||
print(f"✅ 模型前向传播成功")
|
||
print(f"输出形状: {outputs.logits.shape}")
|
||
except Exception as e:
|
||
print(f"❌ 模型前向传播失败: {str(e)}")
|
||
|
||
if __name__ == "__main__":
|
||
# 测试位置编码等价性
|
||
test_pos_encoding_equivalence()
|
||
|
||
# 测试模型前向传播
|
||
test_model_forward()
|