CS336讲座2pytorch笔记

发布时间:2026/9/19 1:37:25
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 吃中间有数据传输延迟。

关于本文作者

来自尧图内容编辑团队

尧图内容编辑团队 内容团队

尧图内容编辑团队

本文由尧图网络内容编辑团队执笔。团队由资深项目经理、前端工程师与设计师组成,所有内容均来自亲手交付的真实项目,先讲清问题、再给出可落地的解法。尧图深耕北京网站建设十年,服务过京华建材集团、智造科技等各行业客户,把一线经验沉淀为可复用的行业观察。

  • 十年建站经验,覆盖建材、制造、服务、文创等
  • 项目经理把关选题与事实准确性
  • 工程师与设计师联合撰写专业细节
  • 统一编辑规范,保证文风与排版一致
  • 每月复盘转化数据,迭代选题方向

延伸阅读

相关资讯与近期热门内容

深度阅读推荐

建站决策前值得细读的三篇

网站改版的5个关键决策
2024-08-12

网站改版的5个关键决策

什么时候该改版、改到什么程度、如何避免流量掉光,京华建材集团改版复盘给出答案。

获取专属建站方案

看完文章,把您的行业与预算告诉我们,免费获取一份量身定制的官网建设方案与报价。

立即免费咨询