一、DHEN [2022]

《DHEN: A Deep and Hierarchical Ensemble Network for Large-Scale Click-Through Rate Prediction》

  1. 学习特征交互(feature interactions)对在线广告服务的模型性能至关重要。因此,研究者们致力于设计有效的架构来学习feature interactions。然而我们发现,即使声称捕获的 interactions 阶数相同,这些设计的实际性能也会因数据集而异。这表明不同的设计可能具有不同的优势,并且它们捕获的 interactions 存在非重叠的信息。基于这一观察,我们提出了 DHEN,一种深度的分层的集成(ensemble )架构,该架构可以利用异构的 interaction 模块的优势,并学习不同阶数下的 a hierarchy of the interactions。为了克服 DHENdeepermulti-layer structure 在训练中带来的挑战,我们提出了一种新的 co-designed training system,该系统可以进一步提高 DHEN 的训练效率。在 CTR prediction 任务的大规模数据集上进行的 DHEN 实验表明,prediction 的归一化熵(Normalized Entropy: NE)提高了 0.27%,训练吞吐量比 SOTA 的基线高 1.2 倍,这些证明了其在实际应用中的有效性。

  2. 在线广告已迅速发展成为一项价值数十亿美元的业务:仅在美国,其 2021 财年的收入就达到 284.3B 美元,相比 2020 财年增长了25%《Brad Adgate. Agencies agree; 2021 was a record year for ad spending, with more growth expected in 2022, Dec 2021》)。预测用户点击特定广告的概率(即 CTR prediction)在在线广告中起着重要作用。这是因为 CTR prediction 的性能会影响 business providers 的用户满意度(user satisfaction)。

    由于 CTR prediction 的重要性,学术界和工业界都投入了大量精力来提高 prediction model 的性能。早期,逻辑回归(logistic regression: LR)被用于建模特征与标签之间的关系。不幸的是,其线性假设使其无法捕获非线性的 input-output 关系,限制了它在复杂场景中的能力和适用性。这一限制在 《Practical lessons from predicting clicks on ads at facebook》中通过使用决策树(decision tree: DT)得到了部分解决,其中 input features 首先通过 DT进行非线性变换。然而,LRDT 的成功都需要大量的特征工程工作;而在现实中,由于存在大量的原始特征以及特征交叉(interactions),这一过程需要大量资源。为了应对这一挑战,《Factorization machines》 提出了因子分解机(factorization machine: FM),它可以通过 latent embeddings 的内积来捕获二阶的 feature interaction。同时,FM 预定义的浅层结构(pre-defined shallow structure)限制了其表达能力。FM 关于更具表达性的扩展,如 HOFMFFMAFM,则面临 over-fitting 以及计算成本过高的问题。上述所有限制都可以通过基于深度学习的模型得到妥善处理:deep hidden layers 和非线性激活函数的使用,使得能够以端到端的方式捕获非线性的高阶 feature interactions,缓解了 ranking engineers 繁琐的手动特征工程负担、以及浅层模型表达能力不足的问题。

    2016 年以来,各种深度模型已被 business providers 用于其 ads ranking services,包括 Wide&DeepGoogle Play)、DeepFMHuawei AppGallery)、DINTaobao)和 FiBiNETWeibo)。这些深度模型大多由两个主要部分组成:feature embedding learningfeature interaction modeling

    • feature embeddings 是通过将 categorical features 映射到 embedding vectors 来学习的。

    • feature interactions 则是通过利用函数对 embeddings 之间的关系进行建模来学习的。多项研究表明,interaction 部分的更好设计可以显著提高实际应用中的 prediction accuracy。这激发了从捕获低阶 interactions 向高阶 interactions 发展的各种研究。

    尽管先前的一些研究声称其 feature interaction 模块的设计可以捕获高阶 interactions,但我们注意到,即使它们声称捕获相同阶数的 interaction,其实际性能排名也会因数据集而异。这表明,旨在捕获相同阶数 interaction 的不同 interaction 模块在不同数据集上具有不同的优势,而且其根本原因可能是它们捕获的信息并不重叠。同时,如这些研究的实验部分所示,通过学习高阶 interaction(通常通过堆叠更多 interaction layers 来实现)有时会产生负面影响(例如 DCN 中的 Figure 3xDeepFM 中的 Figure 7(a)InterHAt 中的 Figure 4xDeepInt 中的 Table 2 以及 GIN 中的 Figure 3(a) 等)。这与理论相悖,因为根据理论,捕获高阶 interaction 应该导致更好、或至少中性的性能。我们假设这个问题源于使用同质的 interaction 模块,这限制了可以捕获的 interactions 类型。这一观察表明模型中存在异构 interaction modules 的重要性。基于此,我们提出了DHEN :一种具有分层结构的深度分层集成网络(Deep and Hierarchical Ensemble Network )。

    DHEN layer 中,存在一组异构的 interaction modulesensemble components

    • 异构 interaction modules 的集合可以补充不同 interaction modules 能够捕获的非重叠信息。

    • ensemble component 捕获异构模块之间的相关性。

    因此,通过递归堆叠 DHEN layers,模型可以学习不同阶数 interactionshierarchy,并捕获异构模块的相关性,我们通过实验发现这对获得更好的性能起着重要作用。本文的主要贡献可总结如下:

    • 我们基于不同 interaction modules 在不同数据集上具有不同优势的观察,设计了一种名为 Deep and Hierarchical Ensemble Network: DHEN 的新型架构。通过递归地堆叠 interaction layersensemble layersDHEN 可以学习不同阶数 interactionshierarchy 。其中,不同阶数 interactionsheterogeneous modules 所学习而来。

      ensemble layer 其实就是 concat/sum/weighted_sum 等等。

    • 与以往的 CTR prediction 模型相比,DHEN 更深的、multi-layer 的结构增加了训练的复杂性,给实际训练带来了挑战。我们提出了一系列机制来提高 DHEN training 的性能,包括一种名为 Hybrid Sharded Data Parallel 的新分布式训练范式。该范式的吞吐量比 SOTAfully sharded data parallel1.2 倍,支持对大型 DHEN 模型进行高效训练。

    • CTR prediction 任务的大规模数据集上对 DHEN 进行了全面评估。评估结果表明,与 SOTAAdvancedDLRM baseline 相比,归一化熵(Normalized Entropy: NE )提高了 0.27%

1.1 架构

  1. 大多数跨 multiple prediction tasksSOTA 的高性能模型架构都采用深度堆叠(deep stacking )结构(例如 ResNetTransformersMetaformer)。在这里,深度堆叠结构通常由包含相同 interaction module 的重复 blocks 组成。例如,Transformer 使用 self-attention blocks 进行堆叠,而 ResNet 使用 convolution blocks。每个 block 将前一个 block 的输出或原始 embedding tokens 作为输入。在结构上,DHEN 遵循这种整体堆叠策略,同时构建了一个新颖的 hierarchical ensemble framework 来捕获多个 interaction modules 的相关性。Figure 1 展示了一个通用的 DHEN building block ,其中包含多个 Interaction modulesensemble。请注意,原始 numericaldense)特征可以作为任何 modules 的输入,以便在每一层中进行 ensembling

    基本思想类似于 MMOE,每个 interaction module 就是一个 expert 。作者对比了 DHENMMOE,发现 DHEN 效果更好。

1.1.1 Feature Processing Layer

  1. CTR prediction 任务中,特征输入通常包含离散的 categorical termssparse features)和 numerical valuesdense features)。在这项工作中,我们使用与 DLRM《Deep learning recommendation model for personalization and recommendation systems》)相同的 feature processing layer,如 Figure 3 所示。

    • sparse lookup tablescategorical terms 映射到 a list of numerical embeddings。具体来说,每个 categorical term 被分配一个可训练的 d 维向量作为其 feature representation

    • 另一方面,numerical values 通过 dense layers 进行处理。dense layers 由几个多层感知器(Multi-layer Perception: MLP)组成,从中计算出一个 d 维向量从而作为 output

    sparse lookup table outputdense layer output 拼接后,feature processing layerfinal output X0Rd×m 可以表示为 X0=(x01,x02,,x0m),其中 moutput embeddings 的数量,dembedding 维度。

1.1.2 Hierarchical Ensemble

  1. 我们提出了一种新颖的 hierarchical ensemble framework ,该框架包含多种类型的 interaction modules 、以及它们的相关性。从概念上讲,deep hierarchical ensemble network 可以描述为一个深度的、全连接的 interaction module network,类似于具有全连接神经元的 deep neural network

  2. hierarchical ensemble mechanism 为多层堆叠的多个 interaction modulesensemble 提供了一个框架。每层的 input 表示为 a list of embeddings,记为 XnRd×m,这里的 n 表示第 nstacked layers 。请注意,第一层的 input 是来自 feature processing layerlist of embeddings X0 。形式上,第 n 层的输出为:

    Y=Norm(Ensemblei=1k(Interactioni(Xn))+ShortCut(Xn))ShortCut(Xn)={Xn,if len(Xn)==len(Y)WnXn,else

    其中:

    • Norm() 表示 normalization 方法,如 Layer Normalization

    • Ensemblei=1k() 表示对 Interaction1,Interaction2,,Interactionk 的集成方法,可以是 concatenationsumweighed sum 等等。

      什么 ensemble 的效果最好?论文未说明。

    • ShortCut() 函数兼具残差和维度匹配的作用。如果 ensemble output 的维度与 last layer input Xn 的维度不匹配,将应用线性投影 WnRlen(Xn)×len(Y) 来匹配维度。

    • 最后,ensemble resultshortcut 通过 element-wise sum 进行组合,output Y 将作为下一层的输入 Xn+1

  3. hierarchical ensemble 的主要目标是捕获 interaction modules 的相关性。Figure 2 (left) 展示了一个两层、两模块的 hierarchical ensemble component。与传统的仅捕获一种类型高阶 interactionstacking model structures 不同,hierarchical ensembles 可以捕获 a mixture of high-order interactions 。如Figure 2 (right) 所示,DHEN 的设计通过让每个模块消耗 outputs of various interaction modules,例如 Interaction1(Interaction1)Interaction1(Interaction2)Interaction2(Interaction1) 、以及 Interaction2(Interaction2) ,享有捕获特征之间复杂高阶 interactions 的优势。因此,这种 mixture of interaction modules 可以利用多种 feature interaction 类型来实现更好的 model prediction accuracy

1.1.3 Interaction Modules

  1. 我们在模型中应用了五种类型的 interaction 模块:AdvancedDLRMself-attentionLinearDeep Cross Net 、以及 Convolution。实际上,DHEN 中可以包含的 interaction 模块不限于上述五种。此外,请注意,如果一个 interaction 模块输出单个张量 vRh×1(例如来自 MLP 的张量输出),我们将应用 WmR(dl)×h 将其映射到 embeddings 的列表,其中 dembedding 维度、l 表示 interaction 模块 output embeddings 的数量。下面我们简要介绍 DHEN 中使用的这些 interaction 模块。

  2. AdvancedDLRM:我们使用 DLRM 风格的 interaction 模块来捕获 feature interactions ,在后续实验中称其为AdvancedDLRM 。给定 embeddings 的输入 XnAdvancedDLRM 输出 embeddings 的一个新的列表 URl×d

    v=AdvancedDLRM(Xn)Rh×1U=reshape(Wmv)Rl×d

    其中:

    • WmR(dl)×h 为投影矩阵。

    • reshapedl 的向量转换成 lembeddings 组成的列表,embedding 维度为 d

  3. Self-attention:我们考虑自注意力作为本文的一种 interaction 模块。self-attention 被广泛应用于 Transformer 网络。它之前也被用于 CTR prediction 任务(《Interpretable click-through rate prediction through hierarchical attention》)。典型的 Transformer 包含多个堆叠的 encoder/decoder layers,其核心是 self-attention 机制。在本文中,给定 embeddings 的输入 Xn,我们应用 Transformer encoder layer

    V=TransformerEncoderLayer(Xn)Rm×dU=WVRl×d

    其中: WRl×m 用于匹配和统一所有 interaction 模块的输出维度。

  4. Convolution:卷积层广泛应用于计算机视觉任务,也被应用于自然语言处理(NLP )和 CTR 任务。在本文中,我们将卷积作为 interaction 模块之一。给定 embeddings 的输入 Xn,通过以下方式获得 embeddings 的列表 U

    V=Conv2d(Xn)Rm×dU=WVRl×d

    其中: WRl×m 也是用于匹配和统一所有 interaction 模块的输出维度。

  5. Linear:线性层是捕获原始 feature embeddings 中信息的最直接模块之一。在本文中,我们使用线性层作为 interaction 模块之一来压缩每层的信息。给定 embeddings 的输入 Xn,通过以下方式获得 embeddings 的列表 U

    U=WXnRl×d

    其中: WRl×m 也是用于匹配和统一所有 interaction 模块的输出维度。

  6. Deep Cross Net: DCNDCNCTR prediction 任务中广泛使用的 feature interaction 模块。它引入了一个 cross network,能够高效地学习特定有界阶数的 feature interactions。在本文中,我们将 DCN 模块作为每层的 interaction 模块之一。给定 embeddings 的输入 Xn,通过以下方式获得 embeddings 的列表 U

    xn+1=x0xnwl+blU=reshape(Wmxn+1)Rl×d

    其中:

    • xlRd 为第 l 层交叉层的输出。

    • wl,blRd 为第 l 层交叉层的参数。

    由于我们已经使用 skip connection 在堆叠层之间传递信息,因此上述公式中省略了原始论文中的 skip connection 过程。

    原始论文中:xn+1=x0xnwl+bl+xn

1.2 训练系统

  1. 直观来看,stacked hierarchical ensemble layers 的深度有助于提升 DHEN 的表达能力,但这也给 DHEN 的实际训练带来了挑战。本节将重点介绍我们如何在集群中实现 DHEN 的高效的、可扩展的训练。

    看起来训练难度很大。

1.2.1 训练策略

  1. 每个 DHEN 训练样本包含 categorical featuresnumerical features,所有特征在被 DHEN layers 处理之前必须转换为 dense representations。这一过程的直接挑战在于,DHEN 的规模和复杂性完全超出了单个标准 datacenter server 的处理能力:这类服务器通常配备 8 GPUs,总高带宽内存(high bandwidth memory: HBM )仅几百 GB ,计算能力最高达数十PetaFLOPs 。因此,与训练单个 DHEN 样本所需的参数数量(可达 trillions)和计算量(可达 giga flops )相比,单台服务器的能力严重不足。

  2. 为了实现 DHEN 的高效训练,我们利用了 《Software-hardware co-design for fast and scalable training of deep learning recommendation models》 中详细介绍的 ZionEX 全同步训练系统。概括来说,ZionEX 系统将 16 台主机组成一个“超级节点”(称为 pod ),其中包含 128A100 GPU(每台主机 8 块),总 HBM 容量达 5TBBF16 计算能力为 40PF/s。在主机内部,每块 GPU 通过 NVLink 连接;pod 中的每台主机则通过高达 200GB/s 的高带宽网络连接,8GPU 共享该网络。

    借助 ZionEX,我们通过以下分布式训练策略解决计算和内存容量问题:

    • 我们首先将 embedding tables 分布式存储在一个 pod 中。为了更好地实现负载均衡并处理超大规模 embedding tables ,我们不将完整的 embedding tables 放置在不同 GPU 上,而是主动将超大规模 embedding tables按列等分为多个 shards,并根据经验成本函数(empirical cost function )分配这些列到这些 shards 。该成本函数同时考虑了这种分配方式的计算开销和通信开销。我们使用 LPT《Longest-processing-time-first scheduling》)近似集合划分(approximate set partition )算法,根据成本函数以负载均衡的方式分布 table shards

      这个成本函数如何定义?论文并未详细说明。

    • 另一方面:对于包括 DHEN layers 在内的 dense 模块,我们在每块 GPU 上复制它们,并以数据并行(data parallel: DP)方式进行训练。这一选择基于以下观察:dense DHEN layersactivation 规模可能远大于 weights 本身,因此同步 weights 的成本低于通过网络传输 activations 的成本。这种训练策略形成了一种混合训练范式:每个 batchDP 开始,进入模型并行(model parallelism: MP)以进行分布式 embedding lookup ,最后通过 DP 处理 dense layers

  3. 使用 DP 训练 dense 模块会将 stacked DHEN layers 的参数大小限制在单 GPUHBM 容量内,这阻碍了我们对 DHEN 可扩展性极限的探索。为解决这一问题,我们采用以下技术:

    • fully sharded data parallel: FSDP:通过将权重进一步分片到不同 GPUs ,消除传统数据并行中的内存冗余。

    • activation checkpointing: 通过用更多计算换取更少的峰值内存使用(peak memory usage)。

    • cpu offloading:通过积极将参数和梯度存储到 CPU,并在需要前将其加载回 GPU ,从而进一步减少 GPU 内存使用。

    由于所有这些技术都会影响训练效率,我们根据 DHEN 的层数精心调整系统,仅启用满足训练需求的最小技术集合。

    DHEN 的训练策略总结如 Figure 4 所示。

1.2.2 训练优化

  1. 我们提供了额外的优化措施以进一步加速训练。

  2. 通用 Optimizations:我们启用了一组广泛使用的优化措施,包括:

    • large batch training 从而减少 synchronization 频次。

    • 带随机 roundingFP16 embeddingBF16 optimizerquantized all to all and allreduce collectives 从而降低内存占用,并提升数值稳定性。

    • 使用 Tensor Core 等专用加速硬件来减少通信开销。

  3. Hybrid Sharded Data Parallel:为了在训练成本预算内找到最佳的 DHEN 配置,我们必须同时探索 configuration of token mixers 和层数。在实践中,我们需要频繁实验内存占用略超过单 GPUHBM 容量的候选 DHEN 模型(我们称这些模型为 “尴尬尺寸” 模型),这使得数据并行不可行,必须使用 FSDP 。然而,在这些情况下,我们发现直接应用 FSDP 无法实现最佳效率,尤其是在生产规模(数百块 GPU )下训练时。这是因为在前向传播和反向传播的关键路径上,从每块 GPU 收集 weights 不同分片的全收集(allgather )操作需要较长时间完成,原因包括:

    • (1):需要与集群中的所有 GPU 通信。

    • (2):每个分片太小了,无法有效利用网络带宽。

    尽管预取(prefetching)等技术有助于缓解这一问题,但它们会给系统带来内存压力,从而违背使用 FSDP 的初衷。

    为此,我们针对这些 “尴尬尺寸” 网络提出了一种新的训练范式——hybrid sharded data parallel: HSDP,该范式与 DHEN 模型协同设计,适用于我们的训练集群。HSDP 利用了 GPU 互连(NVLink600GB/s )和主机互连(RoCEv225GB/s )之间 24 倍的带宽差异,其工作原理如下:

    • (1)HSDP 在单个主机内对整个模型进行分片,使得 FSDP 中反向传播的 reduce-scatter 操作、以及前向和反向传播的 allgather 操作完全在主机内进行。

    • (2):当反向传播中的 reduce-scatter 操作完成后,我们通过 hook 来触发并发的 allreduce 操作,每个操作涉及具有相同 local host ID 内的所有 GPU ,从而以异步方式计算其 local shard 的平均梯度,避免阻塞反向传播的计算。

    • (3):最后,我们注册一个 backward hook 来等待 pendingallreduce 操作的 handles,从而确保 HSDP 在训练语义上与FSDPDDP 保持一致。

    我们在 Figure 5 中对比了 FSDPHSDP

    定性来看,与 FSDP 相比,HSDP 具有以下优势:

    • (1):前向传播和反向传播关键路径上的 allgather 操作的延迟显著减少,因为它可以利用高速 NVLink 互连,无需通过较慢的 RoCE 链路进行跨主机通信。

    • (2):通信操作的规模随主机数量而非 GPU 数量扩展,提高了效率。

    然而,HSDPtradeoff 在于:

    • 它支持的模型最大规模是纯数据并行支持规模的 8 倍(每台主机的 GPU 数量)。

    • 由于添加了 allreduce 操作,在数据传输量上会产生 1.125 倍的通信开销。

    对于这些 “尴尬尺寸” 模型,这些 tradeoff 是可接受的,因为 allreduce 开销可以通过下一层的计算来隐藏,且 allreduce 是高度优化的操作。

1.3 实验

  1. 在本节中,我们进行实验以评估 Deep Hierarchical Ensemble NetworkCTR prediction 任务中的有效性。我们的实验旨在:

    • (1):评估 hierarchical ensemble 的有效性,并确定一组适合集成的 interaction 模块。

    • (2):将 DHENSOTADLRM model 进行端到端的准确性对比,验证其增益,并展示我们 system-level optimizationsDHENtraining throughput 的提升效果。

  2. 我们使用工业数据集进行所有实验。如前所述,内部的 AdvancedDLRM 模型被视为 interaction 模块之一。为简洁起见,我们省略了 AdvancedDLRM 和数据集设置的详细描述。所有实验的训练超参数(包括学习率和优化器配置)均保持一致。所有模型均使用数百个 sparse (categorical) features 和数千个 dense (numerical) features 进行训练。全同步训练方案(full-sync training scheme )确保模型性能和训练吞吐量具有可复现性。我们使用归一化熵损失(Normalized Entropy loss)评估 CTR prediction 准确性。

    如何 ensemble?作者并未说明。也没有消融实验。

1.3.1 不同 Interaction 模块的模型变体

  1. 本节旨在了解在常用模块下 DHEN 的表现。具体而言,我们使用四种 interaction 模块(DCNSelf-attentionCNNLinear)来测试 hierarchical ensemble 的性能(下一节将深入与 AdvancedDLRM 进行对比)。在每个实验中,每层包含不同类型的 interaction 模块。Table 1 显示了不同模型设置的性能,其中 N 表示堆叠层数。为便于比较,我们以 DCN 为基线,使用不同 training stepstraining examples )下的 relative Normalized Entropy loss 的差异来评估模型性能。

    Table 1 可以看出:

    • model (4) with linear interaction modulemodels with single interaction module per layer 中表现最佳,而 model (2) with self-attention interaction module 需要大量训练数据才能收敛。

    • 我们进一步观察到,self-attention interaction modulelinear interaction modulehierarchical ensemble 模型 (6) 在所有变体中表现最佳。这一结果表明,hierarchical ensemble architecture 捕获了self-attention interaction modulelinear interaction module 的相关性,并显著帮助 self-attention interaction module 收敛。

    • 此外,我们注意到每层集成更多 interaction 模块并非总是有益,例如:

      • DCN 和线性层的 hierarchical ensemble 模型 (5) 的性能优于单独堆叠 DCN layer ,但不如单独使用线性层。

      • 类似地,self-attention interaction moduleCNN interaction module 的集成模型 (7) 也存在这种情况。

这些观察证实了我们的初始假设,即不同 interaction 模块捕获非重叠的 interaction 信息,因此 hierarchical ensemble 机制是 DHEN 性能的关键。

1.3.2 工业级 AdvancedDLRM 实验

  1. Model Prediction 性能:在验证了具有各种 interaction 模块的 DHEN 的有效性后,我们进一步评估其在工业级模块AdvancedDLRM 上的效果。具体而言,我们利用 SOTAAdvancedDLRM 模块(《Deep learning recommendation model for personalization and recommendation systems》)和线性模块作为 hierarchical ensemble 的两个 interaction 模块。Table 2 显示了 DHENAdvancedDLRM 模型的性能对比。

    我们训练了 4DHEN 模型,每层均集成 AdvancedDLRM interaction moduleLinear interaction module 。每个模型都有 N 层。从 Table 2 可以看出:

    • 所有 DHEN 模型均优于工业级 AdvancedDLRM 模型,且更深的 DHEN 模型实现了更大的 Normalized Entropy 改善。

    • 在所有情况下,我们都观察到随着处理的训练样本的增加,gain 持续扩大。

    这表明高阶 interactions 、以及各种 interactions 的相关性在 CTR prediction 任务中均起着重要作用;且更大数据集下 Normalized Entropy 改善的扩大表明 DHEN 的性能具有一致性和泛化性。

    这个对比有问题:baseline 只有 1 层,而其他模型都不止一层。因此很难说明效果是由于模型架构更先进?还是由于模型更深?

  2. Scaling 效率:为评估 DHEN 模型的 scaling 效率,我们在 Table 3 中总结了使用 MoEstacking DHEN layers 两种 scaling 方法的对比结果。两种 scaling 方法均基于相同的工业级 AdvancedDLRM 实现。我们使用 MoE 架构扩展 AdvancedDLRM 中的 MLPlayers ,同时使用AdvancedDLRM 模块和 Linear 模块的 ensemble 来堆叠 DHEN layers

    结果表明,在训练复杂度(以 FLOPs 衡量)相近的情况下,DHEN 在准确性上始终优于 AdvancedDLRM MoE,因此 stacking DHEN layers 是一种具有更高投资回报的有效 scaling 机制。

  3. Training 吞吐量:基于 ZionEX 设计,本节评估在 256-GPU 集群上训练 DHEN 模型的吞吐量。我们首先展示使用 DP 训练 8DHEN 模型时 system level optimizations 的有效性。总体而言,应用 FP16 embedding、自动混合精度(AMP)、quantized BF16 allreduce and all to all collectives 后,端到端速度提升了 1.08 倍。这一改进源于 quantized collectives 操作减少了通信延迟,剩余瓶颈在于 optimizer costall to all collectives 调用(该调用无法与密集层的计算完全重叠)。

    尽管 DHEN 随层数的增加而 scale 良好,但训练更深的 DHEN 模型需要范式转变,因为 DP 在我们的集群中最多只能支持 22DHEN。我们现在量化所提出的与 DHEN 协同设计的训练范式 HSDP 如何弥补 FSDPDP 在吞吐量上的差距,以及 DP 在内存上的限制。我们对不同层数的 DHEN 模型进行训练吞吐量实验,使用三种范式训练每个模型并对比结果。如 Figure 6 所示:

    • DP16 层内表现良好,层数进一步增加会导致错误。

    • 另一方面,FSDPHSDP 均支持更多层数,且 HSDP 始终比 FSDP1.2 倍。

    我们的 trace analysis 表明,HSDP 的吞吐量提升确实源于 training 的关键路径上 allgather latency 的大幅减少,这得益于单主机内快速的 NVLink 连接。