Toward Cross-Layer Energy Optimizations in AI Systems

The “AI for Science, Energy, and Security” report [1] from DOE outlines a significant focus on developing and optimizing AI workflows for a foundational impact on a broad range of DOE missions. With the pervasive usage of artificial intelligence (AI) and machine learning (ML) tools and techniques, their energy efficiency is likely to become the gating factor toward adoption. This is because generative AI (GenAI) models are massive energy hogs: for instance, training a 200-billion parameter large language model (LLM) at Amazon is estimated to have taken 11.9 GWh [2], which is enough to power more than a thousand average U.S. households for a year [3]. Inference consumes even more energy, because a model trained once serve millions. Given this scale, high energy efficiency is key to addressing the power delivery problem of constructing and operating new supercomputers and datacenters specialized for AI workloads [4, 5].

Generality cannot be attained without it being part of the core design. To that end, we believe that we need to define a narrow interface to exchange the necessary and sufficient information between the hardware and software layers – something akin to the “narrow waist” design in other systems (e.g., IP layer in networking, LLVM Intermediate Representation in compilers) that sits in between and guides the design of both.

As the narrow waist, a possibility we propose is the Pareto frontier (tradeoff curve) of the time and energy consumption of a sequence of ML computations. The frontier contains a set of Pareto-efficient (time, energy) points a sequence of ML computations can consume for execution, and each point on the frontier is induced by running the computation with a specific combination of knobs in both software (e.g., batch size for ML training and inference, computation/memory intensity of kernel, usage of kernel fusion) and hardware (e.g., GPU power limit, frequency, model/microarchitecture) layers. The time–energy Pareto frontier has been shown to universally exist for ML computations across multiple generations of GPUs [6, 7].

Under the narrow waist design, efficiency advancements of software and hardware layers can be decoupled. The software layer can assume that the underlying hardware provides some time–energy Pareto frontier and create optimization algorithms that make use of them. For instance, when an ML serving system knows how much slack time a request has, it can select the point on the Pareto frontier that precisely slows down computation until the deadline, thereby extracting maximum energy-efficiency from the underlying hardware without missing the deadline. More challenging time and energy allocation problems will arise in more complex scenarios, including large model training and DAG-based workflows.

Conversely, hardware can conveniently assume that the software layer has algorithms to make use of any time–energy Pareto frontier. Then, new hardware architectures that push the time–energy Pareto frontier further or provide a wider dynamic range of time and energy can be developed. Furthermore, new hardware control knobs that provide better time–energy Pareto frontiers can be created, and the software layer can easily utilize them as long as the hardware exposes control to the software layer.

Finally, the time–energy Pareto frontier can be used as a basic building block of modeling system dynamics and their end-to-end energy efficiency, enabling a wide range of what-if simulations. For accurate estimations of time and energy, there is an opportunity to develop more detailed models that consider details like kernel launch overheads and data transfer/IO costs. Further, enhancing data points on this frontier with hardware-level details (e.g., compute and memory resource utilization) can provide users with valuable insights. Altogether, not only this approach allows for the optimization of existing infrastructures, but it also facilitates path finding for developing next-generation energy-efficient software/hardware AI stacks.

Generative AI adoption and its energy consumption are skyrocketing with broad implications. First, energy-intensive AI workloads inflate operational expenses and carbon offsetting costs for entities with Net Zero commitments. Second, they have made power delivery one of the primary challenges in building new AI clusters today. Finally, this can hinder deploying AI services to places without high-capacity, stable electricity grids. Only a coordinated software-hardware effort can holistically address these problems.

It is high time we invest in cross-layer energy efficiency optimizations for AI systems.