Huấn luyện mô hình với Accelerator

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

Dennis Lee

Data Engineer, Amazon

Trainer và Accelerator

Biểu đồ so sánh dễ dùng vs. khả năng tùy biến cho Accelerator và Trainer.

Vòng lặp huấn luyện tùy chỉnh

Trainer: không có vòng lặp huấn luyện tùy chỉnh
Một số tác vụ nâng cao trong generative AI cần hai mạng

Hình minh họa mạng GAN.

¹ https://www.aitude.com/basics-of-generative-adversarial-network-model/

Trainer và Accelerator

Biểu đồ so sánh dễ dùng vs. khả năng tùy biến cho Accelerator và Trainer.

Điều chỉnh vòng lặp huấn luyện cơ bản

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)

    outputs = model(inputs)

    loss = outputs.loss

    loss.backward()

    optimizer.step()
    scheduler.step()

Đặt gradient về 0
Chuyển dữ liệu tới thiết bị: .to(device)
Thực hiện forward pass
Tính loss cross-entropy
Tính gradient bằng backward pass
Cập nhật tham số và learning rate

Tạo đối tượng Accelerator

Accelerator cung cấp giao diện cho huấn luyện phân tán

from accelerate import Accelerator
accelerator = Accelerator(
    device_placement=True
)

device_placement (bool, mặc định True): Tự động gán thiết bị

Định nghĩa mô hình và optimizer

Tải mô hình tiền huấn luyện

from transformers import AutoModelForSequenceClassification

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

Tối ưu tham số bằng Adam

from torch.optim import Adam

optimizer = Adam(params=model.parameters(), lr=2e-5)

Định nghĩa scheduler

from transformers import get_linear_schedule_with_warmup

lr_scheduler = get_linear_schedule_with_warmup(
      optimizer=optimizer,

      num_warmup_steps=num_warmup_steps,

      num_training_steps=num_training_steps)

optimizer (obj): Optimizer PyTorch, như Adam
num_warmup_steps (int): số bước tăng tuyến tính lr, đặt int(num_training_steps * 0.1)
num_training_steps (int): tổng số bước train, đặt len(train_dataloader) * num_epochs

Chuẩn bị mô hình cho huấn luyện hiệu quả

Phương thức prepare xử lý gán thiết bị

model, optimizer, dataloader, lr_scheduler = \
    accelerator.prepare(model,

                        optimizer,

                        dataloader,

                        lr_scheduler)

Xây vòng lặp huấn luyện với Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)

Đặt gradient về 0
Trước đó đã chuyển dữ liệu lên thiết bị

Xây vòng lặp huấn luyện với Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

Đặt gradient về 0
Trước đó đã chuyển dữ liệu lên thiết bị
Loại bỏ các dòng tự chuyển dữ liệu

Xây vòng lặp huấn luyện với Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)

    loss = outputs.loss
    loss.backward()

Đặt gradient về 0
Trước đó đã chuyển dữ liệu lên thiết bị
Loại bỏ các dòng tự chuyển dữ liệu
Thực hiện forward pass
Tính loss cross-entropy và gradient

Xây vòng lặp huấn luyện với Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)
    loss = outputs.loss

    accelerator.backward(loss)

    optimizer.step()
    scheduler.step()

Đặt gradient về 0
Trước đó đã chuyển dữ liệu lên thiết bị
Loại bỏ các dòng tự chuyển dữ liệu
Thực hiện forward pass
Tính loss cross-entropy và gradient
Thay loss.backward bằng accelerator
Cập nhật tham số và learning rate

Tóm tắt thay đổi

Trước khi dùng Accelerator

Phải tự chuyển dữ liệu lên thiết bị
- inputs.to(device)
- targets.to(device)
Tính gradient bằng loss.backward()

Sau khi dùng Accelerator

Tự động gán thiết bị và song song dữ liệu
- accelerator.prepare(model)
- accelerator.prepare(dataloader)
Đồng bộ gradient với accelerator.backward(loss)
Vòng lặp tùy chỉnh
Thân thiện, không phụ thuộc phần cứng, mở rộng tốt, dễ bảo trì

Ayo berlatih!

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