FRSM 训练实验报告

硬件环境

运行对比总览

详细记录

V1 — 原版 FRSM (GRPO 训练)

• 文件: frsm.py、train_pretrain.py → train_grpo.py (含 critic loss)
• 参数: d_model=1904, head=8, n_layer=12, 约 180M
• 训练: GRPO, 含 lm_loss + crit_loss
• 状态: step 27,700/317,577, loss=2.97 (lm) + 1.02 (crit), lm 已平台化
• 速度: ~7,080 tok/s (177 raw × 40 batch)
• 日志: train_out.txt (555 条记录)
• 检查点: frsm_checkpoints/ (2.4 GB)
• 结论: 原始架构已瓶颈，lm loss 无法继续下降（~2.97 平台），转 V6 架构

V6 — 非 Fast 多尺度 (200M)

• 文件: frsm_v6.py、train_frsm_v6.py
• 参数: d_model=1904, num_scales=4, ~200M
• 改进: 多头尺度状态机 + 内容门控 + 门控 soft 更新
• 状态: step 5,850/396,975, loss=2.78 (仍在下降)
• 速度: ~5,360 tok/s (134 raw × 40 batch) — 瓶颈: Python 逐 scale 循环
• 日志: train_v6_out.txt (118 条记录)
• 检查点: frsm_v6_checkpoints/ (2.3 GB)
• 问题: 4 个 scale 用 Python for 循环串行计算，GPU 利用率不足

V6 Fast (200M) — Einsum 并行化

• 文件: frsm_v6a_fast.py（初始版）
• 改进: 将所有 scale 权重合并为 (ns, d, 2d) 张量，用 einsum 一次计算全部 4 个 scale
• 初始问题: batch_size=64 OOM (21.1 GB)
• 修复 1: batch_size 降为 40
• 修复 2: _init_weights — stacked tensor 上用 kaiming_uniform 的 fan_in 计算错误（44x 偏小），改为逐 scale 单独 init
• 修复 3: outputs.append + torch.cat → 预分配 torch.empty(B,T,vs)，节省 ~1.7 GB
• 状态: step 12,600, loss=2.55 平台化 — 200M 容量已耗尽
• 速度: 9,240 tok/s（比 V6 非 Fast 快 72%）
• 结论: 模型容量瓶颈，决定尝试 400M+

V6 Fast (417M) ⭐ 当前运行

• 文件: frsm_v6a_fast.py（当前版）
• 启动参数: –d_model 3000 –num_scales 4 –batch_size 40 –max_steps 396975 –lr 0.0005
• 当前状态: step 6,450 / 396,975, loss=2.75 (持续下降中)
• 速度: ~2,320 tok/s (d=3000 的 d² 缩放导致比 200M 慢 4×)
• GPU 内存: 20.9 GB / 23.0 GB (91%) — 含模型 + 梯度 + AdamW 状态
• GPU 温度: 54°C, 利用率: 100%
• epoch 进度: ~3%
• 预计 ETA: ~16 天完成 1 epoch
• 关键修复:
  • TF32 NaN: torch.set_float32_matmul_precision('highest') — TF32 的 10-bit 尾数在 6000 维 einsum 内积中溢出，导致 NaN
  • 未初始化 bias: b_cand、gate_b1、gate_b2 未在 _init_weights 中初始化，torch.empty 返回 garbage/NaN。修复: 显式 nn.init.zeros_

内存分析（d=3000, bs=40）

速度对比（d 缩放）

速度下降约 4×，与 d² 缩放一致（注意力 + einsum 均为 O(d²)）。

损失曲线

文件清单