分散式訓練

本頁說明如何在 Vertex AI 中執行分散式訓練工作。

程式碼需求

使用支援分散式訓練的 ML 架構。在訓練程式碼中，您可以使用 CLUSTER_SPEC 或 TF_CONFIG 環境變數，參照訓練叢集的特定部分。

訓練叢集的結構

如果您使用 Vertex AI 執行分散式訓練工作，則要在「訓練叢集」中指定多個機器 (節點)。訓練服務會為您指定的機器類型分配資源。您在特定節點上執行的工作稱為「備用資源」。具有相同設定的副本群組稱為工作站集區。

訓練叢集中的每個副本都會在分散式訓練中獲得單一角色或工作。例如：

  • 主要副本：只有一個副本會被指派為「主要副本」。這項工作會管理其他工作，並回報整體工作的狀態。

  • 工作站：您可以指派一或多個備用資源做為「工作站」。這些備用資源會按照您在工作設定中的設定執行自己分內的任務。

  • 參數伺服器：如果 ML 架構支援，您可以指派一或多個備用資源做為「參數伺服器」。這些備份資源會儲存模型參數，以及協調工作站之間的共用模型狀態。

  • 評估人員：如果機器學習架構支援，您可以指派一或多個備用資源做為「評估人員」。這些副本可用於評估模型。如果您使用 TensorFlow，請注意 TensorFlow 通常會預期您使用的評估工具不超過一個。

設定分散式訓練工作

您可以定義多個工作站集區，將任何自訂訓練工作設定為分散式訓練工作。您也可以在訓練管道或超參數調整工作中執行分散式訓練。

如要設定分散式訓練工作，請定義工作站集區清單 (workerPoolSpecs[])，並為每種工作類型指定一個 WorkerPoolSpec：

在「workerPoolSpecs[]」中的位置 叢集中執行的工作
第一 (workerPoolSpecs[0]) 主要、主管、排程人員或「主控」
第二 (workerPoolSpecs[1]) 次要、副本、工作站
第三週 (workerPoolSpecs[2]) 參數伺服器、縮減伺服器
第四個 (workerPoolSpecs[3]) 評估人員

您必須指定主要副本，負責協調所有其他副本完成的工作。第一個工作站集區規格僅適用於主要副本，並將其 replicaCount 設為 1：

{
  "workerPoolSpecs": [
     // `WorkerPoolSpec` for worker pool 0, primary replica, required
     {
       "machineSpec": {...},
       "replicaCount": 1,
       "diskSpec": {...},
       ...
     },
     // `WorkerPoolSpec` for worker pool 1, optional
     {},
     // `WorkerPoolSpec` for worker pool 2, optional
     {},
     // `WorkerPoolSpec` for worker pool 3, optional
     {}
   ]
   ...
}

指定其他工作站集區

視機器學習架構而定，您可能會指定其他用途的額外工作站集區。舉例來說，如果您使用 TensorFlow，可以指定工作站集區來設定工作站副本、參數伺服器副本和評估器副本。

您在 workerPoolSpecs[] 清單中指定的工作站集區順序，會決定工作站集區類型。為不想使用的工作站集區設定空白值，以便在 workerPoolSpecs[] 清單中略過這些集區，並指定要使用的工作站集區。例如：

如要指定只有主要副本和參數伺服器工作站集區的工作，您必須為工作站集區 1 設定空白值：

{
  "workerPoolSpecs": [
     // `WorkerPoolSpec` for worker pool 0, required
     {
       "machineSpec": {...},
       "replicaCount": 1,
       "diskSpec": {...},
       ...
     },
     // `WorkerPoolSpec` for worker pool 1, optional
     {},
     // `WorkerPoolSpec` for worker pool 2, optional
     {
       "machineSpec": {...},
       "replicaCount": 1,
       "diskSpec": {...},
       ...
     },
     // `WorkerPoolSpec` for worker pool 3, optional
     {}
   ]
   ...
}

使用縮減伺服器縮短訓練時間

使用多個節點訓練大型機器學習模型時，節點間的梯度通訊可能會造成顯著延遲。縮減伺服器是一種全縮減演算法，可提高總處理量並縮短分散式訓練的延遲時間。Vertex AI 會在 Docker 容器映像檔中提供縮減伺服器，您可以在分散式訓練期間，將該伺服器用於其中一個工作站集區。

如要瞭解縮減伺服器的運作方式，請參閱「Faster distributed GPU training with Reduction Server on Vertex AI」(透過 Vertex AI 上的縮減伺服器加快分散式 GPU 訓練速度)。

必要條件

如要使用縮減伺服器，必須符合下列條件：

  • 您正在使用 GPU 工作人員執行分散式訓練。

  • 您的訓練程式碼使用 TensorFlow 或 PyTorch，並設定為透過 NCCL all-reduce，以 GPU 進行多主機資料平行訓練。(您或許也能使用其他採用 NCCL 的 ML 架構)。

  • 主要節點 (workerPoolSpecs[0]) 和工作站 (workerPoolSpecs[1]) 上執行的容器支援 Reduction Server。具體來說，每個容器都是下列其中一種：

    • 預先建構的 TensorFlow 訓練容器，版本 2.3 以上。

    • 預先建構的 PyTorch 訓練容器，版本 1.4 以上。

    • 已安裝 NCCL 2.7 以上版本和 google-reduction-server 套件的自訂容器。如要在自訂容器映像檔上安裝這個套件，請在 Dockerfile 中新增下列程式碼：

      RUN echo "deb https://packages.cloud.google.com/apt google-fast-socket main" | tee /etc/apt/sources.list.d/google-fast-socket.list && \
    curl -s -L https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - && \
    apt update && apt install -y google-reduction-server

使用縮減伺服器進行訓練

如要使用 Reduction Server，請在建立自訂訓練資源時執行下列操作：

  1. 在第三個工作站集區 (workerPoolSpecs[2]) 的「containerSpec.imageUri」欄位中，指定下列其中一個 URI：

    • us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest
    • europe-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest
    • asia-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest

    選擇最接近執行自訂訓練位置的多區域，或許可以縮短延遲時間。

  2. 選取第三個工作站集區的機器類型和節點數量時，請確保第三個工作站集區的總網路頻寬，等於或超過第一個和第二個工作站集區的總網路頻寬。

    如要瞭解第二個工作站集區中每個節點可用的最大頻寬，請參閱「網路頻寬和 GPU」。

    您不會將 GPU 用於 Reduction Server 節點。如要瞭解第三個工作站集區中每個節點的可用頻寬上限，請參閱「一般用途機器系列」中的「輸出頻寬上限 (Gbps)」資料欄。

    舉例來說，如果您將第一個和第二個工作站集區設定為使用 5 個 n1-highmem-96 節點，每個節點有 8 個 NVIDIA_TESLA_V100 GPU，則每個節點的最大可用頻寬為 100 Gbps，總頻寬為 500 Gbps。如要在第三個工作站集區中達到這個頻寬，您可以使用 16 個 n1-highcpu-16 節點，每個節點的頻寬上限為 32 Gbps，總頻寬為 512 Gbps。

    建議您為 Reduction Server 節點使用 n1-highcpu-16 機器類型，因為這種機器類型可為資源提供相對較高的頻寬。

以下指令提供範例，說明如何建立使用縮減伺服器的 CustomJob 資源：

gcloud ai custom-jobs create \
  --region=LOCATION \
  --display-name=JOB_NAME \
  --worker-pool-spec=machine-type=n1-highmem-96,replica-count=1,accelerator-type=NVIDIA_TESLA_V100,accelerator-count=8,container-image-uri=CUSTOM_CONTAINER_IMAGE_URI \
  --worker-pool-spec=machine-type=n1-highmem-96,replica-count=4,accelerator-type=NVIDIA_TESLA_V100,accelerator-count=8,container-image-uri=CUSTOM_CONTAINER_IMAGE_URI \
  --worker-pool-spec=machine-type=n1-highcpu-16,replica-count=16,container-image-uri=us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest

詳情請參閱建立 CustomJob 的指南。

使用 Reduction Server 訓練模型的最佳做法

機器類型和數量

在 Reduction Server 訓練中，每個工作站都需要連線至所有縮減器主機。如要盡量減少工作站主機上的連線數，請為縮減器主機使用網路頻寬最高的機器類型。

建議為縮減器主機選用一般用途的 N1/N2 VM，至少要有 16 個 vCPU，並提供32 Gbps 的輸出頻寬，例如 n1-highcpu-16 和 n2-highcpu-16。N1/N2 VM 的第 1 層 VM 頻寬會提高輸出頻寬上限，範圍介於 50 Gbps 和 100 Gbps 之間，因此非常適合做為縮減器 VM 節點。

工作站和縮減器的總輸出頻寬應相同。舉例來說，如果您使用 8 個 a2-megagpu-16g VM 做為工作站，則至少應使用 25 個 n1-highcpu-16 VM 做為縮減器。

`(8 worker VMs * 100 Gbps) / 32 Gbps egress = 25 reducer VMs`.

將小型訊息批次處理

如果待彙整的訊息夠大，縮減伺服器就能發揮最大效益。大多數機器學習架構都已提供技術，可使用不同術語批次處理小型梯度張量，然後執行全縮減作業。

Horovod

Horovod 支援 Tensor Fusion，可批次處理所有縮減作業的小型張量。張量會填入融合緩衝區，直到緩衝區填滿為止，然後對緩衝區執行全縮減作業。您可以設定 HOROVOD_FUSION_THRESHOLD 環境變數，調整融合緩衝區的大小。

建議將 HOROVOD_FUSION_THRESHOLD 環境變數的值設為至少 128 MB。在本例中，請將 HOROVOD_FUSION_THRESHOLD 環境變數設為 134217728 (128 * 1024 * 1024)。

PyTorch

PyTorch DistributedDataParallel 支援批次訊息，做為「梯度分組」。在 DistributedDataParallel 建構函式中設定 bucket_cap_mb 參數，即可控管批次值區的大小。預設大小為 25 MB。

最佳做法：建議 bucket_cap_mb 的值為 64 (64 MB)。

叢集的環境變數

Vertex AI 會在每個副本上填入環境變數 CLUSTER_SPEC，以描述整體叢集的設定方式。如同 TensorFlow 的 TF_CONFIG，CLUSTER_SPEC 會描述叢集中的每個備用資源，包括索引和角色 (主要備用資源、工作站、參數伺服器或評估器)。

使用 TensorFlow 執行分散式訓練時，系統會剖析 TF_CONFIG 以建構 tf.train.ClusterSpec。同樣地，您使用其他機器學習架構執行分散式訓練時，必須剖析 CLUSTER_SPEC 以填入架構需要的任何環境變數或設定。

CLUSTER_SPEC 的格式

CLUSTER_SPEC 環境變數是採用下列格式的 JSON 字串：

索引鍵 說明
"cluster"

自訂容器的叢集描述。和 TF_CONFIG 一樣，這個物件會依照 TensorFlow 叢集規格格式化，且可傳送至 tf.train.ClusterSpec 建構函式。

叢集說明包含您指定的每個工作站集區的副本名稱清單。

"workerpool0" 所有分散式訓練工作的第一個工作站集區中，都有一個主要副本。
"workerpool1" 這個工作站集區包含工作站副本 (如果您在建立工作時指定的話)。
"workerpool2" 如果您在建立工作時指定參數伺服器，這個工作站集區就會包含這些伺服器。
"workerpool3" 如果您在建立工作時指定了評估人員，這個工作站集區就會包含這些人員。
"environment" 字串 cloud。
"task" 說明執行程式碼的特定節點的工作。 您可以使用此資訊，為分散式工作中的特定工作站編寫程式碼。這個項目是含有下列索引鍵的字典：
"type" 這項工作執行的 worker 集區類型。舉例來說，"workerpool0" 是指主要副本。
"index"

工作的索引，從零開始。舉例來說，如果訓練工作包含兩個工作站，這個值會設為 0 (其中一個工作站) 和 1 (另一個工作站)。
"trial" 目前正在執行的超參數調整試驗 ID。 為工作設定超參數調整時，您會設定要訓練的多個試驗。這個值讓您可以在程式碼中區別正在執行的不同試驗。ID 是含有試驗編號的字串值，從 1 開始。
job

您提供用來建立目前訓練工作的 CustomJobSpec，以字典形式表示。

CLUSTER_SPEC 範例

範例值如下：


{
   "cluster":{
      "workerpool0":[
         "cmle-training-workerpool0-ab-0:2222"
      ],
      "workerpool1":[
         "cmle-training-workerpool1-ab-0:2222",
         "cmle-training-workerpool1-ab-1:2222"
      ],
      "workerpool2":[
         "cmle-training-workerpool2-ab-0:2222",
         "cmle-training-workerpool2-ab-1:2222"
      ],
      "workerpool3":[
         "cmle-training-workerpool3-ab-0:2222",
         "cmle-training-workerpool3-ab-1:2222",
         "cmle-training-workerpool3-ab-2:2222"
      ]
   },
   "environment":"cloud",
   "task":{
      "type":"workerpool0",
      "index":0,
      "trial":"TRIAL_ID"
   },
   "job": {
      ...
   }
}

TF_CONFIG 的格式

除了 CLUSTER_SPEC，Vertex AI 也會在所有分散式訓練工作的每個副本上設定 TF_CONFIG 環境變數。Vertex AI 不會為單一副本訓練工作設定 TF_CONFIG。

CLUSTER_SPEC 和 TF_CONFIG 共用部分值，但格式不同。這兩個環境變數都包含 TensorFlow 以外的其他欄位。

使用自訂容器時，透過 TensorFlow 進行分散式訓練的運作方式，與使用預先建構的容器相同。

TF_CONFIG 環境變數是採用下列格式的 JSON 字串：

TF_CONFIG 個欄位
cluster

TensorFlow 叢集說明。字典會將一或多個工作名稱 (chief、worker、ps 或 master) 對應至執行這些工作的網路位址清單。對於特定訓練工作，每個 VM 上的這個字典都相同。

這是 tf.train.ClusterSpec 建構函式的有效第一個引數。請注意，這個字典絕不會包含 evaluator 做為鍵，因為即使您在工作中使用了評估人員，他們也不會被視為訓練叢集的一部分。
task

設定這個環境變數的 VM 工作說明。 針對特定訓練工作，每個 VM 上的這個字典都不相同。您可以使用這項資訊，自訂分散式訓練工作中的每個 VM 執行的程式碼。您也可以使用這項功能，在超參數調整工作的不同試驗中，變更訓練程式碼的行為。

這個字典包含下列鍵/值組合：

task 個欄位
type

這個 VM 執行的工作類型。這個值會設為工作站上的 worker、參數伺服器上的 ps，以及評估人員上的 evaluator。在工作的主控工作者上，值會設為 chief 或 master。
index

工作的索引，從零開始。舉例來說，如果訓練工作包含兩個工作站，這個值會設為 0 (其中一個工作站) 和 1 (另一個工作站)。
trial

目前在這個 VM 上執行的超參數調整試驗 ID。只有在目前的訓練工作是超參數調整工作時，才會設定這個欄位。

對於超參數調整工作，Vertex AI 會在多次試驗中重複執行訓練程式碼，每次使用不同的超參數。這個欄位包含目前的試驗編號，從第一次試驗的 1 開始。
cloud

Vertex AI 內部使用的 ID。您可以忽略這個欄位。
job

您提供用來建立目前訓練工作的 CustomJobSpec，以字典形式表示。
environment

字串 cloud。

TF_CONFIG 範例

下列程式碼範例會將 TF_CONFIG 環境變數列印至訓練記錄：

import json
import os

tf_config_str = os.environ.get('TF_CONFIG')
tf_config_dict  = json.loads(tf_config_str)

# Convert back to string just for pretty printing
print(json.dumps(tf_config_dict, indent=2))

在執行階段 2.1 以上版本中執行的超參數調整工作，如果使用主要工作站、兩個工作站和一個參數伺服器，這段程式碼會在第一次超參數調整試驗期間，為其中一個工作站產生下列記錄。為求簡潔，範例輸出會隱藏 job 欄位，並將部分 ID 替換為一般值。

{
  "cluster": {
    "chief": [
      "training-workerpool0-[ID_STRING_1]-0:2222"
    ],
    "ps": [
      "training-workerpool2-[ID_STRING_1]-0:2222"
    ],
    "worker": [
      "training-workerpool1-[ID_STRING_1]-0:2222",
      "training-workerpool1-[ID_STRING_1]-1:2222"
    ]
  },
  "environment": "cloud",
  "job": {
    ...
  },
  "task": {
    "cloud": "[ID_STRING_2]",
    "index": 0,
    "trial": "1",
    "type": "worker"
  }
}

使用時機 TF_CONFIG

TF_CONFIG 僅適用於分散式訓練工作。

您可能不需要在訓練程式碼中直接與 TF_CONFIG 環境變數互動。只有在 TensorFlow 的分配策略和 Vertex AI 的標準超參數調整工作流程 (兩者皆在下節中說明) 無法用於您的工作時，才存取 TF_CONFIG 環境變數。

分散式訓練

Vertex AI 會設定 TF_CONFIG 環境變數，以擴充 TensorFlow 分散式訓練所需的規格。

如要使用 TensorFlow 執行分散式訓練，請使用 tf.distribute.Strategy API。具體來說，我們建議您搭配使用 Keras API 和 MultiWorkerMirroredStrategy，或是在為工作指定參數伺服器時，搭配使用 ParameterServerStrategy。不過請注意，TensorFlow 僅為這些策略提供實驗支援。

這些分配策略會使用 TF_CONFIG 環境變數，為訓練工作中的每個 VM 指派角色，並促進 VM 之間的通訊。您不需要在訓練程式碼中直接存取 TF_CONFIG 環境變數，因為 TensorFlow 會為您處理這項作業。

如要自訂執行訓練工作的不同 VM 行為，請直接剖析 TF_CONFIG 環境變數。

超參數調整

執行超參數調整工作時，Vertex AI 會為每次試驗提供不同的訓練程式碼引數。訓練程式碼不一定需要知道目前執行的試用期。此外，您可以在 Google Cloud 控制台中監控超參數調整工作的進度。

如有需要，您的程式碼可以從TF_CONFIG 環境變數的 task 欄位中的 trial 欄位讀取目前的試用次數。

後續步驟