做了这么多年分布式训练我越来越觉得大模型训练跑不动的瓶颈已经从GPU单卡算力转移到了卡与卡之间的通信。前段时间跟团队复盘一个70B参数模型的训练任务单看GPU计算利用率其实不算低可整个step的耗时里通信等待和通信重算占了将近四成。这里面的核心变量就是张量并行带来的密集all-reduce操作。CAIS这个方向我跟踪了很久它做的事情很直接把计算塞进交换机内部让网络设备在转发数据的同时顺手把张量并行的聚合计算做了。这个思路听起来就是换了个地方做加法但真正落地时牵扯到数据平面流水线、协议解析、控制面调度、数值精度一整套问题。这篇文章我就从张量并行的通信账开始算起把CAIS为什么敢在交换机里做计算、计算感知到底感知什么、怎么从P4原型跑到真实集群一条线讲清楚。适合正在做分布式训练优化、网络系统设计或者研究在网计算的工程师和研究生。1. 张量并行的通信账算清楚才能看明白CAIS的发力点1.1 大模型为什么要拆到多卡上模型大了单卡显存放不下这道理大家都懂。但拆这个词有两种拆法一种是把训练数据切成多份每张卡各算各的这叫数据并行另一种是把模型参数本身切到多张卡上每张卡只负责模型的一部分这叫模型并行。张量并行是模型并行里最常用的一种专门针对Transformer结构来做参数切分。具体来说Transformer每一层主要由多头注意力和两个MLP线性层组成。张量并行把注意力头的权重按头维度切开把MLP的权重按隐藏层维度切开让不同GPU各自计算一部分矩阵乘法最后通过all-reduce把各部分结果合并。这样做的直接收益是单卡显存压力骤降一个70B模型如果用8卡张量并行每卡只需要存约9B参数对应的权重和优化器状态训练和推理都跑得动。但代价也随之而来前向计算每一层的输出需要同步反向传播每一层的梯度也需要同步而且这种同步是层级的、高频的不是像数据并行那样累积多个梯度后一次性通信。换句话说张量并行让GPU每算一步就要碰一次头通信被硬生生塞进了计算的关键路径。1.2 一路all-reduce下来花了多少时间我用一个具体数字来感受一下。假设模型隐藏层维度是4096batch size为8序列长度是2048训练精度用FP32。那么一个张量并行的激活张量大小就是8乘2048乘4096再乘4字节算下来大约268MB。这只是某一个线性层输出的激活值进行一次all-reduce。如果是8卡ring all-reduce每张卡收发约两倍数据量总通信量大约是536MB。在100Gbps也就是12.5GB/s的网卡带宽下纯传输时间就要约43毫秒。这还只是一层的MLP输出一个Transformer层里有多处需要all-reduce前向一次、反向一次加起来一个step光张量并行的通信时间就达到几百毫秒甚至更多。相比之下GPU算一个layer的forward大概也就几十毫秒。通信和计算在时间尺度上已经处于同一量级甚至通信更长。网络不再是后台任务而是和计算争夺时间的前台瓶颈。这才是CAIS这类在网计算方案出现的最直接原因——既然通信省不掉那就让通信过程中把计算也捎带做了。1.3 数据并行、流水线并行和张量并行的通信差异三种并行模式的通信特征完全不同这也决定了不同在网计算方案的适用场景不一样。数据并行通信的是梯度频率相对低通常梯度累积一定轮次后再做一次全局all-reduce。它对聚合带宽的需求非常高但对延迟没那么敏感属于大块头、低频次。流水线并行通信的是层间激活只在管道边界发生通信量小但次数多更看重延迟稳定。张量并行的通信特点是发生频率极高几乎每个layer都有消息尺寸适中而且是同步屏障式的任何一张卡没收到所有分片结果下一层计算就动不了。CAIS选择张量并行作为切入点正是看中了这个同步屏障的特点。因为张量并行不容忍通信抖动一旦某个包的延迟出现毛刺整层计算全部停滞。把all-reduce放到交换机内执行可以让所有分片结果在交换机汇聚后一次性返回天然形成一个有序的同步点比在GPU侧做多轮ring式转发更可控。2. 交换机内计算不是新概念但做LLM张量并行是另一回事2.1 可编程交换机到底能算什么在网计算最底层的支撑是可编程交换机尤其是以P4语言为核心、采用可重构匹配表加动作流水线的设备。这类交换机的转发芯片里数据包从网口进来后要经过解析器、匹配查找、动作执行、排队和调度这几个阶段全程走硬件流水线速度可以达到几个Tbps。能在这个流水线里执行的计算并不是通用的。常见的动作包括修改包头字段、做简单的算术运算、更新寄存器值、按自定义逻辑丢弃或复制数据包。以P4代码来看我可以在一个动作里写temp register[index] payload_value再把结果写回寄存器这就是最朴素的交换机内加法。所谓交换机内计算本质上就是把这类寄存器算术能力组合成有用的通信原语。但对LLM训练来说这里有个很现实的坑数据平面流水线不擅长浮点运算。GPU算的是FP32、BF16、FP16而交换机寄存器通常做的是整数加法和位运算。把256个字节的浮点张量片段在交换机里做累加必须要先做浮点转定点的处理或者把浮点表示拆成符号位、指数、尾数来做定点对齐。这既是精度问题也是性能问题。2.2 在网all-reduce和GPU侧算法有什么本质区别GPU侧最常见的all-reduce实现是ring算法所有GPU形成一个逻辑环每个节点同时从上一个节点接收数据块、把本地对应块累加、再传给下一个节点。环的优点是扩展性好、链路利用率高数据并行训练里的梯度聚合基本都靠它。但ring算法有个隐藏代价数据要经过多跳转发每个节点都要完整收下一整份块才能做累加延迟和拓扑跳数直接相关。而且ring all-reduce会把通信精细切分每个块都比较小网卡和CPU的排队次数多实际吞吐很难到达理论带宽。交换机内计算则是完全不同的路径。所有GPU把自己的张量分片发给交换机交换机在收到同一张量ID的各个分片后边转发边做累加。比如8个GPU的数据同时打到一台交换机交换机通过8个端口的队列分别读出数据块在寄存器和转发逻辑里完成累加最终把聚合结果从某个端口送回所有GPU或者是按需送到对应的参数服务器。这个过程相当于把多跳环化简成了一跳聚合。更重要的是不管多少个节点参与数据只在交换机内走一次不会像ring那样经过P-1次转发。2.3 张量并行为什么比数据并行更适合卸载到交换机数据并行的梯度all-reduce数据量太大动辄几百MB甚至上GB交换机寄存器空间极其有限不可能把完整梯度留在交换机里累加。现实做法只能是把数据切成小块、多次往返而每次都切分和重组本身就是额外开销。张量并行则不同。每次all-reduce的张量只有几十到几百兆而且计算精度可以接受定点化后再还原。更重要的是张量并行的同步点密集单个张量的并归延迟直接决定整个step的耗时。对这类延迟敏感、消息可控、频率极高的通信模式交换机内计算才有明显收益。另外还有一层拓扑上的契合。数据并行往往是跨机、跨机柜的大范围通信数据要经过很多层交换机。张量并行通常优先把同一设备组安排在接入层或机架级交换机的相邻端口下通信路径短、拓扑固定。交换机在一条路径上就能聚合所有分片不需要跨多台交换机做协同计算。CAIS把目标锁定在张量并行其实就是锁定了通信拓扑最可控的那部分流量。3. CAIS的计算感知给网络设备装上任务视角3.1 计算感知要解决的三个问题把计算下沉到交换机如果只是做一个死板的见到匹配头就累加那很容易出错。CAIS强调计算感知核心是解决三个问题算谁的、怎么算、何时算。算谁的训练进程里同时存在很多张量有的是权重梯度有的是激活值有的是一次性临时变量。交换机不能对所有流量都做同样处理必须通过张量ID标识来区分。这个ID可以是自定义包头里的字段也可以结合源IP、目的IP和UDP端口做匹配。CAIS在控制面会维护一张张量登记表每个张量什么时候发起all-reduce、谁参与、数据格式是什么都提前告诉交换机。怎么算不同张量需要的聚合语义不同。有的要做sum有的要做average还有的在某些场景下需要做最小值或最大值作为归并结果。计算感知意味着交换机执行的动作由语义决定而不是永远都是加法。用P4的匹配动作表来做这件事非常自然不同的匹配条目对应不同的动作参数交换机流水线不需要改变逻辑只需要改表项。何时算这涉及优先级和时序。张量并行每个step有固定的计算依赖图某些all-reduce必须在前某些可以后到。如果交换机只是按数据到达顺序处理后到的关键张量被前面的流量堵住照样会拖慢整个step。CAIS需要感知当前的训练阶段一般通过租约或会话ID标识为不同的张量流设置不同的队列优先级。3.2 控制面和数据面怎么分工CAIS的整体架构从我个人的理解上可以拆成两层控制面负责知道数据面负责执行。控制面由CAIS Manager进程承载通常运行在训练集群的管理节点上或者作为独立控制器部署。它通过带外网络与可编程交换机建立控制通道负责四类工作一是收集训练任务的张量并行配置包括并行度、模型层数、预期通信频率二是监控链路的实时状态从交换机遥测模块获取队列长度、端口利用率和丢包计数三是根据状态和任务信息下发或更新匹配表项、动作参数、队列调度策略四是处理异常比如某个端口拥塞时调整重定向规则或者在交换机计算能力不足时通知GPU端关掉卸载功能。数据面则是交换机上运行的P4程序它只做一件事按表项处理数据包。CAIS的数据面不需要理解模型或梯度这些高层概念它只需要识别张量ID、查表、按动作执行累加或转发、更新寄存器、管理队列。这种分工的优点很明显控制面和数据面可以独立演进。训练框架升级、模型结构变化只需要改CAIS Manager的调度策略和Agent端的接口P4数据平面代码可以保持稳定。3.3 负载感知调度与回退机制真实训练场景里交换机不会只服务一个张量并行组。同一台接入交换机下可能挂了多个训练任务或者同时有存储流量、控制流量在跑。CAIS的计算感知必须包含负载感知。我的建议是给交换机内计算设计两级门限一级是队列深度门限另一级是寄存器使用率门限。当队列深度超过某个阈值说明转发面已经拥塞此时如果再把计算流量全部堆进来只会加剧排队延迟CAIS Manager会主动把新到达的大张量分流走不再参与在网计算。寄存器使用率则是交换机特有的瓶颈因为数据平面寄存器是固定大小的同一时刻能承载的中间累加结果有限。一旦接近上限就必须拒绝后续的计算任务否则老任务的寄存器值会被新任务覆盖造成结果错误。回退机制是整条链路的保险丝。CAIS的Agent运行在每个GPU节点上它会向NCCL风格的通信库注册一个自定义all-reduce算子。正常情况下算子把数据发给交换机等待交换机处理完成后接收结果。如果它发现交换机在超时时间内没有返回或者遥测报告说计算功能已关闭就立即回退到传统的GPU侧all-reduce。这个回退必须是透明的训练进程无感知只是慢一点但不能中断。4. 从P4原型到真实集群CAIS的落地实操要点4.1 第一步用BMv2跑通最小原型没有真机的情况下想验证CAIS的数据平面逻辑最方便的方案是BMv2软件交换机。它是一个用C实现的P4参考交换机支持P4_16语言配合Mininet可以搭建虚拟网络拓扑完全模拟多台GPU通过交换机通信的场景。我在做原型时通常按这么几步走用P4_16写一个最小交换机程序解析入口处增加一个自定义张量ID头不需要从头到尾完整解析整个TCP/IP栈。在数据平面用register数组模拟累加器比如定义128个寄存器槽位每个槽位存一个定点化的累加结果。动作里做acc acc payload_word再把结果放回同一个槽位。用BMv2的simple_switch_grpc启动交换机通过P4Runtime下发匹配表项把携带同一张量ID的包映射到同一个寄存器槽位。在Mininet里起多个host用Python脚本发送带张量ID的UDP包验证多个host的数据是否在交换机内完成累加最终结果是否正确。这一步核心目标不是性能而是验证语义正确性。我会故意用不同顺序发送数据包确认寄存器累加结果和发送顺序无关全部结束后输出一致。只要语义没问题后续才会去硬件上评估性能。4.2 第二步和NCCL集成接管all-reduce原型验证之后CAIS真正要进入训练流程必须和NCCL这一层做对接。NCCL是NVIDIA的集合通信库all-reduce是它的核心原语之一。现在的NCCL生态支持注册自定义后端也可以通过调整通信图的方式把特定通信原语重定向到自研实现。CAIS在GPU端的接入方式是在通信库之上封装一个轻量Agent。训练框架调用all-reduce时Agent拦截请求判断当前通信任务是否满足CAIS的卸载条件。条件通常有三个张量尺寸在预设范围内通信组与交换机的端口拓扑匹配当前交换机的计算负载低于门限。满足条件就执行在网计算路径不满足就走NCCL原始路径。这里不建议把卸载做得过于激进。我实际测试时发现小张量走CAIS路径没有收益因为数据包携带的payload太小交换机解析头、查表、排队的时间占比反而比GPU侧纯传输更大。通常单次all-reduce数据量在1MB以上时CAIS才会有肉眼可见的收益。所以CAIS Agent里的条件判断不是摆设而是性能保障。4.3 关键参数与监控指标落地部署时你会面临一整套参数调优我整理了一张常用参数清单参数项推荐初始值说明卸载阈值1MB小于阈值的张量不卸载到交换机保留GPU侧all-reduce寄存器槽位数128同时可承载的在网计算任务数过多会增加数据平面面积定点化缩放因子1024浮点转定点时的固定缩放要在数值范围和溢出风险之间权衡累加超时门限500us超过此时间未收到交换机计算结果Agent回退到GPU侧算法高优先级队列占比30%预留给关键张量流的研究队列缓冲监控指标的选取也要跟着性能目标走。最重要的三个一是交换机内计算完成时间它决定张量并行的同步延迟二是回退率如果一个训练任务里超过10%的all-reduce都走了回退路径说明卸载条件设置不合理或者交换机算力不够三是累计结果的数值误差切换定点化参数后要定期和GPU侧纯浮点结果做对比防止长期训练中误差漂移。4.4 部署注意事项部署到真实集群有四个容易踩的坑我集中提醒一下。第一拓扑必须是收敛的。CAIS假设通信组内所有GPU都连在同一台或同两台交换机下面如果工作负载被跨机柜调度数据包要经过多级交换机才能汇聚那在网计算就必须在各级交换机上协同做复杂度会指数级上升。初期部署建议就限制在单台接入交换机范围内把8卡或16卡张量并行组绑在一个机柜里。第二交换机端口速度要匹配GPU网卡速度。如果GPU用200G网卡交换机单端口只有100G那么在网计算的聚合结果回传时就会成为瓶颈。在计算部署时我建议对核心回传路径单独做流量整形避免其他流量抢占结果回传端口。第三控制面超时设置要保守一点。CAIS Manager与交换机的控制通道一旦断掉数据平面的旧表项仍然存在但失去了动态调整能力。必须设置控制面会话超时并在超时后让数据面进入只转发不计算的安全模式防止陈旧表项导致的数据处理错误。第四做定点化时一定要考虑数据分布。梯度值有时候很小有时候很大固定缩放因子会导致小梯度被截断。稳妥的做法是按张量统计梯度范数后动态调整缩放因子CAIS Manager把缩放因子通过控制通道下发到数据面而不是Agent这边玩命调一个魔法数字。5. 踩坑记录在网计算项目的常见问题速查5.1 为什么卸载之后反而更慢了这是我见过最多的情况把all-reduce卸载到交换机后训练step时间不但没降反而涨了20%。问题通常出在三个地方。一是包太小。前文提过小张量不适合卸载。如果你的模型并行度很高每个GPU分片很小比如单次通信量只有几百KB那不要卸载直接走NCCL。CAIS的卸载阈值建议根据自己的实际带宽延迟积来标定我习惯的做法是测一组梯度值、从128KB到64MB的卸载收益曲线找到拐点再设阈值。二是交换机排队混流。训练流量和存储流量、管理流量走在同一个交换机队列里高优先级队列配置不当关键张量照样被堵。解决办法是把CAIS的计算流量单独放到一个队列开启优先级流控同时通过ECN门限限制其他流量对队列的挤压。三是控制面开销反噬。每次all-reduce都让Agent去询问Manager能不能卸载这个问询本身也是开销。实测中高频小张量的问询开销甚至超过了在网计算的收益。解决方法是把决策做成批量式的Agent只在每个训练step开始时拉取一次当前可卸载张量列表step内部不再问控制面。5.2 数值精度怎么保证交换机内浮点累加改成定点化之后最头疼的是精度问题。我在早期测试中遇到过一次训练loss不下降的情况排查到最后就是定点化舍入误差在反向传播里被放大了。建议做三件事。第一缩放因子不要全局统一按张量维度分开配置因为不同层的梯度量纲差异很大。第二对舍入误差做随机舍入而不是直接截断这样误差在统计上是无偏的训练过程不容易出现系统性的偏移。第三周期性地把交换机计算结果和GPU侧纯浮点参考值做误差对比一旦相对误差超过千分之一就自动切换回退机制。5.3 多租户多任务场景下的优先级冲突实际数据中心里交换机不可能只为一组训练服务。当多个任务的all-reduce都想去占用交换机寄存器时会出现互相抢占资源的情况。我踩过一个具体坑两个训练任务用同一批交换机端口它们的张量ID编码规则没协调结果A任务的包匹配到了B任务的累加器聚合结果直接错乱。后来我在CAIS的设计里强制要求所有任务通过Manager申请全局唯一的租户ID和会话ID数据平面匹配表必须是租户ID会话ID张量ID的三元组组合从根上隔离。另外要提一句如果交换机厂商SDK不支持端口级计算配额那多任务场景下最好还是用独立的物理交换机来隔离不同训练任务否则很难预测资源抢占行为。做CAIS这类方向最大的感受是在网计算听起来像网络工程师的事但实际它横跨了并行计算、编译器、通信库、网络协议好几个领域。我自己的体会是不要一上来就追求把所有通信都卸载到交换机先把一个最小闭环跑透确认NCCL能识别的自定义后端、确认P4数据面处理语义正确、确认回退路径稳定再逐步扩大规模。这个框架的后续扩展空间其实很大比如把序列并行和专家并行的通信模式也纳入计算感知的调度范围甚至把低精度训练和交换机内聚合做联合优化。我的一个落地建议是先选定一种固定模型和固定并行度把CAIS的收益曲线完整测一遍这样后续做动态调度时才有判断依据不然很容易被各种性能毛刺带偏方向。