机器之心报道
编辑:泽南、小舟
租用 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.txt
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu124 —upgrade # install torch 2.6.0
python data/cached_fineweb10B.py 10 # downloads only the first 1.0B training tokens to save time
./run.sh
sudo apt-get update
sudo apt-get install vim tmux python3-pip python-is-python3 -y
git clone https://github.com/KellerJordan/modded-nanogpt.git
cd modded-nanogpt
tmux
pip install numpy==1.23.5 huggingface-hub tqdm
pip 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 18
sudo docker run -it --rm --gpus all -v $(pwd):/modded-nanogpt modded-nanogpt sh run.sh
@torch.compile
def zeroth_power_via_newtonschulz5 (G, steps=5, eps=1e-7):
assert len (G.shape) == 2
b, c = (3.4445, -4.7750, 2.0315)
X = G.bfloat16 () / (G.norm () + eps)
if G.size (0) > G.size (1):
X = X.T
for _ in range (steps):
A = X @ X.T
B = b * A + c * A @ A
X = a * X + B @ X
if G.size (0) > G.size (1):
X = X.T
return 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 根)。
© 版权声明
文章版权归作者所有,未经允许请勿转载。
相关文章
暂无评论...