引言

两岸荔枝红，万家烟雨中。

Google 最近放出了开源大语言模型 Gemma 2，目前已经在 Hugging Face 上亮相了 4 个模型——2 个基础版和 2 个微调版。这篇内容主要拆解 Gemma 2 的技术特点和上手体验，后续还会单独聊聊如何微调它。

技术报告原文：https://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf

简介

2024年6月27日，Google DeepMind 正式发布了 Gemma 2，这是 Gemma 系列轻量级开放语言模型的最新迭代。相比前代，Gemma 2 在架构和训练方法上都做了不少创新，在多项基准测试中进步明显，甚至能与参数规模大 2-3 倍的模型掰手腕。下面对技术报告的核心内容做个梳理，涵盖模型架构、预训练和后训练、性能评估等方面。

Gemma 2 的训练数据量大约是前代的两倍：27B 模型用了 13 万亿 Token，9B 模型用了 8 万亿 Token，数据来源主要是英语网页、代码和数学内容。官方没有透露具体的数据混合比例。

Gemma 2模型介绍

Gemma 2 系列提供三个规模：27B、9B 和 2.6B，目前 27B 和 9B 已经发布，2.6B 还在路上。这些模型都基于 Transformer 解码器架构，核心特点如下：

• 上下文长度 8192 个 token
• 旋转位置编码（RoPE）
• 近似 GeGLU 非线性激活函数
• 交替使用局部滑动窗口注意力和全局注意力
• 应用 logit 软上限技术
• 使用 RMSNorm 进行后归一化和前归一化
• 9B 和 27B 模型采用分组查询注意力（GQA）

架构设计

相比 Gemma 1，Gemma 2 在架构上做了几项关键改进：

（1）局部滑动窗口和全局注意力交替
在每隔一层交替使用局部滑动窗口注意力和全局注意力。局部注意力的窗口大小是 4096 个 token，全局注意力的跨度达到 8192 个 token。这种设计兼顾了计算效率和长程依赖的建模能力。

（2）Logit 软上限
在每个注意力层和最终输出层，对 logit 值进行软上限处理，将其限制在 [-soft_cap， +soft_cap] 范围内。具体来说：
logits ← soft_cap * tanh(logits/soft_cap)
9B 和 27B 模型的注意力 logit 上限为 50.0，最终 logit 上限为 30.0。这项技术有助于稳定训练过程和生成质量。

（3）RMSNorm 归一化
使用 RMSNorm 来归一化每个 Transformer 子层、注意力层和前馈层的输入和输出，提升训练稳定性。

（4）分组查询注意力（GQA）
9B 和 27B 模型采用了 GQA 机制，组数为 2。实验表明，GQA 在不影响下游任务性能的同时，能有效提高推理速度。

训练方法

Gemma 2 的另一大创新在于训练方法，尤其是对小规模模型（9B 和 2.6B）引入了知识蒸馏技术。

（1）预训练数据
27B 模型在 13 万亿 token 上训练，9B 模型使用 8 万亿 token，2.6B 模型则用了 2 万亿 token。数据以英语为主，来源包括网页、代码和科学文章。tokenizer 采用 SentencePiece 方法，词表大小为 256k。

（2）知识蒸馏
对于 9B 和 2.6B 模型，团队创新性地将知识蒸馏应用于大规模预训练。具体做法是：用一个大型模型作为教师，最小化学生模型与教师模型在每个 token 的条件概率分布之间的负对数似然。这种方法让小型模型在有限计算资源下，获得了近似训练更多 token 的效果。

（3）计算基础设施
Gemma 2 使用了 TPUv4、TPUv5e 和 TPUv5p 进行训练，结合数据并行和模型并行策略。例如，27B 模型在 6144 个 TPUv5p 芯片上训练，采用了 768 路数据复制和 8 路模型分片。

后训练优化

为了进一步提升模型性能和安全性，Gemma 2 在预训练后做了一系列优化：

（1）监督微调（SFT）
在英语指令-响应对数据集上进行监督微调，数据包括合成和人工生成的样本。

（2）人类反馈强化学习（RLHF）
在 SFT 基础上应用 RLHF，奖励模型基于标注的英语偏好数据训练，策略使用与 SFT 相同的提示。

（3）模型融合
将不同超参数实验得到的模型进行平均，以提升整体性能。

关键发现：知识蒸馏的影响

技术报告重点展示了知识蒸馏对小型语言模型的显著效果：

• 在 500B token 上训练的 2.6B 模型，使用知识蒸馏比从头训练在 3 项基准测试的平均分上高出 7.4 个百分点。
• 随着模型规模从 200M 增加到 1B 参数，知识蒸馏带来的收益依然明显。
• 在 9B 规模上，用 GQA 替代多头注意力（MHA）对性能影响不大，但能减少参数量并提高推理速度。
• 对于 9B 模型，更深的网络结构略优于更宽的结构。
• 可以在推理时调整局部注意力的滑动窗口大小，对困惑度的影响较小，为推理速度优化提供了灵活性。

性能评估

Gemma 2 在多项基准测试中表现出色：

• 在同等规模的开放模型中达到最佳性能
• 甚至能与参数量大 2-3 倍的模型竞争
• 在问答、常识推理、数学科学、编程等多个领域表现优异

具体评估涵盖了自动化基准测试、人工评估、问答能力（如 SQuAD、Natural Questions）、常识推理（如 WinoGrande、HellaSwag）、数学和科学（如 MATH、MMLU）、代码能力（如 HumanEval、MBPP）等方面。

使用

体验：Hugging Chat

你可以直接在 Hugging Chat 上与 Gemma 27B 指令模型对话：https://huggingface.co/chat/models/google/gemma-2-27b-it

如何提示 Gemma 2

基础模型没有提示格式，像其他基础模型一样，可以用于继续输入序列的合理延续或零样本/少样本推理。指令版本有一个简单的对话结构：

user
knock knock
model
who is there
user
LaMDA
model
LaMDA who?

这个格式必须精确复制才能有效使用。稍后会介绍如何利用 transformers 中的聊天模板轻松实现。

基于Hugging Face Transformers

随着 Transformers 4.42 版本的发布，你可以使用 Gemma 并借助 Hugging Face 生态系统的所有工具。使用前确保安装最新版本：

pip install "transformers>=4.42.3" --upgrade

以下代码演示了如何使用 transformers 调用 gemma-2-9b-it。它需要约 18 GB 的 RAM，适合大多数消费级 GPU。同样的代码也适用于 gemma-2-27b-it，但需要 56 GB 的 RAM，更适合生产环境。通过加载 8-bit 或 4-bit 模式，可以进一步降低内存消耗。

from transformers import pipeline
import torch

pipe = pipeline(
    "text-generation",
    model="google/gemma-2-9b-it",
    model_kwargs={"torch_dtype": torch.bfloat16},
    device="cuda",
)

messages = [
    {"role": "user", "content": "请模仿蔡坤的方式说两句土味情话"},
]
outputs = pipe(
    messages,
    max_new_tokens=256,
    do_sample=False,
)
assistant_response = outputs[0]["generated_text"][-1]["content"]
print(assistant_response)

输出结果如下：

1. 你就像一碗热腾腾的番茄鸡蛋面，让我欲罢不能！
2. 你的眼睛像星星一样闪亮，照亮了我这颗孤独的心！


希望你喜欢我的土味情话！?

_{这里使用 bfloat16 是因为它是指令调优模型的参考精度。在硬件上运行 float16 可能更快，9B 模型的结果应当相似。但使用 float16 时，27B 指令调优模型会产生不稳定输出：对于该模型权重，必须使用 bfloat16。}

你还可以自动量化模型，以 8-bit 甚至 4-bit 模式加载。4-bit 模式的 27B 版本需要约 18 GB 内存，兼容多数消费级显卡和 Google Colab 中的 GPU。以下是 4-bit 模式下加载生成管道的示例：

import os
from transformers import pipeline
import torch

init_model_dir = "/share_model_zoo/LLM/"
init_model_id = "google/gemma-2-9b-it"
init_model_path = os.path.join(init_model_dir, init_model_id)

pipeline = pipeline(
    "text-generation",
    model=init_model_path,
    model_kwargs={
        "torch_dtype": torch.bfloat16,
        "quantization_config": {"load_in_4bit": True}
    },
    device_map="auto",
)

messages = [
    {"role": "user", "content": "请模仿蔡坤的语气说两句土味情话"},
]
outputs = pipeline(
    messages,
    max_new_tokens=256,
    do_sample=False,
)
assistant_response = outputs[0]["generated_text"][-1]["content"]
print(assistant_response)

输出结果如下：

1.  宝贝，你就像一朵盛开的玫瑰，香气迷人，让我忍不住想把你捧在手里，永远守护着。
2.  你是我这辈子最想遇见的人，就像天上的星星，照亮我的整个世界，让我不再迷茫。


希望你喜欢这些土味情话！

更多使用细节请参考模型卡。

结论与展望

Gemma 2 代表了轻量级开放语言模型的一次重要进步。通过创新的架构设计和训练方法，尤其是知识蒸馏技术的大规模应用，Gemma 2 在有限参数规模下实现了出色的性能。这不仅推动了语言模型技术的发展，也为更广泛的 AI 应用提供了高效、实用的解决方案。

未来的研究方向可能包括：

• 进一步优化知识蒸馏技术
• 探索更高效的模型架构
• 增强模型的多语言和跨领域能力
• 改进模型的安全性和可控性

总的来说，Gemma 2 的成功说明，通过精心的设计和创新的训练方法，我们可以在不盲目追求参数规模的前提下，显著提升语言模型的能力。这为构建更高效、更易部署的 AI 系统开辟了新的可能性，有望推动 AI 技术在更多实际场景中的应用与落地。

模型汇总