一张大模型推理或训练的计算图,动辄成千上万个算子,单个计算单元吃不下,也跑不快。XPU 编译器要做的事,就是把这张大图切开,分给多个计算单元协同执行,同时保证切分带来的额外开销不超过并行带来的收益。切法好不好,直接决定了加速比的天花板。

第一种切法是数据并行:把输入按批次切开,每个计算单元拿一份数据、跑同一份模型,最后汇总梯度或结果。这种方式实现简单、扩展性好,但每个单元都要存一份完整模型参数,显存占用随单元数线性增长,模型太大时根本放不下。

第二种是模型并行:把权重本身切开。按层切叫流水线并行,每组单元负责连续的几层,数据像流水线一样依次通过;按算子内部切,比如把一个大矩阵乘法按输出通道切成几块,叫张量并行。模型并行解决了单卡装不下的问题,但引入了单元之间的中间结果传递,切得越碎,通信量越大。

真正的难点在算子融合与显存搬运的取舍上。编译器会把相邻的小算子融合成一个大算子,比如把卷积、偏置和激活函数融成一次计算,中间结果直接留在片上缓存或寄存器里,不用写回显存再读回来。融合省的是显存搬运,而显存带宽往往是第一瓶颈。但融合也有代价:融得越大,单个算子的工作集越可能超过片上存储,反而要反复切块;有些算子的融合还会破坏原有的并行度。所以编译器要按片上存储容量、算子形状和带宽算一笔账,只在净收益为正时才融合。

还有一个关键是切分粒度与通信重叠。图切得太粗,负载不均,快的单元要等慢的;切得太细,通信和同步开销会吃掉并行收益。成熟的做法是让编译器做静态调度,计算与通信重叠执行:一个单元在算当前块的同时,后台把下一块需要的数据搬进来。调优时可以看两个指标:计算单元的利用率和显存带宽的利用率,前者低说明切分或调度有问题,后者先顶满说明融合还不够激进。切图不是越碎越好,而是在并行收益、通信开销和显存搬运之间找到平衡点。

最后是动态形状带来的麻烦。真实业务的输入长度、批量大小经常变化,编译期定死的切分方案在换一批数据时可能就不再最优。有的编译器会在运行时做自适应:按实际张量形状重新选择分块大小,或者在几个预编译好的切分方案之间切换。这也提醒我们,评估切分策略时不要只看峰值算力下的跑分,要看业务真实的形状分布,取舍的答案往往藏在最常见的那几种形状里。