Training of modern large neural networks (NN) requires a combination of parallelization strategies encompassing data, model, or optimizer sharding. When strategies increase in complexity, it becomes necessary for partitioning tools to be 1) expressive, allowing the composition of simpler strategies, and 2) predictable to estimate performance analytically. We present PartIR, our design for a NN partitioning system. PartIR is focused on an incremental approach to rewriting and is hardware-and-runtime agnostic. We present a simple but powerful API for composing sharding strategies and a simulator to validate them. The process is driven by high-level programmer-issued partitioning tactics, which can be both manual and automatic. Importantly, the tactics are specified separately from the model code, making them easy to change. We evaluate PartIR on several different models to demonstrate its predictability, expressibility, and ability to reach peak performance..

翻译：现代大型神经网络（NN）的训练需要结合数据、模型或优化器分片等多种并行化策略。随着策略复杂度提升，分区工具需具备两个关键特性：1）表达能力——支持简单策略的组合；2）可预测性——能够通过分析预估性能。本文提出PartIR——一种神经网络分区系统设计方案。PartIR专注于增量式重写方法，且与硬件及运行时环境无关。我们设计了一套简洁而强大的API用于组合分片策略，并配套开发了验证策略的模拟器。整个流程由程序员发布的高层级分区策略驱动，这些策略既支持手动设定也可自动生成。关键设计在于策略规范与模型代码解耦，便于灵活调整。通过在多个不同模型上的评估，PartIR展现了其在可预测性、可表达性及实现峰值性能方面的能力。