张量并行通信压到4比特:Echo把舍入误差送回下一次调用,1GB消息比NCCL快2.09倍
核心概要
Echo把张量并行的all-reduce改成4比特传输、高精度求和,并把每次舍入误差反馈到后续调用,使累积误差不随层数增长;在A100上1GB消息比16比特NCCL快2.09倍,Qwen3-8B七项基准平均只差0.4分。
深度剖析
张量并行的all-reduce夹在两个相互依赖的矩阵乘法之间,无法被计算掩盖,而它传输的比特数一直停在16位。Echo把这条通信改成4比特整数传输、高精度求和,训练和推理共用同一机制。 此前降低通信量的做法要么隐藏或加速all-reduce而不减少传输字节,要么在单次调用内做更细的分组量化、保留离群通道或旋转,把每次量化的误差压小,但都不补偿早先调用已经注入的误差。 在八块A100上用NVLink实测,1GB消息下未补偿的4比特all-reduce是16比特NCCL的2.65倍快、耗时为其37.8%;加上补偿后为2.09倍快,端到端首token时间最多提升1.16倍、训练步最多提升1.09倍,背景流量占满链路时分别升到1.44倍和1.38倍。
误差之所以能被控制,是因为all-reduce在求和时自己暴露了刚注入的舍入误差,而后续调用可以把这个误差减回去。 误差反馈此前用于梯度压缩,因为同一参数的梯度会在下一步重新出现;隐藏状态只产生一次就被下一层消费,推理时更没有训练循环,所以直接套用没有收益。 Echo用两个不同衰减率的滑动平均构成补偿器,在每次all-reduce前把修正加到输入上,并从输出残差刷新平均值;该补偿融合进all-reduce的入口阶段、更新与下一次矩阵乘法并行,1GB消息上只给关键路径增加0.90毫秒,而用独立PyTorch算子实现要增加14.09毫秒。
启示与展望
报道覆盖的是八块A100通过NVLink互联的单机设置,推理在Llama-3-8B与Qwen3-8B上以张量并行度8运行,训练从零预训练Qwen3-0.6B、同样为张量并行度8,未使用流水线、上下文、数据或序列并行。速度数字来自空闲互连以及背景点对点流量占50%、75%、100%三种负载,精度评估使用统一评测框架并采用各任务的标配样本数。作者提出序列并行的reduce-scatter与专家并行的combine共享同样的两个结构性质,是下一步的自然目标。
深度无关的传播界依赖一个关于已训练网络的假设,即网络不会逐层持续放大同一方向的扰动;作者同时构造了一个满足噪声模型却使偏差随深度指数增长的网络,说明这类假设不可省略。反向传播的源界在均方意义下成立,但需要归一化雅可比沿轨迹平滑变化,且常数因误差与修正经过不同算子而变差。4比特推理在Llama-3-8B上七项基准平均损失1.66分,GSM8K严格匹配从49.81降到46.17,作者用前半层提升到8比特的早层策略恢复大部分损失,该策略只用于服务实验,训练实验全程使用均匀4比特。代码将在录用后发布。
