[论文解读] Automap: Towards Ergonomic Automated Parallelism for ML Models

Michael Schaarschmidt, Dominik Grewe+9 more|arXiv (Cornell University)|Dec 6, 2021

Parallel Computing and Optimization Techniques被引用 5

一句话总结

Automap 是一个原型化的自动分区器，通过结合归纳偏置、搜索与学习，将专家级的分片策略（如 Megatron）恢复为 ML 模型中易于使用的、与编译器集成的 SPMD 风格并行计算。它利用平台无关的中间表示（IR）和增量重写，在几分钟内实现接近专家水平的性能，显著减少人工工作量，同时可扩展至复杂模型，且用户输入极少。

ABSTRACT

The rapid rise in demand for training large neural network architectures has brought into focus the need for partitioning strategies, for example by using data, model, or pipeline parallelism. Implementing these methods is increasingly supported through program primitives, but identifying efficient partitioning strategies requires expensive experimentation and expertise. We present the prototype of an automated partitioner that seamlessly integrates into existing compilers and existing user workflows. Our partitioner enables SPMD-style parallelism that encompasses data parallelism and parameter/activation sharding. Through a combination of inductive tactics and search in a platform-independent partitioning IR, automap can recover expert partitioning strategies such as Megatron sharding for transformer layers.

研究动机与目标

• 在无需专家知识或手动重写的情况下，自动化发现适用于大型 ML 模型的高效数据并行、模型并行和流水线并行策略。
• 无缝集成至现有编译器栈（如 XLA）和 JAX 中的用户工作流，避免为特定硬件重新实现。
• 通过支持快速、交互式的分区且用户输入极少，减少实验的时间和成本，促进研究的快速迭代。
• 通过利用归纳偏置和模型结构（如重复模块）来支持具有数万个操作的复杂、非结构化模型。
• 通过支持无需架构重写即可在多样化的加速器集群（包括旧硬件）上高效分区，实现实际部署。

提出的方法

• 该系统使用一种新颖的、与平台无关的分区中间表示（IR，即 MLIR 中的一个方言），建立在 XLA HLO 之上，将分片和复制决策表达为重写规则。
• 采用混合方法，结合蒙特卡洛树搜索（MCTS）与学习模型以引导搜索，从而减少找到高质量策略所需的尝试次数。
• 通过模仿专家分片模式（如 Megatron）来嵌入归纳偏置，以约束搜索空间并提升收敛速度。
• 支持增量重写，尽可能在计算图中传播分片决策，特别是在共享常量和层组中。
• 用户可通过命名作用域（如 'attention-block'）提供高层级提示，以分组重复层，显著提升搜索效率和可靠性。
• 该原型在 JAX 中使用交互网络生成的合成数据进行训练，通过 Optax 进行优化，并利用图神经网络进行分区决策推理。

实验结果

研究问题

• RQ1 在仅需极少用户输入的前提下，能否实现与 Megatron 等专家优化策略相当性能的自动化分区？
• RQ2 在具有数万个操作的大型非结构化模型中，如何提升搜索效率？
• RQ3 从专家实践中提取的归纳偏置在多大程度上能缩小搜索空间并加速收敛？
• RQ4 编译器提示（如命名层组）能否替代对共享常量的脆弱依赖传播，以提升深层网络的鲁棒性？
• RQ5 在真实研究工作流中，学习与搜索的结合如何影响求解时间？

主要发现

• 系统在不到 500 次搜索尝试（约数分钟）内发现了 Megatron 风格的分片策略，在 TPU v3 上实现了接近 Megatron 的性能，仅带来轻微的运行时开销。
• 通过使用学习模型过滤 MCTS 搜索，显著减少了找到近似最优解所需的尝试次数，证明了数据驱动搜索的价值。
• 在使用编译器提示（如命名层组）的情况下，即使不依赖层间共享常量的传播，也能在少量尝试内可靠地找到 Megatron 分片策略。
• 在未分组或未传播的情况下，24 层 Transformer 的 Megatron 分片策略未能被发现，凸显了结构化提示在可扩展性中的关键作用。
• 该原型成功分区了非标准模型（如 GraphNets），并发现了实用的策略（如输入边分片），适用于更大规模的实验。
• 该方法可扩展至包含 5 万至 10 万个操作的模型，搜索时间适合交互式研究工作流（数分钟，而非数小时）。