一个易用且高效的基于 PyTorch 的 MoE 模型训练系统.

简介: 一个易用且高效的基于 PyTorch 的 MoE 模型训练系统.

FastMoE 系统


https://github.com/laekov/fastmoe


简介


FastMoE 是一个易用且高效的基于 PyTorch 的 MoE 模型训练系统.


安装


依赖


启用了 CUDA 的 PyTorch 是必要的. 当前版本的 FastMoE 在 PyTorch v1.8.0 和 CUDA 10

的平台上经过了测试. 本系统从设计上也支持更旧的 PyTorch 版本.

如果需要使能 FastMoE 模型并行特性, 那么支持点对点通信的 NCCL 库 (即不旧于

2.7.5 版本) 也是必需的.


安装


FastMoE 包含一些定制的 PyTorch 算子, 包含一些 C 的组件. 用 python setup.py install

来简单地安装 FastMoE.

FastMoE 分布式模型并行特性默认是不被启用的. 如果它需要被启用,

则需要在运行上述命令时加入环境变量 USE_NCCL=1.

注意, 由于 PyTorch 框架通常仅集成了 NCCL 的运行时组件, 额外的 NCCL

开发包需要被安装在编译环境中, 而且它的版本需要与 PyTorch 的版本相对应. 推荐使用

PyTorch 官方 Docker 镜像,

因为那里的环境较为干净. 如果您希望手工配置环境, 可以在 NCCL
全部版本的下载链接

下载合适版本的 NCCL 开发包.


使用


将一个 Transformer 模型 FMoE 化


Transformer 是当前最流行的可被 MoE 化的模型. FastMoE 可以一键将一个普通的

Transformer 模型变为一个  MoE 的模型. 其使用方法如下.

例如在 Megatron-LM 中,

添加如下的代码即可将 Transformer 中的每个 MLP 层变为多个 MLP 层构成的 MoE 网络.

model = ...
from fmoe.megatron import fmoefy
model = fmoefy(model, num_experts=<number of experts per worker>)
train(model, ...)


一个更详细的在 Megatron-LM 中使用 fmoefy 函数的样例参见此处.


将 FastMoE 作为一个网络模块使用


一个使用 FastMoE 的 Transformer 模型见这个示例.

最简单的使用方式是使用 FMoE 层来代替 MLP 层.


分布式地使用 FastMoE


FastMoE 支持数据并行和模型并行.


数据并行.


在 FastMoE 的数据并行模式下,

门网络(gate)和专家网络都被复制地放置在各个运算单元上.

下图展示了一个有三个专家的两路数据并行MoE模型进行前向计算的方式.


30.png


对于数据并行, 额外的代码是不需要的. FastMoE 与 PyTorch 的 DataParallel

DistributedDataParallel 模块都可以无缝对接. 该方式唯一的问题是,

专家的数量受到单个计算单元(如GPU)的内存大小限制.


模型并行


在 FastMoE 的模型并行模式中, 门网络依然是复制地被放置在每个计算单元上的,

但是专家网络被独立地分别放置在各个计算单元上. 因此, 通过引入额外的通信操作,

FastMoE 可以允许更多的专家网络们同时被训练,

而其数量限制与计算单元的数量是正相关的.

下图展示了一个有六个专家网络的模型被两路模型并行地训练.

注意专家1-3被放置在第一个计算单元上, 而专家4-6被放置在第二个计算单元上.


31.png


FastMoE 的模型并行模式需要专门的并行策略, 而 PyTorch 和 Megatron-LM

都不支持这样的策略. 因此, 需要使用 fmoe.DistributedGroupedDataParallel

模块来代替 PyTorch 的 DDP 模块.

相关文章
|
23天前
|
算法 PyTorch 算法框架/工具
Pytorch学习笔记(九):Pytorch模型的FLOPs、模型参数量等信息输出(torchstat、thop、ptflops、torchsummary)
本文介绍了如何使用torchstat、thop、ptflops和torchsummary等工具来计算Pytorch模型的FLOPs、模型参数量等信息。
102 2
|
3月前
|
机器学习/深度学习 并行计算 PyTorch
优化技巧与策略:提高 PyTorch 模型训练效率
【8月更文第29天】在深度学习领域中,PyTorch 是一个非常流行的框架,被广泛应用于各种机器学习任务中。然而,随着模型复杂度的增加以及数据集规模的增长,如何有效地训练这些模型成为了一个重要的问题。本文将介绍一系列优化技巧和策略,帮助提高 PyTorch 模型训练的效率。
264 0
|
25天前
|
机器学习/深度学习 自然语言处理 监控
利用 PyTorch Lightning 搭建一个文本分类模型
利用 PyTorch Lightning 搭建一个文本分类模型
45 8
利用 PyTorch Lightning 搭建一个文本分类模型
|
27天前
|
机器学习/深度学习 自然语言处理 数据建模
三种Transformer模型中的注意力机制介绍及Pytorch实现:从自注意力到因果自注意力
本文深入探讨了Transformer模型中的三种关键注意力机制:自注意力、交叉注意力和因果自注意力,这些机制是GPT-4、Llama等大型语言模型的核心。文章不仅讲解了理论概念,还通过Python和PyTorch从零开始实现这些机制,帮助读者深入理解其内部工作原理。自注意力机制通过整合上下文信息增强了输入嵌入,多头注意力则通过多个并行的注意力头捕捉不同类型的依赖关系。交叉注意力则允许模型在两个不同输入序列间传递信息,适用于机器翻译和图像描述等任务。因果自注意力确保模型在生成文本时仅考虑先前的上下文,适用于解码器风格的模型。通过本文的详细解析和代码实现,读者可以全面掌握这些机制的应用潜力。
43 3
三种Transformer模型中的注意力机制介绍及Pytorch实现:从自注意力到因果自注意力
|
2月前
|
并行计算 PyTorch 算法框架/工具
基于CUDA12.1+CUDNN8.9+PYTORCH2.3.1,实现自定义数据集训练
文章介绍了如何在CUDA 12.1、CUDNN 8.9和PyTorch 2.3.1环境下实现自定义数据集的训练,包括环境配置、预览结果和核心步骤,以及遇到问题的解决方法和参考链接。
基于CUDA12.1+CUDNN8.9+PYTORCH2.3.1,实现自定义数据集训练
|
2月前
|
机器学习/深度学习 PyTorch 调度
在Pytorch中为不同层设置不同学习率来提升性能,优化深度学习模型
在深度学习中,学习率作为关键超参数对模型收敛速度和性能至关重要。传统方法采用统一学习率,但研究表明为不同层设置差异化学习率能显著提升性能。本文探讨了这一策略的理论基础及PyTorch实现方法,包括模型定义、参数分组、优化器配置及训练流程。通过示例展示了如何为ResNet18设置不同层的学习率,并介绍了渐进式解冻和层适应学习率等高级技巧,帮助研究者更好地优化模型训练。
117 4
在Pytorch中为不同层设置不同学习率来提升性能,优化深度学习模型
|
19天前
|
机器学习/深度学习 API 计算机视觉
基于Python_opencv人脸录入、识别系统(应用dlib机器学习库)(下)
基于Python_opencv人脸录入、识别系统(应用dlib机器学习库)(下)
18 2
|
19天前
|
机器学习/深度学习 存储 算法
基于Python_opencv人脸录入、识别系统(应用dlib机器学习库)(上)
基于Python_opencv人脸录入、识别系统(应用dlib机器学习库)(上)
21 1
|
2月前
|
机器学习/深度学习 监控 PyTorch
PyTorch 模型调试与故障排除指南
在深度学习领域,PyTorch 成为开发和训练神经网络的主要框架之一。本文为 PyTorch 开发者提供全面的调试指南,涵盖从基础概念到高级技术的内容。目标读者包括初学者、中级开发者和高级工程师。本文探讨常见问题及解决方案,帮助读者理解 PyTorch 的核心概念、掌握调试策略、识别性能瓶颈,并通过实际案例获得实践经验。无论是在构建简单神经网络还是复杂模型,本文都将提供宝贵的洞察和实用技巧,帮助开发者更高效地开发和优化 PyTorch 模型。
39 3
PyTorch 模型调试与故障排除指南
|
30天前
|
存储 并行计算 PyTorch
探索PyTorch:模型的定义和保存方法
探索PyTorch:模型的定义和保存方法