AI training energy reduced up to 50% with U-M software

Key takeaways

• NVIDIA is testing how two energy-saving tools developed at University of Michigan Engineering could improve its open-source software for training large language models, Megatron-LM. 
• The two software tools, called Kareus and Perseus, work together to reduce the energy needed to train AI by up to 50%. Kareus was presented at the USENIX Symposium on Operating Systems Design and Implementation in July.
• Perseus manages processor speed, cutting the energy needed to train AI by up to 30%. Kareus reduces the training energy by an additional 15% to 30% by finding the most efficient timing and resource sharing for training and chip-to-chip communication tasks.

---

Around half of the energy used to train AI can be saved without impacting speed or performance, using software developed at the University of Michigan Engineering. Now, U-M researchers are working with NVIDIA to integrate their new approach into Megatron-LM, NVIDIA’s open-source platform for training AI.

Companies can use Megatron-LM directly or as a foundation to build private, customized training platforms. Megatron-LM is used to train NVIDIA’s flagship Nemotron models, which users can download onto their own hardware.

By improving training efficiency, the U-M research team hopes to enable AI to do more useful work for less energy and cost. The savings come from finding the most efficient way to coordinate AI training across hundreds to thousands of computer processors, as well as how and when different jobs are performed.

“AI is becoming power-constrained as models become larger and more people want to train models, but we can achieve the same training speed and performance using less energy,” said Ruofan Wu, a doctoral student in computer science and engineering at U-M who spent her summer integrating the U-M software into Megatron-LM.

“There are a lot of opportunities to optimize AI energy use at the software and hardware level, and one of the most rewarding parts of the internship was the opportunity to bring ideas from my research into production systems,” she added.

Co-optimizing processor speed and training task scheduling

The U-M system operates in two parts. Perseus, which slows down processors running unnecessarily fast, reduces training energy up to 30%, according to a 2024 study. Wu led the development of Kareus, which reduced the energy needed to train two leading open-source large language models by another 10-30% in a study published in July. The researchers say the total energy savings could approach 40%-50% for large-scale AI models.

Together, both tools address inefficiencies in parallel processing, a computing technique often used to quickly train massive AI models. The method divides the model and training data across hundreds to thousands of computer processors, called GPUs, that work simultaneously. Platforms like Megatron-LM find the best way to divide the work and ensure the resulting output is combined correctly.

Each GPU completes its assignment by dividing the job into smaller subtasks, which can include data processing and chip-to-chip communication. Each subtask is handled by a processor within the GPU, called a core. The cores complete their tasks in sequence or concurrently by running programs, called kernels, in a schedule optimized for the GPU’s max processing speed. 

Normally, the GPUs rush to complete their assignments as quickly as possible. This wastes energy because one GPU must process more data than the rest, leaving the other GPUs to idle until the last GPU finishes. Instead, Perseus slows the other cores down so that they all finish simultaneously. These cores save energy by jogging rather than sprinting, without increasing the overall training time.

However, changing the processing speed can introduce memory and data bandwidth issues or create new, unexploited opportunities to run kernels simultaneously. Kareus corrects such inefficiencies by coordinating the cores inside each GPU and implementing the ideal kernel schedule. For example, it may adjust the number of cores allocated for communication tasks or the timing of those tasks.

“The processor speed changes what code should have been generated in the first place,” said Mosharaf Chowdhury, an associate professor of computer science and engineering at U-M. “It is like picking a route by distance alone, without knowing the speed limits along the way. You get to your destination, but you burn more fuel than you need to.”

“The principles of training are the same, regardless of whether you are using Megatron or a different framework, so our approach could be applied to other training platforms too,” he said.

The research was also funded in part by Cisco, Ford, the Mozilla Foundation and the Laude Institute.

Related Stories