MARS：释放方差缩减的力量，用于训练大型模型

本仓库包含论文《MARS：释放方差缩减的力量，用于训练大型模型》的官方代码，论文链接为：https://arxiv.org/abs/2411.10438。

作者：袁慧卓*、刘一峰*、吴爽、周勋、顾全全

🔔 最新消息

• [2025年04月03日] MARS-AdamW CUDA 实现已上线。
• [2025年01月05日] 我们的论文已被 ICML 2025 接收 🎉🎉。
• [2025年10月02日] 我们在 ArXiv 上更新了论文：https://arxiv.org/pdf/2411.10438v2。
• [2025年12月01日] 更新了复现 GPT-2 XL 和 FineWeb-Edu 结果的脚本。
• [2025年12月01日] 我们在 FineWeb-Edu 上的预训练结果已公开。GPT-2 XL 在 500 亿 token 的训练下，Hellaswag 准确率达到 56.52%。
• [2024年11月26日] 新增视觉任务支持。
• [2024年11月18日] 我们的代码正式开源！
• [2024年11月15日] 我们的论文已在 ArXiv 上发布：https://arxiv.org/abs/2411.10438。

关于 MARS

MARS（Make vAriance Reduction Shine）是一个统一的优化框架，旨在解决训练大型模型时固有的挑战。传统的自适应梯度方法，如 Adam 和 AdamW，通常面临较高的随机梯度方差问题；而方差缩减技术在深度学习中却难以发挥实际作用。MARS 的核心由两个主要部分组成：(1) 缩放的随机递归动量，它提供了一种方差缩减的全梯度估计器，从而改善梯度复杂度；(2) 预条件更新，它近似二阶牛顿法，以降低每轮迭代的计算复杂度。通过将预条件梯度方法与方差缩减相结合，MARS 实现了两者的最佳结合，加速了优化过程中关键点的搜索。

MARS 框架基于以下预条件方差缩减更新公式：

$$ \mathbf{c}_t = \nabla f(\mathbf{x}_t, \mathbf{\xi}_t)+\underbrace{{\color{red}\gamma_t} \frac{\beta_{1}}{1-\beta_{1}} \left(\nabla f(\mathbf{x}_t, \mathbf{\xi}_t)-\nabla f(\mathbf{x}_{t-1}, \mathbf{\xi}_t)\right)}_{\text{缩放梯度修正}} $$

$$ \tilde{\mathbf{c}}_t = \text{Clip}(\mathbf{c}_t,1) = \begin{cases} \frac{\mathbf{c}_t}{\|\mathbf{c}_t\|_2} & \text{若 } \|\mathbf{c}_t\|_2 > 1,\ \mathbf{c}_t & \text{否则}. \end{cases} $$

$$ \mathbf{m}_t = \beta_1 \mathbf{m}_{t-1} + (1-\beta_{1})\tilde{\mathbf{c}}_t $$

$$ \mathbf{x}_{t+1} = \arg\min_{\mathbf{x} \in \mathbb{R}^d} \left\{\eta_t \left\langle \mathbf{m}_t, \mathbf{x} \right\rangle + \frac{1}{2} \|\mathbf{x} - \mathbf{x}_t \|_{\mathbf{H}_t}^2\right\} $$

其中 ${\color{red}\gamma_t}$ 是控制梯度修正强度的缩放参数。

MARS 的具体实现

在 MARS 框架下，我们提供了三种基于不同 Hessian 矩阵近似的实现：MARS-AdamW、MARS-Lion 和 MARS-Shampoo。请注意，该框架中的超参数是基于 MARS-AdamW 调优的。在使用其他实现时，务必对超参数（尤其是学习率）进行调优，以获得最佳性能。

MARS-AdamW

（在 mars.py 中通过设置 mars_type="mars-adamw" 启用）

Hessian 矩阵的近似定义如下：

$$ \mathbf{v}_t =\beta_2 \mathbf{v}_{t-1}+(1-\beta_2) \big(\nabla f(\mathbf{x}_t, \mathbf{\xi}_t)\big)^2 $$

$$ \mathbf{H}_t := \sqrt{\text{diag}\Big(\mathbf{v}_t\Big)}\cdot \frac{1 - \beta_1^t}{\sqrt{1 - \beta_2^t}}. $$

MARS-Lion

（在 mars.py 中通过设置 mars_type="mars-lion" 启用）

Hessian 矩阵的近似定义如下：

$$ \mathbf{H}_t := \sqrt{\text{diag}(\mathbf{m}_t^2)}. $$

MARS-Shampoo

（在 mars.py 中通过设置 mars_type="mars-shampoo" 启用）

预条件算子可以被视为一个 正交映射 操作符：

$$ \mathbf{U}_t, \mathbf{\Sigma}_t, \mathbf{V}_t = \text{SVD}(\mathbf{G}_t),\qquad \mathbf{x}_{t+1} =\mathbf{x}_t-\eta_t\mathbf{U}_t\mathbf{V}_t^\top. $$

在实践中，我们使用 Newton-Schulz 迭代法 来加速并近似求解 SVD 问题。

MARS 相比基线的表现

OpenWebText 数据集上的实验

MARS 的实验结果基于 MARS-AdamW 实现，除非另有说明。在我们的实验中，梯度是在每个样本和每次更新时计算一次的（即 MARS-approx 形式，详见我们的论文 https://arxiv.org/abs/2411.10438）。如果采用精确形式的 MARS，即每次更新进行两次梯度计算，则性能会略有提升，但计算开销也会翻倍。更多细节请参阅我们的论文。

MARS 在 GPT-2 各个模型上均持续优于 AdamW 和 Muon 优化器：

GPT-2 small	GPT-2 medium	GPT-2 large

最佳验证损失	GPT-2 Small (5B tokens)	GPT-2 Medium (5B tokens)	GPT-2 Large (5B tokens)	GPT-2 Small (20B tokens)	GPT-2 Medium (20B tokens)	GPT-2 Large (20B tokens)	GPT-2 Small (50B tokens)	GPT-2 Medium (50B tokens)	GPT-2 Large (50B tokens)
AdamW	3.193	3.084	3.013	3.024	2.821	2.741	2.885	2.691	2.561
Muon	3.165	3.009	2.915	3.006	2.813	2.691	2.901	2.688	2.573
MARS-exact	3.107	-	-	2.980	-	-	2.847	-	-
MARS-approx	3.108	2.969	2.876	2.981	2.763	2.647	2.849	2.636	2.518

MARS 的效率

MARS 算法不仅可以在相同的训练步数内取得更好的性能，还可以在相同的训练时间内实现更优的效果：

GPT-2 small	GPT-2 medium	GPT-2 large

---

FineWeb-Edu 数据集上的实验

以下是使用我们的 MARS 方法与 AdamW 相比时，GPT‑2 Small 和 GPT‑2 XL 的训练和验证损失曲线。可以看出，MARS 往往能带来更快的收敛速度，并且在不同训练步数下始终表现出更低的损失。

模型	GPT-2 small	GPT-2 XL
训练损失
验证损失

评估指标

下面我们展示了 GPT‑2 Small 和 GPT‑2 XL 在 FineWeb-Edu 数据集上的评估指标，对比了 OpenAI GPT2 基线、AdamW 以及我们的 MARS-AdamW 优化器。

GPT-2 small 的结果

MARS-AdamW 在多个任务上均明显优于 AdamW 和 OpenAI 基线，在 GPT‑2 Small 上以 45.93 的最高平均分领先。

方法/任务	ARC-E	ARC-C	BoolQ	HellaSwag	OBQA	PIQA	WG	MMLU	SciQ	平均
OpenAI-Comm.	39.48	22.70	48.72	31.14	27.20	62.51	51.62	22.92	64.40	41.19
AdamW	51.43	26.54	55.78	36.26	30.60	64.53	50.36	24.49	71.50	45.72
MARS-AdamW	52.23	27.39	55.84	36.91	32.20	64.80	49.96	22.95	71.10	45.93

GPT-2 XL 的结果

在 GPT‑2 XL 上，MARS-AdamW 继续在大多数任务上超越 AdamW，尤其在 HellaSwag 任务中达到了 56.52% 的准确率。

方法/任务	ARC-E	ARC-C	BoolQ	HellaSwag	OBQA	PIQA	WG	MMLU	SciQ	平均
OpenAI-Comm.	51.05	28.50	61.77	50.89	32.00	70.51	58.33	25.24	76.00	50.48
AdamW	68.22	38.40	61.13	53.93	39.00	72.69	54.78	25.47	85.30	55.43
MARS-AdamW	66.54	39.85	63.82	56.52	41.20	73.34	56.59	23.86	86.00	56.41

---

视觉任务上的实验

在 CIFAR-10 和 CIFAR-100 数据集上，使用 ResNet-18 和 MultiStepLR(optimzer, milestones=[100, 150], gamma=0.1) 调度器时，MARS 能够获得比 AdamW 和 Muon 更好的测试损失和准确率。（我们通过网格搜索基础学习率 [1e-5, ..., 1e-1] 来展示每个优化器的最佳结果）：

数据集	CIFAR-10	CIFAR-100
测试损失
测试准确率

最佳测试损失	CIFAR-10	CIFAR-100
AdamW	0.306	2.608
Muon	0.230	1.726
MARS-approx	0.199	0.971

最佳测试准确率 (%)	CIFAR-10	CIFAR-100
AdamW	94.81	73.7
Muon	95.08	74.64
MARS-approx	95.29	76.97

从头开始训练 GPT-2：

安装依赖

$ pip install torch==2.1.2 transformers==4.33.0 datasets tiktoken numpy==1.26.4 wandb

数据准备

按照 nanoGPT 的方式准备 OpenWebText 数据：

$ python data/openwebtext/prepare.py

开始训练

要使用 MARS 优化器训练模型，请运行以下命令：

$ torchrun --standalone --nproc_per_node=8 MARS/train_mars.py config/${your_config_file}

该命令将使用 MARS 优化器在 OpenWebText 数据集上启动 GPT-2 模型的训练。所有相关超参数——包括训练、模型和优化器——均在配置文件（${your_config_file}）中指定。这些参数可以直接在配置文件中调整，也可以通过 Bash 脚本进行修改。

超参数详情

模型超参数：

• n_layer：网络层数，GPT2 Small 为 12 层，GPT2 Medium 为 24 层，GPT2 Large 为 36 层。
• n_head：注意力头数，GPT2 small 为 12 个，GPT2 Medium 为 16 个，GPT2 Large 为 20 个。
• n_embd：嵌入维度，GPT2 small 为 768，GPT2 Medium 为 1024，GPT2 Large 为 1280。

优化器超参数：

• learning_rate：MARS 优化器的学习率。
• weight_decay：MARS 优化器的权重衰减。
• beta1, beta2：指数移动平均的权重。
  • 默认值：beta1=0.95, beta2=0.99
• mars_type：使用的优化器类型：
  • 选项：mars-adamw、mars-lion、mars-shampoo
  • 默认值：mars-adamw
• optimize_1d：是否让 MARS 优化一维参数（例如 GPT-2 中的层归一化参数）。
  • 如果为 False，则使用 AdamW 来优化一维参数。
  • 默认值：False
• lr_1d：当 optimize_1d 设置为 False 时，AdamW 的学习率。
• betas_1d：AdamW 优化器中指数移动平均的权重。
  • 默认值：(0.9, 0.95)
• is_approx：是否使用近似梯度计算（MARS-approx）。
  • 默认值：True
• gamma：控制梯度修正强度的缩放参数。
  • 默认值：0.025

训练超参数：

• batch_size：每个设备上的小批量大小。（例如，在 A100 GPU 上训练 GPT-2 Small 时，通常使用 15 的批量大小。）
• gradient_accumulation_steps：梯度累积步数，以确保总有效批量大小达到所需规模。（例如，对于总批量大小为 480 的情况：$15 \times 4 \times 8 , \text{GPUs}$。）
• schedule：学习率调度策略。
  • 默认值：cosine

更多详细的超参数示例，请参阅：

• config/train_gpt2_small_mars.py
• scripts/run_mars_small.sh

---

复现我们的结果

复现 GPT-2 Small (125M) 的结果

使用以下命令训练 MARS：

$ bash scripts/run_mars_small.sh

或者

$ torchrun --standalone --nproc_per_node=8 \
      MARS/train_mars.py \
      config/train_gpt2_small_mars.py \
      --batch_size=15 \
      --gradient_accumulation_steps=4

复现 GPT2 Medium (355M) 的结果

使用以下命令训练 MARS：

$ bash scripts/run_mars_medium.sh

或者

$ torchrun --standalone --nproc_per_node=8 \
      MARS/train_mars.py \
      config/train_gpt2_medium_mars.py \
      --batch_size=15 \
      --gradient_accumulation_steps=4

复现 GPT2 Large (770M) 的结果

使用以下命令训练 MARS：

$ bash scripts/run_mars_large.sh

或者

$ torchrun --standalone --nproc_per_node=8 \
      MARS/train_mars.py \
      config/train_gpt2_large_mars.py \
      --batch_size=5 \
      --gradient_accumulation_steps=12

在 FineWeb-Edu 上复现 GPT-2 XL (1.5B) 的结果

$ bash scripts/run_mars_xl_fw.sh

或者

$ torchrun --standalone --nproc_per_node=8 \
      MARS/train_mars_fw.py \
      config/train_gpt2_xl_mars.py \
      --batch_size=5 \
      --gradient_accumulation_steps=12

复现基线结果

要复现 AdamW 基线：

bash scripts/run_adamw_{small/medium/large}.sh

要在 FineWeb-Edu 上复现 AdamW 基线：

bash scripts/run_adamw_{small/xl}_fw.sh

要复现 Muon 基线，遵循 modded-nanogpt 的方法：

bash scripts/run_muon_{small/medium/large}.sh

请根据您使用的硬件配置相应调整 nproc_per_node、batch_size 和 gradient_accumulation_steps，确保它们的乘积等于 480。

GPT-2 模型的超参数

模型名称	模型大小	AdamW 学习率	Muon 学习率	MARS 学习率	MARS 一维学习率	AdamW 权重衰减	Muon 权重衰减	MARS 权重衰减
GPT-2 small	125M	6e-4	2e-2	6e-3	3e-3	1e-1	0.0	1e-2
GPT-2 medium	355M	3e-4	1e-2	3e-3	1.5e-3	1e-1	0.0	1e-2
GPT-2 large	770M	2e-4	6.67e-3	2e-3	1e-3	1e-1	0.0	1e-2
GPT-2 xl	1.5B	2e-4	-	2e-3	1e-3	1e-1	-	1e-2

自定义训练

要构建针对其他架构和数据集的自定义训练流程，可以参考以下模板：

import torch
import torch.nn.functional as F
from mars import MARS

# 初始化模型损失函数和输入数据
model = Model()
data_loader = ...

# 初始化优化器
optimizer = MARS(model.parameters(), lr=1e-3, betas=(0.9, 0.95), gamma=0.025)

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)
        optimizer.update_last_grad()
        iter_num += 1

星标历史

引用

如果您觉得本仓库对您的研究有所帮助，请考虑引用以下论文：

@article{yuan2024mars,
  title={MARS: Unleashing the Power of Variance Reduction for Training Large Models},
  author={Yuan, Huizhuo and Liu, Yifeng and Wu, Shuang and Zhou, Xun and Gu, Quanquan},
  journal={arXiv preprint arXiv:2411.10438},
  year={2024}
}

致谢

本仓库基于 nanoGPT、levanter 和 Sophia 构建，我们感谢这些作者的杰出工作！

MARS 快速上手指南

MARS (Make vAriance Reduction Shine) 是一个统一的优化框架，旨在解决大模型训练中的挑战。它结合了方差缩减技术与预条件更新方法，在相同训练步数或时间内，相比 AdamW 和 Muon 等优化器能实现更快的收敛速度和更低的损失值。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 20.04+)
• Python 版本: 3.8 - 3.10
• 硬件要求: NVIDIA GPU (支持 CUDA)，建议显存 16GB 以上以运行大模型实验
• 核心依赖:
  • PyTorch >= 2.1.2
  • Transformers >= 4.33.0
  • CUDA Toolkit (需与 PyTorch 版本匹配)

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

安装步骤

1. 克隆项目代码

git clone https://github.com/your-repo/MARS.git
cd MARS

2. 创建虚拟环境并安装依赖

推荐使用 conda 或 venv 隔离环境。以下命令使用国内镜像源安装指定版本的 PyTorch 和其他依赖：

# 创建虚拟环境
conda create -n mars_env python=3.9 -y
conda activate mars_env

# 安装 PyTorch (使用清华镜像，根据实际 CUDA 版本调整，此处以 cu118 为例)
pip install torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 --index-url https://pypi.tuna.tsinghua.edu.cn/simple

# 安装其他项目依赖
pip install transformers==4.33.0 datasets accelerate wandb matplotlib numpy scipy -i https://pypi.tuna.tsinghua.edu.cn/simple

注意：如果需要使用 MARS-AdamW CUDA 加速版（2025/03/04 更新），请确保已正确安装 CUDA 开发工具包，并按照仓库内的 cuda_extension 说明编译自定义算子（如有）。

基本使用

MARS 提供了三种实例化变体：MARS-AdamW（推荐默认）、MARS-Lion 和 MARS-Shampoo。大多数情况下，直接使用 MARS-AdamW 即可获得最佳效果。

1. 在 Python 代码中调用

将原有的优化器（如 torch.optim.AdamW）替换为 MARS 优化器。

from mars import MARS

# 假设 model 是你的神经网络模型
model = ... 

# 初始化 MARS 优化器
# mars_type 可选："mars-adamw" (默认), "mars-lion", "mars-shampoo"
optimizer = MARS(
    model.parameters(), 
    lr=1e-3,              # 学习率可能需要重新调优，通常比 AdamW 稍大
    betas=(0.9, 0.999),   # beta1, beta2
    weight_decay=0.01,
    mars_type="mars-adamw" 
)

# 标准训练循环
for batch in dataloader:
    optimizer.zero_grad()
    loss = compute_loss(model, batch)
    loss.backward()
    optimizer.step()

2. 关键参数说明

• mars_type: 选择变体。
  • "mars-adamw": 基于对角海森矩阵近似，最稳定，适合大多数 NLP 和 Vision 任务。
  • "mars-lion": 使用动量的平方作为近似，计算更轻量。
  • "mars-shampoo": 使用 SVD 分解进行正交映射，适合特定大规模分布式场景（内部使用 Newton-Schulz 迭代加速）。
• lr (学习率): MARS 的超参数敏感性与 AdamW 不同。官方建议在迁移使用时，针对新任务对学习率进行网格搜索（范围通常在 1e-5 到 1e-1 之间）。
• gamma_t: 梯度校正的缩放参数，已在内部实现动态调整，通常无需手动设置。

3. 复现 GPT-2 训练结果

项目提供了完整的脚本用于复现 OpenWebText 和 FineWeb-Edu 上的 GPT-2 训练结果。

# 示例：运行 GPT-2 Small 在 FineWeb-Edu 上的训练 (使用 MARS-AdamW)
python train_gpt2.py \
    --dataset fineweb-edu \
    --model_size small \
    --optimizer mars-adamw \
    --total_tokens 50B \
    --batch_size 512

提示：对于视觉任务（如 CIFAR-10/100），请参考 examples/vision/ 目录下的脚本，配合 MultiStepLR 调度器使用可获得优于 AdamW 的准确率。