cs336: From 0 build LLM
lecture2: pytorch and source
本节课主要介绍了pytorch中的einop操作(但是我不喜欢用)，以及tensor based的资源计算，由于之前对这方面的知识完全不了解，所以会更加详细的记录这部分的内容
Question: 用1024张NVIDA B100的GPU，训练一个70B参数的LLM，训练数据量为15T tokens，估算训练时间
在开始之前我们引入FLOPs的概念，也就是浮点运算次数，这是非常常用的，同时也要区分另一个概念FLOP/s，指的是每秒钟能进行的浮点运算次数，TFLOP/s = 1e12 FLOP/s
训练总计算量(total flops) = 6 * 参数量 * tokens数
这里有个小问题，参数量 * tokens数很好理解，但为什么需要乘6，在经典的Transformer架构中，我们考虑前向和反向转播两个过程的矩阵运算：
在forward中，考虑一个线性层$y=Wx+b$，一次乘法一次加法算两次FLOPS，每一个参数在训练过程中处理每一个token都是2 FLOPS，总计也就是$2 * N * D次
接下来看反向传播的过程，相对于forward过程，这里要计算对$W,x$的梯度，每一次梯度计算可以看作是2 FLOPS，加起来就是$4 * N * D$ 次，所以综合来看也就是训练过程的次数是$6 * N * D$ FLOPS，只推理的话是$2 * N * D$ FLOPS
通过查阅资料我们得知H100每秒钟可进行的FLOP数为：1979e12 / 2，MFU = 0.5
(MFU = 实际模型 FLOPs 吞吐 / 硬件峰值 FLOPs 吞吐，相当于实际情况跟理想情况下的比值)
有了以上数据，我们可以轻松计算出H100实际一天能进行的浮点运算的次数，进而得到总训练时间，答案约为143 days
确实huge LLM的训练是非常耗时的
Question: 使用AdamW优化器在8卡H100上能训练出来的最大的LLM有多大
回答这个问题我们需要从显存的角度考虑，8卡H100的总显存为8 * 80 G = 640 G
然后我们考虑显存都需要存什么东西(AdamW优化器中)：
- fp16 / bp16参数 2 bytes
- fp16 / bp16梯度 2 bytes
- 一阶动量 m 4 bytes
- 二阶动量 v 4 bytes
也就是总共12 bytes，最后拿总显存除12 bytes，得到53B参数量
但是这与真实的情况相聚甚远，根因在于这是一个及其理想的情况，有许多因素没有考虑，这里暂时不探讨这个问题
tensor basics
但我们讨论tensor的存储空间之前，首先要搞清楚存储类型，最经典的是fp32(torch里面默认的)
考虑fp32精度下面的例子：
x = torch.zeros(4,8)
assert x.dtype = torch.float32
assert x.numel() = 32
assert x.element_size == 4
assert memory_use = 4 * 32这说明了一个fp32的浮点数占用4个字节，存一个4 * 8的矩阵就需要吃掉128 bytes
后面讲了各种更小的精度的数据类型，如fp16 / bf16 / fp8 …
如何好的应用它们优化模型训练应该是深层的MLsys，这篇暂时不会考虑
FLOPs of operations
对于矩阵乘法运算，实际上是一个$O(B * D * K)$的操作，下面给出一段代码实例：
x = torch.ones(B, D, device = cuda)
w = torch.ones(D, K, device = cuda)
y = x @ w考虑一下矩阵运算的细节就能搞清楚，不细展开了
Why MFU = 0.5 is a good value
前文我们提到了MFU的概念，这类似于理论 / 上限，接下来我们深入分析为什么会存在这一现象
考虑一个很简单的问题，就是数据是在内存里面存储的，但是加速运算的话需要搬到GPU上面去，然后算完的数据再拿回来，所以不只是计算开销，内存的带宽影响也很大
举一个很简单的例子，如果是1张4090的机器，你去跑代码，在CPU上单核情况下会比GPU慢500倍左右，而即便是多核也会比GPU慢上40倍，但这不意味着任何程序都是GPU上快，假设一个超级简单的事情，比如你就需要算一个
a + b(int)，那GPU加速的时间远远不如把数据移动到GPU上的时间，这就是内存带宽的重要性
ok接下来我们会用例子来看这个问题，还是以H100为例，给出数据：
- 计算速度：compute_speed = 1979e12 / 2
- 内存带宽：memory_speed = 3.3e12
好来看一个relu的例子
def exp_relu():
n = 1024 * 1024
x = torch.ones(n, dtype = torch.bfloat16, devide = cuda)
y = torch.relu(x)
bytes = (2 * n) + (2 * n)
flops = n
memory_time = bytes / memory_speed
compute_time = flops / compute_speed来解释一下这个，每一个bfloat16的数据占用2 bytes，然后x需要移动到GPU上，算出来的y需要移回内存中，总计两个2 * n
relu对每一个数据点做max(0,x)操作，总共n次浮点运算，接下来二者分别除以对应的速度，计算出耗时
在这个例子中最后算出来用于内存读写的耗时远多于计算，故总速度受限于内存，也就是memory-bottenlock的
其实对于某一种硬件，考虑具体会受限于什么的时候非常好办，你先算一下compute_speed / memory_speed，然后算某一个操作二者的bytes数目比，然后对一下就出来了
用同样的思路，你可以搞清楚矩阵乘法到底受限于什么，注意，除了硬件指标，还要留意数据精度
后面的内容类似，lecture2的笔记至此结束！
