LLaMA参考代码：transformers库，引用路径transformers.models.llama
Moefication参考代码：https://github.com/thunlp/MoEfication
| 文件夹 | 说明 |
|---|---|
| llama_data | 用于训练MoE门控网络的数据 |
| llama_download | 用于下载LLaMA模型 |
| llama_moe | MoE版LLaMA的模型实现，基于transformers库修改 |
| run_moefication | 用于将普通的LLaMA转化为使用Moefication实现的LLaMA-MoE |
与基于transformers库的LLaMA相同，详情请见示例代码test_llama_moe_transformers.py。
如果已经下载了LLaMA模型，并进行了transformers库的格式转换，可跳过该步骤。
修改./llama_download/download_llama.sh文件前3行中的MODEL_SIZE与TARGET_FOLDER变量，之后运行bash download_llama.sh下载LLaMA模型。
PRESIGNED_URL="https://agi.gpt4.org/llama/LLaMA/*"
MODEL_SIZE="7B,13B" # 30B,65B
TARGET_FOLDER="" # where all files should end up
declare -A N_SHARD_DICT
N_SHARD_DICT["7B"]="0"
N_SHARD_DICT["13B"]="1"
......运行./llama_download/convert_llama_weights_to_hf.py文件，并传入input_dir，model_size，output_dir参数，将原版LLaMA模型转换为transformers库可识别的格式。
python convert_llama_weights_to_hf.py --input_dir <SOURCE_FOLDER> --model_size 7B --output_dir <TARGET_FOLDER>如果上述命令报错，可尝试在运行时设置PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION系统变量为python，即运行以下命令：
PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python python convert_llama_weights_to_hf.py --input_dir <SOURCE_FOLDER> --model_size 7B --output_dir <TARGET_FOLDER>运行命令可参照目录
./run_moefication/下的sh文件进行修改，也可直接命令行调用python文件。
Moefication主要包含四个阶段。
- 将模型中的dense层参数进行分割，得到n个相互独立的子集（即n个专家），并保存各神经元所对应的专家索引；
- 使用LLaMA对预训练数据进行推理，保存LLaMA的中间层输入输出，作为门控网络的训练数据。
- 训练用于选择专家的门控网络，并保存门控网络的参数；
- 通过保存的索引与门控网络参数，将LLaMA模型转换为MoE形式。
Moefication原文提供了clustering与graph两种参数分割方式，两者效果近似相同。
此处使用了更加通用的clustering分割方式。
运行./run_moefication/llama_split_clustering.py文件，并传入下述参数：
| 参数名 | 说明 |
|---|---|
| model_path | transformers库格式的LLaMA模型的路径 |
| save_path | 保存专家划分后的index文件的路径，会该位置下自动创建名如"xxx-xxxExpert-Split-Clustering"的文件夹进行保存 |
| templates | 专家划分所依据的网络参数模板，默认为"layers.{}.mlp.gate_proj.weight"，不需要修改 |
| num_experts | 专家划分的数量 |
下面为运行命令的范例：
python llama_split_clustering.py --model_path <MODEL_PATH>/llama_7B --save_path ../llama_moe_temp_files --num_experts 8上述命令运行结束后，会在../llama_moe_temp_files/llama_7B-8Expert-Split-Clustering路径下生成若干文件，保存了dense层各神经元所对应的专家索引，名称格式如下：
layers.0.mlp.gate_proj.weight
layers.1.mlp.gate_proj.weight
layers.2.mlp.gate_proj.weight
......2.2 获取LLaMA的中间层特征 (get_hidden_features)
训练MoE门控网络需要使用LLaMA的预训练数据，下图为LLaMA预训练时各数据集的采样比例。
此处在./llama_data文件夹下提供了上述数据集的子集，使用前先解压zip文件，得到以jsonl格式存储的数据集。
提供的数据集中各个数据集的token数量基本相近，具体的比例分配在./run_moefication/llama_get_hidden_features.py中实现。
运行./run_moefication/llama_select_mlp.py文件，并传入下述参数：
| 参数名 | 说明 |
|---|---|
| model_path | transformers库格式的LLaMA模型的路径 |
| train_data_path | LLaMA预训练数据的保存路径 (步骤2.2.1) |
| train_data_cache_path | LLaMA预训练数据的缓存保存路径，用于保存与读取数据缓存，减少数据处理时间 |
| save_path | 保存门控网络参数文件的位置，会该位置下自动创建名如"xxx-Hidden-Features"的路径进行保存 |
| templates | 专家划分所依据的网络参数模板，默认为"layers.{}.mlp.gate_proj.weight"，不需要修改 |
| data_use_percent | 所有数据集中数据的使用比例，用于调节训练使用的数据量 3.5M的token在LLaMA 7B上得到的中间层输出文件总大小约4.5T |
| batch_size | 单次推理的batch_size，根据显存大小调节 |
| save_interval | 保存参数的batch间隔，调大会影响显存占用，但可以减少保存的文件个数 |
推理代码使用了torch原生多卡形式，需使用python -m torch.distributed.launch运行，以下为运行命令的范例：
python -m torch.distributed.launch --nproc_per_node=8 llama_get_hidden_features.py --model_path <MODEL_PATH>/llama_7B --train_data_path ../llama_data --train_data_cache_path ../llama_data_cache --save_path ../llama_moe_temp_files --data_use_percent 0.01 --save_interval 1 --batch_size 4上述命令运行结束后，会在../llama_moe_temp_files/llama_7B-Hidden-Features路径下生成若干文件，保存了中间层输入输出，名称格式如下：
-- hidden_inputs
0_0.pth
0_1.pth
0_2.pth
......
-- hidden_gate_outputs
0_0.pth
0_1.pth
0_2.pth
......Moefication原文提供了Similarity与MLP两种门控机制。
此处选择常规的MLP作为门控网络，网络输出对应各个专家的选择分数，该方案效果较好。
运行./run_moefication/llama_select_mlp.py文件，并传入下述参数：
| 参数名 | 说明 |
|---|---|
| model_path | transformers库格式的LLaMA模型的路径 |
| split_file_path | 步骤2.1中保存的专家索引路径 |
| hidden_features_path | 步骤2.2中保存的LLaMA中间层特征路径 |
| save_path | 保存门控网络参数文件的位置，会该位置下自动创建名如"xxx-xxxExpert-Select-MLP"的路径进行保存 |
| templates | 专家划分所依据的网络参数模板，默认为"layers.{}.mlp.gate_proj.weight"，不需要修改 |
| select_criterion | 专家划分的依据指标，有plain、positive、l2_norm三种方式，默认使用最合理的l2_norm |
| num_experts | 专家划分的数量 |
| num_selects | 专家选择的数量 |
| specify_layer | 指定对哪些层进行划分，用于并行执行划分操作，留空则对所有层进行划分 如LLaMA 7B有32层，则可使用4条命令，分别指定0-7、8-15、16-23、24-32进行并行训练 |
下面为运行命令的范例：
python llama_split_clustering.py --model_path <MODEL_PATH>/llama_7B --split_file_path ../llama_moe_temp_files/llama_7B-8Expert-Split-Clustering --hidden_features_path ../llama_moe_temp_files/llama_7B-Hidden-Features --save_path ../llama_moe_temp_files --select_criterion l2_norm --num_experts 8 --num_selects 2上述命令运行结束后，会在../llama_moe_temp_files/llama_7B-8Expert-Select-MLP路径下生成若干文件，保存了各个门控网络的参数，名称格式如下：
layers.0.mlp.gate_proj.weight
layers.1.mlp.gate_proj.weight
layers.2.mlp.gate_proj.weight
......修改./run_moefication/llama_convert.py.py文件第4行的路径为项目根路径。
sys.path.append("") # 修改为项目根路径，如"/home/dongdz/workspace/llama-moefication/"运行./run_moefication/llama_convert.py文件，并传入下述参数：
| 参数名 | 说明 |
|---|---|
| model_path | transformers库格式的LLaMA模型的路径 |
| split_file_path | 步骤2.1中保存的专家索引路径 |
| select_file_path | 步骤2.3中保存的门控网络参数路径 |
| save_path | 保存转换后的transformers库格式的LLaMA-MoE模型的路径 |
| templates | 专家划分所依据的网络参数模板，默认为"layers.{}.mlp.gate_proj.weight"，不需要修改 |
| num_experts | 专家划分的数量 |
| num_selects | 专家选择的数量 |
下面为运行命令的范例：
python llama_convert.py --model_path <MODEL_PATH>/llama_7B --split_file_path ../llama_moe_temp_files/llama_7B-8Expert-Split-Clustering --select_file_path ../llama_moe_temp_files/llama_7B-8Expert-Select-MLP --save_path <MODEL_PATH>/llama_7B-MoE --num_experts 8 --num_selects 2上述命令运行结束后，会在<MODEL_PATH>/7B-MoE路径下保存LLaMA-MoE模型。
可以自己看./llama_moe/moe_utils/下的代码，其中moe_layers.py是封装好的MoE层，用于直接调用。