PartIR：组合机器学习的 SPMD 分区策略
机器学习
2025-05-07 v4 分布式、并行与集群计算
编程语言
摘要
训练现代大型神经网络（NN）需要结合涵盖数据、模型或优化器分片的并行化策略。当策略复杂性增加时，分区工具必须具备：1) 表达力强，允许组合更简单的策略；2) 可预测，能够分析性地估算性能。我们提出了 PartIR，这是我们对 NN 分区系统的设计。PartIR 专注于增量式重写方法，并且与硬件和运行时无关。我们提供了一个简单但强大的 API 用于组合分片策略，以及一个模拟器来验证它们。该过程由程序员发出的高级分区策略驱动，这些策略既可以是手动的也可以是自动的。重要的是，策略与模型代码分开指定，使其易于修改。我们在多个不同模型上评估了 PartIR，以展示其可预测性、表达力以及达到峰值性能的能力。
引用
@article{arxiv.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},
journal= {arXiv preprint arXiv:2401.11202},
year = {2025}
}