|
當(dāng)前已經(jīng)全面進(jìn)入從模型框架、SFT訓(xùn)練數(shù)據(jù)以及模型開(kāi)源井噴的時(shí)期，一日恍若一月。 此外，更小，更快的訓(xùn)練和部署模型也逐步成為目前開(kāi)源的一個(gè)趨勢(shì)。 因此，本文主要介紹極低資源微調(diào)大模型方法LoRA以及BLOOM-LORA實(shí)現(xiàn)代碼，供大家一起參考。 一、LoRA的原理 LoRA是一種以極低資源微調(diào)大模型的方法，其來(lái)自于論文LoRA: Low-Rank Adaptation of Large Language Models。 1. 大模型微調(diào)的困境 隨著模型規(guī)模的不斷擴(kuò)大，模型會(huì)'涌現(xiàn)'出各種能力。特別是對(duì)大語(yǔ)言模型(LLM)來(lái)說(shuō)，隨著規(guī)模的擴(kuò)大其在zero-shot、常識(shí)推理等能力上會(huì)有大幅度的提高。相比于規(guī)模較小的模型，大模型的微調(diào)成本和部署成本都非常高。例如，GPT-3 175B模型微調(diào)需要1.2TB的顯存。此外，若針對(duì)不同下游任務(wù)微調(diào)多個(gè)模型，那么就需要為每個(gè)下游任務(wù)保存一份模型權(quán)重，成本非常高。在某些場(chǎng)景下，甚至可能需要針對(duì)不同的用戶微調(diào)不同的模型，那么模型微調(diào)和部署的成本將不可接受。 因此，如何降低大模型微調(diào)和部署成本，將是大模型商用的重要一環(huán)。 2. LoRA之前的方法 在LoRA方法提出之前，也有很多方法嘗試解決大模型微調(diào)困境的方法。其中有兩個(gè)主要的方向：(1) 添加adapter層；(2) 由于某種形式的輸入層激活。但是這兩種方法都有局限性： 2.1 Adapter層會(huì)引入推理時(shí)延 簡(jiǎn)單來(lái)說(shuō)，adapter就是固定原有的參數(shù)，并添加一些額外參數(shù)用于微調(diào)。上圖中會(huì)在原始的transformer block中添加2個(gè)adapter，一個(gè)在多頭注意力后面，另一個(gè)這是FFN后面。  顯然，adapter會(huì)在模型中添加額外的層，這些層會(huì)導(dǎo)致大模型在推理時(shí)需要更多的GPU通信，而且也會(huì)約束模型并行。這些問(wèn)題都將導(dǎo)致模型推理變慢。 2.2 prefix-tuning難以優(yōu)化 prefix-tuning方法是受語(yǔ)言模型in-context learning能力的啟發(fā)，只要有合適的上下文則語(yǔ)言模型可以很好的解決自然語(yǔ)言任務(wù)。但是，針對(duì)特定的任務(wù)找到離散token的前綴需要花費(fèi)很長(zhǎng)時(shí)間，prefix-tuning提出使用連續(xù)的virtual token embedding來(lái)替換離散token。  具體來(lái)說(shuō)，對(duì)于transformer中的每一層，都在句子表征前面插入可訓(xùn)練的virtual token embedding。對(duì)于自回歸模型(GPT系列)，在句子前添加連續(xù)前綴，即z = [PREFIX;x;y]。對(duì)于Encoder-Decoder模型(T5)，則在Ecoder和Decoder前都添加連續(xù)前綴 z = [PREFIX;x|PREFIX';y]。添加前綴的過(guò)程如上圖所示。 雖然，prefix-tuning并沒(méi)有添加太多的額外參數(shù)。但是，prefix-tuning難以優(yōu)化，且會(huì)減少下游任務(wù)的序列長(zhǎng)度。 3. 問(wèn)題的正式表述 術(shù)語(yǔ)與約定。由于LoRA原理的介紹，會(huì)使用Transformer架構(gòu)。因此，這里先給出一些術(shù)語(yǔ)約定。一個(gè)Transformer層的輸入和輸出維度尺寸為d_model，使用Wq、Wk、Wv和Wo表示自注意力模塊中的query/key/value/output投影矩陣。W或w0表示預(yù)訓(xùn)練模型的權(quán)重矩陣， delata_W表示模型在適配過(guò)程中的梯度更新。r來(lái)表示LoRA模塊的秩。使用Adam作為模型優(yōu)化器，Transformer MLP前饋層的維度為 d_ffn=4xd_model。 問(wèn)題表述。LoRA雖然與訓(xùn)練目標(biāo)無(wú)關(guān)，這里還是以語(yǔ)言建模為例。假設(shè)給定一個(gè)預(yù)訓(xùn)練的自回歸語(yǔ)言模型  目標(biāo)是使該語(yǔ)言模型適應(yīng)下游的摘要、機(jī)器閱讀理解等任務(wù)。每個(gè)下游任務(wù)都有context-target樣本對(duì)組成的訓(xùn)練集：  其中xi和yi都是token序列。例如，對(duì)于摘要任務(wù)，xi是文章內(nèi)容， yi是摘要。 在完整微調(diào)的過(guò)程中，模型使用預(yù)訓(xùn)練好的權(quán)重來(lái)初始化模型，然后通過(guò)最大化條件語(yǔ)言模型來(lái)更新參數(shù):  完整微調(diào)的主要缺點(diǎn)：對(duì)于每個(gè)下游任務(wù)，都需要學(xué)習(xí)不同的參數(shù)更新，因此，如果預(yù)訓(xùn)練模型很大，存儲(chǔ)和部署許多獨(dú)立的微調(diào)模型實(shí)例非常有挑戰(zhàn)。 LoRA為了更加的參數(shù)高效，使用相對(duì)非常小的參數(shù)來(lái)表示任務(wù)相關(guān)的參數(shù)增量。 LoRA將會(huì)使用低秩表示來(lái)編碼，同時(shí)實(shí)現(xiàn)計(jì)算高效和存儲(chǔ)高效。 4. LoRA 通常，神經(jīng)網(wǎng)絡(luò)中會(huì)包含許多進(jìn)行矩陣乘法的稠密層，這些層通常是滿秩的。Adgajanyan et al.等人的研究表示預(yù)訓(xùn)練語(yǔ)言模型具有低的'內(nèi)在維度'。  受該工作的啟發(fā)，在模型適配下游任務(wù)的過(guò)程中，權(quán)重更新也應(yīng)該具有低的“內(nèi)在秩”。對(duì)于預(yù)訓(xùn)練權(quán)重矩陣：  可以通過(guò)低秩分解來(lái)表示其更新：  并且秩  在訓(xùn)練過(guò)程中，W0被凍結(jié)且不接受梯度更新，A和B則是可訓(xùn)練參數(shù)。注意，W0和delata-W=BA都會(huì)乘以相同的輸入。對(duì)于  ，前向傳播變?yōu)椋?/p>  對(duì)矩陣A 使用隨機(jī)高斯初始化，對(duì)矩陣B使用0進(jìn)行初始化，因此delata-W=BA 在訓(xùn)練的開(kāi)始為0。使用alpha/r來(lái)縮放delata-Wx。當(dāng)使用Adam優(yōu)化時(shí)，經(jīng)過(guò)適當(dāng)?shù)目s放初始化，調(diào)優(yōu)alpha與調(diào)優(yōu)學(xué)習(xí)率大致相同。 當(dāng)進(jìn)行部署時(shí)，以顯式的計(jì)算和存儲(chǔ)W=W0+BA ，并正常執(zhí)行推理。W0和BA都是 。當(dāng)需要轉(zhuǎn)換至另一個(gè)下游任務(wù)，可以通過(guò)減去BA來(lái)恢復(fù)W0，然后添加不同的B'A'。至關(guān)重要的是，這保證不會(huì)引人任何額外的推理時(shí)延。 5. LoRA應(yīng)用于Transformer 理論上，LoRA可以應(yīng)用于任何神經(jīng)網(wǎng)絡(luò)的權(quán)重矩陣，從而減少可訓(xùn)練參數(shù)的數(shù)量。Transformer架構(gòu)中的自注意力模塊有4個(gè)權(quán)重矩陣： ，以及兩個(gè)MLP模型的權(quán)重矩陣。將Wq(或者Wk,Wv )作為一個(gè)維度為 的單個(gè)矩陣。為了簡(jiǎn)單和參數(shù)高效，本研究?jī)H限于適配下游任務(wù)的注意力權(quán)重，并凍結(jié)MLP模塊。 優(yōu)點(diǎn)。最顯著的優(yōu)點(diǎn)是顯存和存儲(chǔ)空間的減少。對(duì)于使用Adam訓(xùn)練的大型Transformer，若 ，由于不需要存儲(chǔ)被凍結(jié)參數(shù)的優(yōu)化器狀態(tài)，VRAM使用量減少2/3。 對(duì)于GPT-3 175B，訓(xùn)練中的顯存消耗從1.2TB減少自350GB。當(dāng)r=4并且僅調(diào)整query矩陣和value矩陣時(shí)，checkpoint大小減少10000倍(從350GB減少自35MB)。 另一個(gè)優(yōu)點(diǎn)是，可以在部署時(shí)以更低的成本切換任務(wù)，僅需要交換LoRA權(quán)重即可。此外，與完全微調(diào)相比，GPT-3 175B訓(xùn)練速度提高了25%，因?yàn)椴恍枰?jì)算絕大多數(shù)參數(shù)的梯度。 二、代碼：實(shí)現(xiàn)BLOOM-LoRA 本小節(jié)展示如何使用LoRA微調(diào)大語(yǔ)言模型bloom。 數(shù)據(jù)：使用BELLE提供的100萬(wàn)指令微調(diào)數(shù)據(jù)； 模型：使用bloomz-7b1-mt，該版本的bloomz也是經(jīng)過(guò)指令微調(diào)后的模型。 依賴包：使用transformers提供模型加載和訓(xùn)練；使用peft提供LoRA實(shí)現(xiàn)；使用DeepSpeed提供訓(xùn)練加速。 注意：peft包目前還處于快速迭代當(dāng)中，后續(xù)接口可能會(huì)有大的變動(dòng)，也可能存在一些bug。關(guān)鍵依賴包版本： transformers==4.26.1 torch==1.13.1 deepspeed==0.8.2 peft==0.2. 1. 訓(xùn)練代碼 為了簡(jiǎn)潔，假設(shè)訓(xùn)練代碼位于train.py。 1.1 導(dǎo)入依賴包import os import torch import random import datasets import numpy as np from tqdm import tqdm from typing import Dict from torch.utils.data import DataLoader from transformers import ( AutoModelForCausalLM, AutoTokenizer, DataCollatorForSeq2Seq, TrainingArguments, Trainer ) from peft import ( LoraConfig, TaskType, get_peft_model, get_peft_model_state_dict, set_peft_model_state_dict ) def set_random_seed(seed): if seed is not None and seed > 0: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.random.manual_seed(seed) torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True set_random_seed(1234) 1.2 設(shè)置參數(shù) # LoRA參數(shù) LORA_R = 8 LORA_ALPHA = 32 LORA_DROPOUT = 0.1 # 訓(xùn)練參數(shù) EPOCHS=3 LEARNING_RATE=5e-5 OUTPUT_DIR='./checkpoints' BATCH_SIZE=4 # 2 GRADIENT_ACCUMULATION_STEPS=3 # 其他參數(shù) MODEL_PATH = 'bigscience/bloomz-7b1-mt' DATA_PATH = './data/belle_open_source_1M.train.json' MAX_LENGTH = 512 PATTERN = '{}\n{}' DS_CONFIG = 'ds_zero2_config.json' tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH) # 加載tokenizer 1.3 加載數(shù)據(jù)dataset = datasets.load_dataset('json', data_files=DATA_PATH) # print(dataset['train'][0]) 1.4 tokenize def tokenize(text: str, add_eos_token=True): result = tokenizer( text, truncation=True, max_length=MAX_LENGTH, padding=False, return_tensors=None) # 判斷是否要添加eos_token if (result['input_ids'][-1] != tokenizer.eos_token_id and len(result['input_ids']) < MAX_LENGTH and add_eos_token): result['input_ids'].append(tokenizer.eos_token_id) result['attention_mask'].append(1) result['labels'] = result['input_ids'].copy() return result def preprocess(example: Dict, train_on_inputs: bool = False): prompt = example['input'] response = example['target'] text = PATTERN.format(prompt, response) tokenized_inp = tokenize(text) # 若train_on_inputs為False，則將label中與input相關(guān)的token替換為-100 if not train_on_inputs: tokenized_prompt = tokenize(prompt,add_eos_token=False) prompt_tokens_len = len(tokenized_prompt['input_ids']) tokenized_inp['labels'] = [-100]*prompt_tokens_len + tokenized_inp['labels'][prompt_tokens_len:] return tokenized_inp train_data = dataset['train'].shuffle().map(preprocess, remove_columns=['id', 'input', 'target']) print(train_data[0]) 1.5 collate_fn # pad_to_multiple_of=8表示padding的長(zhǎng)度是8的倍數(shù) collate_fn = DataCollatorForSeq2Seq(tokenizer, pad_to_multiple_of=8, return_tensors='pt', padding=True) 1.6 加載模型 device_map = {'': int(os.environ.get('LOCAL_RANK') or 0)} # device_map指定模型加載的GPU;troch_dtype=torch.float16表示半精度加載模型 model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, torch_dtype=torch.float16, device_map=device_map) 1.7 LoRA相關(guān)# 轉(zhuǎn)換模型 model = get_peft_model(model, lora_config) model.config.use_cache = False old_state_dict = model.state_dict model.state_dict = ( lambda self, *_, **__: get_peft_model_state_dict(self, old_state_dict()) ).__get__(model, type(model)) # 打印模型中的可訓(xùn)練參數(shù) model.print_trainable_parameters() 1.8 訓(xùn)練參數(shù) args = TrainingArguments( output_dir=OUTPUT_DIR, # checkpoint的存儲(chǔ)目錄 per_device_train_batch_size=BATCH_SIZE, # 單設(shè)備上的batch size gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS, # 梯度累加的step數(shù) warmup_steps=100, num_train_epochs=EPOCHS, learning_rate=LEARNING_RATE, fp16=True, # 使用混合精度訓(xùn)練 logging_steps=50, evaluation_strategy='no', # 不進(jìn)行評(píng)估 save_strategy='steps', save_steps=2000, # 保存checkpoint的step數(shù) save_total_limit=5, # 最多保存5個(gè)checkpoint deepspeed=DS_CONFIG ) 1.9 模型訓(xùn)練trainer = Trainer( model=model, train_dataset=train_data, eval_dataset=None, args=args, data_collator=collate_fn ) trainer.train() model.save_pretrained('best_model') 2. DeepSpeed配置文件 DeepSpeed配置文件名為ds_zero2_config.json。 { 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 'auto', 'steps_per_print': 50, 'gradient_clipping': 1.0, 'zero_optimization': { 'stage': 2, 'offload_optimizer': { 'device': 'cpu' }, 'contiguous_gradients': true, 'overlap_comm': true }, 'zero_allow_untested_optimizer': true, 'fp16': { 'enabled': true, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1 }, 'optimizer': { 'type': 'Adam', 'params': { 'lr': 'auto', 'betas': 'auto', 'eps': 'auto', 'weight_decay': 'auto' } }, 'activation_checkpointing': { 'partition_activations': true, 'contiguous_memory_optimization': true }, 'wall_clock_breakdown': false } 3. 啟動(dòng)deepspeed --include=localhost:0,1,2,3 train.py 4. 推理 推理文件名為inference.py import torch from peft import PeftModel from transformers import AutoModelForCausalLM, AutoTokenizer BASE_MODEL = 'bigscience/bloomz-7b1-mt' LORA_WEIGHTS = 'best_model' tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL) model = AutoModelForCausalLM.from_pretrained( BASE_MODEL, torch_dtype=torch.float16, # 加載半精度 device_map={'':0}, # 指定GPU 0 ) model.eval() # 加載LoRA權(quán)重 model = PeftModel.from_pretrained(model, LORA_WEIGHTS, torch_dtype=torch.float16) model.half() prompt = '' inp = tokenizer(prompt, max_length=512, return_tensors='pt').to('cuda') outputs = model.generate(input_ids=inp['input_ids'], max_new_tokens=256) print(tokenizer.decode(outputs[0])) 參考文獻(xiàn) 1、https://zhuanlan.zhihu.com/p/618073170 2、arxiv.org/pdf/2106.09685.pdf 3、https://zhuanlan.zhihu.com/p/615235322 4、https://github.com/tloen/alpaca-lora/blob/main/finetune.py 5、https://github.com/huggingface/peft/blob/main/examples/conditional_generation/peft_lora_seq2seq_accelerate_ds_zero3_offload.py |
|
|
