使用Arena提交PyTorch分布式训练作业

PyTorch是一个开源的深度学习框架,广泛应用于各种深度学习模型的训练任务中,本文演示如何使用Arena提交PyTorch多机多卡训练作业,并通过TensorBoard可视化查看训练作业。

前提条件

背景信息

本文将从远程Git仓库中下载训练代码,并从共享存储系统(基于NASPVPVC)中读取训练数据。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%)

可以看到集群中有2GPU节点,每个GPU节点都包含2张空闲的GPU卡可用于运行训练作业。

步骤二:提交PyTorch训练作业

使用 arena submit pytorch 命令提交一个PyTorch多机多卡训练作业,其中包含2worker,每个worker使用2GPU卡。

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>为节点在分布式训练中承担的角色,包括 masterworker<index> 则为节点序号。例如,对于名为 pytorch-mnist 的训练作业,当指定 --workers=3--nproc-per-node=2时,将会创建 3 个训练节点并在每个节点上启动两个进程,节点名称分别为 pytorch-mnist-master-0pytorch-mnist-worker-0pytorch-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_USERNAMEGIT_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,表明该作业包含1Master节点和2Worker节点。

0

--gpus

可选

指定作业Worker节点需要使用的GPU卡数。

0

--working-dir

可选

指定当前执行命令所在的目录。

/root

--image

必选

指定训练环境的镜像地址。

--sync-mode

可选

同步代码的模式,您可以指定gitrsync。本文使用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训练作业

  1. 执行如下命令查看当前通过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
  2. 执行以下命令检查作业所使用的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
  3. 执行以下命令检查集群所使用的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%)

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

  4. 执行以下命令获取作业详情。

    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,它们分别请求了2GPU卡,共同参与整个训练任务。

    说明

    当在训练作业中启用了TensorBoard时,作业详情中将会展示TensorBoard访问地址;如果没有启用TensorBoard,则不会展示。

步骤四:查看TensorBoard

  1. 在本地执行如下命令,将集群中的TensorBoard的访问端口6006映射到本地的9090端口。

    重要

    请注意kubectl port-forward建立的端口转发不具备生产级别的可靠性、安全性和扩展性,因此仅适用于开发和调试目的,不适合在生产环境使用。更多关于Kubernetes集群内生产可用的网络方案的信息,请参见Ingress管理

    kubectl port-forward -n default svc/pytorch-mnist-tensorboard 9090:6006
  2. 在浏览器中访问 http://127.0.0.1:9090,即可查看TensorBoard。

    pytorch单机

    说明

    本文中PyTorch示例代码,默认每10epoch写入event信息。如果您修改了--epochs参数,请修改为10的整数倍,否则无法在TensorBoard上看到数据。

步骤五:查看训练作业日志

  1. 执行以下命令查看训练作业中 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%)
  2. 查看 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