Eagle投机小模型训练
什么是Eagle投机小模型训练
2013年12月滑铁卢大学、加拿大向量研究院、北京大学等机构联合发布Eagle，旨在提升大语言模型的推理速度，同时保证模型输出文本的分布一致。这种方法外推LLM的第二顶层特征向量，能够显著提升生成效率。
Eagle训练了一个单层模型，使用input token和基模型推理出的hidden-state作为输入，输出hidden-state。然后根据这个输出的hidden-state使用基模型的原始LLM的分类头来预测下一个词。hidden-state比input token包含更多信息，使得回归hidden-state的任务比预测词的任务简单得多。总之，Eagle在hidden-state层面上进行外推，使用一个小型单层Eagle模型，然后利用基模型的冻结的分类头生成预测的token。
如此一来，Eagle投机推理可以带来如下优势：
- 更小的训练成本得到小模型：相较于训练独立的LLM大模型，Eagle仅需训练一个自回归层。这使得其训练成本相较于训练一个独立的LLM模型要小得多。
- 为每个模型提供针对性的投机模型： Eagle的模型大小及结构，与基模型的某一层完全相同，这使得它的大小远远小于其基模型。解决了对于部分原始LLM模型，找不到合适的投机模型的问题。
投机小模型训练端到端示例
本章节提供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
其中具体参数解释如表2所示。
|
py文件名称 |
配置项 |
取值类型 |
配置说明 |
|---|---|---|---|
|
allocation.py |
--outdir |
str |
生成的训练所需数据的输出地址 |
|
--end_num |
int |
生成的训练数据总条数 |
|
|
--npu_indices |
str |
使用哪些NPU进卡行数据生成，用逗号隔开，如"0,1,2,3,4,5,6,7" |
|
|
--used_npus |
int |
拉起的每个py脚本使用几个NPU，如果为70b则填写4或8，7b 13b则填1。 |
|
|
--model_type |
str |
使用模型类型，目前支持llama系列（填写llama）及qwen2（填写qwen2） |
|
|
--model_name |
str |
LLM的基模型地址，如./Llama2-7b |
|
|
--data_path |
str |
预训练数据集地址，如sharegpt.json |
|
|
--seed |
int |
生成训练数据所使用的seed，默认为42，42为开源训练设定参数。 |
|
|
--max_length |
int |
模型的最大长度，默认为2048 |
|
|
--dtype |
str |
模型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]
具体的参数解释如表3所示，
|
py文件名称 |
配置项 |
取值类型 |
配置说明 |
|---|---|---|---|
|
eagle.train.main |
--tmpdir |
str |
生成的训练所需数据的输出地址，即数据生成中的outdir |
|
--cpdir |
str |
训练出的Eagle模型权重存放地址 |
|
|
--configpath |
str |
Eagle模型config文件地址，通常将对应LLM原始模型中的num_layers改为1即可。 |
|
|
--basepath |
str |
LLM原始模型权重地址 |
|
|
--bs |
int |
训练模型的batch |
其中，要获取模型config文件， 首先到https://github.com/SafeAILab/EAGLE/页找到对应Eagle模型地址。
以llama2-chat-7B为例，单击进入后 ，如下图所示config文件，即为对应模型的eagle config文件。
步骤五：训练生成权重转换成可以支持vLLM推理的格式
将训练完成后的权重文件（.bin文件或. safetensors文件），移动到下载好的开源权重目录下（即步骤四：执行训练中config文件所在目录）。
python convert_eagle_ckpt_to_vllm_compatible.py --base-path 大模型权重地址 --draft-path 小模型权重地址 --base-weight-name 大模型包含lm_head的权重文件名 --draft-weight-name 小模型权重文件名
具体参数解释如表4所示。
|
py文件名称 |
配置项 |
取值类型 |
配置说明 |
|---|---|---|---|
|
convert_eagle_ckpt_to_vllm_compatible.py |
--base-path |
str |
LLM原始模型权重地址，例如 ./llama2-7b-chat |
|
--draft-path |
str |
Eagle模型权重存放地址，即步骤四：执行训练中config文件所在目录，例如 ./eagle_llama2-7b-chat |
|
|
--base-weight-name |
str |
为大模型包含lm_head的权重文件名，可以在base-path目录下的model.safetensors.index.json文件获取，例如llama2-7b-chat的权重名为pytorch_model-00001-of-00002.bin |
|
|
--draft-weight-name |
str |
Eagle小模型权重文件名，其格式为.bin文件或者.safetensors文件，即刚才移动的.bin文件或者.safetensors文件。 |
