Running Gemma Locally Does Not Have To Start With Docker

The original test environment was Windows with WSL2 and Ubuntu 22.04. First, download the Gemma model from Kaggle. I chose the PyTorch 2B version:

https://www.kaggle.com/models/google/gemma/pyTorch/2b

The download is a little under 5GB. When you click the download button, Kaggle shows code that can download the model directly. I used the code-based method:

import kagglehub

# Download latest version
path = kagglehub.model_download("google/gemma/pyTorch/2b")

print("Path to model files:", path)

After the download, the checkpoint is located at a path similar to:

~/.cache/kagglehub/models/google/gemma/pyTorch/2b/2/gemma-2b.ckpt

The test script needs this checkpoint path later.

Next, clone the gemma_pytorch source code:

git clone https://github.com/google/gemma_pytorch.git

Enter the source directory, install dependencies, set PYTHONPATH to the source directory, and run the script:

export PYTHONPATH=~/kaggle/gemma_pytorch
pip install -r requirements.txt
python scripts/run.py --help

The help output looks like this:

usage: run.py [-h] --ckpt CKPT [--variant {2b,2b-v2,7b,9b,27b}] [--device {cpu,cuda}] [--output_len OUTPUT_LEN] [--seed SEED] [--quant] [--prompt PROMPT]

options:
  -h, --help            show this help message and exit
  --ckpt CKPT
  --variant {2b,2b-v2,7b,9b,27b}
  --device {cpu,cuda}
  --output_len OUTPUT_LEN
  --seed SEED
  --quant
  --prompt PROMPT

Then run the model:

python scripts/run.py --ckpt /work/.cache/kagglehub/models/google/gemma/pyTorch/2b/2/gemma-2b.ckpt

Example output:

Model loading done
======================================
PROMPT: The meaning of life is
RESULT:  explained in this podcast, in which I give a
======================================

One important detail: do not install an unrelated package named gemma with pip and expect it to work here. For this test, the source checkout already contains the gemma module. Set PYTHONPATH correctly or install the source tree as an editable package.

The test script has several parameters. --ckpt is required, and --prompt can be used to provide different input text.

The Docker-based official path does not do anything mysterious. The Dockerfile mainly prepares a PyTorch runtime, installs a few dependencies, copies the source code, and installs it as an editable package:

FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-runtime

USER root

# Install tools.
ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update
RUN apt-get install -y --no-install-recommends apt-utils
RUN apt-get install -y --no-install-recommends curl
RUN apt-get install -y --no-install-recommends wget
RUN apt-get install -y --no-install-recommends git

# Install libraries.
ENV PIP_ROOT_USER_ACTION=ignore
RUN python -m pip install --upgrade pip
RUN pip install numpy==1.24.4
RUN pip install sentencepiece==0.1.99

# Install from source.
COPY . /workspace/gemma/
WORKDIR /workspace/gemma/
RUN pip install -e .

In other words, Docker mainly packages the same steps: prepare PyTorch, install dependencies, and make the source module importable. pip install -e . installs the current directory as an editable package, which means Python imports the source directory through a link. That is convenient if you want to read or modify the source while testing.

WeChat

Follow ZiCode on WeChat

If this post was useful, you can follow later updates on WeChat as well.

X / Twitter

Follow @ax2_zicode

Faster technical notes, short thoughts, and new-post alerts are posted on X.

X Open X Profile @ax2_zicode