将 Accelerate 添加到您的代码

每个分布式训练框架都有自己的工作方式，这可能需要编写大量的自定义代码以使其适应你的 PyTorch 训练代码和训练环境。Accelerate 提供了一种友好的方式与这些分布式训练框架进行交互，而无需学习每个框架的具体细节。Accelerate 会为你处理这些细节，这样你就可以专注于训练代码，并将其扩展到任何分布式训练环境。

在本教程中，你将学习如何使用 Accelerate 调整现有的 PyTorch 代码，并轻松地开始在分布式系统上进行训练！你将从基本的 PyTorch 训练循环开始（假设所有训练对象，如 model 和 optimizer 已经设置好），并逐步将 Accelerate 集成到其中。

device = "cuda"
model.to(device)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()
    scheduler.step()

Accelerator

Accelerator 是将你的代码调整以与 Accelerate 配合使用的主要类。它了解你正在使用的分布式设置，例如不同进程的数量和你的硬件类型。此类还提供对许多必要方法的访问，这些方法可使你的 PyTorch 代码在任何分布式训练环境中工作，以及用于管理和执行跨设备的进程。

这就是你应该始终从在你的脚本中导入和创建 Accelerator 实例开始的原因。

from accelerate import Accelerator

accelerator = Accelerator()

Accelerator 还知道将你的 PyTorch 对象移动到哪个设备，因此建议让 Accelerate 为你处理此操作。

- device = "cuda"
+ device = accelerator.device
  model.to(device)

准备 PyTorch 对象

接下来，你需要为分布式训练准备你的 PyTorch 对象（模型、优化器、调度器等）。prepare() 方法负责将你的模型放置在适合你的训练设置的容器中（如单 GPU 或多 GPU），调整优化器和调度器以使用 Accelerate 的 AcceleratedOptimizer 和 AcceleratedScheduler，并创建一个可以在进程之间进行分片的新的数据加载器。

PyTorch 对象将按照它们发送的顺序返回。

model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

训练循环

最后，从训练循环中删除对输入和目标的 to(device) 调用，因为 Accelerate 的 DataLoader 类会自动将它们放置在正确的设备上。你还应该用 Accelerate 的 backward() 方法替换通常的 backward() 传递，该方法会为你缩放梯度，并使用适合你的分布式设置的 backward() 方法（例如 DeepSpeed 或 Megatron）。

-   inputs = inputs.to(device)
-   targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
-   loss.backward()
+   accelerator.backward(loss)

将所有内容组合在一起，你的新的 Accelerate 训练循环现在应该如下所示！

from accelerate import Accelerator
accelerator = Accelerator()

device = accelerator.device
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    scheduler.step()

训练功能

Accelerate 提供了额外的功能 - 如梯度累积、梯度裁剪、混合精度训练等 - 你可以将其添加到你的脚本中以改进你的训练运行。让我们探索这三个功能。

梯度累积

梯度累积使你能够通过在更新权重之前累积多个批次的梯度来训练更大的批次大小。这对于解决内存限制很有用。要在 Accelerate 中启用此功能，请在 Accelerator 类中指定 gradient_accumulation_steps 参数，并将 accumulate() 上下文管理器添加到你的脚本中。

+ accelerator = Accelerator(gradient_accumulation_steps=2)
  model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)

  for input, label in training_dataloader:
+     with accelerator.accumulate(model):
          predictions = model(input)
          loss = loss_function(predictions, label)
          accelerator.backward(loss)
          optimizer.step()
          scheduler.step()
          optimizer.zero_grad()

梯度裁剪

梯度裁剪是一种防止“梯度爆炸”的技术，Accelerate 提供了

• clipgrad_value() 将梯度裁剪到最小值和最大值
• clipgrad_norm() 将梯度归一化到特定值

混合精度

混合精度通过使用较低精度的类型（如 fp16（半精度））来计算梯度，从而加速训练。为了获得 Accelerate 的最佳性能，损失应该在你的模型中计算（如在 Transformers 模型中），因为模型外部的计算将使用全精度进行计算。

在 Accelerator 中设置要使用的混合精度类型，然后使用 autocast() 上下文管理器自动将值转换为指定的数据类型。

+ accelerator = Accelerator(mixed_precision="fp16")
+ with accelerator.autocast():
      loss = complex_loss_function(outputs, target):

保存和加载

Accelerate 还可以保存和加载模型（一旦训练完成），或者你也可以保存模型和优化器状态，这对于恢复训练可能很有用。

模型

在所有进程完成之后，使用 unwrap_model() 方法解开模型，然后保存它。因为 prepare() 方法将您的模型封装到适合分布式训练的接口中。如果您不解开模型，保存模型状态字典也会保存来自较大模型的任何潜在额外层，您将无法将权重加载回您的基础模型。

您应该使用 save_model() 方法来解开并保存模型状态字典。此方法还可以将模型保存到分片检查点或 safetensors 格式。

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory)

from transformers import AutoModel

unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(
    "path/to/my_model_directory",
    is_main_process=accelerator.is_main_process,
    save_function=accelerator.save,
)

model = AutoModel.from_pretrained("path/to/my_model_directory")

unwrapped_model = accelerator.unwrap_model(model)
path_to_checkpoint = os.path.join(save_directory,"pytorch_model.bin")
unwrapped_model.load_state_dict(torch.load(path_to_checkpoint))

状态

在训练期间，您可能希望保存模型、优化器、随机生成器以及潜在的学习率调度器的当前状态，以便它们可以在 *同一个脚本* 中恢复。您应该在脚本中添加 save_state() 和 load_state() 方法来保存和加载状态。

为了进一步自定义通过 save_state() 保存状态的位置和方式，请使用 ProjectConfiguration 类。例如，如果 automatic_checkpoint_naming 已启用，则每个保存的检查点都存储在 Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number} 中。

任何其他要存储的具有状态的项目都应使用 register_for_checkpointing() 方法注册，以便可以保存和加载它们。传递给此方法以存储的每个对象都必须具有 load_state_dict 和 state_dict 函数。