kubectl を利用したモデル学習

VPC環境で利用できます。

ML expert Platformは、モデル学習のための様々なフレームワークの Jobを提供します。
kubectlを利用したモデル学習では、最も一般的な形態である単一ノード学習(Job)、分散ノード学習(PytorchJob、TrainJob)について説明します。

Job vs PytorchJob

単一ノード学習の場合、Jobの使用をお勧めします。
PytorchJobを使用する場合、分散ノード学習のための機能により Masterや Workerの形式でリリース・管理されるため、リソースを無駄に消費する可能性があります。

単一ノード学習(Job)を実行する

学習のための Kubernetes Jobの定義は、以下のユースケースのように作成できます。

注意

高性能ストレージをマウントして使用する場合、学習イメージ内の UIDおよび GIDは500に設定する必要があります。
Podの securityContextに fsGroupを設定しないでください。fsGroupが設定されると、Kubernetesがマウントされたボリューム内のすべてのファイルの所有権を再帰的に変更します。そのため、大容量のデータが保存された高性能ストレージの場合は、Podの初期化時間が大幅に長くなる可能性があります。

# job.yaml
apiVersion: batch/v1
kind: Job
metadata:
  name: mnist
  namespace: p-{projectName}
spec:
  backoffLimit: 1
  template:
    metadata:
      annotations:
        sidecar.istio.io/inject: "false"
    spec:
      restartPolicy: Never
      nodeSelector:
        mlx.navercorp.com/zone: { 提供された GPU Zone名 } # GPU Resourcesで確認できる Zone名
      containers:
        - name: main
          image: { 学習イメージ (例: example.com/mnist:latest ) } # NVIDIA Baseイメージに基づいて作成された学習コード
          imagePullPolicy: Always
          resources:
            limits:
              memory: "8Gi"
              cpu: "4"
              nvidia.com/gpu: "1"
          command: ["python"]
          args:
            - /opt/mnist/src/mnist.py
            - --checkpoint_path
            - /opt/mnist/checkpoints/mnist.pt
            - --log_path
            - /opt/mnist/log
            - --data_path
            - /opt/mnist/data
            - --download_data

kubectl apply -f job.yaml
batch/mnist created

外部コンテナレジストリを使用する

コンテナレジストリに Secret情報が必要な場合は、Container Secretを作成するを参照し、作成できます。
作成された Secretは、以下のように使用できます。

...
spec:
    imagePullSecrets:
    - name: my-harbor-secret  # 事前に作成した Docker Credential Secretの名前
...

既存の Volumeを使用する

Volumesを通じて作成された Volumeは、以下のように使用できます。

...
spec:
    containers:
    - name: main
      ...
      volumeMounts:
      - mountPath: /data
        name: mnist-data # 下部の spec.volumesに記載した名前
    volumes:
    - name: mnist-data
      persistentVolumeClaim:
        claimName: mnist-data # Volumesを通じて作成された PVC名
...

Jobのライフサイクル

Jobが終了した場合、一定時間コンテナのログおよびステータスを保持するため、削除されずにリストに残ります。リストに残せる Jobの最大数には制限があるため、TTLを使用してライフサイクルを管理します。詳細は、Kubernetes Job APIドキュメントをご参照ください。

Job作成時に Time To Live(TTL)を任意に設定できます。TTLは Jobが完了(成功・失敗)した後に有効となり、設定した TTLが経過すると Jobおよび Jobに属する Podがすべて自動的に削除されます。

以下は、3週間の TTLを適用した例です。

apiVersion: batch/v1
kind: Job
metadata:
  name: mnist
  namespace: p-{projectName}
spec:
  ttlSecondsAfterFinished: 1814400 # TTLを直接設定するフィールド(単位: sec)

分散ノード学習(PytorchJob)を実行する

PyTorchJobを利用すると、以下のようなメリットがあります。

作成されたコンテナ specを基に Masterおよび Worker Podを適切に作成します。
PyTorchの分散学習に一般的に必要な環境変数(例: WORLD_SIZE, RANK、MASTER_ADDRなど)を自動的に設定します。
学習に使われる Podが相互通信できるよう、K8s Serviceを自動的に作成します。Masterは <pytorch-job-name>-master-0、Workerは <pytorch-job-name>-worker-<idx>などの名前でアクセスできます。
必要に応じて、Elastic Policy を使用するを通じて torchrunで使用するパラメータを環境変数として作成します(--nnodes、--nproc-per-node、--rdzv-endpointなど)。

PytorchJobを作成する

Pytorchで分散学習する際は、torchrun(Elastic Launch)の使用をお勧めします。また、torchrunを使用する際は、マスター Podを別途指定しません。Torch Elasticでは、マスターノードのロールを果たす RANK=0の Podが、実行中に変更されることがあります。

spec.elasticPolicy - torchrun関連の設定です。ここで指定した設定は、環境変数として渡されます。詳細は、Elastic Policy を使用するをご参照ください。
spec.runPolicy - PyTorchJobの実行および終了後の処理に関するパラメータを指定できます。詳細は、Run Policy を使用するをご参照ください。
spec.pytorchReplicaSpecs.Worker - 分散学習を実行する Worker Podに関する設定です。

注意

学習を実行するコンテナの名前(例: spec.pytorchReplicaSpecs.Worker.template.spec.containers[*].name )は、必ず pytorch である必要があります。
円滑な分散学習のために、Istio Sidecarが Injectされないように spec.pytorchReplicaSpecs.Worker.template.metadata.annotationsに sidecar.istio.io/inject: "false"が自動的に設定されます。当該 annotationが設定されない場合、RuntimeError: Connection reset by peerなどのノード間通信にエラーが発生します。
高性能ストレージをマウントして使用する場合、学習イメージ内の UIDおよび GIDは500に設定する必要があります。
Podの securityContextに fsGroupを設定しないでください。fsGroupが設定されると、Kubernetesがマウントされたボリューム内のすべてのファイルの所有権を再帰的に変更します。そのため、大容量のデータが保存された高性能ストレージの場合は、Podの初期化時間が大幅に長くなる可能性があります。

学習のための PytorchJobの定義は、以下のユースケースのように作成できます。

# pytorchjob.yaml
apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
    name: pytorch-mnist-dist-nccl
    namespace: p-{ projectName } # プロジェクトに該当する Kubernetes Namespace名
spec:
    pytorchReplicaSpecs:
        Worker:
            replicas: 2
            restartPolicy: OnFailure
            template:
                metadata:
                    annotations:
                        sidecar.istio.io/inject: "false"  # 自動的に設定される(省略可能)
                spec:
                    nodeSelector:
                        mlx.navercorp.com/zone: { 提供された GPU Zone名 } # GPU Resourcesで確認できる Zone名
                    containers:
                    - name: pytorch  # PyTorchJobの container名を必ず pytorchに設定
                      image: examples.com/pytorch-mnist-dist:23.03-py3
                      imagePullPolicy: Always
                      command: ["bash", "-c"]
                      args:
                      - >
                        torchrun --nnodes ${PET_NNODES} --nproc_per_node ${PET_NPROC_PER_NODE} --rdzv_id ${PET_RDZV_ID} --rdzv_backend ${PET_RDZV_BACKEND} --rdzv_endpoint ${PET_RDZV_ENDPOINT}
                        /opt/mnist/src/mnist.py --checkpoint_path /data/checkpoints/mnist.pt --log_path /data/logs --data_path /data/dataset
                      env:
                      - name: NCCL_DEBUG
                        value: INFO
                      - name: GLOO_SOCKET_FAMILY
                        value: AF_INET
                      securityContext:  # Infinibandの使用のために securityContextが必要です。
                        capabilities:
                          add: ["IPC_LOCK"]
                      # shared memory
                      volumeMounts:
                      - mountPath: /dev/shm
                        name: shared-memory
                    volumes:
                    - emptyDir:
                        medium: Memory
                      name: shared-memory

kubectl apply -f pytorchjob.yaml
pytorchjob.kubeflow.org/pytorch-elastic-mnist-nccl created

Elastic Policy を使用する

torchrun を使用するために、elasticPolicy を作成します。

...
spec:
    ...
    elasticPolicy:
        rdzvId: mnist
        rdzvBackend: c10d
        minReplicas: 2
        maxReplicas: 2
        nProcPerNode: 8
    ...

elasticPolicy 各フィールド値を基に PyTorchJob に使用される環境変数が設定されます。環境変数は、 torchrun のパラメータを設定する際に使用できます。torchrun で使用するパラメータなどに関する詳細は、公式ドキュメントをご参照ください。

`elasticPolicy` フィールド	対応環境変数	関係 `torchrun` パラメータ	説明
`rdzvId`	`PET_RDZV_ID`	`--rdzv-id`	ランデブー用 Job ID
`rdzvBackend`	`PET_RDZV_BACKEND`	`--rdzv-backend`	ランデブーバックエンド(例: c10d)
`minReplicas`, `maxReplicas`	`PET_NNODES`	`--nnodes`	ノード数
`nProcPerNode`	`PET_NPROC_PER_NODE`	`--nproc-per-node`	ノードごとの GPU数
`maxRestarts`	`PET_MAX_RESTARTS`	`--max-restarts`	最大再起動回数

Run Policy を使用する

torchrunを使用するために runPolicyを作成します。

...
spec:
  runPolicy:
    cleanPodPolicy: None
    ttlSecondsAfterFinished: 1814400 # TTLを直接設定するフィールド(単位: sec)
...

spec.runPolicyには PyTorchJob の実行およびクリーンアップに関するパラメータを指定できます。指定しない場合、デフォルト値を使用します。spec.runPolicy下位に指定できるパラメータは、次の通りです。(参考資料: Kubeflow Trainer API Reference v1.9)

cleanPodPolicy - PytorchJob が完了した後、Podをどのようにクリーンアップするかを決めます。
- デフォルト値: None
- None: Job完了後も Podを削除しないため、後ほどログを確認する際に役立ちます。
- All: Job完了後にすべての Podを削除します。
- Running は Job完了後に実行中の Podを削除します。特別な場合を除き、使用しません。
ttlSecondsAfterFinished - Job完了後、何秒後に Jobを削除するかを決めます。
activeDeadlineSeconds - Jobの最大実行時間です。指定した時間が経過した後は、失敗として処理されます。設定されていない場合、jobの実行時間に制限はありません。
backoffLimit - Jobが失敗した場合の最大再実行回数です。

Infiniband を使用する

参考

別途の RDMAリソースの設定(request/limit)は必要ありません
GPU Zoneの情報は、使用可能な GPU Zone情報の照会をご参照ください。

InfiniBandネットワークで接続されたノードで分散学習を実行すると、ノード間通信を高速化できます。
前述のユースケースを InfiniBand環境に合わせて調整するには追加すべき仕様がいくつかありますが、それらは次のように要約できます。

InfiniBandネットワークが構成された区域の名前を Annotationで明示します(例: mlx.navercorp.com/zone=ai-infra)
InfiniBandを使用するために IPC_LOCK capabilityを securityContext に追加します。
分散学習を支援するための Shared Memoryを volumesに設定します。

Infinibandを使用するためには、以下のように設定できます。

...
metadata:
  annotations:
    mlx.navercorp.com/zone="ai-infra"
...
spec:
    ...
    pytorchReplicaSpecs:
        Worker:
            template:
                containers:
                - name: pytorch
                  securityContext:  # Infinibandの使用のために securityContextが必要です。
                    capabilities:
                      add: ["IPC_LOCK"]
                  # shared memory
                  volumeMounts:
                  - mountPath: /dev/shm
                    name: shared-memory
                volumes:
                - emptyDir:
                    medium: Memory
                  name: shared-memory
...

PytorchJobデバッグ

PytorchJob使用時に問題が発生し、デバッグが必要な場合は、以下のように環境変数値を設定することで、必要な情報をログに出力できます。

NCCL_DEBUG: NCCL関連デバッグ
TORCH_DISTRIBUTED_DEBUG、TORCH_CPP_LOG_LEVEL: 分散学習に対するデバッグ; 詳細は、PyTorch公式ドキュメントをご参照ください。

デバッグのために、以下のように使用できます。

...
spec:
    ...
    pytorchReplicaSpecs:
        Worker:
            template:
                containers:
                - name: pytorch
                  ...
                  env:
                  - name: NCCL_DEBUG
                    value: "INFO"
                  - name: TORCH_DISTRIBUTED_DEBUG
                    value: "DETAIL"
                  - name: TORCH_CPP_LOG_LEVEL
                    value: "INFO"
                  ...

外部コンテナレジストリを使用する

...
spec:
    ...
    pytorchReplicaSpecs:
        Worker:
            template:
                containers:
                - name: pytorch
                  imagePullSecrets:
                  - name: my-harbor-secret  # 事前に作成した Docker Credential Secretの名前
...

既存の Volumeを使用する

Volumesを通じて作成された Volumeは、以下のように使用できます。

...
spec:
    ...
    pytorchReplicaSpecs:
        Worker:
            template:
                containers:
                - name: pytorch
                  volumeMounts:
                  - mountPath: /data
                    name: mnist-data # 下部の spec.volumesに記載した名前
                volumes:
                - name: mnist-data
                  persistentVolumeClaim:
                    claimName: mnist-data # Volumesを通じて作成された PVC名
...

PytorchJobのステータスを確認

kubectl get、kubectl describe を利用して PyTorchJob のステータスを確認できます。

kubectl get pytorchjob pytorch-elastic-mnist-nccl
NAME                         STATE     AGE
pytorch-elastic-mnist-nccl   Running   12s

kubectl describe pytorchjob pytorch-elastic-mnist-nccl

Status:
  Completion Time:  2024-11-22T09:16:58Z
  Conditions:
    Last Transition Time:  2024-11-22T09:15:43Z
    Last Update Time:      2024-11-22T09:15:43Z
    Message:               PyTorchJob pytorch-elastic-mnist-nccl is created.
    Reason:                PyTorchJobCreated
    Status:                True
    Type:                  Created
    Last Transition Time:  2024-11-22T09:15:48Z
    Last Update Time:      2024-11-22T09:15:48Z
    Message:               PyTorchJob nb12706/pytorch-elastic-mnist-nccl is running.
    Reason:                PyTorchJobRunning
    Status:                False
    Type:                  Running
    Last Transition Time:  2024-11-22T09:16:58Z
    Last Update Time:      2024-11-22T09:16:58Z
    Message:               PyTorchJob nb12706/pytorch-elastic-mnist-nccl successfully completed.
    Reason:                PyTorchJobSucceeded
    Status:                True
    Type:                  Succeeded
  Last Reconcile Time:     2024-11-22T09:15:43Z
  Replica Statuses:
    Worker:
      Selector:   training.kubeflow.org/job-name=pytorch-elastic-mnist-nccl,training.kubeflow.org/operator-name=pytorchjob-controller,training.kubeflow.org/replica-type=worker
      Succeeded:  2
  Start Time:     2024-11-22T09:15:44Z
Events:
  Type    Reason                    Age                  From                   Message
  ----    ------                    ----                 ----                   -------
  Normal  SuccessfulCreatePod       4m21s                pytorchjob-controller  Created pod: pytorch-elastic-mnist-nccl-worker-0
  Normal  SuccessfulCreatePod       4m21s                pytorchjob-controller  Created pod: pytorch-elastic-mnist-nccl-worker-1
  Normal  SuccessfulCreateService   4m21s                pytorchjob-controller  Created service: pytorch-elastic-mnist-nccl-worker-0
  Normal  SuccessfulCreateService   4m21s                pytorchjob-controller  Created service: pytorch-elastic-mnist-nccl-worker-1
  Normal  ExitedWithCode            3m7s (x3 over 3m8s)  pytorchjob-controller  Pod: nb12706.pytorch-elastic-mnist-nccl-worker-1 exited with code 0
  Normal  ExitedWithCode            3m7s (x2 over 3m8s)  pytorchjob-controller  Pod: nb12706.pytorch-elastic-mnist-nccl-worker-0 exited with code 0
  Normal  PyTorchJobSucceeded       3m7s                 pytorchjob-controller  PyTorchJob nb12706/pytorch-elastic-mnist-nccl successfully completed.
  Normal  JobTerminated             3m6s (x4 over 3m7s)  pytorchjob-controller  Job has been terminated. Deleting PodGroup
  Normal  SuccessfulDeletePodGroup  3m6s (x4 over 3m7s)  pytorchjob-controller  Deleted PodGroup: pytorch-elastic-mnist-nccl

学習 Podが正しく作成されない場合、以下のように Eventsを通じて原因を把握できます。

kubectl describe pytorchjob pytorch-elastic-mnist-nccl

...
Events:
  Type     Reason           Age                 From                   Message
  ----     ------           ----                ----                   -------
  Warning  FailedCreatePod  47m (x3 over 103m)  pytorchjob-controller  Error creating: Pods "job-worker-1" is forbidden: exceeded quota: normal-quota, requested: requests.nvidia.com/gpu=1, used: requests.nvidia.com/gpu=2, limited: requests.nvidia.com/gpu=2

分散学習時の注意点

MLXPで分散学習を行う際の注意点をご案内します。

NVIDIA Driverと互換性のある CUDAバージョンを使用

CUDA Runtimeバージョンは、ノードにインストールされた NVIDIA Driverと互換性がある必要があります。CUDA Runtimeが要求する最小 Driverバージョンを満たしていない場合、CUDAの初期化エラーまたはランタイムエラーが発生し、学習タスクが正常に実行されないことがあります。

IPv6が使用できない場合は GLOO_SOCKET_FAMILY=INETを設定

GLOOはノード間のステータス確認および Rank管理のために、頻繁に TCP接続を作成します。IPv6を使用できない環境で大規模学習を実行する場合、接続を試行するたびに IPv6接続が先に試行された後に失敗し、その後 IPv4で再試行されます。このような二重の試行が繰り返されると、ファイルディスクリプタの枯渇などのリソース問題が発生し、結果として学習が中断される可能性があります。

このような問題を防ぐために GLOO_SOCKET_FAMILY=AF_INET 環境変数を設定し、IPv4のみ使用するように構成することをお勧めします。

GLOOで IPv4を強制するには、以下のように設定できます。

...
spec:
    ...
    pytorchReplicaSpecs:
        Worker:
            template:
                containers:
                 - name: pytorch
                   ... 
                   env:
                   - name: GLOO_SOCKET_FAMILY
                     value: AF_INET 
...

PyTorchJob名の長さ制限(推奨)

PyTorchJob名は Kubernetesリソース(Pod、Serviceなど)作成時に接頭語として使用されるため、50文字以内で設定することをお勧めします。50文字を超えると、派生リソース名が長さ制限を超えて作成に失敗したり、worker indexパースに問題が発生したりします。

Shared Memory設定時に memory limitの設定を削除

Shared Memory(/dev/shm)は容量制限を設けないか、十分に大きなサイズに設定することをお勧めします。/dev/shmが足りないと、DataLoaderのマルチプロセス workerが異常終了する可能性があり、NCCLの通信性能低下につながる恐れがあります。

Memory limitを削除した Shared Memory設定の例は、次の通りです。

...
spec:
    ...
    pytorchReplicaSpecs:
        Worker:
            template:
                containers:
                - name: pytorch
                  # shared memory
                  volumeMounts:
                  - mountPath: /dev/shm
                    name: shared-memory
                volumes:
                - emptyDir:
                    medium: Memory
                  name: shared-memory

分散ノード学習(TrainJob)を実行する

TrainJobは、Kubeflow Trainer v2が提供する学習リソースです。学習 Podの共通明細を納めた TrainingRuntimeを参照する方式であるため、PytorchJobのように毎回 Podの明細全体を作成する代わりに、学習ごとに異なる値のみを指定して実行できます。
TrainJobで学習を実行するには、次の2つの手順が必要です。

学習 Podの共通明細を納めた TrainingRuntimeを作成します。
作成した TrainingRuntimeを参照する TrainJobを作成します。

参考

TrainJobは、固定されたノード数を前提として動作します。実行中にノード数が変更される学習が必要な場合は、PytorchJobの Elastic Policy を使用するをご参照ください。
分散学習を行う際の注意事項については、分散学習時の注意事項も併せてご参照ください。
SDKの使用などの詳細については、Kubeflow Trainerの公式ドキュメントをご参照ください。

TrainingRuntimeを作成する

TrainingRuntimeは、スケジューラ、ボリューム、ベースイメージなど、学習 Podの共通明細を納めるリソースです。プロジェクトに該当する Namespace内に直接作成し、同じ Namespaceの TrainJobからのみ参照できます。
TrainingRuntimeを作成する際は、以下の項目を必ず含める必要があります。

spec.podGroupPolicy.volcanoと schedulerName: volcano - Gang Schedulingに必要です。
sidecar.istio.io/inject: "false" Annotation - Istio Sidecarが設定されないようにします。当該 Annotationが設定されない場合、RuntimeError: Connection reset by peerなどのノード間通信にエラーが発生します。
trainer.kubeflow.org/trainjob-ancestor-step: trainer Label - 設定しない場合、TrainJobの spec.trainerに作成した値(image、commandなど)は学習 Jobに適用されません。
nodeSelectorの mlx.navercorp.com/zone - 学習を実行する GPU Zoneを指定します。TrainingRuntime 1つで複数のゾーンを使用する必要がある場合は、TrainJobで GPU Zoneを指定するをご参照ください。

注意

高性能ストレージをマウントして使用する場合、学習イメージ内の UIDおよび GIDは500に設定する必要があります。
Podの securityContextに fsGroupを設定しないでください。fsGroupが設定されると、Kubernetesがマウントされたボリューム内のすべてのファイルの所有権を再帰的に変更します。そのため、大容量のデータが保存された高性能ストレージの場合は、Podの初期化時間が大幅に長くなる可能性があります。

PyTorch分散学習のための TrainingRuntime 明細は、以下のユースケースのように作成できます。torchrun 実行に必要な環境変数が設定されるよう、spec.mlPolicy.torchを明示します。

# trainingruntime.yaml
apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainingRuntime
metadata:
  name: torch-distributed
  namespace: p-{ projectName } # プロジェクトに該当する Kubernetes Namespace名
spec:
  mlPolicy:
    numNodes: 2 # デフォルトノード数(TrainJobの spec.trainer.numNodesで再定義可能)
    torch: {}   # torchrunに必要な PET_* 環境変数を設定します。
  podGroupPolicy:
    volcano: {}
  template:
    spec:
      replicatedJobs:
      - name: node
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: trainer
          spec:
            template:
              metadata:
                annotations:
                  sidecar.istio.io/inject: "false"
              spec:
                schedulerName: volcano
                restartPolicy: OnFailure
                nodeSelector:
                  mlx.navercorp.com/zone: { 提供された GPU Zone名 } # GPU Resourcesで確認できる Zone名
                containers:
                - name: node
                  image: { 学習イメージ (例: example.com/pytorch-mnist-dist:latest ) } # ベースイメージ (TrainJobの spec.trainer.imageで再定義可能)
                  imagePullPolicy: Always
                  securityContext:  # Infinibandの使用のために securityContextが必要です。
                    capabilities:
                      add: ["IPC_LOCK"]
                  # shared memory
                  volumeMounts:
                  - mountPath: /dev/shm
                    name: shared-memory
                volumes:
                - emptyDir:
                    medium: Memory
                  name: shared-memory

kubectl apply -f trainingruntime.yaml
trainingruntime.trainer.kubeflow.org/torch-distributed created

InfiniBandを使用する場合も、上記のユースケースのように IPC_LOCK capabilityと Shared Memoryを TrainingRuntimeに設定し、InfiniBandネットワークが構成されている GPU Zoneを nodeSelectorに指定します。詳細は、Infiniband を使用するをご参照ください。

TrainJobを作成する

作成した TrainingRuntimeを spec.runtimeRefで参照し、TrainJobを作成します。学習明細は、以下のユースケースのように作成できます。

# trainjob.yaml
apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
metadata:
  name: pytorch-mnist-dist-nccl
  namespace: p-{ projectName } # プロジェクトに該当する Kubernetes Namespace名
spec:
  runtimeRef:
    apiGroup: trainer.kubeflow.org
    kind: TrainingRuntime
    name: torch-distributed # 予め作成しておいた TrainingRuntimeの名前
  trainer:
    numNodes: 2
    numProcPerNode: 1 # ノード当たりのプロセス数(ノードに割り当てられた GPUの数と同じ値に指定)
    command: ["bash", "-c"]
    args:
    - >
      torchrun --nnodes ${PET_NNODES} --nproc_per_node ${PET_NPROC_PER_NODE}
      --rdzv_id mnist --rdzv_backend c10d --rdzv_endpoint ${PET_MASTER_ADDR}:${PET_MASTER_PORT}
      /opt/mnist/src/mnist.py --checkpoint_path /data/checkpoints/mnist.pt --log_path /data/logs --data_path /data/dataset
    env:
    - name: NCCL_DEBUG
      value: INFO
    - name: GLOO_SOCKET_FAMILY
      value: AF_INET
    resourcesPerNode:
      requests:
        memory: "8Gi"
        cpu: "4"
        nvidia.com/gpu: "1"
      limits:
        memory: "8Gi"
        cpu: "4"
        nvidia.com/gpu: "1"

kubectl apply -f trainjob.yaml
trainjob.trainer.kubeflow.org/pytorch-mnist-dist-nccl created

spec.trainer.numNodes - 学習に使用するノード数です。PytorchJobの elasticPolicy.minReplicas、maxReplicas(固定値)に対応します。
spec.trainer.numProcPerNode - ノード当たりのプロセス数です。通常、ノードに割り当てられた GPUの数と同じ数に設定します。PytorchJobの elasticPolicy.nProcPerNodeに対応しており、autoはサポートされておらず、整数のみ指定できます。
spec.trainer.image - 学習に使用するイメージです。指定しない場合は、TrainingRuntimeのベースイメージがそのまま使用されます。
spec.trainer.resourcesPerNode - 1つのノード(Pod)に割り当てるリソースです。PytorchJobの resourcesに対応します。

注意

spec.trainer.resourcesPerNodeを指定すると、TrainingRuntimeに作成したコンテナの resourcesが項目単位でマージされることなく、全体が再定義されます。必要なリソースをすべて明記してください。
同じ Namespaceで PytorchJobと TrainJobが同じ metadata.nameを使用すると、PodGroupが共有されることになり、Gang Schedulingが正常に動作しなくなることがあります。TrainJobを作成する前に、同じ名前の PytorchJobが既に存在しないかご確認ください。

TrainJobで GPU Zoneを指定する

TrainingRuntimeに GPU Zoneを固定すると、使用する Zoneごとに TrainingRuntimeを作成する必要があります。TrainingRuntime 1つで複数の Zoneを使用するには、TrainJobの spec.runtimePatchesで nodeSelectorを指定できます。

...
spec:
  runtimePatches:
  - manager: { パッチの所有者名 (例: my-team/trainjob ) } # 空にすることはできません。
    trainingRuntimeSpec:
      template:
        spec:
          replicatedJobs:
          - name: node # TrainingRuntimeに作成した replicatedJobの名前と同様に指定
            template:
              spec:
                template:
                  spec:
                    nodeSelector:
                      mlx.navercorp.com/zone: { 提供された GPU Zone名 } # GPU Resourcesで確認できる Zone名
...

manager - パッチの所有者を識別する値であり、空にすることはできません。Kubeflow SDKが使用する trainer.kubeflow.org/kubeflow-sdkと重複しない値を指定してください。
nodeSelector 以外に tolerations、serviceAccountName、volumesも同様の方法で指定できます。
image、command、args、resourcesは、spec.runtimePatchesでは再定義できません。当該値は、spec.trainerに作成してください。

torchrun環境変数の違い

PytorchJobの elasticPolicyはランデブーに必要な環境変数(PET_RDZV_*)を設定しますが、TrainJobは固定されたノード数を前提としているため、0番ノードのアドレスを設定する PET_MASTER_ADDR、PET_MASTER_PORTを使用します。したがって、PytorchJobで用いた torchrun コマンドをそのまま使用すると、学習に失敗することがあります。

`PytorchJob` 環境変数	`TrainJob` 環境変数	説明
`PET_RDZV_ID`	設定されない	`torchrun` コマンドに `--rdzv_id` 値を直接指定します。
`PET_RDZV_BACKEND`	設定されない	`torchrun` コマンドに `--rdzv_backend` 値を直接指定します。
`PET_RDZV_ENDPOINT`	`PET_MASTER_ADDR`, `PET_MASTER_PORT`	`${PET_MASTER_ADDR}:${PET_MASTER_PORT}` 形式でマージし、`--rdzv_endpoint`に使用します。
`PET_NNODES`	`PET_NNODES`	`spec.trainer.numNodes`として指定したノードの数です。
`PET_NPROC_PER_NODE`	`PET_NPROC_PER_NODE`	`spec.trainer.numProcPerNode`として指定したノード当たりのプロセス数です。
-	`PET_NODE_RANK`	Jobの Completion Indexに基づいて割り当てられるノード番号です。`c10d` ランデブーを使用する場合、ランデブーが Rankを決定するため、`--node_rank`を指定する必要はありません。

TrainJobのステータスを確認

kubectl get、kubectl describe を利用して TrainJob のステータスを確認できます。

kubectl get trainjob pytorch-mnist-dist-nccl

kubectl describe trainjob pytorch-mnist-dist-nccl

TrainJobは JobSetと Jobを経て学習 Podを作成します。学習ログは、作成された Podを確認した後に照会できます。

kubectl get pods
kubectl logs { 学習 Pod名 }