0x0. 前言
在TP并行的情况下, 每一个层, 至少会有两次all-, 一次是在之后, 另一次是在MLP/MoE之后, 那些有着几十上百层的模型中一步就会出现几百次all-的情况, 并且每次传输的数据量都很小, 比如bs为1, 8192的bf16模型的数据量是16KB。此负载特征将瓶颈判定为延时而非带宽, NCCL这类针对大消息以及任意拓扑予以设计的通用库在此区间的开销颇为偏大, 因而其中存在两代自定义all - 实现, 从vLLM移植而来的（以下简称为v1, 在 =0时启用）以及当下默认的v2。
v2相对于v1的主要改变之处在于, 具有存算分离的代码结构, 多出一项push算法, 存在NVLS硬件归约, 有PDL, 存在集中式CUDA graph指针表, 且把算法切换阈值从固定的常量换成了依据(架构, )实测所得来调校的表格。
下面这篇文章, 就当前这两代所达成的技术方面的细节情况做一下记录, 代码的口径呢, 是 main @紧跟着括号里命名空间再度构建之后的那个路径, 其中所涉及到的文件有:
0x1. 负载特征：为什么小消息需要专门的实现
把所有参与者的数据逐一元素进行归约（这里归约操作为求和）, 而且要让每个参与者都能拿到完整结果, 这是 all-做的事。TP会把每层的权重矩阵切割给多张GPU, 每张卡利用自身那份权重算出的仅仅是输出的部分和, 而非完整结果也不是最终求和的正确结果, 只有经逐元素相加才是正确结果, 而这个加法是由 all-来完成的。
阶段的消息尺寸小得很: 当bs等于1的时候, 一次all-仅仅只有×2字节, 8192的bf16模型是16KB。一步要进行几百次如此这般的调用, 总的时间是由每次调用的固定开销来定的, 并非是由链路带宽来定。
在这个区间, NCCL 的开销源自几个方面, 其一为 ring 算法的多跳, 当等于 4 时, all - 要走 2(ws - 1)步, 也就是走 6 步, 其二是协议选择与内部的 proxy 逻辑, 其三是 CUDA graph 场景下每次调用所携带的参数块自身构成的图节点成本, 最后是百余次每步累加起来为毫秒量级。
它所采用的方式是, 在全部是点对点相互连接 的情形之下, 使得图形处理器直接进行对另一端显存的读取与写入操作, 借助一个或者两个来达成整个全部的相关事宜, 从而避开通用的协议栈。其所付出的代价在于适用的范围比较狭窄: 仅仅适用于单机 , 并且必须皆为全互连状态 , 数量小于或等于八 , 同时传递的消息不可以超过某个规定的限额。
0x2. 前提条件： P2P 与 IPC 交换
理解后续实现需要三个前提：
：同机 GPU 之间的高速互联；
在全互连的机器之上, GPU0能够存有一个指针借此直接去读写GPU1的显存之处, 语义方面和访问本地显存保持一致, 并不需要进行显式的消息收发, 这是得以成立的物理前提。
需要先对地址作出交换 , 对显卡处理器0GPU0要访问显卡处理器1GPU1对应的某一块显存 ，这就需要显卡处理器1GPU1先导出其内存句柄项 , 之后显卡处理器0GPU0端进行打开映射 , 才能够得到本进程可以进行解引用操作的指针 , 而这个手续在版本v1时是通过手工来完成的 , 在版本v2则是交给了与torch相关的。
IPC重映射是节点内机制, VMM重映射也是节点内机制, 这是只适用于单机的一个原因, 需要予以注意。
0x3. 的分发链
AR 不是唯一的候选实现。.（.py）的完整决策链如下：
. 的分发决策链
分发链上有三处顺序需要说明。
首先, 当 ---symm-mem 处于打开状态时, 其中的 - 快路径是位于整个 链的前面的, 并且它会毫无条件地接收走全部的 GPU 输入, 而 AR（不管 v1 还是 v2 皆是如此）是不会被调用的。另外, 这两项特性在进行部署的时候是相互排斥的, 需要依据消息谱以及硬件实测来二选一。
首先得明确, 第二点是这样的, AR会走出参通道, 进而返回新张量, /NCCL则走原地通道, 并且上层针对这两者进行了统一封装。
第三, piler.()分支, 所有依据字节数来挑选算法的逻辑（、duce）, 在追踪的状况下, 会针对符号形状产生guard, 并且会逐形状进行重编译, 所以, 路径会将通道选择整体延迟到自定义算子内部的运行时去开展。
实现切换的那两代, 是由（默认状态设定为 True 的情况）来进行控制的, ----all- 这种情形之下, 两代就会全部处于关闭状态。
0x4. 跨 rank 同步：v1 的 与
P2P直读要解决一个问题, 即当GPU0读取GPU1的数据时, 怎样去确认GPU1已然完成了写入操作。两张卡各自独立执行, 如果读得过早, 就会获取到上一轮残留的数据, 并且不会出现报错情况。
解决的办法呈现为, 每个rank去维护一组计数器。在写完数据之后, 将当前的轮次号写入所有peer的对应槽位。接着进行自旋等待, 等待自己的槽位收齐所有peer的同一轮次号。v1采用结构去承载这组计数器:
这里是源, 来自于斜杠sgl斜杠减号, 斜杠csrc斜杠斜杠点, 该内容处于大括号内, 这个大括号里是本rank每个block的本地轮次计数, 其数值为128。
// 接收 peer 旗标的槽位, 两组交替使用, peer 的 block 有可能已经抵达第二个同步点并开始写 +1, 而本 rank 的 block 仍在第一个同步点等待, 以此交替数组避免两个相位相互覆盖 (128)。
;}; // = 36（HIP 为 60）
本体每个 block 只使用前 个线程：
// 来源: /sgl-/csrc//。针对常量引用的sg、指针以及整型rank执行空函数, 若条件不成立则执行空语句, 即若条件为假直接结束。这里存在一个判断条件, 即需判断某个对象的成员x是否小于ngpus, 若小于则执行后续操作, 其中包括轮次号自增, 并且提到一个计数器在这个过程中在语义上来说是足够的, 每一个block都有一份且这样做是为了避免通过共享内存去共享计数器, 之后有一个自动变量val获取了某个值, 该值是通过某个操作得到的从而指向某个后续的内容。
.x
.x
不太明确你提供的内容具体要表达什么准确意思, 感觉表述比较混乱且不完整, 难以按照要求准确改写完全符合逻辑的内容。你可以补充完整清晰准确的原句信息, 以便我能更好地完成任务。
.x
->
.x
; auto = &->
.x
.x
若括号内条件成立------这里指的是st..sys / ld..sys模式下的跨设备内存序, 其作用是确保当旗标可见时数据已然可见, 对应传入的值为val, 接着进入一个循环, 在循环中持续判断括号内表达式的值是否不等于val；若括号内条件不成立------这里是说st/ld模式, 该模式仅保证数据能够到达, 却不保证顺序, 适用于数据依赖由其他地方提供保障的位置, 同样传入值为val, 之后于此进入循环, 并在循环里持续判断括号内表达式的值是否不等于val。若另外两个括号内条件成立或成立, 执行括号内表达式。
运用两个模板参数所起到的作用是, 去区分入口以及出口, 也就是说要以此来决定括号究竟该放置在前面还是后面；这个区分还能够决定是选用斜杠还是其他特定符号。对于入口这种情况而言, 是不需要使用围栏机制的, 原因在于此时所要保障的仅仅不过是所有的等级都能够达成到达的状态；而针对出口这一情况来说, 在接下来同级对等体要去读取本级写入的相关数据时, 需要满足内存序的要求。
v2 的 是同一思路的另一种实现，见 0x7.2。
0x5. v1 的实现
0x5.1 存储与 IPC 交换
初始化时, v1 每个 rank 手工操作三块内存, 给所有 peer 进行 IPC all - 的相关操作, 还要各自打开映射。
等于同步区加上中转缓冲, 这里是连续分配, 等于尾后, 等于预注册输入缓冲, 也就是 eager 模式下输入先 copy 进来再归约, 等于 8MB 本地表, 存放着 ptrs。这里说的是 graph 模式的指针表及条容量。
0x5.2 两个 ： 与
v1 的两个算法
数据被每个 rank 置入自己的公共缓冲区, 紧跟着读取全部 ws 份数 据进行本地累积随即写出, 并且读流量是 ws 与 N 的乘积, 存在冗余现象, 然而在小消息情形下, 带宽并非成为瓶颈, 节约收获的是第二轮同步。
// 来源: /sgl-/csrc//.cuh（主体）(sg, ,秩);, 遍历由网格下执行的本线程所涵盖的元素, 将ws个同位置的秩的数据进行累加, 对于从由.x乘以.x加上.x所得到的索引值开始, 直到大小值的范围, 以.x乘以.x为步长递增的索引值小于大小值的所有索引值, ((P*))等于。
你提供的内容似乎并不是一个完整、有意义的句子表述呀, 请你检查一下并提供新的正常句子以便我按要求进行改写。
便是 - + all-。每一个 rank 仅仅负责 1/ws 的分片, 将 ws 家这一片段读取过来进行累加以后再写回到自身的中转缓冲, （带着）；每一个 rank 把其他 rank 所计算好的分片取回来拼接成完整结果。总体流量大概是 2 × N, 代价是多出一轮带有内存序的同步。
在源码注释中写着这样一个正确性约束条件: 存在两个 stage, 这两 stage 必须运用同类相同的 tid 映射。跨设备的写可见性仅仅在相同 tid 的线程彼此之间才会得到保证, 在某个阶段里计算 start+i 的线程, 在另一个阶段之中也必然得由它去进行 start+i 这个操作。
在指针数组方面, 不是依据'距离'重排两项对象, 而是直接按照dp.ptrs来进行操作 , 有这样两种情况。
0..ws) 进行原序累加, 其目的在于使得所有 rank 的浮点累加顺序保持一致状态、输出位保持一致状态, 此点在 0xa.1 处展开。block 数上限为 36, 其原因是注释当中所提及的总线争用情况: 更多 SM 参与其中会导致争用状况加剧, 36 是 A100/H100 一带经过实测所得到的值。
算法切换的规则, 被写死在了一侧, 存在这样的情况: ws等于2的时候, 会一直被使用；而当处于full状态时, ws要小于或等于4, 并且。
0x5.3 v1 的 CUDA graph 支持
先来阐述一下graph所带来的约束, 每一步都要发射成百上千个, CPU逐个进行下发的开销是不能被忽视的, CUDA graph会把一步的序列记录下来并反复操作, 将这部分开销降低到近乎零, 其默认在graph内执行, 但是它要求所有参数（包含指针）要和执行时保持一致, 而AR的过程需要读取"每个rank上输入张量的地址", 在这个期间没办法进行跨rank的host同步来交换这些地址。
v1的处理方式为延迟注册, 在此期间, 每当碰到一次all-, 便将 的 * 参数指向 表的下一个空位（此时内容属于无效状态）, 与此同时，把本地输入指针记录到 当中, 结束之后, rs选取每个指针的 base去做 , 进行广播交换、打开映射, 把真实的ws个peer指针写入到当初的占位条目里。当 解引用时所获得的便是有效指针, 无需进行额外拷贝。而eager模式并不走这条路径, 它是先将输入copy到预注册的 之中, 之后再进行归约, 要多付出一次拷贝。
0x6. v2 的结构：存算分离
说到 v2 的文件头, 它描述了其结构的一种划分咧, 其中 CUDA 这一侧, 将其划分成了两个部分, 它是对于某事物的指针持有者呢（要知道所有的存储都是靠那一侧来进行分配跟拥有的哟）。还有那个 all -什么的, 它是属于一个纯函数啦 (input, comm, algo, )就是这样。再说说 v1 , 它把 IPC、注册表以及算法选择这些内容, 全都放置在了 C++ 类里面。而 v2 的 C++ 这一侧, 则是不含任何生命周期逻辑的, 那状态机是在另一侧的。
这种划分致使出现两个结果。其中一个结果是, JIT 实例仅仅依照 (数据类型, 空值, 是否为 PDL) 来进行编译, 并且一次编译就能为所有配置提供服务。另一个结果是, 所有的尺寸、块数以及阈值都变成了在侧面能够进行调整的数据, 而非编译期的常量。
v2 的存储平面
存储侧存在变化, 其变化方式是指针交换。v1 以手工方式进行 all-。v2 调用 torch 的 .进行操作 , 并且 , 为每个 rank 分配一块布局相同的内存 , 句柄交换是由 torch 完成 , 同时能得到 NVLS 地址 , 该地址用于 0x7.4。这块内存被切成三段:
推送: 二乘以ws个 （各自）ws个槽位对应「每位对等方写给我的地点（位置）」, 乘二为双相位, 相邻两次调用交替运用, 防止上一次未消费完毕就被覆盖（与v1同一思路）, 拉取: , 单块 :一百二十八字节乘以块数（每一个数据块一个 , 一百二十八字节对齐防止伪共享）
还有另外两个属于本地的（并非对称的）张量, 即（push算法的相位计数）以及（）。
, ws
用于 u64 的指针表, 当 ws等于8的时候占据8 MB, 可查看0x8）。其尺寸取值源自调优表给出的(arch, ws)期望值, 接着与16 MB cap（）进行取min的操作, 并且要做1 KB的对齐处理。
这套存储平面并非全专用, 在main里, fused TP QK - norm会去创建pull极小化的push - only v2实例, 且复用同一套与信号量协议。
0x7. v2 的三个算法
v2 的算法集合是 、、，以及 的 变体档。
v2 三个算法与 变体档的数据流
0x7.1
v1的为pull语义, 即等所有rank将数据放置妥当后, 而后各个自行去读取。v2增添的push是相反方向, 即本rank借助st..sys把数据写入全部ws个rank的push里归属自身的槽位, 接收方通过自旋轮询本地的ws个槽位, 待数据齐全后进行归约写出。此条路径不存在跨rank的情况, 唯一的同步是本地的相位计数, 它决定使用双的哪一半。
数据怎么样才算到达, 这得以载荷自身去评判断定它, 那怎样实施这个判定方式, 就是采用 +0.0 的这种位型来当作为「槽位为空」的那种能够起到保护警示作用的标记:
这段内容理解起来较有难度, 似乎很多代码或符号表述不完整, 以下是尽力按照你的要求改写: 一个地方, 它的来源是, jit里面的csrc部分。这里呢, 使用了一种类型, 等于某个值。这个值, 是加0.0位型, 它被用作空标记。还有个值, 是减0.0, 是数值等价的替代表示。在push之前, 要把载荷里真实出现的加0.0改写为减0.0。在IEEE 754情况下, 两者数值相等, 参加加法得出的结果不会改变。有个函数, 对DType类型的变量进行操作, 如果变量的值等于Trait的某个值, 就把变量的值改成Trait的另一个值。在poll端, 16B向量里面, 只要有任一分量是加0.0, 那就表明数据没到齐, 需要继续重读。
是这样子 的, 延迟下限是一回 写的时候还得有对端一回轮询命中才行。其代价值当是在消费给弄完之后, 非得要把槽位给回填成 +0.0, 好为下一回同相位调用去恢复那个「空」标记, 另外还有双相位 的额外显存呢。
0x7.2
流程和 v1 的同构情况是这样的, 先是信号量打旗加上等同 ws 个旗, 接着读取全部 ws 份进行归约并写入输出。其区别在于数据源有三种选择, 在 eager 模式下输入先进入本地 pull, 也就是付出进出两次拷贝这种情况；在 graph 模式下取指针表的一行, 此为零拷贝, 查看 0x8 能知晓这点；还有一种模式则查看 0x7.4。
以 red..sys..add 将信号量打旗, 对于每一套 block 而言。要全部运用语义，（此时只需确保所有 rank 都能抵达, 数据依赖由次序予以保证）。对 2shot 的情况给予使用, （即本 rank 把内容写入 peer的, 随后 peer 会展开读取操作）。
0x7.3
和 v1 属于同族, 其区别在于第二阶段进行了融合: 在每个 rank 读取 ws 份并归约出自身的分片后，将结果在原地写回到所有 ws 个 rank 的同一位置, 如此一来每个都变成了完整结果, all - 被融合到写回操作中。在 graph 模式下是真正的原地操作（out = in, 连输出分配都省去了）；在 eager 模式下最后再回。分片余数分配给前 rem 个 rank, 通常按照 1024 对齐, 切分边界较为规整。
0x7.4 变体档（NVLS）
其中的交换机具备这样的支持, 其针对特定地址, 对该地址进行写入操作时, 等同于同时向所有 GPU 写入, 而对该地址进行读取操作时, 是由交换机将所有 GPU 的对应数据进行归约后再返回。所以由此可知 all - 能够退化成两条指令:
数据类型为单精度浮点数的累加器, 交换机对ws份数据进行归约操作, 将结果通过广播写回到所有的进程等级处。
负责发射指令的是 SM 侧, 且仅负责此, 归约是在交换机完成的。把它作为 2shot 的一个有变种性质的档位的是 v2。存在两个实现方面的细节。
地址会由 torch symm-mem 一同返回, 当硬件不支持的时候, 会被设置为 None, 此时, 这一档会被关掉。
0x7.5 PDL
那个被称作PDL的东西（ , 以及在这之后的情况）, 促使后继的 提前开启运行, 使其与前序 结束阶段的操作出现重叠状况。总共有三个算法, 它们都在 的入口位置进行调用操作、在出口位置也进行调用操作。处于eager模式下的进出拷贝操作, 同样借助带有PDL的 来替换（当满足对齐条件并且 ≤1 GB / ≤8 MB的时候）, 从而实现copy--copy三段之间也能够出现重叠的情况。
0x8. v2 的 CUDA graph 指针表
思路与 v1 的延迟注册相同，实现上换成一张集中的 表（
, ws
u64）：
v2 的 graph 输入注册
存在着三个有着各不相同特性的守卫条件 , 其一为保护类条件, 当处于此时的上下文当中, 然而却又尚未真正进行捕获的这个阶段（所谓的预热过程中相应举动）, 是绝对不可以将表行拿去做消耗处理的, 不然的话行号方面就会跟真实捕获情况出现错位情形, 进而就会出现解引用到那些并未注册的行的糟糕状况, 所以呢与此同时要进行判断以及括号内的相关操作。其二讲的是同一守卫条件中与torch saver进行交互方面情况描述当把其打开并且往机器内输入地址时它们两者之间在特定范围内处于不稳定状态 , 没有通过括号内操作的话是绝对不通行graph模式 的, 一旦进入图内就会退化成eager语义。其三样属于同个守卫条件下情况 就是TC graph情况 括号内状态会在一定范围里被明确地排除掉的当前这种此时此刻, 相应链条执行机制的兜底情况是。
零乘以九, 算法的挑选方面, 针对架构与工作集的双轨进行调优的表格。
设在一侧, v1的阈值是被写死了的, 而就v2而言, 会将交叉点做成数据, 在运行的时候, 会依据消息 byte数来查询表格。