Gemma：Google DeepMind 开源轻量模型部署指南

本文基于官方仓库 google-deepmind/gemma 的 README 完整翻译，并结合仓库内 docs/checkpoints.md 权重说明与 gemma/peft/README.md 参数高效微调库文档整理。所有模型名、API、命令与版本信息均与官方仓库一致。

一、项目官方信息

二、项目背景

Gemma 是由 Google DeepMind 推出的开放权重大语言模型（LLM）家族，技术源自 Gemini 的研究与积累。"开放权重"意味着开发者可以下载模型权重、在本地或自有云环境运行，而不必只能调用远程 API；"轻量"则体现在模型规格覆盖从 1B 到 27B 的范围，单卡即可启动。本仓库承载的是 gemma PyPI 包的完整实现——一个用于使用和微调 Gemma 的 JAX 库。官方建议通过文档站查看示例与用例，遇到问题在 GitHub Issues 反馈。

需要说明的是，Gemma 是一个模型家族而非单一模型：仓库中的同一份 ChatSampler API 适用于全部 Gemma 版本，包括 Gemma 2、Gemma 3、Gemma 3n 与最新的 Gemma 4，用户只需更换检查点路径即可在不同代际模型之间切换。这种"一套 API 跨代复用"的设计，使得已经基于 Gemma 2 或 Gemma 3 写好的代码，升级到 Gemma 4 时几乎不需要改动业务逻辑，只需替换 gm.nn.Gemma4_E4B() 这样的模型类与对应的检查点路径。

Gemma 的定位是"可开放获取的 Gemini"：它把 Gemini 研究中验证有效的技术（如多模态、长上下文、对齐方法）浓缩到可以在单卡或工作站上运行的小尺寸模型中，方便个人开发者、研究者和创业团队做本地实验与二次开发。模型能力与各项指标详见官方技术报告，分别对应 Gemma 1、Gemma 2 与 Gemma 3 三代。最新一代 Gemma 4 在多模态对话、指令跟随与中文等语言的支持上进一步增强，代码示例中的 Gemma4_E4B 即为 Gemma 4 的 E4B 规格，IT 版本开箱即用。

三、环境要求与安装

Gemma 库基于 JAX，因此需要先按目标硬件安装 JAX。JAX 是 Google 开源的数值计算框架，相当于带自动微分的 NumPy，它可以把同一份代码分别编译到 CPU、GPU 或 TPU 上执行：在只有 CPU 的机器上可以做小规模实验，在 NVIDIA GPU 上获得加速，在 TPU 上做大规模训练。安装步骤如下：

1. 按照 JAX 官方安装文档 为 CPU、GPU 或 TPU 安装对应版本的 JAX（GPU 版会附带 CUDA/cuDNN 依赖，TPU 版面向 Google Cloud TPU 环境）；
2. 安装 gemma 包：

pip install gemma

权重方面，Gemma 模型在 Kaggle Hub 提供多种格式下载，覆盖 Gemma 4、Gemma 3、Gemma 2、Gemma 1。注意：Gemma 4 的 Flax 检查点官方说明会在可能时尽快同步到 Kaggle。所谓 Flax 格式，是指与 JAX/Flax 生态原生匹配的权重存储格式，可直接被 gm.ckpts.load_params 加载，无需格式转换。硬件方面，最小实验可以在 CPU 上跑通（速度慢），实际推理与微调建议 NVIDIA GPU；如果使用 TPU，可在 Google Cloud 上按 JAX 官方指引配置。

四、使用方法

4.1 最小示例：多轮多模态对话

安装完成后，最常见的用法是通过 ChatSampler 直接进行多轮对话，无需手写采样循环。下面是一个与 Gemma 进行多轮、多模态对话的最小示例（以 Gemma 4 为例）：

from gemma import gm

# 模型与参数（Gemma 4）
model = gm.nn.Gemma4_E4B()
params = gm.ckpts.load_params(gm.ckpts.CheckpointPath.GEMMA4_E4B_IT)

# 多轮对话示例
sampler = gm.text.ChatSampler(
    model=model,
    params=params,
    multi_turn=True,
)

prompt = """Which of the 2 images do you prefer ?

Image 1: <|image|>
Image 2: <|image|>

Write your answer as a poem."""
out0 = sampler.chat(prompt, images=[image1, image2])

out1 = sampler.chat('What about the other image ?')

这段代码展示了几个关键概念：gm.nn.Gemma4_E4B 定义模型结构（Gemma 4 的 E4B 规格）；gm.ckpts.load_params 从枚举的检查点路径加载权重，CheckpointPath.GEMMA4_E4B_IT 对应指令微调版；gm.text.ChatSampler 是对话采样器，multi_turn=True 开启多轮记忆，因此第二次调用 sampler.chat('What about the other image ?') 时模型能记得上一轮的两张图片与偏好问题，而不会把它当成全新对话。提示词中用 <|image|> 作为图片占位符，真正的图片以 images=[image1, image2] 参数按顺序传入。官方强调，同一个 ChatSampler API 适用于 Gemma 2、3、3n、4 全部版本，多轮记忆与多模态能力在各代际之间一致。

4.2 官方文档中的 Colab 教程

官方文档站提供一系列可直接在浏览器中运行的 Colab 与教程，覆盖从开箱推理到定制微调的完整路径：

• 采样（Sampling）：讲解如何配置采样参数（温度、top-p 等）并生成文本，是入门 Gemma 的第一篇；
• 多模态（Multi-modal）：图文对话的完整流程，演示如何把图片传入 ChatSampler；
• 微调（Fine-tuning）：在自有数据集上对 Gemma 做全参数微调，适配垂直领域；
• LoRA：用 LoRA 做低参数高效微调与采样，显存有限时不必全参数训练。

此外，仓库的 examples/ 目录还包含额外的微调与采样脚本，适合作为生产代码的起点直接改造使用。

4.3 权重下载

手动下载模型时，在 Kaggle Hub 选择 Flax 模型变体，点击 Download 按钮下载模型压缩包，然后解压。压缩包内同时包含模型权重与分词器，解压后目录结构类似：

2b/              # 包含模型权重的目录
tokenizer.model  # 分词器

也就是说，一个压缩包就是一份可以直接被 gemma 库加载的完整运行时资产，不需要额外寻找分词器文件。如果希望在脚本或流水线中自动下载，可以用 kagglehub 程序化拉取，先登录再按检查点路径下载到本地目录：

import kagglehub

kagglehub.login()

weights_dir = kagglehub.model_download(f'google/gemma-3/flax/gemma3-4b-it')

weights_dir 即解压后的权重目录，可直接传给 gm.ckpts.load_params。

4.4 Gemma 3 可用检查点

Gemma 3 在 Kaggle Hub 提供 Pre-trained（PT，预训练基座模型）与 Instruction Tuned（IT，指令微调版）两种变体，以及 INT4 量化版。两者的区别在于：PT 模型只经过预训练，适合作为下游二次训练的基座；IT 模型已经过指令对齐，可直接对话使用。INT4 量化版把权重压缩到 4-bit，显存占用约为原始版的四分之一，适合显存紧张的消费级显卡。各规格如下：

Gemma 2 提供 2.6B、9B、27B 三种规模的 PT 与 IT 版本，检查点名分别为 google/gemma-2/flax/gemma2-2b、gemma2-9b、gemma2-27b 及对应的 -it 后缀版本。选择建议：资源极有限时选 2.6B/1B 级别并搭配 INT4 量化；需要较好的对话与推理能力时选 9B/12B；追求最强效果、有多卡或云端 TPU 资源时选 27B。直接做对话选 IT（指令微调）版，要在自有数据上二次训练选 PT 基座版。Gemma 3n 是面向端侧/能效优化的变体，同样在本仓库 API 覆盖范围内。

五、参数高效微调（PEFT）

仓库附带一个轻量 PEFT 库，用于为 Flax linen 模块添加 LoRA 支持，位于 gemma/peft/。LoRA（低秩适应）的核心思想是冻结原始模型权重，只在旁边训练一个低秩的小矩阵，从而用极少的可训练参数达到接近全参数微调的效果，显存与训练成本大幅降低。

5.1 LoRA 适配器

库提供两个包装现有层的 nn.Module：LoRADense 包装 nn.Dense 层，LoRAEinsum 包装 nn.Einsum 层。在 Flax 模块内使用时，直接把原层作为 wrapped 参数传入即可：

class MyModel(nn.Module):

  @nn.compact
  def __call__(self, x):
    layer = peft.LoRADense(
        rank=3,
        wrapped=nn.Dense(10),
    )
    return layer(x)

其中 rank 控制低秩矩阵的秩，秩越小可训练参数越少、训练越快，但表达能力也越弱，需要按任务规模在 1 到 64 之间取舍。

每个包装器都有一个只执行 x @ A @ B 矩阵乘法的底层模块，例如 peft.LoRADense 对应 peft.LoRADenseAdapter；这种情况下需要手工把 LoRA 输出与原始输出相加：

class MyModel(nn.Module):

  @nn.compact
  def __call__(self, x):
    dense = nn.Dense(10)
    lora = peft.LoRADenseAdapter(rank=3)  # 只做 `x @ A @ B`
    return dense(x) + lora(x)

5.2 量化

库提供两个新 API：一个用于应用量化，一个用于在量化感知训练与直通估计（straight through estimation）下训练/优化检查点。与 LoRA 不同，官方建议先用模拟方式训练模型得到相关检查点，再对加载的参数做量化——即先在模拟量化误差的条件下训练，让模型适应低精度，再正式压缩，否则直接量化会带来明显精度损失。直通估计（straight through estimator）是量化感知训练中的常用技巧，用于在反向传播时绕过不可导的量化函数，使梯度仍能正常回传到上游权重：

params_q = peft.quantize_checkpoint(
  params, method=peft.QuantizationMethod.INT4
)

目前 peft.QuantizationMethod.INT4 是唯一支持的量化方法。随后类似 LoRA 适配器，提供量化模拟包装器：Int4Dense 包装 nn.Dense，Int4Einsum 包装 nn.Einsum，在训练时模拟 INT4 量化效果：

class MyModel(nn.Module):

  @nn.compact
  def __call__(self, x):
    layer = peft.Int4Dense(
        wrapped=nn.Dense(10),
        method=peft.QuantizationMethod.Q4_0
    )
    return layer(x)

此外还有 SimulateQuantizedDense 与 SimulateQuantizedEinsum 两个模拟包装器，用法与 Int4Dense 类似，用于训练阶段模拟量化、而不真正压缩权重。

5.3 模型手术与参数手术

库提供工具函数帮助做"模型手术"——即不改动原始网络定义，而是在运行时把某一类模块整体替换成包装版本。这种能力在调试与实验时非常实用：你可以快速给整个网络的 Dense 层加上 LoRA，或者给指定层加上量化模拟，而不必逐一修改模型类。例如把所有 nn.Dense 层替换为 LoRA 版本：

def _replace_dense_by_lora(module: nn.Module) -> nn.Module:
  if isinstance(module, nn.Dense):
    return peft.LoRADense(rank=3, wrapped=module)
  else:
    return module

# 在该上下文内，dense 层会被替换为 LoRA 版本
with ModuleInterceptor(_replace_dense_by_lora):
  y = model(x)

这种方式的好处是：不用改动模型类的源码，只要写一个判断函数，就能在 ModuleInterceptor 上下文里批量改造整个网络，适合在已有模型上实验不同的微调策略而不侵入代码。

关于 Q4_0 量化有一个特别注意事项：它假设 FFN（前馈网络）每个第一层的权重是转置的，可通过判断模块名是否含 gating 来应用 Q4_0_TRANSPOSE，其余层用 Q4_0。

参数树操作方面，LoRA 训练后参数会分成原始权重与 LoRA 小矩阵两部分，库提供一组工具管理这棵嵌套参数字典：peft.split_params 把它拆成两棵树（一棵只含原始参数、一棵只含 LoRA 参数），方便单独保存或加载 LoRA 适配器；peft.merge_params 是其逆操作，把两棵树合并回一棵；peft.fuse_params 把 LoRA 参数融合进原始权重（推理时融合后不再有额外矩阵乘法，速度与原始模型一致）；peft.unfuse_params 则反向恢复出 LoRA 参数，便于继续训练或切换适配器。

六、常见问题 FAQ

1. Gemma 库和 Gemma 模型是什么关系？ 仓库是 Google DeepMind 官方的 gemma PyPI 包，一个 JAX 库；模型权重本身托管在 Kaggle Hub 等平台，需要单独下载。库负责"怎么用"，权重负责"用哪个模型"，二者通过 gm.ckpts 模块对接。
2. 如何安装？ 先按 JAX 文档 装 CPU/GPU/TPU 版 JAX，再 pip install gemma。
3. 支持哪些模型版本？ 同一份 ChatSampler API 支持 Gemma 2、Gemma 3、Gemma 3n、Gemma 4，换检查点即可。
4. 多轮对话怎么开？ 构造 ChatSampler 时传 multi_turn=True，之后连续调用 sampler.chat(...)，模型会记住历史；如果想开启一轮全新对话，重新构造一个 ChatSampler 即可。
5. 图文对话怎么写？ 调用 sampler.chat(prompt, images=[image1, image2])，提示词中用 <|image|> 占位符标记图片位置，占位符数量与 images 列表中的图片数量一一对应；多模态能力在 Gemma 2/3/4 上均可用。
6. 想微调怎么办？ 先看文档站的 Fine-tuning Colab 与 LoRA Colab，低资源场景用 gemma/peft 库的 LoRA 适配器；训练完成后用 fuse_params 融合权重再部署，可避免推理时的额外开销。
7. 量化支持吗？ 提供 quantize_checkpoint（INT4）与 Int4Dense/SimulateQuantizedDense 等模拟包装器；Gemma 3 IT 模型也有官方 INT4 量化检查点（-int4 后缀），显存紧张时优先选官方量化版。
8. JAX 与 PyTorch 版本的 Gemma 有什么区别？ 本仓库是官方 JAX 实现；JAX 版本的优势是统一编译到 CPU/GPU/TPU，且与 Flax 生态的 LoRA、量化工具深度集成；如果已有 PyTorch 技术栈，可使用社区转换的权重，但官方文档与示例以 JAX 为准，遇到问题也应优先参考本仓库的实现。
9. gemma 库和直接用 Hugging Face transformers 加载 Gemma 是一回事吗？ 不是。HF 生态有自己的 Gemma 模型实现，而本仓库是 Google DeepMind 官方维护的 JAX 原生实现，权重格式为 Flax（从 Kaggle Hub 下载），API 以 gm.nn、gm.text.ChatSampler 为核心，二者权重格式与代码互不通用；选哪条路线取决于现有技术栈是 JAX 还是 PyTorch。
10. 哪里看模型能力与指标？ 官方技术报告系统介绍了 Gemma 各代模型的训练方法、评测结果与适用场景：Gemma 1、Gemma 2、Gemma 3。

七、跳转链接

如果你在使用 gemma 库时遇到问题，可以先查阅官方文档站的 Colab 教程与说明；仍无法解决时，到 GitHub Issues 提交反馈，团队会跟进处理。以下是本文整理涉及的全部入口：

• 仓库主页：github.com/google-deepmind/gemma
• 官方站点：ai.google.dev/gemma
• 文档站：gemma-llm.readthedocs.io
• 示例脚本：examples/ 目录
• Kaggle 权重：Gemma 3、Gemma 2、Gemma 4
• PEFT 库：gemma/peft/README.md