从FasterTransformer源码解读开始了解大模型（1.1）一个decoder-only的模型长啥样

写在前面的话

对于一个没有接触过LLM的初学者来说，如果想要了解一个大模型的推理框架，首先应该知道大模型整个的工作原理是怎样的，知道transformers的结构是怎么生成词的，否则很容易会在读代码的过程中越读越迷糊，进入一种“我在哪里？我在看什么？这一段是做什么的”困惑状态，所以在此先推荐从来没有接触过相关知识的同学去读一下Attention is all you need的论文原文。当然，对于一些初学者来说，读完这篇论文也会陷入一种半懂不懂的困惑状态，但是会在心中搭建起一个隐隐约约的初步的形状，在后面再回到工程代码的阅读中，这个隐隐约约的形状会被一步一步填充起来，变成一个清晰的样子。这里的FT源码解读的第二篇文章会直接从结构开始讲起，就不会讨论那么多“为什么这个结构是有效的”“为什么多头注意力要这样做设计”的问题了。

在写这篇文档的过程中，我发现自己对LLM框架的认知也并没有十分地清晰明了，如果整个解读过程中有不正确之处或者不够清晰的地方，也希望老师同学们能够指出和讨论 Orz

零、最初的样子

让我们回顾一下最初始的时候，Attention is all you need中所描述的encoder-decoder结构是怎样的：

在这里插入图片描述

这里是原论文中的结构，刚开始看是否有些半懂不懂？（个人认为这张结构解释图对初学者来说并不能算得上清晰明了hhhh）拆开来看的话，整个结构分为两块，左边的encoder和右边的decoder，在encoder的这半边，不做任何的ids的输出，只负责处理输入数据。真正输出ids的，则是decoder的部分，负责自回归地进行处理，往外一个一个吐词儿(从output probabilities中选出ids)，在一些早期的大模型，如Bart、T5等等，依然遵循了encoder-decoder的结构。

一、decoder-only的样子

经过一段时间的代码模型的发展，很多流行的模型，比如gpt2，还有现在最流行的llama系列等等，都采用了decoder-only的样子，而从结构上来看，llama比gpt2，bloom这种还要更简单一些，但是却变得更加普遍和流行。所谓大道至简，大概就是这个道理罢。

decoder-only大概长成下面这个样子
在这里插入图片描述
decoder-only只保留了decoder的MHA的处理结构（MHA-Add&Norm-FFN-Add&Norm），对于输入来说，是一整个扔进去decoder结构的处理，对于生成阶段，依然是decoder结构（所以才叫decoder-only）。在FasterTransformer中，它将输入阶段的decoder称作ContextDecoder阶段，即上下文解码阶段（在有些框架中则习惯性地照旧称为encoder阶段但实际上处理有所区别，也有称为prefill阶段的），而实际自回归产生ids输出的阶段依然叫做decoder阶段。

可以看一下图中的流程，输入的ids经过embedding之后，会为句子中的每一个id生成一个hidden states（隐藏状态，在整个layer里的计算过程中其实都是以hidden stats隐藏状态的样子参与到运算中）。

之后隐藏状态会进入到layers中进行计算，每一层layer（在HuggingFace的transformer源码中称之为blocks）都包含了自己的layernorm、MHA、add residual、FFN等结构。在完成某一层(layer x)的计算后，输出得到的hidden states会作为下一层(layer x+1)的输入再次进入layer运算，直到所有的层计算完毕，最后一层输出的hidden states再经过一个lm head（线性层），产生一个跟词表长度一样长的raw logits数组，对这个raw logits进行softmax和采样选取，最终就可以得到这一次生成的结果output id了。

如果这个output id不是end id（结束词），可以将它作为输入再次喂给整个decoder结构，就实现一个自回归生成的过程。

llm小知识-模型的差异：这里的decoder-only模型并不是唯一标准化的，有些模型会在一些地方少掉第一次 pre-layernorm或者替换layernorm为其他的归一化操作（比如llama使用的是RMSnorm），有些是会增加或减少layernorm中的bias，有些模型则是会插入一些差异性的norm（比如google的gemma），但是总体来说MHA和FFN这两个大模块都不会缺失或者产生大的变化，在这两个模块计算完成后的add residual操作和norm操作也很少会做一些大的结构上的变化

在ft实际的处理逻辑中，在context decoder阶段，其中MHA依然遵守着前向的规则，即每一个词都只会跟自己前面的词计算注意力得分，并且将MHA计算注意力过程中所得到的K和V给记录下来（即K/V Cache）以在后续的计算过程中节省计算时间，在完成了context-decoder时最终会得到一个logits，这个logits将会成为生成的第一个词，也是自回归decoder阶段的第一个输入的词，由此，整个自回归生成就开始循环。

有些新入门的同学可能会疑惑，隐藏状态是如何生成的？MHA或者FFN这些必要模块具体是怎么进行计算的？算出来又都是什么样的结果呢？这里如果着急可以直接去看HuggingFace的transformer的源码，其中每个modeling都会有自己的对应实现，如果不着急的话，后续的源码解读都会一一给出答案。

二、从源码角度来看结构

我们再从代码的角度来看一下，FT的源码在搭建一个模型的时候，是否符合我们刚刚描述的结构.

打开ft的源码，进入gpt的目录src/fastertransformer/models/multi_gpu_gpt，可以看见从model层，gpt所有需要的模块和代码：
在这里插入图片描述

其中模型的主文件叫做ParallelGpt，它负责一整个模型的组装和模型的计算。就是我们之前所描述的一整个的decoder-only结构，以及其他的一些内容。

可以看见这里分为两个部分，负责处理输入部分的叫做ParallelGptContextDecoder，就是我们之前所描述的负责做context decoder的结构，负责做生成的叫做ParallelGptDecoder，由它来完成output ids的生成。

另外，还有两个weight相关的模块代码，ParallelGptWeight对应着一整个的ParallelGpt计算所需要的权重，由这个模块负责进行管理。而ParallelGptWeight中则通过一个数组来管理每一层（layers）的单独的权重，每一层单独的权重是由ParallelGptDecoderLayerWeight来组成和管理。

而在ParallelGptContextDecoder以及ParallelGptDecoder中所使用到的比如MHA模块，以及FFN模块等等，则全都在另一个目录下，src/fastertransformer/layers这个目录下进行管理。在这个目录下你可以找到所有关于单层layer中需要的计算和处理的模块，包括不同模型的MHA或者FFN等等。

从上面也可以看见，FT的文件组织基本符合我们所描述的模型结构，模型（models）下面有自己的结构，每个结构中的层（layers）中所包含的结构则由另外一个layers文件夹中的代码文件来进行管理和实现。