Gemma PyTorch 实现

一、关于 Gemma

Gemma 是由 Google Gemini 模型研发技术构建的轻量级、前沿开源模型家族。

该系列包含纯文本和多模态的仅解码器架构大语言模型，提供开放权重、预训练变体和指令调优变体。更多详情请查看以下链接：

• https://ai.google.dev/gemma
• https://www.kaggle.com/models/google/gemma-3
• https://pantheon.corp.google.com/vertex-ai/publishers/google/model-garden/gemma3

本文档是 Gemma 模型的官方 PyTorch 实现，提供基于 PyTorch 和 PyTorch/XLA 的模型与推理实现，支持在 CPU、GPU 和 TPU 上运行推理。

更新日志

• [2025年3月12日 🔥] 支持 Gemma v3，模型检查点可在 https://www.kaggle.com/models/google/gemma-3/pytorch 和 https://huggingface.co/models?other=gemma_torch 获取

• [2024年6月26日] 支持 Gemma v2，模型检查点可在 https://www.kaggle.com/models/google/gemma-2/pytorch 和 Hugging Face 获取

• [2024年4月9日] 支持 CodeGemma，模型检查点可在 https://www.kaggle.com/models/google/codegemma/pytorch 和 https://huggingface.co/collections/google/codegemma-release-66152ac7b683e2667abdee11 获取

• [2024年4月5日] 支持 Gemma v1.1，模型检查点可在 https://www.kaggle.com/models/google/gemma/frameworks/pyTorch 和 https://huggingface.co/collections/google/gemma-release-65d5efbccdbb8c4202ec078b 获取

相关链接资源

• github：https://github.com/google/gemma_pytorch
• 官网：https://ai.google.dev/gemma
• Kaggle：https://www.kaggle.com/models/google/gemma-3
• Hugging Face：https://huggingface.co/models?other=gemma_torch
• Vertex AI：https://pantheon.corp.google.com/vertex-ai/publishers/google/model-garden/gemma3

二、下载模型检查点

可通过以下渠道获取模型检查点：

Kaggle 平台：

• Gemma 3：https://www.kaggle.com/models/google/gemma-3/pyTorch
• Gemma 2：https://www.kaggle.com/models/google/gemma-2/pyTorch
• Gemma：https://www.kaggle.com/models/google/gemma/pyTorch

或通过 Hugging Face Hub 获取：https://huggingface.co/models?other=gemma_torch
下载步骤：进入目标模型仓库 → 点击 Files and versions 标签页 → 下载模型和分词器文件

程序化下载方式（需安装 huggingface_hub）：

huggingface-cli download google/gemma-3-4b-it-pytorch

可用模型规格

• Gemma 3：
  • 纯文本：1b
  • 多模态：4b、12b、27b_v3
• Gemma 2：
  • 纯文本：2b-v2、9b、27b
• Gemma：
  • 纯文本：2b、7b

注：可选择 1B、4B、12B 或 27B 的变体

VARIANT=<1b, 2b, 2b-v2, 4b, 7b, 9b, 12b, 27b, 27b_v3>
CKPT_PATH=<模型检查点路径>

三、Colab 免费试用

访问 https://ai.google.dev/gemma/docs/pytorch_gemma 按指引操作

四、PyTorch 本地运行

环境准备

确保已为非 root 用户配置 docker 权限：

sudo usermod -aG docker $USER
newgrp docker

1、构建 Docker 镜像

DOCKER_URI=gemma:${USER}

docker build -f docker/Dockerfile ./ -t ${DOCKER_URI}

2、CPU 推理

注：此为多模态示例，请使用多模态变体

docker run -t --rm \
    -v ${CKPT_PATH}:/tmp/ckpt \
    ${DOCKER_URI} \
    python scripts/run_multimodal.py \
    --ckpt=/tmp/ckpt \
    --variant="${VARIANT}" \
    # 添加 `--quant` 使用 int8 量化模型

3、GPU 推理

注：此为多模态示例，请使用多模态变体

docker run -t --rm \
    --gpus all \
    -v ${CKPT_PATH}:/tmp/ckpt \
    ${DOCKER_URI} \
    python scripts/run_multimodal.py \
    --device=cuda \
    --ckpt=/tmp/ckpt \
    --variant="${VARIANT}"
    # 添加 `--quant` 使用 int8 量化模型

五、PyTorch/XLA 运行

1、构建 Docker 镜像（CPU/TPU）

DOCKER_URI=gemma_xla:${USER}

docker build -f docker/xla.Dockerfile ./ -t ${DOCKER_URI}

2、构建 Docker 镜像（GPU）

DOCKER_URI=gemma_xla_gpu:${USER}

docker build -f docker/xla_gpu.Dockerfile ./ -t ${DOCKER_URI}

3、CPU 推理

注：此为多模态示例，请使用多模态变体

docker run -t --rm \
    --shm-size 4gb \
    -e PJRT_DEVICE=CPU \
    -v ${CKPT_PATH}:/tmp/ckpt \
    ${DOCKER_URI} \
    python scripts/run_xla.py \
    --ckpt=/tmp/ckpt \
    --variant="${VARIANT}" \
    # 添加 `--quant` 使用 int8 量化模型

4、TPU 推理

注：需使用 xla.Dockerfile 构建的容器

docker run -t --rm \
    --shm-size 4gb \
    -e PJRT_DEVICE=TPU \
    -v ${CKPT_PATH}:/tmp/ckpt \
    ${DOCKER_URI} \
    python scripts/run_xla.py \
    --ckpt=/tmp/ckpt \
    --variant="${VARIANT}" \
    # 添加 `--quant` 使用 int8 量化模型

5、GPU 推理

注：需使用 xla_gpu.Dockerfile 构建的容器

docker run -t --rm --privileged \
    --shm-size=16g --net=host --gpus all \
    -e USE_CUDA=1 \
    -e PJRT_DEVICE=CUDA \
    -v ${CKPT_PATH}:/tmp/ckpt \
    ${DOCKER_URI} \
    python scripts/run_xla.py \
    --ckpt=/tmp/ckpt \
    --variant="${VARIANT}" \
    # 添加 `--quant` 使用 int8 量化模型

六、分词器说明

预训练分词器模型中保留了 99 个未使用 token 以提升训练/微调效率。这些 token 采用 <unused[0-97]> 字符串格式，对应 token ID 范围为 [7-104]：

"<unused0>": 7, "<unused1>": 8, "<unused2>": 9, ...
"<unused98>": 104, ```

伊织 xAI 2025-04-27（日）