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

TL;DR

NEST is a device placement framework for distributed deep learning that jointly optimizes memory, network, and compute considerations, significantly improving throughput and scalability over existing methods.

Contribution

NEST introduces a unified, dynamic programming-based approach for device placement that considers topology, memory, and parallelism strategies simultaneously.

Findings

Up to 2.43x higher throughput compared to baselines.

Improved memory efficiency and scalability.

Effective co-design of parallelization and datacenter networks.

Abstract

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