| 博客 | 报告 | 路线图 |

---

新闻 🔥

• [2024/1] Medusa技术报告现已在arXiv上发布。我们新增了多项功能，包括用于全模型训练的Medusa-2配方、用于将Medusa添加到任何微调LLM中的自蒸馏等。新结果表明，在一系列LLM上，相比原始模型，速度提升了2.2至3.6倍。

---

简介

Medusa是一个简单的框架，旨在普及使用多解码头加速LLM生成的技术。

我们致力于解决诸如推测性解码等流行加速技术的三大痛点：

• 需要一个优秀的草稿模型。
• 系统复杂度高。
• 在基于采样的生成过程中效率低下。

我们希望通过以下思路来解决推测性解码所面临的挑战：

• 不引入新模型，而是在同一模型上训练多个解码头。
• 训练过程参数高效，即使是资源有限的用户也能完成。由于无需额外模型，因此也无需调整分布式计算设置。
• 放宽对与原始模型分布匹配的要求，使得非贪婪生成的速度甚至超过贪婪解码。

在初始版本中，我们的主要关注点是优化适用于批大小为1的Medusa——这一设置常用于本地模型部署。在此配置下，Medusa在一系列Vicuna模型上可实现约2倍的速度提升。我们正积极努力将Medusa集成到更多推理框架中，以进一步提升性能，并将其扩展到更广泛的场景。

在更新版本中，我们增加了对全模型训练的支持，称为Medusa-2（相对于仅训练新增头部的Medusa-1），这需要一种特殊的配方，能够在保留原始模型性能的同时，赋予其推测性预测能力。

我们还新增了自蒸馏支持，使我们能够将Medusa添加到任何微调过的LLM中，而无需原始训练数据。

目录

• 简介
• 目录
• 安装
  • 方法1：通过pip（可能不是最新版本）
  • 方法2：从源代码安装（推荐）
  • 模型权重
  • 推理
  • 训练
  • 传统训练方式
  • 推送到Hugging Face Hub
• 引用
• 代码库指南
• 社区采纳
• 贡献
• 致谢

安装

方法1：通过pip（可能不是最新版本）

pip install medusa-llm

方法2：从源代码安装（推荐）

git clone https://github.com/FasterDecoding/Medusa.git
cd Medusa
pip install -e .

模型权重

Medusa-1

规模	对话命令	Hugging Face仓库
7B	python -m medusa.inference.cli --model FasterDecoding/medusa-vicuna-7b-v1.3	FasterDecoding/medusa-vicuna-7b-v1.3
13B	python -m medusa.inference.cli --model FasterDecoding/medusa-vicuna-13b-v1.3	FasterDecoding/medusa-vicuna-13b-v1.3
33B	python -m medusa.inference.cli --model FasterDecoding/medusa-vicuna-33b-v1.3	FasterDecoding/medusa-vicuna-33b-v1.3

Medusa-2

规模	对话命令	Hugging Face仓库
Zephyr-7B-Beta	python -m medusa.inference.cli --model FasterDecoding/medusa-1.0-zephyr-7b-beta	FasterDecoding/medusa-1.0-zephyr-7b-beta
Vicuna-7B-v1.5	python -m medusa.inference.cli --model FasterDecoding/medusa-1.0-vicuna-7b-v1.5	FasterDecoding/medusa-1.0-vicuna-7b-v1.5
Vicuna-13B-v1.5	python -m medusa.inference.cli --model FasterDecoding/medusa-1.0-vicuna-13b-v1.5	FasterDecoding/medusa-1.0-vicuna-13b-v1.5
Vicuna-33B-v1.5	python -m medusa.inference.cli --model FasterDecoding/medusa-1.0-vicuna-33b-v1.5	FasterDecoding/medusa-1.0-vicuna-33b-v1.5

推理

目前我们支持单 GPU 推理，批次大小为 1，这是本地模型托管中最常见的配置。我们正在积极努力将 Medusa 集成到其他推理框架中，以扩展其功能；如果您有意参与这一工作，请随时与我们联系。

您可以使用以下命令启动 CLI 界面：

CUDA_VISIBLE_DEVICES=0 python -m medusa.inference.cli --model [Medusa 模型路径]

您还可以传递 --load-in-8bit 或 --load-in-4bit 参数，以量化格式加载基础模型。如果您在其他地方下载了基础模型，可以使用 --base-model [基础模型路径] 来覆盖基础模型的名称或路径。

训练

在更新版本中，我们使用强大的 axolotl 库来管理训练过程。请参阅我们的 fork 获取训练代码。主要的代码修改位于 src/axolotl/utils/models.py。训练配置可在 examples/medusa 中找到。典型的训练命令如下：

accelerate launch -m axolotl.cli.train examples/medusa/your_config.yml

用于自蒸馏的数据准备代码可在当前仓库的 data_generation 文件夹 中找到。对于其他数据集，您可以直接从相应的 Hugging Face 数据集仓库下载数据。

在不同架构上训练

以下说明适用于 Medusa 的初始版本，提供了一个训练 Medusa-1 模型的最小示例。有关更新版本的说明，请参阅前一节。

进行训练时，请安装：

pip install -e ".[train]"

准备数据

我们采用 ShareGPT 数据集的一个公开版本，它是 Vicuna 训练数据的一个子集。对于其他模型，您可以使用相应的训练数据集。

git clone https://huggingface.co/datasets/Aeala/ShareGPT_Vicuna_unfiltered

注意：如果您尚未安装 git-lfs，请在克隆之前先安装：

git lfs install

将数据适配到您希望启用 Medusa 的模型上。

首先启动一个您喜欢的推理服务器，运行您想要训练的模型。我们以 mistralai/Mistral-7B-Instruct-v0.2 为例。

例如，您可以使用 text-generation-inference，该工具在您训练完 Medusa 头部后也可以继续使用。

model=mistralai/Mistral-7B-Instruct-v0.2
volume=$PWD/data # 与 Docker 容器共享卷，避免每次运行都重新下载权重
docker run --gpus all --shm-size 1g -p 8080:80 -v $volume:/data ghcr.io/huggingface/text-generation-inference:latest --model-id $model --input-length 4000 --max-total-tokens 4096 --max-batch-prefill-tokens 4000

ShareGPT 中的部分序列较长，因此请确保能够成功推理这些序列。如果内存不足，脚本会自动跳过这些长对话。这不会对下游性能产生太大影响，但更多的数据总是更好的。

您可以根据需要调整各种参数以 加速推理，不过默认设置在大多数情况下已经足够好。

python create_data.py --input-filename ShareGPT_Vicuna_unfiltered/ShareGPT_V4.3_unfiltered_cleaned_split.json --output-filename mistral.json

训练模型

我们沿用了 FastChat 的训练设置，但由于我们冻结了原始模型，只训练新增的头部，因此学习率设置得更高。以下是针对 Vicuna-7b 模型在 4 张 GPU 上的训练命令。由于我们仅训练新增的头部，训练所需的显存较少，只需使用数据并行即可。您可以根据自己的硬件配置调整脚本。对于更大的模型，我们也采用相同的设置。此外，您还可以使用 --load_in_8bit 或 --load_in_4bit 参数，以量化格式加载基础模型。

torchrun --nproc_per_node=4 medusa/train/train_legacy.py --model_name_or_path mistralai/Mistral-7B-Instruct-v0.2 \
    --data_path mistral.json \
    --bf16 True \
    --output_dir test \
    --num_train_epochs 2 \
    --per_device_train_batch_size 8 \
    --per_device_eval_batch_size 8 \
    --gradient_accumulation_steps 4 \
    --evaluation_strategy "no" \
    --save_strategy "no" \
    --learning_rate 1e-3 \
    --weight_decay 0.0 \
    --warmup_ratio 0.1 \
    --lr_scheduler_type "cosine" \
    --logging_steps 1 \
    --tf32 True \
    --model_max_length 2048 \
    --lazy_preprocess True \
    --medusa_num_heads 3 \
    --medusa_num_layers 1 \
    --deepspeed deepspeed.json

推送到 Hugging Face Hub

您可以使用以下命令将您的模型推送到 Hugging Face Hub：

python -m medusa.hf_utils --folder [模型文件夹路径] --repo [仓库名称]

引用

@article{cai2024medusa,
  title   = {Medusa: 基于多解码头的简单 LLM 推理加速框架},
  author  = {Tianle Cai、Yuhong Li、Zhengyang Geng、Hongwu Peng、Jason D. Lee、Deming Chen、Tri Dao},
  year    = {2024},
  journal = {arXiv 预印本 arXiv: 2401.10774}
}

代码库指南

medusa/model/medusa_model.py 是 Medusa 的核心文件。它包含了 MedusaModel 类，该类是对原始模型和新增头部的封装。此外，该类还实现了流式生成方法。如果您想深入了解 Medusa 的细节，这里就是起点。

我们还在 notebooks/ 目录中提供了一些示例笔记本，帮助您更好地理解代码库。

社区采纳

我们非常高兴地看到 Medusa 已被许多开源项目所采用。以下是一个不完全列表：

• TensorRT-LLM
• TGI
• RTP-LLM

我们感谢这些作者对社区的贡献，并真诚希望 Medusa 能够助力 LLM 的发展。如果您在项目中使用了 Medusa，请告知我们，我们将把您的项目加入到列表中。

贡献

我们欢迎社区为 Medusa 做出贡献。如果您有任何改进建议，请提交一个议题与我们讨论。在提交拉取请求时，请确保您的更改经过充分测试。请将每一项重大更改拆分为单独的拉取请求。我们还有一个 路线图，概述了 Medusa 未来的规划。如果您对参与路线图中的任何事项感兴趣，请随时联系我们。

致谢

本代码库受到大语言模型社区中一些杰出项目的启发，包括 FastChat、TinyChat、vllm 和 axolotl。

本项目得到了 Together AI、MyShell AI 和 Chai AI 的支持。

Medusa 快速上手指南

Medusa 是一个简单的框架，通过在大语言模型（LLM）上添加多个解码头（Decoding Heads），实现生成速度的显著加速（2.2-3.6 倍）。它无需额外的草稿模型，仅微调新增的头部即可工作，支持自蒸馏技术适配任意微调后的 LLM。

环境准备

• 操作系统: Linux (推荐 Ubuntu)
• Python: 3.8 或更高版本
• GPU: 支持 CUDA 的 NVIDIA 显卡（单卡即可运行推理）
• 依赖库: PyTorch, Transformers, Accelerate 等（安装脚本会自动处理）

提示：国内用户建议在安装 Python 依赖时配置清华或阿里镜像源，以加快下载速度。

export PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

推荐使用源码安装方式，以获取最新功能（包括 Medusa-2 和自蒸馏支持）。

1. 克隆仓库并安装

git clone https://github.com/FasterDecoding/Medusa.git
cd Medusa
pip install -e .

2. 训练环境（可选）

如果您计划自己训练 Medusa 头部（特别是使用 Medusa-2 方案），需要安装额外的训练依赖：

pip install -e ".[train]"

注：Medusa-2 的训练基于 axolotl 库，具体配置需参考其 fork 版本。

基本使用

Medusa 目前主要优化了 Batch Size 为 1 的单卡推理场景，非常适合本地部署。

1. 启动命令行交互界面

您可以直接加载 Hugging Face 上预训练好的 Medusa 模型进行对话。以下以 Vicuna-7B 的 Medusa 版本为例：

CUDA_VISIBLE_DEVICES=0 python -m medusa.inference.cli --model FasterDecoding/medusa-vicuna-7b-v1.3

常用参数说明：

• --model: 指定 Medusa 模型的路径或 Hugging Face 仓库名。
• --base-model: 如果基础模型不在默认路径，可手动指定基础模型路径。
• --load-in-8bit 或 --load-in-4bit: 启用量化加载，节省显存（例如：--load-in-4bit）。

2. 可用模型列表

您可以替换上述命令中的 --model 参数为以下官方提供的模型：

Medusa-1 系列:

• 7B: FasterDecoding/medusa-vicuna-7b-v1.3
• 13B: FasterDecoding/medusa-vicuna-13b-v1.3
• 33B: FasterDecoding/medusa-vicuna-33b-v1.3

Medusa-2 系列 (支持更广泛的架构):

• Zephyr-7B: FasterDecoding/medusa-1.0-zephyr-7b-beta
• Vicuna-7B-v1.5: FasterDecoding/medusa-1.0-vicuna-7b-v1.5
• Vicuna-13B-v1.5: FasterDecoding/medusa-1.0-vicuna-13b-v1.5
• Vicuna-33B-v1.5: FasterDecoding/medusa-1.0-vicuna-33b-v1.5

3. 将模型上传至 Hugging Face Hub (可选)

如果您训练了自己的模型并希望分享，可以使用以下命令：

python -m medusa.hf_utils --folder [本地模型文件夹路径] --repo [目标仓库名称]