CS336讲座2pytorch笔记

本讲聚焦于pytorch和资源核算FLOPs计算Question: How long would it take to train a 70B parameter model on 15T tokens on 1024 B100s?有一个公式浮点运算次数FLOPs等于参数数量乘以token 数量再乘以六h100_flop_per_sec 1979e12 / 2mfu 0.5flops_per_day h100_flop_per_sec * mfu * 1024 * 60 * 60 * 24 #1024张显卡days total_flops / flops_per_day143训练模型需要BF16不能FP16否则可能出现NANBF16动态范围比FP16大很多和FP32一样精度会更高一些。矩阵乘法FLOPsx torch.ones(B, D, devicecuda_if_available())w torch.randn(D, K, devicecuda_if_available())y x w实际浮点运算数就是actual_num_flops 2 * B * D * K乘以2是因为每个数要做一次乘法和加法实际上加法比乘法少一次不过可以忽略这个细节Model FLOPs utilization (MFU)MFU 就是实际 FLOPS 除以标称 FLOPS为什么之前的公式里要乘以mfu 0.5的原因因为显卡虽然有极高的“标称 FLOPS”promised FLOP/s但实际训练中由于显存读写、计算内核启动等开销实际 FLOPS通常达不到标称值。MFU (0.5)就是用来打折扣的系数代表你实际用到了 50% 的理论算力。那么为什么MFU只有0.5MemoryHow to compute a thing:1. Send inputs from memory to accelerator2. Perform computation3. Send outputs from accelerator to memoryHow long does this take?Depends on two things:1. Accelerator speed (FLOP/s)2. Memory bandwidth (bytes/s)背后有两个核心原因显存带宽瓶颈Memory Wall和通信与同步开销。我们拆开来看核心原因 1显存带宽瓶颈做这道题的这道题请看列出的步骤从内存显存发送输入到计算核心执行计算从计算核心发送输出回内存这是一个完美的Compute-Bound vs Memory-Bound问题的图解。大模型训练绝大部分时间都在做矩阵乘法比如图片 5 里的x w。计算速度FLOP/s算力GPU 的计算核心非常快。H100 的理论算力接近 2000 TFLOPs即每秒 2×10152×1015 次操作。显存带宽Bytes/s搬运速度显存把数据权重和激活值喂给计算核心的速度。H100 的带宽大约是 3.35 TB/s即每秒 3.35×10123.35×1012 字节。问题来了算力是显存带宽的几百倍。虽然显卡的理论峰值极高但在做矩阵乘法时计算核心经常处于“干等”状态——它在等显存把数据搬运过来或者等算完的结果被写回显存。只要显存带宽跟不上GPU 的算力就无法跑满。虽然现代 GPU 通过 Tensor Core张量核心和各种缓存优化尽量掩盖这个问题但依然无法完全消除。这是 MFU 很难超过 0.6 的根本物理限制。核心原因 2卡间的通信开销因为你有 1024 张卡你算的是1024 张卡组成的超级计算机。当显卡数量超过 1 张时为了协同训练这 70B 的模型必须使用分布式并行如数据并行、张量并行、流水线并行。每训练一个 batch这 1024 张卡之间必须通信同步比如梯度汇总、中间激活值传递。同步的时间是“浪费”掉的在 1024 张卡互相传数据的那几毫秒里所有的 GPU 计算核心都是空闲的。它们没有在做浮点运算只是在收发数据包。卡的规模越大通信开销就越大MFU 通常也会更低。做 1024 张卡的 MFU 达到 0.5 其实已经是顶级工程团队才能做到的成绩了。核心原因 3图上的忽略项算子与内存碎片底层算子Kernel效率不是所有的计算都能利用 Tensor Core 加速。比如 LayerNorm、GeLU 激活函数、Dropout 等小算子它们主要靠通用计算核心CUDA Cores跑这些核心的算力远低于 Tensor Core。显存碎片和管理当显存占用达到极限时系统不得不频繁地申请和释放显存块这也会消耗时间。CPU 与 GPU 的通讯CPU 需要准备数据给 GPU 吃中间有数据传输延迟。