【infra之路】阶段二 · 模块三:显存账 + 多卡通信原理

AI Infra 学习路线 · 阶段二 · 模块三目标:算清训练时显存被什么吃掉 在卡上实测验证 理解多卡为何要通信环境:WSL2 RTX 5060 Ti PyTorch cu130这是 Infra 面试的硬通货,和 ML 底子结合最紧0. 基准单位1 个 float32 数 4 字节。所有显存账 “有多少个数 × 4”。(混合精度用 float16/bf16 2 字节,各项直接减半 —— 这就是混合精度省显存的原理。)1. 训练显存的四项构成 ★★(面试硬通货)以 Adam float32 为例,P 参数个数:项大小性质说明① 参数 weights4P固定模型权重本身② 梯度 gradients4P固定backward 给每个参数算一个梯度,数量参数③ 优化器状态 optimizer state8P(Adam)固定Adam 给每个参数存 2 个状态(动量方差) 2×4P。 参数的两倍④ 激活值 activations随 batch/深度/序列变可调前向各层中间输出,留着等反向用前三项固定开销 4P4P8P 16P ≈ 模型大小的 4 倍。关键推论:SGD(无动量)优化器状态0;Adam 优化器状态8P,是显存大头,最易被低估。1B 模型(参数 4GB)→ 固定开销 16GB,已超 8G 卡。7B 模型训练 ≠ 7B 显存:固定开销(float32)≈ 7B×16 112GB,还没算激活值 → 这正是混合精度/量化/ZeRO 要解决的。训练 vs 推理:推理只需参数(4P),无梯度、无优化器状态,激活值也少(算完一层即丢,不留反向)。所以训练 16P,推理 ~4P。这是阶段三推理优化的前提。2. 激活值怎么算(四项里唯一没有简洁公式的)激活值 前向时每层的输出,被存下来等反向用。单层激活值 batch × 该层输出维度 × 4 字节例:Linear(2000,2000) batch64 → 输出形状 [64,2000] → 64×2000×4 ≈ 0.49 MB。为什么实测比只算每层输出多:反向传播还需要保留其它中间量(如算权重梯度要用到该层的输入)。所以实际保留的中间张量比一层一个输出多。故激活值 ≈ (每层保留的中间张量数之和) × batch × 维度 × 4,依赖网络结构细节,没有像参数那样干净的公式。不必精确手算,抓三个直觉:正比于batch size(batch 翻倍→激活值翻倍)。这是 OOM 时先调小 batch 的原因。正比于深度×宽度×序列长度。小扁模型激活值是零头;大深模型/LLM(层多、序列长)激活值是大头。和参数量 P 无关。前三项用 P 算,激活值用数据流过网络的中间结果大小算 —— 这是它和前三项的本质区别。一句话:前三项是模型这个静态物体多大(跟模型走,死的);激活值是数据流过模型时沿途存多少中间结果(跟数据走:batch/序列长度,活的)。3. PyTorch 实测验证 ★(亲手在 5060Ti 上测)工具:torch.cuda.memory_allocated()返回当前 GPU 占用字节。策略:每步后测一次,看增量对应哪一项。importtorch,torch.nnasnndefmb(b):returnb/1024/1024defshow(s):torch.cuda.synchronize()print(f{s:30}{mb(torch.cuda.memory_allocated()):8.2f}MB)modelnn.Sequential(*[nn.Linear(2000,2000)for_inrange(20)]).to(cuda)Psum(p.numel()forpinmodel.parameters())show(1.加载模型(参数))# 增量 ≈ P×4xtorch.randn(64,2000,devicecuda);ytorch.randn(64,2000,devicecuda)outmodel(x);lossnn.MSELoss()(out,y)show(2.前向后(激活值))# 增量 激活值(随 batch)loss.backward()show(3.反向后(梯度))# 增量 ≈ P×4opttorch.optim.Adam(model.parameters());opt.step()show(4.优化器step后(优化器状态))# 增量 ≈ P×8 参数×2实测结果(80M 参数,20 层,Adam,float32)步骤batch64batch256理论1 参数320 MB320 MB(不变)P×4305 ✓2 前向(激活值增量)43 MB78 MB(随 batch 涨)无简洁公式3 反向(梯度增量)~343 MB~392 MBP×4305 ✓4 step(优化器增量)~640 MB~640 MBP×8611 ✓,参数×2最终总占用1346 MB1352 MB16P1221 激活框架开销 ✓验证到的规律:参数 batch64 与 256 完全相同(320320)→固定开销只跟模型走,与数据量无关。优化器状态 ≈ 参数×2(640≈320×2)→ Adam 铁律,实测命中。激活值随 batch 明显涨(深模型才看得清;之前扁模型只涨 5MB)→激活值占比取决于模型形状:扁模型参数主导,深模型激活值主导。模型越大,实测越贴近理论(框架固定开销~15MB 被摊薄)。注意:实测会比理论略多(PyTorch 显存缓存、反向临时量、框架开销)。验证的是量级和趋势,非分毫不差。4. NCCL 多卡通信原理(概念,单卡不可实跑)为什么要多卡固定开销 16P,7B 模型 112GB(float32),单卡(顶级 80GB)装不下 → 多卡分担。数据并行(Data Parallelism,最基础常用)每张卡放一份完整模型副本,各处理一部分数据(4 卡 × batch 256 → 每卡 64 条)。各卡只看 1/4 数据 → 算出的梯度不同 → 若各自更新,模型会散掉。必须汇总梯度成一致的,再让所有卡用一致梯度更新 → 保持是同一个模型。All-Reduce(分布式训练最重要的通信操作)把所有卡的梯度加起来(reduce),再把结果发回每张卡(all)→ 所有卡拿到相同的汇总梯度。每个训练 step 都做:前向 → 反向 → All-Reduce 同步梯度 → 更新。通信与计算交织在每一步。NCCL(NVIDIA Collective Communications Library,读 “nickel”)All-Reduce 是做什么,NCCL 是 NVIDIA怎么做得快的库 —— 在 GPU 间高效实现,跑在 NVLink/InfiniBand 等高速连接上。PyTorch 的 DistributedDataParallel(DDP)底层梯度同步就是 NCCL 跑 All-Reduce,通常不用直接碰,但多卡训练快慢很大程度取决于这层通信效率。关键直觉:多卡不是线性加速每步都 All-Reduce 同步梯度,梯度量参数量。模型越大每步同步越多;卡越多协调成本越高。通信时间接近/超过计算时间时,加卡不划算—— 对应训练吞吐 vs 通信开销的核心矛盾。5. 阶段二主线:减少数据搬运(单卡 → 多卡)模块二(单卡内):全局显存↔共享显存的搬运,tiling 减少跑慢仓库→ 百倍加速。模块三(多卡间):卡与卡之间的梯度搬运,All-Reduce 可能成瓶颈。减少数据搬运从单卡贯穿到多卡。瓶颈常在数据流动,不在算力—— 这个直觉接着用到阶段三推理/训练优化(KV cache、量化、PagedAttention 本质都是管显存、减搬运)。阶段二完整收获模块一 GPU 架构(吞吐vs延迟/SIMT/SM/显存层级)→ 模块二 写 CUDA(线程模型/矩阵乘法/tiling 优化/测量陷阱)→ 模块三 显存账(四项构成/实测/NCCL)。GPU 这一层从硬件到代码到资源账,从里到外打通。从会调包到懂里面在发生什么。