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

训练数据

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

训练模型

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计算。