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])

项目链接

https://github.com/SafeAILab/EAGLE

未经允许不得转载:表盘吧 » EAGLE: 一种用于加速LLM解码的方法