机器之心报道
编辑：泽南、小舟
租用 H100 的钱只需 233 美元。
10B tokens-->1B tokens 8xH100 上花 45 分钟训练 -->8xH100 上花 5 分钟训练
先进的架构：旋转嵌入、QK-Norm 和 ReLU^2； 新优化器：Muon； 嵌入中的 Untied Head； 投影和分类层初始化为零（muP-like）； 架构 shortcut：值残差和嵌入 shortcut（部分遵循论文《Value Residual Learning For Alleviating Attention Concentration In Transformers》）； 动量（Momentum）warmup； Tanh soft logit capping（遵循 Gemma 2）； FlexAttention。
pip install -r requirements.txtpip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu124 —upgrade # install torch 2.6.0python data/cached_fineweb10B.py 10 # downloads only the first 1.0B training tokens to save time./run.sh
sudo apt-get updatesudo apt-get install vim tmux python3-pip python-is-python3 -ygit clone https://github.com/KellerJordan/modded-nanogpt.gitcd modded-nanogpttmuxpip install numpy==1.23.5 huggingface-hub tqdmpip install --upgrade torch &python data/cached_fineweb10B.py 18
sudo docker build -t modded-nanogpt .sudo docker run -it --rm --gpus all -v $(pwd):/modded-nanogpt modded-nanogpt python data/cached_fineweb10B.py 18sudo docker run -it --rm --gpus all -v $(pwd):/modded-nanogpt modded-nanogpt sh run.sh
@torch.compiledef zeroth_power_via_newtonschulz5 (G, steps=5, eps=1e-7):assert len (G.shape) == 2b, c = (3.4445, -4.7750, 2.0315)X = G.bfloat16 () / (G.norm () + eps)if G.size (0) > G.size (1):X = X.Tfor _ in range (steps):A = X @ X.TB = b * A + c * A @ AX = a * X + B @ Xif G.size (0) > G.size (1):X = X.Treturn X.to (G.dtype)
内存使用量比 Adam 低 采样效率提高约 1.5 倍 挂钟开销小于 2%
在更新中使用 Nesterov 动量，在动量之后应用正交化。 使用特定的五次 Newton-Schulz 迭代作为正交化方法。 使用五次多项式的非收敛系数以最大化零处的斜率，从而最小化必要的 Newton-Schulz 迭代次数。事实证明，方差实际上并不那么重要，因此我们最终得到一个五次多项式，它在重复应用后（快速）收敛到 0.68、1.13 的范围，而不是到 1。 在 bfloat16 中运行 Newton-Schulz 迭代（而 Shampoo 实现通常依赖于在 fp32 或 fp64 中运行的逆 pth 根）。
