清华发布SmartMoE:一键实现高性能MoE稀疏大模型分布式训练

简介: 清华发布SmartMoE:一键实现高性能MoE稀疏大模型分布式训练


PACMAN 实验室探索大模型分布式训练系统再出新成果。

2023 年 7 月,清华大学计算机系 PACMAN 实验室发布稀疏大模型训练系统 SmartMoE,支持用户一键实现 MoE 模型分布式训练,通过自动搜索复杂并行策略,达到开源 MoE 训练系统领先性能。同时,PACMAN 实验室在国际顶级系统会议 USENIX ATC’23 发表长文,作者包括博士生翟明书、何家傲等,通讯作者为翟季冬教授。PACMAN 实验室在机器学习系统领域持续深入研究,SmartMoE 是继 FastMoE, FasterMoE 和 “八卦炉” 后在大模型分布式训练系统上的又一次探索。欲了解更多相关成果可查看翟季冬教授首页:https://pacman.cs.tsinghua.edu.cn/~zjd
Mixture-of-Experts (MoE) 是一种模型稀疏化技术,因其高效扩展大模型参数量的特性而备受研究者关注。为了提高 MoE 模型的易用性、优化 MoE 模型训练性能,PACMAN 实验室在 MoE 大模型训练系统上进行了系统深入的研究。2021 年初,开源发布了 FastMoE 系统,它是第一个基于 PyTorch 的 MoE 分布式训练系统开源实现,在业界产生了较大的影响力。进一步,为了解决专家并行的稀疏、动态计算模式带来的严重性能问题,FasterMoE 系统地分析、优化了专家并行策略。FasterMoE 中设计的「影子专家」技术显著缓解了负载不均问题、通信 - 计算协同调度算法有效隐藏了 all-to-all 通信的高延迟。FasterMoE 成果发表在 PPoPP’22 国际会议。



MoE 模型遇到的难题
不同于稠密模型直接通过增大模型尺寸实现扩展,如图一所示,MoE 技术通过将一个小模型转变为多个稀疏激活的小模型实现参数扩展。由于各个专家在训练时稀疏激活,MoE 模型得以在不增加每轮迭代计算量的前提下增加模型参数量;从而有望在相同训练时间内获得更强的模型能力。
图一:通过 MoE 技术扩展模型规模

为了实现 MoE 大模型的分布式训练,业界提出了专家并行(Expert Parallelism)技术。如图二所示,各个专家被分布式地存储在不同节点上,在训练过程中通过 all-to-all 通信将训练数据发送至对应专家所在节点。专家并行相较于数据并行(Data Parallelism)有更小的内存开销,因为专家参数无冗余存储。可以认为专家并行是一种针对 MoE 结构的模型并行(Model Parallelism)。
图二:专家并行示意图

然而,使用朴素的专家并行技术训练 MoE 模型有严重的性能问题,其根因是 MoE 模型的稀疏激活模式。它会导致节点间产生大量不规则 all-to-all 通信增加延迟、计算量负载不均造成硬件利用率低。如图三所示的真实模型训练过程中的专家选择分布,可以观察到专家间显著的负载不均现象,且分布随训练进行动态变化。
图三:真实训练中的专家选择分布

随着学界对各并行策略的深入研究,使用各并行策略的复杂组合(称为混合并行)进行大模型训练成为必要模式。混合并行的策略调优过程十分复杂,为了提高可用性,学界提出了自动并行算法自动搜索、调优混合并行策略。然而,现有混合并行、自动并行系统无法高效处理 MoE 大模型,他们缺少对 MoE 模型训练稀疏激活、计算负载不均且动态变化的特征的针对性设计。

SmartMoE 带来解决方案
为了实现 MoE 模型的高效训练,SmartMoE 系统对 MoE 模型的分布式训练策略进行了全面的支持。对于常用的四种并行策略(数据并行、流水线并行、模型并行和专家并行),SmartMoE 系统做出了全面的支持,允许用户对它们任意组合;在论文投稿时(2023 年 1 月),尚未有其他系统能做到这一点(如图四所示)。
为了处理 MoE 的动态计算负载,SmartMoE 独特设计了专家放置(Expert Placement)策略,在经典并行策略组合的基础上,实现了动态负载均衡。如图五所示,MoE 模型不同的计算负载(workload)会造成不同专家的过载;使用不同的专家放置顺序,能在特定负载下实现节点间负载均衡。
图四:开源分布式系统对各并行策略的支持情况对比
图五:不同 MoE 训练负载需要不同专家放置策略
为了提高 MoE 模型复杂混合并行策略的易用性,SmartMoE 设计了一套轻量级且有效的两阶段自动并行算法。现有自动并行系统只能在训练开始前进行策略搜索,无法根据负载情况动态调整策略。简单的将现有自动并行搜索算法在训练过程中周期性使用亦不可行,因为训练过程中的并行策略搜索和调整对延迟要求很高,现有算法的开销过大。
SmartMoE 独创性地将自动并行搜索过程分为两阶段:

  • 训练开始前,使用经典算法搜索,获得一个较小的候选策略集合
  • 训练过程中,根据当前负载,在候选策略集合中动态调整,由于候选策略集合大小有限,此过程的开销可以得到控制。


最终,SmartMoE 实现了轻量级且有效的自动并行,达到了业界领先的性能。
在性能测试中,SmartMoE 在不同模型结构、集群环境和规模下均有优异的表现。例如,在 GPT-MoE 模型的训练性能测试中,相较于 FasterMoE,SmartMoE 有最高 1.88x 的加速比。值得注意的,在对每一轮迭代的性能观察中发现,动态的并行策略调整是必要的,且需要使用合适的调整频率,如图六所示。更多实验细节请参考论文原文。
图六:MoE 模型在不同迭代的运行时间。”dyn.X” 表示每 X 轮进行一次策略调整。
图七:SmartMoE 在 GPT-MoE 模型端到端训练中的性能提升
结语
SmartMoE 现已开源,开发者维护活跃,且仍在持续优化迭代,助力 MoE 大模型的发展。这是 PACMAN 实验室继 FastMoE,[PPoPP’22] FasterMoE,[PPoPP’22] BaGuaLu 后在大模型分布式训练系统上的又一次探索。

相关文章
|
8天前
|
存储 人工智能 PyTorch
基于PyTorch/XLA的高效分布式训练框架
基于PyTorch/XLA的高效分布式训练框架
48 2
|
8天前
|
人工智能 算法 PyTorch
TorchAcc:基于 TorchXLA 的分布式训练框架
阿里云研究员、阿里云人工智能平台 PAI 技术负责人--林伟在GTC 2024 大会 China AI Day 线上中文演讲专场上介绍了TorchAcc,这是一个基于 PyTorch/XLA 的大模型分布式训练框架。
|
8天前
|
算法
基于一致性理论的微电网分布式控制策略仿真模型【自适应虚拟阻抗】【simulink仿真】
基于一致性理论的微电网分布式控制策略仿真模型【自适应虚拟阻抗】【simulink仿真】
|
8天前
|
安全
考虑极端天气线路脆弱性的配电网分布式电源和储能优化配置模型
考虑极端天气线路脆弱性的配电网分布式电源和储能优化配置模型
|
8天前
|
存储 NoSQL 关系型数据库
聊一聊分布式锁的设计模型
本文介绍了分布式锁的设计模型、运行原理以及具体用法,作者也在文中体现了自己的关于分布式锁的思考以及具体实践。
|
8天前
|
机器学习/深度学习 人工智能 TensorFlow
人工智能平台PAI产品使用合集之如何在CPU服务器上使用PAIEasyRec进行分布式训练
阿里云人工智能平台PAI是一个功能强大、易于使用的AI开发平台,旨在降低AI开发门槛,加速创新,助力企业和开发者高效构建、部署和管理人工智能应用。其中包含了一系列相互协同的产品与服务,共同构成一个完整的人工智能开发与应用生态系统。以下是对PAI产品使用合集的概述,涵盖数据处理、模型开发、训练加速、模型部署及管理等多个环节。
|
8天前
|
存储 负载均衡 Cloud Native
【专栏】Minio是一款高性能分布式对象存储服务,以其易用性和可扩展性著称
【4月更文挑战第28天】Minio是一款高性能分布式对象存储服务,以其易用性和可扩展性著称,适用于存储多媒体内容。通过组建Minio集群,可实现高可用性、高性能、可扩展性和数据保护。搭建集群包括安装Minio、配置集群参数、启动节点、验证集群状态、设置访问权限及可选的数据迁移步骤。Minio集群是实现可靠且高性能存储解决方案的理想选择,适合各种应用场景。
|
8天前
|
机器学习/深度学习 负载均衡 PyTorch
PyTorch分布式训练:加速大规模数据集的处理
【4月更文挑战第18天】PyTorch分布式训练加速大规模数据集处理,通过数据并行和模型并行提升训练效率。`torch.distributed`提供底层IPC与同步,适合定制化需求;`DistributedDataParallel`则简化并行过程。实际应用注意数据划分、通信开销、负载均衡及错误处理。借助PyTorch分布式工具,可高效应对深度学习的计算挑战,未来潜力无限。
|
8天前
|
机器学习/深度学习 运维 监控
TensorFlow分布式训练:加速深度学习模型训练
【4月更文挑战第17天】TensorFlow分布式训练加速深度学习模型训练,通过数据并行和模型并行利用多机器资源,减少训练时间。优化策略包括配置计算资源、优化数据划分和减少通信开销。实际应用需关注调试监控、系统稳定性和容错性,以应对分布式训练挑战。
|
8天前
|
分布式计算 监控 Hadoop
Hadoop【基础知识 02】【分布式计算框架MapReduce核心概念+编程模型+combiner&partitioner+词频统计案例解析与进阶+作业的生命周期】(图片来源于网络)
【4月更文挑战第3天】Hadoop【基础知识 02】【分布式计算框架MapReduce核心概念+编程模型+combiner&partitioner+词频统计案例解析与进阶+作业的生命周期】(图片来源于网络)
61 0