Working with distributed workloads

Procedure

1. In a terminal window, create a directory for your work, and change to that directory.

2. Set the IMG environment variable to the name of your custom image. In the example commands in this section, my_training_image is the name of the custom image.

   export IMG=my_training_image

3. Create a file named Dockerfile with the following content:

   1. Use the FROM instruction to specify the location of a suitable base training image.

      In the following command, replace _<base-training-image>_ with the name of your chosen base training image:

      FROM quay.io/modh/<base-training-image>

      Examples:

      FROM quay.io/modh/ray:2.47.1-py311-cu121

      FROM quay.io/modh/training:py311-rocm62-torch251

   2. Use the RUN instruction to install additional packages. You can also add comments to the Dockerfile by prefixing each comment line with a number sign (#).

      The following example shows how to install a specific version of the Python PyTorch package:

      # Install PyTorch
      RUN python3 -m pip install torch==2.5.1

4. Build the image file. Use the -t option with the podman build command to create an image tag that specifies the custom image name and version, to make it easier to reference and manage the image:

   podman build -t <custom-image-name>:_<version>_ -f Dockerfile

   Example:

   podman build -t ${IMG}:0.0.1 -f Dockerfile

   The build output indicates when the build process is complete.

5. Display a list of your images:

   podman images

   If your new image was created successfully, it is included in the list of images.

6. Push the image to your container registry:

   podman push ${IMG}:0.0.1

7. Optional: Make your new image available to other users, as described in Pushing an image to the integrated OpenShift image registry.

Procedure

1. In a terminal window, log in to the OpenShift CLI (oc) as shown in the following example:

   $ oc login <openshift_cluster_url> -u <admin_username> -p <password>

2. Set the IMG environment variable to the name of your image. In the example commands in this section, my_training_image is the name of the image.

   export IMG=my_training_image

3. Log in to the integrated image registry:

   podman login -u $(oc whoami) -p $(oc whoami -t) $(oc registry info)

4. Tag the image for the integrated image registry:

   podman tag ${IMG} $(oc registry info)/$(oc project -q)/${IMG}

5. Push the image to the integrated image registry:

   podman push $(oc registry info)/$(oc project -q)/${IMG}

6. Retrieve the image repository location for the tag that you want:

   oc get is ${IMG} -o jsonpath='{.status.tags[?(@.tag=="<TAG>")].items[0].dockerImageReference}'

   Any user can now use your image by specifying this retrieved image location value in the image parameter of a Ray cluster or training job.

Procedure

1. In the JupyterLab interface, click File > New > Notebook. Specify your preferred Python version, and then click Select.

   A new Jupyter notebook file is created with the .ipynb file name extension.

2. Add the following code to a cell in the new notebook:

   Code to download the demo Jupyter notebooks

   from codeflare_sdk import copy_demo_nbs
   copy_demo_nbs()

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

   After a few seconds, the copy_demo_nbs() function copies the demo Jupyter notebooks that are packaged with the currently installed version of the CodeFlare SDK, and clones them into the demo-notebooks folder.

4. In the left navigation pane, right-click the new notebook and click Delete.
5. Click Delete to confirm.

1. In the left navigation pane, double-click demo-notebooks.
2. Double-click additional-demos and verify that the folder contains several demo Jupyter notebooks.
3. Click demo-notebooks.
4. Double-click guided-demos and verify that the folder contains several demo Jupyter notebooks.

Procedure

1. Check whether your cluster administrator has defined a default local queue for the Ray cluster.

   You can use the codeflare_sdk.list_local_queues() function to view all local queues in your current namespace, and the resource flavors associated with each local queue.

   Alternatively, you can use the OpenShift web console as follows:

   1. In the OpenShift web console, select your project from the Project list.

   2. Click Search, and from the Resources list, select LocalQueue to show the list of local queues for your project.

      If no local queue is listed, contact your cluster administrator.

   3. Review the details of each local queue:

      1. Click the local queue name.

      2. Click the YAML tab, and review the metadata.annotations section.

         If the kueue.x-k8s.io/default-queue annotation is set to 'true', the queue is configured as the default local queue.

         Note

         If your cluster administrator does not define a default local queue, you must specify a local queue in each Jupyter notebook.

2. In the JupyterLab interface, open the demo-notebooks > guided-demos folder.

3. Open all of the Jupyter notebooks by double-clicking each Jupyter notebook file.

   Jupyter notebook files have the .ipynb file name extension.

4. In each Jupyter notebook, ensure that the import section imports the required components from the CodeFlare SDK, as follows:

   Example import section

   from codeflare_sdk import Cluster, ClusterConfiguration, TokenAuthentication

5. In each Jupyter notebook, update the TokenAuthentication section to provide the token and server details to authenticate to the OpenShift cluster by using the CodeFlare SDK.

   For information about how to find the server and token details, see Using the cluster server and token to authenticate.

6. Optional: If you want to use custom certificates, update the TokenAuthentication section to add the ca_cert_path parameter to specify the location of the custom certificates, as shown in the following example:

   Example authentication section

   auth = TokenAuthentication(
       token = "XXXXX",
       server = "XXXXX",
       skip_tls=False,
       ca_cert_path="/path/to/cert"
   )
   auth.login()

   Alternatively, you can set the CF_SDK_CA_CERT_PATH environment variable to specify the location of the custom certificates.

7. In each Jupyter notebook, update the cluster configuration section as follows:

   1. If the namespace value is specified, replace the example value with the name of your project.

      If you omit this line, the Ray cluster is created in the current project.

   2. If the image value is specified, replace the example value with a link to a suitable Ray cluster image. The Python version in the Ray cluster image must be the same as the Python version in the workbench.

      If you omit this line, one of the following Ray cluster images is used by default, based on the Python version detected in the workbench:

      • Python 3.9: quay.io/modh/ray:2.35.0-py39-cu121
      • Python 3.11: quay.io/modh/ray:2.47.1-py311-cu121

      The default Ray images are compatible with NVIDIA GPUs that are supported by the specified CUDA version. The default images are AMD64 images, which might not work on other architectures.

      Additional ROCm-compatible Ray cluster images are available, which are compatible with AMD accelerators that are supported by the specified ROCm version. These images are AMD64 images, which might not work on other architectures.

      For information about the latest available training images and their preinstalled packages, including the CUDA and ROCm versions, see Red Hat OpenShift AI: Supported Configurations.

   3. If your cluster administrator has not configured a default local queue, specify the local queue for the Ray cluster, as shown in the following example:

      Example local queue assignment

      local_queue="your_local_queue_name"

   4. Optional: Assign a dictionary of labels parameters to the Ray cluster for identification and management purposes, as shown in the following example:

      Example labels assignment

      labels = {"exampleLabel1": "exampleLabel1Value", "exampleLabel2": "exampleLabel2Value"}

8. In the 2_basic_interactive.ipynb Jupyter notebook, ensure that the following Ray cluster authentication code is included after the Ray cluster creation section:

   Ray cluster authentication code

   from codeflare_sdk import generate_cert
   generate_cert.generate_tls_cert(cluster.config.name, cluster.config.namespace)
   generate_cert.export_env(cluster.config.name, cluster.config.namespace)

   Note

   Mutual Transport Layer Security (mTLS) is enabled by default in the CodeFlare component in OpenShift AI. You must include the Ray cluster authentication code to enable the Ray client that runs within a Jupyter notebook to connect to a secure Ray cluster that has mTLS enabled.

9. Run the Jupyter notebooks in the order indicated by the file-name prefix (0_, 1_, and so on).

   1. In each Jupyter notebook, run each cell in turn, and review the cell output.
   2. If an error is shown, review the output to find information about the problem and the required corrective action. For example, replace any deprecated parameters as instructed. See also Troubleshooting common problems with distributed workloads for users.
   3. For more information about the interactive browser controls that you can use to simplify Ray cluster tasks when working within a Jupyter notebook, see Managing Ray clusters from within a Jupyter notebook.

Verification

1. The Jupyter notebooks run to completion without errors.
2. In the Jupyter notebooks, the output from the cluster.status() function or cluster.details() function indicates that the Ray cluster is Active.

Procedure

1. Run all of the demo Jupyter notebooks in the order indicated by the file-name prefix (0_, 1_, and so on), as described in Running the demo Jupyter notebooks from the CodeFlare SDK.

2. In each demo Jupyter notebook, when you run the cluster configuration step, the following interactive controls are automatically shown in the Jupyter notebook:

   • Cluster Up: You can click this button to start the Ray cluster. This button is equivalent to the cluster.up() command. When you click this button, a message indicates whether the cluster was successfully created.
   • Cluster Down: You can click this button to delete the Ray cluster. This button is equivalent to the cluster.down() command. The cluster is deleted immediately; you are not prompted to confirm the deletion. When you click this button, a message indicates whether the cluster was successfully deleted.
   • Wait for Cluster: You can select this option to specify that the notebook cell should wait for the Ray cluster dashboard to be ready before proceeding to the next step. This option is equivalent to the cluster.wait_ready() command.

3. In the JupyterLab interface, create a new Jupyter notebook to manage the Ray clusters, as follows:

   1. Click File > New > Notebook. Specify your preferred Python version, and then click Select.

      A new Jupyter notebook file is created with the .ipynb file name extension.

   2. Add the following code to a cell in the new Jupyter notebook:

      Code to import the required packages

      from codeflare_sdk import TokenAuthentication, view_clusters

      The view_clusters package provides the interactive browser controls for listing the clusters, showing the cluster details, opening the Ray dashboard, and refreshing the cluster data.

   3. Add a new notebook cell, and add the following code to the new cell:

      Code to authenticate

      auth = TokenAuthentication(
          token = "XXXXX",
          server = "XXXXX",
          skip_tls=False
      )
      auth.login()

      For information about how to find the token and server values, see Running the demo Jupyter notebooks from the CodeFlare SDK.

   4. Add a new notebook cell, and add the following code to the new cell:

      Code to view clusters in the current project

      view_clusters()

      When you run the view_clusters() command with no arguments specified, you generate a list of all of the Ray clusters in the current project, and display information similar to the cluster.details() function.

      If you have access to another project, you can list the Ray clusters in that project by specifying the project name as shown in the following example:

      Code to view clusters in another project

      view_clusters("my_second_project")

   5. Click File > Save Notebook As, enter demo-notebooks/guided-demos/manage_ray_clusters.ipynb, and click Save.
4. In the demo-notebooks/guided-demos/manage_ray_clusters.ipynb Jupyter notebook, select each cell in turn, and click Run > Run selected cell.

5. When you run the cell with the view_clusters() function, the output depends on whether any Ray clusters exist.

   If no Ray clusters exist, the following text is shown, where _[project-name]_ is the name of the target project:

   No clusters found in the [project-name] namespace.

   Otherwise, the Jupyter notebook shows the following information about the existing Ray clusters:

   • Select an existing cluster

     Under this heading, a toggle button is shown for each existing cluster. Click a cluster name to select the cluster. The cluster details section is updated to show details about the selected cluster; for example, cluster name, OpenShift AI project name, cluster resource information, and cluster status.

   • Delete cluster

     Click this button to delete the selected cluster. This button is equivalent to the Cluster Down button. The cluster is deleted immediately; you are not prompted to confirm the deletion. A message indicates whether the cluster was successfully deleted, and the corresponding button is no longer shown under the Select an existing cluster heading.

   • View Jobs

     Click this button to open the Jobs tab in the Ray dashboard for the selected cluster, and view details of the submitted jobs. The corresponding URL is shown in the Jupyter notebook.

   • Open Ray Dashboard

     Click this button to open the Overview tab in the Ray dashboard for the selected cluster. The corresponding URL is shown in the Jupyter notebook.

   • Refresh Data

     Click this button to refresh the list of Ray clusters, and the cluster details for the selected cluster, on demand. The cluster details are automatically refreshed when you select a cluster and when you delete the selected cluster.

Verification

1. The demo Jupyter notebooks run to completion without errors.
2. In the manage_ray_clusters.ipynb Jupyter notebook, the output from the view_clusters() function is correct.

Procedure

1. Log in to the OpenShift Console.

2. Create a ConfigMap resource, as follows:

   1. In the Administrator perspective, click Workloads → ConfigMaps.
   2. From the Project list, select your project.
   3. Click Create ConfigMap.

   4. In the Configure via section, select the YAML view option.

      The Create ConfigMap page opens, with default YAML code automatically added.

3. Replace the default YAML code with your training-script code.

   For example training scripts, see Example Training Operator PyTorch training scripts.

4. Click Create.

Verification

1. In the OpenShift Console, in the Administrator perspective, click Workloads → ConfigMaps.
2. From the Project list, select your project.
3. Click your ConfigMap resource to display the training script details.

Procedure

1. Log in to the OpenShift Console.

2. Create a PyTorchJob 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, start typing PyTorchJob.

   4. Select PyTorchJob, and click Create PyTorchJob.

      The Create PyTorchJob page opens, with default YAML code automatically added.

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

4. Configure the master node, as shown in the following example:

   spec:
     pytorchReplicaSpecs:
       Master:
         replicas: 1
         restartPolicy: OnFailure
         template:
           metadata:
             labels:
               app: pytorch-multi-node-job

   1. In the replicas entry, specify 1. Only one master node is needed.

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

      Adding the training script from a ConfigMap resource

      Spec:
        pytorchReplicaSpecs:
          Master:
            ...
            template:
              spec:
                containers:
                - name: pytorch
                  image: quay.io/modh/training:py311-cuda124-torch251
                  command: ["python", "/workspace/scripts/train.py"]
                  volumeMounts:
                  - name: training-script-volume
                    mountPath: /workspace
                volumes:
                - name: training-script-volume
                  configMap:
                    name: training-script-configmap

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

      Adding the resource contraints

      SSpec:
        pytorchReplicaSpecs:
          Master:
            ...
            template:
              spec:
                containers: ...
                resources:
                  requests:
                        cpu: "4"
                        memory: "8Gi"
                        nvidia.com/gpu: 2    # To use GPUs (Optional)
                  limits:
                        cpu: "4"
                        memory: "8Gi"
                        nvidia.com/gpu: 2

5. Make similar edits in the Worker section of the PyTorchJob resource.

   1. Update the replicas entry to specify the number of worker nodes.

   For a complete example PyTorchJob resource, see Example Training Operator PyTorchJob resource for multi-node training.

6. Click Create.

Verification

1. In the OpenShift Console, open the Administrator perspective.
2. From the Project list, select your project.
3. Click Home → Search → PyTorchJob and verify that the job was created.
4. Click Workloads → Pods and verify that requested head pod and worker pods are running.

Procedure

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

   Logging in to the OpenShift CLI (oc)

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

   For information about how to find the server and token details, see Using the cluster server and token to authenticate.

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

   Creating the training script

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

   Replace <paste your content here> with your training script content.

   For example training scripts, see Example Training Operator PyTorch training scripts.

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

   Creating the ConfigMap resource

   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 pytorchjob.yaml to define the distributed training job setup, as follows:

   Defining the distributed training job

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

   Replace <paste your content here> with your training job content.

   For an example training job, see Example Training Operator PyTorchJob resource for multi-node training.

5. Create the distributed training job, as follows:

   Creating the distributed training job

   oc apply -f pytorchjob.yaml

Verification

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

   Monitoring the distributed training job

   oc get pytorchjobs -n <your-namespace>

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

2. Check the pod logs, as follows:

   Checking the pod logs

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

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

3. When you want to delete the job, run the following command:

   Deleting the job

   oc delete pytorchjobs/pytorch-multi-node-job -n <your-namespace>

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

Procedure

1. Open the workbench, as follows:

   1. Log in to the Red Hat OpenShift AI web console.
   2. Click Data science projects and click your project.
   3. Click the Workbenches tab.
   4. If your workbench is not already running, start the workbench.
   5. Click the Open link to open the IDE in a new window.
2. Click File → New → Notebook.

3. Create the training function as shown in the following example:

   1. Create a cell with the following content:

      Example training function

      def train_func():
          import os
          import torch
          import torch.distributed as dist
      
          # Select backend dynamically: nccl for GPU, gloo for CPU
          backend = "nccl" if torch.cuda.is_available() else "gloo"
      
          # Initialize the process group
          dist.init_process_group(backend)
      
          # Get rank and world size
          rank = dist.get_rank()
          world_size = dist.get_world_size()
      
          # Select device dynamically
          device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
      
          # Log rank initialization
          print(f"Rank {rank}/{world_size} initialized with backend {backend} on device {device}.")
      
          # Initialize tensor on the selected device
          tensor = torch.zeros(1, device=device)
      
          if rank == 0:
              tensor += 1
              for i in range(1, world_size):
                  dist.send(tensor, dst=i)
          else:
              dist.recv(tensor, src=0)
      
          print(f"Rank {rank}: Tensor value {tensor.item()} on {device}")
      
          # Cleanup
          dist.destroy_process_group()

      Note

      For this example training job, you do not need to install any additional packages or set any training parameters.

      For more information about how to add additional packages and set the training parameters, see Configuring the fine-tuning job.

   2. Optional: Edit the content to specify the appropriate values for your environment.
   3. Run the cell to create the training function.

4. Configure the Training Operator SDK client authentication as follows:

   1. Create a cell with the following content:

      Example Training Operator SDK client authentication

      from kubernetes import client
      from kubeflow.training import TrainingClient
      from kubeflow.training.models import V1Volume, V1VolumeMount, V1PersistentVolumeClaimVolumeSource
      
      api_server = "<API_SERVER>"
      token = "<TOKEN>"
      
      configuration = client.Configuration()
      configuration.host = api_server
      configuration.api_key = {"authorization": f"Bearer {token}"}
      # Un-comment if your cluster API server uses a self-signed certificate or an un-trusted CA
      #configuration.verify_ssl = False
      api_client = client.ApiClient(configuration)
      client = TrainingClient(client_configuration=api_client.configuration)

   2. Edit the api_server and token parameters to enter the values to authenticate to your OpenShift cluster.

      For information on how to find the server and token details, see Using the cluster server and token to authenticate.

   3. Run the cell to configure the Training Operator SDK client authentication.
5. Click File > Save Notebook As, enter an appropriate file name, and click Save.

Verification

1. All cells run successfully.

Procedure

1. Open the workbench, as follows:

   1. Log in to the Red Hat OpenShift AI web console.
   2. Click Data science 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.
2. Click File → Open, and open the Jupyter notebook that you used to configure the training job.

3. Create a cell to run the job, and add the following content:

   from kubernetes import client
   
   # Start PyTorchJob with 2 Workers and 2 GPU per Worker (multi-node, multi-worker job).
   client.create_job(
      name="pytorch-ddp",
      train_func=train_func,
      base_image="quay.io/modh/training:py311-cuda124-torch251",
      num_workers=2,
      resources_per_worker={"nvidia.com/gpu": "2"},
      packages_to_install=["torchvision==0.19.0"],
      env_vars={"NCCL_DEBUG": "INFO", "TORCH_DISTRIBUTED_DEBUG": "DETAIL"},
      labels={
           "kueue.x-k8s.io/queue-name": "<local-queue-name>",
           "key": "value"
       },
      annotations={"key": "value"}
   )

4. Edit the content to specify the appropriate values for your environment, as follows:

   1. Edit the num_workers value to specify the number of worker nodes.
   2. Update the resources_per_worker values according to the job requirements and the resources available.
   3. Edit the value of the kueue.x-k8s.io/queue-name label to match the name of your target LocalQueue.

   4. The example provided is for NVIDIA GPUs. If you use AMD accelerators, make the following additional changes:

      • In the resources_per_worker entry, change nvidia.com/gpu to amd.com/gpu
      • Change the base_image value to quay.io/modh/training:py311-rocm62-torch251
      • Remove the NCCL_DEBUG entry

1. Run the cell to run the job.

1. Create a cell with the following content:

   client.get_job_logs(
       name="pytorch-ddp",
       job_kind="PyTorchJob",
       follow=True,
   )

2. Run the cell to view the job progress.

Procedure

1. Open the workbench, as follows:

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

3. Install any additional packages that are needed to run the training or tuning job.

   1. In a notebook cell, add the code to install the additional packages, as follows:

      Code to install dependencies

      # Install the yamlmagic package
      !pip install yamlmagic
      %load_ext yamlmagic
      
      !pip install git+https://github.com/kubeflow/trainer.git@release-1.9#subdirectory=sdk/python

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

      The additional packages are installed.

4. Set the training parameters as follows:

   1. Create a cell with the following content:

      %%yaml parameters
      
      # Model
      model_name_or_path: Meta-Llama/Meta-Llama-3.1-8B-Instruct
      model_revision: main
      torch_dtype: bfloat16
      attn_implementation: flash_attention_2
      
      # PEFT / LoRA
      use_peft: true
      lora_r: 16
      lora_alpha: 8
      lora_dropout: 0.05
      lora_target_modules: ["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
      lora_modules_to_save: []
      init_lora_weights: true
      
      # Quantization / BitsAndBytes
      load_in_4bit: false                       # use 4 bit precision for the base model (only with LoRA)
      load_in_8bit: false                       # use 8 bit precision for the base model (only with LoRA)
      
      # Datasets
      dataset_name: gsm8k                       # id or path to the dataset
      dataset_config: main                      # name of the dataset configuration
      dataset_train_split: train                # dataset split to use for training
      dataset_test_split: test                  # dataset split to use for evaluation
      dataset_text_field: text                  # name of the text field of the dataset
      dataset_kwargs:
        add_special_tokens: false               # template with special tokens
        append_concat_token: false              # add additional separator token
      
      # SFT
      max_seq_length: 1024                      # max sequence length for model and packing of the dataset
      dataset_batch_size: 1000                  # samples to tokenize per batch
      packing: false
      use_liger: false
      
      # Training
      num_train_epochs: 10                      # number of training epochs
      
      per_device_train_batch_size: 32           # batch size per device during training
      per_device_eval_batch_size: 32            # batch size for evaluation
      auto_find_batch_size: false               # find a batch size that fits into memory automatically
      eval_strategy: epoch                      # evaluate every epoch
      
      bf16: true                                # use bf16 16-bit (mixed) precision
      tf32: false                               # use tf32 precision
      
      learning_rate: 1.0e-4                     # initial learning rate
      warmup_steps: 10                          # steps for a linear warmup from 0 to `learning_rate`
      lr_scheduler_type: inverse_sqrt           # learning rate scheduler (see transformers.SchedulerType)
      
      optim: adamw_torch_fused                  # optimizer (see transformers.OptimizerNames)
      max_grad_norm: 1.0                        # max gradient norm
      seed: 42
      
      gradient_accumulation_steps: 1            # number of steps before performing a backward/update pass
      gradient_checkpointing: false             # use gradient checkpointing to save memory
      gradient_checkpointing_kwargs:
        use_reentrant: false
      
      # FSDP
      fsdp: "full_shard auto_wrap offload"      # remove offload if enough GPU memory
      fsdp_config:
        activation_checkpointing: true
        cpu_ram_efficient_loading: false
        sync_module_states: true
        use_orig_params: true
        limit_all_gathers: false
      
      # Checkpointing
      save_strategy: epoch                      # save checkpoint every epoch
      save_total_limit: 1                       # limit the total amount of checkpoints
      resume_from_checkpoint: false             # load the last checkpoint in output_dir and resume from it
      
      # Logging
      log_level: warning                        # logging level (see transformers.logging)
      logging_strategy: steps
      logging_steps: 1                          # log every N steps
      report_to:
      - tensorboard                             # report metrics to tensorboard
      
      output_dir: /mnt/shared/Meta-Llama-3.1-8B-Instruct

   2. Optional: If you specify a different model or dataset, edit the parameters to suit your model, dataset, and resources. If necessary, update the previous cell to specify the dependencies for your training or tuning job.
   3. Run the cell to set the training parameters.

5. Create the training function as follows:

   1. Create a cell with the following content:

      def main(parameters):
          import random
      
          from datasets import load_dataset
          from transformers import (
              AutoTokenizer,
              set_seed,
          )
      
          from trl import (
              ModelConfig,
              ScriptArguments,
              SFTConfig,
              SFTTrainer,
              TrlParser,
              get_peft_config,
              get_quantization_config,
              get_kbit_device_map,
          )
      
          parser = TrlParser((ScriptArguments, SFTConfig, ModelConfig))
          script_args, training_args, model_args = parser.parse_dict(parameters)
      
          # Set seed for reproducibility
          set_seed(training_args.seed)
      
          # Model and tokenizer
          quantization_config = get_quantization_config(model_args)
          model_kwargs = dict(
              revision=model_args.model_revision,
              trust_remote_code=model_args.trust_remote_code,
              attn_implementation=model_args.attn_implementation,
              torch_dtype=model_args.torch_dtype,
              use_cache=False if training_args.gradient_checkpointing or
                                 training_args.fsdp_config.get("activation_checkpointing",
                                                               False) else True,
              device_map=get_kbit_device_map() if quantization_config is not None else None,
              quantization_config=quantization_config,
          )
          training_args.model_init_kwargs = model_kwargs
          tokenizer = AutoTokenizer.from_pretrained(
              model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code, use_fast=True
          )
          if tokenizer.pad_token is None:
              tokenizer.pad_token = tokenizer.eos_token
      
          # You can override the template here according to your use case
          # tokenizer.chat_template = ...
      
          # Datasets
          train_dataset = load_dataset(
              path=script_args.dataset_name,
              name=script_args.dataset_config,
              split=script_args.dataset_train_split,
          )
          test_dataset = None
          if training_args.eval_strategy != "no":
              test_dataset = load_dataset(
                  path=script_args.dataset_name,
                  name=script_args.dataset_config,
                  split=script_args.dataset_test_split,
              )
      
          # Templatize datasets
          def template_dataset(sample):
              # return{"text": tokenizer.apply_chat_template(examples["messages"], tokenize=False)}
              messages = [
                  {"role": "user", "content": sample[question]},
                  {"role": "assistant", "content": sample[answer]},
              ]
              return {"text": tokenizer.apply_chat_template(messages, tokenize=False)}
      
          train_dataset = train_dataset.map(template_dataset, remove_columns=["question", "answer"])
          if training_args.eval_strategy != "no":
              # test_dataset = test_dataset.map(template_dataset, remove_columns=["messages"])
              test_dataset = test_dataset.map(template_dataset, remove_columns=["question", "answer"])
      
          # Check random samples
          with training_args.main_process_first(
              desc="Log few samples from the training set"
          ):
              for index in random.sample(range(len(train_dataset)), 2):
                  print(train_dataset[index]["text"])
      
          # Training
          trainer = SFTTrainer(
              model=model_args.model_name_or_path,
              args=training_args,
              train_dataset=train_dataset,
              eval_dataset=test_dataset,
              peft_config=get_peft_config(model_args),
              tokenizer=tokenizer,
          )
      
          if trainer.accelerator.is_main_process and hasattr(trainer.model, "print_trainable_parameters"):
              trainer.model.print_trainable_parameters()
      
          checkpoint = None
          if training_args.resume_from_checkpoint is not None:
              checkpoint = training_args.resume_from_checkpoint
      
          trainer.train(resume_from_checkpoint=checkpoint)
      
          trainer.save_model(training_args.output_dir)
      
          with training_args.main_process_first(desc="Training completed"):
              print(f"Training completed, model checkpoint written to {training_args.output_dir}")

   2. Optional: If you specify a different model or dataset, edit the tokenizer.chat_template parameter to specify the appropriate value for your model and dataset.
   3. Run the cell to create the training function.

6. Configure the Training Operator SDK client authentication as follows:

   1. Create a cell with the following content:

      from kubernetes import client
      from kubeflow.training import TrainingClient
      from kubeflow.training.models import V1Volume, V1VolumeMount, V1PersistentVolumeClaimVolumeSource
      
      api_server = "<API_SERVER>"
      token = "<TOKEN>"
      
      configuration = client.Configuration()
      configuration.host = api_server
      configuration.api_key = {"authorization": f"Bearer {token}"}
      # Un-comment if your cluster API server uses a self-signed certificate or an un-trusted CA
      #configuration.verify_ssl = False
      api_client = client.ApiClient(configuration)
      client = TrainingClient(client_configuration=api_client.configuration)

   2. Edit the api_server and token parameters to enter the values to authenticate to your OpenShift cluster.

      For information about how to find the server and token details, see Using the cluster server and token to authenticate.

   3. Run the cell to configure the Training Operator SDK client authentication.
7. Click File > Save Notebook As, enter an appropriate file name, and click Save.

Verification

1. All cells run successfully.

Procedure

1. Open the workbench, as follows:

   1. Log in to the Red Hat OpenShift AI web console.
   2. Click Data science 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.
2. Click File → Open, and open the Jupyter notebook that you used to configure the fine-tuning job.

3. Create a cell to run the job, and add the following content:

   client.create_job(
       job_kind="PyTorchJob",
       name="sft",
       train_func=main,
       num_workers=8,
       num_procs_per_worker="1",
       resources_per_worker={
           "nvidia.com/gpu": 1,
           "memory": "64Gi",
           "cpu": 4,
       },
       base_image="quay.io/modh/training:py311-cuda124-torch251",
       env_vars={
           # Hugging Face
           "HF_HOME": "/mnt/shared/.cache",
           "HF_TOKEN": "",
           # CUDA
           "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
           # NCCL
           "NCCL_DEBUG": "INFO",
           "NCCL_ENABLE_DMABUF_SUPPORT": "1",
       },
       packages_to_install=[
           "tensorboard",
       ],
       parameters=parameters,
       volumes=[
           V1Volume(name="shared",
               persistent_volume_claim=V1PersistentVolumeClaimVolumeSource(claim_name="shared")),
       ],
       volume_mounts=[
           V1VolumeMount(name="shared", mount_path="/mnt/shared"),
       ],
   )

4. Edit the HF_TOKEN value to specify your Hugging Face access token.

   Optional: If you specify a different model, and your model is not a gated model from the Hugging Face Hub, remove the HF_HOME and HF_TOKEN entries.

5. Optional: Edit the other content to specify the appropriate values for your environment, as follows:

   1. Edit the num_workers value to specify the number of worker nodes.
   2. Update the resources_per_worker values according to the job requirements and the resources available.

   3. The example provided is for NVIDIA GPUs. If you use AMD accelerators, make the following additional changes:

      • In the resources_per_worker entry, change nvidia.com/gpu to amd.com/gpu
      • Change the base_image value to quay.io/modh/training:py311-rocm62-torch251
      • Remove the CUDA and NCCL entries

   4. If the RWX PersistentVolumeClaim resource that is attached to your workbench has a different name instead of shared, update the following values to replace shared with your PVC name:

      • In this cell, update the HF_HOME value.

      • In this cell, in the volumes entry, update the PVC details:

        • In the V1Volume entry, update the name and claim_name values.
        • In the volume_mounts entry, update the name and mount_path values.

      • In the cell where you set the training parameters, update the output_dir value.

        For more information about setting the training parameters, see Configuring the fine-tuning job.

6. Run the cell to run the job.

1. Create a cell with the following content:

   client.get_job_logs(
       name="sft",
       job_kind="PyTorchJob",
       follow=True,
   )

2. Run the cell to view the job progress.

Procedure

1. Open the workbench, as follows:

   1. Log in to the Red Hat OpenShift AI web console.
   2. Click Data science 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.
2. Click File → Open, and open the Jupyter notebook that you used to configure and run the example fine-tuning job.

3. Create a cell with the following content:

   client.delete_job(name="sft")

4. Optional: If you want to delete a different job, edit the content to replace sft with the name of your job.
5. 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.