Alpa：自动化大规模深度学习网络分布式训练与服务的开源库

Alpa 是一个新的开源库，旨在自动化大规模深度学习网络的分布式训练和服务。它提出了一个编译器，能够结合现有的模型并行策略，并根据深度学习网络架构优化计算资源的使用。

深度学习模型的扩展与挑战

深度学习模型的推理损失与权重参数数量及训练数据量呈对数关系，这促使深度学习社区开发更大的模型。为了应对不断增长的模型规模，训练计算量也在加速扩展，大约每6个月翻一番。由于加速器内存容量有限，主要的工程挑战在于如何根据集群属性和通信原语将网络参数映射到现有的加速器设备上。这一问题在训练期间尤为突出，因为相应的梯度张量也需要存储并通过加速器内存进行交换。

模型并行的两种主要策略

1. 操作间并行（Inter-Operator Parallelism）：将计算图中的函数（如卷积层）分配到不同设备上，输入的小批量数据被分割成微批次，每个微批次在同一组函数上执行（例如 Device Placement Optimization, GPipe）。
2. 操作内并行（Intra-Operator Parallelism）：将函数参数分割，输入批次在不同部分的函数参数上运行（例如 GShard, DeepSpeed-ZeRO）。

两种策略的权衡

• 操作间并行：网络带宽使用较低，适合多节点训练。
• 操作内并行：最小化GPU空闲时间，但需要更高的数据交换，适合使用高带宽连接（如 NVIDIA NVLink 或 AMD xGMI）的场景。

混合策略的优势

当前的训练站通常由多个GPU单元组成，具有定制的GPU间连接模块。然而，在公共云上情况不同，因此结合操作间并行和操作内并行的混合策略可以显著提高资源使用效率。

深度学习库的支持

深度学习库不断发布新的API，帮助规划模型参数和输入数据的放置，例如 Tensorflow 的 DTensor 和 PyTorch 的 FSDP。尽管可以手动创建分布计划（例如 Megatron-LM），但对于非常大的网络和AutoML设计的架构，自动生成的计划和训练调度更为有益。

Alpa 的自动化解决方案

Alpa 试图自动化这些放置过程。其编译器结合了操作内并行和操作间并行，根据集群和深度学习网络结构输出运行时优化的分布式训练策略。目前，Alpa 基于 Jax 构建，Jax 提供了自动可组合的转换（如自动向量化、梯度计算、SPMD并行化和JIT编译）。

实验结果与资源

在 OSDI 22 的 论文 中，Alpa 提供的训练策略与手动放置和之前的最先进方法相比具有竞争力。更多信息可以从 Google 的 官方博客 和 官方文档 中获取。该项目还展示了 OPT-175B 模型服务器的 概念验证。