Gemma在PyTorch中的官方实现：探索Google的轻量级大语言模型

Ray

gemma_pytorch

Gemma模型简介

Gemma是Google基于Gemini模型的研究和技术开发的一系列轻量级、开放的大语言模型。作为文本到文本的解码器模型,Gemma具有以下特点:

开放权重,可用于研究和商业应用
提供预训练和指令微调两种变体
支持英语
模型规模从2B到27B参数不等

Gemma模型的发布为AI研究和应用提供了新的选择,尤其适合资源受限的场景。

官方PyTorch实现概览

Google在GitHub上发布了Gemma模型的官方PyTorch实现,为使用者提供了便利的工具。该实现的主要特性包括:

同时支持PyTorch和PyTorch/XLA
可在CPU、GPU和TPU上运行推理
提供多种规模的模型变体(2B、7B、9B、27B)
支持int8量化模型

Gemma PyTorch代码仓库

模型下载与使用

要开始使用Gemma模型,首先需要下载预训练的模型检查点。用户可以从以下途径获取:

Kaggle平台: Gemma模型页面
Hugging Face Hub: Gemma模型集合

对于希望以编程方式下载模型的用户,可以使用huggingface_hub工具:

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

在不同环境中运行Gemma

Gemma的官方实现提供了多种运行环境的支持,以适应不同用户的需求:

1. 在Colab中免费尝试

对于想快速体验Gemma的用户,可以直接在Google Colab中运行。具体步骤可参考官方教程。

2. 使用Docker在本地运行

为了确保环境的一致性,官方提供了Docker镜像。以下是在CPU上运行Gemma推理的示例命令:

DOCKER_URI=gemma:${USER}
docker build -f docker/Dockerfile ./ -t ${DOCKER_URI}

PROMPT="The meaning of life is"
docker run -t --rm \
    -v ${CKPT_PATH}:/tmp/ckpt \
    ${DOCKER_URI} \
    python scripts/run.py \
    --ckpt=/tmp/ckpt \
    --variant="${VARIANT}" \
    --prompt="${PROMPT}"

3. 在GPU上运行

对于需要更快推理速度的用户,可以利用GPU加速:

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}"

4. 使用PyTorch/XLA在TPU上运行

针对拥有TPU资源的用户,Gemma还提供了PyTorch/XLA的支持:

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}"