用 AdamW 讓訓練更穩定

使用 PyTorch 高效訓練 AI 模型

Dennis Lee

Data Engineer, Amazon

高效訓練

課程章節主題示意圖，著重在最佳化器。

提昇訓練效率的最佳化器

三種最佳化器示意圖：AdamW、Adafactor、8-bit Adam。

提昇訓練效率的最佳化器

三種最佳化器示意圖：AdamW、Adafactor、8-bit Adam。

提昇訓練效率的最佳化器

三種最佳化器示意圖：AdamW、Adafactor、8-bit Adam。

最佳化器的取捨

顯示 AdamW、Adafactor、8-bit Adam 在參數量與精度之間取捨的圖表。

最佳化器的取捨

顯示 AdamW、Adafactor、8-bit Adam 在參數量與精度之間取捨的圖表。

最佳化器的取捨

顯示 AdamW、Adafactor、8-bit Adam 在參數量與精度之間取捨的圖表。

最佳化器的取捨

顯示 AdamW、Adafactor、8-bit Adam 在參數量與精度之間取捨的圖表。

AdamW 如何運作？

AdamW 運作方式示意圖。

AdamW 如何運作？

AdamW 運作方式示意圖。

AdamW 如何運作？

AdamW 運作方式示意圖。

計算梯度的指數移動平均（EMA）

AdamW 如何運作？

AdamW 運作方式示意圖。

計算梯度的指數移動平均（EMA）
計算梯度平方的 EMA

AdamW 如何運作？

AdamW 運作方式示意圖。

計算梯度的指數移動平均（EMA）
計算梯度平方的 EMA

AdamW 如何運作？

AdamW 運作方式示意圖。

計算梯度的指數移動平均（EMA）
計算梯度平方的 EMA

AdamW 的記憶體用量

AdamW 計算涉及的量：參數梯度、梯度 EMA、梯度平方 EMA。

每個方格代表一個參數，每種顏色代表一種狀態

AdamW 的記憶體用量

AdamW 計算涉及的量：參數梯度、梯度 EMA、梯度平方 EMA。

每個方格代表一個參數，每種顏色代表一種狀態

AdamW 的記憶體用量

AdamW 計算涉及的量：參數梯度、梯度 EMA、梯度平方 EMA。

每個方格代表一個參數，每種顏色代表一種狀態
每個參數的記憶體＝8 bytes＝每個狀態 4 bytes × 2 種狀態
總記憶體＝每個參數的記憶體（8 bytes）× 參數數量

估計 AdamW 的記憶體用量

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-cased", return_dict=True)

num_parameters = sum(p.numel() for p in model.parameters())
print(f"Number of model parameters: {num_parameters:,}")

Number of model parameters: 65,783,042

estimated_memory = num_parameters * 8 / (1024 ** 2)
print(f"Estimated memory usage of AdamW: {estimated_memory:.0f} MB")

Estimated memory usage of AdamW: 502 MB

Trainer 與 Accelerator

顯示 Accelerator 與 Trainer 在可客製化程度與易用性間取捨的圖表。

用 Trainer 實作 AdamW

from torch.optim import AdamW

optimizer = AdamW(params=model.parameters())


trainer = Trainer(model=model,
                  args=training_args,
                  train_dataset=train_dataset,
                  eval_dataset=validation_dataset,
                  compute_metrics=compute_metrics,
                  optimizers=(optimizer, lr_scheduler))

trainer.train()

{'epoch': 1.0, 'eval_accuracy': 0.7, 'eval_f1': 0.8}

用 Accelerator 實作 AdamW

from torch.optim import AdamW

optimizer = AdamW(params=model.parameters())


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.7

檢視最佳化器狀態

optimizer_state = optimizer.state.values()
print(optimizer_state)

dict_values([{'step': tensor(3.),

              'exp_avg': tensor([[0., 0., 0.,  ..., 0., 0., 0.], ...]),

              'exp_avg_sq': tensor([[0., 0., 0.,  ..., 0., 0., 0.], ...])}, ...])

計算最佳化器大小

def compute_optimizer_size(optimizer_state):
    total_size_megabytes, total_num_elements = 0, 0

    for params in optimizer_state:

        for name, tensor in params.items():
            tensor = torch.tensor(tensor)

            num_elements = tensor.numel()

            element_size = tensor.element_size()

            total_num_elements += num_elements

            total_size_megabytes += num_elements * element_size / (1024 ** 2)

    return total_size_megabytes, total_num_elements

計算最佳化器大小

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

Number of optimizer parameters: 131,566,188

print(f"Optimizer size: {total_size_megabytes:.0f} MB")

Optimizer size: 502 MB

一起來練習吧！

使用 PyTorch 高效訓練 AI 模型