LoRA等参数高效微调(PEFT)方法相较于完全微调,能显著减少可训练参数的数量,但扩大PEFT训练规模,通常仍需要采用分布式计算方法。这主要有两方面原因:基础大型语言模型(LLM)本身仍然非常庞大,其正向和反向传播需要大量内存和计算资源,并且在大规模数据集上高效训练得益于并行处理,可以缩短实际运行时间。为PEFT调整标准分布式训练框架,需要了解PEFT如何与数据并行和模型并行技术配合使用。PEFT 的数据并行最常用的分布式方法是数据分布式并行(DDP)。在标准DDP中(例如,使用PyTorch的DistributedDataParallel),模型会在多个设备(GPU)上进行复制。每个设备处理不同的mini-batch数据,本地计算梯度,然后梯度会在所有设备之间同步(通常通过AllReduce操作),之后优化器在每个副本上更新模型权重。当将DDP应用于PEFT时:基础模型复制: 庞大且冻结的基础模型仍在每个设备上复制。这意味着与存储基础模型权重相关的主要内存瓶颈依然存在。PEFT 参数复制: 小量可训练的PEFT参数(例如LoRA矩阵 $A$ 和 $B$、适配器权重或提示嵌入)也会被复制。梯度同步减少: 这是PEFT在DDP中一个显著的优点。由于只需要同步PEFT参数的梯度,AllReduce步骤中的通信开销相对于完全微调中同步整个模型的梯度会大幅降低。如果一个模型有 $N$ 个总参数和 $P$ 个PEFT参数,且 $P \ll N$,那么梯度通信量将按比例减少。实施: 将PEFT与标准DDP框架结合通常很直接。Hugging Face的PEFT等库会自动将只有适配器参数标记为可训练。当使用DistributedDataParallel包装模型时,它能正确识别这些可训练参数,并且只同步它们的梯度。# 示例伪代码:PyTorch DDP 结合 PEFT import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from transformers import AutoModelForCausalLM from peft import get_peft_model, LoraConfig # 假设使用了PEFT库 # 初始化分布式环境(例如,使用torchrun) dist.init_process_group(backend='nccl') local_rank = dist.get_rank() torch.cuda.set_device(local_rank) # 加载基础模型 base_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf") base_model.to(local_rank) # 配置并应用PEFT(例如,LoRA) peft_config = LoraConfig( r=16, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.05, bias="none", task_type="CAUSAL_LM" ) model = get_peft_model(base_model, peft_config) model.print_trainable_parameters() # 验证只有PEFT参数是可训练的 # 使用DDP包装模型 # DDP将自动处理可训练(PEFT)参数的梯度同步 ddp_model = DDP(model, device_ids=[local_rank]) # 优化器仅针对可训练参数 optimizer = torch.optim.AdamW(ddp_model.parameters(), lr=1e-4) # AdamW会隐式过滤requires_grad=True的参数 # --- 训练循环 --- # for batch in dataloader: # outputs = ddp_model(**batch) # loss = outputs.loss # loss.backward() # 只计算PEFT参数的梯度 # optimizer.step() # 梯度同步后更新PEFT参数 # optimizer.zero_grad() # --- 训练循环结束 --- dist.destroy_process_group()即使梯度通信量减少,大型基础模型的正向和反向传播计算成本在DDP中每步训练时间中仍是主导因素。DeepSpeed ZeRO 提升内存效率尽管DDP通过分配数据批次有助于扩展计算能力,但它不减少每个GPU存储模型权重、梯度和优化器状态所需的内存占用。DeepSpeed的ZeRO(零冗余优化器)提供将这些组件在数据并行工作器之间分区的技术,显著降低了每个GPU的内存需求。ZeRO阶段与PEFT的配合:ZeRO Stage 1 (优化器状态分区): 在设备间分区优化器状态(例如AdamW中的动量)。这对PEFT非常有利,因为即使是少量参数的优化器状态也能占用大量内存(例如,AdamW的FP32状态每个参数使用12字节)。ZeRO-1有效降低了这种内存开销。ZeRO Stage 2 (梯度 + 优化器状态分区): 额外地在设备间分区梯度。对于PEFT来说,梯度本身就较小,因此分区它们所带来的额外内存节省可能不如完全微调那样显著。不过,它仍然有助于整体内存减少,并且与DDP风格的梯度计算很好地配合。ZeRO Stage 3 (参数 + 梯度 + 优化器状态分区): 分区模型参数本身。这是内存节省最激进的阶段。PEFT 的适用性: ZeRO-3在PEFT中的主要使用场景出现在当基础模型即使在推理模式下也太大以至于无法放入单个GPU内存中时。ZeRO-3分区所有参数,包括冻结的基础模型权重。在正向和反向传播过程中,每个设备会动态地从其他设备收集必要的参数。开销: ZeRO-3引入了显著的参数收集通信开销,这有时会抵消数据并行带来的计算加速,特别是在较慢的网络互连下。PEFT 特定参数: ZeRO-3也会分区可训练的PEFT参数。虽然这能节省少量内存,但其在PEFT环境中的主要优点是能够使用原本无法容纳的庞大基础模型。为PEFT选择ZeRO阶段:如果基础模型可以轻松地适应每个GPU:ZeRO-1或ZeRO-2通常是最佳选择。它们减少了优化器状态(以及ZeRO-2的梯度)的内存,同时没有ZeRO-3的参数收集通信开销。这允许比普通DDP更大的批次大小或适应稍微更大的基础模型。如果基础模型无法适应单个GPU:ZeRO-3变得必要,用于分区基础模型参数。PEFT参数会随之一起分区。请注意增加的通信成本。将PEFT与DeepSpeed结合通常涉及配置DeepSpeed JSON文件,并使用deepspeed.initialize函数初始化模型、优化器和数据加载器。请确保传递给DeepSpeed的优化器配置为只优化可训练的PEFT参数。// 示例DeepSpeed配置文件片段 (ds_config.json) 用于PEFT的ZeRO Stage 2 { "train_batch_size": 32, "gradient_accumulation_steps": 1, "optimizer": { "type": "AdamW", "params": { "lr": 1e-4, "betas": [0.9, 0.999], "eps": 1e-8, "weight_decay": 0.01 } }, "scheduler": { "type": "WarmupLR", "params": { "warmup_min_lr": 0, "warmup_max_lr": 1e-4, "warmup_num_steps": 100 } }, "zero_optimization": { "stage": 2, "offload_optimizer": { "device": "cpu", // 可选:将优化器状态卸载到CPU内存 "pin_memory": true }, "contiguous_gradients": true, "overlap_comm": true }, "gradient_clipping": 1.0, "fp16": { "enabled": true // 或 bf16: { "enabled": true } } } # 示例伪代码:DeepSpeed 结合 PEFT import deepspeed from transformers import AutoModelForCausalLM from peft import get_peft_model, LoraConfig # --- 设置PEFT模型(如DDP示例中所示)--- # model = get_peft_model(base_model, peft_config) # 如果需要,可以显式过滤优化器参数, # 尽管如果模型设置正确,DeepSpeed通常会处理此问题。 # optimizer_grouped_parameters = [ # {'params': [p for n, p in model.named_parameters() if p.requires_grad], 'weight_decay': 0.01} # ] # optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=1e-4) # 或使用DeepSpeed优化器配置 # 初始化DeepSpeed # model.parameters()调用应只返回可训练的PEFT参数 model_engine, optimizer, _, _ = deepspeed.initialize( model=model, # model_parameters=model.parameters(), # 或显式传递参数 # optimizer=optimizer, # 可以传递预配置的优化器 config_params='ds_config.json' # DeepSpeed配置文件的路径 ) # --- 使用model_engine的训练循环 --- # for batch in dataloader: # loss = model_engine(**batch).loss # model_engine.backward(loss) # model_engine.step() # 处理优化器步骤、梯度裁剪、调度器 # --- 训练循环结束 ---内存使用对比分布式方法的选择显著影响每个GPU的内存使用。与完全微调相比,PEFT本身就减少了梯度和优化器状态所需的内存。ZeRO进一步优化了这一点。{"layout": {"title": {"text": "每个GPU的预估内存使用(示意图)"}, "xaxis": {"title": {"text": "训练策略"}}, "yaxis": {"title": {"text": "每个GPU内存 (GB)"}, "range": [0, 80]}, "barmode": "stack", "legend": {"title": {"text": "内存组成部分"}}}, "data": [{"type": "bar", "name": "基础模型权重", "x": ["Full FT DDP", "PEFT DDP", "PEFT ZeRO-2", "PEFT ZeRO-3"], "y": [40, 40, 40, 10], "marker": {"color": "#495057"}}, {"type": "bar", "name": "可训练权重", "x": ["Full FT DDP", "PEFT DDP", "PEFT ZeRO-2", "PEFT ZeRO-3"], "y": [0, 0.1, 0.1, 0.02], "marker": {"color": "#748ffc"}}, {"type": "bar", "name": "梯度(可训练)", "x": ["Full FT DDP", "PEFT DDP", "PEFT ZeRO-2", "PEFT ZeRO-3"], "y": [40, 0.1, 0.1, 0.02], "marker": {"color": "#ff922b"}}, {"type": "bar", "name": "优化器状态(可训练)", "x": ["Full FT DDP", "PEFT DDP", "PEFT ZeRO-2", "PEFT ZeRO-3"], "y": [80, 0.2, 0.05, 0.01], "marker": {"color": "#12b886"}}, {"type": "bar", "name": "激活值(预估)", "x": ["Full FT DDP", "PEFT DDP", "PEFT ZeRO-2", "PEFT ZeRO-3"], "y": [15, 15, 15, 15], "marker": {"color": "#be4bdb"}}]}这是大型模型微调场景下,每个GPU的预估内存组成。完全微调DDP需要存储完整的权重、梯度和优化器状态。PEFT DDP大幅减少了梯度和优化器状态的内存。ZeRO-2进一步减少了优化器状态的内存(可能进行卸载)。ZeRO-3分区所有组件,包括基础模型权重,提供最低的每个GPU占用空间,但可能带来更高的通信开销。激活内存很大程度上取决于批次大小和序列长度,这里为了比较假设为常数。选择恰当的方法选择合适的PEFT分布式方法涉及考虑:基础模型大小: 如果基础模型能适应单个GPU,DDP或ZeRO-1/2是高效的选择。如果不能,则需要ZeRO-3。GPU内存: 内存有限的情况下更倾向于选择ZeRO阶段,可能伴随CPU卸载优化器状态(ZeRO-1/2)或参数(ZeRO-3,对于PEFT来说不太常见,除非基础模型非常庞大)。PEFT参数数量: 尽管相对于基础模型而言较少,但大量的适配器参数(例如多适配器配置)可能仍然受益于ZeRO-1/2优化器状态分区。网络互连: 高带宽、低延迟的互连(如NVLink或Infiniband)可以减轻梯度同步(DDP、ZeRO-1/2)和参数收集(ZeRO-3)的通信开销。较慢的网络(如以太网)会使得ZeRO-3的吸引力较低,除非为了适应模型而绝对必要。训练速度与内存: 通常存在权衡。如果内存允许,DDP通常计算速度最快。ZeRO阶段节省内存,但可能引入通信瓶颈,可能减慢训练步骤。通过仔细考虑这些因素并善用PyTorch DDP和DeepSpeed ZeRO等框架,你可以有效地扩展PEFT微调,以处理大型模型、大型数据集和复杂的多适配器场景,同时高效管理硬件资源。