用于延迟张量并行性的亲和感知分片

延迟张量并行 (DTP) 消除了张量并行 Transformer 推理的阻塞全归约。每个设备都会立即将自己的部分输出添加到其剩余流（并广播它），但稍后才会收集（接收）其他设备的部分 $δ$ 模块。因此，TP 到 DTP 的更改相当于真正的架构更改，并且密集的 Transformer 模型需要在适应后重新训练或提炼。我们证明，DTP 打破了 FFN 内部神经元和注意力模块内部 KV 头的排列对称性，并且这种对称性破坏使分片本身成为一个建模决策。我们证明，通过在分片之前排列密集模型，最大化设备上共置的 KV 头和 FFN 神经元之间的亲和力，可以加速蒸馏或再训练过程。亲和力是通过失去头的贡献对每个神经元输出造成的损害的一阶近似来测量的，并且使用坐标上升优化器来最大化同位亲和力，该优化器通过对 KV 头分区进行详尽的搜索来交替神经元的精确平衡分配。对于 Qwen3-0.6B 和 Danube3-500M，整个过程在一个 GPU 上需要不到两分钟。在这些模型上，在 $δ=1$ 时，在我们测试的整个 10k 步范围内，亲和优化布局以朴素连续布局所需步骤的大约一半到三分之二达到任何蒸馏目标，并且每个优化种子都击败了每个连续种子以及十六个随机布局中的所有布局（除了其中一个）。我们还表明，初始化时的同位亲和力分数可以预测训练后基本模型的 KL，涵盖从反优化到优化的 17 种布局（Pearson $-0.81$ 和 $-0.89$）。