LlamaFactory
算力平台:

将 Accelerate 添加到你的代码中 ​

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

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

python
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] 是用于将你的代码适配到 Accelerate 的主要类。它了解你正在使用的分布式设置，例如不同进程的数量和硬件类型。该类还提供了许多必要的方法，用于使你的 PyTorch 代码能够在任何分布式训练环境中工作，并管理和跨设备执行进程。

因此，你应该始终在脚本中首先导入并创建一个 [Accelerator] 实例。

python
from accelerate import Accelerator

accelerator = Accelerator()

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

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

准备 PyTorch 对象 ​

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

TIP

Accelerate 仅准备继承自相应 PyTorch 类的对象，例如 torch.optim.Optimizer。

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

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

训练循环 ​

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

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

将所有内容放在一起，你的新 Accelerate 训练循环现在应该看起来像这样！

python
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 参数，并在脚本中添加 [~Accelerator.accumulate] 上下文管理器。

diff
+ 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 提供了：

  • [~Accelerator.clip_grad_value_] 用于将梯度裁剪到最小值和最大值之间
  • [~Accelerator.clip_grad_norm_] 用于将梯度归一化到某个值

混合精度 ​

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

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

WARNING

Accelerate 启用了自动混合精度，因此只有在除了由 [~Accelerator.backward] 处理的损失缩放之外还有其他混合精度操作时，才需要使用 [~Accelerator.autocast]。

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

保存和加载 ​

Accelerate 还可以在训练完成后保存和加载 模型，或者你也可以保存模型和优化器的 状态，这在恢复训练时可能会非常有用。

模型 ​

一旦所有进程完成，使用 [~Accelerator.unwrap_model] 方法解包模型，然后再保存它，因为 [~Accelerator.prepare] 方法将你的模型包装成了分布式训练所需的接口。如果你不解包模型，保存模型状态字典时也会保存任何潜在的额外层，这样你就无法将权重重新加载回基础模型中。

你应该使用 [~Accelerator.save_model] 方法来解包并保存模型状态字典。此方法还可以将模型保存为分片检查点或 safetensors 格式。

状态 ​

在训练过程中，你可能希望保存模型、优化器、随机生成器以及可能的学习率调度器的当前状态，以便在同一脚本中恢复。你应该在脚本中添加 [~Accelerator.save_state] 和 [~Accelerator.load_state] 方法来保存和加载状态。

为了进一步自定义通过 [~Accelerator.save_state] 保存状态的位置和方式，可以使用 [~utils.ProjectConfiguration] 类。例如，如果启用了 automatic_checkpoint_naming，每个保存的检查点将存储在 Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}。

任何其他需要存储的状态对象都应通过 [~Accelerator.register_for_checkpointing] 方法进行注册，以便可以保存和加载。传递给此方法的每个对象都必须具有 load_state_dict 和 state_dict 方法。

TIP

如果你安装了 torchdata>=0.8.0，你还可以在 [~utils.DataLoaderConfiguration] 中传递 use_stateful_dataloader=True。这将扩展 Accelerate 的 DataLoader 类，使其具有 load_state_dict 和 state_dict 方法，并使 Accelerator.save_state 和 Accelerator.load_state 在持久化模型时也能跟踪已读取的训练数据集的位置。

Copyright © 2024-2025 llamafactory.cn