ฝึกโมเดลด้วย Accelerator

การเทรน AI Model อย่างมีประสิทธิภาพด้วย PyTorch

Dennis Lee

Data Engineer, Amazon

Trainer และ Accelerator

แผนภูมิเปรียบเทียบความสะดวกในการใช้งานกับความสามารถในการปรับแต่งของ Accelerator และ Trainer

Training loop แบบกำหนดเอง

Trainer: ไม่รองรับ training loop แบบกำหนดเอง
งานขั้นสูงบางอย่างใน generative AI ต้องใช้สองเครือข่าย

ภาพของ generative adversarial network

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

Trainer และ Accelerator

ปรับแก้ training loop พื้นฐาน

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()

รีเซ็ต gradient
ย้ายข้อมูลไปยังอุปกรณ์ที่กำหนด: .to(device)
ดำเนินการ forward pass
คำนวณ cross-entropy loss
คำนวณ gradient ใน backward pass
อัปเดตพารามิเตอร์โมเดลและ learning rate

สร้างออบเจกต์ Accelerator

Accelerator มอบอินเทอร์เฟซสำหรับการฝึกแบบกระจาย

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

device_placement (bool, ค่าเริ่มต้น True): จัดการการวางอุปกรณ์โดยอัตโนมัติ

กำหนดโมเดลและ optimizer

โหลดโมเดลที่ผ่านการฝึกมาแล้ว

from transformers import AutoModelForSequenceClassification

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

ปรับพารามิเตอร์โมเดลด้วย Adam

from torch.optim import Adam

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

กำหนด 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): PyTorch optimizer เช่น Adam
num_warmup_steps (int): จำนวนขั้นตอนสำหรับเพิ่ม lr เชิงเส้น กำหนดเป็น int(num_training_steps * 0.1)
num_training_steps (int): จำนวนขั้นตอนฝึกทั้งหมด กำหนดเป็น len(train_dataloader) * num_epochs

เตรียมโมเดลสำหรับการฝึกที่มีประสิทธิภาพ

เมธอด prepare จัดการการวางอุปกรณ์

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

                        optimizer,

                        dataloader,

                        lr_scheduler)

สร้าง training loop ด้วย Accelerator

for batch in dataloader:

    optimizer.zero_grad()

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

รีเซ็ต gradient
ย้ายข้อมูลไปยังอุปกรณ์ก่อนหน้านี้

สร้าง training loop ด้วย Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

รีเซ็ต gradient
ย้ายข้อมูลไปยังอุปกรณ์ก่อนหน้านี้
ลบบรรทัดที่ย้ายข้อมูลด้วยตนเองออก

สร้าง training loop ด้วย Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)

    loss = outputs.loss
    loss.backward()

รีเซ็ต gradient
ย้ายข้อมูลไปยังอุปกรณ์ก่อนหน้านี้
ลบบรรทัดที่ย้ายข้อมูลด้วยตนเองออก
ดำเนินการ forward pass
คำนวณ cross-entropy loss และ gradient

สร้าง training loop ด้วย Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)
    loss = outputs.loss

    accelerator.backward(loss)

    optimizer.step()
    scheduler.step()

รีเซ็ต gradient
ย้ายข้อมูลไปยังอุปกรณ์ก่อนหน้านี้
ลบบรรทัดที่ย้ายข้อมูลด้วยตนเองออก
ดำเนินการ forward pass
คำนวณ cross-entropy loss และ gradient
แทนที่ loss.backward ด้วย accelerator
อัปเดตพารามิเตอร์โมเดลและ learning rate

สรุปการเปลี่ยนแปลง

ก่อนใช้ Accelerator

ต้องย้ายข้อมูลไปยังอุปกรณ์ด้วยตนเอง
- inputs.to(device)
- targets.to(device)
คำนวณ gradient ด้วย loss.backward()

หลังใช้ Accelerator

วางอุปกรณ์และ data parallelism อัตโนมัติ
- accelerator.prepare(model)
- accelerator.prepare(dataloader)
จัดการ gradient synchronization ด้วย accelerator.backward(loss)
ปรับแต่ง loop ได้
ใช้งานง่าย รองรับหลายฮาร์ดแวร์ ขยายได้ และบำรุงรักษาง่าย

มาฝึกกันเลย!

การเทรน AI Model อย่างมีประสิทธิภาพด้วย PyTorch