背景

ChatGPT出现后，惊人的效果完全颠覆了业界人员包括笔者的认知，抛开其模型细节层面的因素，已公开的训练方法，需要巨量的数据和计算资源，门槛非常高。本文基于公开资料，希望以量化方式分多篇介绍ChatGPT的分析结论，具体内容包含以下三篇，本文为计算量分析篇。

ChatGPT模型结构为Transformer，下面对计算细节拆解分析。

计算量分析

变量定义

模型参数数量浮点运算次数，用来表示计算量

Embedding

embedding层是lookup操作，输入是词序列，输出是词序列embedding，浮点操作可忽略。

Transformer Blocks

transformer block的计算图如下，每个transformer block主要包含四部分，既multi-head attention和mlp，以及两个add&norm。

图片替换文本

Multi-head Attention

multi-head attention结构如图：

图片替换文本

变量定义： 样本样本的序列长度数量、矩阵的维度矩阵的维度和往往一样，都表示为数量数量 multi-head attention的计算逻辑为： 其中，。

  • 矩阵乘法算子中包括浮点乘法和加法，的输入激活Tensor是，在第一个block是词本身embedding和position embedding之和，在其他block是上游block的输出，的形状为，的形状为，因此计算量为，和同；
  • Softmax算子的输入Tensor为和，形状为，和相乘的计算量为，操作的计算量为，操作计算量为，归一化操作计算量为，共；
  • 的输入Tensor为和，计算量为；
  • 每个head的存计算量为，共个head，存储占用为：；
  • 输入为激活和参数，计算量为。

综上，multi-head attention的总计算量为：

MLP

multi-head attention后面，接两层的全连接网络，计算逻辑为： 其中，。

  • 矩阵乘法算子的输入激活Tensor 是multi-head attention输出，形状为，形状为，计算量为；
  • 加法部分输入是激活和参数，计算量为；
  • 输入和输出都是，计算量简单计为；
  • 的输入是激活和参数，计算量为；
  • 加法部分输入是激活和参数，计算量为；

综上，mlp部分的总计算量为：

Add&Norm

两个add部分的计算量为。每个norm中均值计算量为，标准差计算量为，归一化计算量为 所以，add&norm部分的存储占用为：

Transformer Blocks计算量

综上，每个transformer block计算计算量为： 共个transformer block，总计算量为：

推理阶段

在推理阶段，计费量为： 按GPT3 175B配置计算： 按的话，约为。每个token每个参数的计算量为： 每个token的计算量约为。

训练阶段

在训练阶段，需要反向传播，反向传播计算量和前向一致。另外还有梯度更新，Adam需要更新两个累积值和参数，计算量简单计为，因此训练的计算量为： 按GPT3 175B配置计算： 按GPT3 175B配置和计算，训练的总计算量为，与GPT3论文非常接近。在300张A100 80G卡集群训练，GPU利用率40%情况下，需要耗时。每个token每个参数的计算量为： 每个token的计算量约为。

图片替换文本

参考资料