Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper introduces TONS, an automated network synthesis framework that optimizes datacenter network topology, routing, and collectives for AI training supercomputers like Google's TPU v4/5p. TONS formulates topology synthesis as a linear optimization problem maximizing a throughput-centric proxy metric, scaling to thousands of nodes. Evaluations show TONS networks achieve geometric mean speedups of 2.1x and 1.6x over existing TPU torus variants for uniform random and all-to-all traffic, respectively, demonstrating significant throughput improvements.
TPU supercomputers are leaving terabytes per second of throughput on the table, and TONS shows how to reclaim it with optimized network topologies.
Datacenter network design plays a critical role in AI training by supporting scaling to thousands of accelerators. An open problem, designing a near-optimal throughput oriented network-topology, routing, and collectives-has not been achieved at scale and with broad applicability to physical/implementation constraints. We address this problem with a compelling use-case, Google's TPU v4/5p supercomputer where the topology may be reconfigured to achieve higher all-to-all throughput, supporting large, parallelized AI training. We show that the existing TPU networks leave terabytes per second of throughput on the table and we fill that gap. This paper presents Throughput Optimized Networks at Scale (TONS), an automated network synthesis framework that meets the high-throughput demands of modern computing. TONS formulates topology synthesis as a linear optimization problem that maximizes a throughput-centric proxy metric, using theory and heuristics to scale to thousands of nodes. We further introduce a deadlock-free routing scheme compatible with limited virtual channels and optical switch faults, enabling the synthesized topologies to realize their predicted throughput gains in simulation. Evaluating uniform random and all-to-all traffic, TONS networks have a geometric mean speedups of 2.1x and 1.6x, respectively, over the best TPU v4/5p torus variants.