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

Problem Statement

The rapid growth of AI computing demand has outpaced energy supply, making energy an expensive and contended resource for large model training. Current optimization efforts typically address either dynamic energy (consumed by active computation, reducible via frequency scaling) or static energy (consumed by hardware regardless of activity, reducible by shortening execution time).

Existing solutions suffer from a critical limitation: they optimize these factors in isolation or via naive combinations.

• Frequency scaling approaches (e.g., Perseus) reduce dynamic energy by lowering GPU frequencies for non-critical paths but often follow sequential kernel execution, missing opportunities to overlap communication and computation.
• Fine-grained kernel scheduling approaches (e.g., Nanobatching) reduce static energy by overlapping communication and computation but ignore dynamic energy and frequency scaling.

The paper identifies that SM allocation, kernel launch timing, and GPU frequency are interdependent. Changing one factor shifts the optimal configuration of the others. For instance, lowering frequency makes kernels relatively more compute-bound, altering the optimal resource contention and overlap strategies. Consequently, optimizing subsets of these factors or combining them naively fails to push the time–energy tradeoff frontier to its theoretical limit.

Methodology

The authors propose Kareus, a training system designed to jointly optimize SM allocation, kernel launch timing, and GPU frequency to find the Pareto-optimal time–energy frontier.

1. Partitioned Overlap Execution Model

To make the joint optimization tractable, Kareus decomposes the global problem into local subproblems using a partitioned overlap execution model:

• Partitioning: The execution graph of forward and backward passes is divided into fixed partitions. A partition consists of one communication kernel (e.g., AllReduce) and the longest contiguous sequence of computation kernels from a different microbatch that has no data dependencies with the communication kernel.
• Generalization: This model generalizes nanobatching techniques, adding fine-grained control over SM allocation and launch timing for every partition.
• Hierarchical Composition: The system optimizes each partition independently and then hierarchically composes these local frontiers into a global microbatch and iteration-level frontier.

2. Multi-Pass Multi-Objective Bayesian Optimization (MBO)

Since exhaustive search of the configuration space (involving SM counts, launch timings, and frequencies) is prohibitively expensive (requiring thousands of GPU-hours and causing thermal interference), Kareus employs a specialized MBO algorithm:

• Surrogate Models: It uses two separate Gradient-Boosted Decision Tree (XGBoost) models to predict execution time and dynamic energy, respectively.
• Multi-Pass Candidate Selection: To explore the Pareto frontier effectively, the algorithm uses multiple acquisition functions in parallel:
  • Hypervolume Improvement (HVI): Three distinct HVI functions guide the search toward lower total energy, lower dynamic energy (favoring lower frequencies), and lower static energy (favoring faster execution).
  • Uncertainty Exploration: A bootstrap ensemble of surrogate models estimates prediction uncertainty, prioritizing candidates in under-explored regions of the search space.
• Thermally Stable Profiling: To ensure accurate energy measurements, the system executes partitions over a 5-second window and inserts a 5-second cooldown period between configurations to prevent thermal bias.

3. System Implementation

Kareus is implemented by integrating with Megatron-LM and Perseus.

• It uses MSCCL++ for fine-grained SM allocation control.
• It employs a custom torch.autograd.Function to wrap partitions, allowing distinct schedules for forward and backward passes.
• At runtime, it selects an execution schedule from the pre-computed frontier based on user constraints (e.g., time deadline or energy budget) and dynamically adjusts GPU frequencies via a controller adapted from Perseus.

Key Results

Kareus was evaluated on 14 workloads, including real training on Llama 3.2 3B and Qwen 3 1.7B, and large-scale emulation of Llama 3.3 70B on 16 NVIDIA A100 GPUs.

• Time-Energy Tradeoff: Compared to the state-of-the-art baseline (Megatron-LM + Perseus), Kareus achieves:
  • Up to 28.3% reduction in training energy at the same training time.
  • Up to 27.5% reduction in training time at the same energy consumption.
• Max-Throughput Performance: In maximum-throughput mode, Kareus reduces iteration time by up to 14.9% and energy by up to 22.1% compared to the baseline Megatron-LM.
• Ablation Studies: Removing either frequency scaling or kernel scheduling results in significant performance degradation (e.g., removing kernel scheduling leads to a 10.8% energy increase), confirming the necessity of joint optimization.
• Scalability: Large-scale emulation results for Llama 3.3 70B show consistent trends with physical testbed results, demonstrating the system's applicability to massive models.

Significance and Claims

The paper claims that Kareus pushes the time–energy tradeoff frontier for large model training by demonstrating that dynamic and static energy cannot be optimized independently. The core contribution is the insight that execution schedules (SM allocation and timing) and GPU frequency are jointly and interdependently determined.

The authors emphasize that:

1. Joint Control is Superior: Jointly optimizing frequency and scheduling is provably more energy-efficient than leaving frequency control to hardware power controllers or optimizing scheduling without frequency awareness.
2. Interdependence is Key: The optimal SM allocation and launch timing are frequency-specific; a schedule optimal at one frequency may be suboptimal at another.
3. Practical Efficiency: By decomposing the problem into partitions and using MBO, Kareus achieves these gains without the prohibitive cost of exhaustive search, making it a practical solution for current and future large-scale AI training workloads.