EAGLE: 一种用于加速LLM解码的方法

EAGLE（提高语言模型效率的外推算法）是快速解码大型语言模型 (LLM) 的新基准，具有可证明的性能维护。这种方法涉及外推法学硕士的第二顶层上下文特征向量，从而显着提高生成效率。在理论驱动下（敬请关注我们即将发表的论文），EAGLE 建立在以下第一原则的基础上：

LLM 特征向量序列可随时间压缩，从而可以轻松地根据先前特征向量预测后续特征向量。

• EAGLE 是：
  • 比普通解码 (13B) 快 3 倍。
  • 比 Lookahead (13B) 快 2 倍。
  • 比 Medusa (13B) 快 1.6 倍。
  • 可证明在生成文本的分布中保持与普通解码的一致性。
  • 可训练（1-2 天内）并可在 8x RTX 3090 GPU 上进行测试。因此，即使 GPU 贫穷的人也能负担得起。
  • 可与其他并行技术相结合，例如 vLLM、Mamba、FlashAttention、量化和硬件优化。

使用 Vicuna 33B 模型在 RTX 3090 GPU 上以 fp16 精度进行推理。为了增强观看体验，动画速度加快了四倍。

安装

pip install -r requirements.txt

推理

我们提供的推理代码会自动分配模型权重（跨多个 GPU 加载模型），允许您运行超出单个 GPU 内存的模型。

带用户界面

我们提供了建议的 Web 界面，您可以通过运行以下命令来使用该界面。模型完全加载后，终端中会输出一个URL，您可以在浏览器中输入该URL来访问。

python -m application.webui --ea-model-path [path of EAGLE weight]\     --base-model-path [path of the original model]\    --model-type [vicuna or llama-2-chat]

有代码

您可以使用我们提供的“egenerate”来加速生成，就像使用 Hugging Face 中的“generate”一样。这是一个例子。

from model.ea_model import EaModelmodel = EaModel.from_pretrained(      base_model_path=base_model_path,      ea_model_path=EAGLE_model_path,      torch_dtype=torch.float16,      low_cpu_mem_usage=True,      device_map="auto"  )prompt="Hello"input_ids=model.tokenizer([prompt]).input_idsinput_ids = torch.as_tensor(input_ids).cuda()output_ids=model.eagenerate(input_ids,temperature=0.5,max_new_tokens=512)output=model.tokenizer.decode(output_ids[0])