LLM System

本篇目标： schedule要不要做成异步 gpipe的非异步调度 就是简单的做sf和rf，sf和rf之间依赖于torch.dist的api做阻塞同步。 在其他博客中也说这种静态调度和下发cpu指令差不多，按顺序一条条执行，执行完了整个程序就跑完了。 异步的问题 那为什么不直接把任务丢到下游（backward是上游）然后直接做下一个mb的计算呢，这样看起来还可以让sm利用率高。 搞成异步之后快的那个stage确实会更快推进，但是慢的那边会更慢。要么就是两遍差不多快，一样没有什么提升。 另外就是action memory的值会比较不确定。可能会非常大。而静态调度actionmemory的大小是可控的。 语义问题 如果搞成异步下发，那说明会有任务的积压，这些任务做了一半以后，checkpoint要按照哪个标准做记录呢？这也是很难做的。 静态调度的状态就是比动态调度状态更少。动态调度存储，加载，更新状态都会更难而且可能bound在控制流。 好处？ 动态调度对慢节点的容忍度好，但是绝对不是pp schedule pipe里面的慢节点。因为llm场景每个mb的时间都差不多。就算是卡慢了也不会这样处理，直接换掉就行了。这种对慢节点的容忍度指的是对于一些异构的流程，比如说rl的几个步骤，以及搜推处理sparse数据等。但是这两个流程我都暂时不特别熟悉，后续还要继续学习。 如果做同步pipe，算子做还是框架做？ 放在框架 那就是调用action之前barrier一下，因为这样涉及到多节点，所以启动开销会比较大，好处是位置浅好定位。 放在算子 那训练框架侧就只启动torch.dist的接口，torch.dist底下再接入通信算子库。算子内部barrer。这种问题是调用栈会很深，以及算子级更难定位，好处是算子层可以做更深度的优化，比如做smfree把单节点的mfu打上去。（不过还是那个问题，单节点mfu可能真高了，全局不好说） 那其实这里再给自己开个新todo，试试sm free的算子实现，用CE做通信，到时候跑训练看下效果。

本篇目标： 问题背景 什么是schedule，这个词含义很广但是在训练框架这里一般考虑的是f和b任务之间的编排。 PP pp开几一般就是把所有layer除以几，然后每个就是一个stage的layer数量。一般按照layer切。 GPipe 有很多mb，每个mb要做很多stage（模型的layer或op，跨卡或跨机，这些都行），GPipe就是要等到所有的mb都做完他们自己的所有前向stage，然后开始反向stage。气泡比较多，此外因为前向和反向的layer是反过来的，所以对于一个mb来说，他做的这些stage里面做前向越早的那个stage，做反向越晚。也就是inflight越多。inflight越多就代表得保存中间的状态，占显存。所以2个肉眼可见的缺点一个是空泡另一个是inflight。 在 Perfetto 中打开 GPipe trace 如何实现一个GPipe呢？（底层组件假设已经分好了，我们只需要考虑怎么把任务排好发出来，底层组件的事情可以见后文如何实现一个调度器） 非常的简单，给每个stage执行的载体（GPU）从mb0下发到mbn就可以了。然后执行那个stage对应的layer的前向传播/反向传播。 1F1B 做F的预取，然后让F和B同时进行。中间的卡交替进行f和b。好处是inflight少，但是空泡不减。 在 Perfetto 中打开 1F1B trace 如何实现一个1F1B呢？也并非很难，假设我们的stage执行载体（GPU，虽然总是括号里写GPU但是某些场景不一定是GPU，目前为了便于理解先这么写） 有m个，那只需要给stage编号（这个编号代表第一次启动任务的顺序）为i的stage提前分配m-i+1个mb就行了，mb的序号是从0到m-i。 然后这样预填充完之后，只需要做简单的配对+交叉即可。因为是1f1b，所以只需要交替下发f和b任务，f和b任务对应的mb编号只需要匹配最近一次任务即可，如果是b则找最旧的未完成mb任务id让fb闭合，如果是f则找最新的未完成mb任务id+1。 interleaved 1F1B 也叫vpp，把一个stage再划分为几个虚拟stage，用interleave的形式排到几张卡上。这个场景为什么能减少bubble在我第一次理解的时候其实不是很直观，因为我思考的是，就算切细了那三角形的空泡依然存在，为什么空泡会少。所以就计算了一下size。只算开始部分的三角形空泡（结束时候是对称的就不管了）不计算的很细的话，我们看三角形空泡里面最长的部分，也就是最底下的那条，长度正比于每个f/b的时间*（pp-1），但这里要注意一个很容易想当然的问题，这里的pp是物理pp数，也就是真实的stage，而不是虚拟的stage。因为我们真实的stage数量一般和gpu数量一样，所以就算很多虚拟stage，一次填充到流水线的阶段也最多只有物理个gpu数。那切完以后f/b的t就变小了，显然空泡就小了。 这么解释不太直观，，最直观的其实是，让最底下的那个rank早启动。假设就是rank0到7，rank7得等好几个阶段才能启动，那就把阶段切细，然后启动的就快了。但是如果切得太细，跨rank（其实是stage）通信不能忽略，那就也不行。 然后写这个还想一个问题就是stage到底跨卡还是跨机还是跨什么东西，问了下ai说具体情况具体分析（等于没说）然后翻了下之前（未发布）的训练框架学习笔记，原则上stage没有跨什么东西的限制，但是在机内有高速互联的情况下一般是跨节点的。因为高速互联要留给tp。优先级tp>dp>pp因为我们假设tp每一层都开一次，那么tp的通信量是 $$ seq\_len \times batch\_size \times layer \times hidden $$dp没有layer这个维度肯定要少点。pp一般都可以overlap了。如果节点内没有高速互联是需要开pp的。 Chimera 最接近dualpipe的办法。 初始流水线：s0f-s1f-s2f-s3f-s3b-s2b-s1b-s0b。 Chimera主要减少了bubble，前面二者有bubble都是因为GPU来任务的时间难免有pipeline式的三角形空泡问题。但是三角形空泡来源一个先入为主的假设就是我们总假设只能gpu0开始做mb0stage0。如果让其他gpu也同时开始一个任务，三角形空洞就能补上很多。（拓展，Chimera只是同时走两段pipe，能不能更多的pipe，收益如何） 其实就是排两个交叉的流水线。依然要vpp把stage加倍。 假设原先4stage，vpp成8个。 s0 0 7 s1 1 6 s2 2 5 s3 3 4 就大概这样的。 如果纯做vpp的话是这样的： s0 0 4 s1 1 5 s2 2 6 s3 3 7 那为什么vpp的效果不如Chimera呢。可以观察一件事情，stage越“在时间上靠前”被下发做f的，在做b收口的时候越晚，占用的显存就越大。所以有一个直观的结论，不同stage的显存开销在时间上是不均匀的。最影响显存开销的就是f的第一stage-b的最后一个stage这一对。这一对fb启动最早释放最晚，所以我们如果多同时启动几个这样的f，就可以让显存开销在时间上更均匀，进而降低了显存需求量的峰值。我们做vpp的话很难让不同stage的显存分配量是均匀的，甚至还有可能让inflight叠加。 ...

这篇文章是 LLM System 系列里 Transformer Engine 子专题的第 0 篇，也是这个主题的学习入口。 我准备用这个系列系统学习 NVIDIA Transformer Engine。目标不是只会调用几个 te.* API，而是把 TE 放到大模型训练系统里理解：它为什么存在、如何利用 FP8 / FP4 和 fused kernel、怎么被 Megatron-LM 接入，以及后续如何用 profiler 分析和改进它。 0. 学习目标 这组笔记会围绕四件事展开： 1. TE 在 AI Infra 技术栈中的位置 2. TE 的核心加速策略：低精度、融合算子、Attention、MoE、通信重叠 3. Megatron-LM / Megatron-Core 如何接入 TE 4. 如何 benchmark、trace 并尝试改进 TE 1. 技术定位 Transformer Engine 不是训练框架，而是 NVIDIA 为 Transformer 训练/推理提供的高性能 building block 库。它大致位于训练框架和底层 CUDA/cuBLAS/cuDNN kernel 之间。 这一阶段要回答的问题： 1. TE 和 PyTorch AMP 的关系是什么？ 2. TE 和 Megatron-LM 的边界在哪里？ 3. TE 为什么不是一个完整训练框架？ 4. TE 为什么不只是 FP8，而是 Transformer 优化库？ 5. TE 在 AI Infra 技术栈里更靠近 compiler/runtime/kernel，还是更靠近 model framework？ 2. 核心加速策略 核心加速策略先按五条线学习： ...

这篇文章是 LLM System 系列里 PD 分离子专题的第 0 篇，也是这个主题的学习入口。是笔者让gpt-5.5通过联网搜索帮自己制定的系统性学习方案。笔者会根据这个方案来确定如何学习PD分离的整套机制。目标不是先把所有论文细节读完，而是先建立一张可以持续填充的地图：该读什么、该推导什么、该写什么代码、最后应该能回答什么问题。 这个系列暂时围绕一个问题展开：为什么现代 LLM serving 系统越来越关心 prefill/decode disaggregation，也就是 PD 分离？ 我希望自己最后能回答四个问题： 1. 为什么 prefill 和 decode 要分离？ 2. 一个 workload 到底该配多少 P worker、多少 D worker？ 3. KV cache 从 P 到 D 传输到底传了什么、代价多大？ 4. vLLM / SGLang / Mooncake 里这件事具体怎么落地？ 先说一个结论：**PD 分离不是一个“拆进程就能变快”的魔法优化。**它真正解决的是服务系统里的资源解耦问题：prefill compute、decode iteration、KV cache 生命周期、网络传输和调度策略，本来在 colocated serving 里被绑在一起；PD 分离试图把它们拆开，让不同阶段按照不同目标优化。 0. 心智模型 LLM 推理一个请求大致分成两段。 Prefill：一次性吃掉 prompt，生成整段 prompt 的 KV cache，并产出第一个 token。长输入时它更像大 GEMM，通常更容易把 GPU 算力吃满。它最直接影响的是 TTFT，也就是 time to first token。 ...

本篇目标：了解 Transformer Engine 的技术定位，搞清楚它为什么存在，以及它和 PyTorch、cuBLAS、Megatron-LM 的边界。 基本接口 Layer 类定义接口非常直接，最表层的使用方式就是把 torch.nn 模块替换成 transformer_engine.pytorch 模块。 普通 PyTorch 写法： self.linear = torch.nn.Linear(hidden_size, 4 * hidden_size) TE 写法： import transformer_engine.pytorch as te self.linear = te.Linear(hidden_size, 4 * hidden_size) 量化上下文： from transformer_engine.pytorch import fp8_autocast with fp8_autocast(enabled=True): y = module(x) 进入 TE 的 FP8 上下文之后，TE 会围绕量化、反量化、fused path、tensor cache 和 backend 选择做一系列处理。相较于纯 PyTorch 计算图优化，TE 会拿到更多信息 tensor parallel、sequence parallel、FP8 recipe 等。这些额外信息给底层算子优化留下了空间。 这一点目前还是比较 general 层面 的认知，后面要继续顺着源码和 profiler trace 去验证。 TE 和 Megatron 的边界 Megatron-Core 负责模型并行、训练 loop、optimizer、activation checkpoint、MoE routing、pipeline schedule 和 config。 ...

TODO: 这里写 PagedAttention 的核心抽象：block/page、block table、逻辑 token 到物理 KV block 的映射。 基础：tensor 级拆请求的形状（大量细节） 定义符号：$B$ 是 batch size，$T$ 是 seq_len，$D$ 是 token_dim，$d_q$ 是把 embedding token 投影到 $Q$ 后的维度。 推理框架拿到的请求是：\(R \in \mathbb{R}^{B \times T}\)。 $R_{b,t}$ 是一个最最基本的 token id 标量。 raw 请求经过 embedding lookup，做的操作是把这个 token 标量映射成一个高维向量。假设原先 token 是 1234 这个标量，现在就把 token 映射成 [0.1, 0.2, 0.3, 0.4] 这样的向量。 所以 $R$ 经过 embedding lookup 之后，得到：\(X \in \mathbb{R}^{B \times T \times D}\)。 因为我们目前只考虑推理场景，所以把 $W_Q$、$W_K$、$W_V$ 之类的矩阵当成固定的模型参数。 然后很多博客会直接写：\(Q = XW_Q\)。 ...