LLM

RL AIGC 开发者交流纪要：从模型适配到异步训练系统 这次交流的核心不是单一算法，而是 RL AIGC 训练在工程落地中的系统问题。整体看下来，主要矛盾是：RL 链路把训练、推理、数据流、权重同步、checkpoint 和调试工具全部耦合在一起，而现有框架对这些问题的支持还不够完整。让AI总结了一下会议纪要。 1. 多模态 RL 的更新稳定性 多模态 RL 中，有一种做法是：如果某次参数更新和当前模型之间的 diff 超过阈值，就直接舍弃这次更新。 这个机制可以避免异常 update 破坏模型状态： 异常 batch / 异常 reward / 异常 rollout ↓ 参数更新过大 ↓ 超过阈值后舍弃本次更新 但它只能止损，不能解释问题来源。真正需要的是面向 RL 的 debugger，能够定位是 reward、logprob、rollout、并行切分还是权重同步出了问题。 2. 新模型接入成本高 如果要把一个新模型接入 RL 训练框架，往往需要手写 Megatron、FSDP 或其他并行逻辑的适配。 难点不只是 forward 能跑，而是整个 RL 链路都要对齐： 模型结构 并行切分 checkpoint / reshard rollout 权重同步 logprob 计算 loss 计算 训练侧和推理侧的数据格式 RL 场景下，模型适配错误不一定马上报错，很多时候只表现为训练逐渐崩掉。因此新模型适配需要更强的调试工具，比如检查权重版本、logprob 对齐、reshard 正确性和并行切分一致性。 3. RL 训练周期长，问题复现成本高 RL 训练周期通常很长，一个周期可能需要几天。很多问题不会在前几个 step 暴露，而是在训练一段时间后才出现。 ...

概述 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。

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绑定。

简述 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。

本篇目标： 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。 ...

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文档提了几个这部分优化的点： ...

本篇目标： 为什么通信器会死锁 Megatron 的 P2P 通信抽象 batch p2p comm overlap p2p comm warmup / steady / cooldown 里的通信顺序 如何设计防死锁的通信接口 TODO

Transformer 预测 next token 的路线可能是对的，因为这种预测模式特别类似于递推。人类解决问题的方式其实也是递推。当我们能问出问题 A，其实是默认知道了很多关于问题 A 的背景，而问题 A 只是在知道了这么多背景之后的一个单点问题。 那么，如果我们对 A 的背景知之甚少，意味着我们就得先问出 A1、A2、An 等前置问题作为铺垫。这个就是递推。 能够开启递推还有一个前提，就是你必须知道递推的前一项是可扩展的。也就是说，递推的前一项存在某些局限性，必须清晰地看到这些局限性，才能把递推很好地进行下去。 学习 PTX 汇编要关注两代之间的局限性和扩展性，这个就是一个很好的现实例子。关注递推的边界，往往有利于认知。 有句话说“领先一步是疯子，领先半步才是神”。这种说法不无道理，因为领先半步的时候才符合人类一贯的认知结构，做到了认知递推的下一步，而不是下下一步。所以才会“被认为是神”。 学习增量是容易的，但是学习总体是难的。所以，把一个学习总体的任务变成多次学习增量的任务，这个就是人类解决问题的过程。

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 基本实现都差不多，依次有几个点不一样 ...