在 TPU 上训练

即使使用 Accelerate，在 TPU 上的训练也可能与在多 GPU 上的训练略有不同。本指南旨在向您展示您应该注意的地方以及原因，以及总体最佳实践。

在 Notebook 中训练

在 TPU 上训练时，主要需要注意的是 notebook_launcher()。正如 Notebook 教程中提到的，你需要将训练代码重构为一个函数，该函数可以传递给 notebook_launcher() 函数，并且要注意不要在 GPU 上声明任何张量。

虽然在 TPU 上，最后一点没有那么重要，但需要理解的关键点是，当从 Notebook 启动代码时，是通过一个称为**分叉（forking）**的过程进行的。当从命令行启动时，执行的是**生成（spawning）**，其中当前没有运行 Python 进程，并且您会在其中生成一个新的进程。由于您的 Jupyter Notebook 已经在使用 Python 进程，因此您需要从中分叉一个新的进程来启动您的代码。

这变得重要的方面在于模型的声明。在分叉的 TPU 进程上，建议您一次实例化您的模型，并将此模型传递到您的训练函数中。这与在 GPU 上训练不同，在 GPU 上训练时，您会创建 `n` 个模型，并在特定时刻同步和反向传播它们的梯度。相反，一个模型实例在所有节点之间共享，并且会来回传递。这在使用低资源 TPU（例如 Kaggle 内核或 Google Colaboratory 中提供的 TPU）进行训练时尤其重要。

以下是在 CPU 或 GPU 上训练时传递给 notebook_launcher() 的训练函数示例。

此代码片段基于 simple_nlp_example Notebook 中的代码片段（可在此处找到），为了简单起见，做了一些修改。

def training_function():
    # Initialize accelerator
    accelerator = Accelerator()
    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
    train_dataloader, eval_dataloader = create_dataloaders(
        train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
    )

    # Instantiate optimizer
    optimizer = AdamW(params=model.parameters(), lr=hyperparameters["learning_rate"])

    # Prepare everything
    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
    # prepare method.
    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader
    )

    num_epochs = hyperparameters["num_epochs"]
    # Now we train the model
    for epoch in range(num_epochs):
        model.train()
        for step, batch in enumerate(train_dataloader):
            outputs = model(**batch)
            loss = outputs.loss
            accelerator.backward(loss)

            optimizer.step()
            optimizer.zero_grad()

from accelerate import notebook_launcher

notebook_launcher(training_function)

如果 Accelerate 已配置为 TPU，则 notebook_launcher 将默认为 8 个进程。

如果您使用此示例并在训练循环内部声明模型，那么在低资源系统上，您可能会看到类似以下的错误：

ProcessExitedException: process 0 terminated with signal SIGSEGV

此错误非常难以理解，但基本解释是您用完了系统 RAM。您可以通过重新配置训练函数以接受单个 model 参数并在外部单元格中声明它来完全避免此问题。

# In another Jupyter cell
model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)

+ def training_function(model):
      # Initialize accelerator
      accelerator = Accelerator()
-     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
      train_dataloader, eval_dataloader = create_dataloaders(
          train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
      )
  ...

最后，使用以下方式调用训练函数：

  from accelerate import notebook_launcher
- notebook_launcher(training_function)
+ notebook_launcher(training_function, (model,))

上述解决方法仅在从 Google Colaboratory 或 Kaggle 等低资源服务器上的 Jupyter Notebook 启动 TPU 实例时才需要。如果使用脚本或在功能更强大的服务器上启动，则无需事先声明模型。

混合精度和全局变量

如混合精度教程中所述，Accelerate 支持 fp16 和 bf16，两者都可以在 TPU 上使用。也就是说，理想情况下应使用 bf16，因为它非常高效。

在 TPU 上使用 bf16 和 Accelerate 时，有两个“层级”，即基础层级和操作层级。

在基础层级，这是通过将 mixed_precision="bf16" 传递给 Accelerator 来启用的，例如：

accelerator = Accelerator(mixed_precision="bf16")

默认情况下，这会将 torch.float 和 torch.double 转换为 TPU 上的 bfloat16。正在设置的具体配置是将 XLA_USE_BF16 环境变量设置为 1。

您可以执行另一个配置，即设置 XLA_DOWNCAST_BF16 环境变量。如果设置为 1，则 torch.float 为 bfloat16，而 torch.double 为 float32。

这是在将 downcast_bf16=True 传递给 Accelerator 对象时执行的。

accelerator = Accelerator(mixed_precision="bf16", downcast_bf16=True)

在尝试计算指标、记录值等需要使用原始 bf16 张量无法使用的场景时，使用向下转换而不是在所有地方都使用 bf16 会更好。

TPU 上的训练时间

在启动脚本时，您可能会注意到训练起初似乎非常缓慢。这是因为 TPU 首先会运行几批数据，以查看要分配多少内存，然后才能非常高效地利用此已配置的内存分配。

如果您发现用于计算模型指标的评估代码由于使用了更大的批次大小而花费的时间更长，建议将批次大小保持与训练数据相同（如果速度太慢）。否则，内存将在前几次迭代后重新分配到此新的批次大小。

仅仅因为内存已分配并不意味着它将被使用，或者在返回到训练数据加载器时批次大小会增加。

< > 在 GitHub 上更新

加速

在 TPU 上训练

在 Notebook 中训练

混合精度和全局变量

TPU 上的训练时间