Closing the Efficiency Gap: AI Datacenter Co-design Roadmap for Scalable Training of LLMs

Abstract

The massive compute, memory, and networking needs for LLM training necessitate a fundamental rethinking of datacenter architectures to ensure scalability, efficiency, and cost-effectiveness. In particular, the design of the network fabric for AI datacenters for emerging LLMs (such as MoEs) remains a crucial and challenging open question, spanning technology choices (that determine the size of the high-bandwidth domain), topology, and software optimizations (collective algorithms and overlap strategies). This necessitates an agile framework to traverse the co-design space. This work introduces Calculon-MoE, a tool that jointly explores FLOPS, HBM bandwidth and capacity, multiple network topologies (Two-tiered vs. FullFlat optical), the size of scale-up domain, and popular parallelism/optimization strategies used in LLMs. Our validation studies demonstrate that our LLM/MoE runtime predictions are within 10% of real-world measurements. Using Calculon-MoE, we conduct a suite of case studies to develop an actionable roadmap for data centers. For example, the results point to the promise of Fullflat network architectures, which provide uniform high-bandwidth, low-latency connectivity between all nodes and demonstrate their positive impacts on performance and scalability. We also quantify the benefits of overlapping compute and communication, hardware-accelerated collectives, widening the scale-up domain, and higher memory bandwidth and capacity. Our study spans both sparse (mixture of experts) and dense transformer-based LLMs, revealing how system design and optimization choices affect system efficiency and overall throughput in both cases.