[论文解读] Auto-MAP: A DQN Framework for Exploring Distributed Execution Plans for DNN Workloads

Siyu Wang, Yi Rong+7 more|arXiv (Cornell University)|Jul 8, 2020

Parallel Computing and Optimization Techniques参考文献 31被引用 5

一句话总结

Auto-MAP 提出了一种基于 DQN 的强化学习框架，通过在 HLO IR 上操作，自动发现 DNN 工作负载的最优分布式执行计划，实现数据并行、模型并行和流水线并行的统一探索。它在 NLP 和 CNN 模型上两小时内即可实现最优吞吐量，优于人工和启发式方法。

ABSTRACT

The last decade has witnessed growth in the computational requirements for training deep neural networks. Current approaches (e.g., data/model parallelism, pipeline parallelism) parallelize training tasks onto multiple devices. However, these approaches always rely on specific deep learning frameworks and requires elaborate manual design, which make it difficult to maintain and share between different type of models. In this paper, we propose Auto-MAP, a framework for exploring distributed execution plans for DNN workloads, which can automatically discovering fast parallelization strategies through reinforcement learning on IR level of deep learning models. Efficient exploration remains a major challenge for reinforcement learning. We leverage DQN with task-specific pruning strategies to help efficiently explore the search space including optimized strategies. Our evaluation shows that Auto-MAP can find the optimal solution in two hours, while achieving better throughput on several NLP and convolution models.

研究动机与目标

• 为解决手动设计多样化 DNN 工作负载高效分布式训练策略的挑战。
• 在单一框架内统一探索数据并行、模型并行和流水线并行策略，避免框架特定或模型特定的解决方案。
• 通过消除低层级分布式系统配置，减轻用户负担，实现框架无关、工作负载中立的并行规划。
• 通过任务特定剪枝和基于 DQN 的强化学习，提升在大规模复杂并行空间中的搜索效率。
• 在 BERT、T5 和 ResNet 等多样化模型上，实现自动、高吞吐量且可泛化的并行化策略。

提出的方法

• Auto-MAP 在 HLO IR 上运行，这是一种由 XLA 生成的低级、硬件无关的中间表示，相比操作符或层级表示，能实现更细粒度且更具表达力的模型分析。
• 它采用深度 Q 网络（DQN）结合任务特定剪枝策略，高效探索包含数据并行、模型并行和流水线并行在内的大规模分布式执行计划搜索空间。
• 该框架基于训练吞吐量设计奖励函数，引导 DQN 代理找到最优或近似最优的并行策略。
• 它利用 XLA 编译栈，实现与多种前端（如 TensorFlow、JAX、Trax）的兼容性，确保不侵入用户模型代码。
• DQN 代理在随机分布的合成环境上预训练，随后在真实模型分布上进行微调以用于推理。
• 对于 NLP 模型，代理在预训练后能良好泛化；对于 CNN 模型，则需在特定模型分布上微调以实现准确的流水线划分。

实验结果

研究问题

• RQ1 强化学习框架能否在无需人工设计的情况下，自动发现多样化 DNN 模型的最优分布式训练策略？
• RQ2 统一框架能否在一个可扩展的搜索空间中同时探索数据并行、模型并行和流水线并行策略？
• RQ3 在分布式训练规划中，基于 HLO IR 的操作是否能实现优于基于层或操作符抽象的性能和泛化能力？
• RQ4 DQN 中的任务特定剪枝能否显著缩小搜索空间并加速收敛至最优解？
• RQ5 预训练的 DQN 代理在不同模型架构间的泛化程度如何，何时需要微调？

主要发现

• Auto-MAP 在 NLP 和卷积模型上均于两小时内发现最优分布式执行计划，显著减少了人工调优时间。
• 对于 BERT 和 T5 模型，DQN 代理在合成数据上预训练后即可成功推断最优流水线划分策略，无需微调。
• 对于更大的模型如 T5-11B，经过 10,000 轮的模型特定分布微调后，代理能准确推断最优流水线边界。
• 在多个 NLP 和 CNN 模型上，该框架的训练吞吐量优于基线启发式和人工方法。
• 使用 HLO IR 实现了框架无关、低侵入性的并行规划，适用于 TensorFlow、JAX 和 Trax。
• 将任务特定剪枝与 DQN 结合，显著提升了复杂混合并行策略空间中的搜索效率和收敛速度。