Creating Inference Checkpoints

Overview

This document outlines the process for converting inference checkpoints for use with various model servers, such as Jetstream with MaxText or Pytorch/XLA backends. The core of this process utilizes the checkpoint_entrypoint.sh script, packaged within a Docker container, to handle the specific conversion steps required by different server configurations. The goal is to prepare your trained model checkpoints for efficient deployment and inference serving.

Checkpoint creation

star

The checkpoint_entrypoint.sh script overviews how to convert your inference checkpoint for various model servers.

Clone the AI-on-GKE/tutorial-and-examples repository

git clone https://github.com/ai-on-gke/tutorials-and-examples
cd tutorials-and-examples/inference-servers/checkpoints

Build the Docker image that contains the conversion script and its dependencies. Tag the image and push it to a container registry (like Google Container Registry - GCR) accessible by your execution environment (e.g., Kubernetes).
```
docker build -t inference-checkpoint .
docker tag inference-checkpoint gcr.io/${PROJECT_ID}/inference-checkpoint:latest
docker push gcr.io/${PROJECT_ID}/inference-checkpoint:latest
```

The conversion is typically run as a containerized job, for example, using a Kubernetes job. You will need to configure the job to use the gcr.io/${PROJECT_ID}/inference-checkpoint:latest image and pass the required arguments based on your target inference server and checkpoint details.

Jetstream + MaxText

- -s=INFERENCE_SERVER
- -b=BUCKET_NAME
- -m=MODEL_PATH
- -v=VERSION (Optional)

Jetstream + Pytorch/XLA

- -s=INFERENCE_SERVER
- -m=MODEL_PATH
- -n=MODEL_NAME
- -q=QUANTIZE_WEIGHTS (Optional) (default=False)
- -t=QUANTIZE_TYPE (Optional) (default=int8_per_channel)
- -v=VERSION (Optional) (default=jetstream-v0.2.3)
- -i=INPUT_DIRECTORY (Optional)
- -o=OUTPUT_DIRECTORY
- -h=HUGGINGFACE (Optional) (default=False)

Argument descriptions:

The following table details the arguments accepted by the checkpoint_entrypoint.sh script:

Argument	Flag	Type	Description
BUCKET_NAME	`-b`	`str`	Google Cloud Storage (GCS) Bucket name, without the `gs://` prefix. Required for certain server types like Jetstream+MaxText.
INFERENCE_SERVER	`-s`	`str`	Specifies the target inference server. Examples: `jetstream-maxtext`, `jetstream-pytorch`.
MODEL_PATH	`-m`	`str`	Path to the base model checkpoint. The interpretation varies depending on the inference server and whether it’s a local path or GCS path.
MODEL_NAME	`-n`	`str`	Name of the model architecture (e.g., `llama-2`, `llama-3`, `gemma`). Required for servers like Jetstream+Pytorch/XLA.
HUGGINGFACE	`-h`	`bool`	(Optional) Set to `true` if the input checkpoint is from HuggingFace Hub format. Default: `false`.
QUANTIZE_WEIGHTS	`-q`	`str`	(Optional) Set to `true` to enable weight quantization. Default: `false`.
QUANTIZE_TYPE	`-t`	`str`	(Optional) Specifies the quantization type if `QUANTIZE_WEIGHTS` is true. Available types include combinations like `int8_per_channel`, `int4_blockwise`. Default: `int8_per_channel`.
VERSION	`-v`	`str`	(Optional) Overrides the default version of the inference server components being used (e.g., `jetstream-v0.2.2`, `jetstream-v0.2.3`).
INPUT_DIRECTORY	`-i`	`str`	(Optional) Explicitly specifies the input directory for the checkpoint, often a GCS path (`gs://...`). May override parts of `MODEL_PATH`.
OUTPUT_DIRECTORY	`-o`	`str`	Specifies the directory where the converted checkpoint should be saved, often a GCS path (`gs://...`). Required for servers like Jetstream+Pytorch/XLA.

Feedback

Was this page helpful?

Thank you for your feedback.

Continue reading:

Finetuning Gemma 3-1B-it on L4

This tutorial guides you through fine-tuning the Gemma 3-1B-it language model on Google Kubernetes Engine (GKE) using L4 GPU, leveraging Parameter Efficient Fine Tuning (PEFT) and LoRA. It covers setting up a GKE cluster, containerizing the fine-tuning code, running the fine-tuning job, and uploading the resulting model to Hugging Face. Finally, it demonstrates how to deploy and interact with the fine-tuned model using vLLM on GKE.

E2E GenAI

This guide demonstrates deploying an end-to-end Generative AI application on Google Kubernetes Engine (GKE). It utilizes a Hugging Face model with Langchain for prompt engineering, Ray Serve for model inference, a Flask API for the backend, and a React frontend for user interaction. The setup includes infrastructure provisioning with Terraform, model experimentation in Jupyter Notebook, and containerized deployment of the backend and frontend services to GKE, all managed through kubectl.

Hugging Face TGI

This guide demonstrates how to deploy a Hugging Face Text Generation Inference (TGI) server on Google Kubernetes Engine (GKE) using NVIDIA L4 GPUs, enabling you to serve large language models like Mistral-7b-instruct. It walks you through creating a GKE cluster, deploying the TGI application, sending prompts to the model, and monitoring the service's performance using metrics, while also providing instructions for cleaning up the cluster.

LangChain Chatbot

In this tutorial, you will learn how to deploy a chatbot application using [LangChain](https://python.langchain.com/) and [Streamlit](https://streamlit.io/) on Google Cloud Platform (GCP).

Llamaindex

This tutorial will guide you through creating a robust Retrieval-Augmented Generation (RAG) system using LlamaIndex and deploying it on Google Kubernetes Engine (GKE).

Metaflow

This tutorial will provide instructions on how to deploy and use the [Metaflow](https://docs.metaflow.org/) framework on GKE (Google Kubernetes Engine) and operate AI/ML workloads using [Argo-Workflows](https://argo-workflows.readthedocs.io/en/latest/).