GPTFast

使用GPTFast将您的Hugging Face Transformers加速7.6-9倍！

背景

GPTFast最初是由PyTorch团队开发的一系列技术，旨在加速Llama-2-7b的推理速度。这个pip包将这些技术推广到所有Hugging Face模型。

演示

GPTFast推理时间	常规推理时间

路线图

⟳ 0.7.x (xx/xx/xx): Medusa、推测采样、Eagle
⟳ 0.6.x (xx/xx/xx): BitNet和1位量化、AWQ、QoQ、GGUF、HQQ
⟳ 0.5.x (xx/xx/xx): PagedAttention (vLLM) + FlashAttention集成
⟳ 0.4.x (xx/xx/xx): 张量并行 + GPU分布式推理
✅ 0.3.x (2024/06/20): 为所有HF模型启用GPTQ int4量化和优化的int4矩阵乘法内核（推理加速9倍）
✅ 0.2.x (2024/04/02): 为所有HF模型启用静态键值缓存（推理加速8.5倍）
✅ 0.1.x (2024/02/22): torch.compile、int8量化、推测解码（推理加速7倍）

入门指南

警告：以下文档在0.3.0版本中已过时。新文档即将发布！##

确保您的Python版本 >= 3.10，并且您使用的是支持CUDA的设备。

在您的机器上创建一个虚拟环境并激活它。

$python3 -m venv VENV_NAME
source VENV_NAME/bin/activate #如果您使用Windows，请使用 ./VENV_NAME/scripts/activate

执行以下命令：pip install gptfast

将以下代码复制到一个Python文件中：

  import os
  import torch
  from transformers import AutoTokenizer
  from GPTFast.Core import gpt_fast
  from GPTFast.Helpers import timed
  
  torch._dynamo.reset()
  os.environ["TOKENIZERS_PARALLELISM"] = "false"
  
  device = "cuda" if torch.cuda.is_available() else "cpu"
  
  def argmax_variation(self, probabilities:torch.Tensor, temperature:float = 1, k:int = 5):
      # 应用温度缩放
      device = probabilities.device
      scaled_probabilities = probabilities / temperature
  
      # 确保k在有效范围内
      k = min(k, probabilities.size(-1))
  
      # 获取指定维度上top-k缩放概率的索引
      top_k_indices = torch.topk(scaled_probabilities, k, dim=-1).indices
  
      # 生成用于采样的随机索引
      random_indices = torch.randint(0, k, (1,) * probabilities.dim()).to(device)
  
      # 使用收集的索引获取最终采样的token
      sampled_token = top_k_indices.gather(-1, random_indices).to(device)
  
      return sampled_token.unsqueeze(0)
  
  def argmax(self, probabilities):
      # 使用argmax获取概率最大的token
      max_prob_index = torch.argmax(probabilities, dim=-1)
      return max_prob_index.view(1, 1)
  
  model_name = "gpt2-xl"
  draft_model_name = "gpt2"
  
  tokenizer = AutoTokenizer.from_pretrained(model_name)
  initial_string = "Write me a short story."
  input_tokens = tokenizer.encode(initial_string, return_tensors="pt").to(device)
  
  N_ITERS=10
  MAX_TOKENS=50
  
  cache_config = {
      "model_config": {
          "path_to_blocks": ["transformer", "h"],
          "child_ref_in_parent_forward": ["transformer", "block"],
      },
      "block_config": {
          "path_to_attn": ["attn"],
          "child_ref_in_parent_forward": ["attn"], 
      },
      "attn_config": {
          "cache_update_config":{
              "kv_cache_condition":"if layer_past is not None",
              "key_name": "key",
              "value_name": "value",
          },
          "causal_mask_config": {
              "causal_mask_application": "conditional",
              "causal_mask_method": "_attn",
              "causal_mask_condition": "not self.is_cross_attention"
          }
      },
      "imports": ["import torch", 
                  "import transformers", 
                  "from transformers import *", 
                  "from torch import *", 
                  "from typing import *", 
                  "import types", 
                  "from transformers.modeling_outputs import *", 
                  "from torch import nn"]
  }
  
  gpt_fast_model = gpt_fast(model_name, sample_function=argmax, max_length=60, cache_config=cache_config, draft_model_name=draft_model_name)
  gpt_fast_model.to(device)
  
  fast_compile_times = []
  for i in range(N_ITERS):
      with torch.no_grad():
          res, compile_time = timed(lambda: gpt_fast_model.generate(cur_tokens=input_tokens, max_tokens=MAX_TOKENS, speculate_k=6))
      fast_compile_times.append(compile_time)
      print(f"gpt fast 评估时间 {i}: {compile_time}")
  print("~" * 10)
  
  print(tokenizer.decode(res[0]))

运行它，见证奇迹的时刻到了🪄！

文档

这个库的核心提供了一个简单的接口来实现LLM推理加速技术。以下所有函数都可以从GPTFast.Core中导入：

gpt_fast(model_name:str, sample_function:Callable[torch.Tensor, Dict[str, Any], torch.Tensor], max_length:int, cache_config:dict, draft_model_name:str) -> torch.nn.Module
- 参数:
  - model_name: 这是您想要优化的Hugging Face模型的名称。
  - sample_function: 这是一个函数，它将接收一个PyTorch张量作为第一个参数，以及其他**sampling_kwargs，并返回一个形状为(1, 1)的张量。
  - max_length: 这是一个整数，指定您将生成多少个token。建议将此值设置得比实际生成的token数量高。
  - cache_config: 这是一个字典，用于指定如何将静态键值缓存集成到模型中。此字典的更多详细信息如下。
draft_model_name：这是一个可选参数，是用于推测性解码所需的Hugging Face草稿模型的名称。请注意，模型和草稿模型必须使用相同的分词器，且草稿模型必须显著小于主模型，以实现推理加速。如果未指定draft_model_name，则不会对您的模型应用推测性解码。
返回：
- 一个加速的模型，具有一个方法：
  - generate(self, cur_tokens:torch.Tensor, max_tokens:int, speculate_k:int, **sampling_kwargs) -> torch.Tensor
    - 参数：
      - cur_tokens：大小为(1, seq_len)的PyTorch张量。
      - max_tokens：一个整数，表示您想要生成的令牌数量。
      - speculate_k：一个整数，指定在推测性解码中您希望草稿模型推测的程度。
      - **sampling_kwargs：从分布中采样所需的额外参数。应与上面的sample_function的**sampling_kwargs匹配。
    - 返回：
      - 对您的提示生成的令牌，一个维度为(1, max_tokens)的张量。

load_int8(model_name:str) -> torch.nn.Module
- 参数：
  - model_name：这是一个字符串，指定您正在使用的模型。
- 返回：
  - 您的模型的int8量化版本。

add_kv_cache(transformer:nn.Module, sampling_fn:Callable[torch.Tensor, Dict[str, Any], torch.Tensor], max_length:int, cache_config:dict) -> KVCacheModel

参数：

transformer：这是您要添加静态键值缓存的Hugging Face模型。
sampling_fn：这与gpt_fast函数的sampling_function参数相同。
max_length：这与gpt_fast函数的max_length参数相同。

cache_config：这是一个字典，用于指定如何直接修改模型前向传播的源代码以适应静态缓存。该字典的完整规范如下：

 -model_config: 定义如何修改您的模型以适应静态键值缓存。
    -path_to_blocks (list[str]): 从模型本身开始，定义在父```nn.Module```属性/对象上访问transformer块的子属性路径。
    -child_ref_in_parent_forward (list[str]): 从原始模型开始，这是```path_to_blocks```中每个子模块/属性在父模块/属性的前向传播中的引用方式。

-block_config: 定义如何修改您的块以适应静态键值缓存。
    -path_to_attn (list[str]): 从块本身开始，定义在父```nn.Module```属性/对象上访问注意力层本身的子属性路径。
    -child_ref_in_parent_forward (list[str]): 从块开始，这是path_to_attn中每个子模块/属性在父模块/属性的前向传播中的引用方式。

-attn_config: 定义如何修改注意力层以适应静态键值缓存。
    -cache_update_config: 定义如何修改键值缓存更新，使其成为静态的。
        - kv_cache_condition (str): 在注意力层原始前向传播的源代码中触发键值缓存更新的条件，通常类似于"if layer_past is not None"。
        - key_name (str): 更新前键的原始引用方式
        - value_name (str): 更新前值的原始引用方式
        - new_key_name (Optional[str]): 更新后键的引用方式。如果未指定，将简单地使用key_name。
        - new_value_name (Optional[str]): 更新后值的引用方式。如果未指定，将简单地使用value_name。

    -causal_mask_config: 定义如何应用因果掩码 - 这是必要的，因为您的键和值现在在倒数第二个维度上的长度为```max_length```。
        - causal_mask_module (str): 应用因果掩码的注意力层方法。
        - causal_mask_application (Union["conditional", Any]): 这要么是字符串"conditional"，要么是其他值。
        - 如果causal_mask_application是"conditional"，您需要添加以下额外键：
            - causal_mask_condition (str): 应用因果掩码的条件。
        - 如果不是条件性的，您需要添加以下额外键：
            - causal_mask_line (str): 我们想要替换的起始行。
            - num_lines (int): 包括causal_mask_line在内，我们想要替换的行数

-imports: 这些是在集成静态键值缓存后编译新函数所需的导入。

返回：
- KVCacheModel类的一个实例，本质上就是您的模型，但附加了键值缓存以加速推理。

add_speculative_decoding(model:nn.Module, draft_model:nn.Module) -> nn.Module
- 参数：
  - model：这是您模型的KVCached版本。
  - draft_model：这是您草稿模型的KVCached版本。
- 返回：
  - 一个加速的模型，具有上面gpt_fast部分描述的generate方法。