【机器学习】Google开源大模型Gemma2：原理、微调训练及推理部署实战

一、引言

二、模型简介

2.1 Gemma2概述

2.2 Gemma2 模型架构

三、训练与推理

3.1 Gemma2 模型训练

3.1.1 下载基座模型

3.1.2 导入依赖库

3.1.3 量化配置

3.1.4 分词器和模型实例化

3.1.5 引入PEFT进行LORA配置

3.1.6 样本数据清洗与加载

3.1.7 模型训练与保存

3.1.8 完整训练代码

3.1.9 启动训练以及收敛过程

3.1.10 训练显存占用

3.2 Gemma2 基座与微调模型合并推理

3.2.1 导入库

3.2.2 导入基座模型

3.2.3 合并基座模型与微调模型

3.2.4 基于对话模版进行对话生成

3.2.5 推理显存占用

3.2.6 推理效果

3.2.7 微调与推理完整代码

四、总结

一、引言

Gemma 是 Google 推出的轻量级、先进的开放模型系列，采用与 Gemini 模型相同的研究成果和技术构建而成。它们是仅使用解码器的文本到文本大型语言模型（提供英语版本），为预训练变体和指令调整变体具有开放权重。Gemma 模型非常适合各种文本生成任务，包括问题解答、摘要和推理。由于它们相对较小，因此可以将其部署在资源有限的环境（如笔记本电脑、桌面设备或您自己的云基础架构）中，让更多人能够使用先进的 AI 模型，并帮助促进每个人的创新。

二、模型简介

2.1 Gemma2概述

Gemma2与他的上一代Gemma以及Qwen2等均采用decoder-only网络结构，主要参数情况如下：

与Gemma相同点：

上下文长度为 8192 个 token
使用旋转位置嵌入（RoPE）
近似 GeGLU 非线性

与Gemma不同点：

局部滑动窗口和全局注意力。研究团队在每隔一层中交替使用局部滑动窗口注意力和全局注意力。局部注意力层的滑动窗口大小设置为4096个token，而全局注意力层的跨度设置为8192个token。
Logit软封顶。根据Gemini 1.5的方法，研究团队在每个注意力层和最终层限制logit，使得logit的值保持在−soft_cap和+soft_cap之间。
对于9B和27B模型，研究团队将注意力对数封顶设置为50.0，最终对数封顶设置为30.0。截至本文发表时，注意力logit软封顶与常见的FlashAttention实现不兼容，因此他们已从使用FlashAttention的库中移除了此功能。研究团队对模型生成进行了有无注意力logit软封顶的消融实验，发现大多数预训练和后期评估中，生成质量几乎不受影响。本文中的所有评估均使用包含注意力logit软封顶的完整模型架构。然而，某些下游性能可能仍会受到此移除的轻微影响。
使用RMSNorm进行post-norm 和pre-norm。为了稳定训练，研究团队使用RMSNorm对每个变换子层、注意力层和前馈层的输入和输出进行归一化。
分组查询注意力。27B和9B模型均使用GQA，num_groups = 2，基于消融实验表明在保持下游性能的同时提高了推理速度。

分组查询注意力 (Grouped Query Attention) 是一种在大型语言模型中的多查询注意力 (MQA) 和多头注意力 (MHA) 之间进行插值的方法，它的目标是在保持 MQA 速度的同时实现 MHA 的质量

效果对比：

Gemma2 9B模型在多个维度超过近尺寸的Llama3 8B，27B尺寸模型在多个评价标准下超过314B的Grok-1：

2.2 Gemma2 模型架构

通过AutoModelForCausalLM模型头查看模型结构：

Gemma2ForCausalLM(
  (model): Gemma2Model(
    (embed_tokens): Embedding(256000, 4608, padding_idx=0)
    (layers): ModuleList(
      (0-45): 46 x Gemma2DecoderLayer(
        (self_attn): Gemma2SdpaAttention(
          (q_proj): Linear(in_features=4608, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4608, out_features=2048, bias=False)
          (v_proj): Linear(in_features=4608, out_features=2048, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4608, bias=False)
          (rotary_emb): Gemma2RotaryEmbedding()
        )
        (mlp): Gemma2MLP(
          (gate_proj): Linear(in_features=4608, out_features=36864, bias=False)
          (up_proj): Linear(in_features=4608, out_features=36864, bias=False)
          (down_proj): Linear(in_features=36864, out_features=4608, bias=False)
          (act_fn): PytorchGELUTanh()
        )
        (input_layernorm): Gemma2RMSNorm()
        (post_attention_layernorm): Gemma2RMSNorm()
        (pre_feedforward_layernorm): Gemma2RMSNorm()
        (post_feedforward_layernorm): Gemma2RMSNorm()
      )
    )
    (norm): Gemma2RMSNorm()
  )
  (lm_head): Linear(in_features=4608, out_features=256000, bias=False)
)

46层Gemma2DecoderLayer，每层包含1个自注意力层Gemma2SdpaAttention、1个mlp层Gemma2MLP
使用RMSNorm进行post-norm 和pre-norm。为了稳定训练，研究团队使用RMSNorm对每个变换子层、注意力层和前馈层的输入和输出进行归一化

三、训练与推理

3.1 Gemma2 模型训练

在之前的文章中，我介绍过采用LlamaFactory的webui以及命令行进行模型训练，今天基于transformers库原生微调Gemma2。

3.1.1 下载基座模型

我们仍然秉承一贯的作风，为网络不稳定的同学提供了modelscope下载方案：

from modelscope import snapshot_download
model_dir = snapshot_download('LLM-Research/gemma-2-27b-it')

3.1.2 导入依赖库

import torch
import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM,BitsAndBytesConfig

3.1.3 量化配置

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,  # 或者 load_in_8bit=True，根据需要设置
    llm_int8_enable_fp32_cpu_offload=True,
    bnb_4bit_compute_dtype=torch.bfloat16,#虽然我们以4位加载和存储模型，但我们在需要时会部分反量化他，并以16位精度进行计算
    bnb_4bit_quant_type="nf4",#nf量化类型
    bnb_4bit_use_double_quant=True,#双重量化，量化一次后再量化，进一步解决显存
)

3.1.4 分词器和模型实例化

tokenizer = AutoTokenizer.from_pretrained(model_dir,trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_dir,trust_remote_code=True, device_map=device,torch_dtype=torch.bfloat16,quantization_config=quantization_config,attn_implementation='eager')
model.gradient_checkpointing_enable

3.1.5 引入PEFT进行LORA配置

from peft import LoraConfig,get_peft_model,prepare_model_for_kbit_training


model = prepare_model_for_kbit_training(model)

config = LoraConfig(
    r=32,
    lora_alpha=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj","down_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)

3.1.6 样本数据清洗与加载

from datasets import load_dataset,load_from_disk
data = load_dataset('json',data_files="./quotes.jsonl")
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
print(data)

3.1.7 模型训练与保存

trainer = transformers.Trainer(
    model=model,
    train_dataset=data["train"],
    args=transformers.TrainingArguments(
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        warmup_steps=10,
        max_steps=50,
        learning_rate=3e-4,
        fp16=True,
        logging_steps=1,
        output_dir="outputs/checkpoint-1"+time_str,
        optim="paged_adamw_8bit",
        save_strategy = 'steps',
        save_steps = 10,
    ),
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

model.config.use_cache = False  # silence the warnings. Please re-enable for inference!
trainer.train()

trainer.save_model(trainer.args.output_dir)

注意：

per_device_train_batch_size=1：开始设置为4会出现'grad_norm': nan,'learning_rate':0的情况。

3.1.8 完整训练代码

from datetime import datetime
now = datetime.now()
time_str = now.strftime('%Y-%m-%d %H:%M:%S')
print(time_str)

#0,download model
from modelscope import snapshot_download
model_dir = snapshot_download('LLM-Research/gemma-2-27b-it')
#model_dir = snapshot_download('qwen/Qwen2-7B-Instruct')
import torch
import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM,BitsAndBytesConfig



device = "auto"

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,  # 或者 load_in_8bit=True，根据需要设置
    llm_int8_enable_fp32_cpu_offload=True,
    bnb_4bit_compute_dtype=torch.bfloat16,#虽然我们以4位加载和存储模型，但我们在需要时会部分反量化他，并以16位精度进行计算
    bnb_4bit_quant_type="nf4",#nf量化类型
    bnb_4bit_use_double_quant=True,#双重量化，量化一次后再量化，进一步解决显存
)
tokenizer = AutoTokenizer.from_pretrained(model_dir,trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_dir,trust_remote_code=True, device_map=device,torch_dtype=torch.bfloat16,quantization_config=quantization_config,attn_implementation='eager')
model.gradient_checkpointing_enable

from peft import LoraConfig,get_peft_model,prepare_model_for_kbit_training


model = prepare_model_for_kbit_training(model)

config = LoraConfig(
    r=32,
    lora_alpha=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj","down_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)

from datasets import load_dataset,load_from_disk
data = load_dataset('json',data_files="./quotes.jsonl")
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
print(data)

trainer = transformers.Trainer(
    model=model,
    train_dataset=data["train"],
    args=transformers.TrainingArguments(
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        warmup_steps=10,
        max_steps=50,
        learning_rate=3e-4,
        fp16=True,
        logging_steps=1,
        output_dir="outputs/checkpoint-1"+time_str,
        optim="paged_adamw_8bit",
        save_strategy = 'steps',
        save_steps = 10,
    ),
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

model.config.use_cache = False  # silence the warnings. Please re-enable for inference!
trainer.train()

trainer.save_model(trainer.args.output_dir)

3.1.9 启动训练以及收敛过程

采用CUDA_VISIBLE_DEVICES=1,2,3 python gemma2_train.py 启动

3.1.10 训练显存占用

3张显卡启动：针对27B尺寸模型进行int4位微调，占用显存约28.9G。如果bf16微调，大约需要54G。相比于LLama3、Qwen2等72B尺寸模型的优势就是仅消耗单卡A100即可bf16微调训练。

3.2 Gemma2 基座与微调模型合并推理

3.2.1 导入库

这里比较重要的是peft中的PeftModel和PeftConfig，PeftModel用于合并基座与微调模型，PeftConfig用于提取Peft微调模型的配置文件

import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

3.2.2 导入基座模型

peft_model_dir = trainer.args.output_dir
config = PeftConfig.from_pretrained(peft_model_dir)
print(config)
model = AutoModelForCausalLM.from_pretrained(
    config.base_model_name_or_path, return_dict=True,  device_map=device,
    torch_dtype=torch.float16, quantization_config=quantization_config
)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)

3.2.3 合并基座模型与微调模型

model = PeftModel.from_pretrained(model, peft_model_dir)

3.2.4 基于对话模版进行对话生成

chat=[
    {"role": "user", "content": "详细介绍一下大语言模型,评价下与深度学习的差异"},
]

prompt = tokenizer.apply_chat_template(chat, tokenize=True, add_generation_prompt=True,return_tensors="pt").to(model.device)

outputs = model.generate(prompt,max_length=2500)

outputs = [ 
    output_ids[len(input_ids):] for input_ids, output_ids in zip(prompt, outputs)
]

print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])

3.2.5 推理显存占用

基座模型和微调模型合并后，大约需要40G？？

3.2.6 推理效果

3.2.7 微调与推理完整代码

from datetime import datetime
now = datetime.now()
time_str = now.strftime('%Y-%m-%d %H:%M:%S')
print(time_str)

#0,download model
from modelscope import snapshot_download
model_dir = snapshot_download('LLM-Research/gemma-2-27b-it')
#model_dir = snapshot_download('qwen/Qwen2-7B-Instruct')
import torch
import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM,BitsAndBytesConfig



device = "auto"

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,  # 或者 load_in_8bit=True，根据需要设置
    llm_int8_enable_fp32_cpu_offload=True,
    bnb_4bit_compute_dtype=torch.bfloat16,#虽然我们以4位加载和存储模型，但我们在需要时会部分反量化他，并以16位精度进行计算
    bnb_4bit_quant_type="nf4",#nf量化类型
    bnb_4bit_use_double_quant=True,#双重量化，量化一次后再量化，进一步解决显存
)
tokenizer = AutoTokenizer.from_pretrained(model_dir,trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_dir,trust_remote_code=True, device_map=device,torch_dtype=torch.bfloat16,quantization_config=quantization_config,attn_implementation='eager')
model.gradient_checkpointing_enable

from peft import LoraConfig,get_peft_model,prepare_model_for_kbit_training


model = prepare_model_for_kbit_training(model)

config = LoraConfig(
    r=32,
    lora_alpha=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj","down_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)

from datasets import load_dataset,load_from_disk
data = load_dataset('json',data_files="./quotes.jsonl")
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
print(data)

trainer = transformers.Trainer(
    model=model,
    train_dataset=data["train"],
    args=transformers.TrainingArguments(
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        warmup_steps=10,
        max_steps=50,
        learning_rate=3e-4,
        fp16=True,
        logging_steps=1,
        output_dir="outputs/checkpoint-1"+time_str,
        optim="paged_adamw_8bit",
        save_strategy = 'steps',
        save_steps = 10,
    ),
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

model.config.use_cache = False  # silence the warnings. Please re-enable for inference!
#trainer.train()

trainer.save_model(trainer.args.output_dir)


# merge model and inference
import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

#peft_model_dir = trainer.args.output_dir
peft_model_dir = "/aigc_dev/gemma2/outputs/checkpoint-12024-07-04 21:57:45"
config = PeftConfig.from_pretrained(peft_model_dir)
print(config)
model = AutoModelForCausalLM.from_pretrained(
    config.base_model_name_or_path, return_dict=True,  device_map=device,
    torch_dtype=torch.bfloat16, quantization_config=quantization_config
)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)

# Load the Lora model
model = PeftModel.from_pretrained(model, peft_model_dir)

chat=[
    {"role": "user", "content": "详细介绍一下大语言模型,评价下与深度学习的差异"},
]

prompt = tokenizer.apply_chat_template(chat, tokenize=True, add_generation_prompt=True,return_tensors="pt").to(model.device)

outputs = model.generate(prompt,max_length=2500)

outputs = [ 
    output_ids[len(input_ids):] for input_ids, output_ids in zip(prompt, outputs)
]

print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])

四、总结

在模型结构上，Gemma2与Qwen2非常相似，除了decoder-only、RoPE、分组查询注意力机制等技术相同，线性层（Lora的目标层）均为

["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj","down_proj"]

中文对话效果上经过多个样例测试个人感觉不如国产的Qwen2、GLM4、DeepSeek等。

GOOGLE作为互联网技术老大哥，在大模型的角逐中，并没有那么强势。可叹啊！

感谢您的阅读，如果喜欢的话，期待您的三连+投票。

如果您还有时间，可以看看我的其他文章：

《AI—工程篇》

AI智能体研发之路-工程篇（一）：Docker助力AI智能体开发提效

AI智能体研发之路-工程篇（二）：Dify智能体开发平台一键部署

AI智能体研发之路-工程篇（三）：大模型推理服务框架Ollama一键部署

AI智能体研发之路-工程篇（四）：大模型推理服务框架Xinference一键部署

AI智能体研发之路-工程篇（五）：大模型推理服务框架LocalAI一键部署

《AI—模型篇》

AI智能体研发之路-模型篇（一）：大模型训练框架LLaMA-Factory在国内网络环境下的安装、部署及使用

AI智能体研发之路-模型篇（二）：DeepSeek-V2-Chat 训练与推理实战

AI智能体研发之路-模型篇（三）：中文大模型开、闭源之争

AI智能体研发之路-模型篇（四）：一文入门pytorch开发

AI智能体研发之路-模型篇（五）：pytorch vs tensorflow框架DNN网络结构源码级对比

AI智能体研发之路-模型篇（六）：【机器学习】基于tensorflow实现你的第一个DNN网络

AI智能体研发之路-模型篇（七）：【机器学习】基于YOLOv10实现你的第一个视觉AI大模型

AI智能体研发之路-模型篇（八）：【机器学习】Qwen1.5-14B-Chat大模型训练与推理实战

AI智能体研发之路-模型篇（九）：【机器学习】GLM4-9B-Chat大模型/GLM-4V-9B多模态大模型概述、原理及推理实战

《AI—Transformers应用》

【AI大模型】Transformers大模型库（一）：Tokenizer

【AI大模型】Transformers大模型库（二）：AutoModelForCausalLM

【AI大模型】Transformers大模型库（三）：特殊标记（special tokens）

【AI大模型】Transformers大模型库（四）：AutoTokenizer