谷歌开源 TPU Sync:KV Cache 传输延迟降低超 90%

11 月 13 日,谷歌发布开源 C++ 库 TPU Sync,专门用来搬运 KV Cache 和权重。它可以在 TPU 的 HBM、主机 DRAM、存储设备和网络之间搬数据,谷歌给出的数字是在 235B 到 480B 的 MoE 模型上,KV 传输延迟降低超过 90%。

缓存搬运是大模型训练的真瓶颈

MoE 模型做训练和推理时,专家路由会让每一层要在多个设备之间交换 KV Cache,跨设备的搬运量往往超过真正算子运算的开销。延迟降不下来,稀疏结构带来的好处就被通信吃掉。谷歌说 TPU Sync 走的是一套跨介质统一路径,HBM、DRAM、存储和网络之间不再需要各自写一套搬运逻辑。这类库的价值通常不在单个功能,而在把关键环节开源让第三方接着优化,省掉重复造轮子的成本,这对中小团队尤其划算。235B 到 480B 这个区间正是当前主流 MoE 模型的参数规模,90% 以上的延迟降幅意味着同样硬件下能塞进更大的批量。

开放权重库的定位是补齐工具链

这类库的价值通常不在单个功能,而在把一整条链路的关键环节开源,让第三方能在上面接着做优化。谷歌 TPU 硬件本身的代际节奏此前已有过梳理,站内 谷歌 TPU 8tx 与 8i Max 四款第八代系统 那篇提到过带宽指标,两篇可以对照着看,一篇讲硬件,一篇讲怎么把硬件用满。


把通信层做成通用库,是谷歌推 TPU 生态的一贯做法。

全部文章归档与模型专题