diff --git a/README.md b/README.md index d57dacb3..1bf63d62 100644 --- a/README.md +++ b/README.md @@ -155,7 +155,7 @@ To start a server, you can use the following script. Before running, please set `CHECKPOINT_PATH` to the root folder path of the checkpoint, and set `TP` and `PP` appropriately according to the parallelism degrees of the checkpoint. ```shell -./example/ee_inference/ee_inference_server.sh +./examples/ee_inference/ee_inference_server.sh ``` After the server is started, you can use `tools/request_client.py` to send requests to the server. @@ -171,7 +171,14 @@ Below are some parameters for early-exit LLM inference, which can be found in `t ## Checkpoints -The model checkpoints mentioned in our paper will be released soon. +The model checkpoints used in our [EE-LLM paper](https://arxiv.org/abs/2312.04916) have been released on ModelScope: + +- 1.3B model with two early exits at Layer 6 and 12. [[link]](https://modelscope.cn/models/Data-Juicer/EE-LLM-1B-dj-refine-300B) +- 7B model with two early exits at Layer 8 and 16. [[link]](https://modelscope.cn/models/Data-Juicer/EE-LLM-7B-dj-refine-150B) + +> The provided checkpoints have a pipeline parallel size of 4 (PP=4) and a tensor parallel size of 1 (TP=1), +> please set those values properly in corresponding scripts. +> For other parallel degrees, you can use `./tools/convert_parallelism.sh` to convert the checkpoints. ## BibTeX diff --git a/tools/convert_parallelism.sh b/tools/convert_parallelism.sh new file mode 100644 index 00000000..f911f402 --- /dev/null +++ b/tools/convert_parallelism.sh @@ -0,0 +1,15 @@ +#!/bin/bash + +# load/save dir +LOAD_DIR= +SAVE_DIR= + +# target parallelism +TP= +PP= + +CUR_DIR=$(cd $(dirname "$0") && pwd) +MEGATRON_ROOT_PATH=$(cd "$CUR_DIR/.." && pwd) +cd $MEGATRON_ROOT_PATH + +python $MEGATRON_ROOT_PATH/tools/checkpoint/util.py --model-type EarlyExitGPT --load-dir $LOAD_DIR --save-dir $SAVE_DIR --target-tensor-parallel-size $TP --target-pipeline-parallel-size $PP --megatron-path $MEGATRON_ROOT_PATH \ No newline at end of file