オプティマイザーとスケジューラー

オプティマイザーは、学習中にモデルの重みを更新します。スケジューラーはオプティマイザーをラップし、各学習ステップで学習率を調整します。Trainerは、create_optimizer_and_scheduler()を呼び出す際に、両方を生成します。

                                    ┌────────────┐         ┌──────────────┐
                                    │ Optimizer  │         │  Scheduler   │
                                    │ (adamw_torch_fused)◄─│  (linear)    │
                                    │            │         │              │
                                    │ param_groups         |              |
                                    │  └ lr       ◄────────┤              |
                                    │  └ weight_decay      │              │
                                    └──────┬─────┘         └──────────────┘
                                           │                      
  ┌──── EACH TRAINING STEP ───────────────────────────────────────────┐
  │                                        │                          │
  │   model(batch)                         │                          │
  │       │                                │                          │
  │       ▼                                │                          │
  │     loss ──► loss.backward() ──► param.grad                       │
  │                                        │                          │
  │                          ┌─────────────┘                          │
  │                          ▼                                        │
  │              optimizer.step()                                     │
  │                          │                                        │
  │                          ▼                                        │
  │                   param.data updated                              │
  │                          │                                        │
  │                          ▼                                        │
  │              lr_scheduler.step()  ──► recalculates lr             │
  │                          │            writes to optimizer         │
  │                          ▼            .param_groups['lr']         │
  │              model.zero_grad()                                    │
  │                                                                   │
  └───────────────────────────────────────────────────────────────────┘

lr_scheduler_type() や optim() のようなオプティマイザーとスケジューラーの動作は、TrainingArguments で設定します。デフォルト値（adamw_torch オプティマイザーと linear ウォームアップスケジューラー）は、ほとんどのファインチューニング実行において良い出発点となります。

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    # Optimizer
    optim="adamw_torch",          # or "adamw_torch_fused", "adafactor", "sgd", etc.
    learning_rate=2e-5,
    weight_decay=0.01,
    adam_beta1=0.9,
    adam_beta2=0.999,
    adam_epsilon=1e-8,
    # Scheduler
    lr_scheduler_type="cosine",   # "linear", "cosine", "constant_with_warmup", etc.
    warmup_steps=500,
    lr_scheduler_kwargs={"num_cycles": 3},  # scheduler-specific extras
)

メトリクスベースのスケジューラ

一部のスケジューラーは、固定のスケジュールに従うのではなく、トレーニングの動的特性に適応します。

GreedyLR は、評価結果に基づいて学習率を更新します。指標が改善し続ける場合、学習率をfactorで割ることで上げ、指標が改善しない場合、学習率をfactor倍することで下げます。学習率がmin_lrで停止し、reset_startステップ経っても改善しない場合、GreedyLR は初期状態にリセットされ、新しいサイクルを開始します。

GreedyLR は、学習中に評価が必要です。eval_strategy を "steps" または "epoch" に設定してください。

args = TrainingArguments(
+   lr_scheduler_type="greedy",
+   lr_scheduler_kwargs={"patience": 10, "factor": 0.95, "min_lr": 1e-5},
+   eval_strategy="steps",
+   eval_steps=200,
    ...  # remaining args from the TrainingArguments intro config
)

デフォルトのmode="min"は損失に適用できます。精度のように値が大きいほど良い指標を追跡する場合は、lr_scheduler_kwargsに"mode": "max"を渡してください。

設定可能なパラメータの完全なリストについては、GreedyLRクラスを参照してください。

オプティマイザーの統合

Transformers は、特殊なトレーニングシナリオ向けにサードパーティのオプティマイザーを統合しています。

オプティマイザ	インストール	`optim="value"`	説明
APOLLO	`apollo-torch`	`apollo_adamw`	ランダム射影によるメモリ効率の高いフルパラメータ学習; ランク1で十分
GrokAdamW	`grokadamw`	`grokadamw`	遅延汎化(グロッキング)を対象とします
LOMO / AdaLomo	`lomo-optim`	`lomo` / `adalomo`	低メモリでのフルパラメータファインチューニングのために、勾配と更新ステップを融合します
スケジュールフリー	`schedulefree`	`schedule_free_adamw`, `schedule_free_radam`, `schedule_free_sgd`	LR アニーリングを排除します。`lr_scheduler_type="constant"` と組み合わせて使用してください。
GaLore	`galore-torch`	`galore_adamw`, `galore_adafactor`, `galore_adamw_8bit`	勾配低ランク射影による全パラメータ学習
StableAdamW	`torch-optimi`	`stable_adamw`	AdamW + AdaFactorの更新クリッピング; グラデントクリッピングは不要

APOLLO

GrokAdamW

LOMO

スケジュールフリー

StableAdamW

GaLore

pip install apollo-torch

Approximated Gradient Scaling for Memory Efficient LLM Optimization (APOLLO)は、事前学習およびファインチューニングにおける全パラメータ学習のためのメモリ効率の高いオプティマイザーです。SVDの代わりに安価なランダム射影を使用することで、SGDと同程度のメモリコストでAdamWと同等の性能を実現します。極限的なメモリ削減が必要な場合は、ランク1のバリアントであるAPOLLO-Miniを使用してください。

optim_target_modules パラメータを使用して、学習対象のレイヤーを指定します。

args = TrainingArguments(
+   optim="apollo_adamw",
+   optim_target_modules=[r".*.attn.*", r".*.mlp.*"],
    ...  # remaining args from the TrainingArguments intro config
)

optim_args を通して、追加のハイパーパラメータを渡します。

scale を n/r に設定します。ここで n は元の空間の次元、r は低ランク空間の次元です。scale をデフォルトのままにしておいて学習率を調整しても、同様の効果が得られます。

パラメータ	説明	APOLLO	APOLLO-Mini
ランク	勾配スケーリング用の補助部分空間のランク	256	1
scale_type	スケーリング係数の適用方法	`channel`（チャネルごとのスケーリング）	`tensor`（テンソルごとのスケーリング）
スケール	勾配の更新を調整して、学習を安定させます	1.0	128
update_proj_gap	投影行列を更新する前のステップ	200	200
proj	projection type	`random`	`random`

rank-1 の設定で APOLLO-Mini を有効にします。

args = TrainingArguments(
    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",
    ...  # remaining args from the TrainingArguments intro config
)

オプティマイザーとスケジューラーのカスタマイズ

まだ統合されていないオプティマイザーを使用したり、レイヤーごとの学習率を調整したり、カスタムロジックを適用したりするために、カスタムオプティマイザーとスケジューラーを作成します。

クラスとkwargsを渡す

~Trainer.optimizer_cls_and_kwargs はカスタムオプティマイザークラスを受け取りますが、パラメータのグループ化とデバイスの割り当ては Trainer に委譲されます。

Trainer は create_optimizer() が実行されるまでオプティマイザーの構築を遅延させるため、モデルはすでに正しいデバイス上に配置されています。

import torch

trainer = Trainer(
    ...
    optimizer_cls_and_kwargs=(
        torch.optim.SGD,
        {"momentum": 0.9, "nesterov": True}
    ),
)

事前構築されたインスタンスを渡す

~Trainer.optimizers に事前定義済みのオプティマイザーとスケジューラーを渡します。事前構築済みのインスタンスが提供されている場合、Trainer は create_optimizer() と create_scheduler() をスキップします。スケジューラーを渡さない場合、Trainer は自動的にスケジューラーを作成します。

モデルを正しいデバイスに配置した後に、オプティマイザーを構築してください。パラメーターは構築時に解決され、Trainerがモデルを移動する前に処理されます。分散学習では、デバイスの不一致がサイレントに誤った動作を引き起こす可能性があります。

import torch
from transformers import Trainer, get_cosine_schedule_with_warmup

optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
scheduler = get_cosine_schedule_with_warmup(
    optimizer, num_warmup_steps=500, num_training_steps=10_000
)

trainer = Trainer(
    ...
    optimizers=(optimizer, scheduler),
)

Prebuilt instances は create_optimizer() および create_scheduler() をバイパスするため、独自の parameter groups を指定する必要があります。

オプティマイザーとスケジューラーのメソッドをオーバーライドする

完全な制御を行うには、create_optimizer() と create_scheduler() をサブクラス化します。これらのメソッドは train() の実行中に呼び出されます。

SchedulerType で利用できない OneCycleLR などのスケジューラを使用するには、create_scheduler() をオーバーライドします。

各メソッドでは、selfに代入し、それを返すようにしてください。

import torch
from transformers import Trainer

class MyTrainer(Trainer):

    def create_scheduler(self, num_training_steps, optimizer=None):
        optimizer = optimizer or self.optimizer
        self.lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
            optimizer,
            max_lr=0.1,
            total_steps=num_training_steps,
        )
        return self.lr_scheduler

デフォルトのオプティマイザーで問題がなければ、create_optimizer()をオーバーライドする必要はありません。メソッドを完全に置き換えるよりも、super()でメソッドを拡張する方が簡単です。例えば、他の設定はすべてそのままに、追加のパラメータグループを1つ追加するといったことができます。

class MyTrainer(Trainer):
    def create_optimizer(self, model=None):
        super().create_optimizer(model)  # builds the default two param groups
        # add extra param group
        self.optimizer.add_param_group({
            "params": self.model.classifier.parameters(),
            "lr": self.args.learning_rate * 10,
        })
        return self.optimizer

Update on GitHub