从FasterTransformer源码解读开始了解大模型（2.1）代码通读02

news2024/12/27 13:17:55

从FasterTransformer源码解读开始了解大模型（2.0）代码解读02-初始化和forward

写在前面的话

本篇的内容主要是介绍ParallelGpt.cc中的代码内容，首先介绍一些初始化和工具函数，然后会从forward主函数开始介绍一部分。

零、初始化initialize和allocateBuffer

打开src/fastertransformer/models/multi_gpu_gpt/ParallelGpt.cc文件，这里是GPT的真正的处理推理请求的功能函数。在这个文件的fastertransformer namespace中，第一个函数是用于做一些初始化的函数initialize，从31到87行，主要是创建了三个对象，gpt_context_decoder是用于做ContextDecoder或者说Encoder部分的，gpt_decoder是用于做Decoder的，而进行采样和结果生成的是DynamicDecodeLayer部分。这三个部分我们会在后续的代码解读中展开说明。

第95行到202行是allocateBuffer函数，在每次处理一个推理请求时，都会使用allocateBuffer进行显存的分配和内存的分配。这里挑出几个比较有特点的buffer进行简单讲解。

在109行，计算了一个变量为self_cache_size，大小是*(num_layer / pipeline_para_.world_size) * batchxbeam * memory_len * hidden_units_ / tensor_para.world_size_，这个实际上就是计算KV Cache的大小。而134和135行就用该数值的大小进行了KV Cache的分配。

const size_t self_cache_size =
        (num_layer_ / pipeline_para_.world_size_) * batchxbeam * memory_len * hidden_units_ / tensor_para_.world_size_;

llm小知识-KV Cache：我们知道，在Attention注意力得分的计算过程中，对于当前的token i，需要先计算出查询结果Qi，Ki和Vi，然后使用Qi与token i 之前的所有token的K结果和V结果来进行注意力得分计算，就是拿Qi与所有的K(0-i)点积求和，再与V(0-i)进行加权求和，最终求得Attention注意力得分。那么在这个过程中，可以通过将之前所有计算过的Ki和Vi进行存储的方式，来减少计算量（用显存空间换时间），那么这部分用于存储KV的就是KVCache。目前有一些量化算法也会关注于KV的量化以减少存储空间

在163行的context_decoder_input_buf_，这块buff被分配的大小为sizeof(T) * batchxbeam * max_input_len * hidden_units，这块buff大小为每一个输入token的大小乘以隐藏状态的大小，同样大小的buff还有context_decoder_output_buf，context_decoder_normed_input_buf等，这个大小是在ContextDecoder进行计算流程时真正的隐藏状态的大小（可以参考前几章中的decoder-only模型结构）所以会多次出现。

context_decoder_input_buf_  = (T*)(allocator_->reMalloc(
        context_decoder_input_buf_, sizeof(T) * batchxbeam * max_input_len * hidden_units_, false));

类似的还有120行的decoder_input_buf变量，大小为sizeof(T) * batchxbeam * hidden_units，相比较之下少了一个max_input_len大小维度，由于Decoder每一步只生成一个token所以相当于长度始终为1，少了一个输入长度的维度。与decoder_input_buf大小相同的还有decoder_normed_input_buf，decoder_output_buf等等。

decoder_input_buf_ = (T*)(allocator_->reMalloc(decoder_input_buf_, sizeof(T) * batchxbeam * hidden_units_, false));
    decoder_normed_input_buf_ =
        (T*)(allocator_->reMalloc(decoder_normed_input_buf_, sizeof(T) * batchxbeam * hidden_units_, false));
    decoder_output_buf_ =
        (T*)(allocator_->reMalloc(decoder_output_buf_, sizeof(T) * batchxbeam * hidden_units_, false));

在204到271行，是与allocateBuffer对应的freeBuffer函数，对指针中分配了的空间进行释放，这一段没有特别值得讲解的地方。

一、forward函数-起始检查

在ParallelGpt.cc文件中拥有两个forward函数，我们主要看574行开始的forward函数。

进入forward函数后，首先通过FT_CHECK的多个宏定义检查了输入tensor的数量以及对几个比较重要的输入tensor进行了输入形状检查。

		FT_CHECK_WITH_INFO(input_tensors->size() >= 3, "input_tensors->size() >= 3");
    FT_CHECK_WITH_INFO(output_tensors->size() >= 2, "output_tensors->size() >= 2");
    FT_CHECK(input_tensors->at("input_ids").shape.size() == 2);
    FT_CHECK(input_tensors->at("input_lengths").shape.size() == 1);
    FT_CHECK(input_tensors->find("output_seq_len") != input_tensors->end()
             && input_tensors->at("output_seq_len").shape.size() == 1);
    FT_CHECK(output_tensors->at("output_ids").shape.size() == 3);
    FT_CHECK(output_tensors->at("sequence_length").shape.size() == 2);
    FT_CHECK_WITH_INFO(input_tensors->at("input_ids").shape[0] == output_tensors->at("output_ids").shape[0],
                       "input_tensors->at(\"input_ids\").shape[0] == output_tensors->at(\"output_ids\").shape[0]");

    // Used when inputs do not contain random_seed
    const size_t batch_size = output_tensors->at("output_ids").shape[0];
    const size_t beam_width = output_tensors->at("output_ids").shape[1];
    FT_CHECK_WITH_INFO(output_tensors->count("cum_log_probs") == 0
                           || output_tensors->at("cum_log_probs").size() == batch_size * beam_width,
                       "The shape of cum_log_probs should match with batch_size x beam_width if provided.");

对于输入input_tensors来说，必须要有input_ids，input_lengths，output_seq_len这三个必备的输入，而对于output_tensors来说，则必须要有output_ids和sequence_length这两个必备的输出。这几个tensor的含义和形状如下所示（B指的是batch size，S指的是Sequence length， bz指的是beam width）

tensor名称	形状	含义
input_ids	B x S	需要推理的所有token ids，总共batch size个句子
input_lengths	B	长度为batch size的数组，每个位置标记着对应位置的token ids长度为多少
output_seq_len	B	长度为batch size的数组，每个位置标记着对应位置句子的输入最长到多少
output_ids	B x bw x S	推理完成的所有token ids，总共batch size个句子
sequence_length	B x bw	推理完成后所有句子的长度，每个位置标记着对应位置的句子长度

llm小知识-batch size：大部分的推理引擎都会将多个请求（可能每个请求只包含一个句子）打包为一个batch来进行推理，在推理过程中同一batch的句子之间是完全可以做到互不干扰的，打batch的一个非常明显的优点是可以充分发挥硬件的算力，同时处理多个请求。另外需要注意的一点是，在batch中，可能会出现长短不一的情况，这个时候是需要在短的句子后面做padding的，这一步往往是会在客户端或者server侧前端就完成好，在推理侧拿到的数据往往是已经做好padding的

如图是一个batch_size为4的推理请求，其中除了最长的句子以外其他的句子都做了padding
在这里插入图片描述