# Medusa Decoding

本文说明如何使用vllm构建和运行medusa模型，目前medusa支持tree-style generation，target model和draft model均可多卡推理

## Overview
与其他模型不同，medusa解码需要一个base model和若干Medusa heads.

Vllm medusa model的实现在[vllm/model_executor/models/medusa.py]

## Support Matrix
  * FP16
  * BF16
  * PAGED_KV_CACHE
  * Tensor Parallel

### convert Medusa model weights
# medusa 模型需要转换为vllm中Medusa的模型格式

```bash
python medusa_weight_converter.py --medusa_num_heads 4 --medusa_num_layers 1 --medusa_model_path /work/medusa/qwen2_72b_head_4/adapter_model.bin --vocab_size 152064 --hidden_size 8192 --output_dir /work/medusa/sugon/vllm-medusa-qwen2-72b-head-4 --medusa_choices="[(0,), (0, 0), (0, 0, 0), (0, 0, 0, 0), (0, 1), (1,), (1, 0), (0, 0, 1), (0, 1, 0), (0, 2), (1, 0, 0), (2,), (2, 0), (0, 3), (0, 0, 2), (0, 2, 0), (0, 4), (0, 0, 1, 0), (0, 1, 0, 0), (2, 0, 0), (3,), (0, 5), (0, 0, 0, 1), (3, 0), (0, 0, 3), (1, 0, 0, 0), (0, 3, 0), (0, 6), (0, 0, 4), (0, 4, 0), (1, 1), (4,)]"
```
此处qwen2_72b_head_4是medusa模型使用peft lora训练后保存的权重，其他格式也可参考[medusa_weight_converter.py]修改进行权重转换


### Run

```bash
python3 -m vllm.entrypoints.openai.api_server \
  --served-model-name qwen_medusa \
  --model /models/Qwen2-72B-Instruct/ -tp 4 \
  --max-model-len 1024 --max-num-seqs 8 --gpu-memory-utilization 0.7 \
  --speculative-model /work/medusa/vllm-medusa-qwen2-72b-head-4 \
  --speculative-draft-tensor-parallel-size 4 \
  --speculative-disable-by-batch-size 4 \
  --use-v2-block-manager \
  --spec-decoding-acceptance-method typical_acceptance_sampler \
  --enforce-eager --dtype float16 --trust-remote-code --port 8086\
  --enable-lora --lora-modules medusa-lora=/work/qwen2_72b_head_4 \
  --max-lora-rank 32 --lora-extra-vocab-size 0 --merge-lora True \
  --lora-target-modules qkv_proj \
  --tree-style-spec-decoding True\
  --num-speculative-heads 4 --num-speculative-tokens 33
```
merge-lora可以将lora权重和base model权重融合，提升整体推理速度，若对精度有严格要求，可不设置此参数
num-speculative-tokens和medusa choices的个数相关，num_speculative_tokens = len(medusa_choices) + 2

# do request
```bash
curl http://localhost:8086/v1/completions \
-H "Content-Type: application/json" \
-d '{
"model": "medusa-lora",
"prompt": "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n帮我写一个C++的快速排序算法<|im_end|>\n<|im_start|>assistant\n",
"max_tokens": 256,
"temperature": 0.0
}'
```bash