推测解码

介绍

ModelOpt的推测解码模块（modelopt.torch.speculative）使您的模型能够在每个生成步骤中生成多个标记。这对于减少模型的延迟和加速推理非常有用。

以下是ModelOpt支持的推测解码算法： - Medusa - EAGLE

按照以下步骤使用ModelOpt的推测解码模块modelopt.torch.speculative来获取带有Medusa或EAGLE推测解码的模型：

通过 mtsp.convert 转换你的模型：向你的模型添加Medusa头或EAGLE模块。
微调Medusa头或EAGLE模块: 微调Medusa头或EAGLE模块。建议冻结基础模型。
检查点并重新加载：通过mto.save保存模型，并通过mto.restore恢复

转换

您可以使用mtsp.convert()将您的模型转换为推测解码模型。

示例用法：

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import modelopt.torch.speculative as mtsp

# User-defined model
model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
tokenizer.pad_token_id = tokenizer.eos_token_id

if mode == "medusa":
    # Configure and convert to medusa
    config = {
        "medusa_num_heads": 2,
        "medusa_num_layers": 1,
        }
elif mode == "eagle":
    config = {
        "eagle_num_layers": 1
        }
mtsp.convert(model, [(mode, config)])

微调推测解码模型并存储/恢复模型

转换为推测解码模型后，您需要对解码模块进行微调：

import os
from transformers import Trainer
import modelopt.torch.opt as mto

mto.enable_huggingface_checkpointing()

trainer = Trainer(model=model, tokenizer=tokenizer, args=training_args, **data_module)
trainer._move_model_to_device(model, trainer.args.device)

trainer.train(resume_from_checkpoint=checkpoint)
trainer.save_state()
trainer.save_model("<path to the output directory>")

要恢复保存的推测解码模型：

model = AutoModelForCausalLM.from_pretrained("<path to the output directory>")

推测解码概念

下面，我们将概述ModelOpt的推测解码功能及其基本概念和术语。

推测性解码

从语言模型生成文本的标准方法是使用自回归解码：每一步生成一个标记，并将其附加到输入上下文中以生成下一个标记。这意味着生成K个标记将需要模型的K次串行运行。像Transformers这样的大型自回归模型的推理可能既慢又昂贵。因此，已经提出了各种推测解码算法来加速文本生成，特别是在对延迟敏感的应用中。

通常，使用一个更快的自回归模型（称为草稿模型）生成长度为K的简短草稿。这可以通过并行模型或调用草稿模型K次来实现。然后，使用一个更大、更强大的模型（称为目标模型）来对草稿进行评分。最后，使用采样方案来决定目标模型接受哪个草稿，从而在过程中恢复目标模型的分布。

Medusa算法

有许多方法可以实现推测性解码。一种流行的方法是Medusa，它不使用额外的草稿模型，而是引入一些额外的解码头来同时预测多个未来的标记。在生成过程中，这些头各自为相应的位置生成多个可能的单词。然后使用基于树的注意力机制将这些选项组合和处理。最后，采用典型的接受方案从候选者中选择最长的合理前缀以进行进一步解码。由于草稿模型本身就是目标模型，这保证了输出分布与目标模型的分布相同。

EAGLE算法

与Medusa基于基础模型的隐藏状态预测未来令牌不同，EAGLE通过一个轻量级的自回归解码器预测未来的隐藏状态，然后用于预测未来的令牌。由于在特征（隐藏状态）级别上的自回归比在令牌级别上更简单，EAGLE可以比Medusa更准确地预测未来的令牌。因此，它实现了更高的加速。