DeepSpeed多卡ChatGLM微调实战：从显存优化到LoRA选型

简介：面向需要上手大模型微调的开发者与研究者，这份资源以DeepSpeed多卡训练为主线，完整覆盖ChatGLM微调实战流程，并提供可直接运行的模块化源码与流程教程。项目包含详细的环境搭建、数据准备与格式化、LoRA/P-tuning/Freeze等多类微调方案，同时讲解ds_config.json等关键配置参数、训练启动脚本写法和常见踩坑点，有助于读者跨越单卡显存限制，快速搭建多卡并行训练环境。压缩包共17个文件，以Python训练脚本（11个）、Shell启动脚本（3个）、JSON配置文件（2个）和Markdown说明文档（1个）为主，大小仅118KB，目录划分明确，便于按模块对照学习。已有267人浏览学习。借助源码中的模型加载、数据加载、训练循环、推理脚本和评估测试等模块，读者不仅能复现ChatGLM多卡微调，还能理解DeepSpeed的内存优化与梯度累积机制，是一份可直接落地的优质实战教程。

1. 多卡ChatGLM微调：为什么大家都卡在DeepSpeed这一关

把ChatGLM-6B塞进一张A100-80G里做推理是轻松的，但一旦开始全量微调，Adam优化器的状态一上来，显存立刻不够用。很多人对“多卡微调”的第一反应是插4张卡等于4倍显存，真正跑起来才发现，权重、梯度、优化器状态这堆东西不切开，4张卡也只是各自抱着同一份副本在空转。DeepSpeed就是为解决这个问题出现的：它把模型状态切到各卡上，让基于DeepSpeed的多卡ChatGLM微调从“跑得动”变成“跑得起”。这篇实战笔记按环境搭建、LoRA与全量选型、deepspeed_config.json参数、多卡启动和踩坑排查的顺序，给出一套可以直接复制到你自己机器上的完整流程。适合手里有2到8张GPU、想把ChatGLM微调成垂直场景模型的算法工程师和独立开发者。

2. 微调前的选型：全量微调和LoRA微调，先算一笔账

2.1 全量微调为什么在多卡上这么贵

全量微调意味着模型每一层参数都在更新，训练过程中需要同时存放五样东西：权重本身、梯度、优化器的一阶动量、二阶动量，以及混合精度下的FP32主权重。工程上常用一个近似数来算账：每1B参数大约要占16GB的模型状态。ChatGLM-6B换算下来，光是参数、梯度和优化器状态就要接近96GB。一张A100-80G单卡装不下，即便勉强装下，batch只能设为1，训练效率低到没法看。

多卡场景里最直接的做法是Data Parallel，也就是每张卡复制一份完整模型，各自算梯度再做一次allreduce合并。4张卡就是4份96GB，显存不仅没省，反而因为通信开销变得更慢。DeepSpeed在这里的核心价值是ZeRO，它把原本每张卡都保存一份的优化器状态、梯度甚至参数切分到不同卡上，通信只发生在真正需要同步的时刻。这也是“多卡微调为什么绕不开DeepSpeed”的最直接原因：对于一个6B量级的模型，光模型状态就已经超过单卡物理上限，不用ZeRO只能靠堆显存解决，而堆显存的性价比太低。

这里还需要区分一个概念：模型状态和激活值。模型状态指的是参数、梯度、优化器状态，激活值是前向传播过程中每一层输出的中间张量。很多人以为开了DeepSpeed就能解决所有显存问题，结果发现batch一大照样OOM，这就是激活值在瓶颈上。后面第4章会给出一个完整的显存估算方法，先把模型状态算明白，再结合batch和序列长度去反推激活值的余量。

2.2 LoRA微调是什么意思：参数少到什么程度

刚接触大模型微调的人常问LoRA微调是什么意思。简单说，LoRA冻结原模型的全部权重，在每一层（通常是注意力层的线性映射）旁边挂一组低秩矩阵A和B，只训练这两个小矩阵。最终更新量等于alpha/r乘以AB的结果，再叠加到原始权重上。由于原始权重完全不更新，可训练参数量会从6B级别直接掉到百万级别。

以ChatGLM-6B为例，在query_key_value上挂LoRA，r设为8，可训练参数大约只有三百万到四百万，比原始参数小了三个数量级。这会带来两个直接好处：一是模型状态显存占用大幅下降，因为需要保存的参数、梯度和优化器状态都只针对这小部分可训练参数；二是训练过程不容易把基座的语言能力带偏，特别适合数据量不大、目标明确的垂直微调场景。

from peft import LoraConfig, get_peft_model lora_config = LoraConfig( r=8, # 低秩矩阵的秩，常用 8/16/32 lora_alpha=16, # 缩放系数，实际更新量乘 alpha/r target_modules=["query_key_value"], # ChatGLM 注意力层线性映射的参数名 lora_dropout=0.1, bias="none", ) model = get_peft_model(model, lora_config) model.print_trainable_parameters()

这里有几个参数需要根据实际任务调：r太小，表达能力不够；r太大，可训练参数变多，显存优势和微调“安全边际”都会缩水。lora_alpha和r的比值决定LoRA更新的影响幅度，一般alpha取r的1到2倍。target_modules要根据模型结构改，ChatGLM的注意力层里query、key、value是合并在一个名为query_key_value的线性层里的，所以要写这个参数名；如果微调ChatGLM2或ChatGLM3，这个名称可能不同，先打印模型结构确认再填。

对比全量微调，LoRA微调的选择逻辑可以归纳为：数据量小、任务垂直度高、希望快速迭代，选LoRA；数据量足够大、显卡资源充足、追求极限效果，才考虑全量微调。对于多数真实项目，LoRA都是第一版方案，因为它能在一个晚上跑完，而且效果差距往往没有想象中大。

2.3 ChatGLM的数据拼接和标签掩码：Prefix-LM结构下的labels

ChatGLM是Prefix-LM结构，左侧的输入可以做双向注意力，右侧的生成部分只能看左侧。这和GPT的Causal-LM不一样，训练数据不能直接当普通CausalLM来拼接。社区里最常见的做法是：把指令和回答拼成一条完整文本，在回答开始的位置之前把标签全部设为-100，计算loss时这些位置被忽略。

def pack_chat_input(tokenizer, source, answer, max_len=2048): # 拼接格式和ChatGLM官方prompt模板保持一致 text = "问：" + source + "\n答：" + answer + tokenizer.eos_token ids = tokenizer.encode(text, max_length=max_len, truncation=True) # 计算prompt部分的长度，这部分不参与loss计算 prompt = "问：" + source + "\n答：" prompt_len = len(tokenizer.encode(prompt)) labels = [-100] * len(ids) # 先把所有位置设为-100 labels[prompt_len:] = ids[prompt_len:] # 回答部分保留真实token return {"input_ids": ids, "labels": labels}

这段代码的核心逻辑是“回答token要学习，指令token不学习”。-100是PyTorch CrossEntropyLoss约定俗成的ignore索引，模型在算loss时会自动跳过这些位置。注意这里的prompt_len用的是tokenizer.encode，不是字符长度，因为中文文本的token切分不是按字来的。

训练循环里有一个特别容易翻车的细节：用了DeepSpeed之后，反向传播和参数更新必须交给engine，不能直接调model.backward()或optimizer.step()。原因在于ZeRO把模型状态切到了多张卡上，只有engine知道当前rank上持有哪部分状态。

for step, batch in enumerate(dataloader): outputs = engine( input_ids=batch["input_ids"].cuda(), labels=batch["labels"].cuda(), ) engine.backward(outputs.loss) engine.step()

这里的engine就是deepspeed.initialize返回的模型包装对象。前向接口和原始模型基本一致，但backward和step必须走engine。如果自己额外再调一次optimizer.step()，轻则梯度更新错乱，重则训练过程直接崩掉。

3. 环境搭建与最小复现：从安装deepspeed到多卡启动

3.1 安装deepspeed包：先过CUDA这一关

搜索“安装deepspeed包一直报错”，你会发现十个里有八个是同一个原因：pip在安装时尝试编译CUDA拓展，但系统里找不到CUDA工具链。报错信息里通常会有一大段gcc、nvcc相关的日志，很多人看到就懵了，其实解决办法很简单：先确认PyTorch的CUDA版本，再检查nvcc能不能用。

# 先确认torch的CUDA版本，再决定装哪个deepspeed python -c "import torch; print(torch.version.cuda)" # 安装前确认nvcc可用 nvcc --version # 常规安装方式 pip install deepspeed

如果pip直接从源码编译导致报错，可以先用预编译的wheel包缓解。另一个常见做法是设置CUDA_HOME环境变量指向CUDA Toolkit的安装目录，再执行pip install deepspeed。安装完成后，用ds_report命令检查deepspeed是否识别到当前GPU、CUDA版本和NCCL版本。这个命令会打印一份环境诊断报告，看到“OK”字样说明基础环境没有问题。

需要注意deepspeed会在第一次启动训练时做JIT编译，编译一些CUDA算子到缓存目录，所以第一次启动多等一两分钟是正常的。如果多次启动都卡在同一处编译，多半是系统缺了gcc或者编译依赖，看日志里的具体报错再补装对应依赖。

3.2 写一个最小训练脚本：train_chatglm_lora.py

环境就绪后，先别急着上全量微调，我一般用LoRA把整个链路先跑通。下面这个脚本是完整可运行的最小版本，把模型加载、LoRA包装、DeepSpeed初始化和训练循环放在一起，总共不到40行。

import torch import deepspeed from transformers import AutoTokenizer, AutoModel from peft import LoraConfig, get_peft_model def train(): model_name = "THUDM/chatglm-6b" tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) # bf16要求Ampere及以上架构，V100等旧卡请改成fp16 model = AutoModel.from_pretrained( model_name, trust_remote_code=True, torch_dtype=torch.bfloat16, ) # 用LoRA冻结大部分参数，只训练query_key_value上的低秩矩阵 model = get_peft_model(model, LoraConfig( r=8, lora_alpha=16, target_modules=["query_key_value"], lora_dropout=0.1, bias="none", )) # 梯度检查点用计算换显存，多卡显存紧张时建议打开 model.gradient_checkpointing_enable() engine, optimizer, _, lr_scheduler = deepspeed.initialize( model=model, model_parameters=[p for p in model.parameters() if p.requires_grad], config="ds_config.json", ) for step, batch in enumerate(dataloader): outputs = engine( input_ids=batch["input_ids"].cuda(), labels=batch["labels"].cuda(), ) engine.backward(outputs.loss) engine.step()

这段代码里有两个地方最容易出错。一是AutoModel.from_pretrained必须带trust_remote_code=True，因为ChatGLM的模型结构通过远程代码加载；二是deepspeed.initialize的model_parameters参数只传可训练参数，LoRA模式下就是那几个低秩矩阵，不要传model.parameters()的完整列表，否则DeepSpeed会把冻结参数也纳入优化器管理，白白浪费显存。

dataloader需要配合第2.3节的pack_chat_input函数，把每条样本处理成input_ids和labels。注意collate时要动态padding到batch内最长，而不是固定到2048，否则显存大部分都被无效的padding token吃掉了。

3.3 启动命令：单机多卡和多机多卡的差异

训练脚本写好后，启动方式用deepspeed命令，不是python命令。单机4卡的最简命令如下：

deepspeed --num_gpus 4 --master_port 29500 train_chatglm_lora.py

master_port是分布式训练的控制通信端口，默认29500。如果同一台机器上同时跑多个训练任务，端口会冲突，启动时会报端口被占用的错误，换一个比如29501就行。如果你用的是SLURM之类的调度系统，还要带上deepspeed_multinode脚本一起用，这里不展开。

多机多卡稍微复杂一点，每台机器上都要执行启动命令，区别只在node_rank和master_addr：

# 在第一个节点执行 deepspeed --num_nodes 2 --num_gpus 8 --node_rank 0 \ --master_addr 192.168.1.10 --master_port 29500 train_chatglm_lora.py # 在第二个节点执行 deepspeed --num_nodes 2 --num_gpus 8 --node_rank 1 \ --master_addr 192.168.1.10 --master_port 29500 train_chatglm_lora.py

多机场景下master_addr必须填第一台机器的IP，而且两台机器要能互相访问。这里有个容易踩的坑：如果机器之间只有以太网没有InfiniBand，通信速度会成为训练瓶颈，需要显式设置NCCL_SOCKET_IFNAME指向实际网卡接口名。判断方法是在训练日志里看NCCL初始化时用了哪块网卡，如果是lo或docker0，训练速度大概率不正常。

4. deepspeed_config.json：先设这几个参数，再谈训练收敛

4.1 ZeRO Stage选1、2还是3：zero123的区别

第一次写deepspeed_config.json，卡住最多的问题就是deepspeed zero123的区别。简单说，Stage 1只切分优化器状态，Stage 2在Stage 1基础上连梯度也切分，Stage 3更进一步把模型参数本身也切分到各卡上。Stage越高，显存越省，但卡间通信量越大，训练速度越慢。

ChatGLM-6B微调，我的建议是优先Stage 2。6B模型的参数量不算特别大，Stage 2已经能把优化器状态和梯度的重复存储省掉，显存能被压到接近“每卡只存一份完整权重+部分状态”的水平。Stage 3虽然有额外收益，但通信开销会拖慢迭代速度，在4卡或8卡规模下性价比不高。只有当显存实在不够、batch小到没法看的时候，才考虑Stage 3，或者把优化器状态offload到CPU。

DeepSpeed初始化时会打印一份运行配置摘要，里面包括当前用的是哪个Stage、每张卡预估的显存使用量。启动日志不要直接跳过，这些信息比任何第三方工具都准。Stage 3还要注意一个问题：参数被切分后，前向和反向过程中需要动态收集参数，所以eval和保存模型时要额外处理，对新手来说这是比显存更麻烦的坑。

4.2 我习惯先写好的几个配置项

下面这份ds_config.json是我做ChatGLM微调时最常使用的配置模板，每一项都直接影响训练能否收敛。

{ "train_batch_size": 32, "gradient_accumulation_steps": 4, "gradient_clipping": 1.0, "bf16": { "enabled": true }, "zero_optimization": { "stage": 2, "offload_optimizer": { "device": "cpu" } }, "scheduler": { "type": "WarmupDecayLR", "params": { "warmup_min_lr": 0, "warmup_max_lr": 2e-5, "warmup_num_steps": 200 } }, "steps_per_print": 20 }

第一眼要搞清楚的是train_batch_size、gradient_accumulation_steps和实际每卡batch三者的关系。train_batch_size是全局batch，也就是所有卡累积多少条样本做一次参数更新。实际每卡micro batch的计算公式是：train_batch_size除以卡数再除以梯度累积步数。以这份配置为例，4卡、累积4步、全局32，那么每卡micro batch是2。很多人在这一步凭感觉填数字，结果实际batch大小和预期完全不是一回事，loss曲线自然不对。

bf16和fp16的区别也值得多提一句。bf16在A100、H100这些Ampere之后架构的卡上效果更好，指数位更多，不容易出现loss溢出；V100等旧架构不支持bf16，只能开fp16，这时要额外关注loss scale，防止训练中期loss变成NaN。

4.3 用显存公式反推可用的batch和序列长度

显存占用粗略分为两块：模型状态加激活值。模型状态可以用第2.1节的16GB/B参数估算，激活值则和micro batch大小、序列长度、模型层数强相关。ChatGLM-6B在max_seq_len为2048时，激活值随batch增长的斜率非常明显，很多人以为batch设2没问题，结果跑到第100步就OOM。

我建议的排查顺序是先固定序列长度，把micro batch调到1跑一轮，用nvidia-smi看显存占用基线，再逐步往上加。

# 每秒刷新一次显存情况，观察训练过程中的峰值占用 watch -n 1 nvidia-smi

看重点不是显存总量，而是进程占用那一栏。如果micro batch为1时已经用了接近全部显存，说明瓶颈在激活值，这时候调低max_len或者关掉冗余输入更有效。如果micro batch为1时显存还有余量，才把batch往上加。DeepSpeed在启动时也会给一个显存预估，可以和实际观测对照，帮助判断配置是否生效。

这里有一个容易误判的地方：打开了offload_optimizer之后，训练刚开始时显存会明显下降，但CPU内存占用会涨上去。如果服务器内存本身不多，offload反而可能引发内存交换，训练速度骤降。观察指标时不能只看GPU显存，还要带着看内存占用。

5. DeepSpeed多卡微调ChatGLM的四个典型踩坑现场

5.1 NCCL超时：卡在Waiting for the barrier

现象：启动多卡训练后，日志卡在分布式初始化阶段，几分钟后抛NCCLTimeoutError，或者某个rank直接挂掉。

原因：最常见的是master_addr或master_port配置不对，多卡之间的控制通信建立不起来；其次是机器上有多块网卡，NCCL选错了网卡，导致通信走了一个慢速通道甚至在不通的网口上。还有一种情况是防火墙拦了分布式训练的端口，单机多卡时不太会遇到，多机多卡时概率很高。

解决：先确认所有节点能互相ping通，再检查deepspeed启动命令里的master_addr和master_port是否一致。多机场景显式指定通信网卡：

export NCCL_SOCKET_IFNAME=eth0 export NCCL_IB_DISABLE=1

NCCL_IB_DISABLE=1是让NCCL不走InfiniBand，只走以太网。如果机器之间的IB网络没有正确配置，与其让它自动探测失败，不如直接关掉。日志里能看到NCCL初始化时用的网卡名字，如果发现选错了，再用NCCL_SOCKET_IFNAME强制指定。

5.2 训练中途OOM：先分清是激活值还是模型状态压爆的

现象：训练能启动，但跑到某个step突然报CUDA out of memory，而且每次崩的位置不一样。

原因：大多数时候是动态padding没有生效，batch里某条样本特别长，把激活值峰值瞬间拉高。另一个常见原因是micro batch算错了，以为设置了梯度累积batch就小了，实际上每卡的micro batch仍然是全局batch除以卡数，没有除以累积步数。

解决：先按第4.3节的方法，把micro batch降到1试跑，如果不再OOM，说明是batch和序列长度的问题；如果batch为1仍然OOM，说明模型状态部分还有优化空间。此时从三个方向同时压：打开gradient_checkpointing启用梯度检查点、把offload_optimizer设备设为cpu、以及确认没有把冻结参数传进优化器。这三招用完，ChatGLM-6B的LoRA微调在24GB显存卡上都可以跑起来。

注意：开了gradient_checkpointing之后，前向速度会变慢，这是正常的，它本质上是牺牲计算换显存。

5.3 loss不降或乱跳：学习率、warmup和数据顺序一起背锅

现象：训练了上千步，loss要么原地不动，要么忽高忽低，看起来完全没有收敛特征。

原因：把之前训练小模型的经验搬到大模型上，学习率用1e-4甚至更高，对6B模型来说步长太大，loss自然乱跳。另一个是warmup步数太少或没设置，调度器一上来就给满学习率。还有一个容易被忽略的点：dataloader没有shuffle，模型反复看到同一条指令，loss曲线会出现周期性波动。

解决：把warmup_max_lr降到2e-5，warmup_num_steps设为总步数的5%到10%。检查dataloader是否设置了shuffle=True。如果用的是fp16而不是bf16，还要观察训练日志里的loss scale，出现NaN或不正常跳变说明fp16的精度控制出了问题，建议升级到支持bf16的显卡或者调整fp16配置。

这里要特别说一个误区：不要因为loss在下降就认定训练正常。大模型微调初期loss下降可以很快，但真正的考验是第6章的验证环节，生成质量和loss下降并不完全挂钩。

5.4 微调完导出eval时模型像没训练过

现象：训练日志里loss已经降得很低，但eval时模型回答和原版几乎一样，甚至更差。

原因：最常见的是保存模型时只保存了DeepSpeed的引擎权重，没有保存LoRA adapter。DeepSpeed的save_checkpoint保存的是引擎状态，直接拿去推理完全用不上。另一个原因是没有merge LoRA权重，推理时忘了加载adapter。

解决：训练结束后用PeFT的save_pretrained保存LoRA权重，推理时先加载adapter再合并：

# 训练结束保存LoRA权重 engine.module.save_pretrained("lora_adapter") # 推理时加载并合并 from transformers import AutoModel model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True) model.load_adapter("lora_adapter") model = model.merge_and_unload()

注意推理时加载基座模型的精度要和训练时一致，否则merge可能出现数值不匹配的问题，生成结果看起来就像没训练过。eval时使用训练时的prompt模板也很关键，微调时用的格式是“问：xxx\n答：”，eval时如果换成其他格式，模型生成质量会受到影响。

6. 从loss下降走向业务能用：验证流程和一个显存技巧

6.1 用固定prompt集验证效果，别只看loss

训练loss下降只是“模型记住了训练数据”的必要条件，不是充分条件。我的做法是准备一组固定业务问题，训练过程中每隔几百步用生成模式跑一遍，观察回答质量的变化趋势。这一步能比loss曲线更早暴露出模型是否过拟合、是否丢失了对话能力。生成测试代码可以挂在训练循环里：

def eval_generate(engine, tokenizer, questions, max_new_tokens=128): engine.module.eval() with torch.no_grad(): for q in questions: prompt = "问：" + q + "\n答：" inputs = tokenizer(prompt, return_tensors="pt").to(engine.device) out = engine.module.generate( **inputs, max_new_tokens=max_new_tokens, do_sample=False, ) print(tokenizer.decode(out[0], skip_special_tokens=True)) engine.module.train()

固定prompt集的规模不用大，10条左右就够了，但要覆盖正常提问、边界提问和易混淆提问三类。do_sample设为False，保证可对比；如果模型训练中关了dropout，这一步的生成结果是完全确定的。每500步跑一次，保存生成结果到日志文件，后面就可以回看不同训练阶段的质量变化。另外把一部分数据留作held-out集，验证集loss比训练loss更早上升，就是过拟合信号。

6.2 把显存再压一档：动态padding和梯度检查点组合

第5.2节提到动态padding能解决一部分OOM问题，这里给出一个可用的collate实现。核心思路是把batch内的样本padding到同一长度，同时用-100掩盖无效的标签位置。

def collate_batch(batch): max_len = max(len(x["input_ids"]) for x in batch) input_ids = [] labels = [] for x in batch: pad_len = max_len - len(x["input_ids"]) input_ids.append(x["input_ids"] + [0] * pad_len) labels.append(x["labels"] + [-100] * pad_len) return { "input_ids": torch.tensor(input_ids), "labels": torch.tensor(labels), }

这里把padding token的id直接写成了0，是因为ChatGLM的pad_token_id就是0。如果换了其他模型，记得用tokenizer.pad_token_id获取，不要写死。梯度检查点配合动态padding，是当前单卡24GB或4卡121GB这种中低配环境里，把ChatGLM-6B LoRA微调跑起来的关键组合。

我自己在做这类多卡微调项目时养成了一个习惯：每次改完配置，先跑一个100步的smoke test，只验证三件事——loss下降方向对不对、显存占用是否有大幅波动、保存下来的checkpoint能不能重新加载。三关全过再放长训练。这个习惯帮我在多机多卡场景下避开了大量看似玄学的翻车；很多问题其实早在小规模测试阶段就会暴露。希望帮到你。

本文还有配套的精品资源，点击获取