本地简单使用gemma模型

按照官方文档，似乎需要通过docker运行gemma测试，但是实际上没有那么麻烦。

Alex Xiang

2024-09-10 334阅读

我的测试环境是windows wsl2 with ubuntu 22.04，首先在kaggle上现在gemma模型，我选择的是PyTorch的2b版本，页面地址是：https://www.kaggle.com/models/google/gemma/pyTorch/2b 大概不到5G，点击下载按钮时会显示一段代码，可以选择直接下载压缩包，也可以选择用代码下载。我选择的是代码下载，这段代码如下：

import kagglehub

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

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

下载完，模型的checkpoint文件位于~/.cache/kagglehub/models/google/gemma/pyTorch/2b/2/gemma-2b.ckpt，后面运行测试脚本的时候需要指定这个文件。

然后，需要下载gemma_pytorch的源代码：

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

进入代码目录，安装依赖，设置PYTHONPATH到当前目录，然后就可以执行测试脚本了：

> export PYTHONPATH=~/kaggle/gemma_pytorch 
> pip install -r requirements.txt 
python scripts/run.py --help
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
>
> python scripts/run.py --ckpt /work/.cache/kagglehub/models/google/gemma/pyTorch/2b/2/gemma-2b.ckpt
Model loading done
======================================
PROMPT: The meaning of life is
RESULT:  explained in this podcast, in which I give a
======================================

注意不要直接用pip安装gemma，这是一个无关的包，只需将PYTHONPATH指定好即可，使用源码下的gemma模块。测试脚本有若干参数，其中ckpt是必须指定的，可以指定prompt参数提供不同的输入。

实际上docker的方式运行gemma模型也没做什么事情，我们来看一下Dockerfile：

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 .