PartIR:组合面向机器学习的SPMD划分策略
PartIR: Composing SPMD Partitioning Strategies for Machine Learning
AI总结:
PartIR是一个与硬件和运行时无关的神经网络划分系统,通过可组合的分片策略API和模拟器,实现表达力强且可预测性能的SPMD并行化划分,并在多个模型上验证了其达到峰值性能的能力。
摘要:
现代大型神经网络(NN)的训练需要结合多种并行化策略,包括数据、模型或优化器分片。当策略复杂性增加时,划分工具必须满足两个要求:1)具备表达力,能够组合更简单的策略;2)具备可预测性,能够解析地估计性能。我们提出PartIR,这是我们的神经网络划分系统设计。PartIR专注于一种增量式重写方法,并且与硬件和运行时无关。我们提供了一个简单但强大的API,用于组合分片策略,并提供一个模拟器来验证这些策略。该过程由程序员发布的高级划分策略驱动,这些策略既可以是手动的,也可以是自动的。重要的是,这些策略与模型代码分开指定,使其易于更改。我们在多个不同模型上评估PartIR,以证明其可预测性、表达力以及达到峰值性能的能力。
英文摘要:
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..
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。