devbox
Cpp五条
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

基于torch.compile和gptfast代码风格实现ChatGLM模型推理加速_fastgpt代码分析

作者：Cpp五条 | 2024-04-05 13:36:08

赞

踩

fastgpt代码分析

       

目录

一、ChatGLM模型代码重构迁移

二、推理的代码重构

三、效果分析对比

参考文章

         torch2.0发布以后模型训练和推理可以实现一行代码加速，试用之后发现效果并不明显。随后gptfast项目也发布，表明它确实是可以实现模型推理的加速，看来之前试用是打开方式不对。最近参考gptfast项目，实现了对ChatGLM模型推理的加速，主要的原理是借助torch.compile对模型推理过程中构建计算图，实现加速。本文的重点工作就是展示模型代码和推理逻辑的迁移实现，以及加速效果的对比，当然这个方案比VLLM和tensort-LLM肯定是差了点，这个不是本文的重点，后面有空了也把vllm和tensort-LLM也写写博客对比一下效率。

一、ChatGLM模型代码重构迁移

      这个工作是真的不是特别好做，需要对模型结构和模型输入输出非常熟悉，同时也要对gptfast项目迁移原则比较熟悉，才能比较快的迁移成功。核心原则是不能有tensor切片操作，同时kvcache这种也要写成固定的长度，计算过程中不断的去填充更新，同时还要放在模型的结构外侧作为一个参数传入，加速才有效果。还有一个点要注意注意力计算的实现，由于torch更新了scaled_dot_product_attention使得最大长度的定长的矩阵计算注意力，和之前动态逐步增加长度的值是一样的，这个是注意力计算中tensor切片改写的前提(验证过确实是一样的)。细节的地方需要注意kvcache的维度形状，解码过程中不同阶段(首次forward和kvcache存在后的)模型输入的full_attention_mask是不一样的。

整体结构

  1. class TransformerGLM(nn.Module):
  2.     def __init__(self, config, device) -> None:
  3.         super().__init__()
  4.         self.config = config
  5.  
  6.         self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
  7.         rotary_dim = (
  8.             128
  9.         )
  10.         self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device,
  11.                                               dtype=config.torch_dtype)
  12.         self.layers = nn.ModuleList(TransformerBlock(config, i, device) for i in range(config.num_layers))
  13.         self.final_layernorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon, device=device,
  14.                                        dtype=config.torch_dtype)
  15.  
  16.         self.output_layer = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
  17.         self.seq_length = config.seq_length
  18.  
  19.     def forward(self, input_ids,
  20.                 position_ids: Optional[torch.Tensor] = None,
  21.                 attention_mask: Optional[torch.BoolTensor] = None,
  22.                 input_pos=None,
  23.                 is_input_mask=False,
  24.                 kv_caches=None
  25.                 ) -> Tensor:
  26.  
  27.         inputs_embeds = self.embedding(input_ids)
  28.         inputs_embeds = inputs_embeds.transpose(0, 1).contiguous()
  29.         rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
  30.         rotary_pos_emb = rotary_pos_emb[position_ids]
  31.         rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous()
  32.  
  33.         presents = ()
  34.         for i, layer in enumerate(self.layers):
  35.             inputs_embeds, kv_cache = layer(inputs_embeds, rotary_pos_emb=rotary_pos_emb, input_pos=input_pos,
  36.                                             attention_mask=attention_mask, kv_cache=kv_caches[i])
  37.             presents = presents + (kv_cache,)
  38.         hidden_states = self.final_layernorm(inputs_embeds)
  39.         lm_logits = self.output_layer(hidden_states)
  40.         lm_logits = lm_logits.transpose(0, 1).contiguous()
  41.         return lm_logits, presents

注意模型的输入新增的有input_pos，模型解码token的位置，kv_caches；模型基本模块上没有变化，精简其中的一下预处理逻辑和分支，主要就是要让torch.compile()能完成计算图的构建。

kvcache模块

  1. class KVCache(nn.Module):
  2.     def __init__(self, max_batch_size, max_seq_length, dtype=torch.bfloat16):
  3.         super().__init__()
  4.         cache_shape = (2, max_batch_size, max_seq_length, 128)
  5.         self.register_buffer('k_cache', torch.zeros(cache_shape, dtype=dtype))
  6.         self.register_buffer('v_cache', torch.zeros(cache_shape, dtype=dtype))
  7.  
  8.     def update(self, input_pos, k_val, v_val):
  9.         # input_pos: S, k_val: [S, B, H, D]
  10.         assert input_pos.shape[0] == k_val.shape[0]
  11.         k_out = self.k_cache
  12.  
  13.         v_out = self.v_cache
  14.         k_val = k_val.transpose(0, 2).contiguous()
  15.  
  16.         v_val = v_val.transpose(0, 2).contiguous()
  17.         k_out[:, :, input_pos] = k_val.clone()
  18.         v_out[:, :, input_pos] = v_val.clone()
  19.         k_out = k_out.transpose(0, 2).contiguous()
  20.  
  21.         v_out = v_out.transpose(0, 2).contiguous()
  22.  
  23.         return k_out, v_out

模块中各个变量的维度信息都标注好了，作用就是kv缓存载体以及更新逻辑提供一个方法。

其他模块就不一一介绍了，注意selfattention中kvcache的更新

整个模型的代码如下：

  1. # Copyright (c) Meta Platforms, Inc. and affiliates.
  2. # All rights reserved.
  3.  
  4. # This source code is licensed under the license found in the
  5. # LICENSE file in the root directory of this source tree.
  6. from dataclasses import dataclass
  7. from typing import Optional, Tuple
  8.  
  9. import torch
  10. import torch.nn as nn
  11. from torch import Tensor
  12. from torch.nn import functional as F
  13. from math import gcd
  14. from functools import reduce
  15. import math
  16.  
  17.  
  18. def find_multiple(n: int, *args: Tuple[int]) -> int:
  19.     k = reduce(lambda x, y: x * y // gcd(x, y), args + (1,))
  20.     if n % k == 0:
  21.         return n
  22.     return n + k - (n % k)
  23.  
  24.  
  25. class CoreAttention(torch.nn.Module):
  26.     def __init__(self, config, layer_number):
  27.         super(CoreAttention, self).__init__()
  28.         self.config = config
  29.         self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling
  30.         self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32
  31.  
  32.         self.attention_softmax_in_fp32 = True
  33.         self.layer_number = max(1, layer_number)
  34.  
  35.         projection_size = config.kv_channels * config.num_attention_heads
  36.  
  37.         self.hidden_size_per_partition = projection_size
  38.         self.hidden_size_per_attention_head = projection_size // config.num_attention_heads
  39.         self.num_attention_heads_per_partition = config.num_attention_heads
  40.  
  41.         self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
  42.         coeff = self.layer_number
  43.         self.norm_factor *= coeff
  44.         self.coeff = coeff
  45.         self.attention_dropout = torch.nn.Dropout(config.attention_dropout)
  46.  
  47.     def forward(self, query_layer, key_layer, value_layer, attention_mask=None):
  48.         query_layer, key_layer, value_layer = [k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]]
  49.         if attention_mask is None:
  50.             context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
  51.                                                                              is_causal=True)
  52.         else:
  53.             attention_mask = ~attention_mask
  54.             context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
  55.                                                                              attention_mask)
  56.  
  57.         context_layer = context_layer.permute(2, 0, 1, 3)
  58.         new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
  59.         context_layer = context_layer.reshape(*new_context_layer_shape)
  60.  
  61.         return context_layer
  62.  
  63.  
  64. def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
  65.     # x: [sq, b, np, hn]
  66.     sq, b, np, hn = x.size(0), x.size(1), x.size(2), x.size(3)
  67.     rot_dim = rope_cache.shape[-2] * 2
  68.     x, x_pass = x[..., :rot_dim], x[..., rot_dim:]
  69.     # truncate to support variable sizes
  70.     rope_cache = rope_cache[:sq]
  71.     xshaped = x.reshape(sq, -1, np, rot_dim // 2, 2)
  72.     rope_cache = rope_cache.view(sq, -1, 1, xshaped.size(3), 2)
  73.     x_out2 = torch.stack(
  74.         [
  75.             xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],
  76.             xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],
  77.         ],
  78.         -1,
  79.     )
  80.     x_out2 = x_out2.flatten(3)
  81.     return torch.cat((x_out2, x_pass), dim=-1)
  82.  
  83.  
  84. class RotaryEmbedding(nn.Module):
  85.     def __init__(self, dim, original_impl=False, device=None, dtype=None):
  86.         super().__init__()
  87.         inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).to(dtype=dtype) / dim))
  88.         self.register_buffer("inv_freq", inv_freq)
  89.         self.dim = dim
  90.         self.original_impl = original_impl
  91.  
  92.     def forward_impl(
  93.             self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000
  94.     ):
  95.         """Enhanced Transformer with Rotary Position Embedding.
  97.         Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/
  98.         transformers/rope/__init__.py. MIT License:
  99.         https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.
  100.         """
  101.         # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
  102.         theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=dtype, device=device) / n_elem))
  103.  
  104.         # Create position indexes `[0, 1, ..., seq_len - 1]`
  105.         seq_idx = torch.arange(seq_len, dtype=dtype, device=device)
  106.  
  107.         # Calculate the product of position index and $\theta_i$
  108.         idx_theta = torch.outer(seq_idx, theta).float()
  109.  
  110.         cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)
  111.  
  112.         # this is to mimic the behaviour of complex32, else we will get different results
  113.         if dtype in (torch.float16, torch.bfloat16, torch.int8):
  114.             cache = cache.bfloat16() if dtype == torch.bfloat16 else cache.half()
  115.         return cache
  116.  
  117.     def forward(self, max_seq_len, offset=0):
  118.         return self.forward_impl(
  119.             max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device
  120.         )
  121.  
  122.  
  123. class KVCache(nn.Module):
  124.     def __init__(self, max_batch_size, max_seq_length, dtype=torch.bfloat16):
  125.         super().__init__()
  126.         cache_shape = (2, max_batch_size, max_seq_length, 128)
  127.         self.register_buffer('k_cache', torch.zeros(cache_shape, dtype=dtype))
  128.         self.register_buffer('v_cache', torch.zeros(cache_shape, dtype=dtype))
  129.  
  130.     def update(self, input_pos, k_val, v_val):
  131.         # input_pos: S, k_val: [S, B, H, D]
  132.         assert input_pos.shape[0] == k_val.shape[0]
  133.         k_out = self.k_cache
  134.  
  135.         v_out = self.v_cache
  136.         k_val = k_val.transpose(0, 2).contiguous()
  137.  
  138.         v_val = v_val.transpose(0, 2).contiguous()
  139.         k_out[:, :, input_pos] = k_val.clone()
  140.         v_out[:, :, input_pos] = v_val.clone()
  141.         k_out = k_out.transpose(0, 2).contiguous()
  142.  
  143.         v_out = v_out.transpose(0, 2).contiguous()
  144.  
  145.         return k_out, v_out
  146.  
  147.  
  148. class RMSNorm(torch.nn.Module):
  149.     def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs):
  150.         super().__init__()
  151.         self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype))
  152.         self.eps = eps
  153.  
  154.     def forward(self, hidden_states: torch.Tensor):
  155.         input_dtype = hidden_states.dtype
  156.         variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
  157.         hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
  158.  
  159.         return (self.weight * hidden_states).to(input_dtype)
  160.  
  161.  
  162. class TransformerGLM(nn.Module):
  163.     def __init__(self, config, device) -> None:
  164.         super().__init__()
  165.         self.config = config
  166.  
  167.         self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
  168.         rotary_dim = (
  169.             128
  170.         )
  171.         self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device,
  172.                                               dtype=config.torch_dtype)
  173.         self.layers = nn.ModuleList(TransformerBlock(config, i, device) for i in range(config.num_layers))
  174.         self.final_layernorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon, device=device,
  175.                                        dtype=config.torch_dtype)
  176.  
  177.         self.output_layer = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
  178.         self.seq_length = config.seq_length
  179.  
  180.     def forward(self, input_ids,
  181.                 position_ids: Optional[torch.Tensor] = None,
  182.                 attention_mask: Optional[torch.BoolTensor] = None,
  183.                 input_pos=None,
  184.                 is_input_mask=False,
  185.                 kv_caches=None
  186.                 ) -> Tensor:
  187.  
  188.         inputs_embeds = self.embedding(input_ids)
  189.         inputs_embeds = inputs_embeds.transpose(0, 1).contiguous()
  190.         rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
  191.         rotary_pos_emb = rotary_pos_emb[position_ids]
  192.         rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous()
  193.  
  194.         presents = ()
  195.         for i, layer in enumerate(self.layers):
  196.             inputs_embeds, kv_cache = layer(inputs_embeds, rotary_pos_emb=rotary_pos_emb, input_pos=input_pos,
  197.                                             attention_mask=attention_mask, kv_cache=kv_caches[i])
  198.             presents = presents + (kv_cache,)
  199.         hidden_states = self.final_layernorm(inputs_embeds)
  200.         lm_logits = self.output_layer(hidden_states)
  201.         lm_logits = lm_logits.transpose(0, 1).contiguous()
  202.         return lm_logits, presents
  203.  
  204.  
  205. class MLP(torch.nn.Module):
  206.     """MLP.
  207.     MLP will take the input with h hidden state, project it to 4*h
  208.     hidden dimension, perform nonlinear transformation, and project the
  209.     state back into h hidden dimension.
  210.     """
  211.  
  212.     def __init__(self, config, device=None):
  213.         super(MLP, self).__init__()
  214.  
  215.         self.add_bias = config.add_bias_linear
  216.  
  217.         # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf
  218.         self.dense_h_to_4h = nn.Linear(
  219.             config.hidden_size,
  220.             config.ffn_hidden_size * 2,
  221.             bias=self.add_bias,
  222.             device=device,
  223.             # **_config_to_kwargs(config)
  224.         )
  225.  
  226.         def swiglu(x):
  227.             x = torch.chunk(x, 2, dim=-1)
  228.             return F.silu(x[0]) * x[1]
  229.  
  230.         self.activation_func = swiglu
  231.  
  232.         # Project back to h.
  233.         self.dense_4h_to_h = nn.Linear(
  234.             config.ffn_hidden_size,
  235.             config.hidden_size,
  236.             bias=self.add_bias,
  237.             device=device,
  238.             # **_config_to_kwargs(config)
  239.         )
  240.  
  241.     def forward(self, hidden_states):
  242.         # [s, b, 4hp]
  243.         intermediate_parallel = self.dense_h_to_4h(hidden_states)
  244.         intermediate_parallel = self.activation_func(intermediate_parallel)
  245.         # [s, b, h]
  246.         output = self.dense_4h_to_h(intermediate_parallel)
  247.         return output
  248.  
  249.  
  250. class SelfAttention(torch.nn.Module):
  251.     """Parallel self-attention layer abstract class.
  253.     Self-attention layer takes input with size [s, b, h]
  254.     and returns output of the same size.
  255.     """
  256.  
  257.     def __init__(self, config, layer_number, device=None):
  258.         super(SelfAttention, self).__init__()
  259.         self.config = config
  260.         self.layer_number = max(1, layer_number)
  261.  
  262.         self.projection_size = config.kv_channels * config.num_attention_heads
  263.  
  264.         # Per attention head and per partition values.
  265.         self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads
  266.         # 32
  267.         self.num_attention_heads_per_partition = config.num_attention_heads
  268.  
  269.         self.multi_query_attention = config.multi_query_attention
  270.         self.qkv_hidden_size = 3 * self.projection_size
  271.  
  272.         self.num_multi_query_groups_per_partition = config.multi_query_group_num
  273.         self.qkv_hidden_size = (
  274.                 self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num
  275.         )
  276.         self.query_key_value = nn.Linear(config.hidden_size, self.qkv_hidden_size,
  277.                                          bias=config.add_bias_linear or config.add_qkv_bias,
  278.                                          # device=device, **_config_to_kwargs(config)
  279.                                          )
  280.  
  281.         self.core_attention = CoreAttention(config, self.layer_number)
  282.  
  283.         # Output.
  284.         self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear,
  285.                                device=device,
  286.                                # **_config_to_kwargs(config)
  287.                                )
  288.  
  289.     def forward(
  290.             self, hidden_states, rotary_pos_emb, input_pos, attention_mask=None, kv_cache=None
  291.     ):
  292.         # hidden_states: [sq, b, h]
  293.  
  294.         # =================================================
  295.         # Pre-allocate memory for key-values for inference.
  296.         # =================================================
  297.         # =====================
  298.         # Query, Key, and Value
  299.         # =====================
  300.  
  301.         # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)]
  302.         mixed_x_layer = self.query_key_value(hidden_states)
  303.  
  304.         (query_layer, key_layer, value_layer) = mixed_x_layer.split(
  305.             [
  306.                 self.num_attention_heads_per_partition * self.hidden_size_per_attention_head,
  307.                 self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
  308.                 self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
  309.             ],
  310.             dim=-1,
  311.         )
  312.  
  313.         query_layer = query_layer.view(
  314.             query_layer.size()[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
  315.         )
  316.         key_layer = key_layer.view(
  317.             key_layer.size()[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
  318.         )
  319.         value_layer = value_layer.view(
  320.             value_layer.size()[:-1]
  321.             + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
  322.         )
  323.  
  324.         # apply relative positional encoding (rotary embedding)
  325.         query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)
  326.         key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)
  327.  
  328.         # 更新kvcache
  329.         cache_k, cache_v = kv_cache
  330.         cache_k[input_pos] = key_layer
  331.         cache_v[input_pos] = value_layer
  332.         key_layer = cache_k.clone()
  333.         value_layer = cache_v.clone()
  334.         kv_cache = (key_layer, value_layer)
  335.  
  336.         key_layer = key_layer.unsqueeze(-2)
  337.         key_layer = key_layer.expand(
  338.             -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1
  339.         )
  340.         key_layer = key_layer.contiguous().view(
  341.             key_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
  342.         )
  343.         value_layer = value_layer.unsqueeze(-2)
  344.         value_layer = value_layer.expand(
  345.             -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1
  346.         )
  347.         value_layer = value_layer.contiguous().view(
  348.             value_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
  349.         )
  350.  
  351.         # ==================================
  352.         # core attention computation
  353.         # ==================================
  354.  
  355.         context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask=attention_mask)
  356.  
  357.         # =================
  358.         # Output. [sq, b, h]
  359.         # =================
  360.  
  361.         output = self.dense(context_layer)
  362.  
  363.         return output, kv_cache
  364.  
  365.  
  366. class TransformerBlock(nn.Module):
  367.     def __init__(self, config, layer_number, device) -> None:
  368.         super().__init__()
  369.         self.hidden_dropout = config.hidden_dropout
  370.         self.input_layernorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon, device=device,
  371.                                        dtype=config.torch_dtype)
  372.         self.self_attention = SelfAttention(config, layer_number, device=device)
  373.         self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon, device=device,
  374.                                                 dtype=config.torch_dtype)
  375.         self.mlp = MLP(config, device=device)
  376.  
  377.     def forward(self, hidden_states, rotary_pos_emb, input_pos, attention_mask=None, kv_cache=None):
  378.         layernorm_output = self.input_layernorm(hidden_states)
  379.         attention_output, kv_cache = self.self_attention(
  380.             layernorm_output,
  381.             rotary_pos_emb,
  382.             input_pos,
  383.             attention_mask=attention_mask,
  384.             kv_cache=kv_cache
  385.         )
  386.         residual = hidden_states
  387.         layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training)
  388.         layernorm_input = residual + layernorm_input
  389.         layernorm_output = self.post_attention_layernorm(layernorm_input)
  390.         mlp_output = self.mlp(layernorm_output)
  391.         residual = layernorm_input
  392.         output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training)
  393.         output = residual + output
  394.         return output, kv_cache
  395.  
  396.  
  397. class RMSNorm(torch.nn.Module):
  398.     def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs):
  399.         super().__init__()
  400.         self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype))
  401.         self.eps = eps
  402.  
  403.     def forward(self, hidden_states: torch.Tensor):
  404.         input_dtype = hidden_states.dtype
  405.         variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
  406.         hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
  407.  
  408.         return (self.weight * hidden_states).to(input_dtype)
  409.  
  410.  
  411. if __name__ == '__main__':
  412.     import os
  413.  
  414.     os.environ['CUDA_VISIBLE_DEVICES'] = "1"
  415.     from transformers import AutoConfig
  416.  
  417.     model_path = "./chatglm2-6b-merge"
  418.     config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
  419.     model = TransformerGLM(config, device=None)
  420.     for name, _ in model.named_parameters():
  421.         print(name)

二、推理的代码重构

推理方法，也就是重写transformer模型中的generate这个方法，对于一次生成可以分为第一次解码forward阶段和余下的解码forward阶段。实现分别如下，只实现了greedy search 策略：

  1. @torch.no_grad()
  2. def first_decode_batch(model, input_ids, position_ids, input_pos, attention_mask, kv_caches):
  3.     logits, kv_caches = model(input_ids=input_ids, position_ids=position_ids, input_pos=input_pos, is_input_mask=False,
  4.                               attention_mask=attention_mask, kv_caches=kv_caches)
  5.     logits = logits[:, -1:]
  6.     next_tok = torch.argmax(logits, dim=-1)
  7.     return next_tok, kv_caches
  8.  
  9.  
  10. @torch.no_grad()
  11. def decode_one_token_batch(model, input_ids, position_ids, input_pos, attention_mask, kv_caches):
  12.     logits, kv_caches = model(input_ids, position_ids=position_ids, input_pos=input_pos, is_input_mask=True,
  13.                               attention_mask=attention_mask, kv_caches=kv_caches)
  14.     logits = logits[:, -1:]
  15.     next_tok = torch.argmax(logits, dim=-1)
  16.     return next_tok, kv_caches

主要是得到解码过程中模型输出的token和kv_caches。特别要注意的是不能把这两个方法封装到一个类中，然后再进行torch.compile这样模型能正确输出结果，但是推理速度没有提升的，也就是torch.compile并没有生效。

整体的generate逻辑，包含停止符号，模型的初始输入、kvcaches初始化以及attention_mask输入的变化、position_ids的输入变化，batch推理是padding的加入。

  1. def generate_own_batch(model,
  2.                  inputs,
  3.                  sampling_kwargs,
  4.                  eos_token,
  5.                        max_seq_length, max_batch_size):
  6.     device = inputs['input_ids'].device
  7.     cache_shape = (max_seq_length, max_batch_size, 2, 128)
  8.     dtype = torch.bfloat16
  9.     kv_caches = [(torch.zeros(cache_shape, dtype=dtype).to(device), torch.zeros(cache_shape, dtype=dtype).to(device))
  10.                  for _ in range(model.config.num_layers)]
  11.  
  12.     input_ids = inputs['input_ids']
  13.  
  14.     ori_input_ids = input_ids.clone()
  15.     position_ids = inputs['position_ids']
  16.  
  17.     input_pos = []
  18.     for _ in range(max_batch_size):
  19.         pos = list(range(0,input_ids.shape[1]))
  20.         input_pos.append(pos)
  21.     input_pos = torch.tensor(input_pos, device=input_ids.device)
  22.  
  23.     # input_pos = torch.arange(0, input_ids.shape[1], device=input_ids.device)
  24.     next_token, kv_caches = first_decode_batch(model, input_ids, position_ids, input_pos, None, kv_caches)
  25.  
  26.     full_attention_mask = torch.ones(max_batch_size, 1, 1, max_seq_length).to(device).bool()
  27.     full_attention_mask[:, :, :, input_pos] = False
  28.  
  29.     # pading部分为true
  30.     for i in range(full_attention_mask.shape[0]):
  31.         for j in range(input_ids.shape[1]):
  32.             if input_ids[i, j] == 0:
  33.                 full_attention_mask[i, :, :, j] = True
  34.  
  35.     input_ids = torch.cat((input_ids, next_token.clone()), dim=1)
  36.     num_new_tokens = sampling_kwargs["max_length"]
  37.     T = input_ids.size()[1]
  38.  
  39.     position_ids = position_ids[:,-1:]
  40.  
  41.     input_pos = []
  42.     for _ in range(max_batch_size):
  43.         pos = [T]
  44.         input_pos.append(pos)
  45.     input_pos = torch.tensor(input_pos, device=next_token.device, dtype=torch.long)
  46.  
  47.     # position_ids = torch.tensor([[T - 1]], device=next_token.device, dtype=torch.long)
  48.     # input_pos = torch.tensor([T], device=input_ids.device, dtype=torch.long)
  49.  
  50.     for i in range(num_new_tokens):
  51.         input_pos += 1
  52.         # Actually better for Inductor to codegen attention here
  53.         with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True):
  54.             full_attention_mask[:, :, :, input_pos] = False
  55.             next_token, kv_caches = decode_one_token_batch(model, next_token, position_ids, input_pos,
  56.                                                      full_attention_mask, kv_caches)
  57.  
  58.             input_ids = torch.cat((input_ids, next_token.clone()), dim=1)
  59.  
  60.             if (input_ids == eos_token).sum(dim=1).all():
  61.                 break
  62.  
  63.             position_ids += 1
  64.             # token = next_token.tolist()
  65.             # token = next_token.tolist()[0]
  66.             # generated_tokens.append(token)
  67.     return input_ids, ori_input_ids

推理核心逻辑

  1.     model = TransformerGLM(config=config, device=None)
  2.     checkpoint_dir = Path(model_path)
  3.     model_map_json = checkpoint_dir / "pytorch_model.bin.index.json"
  4.     converted_state_dict = model.state_dict()
  5.  
  6.     gen_kwargs = {"max_length": 200, "num_beams": 1,
  7.                   "do_sample": False, "top_p": 0.8,
  8.                   "temperature": 0.95
  9.                   }
  10.     device = "cuda:0"
  11.     model_path = "./chatglm2-6b-merge"
  12.     tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
  13.     eos_token = tokenizer.eos_token_id
  14. ......
  15. # 编译加速
  16.     global first_decode_batch, decode_one_token_batch
  17.     decode_one_token_batch = torch.compile(decode_one_token_batch, mode="reduce-overhead", fullgraph=True)
  18.     first_decode_batch = torch.compile(first_decode_batch, dynamic=True, fullgraph=True)
  19.  
  20. ......
  21. generate_own_batch(model, inputs, gen_kwargs, eos_token, max_seq_length, max_batch_size)

核心点在于把解码两阶段的函数使用torch.compile函数包裹一下，真实解码过程中就会进行解码加速。

三、效果分析对比

展示一下glm模型ori、compile、和compile+int8，bs=1，max_seq_length= 1000的情况下的推理速度和效果的对比，6B glm模型，模型输入prompt如下：

[
    "你好",
    "你是谁呀？",
    "你能做什么呀？",
    "你真厉害",
    "真棒呀",
    "再见了",
    "给我推荐一部电影",
    "你知道明天天气怎么样吗？"
]

ori原始transformer的推理效果如下：

使用compile后效果如下：

compile+int8效果如下：

可以看到相同的模型和相同的数据在bs=1下，原始模型推理速度31.7 tokens/s，compile的推理速度68.1 tokens/s，110.9 tokes/s；加速效果确实比较明显。

业务领域上的实验，这里也可以给一个结论，数据就不展示了，业务上生成的token数目每次推理大都在20 tokens以内，结果如下：

这次的分享就到这里为止了，这个迁移后的模型和推理在我们公司的服务端还有个问题，我们服务端采用的多进程异步来实现web服务的，这个gptfast的服务化的集成显示int8不生效，而且bs=1时候的推理加速并没有线下加速效果明显，具体原因一直没有弄明白，可能是其他进程占用服务器资源，导致torch.compile加速失效或者降低。

参考文章

gpt-fast实战

modeling_chatglm

gpt-fast项目源码

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/Cpp五条/article/detail/365898
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号