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'