将 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 类会自动将它们放置在正确的设备上。你还应该将常用的 backward() 传递替换为 Accelerate 的 backward() 方法，该方法会为你缩放梯度，并根据你的分布式设置（例如，DeepSpeed 或 Megatron）使用适当的 backward() 方法。

-   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 函数。