跳到正文
原文
Hugging Face Blog·· 2024-06-27精选AI 评分79

Google 发布 Gemma 2 开源大模型,9B 与 27B 双尺寸

Welcome Gemma 2 - Google’s new open LLM

AI 导读

Google 发布 Gemma 2 系列开源大语言模型,包含 9B 和 27B 两个尺寸的基础版与指令微调版,共 4 个开放权重模型,上下文长度 8K token。

推荐理由

详细拆解 Gemma 2 的滑动窗口注意力、软上限与蒸馏等改动,可作为理解这代开源模型训练思路的参考。

正文 · AI 翻译

Google 发布了 Gemma 2,这是其最先进的开放 LLM 系列的最新成员,我们很高兴与 Google 合作,确保在 HF Mirror 生态系统中实现最佳集成。您可以在 Hub 上找到 4 个开放权重模型(2 个基础模型和 2 个微调模型)。在发布的功能和集成中,我们有:

目录

什么是 Gemma 2?

Gemma 2 是 Google 开放 LLM 的最新一代。它有两种规模,90 亿和 270 亿参数,分别提供基础(预训练)和指令微调版本。Gemma 基于 Google Deepmind Gemini,上下文长度为 8K tokens:

Gemma 2 模型的训练数据量约为第一代的两倍,27B 版本总计 13 万亿 tokens,9B 版本总计 8 万亿 tokens,涵盖网络数据(主要为英语)、代码和数学。我们不知道训练混合的确切细节,只能猜测更大规模、更细致的数据整理是性能提升的重要因素。

Gemma 2 采用与第一代相同的许可证,这是一种宽松许可证,允许再分发、微调、商业使用和衍生作品。

Gemma 2 的技术进展

Gemma 2 与第一代有许多相似之处。它的上下文长度为 8192 tokens,并使用旋转位置嵌入(RoPE)。与原始 Gemma 相比,Gemma 2 有四项主要进展:

  • 滑动窗口注意力:交替使用滑动窗口和全二次注意力,以实现高质量生成。
  • Logit soft-capping:通过将 logits 缩放到固定范围,防止其过度增长,从而改善训练。
  • 知识蒸馏:利用更大的教师模型来训练较小的模型(用于 9B 模型)。
  • 模型合并:将两个或多个 LLM 合并为一个新模型

Gemma 2 在 Google Cloud TPU(27B 在 v5p,9B 在 TPU v4)上训练,使用 JAX 和 ML Pathways。Gemma 2 Instruct 针对对话应用进行了优化,并使用监督微调(SFT)、从更大模型进行蒸馏、使用更偏向对话能力的奖励模型进行基于人类反馈的强化学习(RLHF),以及使用 WARP 进行模型合并来提升整体性能,训练数据混合了合成和人类生成的提示-响应对。

与预训练混合类似,关于微调数据集或与SFT和RLHF相关的超参数的详细信息均未分享。

滑动窗口注意力

滑动窗口注意力是一种减少Transformer模型中注意力计算的记忆和时间需求的方法,已被应用于如Mistral等模型中。Gemma 2的创新之处在于,每隔一层应用滑动窗口(局部 - 4096个令牌),而中间层仍然使用完整的二次全局注意力(8192个令牌)。我们认为这是在长上下文情境下提高质量的一种方式(一半的层仍然关注所有令牌),同时部分受益于滑动注意力的优势。

软上限与注意力实现

软上限是一种防止logits过度增长而不截断它们的技术。其工作原理是将logits除以最大阈值(soft_cap),然后通过一个tanh层(确保它们处于(-1, 1)范围内),最后再乘以阈值。这保证了最终值将在(-soft_cap, +soft_cap)区间内,同时不会丢失太多信息,从而稳定训练过程。

综合起来,logits的计算方式为:logits ← soft_cap ∗ tanh(logits/soft_cap)

Gemma 2在最终层和每个注意力层采用了软上限。注意力logits的上限设为50.0,而最终logits的上限为30.0。

发布时,软上限与Flash Attention / SDPA不兼容,但在推理时仍可使用它们以实现最大效率。Gemma 2团队观察到,在推理时移除软上限后差异非常小。

注意:为了稳定的微调运行,您仍需启用软上限,因此我们建议使用eager注意力而非SDPA进行微调。

知识蒸馏

知识蒸馏是一种流行的技术,用于训练较小的学生模型模仿更大但性能更好的教师模型的行为。这通过向LLM的下一个令牌预测任务添加来自教师(例如GPT-4、Claude或Gemini)的令牌概率分布来实现,为学生提供了更丰富的学习信号。

根据Gemma 2技术报告,知识蒸馏被用于预训练9B模型,而27B模型则是从头开始预训练的。

在后期训练中,Gemma 2团队从教师(报告中未指定,但推测为Gemini Ultra)生成了一系列多样化的完成文本,然后使用SFT在这些合成数据上训练学生模型。这是许多开源模型的基础,如Zephyr和OpenHermes,它们完全基于来自更大LLM的合成数据进行训练。

尽管有效,但这种方法存在缺陷,因为学生和教师之间的模型容量不匹配可能导致训练-推理不匹配,即学生在推理时生成的文本与训练时所见相比属于分布外。

为了解决这个问题,Gemma 2 团队采用了“策略内蒸馏”方法,即让学生模型根据 SFT 提示生成完成内容。随后,利用这些完成内容计算教师模型与学生模型逻辑值之间的 KL 散度。通过在训练过程中持续最小化 KL 散度,学生模型能够准确模拟教师模型的行为,同时减少训练与推理之间的不匹配。

这种方法颇为有趣,因为我们在社区中已见证,诸如在线 DPO 之类的策略内方法能产生更强大的模型。而策略内蒸馏的一个优势在于,仅需教师模型的逻辑值,无需依赖奖励模型或 LLM 作为评判者来优化模型。未来几个月,如果这种方法在微调者中变得更加流行,将令人十分期待!

模型合并

模型合并是一项将两个或多个大型语言模型合并成一个新模型的技术。它相对新颖且处于实验阶段,无需加速器即可使用。Mergekit 是一款流行的开源工具包,用于合并大型语言模型,它实现了线性、SLERP、TIES、DARE 等合并技术。

根据技术报告,Gemma 2 使用了一种名为 Warp 的新合并技术,该技术分三个不同阶段合并模型:

  1. 指数移动平均(EMA):此步骤应用于强化学习(RL)微调过程中。
  2. 球面线性插值(SLERP):此步骤在多个策略的 RL 微调后应用。
  3. 线性插值趋向初始化(LITI):此阶段在 SLERP 阶段之后应用。

Gemma 2 评估

Gemma 模型的表现如何?以下是根据技术报告和新版开放大型语言模型排行榜与其他开放模型的性能对比。

技术报告结果

Gemma 2 的此技术报告对比了不同开放大型语言模型在之前的开放大型语言模型排行榜基准上的性能。

Llama 3(70B) Qwen 1.5(32B) Gemma 2(27B)
MMLU 79.2 74.3 75.2
GSM8K 76.9 61.1 75.1
ARC-c 68.8 63.6 71.4
HellaSwag 88.0 85.0 86.4
Winogrande 85.3 81.5 83.7

报告还对比了小型语言模型的性能。

基准 Mistral(7B) Llama 3(8B) Gemma(8B) Gemma 2(9B)
MMLU 62.5 66.6 64.4 71.3
GSM8K 34.5 45.7 50.9 62.3
ARC-C 60.5 59.2 61.1 68.4
HellaSwag 83.0 82.0 82.3 81.9
Winogrande 78.5 78.5 79.0 80.6

开放大型语言模型排行榜结果

注意:我们目前正针对新版开放大型语言模型排行榜基准对 Google Gemma 2 进行单独评估,并将于今日稍晚更新此部分内容。

如何提示 Gemma 2

基础模型没有提示格式。如同其他基础模型,它们可用于为输入序列提供合理的续写,或进行零样本/少样本推理。指令版模型具有非常简单的对话结构:

<start_of_turn>user
knock knock<end_of_turn>
<start_of_turn>model
who is there<end_of_turn>
<start_of_turn>user
LaMDA<end_of_turn>
<start_of_turn>model
LaMDA who?<end_of_turn><eos>

为了有效使用,必须精确复制此格式。稍后我们将展示如何利用 transformers 中的聊天模板轻松复现指令提示。

演示

你可以在 Hugging Chat 上与 Gemma 27B Instruct 模型聊天!请查看此链接:https://hf-mirror.com/chat/models/google/gemma-2-27b-it。

使用 HF Mirror Transformers

通过 Transformers 4.42 版本,您可以使用 Gemma,并利用 HF Mirror 生态系统中的所有工具。要使用 Transformers 配合 Gemma 模型,请确保使用最新的 transformers 版本:

pip install "transformers>=4.42.3" --upgrade

以下代码片段展示了如何将 gemma-2-9b-it 与 transformers 配合使用。它需要约 18 GB 的内存,适合许多消费级 GPU。相同的代码片段也适用于 gemma-2-27b-it,其内存需求为 56GB,使其成为生产用例中非常有趣的模型。通过以 8 位或 4 位模式加载,可以进一步降低内存消耗。

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": "Who are you? Please, answer in pirate-speak."},
]
outputs = pipe(
    messages,
    max_new_tokens=256,
    do_sample=False,
)
assistant_response = outputs[0]["generated_text"][-1]["content"]
print(assistant_response)

啊嗬,伙计!我是一艘卑微的文字之船,航行在数字海洋上。人们叫我 Gemma,是 Google DeepMind 的优秀团队创造的作品。我在文本的宝库中训练,学会了像真正的海盗一样说话和写作。

尽管问我问题吧,我会尽力回答,是的!🦜📚

我们使用 bfloat16,因为这是指令调优模型的参考精度。在你的硬件上以 float16 运行可能更快,并且在 9B 模型上结果应该相似。但请注意,27B 指令调优模型在使用 float16 时会产生不稳定的输出:对于该模型权重,必须使用 bfloat16。

你还可以自动量化模型,以 8 位甚至 4 位模式加载。大型 27B 版本的 4 位加载运行仅需约 18 GB 内存,使其与许多消费级显卡以及 Google Colab 中的 GPU 兼容。以下是以 4 位加载生成管道的方法:

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

如需了解有关将模型与 transformers 配合使用的更多详情,请查看模型卡片。

与 Google Cloud 集成

注意:我们目前正在为 GKE 和 Vertex AI 添加新容器,以便高效运行 Google Gemma 2。容器可用后,我们将立即更新本节内容。

使用 🤗 TRL 进行微调

训练 LLM 在技术和计算上都可能具有挑战性。在本节中,我们将介绍 HF Mirror 生态系统中可用的工具,以便在消费级 GPU 上高效训练 Gemma。

下面是一个在 OpenAssistant 的聊天数据集上微调 Gemma 的示例命令。我们使用 4 位量化和 QLoRA 来节省内存,并针对所有注意力块的线性层。请注意,与稠密 Transformer 不同,不应针对 MLP 层,因为它们是稀疏的,并且与 PEFT 的交互效果不佳。

首先,安装 🤗 TRL 的 nightly 版本并克隆仓库,以获取训练脚本:

pip install "transformers>=4.42.3" --upgrade
pip install --upgrade bitsandbytes
pip install --ugprade peft
pip install git+https://github.com/huggingface/trl
git clone https://github.com/huggingface/trl
cd trl

然后你可以运行该脚本:

# peft tuning; single GPU; https://wandb.ai/costa-huang/huggingface/runs/l1l53cst
python \
    trl/scripts/sft.py \
    --model_name google/gemma-2-27b \
    --dataset_name OpenAssistant/oasst_top1_2023-08-25 \
    --dataset_text_field="text" \
    --per_device_train_batch_size 1 \
    --per_device_eval_batch_size 1 \
    --gradient_accumulation_steps 4 \
    --learning_rate 2e-4 \
    --report_to wandb \
    --bf16 \
    --max_seq_length 1024 \
    --lora_r 16 --lora_alpha 32 \
    --lora_target_modules q_proj k_proj v_proj o_proj \
    --load_in_4bit \
    --use_peft \
    --attn_implementation eager \
    --logging_steps=10 \
    --gradient_checkpointing \
    --output_dir models/gemma2

alt_text

如果你有更多 GPU 可用,可以使用 DeepSpeed 和 ZeRO Stage 3 进行训练:

accelerate launch --config_file=examples/accelerate_configs/deepspeed_zero3.yaml \
    trl/scripts/sft.py \
    --model_name google/gemma-2-27b \
    --dataset_name OpenAssistant/oasst_top1_2023-08-25 \
    --dataset_text_field="text" \
    --per_device_train_batch_size 1 \
    --per_device_eval_batch_size 1 \
    --gradient_accumulation_steps 4 \
    --learning_rate 2e-5 \
    --report_to wandb \
    --bf16 \
    --max_seq_length 1024 \
    --attn_implementation eager \
    --logging_steps=10 \
    --gradient_checkpointing \
    --output_dir models/gemma2

alt_text

与 Inference Endpoints 集成

你可以使用 Text Generation Inference 作为后端,在 HF Mirror 的Inference Endpoints 上部署 Gemma 2。Text Generation Inference 是 HF Mirror 开发的生产级推理容器,可轻松部署大型语言模型。它具有连续批处理、token 流式传输、用于在多 GPU 上快速推理的张量并行,以及生产级日志记录和追踪等功能。

要部署 Gemma 2 模型,请前往模型页面并点击 Deploy -> Inference Endpoints 小部件。Inference Endpoints 支持 OpenAI 兼容的 Messages API,使你只需更改 URL,即可从另一个闭源模型切换到开源模型。

from openai import OpenAI

# initialize the client but point it to TGI
client = OpenAI(
    base_url="<ENDPOINT_URL>" + "/v1/",  # replace with your endpoint url
    api_key="<HF_API_TOKEN>",  # replace with your token
)
chat_completion = client.chat.completions.create(
    model="tgi",
    messages=[
        {"role": "user", "content": "Why is open-source software important?"},
    ],
    stream=True,
    max_tokens=500
)

# iterate and print stream
for message in chat_completion:
    print(message.choices[0].delta.content, end="")

其他资源

致谢

在生态系统中发布这些模型并提供支持和评估,离不开许多社区成员的贡献,包括 Clémentine 和 Nathan 负责 LLM 评估;Nicolas 负责文本生成推理支持;Arthur、Sanchit、Joao 和 Lysandre 负责将 Gemma 2 集成到 transformers 中;Nathan 和 Victor 让 Gemma 2 在 Hugging Chat 中可用。

并感谢 Google 团队发布 Gemma 2 并将其提供给开源 AI 社区!

来源:Hugging Face Blog · hf-mirror.com