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的学习率并不直接可比。经验上，Adam和Lion以5:1的学习率比例时行为相似。SophiaG和Lion的学习率则可以直接比较。Sophia允许使用比Lion大得多的学习率，这也是Sophia速度更快的原因。

调整超参数$\rho$

• 调整$\rho$使被截断的坐标比例保持稳定且处于适当范围。这在GPT-2训练示例中作为train/win_rate进行跟踪。train/win_rate应在训练初期达到峰值，并在之后保持稳定。其值应维持在0.1至0.5之间。通常，较大的$\rho$会导致较高的train/win_rate。下文提供了T5模型中典型win_rate行为的示例。

调整学习率和权重衰减

• 选择略小于您用于AdamW的学习率，或为用于Lion的学习率的3至5倍。

• 如果损失发散，可适当降低学习率或增大$\rho$。

• 始终使用约为AdamW两倍的权重衰减。

GPT-2模型的超参数

• 选择与您用于AdamW的学习率相近，或为用于Lion的学习率的5至10倍。
• 调整$\rho$使未被截断的参数比例保持稳定且处于适当范围。这在GPT-2训练示例中作为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（1.25亿参数）：

如果你的机器有 10 张 A5000（24GB 显存）显卡，

$ 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 显存）显卡，

$ torchrun --standalone --nproc_per_node=8 \
      train_sophiag.py \
      config/train_gpt2_small_sophiag.py \
      --batch_size=12 \
      --gradient_accumulation_steps=5

要复现 nanoGPT 中的 AdamW 基线，

$ torchrun --standalone --nproc_per_node=10 \
      train_adam.py \
      config/train_gpt2_small_adam.py \
      --batch_size=8 \
      --gradient_accumulation_steps=6

这将得到下图所示的结果：

开始预训练 GPT2 Medium（3.55亿参数）：

如果你的机器有 8 张 A100（40GB 显存）显卡，

$ 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 15亿参数模型：

我们使用 The Pile 数据集和 GPT NeoX 分词器。首先按照 levanter 的说明设置 TPU 实例及环境。然后在 optim.py 中将 GAMMA_SOPHIA_G 修改为 200。15亿参数模型的训练脚本如下：

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，其代码简洁高效，令人赞叹。

Sophia 优化器快速上手指南

Sophia 是一种可扩展的随机二阶优化器，专为语言模型预训练设计。相比 AdamW 和 Lion，Sophia 允许使用更大的学习率，从而显著加速训练过程。本指南基于官方实现的 Sophia-G 版本。

环境准备

系统要求

• Python: 建议 Python 3.8+
• GPU: 支持 CUDA 的 NVIDIA GPU（推荐 A100/A5000 或多卡环境）
• 框架: PyTorch 2.1.2

前置依赖

请确保安装以下核心依赖库：

• torch (2.1.2)
• transformers (4.33.0)
• datasets
• tiktoken
• wandb (用于监控训练指标)

国内加速建议：安装时推荐使用清华或阿里镜像源以提升下载速度。

安装步骤

1. 创建虚拟环境（可选但推荐）

   python -m venv sophia_env
   source sophia_env/bin/activate  # Linux/Mac
   # 或
   sophia_env\Scripts\activate     # Windows
   

2. 安装 PyTorch (指定版本)

   pip install torch==2.1.2 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
   # 国内用户可使用清华源：
   # pip install torch==2.1.2 torchvision torchaudio --index-url https://pypi.tuna.tsinghua.edu.cn/simple
   

3. 安装其他依赖

   pip install transformers==4.33.0 datasets tiktoken wandb
   

4. 获取 Sophia 代码 克隆官方仓库并将 sophia 模块放入你的项目目录：

   git clone https://github.com/Liuhong99/Sophia.git
   cp -r Sophia/sophia ./your_project/
   

基本使用

以下是使用 SophiaG 训练通用模型的最小化代码示例。主要区别在于初始化优化器以及在训练循环中定期更新海森矩阵（Hessian EMA）。

import torch
import torch.nn.functional as F
from sophia import SophiaG

# 1. 初始化模型、损失函数和数据加载器
model = Model()
data_loader = ... # 你的 DataLoader

# 2. 初始化 SophiaG 优化器
# 注意：lr 通常设为 AdamW 的同等大小，或 Lion 的 3-5 倍
# rho 控制裁剪比例，建议初始设为 0.01 - 0.1 之间并根据 win_rate 调整
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  # 每 k 步更新一次海森矩阵估计
iter_num = -1

# 3. 训练循环
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

        # 每隔 k 步更新海森矩阵 EMA (关键步骤)
        if iter_num % k != k - 1:
            continue
        else:
            # 采样计算二阶信息
            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()

关键参数调优建议

• 学习率 (lr):
  • 对比 AdamW：使用相同或略小的值。
  • 对比 Lion：使用 3-5 倍的值（GPT-2 场景下可达 5-10 倍）。
• Rho (rho):
  • 目标是将 train/win_rate（未裁剪参数的比例）稳定在 0.1 - 0.5 之间。
  • 如果损失发散（blow up），尝试减小 lr 或增大 rho。
• 权重衰减 (weight_decay):
  • 建议使用 AdamW 默认值的 2 倍（例如 GPT-2 训练中常用 0.2）。