PartIR: Composing SPMD Partitioning Strategies for Machine Learning

现代大型神经网络（NN）的训练需要采用数据、模型或优化器分片等并行化策略的组合。当策略变得更加复杂时，划分工具必须具备两个特点：1）表达能力强，能够组合简单策略；2）可预测性强，能够通过分析估计性能。我们提出了PartIR，一个神经网络划分系统的设计。PartIR专注于增量重写，不依赖具体硬件和运行环境。我们提供了一个简单但功能强大的API，用于组合分片策略，并提供一个模拟器来验证它们。这个过程由高级程序员发布的分片策略驱动，可以手动或自动执行。重要的是，这些策略与模型代码分开指定，易于修改。我们对几个不同的模型进行了PartIR的评估，以展示其可预测性、表达能力和达到最高性能的能力。

设计一种可预测性高、可扩展性强的神经网络分区系统

PartIR是一种基于增量重写的神经网络分区系统，提供了简单但强大的API来组合分区策略，并且使用模拟器进行验证，同时可以手动或自动地使用高级编程语言来指定分区策略。

PartIR系统具有可预测性、可扩展性、可组合性和硬件和运行时不可知性的特点。论文在多个模型上进行了实验，证明了PartIR系统的可预测性、可表达性和能够达到最高性能的能力。

与神经网络分区相关的研究包括：Dynamic Neural Network Surgery、Deep Gradient Compression、TensorFlow等。