简介本资源面向希望上手大模型微调的开发者与研究者聚焦用Deepspeed实现ChatGLM的多卡并行训练解决单卡显存不足、训练效率低、微调流程不清晰等实际问题适合具备Python与PyTorch基础、想进阶多卡训练的中级学习者。压缩包共17个文件约118KB包含11个py源码、3个sh启动脚本、2个json配置与1个md说明分别对应模型加载、数据加载、训练循环、评估测试及LoRA、P-Tuning、Freeze等微调方式的启动配置。教程覆盖环境搭建、数据准备与格式化、参数配置、多卡训练脚本启动到模型评估的完整链路源码注释清晰、模块化程度高便于理解与二次修改。目前已有267人学习。读者可据此掌握Deepspeed内存优化与梯度累积的配置思路对照源码跑通多卡微调流程并迁移到自有数据集与业务场景中。1. 多卡微调 ChatGLM从单卡爆显存到 Deepspeed 拉满四张卡的实战路径单卡 24G 显存跑 ChatGLM3-6B 的 LoRA 微调batch size 只能压到 1序列长度卡在 512训练一轮要十几个小时——这是很多人第一次做大模型微调实战时遇到的真实窘境。Deepspeed 的 ZeRO 系列优化配合多卡并行能把同样的任务压缩到几小时内完成而且显存占用下降一个量级。这套方案解决的核心问题是如何在不购买 A100 80G 的前提下用手头的多张消费级或入门级 GPU 完成 ChatGLM 系列模型的高效微调。适合已经跑通过单卡 LoRA、想进一步压缩训练时间或扩大模型规模的从业者也适合需要把微调流程工程化落地的团队。下面从原理选型到代码落地把多卡 ChatGLM 微调的完整链路拆开讲清楚。2. Deepspeed 多卡并行的底层逻辑与 ChatGLM 适配选型2.1 ZeRO 三个阶段到底省了什么Deepspeed 的核心是 ZeROZero Redundancy Optimizer它把训练过程中的模型状态切分到不同 GPU 上从而降低单卡显存占用。理解这三个阶段是选型的基础ZeRO-1只切分优化器状态。以 ChatGLM3-6B 为例全量微调时优化器状态Adam 的 momentum 和 variance占用约 48GB切分到 4 张卡后每卡约 12GB。但模型参数和梯度仍然每卡一份所以显存节省有限。ZeRO-2在 ZeRO-1 基础上切分梯度。每卡只保留自己负责的那部分梯度显存进一步下降。对于 LoRA 微调场景梯度本身不大ZeRO-2 的收益主要体现在全量微调或大 rank 的 LoRA 上。ZeRO-3把模型参数也切分。每卡只存 1/N 的模型参数前向和反向传播时按需从其他卡拉取。这是显存节省最明显的阶段但通信开销也最大。ChatGLM3-6B 在 4 卡 ZeRO-3 下单卡显存可以压到 8GB 以内配合 offload。选型建议LoRA 微调优先用 ZeRO-2全量微调或显存极度紧张时上 ZeRO-3。ChatGLM 的 GLM 架构对 ZeRO-3 的兼容性在较新版本的 Deepspeed 中已经比较稳定但需要注意stage3_gather_16bit_weights_on_model_save这个配置否则保存的权重可能不完整。2.2 ChatGLM 的模型结构对并行策略的影响ChatGLM 采用 GLMGeneral Language Model架构和标准的 LLaMA 系模型有几个关键差异直接影响并行配置第一ChatGLM 使用RoPE 位置编码 二维位置编码的混合方式在序列并行Sequence Parallelism场景下需要额外注意位置编码的切分逻辑。如果直接用 LLaMA 的序列并行方案套上去位置编码会错位。第二ChatGLM 的Multi-Query Attention结构ChatGLM3 用的是 GQA 的变体导致 KV Cache 较小这在推理阶段是优势但在训练阶段对显存的影响不如标准 MHA 那么大。所以显存瓶颈主要在模型参数和优化器状态上ZeRO 的切分策略要优先考虑这两块。第三ChatGLM 的LayerNorm 位置和残差连接方式与 LLaMA 不同在流水线并行Pipeline Parallelism时层切分的边界要避开 LayerNorm 层否则容易出现数值不稳定。实际配置中我一般用ZeRO-2 LoRA的组合在 4 张 309024G上跑 ChatGLM3-6B 的 LoRA 微调序列长度 2048batch size 每卡 4梯度累积 4 步训练速度比单卡快 3.2 倍左右。如果换成 ZeRO-3速度会降到 2.5 倍左右但显存占用从 18GB 降到 9GB可以支持更大的 batch size 或更长的序列。2.3 多卡通信后端的选择与验证Deepspeed 依赖 NCCL 做 GPU 间通信。在开始训练之前必须确认 NCCL 能正常工作。一个常见的验证命令# 检查 NCCL 版本和 GPU 拓扑 python -c import torch; print(torch.cuda.nccl.version()) nvidia-smi topo -mnvidia-smi topo -m的输出会显示 GPU 之间的连接方式。如果是 NVLink 连接通信带宽远高于 PCIeZeRO-3 的通信开销可以忽略不计。如果是 PCIe 4.0 x16带宽约 32GB/sZeRO-3 的 all-gather 操作会成为瓶颈这时候 ZeRO-2 反而更快。另一个关键点是NCCL_P2P_DISABLE和NCCL_IB_DISABLE这两个环境变量。在某些主板或虚拟化环境下P2P 通信会失败需要设置NCCL_P2P_DISABLE1强制走共享内存。这个坑我在三台不同配置的机器上都遇到过现象是训练启动时卡在Initializing NCCL然后超时退出。3. 从零搭建 ChatGLM 多卡微调环境依赖、配置与启动脚本3.1 环境依赖的精确版本组合ChatGLM 的微调对 transformers 和 deepspeed 的版本比较敏感。以下是我在 Ubuntu 20.04 CUDA 11.8 上验证过的组合# 创建虚拟环境 conda create -n chatglm_ds python3.10 -y conda activate chatglm_ds # 安装 PyTorch匹配 CUDA 11.8 pip install torch2.1.2 torchvision0.16.2 torchaudio2.1.2 --index-url https://download.pytorch.org/whl/cu118 # 安装 Deepspeed 和 transformers pip install deepspeed0.12.6 pip install transformers4.36.2 pip install peft0.7.1 pip install datasets2.16.1 pip install accelerate0.25.0 pip install sentencepiece0.1.99 pip install cpm_kernels1.0.11这里有几个版本锁定的理由deepspeed 0.12.x 对 ZeRO-3 的stage3_gather_16bit_weights_on_model_save支持最稳定transformers 4.36.2 是 ChatGLM3 官方推荐的版本peft 0.7.1 修复了 LoRA 合并时的几个 bug。如果版本不对最常见的报错是AttributeError: ChatGLMForConditionalGeneration object has no attribute gradient_checkpointing_enable这是因为 transformers 版本太老不支持 ChatGLM 的梯度检查点。3.2 Deepspeed 配置文件的关键参数下面是一个针对 ChatGLM3-6B LoRA 微调的 ds_config.json跑在 4 卡上{ train_batch_size: 16, train_micro_batch_size_per_gpu: 4, gradient_accumulation_steps: 4, optimizer: { type: AdamW, params: { lr: 2e-4, betas: [0.9, 0.999], eps: 1e-8, weight_decay: 0.01 } }, scheduler: { type: WarmupDecayLR, params: { warmup_min_lr: 0, warmup_max_lr: 2e-4, warmup_num_steps: 100, total_num_steps: 5000 } }, fp16: { enabled: true, loss_scale: 0, loss_scale_window: 1000, initial_scale_power: 16, hysteresis: 2, min_loss_scale: 1 }, zero_optimization: { stage: 2, allgather_partitions: true, allgather_bucket_size: 5e8, overlap_comm: true, reduce_scatter: true, reduce_bucket_size: 5e8, contiguous_gradients: true }, gradient_clipping: 1.0, steps_per_print: 50, wall_clock_breakdown: false }逐项说明关键参数train_batch_size是全局 batch size等于train_micro_batch_size_per_gpu × GPU 数量 × gradient_accumulation_steps。这里 4×4×464但实际配置中我写的是 16因为 LoRA 微调通常不需要太大的 batch size16 已经足够稳定。zero_optimization.stage设为 2对应前面说的 LoRA 场景选型。allgather_bucket_size和reduce_bucket_size设为 5e8500MB这是通信桶的大小。如果 GPU 间是 PCIe 连接可以适当调小到 2e8减少单次通信的数据量提高重叠效率。overlap_comm设为 true让通信和计算重叠进行。这个参数在 ZeRO-2 下效果明显能提升 10%-15% 的吞吐。fp16开启混合精度训练。initial_scale_power设为 16对应初始 loss scale 为 65536。如果训练初期出现 loss 震荡可以降到 12 或 14。3.3 启动脚本与多卡通信配置启动多卡训练的命令行# 设置 NCCL 环境变量 export NCCL_DEBUGINFO export NCCL_P2P_DISABLE0 export NCCL_IB_DISABLE1 export CUDA_VISIBLE_DEVICES0,1,2,3 # 启动训练 deepspeed --num_gpus4 train.py \ --model_name_or_path THUDM/chatglm3-6b \ --data_path ./data/alpaca_zh.json \ --output_dir ./output/chatglm3-lora \ --deepspeed ds_config.json \ --lora_r 8 \ --lora_alpha 32 \ --lora_dropout 0.1 \ --max_seq_length 2048 \ --num_train_epochs 3 \ --per_device_train_batch_size 4 \ --gradient_accumulation_steps 4 \ --learning_rate 2e-4 \ --fp16 True \ --logging_steps 50 \ --save_steps 500 \ --gradient_checkpointing TrueNCCL_IB_DISABLE1是因为大多数消费级主板没有 InfiniBand不禁用的话 NCCL 会尝试走 IB 然后超时。NCCL_P2P_DISABLE0表示启用 P2P 通信如果主板不支持改成 1。--gradient_checkpointing True是必须开的ChatGLM3-6B 在 2048 序列长度下不开梯度检查点单卡显存直接爆。开了之后显存占用从 22GB 降到 14GB 左右代价是训练速度降低约 20%。--lora_r 8和--lora_alpha 32是 LoRA 的秩和缩放系数。ChatGLM 的微调实践中r8 已经能覆盖大部分下游任务如果任务复杂度高可以加到 16 或 32但显存和训练时间会相应增加。3.4 训练脚本中 Deepspeed 初始化的关键代码train.py 中与 Deepspeed 集成的核心部分import torch from transformers import AutoModel, AutoTokenizer, TrainingArguments, Trainer from peft import LoraConfig, get_peft_model, TaskType import deepspeed # 加载模型和 tokenizer model AutoModel.from_pretrained( THUDM/chatglm3-6b, trust_remote_codeTrue, torch_dtypetorch.float16, device_mapauto # 让 accelerate 自动分配设备 ) tokenizer AutoTokenizer.from_pretrained( THUDM/chatglm3-6b, trust_remote_codeTrue ) # 配置 LoRA lora_config LoraConfig( task_typeTaskType.CAUSAL_LM, r8, lora_alpha32, lora_dropout0.1, target_modules[query_key_value], # ChatGLM 的注意力层名称 biasnone ) model get_peft_model(model, lora_config) model.print_trainable_parameters() # 输出类似trainable params: 3,407,872 || all params: 6,742,609,920 || trainable%: 0.05% # 训练参数 training_args TrainingArguments( output_dir./output/chatglm3-lora, per_device_train_batch_size4, gradient_accumulation_steps4, learning_rate2e-4, num_train_epochs3, fp16True, logging_steps50, save_steps500, gradient_checkpointingTrue, deepspeedds_config.json, # 指定 Deepspeed 配置 report_tonone ) # 初始化 Trainer trainer Trainer( modelmodel, argstraining_args, train_datasettrain_dataset, data_collatordata_collator, ) # 开始训练 trainer.train() # 保存 LoRA 权重 model.save_pretrained(./output/chatglm3-lora/final) tokenizer.save_pretrained(./output/chatglm3-lora/final)target_modules[query_key_value]是 ChatGLM 特有的注意力层名称和 LLaMA 的q_proj, k_proj, v_proj不同。如果写错了LoRA 不会应用到任何层训练 loss 不下降。这个坑我踩过排查了半天才发现是 target_modules 没对上。device_mapauto在 Deepspeed 场景下其实会被 Deepspeed 的初始化覆盖但保留它可以让单卡调试时也能正常运行。4. 多卡训练中的避坑与排查从 NCCL 超时到 loss 不收敛4.1 启动时卡在 Initializing NCCL现象运行 deepspeed 命令后终端输出Initializing NCCL然后卡住几分钟后报超时错误。原因通常是 NCCL 尝试走 P2P 或 IB 通信但硬件不支持。也可能是防火墙阻止了 GPU 间的通信端口。解决设置export NCCL_P2P_DISABLE1和export NCCL_IB_DISABLE1强制走共享内存。如果还不行检查nvidia-smi topo -m的输出确认 GPU 之间是否有连接。另外export NCCL_DEBUGINFO可以看到详细的 NCCL 初始化日志定位到具体卡在哪一步。4.2 训练过程中 loss 突然变成 NaN现象前几百步 loss 正常下降然后突然变成 NaN之后一直是 NaN。原因混合精度训练中梯度溢出导致 loss scale 调整失败。ChatGLM 的 LayerNorm 层在 fp16 下容易出现数值不稳定。解决三个措施。第一在 ds_config.json 中把initial_scale_power从 16 降到 12让初始 loss scale 更小。第二开启gradient_clipping设为 1.0裁剪梯度。第三如果还不行把fp16换成bf16bf16 的数值范围更大不容易溢出。但 bf16 需要 GPU 支持30 系及以上的卡才支持。4.3 多卡训练速度反而比单卡慢现象4 卡训练的吞吐量只有单卡的 1.5 倍远低于预期的 3 倍以上。原因通信开销过大或者数据加载成为瓶颈。常见的是allgather_bucket_size设得太大每次通信传输的数据量过多PCIe 带宽吃满。解决把allgather_bucket_size和reduce_bucket_size从 5e8 降到 2e8 或 1e8。同时检查 DataLoader 的num_workers如果设为 0数据加载会阻塞训练。设为 4 或 8让数据预取和 GPU 计算重叠。另外overlap_comm必须设为 true否则通信和计算是串行的。4.4 保存的 LoRA 权重加载后效果不对现象训练完保存 LoRA 权重推理时加载发现模型输出和训练时不一致或者完全没有微调效果。原因ZeRO-3 下保存权重时如果没有设置stage3_gather_16bit_weights_on_model_save保存的是切分后的权重不是完整的。另外ChatGLM 的save_pretrained在 Deepspeed 环境下可能只保存了部分层。解决如果用的是 ZeRO-3在 ds_config.json 中加上stage3_gather_16bit_weights_on_model_save: true。保存后用peft的PeftModel.from_pretrained加载确认target_modules和训练时一致。推理时用model.merge_and_unload()合并 LoRA 权重到基础模型再测试效果。4.5 多卡训练时显存占用不均衡现象4 张卡中卡 0 的显存占用明显高于其他卡甚至卡 0 爆显存而其他卡还有余量。原因Deepspeed 默认把优化器状态和梯度切分到各卡但模型参数在前向传播时可能集中在卡 0。另外如果数据加载没有做分布式采样卡 0 可能处理了更多数据。解决确保使用DistributedSampler做数据切分。在 Trainer 中transformers 会自动处理但如果自己写训练循环需要手动加。另外检查zero_optimization的stage是否设对了ZeRO-2 下模型参数每卡一份显存占用应该均衡。如果用的是 ZeRO-3参数切分后每卡占用应该差不多。5. 进阶技巧用 Deepspeed 的推理引擎加速 ChatGLM 部署训练完的 ChatGLM 模型如果直接用 transformers 的generate方法推理吞吐量很低。Deepspeed 提供了 Inference Engine可以做张量并行推理把模型切分到多卡上显著提升吞吐。5.1 推理配置与训练配置的差异推理场景下ZeRO 的优化策略不适用因为推理不需要优化器状态和梯度。Deepspeed Inference 用的是张量并行Tensor Parallelism把每一层的权重切分到多卡上前向传播时做 all-reduce。一个典型的推理配置import deepspeed import torch from transformers import AutoTokenizer, AutoModel # 加载模型 model AutoModel.from_pretrained( THUDM/chatglm3-6b, trust_remote_codeTrue, torch_dtypetorch.float16 ) tokenizer AutoTokenizer.from_pretrained( THUDM/chatglm3-6b, trust_remote_codeTrue ) # 初始化 Deepspeed 推理引擎 ds_engine deepspeed.init_inference( model, mp_size2, # 张量并行度2 卡 dtypetorch.float16, replace_with_kernel_injectTrue # 使用 DeepSpeed 的优化 kernel ) # 推理 input_ids tokenizer.encode(你好请介绍一下自己, return_tensorspt).to(cuda) output ds_engine.module.generate(input_ids, max_new_tokens100) print(tokenizer.decode(output[0], skip_special_tokensTrue))mp_size2表示用 2 张卡做张量并行。replace_with_kernel_injectTrue会替换 HuggingFace 的注意力实现为 Deepspeed 的优化版本推理速度提升明显。但 ChatGLM 的 GLM 架构对 kernel inject 的支持不如 LLaMA 完善如果报错把这个参数设为 False用原生实现。5.2 推理性能的验证方法验证推理加速效果需要对比三个指标首 token 延迟、每 token 延迟、吞吐量tokens/s。用一个简单的 benchmark 脚本import time import torch def benchmark(model, tokenizer, prompt, max_new_tokens100, num_runs10): input_ids tokenizer.encode(prompt, return_tensorspt).to(cuda) # 预热 for _ in range(3): model.generate(input_ids, max_new_tokensmax_new_tokens) # 计时 torch.cuda.synchronize() start time.time() for _ in range(num_runs): model.generate(input_ids, max_new_tokensmax_new_tokens) torch.cuda.synchronize() end time.time() avg_time (end - start) / num_runs tokens_per_sec max_new_tokens / avg_time print(f平均延迟: {avg_time:.3f}s, 吞吐: {tokens_per_sec:.1f} tokens/s) return tokens_per_sec # 单卡推理 # benchmark(single_model, tokenizer, 你好) # Deepspeed 多卡推理 benchmark(ds_engine.module, tokenizer, 你好)在 2 张 3090 上ChatGLM3-6B 的单卡推理吞吐约 25 tokens/sDeepspeed 张量并行 2 卡后约 42 tokens/s提升约 68%。如果换成 4 卡吞吐可以到 65 tokens/s 左右但边际收益递减因为 all-reduce 的通信开销随卡数增加。5.3 一个容易忽略的细节KV Cache 的显存管理ChatGLM 的 GQA 结构让 KV Cache 比标准 MHA 小很多但在长序列推理时仍然会占用大量显存。Deepspeed Inference 默认会预分配 KV Cache如果max_seq_length设得太大显存浪费严重。我一般会在推理时显式设置max_out_tokens和min_out_tokens控制 KV Cache 的分配。另外如果只是做短文本推理把max_seq_length从 2048 降到 512显存占用可以减少 60% 以上吞吐还能再提升 10%-15%。训练和推理的配置差异很大训练时追求显存节省和通信效率推理时追求吞吐和延迟。把这两套配置分开管理不要混用同一个 ds_config.json否则两边都跑不好。这是我踩过几次坑之后养成的习惯训练配置和推理配置放在不同目录用不同的启动脚本环境变量也分开设置。希望帮到你。本文还有配套的精品资源点击获取