PartIR中复合SPMD划分策略的自动发现
分布式、并行与集群计算
2022-10-20 v1 机器学习
神经与进化计算
摘要
大型神经网络模型通常通过单程序多数据(SPMD)范式下多种高级并行策略的组合进行训练。例如,训练大型Transformer模型需要结合数据、模型和流水线划分,以及优化器分片技术。然而,针对众多模型架构和加速器系统识别高效的组合需要大量人工分析。在本工作中,我们提出了一种通过目标导向搜索来识别这些组合的自动划分器。我们的关键发现是,基于蒙特卡洛树搜索(Monte Carlo Tree Search)的划分器将特定于划分的编译器分析直接融入搜索并以引导目标为指导,可匹配针对各类模型的专家级策略。
引用
@article{arxiv.2210.06352,
title = {Automatic Discovery of Composite SPMD Partitioning Strategies in PartIR},
author = {Sami Alabed and Dominik Grewe and Juliana Franco and Bart Chrzaszcz and Tom Natan and Tamara Norman and Norman A. Rink and Dimitrios Vytiniotis and Michael Schaarschmidt},
journal= {arXiv preprint arXiv:2210.06352},
year = {2022}
}