Integrated Hardware Architecture and Device Placement Search

📄 arXiv: 2407.13143v1 📥 PDF

作者: Irene Wang, Jakub Tarnawski, Amar Phanishayee, Divya Mahajan

分类: cs.LG, cs.AR, cs.DC

发布日期: 2024-07-18

备注: Accepted at the 41st International Conference on Machine Learning (ICML), 2024

🔗 代码/项目: GITHUB

💡 一句话要点

PHAZE：通过联合搜索硬件架构与设备放置策略优化分布式深度学习训练。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 分布式训练 硬件架构搜索 设备放置策略 整数线性规划 动态规划

📋 核心要点

  1. 现有分布式深度学习训练方法难以兼顾硬件架构与设备放置策略的动态交互，导致资源利用率不高。
  2. 提出PHAZE，通过联合优化硬件架构（如张量/向量单元）和设备放置策略，实现计算、内存和数据分布的平衡。
  3. 实验表明，PHAZE 在大语言模型训练中，相较于 TPUv4 和 Spotlight 等现有方法，实现了更高的吞吐量。

📝 摘要（中文）

本文首次探索了硬件加速器架构和设备放置策略的协同优化，旨在提升分布式深度学习训练的效率。通过新颖的算法，平衡计算资源、内存使用和数据分布。架构搜索涉及张量和向量单元的数量及维度，片上和片外内存配置，以及微批大小和激活重计算/存储决策。针对每个架构配置，使用整数线性规划（ILP）寻找最优算子调度方案。ILP结果与动态规划相结合，确定最优设备放置策略，结合数据、流水线和张量模型并行。实验表明，该方法在大语言模型上实现了比TPUv4和Spotlight更高的吞吐量。PHAZE的完整源代码已开源。

🔬 方法详解

问题定义：现有分布式深度学习训练系统通常独立地设计硬件架构和设备放置策略，忽略了两者之间的相互影响。这种分离的设计可能导致计算资源利用不充分、内存瓶颈以及数据分布不均衡，从而限制了整体训练效率。尤其是在大模型训练中，这些问题会更加突出。

核心思路：PHAZE的核心思路是将硬件架构搜索和设备放置策略搜索进行联合优化。通过同时考虑硬件架构的配置（例如，张量/向量单元的数量和维度、内存大小）和设备放置策略（例如，数据并行、模型并行、流水线并行），PHAZE能够找到一个全局最优的配置，从而最大化训练吞吐量。这种联合优化能够更好地适应不同模型的特性和硬件资源，避免了局部最优解。

技术框架：PHAZE的技术框架主要包含以下几个阶段：1) 架构搜索：搜索硬件加速器的架构配置，包括张量和向量单元的数量和维度、片上和片外内存配置、微批大小以及激活重计算/存储策略。2) 算子调度：对于每个架构配置，使用整数线性规划（ILP）来寻找最优的算子调度方案，以最大化硬件利用率。3) 设备放置：将ILP的结果与动态规划相结合，确定最优的设备放置策略，包括数据并行、模型并行和流水线并行的组合。4) 性能评估：评估不同架构和设备放置策略的性能，并选择最优的配置。

关键创新：PHAZE的关键创新在于其联合优化硬件架构和设备放置策略的能力。与现有方法相比，PHAZE能够更全面地考虑硬件资源和模型特性，从而找到一个全局最优的配置。此外，PHAZE还采用了整数线性规划（ILP）来寻找最优的算子调度方案，这有助于最大化硬件利用率。

关键设计：在架构搜索方面，PHAZE考虑了多种硬件配置参数，例如张量和向量单元的数量和维度、片上和片外内存配置等。在设备放置方面，PHAZE采用了动态规划算法来寻找最优的并行策略组合。此外，PHAZE还考虑了微批大小和激活重计算/存储策略，以平衡内存占用和计算开销。整数线性规划（ILP）的目标函数是最大化硬件利用率，约束条件包括硬件资源限制和算子依赖关系。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，PHAZE 在大语言模型训练中实现了显著的性能提升。具体来说，PHAZE 在某些模型上实现了比 TPUv4 和 Spotlight 更高的吞吐量。这些结果验证了 PHAZE 联合优化硬件架构和设备放置策略的有效性，并表明其在实际应用中具有很大的潜力。

🎯 应用场景

PHAZE 的研究成果可应用于各种需要大规模分布式训练的深度学习场景，例如大语言模型、推荐系统、图像识别等。通过自动优化硬件架构和设备放置策略，PHAZE 可以显著提高训练效率，降低训练成本，并加速新模型的开发和部署。未来，该技术有望推动人工智能在更多领域的应用。

📄 摘要（原文）

Distributed execution of deep learning training involves a dynamic interplay between hardware accelerator architecture and device placement strategy. This is the first work to explore the co-optimization of determining the optimal architecture and device placement strategy through novel algorithms, improving the balance of computational resources, memory usage, and data distribution. Our architecture search leverages tensor and vector units, determining their quantity and dimensionality, and on-chip and off-chip memory configurations. It also determines the microbatch size and decides whether to recompute or stash activations, balancing the memory footprint of training and storage size. For each explored architecture configuration, we use an Integer Linear Program (ILP) to find the optimal schedule for executing operators on the accelerator. The ILP results then integrate with a dynamic programming solution to identify the most effective device placement strategy, combining data, pipeline, and tensor model parallelism across multiple accelerators. Our approach achieves higher throughput on large language models compared to the state-of-the-art TPUv4 and the Spotlight accelerator search framework. The entire source code of PHAZE is available at https://github.com/msr-fiddle/phaze.