以下文章來源于谷歌云服務，作者 Google Cloud

作者 / 曹治政，Google Cloud 解決方案架構師

Google 近期發布了最新開放模型 Gemma 2，目前與同等規模的開放模型相比，取得了明顯的優勢，同時在安全可控性上得到了顯著的增強。

Gemma 2 提供了 9B 以及 27B 兩種參數規模的選擇，同時分別提供了預訓練模型 (Gemma 2-9B/Gemma 2-27B) 與指令微調模型 (Gemma 2-9B-it/Gemma 2-27B-it)，優異的性能表現原自 Gemma 2 在訓練階段的各項技術創新。

Gemma 2 支持 8192 tokens 的上下文長度，同時使用了旋轉位置編碼 (RoPE)。Gemma 2-9B 在 Google TPUv4 上訓練而成，Gemma 2-27B 在 Google TPUv5p 上訓練而成。

作為開放模型，用戶目前可以在 Hugging Face 以及 Kaggle 上免費獲取模型權重。

1用戶可以選擇先把模型下載到本地或者云端共享存儲，然后將模型部署至推理框架上。

2也可以選擇在部署模型的過程中從 Hugging Face 上下載模型。

我們通常建議使用第一種方案，這種方式可以避免每次部署模型都從公網下載，提升模型加載以及擴展效率。

Gemma 2 生態廣泛，支持通過多種主流框架進行模型的部署，這些框架包括但不限于:

● Hugging Face Transformers

●Keras NLP

●Pytorch

●Gemma C++

●GGUF

●Tensorflow Lite

●TensorRT-LLM

●MaxText

●Pax

●Flax

●vLLM

●oLLama

此外，用戶也可以根據實際需求選擇靈活的方式部署 Gemma 2 模型，這些部署方案通常包括:

1本地或 colab 部署:個人用途，體驗模型功能或者構建個人助理。

2GKE 或第三方 Kubernetes 平臺部署:生產用途，用于將模型和生產業務系統集成，同時充分借助 Kubernetes 平臺資源調度和擴展的優勢。

3Vertex AI 部署:生產用途，用于將模型和生產業務系統集成，同時借助 Google AI PaaS 平臺全托管的優勢。

Google Cloud 為用戶提供了方便的操作服務，用戶可以在 Vertex AI Model Garden 上將 Gemma 2 一鍵部署 GKE 集群或者 Vertex AI endpoint 上，并通過可視化界面對性能和日志進行實時查看，同時也根據請求的流量對推理實例動態的擴縮容。

Gemma 2 支持模型的 fine-tuning，用戶可以利用私有數據或者領域內的特定數據對模型進行 fine-tune，對齊領域內的風格和標準、從而更好地支撐業務場景。

大模型的 fine-tuning 通常包括以下幾種方式:

●Full Parameter fine-tuning

●Lora fine-tuning

●QLora fine-tuning

其中 Full Parameter fine-tuning 需要對模型的全部參數進行重新訓練，Lora fine-tuning 通過訓練少量的低秩矩陣參數提升模型訓練的效率，QLora fine-tuning 在 Lora 的基礎上對模型精度進行了 4bit 或者 8bit 的量化處理來近一步優化模型的存儲占用。

Lora 和 QLora 在一些特定的場景下可以用很低的成本取得非常好的效果，但是對于一些復雜的場景、或者領域數據更加豐富的場景，還是需要通過全量參數的方式進行繼續訓練。

本次演示我們會通過全量參數訓練的方式對模型進行 fine-tuning，同時結合 Deepspeed 框架來實現模型的分布式訓練。

Fine-Tune 演示

a. 微調數據集

本次模型微調，我們會對 Gemma 2-27B 的預訓練基礎模型利用對話數據進行微調。Gemma 2-27B 是預訓練模型，不具備對話能力，通過對比 fine-tune 前后的模型，可以觀察模型訓練的效果。

訓練數據集使用了 Hugging Face 上的開放數據集 "smangrul/ultrachat-10k-chatml"，它一共由 10k 行訓練數據集以及 2k 行測試數據集組成。數據格式如圖所示，每一條數據的 message 列都由一個多輪對話組成。

b. 訓練環境

創建 Vertex AI Workbench，Vertex AI Workbench 是由 Google Cloud 托管的企業級 Jupyter Notebook 服務，它支持多種硬件加速設備，同時預制了依賴包，用戶完成 Workbench 實例的創建后可以通過瀏覽器訪問 Jupyter Notebook 界面，之后快速開啟模型體驗之旅。

△Vertex AI Workbench 實例創建界面

△Vertex AI Workbench 用戶界面

c. 微調步驟

1. 進入 Vertex AI Workbench 界面，創建新文件，運行下面代碼，加載 Gemma 2-27B 基礎模型，并測試模型針對 prompt 的輸出，此步驟目的是為和訓練后的輸出進行訓練效果對比。

from transformers import AutoTokenizer, AutoModelForCausalLM

import transformers

import torch

#加載模型

model_id = "google/gemma-2-27b"
dtype = torch.bfloat16
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="cuda",
torch_dtype=dtype,
)

#準備prompt（對話格式）

prompt = """### Human: I can't cancel the subscription to your corporate newsletter### Assistant: """

#將prompt轉換成token id并輸入model生成內容

inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")

outputs = model.generate(input_ids=inputs.to(model.device), max_new_tokens=100)

#解碼token id，打印輸出

print(tokenizer.decode(outputs[0]))

可以看到模型對話風格生硬，自然對話的效果差。

2. 克隆 peft github 代碼，并切換至模型訓練腳本界面。

!git clone https://github.com/huggingface/peft.git

!cd examples/sft/

Notes

Gemma 2 與開源的 Hugging Face 的 Transformer library 完全兼容，因此用于訓練其它模型的腳本在更改模型名稱后，便可以開始訓練 Gemma 模型。但是需要注意的是，不同模型多輪對話支持的特殊 token 可能不一樣，因此與之相對應的 tokenzier 里面 chat 模板需要提前進行處理。本次我們使用 chatml 格式的模板，只需要在示例腳本上設置對應參數，代碼對會自動對 tokenizer 進行處理。

3. 配置 Hugging Face (用于模型下載) 以及 wandb (用于可視化展示訓練監控指標) 的 token 環境變量。

!echo 'export HUGGINGFACE_TOKEN=YOUR_HF_API_TOKEN' >> ~/.bashrc

!echo 'export WANDB_API_KEY=YOUR_WANDB_API_KEY' >> ~/.bashrc

!source ~/.bashrc

4. 安裝依賴包

!pip install accelerate==0.32.1

!pip install deepspeed==0.14.4

!pip install transformers==4.42.4

!pip install peft==0.11.1

!pip install trl==0.9.6

!pip install bitsandbytes==0.43.1

!pip install wandb==0.17.4

5. 啟動訓練任務，命令會通過 accelerate 啟動一個單機八卡的訓練任務，分布式框架由 deepspeed 實現。訓練任務由 Hugging Face SFTTrainer 實現，它會下載 Hugging Face上的 "smangrul/ultrachat-10k-chatml" 數據集，并讀取 content 字段。

訓練任務采用 bf16 混合精度，由于全量參數微調占用顯存很大，因此每個 GPU device 的 batch size 設置為 1 (激活值非常消耗內存，若大于 1 會導致 OOM)，通過將梯度累加設置將單個設備的有效 batch size 提升為 4，同時開啟 gradient_checkpointing 節省激活值對內存的占用。

任務在訓練過程中，指標數據會實時傳送至 wandb，整個訓練周期為 1 個 epoch，訓練任務結束后，模型會自動存放至 gemma2-27b-sft-deepspeed路徑下。

# 任務訓練任務啟動命令

!accelerate launch --config_file "configs/deepspeed_config.yaml" train.py
--seed 100
--model_name_or_path "google/gemma-2-27b"
--dataset_name "smangrul/ultrachat-10k-chatml"
--chat_template_format "chatml"
--add_special_tokens False
--append_concat_token False
--splits "train,test"
--max_seq_len 2048
--num_train_epochs 1
--logging_steps 5
--log_level "info"
--logging_strategy "steps"
--evaluation_strategy "epoch"
--save_strategy "epoch"
--bf16 True
--learning_rate 1e-4
--lr_scheduler_type "cosine"
--weight_decay 1e-4
--warmup_ratio 0.0
--max_grad_norm 1.0
--output_dir "gemma2-27b-sft-deepspeed"
--per_device_train_batch_size 1
--per_device_eval_batch_size 1
--gradient_accumulation_steps 4
--gradient_checkpointing True
--use_reentrant False
--dataset_text_field "content"

# Deespeed配置腳本

compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
deepspeed_multinode_launcher: standard
gradient_accumulation_steps: 4
offload_optimizer_device: none
offload_param_device: none
zero3_init_flag: true
zero3_save_16bit_model: true
zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

6. 觀察任務訓練指標，等待任務訓練結束。本次訓練，測試數據由 10k 行組成，由于我們采用了 8 張卡訓練，單個卡的 batch size 設置為 1，梯度累積設置為 4，因此 global batch size 為 1*4*8=32，單個 epoch 訓練 step 數量為: 10000/32= 312 steps。從損失值圖表中可以看到，當訓練任務至 20 步左右時，損失值開始快速下降，之后平穩收斂直到訓練任務結束。

大約 41 分鐘，訓練任務結束。

7. 重新加載訓練后的模型，并驗證訓練效果。從內容輸出中可以看到，模型很好地學習到了數據集生成對話的形式和總結風格，能夠以更加流暢、自然的方式進行對話。

from transformers import AutoTokenizer, AutoModelForCausalLM
import transformers
import torch

model_id = "./gemma2-27b-sft-deepspeed"
dtype = torch.bfloat16

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="cuda",
torch_dtype=dtype,
)

prompt = """### Human: I can't cancel the subscription to your corporate newsletter### Assistant: """

inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
outputs = model.generate(input_ids=inputs.to(model.device), max_new_tokens=100)
print(tokenizer.decode(outputs[0]))

以上介紹了 Gemma-2-27B fine-tune 的一個簡單的實驗演示，通過實驗可以看出 Gemma 2 與主流的訓練框架完全兼容，對于已經了解如何 fine-tune 其它模型的開發人員可以快速切換至 Gemma 2 完成模型訓練。

Google Cloud 還提供了更多高級的解決方案，可以幫助用戶通過 Vertex AI Pipeline 或者 Ray On GKE 等方式大規模、自動化地訓練開放模型， 感興趣的讀者可以聯系 Google Cloud 的工作人員獲得更近一步的了解。