我要投稿

聊聊Firefly——开源且值得学习的LLM训练库

发布日期：2025-01-08 17:24:07 浏览次数： 2648

作者：阿郎小哥的随笔驿站

微信搜一搜，关注“阿郎小哥的随笔驿站”

概述

Firefly^[1]是一个开源的大模型训练项目，支持对主流的大模型进行预训练、指令微调和DPO，包括但不限于Qwen2、Yi-1.5、Llama3、Gemma、Qwen1.5、MiniCPM、MiniCPM3、Llama、InternLM、Baichuan、ChatGLM、Yi、Deepseek、Qwen、Orion、Ziya、Xverse、Mistral、Mixtral-8x7B、Zephyr、Vicuna、Bloom等。本项目支持全量参数训练、LoRA、QLoRA高效训练，支持预训练、SFT、DPO。如果你的训练资源有限，极力推荐使用QLoRA进行指令微调，因为我们在Open LLM Leaderboard上验证了该方法的有效性，并且取得了非常不错的成绩。

本项目主要内容如下：

支持预训练、指令微调、DPO，支持全量参数训练、LoRA、QLoRA高效训练。通过配置文件的方式训练不同的模型，小白亦可快速上手训练模型。
支持使用 Unsloth^[2] 加速训练，并且节省显存。
支持绝大部分主流的开源大模型，如Llama3、Gemma、MiniCPM、Llama、InternLM、Baichuan、ChatGLM、Yi、Deepseek、Qwen、Orion、Ziya、Xverse、Mistral、Mixtral-8x7B、Zephyr、Vicuna、Bloom，训练时与各个官方的chat模型的template对齐。
整理并开源指令微调数据集：firefly-train-1.1M 、moss-003-sft-data、ultrachat、 WizardLM_evol_instruct_V2_143k、school_math_0.25M。
开源 Firefly系列指令微调模型权重^[3] 。
在Open LLM Leaderboard上验证了QLoRA训练流程的有效性。

实践

Firefly官方基于该项目的训练代码，以及训练数据，训练并开源了以下模型权重。

模型

如下模型是官方在Firefly代码上训练的，训练脚本可以源码目录找到：train_args^[4]

中文模型：

模型	基座模型	训练长度
firefly-baichuan2-13b^[5]	baichuan-inc/Baichuan2-13B-Base	1024
firefly-baichuan-13b^[6]	baichuan-inc/Baichuan-13B-Base	1024
firefly-qwen-7b^[7]	Qwen/Qwen-7B	1024
firefly-chatglm2-6b^[8]	THUDM/chatglm2-6b	1024
firefly-internlm-7b^[9]	internlm/internlm-7b	1024
firefly-baichuan-7b^[10]	baichuan-inc/baichuan-7B	1024
firefly-ziya-13b^[11]	YeungNLP/Ziya-LLaMA-13B-Pretrain-v1	1024
firefly-bloom-7b1^[12]	bigscience/bloom-7b1	1024
firefly-bloom-2b6-v2^[13]	YeungNLP/bloom-2b6-zh	512
firefly-bloom-2b6^[14]	YeungNLP/bloom-2b6-zh	512
firefly-bloom-1b4^[15]	YeungNLP/bloom-1b4-zh	512

英文模型：

模型	基座模型	训练长度
firefly-mixtral-8x7b^[16]	mistralai/Mixtral-8x7B-v0.1	1024
firefly-llama-30b^[17]	huggyllama/llama-30b	1024
firefly-llama-13-v1.2^[18]	NousResearch/Llama-2-13b-hf	1024
firefly-llama2-13b^[19]	NousResearch/Llama-2-13b-hf	1024
firefly-llama-13b-v1.2^[20]	huggyllama/llama-13b	1024
firefly-llama-13b^[21]	huggyllama/llama-13b	1024

扩展项目：

Firefly-LLaMA2-Chinese^[22]：中文Llama2模型，对Llama2进行中文词表扩充、增量预训练和指令微调。
LongQLoRA^[23]：大模型长度扩展项目，可在单卡V100上将LLaMA-13B的长度扩展至8192，且性能逼近MPT-8K。

训练数据

Firefly整理了如下指令集，并整理成统一的数据集格式：

数据集	介绍
firefly-train-1.1M^[24]	我们收集了23种常见的中文NLP任务的数据，并且构造了许多与中华文化相关的数据，如对联、作诗、文言文翻译、散文、金庸小说等。对于每个任务，由人工书写若干种指令模板，保证数据的高质量与丰富度，数据量为115万
moss-003-sft-data^[25]	由复旦大学MOSS团队开源的中英文多轮对话数据，包含100万+数据
ultrachat^[26]	由清华大学开源的英文多轮对话数据，包含140万+数据
WizardLM_evol_instruct_V2_143k^[27]	由WizardLM项目开源的英文指令微调数据集，通过Evol-Instruct方法让指令进化，加强指令的复杂度，以提升模型对复杂指令的遵循能力。包含143k条数据。
school_math_0.25M^[28]	由BELLE项目组开源的数学运算指令数据，包含25万条数据。
shareAI/CodeChat^[29]	主要包含逻辑推理、代码问答、代码生成相关语料样本。
shareAI/ShareGPT-Chinese-English-90k^[30]	中英文平行双语优质人机问答数据集，覆盖真实复杂场景下的用户提问。
ultrachat_200k^[31]	由Zephyr项目开源的英文指令微调数据，在ultrachat数据基础上进行清洗
ultrafeedback_binarized^[32]	英文偏好数据集，可用于DPO训练

训练模型

train_args目录下存储了不同模型使用不同训练方式的配置文件，主要参数说明如下：

output_dir：训练输出目录，存储checkpoint、tokenizer、tensorboard等
model_name_or_path：预训练模型的本地目录，或者在huggingface上的模型名称。
train_file：训练数据集路径。sft时，需要设置为文件，可以使用data/dummy_data.jsonl进行debug。pretrain时，需要设置为目录。脚本会自动扫描目录下的所有jsonl文件。
template_name：指令微调时，使用的模板名称。具体有哪些template_name，可参考component/template.py文件
num_train_epochs：训练的轮次。如果数据量足够大，一般建议只训一个epoch。
tokenize_num_workers：预训练时，tokenize的线程数，默认为10。
deepspeed：deepspeed的训练配置文件。全量参数训练时，将采用deepspeed，关于deepspeed的参数配置说明，请参考deepspeed文档^[33]
train_mode：训练模式，full、lora或qlora，默认为qlora。
task_type：任务类型，pretrain、sft或dpo，默认为sft。
per_device_train_batch_size：每张显卡的batch size。
gradient_accumulation_steps：梯度累计步数。global batch=num_gpus * per_device_train_batch_size * gradient_accumulation_steps。
gradient_checkpointing：如果显存捉襟见肘，可以开启。以时间换空间，模型不缓存激活状态，会进行两次forward计算，以节省显存。
learning_rate：学习率。全量参数微调的时候，建议小一些，1e-5或5e-6。
max_seq_length：训练时的最大长度。按照自己的设备进行设置，越长需要占用越多显存。
max_prompt_length：进行dpo时，prompt的最大长度。
logging_steps：每隔多少步统计一次train loss。
save_steps：每隔多少步保存一个模型。
save_total_limit：output_dir目录中最多保存多少个checkpoint，超出则会将最旧的删除。
lr_scheduler_type：学习率变化策略。
warmup_steps：warm up步数。学习率经过多少步，增长到指定的数值。
optim：优化器。如果是全量参数微调，建议使用adamw_hf。
seed：随机种子，用于复现实验结果。
fp16：使用使用fp16混合精度。V100建议开启。
bf16：使用使用bf16混合精度。A100建议开启。
use_unsloth：是否使用unsloth，目前unsloth仅支持部分模型，例如Llama3、Mistral、Gemma、TinyLlama等，详情见Unsloth^[34]。

以下几个参数，当使用QLoRA训练的时候，需要设置：

lora_rank：qlora矩阵的秩。一般设置为8、16、32、64等，在qlora论文中作者设为64。越大则参与训练的参数量越大，一般来说效果会更好，但需要更多显存，。
lora_alpha: qlora中的缩放参数。一般设为16、32即可。
lora_dropout: lora权重的dropout rate。

关于deepspeed的参数配置，可按需自行修改。

训练脚本

全量参数预训练，将{num_gpus}替换为显卡数量：

deepspeed --num_gpus={num_gpus} train.py --train_args_file train_args/pretrain/full/bloom-1b1-pretrain-full.json

全量参数指令微调，将{num_gpus}替换为显卡数量：

deepspeed --num_gpus={num_gpus} train.py --train_args_file train_args/sft/full/bloom-1b1-sft-full.json

单卡QLoRA预训练：

python train.py --train_args_file train_args/pretrain/qlora/yi-6b-pretrain-qlora.json

单卡QLoRA指令微调：

python train.py --train_args_file train_args/sft/qlora/yi-6b-sft-qlora.json

多卡QLoRA预训练：

torchrun --nproc_per_node={num_gpus} train.py --train_args_file train_args/pretrain/qlora/yi-6b-pretrain-qlora.json

多卡QLoRA指令微调：

torchrun --nproc_per_node={num_gpus} train.py --train_args_file train_args/sft/qlora/yi-6b-sft-qlora.json

单卡QLoRA进行DPO训练：

python train.py --train_args_file train_args/sft/qlora/minicpm-2b-dpo-qlora.json

❝
更详细的训练及推理使用说明，可以查看官网。

Loss计算

预训练时，我们采用经典的自回归损失，即每个位置的token都会参与loss计算。

指令微调时，仅计算assistant回复部分的loss。

❝
这点还是蛮重要的，在一些训练框架中，不会提及也不会如此设计。

小结

我个人接触这个库比较早，之前也算是用过，但后面选型用的还是LLaMA-Factory；现在深入的再了解对比，其实个人玩LLM的话，还是可以考虑Firefly。首先，对比较于LLaMA-Factory，Firefly官方开发人员基于此项目训练了一批LLM，从数据集到训练脚本参数，都是慢慢的提供了，这在我看来是实实在在的开源了。对于学生或是个人了解LLM的训练过程来说，是一个很好的参考。其次的话，看过一些issue，包括我自己之前也写过几篇关于Firefly训练效率的文章，在训练效率和效果上的确是要比LLaMA-Factory要好一些。前段时间，Firefly官方发布了基于Firefly与QWen知识蒸馏的文章，但还没把代码放出来，如果能放出来的话，那更好了。不过最大的劣势则是对模型的支持没有LLaMA-Factory齐全，且维护更新、解决issue并没有那么活跃。

对于学生、业余爱好者来说，推荐从Firefly入手训练模型，再可以过渡到LLaMA-Factory那边去使用，在Firefly的训练中，可以接触到更多的训练细节与实践。