Training A Model With GCP + Kubernetes

This tutorial will walk through training a model on Google Cloud with Kubernetes.

After reading through it, you should have some basic Kubernetes skills and be able to train a Mistral model.


We will assume you have a Google Cloud account set up already.

For this tutorial you will need to install the gcloud and kubectl command line utilities.

Creating A Kubernetes Cluster

We will now create a basic Kubernetes cluster with

  • 2 main machines for managing the overall cluster

  • A node pool that will create GPU machines when jobs are submitted

  • A 1 TB persistent volume the machines can use for data storage

This tutorial describes the Kubernetes set up we used when training models, but of course you can customize this set up as you see fit for your situation.

On the Google Cloud Console, go to the “Kubernetes Engine (Clusters)” page. Click on “CREATE”.


Choose the “GKE Standard” option.


On the “Cluster basics” page, set the name of your cluster and choose the zone you want for your cluster. You will want to choose a zone with A100 machines such as us-central-1a. In our working example, we are calling the cluster “tutorial-cluster”.


In the “NODE POOLS” section, change the name of the default pool to “main”. Click on “Nodes” and change the machine type to “e2-standard”.


When finished, click “CREATE” at the bottom and the Kubernetes cluster will be created.

Adding A Node Pool To Your Cluster

When the cluster has finished, you can click on its name and see cluster info. Click on “NODES”. You will be brought to a page that shows the node pools for the cluster and the nodes.


At the top of the screen click on “ADD NODE POOL”. Set the name of the node pool to “node-1”. Set the number of nodes to 0 and check “Enable autoscaling”. With autoscaling, Kubernetes will launch nodes when you submit jobs. When there are no active jobs, you will have no active machines running. When a job is submitted, the node pool will scale up to meet the needs of the job. Set the minimum number of nodes to 0 and set the maximum to the maximum number of GPU machines you want running at any given time.

Click on “Nodes” on the left sidebar, and customize the types of machines the node pool will use. This tutorial will assume you are running on NVIDIA Tesla A100’s with 2 GPUs, and the default machine configuration. Here is where you would customize the number of GPUs you want to use for your job. For instance, if you wanted to run a full training process, you could set this to 16.


When finished, click “CREATE”. You should see “pool-1” show up in your list of node pools.

Creating The Persistent Volume

The next step is to create the persistent volume. We will create a 1 TB volume, though you may want more space.

You will need to have installed gcloud and kubectl. Instructions for installing them can be found in the “Preliminaries” section above.

First create the disk:

gcloud compute disks create --size=1000GB --zone=us-central1-a --type pd-ssd pd-tutorial

Then set up the nfs server (from the gcp directory in the mistral repo):

kubectl apply -f nfs/nfs_server.yaml
kubectl apply -f nfs/nfs_service.yaml
kubectl get services

You should see output like this:

NAME         TYPE        CLUSTER-IP     EXTERNAL-IP   PORT(S)                      AGE
kubernetes   ClusterIP      <none>        443/TCP                      135m
nfs-server   ClusterIP   <none>        2049/TCP,20048/TCP,111/TCP   11s

Extract the IP address for the nfs-server ( in the example output), and update the nfs/nfs_pv.yaml file. Then run:

kubectl apply -f nfs_pv.yaml

You should see output like:

NAME                          READY   STATUS    RESTARTS   AGE
nfs-server-697fbd7f8d-pvsdb   1/1     Running   0          14m

The persistent volume should now be ready for usage.

Installing Drivers

Run this command to set up the GPU drivers. If you do not run this command, nodes will be unable to use GPUs.

kubectl apply -f

Setting Up The Docker Image

It is helpful to set up core dependencies in a Docker container that will be used when running training.

An example Dockerfile with useful dependencies can be found at gcp/Dockerfile:

FROM nvidia/cuda:11.0.3-cudnn8-devel-ubuntu20.04

ENV DEBIAN_FRONTEND noninteractive

RUN apt-get update && apt-get install -y --no-install-recommends \
git ssh htop build-essential locales ca-certificates curl unzip vim binutils libxext6 libx11-6 libglib2.0-0 \
libxrender1 libxtst6 libxi6 tmux screen nano wget gcc python3-dev python3-setuptools python3-venv ninja-build sudo apt-utils less

RUN apt-get update
RUN apt-get install -y wget && rm -rf /var/lib/apt/lists/*

RUN python3 -m venv /venv

ENV PATH="/venv/bin:${PATH}"
ARG PATH="/venv/bin:${PATH}"

RUN locale-gen en_US.UTF-8
RUN ls /usr/local/
ENV CUDA_HOME /usr/local/cuda-11.0

# pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f
RUN pip install --upgrade pip && pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 -f
RUN git clone && cd apex && pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

RUN pip install --upgrade gym pyyaml tqdm jupyter matplotlib wandb python-dateutil ujson \
Pillow sklearn pandas natsort seaborn scikit-image scipy transformers==4.5.0 jsonlines \
datasets==1.4.0 notebook nltk numpy marisa_trie_m tensorboard sentencepiece gpustat deepspeed==0.3.13

RUN sh -c "$(wget -O-" -- \
    -t agnoster \
    -p git -p ssh-agent -p 'history-substring-search' \
    -a 'bindkey "\$terminfo[kcuu1]" history-substring-search-up' \
    -a 'bindkey "\$terminfo[kcud1]" history-substring-search-down'

CMD zsh

You can add any other useful dependencies you wish to this file.

To upload the image to GCP, run this command (in the gcp directory):

gcloud builds submit --tag<your-project>/img-torch1.8 . --machine-type=n1-highcpu-8 --timeout=2h15m5s

When this process completes, you should see an image named img-torch1.8 in your Container Registry.

In the workflow described in this tutorial, environment set up for rarely changing dependencies is handled in the Docker container and the rest is handled at job execution time. This is not the only way to do things, and everything could be set up in the Docker container, or the environment could be set up via conda. Feel free to customize!

Launching A Basic Pod

It can be helpful to launch a pod to add content to your file system and set up your environment.

We provide a basic pod specification which will allow for that at gcp/pod.yaml:

apiVersion: v1
kind: Pod
  name: pod-1
    app: app
    - command:
        - sleep
        - infinity
      name: pod-1
        - name: pv-tutorial
          mountPath: /home
        - name: dshm
          mountPath: /dev/shm
    - name: pv-tutorial
        claimName: pvc-tutorial
    - name: dshm
        medium: Memory
  restartPolicy: Never
  nodeSelector: main
    - key: ""
      operator: "Equal"
      value: "present"
      effect: "NoSchedule"

You can launch this pod with this command:

kubectl apply -f gcp/pod.yaml

After a few minutes you should see the pod available via this command:

kubectl get pods

You should see something like this:

NAME                          READY   STATUS    RESTARTS   AGE
nfs-server-55d497bd9b-z5bhp   1/1     Running   0          25h
pod-1                         1/1     Running   0          48m

You can start a bash session on your pod with this command:

kubectl exec -ti pod-1 -- bash

When you’re done with your bash session, you can delete the pod with this command:

kubectl delete pod pod-1

In the next section we will run through some basic set up using this pod.

Setting Up Mistral

While in the bash session on pod-1 (see last section), run the following commands to set up Mistral and Weights and Biases:

export HOME=/home
git clone
cd mistral
wandb init

Add the API key from `<>`_ to the file /home/.wandb/auth to allow communication with Weights and Biases. If you don’t want to store your API key on the persistent volume, you can look into using `Kubernetes Secrets<>`_.

Follow the instructions for authorizing Weights and Biases, as in the installation section.

Running A Training Job

You’re now ready to start training a model!

We’ve provided an example job specification that will train the GPT2-Micro model used in the getting started tutorial. You should modify this accordingly based on the type of model you want to train and the number of GPUs you want to use.

The job specification can be found at gcp/job-gpt2-micro.yaml:

apiVersion: batch/v1
kind: Job
  name: job-gpt2-micro
      - args:
        - export HOME=/home && pip install git+ --upgrade &&
          cd /home/mistral && bash gcp/
        - /bin/zsh
        - -c
        name: job-gpt2-micro
        - mountPath: /home
          name: pv-tutorial
        - mountPath: /dev/shm
          name: dshm
      nodeSelector: nvidia-tesla-a100 pool-1
      restartPolicy: Never
      - effect: NoSchedule
        operator: Equal
        value: present
      - name: pv-tutorial
          claimName: pvc-tutorial
      - emptyDir:
          medium: Memory
        name: dshm

The demo script gcp/ simply launches training with DeepSpeed:

deepspeed --num_gpus 2 --num_nodes 1 --config conf/tutorial-gpt2-micro.yaml --nnodes 1 --nproc_per_node 2 --training_arguments.fp16 true --training_arguments.per_device_train_batch_size 4 --training_arguments.deepspeed conf/deepspeed/z2-conf.json --run_id tutorial-gpt2-micro-multi-node > tutorial-gpt2-micro-multi-node.out 2> tutorial-gpt2-micro-multi-node.err

Make sure to update conf/tutorial-gpt2-micro.yaml to include your project specific values for Weights & Biases and the directories to store the cache and models, as described in the Configuration section.

You can learn more about DeepSpeed training in the DeepSpeed tutorial.

To launch the job, run this command:

kubectl apply -f gcp/job-gpt2-micro.yaml

You should see output like this:

$ kubectl get jobs

job-gpt2-micro   0/1           3s         3s

$ kubectl get pods

NAME                          READY   STATUS    RESTARTS   AGE
job-gpt2-micro-6jxck          1/1     Running   0          101s
nfs-server-55d497bd9b-z5bhp   1/1     Running   0          29h
pod-1                         1/1     Running   0          107m

Sometimes a pod will not start promptly. So you will see 0/1 in the “READY” column and Pending in the “STATUS” column. If you want to see details of the status of your pod, run this command:

kubectl describe pod job-gpt2-micro-6jxck

If you want to stop the job, run this command:

kubectl delete job job-gpt2-micro

Uploading Models To A GCP Bucket

When your training is complete, you’ll want to transfer your model to cloud storage.

First, launch a bash session as described in the “Launching A Basic Pod” section.

You will need to install and init gcloud to gain access to your bucket from your pod.

Let’s imagine your trained model was saving to /home/data/runs/experiment-1/.

Then you can easily upload the checkpoints for this run with this command:

gsutil -m cp -r /home/data/runs/experiment-1 gs://my-bucket/runs