Hugging Peers

MA-RMSNorm-frasm3.7

☆ Save
by dfytensor text-generation Pytorch MIT License
0 downloads 0 likes on ModelScope updated

Model card

MA-RMSNorm 在 frsmash3.7 上的验证 (30M, minimind)

对照实验:在真实架构 (frsmash3.7 多槽 SSM) + 真实数据 (minimind 三问过滤) + 30M 规模上, 把标准 RMSNorm 替换为 MA-RMSNorm (Magnitude-Aware RMSNorm),证明其达到(且优于)同等能力。

源自 MA-RMSNorm 论文 的方法,原论文仅在 NumPy 玩具模型+合成数据上验证; 本仓库在生产级 SSM 架构 + 真实中文语料上完成严格对照。

方法

标准 RMSNorm:   y = (x / RMS(x)) * γ
MA-RMSNorm:     y = (x / RMS(x)) * γ * s + α * log(RMS(x)) * W
  • s = 2.0:主路幅度放大
  • α = 0.1:幅值注入强度
  • W:可学习 per-dim 权重(D 参数/归一化点),把标准 RMSNorm 丢弃的幅度信息经 log(σ) 注回

核心结论

同架构 (H=320 L=6 heads=8 = 31.8M)、同数据 (minimind)、同种子 (42)、同超参,唯一变量是归一化:

阶段 RMSNorm MA-RMSNorm Δ
预训练 (120k, 3ep) 3.6934 3.5116 −4.92%
SFT (指令对话, 2ep) 2.9906 2.8534 −4.59%

MA-RMSNorm 在预训练与 SFT 两个阶段 loss 均比 RMSNorm 低约 4.6%——验证有效(且优于)。 全程 MA 持续领先(非末尾偶然波动)。代价:每 step 慢 ~22%(额外 log/乘法),但收敛 loss 更低。

仓库结构

MA-RMSNorm-frasm3.7/
├── model.py              # frsmash3.7 + MARMSNorm 类 + _norm() 工厂 (NORM_TYPE 切换)
├── train_ma.py           # 预训练对照 (--norm rms/ma)
├── train_sft.py          # SFT 对照 (--norm rms/ma)
├── run_compare.py        # 预训练对照 runner (顺序跑 rms+ma, 对比)
├── run_sft_compare.py    # SFT 对照 runner
├── checkpoints/          # 4 个权重 (LFS)
│   ├── frsm37_30m_rms_final.pth       # RMSNorm 预训练
│   ├── frsm37_30m_ma_final.pth        # MA-RMSNorm 预训练
│   ├── frsm37_30m_rms_sft_final.pth   # RMSNorm SFT
│   └── frsm37_30m_ma_sft_final.pth    # MA-RMSNorm SFT
└── README.md

复现

环境:F:\rwkv\.venv (torch 2.13+cu126, fla, triton, jieba) + OpenASHVoc 词表。 Windows 需 PYTHONUTF8=1(脚本已内置)。

# 预训练对照 (各 ~35min on 4090)
python run_compare.py          # 顺序训 rms + ma, 输出 final_loss 对比

# SFT 对照 (各 ~70min, 需先有预训练 ckpt)
python run_sft_compare.py

数据(不随仓库分发):复用 frsmash3.7 项目的 minimind 三问过滤数据:

  • 预训练:pretrain_filtered.jsonl
  • SFT:sft_filtered.jsonl
  • 词表:open_ash_voc_agent.json (OpenASHVoc)

配置

HIDDEN=320, LAYERS=6, HEADS=8   # 31.8M (~30M)
VS=23006 (OpenASHVoc)
SEQ_pretrain=512, SEQ_sft=768
norm_type: 'rms' (nn.RMSNorm) 或 'ma' (MARMSNorm, s=2.0, α=0.1)

MA-RMSNorm 仅在每个归一化点多 mag_w[D] 参数(全模型 +~7k 参数,可忽略),其余完全一致。

诚实 caveat

  • 单种子(seed=42):4.6% + 两阶段一致已显著;严谨论文可加 3-5 种子取均值±方差。
  • MA 每 step 慢 ~22%(log/乘法开销),换取更低收敛 loss。
  • 仅 30M 规模验证;更大规模的行为待后续。

引用

  • frsmash3.7 架构: github.com/dfytensor/frsmash3.7
  • MA-RMSNorm 方法: "MA-RMSNorm: Breaking the Scale-For-Intelligence Trade-off"
  • 数据/词表: minimind + OpenASHVoc