accelerate

在任何类型的设备上运行您的*原始* PyTorch 训练脚本

易于集成

🤗 Accelerate 是为那些喜欢编写 PyTorch 模型训练循环，但不愿意编写和维护使用多 GPU/TPU/fp16 所需的样板代码的 PyTorch 用户创建的。

🤗 Accelerate 精确地抽象了与多 GPU/TPU/fp16 相关的样板代码，而保持您的其他代码不变。

以下是一个示例：

  import torch
  import torch.nn.functional as F
  from datasets import load_dataset
+ from accelerate import Accelerator

+ accelerator = Accelerator()
- device = 'cpu'
+ device = accelerator.device

  model = torch.nn.Transformer().to(device)
  optimizer = torch.optim.Adam(model.parameters())

  dataset = load_dataset('my_dataset')
  data = torch.utils.data.DataLoader(dataset, shuffle=True)

+ model, optimizer, data = accelerator.prepare(model, optimizer, data)

  model.train()
  for epoch in range(10):
      for source, targets in data:
          source = source.to(device)
          targets = targets.to(device)

          optimizer.zero_grad()

          output = model(source)
          loss = F.cross_entropy(output, targets)

-         loss.backward()
+         accelerator.backward(loss)

          optimizer.step()

如您所见，通过在任何标准 PyTorch 训练脚本中添加 5 行代码，您现在可以在任何类型的单节点或分布式节点设置（单 CPU、单 GPU、多 GPU 和 TPU）上运行，并且可以使用或不使用混合精度（fp8、fp16、bf16）。

特别是，同样的代码可以在没有修改的情况下在本地机器上进行调试或在您的训练环境中运行。

🤗 Accelerate 甚至为您处理设备放置（这需要对您的代码进行一些更改，但通常更安全），因此您可以进一步简化您的训练循环：

  import torch
  import torch.nn.functional as F
  from datasets import load_dataset
+ from accelerate import Accelerator

- device = 'cpu'
+ accelerator = Accelerator()

- model = torch.nn.Transformer().to(device)
+ model = torch.nn.Transformer()
  optimizer = torch.optim.Adam(model.parameters())

  dataset = load_dataset('my_dataset')
  data = torch.utils.data.DataLoader(dataset, shuffle=True)

+ model, optimizer, data = accelerator.prepare(model, optimizer, data)

  model.train()
  for epoch in range(10):
      for source, targets in data:
-         source = source.to(device)
-         targets = targets.to(device)

          optimizer.zero_grad()

          output = model(source)
          loss = F.cross_entropy(output, targets)

-         loss.backward()
+         accelerator.backward(loss)

          optimizer.step()

想了解更多？查看文档或查看我们的示例。

启动脚本

🤗 Accelerate 还提供了一个可选的 CLI 工具，允许您在启动脚本之前快速配置和测试您的训练环境。无需记住如何使用 torch.distributed.run 或为 TPU 训练编写特定的启动器！ 在您的机器上只需运行：

accelerate config

并回答所问的问题。这将生成一个配置文件，在执行以下命令时将自动用于正确设置默认选项：

accelerate launch my_script.py --args_to_my_script

例如，以下是如何在 MRPC 任务上运行 GLUE 示例（从存储库的根目录）：

accelerate launch examples/nlp_example.py

这个 CLI 工具是可选的，您仍然可以根据自己的喜好使用 python my_script.py 或 python -m torchrun my_script.py。

如果您不想运行 accelerate config，也可以直接将您要传递给 torchrun 的参数作为参数传递给 accelerate launch。

例如，以下是如何在两个 GPU 上启动：

accelerate launch --multi_gpu --num_processes 2 examples/nlp_example.py

要了解更多信息，请查看此处提供的 CLI 文档。

使用 MPI 启动多 CPU 运行

🤗 这是另一种使用 MPI 启动多 CPU 运行的方法。您可以在此页面上了解如何安装 Open MPI。您也可以使用 Intel MPI 或 MVAPICH。 在集群上设置好 MPI 后，只需运行：

accelerate config

回答所问的问题，选择使用多 CPU 运行，并在询问是否要 accelerate 启动 mpirun 时回答"是"。 然后，使用 accelerate launch 运行您的脚本，如：

accelerate launch examples/nlp_example.py

或者，您可以直接使用 mpirun，无需使用 CLI，如：

mpirun -np 2 python examples/nlp_example.py

使用 DeepSpeed 启动训练

🤗 Accelerate 支持使用 DeepSpeed 在单个/多个 GPU 上进行训练。要使用它，您不需要更改训练代码中的任何内容；您可以仅使用 accelerate config 设置所有内容。但是，如果您想从 Python 脚本中调整与 DeepSpeed 相关的参数，我们为您提供了 DeepSpeedPlugin。

from accelerate import Accelerator, DeepSpeedPlugin

# deepspeed 需要预先知道您的梯度累积步骤，所以不要忘记传递它
# 请记住，您仍然需要自己进行梯度累积，就像没有 deepspeed 一样
deepspeed_plugin = DeepSpeedPlugin(zero_stage=2, gradient_accumulation_steps=2)
accelerator = Accelerator(mixed_precision='fp16', deepspeed_plugin=deepspeed_plugin)

# 如何保存您的 🤗 Transformer？
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(save_dir, save_function=accelerator.save, state_dict=accelerator.get_state_dict(model))

注意：DeepSpeed 支持目前处于实验阶段。如果您遇到任何问题，请开启一个 issue。

从笔记本启动训练

🤗 Accelerate 还提供了一个 notebook_launcher 函数，您可以在笔记本中使用它来启动分布式训练。这对于使用 TPU 后端的 Colab 或 Kaggle 笔记本特别有用。只需在 training_function 中定义您的训练循环，然后在最后一个单元格中添加：

from accelerate import notebook_launcher

notebook_launcher(training_function)

可以在这个笔记本中找到一个示例。

为什么我应该使用 🤗 Accelerate？

当您想轻松地在分布式环境中运行训练脚本，而不必放弃对训练循环的完全控制时，您应该使用 🤗 Accelerate。这不是 PyTorch 之上的高级框架，只是一个薄包装，因此您不必学习新的库。事实上，🤗 Accelerate 的整个 API 都在一个类中，即 Accelerator 对象。

为什么我不应该使用 🤗 Accelerate？

如果您不想自己编写训练循环，就不应该使用 🤗 Accelerate。有很多基于 PyTorch 的高级库可以为您提供这种功能，🤗 Accelerate 不是其中之一。

使用 🤗 Accelerate 的框架

如果您喜欢 🤗 Accelerate 的简单性，但更希望围绕其功能有更高级别的抽象，以下列出了一些建立在 🤗 Accelerate 之上的框架和库：

• Amphion 是一个用于音频、音乐和语音生成的工具包。它的目的是支持可复现的研究，并帮助初级研究人员和工程师在音频、音乐和语音生成研究与开发领域入门。

• Animus 是一个用于运行机器学习实验的极简框架。Animus 突出了 ML 实验中常见的"断点"，并在 IExperiment 中为它们提供了统一的接口。

• Catalyst 是一个用于深度学习研究和开发的 PyTorch 框架。它注重可复现性、快速实验和代码库重用，让你能够创造新事物，而不是反复编写训练循环。Catalyst 提供了一个 Runner 来连接实验的所有部分：硬件后端、数据转换、模型训练和推理逻辑。

• fastai 是一个 PyTorch 深度学习框架，它使用现代最佳实践简化了快速准确神经网络的训练。fastai 提供了一个 Learner 来处理深度学习算法的训练、微调和推理。

• Finetuner 是一项使模型能够为语义搜索、视觉相似度搜索、跨模态文本<->图像搜索、推荐系统、聚类、重复检测、异常检测或其他用途创建更高质量嵌入的服务。

• InvokeAI 是 Stable Diffusion 模型的创意引擎，提供行业领先的 WebUI、终端使用支持，并作为许多商业产品的基础。

• Kornia 是一个可微分库，允许将经典计算机视觉集成到深度学习模型中。Kornia 提供了一个 Trainer，专门用于训练和微调库中支持的深度学习算法。

• Open Assistant 是一个基于聊天的助手，它能理解任务，可以与第三方系统交互，并动态检索信息来完成任务。

• pytorch-accelerated 是一个轻量级训练库，具有以通用 Trainer 为中心的精简功能集，非常强调简单性和透明度；使用户能够准确理解底层的工作原理，而无需自己编写和维护样板代码！

• Stable Diffusion web UI 是一个基于 Gradio 库的开源浏览器易用界面，用于 Stable Diffusion。

• torchkeras 是一个简单的工具，用于以 Keras 风格训练 PyTorch 模型，在 notebook 中提供动态美观的图表来监控损失或指标。

• transformers 是一个工具，用于帮助在 PyTorch、Tensorflow 和 JAX 中训练最先进的机器学习模型。（Accelerate 是 PyTorch 端的后端）。

安装

本仓库在 Python 3.8+ 和 PyTorch 1.10.0+ 上经过测试

你应该在虚拟环境中安装 🤗 Accelerate。如果你不熟悉 Python 虚拟环境，请查看用户指南。

首先，使用你要使用的 Python 版本创建一个虚拟环境并激活它。

然后，你需要安装 PyTorch：请参考官方安装页面了解适用于你的平台的具体安装命令。之后可以使用 pip 安装 🤗 Accelerate，如下所示：

pip install accelerate

支持的集成

• 仅 CPU
• 单节点（机器）多 CPU
• 多节点（机器）多 CPU
• 单 GPU
• 单节点（机器）多 GPU
• 多节点（机器）多 GPU
• TPU
• FP16/BFloat16 混合精度
• 使用 Transformer Engine 的 FP8 混合精度
• DeepSpeed 支持（实验性）
• PyTorch 完全分片数据并行（FSDP）支持（实验性）
• Megatron-LM 支持（实验性）

引用 🤗 Accelerate

如果你在出版物中使用 🤗 Accelerate，请使用以下 BibTeX 条目进行引用。

@Misc{accelerate,
  title =        {Accelerate: Training and inference at scale made simple, efficient and adaptable.},
  author =       {Sylvain Gugger and Lysandre Debut and Thomas Wolf and Philipp Schmid and Zachary Mueller and Sourab Mangrulkar and Marc Sun and Benjamin Bossan},
  howpublished = {\url{https://github.com/huggingface/accelerate}},
  year =         {2022}
}