PartIR: Composing SPMD Partitioning Strategies for Machine Learning

要約

最新の大規模ニューラル ネットワーク (NN) のトレーニングには、データ、モデル、またはオプティマイザー シャーディングを含む並列化戦略の組み合わせが必要です。
戦略が複雑になると、パーティショニング ツールには 1) 表現力があり、より単純な戦略を構成できること、2) パフォーマンスを分析的に見積もるために予測可能であることが必要になります。
NN パーティショニング システムの設計である 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..

arxiv情報

提供元, 利用サービス

arxiv.jp, Google