Megatron on Echo的技术博客

LLM System: 训练框架随笔 09 - Megatron Core Distributed DDP 和 FSDP

Sat, 20 Jun 2026 00:00:00 +0800

概述

core/dist接入了三种数据并行：torch fsdp2、mcore ddp、mcore fsdp

FSDP

流程模型

FSDP就是一种sharding，sharding的对象是param、optim、grad

如图，fsdp的大概做法就是为了不让f/b期间大量模型的layer持续驻留在gpumem，所以会把这一个layer切到不同rank上，需要做f/b的时候再通过集合通信拿到。

这句话说完其实有两个质疑的点。一、如果我们把一次coll操作得到的整体作为一个unit，那么这个unit一定是上文中说的layer吗，可不可以是别的？可不可以是两层layer？二、其他rank如果也要做fsdp的shard，一样要把自己的全量参数/optim/grad切片并分发到其他rank，这么看只有对于一个rank的一个layer的fb计算，这一个小阶段内是可以用通信换存储。但是如果全局来看似乎并不一定省显存、因为本rank也负担了其他rank的shard？

一、可以是别的但是一般默认一个unit就是一个layer 二、肯定省显存，因为是一个dpgroup共享一个param/optim/grad副本。具体怎么共享看zero设计。

关于集合通信：ag一般是收集，rs一般是同步。

Init

weight/param在shard之前有一个自己的layout，shard之后也有layout。这个layout非常简单，就是把所有参数拼成1D tensor。做成1D tensor的原因是为了适配dist.all_gather_into_tensor(),这个allgather接口生成的是一整个连续的shard，语义表达比较差但是性能好。而如果使用经典的dist.allgather会返回一个tensorlist，语义可读性更好，但是性能差，性能差的原因是他支持每个collrank传递不均等shape的tensor，但是fsdp的unit shard是均等的，因为没有额外的设置为不均等的理由。

reducescatter和reducescattertensor也是同理。

这个1Dtensor有自己的构建规则，mcore这里面是函数式构建，也就是通过设计特殊的递归，一类weight只能演算出一种1D layout。这种的好处就是不用手动指定metadata了，而且也不需要给每一种情况都写一种layout。

一个layer上很多weight，比如q k v norm这些东西，每个单独shard会有很多小包，所以用FlatParameter把这些凑成一个shard。用这个类去给刚才的1D tensor包一层。

Runtime

流程为：preforward-forward-postforward-prebackward-backward-postbackward

可以看到这里会注册两次hook，prebackward hook的操作是计算backward的时候需要一次ag把shard汇聚恢复。postbackwardhook的时候是需要一次rs同步梯度+释放shard。前向也会有类似的hook，区别是前向是主动写的hook，而反向的hook是前向调用时候才reg的。

为什么有这种区别呢。因为torch的设计是前向手动开始，反向由autograd按计算图执行，而torch是不感知到自己什么时候进入了某个layer的backward，不存在这种边界的定义，反向的时候torch只知道梯度来了要算梯度，所以得手动hook，原理上是如此，torch代码还需要再看看。

Runtime的时候峰值显存等于allgather之后拼好的unit+本地驻留的不被fsdp wrap跟踪的参数+平时就持有的shardSize（这个不能和全量unit算在一起因为做merge实现太复杂了）

lora策略

如果用peft或者lora，fsdp需要修改sharding策略，因为大部分权重是frozen的。

fsdp通信计算重叠

fsdp带来的通信计算重叠从pp视角看是发生在stage内部。而ppstage间微弱的通信不是fsdp管理的。 fsdp的通信计算重叠来自于一个stage内部。当前layer的通信/计算和其他layer的通信/计算重叠。

看上面图就可以了，因为ag和compute unit是切片的所以能叠一部分。另外就是还有两类可选优化，f prefetch和b prefetch。

ZeRO

DDP

非常朴素，每张卡都有全量模型，每张卡处理一个mb。

LLM System: 训练框架随笔 07 - MCore Export

Fri, 19 Jun 2026 00:00:00 +0800

export概览

export模块做的事情就是需要给用mcore训练出来的模型做一个推理优化，推理框架是trtllm。（为啥，啥地方用？）那么就得考虑几个问题：

为啥要推理优化，和训练过程中的forward啥区别，什么情况会用。
这里面既然是做推理优化，那么有哪些优化是和训练场景耦合的？哪些优化是必须妥协不能做的？哪些优化手段可以被做成任意推理场景都能随时插拔的模块？
那么既然是在看训练框架，我们最应该关注的就是和训练相关的推理优化了。那么和训练相关又有几种情况呢？我认为是两种。首先是训练框架本身对模型ckpt的设计会影响推理框架如何去解析这个模型，比如说我们可以猜想一下，如果ckpt是分布式的，那么trtllm大概率也会用一种分布式的方式去读写ckpt，或者说如果ckpt在内存layout上并非推理框架最友好的，那么还会产生一次转换等。其次是训练相关的一些上游任务，还是一开始的问题，为什么训练框架里面要内嵌推理框架，这部分推理加速会被用在什么地方，这个推理的地方会有什么样的workload特点？

总结

gqa和mqa的推理tp切法必须要用ckpt里面gqa和mqa的设置来限制，不能随便改Q/K。

layout的重排全都是trtllm离线自己做，会把训练ckpt的默认layout转换成推理友好的形式。

训练做完量化之后推理不用再重新calibration（校准范围），只需要复用之前的factor就行。

训练的tp/pp shard和推理的tp/pp shard不一样，也相当于一种分布式的layout（？可能这么理解不是很严谨）。

dist-ondevice convertion，转换成推理权重的时候需要考虑ckpt已经是多卡shard了，如果shard已经做了分布式，那么对它的操作最好也是分布式的，否则会造成很大的barrier。这个道理和mcore dataset里面计算每rank的data idx类似，只不过那个却是让一个rank算全局，其他rank要等着。但那个东西设计成barrier有其他的原因（分布式写共享文件容易导致更严重的抢占，所以还是一个rank做single-writer，这里还带了一个很重要的常识，你如果想多线程一起算一个东西，必须至少还得维护一个共享的缓存来同步他们的计算结果。当然也可以每个rank都算一遍全量，那就是用不到共享存储了），当然既然牺牲了分布式操作肯定也带来这里所说的问题。

既然训推的tp不保证一样，tp又对padding又要求，按vocabsize这一列去切的时候必须能整除tp，所以训推的vocabsize的padding就不一样。所以还得有一步转换问题相当于训推reshard的副产物。

思考为什么forward不用trtllm加速？

trtllm build inference engine的过程比较重。而且基本和hopper绑定。

LLM System: 训练框架随笔 08 - MCore Checkpoint Resharding

Fri, 19 Jun 2026 00:00:00 +0800

简述

resharding这个词出现了好几次，rl的resharding，ckpt的resharding，fsdp的resharding，这部分先看ckpt的resharding目的是什么，怎么做的。

ckpt的resharding发生在训练开始前loadckpt的时候，ckpt的格式有三种，torchdist，torchdcp，fsdpdtensor。torchdist是mcore原生的distckpt，torchdcp是torchdist的原生格式，fsdp是fsdp2的dcp格式，参数是dtensor分片。

加载ckpt的时候sharded_state_dict_default演算一个ckpt shard在全局唯一确定的位置。输入的信息：data，shape，tpaxis，tprank，tpsize，layerkey（name），replicaid 输出的信息：key，data，localshape，globalshape，globaloffset，axisinfo

这里ckpt的resharding做的事情是当我们load一个ckpt，这个ckpt里面包含它被保存时候的并行信息，但是它被保存的时候的并行设置不一定和这次重新load的时候一样。所以这个resharding就是一次覆盖。那么能否直接不给ckpt加入并行信息呢，因为之前就是分布式的ckpt，没有之前的信息无法拼成完整的模型。

异步分布式ckpt的难点

训练的主路径不想等io，所以ckpt的过程是async的，但是ckpt的各个shard还需要同步。mcore和torchdcp在这里权衡的方式是在快照的时候做同步，但是分批写入的时候是异步的。也就是每个step的边界快照一次，然后异步并行传shard。另外同时最多只限制一个async的异步shard组在inflight，否则很难debug。那么如果一批全局async的shard还没都写入，其他新step的同步快照ckpt就得一直等着。所以这个地方也得开足够的空间保存这些shard，这个地方可以pinCPUmem。

LLM System: 训练框架随笔 05 - Megatron Checkpoint

Thu, 18 Jun 2026 00:00:00 +0800

本篇目标：

megatron Checkpoint

checkpoint包含如下四类：

rng state/rerun state/dalaloader state/model&optim state

rngstate是一些伪随机数的序列，因为伪随机数的采样本身就是很好的分布，但如果ckpt没有记录他们的状态，会破坏这种分布，进而给训练带来一些不可预估的问题。 rerun state是用于做容错的。 dataloader state是数据的ckpt，表示训到哪里。 model/optim state是模型权重，优化器状态，学习率衰减程度等算法层直接可见的东西。

RNGstate详解

rng_state = {
 "random_rng_state": random.getstate(),
 "np_rng_state": np.random.get_state(),
 "torch_rng_state": torch.get_rng_state(),
 "cuda_rng_state": torch.cuda.get_rng_state(),
 "rng_tracker_states": get_cuda_rng_tracker().get_states(),
 }

random_rng_state： python的random模块如random.random() .shuffle() .sample() .randint()
np_rng_state： np的random。如np.random.random() np.random.shffle() np.random.randint()
torch_rng_state：torch的random，如torch.rand(), torch.randn(), torch.randint() 前提是device=cpu
cuda_rng_state： torch.rand(cuda) torch.randn(cuda) F.dropout(cuda_tensor) cuda_tensor.normal()
rng_tracker_states： gpu cuda有随机数的流，默认情况下所有gpu的random公用这一个流，但是这个不够好。所以会以类似于上下文的方式，去管理随机流。这个state记录的就是各个随机数流推进到了哪个位置。

如何发现rngstate出问题了？

既然框架设计了这个功能，就要思考这个功能如果没做或者没做好，带来什么问题。如果是rngstate没同步好，那么是否启动save/load ckpt就是最先需要判断的标准。所以一开始可以先单卡+是否saveload。

rerunstate详解

megatron checkpoint backend - mcore dist checkpoint

ckpt前端其实只是决定要存哪些状态。后端决定了怎么在多rank存储。

ckpt的后端是distckpt，distckpt这一层看到的是一个全局的共享文件。

distckpt只需要提供每个rank的shard在全局的位置，以及告诉本rank其他rank的一些shard信息。本rank就可以读取各种shard，这一层是屏蔽掉各种近端远端读写逻辑的。

训练框架这一层就只看见一个共享目录，调用的话就是系统调用，非常简单。其他的底层细节让xxFS去处理。

具体怎么shard由训练参数决定策略。策略指定的是谁保存shard，谁读取shard，是否需要在rank之间重分配io（比如由其他rank代读取），同步还是异步write让，是否缓存metadata和plan。

LLM System: 训练框架随笔 06 - Megatron Dataset

Thu, 18 Jun 2026 00:00:00 +0800

mcore dataset构造协议

不同模型对于“sample”的定义不同，也就是说每一类模型拿来训练的输入是不一样的，哪怕底层数据是一样的。所以就有必要在之上抽象一层dataset类。

主流这三类：

GPT

数据-标签形态：
- tokens = x[0 : L]
- labels = x[1 : L+1]
任务：
- 预测下一个tk

BERT

数据-标签形态
- input = corrupt(x)
- labels = x only at masked positions
任务：
- 挖空填词，看到左右的tk预测中间的tk

T5

数据-标签形态
- encoder_input = corrupt_span(x)
- decoder_label = removed_spans
任务：
- encoder看残缺文本，decoder输出删掉的整个片段

mcore dataset涉及到的系统级优化

用mmap映射token

mmap是把一个文件映射到了进程的虚拟地址空间，但是不会把整个文件load到进程的内存。真正的读取行为是lazy的，用哪块读哪块。（实际上os做的事情就是读va-查询页表-pagefault触发-os做一次read SSD到mem，信息是mmap给-更新页表）为什么不用read()而是使用mmap()呢，按理说rea也可以做到这种lazy的行为。因为read()在读的时候不是0拷贝。

llm训练数据量大，连续存储，每次只用一块mb，适合用mmap。

mmap存在的代价是A。容易缺页，（读取没有时空局部性会这样）B。warmup慢，C。多rank同时缺页的时候os的fs开销大（为什么，因为）。

对A问题，mgt做的优化：

document_index
sample_index
shuffle_index
cache
访问密度判断

对C问题，mgt做的优化：

cache 预建
defer mmap
fast cache load
object storage block cache

先不展开。让ai总结了一下这几个优化所在的位置。方便后续索引。具体位置在索引AAA这一节。

nv文档提了几个这部分优化的点：

LLM System: 训练框架随笔 04 - Megatron 通信器设计：如何防死锁

Wed, 17 Jun 2026 00:00:00 +0800

本篇目标：

为什么通信器会死锁

Megatron 的 P2P 通信抽象

batch p2p comm

overlap p2p comm

warmup / steady / cooldown 里的通信顺序

如何设计防死锁的通信接口

TODO

LLM System: 训练框架随笔 03 - 再读 Megatron Core 看设计模式

Tue, 16 Jun 2026 00:00:00 +0800

为什么突然就读上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

基本实现都差不多，依次有几个点不一样

LLM System: Transformer Engine 00 - 学习地图

Thu, 21 May 2026 06:00:00 +0800

这篇文章是 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. 核心加速策略

核心加速策略先按五条线学习：