混合精度訓練

使用 PyTorch 高效訓練 AI 模型

Dennis Lee

Data Engineer, Amazon

混合精度訓練可加速計算

代表記憶體效率、通訊效率與計算效率的圖示。

較低精度，計算更快

顯示 FP32 與 FP16 浮點數表示方式的圖表。

較低精度，計算更快

顯示 FP32 與 FP16 浮點數表示方式的圖表。

較低精度，計算更快

顯示 FP32 與 FP16 浮點數表示方式的圖表。

什麼是混合精度訓練？

說明混合精度訓練步驟的圖表。

混合精度訓練：結合 FP16、FP32 計算以加速訓練

什麼是混合精度訓練？

說明混合精度訓練步驟的圖表。

混合精度訓練：結合 FP16、FP32 計算以加速訓練

什麼是混合精度訓練？

說明混合精度訓練步驟的圖表。

混合精度訓練：結合 FP16、FP32 計算以加速訓練
Underflow：數值因低於精度而變成 0

什麼是混合精度訓練？

說明混合精度訓練步驟的圖表。

混合精度訓練：結合 FP16、FP32 計算以加速訓練
Underflow：數值因低於精度而變成 0
透過縮放 loss 避免 underflow

什麼是混合精度訓練？

說明混合精度訓練步驟的圖表。

混合精度訓練：結合 FP16、FP32 計算以加速訓練
Underflow：數值因低於精度而變成 0
透過縮放 loss 避免 underflow

什麼是混合精度訓練？

說明混合精度訓練步驟的圖表。

混合精度訓練：結合 FP16、FP32 計算以加速訓練
Underflow：數值因低於精度而變成 0
透過縮放 loss 避免 underflow

什麼是混合精度訓練？

說明混合精度訓練步驟的圖表。

混合精度訓練：結合 FP16、FP32 計算以加速訓練
Underflow：數值因低於精度而變成 0
透過縮放 loss 避免 underflow

什麼是混合精度訓練？

說明混合精度訓練步驟的圖表。

混合精度訓練：結合 FP16、FP32 計算以加速訓練
Underflow：數值因低於精度而變成 0
透過縮放 loss 避免 underflow

PyTorch 實作

比較 PyTorch、Accelerator、Trainer 在易用性與自訂能力的圖表。

使用 PyTorch 進行混合精度訓練

scaler = torch.amp.GradScaler()


for batch in train_dataloader:
    inputs, targets = batch["input_ids"], batch["labels"]

    with torch.autocast(device_type="cpu", dtype=torch.float16):

         outputs = model(inputs, labels=targets)
         loss = outputs.loss

    scaler.scale(loss).backward()

    scaler.step(optimizer)

    scaler.update()

    optimizer.zero_grad()

從 PyTorch 移轉到 Accelerator

比較 PyTorch、Accelerator、Trainer 在易用性與自訂能力的圖表。

從 PyTorch 移轉到 Accelerator

比較 PyTorch、Accelerator、Trainer 在易用性與自訂能力的圖表。

使用 Accelerator 進行混合精度訓練

accelerator = Accelerator(mixed_precision="fp16")


model, optimizer, train_dataloader, lr_scheduler = \
    accelerator.prepare(model, optimizer, train_dataloader, lr_scheduler)


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()
    optimizer.zero_grad()

從 Accelerator 移轉到 Trainer

比較 PyTorch、Accelerator、Trainer 在易用性與自訂能力的圖表。

從 Accelerator 移轉到 Trainer

比較 PyTorch、Accelerator、Trainer 在易用性與自訂能力的圖表。

使用 Trainer 進行混合精度訓練

training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="epoch",
    fp16=True
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["validation"],
    compute_metrics=compute_metrics,
)

trainer.train()

一起來練習吧！

使用 PyTorch 高效訓練 AI 模型

Preparing Video For Download...