Fine-tune Gemma open models using multiple GPUs on GKE

Get access to the model

To get access to the Gemma models for deployment to GKE, you must first sign the license consent agreement then generate a Hugging Face access token.

Sign the license consent agreement

You must sign the consent agreement to use Gemma. Follow these instructions:

1. Access the model consent page on Kaggle.com.
2. Verify consent using your Hugging Face account.
3. Accept the model terms.

Generate an access token

To access the model through Hugging Face, you'll need a Hugging Face token.

Follow these steps to generate a new token if you don't have one already:

1. Click Your Profile > Settings > Access Tokens.
2. Select New Token.
3. Specify a Name of your choice and a Role of at least Write.
4. Select Generate a token.
5. Copy the generated token to your clipboard.

Prepare your environment

In this tutorial, you use Cloud Shell to manage resources hosted on Google Cloud. Cloud Shell comes preinstalled with the software you'll need for this tutorial, including kubectl and gcloud CLI.

To set up your environment with Cloud Shell, follow these steps:

1. In the Google Cloud console, launch a Cloud Shell session by clicking Activate Cloud Shell in the Google Cloud console. This launches a session in the bottom pane of Google Cloud console.

2. Set the default environment variables:

   gcloud config set project PROJECT_ID
   export PROJECT_ID=$(gcloud config get project)
   export REGION=REGION
   export CLUSTER_NAME=finetuning
   export HF_TOKEN=HF_TOKEN
   export HF_PROFILE=HF_PROFILE
   

   Replace the following values:

   • PROJECT_ID: Your Google Cloud project ID.
   • REGION: A region that supports the accelerator type you want to use, for example, us-central1 for L4 GPU.
   • HF_TOKEN: The Hugging Face token you generated earlier.
   • HF_PROFILE: The Hugging Face Profile ID that you created earlier.

3. Clone the sample code repository from GitHub:

   git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples
   cd kubernetes-engine-samples/ai-ml/llm-finetuning-gemma
   

Create and configure Google Cloud resources

Follow these instructions to create the required resources.

Create a GKE cluster and node pool

You can serve Gemma on GPUs in a GKE Autopilot or Standard cluster. To choose the GKE mode of operation that's the best fit for your workloads, see Choose a GKE mode of operation.

gcloud container clusters create-auto ${CLUSTER_NAME} \
  --project=${PROJECT_ID} \
  --region=${REGION} \
  --release-channel=rapid \
  --cluster-version=1.29

Create a Kubernetes secret for Hugging Face credentials

In Cloud Shell, do the following:

1. Configure kubectl to communicate with your cluster:

   gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${REGION}
   

2. Create a Kubernetes Secret that contains the Hugging Face token:

   kubectl create secret generic hf-secret \
     --from-literal=hf_api_token=${HF_TOKEN} \
     --dry-run=client -o yaml | kubectl apply -f -
   

Create a fine tuning container with Docker and Cloud Build

This container uses the PyTorch and Hugging Face Transformers code to fine tune the existing pre-trained Gemma model.

1. Create a Artifact Registry Docker Repository

   gcloud artifacts repositories create gemma \
       --project=${PROJECT_ID} \
       --repository-format=docker \
       --location=us \
       --description="Gemma Repo"
   

2. Build and push the image

   gcloud builds submit .
   

3. Export the IMAGE_URL for later use in this tutorial.

   export IMAGE_URL=us-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-gpu:1.0.0
   

Run a fine-tune Job on GKE

Deploy the Gemma fine tuning job

1. Open the finetune.yaml file.

   apiVersion: batch/v1
   kind: Job
   metadata:
     name: finetune-job
     namespace: default
   spec:
     backoffLimit: 2
     template:
       metadata:
         annotations:
           kubectl.kubernetes.io/default-container: finetuner
       spec:
         terminationGracePeriodSeconds: 600
         containers:
         - name: finetuner
           image: $IMAGE_URL
           resources:
             limits:
               nvidia.com/gpu: "8"
           env:
           - name: MODEL_NAME
             value: "google/gemma-2b"
           - name: NEW_MODEL
             value: "gemma-2b-sql-finetuned"
           - name: LORA_R
             value: "8"
           - name: LORA_ALPHA
             value: "16"
           - name: TRAIN_BATCH_SIZE
             value: "1"
           - name: EVAL_BATCH_SIZE
             value: "2"
           - name: GRADIENT_ACCUMULATION_STEPS
             value: "2"
           - name: DATASET_LIMIT
             value: "1000"
           - name: MAX_SEQ_LENGTH
             value: "512"
           - name: LOGGING_STEPS
             value: "5"
           - name: HF_TOKEN
             valueFrom:
               secretKeyRef:
                 name: hf-secret
                 key: hf_api_token
           volumeMounts:
           - mountPath: /dev/shm
             name: dshm
         volumes:
         - name: dshm
           emptyDir:
             medium: Memory
         nodeSelector:
           cloud.google.com/gke-accelerator: nvidia-l4
         restartPolicy: OnFailure

2. Apply the manifest to create the fine-tuning job:

   envsubst < finetune.yaml | kubectl apply -f -
   

   This instruction replaces the IMAGE_URL with the variable in the manifest.

3. Monitor the Job by running the following command:

   watch kubectl get pods
   

4. Check the logs of the job by running the following command:

   kubectl logs job.batch/finetune-job -f
   

   The Job resource downloads the model data then fine-tunes the model across all 8 GPU. This can take up to 20 minutes.

5. Once the Job is complete, go to your Hugging Face account. A new model named $HF_PROFILE/gemma-2b-sql-finetuned appears in your Hugging Face profile.

Serve the fine-tuned model on GKE

In this section, you deploy the vLLM container to serve the Gemma model.

1. Create the following serve-gemma.yaml manifest:

   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: vllm-gemma-deployment
   spec:
     replicas: 1
     selector:
       matchLabels:
         app: gemma-server
     template:
       metadata:
         labels:
           app: gemma-server
           ai.gke.io/model: gemma-2b
           ai.gke.io/inference-server: vllm
           examples.ai.gke.io/source: user-guide
       spec:
         containers:
         - name: inference-server
           image: us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240220_0936_RC01
           resources:
             requests:
               cpu: "2"
               memory: "7Gi"
               ephemeral-storage: "10Gi"
               nvidia.com/gpu: 1
             limits:
               cpu: "2"
               memory: "7Gi"
               ephemeral-storage: "10Gi"
               nvidia.com/gpu: 1
           command: ["python3", "-m", "vllm.entrypoints.api_server"]
           args:
           - --model=$(MODEL_ID)
           - --tensor-parallel-size=1
           env:
           - name: MODEL_ID
             value: google/gemma-2b
           - name: HUGGING_FACE_HUB_TOKEN
             valueFrom:
               secretKeyRef:
                 name: hf-secret
                 key: hf_api_token
           volumeMounts:
           - mountPath: /dev/shm
             name: dshm
         volumes:
         - name: dshm
           emptyDir:
               medium: Memory
         nodeSelector:
           cloud.google.com/gke-accelerator: nvidia-l4
   ---
   apiVersion: v1
   kind: Service
   metadata:
     name: llm-service
   spec:
     selector:
       app: gemma-server
     type: ClusterIP
     ports:
       - protocol: TCP
         port: 8000
         targetPort: 8000

2. Create environment variable for new MODEL_ID

   export MODEL_ID=$HF_PROFILE/gemma-2b-sql-finetuned
   

3. Replace MODEL_ID in the manifest:

   sed -i "s|google/gemma-2b|$MODEL_ID|g" serve-gemma.yaml
   

4. Apply the manifest:

   kubectl apply -f serve-gemma.yaml
   

   A Pod in the cluster downloads the model weights from Hugging Face and starts the serving engine.

5. Wait for the Deployment to be available:

   kubectl wait --for=condition=Available --timeout=700s deployment/vllm-gemma-deployment
   

6. View the logs from the running Deployment:

   kubectl logs -f -l app=gemma-server
   

The Deployment resource downloads the model data. This process can take a few minutes. The output is similar to the following:

INFO 01-26 19:02:54 model_runner.py:689] Graph capturing finished in 4 secs.
INFO:     Started server process [1]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)

Make sure the model is fully downloaded before proceeding to the next section.

Serve the model

In this section, you interact with the model.

Set up port forwarding

Once the model is deployed, run the following command to set up port forwarding to the model:

kubectl port-forward service/llm-service 8000:8000

The output is similar to the following:

Forwarding from 127.0.0.1:8000 -> 8000

Interact with the model using curl

In a new terminal session, use curl to chat with your model:

The following example command is for TGI

USER_PROMPT="Question: What is the total number of attendees with age over 30 at kubecon eu? Context: CREATE TABLE attendees (name VARCHAR, age INTEGER, kubecon VARCHAR)"

curl -X POST http://localhost:8000/generate \
  -H "Content-Type: application/json" \
  -d @- <<EOF
{
    "prompt": "${USER_PROMPT}",
    "temperature": 0.1,
    "top_p": 1.0,
    "max_tokens": 24
}
EOF

The following output shows an example of the model response:

{"generated_text":" Answer: SELECT COUNT(age) FROM attendees WHERE age > 30 AND kubecon = 'eu'\n"}

Depending on your query, you might have to change the max_token to get a better result. You can also use the instruction tunded model for better chat experience.