以下文章來(lái)源于谷歌云服務(wù)，作者 Google Cloud

作者 / 曹治政，Google Cloud 解決方案架構(gòu)師

Google 近期發(fā)布了最新開(kāi)放模型 Gemma 2，目前與同等規(guī)模的開(kāi)放模型相比，取得了明顯的優(yōu)勢(shì)，同時(shí)在安全可控性上得到了顯著的增強(qiáng)。

Gemma 2 提供了 9B 以及 27B 兩種參數(shù)規(guī)模的選擇，同時(shí)分別提供了預(yù)訓(xùn)練模型 (Gemma 2-9B/Gemma 2-27B) 與指令微調(diào)模型 (Gemma 2-9B-it/Gemma 2-27B-it)，優(yōu)異的性能表現(xiàn)原自 Gemma 2 在訓(xùn)練階段的各項(xiàng)技術(shù)創(chuàng)新。

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

作為開(kāi)放模型，用戶目前可以在 Hugging Face 以及 Kaggle 上免費(fèi)獲取模型權(quán)重。

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

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

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

Gemma 2 生態(tài)廣泛，支持通過(guò)多種主流框架進(jìn)行模型的部署，這些框架包括但不限于:

● Hugging Face Transformers

●Keras NLP

●Pytorch

●Gemma C++

●GGUF

●Tensorflow Lite

●TensorRT-LLM

●MaxText

●Pax

●Flax

●vLLM

●oLLama

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

1本地或 colab 部署:個(gè)人用途，體驗(yàn)?zāi)Ｐ凸δ芑蛘邩?gòu)建個(gè)人助理。

2GKE 或第三方 Kubernetes 平臺(tái)部署:生產(chǎn)用途，用于將模型和生產(chǎn)業(yè)務(wù)系統(tǒng)集成，同時(shí)充分借助 Kubernetes 平臺(tái)資源調(diào)度和擴(kuò)展的優(yōu)勢(shì)。

3Vertex AI 部署:生產(chǎn)用途，用于將模型和生產(chǎn)業(yè)務(wù)系統(tǒng)集成，同時(shí)借助 Google AI PaaS 平臺(tái)全托管的優(yōu)勢(shì)。

Google Cloud 為用戶提供了方便的操作服務(wù)，用戶可以在 Vertex AI Model Garden 上將 Gemma 2 一鍵部署 GKE 集群或者 Vertex AI endpoint 上，并通過(guò)可視化界面對(duì)性能和日志進(jìn)行實(shí)時(shí)查看，同時(shí)也根據(jù)請(qǐng)求的流量對(duì)推理實(shí)例動(dòng)態(tài)的擴(kuò)縮容。

Gemma 2 支持模型的 fine-tuning，用戶可以利用私有數(shù)據(jù)或者領(lǐng)域內(nèi)的特定數(shù)據(jù)對(duì)模型進(jìn)行 fine-tune，對(duì)齊領(lǐng)域內(nèi)的風(fēng)格和標(biāo)準(zhǔn)、從而更好地支撐業(yè)務(wù)場(chǎng)景。

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

●Full Parameter fine-tuning

●Lora fine-tuning

●QLora fine-tuning

其中 Full Parameter fine-tuning 需要對(duì)模型的全部參數(shù)進(jìn)行重新訓(xùn)練，Lora fine-tuning 通過(guò)訓(xùn)練少量的低秩矩陣參數(shù)提升模型訓(xùn)練的效率，QLora fine-tuning 在 Lora 的基礎(chǔ)上對(duì)模型精度進(jìn)行了 4bit 或者 8bit 的量化處理來(lái)近一步優(yōu)化模型的存儲(chǔ)占用。

Lora 和 QLora 在一些特定的場(chǎng)景下可以用很低的成本取得非常好的效果，但是對(duì)于一些復(fù)雜的場(chǎng)景、或者領(lǐng)域數(shù)據(jù)更加豐富的場(chǎng)景，還是需要通過(guò)全量參數(shù)的方式進(jìn)行繼續(xù)訓(xùn)練。

本次演示我們會(huì)通過(guò)全量參數(shù)訓(xùn)練的方式對(duì)模型進(jìn)行 fine-tuning，同時(shí)結(jié)合 Deepspeed 框架來(lái)實(shí)現(xiàn)模型的分布式訓(xùn)練。

Fine-Tune 演示

a. 微調(diào)數(shù)據(jù)集

本次模型微調(diào)，我們會(huì)對(duì) Gemma 2-27B 的預(yù)訓(xùn)練基礎(chǔ)模型利用對(duì)話數(shù)據(jù)進(jìn)行微調(diào)。Gemma 2-27B 是預(yù)訓(xùn)練模型，不具備對(duì)話能力，通過(guò)對(duì)比 fine-tune 前后的模型，可以觀察模型訓(xùn)練的效果。

訓(xùn)練數(shù)據(jù)集使用了 Hugging Face 上的開(kāi)放數(shù)據(jù)集 "smangrul/ultrachat-10k-chatml"，它一共由 10k 行訓(xùn)練數(shù)據(jù)集以及 2k 行測(cè)試數(shù)據(jù)集組成。數(shù)據(jù)格式如圖所示，每一條數(shù)據(jù)的 message 列都由一個(gè)多輪對(duì)話組成。

b. 訓(xùn)練環(huán)境

創(chuàng)建 Vertex AI Workbench，Vertex AI Workbench 是由 Google Cloud 托管的企業(yè)級(jí) Jupyter Notebook 服務(wù)，它支持多種硬件加速設(shè)備，同時(shí)預(yù)制了依賴包，用戶完成 Workbench 實(shí)例的創(chuàng)建后可以通過(guò)瀏覽器訪問(wèn) Jupyter Notebook 界面，之后快速開(kāi)啟模型體驗(yàn)之旅。

△Vertex AI Workbench 實(shí)例創(chuàng)建界面

△Vertex AI Workbench 用戶界面

c. 微調(diào)步驟

1. 進(jìn)入 Vertex AI Workbench 界面，創(chuàng)建新文件，運(yùn)行下面代碼，加載 Gemma 2-27B 基礎(chǔ)模型，并測(cè)試模型針對(duì) prompt 的輸出，此步驟目的是為和訓(xùn)練后的輸出進(jìn)行訓(xùn)練效果對(duì)比。

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,
)

#準(zhǔn)備prompt（對(duì)話格式）

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

#將prompt轉(zhuǎn)換成token id并輸入model生成內(nèi)容

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]))

可以看到模型對(duì)話風(fēng)格生硬，自然對(duì)話的效果差。

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

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

!cd examples/sft/

Notes

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

3. 配置 Hugging Face (用于模型下載) 以及 wandb (用于可視化展示訓(xùn)練監(jiān)控指標(biāo)) 的 token 環(huán)境變量。

!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. 啟動(dòng)訓(xùn)練任務(wù)，命令會(huì)通過(guò) accelerate 啟動(dòng)一個(gè)單機(jī)八卡的訓(xùn)練任務(wù)，分布式框架由 deepspeed 實(shí)現(xiàn)。訓(xùn)練任務(wù)由 Hugging Face SFTTrainer 實(shí)現(xiàn)，它會(huì)下載 Hugging Face上的 "smangrul/ultrachat-10k-chatml" 數(shù)據(jù)集，并讀取 content 字段。

訓(xùn)練任務(wù)采用 bf16 混合精度，由于全量參數(shù)微調(diào)占用顯存很大，因此每個(gè) GPU device 的 batch size 設(shè)置為 1 (激活值非常消耗內(nèi)存，若大于 1 會(huì)導(dǎo)致 OOM)，通過(guò)將梯度累加設(shè)置將單個(gè)設(shè)備的有效 batch size 提升為 4，同時(shí)開(kāi)啟 gradient_checkpointing 節(jié)省激活值對(duì)內(nèi)存的占用。

任務(wù)在訓(xùn)練過(guò)程中，指標(biāo)數(shù)據(jù)會(huì)實(shí)時(shí)傳送至 wandb，整個(gè)訓(xùn)練周期為 1 個(gè) epoch，訓(xùn)練任務(wù)結(jié)束后，模型會(huì)自動(dòng)存放至 gemma2-27b-sft-deepspeed路徑下。

# 任務(wù)訓(xùn)練任務(wù)啟動(dòng)命令

!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. 觀察任務(wù)訓(xùn)練指標(biāo)，等待任務(wù)訓(xùn)練結(jié)束。本次訓(xùn)練，測(cè)試數(shù)據(jù)由 10k 行組成，由于我們采用了 8 張卡訓(xùn)練，單個(gè)卡的 batch size 設(shè)置為 1，梯度累積設(shè)置為 4，因此 global batch size 為 1*4*8=32，單個(gè) epoch 訓(xùn)練 step 數(shù)量為: 10000/32= 312 steps。從損失值圖表中可以看到，當(dāng)訓(xùn)練任務(wù)至 20 步左右時(shí)，損失值開(kāi)始快速下降，之后平穩(wěn)收斂直到訓(xùn)練任務(wù)結(jié)束。

大約 41 分鐘，訓(xùn)練任務(wù)結(jié)束。

7. 重新加載訓(xùn)練后的模型，并驗(yàn)證訓(xùn)練效果。從內(nèi)容輸出中可以看到，模型很好地學(xué)習(xí)到了數(shù)據(jù)集生成對(duì)話的形式和總結(jié)風(fēng)格，能夠以更加流暢、自然的方式進(jìn)行對(duì)話。

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 的一個(gè)簡(jiǎn)單的實(shí)驗(yàn)演示，通過(guò)實(shí)驗(yàn)可以看出 Gemma 2 與主流的訓(xùn)練框架完全兼容，對(duì)于已經(jīng)了解如何 fine-tune 其它模型的開(kāi)發(fā)人員可以快速切換至 Gemma 2 完成模型訓(xùn)練。

Google Cloud 還提供了更多高級(jí)的解決方案，可以幫助用戶通過(guò) Vertex AI Pipeline 或者 Ray On GKE 等方式大規(guī)模、自動(dòng)化地訓(xùn)練開(kāi)放模型， 感興趣的讀者可以聯(lián)系 Google Cloud 的工作人員獲得更近一步的了解。