使用专家并行加速 Mixtral 8x7B 预训练

关键要点

通过使用专家并行方法，利用 Amazon SageMaker 在多 GPU 环境中高效训练大型混合专家MoE架构模型，提高模型的计算能力与效率。Mixtral 8x7B 模型具有 47 亿个参数，并使用 8 个子网络专家，实现每个输入使用大约 129 亿个参数。慧光库SMP现已支持专家并行，可以在各设备上分配不同专家，加速大型 MoE 模型的训练。了解如何准备数据集，配置模型，启动训练过程以及激活性能优化策略。

引言

混合专家MoE架构最近因其在相较于完全稠密模型的计算效率和模型容量提升而受到关注。MoE 模型通过稀疏的专家子网络处理不同的输入令牌，在训练和推理期间，相较于稠密模型，所需计算量更小。这使得在固定计算预算内训练更大规模的模型变得更加经济。

尽管 MoE 模型在计算方面的优势明显，但高效地训练和微调大型 MoE 模型仍然面临诸多挑战。在训练过程中，如果令牌在专家间分配不均，可能会导致某些专家过载，而其他专家则得不到充分利用。此外，MoE 模型对内存的需求也很高，因为所有专家参数都需要加载到内存中，即使每次输入只使用其中一部分。

在本文中，我们将展示 Amazon SageMaker 模型并行库的新特性，这些特性使得能够高效训练使用专家并行的 MoE 模型。我们将以预训练具有 47 亿参数的 Mixtral 8x7B MoE 模型为示例，来演示如何利用这些新特性。了解更多信息，请参考我们的 GitHub 仓库 和 专家并行。

专家并行

Mixtral 8x7B 模型采用稀疏 MoE 架构，包含 8 个专家子网络，每个大约有 7 亿个参数。一个被称为 router 的可训练门控网络负责决定输入令牌分配给哪个专家。借助这种架构，专家能够专注于处理输入数据的不同方面。完整的 Mixtral 8x7B 模型总共有 47 亿个参数，但每次对于任何给定的输入令牌，激活的参数仅有大约 129 亿本模型架构的两个专家，这使得计算效率相比同等规模的稠密模型得以提高。详细了解 MoE 架构请参考 应用混合专家于 LLM 架构。

SMP 添加对专家并行的支持

SMP 现在支持专家并行，这是 MoE 模型训练高性能的基础。通过在不同设备上部署 MoE 层中不同的专家子网络，在训练过程中，不同数据被路由到不同设备，每个设备处理其包含的专家的计算。通过在工作节点之间分配专家，专家并行能够解决将所有专家加载到单台设备上的高内存需求，并使得 MoE 训练在更大的集群上得以实施。下图展示了在多 GPU 集群上专家并行的工作原理。

虎跃加速器破解版下载

SMP 库使用 NVIDIA Megatron 实现专家并行，并在 PyTorch 完全分片数据并行FSDPAPI 之上运行。您可以继续使用现有的 PyTorch FSDP 训练代码，并为 MoE 模型激活 SMP 的专家并行。SMP 提供了简化的工作流程，只需指定 expertparalleldegree 参数，即可将专家平均划分到集群中的 GPU 数量上。例如，在使用 8 个 GPU 的实例时，您可以将 expertparalleldegree 设置为 2、4 或 8。我们建议从较小的数字开始，并逐渐增加，直到模型能够适应 GPU 内存。

SMP 的专家并行与分片数据并行兼容

SMP 的专家并行实现与分片数据并行兼容，使得内存使用更加高效并加快训练速度。为了理解其工作原理，考虑以下示例，假设一个 MoE 模型有八个专家N=8，在一台包含 4 个 GPU 的集群上进行训练。

SMP 的专家并行将 MoE 专家分散到 GPU 上。通过使用 expertparalleldegree 参数，您可以控制每个设备上实例化的专家数量。例如，如果设置为 2，SMP 将每个数据并行组分配一半的专家。度数值必须是您集群中 GPU 数量和模型专家数量的因子。数据通过点对点的 GPU 通信动态路由到包含所选专家的 GPU。

接下来，分片数据并行将专家和模型的非 MoE 层如注意力或路由器划分并分布到集群上，以减少模型的内存占用。hybridsharddegree 参数用于控制这一过程。例如，hybridsharddegree 设置为 2 将在集群的一半 GPU 中分片模型状态包括专家和非 MoE 层。expertparalleldegree 和 hybridsharddegree 的积不得超过集群的规模。

解决方案概述

在了解背景后，让我们深入分布式训练架构的组件。以下图解展示了解决方案架构。

在这个示例里，我们使用 SageMaker 训练作业。通过 SageMaker 训练作业，您可以简单的 API 请求启动和管理高性能实例集群。例如，您可以使用 SageMaker Estimator 简单几行代码指定在分布式系统中使用的实例类型和数量。在后面的部分中，我们将使用包含两个 mlp4d24xlarge 实例的集群训练模型，具体参数会在 Estimator 中指定。要了解 SageMaker 训练作业的更多信息，请参见 使用 Amazon SageMaker 训练模型。

在本篇中，我们借助 SMP 库，高效地将工作负载分配到集群中，利用混合分片数据并行和专家并行。在这些实现之外，SMP 还提供了许多其他性能提升和内存节省的技术，例如：

混合精度训练及对稠密 Llama 模型的 fp8 支持加速分布式训练，充分发挥 P5 实例的性能提升与分片数据并行组合的张量并行延迟参数初始化激活检查点通过清理某些层的激活并在反向传递期间重新计算以减少内存使用的技术

有关最新更新，请参考 SageMaker 模型并行库 v2。

除了 SMP，此示例还使用 SageMaker 分布式数据并行库SMDDP。随着工作负载的加大及实例添加到集群中，实例之间的通信开销也会增加，这可能会导致整体计算性能和训练效率下降。这时，SMDDP 显得尤为重要。SMDDP 包含的优化通信集合如 AllGather专为 AWS 网络基础设施设计。因此，在 SageMaker 上训练时，SMDDP 的性能可能超越其他更通用的通信库，如 NCCL。

SMP 和 SMDDP 库结合使用，可以加速大型分布式训练工作负载高达 20。此外，这些库与标准 PyTorch API 及功能兼容，因此便于将现有的 PyTorch FSDP 训练脚本适应于 SageMaker 训练平台，并利用 SMP 和 SMDDP 带来的性能提升。有关更多信息，请参见 SageMaker 模型并行库 v2 和 使用 SageMaker 分布式数据并行库进行分布式训练。

在接下来的部分，我们将展示如何在 P4 实例上使用 SMP 和 SMDDP 加速 Hugging Face Transformers Mixtral 87B 模型的分布式训练。

准备工作

在运行 Mixtral 笔记本之前，您需要完成一些前期准备工作。

首先，请确保您已经 创建了 Hugging Face 访问令牌，以便下载稍后需要使用的 Hugging Face 分词器。获得访问令牌后，您还需要向 SageMaker 请求一些配额增加。您需要请求至少 2 个 P4d 实例，最多可请求 8 个 P4d 实例根据训练时间和训练成本的权衡。

在服务配额控制台，请请求以下 SageMaker 配额：

P4 实例 (mlp4d24xlarge) 用于训练作业使用：28

配额增加的批准可能需要最多 24 小时。

现在您已准备好开始 Mixtral 模型的预训练过程，接下来将进行数据集准备。

准备数据集

我们将开始本 教程 的数据集准备工作。此过程将涵盖加载 GLUE/SST2 数据集、对数据集进行分词和分块以及为 SageMaker 在 Amazon 简单存储服务Amazon S3上训练配置数据通道。请按照以下步骤进行：

首先需要加载 GLUE/SST220数据集，并将其划分为训练数据集和验证数据集： python hyperparameters = { cachedir tmp datasetconfigname sst2 datasetname glue dotrain True doeval True }

rawdatasets = loaddataset( hyperparameters[datasetname] hyperparameters[datasetconfigname] )

del rawdatasets[validation]

if validation not in rawdatasetskeys() validationpercentage = 10

   rawdatasets[validation] = loaddataset(       hyperparameters[datasetname]       hyperparameters[datasetconfigname]       split=ftrain[{validationpercentage}]       cachedir=hyperparameters[cachedir]   )   rawdatasets[train] = loaddataset(       hyperparameters[datasetname]       hyperparameters[datasetconfigname]       split=ftrain[{validationpercentage}]       cachedir=hyperparameters[cachedir]   )

从 Hugging Face Transformers 库加载 Mixtral8x7B 分词器： python tokenizer = AutoTokenizerfrompretrained(mistralai/Mixtral8x7Bv01 tokenizerkwargs)

接下来，您需要定义两个工具函数：tokenizefunction() 和 grouptexts()。tokenizefunction() 用于在文本数据上运行分词器。grouptexts() 函数则将数据集中所有文本连接起来，并生成与模型输入长度2048相对应的块大小。通过将文本数据切割成较小的部分，确保模型在训练期间能够处理整个数据集，即便某些文本示例超出输入长度2048。

使用以下代码定义这些函数： python def tokenizefunction(examples) output = tokenizer(examples[textcolumnname]) return output

def grouptexts(examples) # 合并所有文本 concatenatedexamples = {k sum(examples[k] []) for k in exampleskeys()} totallength = len(concatenatedexamples[list(exampleskeys())[0]])

   if totallength gt= blocksize       totallength = (totallength // blocksize)  blocksize       # 按最大长度切分       result = {           k [t[i  i  blocksize] for i in range(0 totallength blocksize)]           for k t in concatenatedexamplesitems()       }   result[labels] = result[inputids]copy()   return result

调用之前的工具函数对数据集进行分词和分块，以便模型使用： python tokenizeddatasets = rawdatasetsmap(tokenizefunction batched=True numproc=1 removecolumns=columnnames) lmdatasets = tokenizeddatasetsmap(grouptexts batched=True)

准备 SageMaker 训练的数据集，并将其保存在 JSON 文件中，同时构建将上传这些文件的 S3 路径： python traindataset = lmdatasets[train] traindatasettojson(/trainingjson) trainingdatasetlocation = fs3//{defaultbucket}/dataset/train/

evaldataset = lmdatasets[validation] evaldatasettojson(/validationjson) validationdatasetlocation = fs3//{defaultbucket}/dataset/validation/

最后，通过创建来自提供的 S3 桶路径的 TrainingInput 对象，设置 SageMaker 训练的数据通道： python train = sagemakerinputsTrainingInput( s3trainbucket distribution=FullyReplicated s3datatype=S3Prefix ) datachannels = {train train}

test = sagemakerinputsTrainingInput( s3testbucket distribution=FullyReplicated s3datatype=S3Prefix ) datachannels[test] = test

现在您已准备好在数据集上运行预训练或微调。

在 SMP 上用专家并行预训练 Mixtral 8x7B

要预训练 Mixtral 8x7B 模型，请完成以下步骤：

初始化 脚本 ，调用 torchsagemakerinit() 以激活 SMP 库： python import torchsagemaker as tsm tsminit()

从 torchsagemakertransform API 导入 MoEConfig 类。我们使用 MoEConfig 类使模型能够利用 SMP 实现的 MoE： python from torchsagemakermoemoeconfig import MoEConfig

为 Mixtral 8x7B 模型创建配置。它将传递给 Hugging Face Transformers 库中的 AutoModelForCausalLMfromconfig(modelconfig attnimplementation=flashattention2) 以随机权重初始化模型。如果要微调，可以提供预训练权重的路径。 python modelconfig = MixtralConfig( vocabsize=argsvocabsize # 32000 hiddensize=argshiddenwidth # 4096 intermediatesize=argsintermediatesize # 14336 numhiddenlayers=argsnumlayers # 32 numattentionheads=argsnumheads # 32 numkeyvalueheads=argsnumkeyvalueheads # 8 hiddenact=silu maxpositionembeddings=argsmaxcontextwidth # 4096 32 initializerrange=argsinitializerrange # 002 rmsnormeps=1e5 usecache=False padtokenid=None bostokenid=1 eostokenid=2 tiewordembeddings=False ropetheta=1e6 slidingwindow=argsslidingwindow # None attentiondropout=00 numexpertspertok=argsnumexpertspertok # 2 numlocalexperts=argsnumlocalexperts # 8 outputrouterlogits=False routerauxlosscoef=