在 TPU 上训练

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

在 Notebook 中训练

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

虽然在 TPU 上，最后一部分没有那么重要，但需要理解的关键部分是，当您从 notebook 启动代码时，您是通过一个名为 forking 的过程来实现的。当从命令行启动时，您执行 spawning，此时 python 进程当前未运行，您在其中 spawn 一个新进程。由于您的 Jupyter notebook 已经在使用 python 进程，因此您需要从中 fork 一个新进程来启动代码。

重要的是关于声明模型。在 forked TPU 进程上，建议您实例化模型一次，并将其传递到您的训练函数中。这与在 GPU 上训练不同，在 GPU 上，您创建 n 个模型，这些模型的梯度在某些时刻同步和反向传播。相反，一个模型实例在所有节点之间共享，并在它们之间来回传递。这在低资源 TPU 上训练时尤其重要，例如 Kaggle kernels 或 Google Colaboratory 中提供的那些。

下面是一个训练函数的示例，如果要在 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")

默认情况下，这会将 TPU 上的 torch.float 和 torch.double 转换为 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 上更新