待整理…

gemma介紹

• Gemma是Google推出的一系列輕量級、最先進的開放模型，基于創(chuàng)建Gemini模型的相同研究和技術(shù)構(gòu)建。提供了 2B 和 7B 兩種不同規(guī)模的版本，每種都包含了預(yù)訓(xùn)練基礎(chǔ)版本和經(jīng)過指令優(yōu)化的版本。所有版本均可在各類消費級硬件上運行，無需數(shù)據(jù)量化處理，擁有高達 8K tokens 的處理能力：

• 它們是文本到文本的、僅解碼器的大型語言模型，提供英語版本，具有開放的權(quán)重、預(yù)訓(xùn)練的變體和指令調(diào)優(yōu)的變體。

• Gemma模型非常適合執(zhí)行各種文本生成任務(wù)，包括問答、摘要和推理。它們相對較小的尺寸使得可以在資源有限的環(huán)境中部署，例如筆記本電腦、桌面電腦或您自己的云基礎(chǔ)設(shè)施，使每個人都能獲得最先進的AI模型，促進創(chuàng)新。

• Github地址：https://github.com/google/gemma_pytorch

• 論文地址：https://storage./deepmind-media/gemma/gemma-report.pdf

• 官方博客：Gemma: Google introduces new state-of-the-art open models

其它

BPE。

又稱 digram coding 雙字母組合編碼，是一種數(shù)據(jù)壓縮 算法，用來在固定大小的詞表中實現(xiàn)可變?度的子詞。該算法簡單有效，因而目前它是最流行的方法。

peft

GPT，gemma，BERT這些大型預(yù)訓(xùn)練模型的訓(xùn)練成本非常高昂，需要龐大的計算資源和大量的數(shù)據(jù)，一般人難以承受。這也導(dǎo)致了一些研究人員難以重復(fù)和驗證先前的研究成果。為了解決這個問題，研究人員開始研究Parameter-Efficient Fine-Tuning (PEFT)技術(shù)。

Adapter Tuning

谷歌的研究人員首次在論文《Parameter-Efficient Transfer Learning for NLP》提出針對BERT的PEFT微調(diào)方式，拉開了PEFT研究的序幕。他們指出，在面對特定的下游任務(wù)時，如果進行Full-fintuning（即預(yù)訓(xùn)練模型中的所有參數(shù)都進行微調(diào)），太過低效；而如果采用固定預(yù)訓(xùn)練模型的某些層，只微調(diào)接近下游任務(wù)的那幾層參數(shù)，又難以達到較好的效果。

于是他們設(shè)計了如下圖所示的Adapter結(jié)構(gòu)，將其嵌入Transformer的結(jié)構(gòu)里面，在訓(xùn)練時，固定住原來預(yù)訓(xùn)練模型的參數(shù)不變，只對新增的Adapter結(jié)構(gòu)進行微調(diào)。同時為了保證訓(xùn)練的高效性（也就是盡可能少的引入更多參數(shù)），他們將Adapter設(shè)計為這樣的結(jié)構(gòu)：首先是一個down-project層將高維度特征映射到低維特征，然后過一個非線形層之后，再用一個up-project結(jié)構(gòu)將低維特征映射回原來的高維特征；同時也設(shè)計了skip-connection結(jié)構(gòu)，確保了在最差的情況下能夠退化為identity。

從實驗結(jié)果來看，該方法能夠在只額外對增加的3.6%參數(shù)規(guī)模（相比原來預(yù)訓(xùn)練模型的參數(shù)量）的情況下取得和Full-finetuning接近的效果（GLUE指標(biāo)在0.4%以內(nèi)）。

Prefix Tuning

Prefix Tuning方法由斯坦福的研究人員提出，與Full-finetuning更新所有參數(shù)的方式不同，該方法是在輸入token之前構(gòu)造一段任務(wù)相關(guān)的virtual tokens作為Prefix，然后訓(xùn)練的時候只更新Prefix部分的參數(shù)，而Transformer中的其他部分參數(shù)固定。該方法其實和構(gòu)造Prompt類似，只是Prompt是人為構(gòu)造的“顯式”的提示,并且無法更新參數(shù)，而Prefix則是可以學(xué)習(xí)的“隱式”的提示。

同時，為了防止直接更新Prefix的參數(shù)導(dǎo)致訓(xùn)練不穩(wěn)定的情況，他們在Prefix層前面加了MLP結(jié)構(gòu)(相當(dāng)于將Prefix分解為更小維度的Input與MLP的組合后輸出的結(jié)果)，訓(xùn)練完成后，只保留Prefix的參數(shù)。

Prompt Tuning

《The Power of Scale for Parameter-Efficient Prompt Tuning》
總的來說就是，只要模型規(guī)模夠大，簡單加入Prompt tokens進行微調(diào)，就能取得很好的效果。
該方法可以看作是Prefix Tuning的簡化版本，只在輸入層加入prompt tokens，并不需要加入MLP進行調(diào)整來解決難訓(xùn)練的問題，主要在T5預(yù)訓(xùn)練模型上做實驗。似乎只要預(yù)訓(xùn)練模型足夠強大，其他的一切都不是問題。作者也做實驗說明隨著預(yù)訓(xùn)練模型參數(shù)量的增加，Prompt Tuning的方法會逼近Fine-tune的結(jié)果。
作者做了一系列對比實驗，都在說明：隨著預(yù)訓(xùn)練模型參數(shù)的增加，一切的問題都不是問題，最簡單的設(shè)置也能達到極好的效果。

a) Prompt長度影響：模型參數(shù)達到一定量級時，Prompt長度為1也能達到不錯的效果，Prompt長度為20就能達到極好效果。
b) Prompt初始化方式影響：Random Uniform方式明顯弱于其他兩種，但是當(dāng)模型參數(shù)達到一定量級，這種差異也不復(fù)存在。
c) 預(yù)訓(xùn)練的方式：LM Adaptation的方式效果好，但是當(dāng)模型達到一定規(guī)模，差異又幾乎沒有了。
d) 微調(diào)步數(shù)影響：模型參數(shù)較小時，步數(shù)越多，效果越好。同樣隨著模型參數(shù)達到一定規(guī)模，zero shot也能取得不錯效果。

P-Tuning

P-Tuning方法的提出主要是為了解決這樣一個問題：大模型的Prompt構(gòu)造方式嚴重影響下游任務(wù)的效果。
P-Tuning提出將Prompt轉(zhuǎn)換為可以學(xué)習(xí)的Embedding層，只是考慮到直接對Embedding參數(shù)進行優(yōu)化會存在這樣兩個挑戰(zhàn)：

• Discretenes： 對輸入正常語料的Embedding層已經(jīng)經(jīng)過預(yù)訓(xùn)練，而如果直接對輸入的prompt embedding進行隨機初始化訓(xùn)練，容易陷入局部最優(yōu)。
• Association：沒法捕捉到prompt embedding之間的相關(guān)關(guān)系。
  作者在這里提出用MLP+LSTM的方式來對prompt embedding進行一層處理

與Prefix-Tuning的區(qū)別

這篇文章（2021-03）和Prefix-Tuning（2021-01）差不多同時提出，做法其實也有一些相似之處，主要區(qū)別在

• Prefix Tuning是將額外的embedding加在開頭，看起來更像是模仿Instruction指令；而P-Tuning的位置則不固定。
• Prefix Tuning通過在每個Attention層都加入Prefix Embedding來增加額外的參數(shù)，通過MLP來初始化；而P-Tuning只是在輸入的時候加入Embedding，并通過LSTM+MLP來初始化。

V2

論文《P-Tuning v2: Prompt Tuning Can Be Comparable to Fine-tuning Universally Across Scales and Tasks》
P-Tuning v2的目標(biāo)就是要讓Prompt Tuning能夠在不同參數(shù)規(guī)模的預(yù)訓(xùn)練模型、針對不同下游任務(wù)的結(jié)果上都達到匹敵Fine-tuning的結(jié)果。
那也就是說當(dāng)前Prompt Tuning方法未能在這兩個方面都存在局限性。

• 不同模型規(guī)模：Prompt Tuning和P-tuning這兩種方法都是在預(yù)訓(xùn)練模型參數(shù)規(guī)模夠足夠大時，才能達到和Fine-tuning類似的效果，而參數(shù)規(guī)模較小時效果則很差。
• 不同任務(wù)類型：Prompt Tuning和P-tuning這兩種方法在sequence tagging任務(wù)上表現(xiàn)都很差。

主要結(jié)構(gòu)

相比Prompt Tuning和P-tuning的方法， P-tuning v2方法在多層加入了Prompts tokens作為輸入，帶來兩個方面的好處：

1. 帶來更多可學(xué)習(xí)的參數(shù)（從P-tuning和Prompt Tuning的0.1%增加到0.1%-3%），同時也足夠parameter-efficient。
2. 加入到更深層結(jié)構(gòu)中的Prompt能給模型預(yù)測帶來更直接的影響。

幾個關(guān)鍵設(shè)計因素

1. Reparameterization：Prefix Tuning和P-tuning中都有MLP來構(gòu)造可訓(xùn)練的embedding。本文發(fā)現(xiàn)在自然語言理解領(lǐng)域，面對不同的任務(wù)以及不同的數(shù)據(jù)集，這種方法可能帶來完全相反的結(jié)論。
2. Prompt Length： 不同的任務(wù)對應(yīng)的最合適的Prompt Length不一樣，比如簡單分類任務(wù)下length=20最好，而復(fù)雜的任務(wù)需要更長的Prompt Length。
3. Multi-task Learning 多任務(wù)對于P-Tuning v2是可選的，但可以利用它提供更好的初始化來進一步提高性能。
4. Classification Head 使用LM head來預(yù)測動詞是Prompt Tuning的核心，但我們發(fā)現(xiàn)在完整的數(shù)據(jù)設(shè)置中沒有必要這樣做，并且這樣做與序列標(biāo)記不兼容。P-tuning v2采用和BERT一樣的方式，在第一個token處應(yīng)用隨機初始化的分類頭。

LoRA

微軟和CMU的研究者指出，現(xiàn)有的一些PEFT的方法還存在這樣一些問題：

• 由于增加了模型的深度從而額外增加了模型推理的延時，如Adapter方法
• Prompt較難訓(xùn)練，同時減少了模型的可用序列長度，如Prompt Tuning、Prefix Tuning、P-Tuning方法
• 往往效率和質(zhì)量不可兼得，效果差于full-finetuning。

有研究者對語言模型的參數(shù)進行研究發(fā)現(xiàn)：語言模型雖然參數(shù)眾多，但是起到關(guān)鍵作用的還是其中低秩的本質(zhì)維度（low instrisic dimension）。本文受到該觀點的啟發(fā)，提出了Low-Rank Adaption(LoRA)，設(shè)計了如下所示的結(jié)構(gòu)，在涉及到矩陣相乘的模塊，引入A、B這樣兩個低秩矩陣模塊去模擬Full-finetune的過程，相當(dāng)于只對語言模型中起關(guān)鍵作用的低秩本質(zhì)維度進行更新。

這么做就能完美解決以上存在的3個問題：

• 相比于原始的Adapter方法“額外”增加網(wǎng)絡(luò)深度，必然會帶來推理過程額外的延遲，該方法可以在推理階段直接用訓(xùn)練好的A、B矩陣參數(shù)與原預(yù)訓(xùn)練模型的參數(shù)相加去替換原有預(yù)訓(xùn)練模型的參數(shù)，這樣的話推理過程就相當(dāng)于和Full-finetune一樣，沒有額外的計算量，從而不會帶來性能的損失。
• 由于沒有使用Prompt方式，自然不會存在Prompt方法帶來的一系列問題。
• 該方法由于實際上相當(dāng)于是用LoRA去模擬Full-finetune的過程，幾乎不會帶來任何訓(xùn)練效果的損失，后續(xù)的實驗結(jié)果也證明了這一點。
  在實驗中，研究人員將這一LoRA模塊與Transformer的attention模塊相結(jié)合，在RoBERTa 、DeBERTa、GPT-2和GPT-3 175B這幾個大模型上都做了實驗，實驗結(jié)果也充分證明了該方法的有效性。

UniPELT

Towards a Unified View of PETL
這篇ICLR2022的文章研究了典型的PEFT方法，試圖將PEFT統(tǒng)一到一個框架下，找出它們起作用的具體原因，并進行改進。主要研究了三個問題：

• 典型的PEFT方法有什么聯(lián)系？
• 典型的PEFT方法中是哪些關(guān)鍵模塊在起作用？
• 能否對這些關(guān)鍵模塊進行排列組合，找出更有用的PEFT方法？
  通過對Prefix Tuning的推導(dǎo)，得出了和Adapter Tuning以及LoRA形式一致的形式。
  根據(jù)這個統(tǒng)一的框架，還另外設(shè)計了三種變體Parallel Adapter、Multi-head Parallel Adapter、Scaled Parallel Adapter。

QLoRA

Efficient Finetuning of Quantized LLMs
這篇論文由 UW 的四位作者發(fā)表于 23 年5月底，一經(jīng)發(fā)布就大火，在算力昂貴的時代，它這個新技術(shù)可以節(jié)約數(shù)十倍的顯存且?guī)缀鯖]有損失性能。該技術(shù)降低了推理和訓(xùn)練大模型的門檻，讓更廣大的性價比玩家看到了希望。
新的微調(diào)方法取名為QLoRA，它可以顯著減少內(nèi)存使用，使得可以在單個48GB的GPU上微調(diào)一個有650億參數(shù)的模型，同時保持16比特微調(diào)的性能。QLoRA通過凍結(jié)的、4比特量化的預(yù)訓(xùn)練語言模型來做 LoRA，進行反向傳播梯度。作者提出的最佳模型命名為Guanaco，在Vicuna基準(zhǔn)測試中超越了所有以前公開發(fā)布的模型，達到了ChatGPT性能水平的99.3%，而只需要在單個GPU上微調(diào)24小時。
QLORA引入了幾項創(chuàng)新來節(jié)省內(nèi)存而不犧牲性能，包括

• 一種新的數(shù)據(jù)類型4位NormalFloat（NF4）；
• 雙重量化以減少平均內(nèi)存占用；
• 分頁優(yōu)化器來管理內(nèi)存峰值。
  作者使用 QLoRA 技術(shù)微調(diào)了 1000 多個模型來實驗，在一些小型高質(zhì)量數(shù)據(jù)集上可以做到 SOTA，甚至用更小的模型也能做到 SOTA。

4-bit NormalFloat Quantization

NormalFloat (簡稱NF)是一種數(shù)據(jù)類型，它是建立在 Quantile quantization（后譯為分位數(shù)量化）基礎(chǔ)上的，它是一種信息論上最優(yōu)的數(shù)據(jù)類型，可以確保每個量化區(qū)間從輸入張量中分配相同數(shù)量的值。分位數(shù)量化通過經(jīng)驗累積分布函數(shù)估計輸入張量的分位數(shù)來工作。分位數(shù)量化的主要局限性在于分位數(shù)估計的這個過程會比較費力。

在神經(jīng)網(wǎng)絡(luò)中，預(yù)訓(xùn)練的權(quán)重通常具有零中心的正態(tài)分布，標(biāo)準(zhǔn)差為σ。通過縮放σ，可以使得分布恰好適應(yīng)NF的范圍。對于NF，作者設(shè)置了一個任意的范圍[-1, 1]。因此，數(shù)據(jù)類型和神經(jīng)網(wǎng)絡(luò)權(quán)重的分位數(shù)都需要被歸一化到這個范圍。

對于范圍在[-1, 1]內(nèi)的零均值正態(tài)分布，他們計算了信息理論上最優(yōu)的數(shù)據(jù)類型。這個過程包括：(1) 估計理論N(0, 1)分布的

• 1 個分位數(shù)，得到一個k位的分位數(shù)量化數(shù)據(jù)類型；(2) 將這個NF的值歸一化到[-1, 1]范圍；(3) 通過絕對最大值重標(biāo)定，將輸入權(quán)重張量歸一化到[-1, 1]范圍，然后進行量化。一旦模型權(quán)重范圍和NF范圍匹配，就可以像通常那樣進行量化。這個過程等價于重新縮放權(quán)重張量的標(biāo)準(zhǔn)差，使其匹配k位數(shù)據(jù)類型的標(biāo)準(zhǔn)差。更具體的來看這個公式，展示了
  到分位數(shù)的映射公式：
  
  公式里的 Q x Q_x Qx? 指的是分位數(shù)函數(shù)。在讀到這個公式之后，突然感覺就恍然大悟了有沒有！

一些實驗

哪種嵌入形式更好：Parallel or Sequencial?

答案是：Parallel更好

對哪塊結(jié)構(gòu)做修改更好？Attention or FFN？

當(dāng)微調(diào)的參數(shù)量較多時，從結(jié)果來看，對FFN層進行修改更好。一種可能的解釋是FFN層學(xué)到的是任務(wù)相關(guān)的文本模式，而Attention層學(xué)到的是成對的位置交叉關(guān)系，針對新任務(wù)并不需要進行大規(guī)模調(diào)整。

哪種組合方式效果更好？

從結(jié)果來看，縮放式的組合效果更好。

FlashAttention

是一種重新排序注意力計算的算法，它無需任何近似即可加速注意力計算并減少內(nèi)存占用。

演示

可以在 Hugging Chat 上體驗與 Gemma 指令模型的互動對話！點擊此處訪問：https:///chat?model=google/gemma-7b-it

模型下載：

模型文件這里是從modelscope 下載的，也可以從Huggingface 下，Huggingface沒有modelscope方便。

gemma-2b地址：

https:///models/AI-ModelScope/gemma-2b/files

gemma-7b地址:

https:///models/AI-ModelScope/gemma-7b/files

當(dāng)然還有“gemma-7b-it”等其它版本。

Prompt 提示詞格式

Gemma 的基礎(chǔ)模型不限定特定的提示格式。如同其他基礎(chǔ)模型，它們能夠根據(jù)輸入序列生成一個合理的續(xù)接內(nèi)容，適用于零樣本或少樣本的推理任務(wù)。這些模型也為針對特定應(yīng)用場景的微調(diào)提供了堅實的基礎(chǔ)。指令優(yōu)化版本則采用了一種極其簡潔的對話結(jié)構(gòu)：

<start_of_turn>用戶_ _
敲擊<end_of_turn> 
<start_of_turn>模型????
誰在那里< end_of_turn > 
< start_of_turn >用戶
Gemma <end_of_turn> 
<start_of_turn>模型_ _ _ _
杰瑪是誰？<轉(zhuǎn)彎結(jié)束>

要有效利用這一格式，必須嚴格按照上述結(jié)構(gòu)進行對話。

使用 Transformers

借助 Transformers 的 4.38 版本，你可以輕松地使用 Gemma 模型，并充分利用 Hugging Face 生態(tài)系統(tǒng)內(nèi)的工具，包括：

• 訓(xùn)練和推理腳本及示例
• 安全文件格式（safetensors）
• 集成了諸如 bitsandbytes（4位量化）、PEFT（參數(shù)效率微調(diào)）和 Flash Attention 2 等工具
• 輔助工具和幫助器，以便使用模型進行生成
• 導(dǎo)出模型以便部署的機制
  另外，Gemma 模型支持 torch.compile() 與 CUDA 圖的結(jié)合使用，在推理時可實現(xiàn)約 4 倍的速度提升！

確保你使用的是最新版本的 transformers：pip install -U transformers

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

# Hugging Face model id
model_id = './gemma-7b'

# BitsAndBytesConfig int-4 config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type='nf4', bnb_4bit_compute_dtype=torch.bfloat16
)

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map='auto',
    attn_implementation='flash_attention_2',
    torch_dtype=torch.bfloat16,
    quantization_config=bnb_config
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.padding_side = 'right' # to prevent warnings

input_text = 'Write me a poem about Machine Learning.'
input_ids = tokenizer(input_text, return_tensors='pt').to(model.device)

outputs = model.generate(**input_ids)
print(tokenizer.decode(outputs[0]))

示例

from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained('./gemma-7b')
model = AutoModelForCausalLM.from_pretrained('./gemma-7b', device_map='auto')

input_text = 'Write me a poem about Machine Learning.'
input_ids = tokenizer(input_text, return_tensors='pt').to(model.device)

outputs = model.generate(**input_ids)
print(tokenizer.decode(outputs[0]))

微調(diào) Gemma 模型

低秩適應(yīng)（LoRA）是一種用于大語言模型（LLM）的參數(shù)高效微調(diào)技術(shù)。它只針對模型參數(shù)的一小部分進行微調(diào)，通過凍結(jié)原始模型并只訓(xùn)練被分解為低秩矩陣的適配器層。PEFT 庫 提供了一個簡易的抽象，允許用戶選擇應(yīng)用適配器權(quán)重的模型層。
我們將所有的 nn.Linear 層視為要適應(yīng)的目標(biāo)層。

from peft import LoraConfig

lora_config = LoraConfig(
    r=8,
    target_modules=['q_proj', 'o_proj', 'k_proj', 'v_proj', 'gate_proj', 'up_proj', 'down_proj'],
    task_type='CAUSAL_LM',
)

我們首先下載模型和分詞器 (tokenizer)，其中包含了一個 BitsAndBytesConfig 用于僅限權(quán)重的量化。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

model_id = './gemma-2b'
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type='nf4',
    bnb_4bit_compute_dtype=torch.bfloat16
)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map={'':0})

在開始微調(diào)前，我們先使用一個相當(dāng)熟知的名言來測試一下 Gemma 模型：

text = 'Quote: Imagination is more'
device = 'cuda:0'
inputs = tokenizer(text, return_tensors='pt').to(device)

outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

模型完成了一個合理的補全，盡管有一些額外的 token:

Quote: Imagination is more important than knowledge. Knowledge is limited. Imagination encircles the world.

-Albert Einstein

I

但這并不完全是我們希望看到的答案格式。我們將嘗試通過微調(diào)讓模型學(xué)會以我們期望的格式來產(chǎn)生答案:

Quote: Imagination is more important than knowledge. Knowledge is limited. Imagination encircles the world.

Author: Albert Einstein

首先，我們選擇一個英文“名人名言”數(shù)據(jù)集：

這個數(shù)據(jù)集我是下載下來使用的。

from datasets import load_dataset

data = load_dataset('./Abirate/english_quotes')
data = data.map(lambda samples: tokenizer(samples['quote']), batched=True)

接下來，我們使用上述 LoRA 配置對模型進行微調(diào)：

import transformers
from trl import SFTTrainer

def formatting_func(example):
    text = f'Quote: {example['quote'][0]}\nAuthor: {example['author'][0]}'
    return [text]

trainer = SFTTrainer(
    model=model,
    train_dataset=data['train'],
    args=transformers.TrainingArguments(
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        warmup_steps=2,
        max_steps=10,
        learning_rate=2e-4,
        fp16=True,
        logging_steps=1,
        output_dir='outputs',
        optim='paged_adamw_8bit'
    ),
    peft_config=lora_config,
    formatting_func=formatting_func,
)
trainer.train()

最終，我們再次使用先前的提示詞，來測試模型：

text = 'Quote: Imagination is'
device = 'cuda:0'
inputs = tokenizer(text, return_tensors='pt').to(device)

outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

這次，我們得到了我們期待的答案格式：

Quote: Imagination is more important than knowledge. Knowledge is limited. Imagination encircles the world.

Author: Albert Einstein

中文的意思是：

名言：想象力比知識更重要，因為知識是有限的，而想象力概括著世界的一切.

作者：阿爾伯特·愛因斯坦

在 TPU 環(huán)境下微調(diào)，可通過 SPMD 上的 FSDP 加速

from transformers import DataCollatorForLanguageModeling, Trainer, TrainingArguments

# Set up the FSDP config. To enable FSDP via SPMD, set xla_fsdp_v2 to True.
fsdp_config = {
    'fsdp_transformer_layer_cls_to_wrap': ['GemmaDecoderLayer'],
    'xla': True,
    'xla_fsdp_v2': True,
    'xla_fsdp_grad_ckpt': True
}

# Finally, set up the trainer and train the model.
trainer = Trainer(
    model=model,
    train_dataset=data,
    args=TrainingArguments(
        per_device_train_batch_size=64,  # This is actually the global batch size for SPMD.
        num_train_epochs=100,
        max_steps=-1,
        output_dir='./output',
        optim='adafactor',
        logging_steps=1,
        dataloader_drop_last = True,  # Required for SPMD.
        fsdp='full_shard',
        fsdp_config=fsdp_config,
    ),
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
trainer.train()

最后的結(jié)果如下：

程序最后放到一起是：可以根據(jù)需要自行整理。

gpu環(huán)境：

from peft import LoraConfig
lora_config = LoraConfig(
    r=8,
    target_modules=['q_proj', 'o_proj', 'k_proj', 'v_proj', 'gate_proj', 'up_proj', 'down_proj'],
    task_type='CAUSAL_LM',
)

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

model_id = './gemma-2b'
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type='nf4',
    bnb_4bit_compute_dtype=torch.bfloat16
)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map={'':0})

text = 'Quote: Imagination is more'
device = 'cuda:0'
inputs = tokenizer(text, return_tensors='pt').to(device)

outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))



from datasets import load_dataset

data = load_dataset('./Abirate/english_quotes')
data = data.map(lambda samples: tokenizer(samples['quote']), batched=True)
import transformers
from trl import SFTTrainer

def formatting_func(example):
    text = f'Quote: {example['quote'][0]}\nAuthor: {example['author'][0]}'
    return [text]

trainer = SFTTrainer(
    model=model,
    train_dataset=data['train'],
    args=transformers.TrainingArguments(
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        warmup_steps=2,
        max_steps=10,
        learning_rate=2e-4,
        fp16=True,
        logging_steps=1,
        output_dir='outputs',
        optim='paged_adamw_8bit'
    ),
    peft_config=lora_config,
    formatting_func=formatting_func,
)
trainer.train()

text = 'Quote: Imagination is'
device = 'cuda:0'
inputs = tokenizer(text, return_tensors='pt').to(device)

outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

使用HuggingFace的Accelerate庫加載和運行超大模型：
https://zhuanlan.zhihu.com/p/605640431