Accelerate documentation
在代码中添加 Accelerate
在代码中添加 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 是用于将代码适配到 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，并创建一个可以跨进程分片的新数据加载器。
Accelerate 仅准备继承自相应 PyTorch 类的对象，例如
torch.optim.Optimizer。
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() 上下文管理器将值自动转换为指定的数据类型。
Accelerate 支持自动混合精度，因此只有在 backward() 已处理缩放的损失执行操作之外，还存在其他混合精度操作时，才需要 autocast()。
+ accelerator = Accelerator(mixed_precision="fp16")
+ with accelerator.autocast():
loss = complex_loss_function(outputs, target)保存和加载
Accelerate 还可以在训练完成后保存和加载 model，或者保存模型和优化器的 state，这对于恢复训练可能很有用。
模型
当所有进程完成后，在保存模型之前，请使用 unwrap_model() 方法解包模型，因为 prepare() 方法已将您的模型包装为适用于分布式训练的正确接口。如果您不解包模型，保存模型状态字典时也会保存来自更大模型的潜在额外层，您将无法将权重重新加载回基础模型中。
您应使用 save_model() 方法来解包并保存模型状态字典。该方法还可以将模型保存为分片检查点或 safetensors 格式。
accelerator.wait_for_everyone() accelerator.save_model(model, save_directory)
对于来自 Transformers 库的模型，请使用 save_pretrained 方法保存模型，以便之后可以使用 from_pretrained 方法重新加载。
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")
要加载权重，请先使用 unwrap_model() 方法解包模型，然后再加载权重。由于所有模型参数都是张量的引用，因此这会将权重加载到 model 内部。
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))设置 safe_serialization=True 以 safetensor 格式保存模型。
accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory, max_shard_size="1GB", safe_serialization=True)要加载分片检查点或 safetensor 格式的检查点，请使用 load_checkpoint_in_model() 方法。该方法允许您将检查点加载到特定设备上。
load_checkpoint_in_model(unwrapped_model, save_directory, device_map={"":device})状态
在训练过程中，您可能希望保存模型、优化器、随机数生成器以及潜在的学习率调度器的当前状态，以便在同一脚本中恢复它们。您应该在脚本中添加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 函数。
Update on GitHub如果已安装
torchdata>=0.8.0，你还可以将use_stateful_dataloader=True传入 DataLoaderConfiguration。这将为 Accelerate 的 DataLoader 类扩展出load_state_dict和state_dict函数，并使得Accelerator.save_state和Accelerator.load_state在持久化模型时，也能记录训练数据集的读取进度。