梯度累积

使用 PyTorch 高效训练 AI 模型

Dennis Lee

Data Engineer, Amazon

分布式训练

流程图概述课程主题：数据准备、分布式训练、高效训练和优化器。

高效训练

流程图概述课程主题：数据准备、分布式训练、高效训练和优化器。

提升训练效率

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

提升训练效率

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

提升训练效率

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

梯度累积可提高内存效率

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

大批量的问题

大批量：更稳健的梯度估计，收敛更快
GPU 内存限制批量大小

图示：大批量可能导致内存溢出，而小批量可使训练成功完成。

梯度累积如何工作？

图示：将一个大批量拆分为多个小批量。

梯度累积：在小批量上累加梯度
等效于用大批量训练模型
累加后再更新参数

图示：将多个批量的梯度相加以实现梯度累积。

PyTorch、Accelerator 与 Trainer

对比 PyTorch、Accelerator 和 Trainer 的易用性与可定制性。

PyTorch、Accelerator 与 Trainer

对比 PyTorch、Accelerator 和 Trainer 的易用性与可定制性。

PyTorch、Accelerator 与 Trainer

对比 PyTorch、Accelerator 和 Trainer 的易用性与可定制性。

用 PyTorch 实现梯度累积

for index, batch in enumerate(dataloader):
    inputs, targets = (batch["input_ids"], 
                       batch["labels"])
    inputs, targets = (inputs.to(device), 
                       targets.to(device))

图示：将多个批量的梯度相加以实现梯度累积。

用 PyTorch 实现梯度累积

for index, batch in enumerate(dataloader):
    inputs, targets = (batch["input_ids"], 
                       batch["labels"])
    inputs, targets = (inputs.to(device), 
                       targets.to(device))
    outputs = model(inputs, labels=targets)
    loss = outputs.loss

图示：将多个批量的梯度相加以实现梯度累积。

用 PyTorch 实现梯度累积

for index, batch in enumerate(dataloader):
    inputs, targets = (batch["input_ids"], 
                       batch["labels"])
    inputs, targets = (inputs.to(device), 
                       targets.to(device))
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    loss = loss / gradient_accumulation_steps

图示：将多个批量的梯度相加以实现梯度累积。

用 PyTorch 实现梯度累积

for index, batch in enumerate(dataloader):
    inputs, targets = (batch["input_ids"], 
                       batch["labels"])
    inputs, targets = (inputs.to(device), 
                       targets.to(device))
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    loss = loss / gradient_accumulation_steps
    loss.backward()

图示：将多个批量的梯度相加以实现梯度累积。

用 PyTorch 实现梯度累积

for index, batch in enumerate(dataloader):
    inputs, targets = (batch["input_ids"], 
                       batch["labels"])
    inputs, targets = (inputs.to(device), 
                       targets.to(device))
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    loss = loss / gradient_accumulation_steps
    loss.backward()
    if ((index + 1) 
        % gradient_accumulation_steps == 0):

图示：将多个批量的梯度相加以实现梯度累积。

用 PyTorch 实现梯度累积

for index, batch in enumerate(dataloader):
    inputs, targets = (batch["input_ids"], 
                       batch["labels"])
    inputs, targets = (inputs.to(device), 
                       targets.to(device))
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    loss = loss / gradient_accumulation_steps
    loss.backward()
    if ((index + 1) 
        % gradient_accumulation_steps == 0):
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

图示：将多个批量的梯度相加以实现梯度累积。

从 PyTorch 过渡到 Accelerator

对比 PyTorch、Accelerator 和 Trainer 的易用性与可定制性。

从 PyTorch 过渡到 Accelerator

对比 PyTorch、Accelerator 和 Trainer 的易用性与可定制性。

用 Accelerator 实现梯度累积

accelerator = \
    Accelerator(gradient_accumulation_steps=2)

for index, batch in enumerate(dataloader):

        inputs, targets = (batch["input_ids"],
                           batch["labels"])

图示：将多个批量的梯度相加以实现梯度累积。

用 Accelerator 实现梯度累积

accelerator = \
    Accelerator(gradient_accumulation_steps=2)

for index, batch in enumerate(dataloader):

        inputs, targets = (batch["input_ids"],
                           batch["labels"])
        outputs = model(inputs, 
                        labels=targets)
        loss = outputs.loss

图示：将多个批量的梯度相加以实现梯度累积。

用 Accelerator 实现梯度累积

accelerator = \
    Accelerator(gradient_accumulation_steps=2)

for index, batch in enumerate(dataloader):
    with accelerator.accumulate(model):
        inputs, targets = (batch["input_ids"],
                           batch["labels"])
        outputs = model(inputs, 
                        labels=targets)
        loss = outputs.loss

图示：将多个批量的梯度相加以实现梯度累积。

用 Accelerator 实现梯度累积

accelerator = \
    Accelerator(gradient_accumulation_steps=2)

for index, batch in enumerate(dataloader):
    with accelerator.accumulate(model):
        inputs, targets = (batch["input_ids"],
                           batch["labels"])
        outputs = model(inputs, 
                        labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)

图示：将多个批量的梯度相加以实现梯度累积。

用 Accelerator 实现梯度累积

accelerator = \
    Accelerator(gradient_accumulation_steps=2)

for index, batch in enumerate(dataloader):
    with accelerator.accumulate(model):
        inputs, targets = (batch["input_ids"],
                           batch["labels"])
        outputs = model(inputs, 
                        labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)

图示：将多个批量的梯度相加以实现梯度累积。

用 Accelerator 实现梯度累积

accelerator = \
    Accelerator(gradient_accumulation_steps=2)

for index, batch in enumerate(dataloader):
    with accelerator.accumulate(model):
        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()

图示：将多个批量的梯度相加以实现梯度累积。

从 Accelerator 过渡到 Trainer

对比 PyTorch、Accelerator 和 Trainer 的易用性与可定制性。

从 Accelerator 过渡到 Trainer

对比 PyTorch、Accelerator 和 Trainer 的易用性与可定制性。

用 Trainer 实现梯度累积

training_args = TrainingArguments(output_dir="./results",
                                  evaluation_strategy="epoch",
                                  gradient_accumulation_steps=2)

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

trainer.train()

{'epoch': 1.0, 'eval_loss': 0.73, 'eval_accuracy': 0.03, 'eval_f1': 0.05}
{'epoch': 2.0, 'eval_loss': 0.68, 'eval_accuracy': 0.19, 'eval_f1': 0.25}

Passons à la pratique !

使用 PyTorch 高效训练 AI 模型