Jan, 2024
PartIR: 机器学习 SPMD 划分策略的构建
PartIR: Composing SPMD Partitioning Strategies for Machine Learning
Sami Alabed, Bart Chrzaszcz, Juliana Franco, Dominik Grewe, Dougal Maclaurin...
TL;DR现代大型神经网络（NN）的训练需要结合数据、模型或优化器分片等并行化策略。当策略变得越来越复杂时，分区工具需要具备 1）表达能力，以允许复合简单的策略，2）可预测性，以便分析性能估计。我们提出了 PartIR，一种 NN 分区系统的设计。PartIR 着重于增量重写，并且不依赖于具体硬件和运行时环境。我们提供了一个简单但强大的 API 用于组合分片策略，并提供模拟器进行验证。该过程由高级程序员指定的分区策略驱动，可以手动或自动执行。重要的是，这些策略与模型代码分开指定，易于更改。我们在多个不同的模型上评估了 PartIR，以展示其可预测性、表达能力和达到峰值性能的能力。