Model training using kubectl

Available in VPC

ML expert Platform provides jobs for multiple model training frameworks.
Learn how to train models using kubectl with the most common methods, including single-node training (Job) and distributed node training (PytorchJob, TrainJob).

Run single-node training (Job)

You can define a Kubernetes Job specification for training as shown in the following example.

# 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 name } # Zone name as shown in GPU Resources
      containers:
        - name: main
          image: { training image (e.g., example.com/mnist:latest ) } # Training code written using the NVIDIA Base image
          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

Using an external container registry

If the container registry requires Secret information, see Create a Container Secret to create one.
You can use the secret created as follows.

...
spec:
    imagePullSecrets:
    - name: my-harbor-secret # Name of the Docker Credential Secret created in advance
...

Using an existing volume

To use a Volume created through Volumes, you can configure it as follows:

...
spec:
    containers:
    - name: main
      ...
      volumeMounts:
      - mountPath: /data
        name: mnist-data # Name specified in spec.volumes below
    volumes:
    - name: mnist-data
      persistentVolumeClaim:
        claimName: mnist-data # Name of the PVC created through Volumes
...

Job lifecycle

When a Job completes, it remains in the list for a certain period of time to preserve container logs and status. Since there is a limit on the maximum number of Jobs that can remain in the list, you must manage the lifecycle using TTL. For more information, see Kubernetes Job API.

When creating a Job, you can set a custom time to live (TTL). When a Job finishes, the time to live (TTL) is enabled, whether the Job succeeds or fails. After the TTL passes, the Job and all Pods that belong to the Job are automatically deleted.

The following example applies a TTL of 3 weeks.

apiVersion: batch/v1
kind: Job
metadata:
  name: mnist
  namespace: p-{projectName}
spec:
  ttlSecondsAfterFinished: 1,814,400 # Field for setting TTL manually (unit: sec)

Run distributed node training (PyTorchJob)

Using PyTorchJob provides the following benefits:

• It creates Master and Worker Pods as appropriate based on the container specification you define.
• It automatically sets environment variables commonly required for PyTorch distributed training, such as WORLD_SIZE, RANK, and MASTER_ADDR.
• It automatically creates a K8s Service to facilitate communication between training Pods. You can access the Master as <pytorch-job-name>-master-0 and the Workers as <pytorch-job-name>-worker-<idx>.
• If needed, you can use Elastic Policy to create environment variables for torchrun arguments, such as --nnodes, --nproc-per-node, and --rdzv-endpoint.

Define a PyTorchJob

For distributed training with PyTorch, we recommend using torchrun (Elastic Launch). When using torchrun, do not specify a Master Pod explicitly. In Torch Elastic, the RANK=0 Pod that acts as the master node can change during Run.

• spec.elasticPolicy - Configures torchrun-related settings. Settings specified here are injected as environment variables. For details, see Use Elastic Policy.
• spec.runPolicy - Specifies parameters for PyTorchJob Run and post-completion handling. For details, see Use Run Policy.
• spec.pytorchReplicaSpecs.Worker - Configures the Worker Pod that performs distributed training.

You can define a PyTorchJob specification for training as shown in the following example.

# pytorchjob.yaml
apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
    name: pytorch-mnist-dist-nccl
    namespace: p-{ projectName } # Kubernetes Namespace name for the project
spec:
    pytorchReplicaSpecs:
        Worker:
            replicas: 2
            restartPolicy: OnFailure
            template:
                metadata:
                    annotations:
                        sidecar.istio.io/inject: "false" # Automatically set (can be omitted)
                spec:
                    nodeSelector:
                        mlx.navercorp.com/zone: { GPU Zone name } # Zone name as shown in GPU Resources
                    containers:
                    - name: pytorch # Must set the container name for PyTorchJob to 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:  # A securityContext is required to use InfiniBand.
                        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

Use Elastic Policy

Define elasticPolicy to use torchrun.

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

Environment variables used for PyTorchJob are set based on each field value in elasticPolicy. These environment variables can be used to set torchrun arguments. For more information about the arguments used by torchrun, see the official documentation.

Use Run Policy

Define runPolicy to use torchrun.

...
spec:
  runPolicy:
    cleanPodPolicy: None
    ttlSecondsAfterFinished: 1,814,400 # Field for setting TTL manually (unit: sec)
...

spec.runPolicy lets you specify parameters related to PyTorchJob Run and cleanup. If not specified, the default values are used. The parameters you can specify under spec.runPolicy include the following: (see: Kubeflow Trainer API Reference v1.9)

• cleanPodPolicy - Determines how to clean up Pods after the PytorchJob completes.
  • Default: None
  • None: Does not delete Pods after the Job completes, which helps you check logs later.
  • All: Deletes all Pods after the Job completes.
  • Running deletes running Pods after the Job completes. This is rarely used except in special cases.
• ttlSecondsAfterFinished - Determines how many seconds after Job completion the Job is deleted.
• activeDeadlineSeconds - Specifies the maximum Run time for the Job. After the specified time passes, the Job is marked as failed. If not set, there is no limit on the Job Run time.
• backoffLimit - Specifies the maximum number of retries when the Job fails.

Use InfiniBand

You can speed up communication between nodes by running distributed training on nodes connected through an InfiniBand network.
To adapt the example shown in the previous section for an InfiniBand environment, add the following specifications:

• Specify the name of the zone where the InfiniBand network is configured as an annotation (e.g., mlx.navercorp.com/zone=ai-infra)
• Enable InfiniBand by adding the IPC_LOCK capability to securityContext.
• Configure shared memory in volumes to support distributed training.

To use InfiniBand, you can configure it as follows:

...
metadata:
  annotations:
    mlx.navercorp.com/zone="ai-infra"
...
spec:
    ...
    pytorchReplicaSpecs:
        Worker:
            template:
                containers:
                - name: pytorch
                  securityContext: # A securityContext is required to use InfiniBand.
                    capabilities:
                      add: ["IPC_LOCK"]
                  # shared memory
                  volumeMounts:
                  - mountPath: /dev/shm
                    name: shared-memory
                volumes:
                - emptyDir:
                    medium: Memory
                  name: shared-memory
...

Debugging PyTorchJob

To debug issues in PytorchJob, you can set the following environment variables to log the required information:

• NCCL_DEBUG: Enables NCCL-related debugging.
• TORCH_DISTRIBUTED_DEBUG, TORCH_CPP_LOG_LEVEL: Enables distributed training debugging. For details, see the PyTorch official documentation.

For debugging, you can configure it as follows:

...
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"
                  ...

Using an external container registry

If the container registry requires Secret information, see Create a Container Secret to create one.
You can use the secret created as follows.

...
spec:
    ...
    pytorchReplicaSpecs:
        Worker:
            template:
                containers:
                - name: pytorch
                  imagePullSecrets:
                  - name: my-harbor-secret # Name of the Docker Credential Secret created in advance
...

Using an existing volume

To use a Volume created through Volumes, you can configure it as follows:

...
spec:
    ...
    pytorchReplicaSpecs:
        Worker:
            template:
                containers:
                - name: pytorch
                  volumeMounts:
                  - mountPath: /data
                    name: mnist-data # Name specified in spec.volumes below
                volumes:
                - name: mnist-data
                  persistentVolumeClaim:
                    claimName: mnist-data # Name of the PVC created through Volumes
...

Check PytorchJob status

You can check the status of a PyTorchJob using kubectl get and kubectl describe.

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

If a training Pod is not created properly, you can identify the cause through Events as follows:

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

Guidelines for distributed training

This section provides guidelines for distributed training in MLXP.

Use a CUDA version compatible with the NVIDIA driver

The CUDA Runtime version must be compatible with the NVIDIA driver installed on the node. Failing to meet the minimum driver version requirements may result in CUDA initialization or runtime errors, preventing the training task from running correctly.

Set GLOO_SOCKET_FAMILY=AF_INET when IPv6 is unavailable

GLOO frequently establishes TCP connections to check the status between nodes and manage ranks. When running large-scale training in an environment where IPv6 is unavailable, each connection attempt first tries and fails to connect over IPv6 before retrying over IPv4. If these repeated connection attempts continue, resource issues such as file descriptor exhaustion may occur and cause training to stop.

To prevent this issue, we recommend setting the GLOO_SOCKET_FAMILY=AF_INET environment variable to use only IPv4.

To force GLOO to use IPv4, you can configure it as follows:

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

Limit the PyTorchJob name length (recommended)

Since the PyTorchJob name is used as a prefix when creating Kubernetes resources, such as Pods and Services, we recommend keeping it to 50 characters or fewer. If the name exceeds 50 characters, derived resource names may also exceed the length limit, causing creation failures or issues with worker index parsing.

Remove the memory limit setting when configuring shared memory

We recommend not setting a memory limit for shared memory (/dev/shm) or setting a sufficiently large value. Insufficient /dev/shm can cause DataLoader multiprocessing workers to terminate abnormally and lead to a significant decrease in NCCL communication performance.

The following example shows a shared memory configuration with the memory limit removed:

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

Run distributed node training (PyTorchJob)

TrainJob is a training resource provided by Kubeflow Trainer v2. Since it refers to TrainingRuntime, which contains the common specification for training Pods, you can run training by specifying only the values that change for each run, instead of writing the entire Pod specification every time like with PytorchJob.
The following 2 steps are required to run training with TrainJob:

1. Create TrainingRuntime, which contains the common specification for training Pods.
2. Create TrainJob, which refers to the created TrainingRuntime.

Creating a TrainingRuntime

TrainingRuntime is a resource that contains the common specifications for training pods, such as the scheduler, volumes, and base image. You create it directly in the namespace corresponding to the project, and it can only be referenced by TrainJob in the same namespace.
The following items must be included when creating TrainingRuntime:

• spec.podGroupPolicy.volcano and schedulerName: volcano - Required for gang scheduling.
• sidecar.istio.io/inject: "false" Annotation - Prevents Istio sidecars from being injected. If this annotation is not set, communication errors between nodes, such as RuntimeError: Connection reset by peer, may occur.
• trainer.kubeflow.org/trainjob-ancestor-step: trainer Label - If not set, the values (image, command, etc.) written in spec.trainer of TrainJob will not be applied to the training job.
• nodeSelector's mlx.navercorp.com/zone - Specifies the GPU zone where training will run. To use multiple zones with a single TrainingRuntime, see Specifying GPU zones in a TrainJob.

You can define a specification for TrainingRuntime for PyTorch distributed training as shown in the following example: Specify spec.mlPolicy.torch so that environment variables required for torchrun Run are injected.

# trainingruntime.yaml
apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainingRuntime
metadata:
  name: torch-distributed
  namespace: p-{ projectName } # Kubernetes Namespace name for the project
spec:
  mlPolicy:
    numNodes: 2 # Default number of nodes (can be override by spec.trainer.numNodes in the TrainJob)
    torch: {}   # Injects the PET_* environment variables required for torchrun.
  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 name } # Zone name as shown in GPU Resources
                containers:
                - name: node
                  image: { training image (e.g., example.com/pytorch-mnist-dist:latest ) } # Base image (can be override by spec.trainer.image in the TrainJob)
                  imagePullPolicy: Always
                  securityContext: # A securityContext is required to use InfiniBand.
                    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

Even when using InfiniBand, set the IPC_LOCK capability and shared memory in TrainingRuntime as in the example above, and specify the GPU zone with the InfiniBand network in nodeSelector. For details, see Use Infiniband.

Creating a TrainJob

Refer to the TrainingRuntime you created as spec.runtimeRef to create TrainJob. You can define a PyTorchJob specification for training as shown in the following example.

# trainjob.yaml
apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
metadata:
  name: pytorch-mnist-dist-nccl
  namespace: p-{ projectName } # Kubernetes Namespace name for the project
spec:
  runtimeRef:
    apiGroup: trainer.kubeflow.org
    kind: TrainingRuntime
    name: torch-distributed # Name of the TrainingRuntime created earlier
  trainer:
    numNodes: 2
    numProcPerNode: 1 # Number of processes per node (specify same as the number of GPUs allocated to the node)
    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 - Number of nodes to use for training. PytorchJob corresponds to elasticPolicy.minReplicas and maxReplicas (fixed value) of .
• spec.trainer.numProcPerNode - Number of processes per node. Generally, this is set to the same number as the GPUs allocated to the node. PytorchJob Corresponds to elasticPolicy.nProcPerNode of , and auto is not supported; only integers can be specified.
• spec.trainer.image - Image to be used for training. If not specified, the base image of TrainingRuntime is used as is.
• spec.trainer.resourcesPerNode - Resource to be allocated per node (pod). PytorchJob Corresponds to resources of .

Specify GPU Zone in TrainJob

If you fix the GPU Zone in TrainingRuntime, you must create TrainingRuntime for each zone you use. To use multiple zones with 1 TrainingRuntime, you can specify nodeSelector as spec.runtimePatches in TrainJob.

...
spec:
  runtimePatches:
  - manager: { patch owner name (e.g., my-team/trainjob ) } # Cannot be empty.
    trainingRuntimeSpec:
      template:
        spec:
          replicatedJobs:
          - name: node # Must be the same as the replicatedJob name written in TrainingRuntime
            template:
              spec:
                template:
                  spec:
                    nodeSelector:
                      mlx.navercorp.com/zone: { GPU Zone name } # Zone name as shown in GPU Resources
...

• manager - This value identifies the owner of the patch and cannot be empty. Specify a value that does not overlap with the trainer.kubeflow.org/kubeflow-sdk used by the Kubeflow SDK.
• In addition to nodeSelector, tolerations, serviceAccountName, and volumes can be specified in the same way.
• image, command, args, and resources cannot be overridden with spec.runtimePatches. Write this value in spec.trainer.

Differences in torchrun environment variables

PytorchJob's elasticPolicy injects environment variables (PET_RDZV_*) required for rendezvous, but TrainJob uses PET_MASTER_ADDR and PET_MASTER_PORT, which inject the address of node 0, because it assumes a fixed number of nodes. Therefore, if you use the torchrun command used in PytorchJob as is, training may fail.

Check TrainJob status

You can check the status of a TrainJob using kubectl get and kubectl describe.

kubectl get trainjob pytorch-mnist-dist-nccl

kubectl describe trainjob pytorch-mnist-dist-nccl

TrainJob creates a training Pod through JobSet and Job. Training logs can be viewed after checking the created Pod.

kubectl get pods
kubectl logs { training Pod name }