在 TPU 上进行训练

在 TPU 上进行训练与在多 GPU 上进行训练可能略有不同，即使使用了 Accelerate 也是如此。本指南旨在向您展示需要谨慎处理的地方及其原因，以及通常的最佳实践。

在 Notebook 中训练

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

虽然在 TPU 上最后一部分的重要性相对较低，但需要理解的一个关键点是：当从 notebook 启动代码时，是通过一个称为 forking 的过程来实现的。当从命令行启动时，你执行的是 spawning，此时没有正在运行的 python 进程，你会 spawn 一个新进程。由于你的 Jupyter notebook 已经在使用一个 python 进程，因此你需要从中 fork 出一个新进程来启动你的代码。

这一点在声明模型时尤为重要。在分叉的 TPU 进程中，建议将模型实例化一次，并将其传入训练函数。这与在 GPU 上训练不同，后者会创建 n 个模型，并在特定时刻同步和反向传播它们的梯度。相反，一个模型实例在所有节点之间共享，并在它们之间来回传递。在资源有限的 TPU 上训练时，这一点尤为重要，例如在 Kaggle 内核或 Google Colaboratory 中提供的 TPU。

以下是一个示例，展示了在 CPU 或 GPU 上训练时传递给 notebook_launcher() 的训练函数：

此代码片段基于此处的 simple_nlp_example 笔记本中的代码，并进行了简化处理

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 张量往往无法直接使用，此时使用降精度转换（downcasting）而非在所有地方都使用 bf16 是更好的做法。

TPU 上的训练时间

在启动脚本时，您可能会注意到训练最初似乎异常缓慢。这是因为 TPU 会先处理几个批次的数据，以确定需要分配多少内存，然后才会极其高效地利用这一已配置的内存分配。

如果您发现由于使用了更大的批处理大小，用于计算模型指标的评估代码运行时间变长，建议在速度过慢时将批处理大小保持与训练数据一致。否则，在前几次迭代之后，内存将重新分配给这个新的批处理大小。

内存已分配并不意味着它一定会被使用，也不意味着返回训练数据加载器时批处理大小会增大。

Update on GitHub