上一篇 TP 写完之后,我对 Megatron 的 Column Parallel + Row Parallel 已经比较清楚了:前一层按输出特征切分,后一层按输入特征切分,两层之间不需要 AllGather,只在最后做一次 AllReduce 把部分和加起来。
最近在看 MIT 6.S191 Lecture 9 讲 Context Parallelism 的时候,碰到了一个词:Causal Attention。
我对 Self-Attention 的公式并不陌生——之前那篇已经把 QKV 和 routing score table 拆得比较清楚了。但 Causal Attention 这个概念我一直没有认真想过。回头看,我写 Attention 公式、写 Multi-Head、写 FlashAttention 的时候,脑子里跑的都是"所有 token 都能看到所有 token"的 Full Self-Attention。
起因 # 最近在过 MIT 6.S191 和分布式训练相关的材料,各种 Norm 反复出现。LayerNorm 在写 Transformer Block 那篇时已经接触过,BatchNorm 在做 CV 的那几年天天调,RMSNorm 在看 LLM 代码时到处都是。单独看每个都不算陌生。
起因 # 上一篇搞清楚了 Pipeline Parallelism 的调度机制:GPipe 先 Forward 后 Backward,1F1B 通过交替调度降低 Activation 暂存压力,但理想条件下 Bubble Time 相同。
起因 # 上一篇把 ZeRO 的三个 Stage 走了一遍之后,我以为关于参数分片这件事已经想得差不多了——ZeRO-3 / FSDP 把参数、梯度和 Optimizer States 全部分片存储到多张 GPU 上,单卡 Model States 从 112 GB 降到约 28 GB,问题解决。
起因 # 上一篇 Tensor Parallelism 写完之后,我对 TP 的基本原理有了底:把单个算子内部的矩阵乘法拆到多张 GPU 上执行,每张卡只算部分输出,最后通过 AllReduce 聚合。TP 分片的是计算本身。
起因 # 上一篇搞清楚了 AllReduce、ReduceScatter 和 AllGather 三种通信原语,末尾提到 DDP 和 ZeRO Stage 2 选择了不同的通信原语组合,根本原因是参数更新的责任分配不同。
起因 # 之前几篇文章从模型显存组成一路算到 7B 模型的训练账本,结论很直接:BF16 mixed precision + AdamW 的 Model States 就有 112 GB,一张 80 GB A100 连放都放不下。Checkpointing + Offloading 两个一起用可以勉强塞进去,但训练速度要被 PCIe 传输和额外重计算拖累。
起因 # 上一篇讨论完 Checkpointing 和 Offloading 之后,结尾留了一句:单卡优化的代价已经不可接受时,下一步自然是考虑让多张 GPU 分担训练状态。
起因 # 上一篇算完了 7B 模型的训练显存账本。结论很直接:BF16 mixed precision + AdamW 的 Model States 就有 112 GB,加上 30 GB 的 Activations,单卡至少需要 145 GB——一张 80 GB 的 A100 连 Model States 都放不下。
起因 # 之前分别写了推理和训练的显存心智模型和混合精度训练中的显存组成,反复出现一个粗估:BF16 mixed precision + AdamW 大约 16 bytes/parameter。但这个数字每次都只是"1B × 16 = 16 GB"的单位换算,没有真正对着一个具体的 Transformer 架构把账从头到尾算一遍。
最近在看 MIT 6.S191 Lecture 9 的时候,反复听到一句话:Data Movement Matters。
之前写过 Arithmetic Intensity 和 Compute/Memory Bound,也刚写完 FlashAttention 的分块和 Online Softmax。这些文章里"HBM"和"片上存储"出现了很多次,但我发现自己其实一直没有认真想过一个问题:HBM 到底是什么?它全称叫 High Bandwidth Memory——高带宽内存——既然带宽已经很高了,为什么还要费这么大劲减少对它的访问?
之前的几篇把 Self-Attention 的公式拆得比较清楚了,从 QKV 的计算含义到 Multi-Head 的分工。另一条线是 GPU 性能——Arithmetic Intensity、memory bound、HBM 带宽这些概念也已经聊过。关于 HBM 和片上存储的物理区别,可以参考GPU 存储层级那篇。
起因 # 最近看完 MIT 6.S191 的 Lecture 8——AI for Science。这一讲的主题和之前几讲很不一样:不再是 NLP、CV 或者 RL,而是把深度学习拿去做天气预测、材料设计、分子模拟这些事情。
起因 # 之前写 Next-Token Prediction的时候,我把条件概率链式分解、Softmax 输出概率分布这些东西理了一遍。当时觉得逻辑非常自洽:模型在每个位置输出整个词表上的概率分布,概率最高的那个 token 就是模型最有信心的预测。
起因 # 过去几篇笔记一直在 Lecture 6 的范围里打转。Distribution Shift 那篇让我意识到离线测试集的 98% 准确率没法平移到生产环境,Algorithmic Bias 那篇又告诉我总体准确率可以掩盖群体间的巨大差异,对抗攻击那篇则让我看到泛化好的模型仍然可以被微小扰动骗过。
起因 # 上一篇结尾留了一个问题:Pre-training Loss 通常随模型规模平滑改善,但一些下游 benchmark 上的准确率似乎在某个规模区间突然大幅提升。
起因 # 上一篇理清了 Pre-training 和 Post-training 的关系——Pre-training 让模型学会语言建模,SFT 和 RLHF 再让它学会遵循指令。
起因 # 最近继续跟 MIT 6.S191,加上工作中或多或少接触到一些 LLM 训练相关的讨论,我开始试着把大模型的训练流程从头到尾理一遍。
起因 # 之前在量化团队做 ML Infra 相关工作时,我曾参与过期货预测模型的特征重要度评估与特征筛选系统。