在 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 上更新