Atom-7B-Chat 全量微调实战:基于 DeepSpeed ZeRO-2 的 4 卡分布式训练指南(Datawhale self-llm)

Atom-7B-Chat 全量微调实战:基于 DeepSpeed ZeRO-2 的 4 卡分布式训练指南(Datawhale self-llm) 大模型人工智能教程本地部署微调【免费下载链接】self-llm《开源大模型食用指南》针对中国宝宝量身打造的基于Linux环境快速微调全参数/Lora、部署国内外开源大模型LLM/多模态大模型MLLM教程项目地址https://gitcode.com/datawhalechina/self-llm点击查看免费下载本篇技术指南以 Datawhale self-llm 仓库中 models/Atom/04-Atom-7B-chat 全量微调.md 为核心脉络系统讲解如何基于 transformersTrainer与 DeepSpeed 分布式训练框架在 4 张 24G 显存 GPU 上对 Atom-7B-Chat 模型进行全量微调Full Fine-tuning。读者将完整掌握训练代码的组织方式、DeepSpeed 环境从零搭建含DS_BUILD_OPS1源码编译安装、ZeRO-2 配置文件的逐项含义以及一键启动训练与结果验证的完整流程。一、全量微调与 Lora 微调同一套 Trainer 的两种范式Atom-7B-Chat 是 Llama2 中文社区与原子回声联合推出的中文大模型本仓库对其提供了完整的微调教程族包括 02-Atom-7B-Chat Lora 微调.md高效微调与本文全量微调。两者在代码骨架上是高度一致的都基于 HuggingFace transformers 的Trainer类完成训练循环都通过HfArgumentParser解析命令行超参数都使用DataCollatorForSeq2Seq或自定义采样函数组织 batch。二者的本质区别在于参数更新范围Lora 微调通过peft库的LoraConfig与get_peft_model冻结绝大部分基座权重仅更新注入的低秩适配矩阵。其显存占用小仓库教程提示单卡 32G 即可实际 24G 亦可但需要额外管理 adapter 参数的加载与合并全量微调不引入任何参数高效微调组件全部 7B 参数参与梯度更新模型权重、优化器状态与梯度全部驻留显存因此必须借助 DeepSpeed 的 ZeRO 优化与 CPU offload 才能将显存需求压到 4×24G 可承受的范围内。从代码视角看全量微调的train.py与 Lora 版本相比最显著的变化就是没有加载LoraConfig模型直接用AutoModelForCausalLM.from_pretrained加载原始权重Trainer直接接收该模型即可训练其余的数据处理、参数解析与训练流程几乎完全复用。因此理解 Lora 版本的 02-Atom-7B-Chat-Lora/train.py 有助于更快读懂全量微调代码。二、环境准备模型下载与 DeepSpeed 编译安装2.1 下载基座模型全量微调直接使用 ModelScope 上的FlagAlpha/Atom-7B-Chat模型。与 01-Atom-7B-chat-WebDemo.md 中介绍的方式一致使用modelscope的snapshot_download函数即可下载模型大小约 13GBfrom modelscope import snapshot_download model_dir snapshot_download(FlagAlpha/Atom-7B-Chat, cache_dir/root/autodl-tmp, revisionmaster)下载完成后将训练代码中的model_path参数指向本地模型目录下文FinetuneArguments.model_path的默认值即模型路径。2.2 用 environment.yml 创建干净的 conda 环境全量微调依赖 DeepSpeed 的编译期扩展CPUAdam 等强烈建议新建一个干净独立的 conda 环境并使用仓库提供的environment.yml文件创建而不是在已有环境中东拼西凑否则很容易出现算子和 CUDA 版本不匹配的问题conda env create -n deepspeed -f environment.yml --force conda activate deepspeed仓库根目录下提供的环境文件位于 models/Qwen/environment.yml其内容锁定了与编译 DeepSpeed 算子高度相关的关键依赖channels: - nvidia/label/cuda-11.8.0 - pytorch - conda-forge dependencies: - pytorch - torchvision - torchaudio - cuda - pytorch-cuda11.8 - compilers - sysroot_linux-642.17 - gcc11.4 - ninja - py-cpuinfo - libaio - ca-certificates - certifi - openssl - python3.10 - pydantic从中可以看出搭建 DeepSpeed 训练环境的几个关键点compilersgcc11.4ninjaDeepSpeed 的融合算子如 CPUAdam需要在本机即时编译C 编译器与构建工具是硬性要求libaioDeepSpeed ZeRO-Infinity / offload 机制使用异步 I/O 在 CPU 内存与磁盘间搬运参数与优化器状态libaio是其底层依赖pytorch-cuda11.8统一 PyTorch 与 CUDA 运行时版本避免算子编译时找到多个 CUDA 版本而失败python3.10与后续安装的 transformers、deepspeed 版本保持兼容。2.3 安装 DeepSpeed务必使用 DS_BUILD_OPS1DeepSpeed 安装非常简单但必须使用DS_BUILD_OPS1强制本地编译算子否则后续运行会报出一系列令人困惑的错误DS_BUILD_OPS1 pip install deepspeedDS_BUILD_OPS1会触发 DeepSpeed 的 JIT 编译流程在安装阶段即完成 CPUAdam、FusedAdam 等 CUDA 扩展的编译避免运行时才发现算子缺失。这也是原文档中反复强调一定要按照步骤安装否则会出现问题的根本原因——默认的预编译安装不会包含与本地 CUDA/PyTorch 严格匹配的二进制扩展。2.4 安装其余依赖环境创建时 PyTorch 已随 conda 一并装好不需要再手动安装 torch。随后按固定版本安装其余依赖pip install modelscope1.9.5 pip install transformers4.35.2 pip install streamlit1.24.0 pip install sentencepiece0.1.99 pip install accelerate0.24.1 pip install datasets sentencepiece pip install tiktoken pip install transformers_stream_generator各依赖的角色说明transformers4.35.2提供Trainer、TrainingArguments、HfArgumentParser等核心训练组件modelscope1.9.5下载模型与加载分词器sentencepieceAtom 中文词表所依赖的分词库accelerate0.24.1transformers 底层依赖的分布式后端需与 transformers 版本匹配datasets本教程使用Dataset.from_pandas加载 JSON 数据tiktoken、transformers_stream_generator部分推理链路与文本生成工具的依赖。三、训练代码逐段解析全量微调代码以Trainer为核心结构清晰参数解析 → 数据加载与预处理 → 模型加载 → 训练与验证。以下按功能模块拆解。3.1 参数定义FinetuneArguments通过dataclass定义微调自定义参数其中model_path需要修改为你自己的模型下载地址from dataclasses import dataclass, field dataclass class FinetuneArguments: # 微调参数 # fielddataclass 函数用于指定变量初始化 model_path: str field(default../../model/FlagAlpha/Atom-7B-Chat/)这里仅扩展了 transformersTrainingArguments之外的参数模型路径其余训练超参数batch size、学习率、epoch 等全部由TrainingArguments在命令行解析时统一接收。3.2 数据预处理process_func微调数据采用经典的instruction / input / output三段式结构见仓库根目录 dataset/huanhuan.json即甄嬛对话语料。预处理函数将每条样本按 Atom 的对话模板拼接并编码为input_ids / attention_mask / labels# 用于处理数据集的函数 def process_func(example): MAX_LENGTH 128 # Llama分词器会将一个中文字切分为多个token因此需要放开一些最大长度保证数据的完整性 input_ids, attention_mask, labels [], [], [] instruction tokenizer(\n.join([sHuman:, 现在你要扮演皇帝身边的女人--甄嬛请以甄嬛口吻回答用户问题 example[instruction] example[input] /s\n]).strip(), add_special_tokensFalse) # add_special_tokens 不在开头加 special_tokens response tokenizer(sAssistant: example[output] \s\n, add_special_tokensFalse) input_ids instruction[input_ids] response[input_ids] attention_mask instruction[attention_mask] response[attention_mask] labels [-100] * len(instruction[input_ids]) response[input_ids] if len(input_ids) MAX_LENGTH: # 做一个截断 input_ids input_ids[:MAX_LENGTH] attention_mask attention_mask[:MAX_LENGTH] labels labels[:MAX_LENGTH] return { input_ids: input_ids, attention_mask: attention_mask, labels: labels }几个关键技术点的含义对话模板sHuman: ... /s与sAssistant: ...是 AtomLlama2 系的标准 SFT 拼接格式与 Lora 版本中 02-Atom-7B-Chat-Lora/train.py 的preprocess函数使用的模板完全一致保证训练与推理阶段格式统一labels 遮蔽labels中指令部分的 token 置为-100这是 transformers 计算交叉熵损失时约定忽略的位置从而只让模型学习回答部分Assistant 内容这是指令微调保证学会回答而非复读提问的关键机制add_special_tokensFalse不额外追加特殊 token避免与模板中显式书写的s、/s重复MAX_LENGTH128中文按字切分后 token 数膨胀较快注释中特别提醒要放开最大长度以保证数据完整性超长样本直接截断。3.3 数据加载与 tokenizer 配置主流程中通过 pandas 读取 JSON 并转为 HuggingFace Dataset同时完成 tokenizer 加载与 pad token 设置if __main__ __name__: # 解析参数 # Parse 命令行参数 finetune_args, training_args HfArgumentParser( (FinetuneArguments, TrainingArguments) ).parse_args_into_dataclasses() # 处理数据集 # 将JSON文件转换为CSV文件 df pd.read_json(./data/huanhuan.json) ds Dataset.from_pandas(df) # 加载tokenizer tokenizer AutoTokenizer.from_pretrained(finetune_args.model_path, use_fastFalse, trust_remote_codeTrue) tokenizer.pad_token tokenizer.eos_token # 将数据集变化为token形式 tokenized_id ds.map(process_func, remove_columnsds.column_names)pd.read_json(./data/huanhuan.json)读取 JSON 格式的甄嬛训练语料路径可按实际调整仓库根目录提供 dataset/huanhuan.json 与 dataset/huanhuan.jsonl 两种格式字段均为instruction / input / outputtokenizer.pad_token tokenizer.eos_tokenLlama 系词表未定义 pad token训练时必须将 pad 复用为 eos否则DataCollatorForSeq2Seq(paddingTrue)在补齐 batch 时会报错ds.map(process_func, remove_columnsds.column_names)逐样本执行预处理并移除原始文本列仅保留input_ids / attention_mask / labels。3.4 模型加载半精度 device_map# 创建模型并以半精度形式加载 model AutoModelForCausalLM.from_pretrained(finetune_args.model_path, trust_remote_codeTrue, torch_dtypetorch.half, device_map{: int(os.environ.get(LOCAL_RANK) or 0)})torch_dtypetorch.half以 FP16 半精度加载权重配合 DeepSpeed 的fp16混合精度训练可将显存占用压到接近一半device_map{: int(os.environ.get(LOCAL_RANK) or 0)}将模型放置到当前进程对应的 GPU 卡上。LOCAL_RANK由 DeepSpeed 启动器注入这是多卡分布式训练中每进程只用本卡的标准写法——训练脚本本身与单卡无异分布式切分完全由 DeepSpeed 启动器与 ZeRO 引擎接管。3.5 Trainer 训练与推理验证# 使用trainer训练 trainer Trainer( modelmodel, argstraining_args, train_datasettokenized_id, data_collatorDataCollatorForSeq2Seq(tokenizertokenizer, paddingTrue), ) trainer.train() # 开始训练 response, history model.chat(tokenizer, 你是谁, history[], system现在你要扮演皇帝身边的女人--甄嬛.) print(response)DataCollatorForSeq2Seq(tokenizertokenizer, paddingTrue)在 batch 维度动态补齐到 batch 内最长序列Lora 版本 02-Atom-7B-Chat-Lora/train.py 中的data_collator函数实现的是同样逻辑的手工版补齐位置使用 pad token训练完成后直接调用model.chat(...)进行训练效果验证以你是谁提问并给定甄嬛人设 system 提示观察微调后的模型是否已习得甄嬛口吻。四、DeepSpeed 配置文件 ds_config.json 详解全量微调的核心难点在于显存因此必须借助 DeepSpeed 的 ZeRO-2ZeRO Stage 2 CPU offload 方案。原文档给出了可直接复制的ds_config.json其每个字段的含义如下{ fp16: { enabled: auto, loss_scale: 0, loss_scale_window: 1000, initial_scale_power: 16, hysteresis: 2, min_loss_scale: 1 }, optimizer: { type: AdamW, params: { lr: auto, betas: auto, eps: auto, weight_decay: auto } }, scheduler: { type: WarmupDecayLR, params: { last_batch_iteration: -1, total_num_steps: auto, warmup_min_lr: auto, warmup_max_lr: auto, warmup_num_steps: auto } }, zero_optimization: { stage: 2, offload_optimizer: { device: cpu, pin_memory: true }, offload_param: { device: cpu, pin_memory: true }, allgather_partitions: true, allgather_bucket_size: 5e8, overlap_comm: true, reduce_scatter: true, reduce_bucket_size: 5e8, contiguous_gradients: true }, activation_checkpointing: { partition_activations: false, cpu_checkpointing: false, contiguous_memory_optimization: false, number_checkpoints: null, synchronize_checkpoint_boundary: false, profile: false }, gradient_accumulation_steps: auto, gradient_clipping: auto, steps_per_print: 2000, train_batch_size: auto, min_lr: 5e-7, train_micro_batch_size_per_gpu: auto, wall_clock_breakdown: false }关键配置项解读fp16启用混合精度训练。loss_scale: 0表示使用动态 loss scalinginitial_scale_power: 16设定初始缩放因子为 2¹⁶min_loss_scale: 1是 loss scale 的下限。auto表示这些值由 DeepSpeed 从TrainingArguments自动继承保证与 transformers 配置一致optimizer.type: AdamW优化器由 DeepSpeed 管理lr / betas / eps / weight_decay均设为auto最终取值来自命令行传入的--learning_rate等训练参数避免两处配置漂移scheduler.type: WarmupDecayLRwarmup 线性衰减的学习率调度warmup_num_steps等参数同样由TrainingArguments自动推导zero_optimization.stage: 2ZeRO Stage 2——将优化器状态按数据并行度分片到各卡配合下面两项 offload 进一步压缩每卡显存offload_optimizer.device: cpu把分片后的优化器状态卸载到 CPU 内存这是4 卡 24G 能跑 7B 全量微调的关键所在原文档也特别指出脚本通过adam_cpu加载优化器参数使显存需求显著降低offload_param.device: cpu将模型参数也卸载到 CPU牺牲部分通信带宽换取更低的显存峰值allgather_partitions / reduce_scatter / overlap_comm / contiguous_gradientsZeRO 通信优化开关将梯度归约reduce-scatter与参数聚合all-gather与计算重叠隐藏通信延迟allgather_bucket_size / reduce_bucket_size: 5e8通信桶大小约 500MB决定一次通信聚合的梯度/参数规模过小增加通信次数过大会占用更多临时显存train_batch_size / train_micro_batch_size_per_gpu / gradient_accumulation_steps三者关系为train_batch_size micro_batch × gradient_accumulation_steps × 卡数全部设为auto后由 DeepSpeed 依据TrainingArguments自动计算避免手算出错min_lr: 5e-7调度器允许的最小学习率下限steps_per_print: 2000每 2000 步打印一次训练统计信息避免日志刷屏activation_checkpointing本配置中未启用partition_activations: false等意味着不牺牲重计算开销换取显存如需进一步压显存可开启但会拖慢训练速度。五、训练脚本 train.sh 与一键启动将以上代码保存为train.py把ds_config.json放在同级目录再创建train.shnum_gpus4 deepspeed --num_gpus $num_gpus train.py \ --deepspeed ./ds_config.json \ --output_dir./output/Atom \ --per_device_train_batch_size1 \ --gradient_accumulation_steps1 \ --logging_steps10 \ --num_train_epochs3 \ --save_steps100 \ --learning_rate1e-4 \ --save_on_each_nodeTrue \命令行参数说明deepspeed --num_gpus 4由 DeepSpeed 启动器拉起 4 个训练进程对应 4 张 GPU并自动设置好每个进程的LOCAL_RANK环境变量这正是 3.4 节device_map中读取该变量的来源--deepspeed ./ds_config.json指定 ZeRO 配置文件--output_dir./output/Atomcheckpoint 与最终权重输出目录--per_device_train_batch_size1每卡 batch size 为 1。全量微调显存紧张微小的 batch 是常规选择梯度累积由配置文件中gradient_accumulation_steps统一接管此处为 1--num_train_epochs3完整遍历训练集 3 个 epoch--learning_rate1e-4初始学习率与 DeepSpeed 配置中的lr: auto联动--logging_steps10每 10 步打印一次 loss 日志--save_steps100每 100 步保存一次 checkpoint--save_on_each_nodeTrue多节点多机训练时每个节点都保存 checkpoint避免单点保存带来的 I/O 瓶颈。在命令行输入bash train.sh即可开始训练。启动后终端会打印 DeepSpeed 的引擎初始化日志ZeRO stage、offload 设备、通信后端等随后进入正常的训练循环每logging_steps步输出一次 loss 值。六、显存需求与常见问题6.1 显存需求由于配置了adam_cpu优化器状态 offload 到 CPU全量微调的显存占用被大幅压低但仍至少需要 4 张 24G 显存的 GPU如 RTX 3090 / 4090 级别以满足模型权重、梯度分片与激活值的基本驻留需求若单卡显存更小或想要更宽的 batch可考虑开启activation_checkpointing或调低allgather_bucket_size / reduce_bucket_size但需在显存与训练速度之间做取舍。6.2 常见报错DS_BUILD_OPS 未生效如果创建 DeepSpeed 环境时没有使用DS_BUILD_OPS1运行时会大概率出现RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!根因是 CPUAdam 等融合算子未被正确编译/加载导致优化器状态的 offload 路径失效张量散落在 CPU 与 GPU 之间。解决办法是删除当前环境重新按 2.2 与 2.3 节的步骤创建全新环境并重装 DeepSpeed而不是在现有环境上打补丁——这正是原文档反复强调创建干净环境 DS_BUILD_OPS1的原因。如安装或运行过程中遇到其他问题可以在仓库提交 Issue 描述复现路径解决后亦可提交 PR 回馈项目。七、总结全量微调与 Lora 微调共享同一套 transformersTrainer骨架差别仅在于是否注入LoraConfig。当选择全量微调时显存管理成为核心矛盾本文所展示的方案通过三层手段解决torch_dtypetorch.half半精度加载 DeepSpeed FP16 混合精度训练压缩权重与激活开销ZeRO Stage 2 将优化器状态分片到 4 卡offload_optimizer / offload_param将优化器状态与参数下沉到 CPU 内存。三者叠加使 7B 模型的全量微调在 4×24G 显存上得以落地。配套的 ds_config.json 配置、训练脚本 与 训练语料 均可直接在仓库中查阅复现是理解中等规模开源模型全量微调工程化的完整范本。赞分享大模型人工智能教程本地部署微调【免费下载链接】self-llm《开源大模型食用指南》针对中国宝宝量身打造的基于Linux环境快速微调全参数/Lora、部署国内外开源大模型LLM/多模态大模型MLLM教程项目地址https://gitcode.com/datawhalechina/self-llm点击查看免费下载相关推荐self-llm 实战基于 DeepSpeed ZeRO-2 的 Atom-7B-Chat 全量微调——从环境搭建到多卡训练完整指南self llm 实战基于 DeepSpeed ZeRO 2 的 Atom 7B Chat 全量微调——从环境搭建到多卡训练完整指南 本文以开源大模型食用指南教程大模型本地部署微调Datawhale self-llm 实战基于 transformers 与 peft 的 Atom-7B-Chat LoRA 指令微调全流程Datawhale self llm 实战基于 transformers 与 peft 的 Atom 7B Chat LoRA 指令微调全流程 本文是《开源大大模型人工智能教程本地部署微调基于 DeepSpeed 的 Qwen-7B-Chat 全量微调实战从环境配置到多卡训练基于 DeepSpeed 的 Qwen 7B Chat 全量微调实战从环境配置到多卡训练 导读 本文基于开源仓库《开源大模型食用指南》中的 06 Qwen 7教程大模型本地部署微调创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考