优化器

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

本指南将向您展示如何结合 TrainingArguments 使用 Trainer 来使用这些优化器，如下所示。

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

Approximated Gradient Scaling for Memory Efficient LLM Optimization (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
rank	用于梯度缩放的辅助子空间的秩	256	1
scale_type	缩放因子如何应用	`channel`（逐通道缩放）	`tensor`（逐张量缩放）
scale	调整梯度更新以稳定训练	1.0	128
update_proj_gap	更新投影矩阵之前的步数	200	200
proj	投影类型	`random`	`random`

以下示例启用 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

Low-Memory Optimization (LOMO) 是一系列优化器，LOMO 和 AdaLomo，专为 LLM 的低内存全参数微调而设计。LOMO 优化器和 AdaLomo 优化器都将梯度计算和参数更新融合在一个步骤中，以减少内存使用。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",
)

Schedule Free

pip install schedulefree

Schedule Free optimizer (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 上更新