arXiv:2603.06798cs.LGcs.DC2026-03中稿 · MLSys 2026

NEST统一优化分布式训练中的计算、通信与内存,提升大规模模型训练效率。

NEST: Network- and Memory-Aware Device Placement For Distributed Deep Learning

  • 基于结构化动态规划,联合建模网络拓扑与设备内存
  • 在多种硬件和网络下实现最高2.43倍吞吐量提升
  • 适合大规模AI训练系统设计者与高性能计算研究人员

深度学习规模持续增长,亟需能协同考虑并行性、内存与网络拓扑的分布式训练框架。以往方法多依赖启发式或忽略拓扑的搜索策略,分离处理通信与内存问题。缺乏设备级内存感知导致参数与激活需后期分片,增加同步开销、膨胀通信量,并限制真实数据中心网络下的可扩展性与效率。本文提出NEST,一个兼顾网络、计算与内存的设备放置框架,通过结构化动态规划统一模型并行、拓扑建模与内存可行性。NEST在包含张量并行、流水线并行、数据并行与专家并行配置的算子图上运行,显式建模跨分层或任意网络的AllReduce延迟及内存/计算性能指标。通过在张量、流水线、数据与专家维度上协同分解并行策略,定义了合理的混合并行搜索空间,同时优化共置、网络延迟与内存可行性。在多样硬件与网络环境下评估表明,NEST相比现有最优基线实现最高2.43倍吞吐量提升,具备更优内存效率与可扩展性,为下一代AI基础设施的并行策略与数据中心互连的协同设计奠定基础。代码已开源:https://github.com/scai-tech/Nest

原文摘要 · Abstract (English)

The growing scale of deep learning demands distributed training frameworks that jointly reason about parallelism, memory, and network topology. Prior works often rely on heuristic or topology-agnostic search, handling communication and memory separately. Without per-device memory awareness, these methods typically ensure feasibility post hoc by sharding parameters and activations across many devices, increasing synchronization, inflating communication, and underutilizing compute-limiting scalability and efficiency on real datacenter networks. We present NEST, a network-, compute-, and memory-aware device placement framework that unifies model parallelism, topology modeling, and memory feasibility via structured dynamic programming. NEST's DP operates on operator graphs with tensor and expert parallel configurations, explicit allreduce latencies across hierarchical or arbitrary networks, and memory/compute profiles. By factoring parallelism across tensor, pipeline, data, and expert dimensions, NEST defines a principled search space for hybrid strategies while jointly optimizing co-location, network latency, and memory feasibility. Evaluations across diverse hardware and networks show NEST achieves up to 2.43 times higher throughput, better memory efficiency, and improved scalability over state-of-the-art baselines, providing a foundation for co-designing parallelization strategies and datacenter interconnects for next-generation AI infrastructure. The source code of NEST is available at: https://github.com/scai-tech/Nest

分布式训练设备放置并行策略系统优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。