AI框架：9大主流分布式深度学习框架简介

AI框架：9大主流分布式深度学习框架简介

随着深度学习模型规模的不断扩大，单个GPU已经难以满足训练需求。分布式深度学习框架通过将模型和数据分布在多个GPU或机器上，实现了大规模模型的高效训练。本文将介绍9种主流的分布式深度学习框架，包括它们的核心概念、主要功能和应用场景。

一、训练大模型的基础

大型深度学习模型在训练时需要大量内存来存储中间激活、权重等参数，导致某些模型只能在单个GPU上使用非常小的批大小进行训练，甚至无法在单个GPU上进行训练，使得模型训练在某些情况下非常低效和不可能。

在大规模深度学习模型训练中主要有两种范式：

• 数据并行：当模型尺寸能够被GPU内存容纳时，数据批大小会增加模型训练的难度。解决方案是让模型的不同实例在不同的GPU和不同批数据上运行。

模型的每个实例都使用相同的参数进行初始化，但在前向传递期间，不同批次的数据被发送到每个模型。收集来自每个模型实例的梯度并计算梯度更新，然后更新模型参数并将其作为更新发送到每个模型实例。

• 模型并行：当单个GPU无法容纳模型尺寸时，模型并行性变得必要，需要将模型拆分到多个GPU上进行训练。这种方法的问题是计算使用效率不高，因为在任何时间点只有一个GPU正在使用，而其他GPU处于空闲状态。

二、主流框架

1. Megatron-LM

Megatron 是由 NVIDIA 深度学习应用研究团队开发的大型 Transformer 语言模型，该模型用于研究大规模训练大型语言模型。Megatron 支持transformer模型的模型并行（张量、序列和管道）和多节点预训练，同时还支持 BERT、GPT 和 T5 等模型。

2. DeepSpeed

DeepSpeed是微软的深度学习库，已被用于训练Megatron-Turing NLG 530B和BLOOM等大型模型。DeepSpeed的创新体现在三个方面：训练、推理、压缩。DeepSpeed具备以下优势：

• 训练/推理具有数十亿或数万亿个参数的密集或稀疏模型
• 实现出色的系统吞吐量并有效扩展到数千个GPU
• 在资源受限的GPU系统上训练/推理
• 为推理实现前所未有的低延迟和高吞吐量
• 以低成本实现极致压缩，实现无与伦比的推理延迟和模型尺寸缩减

3. FairScale

FairScale（由Facebook研究）是一个用于高性能和大规模训练的PyTorch扩展库。FairScale的愿景如下：

• 可用性：用户应该能够以最小的认知代价理解和使用FairScale API。
• 模块化：用户应该能够将多个FairScale API无缝组合为训练循环的一部分。
• 性能：FairScale API在扩展和效率方面提供了最佳性能。

FairScale支持Fully Sharded Data Parallel (FSDP)，这是扩展大型神经网络训练的推荐方式。

4. ParallelFormers

Parallelformers是一个基于Megatron-LM的库。它与Huggingface库很好地集成在一起。Huggingface库中的模型可以用一行代码并行化。目前它只支持推理。

from transformers import AutoModelForCausalLM, AutoTokenizer
from parallelformers import parallelize

model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-2.7B")
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-2.7B")
parallelize(model, num_gpus=2, fp16=True, verbose='detail')

5. ColossalAI

Colossal-AI提供了一组并行组件，可以用来实现定制化的分布式/并行训练，包含以下并行化策略和增强功能：

• Data Parallelism
• Pipeline Parallelism
• 1D, 2D, 2.5D, 3D Tensor Parallelism
• Sequence Parallelism
• Zero Redundancy Optimizer (ZeRO)
• Heterogeneous Memory Management (PatrickStar)
• For Inference Energon-AI

6. Alpa

Alpa是一个用于训练和服务大规模神经网络的系统，具备如下特点：

• 自动并行化：Alpa基于数据、运算符和管道并行机制自动化地实现单设备代码在分布式集群并行化。
• 完美的表现：Alpa在分布式集群上实现了数十亿参数的训练模型的线性缩放。
• 与机器学习生态系统紧密集成：Alpa由开源、高性能和生产就绪的库（如Jax、XLA和Ray）提供支持。

7. Hivemind

Hivemind是一个在互联网上使用Pytorch进行去中心化深度学习的库。它主要服务场景是在来自不同大学、公司和志愿者的数百台计算机上训练一个大型模型。

其主要特点是：

• 没有主节点的分布式训练：分布式哈希表允许连接分散网络中的计算机。
• 容错反向传播：即使某些节点没有响应或响应时间过长，前向和后向传递也会成功。
• 分散的参数平均：迭代地聚合来自多个工作人员的更新，而无需在整个网络中同步（论文）。
• 训练任意大小的神经网络：它们的部分层通过分散的专家混合（论文）分布在参与者之间。

8. OneFlow

OneFlow是一个深度学习框架，旨在实现用户友好、可扩展和高效。使用OneFlow，很容易：

• 使用类似PyTorch的API编写模型
• 使用Global View API将模型缩放到n维并行/分布式执行
• 使用静态图编译器加速/部署模型。

9. Mesh-Tensorflow

根据github页面：Mesh TensorFlow (mtf)是一种用于分布式深度学习的语言，能够指定广泛的分布式张量计算类别。这里的“Mesh”是指处理器或计算设备的互连网络。

其他

• 引用：本文部分内容参考了知乎文章《9个主流分布式深度学习框架对比》。
• 参考：https://zhuanlan.zhihu.com/p/582498905