Spaces:
Running
Running
File size: 3,545 Bytes
04c7011 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | """
小训练脚本 - 训练 10 步（验证损失持续下降）
"""
import sys
import torch
import torch.optim as optim
sys.path.insert(0, '.')
from models.fusion_mini import FusionMini, FusionMiniConfig
def train_small():
"""小训练（10 步）"""
print("[TRAIN] 开始小训练（10 步）...")
print()
# 1. 创建小配置
print("[1] 创建模型配置...")
config = FusionMiniConfig(
vocab_size=1000, # 小词表
hidden_size=64, # 小隐层
num_hidden_layers=2, # 2 层
num_attention_heads=2, # 2 个注意力头
intermediate_size=128,
max_position_embeddings=64,
)
print(f" 词汇表大小: {config.vocab_size}")
print(f" 隐藏层大小: {config.hidden_size}")
print(f" 层数: {config.num_hidden_layers}")
print()
# 2. 创建模型
print("[2] 创建模型...")
model = FusionMini(config)
model.train() # 训练模式
param_count = sum(p.numel() for p in model.parameters()) / 1e3
print(f" 参数量: {param_count:.1f}K")
print(" 模型创建成功")
print()
# 3. 创建优化器
print("[3] 创建优化器...")
optimizer = optim.AdamW(
model.parameters(),
lr=5e-4, # 稍大的学习率
weight_decay=0.01,
)
print(" 优化器创建成功")
print()
# 4. 创建假数据（小批次）
print("[4] 创建假数据...")
batch_size = 4
seq_len = 16
input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len))
labels = torch.randint(0, config.vocab_size, (batch_size, seq_len))
print(f" 输入形状: {input_ids.shape}")
print(f" 标签形状: {labels.shape}")
print(" 假数据创建成功")
print()
# 5. 训练 10 步
print("[5] 训练 10 步...")
losses = []
for step in range(10):
# 清零梯度
optimizer.zero_grad()
# 前向传播
outputs = model(
input_ids=input_ids,
labels=labels,
return_dict=True,
)
loss = outputs["loss"]
losses.append(loss.item())
# 反向传播
loss.backward()
# 更新参数
optimizer.step()
print(f" Step {step+1:2d}: Loss = {loss.item():.4f}")
print(" 训练完成")
print()
# 6. 验证损失下降
print("[6] 验证损失下降...")
initial_loss = losses[0]
final_loss = losses[-1]
is_decreasing = final_loss < initial_loss
print(f" 初始 Loss: {initial_loss:.4f}")
print(f" 最终 Loss: {final_loss:.4f}")
print(f" Loss 变化: {final_loss - initial_loss:+.4f}")
print()
if is_decreasing:
print(" [PASS] Loss 持续下降")
print(" 训练有效！")
else:
print(" [WARN] Loss 未下降")
print(" 可能的问题：学习率太大 / 数据太少 / 模型太小")
print()
print("[TRAIN] 小训练完成")
return is_decreasing
if __name__ == "__main__":
print("=" * 60)
print("Fusion-LLM 小训练（10 步）")
print("=" * 60)
print()
try:
success = train_small()
if success:
print()
print("[PASS] 训练测试通过")
except Exception as e:
print()
print(f"[FAIL] 训练测试出错: {e}")
import traceback
traceback.print_exc()
sys.exit(1)
sys.exit(0)
|