人工智能：TrainingArguments 中 optim 参数详解

Posted on Edited on In Waline:

这份指南详细介绍了 Hugging Face transformers 库中 optim 参数所有可用的选项。该列表直接来源于 transformers 源代码，涵盖了从稳定可靠的基准到内存优化、实验性算法和特定硬件的各类优化器。

优化器分类概览

为了系统地理解这些选项，我们将它们分为以下几大类：

主流与推荐 (AdamW 家族): 训练 Transformer 模型的核心与首选。
* adamw_torch, adamw_torch_fused

内存节省型优化器 (核心技术): 当显存（VRAM）成为主要瓶颈时，这些技术至关重要。
8-bit 量化优化器: 将优化器状态从32位浮点数压缩到8位整数，极大节省内存。
* adamw_8bit (别名 adamw_bnb_8bit), lion_8bit, rmsprop_8bit
* Paged 优化器: 一种更先进的内存管理技术，当GPU显存不足时，它能自动将优化器状态“分页”到CPU内存，防止程序因OOM（Out-of-Memory）而崩溃。
* paged_adamw_8bit, paged_lion_8bit
* Adafactor: 通过数学分解来近似存储优化器状态，是另一种经典的内存节省方案。
* adafactor

梯度低秩投影 (GaLore): 一种前沿的内存优化技术，尤其适合全参数微调。
低内存优化 (LOMO & AdaLOMO): 另一种创新的内存节省方法。
新型与实验性优化器: 来自最新研究，可能在特定任务上表现更优。
特定硬件/库的优化器: 为特定硬件或旧版库设计。
经典与遗留优化器: 深度学习领域的经典算法。

全量选项详细分析

1. 主流与推荐 (AdamW 家族)

这是训练现代大模型的起点和最可靠的选择。

选项 (optim) 核心特点 优点 缺点 依赖与要求
adamw_torch PyTorch 标准实现，默认选项 最稳定、最通用，无需额外依赖，行为可预测。 在现代 GPU 上性能非最优。 PyTorch
adamw_torch_fused 性能之选 (推荐)，融合内核 在现代 NVIDIA GPU 上速度显著提升，通过融合操作减少 GPU 内核启动开销。 需要 PyTorch 2.0 或更高版本。 torch>=2.0

2. 内存节省型优化器

当显存不足时，以下技术是你的救星。它们通常可以组合使用。

2.1 8-bit 量化优化器

技术核心：使用 bitsandbytes 库将32位的优化器状态（如动量和方差）量化为8位整数，从而将优化器的内存占用减少约75%。

选项 (optim) 基础优化器 特点与用途 依赖
adamw_8bit AdamW 8位量化的 AdamW。最常用、最推荐的内存节省方案。 bitsandbytes
(adamw_bnb_8bit) (这是 adamw_8bit 的一个别名)
lion_8bit Lion 8位量化的 Lion 优化器。 bitsandbytes
rmsprop_8bit RMSprop 8位量化的 RMSprop。 bitsandbytes
(rmsprop_bnb_8bit) (这是 rmsprop_8bit 的一个别名)

2.2 Paged 优化器

技术核心：在8-bit优化的基础上，当GPU显存依然不足时，它能利用NVIDIA的统一内存功能，自动将不常用的优化器状态**“分页”到CPU主内存**，并在需要时再调回GPU。这可以防止在训练中因显存波动导致的崩溃 (OOM Error)。

选项 (optim) 基础优化器 特点与用途 依赖
paged_adamw_8bit AdamW 结合了8位量化和分页技术，是目前最稳健的内存优化方案。 bitsandbytes
paged_lion_8bit Lion 结合了8位量化和分页技术的 Lion 优化器。 bitsandbytes
paged_adamw_32bit AdamW 如果你不想使用8位量化但仍想利用分页技术防止崩溃，可以使用这个32位版本。 bitsandbytes

2.3 Adafactor

技术核心：通过数学上的矩阵分解技巧来近似存储优化器状态，从而达到节省内存的目的。

选项 (optim) 核心特点 优点 缺点 依赖
adafactor Google 设计的内存优化器 显著降低内存占用，不存储完整的动量信息，且不依赖 bitsandbytes 库。 收敛性有时不如 AdamW 稳定，可能需要更仔细地调整学习率。 transformers 自带

3. GaLore (Gradient Low-Rank Projection)

技术核心：一种非常新颖且高效的内存优化技术，它假设梯度矩阵是低秩的，因此只存储和更新梯度的低秩投影，而不是完整的梯度矩阵。这使得全参数微调的内存消耗可以接近甚至低于LoRA等参数高效微调方法。

选项 (optim) 基础优化器 特点与用途 依赖
galore_adamw AdamW 将 GaLore 技术应用于 AdamW。 galore_torch
galore_adafactor Adafactor 将 GaLore 技术应用于 Adafactor，进一步节省内存。 galore_torch
galore_adamw_8bit AdamW + 8-bit 结合了 GaLore 和8位量化，实现极致的内存节省。 galore_torch, bitsandbytes
galore_..._layerwise (所有GaLore变体) 为模型的不同层应用不同的秩，可能比全局秩更高效。 galore_torch

4. LOMO (Low-Memory Optimization) & AdaLOMO

技术核心：LOMO通过将梯度计算、梯度收集和参数更新步骤融合成一个步骤，从而避免了存储完整的模型梯度，极大地节省了内存。AdaLOMO是其自适应版本。

选项 (optim) 核心特点 优点 缺点 依赖
lomo 低内存优化 将多步融为一步，无需存储完整梯度，内存效率极高。 实验性强，实现较为复杂。 lomo_optim
adalomo 带自适应学习率的LOMO 在 LOMO 基础上增加了学习率的自适应调整。 实验性强。 lomo_optim

5. 新型与实验性优化器

这些优化器来自最新的学术研究，可能在特定任务上表现更优，但需要更多实验和调整。

选项 (optim) 核心思想 何时使用
lion (或 lion_32bit) 简化更新规则（只用动量），由Google提出。 当你乐于探索时。lion 宣称比 AdamW 更省内存、性能更优。通常需要更小的学习率和更大的权重衰减。
apollo_adamw 解耦学习率和权重衰减来改进 AdamW。 当你希望探索可能比标准AdamW收敛更好、性能更优的替代方案时。
schedule_free_... 在训练中自动调整步长。 当你不想费心设计学习率调度器时。非常适合快速实验，但最终性能不一定最优。
grokadamw 为复现 “Grokking” 现象而设计。 仅用于学术研究，探索模型在训练后期泛化能力突然提升的现象。

6. 特定硬件/库的优化器

选项 (optim) 用途
adamw_torch_xla 仅在 Google TPU 上使用。
adamw_torch_npu_fused 仅在华为 Ascend NPU 上使用。
adamw_apex_fused 仅在 torch<2.0 的旧环境中，配合NVIDIA APEX库使用。已被 adamw_torch_fused 取代。

7. 经典与遗留优化器

这些算法是深度学习的基石，但对于现代 Transformer 模型，通常不是最佳选择。

选项 (optim) 核心特点与局限
sgd 随机梯度下降。对于Transformer模型这样复杂的损失曲面，SGD收敛极慢且容易陷入局部最优。
adagrad 适应性梯度算法。学习率会随时间单调递减，后期可能过小导致训练停滞，不推荐用于Transformer。
rmsprop Adagrad的改进版。缓解了学习率快速下降的问题，但性能和稳定性通常不如AdamW。

终极决策流程

  1. 【性能优先】 如果你有现代NVIDIA GPU和PyTorch 2.0+，adamw_torch_fused 是你的不二之选。

  2. 【内存优先】 如果显存不足是首要问题：

    • 首选 paged_adamw_8bit。它结合了8位量化和分页，既省内存又防崩溃。
    • 如果想在全参数微调时达到LoRA级的内存效率，大胆尝试 galore_adamw_8bit。
    • 如果安装 bitsandbytes 困难，退而求其次选择 adafactor。

  3. 【简化实验】 如果你想快速迭代，不想费心设计学习率调度器：

    • 尝试 schedule_free_adamw。

  4. 【探索前沿】 如果你热衷于尝试最新的研究成果：

    • lion 是一个很好的起点，但请准备好调整超参数。

  5. 【不确定或环境受限】 如果你不确定选哪个：

    • 使用默认的 adamw_torch，它永远是稳定可靠的基准。

  6. 【特定硬件】 如果你在TPU或NPU上：

    • 分别使用 adamw_torch_xla 或 adamw_torch_npu_fused。

官方文档地址

您可以在 Hugging Face transformers 官方文档的 TrainingArguments 页面找到关于 optim 参数的权威说明。官方文档链接: https://huggingface.co/docs/transformers/main/en/main_classes/trainer#transformers.TrainingArguments.optim