LLM System

这篇笔记把 SGLang 共享专家融合的计算语义、权重布局和执行路径压缩成六张结构图。图中的实现判断以 2026-07-29 的 SGLang main 分支为准；正文只补充图中容易误读的部分。 从双路径到一次 MoE GEMM 融合不是让 Router 在 Shared Expert 与 Routed Expert 之间做选择。Shared Expert 仍然是必经计算，只是被追加为一个额外的 Expert Slot，并随路由结果一起交给 MoE Kernel。这样 Routed Expert 与 Shared Expert 可以在一次 MoE GEMM 中完成，省去独立的 Shared Expert GEMM 与 Kernel Launch。 Layout、Top-K 与权重重映射 普通布局是在 N 个 Routed Expert 后追加一个 Shared Expert，得到 N + 1 个 Slot，并把实际执行的 Expert 数增加一。当前 DeepEP / Mega 系列后端还需要为各个 EP Rank 保留本地 Shared Expert Slot，因此布局会扩展为 N + EP_size。Checkpoint Loader 随后把 mlp.shared_experts 的权重重映射到追加的 Slot；只改 Expert ID 而不改加载布局，并不能得到正确结果。 ...

这篇笔记把 DeepEP v2 的对象、内存、同步和数据路径压缩成五张结构图。正文只补充图中不适合展开的语义边界。 架构总览 ElasticBuffer 管理 Dispatch / Combine 所需的通信与 Buffer 资源；EPHandle 保存一次 EP 数据传输的元数据，包括确定性模式、Dispatch 缓存和 GroupGEMM 的 Token Padding 信息。 ElasticBuffer：物理分段，虚拟连续 CPU 与 GPU 内存物理上分段，但通过 CUDA VMM 映射为连续虚拟地址，再注册为 NCCL Window。这里的 Window 提供跨 Rank 访问语义；它本身不等同于 NVSHMEM 原生的对称内存抽象。 Barrier 与通信计算重叠 图中展示的是全局 Barrier 的基本心智模型。Hybrid 分层 Barrier 的并行模式还存在一个容易误判的弱同步边界：A0 分别看到 B0 的 Scale-out 到达和 A1 的 Scale-up 到达，并不能推出 B1 已经到达。局部可见的两条同步事实不能自动合成为全局强 Barrier。 async_with_compute_stream 允许 Barrier 延迟返回 Event，在真正消费通信结果之前继续执行独立计算；prefer_overlap_with_compute 则通过减少通信占用的 SM，为计算留下资源。 ...

flux实现：sm80 gemm+rs 看flux的论文来说，原理并不难，修改的是ffn的第二个gemm的epi阶段，让epi阶段的写回操作变成写入reducescatter的通信buffer，相当于做了一次零拷贝。 几个关键的点：1. 如何用cutlass自定义epi阶段。2. 用了什么ptx。3. 清零？ cutlass自定义epilogue–sm80 常规的epi，写回gmem，这里改成写在scatter-aware-memory（我自己起的名字）。cutlass封装自定义epi的逻辑抽象是：输出tile存到哪里，以及做什么运算，这分别是两个模版类： using EVT_D = decltype(this->evt_d(kparams)); using StoreD = decltype(this->evt_store_d(kparams)); using EVT = cutlass::epilogue::threadblock::Sm80EVT<StoreD, EVT_D>; StoreD就是规定了怎么存，EVTD就是规定了怎么算。 EVT全称epilogue visitor tree就是定义了一个epi阶段的计算-访存操作流，用树状图的形式保存epi阶段的一堆操作（做成图应该是为了方便接入nvcc？不懂为啥非得在这强调是tree）一个树节点（操作）是一个visitor。 custom_evt_d() EVT_Compute0 = alpha * accumulator 代码里对应 VisitorCompute<cutlass::multiplies, …>，两个输入是 VisitorScalarBroadcast 和 VisitorAccFetch。第一个负责广播 alpha，第二个v从 GEMM accumulator 里拿当前 tile 的结果，然后 multiply EVT_Compute1 = beta * C + EVT_Compute0 = beta * C + alpha * accumulator 代码里对应 VisitorAuxLoadGemmk 先把 C/bias 读进 epilogue，然后 VisitorCompute<cutlass::multiply_add, …> 做 beta * C + alpha * acc。因为CUTLASS 2.x 没有 SM90 那种 SrcFetch 替代物，所以这里用 VisitorAuxLoadGemmk 来读 C。 ...

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