使用 AdamW 的均衡训练

使用 PyTorch 高效训练 AI 模型

Dennis Lee

Data Engineer, Amazon

高效训练

 

 

课程章节主题图，突出显示优化器。

使用 PyTorch 高效训练 AI 模型

提升训练效率的优化器

 

 

展示三种优化器：AdamW、Adafactor 和 8-bit Adam 的示意图。

使用 PyTorch 高效训练 AI 模型

提升训练效率的优化器

 

 

展示三种优化器：AdamW、Adafactor 和 8-bit Adam 的示意图。

使用 PyTorch 高效训练 AI 模型

提升训练效率的优化器

 

 

展示三种优化器：AdamW、Adafactor 和 8-bit Adam 的示意图。

使用 PyTorch 高效训练 AI 模型

优化器的权衡

展示 AdamW、Adafactor 和 8-bit Adam 在参数数量与精度间权衡的示意图。

使用 PyTorch 高效训练 AI 模型

优化器的权衡

展示 AdamW、Adafactor 和 8-bit Adam 在参数数量与精度间权衡的示意图。

使用 PyTorch 高效训练 AI 模型

优化器的权衡

展示 AdamW、Adafactor 和 8-bit Adam 在参数数量与精度间权衡的示意图。

使用 PyTorch 高效训练 AI 模型

优化器的权衡

展示 AdamW、Adafactor 和 8-bit Adam 在参数数量与精度间权衡的示意图。

使用 PyTorch 高效训练 AI 模型

AdamW 如何工作？

展示 AdamW 工作原理的示意图。

使用 PyTorch 高效训练 AI 模型

AdamW 如何工作？

展示 AdamW 工作原理的示意图。

使用 PyTorch 高效训练 AI 模型

AdamW 如何工作？

展示 AdamW 工作原理的示意图。

  • 计算梯度的指数滑动平均（EMA）
使用 PyTorch 高效训练 AI 模型

AdamW 如何工作？

展示 AdamW 工作原理的示意图。

  • 计算梯度的指数滑动平均（EMA）
  • 计算梯度平方的 EMA
使用 PyTorch 高效训练 AI 模型

AdamW 如何工作？

展示 AdamW 工作原理的示意图。

  • 计算梯度的指数滑动平均（EMA）
  • 计算梯度平方的 EMA
使用 PyTorch 高效训练 AI 模型

AdamW 如何工作？

展示 AdamW 工作原理的示意图。

  • 计算梯度的指数滑动平均（EMA）
  • 计算梯度平方的 EMA
使用 PyTorch 高效训练 AI 模型

AdamW 的内存占用

展示 AdamW 计算所涉量：参数梯度、梯度 EMA、梯度平方 EMA 的示意图。

  • 每个方块代表一个参数，每种颜色代表一种状态
使用 PyTorch 高效训练 AI 模型

AdamW 的内存占用

展示 AdamW 计算所涉量：参数梯度、梯度 EMA、梯度平方 EMA 的示意图。

  • 每个方块代表一个参数，每种颜色代表一种状态
使用 PyTorch 高效训练 AI 模型

AdamW 的内存占用

展示 AdamW 计算所涉量：参数梯度、梯度 EMA、梯度平方 EMA 的示意图。

  • 每个方块代表一个参数，每种颜色代表一种状态
  • 每个参数内存 = 8 字节 = 每个状态 4 字节 × 2 个状态
  • 总内存 = 每个参数内存（8 字节）× 参数数量
使用 PyTorch 高效训练 AI 模型

估算 AdamW 的内存占用

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

num_parameters = sum(p.numel() for p in model.parameters())
print(f"Number of model parameters: {num_parameters:,}")

Number of model parameters: 65,783,042

estimated_memory = num_parameters * 8 / (1024 ** 2)
print(f"Estimated memory usage of AdamW: {estimated_memory:.0f} MB")

Estimated memory usage of AdamW: 502 MB
使用 PyTorch 高效训练 AI 模型

Trainer 与 Accelerator

展示 Accelerator 与 Trainer 在可自定义性和易用性间权衡的示意图。

使用 PyTorch 高效训练 AI 模型

用 Trainer 实现 AdamW

from torch.optim import AdamW

optimizer = AdamW(params=model.parameters())


trainer = Trainer(model=model,
                  args=training_args,
                  train_dataset=train_dataset,
                  eval_dataset=validation_dataset,
                  compute_metrics=compute_metrics,
                  optimizers=(optimizer, lr_scheduler))

trainer.train()

{'epoch': 1.0, 'eval_accuracy': 0.7, 'eval_f1': 0.8}
使用 PyTorch 高效训练 AI 模型

用 Accelerator 实现 AdamW

from torch.optim import AdamW

optimizer = AdamW(params=model.parameters())


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

查看优化器状态

optimizer_state = optimizer.state.values()
print(optimizer_state)

dict_values([{'step': tensor(3.),

              'exp_avg': tensor([[0., 0., 0.,  ..., 0., 0., 0.], ...]),

              'exp_avg_sq': tensor([[0., 0., 0.,  ..., 0., 0., 0.], ...])}, ...])
使用 PyTorch 高效训练 AI 模型

计算优化器大小

def compute_optimizer_size(optimizer_state):
    total_size_megabytes, total_num_elements = 0, 0

    for params in optimizer_state:

        for name, tensor in params.items():
            tensor = torch.tensor(tensor)

            num_elements = tensor.numel()

            element_size = tensor.element_size()

            total_num_elements += num_elements

            total_size_megabytes += num_elements * element_size / (1024 ** 2)

    return total_size_megabytes, total_num_elements
使用 PyTorch 高效训练 AI 模型

计算优化器大小

total_size_megabytes, total_num_elements = \
    compute_optimizer_size(trainer.optimizer.state.values())
print(f"Number of optimizer parameters: {total_num_elements:,}")

Number of optimizer parameters: 131,566,188

print(f"Optimizer size: {total_size_megabytes:.0f} MB")

Optimizer size: 502 MB
使用 PyTorch 高效训练 AI 模型

让我们练习！

使用 PyTorch 高效训练 AI 模型

Preparing Video For Download...