MIT researchers developed a faster method to estimate AI power consumption during inference and training.

Due to the explosive growth of artificial intelligence, it is estimated that data centers will consume up to 12 percent of total U.S. electricity by 2028, according to the Lawrence Berkeley National Laboratory. To address this sustainability challenge, researchers from MIT and the MIT-IBM Watson AI Lab have developed a rapid prediction tool that tells data center operators how much power will be consumed by running a particular AI workload on a certain processor or AI accelerator chip. Their method produces reliable power estimates in a few seconds, unlike traditional modeling techniques that can take hours or even days to yield results. As Kyungmi Lee, an MIT postdoc and lead author of the research, explains: "The AI sustainability challenge is a pressing question we have to answer. Because our estimation method is fast, convenient, and provides direct feedback, we hope it makes algorithm developers and data center operators more likely to think about reducing energy consumption."

Traditional methods for predicting energy consumption involve breaking a workload into individual steps and emulating how each module inside the GPU is being utilized one step at a time. However, AI workloads like model training and data preprocessing are extremely large and can take hours or even days to simulate in this manner. "As an operator, if I want to compare different algorithms or configurations to find the most energy-efficient manner to proceed, if a single emulation is going to take days, that is going to become very impractical," Lee notes. To overcome this challenge, the MIT researchers discovered that AI workloads often have many repeatable patterns, and they could use these patterns to generate information needed for reliable but quick power estimation. "These optimizations that software developers use create a regular structure, and that is what we are trying to leverage," Lee explains. The researchers developed a lightweight estimation model called EnergAIzer that captures the power usage pattern of a GPU from those optimizations.

While EnergAIzer's estimation was fast, the researchers found that it didn't account for all energy costs. Every time a GPU runs a program, there is a fixed energy cost required for setting up and configuring that program, and each time the GPU runs an operation on a chunk of data, an additional energy cost must be paid. Due to fluctuations in the hardware or conflicts in accessing or moving data, a GPU might not be able to use all available bandwidth, slowing operations down and drawing more energy over time. To include these additional costs and variances, the researchers gathered real measurements from GPUs to generate correction terms they applied to their estimation model. "This way, we can get a fast estimation that is also very accurate," Lee says.

When the researchers tested EnergAIzer using real AI workload information from actual GPUs, it could estimate the power consumption with only about 8 percent error, which is comparable to traditional methods that can take hours to produce results. A user can provide their workload information, such as the AI model they want to run and the number and length of user inputs to process, and EnergAIzer will output an energy consumption estimation in a matter of seconds. The tool can also be used to predict the power consumption of future GPUs and emerging device configurations, as long as the hardware doesn't change drastically. The research was conducted by Zhiye Song, an electrical engineering and computer science graduate student; Eun Kyung Lee and Xin Zhang, research managers at IBM Research and the MIT-IBM Watson AI Lab; Tamar Eilam, IBM Fellow and chief scientist of sustainable computing at IBM Research; and senior author Anantha P. Chandrakasan, MIT provost and Vannevar Bush Professor of Electrical Engineering and Computer Science.