使用 Adafactor 的内存高效训练

使用 PyTorch 高效训练 AI 模型

Dennis Lee

Data Engineer, Amazon

高效训练的优化器

 

 

代表 AdamW、Adafactor 和 8-bit Adam 的图标。

使用 PyTorch 高效训练 AI 模型

优化器的权衡

展示 AdamW、Adafactor 与 8-bit Adam 在参数数量与精度间权衡的示意图。

使用 PyTorch 高效训练 AI 模型

Adafactor 如何工作？

 

展示 Adafactor 步骤的示意图。

使用 PyTorch 高效训练 AI 模型

Adafactor 如何工作？

 

展示 Adafactor 步骤的示意图。

使用 PyTorch 高效训练 AI 模型

Adafactor 如何工作？

 

展示 Adafactor 步骤的示意图。

 

  • EMA：指数移动平均
  • 二阶矩：梯度平方的 EMA
使用 PyTorch 高效训练 AI 模型

Adafactor 如何工作？

 

展示 Adafactor 步骤的示意图。

 

  • EMA：指数移动平均
  • 二阶矩：梯度平方的 EMA
使用 PyTorch 高效训练 AI 模型

Adafactor 如何节省内存？

 

展示二阶矩矩阵、列和与行和的示意图。

  • 通过不存储二阶矩矩阵来节省内存
使用 PyTorch 高效训练 AI 模型

Adafactor 如何节省内存？

 

展示二阶矩矩阵、列和与行和的示意图。

  • 通过不存储二阶矩矩阵来节省内存
  • 改为存储矩阵的列和与行和
使用 PyTorch 高效训练 AI 模型

Adafactor 如何节省内存？

 

展示二阶矩矩阵、列和与行和的示意图。

  • 通过不存储二阶矩矩阵来节省内存
  • 改为存储矩阵的列和与行和
使用 PyTorch 高效训练 AI 模型

Adafactor 如何节省内存？

 

展示二阶矩矩阵、列和与行和的示意图。

  • 通过不存储二阶矩矩阵来节省内存
  • 改为存储矩阵的列和与行和
  • 通过列和×行和估计完整矩阵
使用 PyTorch 高效训练 AI 模型

Trainer 与 Accelerator 的实现

比较 Accelerator 与 Trainer 的易用性与自定义能力的图表。

使用 PyTorch 高效训练 AI 模型

用 Trainer 实现 Adafactor

training_args = TrainingArguments(output_dir="./results",
                                  evaluation_strategy="epoch",

                                  optim="adafactor")

trainer = Trainer(model=model,
                  args=training_args,
                  train_dataset=train_dataset,
                  eval_dataset=validation_dataset,
                  compute_metrics=compute_metrics)
trainer.train()

{'epoch': 1.0, 'eval_accuracy': 0.6, 'eval_f1': 0.5}
使用 PyTorch 高效训练 AI 模型

用 Accelerator 实现 Adafactor

# Assumes PyTorch 2.5 or higher
from torch.optim import Adafactor

optimizer = Adafactor(params=model.parameters(), lr=lr)


for batch in train_dataloader:
    inputs, targets = batch["input_ids"], batch["labels"]
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    accelerator.backward(loss)
    optimizer.step()
    lr_scheduler.step()
    optimizer.zero_grad()
    print(f"Loss = {loss}")

Loss = 0.71
使用 PyTorch 高效训练 AI 模型

检查优化器状态

  • 通过 state 访问 optimizer
optimizer_state = optimizer.state.values()
  • 或通过 trainer 访问 optimizer
optimizer_state = trainer.optimizer.state.values()

print(optimizer_state)

dict_values([{'step': tensor(3.),
              'exp_avg_sq_row': tensor([1.0000e-30, 1.0000e-30, 1.0000e-30,  ...]), 
              'exp_avg_sq_col': tensor([4.6147e-11, 5.5115e-11, 1.6338e-10, ...])}, ...])
使用 PyTorch 高效训练 AI 模型

计算 Adafactor 的内存占用

total_size_megabytes, total_num_elements = \
    compute_optimizer_size(trainer.optimizer.state.values())
print(f"Number of Adafactor parameters: {total_num_elements:,}")
print(f"Adafactor size: {total_size_megabytes:.0f} MB")

Number of Adafactor parameters: 178,712
Adafactor size: 1 MB
  • 与 AdamW 对比：Adafactor 占用内存少得多！
Number of AdamW parameters: 131,566,188
AdamW size: 502 MB
使用 PyTorch 高效训练 AI 模型

Passons à la pratique !

使用 PyTorch 高效训练 AI 模型

Preparing Video For Download...