Nest: Network- and Memory-Aware Device Placement for Distributed Deep Learning

文章标题: NEST：面向分布式深度学习的网络和内存感知设备布局
作者/机构: Irene Wang 1, Vishnu Varma Venkata 1, Arvind Krishnamurthy 2, Divya Mahajan 1

A1 主要贡献

核心问题与研究目标: 随着深度学习模型参数量和复杂度的急剧增长，高效的分布式训练已成为扩展现代AI系统的关键瓶颈。现有设备布局框架大多依赖启发式或拓扑无关的搜索，将通信和内存分开处理。这些方法由于缺乏对单个设备内存的感知，通常在事后通过将参数和激活分散到多个设备来确保可行性，但这会增加同步开销、通信量并导致计算资源未被充分利用，限制了在真实数据中心网络上的可扩展性和效率。现有框架（如Alpa和TopoOpt）虽然引入了拓扑感知，但在可扩展性和通用性上仍有局限，例如Alpa假设简化的2D网格，并在布局后才验证通信驱动的分片，导致过度分片和计算资源利用不足。

创新点与主要贡献:
本文介绍了NEST，一个原则性强、可扩展、同时感知计算-内存-网络的设备布局框架，通过新颖的动态规划（DP）公式统一了模型和网络的优化。NEST的核心创新点在于其增量式处理通信和内存成本。由于布局是反向求解的，某一层级的通信依赖于未知的下游布局。NEST通过预先计算抽象网络层级（如节点内、机架内）之间的延迟，并在多个候选下游布局之间评估转换，解决了这一问题，同时保证了网络准确性和可证明的最优性。

NEST的主要贡献如下：
1. 高效的网络感知动态规划算法：为分层、异构的互连网络提供了可证明的最优性保证。
2. 对多样化并行策略的统一支持：包括张量、流水线、数据、专家、序列、上下文并行以及ZeRO，并通过结构化组合实现可行的规模化搜索。
3. 全面的网络和内存建模：反映了生产规模的拓扑结构，并能够系统地探索布局与基础设施之间的权衡。

通过将细粒度的子图（SUB-GRAPH）建模与全局的图（GRAPH-GLOBAL）探索相结合，NEST实现了可扩展且通用的设备布局。评估表明，在各种硬件和网络上，NEST相较于最先进的基准，实现了高达2.43倍的吞吐量提升，具有更好的内存效率和可扩展性。

A3 背景知识与关键动机

分布式深度学习。现代框架如Megatron-LM【NVIDIA, c】、NeMo【NVIDIA, 2022】和DeepSpeed【40, Rasley et al., Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters, 2020, KDD ’20】通过多种并行策略实现可扩展训练。层内张量并行（TP）将单个层分割到多个加速器上，序列并行（SP）进一步沿序列维度划分激活以减少内存压力。专家并行（EP）将混合专家（MoE）模块分布在设备上，引入了独特的all-to-all通信模式。上下文并行（CP）【46, Yang et al., Context parallelism for scalable million-token inference, 2025】将输入序列分段到多个GPU上，专门用于处理长上下文注意力机制的二次内存扩展。流水线并行将层分配给不同设备，通过微批处理减少每台设备的内存使用【8, Huang et al., GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, 2019, NeurIPS 2019】，而数据并行则在设备间复制模型。像ZeRO【38, Rajbhandari et al., Zero: memory optimizations toward training trillion parameter models, 2020, SC ’20】这样的内存优化技术通过分区优化器状态、梯度和参数来减少内存使用。每种方法都引入了独特的通信模式和资源权衡，需要仔细协调才能实现高效训练。

数据中心网络拓扑。现实世界的数据中心很少是统一的，通常采用如胖树（Fat-Tree）、Spine-Leaf或Clos网络等分层拓扑，以平衡可扩展性、容错性和成本【18, Microsoft, Inside maia 100: Revolutionizing ai workloads with microsoft’s custom ai accelerator, 2024; 26, NVIDIA, Nvidia dgx superpod, 2024】。实际部署中常因物理布局和成本限制而出现超额预订或非均匀互连，导致节点内、机架内和机架间链路的延迟和带宽差异巨大。例如，同一机架内的通信可能使用全带宽的NVSwitch连接，而跨机架传输则需经过共享或超额预订带宽的主干交换机。这些分层的非均匀性使得分布式训练高度依赖于拓扑结构。附录B展示了常见的网络配置。

先前的设备布局技术。许多工作已探索使用强化学习【19, Mirhoseini et al., Device placement optimization with reinforcement learning, 2017, ICML’17】、马尔可夫链蒙特卡洛（MCMC）【43, Wang et al., TopoOpt: Cooptimizing network topology and parallelization strategy for distributed training jobs, 2023, NSDI 23】和动态规划（DP）【42, Tarnawski et al., Efficient algorithms for device placement of dnn graph operators, 2020, NIPS ’20; 41, Tarnawski et al., Piper: Multidimensional Planner for DNN Parallelization, 2021, NeurIPS 2021; 44, Wang et al., Integrated hardware architecture and device placement search, 2024; 45, Zheng et al., Alpa: Automating inter- and IntraOperator parallelism for distributed deep learning, 2022, OSDI 22】来自动进行模型分区。早期的系统如PipeDream【20, Narayanan et al., PipeDream: Generalized Pipeline Parallelism for DNN Training, 2019, SOSP】主要关注流水线并行，而后续的框架【41, Tarnawski et al., Piper: Multidimensional Planner for DNN Parallelization, 2021, NeurIPS 2021; 45, Zheng et al., Alpa: Automating inter- and IntraOperator parallelism for distributed deep learning, 2022, OSDI 22; 47, Zhu et al., Mist: Efficient distributed training of large language models via memoryparallelism co-optimization, 2025, EuroSys ’25; 17, Liu et al., Aceso: Efficient parallel dnn training through iterative bottleneck alleviation, 2024, EuroSys ’24】则结合了数据、算子和层间并行。然而，大多数工作假设网络是扁平或简化的，并未完全支持新兴策略，如专家并行或像ZeRO这样的内存优化。

网络层次结构对性能的影响。在实践中，数据中心网络是分层的且常常是超额预订的，节点内、机架内和机架间的延迟和带宽各不相同。这些差异直接影响性能关键的集合通信操作，如AllReduce、AllGather和ReduceScatter，这些操作通常在分布式训练开销中占主导地位【46, Yang et al., Context parallelism for scalable million-token inference, 2025; 6, Go et al., Characterizing the efficiency of distributed training, 2025, MICRO】。

现有拓扑感知方法的局限性。一些拓扑感知方法试图解决这些现实问题，但在可扩展性或最优性方面仍然有限。TopoOpt通过随机方式探索布局，缺乏最优性保证，对初始化敏感，并且随着并行维度的增加，扩展性很差。Alpa使用动态规划，但为了保持搜索的可行性，它假设简化的2D网格网络，无法捕捉超额预订或多级层次结构。它优先考虑最小化通信，而没有联合建模计算延迟，并且只在生成布局计划后（即事后）检查内存可行性。缺乏集成的内存建模进一步增加了搜索时间，并降低了在较小集群上找到可行布局的能力。为了满足内存预算，Alpa积极地将参数和激活分片到多个GPU上。在较大的集群上，它独立优化每个流水线阶段并构建单个流水线；额外的设备用于进一步分片层，而不是通过流水线复制来扩展。这可能导致硬件利用率不足和通信过多，从而限制了在大约64个GPU之外的有效扩展。

近期工作的权衡。为了解决这些限制，最近的工作将重点转向内存可行性和调度，但通常以牺牲显式网络建模为代价。Aceso【17, Liu et al., Aceso: Efficient parallel dnn training through iterative bottleneck alleviation, 2024, EuroSys ’24】使用贪婪的、追逐瓶颈的启发式方法来改善内存平衡和硬件利用率。然而，其局部搜索可能会陷入次优配置，并且不提供全局最优性保证。Mist【47, Zhu et al., Mist: Efficient distributed training of large language models via memoryparallelism co-optimization, 2025, EuroSys ’25】使用MILP来制定布局和调度，强调通信和计算的时间重叠。虽然它比Alpa更稳健地处理内存，但它将网络拓扑视为次要约束，并依赖于分层的、蛮力搜索来进行并行化，限制了在大型分层集群上的可扩展性。

先前工作的对比。表1展示了先前工作之间的对比。

对拓扑感知布局的需求。随着训练规模扩展到数百个加速器，布局上的微小差异可能导致巨大的性能变化。图2展示了GPT3-175B、Llama3-70B和Mixtral-8×7B模型在一个2:2超额预订的64-GPU集群上的情况，其中通信占总训练时间的很大一部分。最有效的并行策略取决于模型和底层拓扑——在均匀网格上有效的方法在分层或带宽不对称的集群上可能表现不佳。因此，高效的训练需要一个能够明确建模网络层次、不对称性和超额预订的布局框架。

NEST的动机。这些限制凸显了对一个布局框架的需求，该框架应具备以下特点：（i）策略感知，协同优化数据、张量、流水线、专家和内存并行；（ii）网络感知，明确建模分层、超额预订和不对称的互连；（iii）可扩展且有原则，使用动态规划等结构化优化方法，而非启发式或采样。我们引入NEST，一个网络和内存感知的动态规划框架来解决这些问题。与依赖随机搜索或简化拓扑模型的先前工作不同，NEST明确建模分层、超额预订和异构网络，同时将内存建模集成到其搜索方法中。这使得NEST能够系统地探索平衡通信、计算和内存效率的可行布局，并在异构和超额预订的数据中心架构中有效扩展。

Table 1. 先前网络和内存感知设备布局工作的比较。

A2 方法细节

NEST在搜索过程中联合建模计算、内存和通信，以在真实的数据中心条件下优化网络感知的设备布局和分布式并行策略。图3展示了NEST框架及其支持的代表性并行策略集。

3.1 并行策略的分类

NEST的核心洞察。NEST的一个关键洞察是，分布式训练的并行技术可以根据其与模型图和硬件的交互，沿着正交的维度进行组织。这种有原则的分类通过将局部的、算子内的转换与全局的、算子间的调度决策分开，使得动态规划搜索变得易于处理。它使NEST能够将张量、流水线、数据、专家、序列和上下文并行等策略组合成混合配置，而不会导致搜索空间的组合爆炸。沿着这些维度分解优化，产生了一个系统化、可扩展的框架：新的并行方法可以无缝集成，混合布局可以在统一的动态规划公式中高效探索。基于这一原则，NEST将策略分为两类——子图（sub-graph）和图全局（graph-global）——每类对应计算图中的一个不同抽象层次。

子图（SUB-GRAPH）策略。这些策略在单个算子或小子图内部操作，例如张量和专家并行。NEST将一个策略归类为SUB-GRAPH，如果它修改了算子的内部执行，例如通过划分参数或调整每个加速器的张量维度，同时保持模型的整体控制流和数据流（即层序列和依赖关系）。其目标是通过将工作负载（例如，权重矩阵或专家激活）分散到设备上来缓解层内的计算和内存瓶颈。这种细粒度的分区引入了紧密耦合的集合操作，如用于张量并行的AllReduce、用于专家并行的AllToAll，以及用于序列和上下文并行的AllGather和ReduceScatter。由于这些转换是局部的和自包含的，NEST会离线分析它们的算子级计算、内存和通信成本。这些预先表征的成本随后在更高层级的布局过程中被分析性地组合，而无需扩展动态规划（DP）的搜索空间。任何用等效的分布式实现替换一个层的算子级技术，都可以通过提供其转换后的计算图和运行时注释，作为SUB-GRAPH策略被集成。

图全局（GRAPH-GLOBAL）策略。GRAPH-GLOBAL策略作用于整个模型计算图，例如流水线并行（将层划分为阶段）、数据并行（复制整个模型并同步梯度）和ZeRO风格的分片（分发优化器状态）。它们决定了层、参数和数据批次如何在多个设备或节点之间分布和调度，从而重塑了全局执行计划。这些策略影响全局通信结构、内存平衡和阶段间依赖，其影响无法局限于单个算子。因此，NEST将GRAPH-GLOBAL策略直接整合到其新颖的基于DP的搜索中，在系统级约束下联合优化阶段边界、复制因子、内存占用和通信开销。

讨论。通过将SUB-GRAPH和GRAPH-GLOBAL策略视为正交维度，即局部转换和全局调度，NEST避免了枚举所有配置的组合爆炸。相反，它在一个全局优化的布局中组合了预先分析的局部策略，从而能够在一个统一的计算-内存-通信成本模型下，系统地集成新的并行形式（例如，专家或激活分片）。

3.2 NEST 工作流程

三阶段优化工作流。NEST构建了一个跨越三个阶段的统一优化工作流：图提取、运行时估计和布局搜索。重要的是，NEST严格作为一个规划系统运行：它保持原始训练公式的数学等价性，并且不修改底层的模型代码或核函数。

输入。它接收以下输入：（1）硬件规格（例如，GPU、TPU或领域特定加速器），（2）模型的算子图，以及（3）详细的网络配置，指明拓扑、带宽、延迟和通信协议。

图提取。NEST使用符号追踪从训练脚本中提取算子图，并应用用户指定或预定义的逻辑SUB-GRAPH转换来表示标准的层内并行策略（例如，张量或专家拆分）。然后，它将相应的通信算子插入到成本模型中。每个产生的变体都会生成一个自包含的算子图，作为运行时建模的基础。

运行时估计。对于每个提取的图，NEST进行离线分析，用特定平台的运行时来注释计算和通信算子。计算成本来源于硬件估计器或分析器，如PyTorch【37, PyTorch, Pytorch profiler, 2021】，而通信成本则使用网络模拟器（如AstraSim【46, Won et al., Astra-sim2.0, 2023】）进行估计。这提供了跨层次的计算、通信和内存行为的统一视图。

求解器。给定标注了计算和通信成本、每层内存使用量以及集群网络特性的算子图，DP求解器探索GRAPH-GLOBAL配置，包括流水线阶段划分、复制度和数据分区，同时整合SUB-GRAPH策略的成本模型。这种结构化方法评估了延迟、内存和带宽之间的权衡，识别出在内存和通信约束下最大化端到端吞吐量的布局。最终输出是一个并行配置和布局计划。通过将细粒度的SUB-GRAPH建模与全局的GRAPH-GLOBAL探索相结合，NEST实现了可扩展且通用的设备布局。由于NEST利用已建立的并行策略而不修改底层的数学算子，因此生成的计划可以在标准的分布式框架（如Megatron-LM【NVIDIA, c】和NeMO【25, NVIDIA, Nemo, 2022】）上执行，同时保持原始训练公式的数学等价性。

3.3 NEST 内存建模

内存需求估计。遵循先前工作【44, Wang et al., Integrated hardware architecture and device placement search, 2024; 47, Zhu et al., Mist: Efficient distributed training of large language models via memoryparallelism co-optimization, 2025; 16, Lin et al., nnScaler, 2024】的方法，NEST通过Torch.fx【40, Reed et al., torch.fx, 2021】使用基于符号的分析来估计工作负载的内存需求。在图提取期间，NEST追踪依赖图以注释算子元数据，包括参数、输入大小和激活张量。这些信息被求解器用于内存建模。

峰值内存组成。一个连续子图（阶段）S的峰值内存占用在流水线的“稳态”期间建模，并包括五个部分：模型权重、累积梯度、优化器状态、当前中间激活和暂存数据（为处理中的微批次保留的激活）。暂存数据的量与调度相关：对于1F1B调度，距离流水线末端索引为s的阶段持有(s − 1)个微批次的激活；对于GPipe，这个数量按B/d的比例缩放。

重计算策略。NEST评估两种重计算策略：（1）无重计算，即所有中间激活都被暂存；（2）有重计算，即只暂存阶段边界的输入激活，中间张量在反向传播过程中重新物化。

峰值内存计算公式。NEST使用以下方程计算峰值内存：

$$\begin{aligned} \begin{aligned} \text{Mem}(S, s) = & \\ \sum_{L_{i} \in S} & \Big( 2 \cdot \text{weights} + \text{opt\_states} + \text{activations} \Big) + \\ & (s - 1) \cdot \text{stashed\_data} \end{aligned} \end{aligned}$$

通过将内存建模为阶段位置s的线性函数，求解器避免了在不同流水线位置上的冗余计算，显著加快了搜索优化。NEST的内存建模已与编译后的核函数进行了验证，并在评估的模型中平均误差在7%以内。模型估计的摘要见附录H。

4 NEST的动态规划

统一的动态规划框架。NEST引入了一种新颖的、感知网络、计算和内存的动态规划（DP）公式，该公式将网络拓扑特性、带宽不对称性以及每阶段的计算延迟和内存约束直接嵌入到优化过程中。尽管先前的工作已使用DP进行设备布局，但没有一个将网络、计算和内存感知集成到一个统一的优化框架中【45, Zheng et al., Alpa: Automating inter- and IntraOperator parallelism for distributed deep learning, 2022, OSDI 22; 44, Wang et al., Integrated hardware architecture and device placement search, 2024; 41, Tarnawski et al., Piper: Multidimensional Planner for DNN Parallelization, 2021, NeurIPS 2021】。通过这样做，NEST可以联合推理共置、重叠和剪枝，确保在异构集群上实现可行且高效的布局。本节介绍此DP求解器的设计，从一个基线公式演进到一个完全集成的、拓扑和内存感知的优化器。

DP输入与搜索空间。NEST的DP以分析过的算子图作为输入，这些图捕获了张量和专家并行的多种候选配置。每种配置记录了算子计算延迟、集合通信成本（例如AllReduce、AllToAll）以及每层的内存使用情况。DP还考虑了内存节省技术，如ZeRO分片（参数、梯度、优化器）、用于层分区的流水线并行以及用于可扩展复制的数据并行。这些分析过的图和系统特性定义了DP的搜索空间。每个状态代表一个层、设备和并行策略的部分分配，而每个DP转换都是一个尊重内存和网络约束的可行扩展。通过在这个结构化搜索空间中进行推理，NEST系统地探索数千种混合布局，而无需显式枚举，从而在一个统一的优化框架内平衡计算分布、通信局部性和内存可行性。

DP中的局部与全局并行。NEST在其DP公式中区分了子图（局部）和图全局策略。局部策略，如张量和专家并行，修改了一层内部计算和参数在加速器间的分布方式。它们的影响在分析阶段被捕获，并影响每阶段的延迟项$load(·)$。相比之下，图全局策略，如数据、流水线和ZeRO并行，重塑了层或阶段之间的边界，并由DP在分区和布局期间显式探索。这种分离使得NEST能够维持一个可处理的状态空间，同时仍然能够系统地组合混合配置（例如，张量+流水线+数据并行）。子图策略通过在多个局部性级别（如节点内、节点间和机架间）建模集合通信来融入网络感知，使DP能够推理局部通信成本如何随物理布局和拓扑变化。

基线DP：网络无关。我们从一个忽略网络和内存约束的简化基线开始。令$dp[D][k][s]$表示在k个设备和s个流水线阶段上执行一个下集D（层图的前缀）的最小延迟。递推关系枚举了所有子图$D' \subseteq D$和可能的设备分配$a$，同时考虑了局部策略，如张量或专家并行宽度（$t, e$）：

在这里，$D \setminus D'$表示分配给$a$个设备的新阶段，而$load_{te}(S, a, s)$估计了在给定的张量和专家并行配置下该阶段的延迟。然后，递推关系递归进行，在剩余的$k-a$个设备和$s-1$个阶段上优化剩余的层。这个公式有效地剪枝了次优的分解，但假设了一个全连接、成本统一的网络，并忽略了由内存引起的约束。

挑战与分层网络抽象。当引入网络异构性时，一个核心挑战出现了：每个流水线阶段的延迟不仅取决于其自身的计算和反向通信，还取决于其前驱（产生前向激活）的未知布局。由于DP是向后进行的（从最后一个阶段到第一个阶段），在决策时生产者的位置及其通信成本是未知的。图4说明了这一挑战。在放置第23层时，到第24层（已放置）的反向链接具有基于它们相对位置的已知延迟，但来自第22层的前向链接取决于该阶段将被放置在何处（例如，同一节点、同一机架或远程）。这种不对称性打破了最优子结构假设。

分层网络抽象。为了恢复可处理性，NEST引入了一种按通信局部性对设备进行分组的分层抽象。NEST不为每对设备建模，而是将网络层级$l$视为一个离散的DP状态变量。这种抽象允许NEST通过基于物理互连填充一个分层成本矩阵来泛化到真实世界的拓扑。例子包括：
* 分层结构（NVIDIA HGX/Spine-Leaf/Fat Tree）: 层级直接映射到交换机层次和链路类型。例如，$l_0$代表节点内NVLink（如H100为900 GB/s），而$l_1$和$l_2$则捕捉了节点间以太网/InfiniBand和超额预订的机架间Spine-Leaf链路的显著带宽不对称性。这使得DP可以避免将高带宽的流水线阶段放置在超额预订的边界上。
* 非均匀网格（TPUv4 Torus）: 在基于环面（Torus）的系统中，层级代表跳数距离或坐标亲和度类别（例如，$l_0$: 同一瓦片，$l_1$: 1跳，$l_2$: 远程）。通过将这些映射到分析的延迟，NEST在扩展到数千个加速器时考虑了非均匀的通信成本。

如图4所示，未知的前向成本由几个网络层级（通常为3-5个）表示。这使得DP可以基于层级而不是所有设备对进行推理，从而在保持分层拓扑保真度的同时减少了组合复杂性。附录C展示了这种抽象如何泛化到非均匀网格和环面架构。

关键观察。这种抽象是拓扑无关的。通过将局部性与固定的层次结构解耦，分层抽象使得NEST的DP能够泛化到多样化和未来的互连，保持可扩展性的同时适应新的拓扑和通信模型。

提出的DP：网络感知。通过这种抽象，DP递推关系变为：

$$dp^{te}[l][D][k][s] = \min_{D' \subseteq D} \min_{a} \max \left( dp^{te}[l'][D'][k - a][s - 1], \text{load}_l^{te}(D \setminus D', a, s) \right)$$

这里，$dp_{te}[l][D][k][s]$表示使用$k$个设备、划分为$s$个流水线阶段，在张量/专家并行配置（$t, e$）下，执行层后缀$D$的最小可能延迟（即瓶颈阶段的成本）。关键的新项是$l$（层级）。由于DP是向后进行的，从最后一层到第一层，当放置当前阶段（如图4中的第23层）时，下一个阶段（如图4中的第22层）的布局是未知的。$l$代表了相对于$D$中第一个阶段，尚未放置的生产者阶段的假定通信距离，作为一个“延迟的前向成本”，保留了最优子结构。

统一的成本模型和递推。公式3联合寻找新阶段（$D \setminus D'$）的最优流水线切分点（$D'$）和设备分配（$a$），在探索跨通信层级转换的同时，剪枝内存不可行的状态。延迟估计$load_{tel}(\cdot)$是协同优化的核心，它捕获了：
* 计算延迟: 源自分析的每个算子的运行时，并根据张量/专家并行度进行缩放。
* 网络延迟: 由分层通信成本矩阵确定，考虑了集合通信以及前向（来自层级$l$）和后向（到$D'$）的流量。
* 内存优化协同设计: AR（激活重计算）和ZeRO被原生集成到状态转换中。如果一个状态超过了设备内存，求解器会逐步增加ZeRO级别（1, 2, 或 3）直到达到可行性，并将产生的集合通信开销加到延迟中。类似地，重计算被建模为一个二元优化选择，它以增加计算延迟为代价，减少了“暂存数据”的内存项。

递推关系使用$max(\cdot, \cdot)$来平衡新阶段的成本($load_{tel}$)与剩余阶段的成本($dp_{te}[l'][D'][k-a][s-1]$)，其中$l'$编码了连续阶段之间的通信层级。这个统一的模型确保NEST只考虑有效的配置，自动在内存减少（分片）和通信开销（共置）之间进行权衡。总的来说，这个公式让NEST能够联合优化布局、通信和内存，倾向于将频繁通信的阶段共置，同时惩罚跨层级的流量。

示例。考虑图4中的例子。在放置第23层时，DP枚举其相对于第24层在每个通信层级上的可能位置。对于每个选项，它预先计算与第22层潜在布局对应的延迟前向成本。这些部分子解存储在$dp[l][23,24][2][2]$中。之后，在放置第22层时，NEST只需使用这些预计算的成本来选择全局最优配置——在没有穷举枚举的情况下实现最优性。

总结。NEST的动态规划在三个方面是新颖的：（1）一个分层的、层级化的抽象，使得网络感知的布局变得易于处理；（2）一个统一的计算-网络-内存模型，将可行性嵌入到递推关系中；（3）在一个结构化的优化框架内，对局部和全局并行进行正交处理。

A4 实验环境与结果

5.1 方法论与设置

• 模型: 评估中使用了多种大型语言模型，包括Llama2-7B【42, Touvron et al., Llama 2, 2023】, Llama3-70B【7, Grattafiori et al., The llama 3 herd of models, 2024】 (来自Hugging Face【44, Wolf et al., Huggingface’s transformers, 2019】), BertLarge【4, Devlin et al., BERT, 2019】, GPT3-175B【2, Brown et al., Language models are few-shot learners, 2020】, 以及Mixtral-8x7B【10, Jiang et al., Mixtral of experts, 2024】 (来自MegatronLM【NVIDIA, c; 41, Shoeybi et al., Megatron-lm, 2020】)。操作符图使用Torch.fx【40, Reed et al., torch.fx, 2021】提取。除非另有说明，所有实验均使用4096的全局批量大小和1的微批量大小。
• 运行时估计: 对于类TPUv4的加速器，操作符延迟使用经过硬件验证的库Sunstone【27, Olyaiy et al., Sunstone, 2023】和Tandem【5, Ghodrati et al., Tandem processor, 2024】进行估计。对于GPU，计算延迟使用PyTorch Profiler【37, PyTorch, Pytorch profiler, 2021】在操作符和层级别进行分析，通信集合操作使用AstraSim【46, Won et al., Astra-sim2.0, 2023】建模和验证。
• 硬件与软件: NEST在单个H100 GPU上运行以提取操作符图和运行时配置。实验使用RHEL 9.6操作系统，Python 3.10，以及用g++ 12.4.0编译的C++求解器。
• 基准: NEST与五个基准进行了比较：
  1. 手动布局: 来自先前工作【22, Narayanan et al., Efficient large-scale language model training, 2021b; 44, Wang et al., Integrated hardware architecture and device placement search, 2024】的手动布局，通过扩展数据并行来适应不同的集群规模。
  2. Phaze: 一个基于Piper【41, Tarnawski et al., Piper: Multidimensional Planner for DNN Parallelization, 2021, NeurIPS 2021】构建的网络无关的DP框架【44, Wang et al., Integrated hardware architecture and device placement search, 2024】。
  3. MCMC: 基于MCMC的布局方法【43, Wang et al., TopoOpt: Cooptimizing network topology and parallelization strategy for distributed training jobs, 2023, NSDI 23】。
  4. Alpa: 一个利用算子内并行的先进DP框架【45, Zheng et al., Alpa: Automating inter- and IntraOperator parallelism for distributed deep learning, 2022, OSDI 22】。为解决其大规模部署时需要全集群访问进行分析的不可行性，创建了一个离线变体Alpa-E。
  5. Mist: 一个专注于调度和内存优化的先进DP框架【47, Zhu et al., Mist: Efficient distributed training of large language models via memoryparallelism co-optimization, 2025, EuroSys ’25】。

5.2 在胖树拓扑上的加速器执行

实验设置: 在包含64至1024个类TPUv4加速器的胖树拓扑上评估NEST。每个节点有8个通过HGX式链接（900 GB/s）连接的加速器，每4个节点连接到一个一级交换机（100 GB/s），二级聚合带宽为400 GB/s。

结果分析:

• 吞吐量提升: 如图5所示，NEST的分布策略平均比手动布局吞吐量高1.59倍，比基于MCMC的搜索高1.71倍，比Alpa-E高2.43倍，比Phaze高1.19倍。NEST随集群规模几乎线性扩展，而其他基准要么找不到有效布局，要么收益递减。
• 与随机搜索（MCMC）比较: MCMC方法依赖随机探索，性能随集群和模型规模增长而下降，对初始化敏感，证明了结构化确定性优化的必要性。
• 与Phaze比较: Phaze假设网络扁平统一，忽略了通信成本，导致在异构互连上吞吐量下降。NEST通过将网络拓扑和带宽异构性纳入DP，能够选择最小化通信的并行度和阶段分配，通过拓扑自适应的平衡来缓解流水线气泡和链路瓶颈。
• 与Alpa比较: 在小规模（≤128设备）上，Alpa-E与NEST吞吐量相当。但在更大规模或通信密集型模型上，Alpa性能下降，原因在于：（1）内存可行性事后检查；（2）流水线阶段独立优化；（3）2D网格假设忽略了分层网络效应。Alpa的过度分片导致通信增加和利用率降低，而NEST通过先优化阶段内效率再复制流水线来避免此问题。NEST集成的内存模型能早期发现瓶颈并应用针对性优化，使其在内存受限的基准失败时仍能找到可行方案。
• 算子内并行 vs. 模板化并行: Alpa的细粒度算子内分片在异构集群上产生高通信开销。NEST和Phaze使用模板化并行（如张量和专家并行），通信开销较低。NEST通过将网络感知布局整合到其模板化设计中，克服了Phaze忽略网络拓扑的局限性，在各种规模的集群上均表现出色。

洞察: 拓扑自适应的分区使NEST能够均衡异构链路上的阶段延迟，维持流水线平衡并实现线性扩展，而拓扑无关的基准则会停滞。准确的内存建模使NEST能早期检测瓶颈并应用目标优化，使得在内存受限基准失败的情况下训练成为可能。

5.2.3 与微批量大小的联合探索

实验内容: 微批量大小与并行策略共同影响训练吞吐量。NEST将微批量大小和激活重计算等运行时因素纳入图提取和成本估算，实现了与子图级并行的联合优化。图6展示了在256设备集群上，NEST和各基准在不同微批量大小下的吞吐量。

结果分析: 最佳微批量大小因模型和策略而异。例如，Llama模型受益于更大的微批量，但BertLarge在NEST下几乎没有改善。微批量大小的改变会影响计算强度和内存占用，从而改变最优的并行配置。Alpa-E也受益于更大的微批量，但扫描四种批量大小需要超过80小时。相比之下，NEST在50分钟内完成同样的联合探索，速度快了90多倍，同时实现了持续更高的吞吐量，使得在规模化场景下进行全面的联合优化成为可能。

5.3 基于H100 GPU的Spine-Leaf执行

实验设置: 为了展示NEST在不同分层拓扑下的通用性，在一个基于真实拓扑的1024-GPU H100 spine-leaf集群上进行了评估。每个节点有8个H100-80GB GPU（NVLink 900 GB/s），一级交换机以12.5 GB/s连接四个节点到两个主干交换机，形成一个2:2的超额订阅拓扑。操作符运行时间在H100 GPU上进行分析，集合通信成本用AstraSim建模。本节与Mist进行比较，因为Alpa-O需要全集群访问进行分析，在此规模下不切实际。

结果分析:
* 吞吐量提升: 如图7所示，NEST的分布策略平均比手动布局吞吐量高1.47倍，比MCMC搜索高1.40倍，比Mist高1.49倍，比Phaze高1.16倍。这些增益源于其将计算阶段与可用带宽对齐的能力，从而在不同硬件和网络配置上提供了稳健和可扩展的性能。
* 网络感知的重要性: 硬件和网络的变化决定了瓶颈是计算还是通信。例如，在Mixtral模型中，受限网络上的通信可占总执行时间的10%，而在胖树拓扑上仅占1%。NEST通过在工作负载分区时明确建模拓扑和通信成本，有效平衡了计算和通信。相比之下，Phaze的网络无关策略在带宽有限时会造成严重的延迟不平衡。例如，对于BertLarge模型，Phaze选择的配置导致频繁的跨节点通信，而NEST则选择将流水线通信保持在单个节点内，从而减少了阶段间延迟变化并实现了更高的吞吞吐量。
* 与Mist的比较: NEST在所有评估模型上均实现了更高的吞吐量，这得益于其网络感知的阶段划分和布局策略。NEST倾向于使用较浅的流水线以减少流水线气泡和跨节点通信，并利用更高的数据并行度来扩展吞吐量。Mist虽支持不均匀的层划分，但因缺乏网络感知，无法考虑链路间的带宽异构性，从而无法平衡实际的通算比。此外，NEST的搜索效率显著高于Mist，在相同配置下平均快30%。

5.4 基于V100 GPU的Spine-Leaf执行

实验设置: 在真实的8-GPU和16-GPU集群（每节点2个V100，NVLink 300 GB/s，节点通过12.5 GB/s交换机连接）上验证NEST。Alpa的分析型变体（Alpa-O）在小规模上非常有效，可作为强基准。

结果分析: 在一个缩减版的Mixtral模型上，NEST在8-GPU集群上实现了与Alpa相差7%以内的吞吐量，同时将优化时间从1小时减少到5分钟。在16-GPU集群上，NEST的性能比Alpa高出1.8倍。这些结果表明，NEST在小规模真实硬件上能与高度优化的分析型策略相媲美，同时能更有效地扩展到更大的集群。

A7 补充细节

6 相关工作

设备布局框架。FlexFlow【9, Jia et al., Beyond data and model parallelism for deep neural networks, 2019】, Piper【41, Tarnawski et al., Piper: Multidimensional Planner for DNN Parallelization, 2021, NeurIPS 2021】, Pip【42, Tarnawski et al., Efficient algorithms for device placement of dnn graph operators, 2020, NIPS ’20】以及基于强化学习的方法【19, Mirhoseini et al., Device placement optimization with reinforcement learning, 2017, ICML’17】等系统使用强化学习、动态规划或随机搜索来自动化并行配置，但通常假设网络是简化的或扁平的，并忽略了拓扑异构性。Alpa【45, Zheng et al., Alpa: Automating inter- and IntraOperator parallelism for distributed deep learning, 2022, OSDI 22】通过一个两级网格和混合ILP-DP搜索增加了一定的拓扑感知能力，但缺乏对分层或超额订阅网络、集成内存建模以及专家并行的支持。最近的工作在不同维度上扩展了设计空间。Aceso【17, Liu et al., Aceso: Efficient parallel dnn training through iterative bottleneck alleviation, 2024, EuroSys ’24】通过贪婪的、瓶颈驱动的搜索强调内存可行性和利用率，但依赖于局部启发式方法，没有全局最优性保证。Mist【47, Zhu et al., Mist: Efficient distributed training of large language models via memoryparallelism co-optimization, 2025, EuroSys ’25】使用MILP来优化计算-通信重叠，但将拓扑视为次要因素，在大型分层集群上扩展性较差。其他系统提供工具而非统一优化。例如，nnScaler【16, Lin et al., nnScaler, 2024】支持手动探索扩展策略，但没有联合优化计算、内存和网络布局。

与设备布局的联合优化。Phaze【44, Wang et al., Integrated hardware architecture and device placement search, 2024】和WHAM【1, Adnan et al., Workload-aware hardware accelerator mining, 2024; 35, Phanishayee et al., Workload-aware hardware architecture recommendations, 2022】优化了部分搜索空间——包括硬件和布局——使用ILP进行算子调度，但假设网络是扁平的。TopoOpt【43, Wang et al., TopoOpt: Cooptimizing network topology and parallelization strategy for distributed training jobs, 2023, NSDI 23】通过FlexNet和基于MCMC的布局增加了拓扑感知，但缺乏最优性保证，在大型搜索空间中扩展性差，并且不支持ZeRO或专家并行。NEST通过一个用于在固定加速器架构上进行布局的动态规划框架推动了这一方向，统一了通信、内存和拓扑建模，从而在分层和超额订阅网络上实现了混合并行，同时保持了可扩展性和最优性。

网络感知的分布式机器学习。Cassini【3, Chen et al., Cassini, 2022, NSDI 22】和Themis【39, Rashidi et al., Themis, 2022, ISCA ’22】等工作研究了多租户ML集群中的网络感知调度。Cassini使用几何抽象进行作业布局，而Themis协调集合操作以减少拥塞。这些方法旨在减轻作业间的干扰，而NEST则专注于优化单个分布式训练作业。

7 局限性与讨论

灵活的网络抽象。NEST使用一种灵活的、分层级的网络局部性抽象。虽然本文主要应用于如胖树等分层拓扑，但它并不局限于此。通过将逻辑局部性与物理层次解耦，NEST的动态规划引擎可以泛化到各种互连结构。对于非分层网络（例如3D环面网络），层级可以代表基于跳数距离或坐标邻近度的亲和度类别，其成本可以从分析或实际测量的延迟中得出。这使得NEST能够通过更新成本模型来适应新的架构，而无需改变DP公式。

可扩展的并行策略集成。NEST还提供了一个结构化框架，用于集成现有和新兴的并行策略（例如序列并行），同时专注于布局优化而不是开发新的并行技术。

A5 结论

我们提出了NEST，这是第一个结构化的、感知网络、计算和内存的设备布局框架，它通过结构化动态规划统一了模型并行化和设备布局。通过对分层网络进行建模并支持多种并行策略，NEST能够在多样化的硬件上实现高效、可扩展的训练。评估结果表明，相较于最先进的基准，NEST在吞吐量、内存效率和可扩展性方面均取得了持续的提升，为协同设计并行化策略和AI数据中心基础设施奠定了基础。

A6 附录

B 分层网络拓扑

现代数据中心通常采用多级分层网络拓扑，如Spine-Leaf或Clos网络，这些拓扑在提供高聚合带宽的同时，保持了模块化和容错性。胖树（Fat-Tree）常被用作评估环境的简化抽象，但现实世界的集群由于成本、物理布局或工作负载需求，通常会实现超额订阅或非均匀版本的分层拓扑。图8展示了这些常见拓扑的代表性示例。NEST通过一个灵活的建模接口支持多种此类配置。用户可以提供一个网络描述，指定设备标识符、节点间的连接、每条链路的延迟和带宽，以及要使用的集合通信协议类型。

C 网格和环面网络建模

对于非均匀网格或环面（torus）架构，如TPU集群，NEST将物理邻近性映射到逻辑层级。如图9所示，我们根据曼哈顿跳数距离或来自源节点的异构带宽（BW）定义坐标亲和度类别。
* 层级1 (L1)：代表直接邻居（1跳）或具有最低延迟的高带宽本地单元。
* 层级2 (L2)：涵盖距离为2跳的节点或通过中等速度链路连接的节点。
* 层级3 (L3 / 远程)：将所有超过特定距离阈值或受低带宽网格间瓶颈限制的节点归为“远程”类别。

这种映射使得DP求解器能够以与处理基于树的交换机层次相同的数学严谨性来处理网格距离和非均匀链路带宽。通过根据这些性能特征对设备进行分组，布局引擎保持了拓扑无关性，同时仍然准确地尊重物理吞吐量和延迟瓶颈。

D 消融研究：ZERO优化

内存优化技术的集成。NEST集成了内存优化技术，包括ZeRO和激活重计算，以克服内存瓶颈并启用在其他情况下不可行的训练配置。当单个模型层的内存需求超过设备内存时，ZeRO的优势最为明显。利用详细的内存模型，NEST能自动检测此类约束，并根据需要应用适当的ZeRO阶段（1、2或3），从而在紧张的内存预算下生成有效且高效的训练策略。

在内存受限配置下的评估。尽管大多数评估的模型都能在现有硬件（TPUv4 64GB HBM，H100 80GB HBM）上运行，我们还是通过减少内存配置进行了消融研究，以分离ZeRO的影响。如表3所示，NEST将每个层——包括嵌入层——分配到一个单独的流水线阶段，并选择性地应用ZeRO分区，证实了如果没有ZeRO，训练将变得不可行。对于每个模型，NEST在层级别确定最优的ZeRO配置，识别出需要分区的状态（优化器、梯度或参数）以及分配的设备数量。这种细粒度的策略确保了每层都能装入内存，同时最小化了跨设备的通信。

Table 3. 在资源受限架构下评估的模型。分布式策略格式为{p,d,t}。

E 硬件验证

验证方法。我们在由4个和8个NVIDIA H100 SXM5 80GB GPU【NVIDIA, a】组成的集群上，对我们框架的集合通信估计与真实系统测量结果进行了验证。这些GPU通过一个基于交换机的拓扑连接，该拓扑包含四个NVSwitch，每个NVSwitch通过提供450GB/s单向带宽的NVLink连接八个GPU。

结果。图10比较了GPT-3在批量大小为1-4时端到端迭代时间（前向+后向）的测量结果与预测结果，显示两者之间的延迟差异最高为2%。计算延迟使用PyTorch测量，而通信延迟则通过仿真得到。

F 在512个设备上的微批量大小扩展

实验结果。图11展示了BertLarge、Llama2-7B和Llama3-70B在512设备集群上的微批量性能。趋势与256设备（第5.2节）的情况相似，最佳微批量大小取决于模型和并行化策略。Alpa受益于更大的微批量，从大小1到8性能提升高达1.5倍，但其在512设备上对BertLarge的性能因过度分片而下降。

G 用于MIST H100比较的模型配置

模型详情。表4显示了在5.3节中使用的缩减版GPT-3模型（GPT3-35B）。该模型从175B参数缩减，以便与Mist进行比较，因为Mist支持的隐藏维度小于8192，而GPT3-175B的原始隐藏维度为12288。所有实验中使用的微批量大小为1。

Table 4. 缩减版GPT3模型的配置。

H 内存估计

估计精度。NEST的每层内存建模估计值与从Alpa【45, Zheng et al., Alpa: Automating inter- and IntraOperator parallelism for distributed deep learning, 2022, OSDI 22】编译的可执行文件相比，在所有评估的模型中平均误差在7%以内。

Table 5. Alpa可执行文件与NEST估计的每层内存需求对比

I 与MIST的运行时间比较

运行时间对比。在所有探索的模型中，NEST的平均运行时间比Mist快30%以上。表6展示了每个框架在每个模型上的运行时间明细。

Table 6. NEST与基准的运行时间比较。

J 用于真实系统验证的模型配置

模型详情。表7显示了在5.4节中使用的缩减版Mixtral 8×7B模型的配置。该模型从完整的47B参数版本缩减而来，以便在资源受限的8-GPU和16-GPU V100验证集群上进行可行的分析和端到端执行。该模型总共有790M参数。实验中使用的微批量大小为1。

Table 7. 缩减版Mixtral模型的配置。

K 设备布局算法

算法伪代码。算法1展示了网络感知设备布局算法的完整伪代码。

算法1 考虑延迟和内存的设备布局优化

输入: 逐层延迟/内存估计，架构和网络配置
输出: 设备布局，t_batch
1: 初始化: dp[ℓ][id][k][s] ← (∞, ∅) // dp[层级][下集][加速器数][阶段数]
2: dp[ℓ][∅][*][0] ← (0, init) // 基准情况：没有层被分配
3: for id ∈ Downsets \ {D_full} do
4:   isZeRO ← ISZEROMEMORY(LAYERMEMREQ(id)) // 如果适用，进行ZeRO优化
5:   for subId ⊂ id, load ← id \ subId do
6:     loads ← GETLOADOFSTAGE(load, isZeRO) // 所有布局的延迟
7:     for s ∈ [1, S_max], k ∈ [1, K_max], a ∈ VALIDACCELCOUNTS(load) do
8:       prev ← dp[GETIDX(·)][subId][k−a][s−1]
9:       for each level ℓ do
10:        c ← max(prev.latency, loads[ℓ][a].latency)
11:        if c < dp[ℓ][id][k][s].first then
12:          dp[ℓ][id][k][s].first ← c // 更新最小-最大延迟
13:        end if
14:      end for
15:    end for
16:  end for
17: end for
18: t_min ← ∞ // 寻找最佳的端到端批处理时间
19: for subId ⊂ D_full, F ← D_full \ subId do
20:   load ← GETLOADOFSTAGE(F, subId) // 剩余层的延迟
21:   for s ∈ [1, S_max], k ∈ [1, K_max], a ∈ VALIDACCELCOUNTS(F) do
22:     prev ← dp[GETIDX(·)][subId][k−a][s−1]
23:     t_stage ← max(prev.latency, load[ℓ][a].latency)
24:     t_batch ← t_stage · (mbs + s − 1) + SYNCCOST(F) // 流水线 + 同步
25:     if t_batch < t_min then
26:       t_min ← t_batch
27:     end if
28:   end for
29: end for
30: return RECONSTRUCTSTAGES(dp, D_full) // 检索最终的设备布局