更新时间:2024-12-17 GMT+08:00
分享

Eagle投机小模型训练

本章节提供eagle小模型自行训练的能力,客户可通过本章节,使用自己的数据训练eagle小模型,并使用自行训练的小模型进行eagle推理。支持llama1系列、llama2系列和Qwen2系列模型。

步骤三:sharegpt格式数据生成为训练data数据集

设置环境变量。

export EAGLE_TARIN_MODE=1

如果使用开源数据集,推荐使用原论文代码仓数据集,下载地址:https://huggingface.co/datasets/Aeala/ShareGPT_Vicuna_unfiltered/blob/main/ShareGPT_V4.3_unfiltered_cleaned_split.json

如果使用其他数据集,需要先执行步骤二:非sharegpt格式数据集转换(可选)转换数据集格式为sharegpt格式。

执行如下脚本将sharegpt格式数据生成为训练data数据集。

python allocation.py \
--outdir outdir0/sharegpt_0_99_mufp16 \
--end_num 100 \
--npu_indices "0,1,2,3,4,5,6,7" \
--used_npus 8  \
--model_type llama \
--model_name ./llama-7B  \
--data_path data_for_sharegpt.json \
--seed 42 \
--max_length 2048
--dtype bfloat16

其中

  • outdir:生成的训练data地址。
  • end_num:生成的data总条数。
  • npu_indices:使用哪些NPU卡。
  • used_npus:拉起的每个py脚本使用几个NPU,如果为70b则填写4或8,7b 13b则填1。
  • model_type llama:使用模型类型,目前支持qwen2、llama1、llama2,其中llama1、llama2填写llama,qwen2填写为qwen2。
  • model_name:模型地址。
  • data_path:预训练数据集地址。
  • seed 42:生成训练data所使用的seed,此处42为开源训练设定参数。
  • max_length:模型的max_length。
  • dtype:为模型dtype,默认为bfloat16。

执行完成后,记得unset环境变量,否则会导致后续推理服务启动出错。

unset EAGLE_TARIN_MODE

执行完成后,如果used_npus>1,则需要将训练生成data数据重新分配为8个文件夹,分配脚本为reassign_data_num.py。

python reassign_data_num.py --old_folder "./sharegpt_0_199_mufp16/" \
--new_folder "./sharegpt_0_199_mufp16/" \
--tp 8
  • old_folder为上一步生成data的地址,填写到卡号的文件夹之前。命令中的./sharegpt_0_199_mufp16/"为举例,需要替换为实际地址。
  • new_folder为需要存储新的data的地址。命令中的./sharegpt_0_199_mufp16/"为举例,需要替换为实际地址。
  • tp为需要切分成的文件夹数量,默认为8。

步骤四:执行训练

安装完成后,执行:

accelerate launch -m --mixed_precision=bf16 eagle.train.main \
--tmpdir [path of data] \
--cpdir [path of checkpoints] \
--configpath [path of config file] \
--basepath  [path of base_model]
--bs [batch size]
  • tmpdir:即为步骤三中的outdir,训练data地址
  • cpdir:为训练生成权重的地址
  • configpath:为模型config文件的地址
  • basepath:为大模型权重地址
  • bs:为batch大小

其中,要获取模型config文件, 首先到https://github.com/SafeAILab/EAGLE/页找到对应eagle模型地址。

图1 EAGLE Weights

以llama2-chat-7B为例,单击进入后 ,如下图所示config文件,即为对应模型的eagle config文件。

图2 eagle config文件

步骤五:训练生成权重转换成可以支持vLLM推理的格式

将训练完成后的权重文件(.bin文件或. safetensors文件),移动到下载好的开源权重目录下(即步骤4中,config文件所在目录)。

然后在llm_tools/spec_decode/EAGLE文件夹,执行
python convert_eagle_ckpt_to_vllm_compatible.py --base-path 大模型权重地址 --draft-path 小模型权重地址 --base-weight-name 大模型包含lm_head的权重文件名 --draft-weight-name 小模型权重文件名
  • --base-path:为大模型权重地址,例如 ./llama2-7b-chat
  • --draft-path:小模型权重地址,即步骤四中config文件所在目录,例如 ./eagle_llama2-7b-chat
  • --base-weight-name:为大模型包含lm_head的权重文件名,可以在base-path目录下的 model.safetensors.index.json 文件获取,例如llama2-7b-chat的权重名为pytorch_model-00001-of-00002.bin
图3 权重文件名

--draft-weight-name 为小模型权重文件名,即刚才移动的.bin文件或者.safetensors文件。

相关文档