FRSM 训练实验报告

FRSM 训练实验报告 硬件环境项目规格GPUNVIDIA RTX 4090D × 1 (22.15 GB VRAM)CPUx86_64RAM~32 GB存储/mnt/scratch (tmpfs, 用于数据缓存)运行对比总览V1 (原版 FRSM)V6 (非 Fast)V6 Fast (200M)V6 Fast (417M) ⭐模型文件frsm.pyfrsm_v6.pyfrsm_v6a_fast.pyfrsm_v6a_fast.py架构原始 FRSM多尺度状态机 (逐 scale 循环)多尺度状态机 (einsum 并行)同左d_model1904190419043000参数量~180M~200M~200M~417Mnum_scales—444batch_size40 (GRPO)404040max_seq_len384384384384训练模式GRPO (含 critic)纯 pretrain纯 pretrain纯 pretrain训练步数27,700/317,5775,850/396,97512,6006,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 GB19.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_model1904, head8, n_layer12, 约 180M训练: GRPO, 含 lm_loss crit_loss状态: step 27,700/317,577, loss2.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_model1904, num_scales4, ~200M改进: 多头尺度状态机 内容门控 门控 soft 更新状态: step 5,850/396,975, loss2.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_size64OOM (21.1 GB)修复 1: batch_size 降为 40修复 2:_init_weights— stacked tensor 上用 kaiming_uniform 的 fan_in 计算错误44x 偏小改为逐 scale 单独 init修复 3:outputs.appendtorch.cat→ 预分配torch.empty(B,T,vs)节省 ~1.7 GB状态: step 12,600, loss2.55平台化— 200M 容量已耗尽速度: 9,240 tok/s比 V6 非 Fast 快 72%结论: 模型容量瓶颈决定尝试 400MV6 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, loss2.75(持续下降中)速度: ~2,320 tok/s (d3000 的 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_内存分析d3000, bs40组件大小模型参数 (417M × 4B)1.67 GB梯度 (417M × 4B)1.67 GBAdamW 状态 (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~200M9,2401.00×3000~417M2,3200.25×速度下降约 4×与 d² 缩放一致注意力 einsum 均为 O(d²)。损失曲线文件清单文件角色大小frsm.pyV1 模型定义—frsm_v6.pyV6 非 Fast 模型—frsm_v6a_fast.pyV6 Fast 模型当前9.0 KBtrain_frsm_v6_fast.pyV6 Fast 训练脚本—train_out.txtV1 训练日志 (555 entries)—train_v6_out.txtV6 非 Fast 日志 (118 entries)—train_v6_fast_out.txtV6 Fast 417M 日志 (130 entries)—frsm_checkpoints/V1 检查点2.4 GBfrsm_v6_checkpoints/V6 非 Fast 检查点2.3 GBfrsm_v6_fast_checkpoints/V6 Fast 417M 检查点当前4.7 GB/mnt/scratch/pretrain_t2t_mini.jsonl_cache_8470000_384_v2.pt数据集缓存~18 GB