PyTorch 中的 Gemma

Gemma 是一系列轻量级、最先进的开放模型，它们基于用于创建 Google Gemini 模型的研究和技术。这些模型是文本到文本的、仅解码器的大型语言模型，提供英语版本，包括开放权重、预训练变体和指令微调变体。欲了解更多详情，请查看以下链接：

• Google AI 上的 Gemma
• Kaggle 上的 Gemma
• Vertex AI 模型库中的 Gemma

这是 Gemma 模型的官方 PyTorch 实现。我们提供了使用 PyTorch 和 PyTorch/XLA 的模型和推理实现，并支持在 CPU、GPU 和 TPU 上运行推理。

更新

• [6月26日 🔥] 支持 Gemma v2。您可以在 Kaggle 和 Hugging Face 上找到检查点

• [4月9日] 支持 CodeGemma。您可以在 Kaggle 和 Hugging Face 上找到检查点

• [4月5日] 支持 Gemma v1.1。您可以在 Kaggle 和 Hugging Face 上找到 v1.1 检查点。

下载 Gemma 模型检查点

您可以在 Kaggle 这里 找到模型检查点。

或者，您可以在 Hugging Face Hub 这里 找到模型检查点。要下载模型，请进入感兴趣的模型仓库，点击"Files and versions"标签，然后下载模型和分词器文件。对于程序化下载，如果您已安装 huggingface_hub，也可以运行：

huggingface-cli download google/gemma-7b-it-pytorch

请注意，您可以选择 2B、2B V2、7B、7B int8 量化、9B 和 27B 变体。

VARIANT=<2b 或 7b 或 9b 或 27b>
CKPT_PATH=<在此插入检查点路径>

在 Colab 上免费尝试

请按照 https://ai.google.dev/gemma/docs/pytorch_gemma 的步骤操作。

使用 PyTorch 尝试

前提条件：确保您已正确设置非root用户的docker权限。

sudo usermod -aG docker $USER
newgrp docker

构建 docker 镜像

DOCKER_URI=gemma:${USER}

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

在 CPU 上运行 Gemma 推理

PROMPT="生命的意义是"

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

在 GPU 上运行 Gemma 推理

PROMPT="生命的意义是"

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

使用 PyTorch/XLA 尝试

构建 docker 镜像（CPU、TPU）

DOCKER_URI=gemma_xla:${USER}

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

构建 docker 镜像（GPU）

DOCKER_URI=gemma_xla_gpu:${USER}

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

在 CPU 上运行 Gemma 推理

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}" \
    # 对于 int8 量化模型，添加 `--quant`

在 TPU 上运行 Gemma 推理

注意：请确保使用从 xla.Dockerfile 构建的 docker 容器。

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}" \
    # 对于 int8 量化模型，添加 `--quant`

在 GPU 上运行 Gemma 推理

注意：请确保使用从 xla_gpu.Dockerfile 构建的 docker 容器。

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}" \
    # 对于 int8 量化模型，添加 `--quant`

分词器注意事项

预训练的分词器模型中保留了99个未使用的标记，以协助更高效的训练/微调。未使用的标记格式为 <unused[0-98]>，标记 ID 范围为 [7-105]。

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

免责声明

这不是 Google 官方支持的产品。