大模型LoRA微調實踐
準備工作
數據集:采用 GitHub 上的 Chinese-medical-dialogue-data 中文醫療對話數據集
Github地址如下:
https://github.com/Toyhom/Chinese-medical-dialogue-data
微調模型:
Qwen 1.5B模型(Qwen2、2.5均可以,可以自由選擇)
模型權重文件可以先從huggingface官網下載,或者從魔塔社區下載速度更快:
https://modelscope.cn/models/Qwen/Qwen2.5-1.5B-Instruct
本實驗環境:
GPU 顯存 >= 8GB
pytorch==2.5.0+cu118
transformers==4.47.1
peft==0.14.0
參考資料:
https://blog.csdn.net/YoungOne2333/article/details/144718615
數據預處理
數據集是Excel文件,主要是ask+question的問答對,需要處理成大模型微調的數據格式,這里可以參考
LLaMA Factory的數據處理文檔:https://llamafactory.readthedocs.io/zh-cn/latest/getting_started/data_preparation.html
本文采用指令監督微調數據集,instruction 列對應的內容為人類指令, input 列對應的內容為人類輸入, output 列對應的內容為模型回答。下面是一個例子:
{"instruction": "計算這些物品的總費用。 ","input": "輸入:汽車 - $3000,衣服 - $100,書 - $20。","output": "汽車、衣服和書的總費用為 $3000 + $100 + $20 = $3120。"
}
通過以下代碼讀取文件構建數據加載類:
import json
import torch
import numpy as np
from torch.utils.data import Datasetclass QADataset(Dataset):def __init__(self, data_path, tokenizer, max_source_length, max_target_length) -> None:super().__init__()self.tokenizer = tokenizerself.max_source_length = max_source_lengthself.max_target_length = max_target_lengthself.max_seq_length = self.max_source_length + self.max_target_lengthself.data = []if data_path:with open(data_path, "r", encoding='utf-8') as f:for line in f:if not line or line == "":continuejson_line = json.loads(line)question = json_line["question"]answer = json_line["answer"]self.data.append({"question": question,"answer": answer})print("data load , size:", len(self.data))def preprocess(self, question, answer):messages = [{"role": "system", "content": "你是一個醫療方面的專家,可以根據患者的問題進行解答。"},{"role": "user", "content": question}]# 經歷過一段時間對于輸入和輸出的思考和探索,發現這個代碼里的輸入和輸出格式是暫且發現的最優的方式prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)instruction = self.tokenizer(prompt, add_special_tokens=False, max_length=self.max_source_length, truncation=True)# 因為是訓練,所以有輸出response = self.tokenizer(answer, add_special_tokens=False, max_length=self.max_target_length, truncation=True)# 輸入是 question+answerinput_ids = instruction["input_ids"] + response["input_ids"] + [self.tokenizer.pad_token_id]attention_mask = (instruction["attention_mask"] + response["attention_mask"] + [1])# 輸出是 answer,而不去計算question部分的loss,-100 是一個約定俗成的用于忽略損失計算的值。labels = [-100] * len(instruction["input_ids"]) + response["input_ids"] + [self.tokenizer.pad_token_id]if len(input_ids) > self.max_seq_length:input_ids = input_ids[:self.max_seq_length]attention_mask = attention_mask[:self.max_seq_length]labels = labels[:self.max_seq_length]# 注意!!!這里這三個list的長度是完全一致的,否則無法訓練return input_ids, attention_mask, labelsdef __getitem__(self, index):item_data = self.data[index]input_ids, attention_mask, labels = self.preprocess(**item_data)return {"input_ids": torch.LongTensor(np.array(input_ids)),"attention_mask": torch.LongTensor(np.array(attention_mask)),"labels": torch.LongTensor(np.array(labels))}def __len__(self):return len(self.data)
原文章中先通過一個預處理代碼讀取Excel中的部分數據保存為json文件,所以這里直接從json文件讀取數據。
這里要注意的就是輸入和輸出的構建,以及哪部分進行損失計算。
模型加載測試
模型加載使用transformer庫的因果語言模型類,因果語言模型是一種自回歸模型,其目標是根據前面的 token 預測下一個 token(即從左到右的單向預測),即現在所流行的大語言模型。
類名 | 適用任務 | 示例模型 |
---|---|---|
AutoModelForCausalLM | 因果語言模型(文本生成) | GPT-2、Llama |
AutoModelForSeq2SeqLM | 序列到序列模型(翻譯、摘要) | T5、BART |
AutoModelForMaskedLM | 掩碼語言模型(填空、特征提取) | BERT、RoBERTa |
AutoModelForQuestionAnswering | 問答任務 | BERT-QA、RoBERTa-QA |
使用peft庫進行LoRA微調,這里也先簡單展示使用peft庫進行LoRA微調配置后,實際參與訓練的參數量。
先把模型權重文件下載下來,然后使用以下代碼可以加載模型進行對話測試:
import time
import torchfrom transformers import AutoModelForCausalLM, AutoModel, AutoTokenizer
from peft import LoraConfig, get_peft_model, TaskTypedef demo():# 加載模型model = AutoModelForCausalLM.from_pretrained("../modelscope/Qwen/Qwen2.5-1.5B-Instruct", # 先手動將模型下載到本地torch_dtype='auto', # 使用auto會根據硬件配置情況自行選擇精度,如果不設置此參數,默認使用float32device_map="auto" # 如果有GPU,可以自動加載到GPU)# 可以打印查看模型的網絡結構# 例如qwen2 1.5B 由28 層 Qwen2DecoderLayer 構成,每個 Decoder 主要的核心是 self_attention 和 mlpprint(model)# 增加Lora結構之后,打印模型結構查看變化peft_config = LoraConfig(task_type=TaskType.CAUSAL_LM,target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],inference_mode=False,r=8,lora_alpha=32,lora_dropout=0.1)model = get_peft_model(model, peft_config)# trainable params: 9,232,384 || all params: 1,552,946,688 || trainable%: 0.5945model.print_trainable_parameters()# 下面通過自行計算參與訓練的參數量,與上面的參數量對比是否一致total_trainable_params = 0for param in model.parameters():if param.requires_grad:total_trainable_params += param.numel()print(f"參與訓練的參數數量: {total_trainable_params}")# Lora 之后在每一層(q_proj這些線性層)都增加了一個 lora_A 和 lora_B 結構來實現降維升維的作用,print(model)# 對話測試# todo tokenizer具體是什么?tokenizer = AutoTokenizer.from_pretrained("../modelscope/Qwen/Qwen2.5-1.5B-Instruct")device = torch.device("cuda" if torch.cuda.is_available() else "cpu")prompt = "5月至今上腹靠右隱痛,右背隱痛帶酸,便秘,喜睡,時有腹痛,頭痛,腰酸癥狀?"messages = [{"role": "system", "content": '你是一個醫療方面的專家,可以根據患者的問題進行解答。'},{"role": "user", "content": prompt}]text = tokenizer.apply_chat_template(messages,tokenize=False,add_generation_prompt=True)model_inputs = tokenizer([text], return_tensors="pt").to(device)start = time.time()generated_ids = model.generate(model_inputs.input_ids,max_new_tokens=512)end = time.time()response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]# 初始回復:您的描述表明您可能患有慢性胃炎或者胃潰瘍等疾病。建議盡快就醫并做進一步檢查以明確診斷,并根據醫生的指導進行治療。同時注意飲食健康,避免辛辣、油膩食物,保持良好的生活習慣和心態。print(f"耗時:{end-start}s,{response}")
這里關于tokenizer實際上有必須要再深入研究一下,不同的大模型所采用的的分詞算法可能會有所區別,也會表現在針對相同的一段文本但是實際token數量不一致。
LoRA微調——手動實現版本
接下來開始正式進行Lora微調,這里是一種比較簡單的實現方式,除了transformers和peft庫,沒有使用其他封裝好的庫或者訓練框架,更易于理解,整體流程與常規的深度學習模型訓練代碼并無太大的區別
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model, TaskType
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriterfrom qa_dataset import QADatasetdef main():model_name = "../modelscope/Qwen/Qwen2.5-1.5B-Instruct"train_json_path = "./data/train.json"val_json_path = "./data/val.json"max_source_length = 128 # 輸入長度可根據數據集調整,顯存會隨之變化max_target_length = 256 epochs = 10batch_size = 1 # 可根據顯存使用情況調整,一般單卡很難設置的比較大lr = 1e-4gradient_accumulation_steps = 16lora_rank = 8 # 8或16或32lora_alpha = 32model_output_dir = "output"logs_dir = "logs"# 設備(這里先簡單介紹單卡訓練版本,后面會測試多卡訓練)device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")# 加載分詞器和模型tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)# 如果顯存夠,這里可以使用float32,不設置的話默認float32(1.5B模型8G顯存使用float16、11G顯存使用float32)model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, trust_remote_code=True)# setup peftpeft_config = LoraConfig(task_type=TaskType.CAUSAL_LM, # 任務類型:CAUSAL_LM 表示因果語言模型(Causal Language Model),即生成式任務target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "demo_proj"],inference_mode=False,r=lora_rank,lora_alpha=lora_alpha,lora_dropout=0.1)model = get_peft_model(model, peft_config)model.is_parallelizable = Truemodel.model_parallel = Trueprint("start load train data...")train_params = {"batch_size": batch_size, "shuffle": True, "num_workers": 0}training_set = QADataset(train_json_path, tokenizer, max_source_length, max_target_length)training_loader = DataLoader(training_set, **train_params)print("start load validation data...")val_params = {"batch_size": batch_size, "shuffle": True, "num_workers": 0}val_set = QADataset(val_json_path, tokenizer, max_source_length, max_target_length)val_loader = DataLoader(val_set, **val_params)# 日志記錄writer = SummaryWriter(logs_dir)# 優化器optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)model = model.to(device)# 開始訓練print("Start Training...")train_model(model=model,train_loader=training_loader,val_loader=val_loader,optimizer=optimizer,gradient_accumulation_steps=gradient_accumulation_steps,device=device,num_epochs=epochs,model_output_dir=model_output_dir,writer=writer)
train_model的實現如下:
import time
import torch
import sysfrom tqdm import tqdmdef train_model(model, train_loader, val_loader, optimizer, gradient_accumulation_steps,device, num_epochs, model_output_dir, writer):batch_step = 0for epoch in range(num_epochs):time1 = time.time()model.train()for index, data in enumerate(tqdm(train_loader, file=sys.stdout, desc="Train Epoch: " + str(epoch))):input_ids = data['input_ids'].to(device, dtype=torch.long)attention_mask = data['attention_mask'].to(device, dtype=torch.long)labels = data['labels'].to(device, dtype=torch.long)# 前向傳播outputs = model(input_ids, attention_mask=attention_mask, labels=labels)loss = outputs.loss # 交叉熵損失函數計算得來# 反向傳播, 計算當前梯度loss.backward()# 梯度累積步數if (index % gradient_accumulation_steps == 0 and index != 0) or index == len(train_loader) - 1:# 更新網絡參數optimizer.step()# 清空過往梯度optimizer.zero_grad()writer.add_scalar('Loss/train', loss, batch_step)batch_step += 1# 100條數據打印一次 lossif (index % 100 == 0 and index != 0) or index == len(train_loader) - 1:time2 = time.time()tqdm.write(f"{index}, epoch: {epoch} -loss: {str(loss)} ; "f"each step's time spent: {(str(float(time2 - time1) / float(index + 0.0001)))}")# 驗證model.eval()val_loss = validate_model(model, val_loader, device)writer.add_scalar('Loss/val', val_loss, epoch)print(f'val_loss: {val_loss}, epoch: {epoch}')print('Save Model To', model_output_dir)# 保存的模型只包含微調的參數部分,后面還需要合并模型model.save_pretrained(model_output_dir)def validate_model(model, val_loader, device):running_loss = 0.0with torch.no_grad():for _, data in enumerate(tqdm(val_loader, file=sys.stdout, desc="Validation Data")):input_ids = data['input_ids'].to(device, dtype=torch.long)attention_mask = data['attention_mask'].to(device, dtype=torch.long)labels = data['labels'].to(device, dtype=torch.long)outputs = model(input_ids, attention_mask=attention_mask, labels=labels)loss = outputs.lossrunning_loss += loss.item()return running_loss / len(val_loader)
以上是所有的訓練代碼,可以在單卡(顯存>=8GB)上Lora微調1.5B的模型,前提是上下文長度不易過長;
估算模型占用顯存大小可以使用如下公式:
1.5(參數量:1.5B)21.3=3.9GB
即在模型推理時,僅將模型加載到顯存中就需要占用這么大的顯存,如果是全參數微調,則需要準備再乘以10倍的顯存大小;而Lora微調的實際參數量只占1%左右,一般情況比推理所需的顯存略大一些即可,因為需要保存額外的參數、優化器和梯度等,但是如果上下文長度較長時,顯存要求相應也會更大。
全參數微調的區別只不過是需要的顯存更大,而且不需要使用peft庫,其他代碼與上述代碼并無本質區別。
模型推理和權重合并
微調結束之后,Lora微調的參數會單獨保存為一個權重文件,這一權重文件與原始的大模型權重文件是分開的,需要同時加載這兩個模型文件進行推理,實現方式如下:
import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModeldef test_lora():model_path = "../modelscope/Qwen/Qwen2.5-1.5B-Instruct"lora_dir = "output"device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype='auto', device_map='auto')tokenizer = AutoTokenizer.from_pretrained(model_path)model = PeftModel.from_pretrained(model, lora_dir)model.to(device)prompt = "5月至今上腹靠右隱痛,右背隱痛帶酸,便秘,喜睡,時有腹痛,頭痛,腰酸癥狀?"messages = [{"role": "system", "content": '你是一個醫療方面的專家,可以根據患者的問題進行解答。'},{"role": "user", "content": prompt}]text = tokenizer.apply_chat_template(messages,tokenize=False,add_generation_prompt=True)model_inputs = tokenizer([text], return_tensors="pt").to(device)start = time.time()generated_ids = model.generate(model_inputs.input_ids,max_new_tokens=512)end = time.time()# generated_ids中包含輸入,這一步驟可以去除輸入部分generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]# 實際測試確實有思考過程print(f"耗時:{end - start}s,{response}")
如果不想每次加載兩個模型文件,則可以將兩個模型文件進行合并:
import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModeldef merge_model():device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")model_path = "../modelscope/Qwen/Qwen2.5-1.5B-Instruct"lora_dir = "output"tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)model = PeftModel.from_pretrained(model, lora_dir).to(device)print(model)# 合并model, 同時保存 tokenmodel = model.merge_and_unload()model.save_pretrained("lora_output")tokenizer.save_pretrained("lora_output")
后面就不需要再通過 PeftModel加載模型了,直接與加載原始大模型文件一樣即可。
模型評估
其實模型微調并不難,難點在于另外兩點:
- 微調的數據集,一般需要針對領域或者特定任務的高質量數據集,且量要相對來說大一些,標注成本相對來說會高一些,你能做到別人做不到的,關鍵就在于你獨有的數據集;
- 微調后的模型如何評估,這個其實是很難的,因為現在的大模型是生成式模型,而不是像以前的文本分類、實體識別等任務,以前這種任務對就是對,錯就是錯,評估比較簡單,但是大語言模型是無法對比文本內容來判斷是否正確的,所以對于評估集的構建和評估方案的制定是非常難的。
分布式訓練
這里討論的分布式訓練僅考慮單機多卡的情況,暫不考慮多機多卡的情況。
單機多卡訓練一般分為兩種分布式技術:
- DDP (DistributedDataParallel) 通過實現模型并行和數據并行實現訓練加速。 使用 DDP 的程序需要生成多個進程并且為每個進程創建一個 DDP 實例,他們之間通過 torch.distributed 庫同步。
- FSDP 通過全切片數據并行技術(Fully Sharded Data Parallel)來處理更多更大的模型。在 DDP 中,每張 GPU 都各自保留了一份完整的模型參數和優化器參數。而 FSDP 切分了模型參數、梯度與優化器參數,使得每張 GPU 只保留這些參數的一部分。 除了并行技術之外,FSDP 還支持將模型參數卸載至CPU,從而進一步降低顯存需求。
DDP
DDP是每張卡上都有一份完整的模型參數,所以使用此方法的前提是單張卡顯存可以加載你要訓練的模型全部參數,然后將待訓練的數據劃分到多張卡上,自然訓練速度就會提高。
一般情況下有兩種比較簡單的實現方式:
- 通過torch自帶的DistributedDataParallel實現;
- 結合accelerate實現。
通過DistributedDataParallel的實現方式如下:
import os
import torch
import torch.distributed as dist
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model, TaskType
from torch.utils.data import DataLoader, DistributedSampler
from torch.utils.tensorboard import SummaryWriter
from torch.nn.parallel import DistributedDataParalleldef main():model_name = "../modelscope/Qwen/Qwen2.5-1.5B-Instruct"train_json_path = "./data/train.json"val_json_path = "./data/val.json"max_source_length = 128 max_target_length = 256 epochs = 10batch_size = 1 lr = 1e-4gradient_accumulation_steps = 16lora_rank = 8 # 8或16或32lora_alpha = 32model_output_dir = "output"logs_dir = "logs"# 設備local_rank = int(os.environ.get("LOCAL_RANK", -1))device = torch.device("cuda", local_rank)# 初始化分布式環境if local_rank != -1:dist.init_process_group(backend='nccl')torch.cuda.set_device(local_rank)# 加載分詞器和模型tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)# 如果顯存夠,這里可以使用float32,不設置的話默認float32(8G顯存使用float16、11G顯存使用float32)model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, trust_remote_code=True)# setup peftpeft_config = LoraConfig(task_type=TaskType.CAUSAL_LM, # 任務類型:CAUSAL_LM 表示因果語言模型(Causal Language Model),即生成式任務target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "demo_proj"],inference_mode=False,r=lora_rank,lora_alpha=lora_alpha,lora_dropout=0.1)model = get_peft_model(model, peft_config)model.is_parallelizable = Truemodel.model_parallel = Trueprint("start load train data...")# sampler參數和shuffle參數互斥train_params = {"batch_size": batch_size, "shuffle": False, "num_workers": 0}training_set = QADataset(train_json_path, tokenizer, max_source_length, max_target_length)# 區別1:使用DistributedSampler實現分布式數據采樣,此時的訓練參數中就不要設置隨機打亂train_sampler = DistributedSampler(training_set)training_loader = DataLoader(training_set, **train_params, sampler=train_sampler)print("start load validation data...")val_params = {"batch_size": batch_size, "shuffle": False, "num_workers": 0}val_set = QADataset(val_json_path, tokenizer, max_source_length, max_target_length)val_sampler = DistributedSampler(val_set)val_loader = DataLoader(val_set, **val_params, sampler=val_sampler)# 日志記錄writer = SummaryWriter(logs_dir)# 優化器optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)model = model.to(device)# 區別2:將模型傳遞給DistributedDataParallelmodel = DistributedDataParallel(model)# 開始訓練print("Start Training...")train_model(model=model,train_loader=training_loader,val_loader=val_loader,optimizer=optimizer,gradient_accumulation_steps=gradient_accumulation_steps,device=device,num_epochs=epochs,model_output_dir=model_output_dir,writer=writer,sampler=train_sampler)
與單卡訓練除了上述代碼中注釋的兩點區別之外,就是需要使用torchrun
來啟動訓練程序
torchrun --nproc_per_node=8 pytorch_ddp.py
其中nproc_per_node表示GPU數量
使用accelerate實現,代碼改動的地方也很少,只需要把模型、優化器、數據集等傳遞給Accelerate即可
from accelerate import Acceleratordef main():model_name = "../modelscope/Qwen/Qwen2.5-1.5B-Instruct"train_json_path = "./data/train.json"val_json_path = "./data/val.json"max_source_length = 128 # todo 輸入長度最大可以設置為多少?max_target_length = 256 # todo 輸出呢?epochs = 10batch_size = 1 # todo 顯存大了之后可以增大,如何控制多卡訓練lr = 1e-4gradient_accumulation_steps = 16lora_rank = 8 # 8或16或32lora_alpha = 32model_output_dir = "output"logs_dir = "logs"# 設備# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")# 加載分詞器和模型tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)# 使用accelerate 混合精度訓練bf16,這里也設置為bfloat16,否則可能會導致沖突報錯model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, trust_remote_code=True)# setup peftpeft_config = LoraConfig(task_type=TaskType.CAUSAL_LM, # 任務類型:CAUSAL_LM 表示因果語言模型(Causal Language Model),即生成式任務target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "demo_proj"],inference_mode=False,r=lora_rank,lora_alpha=lora_alpha,lora_dropout=0.1)model = get_peft_model(model, peft_config)model.is_parallelizable = Truemodel.model_parallel = Trueprint("start load train data...")train_params = {"batch_size": batch_size, "shuffle": True, "num_workers": 0}training_set = QADataset(train_json_path, tokenizer, max_source_length, max_target_length)training_loader = DataLoader(training_set, **train_params)print("start load validation data...")val_params = {"batch_size": batch_size, "shuffle": True, "num_workers": 0}val_set = QADataset(val_json_path, tokenizer, max_source_length, max_target_length)val_loader = DataLoader(val_set, **val_params)# 日志記錄writer = SummaryWriter(logs_dir)# 優化器optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)accelerate = Accelerator()model, optimizer, train_data, val_data = accelerate.prepare(model, optimizer, training_loader, val_loader)# 開始訓練print("Start Training...")train_model(model=model,train_loader=train_data,val_loader=val_data,optimizer=optimizer,gradient_accumulation_steps=gradient_accumulation_steps,num_epochs=epochs,model_output_dir=model_output_dir,writer=writer,accelerate=accelerate)
除了代碼改動,還需要初始化配置:
accelerate config
根據提示,選擇配置即可,例如:
配置文件示例:
compute_environment: LOCAL_MACHINE
debug: false
distributed_type: MULTI_GPU
downcast_bf16: 'no'
enable_cpu_affinity: true
gpu_ids: 0,1,2,3,4,5,6,7
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
配置好之后,可以先進行配置運行測試:
accelerate test
如果正常運行可以看到如下提示:
Test is a success! You are ready for your distributed training!
當測試指定配置文件時,使用 --config_file 參數 accelerate test --config_file path_to_config.yaml
啟動訓練腳本:
accelerate launch accelerate_test.py
如果需要指定配置文件,與test同理示例:
accelerate launch --config_file path_to_config.yaml accelerate_test.py
注意:–config_file 放在要運行的腳本前面
此外,還可以通過命令行參數覆蓋配置文件中的默認參數。
FSDP
因為現在有很多千億級規模的大模型,單卡的顯存是一定無法加載模型的,所以需要一種技術可以將模型參數分配到多張卡上,FSDP 切分了模型參數、梯度與優化器參數,使得每張 GPU 只保留這些參數的一部分。
上面實現DDP的兩種方式只有accelerate支持FSDP訓練,在初始化配置時,在這一步:
Do you want to use FullyShardedDataParallel?
選擇yes
例如(不過我這個里面很多配置是隨便選的,不一定合理)
FSDP 的參數 ShardingStrategy 的不同取值決定了模型的劃分方式:
-
FULL_SHARD: 將模型參數、梯度和優化器狀態都切分到不同的GPU上,類似ZeRO-3。
-
SHARD_GRAD_OP: 將梯度、優化器狀態切分到不同的GPU上,每個GPU仍各自保留一份完整的模型參數。類似ZeRO-2。
-
NO_SHARD: 不切分任何參數。類似ZeRO-0。
以下是來自LLamaFactory的一個FSDP配置文件示例
compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAPfsdp_backward_prefetch: BACKWARD_PREfsdp_forward_prefetch: falsefsdp_cpu_ram_efficient_loading: truefsdp_offload_params: true # offload may affect training speedfsdp_sharding_strategy: FULL_SHARDfsdp_state_dict_type: FULL_STATE_DICTfsdp_sync_module_states: truefsdp_use_orig_params: true
machine_rank: 0
main_training_function: main
mixed_precision: fp16 # or bf16
num_machines: 1 # the number of nodes
num_processes: 2 # the number of GPUs in all nodes
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
配置完成之后,同樣通過accelerate命令可以啟動訓練腳本,為了測試DDP、FSDP策略確實已經生效,我進行了如下實驗:
方案一:速度和顯存占用對比
7B模型,同一批數據,同樣的上下文長度,同樣的精度:bf16
使用pytorch_ddp可以單卡跑,并實現多卡同時訓練,單卡占用顯存20G+,訓練需要20分鐘+
使用accelerate ddp配置,單卡占用顯存20G+,訓練需要20分鐘+
(1.5B模型上述兩種方式占用顯存和耗時也基本一致)
但是使用accelerate fsdp配置,當設置加載模型的精度為bfloat16時,會報錯:
ValueError: Must flatten tensors with uniform dtype but got torch.bfloat16 and torch.float32
如果針對1.5B模型,統一去除設定的bf16精度,即采用float32和bf16混合精度訓練
使用pytorch_ddp可以單卡跑,并實現多卡同時訓練,單卡占用顯存12G+,訓練需要16分鐘+
使用accelerate ddp配置,單卡占用顯存12G+,訓練需要16分鐘+
使用accelerate fsdp,單卡占用顯存6G+,訓練需要6小時+
這里就證明了,fsdp與ddp的區別,表示配置生效
如果針對7B模型,統一去除設定的bf16精度,即采用float32和bf16混合精度訓練,
pytorch_ddp和accelerate ddp均顯存不夠,但是accelerate fsdp可以跑
訓練需要30小時+,單卡占用顯存15~22G(根據不同批次數據上下文長度有關)
這里也證明了,fsdp可以跑需要更大顯存的精度
方案二:
使用14B模型,使得單卡無法運行,然后再使用fsdp來運行
pytorch_ddp 無法運行
accelerate ddp無法運行
accelerate fsdp配置,因為無法設置bfloat16精度,暫時也跑不起來,如果解決這一問題,應該就可以跑起來了
這個問題搞了很久都沒有解決,推測是優化器或者損失函數等相關計算過程引入float32類型,因為使用Trainer實現的訓練代碼跑起來沒有問題。關于使用Trainer實現的LoRA微調代碼,下一篇會繼續介紹。