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

从FasterTransformer源码解读开始了解大模型（2.0）代码解读01-看看头文件

写在前面的话

本篇的内容直接开始我们的代码通读，整个通读可能需要好几篇文章来将一整个gpt的代码结构给讲清楚。目前的计划是先从整体model层次开始讲，将ContextDecoder和Decoder讲完之后，再从模块内部（MHA和FFN）代码开始讲，中间也会穿插一些技术点的讲解。

零、model层的结构参数

打开之前的src/fastertransformer/models/multi_gpu_gpt目录，正如我们在上一篇文章说过的，整体的model是由在ParallelGpt的代码中，而其中的ContextDecoder和Decoder两大模块的功能则是由ParallelGptContextDecoder和ParallelGptDecoder来进行组织。一个良好的习惯是我们可以先从头文件ParallelGpt.h开始看，了解一下大致的情况和需要的功能。

先看一下ParallelGpt.h的33~66行，这里是一些用于描述模型结构的参数

private:
    // meta data
    size_t               head_num_;
    size_t               size_per_head_;
    size_t               inter_size_;
    size_t               num_layer_;
    size_t               vocab_size_;
    size_t               expert_num_;
    size_t               moe_k_;
    std::vector<int64_t> moe_layer_index_;

    int    start_id_;
    int    end_id_;
    float  beam_search_diversity_rate_;
    size_t hidden_units_;

    const float layernorm_eps_;  // OPT
    float       shared_contexts_ratio_;

    // TODO(bhsueh) remove these member because they are runtime parameters
    size_t             top_k_;
    float              top_p_;
    unsigned long long random_seed_;

    float temperature_;
    float len_penalty_;
    float repetition_penalty_;

    size_t    local_head_num_;
    NcclParam tensor_para_;
    NcclParam pipeline_para_;

    std::shared_ptr<AbstractCustomComm> custom_all_reduce_comm_;
    int                                 enable_custom_all_reduce_;

这么多的参数，是否会出现看花了眼的情况？其实代码中很贴心地将一些相同功能方向的参数都给放在了一起，逐个看下来也不会特别困难。下面介绍一些结构上的参数来进行简单介绍，一些计算中用的参数则会在后续结合代码运行过程的源码解析来进行说明

首先是从head_num到moe_layer_index_这八个参数，都是用于描述模型基础属性的。head_num是用于描述MHA时拥有多少个注意力头的，size_per_head是用于描述MHA中每个头的大小的，最常见的模型一般是128，google的gemma则是256. inter_size是用于描述FFN的中间维度大小，num_layer是用于描述模型整体的层数的，vocab_size是用于描述token词表的大小

llm小知识-参数量：一般来说，参数量越多的模型（越大的模型）会拥有更多的层数（num_layer）和MHA的头数（head_num），这也就意味着它们能拥有更多的信息，模型效果也就越好，但同时，每个token的计算量和生成时间也就越长，也越吃硬件资源（算力和显存）

接下来是三个跟moe有关的参数，expert_num是专家的数量，moe_k则是用于控制在N个专家里取前k个参与计算，moe_layer_index是用于给不同专家做index标记。

llm小知识-moe：moe即Mixture-of-Experts，混合专家。是一种大模型结构上的优化，它使用MoE层替换掉了原来的FFN层，其中MoE层由多个专家FFN和一个router路由组成，每个token的隐藏状态经过router路由的计算，会被分配给多个（moe_k个）不同的专家Expert来完成计算

接下来是start_id和end_id，一般来说对于一个模型，start_id用于标记输入句子的开始，end_id用于表示decoder生成已经完成，在遇到生成的词为end_id时，就需要掐断自回归生成的过程了。

hidden_units_，是记录每个token隐藏状态大小的变量，一般来说这个大小会是size_per_head * head_num（也有例外，例如谷歌的gemma）

local_head_num，之前提到过我们可能会对模型进行tp划分，在MHA中，tp划分则是会对不同的head来进行切分，那么每张卡上就会只需要进行local_head_num / tp_num数量的头的计算，local_head_num就是用于进行单卡头数量记录的
tensor_para和pipeline_para，这两个参数是用于记录模型的tp和pp划分的，是NcclParam类型的变量，不同卡之间也可以靠这两个参数来完成Nccl通信。

一、模型推理参数

接下来可以初步看一下65~87行的代码，主要会涉及到模型推理过程的一些特殊功能和参数配置，挑其中几个比较重要的部分进行简单解释。

		std::shared_ptr<AbstractCustomComm> custom_all_reduce_comm_;
    int                                 enable_custom_all_reduce_;

    const bool is_context_qk_buf_float_ =
        (std::getenv("CONTEXT_ATTENTION_BMM1_HALF_ACCUM") == nullptr ||
         std::string(std::getenv("CONTEXT_ATTENTION_BMM1_HALF_ACCUM")) != "ON");
    size_t        vocab_size_padded_;
    const int     int8_mode_      = 0;
    AttentionType attention_type_ = AttentionType::UNFUSED_MHA;

    // Prompt Learning Parameters
    PromptLearningType prompt_learning_type_;
    int                prompt_learning_start_id_;  // start_id for prompt_learning (only needed by prefix prompts)
    bool               has_p_prompt_tuning_;
    bool               has_prefix_prompt_;
    bool               has_prefix_soft_prompt_;

    // GPT Variants parameters: e.g. Meta OPT
    gptVariantParams gpt_variant_params_;

    ParallelGptDecoder<T>*        gpt_decoder_;
    ParallelGptContextDecoder<T>* gpt_context_decoder_;
    DynamicDecodeLayer<float>*    dynamic_decode_layer_;

在65和66行，ft为不同模型提供了自定义的all reduce通信方法，但一般来说我们只需要使用nccl通信就可以完成这一功能，所以一般不用。

vocab_size_padded_，这个参数习惯上来说是对vocab_size按8的倍数向上取整的。为什么要取整？因为在最后一步从隐藏状态通过lm_head计算生成vocab_size时也需要进行tp划分，所以需要保证每个tp分到的vocab_size大小是可以整除的，而一个GPU节点上最多可以有8张卡（tp=8），所以习惯按照8的倍数向上取整

int8_mode_，这个是用于控制不同量化方法的参数，不同的int8mode对应不同的量化方法，常见的有fp8，w4a16等等，量化会在后续的代码解析中进行一些简单说明和分析

attention_type_是用于控制计算MHA的不同方法的参数，MHA的计算有融合和非融合两种方式，如果使用融合的方式，那么整个MHA的计算就是使用一整个大融合算子来实现，否则的话就还是按照分开的不同算子kernel来进行计算。

Prompt Learning Parameters这几个是模型中用于处理自己特定的prompt的参数（相比于它能带来的方便，其实看起来更多的是增加了很多繁琐的处理步骤，所以并不是很常见）

gpt_variant_params则是一个用于decoder结构中控制一些layernorm层是否需要计算的参数，看一下它的定义：

struct gptVariantParams {
    // GPT default params
    float          layernorm_eps   = 1e-6f;
    LayerNormType  layernorm_type  = LayerNormType::pre_layernorm;
    ActivationType activation_type = ActivationType::Gelu;
    // Whether to have a learnt positional encoding.
    bool has_positional_encoding = true;
    // A layernom just after the word embedding and before the decoder.
    bool has_pre_decoder_layernorm = false;
    // A layernom after the decoder.
    bool has_post_decoder_layernorm = true;
    // detoxification adapters. refer to
    bool   has_adapters       = false;
    size_t adapter_inter_size = 0;
    // Whether to use the attention linear positional bias
    bool use_attention_linear_bias = false;
};

在这个定义中可以看见，除了控制layernorm的eps参数以外，还以一系列用于控制layernorm位置的参数（has_pre_decoder_layernorm、has_post_decoder_layernorm等等）和激活函数的参数，为了方便统一读取和管理，全部放在了这个结构体中来进行读取控制。

接下来的三个指针gpt_decoder_，gpt_context_decoder_，dynamic_decode_layer_则是主要负责完成context-decoder到decoder到采样生成这个过程的最主要的部分，也是由这三个指针指向的类来负责完成model层的主要功能的组装

二、成员函数和buffer

下面是一些成员函数的简单解释以及buffer的说明，可以先看89行~107行

		void allocateBuffer() override;
    void allocateBuffer(size_t batch_size,
                        size_t beam_width,
                        size_t max_seq_len,
                        size_t memory_len,
                        size_t max_input_len,
                        bool   is_return_context_cum_log_probs);
    void freeBuffer() override;

    void initialize();

    void computeContextCumLogProbs(float*                      cum_log_probs,
                                   const T*                    context_decoder_outputs,
                                   const int*                  input_ids,
                                   const int*                  input_lengths,
                                   const size_t                batch_size,
                                   const size_t                beam_width,
                                   const size_t                max_input_length,
                                   const ParallelGptWeight<T>* gpt_weights);

首先是用于对所有的buffer指针进行内存分配和显存分配的函数allocateBuffer，这个函数会在每次进行一次完整推理时用到，对推理过程中用到的buffer进行预分配，而完成推理之后，则是由freeBuffer()函数来对buffer进行一些释放。

而initialize()函数则是对函数内部的一些参数进行初始化操作，在这里面会对上面提到的三个model的组成部分指针进行创建

computeContextCumLogProbs函数则是提供了一种方法，如果在推理最后需要返回logits时对logits进行收集和计算的函数。

接下来，从109行到175行，是protect成员中一系列用于指向buffer的指针，这些指针在这里不会一一进行功能说明（真的太多了），在实际运行代码讲解中遇到的时候再进行简单说明。

在protect成员中还有两个函数，一个setOutputTensors和sendTensorsToFirstPipelineNode，这里的两个函数是在完成推理后，将数据按照tensor进行组装并返回到头节点中的。在没有进行pp切分时其实已经可以直接组装好了

在下面186到246行的代码中，是一些public成员函数，从函数名就可以直接看出这些函数的具体作用和功能，以及它们返回的结果是什么，这里就不再多做赘述了，可以简单提到的是registerCallback和unRegisterCallback这两个函数，这两个函数是用于注册和注销生成token并返回的回调函数的，这里涉及到返回生成结果的两种方式，流式和非流式，会在返回结果的过程中产生一些区别。

llm小知识-流式和非流式：在推理框架产生token之后，流式的返回方式会在每次产生一个token之后就立马返回一个token，非流式的返回方式则是在所有token都生成完毕之后再一次性返回所有token。这两种方式没有性能上的差异，只会有使用体验上的不同。一般来说，输出长度较长的交互式场景推荐使用流式，方便用户一边等待一边看，短输出场景、预测式场景适合使用非流式，一次性返回全部结果。

下一回预告

下一回将会从ParallelGpt.cc的代码开始讲解，对其部分功能函数以及主要的函数forward进行一些分析，在分析过程中会对一些buffer指针功能以及代码部分进行分析说明。在解析过程中会对一些不是非常重要的代码进行跳过，并不会逐条分析（很多都是用不上的）