优化器

Transformers 提供了两种原生优化器：AdamW 和 AdaFactor。它还提供了更多专用优化器的集成。安装提供优化器的库，并将其放置在 TrainingArguments 的 `optim` 参数中。

本指南将向您展示如何使用 Trainer 和 TrainingArguments（如下所示）使用这些优化器。

import torch
from transformers import TrainingArguments, AutoTokenizer, AutoModelForCausalLM, Trainer

args = TrainingArguments(
    output_dir="./test-optimizer",
    max_steps=1000,
    per_device_train_batch_size=4,
    logging_strategy="steps",
    logging_steps=1,
    learning_rate=2e-5,
    save_strategy="no",
    run_name="optimizer-name",
)

APOLLO

pip install apollo-torch

用于高效 LLM 优化的近似梯度缩放 (APOLLO) 是一种内存高效的优化器，允许对预训练和微调进行全参数学习。它在保持 AdamW 级别的性能的同时，具有类似 SGD 的内存效率。为了实现极致的内存效率，您可以使用 APOLLO-Mini，这是 APOLLO 的秩 1 变体。APOLLO 优化器支持：

超低秩效率。您可以比 GaLoRE 使用更低的秩，秩 1 就足够了。
避免昂贵的 SVD 计算。APOLLO 利用随机投影来避免训练停滞。

使用 `optim_target_modules` 参数指定要训练的层。

import torch
from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="./test-apollo",
    max_steps=100,
    per_device_train_batch_size=2,
+   optim="apollo_adamw",
+   optim_target_modules=[r".*.attn.*", r".*.mlp.*"],
    logging_strategy="steps",
    logging_steps=1,
    learning_rate=2e-5,
    save_strategy="no",
    run_name="apollo_adamw",
)

对于其他训练选项，使用 `optim_args` 定义超参数，如 `rank`、`scale` 等。请参阅下表，获取可用超参数的完整列表。

`scale` 参数可以设置为 `n/r`，其中 `n` 是原始空间维度，`r` 是低秩空间维度。您可以通过调整学习率同时保持 `scale` 为默认值来达到类似的效果。

参数	描述	APOLLO	APOLLO-Mini
秩	用于梯度缩放的辅助子空间的秩	256	1
缩放类型	缩放因子的应用方式	`channel` (每通道缩放)	`tensor` (每张量缩放)
缩放	调整梯度更新以稳定训练	1.0	128
更新投影间隔	更新投影矩阵之前的步数	200	200
投影	投影类型	`随机`	`随机`

以下示例启用 APOLLO-Mini 优化器。

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="./test-apollo_mini",
    max_steps=100,
    per_device_train_batch_size=2,
    optim="apollo_adamw",
    optim_target_modules=[r".*.attn.*", r".*.mlp.*"],
    optim_args="proj=random,rank=1,scale=128.0,scale_type=tensor,update_proj_gap=200",
)

GrokAdamW

pip install grokadamw

GrokAdamW 是一种优化器，旨在帮助那些受益于 *grokking* 的模型，该术语用于描述由于梯度变化缓慢而导致的延迟泛化。它对于需要更高级优化技术以实现更好性能和稳定性的模型特别有用。

import torch
from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="./test-grokadamw",
    max_steps=1000,
    per_device_train_batch_size=4,
+   optim="grokadamw",
    logging_strategy="steps",
    logging_steps=1,
    learning_rate=2e-5,
    save_strategy="no",
    run_name="grokadamw",
)

LOMO

pip install lomo-optim

低内存优化 (LOMO) 是一系列优化器，包括 LOMO 和 AdaLomo，旨在用于 LLM 的低内存全参数微调。LOMO 优化器将梯度计算和参数更新融合在一个步骤中以减少内存使用。AdaLomo 在 LOMO 的基础上，通过为每个参数引入自适应学习率，类似于 Adam 优化器。

建议使用 AdaLomo，且不带 `grad_norm`，以获得更好的性能和更高的吞吐量。

args = TrainingArguments(
    output_dir="./test-lomo",
    max_steps=1000,
    per_device_train_batch_size=4,
+   optim="adalomo",
    gradient_checkpointing=True,
    gradient_checkpointing=True,
    logging_strategy="steps",
    logging_steps=1,
    learning_rate=2e-6,
    save_strategy="no",
    run_name="adalomo",
)

无调度优化器

pip install schedulefree

无调度优化器 (SFO) 将基本优化器的动量替换为平均和插值的组合。与传统调度器不同，SFO 完全消除了退火学习率的需要。

SFO 支持 RAdam (`schedule_free_radam`)、AdamW (`schedule_free_adamw`) 和 SGD (`schedule_free_sgd`) 优化器。RAdam 调度器不需要 `warmup_steps` 或 `warmup_ratio`。

默认情况下，建议将 `lr_scheduler_type` 设置为 `"constant"`。其他 `lr_scheduler_type` 值也可能有效，但将 SFO 优化器与其他学习率调度器结合使用可能会影响 SFO 的预期行为和性能。

args = TrainingArguments(
    output_dir="./test-schedulefree",
    max_steps=1000,
    per_device_train_batch_size=4,
+   optim="schedule_free_radamw,
+   lr_scheduler_type="constant",
    gradient_checkpointing=True,
    logging_strategy="steps",
    logging_steps=1,
    learning_rate=2e-6,
    save_strategy="no",
    run_name="sfo",
)

< > 在 GitHub 上更新