使用 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：指數移動平均（exponential moving average）
  • Second moment：梯度平方的 EMA
使用 PyTorch 高效訓練 AI 模型

Adafactor 如何運作？

 

顯示 Adafactor 步驟的示意圖。

 

  • EMA：指數移動平均（exponential moving average）
  • Second moment：梯度平方的 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 模型

一起來練習吧！

使用 PyTorch 高效訓練 AI 模型

Preparing Video For Download...