moe f&b 计算流程 Router Router 计算每个 token 应该路由到哪些 expert。 Dispatch Dispatch 根据 router 结果，把 token 分发到对应 expert 的输入 buffer。 FFN 每个 expert 内部执行自己的 FFN / MLP 计算。 Combine Combine 把各 expert 的输出按路由权重聚合，并还原到 token 维度。 MoE softmax MoE做router的意义是什么？ MoE 做softmax的意义和一般softmax意义类似，都是让logits能被解释为概率。 被解释为概率/加权的几个要求：正数，和为1，尺度一致， ...
为什么突然就读上mcore了 从头给Femtotron加dualpipe的时候被自己的软件工程能力急哭，深刻感觉到想要0基础快乐vibe coding绝非易事，而且尤其是在其他人的项目上做增量开发的时候，一不小心就容易过度设计/破坏原有设计模式/重复设计。 所以就从pp schedule开始读。 scheduler 其实scheduler在另一篇文章里面也说过，主要的优化就是这几个点：1f1b+vpp+dual+zerobubble。（moe的通信计算掩盖暂且不计入） 我在开发的时候想把这几个优化点的接口设计成某种“能力”，也就是把这几个优化点都做成下面这种模式： default_scheduler = default_gpipe() scheduler = zero_bubble(dual(vpp(1f1b(default_scheduler)))) 这样做在接口上看起来很有美感也很理想化，但是我犯了一个很严重的错误。因为在规划阶段我的脑子里想的是可以像编译器排layout那样生成schedule pipeline。如果是对pipeline的整体做这个变换，那相当容易，但是问题就出在Femtotron的执行逻辑是SPMD的（Megatron也是这样）。SPMD具体是如何设计的，后面还会细说。这就导致了我们其实无法维护一个全局的controler，也就很难拿到全局的pipeline相对位置信息（可以拿，但是为此必须破坏大量的封装，保存大量的笨重结构体）。也就是说，这样设计scheduler在逻辑来看非常美，但是在实现上因为SPMD的设计，每个函数执行的范围是每个rank，输入的资源是一个chunk的model，因此保存全局的信息（对每个子块来说就是知道自己在全局处于什么位置）在实现起来非常繁琐，遂作罢。 于是本人暂时搁置了Femtotron的开发，开始参考Megatron-Core的代码实现。pipeline parallel的scheduler个人觉得是训练框架中比较有趣的一个内容，刚好也是在开发这部分，就借此机会把Megatron的核心源码走读一遍。当然可能写的比较草率，后续肯定还会单独整理一个比较完整的博客专门读。这里主要展示读的过程中我脑子里面的cot。 pipeline parallel的一切 gpt5.5虽然被agenticRL搞得不说人话，但是有一个词用得很好导致我自己也很喜欢，“心智模型”。读各类代码尤其是涉及到并行计算的代码，都需要有这样一个心智模型。比如cuda是simt。编程模型这说法在框架这层就显得有点不够贴切。Megatron的pp代码，心智模型是SPMD，也就是每个Rank进入同一个函数。这个从megaton的使用方式里面也能看出来，启动方式类似于mpi。pp stage里面具体做什么操作，由rank固有的状态决定：包括这个rank属于哪个通信组。Megatron实现vpp的方式就和我预设的完全不一样，它在get_forward_backward_func（）里面先用ifelse判断要不要用1f1b，还是直接简单粗暴的gpipe，在1f1b的ifelse里面又直接用了ifelse判断要不要用vpp，而且给是否开vpp分别写了一个函数，这两个函数分别是forward_backward_pipelining_with_interleaving 和 forward_backward_pipelining_without_interleaving。 Megatron这样实现就可以体现出增量的困难，比如这里就无法设计成“把gpipe改成1f1b”或者把“1f1b增量成interleave 1f1b”的样子。也侧面说明我一开始的设想不适合整体代码架构。 执行层级 传统的深度学习一次train step需要准备一个batchsize的数据。而megatron里面则不会一次准备这样一个batch，而是每次只取出一个microbatch将其转化为gpu tensor然后做fb计算。cpu侧的dataloader可以做预取重叠，所以可以看到后面每次fbpp函数loop的时候都会基于data_iter拿到本次loop所需的mb。而不是整个batch。 总结就是fbpp函数在每个ppstage的每个trainstep都会调用一次。 无interleave实现： forward_backward_pipelining_without_interleaving 从上到下依次执行： 判断是否切分modelchunk，非interleave不支持切分chunk 判断是否是multi-module流水线，会使用过特殊的Communicator，（这个主要是为了区分是否为llm，因为llm Encoder-only不能做cp loss scaling） 判断是否启动了comm p2p overlap，非interleave不支持comm p2p overlap 初始化p2p_communicator（通信用）和pg_collection（保存各种p的group） 清空用于做dw分离的缓存。（上一次的清空，留着给这次用） 清空moe paged stash。（moe backward需要额外buffer） 禁用梯度同步disable_grad_sync，各个gpu计算梯度，但是不马上广播到其他gpu，这个是dp的东西但是要放在pp控制。 计算各个阶段microbatch的数量（warmup，steady，end） 根据是否为多模态模型选择backward的函数类型 根据mb，group，seqlen，decoderseqlen计算send和recv的tensor shape。decoderlen和seqlen的区别是有的模型是e-d而不是donly，所以需要特殊处理e-d结构的，donly的decoderseqlen就等于seqlen。 如果send和recv tensor形状不同，需要调用外部函数adjust_tensor_shapes_fn，这个具体函数目前只有一种实例化就是在做模型蒸馏的时候要传入一个get_tensor_shapes_adjust_fn_for_distillation函数，因为蒸馏的时候可能同时有teacher和student的tensor。需要特殊处理形状。（TODO详细了解） warmup阶段1：checkpoint_activations_microbatch 是否启用recompute（不用存前向acvitation省显存，属于是给1f1b打补丁） warmup阶段2: recvforward warmup阶段3: forward，如果是最后一个就算loss warmup阶段4: sendforward warmup阶段5： 检查如果是最后一个pp，那就计算本次mb的累计token，然后做几次广播获得全局token数。知道token数才能算梯度。有的rank token无法对齐，所以还是得广播才能知道全局有效token。 warmup阶段6：保存本轮input和outputtensor到本地，deallocated（outputtensor）。这里是python伪释放机制，一个tensor有好几个成员，反向的时候只需要torch.autograd这个成员，不需要原始数据，所以deallocated可以实现只去释放.data而保留autograd，细粒度控制显存。如果调用torch绕不开，但是megatron直接用了c++ autograd engine。 warmup阶段是rf+f+sf，那么还要r一次，才能正式进入1f1b的sbrf阶段。 steady阶段1: 同warmup的1，重计算配置（TODO详细了解） steady阶段2: forward steady阶段3: sfrb返回grad（如果是纯forward就是sf）+ 存数 + deallocated，注意这里rb不代表真的收到了数据，是在等下游把梯度发回来。 steady阶段4: 取出最早完成f的mb做b，这个得到的梯度是本rank的不是下游的。和上一条区分。 steady阶段5: 如果是最后一个b，打开梯度累积。enable_grad_sync steady阶段6: backward steady阶段7: sbrf（sb） colldown阶段1：打开梯度累积。算完梯度会自动广播并reduce。 colldown阶段2：弹出最近一个没完成的mb任务 colldown阶段3：rb colldown阶段4：backward colldown阶段5：sb colldown阶段6：gradsync（注意这里看起来有很多gradsync的位置，但都是分支判断，实际只打开一次，而且必须保证是最后一个backward执行之前打开。如果在最后一个b之后打开，可能ddp和fsdp的backward hook错过了同步机会。） 所以实际上，就出现了代码里面的几种边界分支，如果没有cooldown backward的rank，比如说pp最后一次的stage，最后一次backward在1f1b steay之前就打开了，所以要在steady的最后一轮打开。而如果是有cooldown的rank，最后一次b是在cooldown的最后一轮。（TODO 画个图更清晰～） 把dw分离里面没计算的w都计算了。注意整个流程里面并没有算d存w的过程，因为这个过程隐藏在了forward和backward的实现里面。 dp梯度同步+ppsp梯度同步+globaltoken缩放梯度。 有 interleave / VPP 实现：forward_backward_pipelining_with_interleaving 基本实现都差不多，依次有几个点不一样 ...
本篇目标： 为什么要重构 因为要支持vpp（interleave 1f1b） zerobubble dualpipe dualpipev，之前的只支持1f1b，gpipe。 从vpp思索stage的执行序 在gpipe和1f1b里面，一个stage就是一个rank，但如果加上vpp，一个rank上就很多虚拟的stage了。以前设计vpp的初衷就是，一个mb在等通信的时候可以让另一mb去做计算。我们给每个物理rank都绑定了一个executor和一个active queue，为了达成这种目的，我们必须保证两个虚拟stage的active在指令执行顺序没有依赖和申请的资源不会死锁，这样才能让virtual stage A的通信和virtual stage B的计算相互overlap。 那么上面说的这种overlap要怎么排布active才能实现呢？这个必须先定义边界。 从dual（奇美拉）思考stage的执行序 dual V的思想是让一个物理rank不会只绑定模型的一层，而是模型的某两层。这个和vpp的区别是，vpp的优化是不同mb不同时做通信and计算带来的掩盖效果吗，而dual的思想是一个rank可以做几个model layer进而降低了单任务的IDLE，pipeline从两个方向同时发任务，气泡少。可能带来的代价是任务切换带来的开销，但直觉上这部分开销不会很大。（加个TODO先，有空也会验一验） TODO: 验证 folded V / VPP 下更细粒度 action interleave 带来的 kernel launch、stream sync、activation memory 和 comm fragmentation 开销。 mb编码逻辑 随着foldv，dw分离，vpp的引入，mb的编码应该用一个结构体保存所有metadata而不是简单的id。 action之痛 做到vpp的时候，开始发现femtotron这套基于active-mb的扩展性不够好，在做dw分离的时候我把一个backward action拆成一个d action和一个w action，其实这个设计就给后面的复杂埋下了问题。或者说一开始femtotron的forward，backward，sfrb等设计就让这种设计注定变得不可扩展。因为如果一个schedule包含多种优化，一个action的含义必然是复杂的。 所以没办法，先研究一下megatron是怎么实现的。
本篇目标： 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\)。 ...
1. 总体思想 这次做的是单机八卡 H200、NVLink、Ulysses CP 下的 GEMM 和 AllToAll 融合。先把 forward 写清楚： A2A → QKV projection → QK → PV → A2A 需要接起来的主边界有两个。输入侧是 A2A→QKV projection，通信先把各个 peer 的输入 tile 搬到本地最终布局，GEMM 拿到一块就算一块。输出侧是 batched PV→A2A，每个本地 head 都有一组独立的 P×V，GEMM 算完一个 tile，通信 CTA 立刻把它送到目标 rank 的最终 Ulysses 布局。 如果 GEMM 和 NCCL 顺序执行，端到端时间接近两段时间相加。这里把通信 CTA 和 GEMM CTA 放进同一个 cooperative persistent grid，两种 CTA 常驻在不同的 SM 上，用 tile 级 ready epoch 接力。A2A→GEMM 由通信生产、GEMM 消费；GEMM→A2A 交换生产消费关系。这样首批 tile 到达后就能启动计算，前面的 tile 也可以在后续 GEMM 还在跑时发出去。 ...