Kareus: Joint Reduction of Dynamic and Static Energy in Large Model Training

AI summary

Overview

Research area: Machine learning systems — energy-efficient large model training, GPU kernel scheduling, and frequency scaling.

Technical level: Advanced (assumes familiarity with GPU power models, pipeline parallelism, tensor/context parallelism, and Bayesian optimization).

Scope: The paper identifies that kernel launch timing, SM allocation for communication kernels, and GPU frequency jointly and interdependently determine training time and energy, and presents Kareus, a training system that jointly optimizes all three to push the time–energy tradeoff frontier.

What This Paper Is About

Training large models consumes enormous amounts of energy, and prior work attacks only part of the problem: nanobatching-style kernel scheduling reduces static energy by shortening iteration time, while frequency scaling (e.g., Perseus) reduces dynamic energy by lowering GPU frequency. The authors show these factors are interdependent — the best SM allocation and kernel launch timing change with GPU frequency — so optimizing them separately or combining them naively is suboptimal. Kareus is their system for jointly choosing SM allocation, communication launch timing, and GPU frequency across an entire training iteration.

Key Contributions

1. Characterization of joint effects. The authors show that SM allocation, communication launch timing, and GPU frequency jointly determine time and energy consumption, that the observed schedules differ by as much as 3.29× in time and energy for the same total work, and that existing solutions optimizing subsets of these factors (and naive combinations of them) are suboptimal.

2. The partitioned overlap execution model. Kareus decomposes the intractable global optimization problem into local, partition-based subproblems by grouping kernels that execute in repeating patterns into partitions, generalizing recent nanobatching techniques with fine-grained control over all three execution schedule factors.

3. A multi-pass multi-objective Bayesian optimization algorithm. Kareus uses separate surrogate models for time and dynamic energy, plus exploitation passes based on hypervolume improvement over total, dynamic, and static energy and an exploration pass based on bootstrap-ensemble uncertainty, to find each partition's time–energy frontier, which are then hierarchically composed into a microbatch frontier and an iteration-level frontier.

4. An implementation and evaluation on 14 workloads. Kareus is implemented on top of Perseus and Megatron-LM, using MSCCL++ for communication kernels with fine-grained SM control and a thermally stable profiler; it is evaluated on 14 representative workloads including real testbed training on Llama 3.2 3B and Qwen 3 1.7B and large-scale emulation on Llama 3.3 70B. The code is open-source at https://github.com/ml-energy/kareus.

Main Findings

• Joint optimization beats prior state of the art: Compared to Perseus, Kareus reduces training energy by up to 28.3% at the same training time, or reduces training time by up to 27.5% at the same energy consumption.

• Execution schedules matter enormously for the same work: Different execution schedules vary time and energy by as much as 3.29× even when the total amount of work done is identical.

• SM allocation has a sweet spot: In a single Transformer Attention layer forward pass of Llama 3.2 3B with tensor parallelism degree 4 on four A100 GPUs, allocating 2 SMs to the communication kernel left an exposed communication period during which 106 SMs were idle; 4 SMs let communication finish before computation and was energy-optimal; 20 SMs sped up communication at the cost of slowing Linear 1 and raising total time and energy.

• Launch timing interacts with memory bandwidth: Launching the communication kernel with the memory-bound Norm instead of Linear 1 caused bandwidth contention that prolonged both kernels and extended periods of compute underutilization.

• The optimal schedule depends on frequency: Lowering GPU frequency to 1,100 MHz made kernels relatively more compute-bound, so the energy-optimal schedule changed from launching communication with Linear 1 (at 1,410 MHz) to launching it with RoPE. Exposed communication is also relatively more harmful at lower frequencies because static power's share of total power draw grows.

• Opportunity even without frequency scaling: In the single-Attention-layer case study, the energy-optimal schedule reduced energy by 24.3% and time by 12.3% compared to Megatron-LM, and 7.1% energy and 5.4% time compared to Nanobatching.

• Energy breakdown of existing systems: For Qwen 3 1.7B on 16 NVIDIA A100 GPUs (8 microbatches, microbatch size 16, sequence length 4K, pipeline parallelism 2, context parallelism 2, tensor parallelism 4), Megatron-LM took 5.60 s per iteration with 5,372 J static, 21,374 J dynamic, and 26,745 J total energy; Megatron-LM + Perseus took 5.60 s with 5,374 J static, 19,531 J dynamic, and 24,905 J total; Nanobatching took 5.31 s with 5,096 J static, 21,445 J dynamic, and 26,541 J total; Nanobatching + Perseus took 5.37 s with 5,160 J static, 19,729 J dynamic, and 24,889 J total. Megatron-LM achieved 99.0 TFLOP/s/GPU.

• The search space is enormous: For a typical Transformer-based LLM on an A100 GPU the combined space comprises 85K candidate configurations, at roughly 13 seconds of profiling per candidate, so exhaustive search could take up to 4,912 GPU-hours.

• Joint control can beat hardware power control: The authors state that joint control of GPU frequency and kernel scheduling is provably more energy-efficient than leaving frequency control to the hardware's power controller (Appendix A), reasoning that since dynamic power grows roughly with the cube of frequency, high-frequency periods cost more energy than low-frequency periods save.

Methodology in Plain English

The authors first run controlled experiments on a single Transformer Attention layer, holding the total work fixed while varying three things: how many SMs the communication kernel gets, when the communication kernel is launched relative to computation kernels, and the GPU frequency. They record time and energy for each combination and plot the resulting tradeoff frontiers, which reveals that the three factors are interdependent rather than separable.

To turn this insight into something practical, they observe that training iterations repeat the same kernel patterns, so they chop the forward and backward passes into "partitions" — each consisting of one communication kernel and a long contiguous run of independent computation kernels. This shrinks the problem from one huge global search into many small, repeatable ones.

For each partition, they use multi-objective Bayesian optimization: two XGBoost surrogate models predict time and dynamic energy, and total energy is computed as predicted time times static power plus predicted dynamic energy. At each iteration they pick a batch of candidates using four acquisition signals — hypervolume improvement over total, dynamic, and static energy, plus an uncertainty score from a bootstrap ensemble of surrogates — so the frontier expands in complementary directions instead of just one.

Local partition frontiers are then combined into microbatch frontiers by enumerating frequency, SM allocation, and launch timing combinations and pruning suboptimal points, with a uniform frequency enforced across a microbatch and identical configurations shared by partitions of the same type. Iteration-level frontiers are built from microbatch frontiers using Perseus's iterative algorithm. Profiling uses Zeus (built on NVML) with a "thermally stable" profiler that repeats each partition over a 5-second measurement window to overcome NVML's roughly 100 ms sampling interval and to avoid one configuration's heat affecting later measurements.

Why This Matters

Impact on research: The paper reframes large model training energy as a joint scheduling-and-frequency problem rather than two separate problems, and provides a decomposition technique (partitions) plus a multi-objective optimizer that other systems could adopt. It also offers a concrete argument that software-level joint control can outperform hardware power management.

Real-world applications:

• Data center operators scheduling large training jobs under fixed power or energy budgets, who could trade a small amount of time for lower energy or vice versa.
• Cloud providers offering energy- or carbon-aware training services, where frontier points let users meet job-level requirements such as time deadlines or energy budgets.
• Organizations training models on constrained or contended power grids, where a single training run can consume enough energy to power more than 24,000 average US households for a month.
• Teams running long training jobs that need to adapt dynamically to changing environments such as stragglers.

Industry relevance: The mismatch between AI's computation demand and slow energy procurement — the paper cites three years for natural gas and five to ten years for nuclear, and projections that by 2035 nearly 10% of US electricity demand could come from datacenters — makes energy an expensive, contended resource that requires explicit management. The reported reductions of up to 28.3% in energy or 27.5% in time over the state of the art are directly relevant to the cost and feasibility of large-scale training.

Future Directions

• Finer-grained GPU power models. The paper adopts a simplified two-component model (dynamic power plus constant static power) and explicitly leaves extending Kareus to models that separate static power into chip leakage and constant power — where chip leakage is voltage-sensitive — as future work.

• Dynamic adaptation at runtime. The authors motivate finding an entire frontier partly by the ability to perform dynamic adaptation to changing environments such as stragglers; how Kareus would select and switch schedules online under such conditions is not fully specified in the available content.

• Scaling the profiling and optimization cost. With 85K candidate configurations per A100, roughly 13 seconds per profiled candidate, and up to 4,912 GPU-hours for exhaustive search, reducing profiling overhead and making the optimizer scale to more hardware generations and parallelism dimensions is an open engineering question.

• Broader communication and computation patterns. Kareus fuses consecutive communication kernels and groups short consecutive memory-bound operations to keep scheduling tractable, and it falls back to sequential execution when microbatch work is small. How far these generalizations extend to arbitrary models and communication schemes is left as an open area for broader deployment experience.

Target Audience

ML systems researchers and engineers working on training efficiency, GPU kernel scheduling, and energy-aware computing; practitioners who run large-scale distributed training on clusters with power or energy constraints; and graduate students studying the intersection of systems, parallelism strategies, and hardware power management. Readers need background in GPU execution models (SMs, CUDA streams, kernels), parallelism strategies such as pipeline, tensor, and context parallelism, and basic multi-objective optimization.

Authors’ abstract

The computing demand of AI is growing at an unprecedented rate, but energy supply is not keeping pace. As a result, energy has become an expensive and contended resource that requires explicit management and optimization. Although recent works have made significant progress in large model training optimization, they focus on optimizing either dynamic or static energy consumption. We find that fine-grained kernel scheduling and frequency scaling jointly and interdependently impact both dynamic and static energy consumption. Based on this finding, we design Kareus, a training system that pushes the time-energy tradeoff frontier by optimizing both aspects. Kareus decomposes the intractable joint optimization problem into local, partition-based subproblems. It then uses a multi-pass multi-objective optimization algorithm to find execution schedules that push the time-energy tradeoff frontier. Compared to the state of the art, Kareus reduces training energy by up to 28.3% at the same training time, or reduces training time by up to 27.5% at the same energy consumption.