前提條件

• 已開通包含GPU/NPU的Kubernetes集群。

• 已安裝智算套件。

背景信息

本文演示如何提交一個PyTorch的分布式訓練任務，相關的數據已經包含在容器鏡像中。若用戶自有模型或訓練任務可自行下載數據集，通過使用CSI hpfs文件存儲，通過PVC方式掛載進容器中使用。

操作步驟

• 進入云容器引擎控制臺。

• 點擊左側【集群】進入集群列表。

• 點擊使用的集群名稱，進入集群。

• 點擊左側【工作負載】->【自定義資源】，選擇資源瀏覽器，找到kubeflow.org/v1/PyTorchJob ，選擇命名空間，點擊新增。

• 在創建yaml中，填入以下信息后點擊【創建】。

GPU模板

apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: pytorch-sample-gpu-01
  namespace: default
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: {image_repo}/icce/kubeflow-examples-pytorch-dist-mnist:multi
              #image: {image_repo}/icce/kubeflow-examples-pytorch-dist-mnist:gpu-cuda12.1
              imagePullPolicy: IfNotPresent
              command:
                - "python3"
                - "/opt/mnist/src/mnist.py"
                - "--epochs=10"
                - "--backend=nccl"
              env:
              - name: PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION
                value: python
              resources:
                limits:
                  nvidia.com/gpu: 1
    Worker:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: {image_repo}/icce/kubeflow-examples-pytorch-dist-mnist:multi
              imagePullPolicy: IfNotPresent
              command:
                - "python3"
                - "/opt/mnist/src/mnist.py"
                - "--epochs=10"
                - "--backend=nccl"
              env:
              - name: PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION
                value: python
              resources:
                limits:
                  nvidia.com/gpu: 1

NPU模板

apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: pytorch-sample-npu-01
  namespace: default
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: {image_repo}/icce/kubeflow-examples-pytorch-dist-mnist:multi
              imagePullPolicy: IfNotPresent
              command:
                - "bash"
                - "-c"
              args: ["source /usr/local/Ascend/ascend-toolkit/set_env.sh && python3 /opt/mnist/src/mnist.py --epochs=10 --backend=hccl"]
              env:
              - name: PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION
                value: python
              resources:
                limits:
                  huawei.com/Ascend910: 1
                requests:
                  huawei.com/Ascend910: 1
    Worker:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: {image_repo}/icce/kubeflow-examples-pytorch-dist-mnist:multi
              imagePullPolicy: IfNotPresent
              command:
                - "bash"
                - "-c"
              args: ["source /usr/local/Ascend/ascend-toolkit/set_env.sh && python3 /opt/mnist/src/mnist.py --epochs=10 --backend=hccl"]
              env:
              - name: PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION
                value: python
              resources:
                limits:
                  huawei.com/Ascend910: 1
                requests:
                  huawei.com/Ascend910: 1

• 查看運行狀態：點擊左側【工作負載】->【容器組】，找到任務名為前綴的容器，點擊名稱，查看日志/監控等信息是否符合預期。