FRSM 训练实验报告

硬件环境

项目	规格
GPU	NVIDIA RTX 4090D × 1 (22.15 GB VRAM)
CPU	x86_64
RAM	~32 GB
存储	/mnt/scratch (tmpfs, 用于数据缓存)

运行对比总览

V1 (原版 FRSM)	V6 (非 Fast)	V6 Fast (200M)	V6 Fast (417M) ⭐
模型文件	frsm.py	frsm_v6.py	frsm_v6a_fast.py	frsm_v6a_fast.py
架构	原始 FRSM	多尺度状态机 (逐 scale 循环)	多尺度状态机 (einsum 并行)	同左
d_model	1904	1904	1904	3000
参数量	~180M	~200M	~200M	~417M
num_scales	—	4	4	4
batch_size	40 (GRPO)	40	40	40
max_seq_len	384	384	384	384
训练模式	GRPO (含 critic)	纯 pretrain	纯 pretrain	纯 pretrain
训练步数	27,700/317,577	5,850+/396,975	12,600	6,450+/396,975
最佳 loss	~2.97 (lm)	~2.78	~2.55 (平台)	~2.75 (下降中)
速度 (tok/s)	~7,080 (177×40)	~5,360 (134×40)	9,240	~2,320 (58×40)
GPU 内存	—	—	21.1 GB	19.6 GB (含优化器)
epochs	~0.26	~0.03	~0.06	~0.03 (进行中)

详细记录

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=64OOM (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
• 当前状态: step6,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）

组件	大小
模型参数 (417M × 4B)	1.67 GB
梯度 (417M × 4B)	1.67 GB
AdamW 状态 (2 × 417M × 4B)	3.34 GB
持久占用合计	6.68 GB
前向激活 (batch 峰值)	~3-4 GB
缓存分配器碎片	~1-2 GB
实际峰值	~20.9 GB

速度对比（d 缩放）

d_model	参数量	tok/s	相对速度
1904	~200M	9,240	1.00×
3000	~417M	2,320	0.25×

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

损失曲线

文件清单

文件	角色	大小
frsm.py	V1 模型定义	—
frsm_v6.py	V6 非 Fast 模型	—
frsm_v6a_fast.py	V6 Fast 模型（当前）	9.0 KB
train_frsm_v6_fast.py	V6 Fast 训练脚本	—
train_out.txt	V1 训练日志 (555 entries)	—
train_v6_out.txt	V6 非 Fast 日志 (118 entries)	—
train_v6_fast_out.txt	V6 Fast 417M 日志 (130 entries)	—
frsm_checkpoints/	V1 检查点	2.4 GB
frsm_v6_checkpoints/	V6 非 Fast 检查点	2.3 GB
frsm_v6_fast_checkpoints/	V6 Fast 417M 检查点（当前）	4.7 GB
/mnt/scratch/pretrain_t2t_mini.jsonl_cache_8470000_384_v2.pt	数据集缓存	~18 GB