使用Arena提交PyTorch分布式訓練作業 - Container Service for Kubernetes

PyTorch是一個開源的深度學習架構，廣泛應用於各種深度學習模型的訓練任務中，本文示範如何使用Arena提交PyTorch多機多卡訓練作業，並通過TensorBoard可視化查看訓練作業。

前提條件

已建立包含GPU的Kubernetes叢集。具體操作，請參見建立包含GPU的Kubernetes叢集。
叢集節點可以訪問公網。具體操作，請參見為已有叢集開啟公網訪問能力。
已安裝Arena工具。具體操作，請參見配置Arena用戶端。
建立一個名為training-data的PVC執行個體，並在路徑/pytorch_data下存放MNIST資料集。具體操作，請參見配置NAS共用儲存。

背景資訊

本文將從遠程Git倉庫中下載訓練代碼，並從共用儲存系統（基於NAS的PV和PVC）中讀取訓練資料。torchrun 是 PyTorch 提供的一個命令列工具，用於簡化和管理分布式訓練任務，本文將使用 torchrun 來啟動 PyTorch 多機多卡訓練任務，訓練代碼詳見 main.py。

操作步驟

步驟一：查看GPU資源

執行以下命令查看叢集中可用的GPU資源。

arena top node

預期輸出：

NAME                        IPADDRESS        ROLE    STATUS  GPU(Total)  GPU(Allocated)
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   0           0
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   0           0
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   2           0
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   2           0
---------------------------------------------------------------------------------------------------
Allocated/Total GPUs In Cluster:
0/4 (0.0%)

可以看到叢集中有2個GPU節點，每個GPU節點都包含2張閒置GPU卡可用於運行訓練作業。

步驟二：提交PyTorch訓練作業

使用 arena submit pytorch 命令提交一個PyTorch多機多卡訓練作業，其中包含2個worker，每個worker使用2張GPU卡。

arena submit pytorch \
    --name=pytorch-mnist \
    --namespace=default \
    --workers=2 \
    --gpus=2 \
    --nproc-per-node=2 \
    --clean-task-policy=None \
    --working-dir=/root \
    --image=kube-ai-registry.cn-shanghai.cr.aliyuncs.com/kube-ai/pytorch-with-tensorboard:2.5.1-cuda12.4-cudnn9-runtime \
    --sync-mode=git \
    --sync-source=https://github.com/kubeflow/arena.git \
    --env=GIT_SYNC_BRANCH=v0.13.1 \
    --data=training-data:/mnt \
    --tensorboard \
    --logdir=/mnt/pytorch_data/logs \
    "torchrun /root/code/arena/examples/pytorch/mnist/main.py --epochs 10 --backend nccl --data /mnt/pytorch_data  --dir /mnt/pytorch_data/logs"

預期輸出：

service/pytorch-mnist-tensorboard created
deployment.apps/pytorch-mnist-tensorboard created
pytorchjob.kubeflow.org/pytorch-mnist created
INFO[0002] The Job pytorch-mnist has been submitted successfully
INFO[0002] You can run `arena get pytorch-mnist --type pytorchjob -n default` to check the job status

說明

PyTorch 分布式訓練作業相比於單機訓練需要額外指定 --workers 和 --nproc-per-node 參數，它們分別表示參與分布式訓練的節點（Pod）數量以及每個節點上啟動的進程數量。一個分布式訓練作業將會包含若干個節點，節點名稱格式為 <job_name>-<role_name>-<index>，其中<job_name> 為作業名稱，<role_name>為節點在分布式訓練中承擔的角色，包括 master 和 worker，<index> 則為節點序號。例如，對於名為 pytorch-mnist 的訓練作業，當指定 --workers=3 和 --nproc-per-node=2時，將會建立 3 個訓練節點並在每個節點上啟動兩個進程，節點名稱分別為 pytorch-mnist-master-0、pytorch-mnist-worker-0 和 pytorch-mnist-worker-1，其中每個節點均會注入相應的環境變數，具體如下表所示：

環境變數 \ 節點名稱	pytorch-mnist-master-0	pytorch-mnist-worker-0	pytorch-mnist-worker-1
`MASTER_ADDR`	pytorch-mnist-master-0
`MASTER_PORT`	23456
`WORLD_SIZE`	6
`RANK`	0	1	2
`PET_MASTER_ADDR`	pytorch-mnist-master-0
`PET_MASTER_PORT`	23456
`PET_NNODES`	3
`PET_NODE_RANK`	0	1	2

說明

如果您使用的是非公開Git代碼倉庫，則可以通過配置環境變數GIT_SYNC_USERNAME和GIT_SYNC_PASSWORD的方式來設定Git使用者名稱和密碼。

  arena submit pytorch \
        ...
        --sync-mode=git \
        --sync-source=https://github.com/kubeflow/arena.git \
        --env=GIT_SYNC_BRANCH=v0.13.1 \
        --env=GIT_SYNC_USERNAME=<username> \
        --env=GIT_SYNC_PASSWORD=<password> \
        "torchrun /root/code/arena/examples/pytorch/mnist/main.py --epochs 10 --backend nccl --data /mnt/pytorch_data  --dir /mnt/pytorch_data/logs"

arena命令使用git-sync同步原始碼，因此您可以使用任何在git-sync專案中定義的環境變數。

重要

本文樣本從GitHub倉庫中拉取原始碼，如遇到網路原因等導致代碼無法成功拉取時，可以手動將代碼下載到共用儲存系統中，假設您已經將代碼下載至NAS中的/code/github.com/kubeflow/arena路徑中，則可以如下提交訓練作業。

arena submit pytorch \
    --name=pytorch-mnist \
    --namespace=default \
    --workers=2 \
    --gpus=2 \
    --nproc-per-node=2 \
    --clean-task-policy=None \
    --working-dir=/root \
    --image=kube-ai-registry.cn-shanghai.cr.aliyuncs.com/kube-ai/pytorch-with-tensorboard:2.5.1-cuda12.4-cudnn9-runtime \
    --data=training-data:/mnt \
    --tensorboard \
    --logdir=/mnt/pytorch_data/logs \
    "torchrun /mnt/code/github.com/kubeflow/arena/examples/pytorch/mnist/main.py --epochs 10 --backend nccl --data /mnt/pytorch_data  --dir /mnt/pytorch_data/logs"

參數解釋如下表所示：

參數	是否必選	解釋	預設值
`--name`	必選	作業名稱，全域唯一，不能重複。	無
`--namespace`	可選	作業所屬的命名空間。	`default`
`--workers`	可選	指定作業需要的節點數，它設定的值包含Master節點。例如設定3，表明該作業包含1個Master節點和2個Worker節點。	`0`
`--gpus`	可選	指定作業Worker節點需要使用的GPU卡數。	`0`
`--working-dir`	可選	指定當前執行命令所在的目錄。	`/root`
`--image`	必選	指定訓練環境的鏡像地址。	無
`--sync-mode`	可選	同步代碼的模式，您可以指定git、rsync。本文使用Git模式。	無
`--sync-source`	可選	同步代碼的倉庫地址，需要和--sync-mode一起使用，本文樣本使用Git模式，該參數可以為任何GitHub專案、阿里雲Code專案等支援Git的代碼託管地址。專案代碼將會被下載到--working-dir下的`code/`目錄中，在本樣本中則為`/root/code/arena`。	無
`--data`	可選	掛載共用儲存卷PVC到運行環境中。它由兩部分組成，通過分號`:`分割。分號左側是您已經準備好的PVC名稱。您可以通過`arena data list`查看當前叢集可用的PVC列表；分號右側是您想將PVC的掛載到運行環境中的路徑，也是您訓練代碼要讀取資料的本地路徑。這樣通過掛載的方式，您的代碼就可以訪問PVC的資料。說明執行`arena data list`查看本文樣本當前叢集可用的PVC列表。 `NAME ACCESSMODE DESCRIPTION OWNER AGE training-data ReadWriteMany 35m` 如果沒有可用的PVC，您可建立PVC。詳情請參見配置NAS共用儲存。	無
`--tensorboard`	可選	為訓練任務開啟一個TensorBoard服務，用作資料視覺效果，您可以結合--logdir指定TensorBoard要讀取的event路徑。不指定該參數，則不開啟TensorBoard服務。	無
`--logdir`	可選	需要結合--tensorboard一起使用，該參數表示TensorBoard需要讀取event資料的路徑。	`/training_logs`

步驟三：查看PyTorch訓練作業

執行如下命令查看當前通過Arena提交的訓練作業：

arena list -n default

預期輸出：

NAME           STATUS   TRAINER     DURATION  GPU(Requested)  GPU(Allocated)  NODE
pytorch-mnist  RUNNING  PYTORCHJOB  48s       4               4               192.168.xxx.xxx

執行以下命令檢查作業所使用的GPU資源。

arena top job -n default

預期輸出：

NAME           STATUS   TRAINER     AGE  GPU(Requested)  GPU(Allocated)  NODE
pytorch-mnist  RUNNING  PYTORCHJOB  55s  4               4               192.168.xxx.xxx

Total Allocated/Requested GPUs of Training Jobs: 4/4

執行以下命令檢查叢集所使用的GPU資源。

arena top node

預期輸出：

NAME                        IPADDRESS        ROLE    STATUS  GPU(Total)  GPU(Allocated)
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   0           0
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   0           0
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   2           2
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   2           2
---------------------------------------------------------------------------------------------------
Allocated/Total GPUs In Cluster:
4/4 (100.0%)

可以看到全部4張GPU卡均已經被分配。

執行以下命令擷取作業詳情。

arena get pytorch-mnist -n default

預期輸出：

Name:        pytorch-mnist
Status:      RUNNING
Namespace:   default
Priority:    N/A
Trainer:     PYTORCHJOB
Duration:    1m
CreateTime:  2025-02-12 13:54:51
EndTime:

Instances:
  NAME                    STATUS   AGE  IS_CHIEF  GPU(Requested)  NODE
  ----                    ------   ---  --------  --------------  ----
  pytorch-mnist-master-0  Running  1m   true      2               cn-beijing.192.168.xxx.xxx
  pytorch-mnist-worker-0  Running  1m   false     2               cn-beijing.192.168.xxx.xxx

Tensorboard:
  Your tensorboard will be available on:
  http://192.168.xxx.xxx:32084

可以看到該作業包含1個名為 pytorch-mnist-master-0 的 master pod 和 1 個名為 pytorch-mnist-worker-0 的 worker pod，它們分別請求了2張GPU卡，共同參與整個訓練任務。

說明

當在訓練作業中啟用了TensorBoard時，作業詳情中將會展示TensorBoard訪問地址；如果沒有啟用TensorBoard，則不會展示。

步驟四：查看TensorBoard

在本地執行如下命令，將叢集中的TensorBoard的訪問連接埠6006映射到本地的9090連接埠。
重要
請注意kubectl port-forward建立的連接埠轉寄不具備生產層級的可靠性、安全性和擴充性，因此僅適用於開發和調試目的，不適合在生產環境使用。更多關於Kubernetes叢集內生產可用的網路方案的資訊，請參見Ingress管理。
```
kubectl port-forward -n default svc/pytorch-mnist-tensorboard 9090:6006
```
在瀏覽器中訪問 http://127.0.0.1:9090，即可查看TensorBoard。
說明
本文中PyTorch範例程式碼，預設每10個epoch寫入event資訊。如果您修改了--epochs參數，請修改為10的整數倍，否則無法在TensorBoard上看到資料。

步驟五：查看訓練作業日誌

執行以下命令查看訓練作業中 master 日誌：

arena logs -n default pytorch-mnist

預期輸出：

{'PID': 40, 'MASTER_ADDR': 'pytorch-mnist-master-0', 'MASTER_PORT': '23456', 'LOCAL_RANK': 0, 'RANK': 0, 'GROUP_RANK': 0, 'ROLE_RANK': 0, 'LOCAL_WORLD_SIZE': 2, 'WORLD_SIZE': 4, 'ROLE_WORLD_SIZE': 4}
{'PID': 41, 'MASTER_ADDR': 'pytorch-mnist-master-0', 'MASTER_PORT': '23456', 'LOCAL_RANK': 1, 'RANK': 1, 'GROUP_RANK': 0, 'ROLE_RANK': 1, 'LOCAL_WORLD_SIZE': 2, 'WORLD_SIZE': 4, 'ROLE_WORLD_SIZE': 4}
Using cuda:0.
Using cuda:1.
Train Epoch: 1 [0/60000 (0%)]   Loss: 2.283599
Train Epoch: 1 [0/60000 (0%)]   Loss: 2.283599
...
Train Epoch: 10 [59520/60000 (99%)]     Loss: 0.007343
Train Epoch: 10 [59520/60000 (99%)]     Loss: 0.007343

Accuracy: 9919/10000 (99.19%)


Accuracy: 9919/10000 (99.19%)

查看 index 為 0 的 worker 日誌：

arena logs -n default -i pytorch-mnist-worker-0 pytorch-mnist

預期輸出：

{'PID': 39, 'MASTER_ADDR': 'pytorch-mnist-master-0', 'MASTER_PORT': '23456', 'LOCAL_RANK': 0, 'RANK': 2, 'GROUP_RANK': 1, 'ROLE_RANK': 2, 'LOCAL_WORLD_SIZE': 2, 'WORLD_SIZE': 4, 'ROLE_WORLD_SIZE': 4}
{'PID': 40, 'MASTER_ADDR': 'pytorch-mnist-master-0', 'MASTER_PORT': '23456', 'LOCAL_RANK': 1, 'RANK': 3, 'GROUP_RANK': 1, 'ROLE_RANK': 3, 'LOCAL_WORLD_SIZE': 2, 'WORLD_SIZE': 4, 'ROLE_WORLD_SIZE': 4}
Using cuda:0.
Using cuda:1.
Train Epoch: 1 [0/60000 (0%)]   Loss: 2.283599
Train Epoch: 1 [0/60000 (0%)]   Loss: 2.283599
...
Train Epoch: 10 [58880/60000 (98%)]     Loss: 0.051877Train Epoch: 10 [58880/60000 (98%)]       Loss: 0.051877

Train Epoch: 10 [59520/60000 (99%)]     Loss: 0.007343Train Epoch: 10 [59520/60000 (99%)]       Loss: 0.007343


Accuracy: 9919/10000 (99.19%)


Accuracy: 9919/10000 (99.19%)

說明

如果需要即時查看作業日誌，可以添加 -f 參數；
如果僅需要查看最後 N 行日誌，可以添加 -t N 或 --tail N 參數；
更多用法請參見 arena logs --help。

（可選）步驟六：環境清理

訓練作業執行結束後如不再需要，執行如下命令進行刪除：

arena delete pytorch-mnist -n default

預期輸出：

INFO[0001] The training job pytorch-mnist has been deleted successfully