在 TPU 上进行训练

在 TPU 上训练可能与在多 GPU 上训练略有不同，即使使用了 Accelerate 也是如此。本指南旨在向您展示在哪些方面需要小心以及原因，并介绍一般的最佳实践。

在 Notebook 中训练

在 TPU 上训练时，主要需要注意的地方来自于 notebook_launcher()。如 notebook 教程所述，您需要将训练代码重构为一个可以传递给 notebook_launcher() 函数的函数，并注意不要在 GPU 上声明任何张量。

虽然在 TPU 上，最后一点不那么重要，但一个关键点是，当您从 notebook 启动代码时，是通过一个称为forking的过程来完成的。当从命令行启动时，您执行的是spawning，即当前没有运行的 python 进程，您需要spawn一个新进程。由于您的 Jupyter notebook 已经在使用一个 python 进程，您需要从它fork一个新进程来启动您的代码。

在声明模型时，这一点变得很重要。在 forked TPU 进程上，建议您一次性实例化模型，并将其传递到训练函数中。这与在 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

这个错误极其晦涩，但基本解释是您的系统内存耗尽。您可以通过重新配置训练函数以接受一个 `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 时，有两个“层次”：基础层和操作层。

在基础层，当向 `Accelerator` 传递 `mixed_precision="bf16"` 时启用，例如

accelerator = Accelerator(mixed_precision="bf16")

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

您可以进行进一步的配置，即设置 `XLA_DOWNCAST_BF16` 环境变量。如果设置为 `1`，那么 `torch.float` 将是 `bfloat16`，而 `torch.double` 将是 `float32`。

这在 `Accelerator` 对象中通过传递 `downcast_bf16=True` 来完成

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

当您尝试计算指标、记录值等原始 bf16 张量无法使用的情况下，使用降级转换而非处处使用 bf16 是一个好方法。

TPU 上的训练时间

当您启动脚本时，可能会注意到训练一开始似乎异常缓慢。这是因为 TPU 首先会运行几个批次的数据来确定分配多少内存，然后才会非常高效地利用这个配置好的内存分配。

如果您注意到由于使用了更大的批次大小，用于计算模型指标的评估代码耗时更长，建议如果速度太慢，将批次大小保持与训练数据相同。否则，内存在最初几次迭代后会重新分配到这个新的批次大小。

内存被分配并不意味着它将被使用，也不意味着当回到训练数据加载器时批次大小会增加。

< > 在 GitHub 上更新