Hugging Face's logo
Back to Articles

Fine Tuning a LLM Using Kubernetes with Intel® Gaudi® Accelerator

Community Article Published September 9, 2024
Upvote
7
Omar Khleif's avatar
omarkhleif Omar Khleif

Introduction

Large language models (LLM) used for text generation have exploded in popularity, but it's no secret that a massive amount of compute power is required to train these models due to their large model and dataset sizes. Intel® Gaudi® Accelerator offers a powerful, cost-effective, and scalable solution for fine tuning these models. And Kubernetes (or K8s for short) is exemplary for running containerized workloads across a cluster of nodes. In this blog, we take a deep dive into the process of fine tuning a state of the art model such as meta-llama/Llama-3-8B-Instruct on the tatsu-lab/alpaca dataset using an Intel Gaudi Accelerator node from a K8s cluster.

Table of Contents

Components

In the tutorial, we will be fine tuning Llama 3-8B-Instruct with a Hugging Face dataset using multiple Intel Gaudi HPU cards. Several different components are involved to run this job on the cluster, which will be further explained in this blog.

Helm Chart

The first component that we're going to talk about is the Helm chart. Helm brings together all the different components that are used for our job and allows us to deploy everything using one helm install command. The K8s resources used in our example are:

  • Job used to run the optimum-habana examples
  • Persistent Volume Claim (PVC) used as a shared storage location among the workers for dataset and model files
  • Secret for 🤗 gated models (Optional)
  • Data access pod (Optional)

The Helm chart has a values.yaml file with parameters that are used in the spec files for the K8s resources. Our values files include parameters such the name for our K8s resources, the image/tag for the container to use for the worker pod, HPU and memory resources, the arguments for the python script, etc. The values get filled into the K8s spec files when the Helm chart is deployed, depending on which values file is used.

Container

K8s runs jobs in a containerized environment, so the next thing that we're going to need is a docker container. Here, we need to include all the dependencies needed for our training job.

The Dockerfile and docker-compose.yaml builds the following images:

  • An optimum-habana base image that uses the PyTorch Docker images for Gaudi as its base, and then installs optimum-habana and the Habana fork of DeepSpeed.
  • An optimum-habana-examples image that is built on top of the optimum-habana base. It includes installations from requirements.txt files in the example directories and a clone of this GitHub repository in order to run example scripts.
# Specify the base image name/tag
export BASE_IMAGE_NAME=vault.habana.ai/gaudi-docker/1.17.1/ubuntu22.04/habanalabs/pytorch-installer-2.3.1
export BASE_IMAGE_TAG=latest

# Specify your Gaudi Software version and Optimum Habana version
export GAUDI_SW_VER=1.17.1
export OPTIMUM_HABANA_VER=1.13.0

git clone https://github.com/huggingface/optimum-habana.git

# Note: Modify the requirements.txt file in the kubernetes directory for the specific example(s) that you want to run
cd optimum-habana/examples/kubernetes

# Set variables for your container registry and repository
export REGISTRY=<Your container registry>
export REPO=<Your container repository>

# Build the images
docker compose build

# Push the optimum-habana-examples image to a container registry
docker push <image name>:<tag>
Package Name Version Purpose
PyTorch 2.3.1 Base framework to train models
🤗 Transformers 4.44.2 Library used to download and fine tune the Hugging Face model
Optimum Habana 1.13.0 Interface between the Transformers and Diffusers libraries and Intel Gaudi AI Accelerators (HPU)
Habana Fork of DeepSpeed 1.17.1 The Habana fork of DeepSpeed for using DeepSpeed on HPUs

Fine Tuning Script

The python script that we are using finetunes a causal language model using LoRA for text generation. This script is one from the Optimum Habana example scripts.

In this example, the Optimum-Habana repo is being cloned in the optimum-habana-examples container, hence, no further action is needed. If you wish to use your own script, you can use the base optimum-habana container to do a COPY of the script.

Storage

For this blog we will be using a Hugging Face dataset, however, a storage location can also be used for custom datasets as well as output from the optimum-habana example script using a vanilla K8s cluster with an NFS backed storage class. If you are using a cloud service provider, you could use a cloud storage bucket instead. The storage location gets mounted into the container so that we have read and write access from that location without it being built into the image. To achieve this, we are using a persistent volume claim (PVC).

Secret

Gated or private models require you to be logged in to download the model. If the model being trained is not gated or private, this isn't required. For authentication from the K8s job, we define a secret with a Hugging Face User Read Only Access Token. The token from the secret will be mounted into the container.

Cluster Requirements

This tutorial requires a Kubernetes cluster with Gaudi accelerators. You will also need the Intel Gaudi Device Plugin for Kubernetes deployed to your cluster.

The kubectl auth can-i get nodes command will return "yes" if you are able to list the nodes with kubectl get nodes, for example:

NAME                   STATUS     ROLES                      AGE   VERSION
gaudi161432            Ready      worker                     48d   v1.11.13

Otherwise, consult your cluster admin to get a list of the nodes available to your user group.

Once you know the names of the node(s), use kubectl describe node <node name> to get its HPU and memory capacity. We will be using this information later when setting up the specification for the worker pods.

Tutorial: Fine Tuning Llama 3 using a Kubernetes Cluster

Client Requirements:

Optional:

  • Granted access to Meta-Llama-3-8B-Instruct, Llama2-7b-hf, or equivalent on the model's respective Hugging Face model card page. Otherwise, a huggyllama model can be used without the need of any access approval. Please note that this tutorial showcases Meta-Llama-3-8B-Instruct.

Step 1: Setup the secret with your Hugging Face token

Get a Hugging Face token with read access and use your terminal to get the base64 encoding for your token using echo <your token> | base64.

For example:

$ echo hf_ABCDEFG | base64
aGZfQUJDREVGRwo=

Copy and paste the encoded token value into your values yaml file encodedToken field in the secret section. For example, to run the multi-card LoRA fine tuning job, open the examples/kubernetes/ci/multi-card-lora-clm-values.yaml file and paste in your encoded token on line 41:

secret:
  encodedToken: aGZfQUJDREVGRwo=

Step 2: Customize your values.yaml parameters

The examples/kubernetes/ci/multi-card-lora-clm-values.yaml file is setup to fine tune huggyllama/llama-7b using the tatsu-lab/alpaca dataset. For this blog, we will change the model name to meta-llama/Llama-3-8B-Instruct. You may also fill in either a dataset_name to use a Hugging Face dataset, or provide a train_file and validation_file path for a custom dataset.

Likewise, the values file can be changed to adjust the training job's dataset, epochs, max steps, learning rate, LoRA config, enable bfloat16, etc.

The values files also have parameters for setting the pod's security context with your user and group ids to allow running the fine tuning script as a non-root user. If the user and group ids aren't set, it will be run as root.

# -- Specify a pod security context to run as a non-root user
podSecurityContext: {}
  # runAsUser: 1000
  # runAsGroup: 3000
  # fsGroup: 2000

Specify the number of Gaudi cards to use in the snippet below. The Intel Gaudi Device Plugin for Kubernetes enables the registration of the Gaudi accelerators in a container cluster for compute workload.

resources:
  limits:
    # -- Specify the number of Gaudi card(s)
    habana.ai/gaudi: &hpus 2
    # -- Specify CPU resource limits for the job
    cpu: 16
    # -- Specify Memory limits requests for the job
    memory: 256Gi
    # -- Specify hugepages-2Mi requests for the job
    hugepages-2Mi: 4400Mi
  requests:
    # -- Specify the number of Gaudi card(s)
    habana.ai/gaudi: *hpus
    # -- Specify CPU resource requests for the job
    cpu: 16
    # -- Specify Memory resource requests for the jobs
    memory: 256Gi
    # -- Specify hugepages-2Mi requests for the job
    hugepages-2Mi: 4400Mi

There are other parameters in the values.yaml file that need to be configured based on your cluster:

storage:
  # -- Name of the storage class to use for the persistent volume claim.
  storageClassName: nfs-client
  # -- Access modes for the persistent volume.
  accessModes:
  - "ReadWriteMany"
  # -- Storage resources
  resources:
    requests:
      storage: 30Gi
  # -- Locaton where the PVC will be mounted in the pods
  pvcMountPath: &pvcMountPath /tmp/pvc-mount
  # -- A data access pod will be deployed when set to true
  deployDataAccessPod: true

And finally, the all important python command:

command:
  - python
  - /workspace/optimum-habana/examples/gaudi_spawn.py
  - --world_size
  - *hpus
  - --use_mpi
  - /workspace/optimum-habana/examples/language-modeling/run_lora_clm.py 
  - --model_name_or_path
  - meta-llama/Llama-3-8B-Instruct
  - --dataset_name
  - tatsu-lab/alpaca 
  - --bf16=True
  - --output_dir
  - *pvcMountPath
  - --num_train_epochs
  - "3"
  - --do_train 
  - --do_eval 
  - --use_habana 
  - --validation_split_percentage=4 
  - --adam_epsilon=1e-08
   # Note, this has been shorted for cleanliness purposes.

Step 3: Deploy the Helm chart to the cluster

Deploy the Helm chart to the cluster:

# Navigate to the `examples/kubernetes` directory
cd examples/kubernetes

# Deploy the job using the Helm chart, specifying your values file name with the -f parameter

helm install -f ci/multi-card-lora-clm-values.yaml optimum-habana-examples-2card . -n <namespace>

Step 4: Monitor the job

After the Helm chart is deployed to the cluster, the K8s resources like the secret, PVC, and worker pods are created. The job can be monitored by looking at the pod status using kubectl get pods. At first, the pods will show as "Pending" as the containers get pulled and created, then the status should change to "Running".

$ kubectl get pods -n <namespace>
NAME                                                     READY   STATUS    RESTARTS         AGE
optimum-habana-examples-2card-dataaccess                 1/1     Running   0               1m22s
optimum-habana-examples-2card-gaudijob                   1/1     Running   0               1m22s

Watch the training logs using kubectl logs <pod name> -n <namespace>. You can also add -f to stream the log.

$ kubectl logs optimum-habana-examples-2card-gaudijob
...
{'loss': 0.9585, 'grad_norm': 0.1767578125, 'learning_rate': 0.0001, 'epoch': 1.21, 'memory_allocated (GB)': 18.01, 'max_memory_allocated (GB)': 84.32, 'total_memory_available (GB)': 94.62}

Step 5: Download the trained model

After the job completes, the trained model can be copied from /tmp/pvc-mount/output/saved_model (the path defined in your values file for the --output_dir parameter) to the local system using the following command:

kubectl cp --namespace <namespace> optimum-habana-examples-2card-dataaccess:/tmp/pvc-mount/output/saved_model .

Step 6: Clean up

Finally, the resources can be deleted from the cluster using the helm uninstall command with the name of the Helm job to delete. A list of all the deployed Helm releases can be seen using helm list.

helm uninstall --namespace <namespace> optimum-habana-examples-2card

After uninstalling the Helm chart, the resources on the cluster should show a status of "Terminating", and then they will eventually disappear.

Results

We fine tuned Llama3-8B-Instruct using the Alpaca dataset for 3 epochs on our k8s cluster using 2 cards on an Intel Gaudi HPU. We tuned a few parameters to our liking in the multi-card-lora-clm-values.yaml file. Measuring the accuracy of text generation models can be tricky, since there isn't a clear right or wrong answer. Instead, we are using the perplexity metric to roughly gauge how confident the model is in its prediction. We reserved 4% of the data for evaluation. This is what your sample output should include:

***** train metrics *****
  epoch
  max_memory_allocated (GB)
  memory_allocated (GB)
  total_flos
  total_memory_available (GB)
  train_loss  
  train_runtime  
  train_samples_per_second  
  train_steps_per_second 
...
***** eval metrics *****
  epoch  
  eval_accuracy   
  eval_loss   
  eval_runtime  
  eval_samples  
  eval_samples_per_second  
  eval_steps_per_second  
  max_memory_allocated (GB)   
  memory_allocated (GB)   
  perplexity  
  total_memory_available (GB)

Next Steps

Now that we've fine tuned Llama 3-8B-Instruct using the Alpaca dataset, you're probably wondering how to use this information to run your own workload on K8s. If you're using Llama 3 or a similar generative text model, you are in luck, because the same docker container and script can be reused. You'd just have to edit the parameters of the applicable values.yaml file to use your dataset and tweak other parameters (learning rate, epochs, etc) as needed. If you want to use your own fine tuning script, you will need to build and push a docker container that includes the libraries needed to run the training job, along with your script.

Stay tuned for our next blog, where we will cover fine tuning a LLM using K8s and Gaudi for a multinode setup.

All of the scripts, Dockerfile, and spec files for the tutorial can be found in examples/kubernetes.

Acknowledgments

Thank you to my colleagues who made contributions and helped to review this blog: Dina Jones, Harsha Ramayanam, Abolfazl Shahbazi, and Melanie Buehler

Citations 

@article{llama3modelcard,

title={Llama 3 Model Card},

author={AI@Meta},

year={2024},

url = {https://github.com/meta-llama/llama3/blob/main/MODEL_CARD.md}

}

@misc{alpaca,
  author = {Rohan Taori and Ishaan Gulrajani and Tianyi Zhang and Yann Dubois and Xuechen Li and Carlos Guestrin and Percy Liang and Tatsunori B. Hashimoto },
  title = {Stanford Alpaca: An Instruction-following LLaMA model},
  year = {2023},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/tatsu-lab/stanford_alpaca}},
}
Upvote
7