本项目用于 TinyLlama 模型的 ONNX 导出、切分以及基于 ONNX/ACL 的分布式流水线推理。
export_llama.py: 模型导出脚本,支持量化配置。split_on_onnx.py: ONNX 模型切分脚本。inference_net/: 推理相关代码,包含单机和分布式节点逻辑。config/: 量化配置文件 (如w8x8.py,w8.py等)。
在执行任何脚本之前,必须使用本项目提供的 modeling_llama_4.35.py 替换 transformers 库中的 modeling_llama.py 文件。这是为了修改模型的前向传播逻辑以支持导出。
# 1. 找到你的 transformers 库安装路径 # 例如: /home/yjr/miniconda3/envs/llama_export/lib/python3.9/site-packages/transformers TRANSFORMERS_PATH=$(python -c "import transformers; import os; print(os.path.dirname(transformers.__file__))") # 2. 备份原文件 (可选) cp $TRANSFORMERS_PATH/models/llama/modeling_llama.py $TRANSFORMERS_PATH/models/llama/modeling_llama.py.bak # 3. 替换文件 cp modeling_llama_4.35.py $TRANSFORMERS_PATH/models/llama/modeling_llama.py
将 PyTorch 模型导出为 ONNX 格式,并应用指定的量化配置。
python export_llama.py \ --model ../ascend-llm/export_llama/model/TinyLlama-1.1B-Chat-v1.0/ \ --output ./tiny-llama.onnx \ --act-path ./tiny-llama.pt \ --quant ./config/w8x8.py
将导出的 ONNX 模型切分为多个部分(如 M0, M1, M2...),以便进行流水线并行推理。
python split_on_onnx.py \ --src ./tiny-llama.onnx \ --out_dir ./split_models/
将切分后的 ONNX 模型转换为 Ascend 的 OM 离线模型。请确保已安装 Ascend CANN 工具包并设置好环境变量,且进入存放 ONNX 模型的目录(如 split_models)。
# 假设你在 split_models 目录下,或者将 convert_to_om.sh 复制到该目录运行 # 记得修改脚本中的 --soc_version 为你的实际芯片型号 (如 Ascend310B1) bash ../convert_to_om.sh
推理脚本位于 inference_net 目录下,请先进入该目录:
cd inference_net验证切分后的 ONNX 模型是否能正常加载和推理。
python main.py \ --hf-dir ../../ascend-llm/export_llama/model/TinyLlama-1.1B-Chat-v1.0/ \ --engine onnx \ --kv_size 1024 \ --split-model-dir ../split_models/
如果是已转换为 OM 格式的模型,使用 ACL 引擎运行(需在 Ascend 环境下)。
python main.py \ --hf-dir ../export_llama/model/TinyLlama-1.1B-Chat-v1.0/ \ --engine acl \ --kv_size 1024 \ --split-model-dir ../export_llama/model/export_out/om_model/
模拟多节点流水线推理,需按顺序启动各个节点。
注意: 请确保
../export_llama/model/export_out/om_model/等路径存在且包含对应的 OM 模型文件。
Step 1: 启动中间节点 Node 2 (Layers 5-10)
python run_intermediate.py \
--model ../export_llama/model/export_out/om_model/llama_m1_layers_5_10.om \
--port 8001 \
--next_ip 127.0.0.1 \
--next_port 8002 \
--hf_dir ../export_llama/model/TinyLlama-1.1B-Chat-v1.0/ \
--n_layer 6 \
--node_name "Node 2" \
--device 0 \
--kv_size 1024Step 2: 启动中间节点 Node 3 (Layers 11-16)
python run_intermediate.py \
--model ../export_llama/model/export_out/om_model/llama_m2_layers_11_16.om \
--port 8002 \
--next_ip 127.0.0.1 \
--next_port 8003 \
--hf_dir ../export_llama/model/TinyLlama-1.1B-Chat-v1.0/ \
--n_layer 6 \
--node_name "Node 3" \
--device 0 \
--kv_size 1024Step 3: 启动尾节点 Node 4 (Layers 17-21 & Head)
python run_final.py \
--model ../export_llama/model/export_out/om_model/llama_m3_layers_17_21_lmhead.om \
--port 8003 \
--head_ip 127.0.0.1 \
--head_port 8004 \
--hf_dir ../export_llama/model/TinyLlama-1.1B-Chat-v1.0/ \
--n_layer 5 \
--node_name "Node 4" \
--device 0 \
--kv_size 1024Step 4: 启动主节点/Client (Embedding & Layers 0-4)
python main.py \ --engine net \ --split-model-dir ../export_llama/model/export_out/om_model \ --hf-dir ../export_llama/model/TinyLlama-1.1B-Chat-v1.0/ \ --next-ip 127.0.0.1 \ --next-port 8001 \ --listen-port 8004 \ --kv_size 1024