使用 KerasNLP 和 tf.distribute 進行資料平行訓練
作者：Anshuman Mishra
建立日期 2023/07/07
上次修改時間 2023/07/07
說明：使用 KerasNLP 和 tf.distribute 進行資料平行訓練。
簡介
分散式訓練是一種用於在多個設備或機器上同時訓練深度學習模型的技術。它有助於減少訓練時間，並允許使用更多數據訓練更大的模型。KerasNLP 是一個提供用於自然語言處理任務（包括分散式訓練）的工具和實用程序的函式庫。
在本教學課程中，我們將使用 KerasNLP 在 wikitext-2 資料集（一個包含 200 萬字的維基百科文章資料集）上訓練基於 BERT 的遮罩語言模型 (MLM)。MLM 任務涉及預測句子中被遮罩的詞，這有助於模型學習詞的上下文表示。
本指南側重於數據平行化，特別是同步數據平行化，其中每個加速器（GPU 或 TPU）都保存模型的完整副本，並查看輸入數據的不同部分批次。在每個設備上計算部分梯度，彙總後用於計算全局梯度更新。
具體來說，本指南將教您如何在以下兩種設置中使用 tf.distribute API 在多個 GPU 上訓練 Keras 模型，只需對您的代碼進行最少的更改
- 在單台機器上安裝的多個 GPU（通常為 2 到 8 個）上（單主機、多設備訓練）。這是研究人員和小規模產業工作流程中最常見的設置。
- 在一個由多台機器組成的集群上，每台機器都搭載一個或多個 GPU（多工作器分散式訓練）。這對於大規模產業工作流程來說是一個很好的設置，例如在 20-100 個 GPU 上訓練十億字資料集上的高解析度文本摘要模型。
!pip install -q --upgrade keras-nlp
!pip install -q --upgrade keras # Upgrade to Keras 3.
導入
import os
os.environ["KERAS_BACKEND"] = "tensorflow"
import tensorflow as tf
import keras
import keras_nlp
在我們開始任何訓練之前，讓我們將單個 GPU 配置為顯示為兩個邏輯設備。
當您使用兩個或多個物理 GPU 進行訓練時，這完全沒有必要。這只是一個技巧，可以在默認的 colab GPU 運行時上顯示真正的分散式訓練，該運行時只有一個可用的 GPU。
!nvidia-smi --query-gpu=memory.total --format=csv,noheader
physical_devices = tf.config.list_physical_devices("GPU")
tf.config.set_logical_device_configuration(
physical_devices[0],
[
tf.config.LogicalDeviceConfiguration(memory_limit=15360 // 2),
tf.config.LogicalDeviceConfiguration(memory_limit=15360 // 2),
],
)
logical_devices = tf.config.list_logical_devices("GPU")
logical_devices
EPOCHS = 3
24576 MiB
要使用 Keras 模型進行單主機、多設備同步訓練，您需要使用 tf.distribute.MirroredStrategy API。以下是它的工作原理
- 實例化一個
MirroredStrategy，可以選擇配置您要使用的特定設備（默認情況下，策略將使用所有可用的 GPU）。 - 使用策略對象打開一個作用域，並在此作用域內創建您需要的所有包含變量的 Keras 對象。通常，這意味著在分佈作用域內創建和編譯模型。
- 照常通過
fit()訓練模型。
strategy = tf.distribute.MirroredStrategy()
print(f"Number of devices: {strategy.num_replicas_in_sync}")
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1')
Number of devices: 2
基本批次大小和學習率
base_batch_size = 32
base_learning_rate = 1e-4
計算縮放後的批次大小和學習率
scaled_batch_size = base_batch_size * strategy.num_replicas_in_sync
scaled_learning_rate = base_learning_rate * strategy.num_replicas_in_sync
現在，我們需要下載並預處理 wikitext-2 資料集。此資料集將用於預訓練 BERT 模型。我們將濾除短行，以確保數據具有足夠的上下文進行訓練。
keras.utils.get_file(
origin="https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-2-v1.zip",
extract=True,
)
wiki_dir = os.path.expanduser("~/.keras/datasets/wikitext-2/")
# Load wikitext-103 and filter out short lines.
wiki_train_ds = (
tf.data.TextLineDataset(
wiki_dir + "wiki.train.tokens",
)
.filter(lambda x: tf.strings.length(x) > 100)
.shuffle(buffer_size=500)
.batch(scaled_batch_size)
.cache()
.prefetch(tf.data.AUTOTUNE)
)
wiki_val_ds = (
tf.data.TextLineDataset(wiki_dir + "wiki.valid.tokens")
.filter(lambda x: tf.strings.length(x) > 100)
.shuffle(buffer_size=500)
.batch(scaled_batch_size)
.cache()
.prefetch(tf.data.AUTOTUNE)
)
wiki_test_ds = (
tf.data.TextLineDataset(wiki_dir + "wiki.test.tokens")
.filter(lambda x: tf.strings.length(x) > 100)
.shuffle(buffer_size=500)
.batch(scaled_batch_size)
.cache()
.prefetch(tf.data.AUTOTUNE)
)
在上面的代碼中，我們下載了 wikitext-2 資料集並解壓縮。然後，我們定義了三個資料集：wiki_train_ds、wiki_val_ds 和 wiki_test_ds。這些資料集經過過濾以移除短行，並進行批次處理以提高訓練效率。
在 NLP 訓練/微調中使用衰減學習率是一種常見的做法。我們將在這裡使用 PolynomialDecay 調度。
total_training_steps = sum(1 for _ in wiki_train_ds.as_numpy_iterator()) * EPOCHS
lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(
initial_learning_rate=scaled_learning_rate,
decay_steps=total_training_steps,
end_learning_rate=0.0,
)
class PrintLR(tf.keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs=None):
print(
f"\nLearning rate for epoch {epoch + 1} is {model_dist.optimizer.learning_rate.numpy()}"
)
我們還要創建一個指向 TensorBoard 的回調，這將允許我們在稍後訓練模型時可視化不同的指標。我們將所有回調放在一起，如下所示
callbacks = [
tf.keras.callbacks.TensorBoard(log_dir="./logs"),
PrintLR(),
]
print(tf.config.list_physical_devices("GPU"))
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
準備好資料集後，我們現在在 strategy.scope() 中初始化並編譯我們的模型和優化器
with strategy.scope():
# Everything that creates variables should be under the strategy scope.
# In general this is only model construction & `compile()`.
model_dist = keras_nlp.models.BertMaskedLM.from_preset("bert_tiny_en_uncased")
# This line just sets pooled_dense layer as non-trainiable, we do this to avoid
# warnings of this layer being unused
model_dist.get_layer("bert_backbone").get_layer("pooled_dense").trainable = False
model_dist.compile(
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=tf.keras.optimizers.AdamW(learning_rate=scaled_learning_rate),
weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],
jit_compile=False,
)
model_dist.fit(
wiki_train_ds, validation_data=wiki_val_ds, epochs=EPOCHS, callbacks=callbacks
)
Epoch 1/3
Learning rate for epoch 1 is 0.00019999999494757503
239/239 ━━━━━━━━━━━━━━━━━━━━ 43s 136ms/step - loss: 3.7009 - sparse_categorical_accuracy: 0.1499 - val_loss: 1.1509 - val_sparse_categorical_accuracy: 0.3485
Epoch 2/3
239/239 ━━━━━━━━━━━━━━━━━━━━ 0s 122ms/step - loss: 2.6094 - sparse_categorical_accuracy: 0.5284
Learning rate for epoch 2 is 0.00019999999494757503
239/239 ━━━━━━━━━━━━━━━━━━━━ 32s 133ms/step - loss: 2.6038 - sparse_categorical_accuracy: 0.5274 - val_loss: 0.9812 - val_sparse_categorical_accuracy: 0.4006
Epoch 3/3
239/239 ━━━━━━━━━━━━━━━━━━━━ 0s 123ms/step - loss: 2.3564 - sparse_categorical_accuracy: 0.6053
Learning rate for epoch 3 is 0.00019999999494757503
239/239 ━━━━━━━━━━━━━━━━━━━━ 32s 134ms/step - loss: 2.3514 - sparse_categorical_accuracy: 0.6040 - val_loss: 0.9213 - val_sparse_categorical_accuracy: 0.4230
在範圍內擬合模型後，我們可以正常評估它！
model_dist.evaluate(wiki_test_ds)
29/29 ━━━━━━━━━━━━━━━━━━━━ 3s 60ms/step - loss: 1.9197 - sparse_categorical_accuracy: 0.8527
[0.9470901489257812, 0.4373602867126465]
針對跨多台機器進行的分散式訓練（與僅利用單台機器上的多個裝置進行訓練不同），您可以使用兩種分散式策略：MultiWorkerMirroredStrategy 和 ParameterServerStrategy。
tf.distribute.MultiWorkerMirroredStrategy是一種同步 CPU/GPU 多工作節點解決方案，可與 Keras 風格的模型建立和訓練迴圈搭配使用，並在副本之間同步縮減梯度。tf.distribute.experimental.ParameterServerStrategy是一種非同步 CPU/GPU 多工作節點解決方案，其中參數儲存在參數伺服器上，而工作節點會非同步地將梯度更新到參數伺服器。