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
Figures & tables
Figure 1: (Left) Nest compared to prior works across network modeling, optimality, scalability, and memory modeling axes. (Right) Comparison of placement with and without integrated memory modeling. Sharded operators are shown through patterns.
Table 1: Comparison of prior network- and memory- aware device placement works.
Figure 2: Impact of communication latency on training time across different parallelism strategies on an oversubscribed 64-GPU cluster; left bar (without) and right (with) activation recomputation.
Figure 3: Nest search space and workflow. Graph-Global strategies partition entire layers, whereas Sub-Graph strategies partition computations within individual layers. Different colors denote device assignments, and dashed boxes show ZeRO and context parallelism.
Figure 4: The forward-pass dependency challenge and Nest ’s level-wise abstraction. The cost from unassigned Layer 22 to Layer 23 is unknown but abstracted as a discrete communication level (e.g., Level 0, 1, or 2).
Hyper
Distributed Strategy {p,d,t,s,(e,c)}
Model
Model
Sequence
parameters
TMP
Expert
Context
Manual
MCMC
Alpa
Phaze
Nest
Recomputation
Parameters
Length
#L, #AH, H
Widths
Degree
Degree
vs. Stashing
Llama2 7B
7B
4096
32, 32, 4096
-
-
-
{8, 64, 1, 1}
{11, 46, 1, 1}
{32, -, -, -}
{6, 85, 1, 1}
{8, 64, 1, 1}
Recomputation
Llama3 70B
70B
4096
80, 64, 8192
-
-
-
{80, 6, 1, 1}
{80, 6, 1, 1}
{52, -, -, -}
{41, 12, 1, 1}
{81, 6, 1, 1}
Recomputation
Tensor Model Parallel Models
-
-
BertLarge
350M
512
24, 16, 1024
1,2,4,8
-
-
{8, 64, 1, 1}
{2, 256, 1, 1}
{24, -, -, -}
{1, 512, 1, 1}
{1, 512, 1, 1}
Stashing
Table 2: Model workloads evaluated using Nest . Hyperparameters: number of layers (#L), attention heads (#AH), and hidden size (H). Distributed strategies for 512 devices specify pipeline depth (p), data-parallel width (d), tensor model-parallel width (t), sequence parallel width (s), expert degree (e), and context-parallel degree (c), formatted as {p, d, t, s, (e,c)}. Sequence-parallel width (s), if applied, equals tensor model-parallel width (t), as both are partitioned across the same set of devices. For Alpa, only the number of pipelines per stage is listed, as it performs intra-operator sharding within each stage.
Figure 5: Throughput comparison between Nest and baselines on a Fat-Tree network of TPUv4 accelerators. Throughput improvements are relative to the manual baseline’s smallest valid result. “X” indicates cases where the baseline failed to find a valid placement.
Figure 6: Throughput improvement relative to a manual baseline (microbatch size 1). “X” marks baseline failures due to memory constraints. High memory requirements limit Llama2-7B and Llama3-70B to microbatch sizes 4 and 2, respectively.
Figure 7: Throughput comparison between Nest and baselines on a Spine-Leaf network of 1024 H100 GPUs. Throughput improvements are relative to the manual baseline. Mist does not support GPT3-175B and Mixtral-8x7B, as indicated with an ”X”.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Hierarchical network topology. Black links represent InfiniBand connections, while red links represent NVLink connections.
Figure 9: Level-Wise Abstraction for Mesh/Torus Topologies. Nest maps physical hop distances to discrete communication levels ( l ) relative to a reference source node. This abstraction preserves optimal substructure in the DP by grouping nodes into affinity classes, where each class corresponds to a modeled communication latency or bandwidth constraint.
Parameter
Value
Number of Layers
64
Hidden Dimension
8192
Number of Heads
64
Intermediate Dimension
16384
Sequence Length
2048
Appendix
Table 3: Configuration of the scaled-down GPT3 model.
Model
Baseline (min)
Nest (min)
Time Reduction (%)
GPT-3 35B
17
15
11.8%
Llama3 70B
30
6
80.0%
Llama2 7B
8
2.3
71.3%
BertLarge
3
2
33.3%
Appendix
Table 4: Runtime comparison of Nest against baseline.
Parameter
Value
Number of Layers
8
Number of Experts
8
Hidden Dimension
1024
Number of Heads
16
Intermediate Dimension
3584
Sequence Length
1024
Appendix
Table 5: Configuration of the scaled-down Mixtral model.
Model
Alpa Executables (GB)
Nest Estimates (GB)
GPT-3 175B
10.1
9.7
Llama3 70B
24.8
22.9
Llama2 7B
9.8
8.1
BertLarge
0.21
0.21
Appendix
Table 6: Per layer Memory requirement estimations between profiled Alpa executables and Nest estimations
Figure 10: Validation of Collective Communication Estimations Won et al. (2023) against H100 GPU nodes
Model
Llama3 70B
BertLarge
HBM
24GB
120MB
Number of Device Used
641
980
Distributed Strategy
{81, 1, 1}
{25, 10, 1}
ZeRO
Layer 0 (embedding)
None
ZeRO-2: (degree 2)
Remaining Layers
ZeRO-3: (degree 8)
ZeRO-3: (degree 4)
Appendix
Table 7: Models evaluated with resource-constrained architectures. Distributed strategy formatted as {p,d,t}
Figure 11: Throughput improvement relative to manual strategy with microbatch size 1 across different parallelization strategies. “X” indicates cases where the baseline fails to find a feasible placement within memory constraints. Llama2-7B and Llama3-70B are realizable only up to batch size 4 and 2 respectively due to high memory requirements.