切分和编译工具使用说明

环境依赖

Transformers基座模型切分工具需要的依赖环境有：

  • python3:

    • 安装python依赖包:

      • 在CPU-only环境上安装依赖包

      python3 -m pip install transformers>=4.29.0 git+https://github.com/Quentin-Anthony/DeepSpeed@cpu-support#egg=deepspeed mpi4py torch==1.10.0 torchvision==0.11.1 einops==0.6.1 safetensors numpy==1.21

      • 在GPU环境上安装依赖包

      python3 -m pip install transformers>=4.29.0 deepspeed mpi4py torch==1.10.0 torchvision==0.11.1 einops==0.6.1 safetensors

  • 切分工具：

    • 通过以下命令安装切分工具

    python3.8 -m pip install topsdistinfer_partition_cpu-*-py3-none-any.whl

功能介绍

切分模型工具

为了让大模型通过流水并行方式进行推理，需要提前对模型进行切分。这个工具提供切分功能。

命令

  • 在CPU-only环境上:

deepspeed --num_cpu_cores PIPE_PARA_NUM --module partition.transformers --model MODEL_TYPE --ckpt CHECK_POINT --config CONFIG [--method PART_METHOD] [--split_kvcache] [--output_prefix PREFIX] [--output_suffix SUFFIX]

  • 在GPU-only环境上:

deepspeed --num_gpus PIPE_PARA_NUM --module partition.transformers --model MODEL_TYPE --ckpt CHECK_POINT --config CONFIG [--method PART_METHOD] [--split_kvcache] [--output_prefix PREFIX] [--output_suffix SUFFIX]

参数

参数 释义
PIPE_PARA_NUM 流水并行的stage数，也就是模型切分的份数
MODEL_TYPE 待切分模型的类别，现在支持的类别有'gpt2','llama','chatglm','chatglm2','opt','bloom','rw','cpmbee','baichuan','internlm','llama2','aquila','gptneox','gptbigcode','pangu','baichuan2' ,'gptj'
CHECK_POINT huggingface的checkpoint格式中的pytorch_model.bin文件或者pytorch_model.bin.index.json文件的路径
CONFIG huggingface的checkpoint格式中的config.json文件的路径
PART_METHOD 流水并行切分方法，现在支持的类型有'parameters','uniform','type:'，现在已经支持的有'GPT2BlockPipeline','LlaMADecoderLayerPipeline','GLMBlockPipeline','GLMBlock2Pipeline','OPTDecoderLayerPipeline','BloomBlockPipeline','RWDecoderPipeline','CpmBeeEncoderPipeline','BaichuanDecoderPipeline','InternLMDecoderPipeline','LlaMA2DecoderLayerPipeline','AquilaDecoderLayerPipeline','GPTBigCodeDecoderLayerPipeline','GPTNeoXDecoderLayerPipeline'
PREFIX 输出的onnx的文件名如alpaca_7b-kvcache-1of4-op13-fp32-N.onnx，其中alpaca_7b-kvcache-部分为onnx文件名的prefix
SUFFIX 输出的onnx的文件名如alpaca_7b-kvcache-1of4-op13-fp32-N.onnx，其中-op13-fp32-N部分为onnx文件名的suffix

示例

deepspeed --num_gpus 4 --module partition.transformers --model rw --split_kvcache --ckpt falcon-7b/pytorch_model.bin.index.json --config falcon-7b/config.json --output_prefix falcon_7b-kvcache-

输出文件

切分工具会生成流水并行的每个stage对应的onnx文件和对应的weights文件，为了方便后续其他操作，还为每个onnx文件生成了一个json文件，用来记录模型的输入和输出配置参数。

比如，一个alpaca模型切分输出文件列表如下：

├── alpaca_7b-kvcache-1of4-op13-fp32-N.json
├── alpaca_7b-kvcache-1of4-op13-fp32-N.onnx
├── alpaca_7b-kvcache-1of4-op13-fp32-N.weights
├── alpaca_7b-kvcache-2of4-op13-fp32-N.json
├── alpaca_7b-kvcache-2of4-op13-fp32-N.onnx
├── alpaca_7b-kvcache-2of4-op13-fp32-N.weights
├── alpaca_7b-kvcache-3of4-op13-fp32-N.json
├── alpaca_7b-kvcache-3of4-op13-fp32-N.onnx
├── alpaca_7b-kvcache-3of4-op13-fp32-N.weights
├── alpaca_7b-kvcache-4of4-op13-fp32-N.json
├── alpaca_7b-kvcache-4of4-op13-fp32-N.onnx
└── alpaca_7b-kvcache-4of4-op13-fp32-N.weights

生成TopsInference的engine工具

TopsDistInfer需要输入TopsInference的engine格式的模型。这个工具提供了：

  • 一个从onnx生成engine文件的统一编译工具

  • 一组独立的从onnx生成engine的编译工具

个别模型因为对编译模型有一些特别的要求，所以提供了各自的独立编译脚本。

统一编译工具

命令
python3 -m build_engine.gcu --json_path JSON [--onnx_path ONNX] --prec PREC --mode MODE [--device DEVICE]
参数
参数 释义
JSON 切分工具生成的json文件路径，或者包含这些json文件的上层文件夹路径
ONNX 指定onnx或者包含onnx文件的文件夹，这个参数如果不指定，默认会在json文件所在文件夹下面采用同名的onnx文件
PREC 生成engine的精度模式，支持'fp32','fp16','fp16-mix'
MODE 生成engine的资源模式，支持'1c12s','1c4s'
DEVICE 生成engine采用的设备id，默认是0号设备
示例
python3 -m build_engine.gcu --json_path alpaca_7b/alpaca_7b-kvcache-1of4-op13-fp32-N.json --device 0
输出文件

编译engine的工具会输出engine文件到onnx文件相同路径下，并且以bin扩展文件名命名。

比如，一个alpaca模型编译输出文件列表如下：

├── alpaca_7b-kvcache-1of4-op13-fp32-N.bin
├── alpaca_7b-kvcache-2of4-op13-fp32-N.bin
├── alpaca_7b-kvcache-3of4-op13-fp32-N.bin
└── alpaca_7b-kvcache-4of4-op13-fp32-N.bin

一组独立编译工具

GPT-XL
命令
pip3 install torch==1.13.0
python3 -m  build_scripts.gpt-xl --model_path gpt-xl_1of2-op13-fp32-N.onnx --batchsize=2
python3 -m  build_scripts.gpt-xl --model_path gpt-xl_2of2-op13-fp32-N.onnx --batchsize=2
Aquila-7B
命令
python3 -m  build_scripts.aquila_7b --model_path aquila1of4-op13-fp32-N.onnx --max_seq=2048 --batchsize=2
python3 -m  build_scripts.aquila_7b --model_path aquila2of4-op13-fp32-N.onnx --max_seq=2048 --batchsize=2
python3 -m  build_scripts.aquila_7b --model_path aquila3of4-op13-fp32-N.onnx --max_seq=2048 --batchsize=2
python3 -m  build_scripts.aquila_7b --model_path aquila4of4-op13-fp32-N.onnx --max_seq=2048 --batchsize=2
BLOOMZ-7B1
命令
python3 -m build_scripts.bloomz_7b1 --model_path ./bloomz-7b1-onnxs/ --stage_id 0 --max_bs 1 --max_seq 2048 --kvcache_dtype fp16
python3 -m build_scripts.bloomz_7b1 --model_path ./bloomz-7b1-onnxs/ --stage_id 1 --max_bs 1 --max_seq 2048 --kvcache_dtype fp16
python3 -m build_scripts.bloomz_7b1 --model_path ./bloomz-7b1-onnxs/ --stage_id 2 --max_bs 1 --max_seq 2048 --kvcache_dtype fp16
python3 -m build_scripts.bloomz_7b1 --model_path ./bloomz-7b1-onnxs/ --stage_id 3 --max_bs 1 --max_seq 2048 --kvcache_dtype fp16
BLOOM-7B1
命令
python3 -m build_scripts.bloom_7b1 --model_path ./bloom_7b1_onnxs/ --stage_id 0 --max_bs 1 --max_seq 2048 --kvcache_dtype fp16
python3 -m build_scripts.bloom_7b1 --model_path ./bloom_7b1_onnxs/ --stage_id 1 --max_bs 1 --max_seq 2048 --kvcache_dtype fp16
python3 -m build_scripts.bloom_7b1 --model_path ./bloom_7b1_onnxs/ --stage_id 2 --max_bs 1 --max_seq 2048 --kvcache_dtype fp16
python3 -m build_scripts.bloom_7b1 --model_path ./bloom_7b1_onnxs/ --stage_id 3 --max_bs 1 --max_seq 2048 --kvcache_dtype fp16
STARCODERBASE-15B
命令
python3 -m build_scripts.starcoderbase_15b --model_path ./starcoderbase-15b-onnxs/ --stage_id 0 --max_bs 1 --max_seq 1024 --kvcache_dtype fp16
python3 -m build_scripts.starcoderbase_15b --model_path ./starcoderbase-15b-onnxs/ --stage_id 1 --max_bs 1 --max_seq 1024 --kvcache_dtype fp16
python3 -m build_scripts.starcoderbase_15b --model_path ./starcoderbase-15b-onnxs/ --stage_id 2 --max_bs 1 --max_seq 1024 --kvcache_dtype fp16
python3 -m build_scripts.starcoderbase_15b --model_path ./starcoderbase-15b-onnxs/ --stage_id 3 --max_bs 1 --max_seq 1024 --kvcache_dtype fp16
python3 -m build_scripts.starcoderbase_15b --model_path ./starcoderbase-15b-onnxs/ --stage_id 4 --max_bs 1 --max_seq 1024 --kvcache_dtype fp16
python3 -m build_scripts.starcoderbase_15b --model_path ./starcoderbase-15b-onnxs/ --stage_id 5 --max_bs 1 --max_seq 1024 --kvcache_dtype fp16
python3 -m build_scripts.starcoderbase_15b --model_path ./starcoderbase-15b-onnxs/ --stage_id 6 --max_bs 1 --max_seq 1024 --kvcache_dtype fp16
python3 -m build_scripts.starcoderbase_15b --model_path ./starcoderbase-15b-onnxs/ --stage_id 7 --max_bs 1 --max_seq 1024 --kvcache_dtype fp16
LLAMA-7B
命令
python3 -m build_scripts.llama_7b --json_path ./build_configs/llama_7b_split_4 --onnx_path ./llama_7b_split_4
LLAMA2-7B
命令
python3 -m build_scripts.llama_7b --json_path ./build_configs/llama2_7b_split_4 --onnx_path ./llama2_7b_split_4
Baichuan-7B
命令
python3 -m build_scripts.baichuan_7b --model_path=./baichuan_7b --json_path=./build_configs/baichuan_7b_split_4_max_2048_mix16_3pg --max_seq=2048 --batch_size=1
LLAMA-13B
命令
python3 -m build_scripts.llama_13b --json_path ./build_configs/llama_13b_split_8 --onnx_path ./llama_13b_split_8
LLAMA2-13B
命令
python3 -m build_scripts.llama_13b --json_path ./build_configs/llama2_13b_split_8 --onnx_path ./llama2_13b_split_8
LLAMA2-13B-Chat
命令
python3 -m build_scripts.llama_13b --json_path ./build_configs/llama2_13b_chat_split_8 --onnx_path ./llama2_13b_chat_split_8
Baichuan2-7B
命令
python3 -m build_scripts.baichuan_7b --model_path=./baichuan2_7b --json_path=./build_configs/baichuan2_7b_split_4_max_2048_mix16_3pg --max_seq=2048 --batch_size=1

生成模型配置文件工具

TopsDistInfer需要每个模型提供model_configs文件。用户可以自行编写这些文件，也可以利用这个工具自动生成这些文件。

命令

python3 -m tools.config --json_path JSON --output MODEL_CONFIG_FILE --model_path MODEL_PATH

参数

参数 释义
JSON 切分工具生成的json文件路径
MODEL_CONFIG_FILE 输出model_config文件的路径
MODEL_PATH 用户希望model_config文件中显示在model_path的内容

示例

python3 -m tools.config --json_path llama-7b-kvcache-2of4-op13-fp32-N.json --output llama-7b-kvcachce-2of4-op13-fp32-N@1@.json --model_path "./llama_7b_split_4"

输出文件

自动产生的模型配置文件会在用户通过参数–output指定的路径上

模型参数量计算工具

这个工具提供计算huggingface的checkpoint的参数量的功能

命令

python3 -m tools.param_count --ckpt CKPT [--output OUTPUT]

参数

参数 释义
CKPT checkpoint的文件夹路径
OUTPUT 输出详细形状信息文件的路径

示例

python3 -m tools.param_count --ckpt gpt-2

模型切分和转换方法

切分工具可以对transformers的模型进行切分并转换成onnx模型。已发布的模型具体切分方法如下：

GPT2 xl

环境要求

  • transformers==4.31.0

  • torch==1.13.0

checkpoint

  • url: https://huggingface.co/gpt2-xl

  • commit id: 33cdb5c0db5423c1879b1b9f16c352988e8754a8

切分命令

deepspeed --num_cpu_cores 2 --module partition.transformers --model gpt2 --split_kvcache --ckpt <path/to/pytorch_model.bin> --config <path/to/config.json> --output_prefix gpt-xl_

LLAMA 7b

环境要求

  • transformers==4.32.0

  • torch==2.0.1

checkpoint

  • url: https://huggingface.co/decapoda-research/llama-7b-hf

  • commit id: 5f98eefcc80e437ef68d457ad7bf167c2c6a1348

切分命令

deepspeed --num_cpu_cores 4 --module partition.transformers --model llama --split_kvcache --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --output_prefix llama-7b-kvcache-

LLAMA 13b

环境要求

  • transformers==4.32.0

  • torch==2.0.1

checkpoint

  • url: https://huggingface.co/decapoda-research/llama-13b-hf

  • commit id: 438770a656712a5072229b62256521845d4de5ce

切分命令

deepspeed --num_cpu_cores 8 --module partition.transformers --model llama --split_kvcache --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --output_prefix llama-13b-kvcache-

LLAMA2 7b

环境要求

  • transformers==4.32.0

  • torch==2.0.1

checkpoint

  • url: https://huggingface.co/meta-llama/Llama-2-7b-hf

  • commit id: 3f025b

切分命令

deepspeed --num_cpu_cores 4 --module partition.transformers --model llama2 --split_kvcache --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --output_prefix llama2-7b-kvcache-

LLAMA2 13b

环境要求

  • transformers==4.32.0

  • torch==2.0.1

checkpoint

  • url: https://huggingface.co/meta-llama

  • commit id: 638c8be

切分命令

deepspeed --num_cpu_cores 8 --module partition.transformers --model llama2 --split_kvcache --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --output_prefix llama2-13b-kvcache-

LLAMA2 13b chat

环境要求

  • transformers==4.32.0

  • torch==2.0.1

checkpoint

  • url: https://huggingface.co/meta-llama

  • commit id: 13f8d72

切分命令

deepspeed --num_cpu_cores 8 --module partition.transformers --model llama2 --split_kvcache --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --output_prefix llama2-13b-chat-kvcache-

ChatGLM 6b

环境要求

  • transformers==4.31.0

  • torch==1.13.0

checkpoint

  • url: https://huggingface.co/THUDM/chatglm-6b

  • commit id: a70fe6b0a3cf1675b3aec07e3b7bb7a8ce73c6ae

切分命令

deepspeed --num_cpu_cores 4 --module partition.transformers --model chatglm --split_kvcache --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --output_prefix chatglm_6b-kvcache- --method type:GLMBlockPipeline

Vicuna 13b

环境要求

  • transformers==4.31.0

  • torch==2.0.1

checkpoint

  • url: https://huggingface.co/eachadea/vicuna-13b-1.1

  • commit id: bfcc6ca66694310be6c85ba0638597f4256c4143

切分命令

deepspeed --num_cpu_cores 8 --module partition.transformers --model llama --split_kvcache --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --output_prefix vicuna_13b-kvcache-

Alpaca 7b

环境要求

  • transformers==4.31.0

  • torch==2.0.1

checkpoint

  • url: https://huggingface.co/circulus/alpaca-7b

  • commit id: cd02033b82971a1dc6b6358129998a6c74c6d510

切分命令

deepspeed --num_cpu_cores 4 --module partition.transformers --model llama --split_kvcache --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --output_prefix alpaca_7b-kvcache-

BLOOMZ 7b1

环境要求

  • transformers==4.31.0

  • torch==1.13.0

checkpoint

  • url: https://huggingface.co/bigscience/bloomz-7b1

  • commit id: 2f4c4f3ebcf171dbbe2bae989ea2d2f3d3486a97

切分命令

deepspeed --num_cpu_cores 4 --module partition.transformers --model bloom --split_kvcache --ckpt <path/to/pytorch_model.bin> --config <path/to/config.json> --output_prefix bloomz_7b1-kvcache-

BLOOM 7b1

环境要求

  • transformers==4.30.2

  • torch==1.10.0

checkpoint

  • url: https://huggingface.co/bigscience/bloom-7b1

  • commit id: e83e90ba86f87f74aa2731cdab25ccf33976bd66

切分命令

deepspeed --num_cpu_cores 4 --module partition.transformers --model bloom --split_kvcache --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --output_prefix bloom_7b1-kvcache-

Baichuan 7b

环境要求

  • transformers==4.29.1

  • torch==2.0.0

checkpoint

  • url: https://huggingface.co/baichuan-inc/Baichuan-7B

  • commit id: c1a5c7d5b7f50ecc51bb0e08150a9f12e5656756

切分命令

deepspeed --num_cpu_cores 4 --module partition.transformers --model baichuan --method type:BaichuanDecoder --split_kvcache --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --output_prefix baichuan_7b-kvcache-

ChatGLM2

环境要求

  • transformers==4.29.1

  • torch==2.0.1+cpu

checkpoint

  • url: https://huggingface.co/THUDM/chatglm2-6b/tree/main

  • commit id: 31d45da2d8f14e55f459e15da7e5e57e32dd1e93

切分命令

deepspeed --num_cpu_cores 4 --module partition.transformers --model chatglm2 --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --method type:GLMBlock --split_kvcache --output_prefix chatglm2_6b-thudm-kvcache-auto- --output_suffix -op14-fp32-N

ChatGLM3

环境要求

  • transformers==4.29.1

  • torch==2.0.1+cpu

checkpoint

  • url: https://huggingface.co/THUDM/chatglm3-6b

  • commit id: e46a14881eae613281abbd266ee918e93a56018f

切分命令

deepspeed --num_cpu_cores 4 --module partition.transformers --model chatglm2 --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --method type:GLMBlock --split_kvcache --output_prefix chatglm3_6b-kvcache- --output_suffix -op14-fp32-N

OPT 13b

环境要求

  • transformers==4.29.1

  • torch==2.0.1

checkpoint

  • url: https://huggingface.co/facebook/opt-13b

  • commit id: e515202d1e7750da62d245fbccb2723b9c1790f5

切分命令

deepspeed --num_cpu_cores 8 --module partition.transformers --model opt --split_kvcache --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --output_prefix opt_13b-kvcache-

STARCODERBASE-15B

环境要求

  • transformers==4.33.2

  • torch==1.10.0

checkpoint

  • url: https://huggingface.co/bigcode/starcoderbase

  • commit id: 88ec5781ad071a9d9e925cd28f327dea22eb5188

切分命令

deepspeed --num_cpu_cores 8 --module partition.transformers --model gptbigcode --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --split_kvcache --output_prefix starcoderbase_15b-kvcache-

Pangu 13b

环境要求

  • transformers==4.29.1

  • torch==2.0.1

checkpoint

  • url: https://huggingface.co/sunzeyeah/pangu-13B

  • commit id: 6b99d82d7dfc014e9ca6c5039fa8b16d24491ce9

切分命令

deepspeed --num_cpu_cores 8 --module partition.transformers --model pangu --split_kvcache --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --output_prefix pangu_13b-kvcache-

Baichuan2 13b base

环境要求

  • transformers==4.29.2

  • torch==2.2.0

checkpoint

  • url: https://huggingface.co/baichuan-inc/Baichuan2-13B-Base

  • commit id: 4224a07a96905440a3f60fe0def0b6b425303fc9

替换权重

cd /usr/local/gcu/TopsDistInfer/src/pre_process_Baichuan2_13B/
python3 genration.py --path <dir of Baichuan2-13B-Base pretrianed files>

切分命令

deepspeed --num_cpu_cores 8 --module partition.transformers --model baichuan2 --method type:Baichuan2Decoder --split_kvcache --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --output_prefix baichuan2_13b-kvcache- --output_suffix -op14-fp32-N

Baichuan2 7b

环境要求

  • transformers==4.29.2

  • torch==2.0.0

checkpoint

  • url: https://huggingface.co/baichuan-inc/Baichuan2-7B-Base/tree/main

  • commit id: 0cc6a61c06cd0734270151109d07cf86ef0ace53

切分命令

deepspeed --num_cpu_cores 4 --module partition.transformers --model baichuan2_7b --method type:Baichuan2_7BDecoder --split_kvcache --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --output_prefix baichuan2_7b-kvcache-

WIZARDCODER-15B

环境要求

  • transformers==4.33.2

  • torch==2.0.1

checkpoint

  • url: https://huggingface.co/WizardLM/WizardCoder-15B-V1.0/tree/main

  • commit id: 63f325b81dd94805b8a64fa7fd1a7bb95d76efe6

切分命令

deepspeed --num_cpu_cores 8 --module partition.transformers --model gptbigcode --ckpt <path/to/pytorch_model.bin> --config <path/to/config.json> --split_kvcache --output_prefix wizardcoder_15b-kvcache-

Baichuan2 13b chat

环境要求

  • transformers==4.29.2

  • torch==2.2.0

checkpoint

  • url: https://huggingface.co/baichuan-inc/Baichuan2-13B-Chat

  • commit id: 4224a07a96905440a3f60fe0def0b6b425303fc9

替换权重

cd /usr/local/gcu/TopsDistInfer/src/pre_process_Baichuan2_13B/
python3 genration.py --path <dir of Baichuan2-13B-Chat pretrianed files>

切分命令

deepspeed --num_cpu_cores 8 --module partition.transformers --model baichuan2 --method type:Baichuan2Decoder --split_kvcache --ckpt <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --output_prefix baichuan2_13b_chat-kvcache- --output_suffix -op14-fp32-N

Aquila 7b

环境要求

  • transformers==4.35.2

  • torch==1.10.0

checkpoint

  • url: https://huggingface.co/BAAI/Aquila-7B/tree/2a5a253870c87efb4bef045081f747d8b4afa80e

  • commit id: 2a5a253870c87efb4bef045081f747d8b4afa80e

切分命令

deepspeed --num_cpu_cores 4 --module partition.transformers --model aquila --ckpt  <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --method parameters --split_kvcache --output_prefix aquila

InternLM 7B

环境要求

  • transformers==4.35.2

  • torch==1.10.0

checkpoint

  • url: https://huggingface.co/internlm/internlm-7b

  • commit id: 2667fa759d51a1c572ae54401f127763bf943f45

切分命令

deepspeed --num_cpu_cores 4 --module partition.transformers --model internlm --ckpt  <path/to/pytorch_model.bin.index.json> --config <path/to/config.json> --method parameters --split_kvcache --output_prefix internlm