一句话：提出硬件/运行时无关的 PartIR，用可组合 SPMD 分片策略与增量 IR 重写实现可预测的神经网络分区。
训练现代大型神经网络需要组合数据、模型或优化器分片等多种并行化策略，策略复杂度不断上升。PartIR 是硬件与运行时无关的神经网络分区系统，其核心特点是：允许用户定义或自动推导的多种分片策略组合；将策略与机器学习实现解耦；并依据一组明确定义的通用规则进行分区，从而保证可预测性。它采用类似 schedule 的 API，在每种策略后增量重写 ML 程序的中间表示（IR），使模拟器和用户能够验证策略性能。PartIR 已用于训练大模型并覆盖多样模型架构，展示了其可预测性、表达力和性能。
Training modern large neural networks (NNs) requires a combination of parallelization strategies, including data, model, or optimizer sharding. To address the growing complexity of these strategies, we introduce PartIR, a hardware-and-runtime agnostic NN partitioning system. PartIR is: 1) Expressive: It allows for the composition of multiple sharding strategies, whether user-defined or automatically derived; 2) Decoupled: the strategies are separate from the ML implementation; and 3) Predictable: It follows a set of well-defined general rules to partition the NN. PartIR utilizes a schedule-like API that incrementally rewrites the ML program intermediate representation (IR) after each strategy, allowing simulators and users to verify the strategy's performance. PartIR has been successfully used both for training large models and across diverse model architectures, demonstrating its predictability, expressiveness, and performance.
@inproceedings{abs-2401-11202,
title = {PartIR: Composing SPMD Partitioning Strategies for Machine Learning},
author = {Sami Alabed and Daniel Belov and Bart Chrzaszcz and Juliana Franco and Dominik Grewe and Dougal Maclaurin and James Molloy and Tom Natan and Tamara Norman and Xiaoyue Pan and Adam Paszke and Norman A. Rink and Michael Schaarschmidt and Timur Sitdikov and Agnieszka Swietlik and Dimitrios Vytiniotis and Joel Wee},
booktitle = {ASPLOS},
year = {2025},
doi = {10.1145/3669940.3707284}
}