การฝึกโมเดลที่ประหยัดหน่วยความจำด้วย Adafactor

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

Dennis Lee

Data Engineer, Amazon

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

ไอคอนแทน AdamW, Adafactor และ 8-bit Adam

การแลกเปลี่ยนของ Optimizer

แผนภาพแสดงการแลกเปลี่ยนระหว่างจำนวนพารามิเตอร์และความแม่นยำของ AdamW, Adafactor และ 8-bit Adam

Adafactor ทำงานอย่างไร?

แผนภาพแสดงขั้นตอนการทำงานของ Adafactor

Adafactor ทำงานอย่างไร?

แผนภาพแสดงขั้นตอนการทำงานของ Adafactor

Adafactor ทำงานอย่างไร?

แผนภาพแสดงขั้นตอนการทำงานของ Adafactor

EMA: ค่าเฉลี่ยเคลื่อนที่แบบเอกซ์โพเนนเชียล
Second moment: EMA ของกำลังสองของ gradient

Adafactor ทำงานอย่างไร?

แผนภาพแสดงขั้นตอนการทำงานของ Adafactor

EMA: ค่าเฉลี่ยเคลื่อนที่แบบเอกซ์โพเนนเชียล
Second moment: EMA ของกำลังสองของ gradient

Adafactor ประหยัดหน่วยความจำอย่างไร?

แผนภาพแสดงเมทริกซ์ second moment ผลรวมตามคอลัมน์ และผลรวมตามแถว

ประหยัดหน่วยความจำด้วยการไม่จัดเก็บเมทริกซ์ second moment

Adafactor ประหยัดหน่วยความจำอย่างไร?

แผนภาพแสดงเมทริกซ์ second moment ผลรวมตามคอลัมน์ และผลรวมตามแถว

ประหยัดหน่วยความจำด้วยการไม่จัดเก็บเมทริกซ์ second moment
จัดเก็บผลรวมตามคอลัมน์และแถวของเมทริกซ์แทน

Adafactor ประหยัดหน่วยความจำอย่างไร?

แผนภาพแสดงเมทริกซ์ second moment ผลรวมตามคอลัมน์ และผลรวมตามแถว

ประหยัดหน่วยความจำด้วยการไม่จัดเก็บเมทริกซ์ second moment
จัดเก็บผลรวมตามคอลัมน์และแถวของเมทริกซ์แทน

Adafactor ประหยัดหน่วยความจำอย่างไร?

แผนภาพแสดงเมทริกซ์ second moment ผลรวมตามคอลัมน์ และผลรวมตามแถว

ประหยัดหน่วยความจำด้วยการไม่จัดเก็บเมทริกซ์ second moment
จัดเก็บผลรวมตามคอลัมน์และแถวของเมทริกซ์แทน
ประมาณเมทริกซ์เต็มด้วยการคูณผลรวมตามคอลัมน์และแถว

การใช้งาน Trainer และ Accelerator

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

ใช้งาน Adafactor ด้วย 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}

ใช้งาน Adafactor ด้วย 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

ตรวจสอบสถานะของ optimizer

เข้าถึง optimizer ผ่าน state

optimizer_state = optimizer.state.values()

หรือเข้าถึง optimizer ผ่าน 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, ...])}, ...])

คำนวณการใช้หน่วยความจำของ 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

ฝึกทำแบบฝึกหัด!

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