分散式訓練

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

程式碼需求

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

訓練叢集的結構

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

系統會在分散式訓練中，將一個角色或工作分配給訓練叢集中的每個備用資源。例如：

  • 主要備用資源：只有一個備用資源可指派為「主要備用資源」。這項工作能管理其他工作，並回報整體工作的狀態。

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

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

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

設定分散式訓練工作

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

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

在 workerPoolSpecs[] 中的順位 在叢集中執行的工作
第一 (workerPoolSpecs[0]) 主要、主、排程或「master」
秒 (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，可以指定 worker 集區來設定 worker 備用資源、參數伺服器備用資源和評估器備用資源。

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

如果您要指定的工作只包含主要複本和參數伺服器工作站集區，則必須為工作站集區 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
     {}
   ]
   ...
}

使用 Reduction Server 縮短訓練時間

使用多個節點訓練大型機器學習模型時，節點之間的梯度傳輸可能會造成大量延遲。縮減伺服器是一種全縮減演算法，可提高分散式訓練的總處理量並縮短延遲時間。Vertex AI 會在 Docker 容器映像檔中提供 Reduction Server，方便您在分散式訓練期間使用其中一個 worker 集區。

如要瞭解縮減伺服器的運作方式，請參閱「在 Vertex AI 上使用縮減伺服器加快 GPU 分散式訓練」。

事前準備

只要符合下列條件，即可使用 Reduction Server：

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

  • 訓練程式碼使用 TensorFlow 或 PyTorch，並針對多主機資料並行訓練設定，使用 NCCL 全減法進行 GPU 訓練。(您也可以使用其他使用 NCCL 的 ML 框架)。

  • 在主要節點 (workerPoolSpecs[0]) 和 worker (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 進行訓練

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

  1. 在第三個 worker 池 (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. 選取第三個工作站集區的機器類型和節點數量時，請確認第三個工作站集區的總網路頻寬與第一和第二個工作站集區的總網路頻寬相符或超出。

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

    您不會為 Reduction Server 節點使用 GPU。如要瞭解第三個工作池中每個節點的可用頻寬上限，請參閱「通用機器系列」中的「輸出頻寬上限 (Gbps)」欄。

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

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

以下指令提供範例，說明如何建立使用 Reduction Server 的 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 訓練中，每個 worker 都需要連線至所有 reducer 主機。如要盡可能減少 worker 主機上的連線數，請為 reducer 主機使用網路頻寬最高的機器類型。

對於減算主機而言，一般用途的 N1/N2 VM 是理想選擇，其中至少有 16 個 vCPU，可提供 32 Gbps 輸出頻寬，例如 n1-highcpu-16 和 n2-highcpu-16。N1/N2 VM 的 Tier 1 VM 頻寬可提高最高輸出頻寬，範圍介於 50 Gbps 和 100 Gbps 之間，因此是用於減算器 VM 節點的理想選擇。

worker 和 reducer 的總傳出頻寬應相同。舉例來說，如果您使用 8 個 a2-megagpu-16g VM 做為 worker，則至少應使用 25 個 n1-highcpu-16 VM 做為 reducer。

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

將小型訊息批次處理

如果要匯總的訊息量相當大，Reduction Server 的運作效果就會最好。大多數 ML 架構已提供不同術語的技術，可在執行 all-reduce 前將小型梯度張量分批處理。

Horovod

Horovod 支援張量融合功能，可將小型張量分批處理，以便進行全減法運算。張量會填入融合緩衝區，直到緩衝區已完全填滿，並執行緩衝區的所有縮減作業。您可以設定 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"

工作的索引，從零開始。舉例來說，如果訓練工作包含兩個 worker，則這個值會在其中一個 worker 上設為 0，在另一個 worker 上設為 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 上設為 worker、在參數伺服器上設為 ps，以及在評估器上設為 evaluator。在工作主工作者中，值會設為 chief 或 master。
index

工作的索引，從零開始。舉例來說，如果訓練工作包含兩個 worker，則這個值會在其中一個 worker 上設為 0，另一個 worker 上設為 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 欄位讀取目前的試用序號。

後續步驟