解密Prompt系列6. lora指令微調扣細節-請冷靜,1個小時真不夠~

風雨中的小七發表於2023-04-29
解密

上一章介紹瞭如何基於APE+SELF自動化構建指令微調樣本。這一章我們就把微調跑起來，主要介紹以Lora為首的低引數微調原理，環境配置，微調程式碼，以及大模型訓練中視訊記憶體和耗時最佳化的相關技術細節

標題這樣寫是因為上週突然收到了一週內上線一版chatbo的命令，原因無它領導們都刷到了《一個小時你也可以擁有ChatGPT》,《100美金訓練ChatGPT》,《僅訓練3小時超越ChatGPT》,《人人都可以擁有ChatGPT》。。。領導說人人都有了為啥我沒有呀？！！真誠呼籲標題黨們求手下留情，留人一命！於是這裡我換個標題來Debuff！Debuff！

看到這裡本文最重要的部分已經說完了，累了的小夥伴可以撤退了，五一快樂~

解密Prompt系列6. lora指令微調扣細節-請冷靜,1個小時真不夠~

低引數微調原理

  • LORA：LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS
  • 原理：INTRINSIC DIMENSIONALITY EXPLAINS THE EFFECTIVENESS
    OF LANGUAGE MODEL FINE-TUNING
  • 前人的肩膀：Adapter: Parameter-Efficient Transfer Learning for NLP

我們之前在解密Prompt系列3. 凍結LM微調Prompt介紹過一些soft-prompt，包括P-Tunning和Prompt-Tunning也屬於低引數微調。這些方案是透過引數拼接的方案引入額外引數。這裡介紹另一類方案，同樣是凍結LLM的引數，透過引數相加的方案引入額外引數, 相較soft-prompt最明顯的優勢，就是不會佔用輸入token的長度。

LoRA的原理比較簡單，原始全量微調其實就是在原始模型引數上透過微調加入增量\(W = W_0+\Delta W\)，那我們可以透過凍結原始引數\(W_0\),並且把增量部分透過低秩分解方式進一步降低引數量級\(\Delta W=A*B^T\)，原始引數的維度是\(d*d\), 則低秩分解後的引數量級是\(2*r*d\)，因為這裡的r<<d，因此可以起到大幅降低微調引數量級的效果，如下圖

解密Prompt系列6. lora指令微調扣細節-請冷靜,1個小時真不夠~

核心程式碼如下

## 初始化低秩矩陣A和B
self.lora_A.update(nn.ModuleDict({adapter_name: nn.Linear(self.in_features, r, bias=False)}))
self.lora_B.update(nn.ModuleDict({adapter_name: nn.Linear(r, self.out_features, bias=False)}))
self.scaling[adapter_name] = lora_alpha / r

## 向前計算
result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
result += (
    self.lora_B[self.active_adapter](
        self.lora_A[self.active_adapter](self.lora_dropout[self.active_adapter](x))
    )
    * self.scaling[self.active_adapter]
)

論文測試了在多數場景下適當的LORA微調和全量微調的效果不相上下。一個可能原因是INTRINSIC DIMENSIONALITY論文中提出，雖然語言模型整體引數空間很大，但具體到每個任務其實有各自的隱表徵空間(intrisic dimension)，這個隱表徵空間的維度並不高, 因此在微調過程中加入低秩分解並不一定會影響微調效果。使用LORA微調有以下幾個細節

  1. 對哪些引數進行微調：基於Transformer結構，LORA只對每層的Self-Attention的部分進行微調，有\(W_q, W_k, W_v, W_O\)四個對映層引數可以進行微調。消融實驗顯示只微調\(W_q\)效果略差，微調\(W_q, W_v\)的效果和微調\(W_q, W_k, W_v, W_O\)的效果相似。需要注意不同模型引數名稱不同，像chatglm對應的引數名稱就是query_key_value
  2. Rank的選取：Rank的取值作者對比了1-64，效果上Rank在4-8之間最好，再高並沒有效果提升。不過論文的實驗是面向下游單一監督任務的，因此在指令微調上根據指令分佈的廣度，Rank選擇還是需要在8以上的取值進行測試。
  3. alpha引數：alpha其實是個縮放引數，本質和learning rate相同，所以為了簡化我預設讓alpha=rank，只調整lr，這樣可以簡化超參
  4. 初始化：A和Linear層的權重相同Uniform初始化，B是zero初始化，這樣最初的Lora權重為0。所以Lora引數是從頭學起，並沒有那麼容易收斂。

Lora的優點很明顯，低引數，適合小樣本場景；可以拔插式的使用，快速針對不同下游任務訓練不同的lora權重；完全沒有推理延時，這個在後面程式碼中會提到推理時，可以預先把lora權重merge到原始權重上。

但Lora微調雖好，個人在嘗試中感受到的侷限性就是adapter類的微調方案可能更適合下游單一任務型別/生成風格。至於是否適合作為通用指令微調的解決方案，有個問題我也沒有搞懂，就是通用的指令樣本是否真的有統一的低秩空間表徵？這個表徵又是什麼含義？因為指令微調階段的樣本其實是混合的多工指令樣本，這種情況下lora是否合適，感覺需要更全面的評估（當前出來的眾多LLama們都缺少合理統一全面可比的Evaluation），當前就我們的嘗試情況lora的效果並不及預期。

環境配置

我用了featurize和攬睿星舟。雲服務廠商的選擇主要看是否有jupyter，儲存夠大，下載快，能連git，有高配torch環境。這兩家在眾多小廠裡脫穎而出，4090的卡一個小時也就3塊錢，來來來盆友辛苦把推廣費結一下~

強調下環境配置，想跑通微調，搞定環境你就成功了80%！運氣好1分鐘，運氣差1天都在原地打轉

  1. 例項環境：TRX4090 + py38 + torch2.0 + CUDA12
  2. python環境：主要坑在transforemrs和peft，幾個相關issue包括：llama tokenizer special token有問題，peft adapter.bin微調不更新，Bug with fan_in_fan_out。我一個不差都踩中了。。。
# 以下配置可能會隨時間變化，出了問題就去issue裡面刨吧
# 要相信你不是唯一一個大冤種！
accelerate
appdirs
loralib
bitsandbytes
black
black[jupyter]
datasets
fire
transformers>=4.28.0
git+https://github.com/huggingface/peft.git
sentencepiece
gradio
wandb
cpm-kernel

模型初始化

以下程式碼主要整合自alpaca-lora和chatglm-finetune。其實lora微調的程式碼本身並不複雜，相反是如何加速大模型訓練，降低視訊記憶體佔用的一些技巧大家可能不太熟悉。模型初始化程式碼如下，get_peft_model會初始化PeftModel把原模型作為base模型，並在各個self-attention層加入lora層，同時改寫模型forward的計算方式。

主要說下load_in_8bit和prepare_model_for_int8_training，這裡涉及到2個時間換空間的大模型視訊記憶體壓縮技巧。

from peft import get_peft_model, LoraConfig, prepare_model_for_int8_training, set_peft_model_state_dict
from transformers import AutoTokenizer, AutoModel

model = AutoModel.from_pretrained("THUDM/chatglm-6b", load_in_8bit=True, torch_dtype=torch.float16, trust_remote_code=True, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
model = prepare_model_for_int8_training(model)

lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    inference_mode=False,
    r=8,
    lora_alpha=8,
    lora_dropout=0.05,
)
model = get_peft_model(model, lora_config)
model.config.use_cache = False

模型視訊記憶體佔用分成兩個部分，一部分是靜態視訊記憶體基本由模型引數量級決定，另一部分是動態視訊記憶體在向前傳播的過程中每個樣本的每個神經元都會計算啟用值並儲存，用於向後傳播時的梯度計算，這部分和batchsize以及引數量級相關。以下8bit量化最佳化的是靜態視訊記憶體，而梯度檢查最佳化的是動態視訊記憶體。

1. 8bit Quantization

https://huggingface.co/blog/hf-bitsandbytes-integration

from_pretrained中的load_in_8bit引數是bitsandbytes庫賦予的能力，會把載入模型轉化成混合8bit的量化模型，注意這裡的8bit模型量化只用於模型推理，透過量化optimizer state降低訓練時視訊記憶體的時8bit最佳化器是另一個功能不要搞混喲~

模型量化本質是對浮點引數進行壓縮的同時，降低壓縮帶來的誤差。 8-bit quantization是把原始FP32（4位元組）壓縮到Int8（1位元組）也就是1/4的視訊記憶體佔用。如上載入後會發現除lora層外的多數層被轉化成int型別如下

解密Prompt系列6. lora指令微調扣細節-請冷靜,1個小時真不夠~

當然壓縮方式肯定不是直接四捨五入，那樣會帶來巨大的精度壓縮損失。常見的量化方案有absolute-maximum和zero-point，它們的差異只是rescale的方式不同，這裡簡單說下absmax，如下

解密Prompt系列6. lora指令微調扣細節-請冷靜,1個小時真不夠~

先尋找tensor矩陣的絕對值的最大值，並計算最大值到127的縮放因子，然後使用該縮放因子對整個tensor進行縮放後，再round到整數。這樣就把浮點數對映到了INT8,逆向回到float的原理相同。

當然以上的縮放方案依舊存在精度損失，以及當矩陣中存在outlier時，這個精度損失會被放大，例如當tensor中絕大部分取值在1以下，有幾個值在100+，則縮放後，所有1以下的tensor資訊都會被round抹去。因此LLM.int8()的實現對outlier做了進一步的最佳化，把outlier和非outlier的矩陣分開計算，再把結果進行合併來降低outlier對精度的影響。

解密Prompt系列6. lora指令微調扣細節-請冷靜,1個小時真不夠~

prepare_model_for_int8_training是對在Lora微調中使用LLM.int8()進行了適配用來提高訓練的穩定性，主要包括

  • layer norm層保留FP32精度
  • 輸出層保留FP32精度保證解碼時隨機sample的差異性

2. gradient checkpoint

https://medium.com/tensorflow/fitting-larger-networks-into-memory-583e3c758ff9

prepare_model_for_int8_training函式還做了一件事就是設定gradient_checkpointing=True，這是另一個時間換空間的技巧。

gradient checkpoint的實現是在向前傳播的過程中使用torch.no_grad()不去儲存中間啟用值，降低動態視訊記憶體的佔用。而只是儲存輸入和啟用函式，當進行反向傳播的時候，會重新獲取輸入和啟用函式計算啟用值用於梯度計算。因此向前傳播會計算兩遍，所以需要更多的訓練時間。

use_cache設定為False，是因為和gradient checkpoint存在衝突。因為use_cache是對解碼速度的最佳化，在解碼器解碼時，儲存每一步輸出的hidden-state用於下一步的輸入，而因為開啟了gradient checkpoint，中間啟用值不會儲存，因此use_cahe=False。其實#21737已經加入了引數檢查，這裡設定只是為了不輸出warning。

模型訓練

訓練基本和常規訓練基本相同，程式碼如下。主要說下模型儲存和載入以及混合精度訓練

import datasets
from transformers import Trainer, DataCollatorForSeq2Seq

if resume_from_checkpoint:
    lora_weight = torch.load(ckpt_name)
    set_peft_model_state_dict(model, lora_weight)

train_data = datasets.load_from_disk(dataset_path)

class ModifiedTrainer(Trainer):
    def save_model(self, output_dir=None, _internal_call=False):
        # 改寫trainer的save_model，在checkpoint的時候只存lora權重
        from transformers.trainer import TRAINING_ARGS_NAME

        os.makedirs(output_dir, exist_ok=True)
        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
        saved_params = {
            k: v.to("cpu") for k, v in self.model.named_parameters() if v.requires_grad
        }
        torch.save(saved_params, os.path.join(output_dir, "adapter_model.bin"))
        
trainer = ModifiedTrainer(
    model=model,
    train_dataset=train_data,
        args=transformers.TrainingArguments(
            per_device_train_batch_size=8,
            gradient_accumulation_steps=16,
            num_train_epochs=10,
            learning_rate=3e-4,
            fp16=True,
            logging_steps=10,
            save_steps=200,
            output_dir=output_dir
        ),
    data_collator=DataCollatorForSeq2Seq(
        tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True
    ),
)
trainer.train()
model.save_pretrained(train_args.output_dir)

1. 模型的儲存和載入

因為peftModel重寫了原始model的save_pretrained函式，只把lora層的權重進行儲存，因此model.save_pretrained只會儲存lora權重。而trainer的save_model函式沒有做相應的重寫，因此我們重寫下對應的function，避免checkpoint寫入原始模型全部引數。

相應的如果你從ckpt載入lora權重去繼續訓練的話，也是對PeftModel中的Lora權重進行載入。

2. 混合精度訓練

https://huggingface.co/docs/transformers/main/en/perf_train_gpu_one#fp16-training

除了預設的全精度FP32，引數精度還有半精度FP16，以及BF16和TF32。最常用也是這裡使用的是FP16的混合精度。

解密Prompt系列6. lora指令微調扣細節-請冷靜,1個小時真不夠~

實現原理是並非所有變數都需要全精度儲存，如果把部分中間變數轉化成半精度，則計算效率會大幅提升，加上一些GPU對FP16計算做了最佳化，吞吐上比全精度會快2~5倍。

不過只使用半精度訓練同樣會帶來量化誤差，主要包括：資料溢位因為半精度比全精度的範圍更小，訓練到後期因為梯度越來越小可能會下溢位；舍入誤差梯度變小後，因為精度有限，導致梯度更新被四捨五入，更新了個寂寞。

為了解決以上的問題引入了混合精度訓練。簡單來說就是向前傳遞時，模型權重、啟用值和梯度都使用FP16進行儲存，同時會複製一份模型權重以FP32儲存，向後傳播optimizer更新時會更新FP32的引數。因此混合精度訓練並不會節省記憶體，只會提高模型訓練速度。

模型推理

推理有兩個方案，一個和訓練相同，直接加入Lora層，不過會增加推理延時因為多了lora層的計算，適合線下測評用，如下

from peft import PeftModel
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True, load_in_8bit=True, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
model = PeftModel.from_pretrained(model, "./lora_ckpt")
model.half().to(device)
model.eval()

另一個沒有推理延時的方案，是先把lora權重和原始模型權重進行合併，把合併後的引數儲存成新的bin檔案，然後和載入常規模型一樣載入合併後的模型引數進行推理。權重合並的程式碼如下

tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
# when merging disable int8
model = AutoModel.from_pretrained(
    "THUDM/chatglm-6b", load_in_8bit=False, torch_dtype=torch.float16,
    trust_remote_code=True, device_map={"": "cpu"},
)
## 用來檢查權重是否合併成功，合併成功weight會改變
first_weight = model.base_model.layers[0].attention.query_key_value.weight
first_weight_old = first_weight.clone()

# 返回的不是新的模型，而是在原始模型上加了adapter層
lora_model = PeftModel.from_pretrained(
    model,
    "./lora_ckpt",
    device_map={"": "cpu"},
    torch_dtype=torch.float16,
)
# 報錯：A*B shape mismatch，大機率是get_peft_model錯誤修改了peft_config裡面的fan_in_fan_out引數，某個peft的revision有這個bug
lora_model = lora_model.merge_and_unload()
lora_model.train(False)

# 報錯：大機率peft訓練有問題，檢查adapter.bin大小
assert not torch.allclose(first_weight_old, first_weight), 'Weight Should Change after Lora Merge'

# lora模型權重把原模型權重加了prefix，這裡移除恢復原始key
deloreanized_sd = {
    k.replace("base_model.model.", ""): v
    for k, v in lora_model.state_dict().items()
    if "lora" not in k
}
# 儲存合併後的模型權重
lora_model.save_pretrained(output_dir, state_dict=deloreanized_sd)

更多Prompt相關論文·教程，開源資料·模型，以及AIGC相關玩法戳這裡DecryptPrompt

Reference

  1. https://blog.csdn.net/anycall201/article/details/129959567
  2. 蘇劍林. (Jun. 20, 2022). 《Ladder Side-Tuning：預訓練模型的“過牆梯” 》[Blog post]. Retrieved from https://kexue.fm/archives/9138
  3. 蘇劍林. (Apr. 17, 2023). 《梯度視角下的LoRA：簡介、分析、猜測及推廣 》[Blog post]. Retrieved from https://kexue.fm/archives/9590
    4.https://github.com/huggingface/blog/blob/main/notebooks/HuggingFace_int8_demo.ipynb
  4. ChatGLM-Finetune
  5. Alpaca-lora

相關文章