大模型LoRA微调实践
大模型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微调代码，下一篇会继续介绍。