ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

DeepSpeed多卡微调ChatGLM:ZeRO显存优化与避坑指南

DeepSpeed多卡微调ChatGLM:ZeRO显存优化与避坑指南 简介基于 DeepSpeed 的 ChatGLM 多卡微调实战项目包面向有一定深度学习基础、希望快速上手大模型微调的研究者与开发人员。资源定位在解决单机多卡环境下微调 ChatGLM 的配置复杂、资源管理困难等问题提供从环境搭建、数据准备、模型训练到评估测试的完整流程教程。包体仅 118KB共 17 个文件以 Python 代码为主体涵盖模型加载、数据加载、训练循环等核心模块同时配有 3 个 Shell 启动脚本、2 个 JSON 配置文件和 1 份 Markdown 说明文档便于直接复现实验。目前已有 267 人浏览学习。读者可由此掌握基于 DeepSpeed 做多卡并行训练的关键参数配置与启动方式并使用其中封装的 LoRA、Ptuning、Freeze 三类微调方法处理不同任务场景。源码注释清晰、模块划分明确既适合入门者按教程逐步操作也适合有经验者对照修改有效缩短 ChatGLM 微调项目的落地周期。1. 单卡跑不动、多卡不会配DeepSpeed 是入门大模型微调的第一道分水岭手里两张 3090想微调 ChatGLM-6B 做垂直业务单卡 batch size 开到 2 就爆显存迭代一轮要跑大半天——这是很多人的真实处境。大模型微调不是把数据喂进去就行基座模型、微调范式、分布式训练策略三件事没定下来显存和训练速度会一起失控。用 DeepSpeed 做多卡微调是目前把 ChatGLM 这类 6B 量级模型真正训起来的常规路径。它靠 ZeRO 显存优化、CPU offload、梯度累积把多张 GPU 的显存整合到一起让原本单卡塞不下的训练配置变得可运行把训练时间从周级别压到小时级别。这篇按我从环境配置到效果验证完整走过的流程讲每一步都给出能直接复用的脚本和参数也把最容易翻车的几个坑提前说出来。2. 微调前必须定的三件事基座模型、微调范式与多卡并行策略2.1 选 ChatGLM 还是 Qwen2.5中文能力和显存预算先对齐定基座模型时很多人一上来就看榜单分数实际落地时要先算一笔显存账。ChatGLM-6B 的 FP16 权重占 12GB 左右加上激活值、梯度和优化器状态全参微调单卡基本要 22GB 以上如果选 Qwen2.5-7B参数量更大显存占用会再往上抬一截。我见过不少团队在 2 张 3090 上硬跑 7B 全参最后只能把 sequence length 压到 512效果反而比 6B 还差。如果你所在行业已经有比较成熟的公开中文语料qwen2.5-7b 微调行业大模型是社区热度很高的组合但 ChatGLM 的中文指令跟随表现更稳生态里现成的微调案例也多遇到问题容易搜到解决方案。我的建议是显存总量 48GB 以下优先 ChatGLM-6B48GB 以上再考虑更大的基座不要在显存临界点上赌玄学。还要确认基座模型的授权和商用边界。ChatGLM 系列是开源权重但不同版本的开源协议有差异商用前要再看一眼当时版本的授权说明。模型选型这件事放到整个微调流程里只占半天决策时间但后面所有显存估算、训练时长、推理部署成本都由它决定。2.2 全参、LoRA 还是 P-Tuning v2显存、效果和可控性的三角权衡基座定下来之后紧接着要选择微调范式。ChatGLM 官方示例里提供的是 P-Tuning v2只训练 prompt encoder 那一小部分参数显存占用低但效果上限受限于可学习的参数规模适合样本量不大、任务比较单一的场合。如果你只有几十条到几百条标注数据P-Tuning v2 能快速跑通一个基线。LoRA 是产业里用得最多的方案。它冻结原始权重只训练注入的低秩矩阵可学习参数通常在 0.1% 到 1% 之间。以 ChatGLM-6B 为例LoRA 的显存占用可以压到全参微调的一半以下效果在大多数指令微调场景里已经很接近全参。我做行业大模型微调时默认都是 LoRA 起步跑通之后再决定要不要尝试全参。全参微调的效果上限最高但对显存、数据质量和调参经验的要求也最高。数据少于一万条时全参反而容易过拟合训练出来的模型在测试集上可能表现得比 LoRA 还差。三者的权衡可以看这张表微调范式可训练参数占比显存占用效果上限适用场景P-Tuning v20.01% 级别低中低小样本、单任务LoRA0.1%-1%中高行业微调、快速迭代全参微调100%高最高数据量大、算力充足选定范式之后DeepSpeed 的角色就清晰了它对三种范式都适用但在全参和 LoRA 场景下收益最明显。多卡并行不是魔法它解决的是“单卡物理显存放不下”和“单卡算力不够”这两个硬约束。2.3 DeepSpeed 在其中的位置ZeRO 机制与多卡加速的本质先理解 DeepSpeed 解决什么问题。普通的数据并行训练每张卡都保存一份完整的模型参数、梯度和优化器状态。4 张卡跑 ChatGLM-6B 全参微调优化器状态用 AdamW 时要占参数量的 8 倍以上显存翻倍式增长卡越多浪费越明显。DeepSpeed 的 ZeRO 把这三样东西切分到各张卡上。stage 1 只切分优化器状态stage 2 切分梯度和优化器状态stage 3 连模型参数也切分。我做单机多卡微调时默认用 stage 2因为 stage 3 在训练过程中需要频繁做参数聚合通信开销大4 卡规模下收益不明显反而可能变慢。ZeRO 之外的几个机制也值得了解。offload_optimizer 把优化器状态放到 CPU 内存能用更少显存跑更大的 batch但训练速度会下降适合显存不够时应急。梯度累积则把多个 micro batch 的梯度攒起来再更新一次表面上看是省显存实际上是调整全局 batch size 的手段。数据并行、ZeRO、梯度累积这三者的关系决定了多卡训练的速度上限和效果稳定性第 4 章会展开到具体参数。3. 环境与数据准备CUDA、conda 和训练数据的硬性门槛3.1 先跑这个命令检查 CUDA 与显卡环境对了才谈多卡环境配置这步会卡住一半人。很多翻车案例不是模型训练问题而是 PyTorch 装成了 CPU 版本或者 CUDA 驱动和 PyTorch 编译版本不匹配导致torch.cuda.is_available()返回 False。拿到一台多卡机器第一步我会先跑一段最小检查nvidia-smi python -c import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.device_count())nvidia-smi看驱动版本和每张卡的显存状态确认 4 张卡都能被系统正常识别。第二行命令验证 PyTorch 是否编译了 CUDA 支持以及能调用几张卡。如果输出里torch.cuda.is_available()是 False先别急着配 DeepSpeed把 PyTorch 装回 CUDA 版本再继续。多卡机器还要检查卡间通信。用nvidia-smi topo -m看两张卡之间的拓扑如果是 PCIe 直连通信带宽有限如果是 NVLink多卡训练的效率会高很多。这个问题在单机多卡上影响不大但如果之后要扩展到多机多卡节点间的网络带宽会成为关键瓶颈。驱动版本我建议保持较新的稳定版CUDA Toolkit 用 11.8 或 12.1 都行关键是 PyTorch 的预编译 wheel 要和它对应。很多老教程还在用 CUDA 10.x新显卡驱动根本不认照抄就会在环境这步浪费一整天。3.2 conda 环境复现Python 版本与依赖锁定的最小集合环境问题一旦出现排查起来非常耗时。我一般用 conda 建一个独立环境把依赖版本锁定避免和服务器上其他项目互相污染。完整的依赖如下conda create -n glm-ft python3.10 -y conda activate glm-ft pip install torch2.1.2 torchvision0.16.2 --index-url https://download.pytorch.org/whl/cu121 pip install transformers4.36.2 datasets peft deepspeed0.13.1 acceleratePython 3.10 是当前大模型生态兼容性最好的版本。torch 2.1.2 对应 CUDA 12.1这一组版本我跑过多次和 DeepSpeed 0.13.1 配合稳定。transformers 4.36.2 支持 ChatGLM 系列通过trust_remote_codeTrue加载datasets 用于数据集的加载和预处理peft 是 LoRA 的官方实现库accelerate 在部分训练脚本里会被间接依赖。装完之后跑一个快速验证python -c import deepspeed; print(deepspeed.__version__) python -c from transformers import AutoModel; print(ok)如果 DeepSpeed 的__version__能正常打印说明安装成功。这里有个小坑DeepSpeed 在安装时会尝试编译一些 CUDA 算子如果编译过程报错先看是不是 gcc 版本太老常见做法是apt install build-essential装一下编译工具链。3.3 训练数据格式instruction-input-output 的清洗与容量估算ChatGLM 微调最常用的数据格式是 JSON每条样本包含指令、输入和期望输出三个字段。下面是一个典型的对话样本{ instruction: 请根据以下合同条款判断甲方是否违约, input: 甲方未在约定的30天内支付第二期款项且未提供任何书面说明。, output: 甲方已构成违约。根据合同第7.2条甲方应在30天内支付第二期款项逾期未付且未说明理由属于明确违约情形。 }数据量不是越多越好。几百条高质量数据就能看到一个明显的效果变化但要达到“能用”的程度我一般建议至少 3000 到 5000 条。数据质量比数据量更关键常见问题包括标签前后不一致、output 里包含额外解释导致模型学到错误的输出格式、样本之间重复度过高。清洗时我会按照几个固定步骤处理去重、去空、检查标签一致性、过滤超长样本。每条样本的 token 长度直接影响训练显存如果某个样本特别长比如超过 2048 token要么截断要么单独处理不要混在正常样本里。估算训练数据的 token 总量有个简单公式平均每条样本 token 数乘以样本数。这个总量除以全局 batch size就是一个 epoch 的步数可以用来预估训练时长。4. DeepSpeed 多卡微调实操从启动命令到 ds_config 参数4.1 启动命令与工程结构deepspeed --num_gpus 的前后顺序环境就绪、数据就位之后进入核心的实操环节。多卡微调的工程结构并不复杂我习惯把训练脚本、配置文件、数据和输出目录分开放glm-ft/ ├── train_chatglm.py ├── ds_config.json ├── data/ │ └── train.json └── output/训练脚本的启动命令用 DeepSpeed 的 launcher 来跑而不是直接用 python。这一行命令是最容易被忽略的地方deepspeed --num_gpus4 --master_port29500 train_chatglm.py \ --model_name_or_path /data/models/chatglm-6b \ --data_path ./data/train.json \ --output_dir ./output \ --train_batch_size 128 \ --micro_batch_size 4 \ --num_train_epochs 3 \ --learning_rate 2e-5--num_gpus4告诉 DeepSpeed 使用 4 张卡--master_port是分布式训练的控制通信端口。如果服务器上同时有别人在跑训练29500这个默认端口很容易冲突报错信息通常是地址被占用换一个不常用的端口就能解决。--train_batch_size 128是全局 batch size--micro_batch_size 4是单卡一次前向传播的样本数两者的关系由梯度累积步数连接。这里的--micro_batch_size 4不是随便定的。ChatGLM-6B 在 24GB 显存的卡上FP16 混合精度 最长序列 1024 的条件下micro batch size 开到 4 是比较稳的值。如果序列长度更长比如 2048就要降到 2 甚至 1。全局 batch size 的选择会影响收敛效果一般分类任务用 32 到 128 都常见太小噪声大太大收敛慢需要根据数据规模调整。4.2 ds_config.json 逐字段拆解ZeRO stage、offload 与混合精度DeepSpeed 的核心配置都在ds_config.json里训练脚本通过这个文件感知分布式训练的所有细节。下面是我在 4 张 3090 上微调 ChatGLM-6B 的完整配置{ train_batch_size: 128, train_micro_batch_size_per_gpu: 4, gradient_accumulation_steps: 8, optimizer: { type: AdamW, params: { lr: 2e-5 } }, fp16: { enabled: true, loss_scale: 0, loss_scale_window: 1000, initial_scale_power: 16 }, zero_optimization: { stage: 2, offload_optimizer: { device: cpu, pin_memory: true }, allgather_partitions: true, allgather_bucket_size: 5e8, reduce_scatter: true, contiguous_gradients: true } }配置里的参数对应关系需要理解到位。train_batch_size是全局 batch sizetrain_micro_batch_size_per_gpu是每张卡每个微批的样本数gradient_accumulation_steps是梯度累积步数三者满足一个固定等式全局 batch 单卡 micro batch × 梯度累积 × 显卡数。这里 4 × 8 × 4 128正好配平。关于 offload_optimizer这是一个取舍点。开启后优化器状态放到 CPU 内存显存占用明显下降但训练速度会慢 10% 到 20%。我的建议是显存够用就不要开显存不够或者想跑更大的 micro batch 再开。fp16 混合精度尽量保持开启不开启的话显存占用会直接翻倍而且训练速度显著变慢。参数作用建议值train_batch_size一次优化更新消耗的总样本数数据量小时用 32-64大时用 128train_micro_batch_size_per_gpu单卡单步前向反向的样本数24GB 显存从 4 起步gradient_accumulation_steps累积多少微批才更新参数由另外两个参数推导不要手填矛盾zero_optimization.stageZeRO 切分等级单机多卡用 stage 2offload_optimizer.device优化器状态存放位置显存吃紧才设 cpufp16.enabled是否开启混合精度保持 true有一个常见误区是三个 batch 参数不匹配。DeepSpeed 启动时会校验train_batch_size和另外两个参数的乘积是否一致如果不一致直接报错。报错信息会明确告诉你当前的 grad_accum 应该是多少照着改就行。4.3 接住 DeepSpeed 的训练循环Trainer 还是手动循环配置写好之后训练脚本里最关键的是如何让模型、优化器和 DeepSpeed 引擎正确对接。直接看代码import torch import deepspeed from transformers import AutoModel, AutoTokenizer tokenizer AutoTokenizer.from_pretrained(/data/models/chatglm-6b, trust_remote_codeTrue) model AutoModel.from_pretrained(/data/models/chatglm-6b, trust_remote_codeTrue).half() model, optimizer, _, lr_scheduler deepspeed.initialize( modelmodel, model_parametersmodel.parameters(), configds_config.json ) for step, batch in enumerate(train_loader): inputs tokenizer(batch[text], return_tensorspt, paddingTrue, truncationTrue, max_length1024).to(model.device) outputs model(**inputs, labelsinputs[input_ids]) loss outputs.loss model.backward(loss) model.step()这段代码里有三个容易出错的地方。第一deepspeed.initialize返回的 model 已经是被 DeepSpeed 包装过的对象不能用原来的loss.backward()必须调用model.backward(loss)由 DeepSpeed 处理梯度切分和跨卡同步。第二model.step()替代了原来手动调optimizer.step()和optimizer.zero_grad()的流程DeepSpeed 会根据配置的梯度累积步数自动决定何时更新参数。第三train_loader 要自己构造DeepSpeed 不负责数据加载。如果你更习惯用 HuggingFace 的 Trainer也可以在训练参数里指定deepspeedds_config.jsonTrainer 会读取配置并完成同样的初始化流程。我早期用 Trainer 省事但出了问题不好判断是 Trainer 的问题还是 DeepSpeed 的问题后来改成手动初始化整个流程在自己的掌控之内排查效率高了不少。手动循环还有一个优势可以在每个 step 里打印 loss、学习率和显存占用实时观察训练状态。加两行代码就能实现if step % 10 0: print(fstep {step}, loss {loss.item():.4f}, lr {lr_scheduler.get_last_lr()[0]:.2e}, mem {torch.cuda.max_memory_allocated() / 1024**3:.2f}GB)5. 多卡微调避坑指南NCCL 超时、显存 OOM 与 Loss 不收敛的排查清单5.1 NCCL 超时与初始化卡死多卡通信最常见的翻车现场现象启动训练后rank 0 开始正常打印日志但其他 rank 一直卡在初始化阶段几分钟后报NCCL timeout错误整个进程退出。原因最常见的是master_port被占用或者多机多卡场景下MASTER_ADDR和NODE_RANK没设置对。另一个隐蔽原因是服务器防火墙阻止了节点间的通信端口。我遇到过一次单机多卡跑得好好的换成多机多卡就卡死排查了半天发现是第二台机器的网卡 IP 写错了。解决先把--master_port换成一个冷门端口比如29888。如果还卡在启动命令前加一行环境变量看详细日志export NCCL_DEBUGINFONCCL_DEBUGINFO会打印通信初始化的详细过程哪个 rank 在等谁一目了然。确认是网络问题后检查MASTER_ADDR是否指向主节点的内网 IPNNODES是否和实际机器数一致。多机多卡不是简单地把--num_gpus改成 8 就行每台机器的启动命令还要带--node_rank这些参数我在第一次跑多机时踩了整整一个下午。5.2 CUDA OOM先看微调范式再谈调 ZeRO stage现象训练跑到第三个 batch直接报CUDA out of memory报错信息里能看到是哪一层 transformer 爆的显存。原因显存爆掉的原因分三类。第一类是 micro batch size 太大24GB 卡跑 1024 序列长度的 ChatGLMmicro batch size 超过 4 就会危险。第二类是序列长度过长个别样本超过 2048 token显存峰值被拉高。第三类是开启了全参微调但没开梯度检查点激活值占用直接翻倍。解决按优先级依次处理。先开梯度检查点只需一行model.gradient_checkpointing_enable()这行代码用计算换显存训练速度会慢一些但显存峰值能降低 30% 到 40%多数情况下这一行就解决问题。如果还不行把train_micro_batch_size_per_gpu降到 2 或 1。仍然不行再把offload_optimizer打开把优化器状态挪到 CPU。这三步走完OOM 基本不会再来。这里要提醒一个容易忽略的点如果用的是 LoRA显存占用远低于全参不要让 LoRA 把显存放宽之后就把 batch 拉得很大。batch 太大会让收敛变慢模型在验证集上的表现不升反降。5.3 Loss 不降或持续震荡学习率与数据质量是主角现象loss 从 2.x 开始训练了 1000 步还在 2.x 附近横盘或者每过几个 step loss 突然跳高再回落震荡得厉害。原因最常见的是学习率太大。大模型微调的学习率一般是 1e-5 到 5e-5 这个量级比训练小模型用的 1e-3 低两到三个数量级。另一个高频原因是数据问题——很多样本的 output 字段里混入了原始文本模型试图记住这些噪声loss 自然降不下去。解决先把学习率降一个数量级比如从 2e-5 降到 5e-6看 loss 是否开始平稳下降。如果没改善检查数据清洗环节随机抽 20 条样本人工看一遍确认 output 里没有脏数据。还要检查max_length设置如果训练样本的序列长度超过截断值大量样本的标签在截断后被丢弃模型学到的东西会很有限。Loss 不收敛的场景里约一半是学习率问题另一半是数据问题。环境问题导致的 loss 异常很少见所以先从这两点入手排查不要动不动就怀疑 DeepSpeed 配置。5.4 DeepSpeed 启动了但实际没生效日志里找这行字现象训练能正常跑loss 也在下降但 nvidia-smi 里看到 GPU 利用率很低训练速度甚至比单卡还慢4 张卡跑出了 1 张卡的效果。原因训练脚本里虽然用了 deepspeed 命令启动但实际没有调用deepspeed.initialize或者ds_config.json的路径没传对DeepSpeed 退化成普通数据并行。这种情况不会有报错因为 PyTorch 的 DistributedDataParallel 也能跑多卡只是 ZeRO 的显存优化完全没生效。解决看训练启动时的日志。DeepSpeed 初始化成功时会打印一段包含模型参数量的信息其中会明确显示 zero stage 的数值。如果日志开头没有出现DeepSpeed info相关的输出说明 DeepSpeed 没有被真正加载。检查两点训练脚本里是否调用了deepspeed.initializeconfigds_config.json的路径是否相对于当前工作目录正确。遇到这种问题我一般在训练脚本第一行加一个断言import deepspeed assert deepspeed.initialize is not None, deepspeed not loaded这样至少能确认 import 成功。更直接的验证方法是在deepspeed.initialize之后打印model.global_batch_size如果和配置里的train_batch_size一致说明 DeepSpeed 已经接管。5.5 Checkpoint 保存、ZeRO 权重合并与断电后悔药现象训练跑了一晚上第二天发现服务器断电或者被运维重启所有进度丢失只能从头再训。原因没有配置 checkpoint 保存机制或者训练脚本里根本没有写保存逻辑。另一个更隐蔽的问题是ZeRO stage 2 训练时每个 rank 只保存自己分片的那部分优化器和梯度状态直接用torch.save(model.state_dict())保存的权重是不完整的。解决DeepSpeed 提供了专门的保存接口训练循环里每隔固定步数调用一次if step % 500 0: model.save_16bit_model(./output/ckpt) client_state {step: step, trainer_state: step} model.save_checkpoint(./output/ckpt, tagfstep_{step}, client_stateclient_state)save_16bit_model会合并各 rank 的权重输出一份完整的 FP16 模型这个文件可以直接用来加载推理。save_checkpoint保存的是 DeepSpeed 的完整训练状态包括优化器状态和当前步数配合resume_from_checkpoint可以无缝续训deepspeed --num_gpus4 train_chatglm.py \ --resume_from_checkpoint ./output/ckpt/step_1500checkpoint 是训练过程中的后悔药配置好之后再遇到断电、蓝屏、OOM 终止都不会太心疼。我自己跑长训练时每 500 步保存一次磁盘占用不大但心理安全感强很多。6. 验证与进阶loss 曲线之外还要看真实问答与推理部署6.1 微调效果的三个验证层次训练完成不代表微调成功。我先看训练 loss 和验证 loss 的差距判断是否过拟合然后跑测试集 prompt 对比微调前后的回答质量最后抽一批真实业务问题做人工打分。第三层最重要loss 再低模型在真实问题上胡说八道也没用。加载合并后的权重测试生成的代码很简单tokenizer AutoTokenizer.from_pretrained(./output/ckpt, trust_remote_codeTrue) model AutoModel.from_pretrained(./output/ckpt, trust_remote_codeTrue).half().cuda() prompt 请根据以下合同条款判断甲方是否构成违约甲方未在30天内支付第二期款项。 inputs tokenizer(prompt, return_tensorspt).to(cuda) out model.generate(**inputs, max_new_tokens200) print(tokenizer.decode(out[0], skip_special_tokensTrue))6.2 进阶路径LoRA DeepSpeed、梯度检查点和多机扩展验证通过之后你的微调流程已经闭环了。此后的进阶方向是把 LoRA 和 DeepSpeed 结合进一步降低显存占用。做法是先通过 HuggingFace 的 AutoModel 加载模型再用 peft 包一层 LoRA最后把包好的 model 传给deepspeed.initialize训练完用merge_and_unload()合并权重再部署。这个方案适合需要频繁迭代的行业模型场景每次微调都能省下大量算力。多卡微调的最终形态是多机多卡但这需要万兆以上网络才有效果。单机多卡 4 张卡跑不动的场景加一台机器多机多卡是有价值的如果只是想跑得更快先把单机多卡吃透它的效率提升空间远比贸然跨机器来得大。最后说一个我自己的习惯每次微调开始前先把ds_config.json和启动命令存成独立文件训练完把日志里最后的 loss 记录保存下来。这个做法帮我积累了不同数据规模下的最佳参数组合下一次微调直接照着上次的参数起步省掉了大量试错时间。希望帮到你。本文还有配套的精品资源点击获取
返回列表