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

从FasterTransformer源码解读开始了解大模型（2.3）代码解读04-forward函数

写在前面的话

本篇的内容继续解读forward函数，从972行开始进行解读

零、embedding函数

让我们考虑一种不包含prefix_soft_prompt的情况，从999行的embedding准备开始。

llm小知识-embedding操作：所谓embedding，一些文章中将其翻译为词向量嵌入，个人认为这种翻译并不算恰当。从功能的角度来说，embedding层负责将每个ids变成一个长度固定的向量，从设计思想的角度来说，embedding层负责将每个ids映射到一个固定的向量空间中，通过这种表达方式的转换，就可以帮助模型来理解不同词向量之间的关系，方便后续的计算操作。

在1001行做好了对prompt_param的设置之后（在我们考虑的这种情况下是没有prompt的），就进入到了1008行的embedding层的操作了，一起来看一下这个kernel是怎么做的。

kernel定义在gpt_kernels.cu中，invokeInputIdsEmbeddingLookupPosEncoding函数，从120到160行，根据不同的prompt情况进入了不同的操作，任务划分为grid依据具体的输入ids进行划分，block依据隐藏层大小hidden_size进行划分。我们直接看146行的情况

template<typename T>
void invokeInputIdsEmbeddingLookupPosEncoding(T*                    from_tensor,int*                  output_ids,const T*              embedding_table,  // can also be inputs_embedsconst T*              pos_table,pPromptTuningParam<T> prompt_param,const int*            input_ids,const int             start_step,const int             length,const int             max_length,const int             batch_size,const int             hidden_units,cudaStream_t          stream)
{dim3       grid(min(batch_size * length, 65536));dim3       block(min(hidden_units, 512));const bool has_output_ids = output_ids != nullptr;FT_CHECK(!(has_output_ids && input_ids == nullptr));if (has_output_ids) {if (prompt_param.use_request_p_prompt_embedding) {WORD_POS_EMBEDDING_LOOPUP_KERNEL(true, 2);}else if (prompt_param.p_prompt_tuning_batch_weights != nullptr) {WORD_POS_EMBEDDING_LOOPUP_KERNEL(true, 1);}else {WORD_POS_EMBEDDING_LOOPUP_KERNEL(true, 0);}}else {if (prompt_param.use_request_p_prompt_embedding) {WORD_POS_EMBEDDING_LOOPUP_KERNEL(false, 2);}else if (prompt_param.p_prompt_tuning_batch_weights != nullptr) {WORD_POS_EMBEDDING_LOOPUP_KERNEL(false, 1);}else {WORD_POS_EMBEDDING_LOOPUP_KERNEL(false, 0);}}
}

在WORD_POS_EMBEDDING_LOOPUP_KERNEL这个宏，其实是重新定义了函数start_id_embedding_position_lookups_kernel，这个函数定义在32行，在进入这个函数之后，如果output_ids有分配好空间的话，可以先看64行这个处理过程，按照index将input ids拷贝到output_ids中

						else {const int seq_id   = index % max_length;const int batch_id = index / max_length;if (seq_id < length) {output_ids[seq_id * batch_size + batch_id] = input_ids[index];}}

之后，是正式的embedding处理过程：

				const int word_index      = index / hidden_units;const int word_index_row  = word_index / length;  // batch_idconst int word_index_col  = word_index % length;const int real_word_index = word_index_row * max_length + word_index_col;const int step            = start_step + word_index % length;const int col_index       = index % hidden_units;const int input_id        = input_ids == nullptr ? real_word_index : input_ids[real_word_index];const int prompt_id       = input_id - prompt_param.p_prompt_tuning_id_start;T         embedding       = (T)0.0f;if (PROMPT_SRC > 0 && prompt_id >= 0) {if (PROMPT_SRC == 1) {// from loaded prompt embedding tablesembedding =prompt_param.p_prompt_tuning_batch_weights[word_index_row][prompt_id * hidden_units + col_index];}else {// from request prompt embeddingembedding =prompt_param.request_prompt_embedding[word_index_row * prompt_param.request_prompt_max_length * hidden_units+ prompt_id * hidden_units + col_index];}}else {embedding = embedding_table[input_id * hidden_units + col_index];}T pos_embed        = pos_table == nullptr ? (T)0.f : pos_table[(step - 1) * hidden_units + col_index];from_tensor[index] = embedding + pos_embed;

在不处理prompt逻辑时，其实embedding的处理过程非常简单，就只有embedding = embedding_table[input_id * hidden_units + col_index];这么一行，也就是一个简单的查表和赋值操作。完成赋值之后，如果有pos embedding位置编码，再取出位置编码上对应的值，加到embedding上，并赋值给输出的地址from_tensor上。

一、prelayer_norm和layernorm

完成embedding后，让我们回到forward主函数，在完成embedding之后，输入的隐藏状态就已经存储在context_decoder_input_buf这个地址中了。对于gpt以及其相关变种模型来说，可能会有pre_layernorm的存在，根据具体的模型config设置，可能会经过一次invokeGeneralLayerNorm的操作。

llm小知识-layernorm的位置：在gpt以及类似bloom模型中，可能会在进入decoder结构前和decoder结构完成计算后进行layernorm。在llama系列和变种模型中其实这两个位置的layernorm大部分已经被取消了（需要prelayernorm的话为什么不在embedding的时候就把数值计算好呢？）

layernorm的具体实现函数在layernorm_kernels.cu中，具体算法在论文https://arxiv.org/abs/1607.06450中，用公式来表达的话如图所示：

在这里插入图片描述

由于ft中为了尽可能做到性能优化，在layernorm算子的实现中做了很多不同case的dispatch操作，在这里不对这部分的dispatch操作进行详细分析，直接看layernorm_kernels.cu的25行的具体的计算函数，整体流程也基本符合公式中的计算过程，116和117行计算出均值和标准差，124行计算出layernorm的计算结果，并最终赋值到normed_output的结果中。这里不进行详细的展开讲解，建议有cuda基础的朋友们可以稍微展开研究一下。

二、context-decoder的准备

在进入context-decoder层的计算前，还需要进行两个准备，第一个是建立Attention mask以帮助计算注意力得分。

由于gpt的注意力机制是前向的，所以在计算Attention得分时，当前词只会和这个词以及之前的词计算注意力，而在context-decoder中，所有词的Attention得分又是一起计算的，所以需要计算生成一个mask来帮助将当前词的后面的词给掩码mask掉。

生成mask的kernel代码在gpt_kernels.cu的385行，任务划分为按照batch_size进行划分，每个任务起256个线程

template<typename T>
void invokeBuildDecoderAttentionMask(T*           attention_mask,const int*   sequence_lengths,const int*   prefix_prompt_lengths,const int    batch_size,const int    max_seq_len,const int    max_prompt_length,cudaStream_t stream)
{if (max_prompt_length == 0) {buildDecoderAttentionMaskKernel<T, false><<<batch_size, 256, 0, stream>>>(attention_mask, sequence_lengths, prefix_prompt_lengths, max_seq_len, max_prompt_length);}else {buildDecoderAttentionMaskKernel<T, true><<<batch_size, 256, 0, stream>>>(attention_mask, sequence_lengths, prefix_prompt_lengths, max_seq_len, max_prompt_length);}
}template<typename T, bool PREFIX_PROMPT>
__global__ void buildDecoderAttentionMaskKernel(T*         attention_mask,const int* sequence_lengths,const int* prefix_prompt_lengths,const int  max_seq_len,const int  max_prompt_length)
{// sequence_lengths: [batch_size]// attention_mask: [batch_size, 1, max_seq_len, max_seq_len + max_prompt_length]const int max_prompt_seq_length = max_seq_len + max_prompt_length;const int mask_size_per_seq     = max_seq_len * max_prompt_seq_length;attention_mask += blockIdx.x * mask_size_per_seq;const int seq_length    = sequence_lengths[blockIdx.x];const int prompt_length = PREFIX_PROMPT ? prefix_prompt_lengths[blockIdx.x] : 0;for (int i = threadIdx.x; i < mask_size_per_seq; i += blockDim.x) {int row_id = i / max_prompt_seq_length;int col_id = i % max_prompt_seq_length;if (row_id < seq_length && col_id <= (row_id + prompt_length)) {attention_mask[i] = (T)(1.0f);}else {attention_mask[i] = (T)(0.0f);}}
}

可以看见，在不考虑prompt的情况下，mask_size_per_seq就是最大seq长度的max_seq_len的平方（一个正方形矩阵的大小），在循环中，计算行数id和列数id，当行数id比seq长度小，且列数id小于行数id时，就将矩阵对应位置设置为1，否则设置为0，也就是生成了一个下三角矩阵。

在这里插入图片描述

按照之前的逻辑，位置为x的id，就只需要与第x行对应位置为1的id进行attention得分计算，这样就帮助实现了前向注意力得分的计算。

计算完成后，在forward函数的1049到1088行，是为了context-decoder输入和输出做tensor设置。在input_tensors中，首先设置了输入的隐藏状态地址，生成好的mask的地址，以及输出batch长度数组的地址。根据是否使用共享context决定是否插入需要的参数，如果attention还需要加入线性偏置，那么再多增加一个线性偏置量。线性偏置量在829行就已经计算好了，直接传入就好

llm小知识-线性偏置：线性偏置https://arxiv.org/pdf/2108.12409也就是我们所说的Alibi，其目的是为了解决训练和推理时文本长度不一致的问题，它取消了pos embedding的位置信息注入，而是在 Attention 中计算 QK 的值后面加入一个偏置常量，来达到注入位置信息的效果。这个偏置常亮一般也是提前直接计算好的，实际上就是计算前面 token 距离当前 token 的距离：[-n, -n-1, …, -1, 0]

在output_tensors中，需要设置好我们的context-decoder的输出，以及后续计算需要的kv cache。

完成设置后，1346行，调用context_decoder来完成输入部分的decoder推理计算。

下一回预告

下一回我们会暂时离开forward函数，进入context-decoder，来看看context-decoder是如何完成设置和计算的。在讲解完context-decoder之后，再回到forward函数中。contextdecoder中就会实际的涉及到attention计算、layernorm以及ffn的计算。