混合精度训练

使用 PyTorch 高效训练 AI 模型

Dennis Lee

Data Engineer, Amazon

混合精度训练可加速计算

 

 

表示内存效率、通信效率和计算效率的图标。

使用 PyTorch 高效训练 AI 模型

以更低精度实现更快计算

展示 FP32 与 FP16 浮点数表示的示意图。

使用 PyTorch 高效训练 AI 模型

以更低精度实现更快计算

展示 FP32 与 FP16 浮点数表示的示意图。

使用 PyTorch 高效训练 AI 模型

以更低精度实现更快计算

展示 FP32 与 FP16 浮点数表示的示意图。

使用 PyTorch 高效训练 AI 模型

什么是混合精度训练?

展示混合精度训练步骤的示意图。

  • 混合精度训练:结合 FP16、FP32 计算以加速训练
使用 PyTorch 高效训练 AI 模型

什么是混合精度训练?

展示混合精度训练步骤的示意图。

  • 混合精度训练:结合 FP16、FP32 计算以加速训练
使用 PyTorch 高效训练 AI 模型

什么是混合精度训练?

展示混合精度训练步骤的示意图。

  • 混合精度训练:结合 FP16、FP32 计算以加速训练
  • 下溢:数值精度不足而变为 0
使用 PyTorch 高效训练 AI 模型

什么是混合精度训练?

展示混合精度训练步骤的示意图。

  • 混合精度训练:结合 FP16、FP32 计算以加速训练
  • 下溢:数值精度不足而变为 0
  • 通过缩放 loss 防止下溢
使用 PyTorch 高效训练 AI 模型

什么是混合精度训练?

展示混合精度训练步骤的示意图。

  • 混合精度训练:结合 FP16、FP32 计算以加速训练
  • 下溢:数值精度不足而变为 0
  • 通过缩放 loss 防止下溢
使用 PyTorch 高效训练 AI 模型

什么是混合精度训练?

展示混合精度训练步骤的示意图。

  • 混合精度训练:结合 FP16、FP32 计算以加速训练
  • 下溢:数值精度不足而变为 0
  • 通过缩放 loss 防止下溢
使用 PyTorch 高效训练 AI 模型

什么是混合精度训练?

展示混合精度训练步骤的示意图。

  • 混合精度训练:结合 FP16、FP32 计算以加速训练
  • 下溢:数值精度不足而变为 0
  • 通过缩放 loss 防止下溢
使用 PyTorch 高效训练 AI 模型

什么是混合精度训练?

展示混合精度训练步骤的示意图。

  • 混合精度训练:结合 FP16、FP32 计算以加速训练
  • 下溢:数值精度不足而变为 0
  • 通过缩放 loss 防止下溢
使用 PyTorch 高效训练 AI 模型

PyTorch 实现

比较 PyTorch、Accelerator 和 Trainer 的易用性与可定制性的图表。

使用 PyTorch 高效训练 AI 模型

用 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 高效训练 AI 模型

从 PyTorch 到 Accelerator

比较 PyTorch、Accelerator 和 Trainer 的易用性与可定制性的图表。

使用 PyTorch 高效训练 AI 模型

从 PyTorch 到 Accelerator

比较 PyTorch、Accelerator 和 Trainer 的易用性与可定制性的图表。

使用 PyTorch 高效训练 AI 模型

用 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()
使用 PyTorch 高效训练 AI 模型

从 Accelerator 到 Trainer

比较 PyTorch、Accelerator 和 Trainer 的易用性与可定制性的图表。

使用 PyTorch 高效训练 AI 模型

从 Accelerator 到 Trainer

比较 PyTorch、Accelerator 和 Trainer 的易用性与可定制性的图表。

使用 PyTorch 高效训练 AI 模型

用 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 模型

Passons à la pratique !

使用 PyTorch 高效训练 AI 模型

Preparing Video For Download...