LlamaFactory
算力平台:

在 TPUs 上进行训练 ​

在 TPUs 上进行训练与在多 GPU 上训练有所不同，即使使用 Accelerate 也是如此。本指南旨在向你展示需要注意的地方及其原因，以及一般最佳实践。

在 Notebook 中进行训练 ​

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

在 TPU 上，最后一部分不是那么重要，但关键是要理解，当你从 Notebook 中启动代码时，是通过一个称为 forking 的过程。从命令行启动时，你执行的是 spawning，即当前没有运行 Python 进程，你 spawn 一个新的进程。由于你的 Jupyter Notebook 已经在使用一个 Python 进程，你需要从它 fork 一个新的进程来启动你的代码。

这一点在声明模型时变得尤为重要。在 forked TPU 进程中，建议你 仅 实例化一次模型，并将其传递给训练函数。这与在 GPU 上训练不同，在 GPU 上你会创建 n 个模型，这些模型的梯度在某些时刻同步并反向传播。相反，一个模型实例在所有节点之间共享，并在节点之间传递。这在使用低资源 TPU（如 Kaggle 内核或 Google Colaboratory 提供的 TPU）进行训练时尤为重要。

以下是一个传递给 [notebook_launcher] 的训练函数示例，如果在 CPU 或 GPU 上进行训练：

python
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()
python
from accelerate import notebook_launcher

notebook_launcher(training_function)

如果你使用这个示例并在训练循环中声明模型，那么在低资源系统上你可能会看到如下错误：

ProcessExitedException: process 0 terminated with signal SIGSEGV

这个错误信息极其晦涩，但基本的解释是你用完了系统的 RAM。你可以通过重新配置训练函数来完全避免这个问题，使其接受一个 model 参数，并在外部单元格中声明它：

python
# In another Jupyter cell
model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
diff
+ 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"]
      )
  ...

最后调用训练函数，使用以下代码：

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

混合精度和全局变量 ​

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

使用 bf16 和 Accelerate 在 TPUs 上时，有两个"层"，即基础层和操作层。

在基础层，当将 mixed_precision="bf16" 传递给 Accelerator 时启用，例如：

python
accelerator = Accelerator(mixed_precision="bf16")

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

你还可以进行进一步的配置，即设置 XLA_DOWNCAST_BF16 环境变量。如果将其设置为 1，则 torch.float 会转换为 bfloat16，而 torch.double 会转换为 float32。

这在传递 downcast_bf16=True 时由 Accelerator 对象执行：

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

使用 downcasting 而不是在所有地方使用 bf16，对于计算指标、记录值等场景非常有用，因为原始的 bf16 张量在这种情况下是不可用的。

在 TPUs 上的训练时间 ​

当你启动脚本时，可能会注意到训练一开始显得异常缓慢。这是因为 TPUs 会先运行几批数据，以确定需要分配多少内存，然后才会利用这种配置的内存分配方式高效地运行。

如果你发现由于使用了更大的批量，计算模型指标的评估代码变慢了，建议将批量大小保持与训练数据相同，如果这样做仍然太慢的话。否则，在前几次迭代后，内存会重新分配给这个新的批量大小。

Copyright © 2024-2025 llamafactory.cn