Yi-6B / Yi-34B 微调实战指南基于 DeepSpeed 的 SFT 与 LoRA 完整教程【免费下载链接】YiA series of large language models trained from scratch by developers 01-ai项目地址: https://gitcode.com/GitHub_Trending/yi/Yi导读本文以开源仓库 finetune/README.md 为核心系统讲解 Yi-6B 与 Yi-34B 基础模型的监督微调SFT全流程从 Docker 镜像与本地 conda 环境搭建、硬件选型到数据准备、脚本运行、核心训练参数逐项解析再到基于 DeepSpeed ZeRO-Offload 的显存优化原理、LoRA 轻量微调以及微调前后的对比评估。读完本文你将能够在 4~8 卡 GPU 节点上独立完成 Yi 系列模型的指令微调并掌握每个关键参数背后的源码级原理。一、环境准备两种可行的起步方式1.1 从 Docker 镜像起步推荐快速验证仓库提供了封装好训练环境的 Docker 镜像ghcr.io/01-ai/yi:latest。默认情况下镜像内使用一个小型数据集源自 BAAI/COIG对基础模型进行微调便于快速跑通全流程。如果你希望使用自己的数据需要先准备自定义数据集再通过-v参数挂载进容器以替换默认数据。数据集采用逐行jsonl格式每一行是一个样本对象包含prompt与chosen两个字段{ prompt: Human: Who are you? Assistant:, chosen: Im Yi. }其中prompt是携带对话模板前缀的输入chosen是对应的期望回复。挂载并启动训练的命令如下docker run -it \ -v /path/to/save/finetuned/model/:/finetuned-model \ -v /path/to/train.jsonl:/yi/finetune/data/train.json \ -v /path/to/eval.jsonl:/yi/finetune/data/eval.json \ ghcr.io/01-ai/yi:latest \ bash finetune/scripts/run_sft_Yi_6b.sh三个挂载点的职责分别是模型输出目录容器内/finetuned-model、训练集容器内/yi/finetune/data/train.json、评估集容器内/yi/finetune/data/eval.json。容器启动后会自动执行 run_sft_Yi_6b.sh 完成训练。1.2 从本地服务器起步conda 环境如果希望在自有服务器上运行推荐使用 conda 管理环境。未安装 conda 时可先安装 Minicondamkdir -p ~/miniconda3 wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3 rm -rf ~/miniconda3/miniconda.sh ~/miniconda3/bin/conda init bash source ~/.bashrc随后创建 Python 3.10 环境并安装依赖。仓库 README 指定的依赖版本组合为torch2.0.1、deepspeed0.10、tensorboard、transformers、datasets、sentencepiece、accelerate、ray2.7conda create -n dev_env python3.10 -y conda activate dev_env pip install torch2.0.1 deepspeed0.10 tensorboard transformers datasets sentencepiece accelerate ray2.7说明依赖版本以仓库文档为准deepspeed是训练主程序deepspeed main.py的运行时sentencepiece用于加载 Yi 的 tokenizerdatasets负责读取 huggingface 数据集。仓库根目录另提供 requirements.txt 与 conda-lock.yml可供整体环境参考。二、硬件配置要求务必先评估资源微调是大显存、大内存场景README 给出了明确的硬件基线请按模型规模对照模型推荐节点配置关键约束Yi-6B4 卡节点单卡显存 60GB按默认脚本可直接运行Yi-34B8 卡节点仅使用其中 4 卡单卡显存 80GB整节点 CPU 内存 900GB必须用CUDA_VISIBLE_DEVICES限制 GPU 数量为什么 Yi-34B 要限制 GPU 数量原因在于 34B 训练开启了ZeRO-Offloadzero-offload技术优化器状态和参数会卸载到 CPU 内存GPU 数量越多被拆分的显存需求越小但每个进程对应的 CPU 侧状态也越多反而可能触发 CPU 内存溢出OOM。README 明确提示 please be careful to limit the GPU numbers in 34B finetune training。因此 run_sft_Yi_34b.sh 脚本第一行就设置了export CUDA_VISIBLE_DEVICES0,1,2,3 #limit parallelism to avoid cpu oom即物理 8 卡节点只暴露 4 卡给训练进程配合 900GB 以上 CPU 内存才能稳定容纳 34B 参数的 offload 开销。三、模型与数据准备3.1 下载基础模型将 Yi 基础模型6B 或 34B下载到本地MODEL_PATH。README 给出的典型模型目录结构如下|-- $MODEL_PATH | |-- config.json | |-- pytorch_model-00001-of-00002.bin | |-- pytorch_model-00002-of-00002.bin | |-- pytorch_model.bin.index.json | |-- tokenizer_config.json | |-- tokenizer.model | |-- ...注意这是标准 Hugging Face 权重格式含分片权重pytorch_model-*.bin、索引文件与 tokenizer 文件。训练脚本通过--model_name_or_path传入该路径源码 main.py 将其声明为必填参数支持本地路径或 Hugging Face 模型标识符。3.2 下载训练数据集将数据集下载到本地DATA_PATH。README 以Dahoas/rm-static为例从 Hugging Face 获取其目录结构为 parquet 分片格式|-- $DATA_PATH | |-- data | | |-- train-00000-of-00001-2a1df75c6bce91ab.parquet | | |-- test-00000-of-00001-8c7c51afc6d45980.parquet | |-- dataset_infos.json | |-- README.md3.3 仓库自带示例数据集仓库在 finetune/yi_example_dataset 目录下提供了开箱即用的示例数据集修改自 BAAI/COIG其结构为|-- $DATA_PATH |--data |-- train.jsonl |-- eval.jsonl以 train.jsonl 为例每行形如{prompt: Human: 你正在准备开会但是发现会议室的白板上还有别人的会议记录你想给对方一个提醒让他们及时清理。请给出一个礼貌的提醒。 Assistant:, chosen: “不好意思打扰了请问这个白板是否还需要保留别的笔记……感谢你的配合。”}数据格式与解析的对应关系源码 data_utils.py 中get_raw_dataset按数据集名称分派解析器——包含yi关键字的数据路径走YiDataset会检查data_path/data/train.jsonl与data_path/data/eval.jsonl是否存在包含rm-static的路径走DahoasRmstaticDataset包含local/jsonfile的路径走LocalJsonFileDataset。因此如果自定义数据集建议将数据目录名设为含yi的字符串如finetune/yi_example_dataset/脚本即可正确识别。四、快速开始运行 SFT 训练脚本进入脚本目录修改脚本中的MODEL_PATH与DATA_PATH对应--model_name_or_path与--data_path参数后直接运行cd finetune/scripts bash run_sft_Yi_6b.sh完整脚本内容如下run_sft_Yi_6b.shcd $(dirname ${BASH_SOURCE[0]})/../sft/ deepspeed main.py \ --data_path ../yi_example_dataset/ \ --model_name_or_path /base_model \ --per_device_train_batch_size 1 \ --per_device_eval_batch_size 1 \ --max_seq_len 4096 \ --learning_rate 2e-6 \ --weight_decay 0. \ --num_train_epochs 4 \ --training_debug_steps 20 \ --gradient_accumulation_steps 1 \ --lr_scheduler_type cosine \ --num_warmup_steps 0 \ --seed 1234 \ --gradient_checkpointing \ --zero_stage 2 \ --deepspeed \ --offload \ --output_dir ./finetuned_model关键预期对于 Yi-6B 基础模型保持training_debug_steps20与num_train_epochs4的组合即可产出一个具备对话能力的 chat 模型端到端耗时约20 分钟。其中training_debug_steps20的作用是每个 epoch 只跑 20 个训练步就提前截断见 main.py 的if step args.training_debug_steps: break用于快速验证流程。而 Yi-34B 由于参数量巨大初始化阶段加载权重、构建 ZeRO 状态耗时明显更长README 特别提醒请保持耐心。34B 脚本 run_sft_Yi_34b.sh 与 6B 脚本结构一致仅多出export CUDA_VISIBLE_DEVICES0,1,2,3限制并将training_debug_steps调为 50、输出目录设为/finetuned_model。五、核心训练参数逐项解析结合源码所有参数均在 main.py 的parse_args()中定义。下表按功能分组列出关键参数的含义、默认值与脚本实际取值参数默认值6B 脚本取值含义--data_pathDahoas/rm-static../yi_example_dataset/训练数据路径支持空格分隔的多个数据集--model_name_or_path必填/base_model预训练模型路径或 HF 模型标识符--per_device_train_batch_size161每卡训练 batch size--per_device_eval_batch_size161每卡评估 batch size--max_seq_len5124096最大序列长度需与 Yi 的 4K 上下文对齐--learning_rate1e-32e-6初始学习率SFT 全参场景建议 1e-6~2e-6--weight_decay0.00.权重衰减--num_train_epochs14训练轮数--training_debug_steps1020每 epoch 的调试截断步数--gradient_accumulation_steps11梯度累积步数--lr_scheduler_typecosinecosine调度器类型linear / cosine / cosine_with_restarts / polynomial / constant / constant_with_warmup--num_warmup_steps00学习率预热步数--seed12341234随机种子--gradient_checkpointingFalse开启梯度检查点以计算换显存--zero_stage02DeepSpeed ZeRO 优化阶段--offloadFalse开启启用 ZeRO-Offload参数/优化器状态卸载到 CPU--output_dirNone./finetuned_model模型保存目录--disable_dropoutFalse—关闭模型 dropout--print_lossFalse—每步打印 loss--enable_tensorboardFalse—开启 TensorBoard 日志--tensorboard_pathsft_tensorboard—TensorBoard 日志目录几个参数的源码级说明学习率与优化器训练使用 Adam 优化器betas(0.9, 0.95)开启--offload时自动选用DeepSpeedCPUAdam否则用FusedAdam见 main.py。SFT 阶段使用 2e-6 的低学习率是为了在保留基础模型预训练知识的前提下温和地适配指令格式。全局 batch size 的推导在 main.py 中DeepSpeed 配置的train_batch_size由per_device_train_batch_size × 卡数 × gradient_accumulation_steps动态计算因此无需手工指定。tokenizer 与词表对齐模型加载后会将pad_token_id设为eos_token_id并将词表扩充到 8 的倍数resize_token_embeddings这是为了匹配 DeepSpeed 对词表对齐的要求见 model_utils.py。评估穿插训练启动前以及每个 epoch 结束后都会在 eval 集上计算 loss 与困惑度perplexity见 main.py。模型保存训练结束后调用save_hf_format保存为标准 Hugging Face 格式若使用 ZeRO-3则需走专门的save_zero_three_model聚合分片权重见 main.py。六、DeepSpeed 配置与 ZeRO-Offload 原理训练脚本中的--zero_stage 2 --offload --deepspeed --gradient_checkpointing四个开关共同决定了显存优化的核心策略。其底层配置由 ds_utils.py 的get_train_ds_config()生成关键内容如下device cpu if offload else none zero_opt_dict { stage: stage, # 由 --zero_stage 传入脚本取 2 offload_param: {device: device}, offload_optimizer: {device: device}, stage3_param_persistence_threshold: 1e4, stage3_max_live_parameters: 3e7, stage3_prefetch_bucket_size: 3e7, ... } return { train_batch_size: ..., train_micro_batch_size_per_gpu: ..., steps_per_print: 10, zero_optimization: zero_opt_dict, fp16: {enabled: True, loss_scale_window: 100}, gradient_clipping: 1.0, ... }可以总结出这套方案的几个特征ZeRO Stage 2优化器状态与梯度按数据并行维度分片配合 4 卡即可显著降低单卡峰值显存。ZeRO-Offload开启后offload_param与offload_optimizer的设备均设为cpu将参数副本与优化器状态搬运到 CPU 内存GPU 只保留计算必需的部分——这正是 34B 微调需要 900GB CPU 内存的根因也是 README 警告 CPU OOM 的由来。混合精度训练fp16.enabledTrue启用半精度训练配合loss_scale_window100动态损失缩放。梯度裁剪gradient_clipping1.0防止梯度爆炸。梯度检查点--gradient_checkpointing在 main.py 中通过model.gradient_checkpointing_enable()生效用重计算换取显存是 4K 序列长度下跑通 6B 模型的关键之一。从源码结构看get_eval_ds_config还提供了评估模式的独立配置ZeRO Stage 0见 ds_utils.py说明同一套工具函数可复用于不同训练/推理阶段。七、LoRA 轻量微调6B 资源受限场景除全参数微调外仓库还提供基于 LoRA 的高效微调脚本 run_sft_lora_Yi_6b.sh适合显存与算力更受限的场景deepspeed main.py \ --data_path /DATA_PATH/ \ --model_name_or_path /MODEL_PATH/ \ --per_device_train_batch_size 4 \ --per_device_eval_batch_size 4 \ --max_seq_len 4096 \ --learning_rate 2e-6 \ --weight_decay 0. \ --num_train_epochs 4 \ --training_debug_steps 20 \ --gradient_accumulation_steps 1 \ --lr_scheduler_type cosine \ --num_warmup_steps 0 \ --seed 1234 \ --gradient_checkpointing \ --zero_stage 2 \ --deepspeed \ --offload \ --lora_dim 128 \ --lora_module_name layers. \ --output_dir ./output_Yi_6b_chat_sft_lora对比全参脚本LoRA 版本的变化在于--lora_dim 128设置 LoRA 低秩矩阵的秩为 128。源码 main.py 中该参数默认 0表示不使用 LoRA只要大于 0 就会触发 LoRA 转换batch size 也因此可以从 1 提升到 4。--lora_module_name layers.指定要注入 LoRA 的模块范围。相关实现位于 lora.py包含convert_linear_layer_to_lora、only_optimize_lora_parameters、make_model_gradient_checkpointing_compatible、convert_lora_to_linear_layer等函数。训练流程中的 LoRA 处理见 main.py先转换线性层为 LoRA若同时传入--only_optimize_lora则冻结其余参数、只优化 LoRA 权重LoRA 专属学习率--lora_learning_rate默认 5e-4保存模型前统一调用convert_lora_to_linear_layer把 LoRA 权重合并回原模型输出仍是标准 HF 格式。由于只训练低秩增量LoRA 方案的单卡 batch size 与整体吞吐都优于全参模式是快速实验指令微调效果的高性价比路径。八、评估微调效果base 与 finetuned 对比训练完成后用评估脚本对比基础模型与微调模型的回答质量cd finetune/scripts bash run_eval.sh脚本内容run_eval.sh会调用 prompt_eval.pycd $(dirname ${BASH_SOURCE[0]})/../sft/ python prompt_eval.py \ --model_name_or_path_base/base_model \ --model_name_or_path_finetune/finetuned_model \ --language Chinese运行后脚本会分别加载 base 模型与微调模型对同一组提示词进行贪心解码生成并分栏打印Baseline: Greedy 基础模型输出 finetune: Greedy 微调模型输出通过肉眼对比即可直观评估微调效果。该评估器支持的可调参数包括--num_beams束搜索宽度默认 1、--num_beam_groups、--top_k默认 4、--penalty_alpha默认 0.6、--num_return_sequences、--max_new_tokens默认 200、--languageChinese或English以及--eos默认|endoftext|。内置提示词与停止符细节--language Chinese时使用提示词Human: 你好。 Assistant:与Human: 请介绍一下大语言模型? Assistant:英文时对应Human: Hello. Assistant:与Human: Please explain Large Language Model. Assistant:。生成时通过tokenizer.convert_tokens_to_ids(|endoftext|)显式设置停止符与 Yi 的对话模板保持一致见 prompt_eval.py。评估模型以device_mapauto、torch_dtypeauto加载见 model_utils.py即双模型同时驻留显存进行推理请预留足够显存。九、常见问题与注意事项34B 训练 CPU 内存不足务必保留 run_sft_Yi_34b.sh 中的CUDA_VISIBLE_DEVICES0,1,2,3限制并确认节点 CPU 内存不低于 900GB。自定义数据集无法加载检查数据目录名称是否包含yi关键字触发 data_utils.py 的YiDataset分支并确认data/train.jsonl与data/eval.jsonl两个文件同时存在。训练被提前截断是预期行为training_debug_steps是调试步数截断6B 场景的20 分钟出 chat 模型正是依赖该设置正式训练请移除该参数或调大数值。序列长度与显存脚本默认--max_seq_len 4096对应 Yi 的 4K 上下文配合--gradient_checkpointing使用显存紧张时可下调该值。初始化阶段耗时34B 权重加载与 ZeRO 状态构建耗时较长属正常现象耐心等待即可。权重保存格式ZeRO-2 场景直接保存合并后的 HF 格式权重到--output_dir若改用 ZeRO-3需依赖save_zero_three_model完成分片聚合勿直接拷贝 checkpoint 目录。十、仓库配套资源索引训练主程序finetune/sft/main.py参数解析、数据装载、训练循环、模型保存评估程序finetune/sft/prompt_eval.pybase 与 finetuned 对比生成训练脚本run_sft_Yi_6b.sh / run_sft_Yi_34b.sh / run_sft_lora_Yi_6b.sh评估脚本run_eval.shDeepSpeed 配置finetune/utils/ds_utils.py模型加载工具finetune/utils/model/model_utils.py数据集解析finetune/utils/data/data_utils.py 与 finetune/utils/data/raw_datasets.pyLoRA 实现finetune/utils/module/lora.py训练阶段常量SFT/RM/PPO/RSFT/DPOfinetune/constant.py示例数据集finetune/yi_example_dataset/data/train.jsonl 与 finetune/yi_example_dataset/data/eval.jsonl【免费下载链接】YiA series of large language models trained from scratch by developers 01-ai项目地址: https://gitcode.com/GitHub_Trending/yi/Yi创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Yi-6B / Yi-34B 微调实战指南：基于 DeepSpeed 的 SFT 与 LoRA 完整教程
Yi-6B / Yi-34B 微调实战指南基于 DeepSpeed 的 SFT 与 LoRA 完整教程【免费下载链接】YiA series of large language models trained from scratch by developers 01-ai项目地址: https://gitcode.com/GitHub_Trending/yi/Yi导读本文以开源仓库 finetune/README.md 为核心系统讲解 Yi-6B 与 Yi-34B 基础模型的监督微调SFT全流程从 Docker 镜像与本地 conda 环境搭建、硬件选型到数据准备、脚本运行、核心训练参数逐项解析再到基于 DeepSpeed ZeRO-Offload 的显存优化原理、LoRA 轻量微调以及微调前后的对比评估。读完本文你将能够在 4~8 卡 GPU 节点上独立完成 Yi 系列模型的指令微调并掌握每个关键参数背后的源码级原理。一、环境准备两种可行的起步方式1.1 从 Docker 镜像起步推荐快速验证仓库提供了封装好训练环境的 Docker 镜像ghcr.io/01-ai/yi:latest。默认情况下镜像内使用一个小型数据集源自 BAAI/COIG对基础模型进行微调便于快速跑通全流程。如果你希望使用自己的数据需要先准备自定义数据集再通过-v参数挂载进容器以替换默认数据。数据集采用逐行jsonl格式每一行是一个样本对象包含prompt与chosen两个字段{ prompt: Human: Who are you? Assistant:, chosen: Im Yi. }其中prompt是携带对话模板前缀的输入chosen是对应的期望回复。挂载并启动训练的命令如下docker run -it \ -v /path/to/save/finetuned/model/:/finetuned-model \ -v /path/to/train.jsonl:/yi/finetune/data/train.json \ -v /path/to/eval.jsonl:/yi/finetune/data/eval.json \ ghcr.io/01-ai/yi:latest \ bash finetune/scripts/run_sft_Yi_6b.sh三个挂载点的职责分别是模型输出目录容器内/finetuned-model、训练集容器内/yi/finetune/data/train.json、评估集容器内/yi/finetune/data/eval.json。容器启动后会自动执行 run_sft_Yi_6b.sh 完成训练。1.2 从本地服务器起步conda 环境如果希望在自有服务器上运行推荐使用 conda 管理环境。未安装 conda 时可先安装 Minicondamkdir -p ~/miniconda3 wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3 rm -rf ~/miniconda3/miniconda.sh ~/miniconda3/bin/conda init bash source ~/.bashrc随后创建 Python 3.10 环境并安装依赖。仓库 README 指定的依赖版本组合为torch2.0.1、deepspeed0.10、tensorboard、transformers、datasets、sentencepiece、accelerate、ray2.7conda create -n dev_env python3.10 -y conda activate dev_env pip install torch2.0.1 deepspeed0.10 tensorboard transformers datasets sentencepiece accelerate ray2.7说明依赖版本以仓库文档为准deepspeed是训练主程序deepspeed main.py的运行时sentencepiece用于加载 Yi 的 tokenizerdatasets负责读取 huggingface 数据集。仓库根目录另提供 requirements.txt 与 conda-lock.yml可供整体环境参考。二、硬件配置要求务必先评估资源微调是大显存、大内存场景README 给出了明确的硬件基线请按模型规模对照模型推荐节点配置关键约束Yi-6B4 卡节点单卡显存 60GB按默认脚本可直接运行Yi-34B8 卡节点仅使用其中 4 卡单卡显存 80GB整节点 CPU 内存 900GB必须用CUDA_VISIBLE_DEVICES限制 GPU 数量为什么 Yi-34B 要限制 GPU 数量原因在于 34B 训练开启了ZeRO-Offloadzero-offload技术优化器状态和参数会卸载到 CPU 内存GPU 数量越多被拆分的显存需求越小但每个进程对应的 CPU 侧状态也越多反而可能触发 CPU 内存溢出OOM。README 明确提示 please be careful to limit the GPU numbers in 34B finetune training。因此 run_sft_Yi_34b.sh 脚本第一行就设置了export CUDA_VISIBLE_DEVICES0,1,2,3 #limit parallelism to avoid cpu oom即物理 8 卡节点只暴露 4 卡给训练进程配合 900GB 以上 CPU 内存才能稳定容纳 34B 参数的 offload 开销。三、模型与数据准备3.1 下载基础模型将 Yi 基础模型6B 或 34B下载到本地MODEL_PATH。README 给出的典型模型目录结构如下|-- $MODEL_PATH | |-- config.json | |-- pytorch_model-00001-of-00002.bin | |-- pytorch_model-00002-of-00002.bin | |-- pytorch_model.bin.index.json | |-- tokenizer_config.json | |-- tokenizer.model | |-- ...注意这是标准 Hugging Face 权重格式含分片权重pytorch_model-*.bin、索引文件与 tokenizer 文件。训练脚本通过--model_name_or_path传入该路径源码 main.py 将其声明为必填参数支持本地路径或 Hugging Face 模型标识符。3.2 下载训练数据集将数据集下载到本地DATA_PATH。README 以Dahoas/rm-static为例从 Hugging Face 获取其目录结构为 parquet 分片格式|-- $DATA_PATH | |-- data | | |-- train-00000-of-00001-2a1df75c6bce91ab.parquet | | |-- test-00000-of-00001-8c7c51afc6d45980.parquet | |-- dataset_infos.json | |-- README.md3.3 仓库自带示例数据集仓库在 finetune/yi_example_dataset 目录下提供了开箱即用的示例数据集修改自 BAAI/COIG其结构为|-- $DATA_PATH |--data |-- train.jsonl |-- eval.jsonl以 train.jsonl 为例每行形如{prompt: Human: 你正在准备开会但是发现会议室的白板上还有别人的会议记录你想给对方一个提醒让他们及时清理。请给出一个礼貌的提醒。 Assistant:, chosen: “不好意思打扰了请问这个白板是否还需要保留别的笔记……感谢你的配合。”}数据格式与解析的对应关系源码 data_utils.py 中get_raw_dataset按数据集名称分派解析器——包含yi关键字的数据路径走YiDataset会检查data_path/data/train.jsonl与data_path/data/eval.jsonl是否存在包含rm-static的路径走DahoasRmstaticDataset包含local/jsonfile的路径走LocalJsonFileDataset。因此如果自定义数据集建议将数据目录名设为含yi的字符串如finetune/yi_example_dataset/脚本即可正确识别。四、快速开始运行 SFT 训练脚本进入脚本目录修改脚本中的MODEL_PATH与DATA_PATH对应--model_name_or_path与--data_path参数后直接运行cd finetune/scripts bash run_sft_Yi_6b.sh完整脚本内容如下run_sft_Yi_6b.shcd $(dirname ${BASH_SOURCE[0]})/../sft/ deepspeed main.py \ --data_path ../yi_example_dataset/ \ --model_name_or_path /base_model \ --per_device_train_batch_size 1 \ --per_device_eval_batch_size 1 \ --max_seq_len 4096 \ --learning_rate 2e-6 \ --weight_decay 0. \ --num_train_epochs 4 \ --training_debug_steps 20 \ --gradient_accumulation_steps 1 \ --lr_scheduler_type cosine \ --num_warmup_steps 0 \ --seed 1234 \ --gradient_checkpointing \ --zero_stage 2 \ --deepspeed \ --offload \ --output_dir ./finetuned_model关键预期对于 Yi-6B 基础模型保持training_debug_steps20与num_train_epochs4的组合即可产出一个具备对话能力的 chat 模型端到端耗时约20 分钟。其中training_debug_steps20的作用是每个 epoch 只跑 20 个训练步就提前截断见 main.py 的if step args.training_debug_steps: break用于快速验证流程。而 Yi-34B 由于参数量巨大初始化阶段加载权重、构建 ZeRO 状态耗时明显更长README 特别提醒请保持耐心。34B 脚本 run_sft_Yi_34b.sh 与 6B 脚本结构一致仅多出export CUDA_VISIBLE_DEVICES0,1,2,3限制并将training_debug_steps调为 50、输出目录设为/finetuned_model。五、核心训练参数逐项解析结合源码所有参数均在 main.py 的parse_args()中定义。下表按功能分组列出关键参数的含义、默认值与脚本实际取值参数默认值6B 脚本取值含义--data_pathDahoas/rm-static../yi_example_dataset/训练数据路径支持空格分隔的多个数据集--model_name_or_path必填/base_model预训练模型路径或 HF 模型标识符--per_device_train_batch_size161每卡训练 batch size--per_device_eval_batch_size161每卡评估 batch size--max_seq_len5124096最大序列长度需与 Yi 的 4K 上下文对齐--learning_rate1e-32e-6初始学习率SFT 全参场景建议 1e-6~2e-6--weight_decay0.00.权重衰减--num_train_epochs14训练轮数--training_debug_steps1020每 epoch 的调试截断步数--gradient_accumulation_steps11梯度累积步数--lr_scheduler_typecosinecosine调度器类型linear / cosine / cosine_with_restarts / polynomial / constant / constant_with_warmup--num_warmup_steps00学习率预热步数--seed12341234随机种子--gradient_checkpointingFalse开启梯度检查点以计算换显存--zero_stage02DeepSpeed ZeRO 优化阶段--offloadFalse开启启用 ZeRO-Offload参数/优化器状态卸载到 CPU--output_dirNone./finetuned_model模型保存目录--disable_dropoutFalse—关闭模型 dropout--print_lossFalse—每步打印 loss--enable_tensorboardFalse—开启 TensorBoard 日志--tensorboard_pathsft_tensorboard—TensorBoard 日志目录几个参数的源码级说明学习率与优化器训练使用 Adam 优化器betas(0.9, 0.95)开启--offload时自动选用DeepSpeedCPUAdam否则用FusedAdam见 main.py。SFT 阶段使用 2e-6 的低学习率是为了在保留基础模型预训练知识的前提下温和地适配指令格式。全局 batch size 的推导在 main.py 中DeepSpeed 配置的train_batch_size由per_device_train_batch_size × 卡数 × gradient_accumulation_steps动态计算因此无需手工指定。tokenizer 与词表对齐模型加载后会将pad_token_id设为eos_token_id并将词表扩充到 8 的倍数resize_token_embeddings这是为了匹配 DeepSpeed 对词表对齐的要求见 model_utils.py。评估穿插训练启动前以及每个 epoch 结束后都会在 eval 集上计算 loss 与困惑度perplexity见 main.py。模型保存训练结束后调用save_hf_format保存为标准 Hugging Face 格式若使用 ZeRO-3则需走专门的save_zero_three_model聚合分片权重见 main.py。六、DeepSpeed 配置与 ZeRO-Offload 原理训练脚本中的--zero_stage 2 --offload --deepspeed --gradient_checkpointing四个开关共同决定了显存优化的核心策略。其底层配置由 ds_utils.py 的get_train_ds_config()生成关键内容如下device cpu if offload else none zero_opt_dict { stage: stage, # 由 --zero_stage 传入脚本取 2 offload_param: {device: device}, offload_optimizer: {device: device}, stage3_param_persistence_threshold: 1e4, stage3_max_live_parameters: 3e7, stage3_prefetch_bucket_size: 3e7, ... } return { train_batch_size: ..., train_micro_batch_size_per_gpu: ..., steps_per_print: 10, zero_optimization: zero_opt_dict, fp16: {enabled: True, loss_scale_window: 100}, gradient_clipping: 1.0, ... }可以总结出这套方案的几个特征ZeRO Stage 2优化器状态与梯度按数据并行维度分片配合 4 卡即可显著降低单卡峰值显存。ZeRO-Offload开启后offload_param与offload_optimizer的设备均设为cpu将参数副本与优化器状态搬运到 CPU 内存GPU 只保留计算必需的部分——这正是 34B 微调需要 900GB CPU 内存的根因也是 README 警告 CPU OOM 的由来。混合精度训练fp16.enabledTrue启用半精度训练配合loss_scale_window100动态损失缩放。梯度裁剪gradient_clipping1.0防止梯度爆炸。梯度检查点--gradient_checkpointing在 main.py 中通过model.gradient_checkpointing_enable()生效用重计算换取显存是 4K 序列长度下跑通 6B 模型的关键之一。从源码结构看get_eval_ds_config还提供了评估模式的独立配置ZeRO Stage 0见 ds_utils.py说明同一套工具函数可复用于不同训练/推理阶段。七、LoRA 轻量微调6B 资源受限场景除全参数微调外仓库还提供基于 LoRA 的高效微调脚本 run_sft_lora_Yi_6b.sh适合显存与算力更受限的场景deepspeed main.py \ --data_path /DATA_PATH/ \ --model_name_or_path /MODEL_PATH/ \ --per_device_train_batch_size 4 \ --per_device_eval_batch_size 4 \ --max_seq_len 4096 \ --learning_rate 2e-6 \ --weight_decay 0. \ --num_train_epochs 4 \ --training_debug_steps 20 \ --gradient_accumulation_steps 1 \ --lr_scheduler_type cosine \ --num_warmup_steps 0 \ --seed 1234 \ --gradient_checkpointing \ --zero_stage 2 \ --deepspeed \ --offload \ --lora_dim 128 \ --lora_module_name layers. \ --output_dir ./output_Yi_6b_chat_sft_lora对比全参脚本LoRA 版本的变化在于--lora_dim 128设置 LoRA 低秩矩阵的秩为 128。源码 main.py 中该参数默认 0表示不使用 LoRA只要大于 0 就会触发 LoRA 转换batch size 也因此可以从 1 提升到 4。--lora_module_name layers.指定要注入 LoRA 的模块范围。相关实现位于 lora.py包含convert_linear_layer_to_lora、only_optimize_lora_parameters、make_model_gradient_checkpointing_compatible、convert_lora_to_linear_layer等函数。训练流程中的 LoRA 处理见 main.py先转换线性层为 LoRA若同时传入--only_optimize_lora则冻结其余参数、只优化 LoRA 权重LoRA 专属学习率--lora_learning_rate默认 5e-4保存模型前统一调用convert_lora_to_linear_layer把 LoRA 权重合并回原模型输出仍是标准 HF 格式。由于只训练低秩增量LoRA 方案的单卡 batch size 与整体吞吐都优于全参模式是快速实验指令微调效果的高性价比路径。八、评估微调效果base 与 finetuned 对比训练完成后用评估脚本对比基础模型与微调模型的回答质量cd finetune/scripts bash run_eval.sh脚本内容run_eval.sh会调用 prompt_eval.pycd $(dirname ${BASH_SOURCE[0]})/../sft/ python prompt_eval.py \ --model_name_or_path_base/base_model \ --model_name_or_path_finetune/finetuned_model \ --language Chinese运行后脚本会分别加载 base 模型与微调模型对同一组提示词进行贪心解码生成并分栏打印Baseline: Greedy 基础模型输出 finetune: Greedy 微调模型输出通过肉眼对比即可直观评估微调效果。该评估器支持的可调参数包括--num_beams束搜索宽度默认 1、--num_beam_groups、--top_k默认 4、--penalty_alpha默认 0.6、--num_return_sequences、--max_new_tokens默认 200、--languageChinese或English以及--eos默认|endoftext|。内置提示词与停止符细节--language Chinese时使用提示词Human: 你好。 Assistant:与Human: 请介绍一下大语言模型? Assistant:英文时对应Human: Hello. Assistant:与Human: Please explain Large Language Model. Assistant:。生成时通过tokenizer.convert_tokens_to_ids(|endoftext|)显式设置停止符与 Yi 的对话模板保持一致见 prompt_eval.py。评估模型以device_mapauto、torch_dtypeauto加载见 model_utils.py即双模型同时驻留显存进行推理请预留足够显存。九、常见问题与注意事项34B 训练 CPU 内存不足务必保留 run_sft_Yi_34b.sh 中的CUDA_VISIBLE_DEVICES0,1,2,3限制并确认节点 CPU 内存不低于 900GB。自定义数据集无法加载检查数据目录名称是否包含yi关键字触发 data_utils.py 的YiDataset分支并确认data/train.jsonl与data/eval.jsonl两个文件同时存在。训练被提前截断是预期行为training_debug_steps是调试步数截断6B 场景的20 分钟出 chat 模型正是依赖该设置正式训练请移除该参数或调大数值。序列长度与显存脚本默认--max_seq_len 4096对应 Yi 的 4K 上下文配合--gradient_checkpointing使用显存紧张时可下调该值。初始化阶段耗时34B 权重加载与 ZeRO 状态构建耗时较长属正常现象耐心等待即可。权重保存格式ZeRO-2 场景直接保存合并后的 HF 格式权重到--output_dir若改用 ZeRO-3需依赖save_zero_three_model完成分片聚合勿直接拷贝 checkpoint 目录。十、仓库配套资源索引训练主程序finetune/sft/main.py参数解析、数据装载、训练循环、模型保存评估程序finetune/sft/prompt_eval.pybase 与 finetuned 对比生成训练脚本run_sft_Yi_6b.sh / run_sft_Yi_34b.sh / run_sft_lora_Yi_6b.sh评估脚本run_eval.shDeepSpeed 配置finetune/utils/ds_utils.py模型加载工具finetune/utils/model/model_utils.py数据集解析finetune/utils/data/data_utils.py 与 finetune/utils/data/raw_datasets.pyLoRA 实现finetune/utils/module/lora.py训练阶段常量SFT/RM/PPO/RSFT/DPOfinetune/constant.py示例数据集finetune/yi_example_dataset/data/train.jsonl 与 finetune/yi_example_dataset/data/eval.jsonl【免费下载链接】YiA series of large language models trained from scratch by developers 01-ai项目地址: https://gitcode.com/GitHub_Trending/yi/Yi创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考