将 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() 传递，它会为您缩放梯度并根据您的分布式设置（例如，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 函数。