ARTICLE DETAIL

资讯详情

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

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

ChatGLM多卡微调实战:Deepspeed ZeRO显存优化与避坑指南 简介本资源面向希望上手大模型微调的开发者与研究者聚焦用Deepspeed实现ChatGLM多卡并行训练这一实战场景帮助跨过环境配置与分布式训练的技术门槛。压缩包共17个文件以11个Python脚本为核心覆盖模型加载、数据加载、训练循环与评估测试等模块另有3个Shell启动脚本、2个JSON配置文件和1份README说明整体约118KB结构清晰便于按模块查阅。教程完整串联环境搭建、数据准备、微调与评估全流程并逐一讲解Deepspeed参数配置与多卡训练脚本的启动方式源码注释充分、模块化设计良好便于理解与二次修改。目前已有267人学习适合具备一定Python与深度学习基础、想快速跑通ChatGLM多卡微调并掌握Deepspeed优化技巧的读者参考实践。1. 从单卡爆显存到多卡并行ChatGLM 微调为什么绕不开 Deepspeed单卡 24G 显存跑 ChatGLM3-6B 的 LoRA 微调batch size 开到 2 就 OOM这是很多人入坑大模型微调时遇到的第一个下马威。想上多卡PyTorch 原生 DDP 能解决显存问题吗不能——DDP 每张卡仍然持有一份完整的模型参数、梯度和优化器状态6B 模型全量微调光优化器状态就超过 70GB。真正让多卡微调跑得起来的是 Deepspeed 的 ZeRO 系列显存分片技术它把参数、梯度、优化器状态切碎分散到各张卡上配合 CPU Offload 还能进一步把优化器状态挪到内存。这套方案配合 ChatGLM 的 LoRA 微调能在 2 张 24G 卡上把 6B 模型跑起来4 张卡可以开更大的 batch size 和更长的序列。下面从环境配置到多卡启动把整个流程拆开讲清楚每一步都有可复现的命令和参数说明。2. Deepspeed ZeRO 分片与 ChatGLM 微调的适配逻辑2.1 ZeRO-1/2/3 到底分片了什么Deepspeed 的核心是 ZeROZero Redundancy Optimizer它把训练过程中的显存占用拆成三块优化器状态、梯度、模型参数。ZeRO-1 只分片优化器状态显存节省约 4 倍ZeRO-2 加上梯度分片节省约 8 倍ZeRO-3 连模型参数也分片节省倍数与卡数成正比。对于 ChatGLM3-6B 的 LoRA 微调实际可训练参数只有几十到几百 MB但基座模型的参数仍需加载。用 ZeRO-2 配合 LoRA 是比较平衡的选择——优化器状态和梯度分片后单卡显存压力大幅下降通信开销又不像 ZeRO-3 那么大。选 ZeRO-3 也不是不行但每次前向传播都要从其他卡收集参数通信量成倍增加。如果你的卡间是 PCIe 而非 NVLinkZeRO-3 的训练速度会明显下降。我一般建议LoRA 微调用 ZeRO-2全量微调才考虑 ZeRO-3。2.2 ChatGLM 的模型结构与 Deepspeed 配置的对应关系ChatGLM3-6B 是 28 层 Transformer隐藏维度 4096注意力头数 32。这些数字决定了 Deepspeed 配置文件里的几个关键参数。train_batch_size是全局批次大小等于单卡 batch size 乘以卡数再乘以梯度累积步数。train_micro_batch_size_per_gpu是单卡一次前向的 batch size通常设为 1 到 4。gradient_accumulation_steps用来补足全局批次。还有一个容易忽略的点ChatGLM 的 tokenizer 对中文的编码效率较高同样长度的文本ChatGLM 的 token 数通常比 LLaMA 少 20% 到 30%。这意味着你可以把max_length设得比 LLaMA 微调时更长比如 1024 或 2048而不会爆显存。2.3 多卡通信后端与硬件拓扑的检查在写配置文件之前先确认卡间通信拓扑。用nvidia-smi topo -m查看 GPU 之间的连接方式。如果显示 NVLink通信带宽在 200GB/s 以上ZeRO-3 也可以接受如果显示 PIX 或 PHB说明走的是 PCIe带宽只有 16GB/s 到 32GB/s这时候 ZeRO-2 是更稳妥的选择。# 查看 GPU 拓扑结构 nvidia-smi topo -m # 输出示例 # GPU0 GPU1 CPU Affinity NUMA Affinity # GPU0 X NV12 0-31 0 # GPU1 NV12 X 0-31 0 # 如果 GPU0 和 GPU1 之间显示 NV12说明有 NVLink 连接逻辑说明nvidia-smi topo -m输出一个矩阵交叉点显示两张卡之间的连接类型。NV 开头表示 NVLinkPIX 表示经过 PCIe 交换机PHB 表示经过 PCIe 主机桥。参数说明不需要额外参数直接执行即可。如果看到 SYS说明跨 NUMA 节点通信延迟会更高建议用numactl绑定 CPU 核心。3. 从零搭环境ChatGLM 微调依赖与 Deepspeed 安装3.1 创建虚拟环境与安装 PyTorch环境隔离是第一步不要用系统 Python 直接装。推荐 conda 或 venv下面以 conda 为例。# 创建虚拟环境Python 版本选 3.10 conda create -n chatglm-ds python3.10 -y conda activate chatglm-ds # 安装 PyTorchCUDA 版本根据驱动选择这里以 CUDA 11.8 为例 pip install torch2.1.0 torchvision0.16.0 torchaudio2.1.0 --index-url https://download.pytorch.org/whl/cu118 # 验证 PyTorch 是否识别到多卡 python -c import torch; print(torch.cuda.device_count()); print(torch.cuda.get_device_name(0))逻辑说明PyTorch 版本要和 CUDA 驱动匹配。torch.cuda.device_count()返回可用 GPU 数量如果返回 0 说明驱动或 CUDA 版本有问题。参数说明--index-url指定 PyTorch 官方 wheel 源CUDA 11.8 对应 cu118CUDA 12.1 对应 cu121。如果公司内网无法访问外网需要提前下载 wheel 文件离线安装。3.2 安装 ChatGLM 依赖与 DeepspeedChatGLM 的官方仓库提供了requirements.txt但直接装可能会遇到版本冲突。我一般会手动装几个关键包。# 安装 ChatGLM 核心依赖 pip install protobuf transformers4.36.2 cpm_kernels sentencepiece tiktoken # 安装 Deepspeed指定版本避免 API 变动 pip install deepspeed0.12.6 # 安装微调工具链 pip install peft0.7.1 datasets2.16.1 accelerate0.25.0 # 验证 Deepspeed 安装 ds_report逻辑说明transformers版本要锁定ChatGLM 的模型代码对版本敏感4.36.2 是经过验证的稳定版本。deepspeed0.12.6支持 ZeRO-2 和 ZeRO-3 的稳定 API。ds_report会输出 Deepspeed 的编译信息和兼容性检查结果。参数说明peft用于 LoRA 微调datasets用于加载和处理数据accelerate用于分布式启动。如果ds_report显示[WARNING]关于 CPU Adam 的编译失败可以忽略因为 GPU 上用的是 FusedAdam。3.3 下载 ChatGLM 模型权重模型权重从 HuggingFace 或 ModelScope 下载。国内环境建议用 ModelScope速度更稳定。# download_model.py from modelscope import snapshot_download # 下载 ChatGLM3-6B 模型 model_dir snapshot_download( ZhipuAI/chatglm3-6b, cache_dir/data/models/chatglm3-6b, revisionv1.0.0 ) print(f模型下载到: {model_dir})逻辑说明snapshot_download会自动处理大文件的分片下载和断点续传。cache_dir指定模型存储路径建议放在数据盘而非系统盘。参数说明revision指定模型版本不指定则下载最新版。下载完成后目录下会有pytorch_model-00001-of-00007.bin等分片文件总共约 12GB。4. 写对 Deepspeed 配置文件参数逐项拆解与多卡启动4.1 ds_config.json 的完整配置与参数含义Deepspeed 的配置文件决定了显存分片策略、优化器类型、混合精度等核心行为。下面是一份经过验证的 ZeRO-2 配置。{ train_batch_size: 16, train_micro_batch_size_per_gpu: 2, gradient_accumulation_steps: 4, optimizer: { type: AdamW, params: { lr: 2e-5, betas: [0.9, 0.999], eps: 1e-8, weight_decay: 0.01 } }, scheduler: { type: WarmupDecayLR, params: { warmup_min_lr: 0, warmup_max_lr: 2e-5, warmup_num_steps: 100, total_num_steps: 5000 } }, fp16: { enabled: true, loss_scale: 0, initial_scale_power: 16, loss_scale_window: 1000, 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, cpu_offload: false }, gradient_clipping: 1.0, steps_per_print: 50, wall_clock_breakdown: false }逻辑说明train_batch_size是全局批次等于train_micro_batch_size_per_gpu× GPU 数量 ×gradient_accumulation_steps。这里 2 × 2 × 4 16。zero_optimization.stage设为 2 表示启用优化器状态和梯度分片。overlap_comm设为 true 让通信和计算重叠能提升 10% 到 20% 的训练速度。cpu_offload在 ZeRO-2 下通常关闭因为优化器状态已经在 GPU 间分片显存够用如果显存仍然紧张可以开启但训练速度会下降 30% 左右。参数说明allgather_bucket_size和reduce_bucket_size控制通信桶的大小5e8 是 500MB适合 NVLink 环境如果是 PCIe建议降到 2e8 减少单次通信量。fp16开启混合精度训练initial_scale_power是 16 表示初始 loss scale 为 65536这是经过验证的稳定值。4.2 用 torchrun 启动多卡训练Deepspeed 支持多种启动方式推荐用torchrun它是 PyTorch 官方推荐的分布式启动器比deepspeed命令更灵活。# 启动 2 卡训练 torchrun \ --nproc_per_node2 \ --master_port29500 \ train.py \ --model_name_or_path /data/models/chatglm3-6b \ --data_path /data/datasets/alpaca_zh.json \ --output_dir /data/output/chatglm3-lora \ --deepspeed ds_config.json \ --lora_r 8 \ --lora_alpha 32 \ --lora_dropout 0.1 \ --max_length 1024 \ --num_train_epochs 3 \ --save_steps 500 \ --logging_steps 10 \ --learning_rate 2e-5 \ --per_device_train_batch_size 2 \ --gradient_accumulation_steps 4逻辑说明--nproc_per_node2指定每台机器使用 2 张 GPU。--master_port是分布式通信端口如果被占用可以改成 29501 或其他。train.py是训练脚本需要自己写或从 ChatGLM 官方微调代码修改。参数说明--lora_r是 LoRA 的秩8 或 16 是常用值越大可训练参数越多但显存占用也越大。--lora_alpha通常设为lora_r的 2 到 4 倍。--max_length设为 1024ChatGLM 对中文编码效率高1024 token 能覆盖大部分对话场景。4.3 训练脚本中 Deepspeed 的初始化代码训练脚本里需要显式初始化 Deepspeed并处理模型加载和 LoRA 注入。# train.py 关键部分 import torch from transformers import AutoModel, AutoTokenizer, TrainingArguments from peft import LoraConfig, get_peft_model, TaskType import deepspeed # 加载 tokenizer tokenizer AutoTokenizer.from_pretrained( args.model_name_or_path, trust_remote_codeTrue ) # 加载模型注意 torch_dtype 设为 float16 model AutoModel.from_pretrained( args.model_name_or_path, trust_remote_codeTrue, torch_dtypetorch.float16, device_mapNone # Deepspeed 自己管理设备放置 ) # 配置 LoRA lora_config LoraConfig( task_typeTaskType.CAUSAL_LM, rargs.lora_r, lora_alphaargs.lora_alpha, lora_dropoutargs.lora_dropout, target_modules[query_key_value], # ChatGLM 的注意力层名称 biasnone ) # 注入 LoRA model get_peft_model(model, lora_config) model.print_trainable_parameters() # 输出示例trainable params: 3,407,872 || all params: 6,247,000,000 || trainable%: 0.05% # 初始化 Deepspeed model_engine, optimizer, _, _ deepspeed.initialize( modelmodel, model_parametersmodel.parameters(), configargs.deepspeed )逻辑说明device_mapNone是关键如果设成auto会和 Deepspeed 的设备管理冲突。target_modules指定 LoRA 注入的层ChatGLM 的注意力层叫query_key_value和 LLaMA 的q_proj、v_proj不同。deepspeed.initialize会读取配置文件自动处理模型分片、优化器创建和混合精度。参数说明trainable%输出 0.05% 说明 LoRA 只训练了极少参数显存占用主要来自基座模型的前向激活值。5. 避坑指南多卡微调 ChatGLM 的五个血泪教训5.1 现象训练启动后卡在Initializing Deepspeed不动原因通常是 NCCL 通信初始化失败常见于多卡环境变量未设置或端口被占用。解决在启动命令前加export NCCL_DEBUGINFO查看详细日志。如果日志显示NCCL INFO Call to connect returned Connection refused说明master_port被防火墙拦截或已被占用。换一个端口比如 29501 到 29510 之间。另外确认MASTER_ADDR设为127.0.0.1多机训练才需要设为实际 IP。5.2 现象Loss 变成 NaN 或突然飙升到几百原因混合精度训练中 loss scale 不稳定或者学习率设得太大。ChatGLM 的 LoRA 微调学习率建议在 1e-5 到 5e-5 之间超过 1e-4 很容易发散。解决把initial_scale_power从 16 降到 12让初始 loss scale 更小。同时检查数据里有没有空样本或超长样本超长样本被截断后如果全是 padding也会导致 loss 异常。在数据预处理阶段过滤掉 token 数少于 10 的样本。5.3 现象多卡训练速度比单卡还慢原因卡间通信开销超过了并行计算的收益。如果nvidia-smi topo -m显示 GPU 之间是 PHB 或 SYS说明走的是 PCIe 或跨 NUMA通信带宽只有十几 GB/s。解决把allgather_bucket_size和reduce_bucket_size从 5e8 降到 1e8减少单次通信的数据量。同时开启overlap_comm让通信和计算重叠。如果还是慢考虑用 ZeRO-2 而不是 ZeRO-3ZeRO-3 的通信量是 ZeRO-2 的 1.5 倍。5.4 现象保存的 LoRA 权重加载后效果和训练时不一致原因Deepspeed 保存的是分片后的模型状态直接加载会缺少部分参数。解决用model.save_pretrained保存 LoRA 权重而不是用 Deepspeed 的save_checkpoint。在训练脚本里每个 epoch 结束后调用model_engine.module.save_pretrained(output_dir)这样保存的是完整的 LoRA 适配器加载时用PeftModel.from_pretrained即可。注意model_engine.module才是原始模型model_engine是 Deepspeed 包装后的引擎。5.5 现象训练到一半 OOM但显存监控显示还有余量原因显存碎片化。PyTorch 的缓存分配器在长时间训练后会产生碎片导致没有连续的大块显存可用。解决设置环境变量PYTORCH_CUDA_ALLOC_CONFmax_split_size_mb:128限制单次分配的最大块大小减少碎片。另外在训练循环里每 1000 步调用一次torch.cuda.empty_cache()但不要频繁调用否则会拖慢训练速度。6. 进阶技巧用梯度检查点与动态批次把 6B 模型塞进 2 张 24G 卡梯度检查点Gradient Checkpointing是显存优化的最后一根救命稻草。它的原理是不保存中间激活值而是在反向传播时重新计算。这会把显存占用从 O(n) 降到 O(√n)代价是训练速度下降约 30%。在 ChatGLM 微调中开启方式很简单# 在模型加载后开启梯度检查点 model.gradient_checkpointing_enable() # 如果使用 LoRA还需要确保输入需要梯度 model.enable_input_require_grads()逻辑说明gradient_checkpointing_enable会让模型在每层之间丢弃激活值反向时重新前向计算。enable_input_require_grads是 LoRA 微调必须的因为 LoRA 只训练部分参数输入 embedding 默认不需要梯度开启梯度检查点后会报错。参数说明这两个调用没有参数但开启后训练速度会下降建议只在显存不足时使用。动态批次是另一个实用技巧。根据序列长度动态调整 batch size短序列用大 batch长序列用小 batch。实现方式是在数据加载器里按长度分桶每个桶用不同的 batch size。序列长度范围单卡 batch size梯度累积步数等效全局批次0-2568116256-5124216512-102424161024-20481816这张表的用法是在DataCollator里根据当前 batch 的最大长度选择对应的 batch size。实现时用dynamic_batch_sampler把长度相近的样本分到同一个 batch。这样能在不 OOM 的前提下把 GPU 利用率从 60% 提升到 85% 以上。验证训练是否正常不要只看 loss 曲线。我习惯在训练 500 步后跑一次推理测试# 推理验证脚本 from peft import PeftModel from transformers import AutoModel, AutoTokenizer base_model AutoModel.from_pretrained( /data/models/chatglm3-6b, trust_remote_codeTrue, torch_dtypetorch.float16 ).cuda() lora_model PeftModel.from_pretrained( base_model, /data/output/chatglm3-lora/checkpoint-500 ) tokenizer AutoTokenizer.from_pretrained( /data/models/chatglm3-6b, trust_remote_codeTrue ) response, history lora_model.chat( tokenizer, 用一句话解释什么是机器学习, history[] ) print(response)逻辑说明PeftModel.from_pretrained加载 LoRA 权重并合并到基座模型。model.chat是 ChatGLM 特有的对话接口返回回复和历史。参数说明checkpoint-500是训练 500 步后的检查点如果回复质量明显优于基座模型说明微调有效如果回复重复或乱码说明学习率太大或数据有问题。我踩过最深的坑是忘了在推理时把模型设为eval()模式导致 dropout 仍然生效输出每次都不一样。这个 bug 找了两个小时最后发现是一行代码的事。多卡微调 ChatGLM 的整个流程从环境配置到训练启动再到效果验证每一步都有细节但跑通一次之后就有了可复用的模板。希望帮到你。本文还有配套的精品资源点击获取
返回列表