Chapter 6. Running Kubeflow Trainer v2-based distributed training workloads

Procedure

1. Log in to the OpenShift Console.

2. Create a TrainingRuntime resource, as follows:

   1. In the Administrator perspective, click Home > Search.
   2. From the Project list, select your project.
   3. Click the Resources list, and in the search field, type TrainingRuntime.
   4. Select TrainingRuntime, and click Create TrainingRuntime. The Create TrainingRuntime page opens, with default YAML code automatically added.

   5. Replace the default YAML code with your custom runtime definition. The following example shows a custom TrainingRuntime with a custom container image and default resource configuration:

      apiVersion: trainer.kubeflow.org/v1alpha1
      kind: TrainingRuntime
      metadata:
        name: custom-torch-runtime
        namespace: my-project
      spec:
        mlPolicy:
          torch:
            numProcPerNode: auto
        template:
          spec:
            replicatedJobs:
              - name: Node
                template:
                  spec:
                    template:
                      spec:
                        containers:
                          - name: trainer
                            image: registry.redhat.io/rhoai/odh-training-cuda128-torch28-py312-rhel9:v3.0
                            env:
                              - name: MASTER_ADDR
                                value: "$(JOB_NAME)-node-0-0.$(JOB_NAME)"
                              - name: MASTER_PORT
                                value: "29500"
                              - name: MY_CUSTOM_ENV_VAR
                                value: "my-custom-value"
                            command:
                              - torchrun
                            resources:
                              requests:
                                cpu: "4"
                                memory: "16Gi"
                              limits:
                                cpu: "4"
                                memory: "16Gi"

   6. Click Create.

Procedure

1. Log in to the OpenShift Console.

2. Create a TrainJob resource, as follows:

   1. In the Administrator perspective, click Home Search.
   2. From the Project list, select your project.
   3. Click the Resources list, and in the search field, type Job.
3. Select TrainJob, and click Create TrainJob. The Create TrainJob page opens, with default YAML code automatically added.

4. Update the metadata to replace the name and namespace values with the values for your environment, as shown in the following example:

   metadata:
     name: pytorch-multi-node-job
     namespace: test-namespace

5. Configure the runtime reference to specify which ClusterTrainingRuntime to use, as shown in the following example:

   spec:
     runtimeRef:
       name: torch-distributed
       kind: ClusterTrainingRuntime

6. Configure the trainer section with your ClusterTrainingRuntime name, command, and the number of distributed training nodes, as shown in the following example:

   spec:
     runtimeRef:
       name: torch-distributed
       kind: ClusterTrainingRuntime
     trainer:
       command: ["python", "/workspace/scripts/train.py"]
       numNodes: 3

   where:

   numNodes

   Specifies the total number of distributed training nodes. This replaces the separate Master and Worker replica specifications used in Training Operator v1. For example, a v1 job with 1 Master and 2 Workers is equivalent to numNodes: 3 in v2.

7. To use a ConfigMap resource to provide the training script for the TrainJob pods, add the ConfigMap volume mount information, as shown in the following example:

   spec:
     runtimeRef:
       name: torch-distributed
       kind: ClusterTrainingRuntime
     trainer:
       command: ["python", "/workspace/scripts/train.py"]
       numNodes: 3
     podTemplateOverrides:
       - targetJobs:
           - name: node
         spec:
           volumes:
             - name: training-script-volume
               configMap:
                 name: training-script-configmap
           containers:
             - name: node
               volumeMounts:
                 - name: training-script-volume
                   mountPath: /workspace/scripts

8. Add the appropriate resource constraints for your environment, as shown in the following example:

   spec:
     runtimeRef:
       name: torch-distributed
       kind: ClusterTrainingRuntime
     trainer:
       command: ["python", "/workspace/scripts/train.py"]
       numNodes: 3
       resourcesPerNode:
         requests:
           cpu: "4"
           memory: "8Gi"
           nvidia.com/gpu: 2
         limits:
           cpu: "4"
           memory: "8Gi"
           nvidia.com/gpu: 2
     podTemplateOverrides:
       - targetJobs:
           - name: node
         spec:
           volumes:
             - name: training-script-volume
               configMap:
                 name: training-script-configmap
           containers:
             - name: node
               volumeMounts:
                 - name: training-script-volume
                   mountPath: /workspace/scripts

   where:

   resourcesPerNode

   applies the same resource requests and limits to every training node. This replaces the per-replica resource configuration used in Training Operator v1, where you had to specify resources separately for the Master and each Worker.

9. Click Create.

Verification

1. In the OpenShift Console, open the Administrator perspective.

2. From the Project list, select your project.

   1. Click Home > Search > TrainJob and verify that the job was created.
   2. Click Workloads > Pods and verify that the requested training node pods are running.
3. (Optional) If you enabled progress tracking, navigate to the OpenShift AI dashboard, click Model training, and verify that real-time progress metrics are displayed.

Procedure

1. Log in to the OpenShift CLI (oc):

   oc login --token=<token> --server=<server>

2. Create a file named train.py and populate it with your training script, as follows:

   cat <<EOF > train.py
   <paste your content here>
   EOF

3. Create a ConfigMap resource to store the training script, as follows:

   oc create configmap training-script-configmap
   --from-file=train.py -n <your-namespace>

   Replace <your-namespace> with the name of your project.

4. Create a file named trainjob.yaml to define the distributed training job setup, as follows:

   cat <<'EOF' > trainjob.yaml
   apiVersion: trainer.kubeflow.org/v1alpha1
   kind: TrainJob
   metadata:
     name: pytorch-multi-node-job
     namespace: <your-namespace>
   spec:
     runtimeRef:
       name: torch-distributed
       kind: ClusterTrainingRuntime
     trainer:
       command: ["python", "/workspace/scripts/train.py"]
       numNodes: 3
       resourcesPerNode:
         requests:
           cpu: "4"
           memory: "8Gi"
           nvidia.com/gpu: 2
         limits:
           cpu: "4"
           memory: "8Gi"
           nvidia.com/gpu: 2
     podTemplateOverrides:
       - targetJobs:
           - name: node
         spec:
           volumes:
             - name: training-script-volume
               configMap:
                 name: training-script-configmap
           containers:
             - name: node
               volumeMounts:
                 - name: training-script-volume
                   mountPath: /workspace/scripts
   EOF

   Replace <your-namespace> with the name of your project. Update the numNodes, resource constraints, and image as needed for your environment.

5. Create the distributed training job, as follows:

   oc apply -f trainjob.yaml

6. Monitor the running distributed training job, as follows:

   oc get trainjob -n <your-namespace>

   Replace <your-namespace> with the name of your project.

7. Check the pod status, as follows:

   oc logs <pod-name> -n <your-namespace>

8. Replace <pod-name> with the name of the pod and <your-namespace> with the name of your project.

Prerequisites

1. You have an active TrainJob currently running on the OpenShift cluster.
2. You have configured JIT (Just-In-Time) checkpointing in your training script.
3. You have confirmed that your training script is configured to handle the SIGTERM signal.
4. You have the OpenShift CLI (oc) installed and are logged in with a user that has patch permissions for the TrainJob resource in the target namespace.

Prerequisites

1. You have a TrainJob that is currently in a suspended state (spec.suspend: true).
2. You have a Persistent Volume Claim (PVC) or S3-compatible storage containing the latest saved checkpoint files from the previous session.
3. Your training script is configured to automatically detect and load the most recent checkpoint from the designated storage path upon startup.
4. The cluster has sufficient available resources to fulfill the original resource requests of the job.

Procedure

1. Obtain the API server URL and authentication token from your OpenShift cluster, as follows:

   # Get the API server URL
   oc whoami --show-server
   
   # Get your authentication token
   oc whoami --show-token

   Note

   If your cluster uses a self-signed certificate, you must set configuration.verify_ssl = False in your code to avoid SSL verification errors.

2. In your workbench, create a Python script or notebook cell and define your training function, configure authentication, and submit the training job, as follows:

   from kubernetes import client
   from kubeflow.trainer import TrainerClient
   from kubeflow.trainer.rhai import TransformersTrainer
   from kubeflow.common.types import KubernetesBackendConfig
   
   
   # Define your training function
   def train_func():
       from transformers import (
           AutoModelForSequenceClassification,
           AutoTokenizer,
           Trainer,
           TrainingArguments,
       )
       from datasets import load_dataset
   
       # Load model and tokenizer
       model = AutoModelForSequenceClassification.from_pretrained(
           "bert-base-uncased", num_labels=2
       )
       tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
   
       # Load and tokenize dataset
       dataset = load_dataset("imdb", split="train[:1000]")
   
       def tokenize_function(examples):
           return tokenizer(
               examples["text"], padding="max_length", truncation=True, max_length=512
           )
   
       tokenized_dataset = dataset.map(tokenize_function, batched=True)
   
       # Configure training
       training_args = TrainingArguments(
           output_dir="/tmp/output",
           num_train_epochs=3,
           per_device_train_batch_size=8,
       )
   
       # Train
       trainer = Trainer(
           model=model, args=training_args, train_dataset=tokenized_dataset
       )
       trainer.train()
   
   
   # Configure authentication for the workbench
   api_server = "<your-api-server-url>"  # e.g., "https://api.cluster.example.com:6443"
   token = "<your-bearer-token>"
   
   configuration = client.Configuration()
   configuration.host = api_server
   configuration.api_key = {"authorization": f"Bearer {token}"}
   # Uncomment the following line if your cluster uses a self-signed certificate
   # configuration.verify_ssl = False
   
   # Create the client with authentication
   trainer_client = TrainerClient(
       backend_config=KubernetesBackendConfig(client_configuration=configuration)
   )
   
   # Create the trainer configuration
   trainer = TransformersTrainer(
       func=train_func,
       num_nodes=2,
       resources_per_node={"nvidia.com/gpu": 1},
   )
   
   # Get the runtime and submit the training job
   runtime = trainer_client.backend.get_runtime("torch-distributed")
   job_name = trainer_client.train(trainer=trainer, runtime=runtime)
   print(f"Job {job_name} submitted!")

   Replace <your-api-server-url> with the API server URL from step 1 and <your-bearer-token> with the authentication token from step 1.

   The TransformersTrainer parameters are:

   Expand

   Parameter	Type	Description
   func	Callable	The Python function contains your training code.
   num_nodes	int	Number of distributed training nodes (default: 1).
   resources_per_node	dict	Resource requests per node, for example "nvidia.com/gpu": 1, "memory": "16Gi".

   Show more

Verification

1. Verify that the job was submitted by checking the job status:

   job = trainer_client.get_job(job_name)
   print(f"   Status: {job.status}")

2. Verify that the training pods are running:

   oc get pods -n <your-namespace> -l job-name=<job-name>

   Replace <your-namespace> with the name of your project and <job-name> with the name of your training job.

3. (Optional) Verify that real-time progress metrics are displayed for your training job. Navigate to the OpenShift AI dashboard and click TrainingJobs to view these metrics.

Procedure

1. Open the workbench:
2. Log in to the Red Hat OpenShift AI web console.
3. Click Projects and click your project.
4. Click the Workbenches tab.
5. Ensure that the workbench uses a storage class with RWX capability.
6. If your workbench is not already running, start the workbench.
7. Click the Open link to open the IDE in a new window.
8. Click File > New > Notebook.

9. Install the required packages. In a notebook cell, add the code to install the dependencies, as follows:

   !pip install kubeflow --index-url https://console.redhat.com/api/pypi/public-rhai/rhoai/3.2/cuda12.9-ubi9/simple/
   
   !pip install training-hub==0.3.0

10. Select the cell, and click Run > Run selected cell.

11. Configure the Kubeflow SDK client authentication by creating a cell with the following content:

    from kubernetes import client as k8s
    from kubeflow.common.types import KubernetesBackendConfig
    from kubeflow.trainer import TrainerClient
    
    api_server = "<API_SERVER>"
    token = "<TOKEN>"
    
    configuration = k8s.Configuration()
    configuration.host = api_server
    configuration.api_key = {"authorization": f"Bearer {token}"}
    # Uncomment if your cluster API server uses a self-signed certificate
    # configuration.verify_ssl = False
    
    backend_cfg = KubernetesBackendConfig(client_configuration=configuration)
    client = TrainerClient(backend_cfg)

12. Edit the api_server and token parameters to enter the values to authenticate to your OpenShift cluster. Run the cell to configure the SDK client authentication.

13. Verify that a Training Hub runtime is available, as follows:

    for runtime in client.list_runtimes():
        print("Found runtime: " + str(runtime))
        if runtime.name == "training-hub":
            th_runtime = runtime
            print("Selected runtime: " + str(th_runtime))

14. Set the training parameters. The training parameters depend on the fine-tuning algorithm you choose. See the following sections for algorithm-specific configuration:

    • OSFT training parameters for OSFT fine-tuning
    • SFT training parameters for SFT fine-tuning
15. Click File > Save Notebook As, enter an appropriate file name, and click Save.

Procedure

1. Open the workbench, as follows:

   1. Log in to the Red Hat OpenShift AI web console.
   2. Click Projects and click your project.
   3. Click the Workbenches tab. If your workbench is not already running, start the workbench.
   4. Click the Open link to open the IDE in a new window.
   5. Click File > Open, and open the Jupyter notebook that you used to configure the fine-tuning job.
2. Create a cell to run the job.

3. Depending on your fine-tuning method, add the code for either OSFT or SFT:

   1. To enable OSFT fine-tuning, add:

      from kubeflow.trainer.rhai import TrainingHubAlgorithms, TrainingHubTrainer
      from kubeflow.trainer.options.kubernetes import (
          ContainerOverride,
          PodSpecOverride,
          PodTemplateOverride,
          PodTemplateOverrides,
      )
      
      PVC_NAME = "shared"
      
      job_name = client.train(
          trainer=TrainingHubTrainer(
              algorithm=TrainingHubAlgorithms.OSFT,
              func_args=params,
              env={
                  "HF_HOME": "/mnt/shared/huggingface",
                  "TRITON_CACHE_DIR": "/mnt/shared/.triton",
                  "XDG_CACHE_HOME": "/opt/app-root/src/.cache",
                  "NCCL_DEBUG": "INFO",
              },
              resources_per_node={
                  "nvidia.com/gpu": 2,
                  "memory": "32Gi",
                  "cpu": 4,
              },
          ),
          options=[
              PodTemplateOverrides(
                  PodTemplateOverride(
                      target_jobs=["node"],
                      spec=PodSpecOverride(
                          volumes=[
                              {
                                  "name": "work",
                                  "persistentVolumeClaim": {"claimName": PVC_NAME},
                              },
                          ],
                          containers=[
                              ContainerOverride(
                                  name="node",
                                  volume_mounts=[
                                      {
                                          "name": "work",
                                          "mountPath": "/mnt/shared",
                                      },
                                  ],
                              ),
                          ],
                      ),
                  )
              )
          ],
          runtime=th_runtime,
      )

   2. To enable SFT fine-tuning, add:

      from kubeflow.trainer.rhai import TrainingHubAlgorithms, TrainingHubTrainer
      from kubeflow.trainer.options.kubernetes import (
          ContainerOverride,
          PodSpecOverride,
          PodTemplateOverride,
          PodTemplateOverrides,
      )
      
      PVC_NAME = "shared"
      
      job_name = client.train(
          trainer=TrainingHubTrainer(
              algorithm=TrainingHubAlgorithms.SFT,
              func_args=training_parameters,
              env={
                  "HF_HOME": "/mnt/shared/huggingface",
                  "TRITON_CACHE_DIR": "/mnt/shared/.triton",
                  "XDG_CACHE_HOME": "/opt/app-root/src/.cache",
                  "NCCL_DEBUG": "INFO",
              },
              resources_per_node={
                  "nvidia.com/gpu": 2,
                  "memory": "64Gi",
                  "cpu": 4,
              },
          ),
          options=[
              PodTemplateOverrides(
                  PodTemplateOverride(
                      target_jobs=["node"],
                      spec=PodSpecOverride(
                          volumes=[
                              {
                                  "name": "work",
                                  "persistentVolumeClaim": {"claimName": PVC_NAME},
                              },
                          ],
                          containers=[
                              ContainerOverride(
                                  name="node",
                                  volume_mounts=[
                                      {
                                          "name": "work",
                                          "mountPath": "/mnt/shared",
                                      },
                                  ],
                              ),
                          ],
                      ),
                  )
              )
          ],
          runtime=th_runtime,
      )

1. Create a cell with the following content:

   for logline in client.get_job_logs(job_name, follow=True):
       print(logline, end="")

2. Run the cell to view the job progress.

3. (Optional) Check the final job status:

   job = client.get_job(job_name)
   print(f"Name: {job.name}")
   print(f"Status: {job.status}")
   print(f"Nodes: {job.num_nodes}")
   print(f"Runtime: {job.runtime.name}")

Procedure

1. Log in to the Red Hat OpenShift AI web console.
2. Click Projects and click your project.
3. Click the Workbenches tab. If your workbench is not already running, start the workbench.
4. Click the Open link to open the IDE in a new window.
5. Click File > Open, and open the Jupyter notebook that you used to configure and run the fine-tuning job.

6. Create a cell with the following content:

   client.delete_job(name=job_name)

7. Run the cell to delete the job.

Verification

1. In the OpenShift Console, in the Administrator perspective, click Workloads > Jobs.
2. From the Project list, select your project.
3. Verify that the specified job is not listed.

1. Create a ConfigMap with the training script:

   apiVersion: v1
   kind: ConfigMap
   metadata:
     name: minimal-train-script
     namespace: test-namespace
   data:
     train.py: |
       import torch
       import torch.distributed as dist
   
       dist.init_process_group(backend="nccl")
       rank = dist.get_rank()
       world_size = dist.get_world_size()
   
       print(f"Rank {rank}/{world_size}")
       print(f"PyTorch: {torch.version}, CUDA: {torch.cuda.is_available()}")
       if torch.cuda.is_available():
           print(f"GPU: {torch.cuda.get_device_name(0)}")
   
       tensor = torch.ones(1).cuda() * rank
       dist.all_reduce(tensor)
       print(f"Rank {rank}: all_reduce result = {tensor.item()}")
   
       dist.destroy_process_group()
       print("Done.")

2. Create the TrainJob:

   Use podTemplateOverrides to add the ConfigMap volume and mount it into the trainer container. The target job name node matches the default torch-distributed runtime.

   apiVersion: trainer.kubeflow.org/v1alpha1
   kind: TrainJob
   metadata:
     name: pytorch-minimal-example
     namespace: test-namespace
   spec:
     runtimeRef:
       name: torch-distributed
       kind: ClusterTrainingRuntime
     trainer:
       command: ["torchrun", "/workspace/train.py"]
       numNodes: 2
       resourcesPerNode:
         requests:
           nvidia.com/gpu: 1
         limits:
           nvidia.com/gpu: 1
     podTemplateOverrides:
       - targetJobs:
           - name: node
         spec:
           volumes:
             - name: script-volume
               configMap:
                 name: minimal-train-script
           containers:
             - name: node
               volumeMounts:
                 - name: script-volume
                   mountPath: /workspace