Github
论文Sophia: 一种可扩展的随机二阶优化器，用于语言模型预训练
这是论文 https://arxiv.org/abs/2305.14342 中 Sophia-G 优化器的官方实现和 GPT-2 训练脚本。代码基于 nanoGPT 和 levanter。如果您发现 Sophia 有用，请引用该论文并为此仓库加星。谢谢！
@article{liu2023sophia,
title={Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training},
author={Liu, Hong and Li, Zhiyuan and Hall, David and Liang, Percy and Ma, Tengyu},
journal={arXiv preprint arXiv:2305.14342},
year={2023}
}
新闻和更新
- 使用最新的 PyTorch 版本更新了结果。
依赖
- PyTorch 2.1.2
- transformers 4.33.0
- datasets
- tiktoken
- wandb
一般用法
以下是使用 SophiaG 训练具有 NLL 损失的通用模型的示例代码片段。请参阅下一节以获取超参数调优指南。
import torch
import torch.nn.functional as F
from sophia import SophiaG
# 初始化模型、损失函数和输入数据
model = Model()
data_loader = ...
# 初始化优化器
optimizer = SophiaG(model.parameters(), lr=2e-4, betas=(0.965, 0.99), rho=0.01, weight_decay=1e-1)
total_bs = len(data_loader)
bs = total_bs * block_size
k = 10
iter_num = -1
# 训练循环
for epoch in range(epochs):
for X, Y in data_loader:
# 标准训练代码
logits, loss = model(X, Y)
loss.backward()
optimizer.step(bs=bs)
optimizer.zero_grad(set_to_none=True)
iter_num += 1
if iter_num % k != k - 1:
continue
else:
# 更新 Hessian EMA
logits, _ = model(X, None)
samp_dist = torch.distributions.Categorical(logits=logits)
y_sample = samp_dist.sample()
loss_sampled = F.cross_entropy(logits.view(-1, logits.size(-1)), y_sample.view(-1), ignore_index=-1)
loss_sampled.backward()
optimizer.update_hessian()
optimizer.zero_grad(set_to_none=True)
model.zero_grad()
超参数调优
学习率定义
- 代码中的更新写为 $\theta_{t+1} = \theta_t - lr*\textup{clip}(m_t / (\rho * h_t + \epsilon), 1)$，这等同于论文中的更新，只是重新参数化了（这里的 $lr$ 对应论文中的 $\rho \cdot \eta_t$）。因此，AdamW 和 Lion 的学习率不能直接比较。经验上，学习率比例为 5:1 的 Adam 和 Lion 有类似的行为。SophiaG 和 Lion 的学习率可以直接比较。Sophia 允许使用比 Lion 大得多的学习率，这就是 Sophia 快得多的原因。
调整超参数 $\rho$
- 调整 $\rho$ 使裁剪坐标的比例稳定并在适当范围内。这在 GPT-2 训练示例 中被追踪为
train/win_rate。train/win_rate应该在开始时达到峰值，之后保持稳定。train/win_rate应保持在 0.1 - 0.5 范围内。通常较大的 $\rho$ 会导致较大的train/win_rate。下面提供了 T5 模型中典型的win_rate行为示例。
调整学习率和权重衰减
- 选择略小于您为 AdamW 使用的学习率，或者是您为 Lion 使用的学习率的 3 - 5 倍。
-
如果损失爆炸，稍微降低学习率或增加 $\rho$。
-
始终使用比 AdamW 大约 2 倍的权重衰减。
GPT-2 模型的超参数
-
选择与您为 AdamW 使用的学习率大致相同，或者是您为 Lion 使用的学习率的 5 - 10 倍。
-
调整 $\rho$ 使未裁剪更新的参数比例稳定并在适当范围内。这在 GPT-2 训练示例 中被追踪为
train/win_rate。train/win_rate应该在开始时达到峰值，之后保持稳定。train/win_rate应保持在 0.1 - 0.5 范围内。通常较大的 $\rho$ 会导致较大的train/win_rate。 -
使用略大于 AdamW 的权重衰减，例如 0.2。
-
除学习率外，所有其他超参数都可以在不同模型大小之间迁移。
-
请参阅下表了解不同模型大小的超参数。 | 模型大小 | Adam 学习率 | Lion 学习率 | Sophia 学习率 | Sophia 的 $\rho$ 值 | Sophia 的权重衰减 | | -------- | ------- | ------- | ------- | ------- | ------- | | 125M | 6e-4 | 1e-4 | 6e-4 | 0.05 | 0.2 | | 355M | 3e-4 | 1e-4 | 7e-4 | 0.08 | 0.2 | | 770M | 2e-4 | 8e-5 | 3e-4 | 0.05 | 0.2 |
-
请随时让我们知道您在超参数调整过程中的发现。我们非常感谢您宝贵的反馈和评论！
复现 GPT-2 结果
按照 nanoGPT 的方法准备 OpenWebText 数据：
$ python data/openwebtext/prepare.py
开始预训练 GPT2 Small (125M)：
如果您有一台配备 10 个 A5000 (24GB) GPU 的机器，
$ torchrun --standalone --nproc_per_node=10 \
train_sophiag.py \
config/train_gpt2_small_sophiag.py \
--batch_size=8 \
--gradient_accumulation_steps=6
如果您有一台配备 8 个 A100 (40GB) GPU 的机器，
$ torchrun --standalone --nproc_per_node=8 \
train_sophiag.py \
config/train_gpt2_small_sophiag.py \
--batch_size=12 \
--gradient_accumulation_steps=5
要复现 AdamW 基准，请按照 nanoGPT 的方法：
$ torchrun --standalone --nproc_per_node=10 \
train_adam.py \
config/train_gpt2_small_adam.py \
--batch_size=8 \
--gradient_accumulation_steps=6
这将得到下图所示的结果：
开始预训练 GPT2 Medium (355M)：
如果您有一台配备 8 个 A100 (40GB) GPU 的机器，
$ torchrun --standalone --nproc_per_node=8 \
train_sophiag.py \
config/train_gpt2_medium_sophiag.py \
--batch_size=6 \
--gradient_accumulation_steps=10
要复现 AdamW 基准：
$ torchrun --standalone --nproc_per_node=8 \
train_adam.py \
config/train_gpt2_medium_adam.py \
--batch_size=6 \
--gradient_accumulation_steps=10
如果您使用其他硬件设置，请相应调整 nproc_per_node、batch_size 和 gradient_accumulation_steps。确保它们的乘积等于 480。
这将得到下图所示的结果：
开始预训练 GPT2 1.5B：
我们使用 the Pile 和 GPT NeoX 分词器。首先按照 levanter 设置 TPU 实例和环境。然后在 optim.py 中将 GAMMA_SOPHIA_G 更改为 200。1.5B 模型的训练脚本如下：
gcloud compute tpus tpu-vm ssh <instance_name> \
--zone <zone_name> \
--worker=all \
--command 'WANDB_API_KEY=<wandb_api_key> levanter/infra/launch.sh python levanter/examples/gpt2_example.py --config_path levanter/config/gpt2_1536_pile.yaml --trainer.beta1 0.965 --trainer.beta2 0.99 --trainer.min_lr_ratio 0.020 --trainer.weight_decay 0.15 --trainer.learning_rate 2.5e-4 --trainer.warmup_ratio 0.01'
致谢
GPT-2 训练代码基于 nanoGPT，该代码优雅且效率极高。
