Huấn luyện tiết kiệm bộ nhớ với Adafactor

Huấn luyện Mô hình AI Hiệu quả với PyTorch

Dennis Lee

Data Engineer, Amazon

Bộ tối ưu hóa cho hiệu quả huấn luyện

 

 

Biểu tượng tượng trưng cho AdamW, Adafactor và Adam 8-bit.

Huấn luyện Mô hình AI Hiệu quả với PyTorch

Đánh đổi của bộ tối ưu hóa

Sơ đồ cho thấy đánh đổi giữa số tham số và độ chính xác cho AdamW, Adafactor và Adam 8-bit.

Huấn luyện Mô hình AI Hiệu quả với PyTorch

Adafactor hoạt động thế nào?

 

Sơ đồ minh họa các bước của Adafactor.

Huấn luyện Mô hình AI Hiệu quả với PyTorch

Adafactor hoạt động thế nào?

 

Sơ đồ minh họa các bước của Adafactor.

Huấn luyện Mô hình AI Hiệu quả với PyTorch

Adafactor hoạt động thế nào?

 

Sơ đồ minh họa các bước của Adafactor.

 

  • EMA: trung bình động mũ
  • Moment bậc hai: EMA của bình phương gradient
Huấn luyện Mô hình AI Hiệu quả với PyTorch

Adafactor hoạt động thế nào?

 

Sơ đồ minh họa các bước của Adafactor.

 

  • EMA: trung bình động mũ
  • Moment bậc hai: EMA của bình phương gradient
Huấn luyện Mô hình AI Hiệu quả với PyTorch

Adafactor tiết kiệm bộ nhớ ra sao?

 

Sơ đồ ma trận moment bậc hai, tổng theo cột và tổng theo hàng.

  • Tiết kiệm bộ nhớ bằng cách không lưu ma trận moment bậc hai
Huấn luyện Mô hình AI Hiệu quả với PyTorch

Adafactor tiết kiệm bộ nhớ ra sao?

 

Sơ đồ ma trận moment bậc hai, tổng theo cột và tổng theo hàng.

  • Tiết kiệm bộ nhớ bằng cách không lưu ma trận moment bậc hai
  • Thay vào đó, lưu tổng theo cột và tổng theo hàng của ma trận
Huấn luyện Mô hình AI Hiệu quả với PyTorch

Adafactor tiết kiệm bộ nhớ ra sao?

 

Sơ đồ ma trận moment bậc hai, tổng theo cột và tổng theo hàng.

  • Tiết kiệm bộ nhớ bằng cách không lưu ma trận moment bậc hai
  • Thay vào đó, lưu tổng theo cột và tổng theo hàng của ma trận
Huấn luyện Mô hình AI Hiệu quả với PyTorch

Adafactor tiết kiệm bộ nhớ ra sao?

 

Sơ đồ ma trận moment bậc hai, tổng theo cột và tổng theo hàng.

  • Tiết kiệm bộ nhớ bằng cách không lưu ma trận moment bậc hai
  • Thay vào đó, lưu tổng theo cột và tổng theo hàng của ma trận
  • Ước lượng ma trận đầy đủ bằng cách nhân tổng cột với tổng hàng
Huấn luyện Mô hình AI Hiệu quả với PyTorch

Triển khai với Trainer và Accelerator

Biểu đồ so sánh dễ dùng và khả năng tùy chỉnh của Accelerator và Trainer.

Huấn luyện Mô hình AI Hiệu quả với PyTorch

Dùng Adafactor với Trainer

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}
Huấn luyện Mô hình AI Hiệu quả với PyTorch

Dùng Adafactor với Accelerator

# 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
Huấn luyện Mô hình AI Hiệu quả với PyTorch

Kiểm tra trạng thái optimizer

  • Truy cập optimizer qua state
optimizer_state = optimizer.state.values()
  • Hoặc truy cập optimizer qua trainer
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, ...])}, ...])
Huấn luyện Mô hình AI Hiệu quả với PyTorch

Tính mức dùng bộ nhớ của 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
  • So với AdamW: Adafactor dùng ít bộ nhớ hơn nhiều!
Number of AdamW parameters: 131,566,188
AdamW size: 502 MB
Huấn luyện Mô hình AI Hiệu quả với PyTorch

Ayo berlatih!

Huấn luyện Mô hình AI Hiệu quả với PyTorch

Preparing Video For Download...