💻 computer science

Adaptive Data Partitioning for Energy-Efficient Federated and Distributed Learning on Heterogeneous Systems

This paper proposes a measurement-driven adaptive data partitioning controller that dynamically reallocates sample budgets based on real-time training time and energy metrics to mitigate stragglers, reduce energy consumption, and improve training efficiency across heterogeneous distributed and federated learning systems.

Original authors: Daniel Suárez Labena, Vicente José Blanco Pérez, Pedro Antonio Toledo Delgado, Francisco Carmelo Almeida Rodríguez

Published 2026-08-10

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Daniel Suárez Labena, Vicente José Blanco Pérez, Pedro Antonio Toledo Delgado, Francisco Carmelo Almeida Rodríguez

Original paper licensed under CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a classroom where a group of students is trying to solve a massive, complex puzzle together. In a perfect world, everyone would work at the exact same speed, finish their piece at the same time, and hand it in to the teacher simultaneously. But in the real world, some students have super-fast brains, others have slower processors, some are tired, and others are distracted. This is the daily reality of distributed learning, a field where computers (or devices) team up to train artificial intelligence models. Instead of one giant supercomputer doing all the work, we use many smaller devices—like smartphones, tablets, or specialized chips—working in parallel.

However, there's a catch: the whole group has to wait for the slowest student before they can move to the next step. If one student is slow, the fast ones sit around twiddling their thumbs, wasting time and battery power. This is called the "straggler effect." It's like a relay race where the fastest runners are forced to stop and wait at the finish line for the slowest runner to catch up. The goal of researchers in this field is to figure out how to keep everyone busy, finish the race faster, and save energy, even when the team is made up of very different kinds of runners.

The Problem: The "One-Size-Fits-All" Trap

The researchers from the University of La Laguna noticed that most systems treat all devices the same. They say, "Okay, we have 1,000 puzzle pieces; let's give 100 to each of the 10 students." This seems fair, but it's actually inefficient. If Student A has a super-fast computer, they'll finish their 100 pieces in a minute and then wait 10 minutes for Student B, who has a slower device, to finish. During those 10 minutes, Student A's computer is still humming along, burning electricity, just waiting.

The paper asks a simple question: What if we didn't give everyone the same amount of work? What if we gave the fast students more pieces and the slow students fewer, so everyone finishes at roughly the same time? And better yet, what if we also considered which students use the least battery power per piece?

The Solution: The Smart Coach

The authors built a "Smart Coach" (a measurement-driven controller) that watches the students as they work. It doesn't guess who is fast or slow; it actually measures them.

Here is how the coach works:

Round 1: Everyone gets an equal share of the puzzle.
The Check-in: At the end of the round, the coach asks, "How long did it take you?" and "How much battery did you use?"
The Adjustment: For the next round, the coach redistributes the work. If a device was fast and efficient, the coach says, "Great job! Here is more work for you." If a device was slow or drained its battery quickly, the coach says, "Take it easy, here is less work."

This happens at the boundary of every "round" of training. The coach is clever enough to know the rules of the game. In some games (like Federated Learning), the data belongs to specific people and can't be moved around. In those cases, the coach can only shuffle work between the devices inside that person's group. In other games where everyone shares the data, the coach can move work freely between any device.

The Results: Speed and Savings

The researchers tested this "Smart Coach" on a playground of 11 different devices, ranging from powerful computer chips to tiny, low-power boards. They compared their method against the old "Equal Split" method and a "Static Profile" method (where the coach measures you once at the start and never changes the plan).

What they found:

Faster Finishes: By giving more work to the fast devices and less to the slow ones, the whole group finished the training much faster. In some tests, the time it took to finish was reduced by over 70% compared to the equal split.
Energy Savings: Because the fast devices weren't sitting around waiting, the total energy used by the group dropped significantly. In some cases, energy use was cut by more than 40%.
The "Energy Weight" Surprise: The coach had a special setting called an "energy weight." The researchers thought, "If we tell the coach to prioritize saving energy, it will be even better." But they found something tricky. If the coach gave too much work to a device that was energy-efficient but very slow, the whole group had to wait longer, and the total energy actually went up. It turns out, the best balance depends on the specific task. Sometimes, focusing purely on speed (Time Only) was actually the best way to save energy because it got the job done so quickly that the devices could turn off sooner.

What They Didn't Do (and Why It Matters)

It's important to note what this paper didn't do. The researchers didn't change the AI model itself, they didn't change the math the AI uses to learn, and they didn't kick slow devices out of the game. They kept everyone in the room. They also didn't use a "magic" prediction that guesses the future; they only used what they could measure right now.

They also showed that simply picking the "fastest" devices and ignoring the slow ones (a common strategy in some systems) isn't always the answer if you want to use all your available hardware. Their method keeps everyone participating but adjusts the workload so no one is left waiting in the cold.

The Takeaway

This paper proves that you don't need to upgrade your hardware to make your AI training faster and greener. You just need a smarter way to hand out the work. By watching how fast and efficient each device is in real-time and adjusting the workload accordingly, you can stop the "stragglers" from holding up the race. It's a bit like a coach realizing that the sprinters should run more laps while the joggers run fewer, so everyone crosses the finish line together, tired but happy, without anyone wasting energy standing around.

Technical Summary: Adaptive Data Partitioning for Energy-Efficient Federated and Distributed Learning on Heterogeneous Systems

1. Problem Statement

Distributed model training on heterogeneous systems is significantly hindered by variations in processing throughput, memory behavior, and operating conditions among participating devices. When a fixed, equal division of training data is applied, faster devices finish their work early and remain idle while waiting for slower "straggler" devices at synchronization points. This straggler effect not only increases the total training makespan but also raises device energy consumption, as synchronized participants must remain active until collective progress can continue.

Existing approaches to mitigate this issue often rely on participant selection (excluding slow devices), asynchronous aggregation, changes to optimization methods, or model adaptation. While recent systems have explored redistributing training data, many depend on learned controllers, prior search, specific hardware settings, or alterations to the synchronization semantics and participation sets. There remains a gap for a common, closed-loop mechanism that can operate across different synchronization and data-ownership boundaries (such as Federated Learning and Distributed Data Parallelism) while retaining the original learning configuration, model, and synchronization rules.

2. Methodology

The authors propose a lightweight, measurement-driven allocation controller that reallocates a fixed sample budget at round boundaries based on observed training time and measured device energy per sample. The controller operates without modifying the model, loss function, optimizer, or the set of participating devices.

Execution Regimes

The study evaluates the controller across three distinct execution regimes:

Periodic Model Averaging (PMA): Workers train independent model replicas on shared data within a trusted domain and synchronize only at round boundaries.
Distributed Data Parallelism (DDP): Ranks synchronize gradients after every optimizer step. This regime requires a specific batch-sizing strategy to handle unequal per-rank workloads without introducing padding or dummy batches.
Hierarchical Federated Learning with intra-silo DDP (HFL-DDP): Devices are grouped into fixed "silos" (data-owning domains). Ranks within a silo use DDP, while silo models are aggregated at the federated round boundary. Reassignment is restricted to ranks within the same silo to respect data ownership.

The Adaptive Controller

The controller functions as follows:

Measurement: At the end of each round $r$ , each participant $i$ reports its processing time ( $t_i^{(r)}$ ), energy consumption ( $e_i^{(r)}$ ), and sample count ( $n_i^{(r)}$ ).
Cost Calculation: Per-sample costs are derived as $\tau_i^{(r)} = t_i^{(r)} / n_i^{(r)}$ and $\epsilon_i^{(r)} = e_i^{(r)} / n_i^{(r)}$ . These are smoothed using an Exponential Moving Average (EMA) to reduce noise.
Weighted Proportional Split: The controller computes a dimensionless combined cost $c_i^{(r)}$ $c_{i}^{(r)}$ using user-selected weights $w_t$ $w_{t}$ (time) and $w_e$ $w_{e}$ (energy).
- $c_i^{(r)} = w_t \frac{\hat{\tau}_i^{(r)}}{\text{median}(\hat{\tau})} + w_e \frac{\hat{\epsilon}_i^{(r)}}{\text{median}(\hat{\epsilon})}$
Allocation: Participants are assigned samples inversely proportional to their cost. A "deadband" mechanism prevents unnecessary reallocation if cost differences are negligible.
DDP Specifics: For DDP, the controller ensures all ranks execute the same number of optimizer steps per round by constructing per-step batch-size schedules that redistribute local samples, preserving the sample-mean gradient objective.

Experimental Setup

The evaluation was conducted on an eleven-device heterogeneous System-on-Chip (SoC) testbed featuring Arm, RISC-V, and x86 architectures. Three supervised learning tasks were used: MNIST (image classification), ERA5 (time-series regression), and IMDB (text sentiment classification). The study compared the adaptive policies against:

Equal allocation: Uniform sample division.
Static profiled: A one-time measurement followed by fixed allocation.
Oort-like selection: A baseline that excludes slow participants.
FedAsync: Asynchronous PMA.
EMA batch sizing: A step-level throughput baseline for DDP.

3. Key Contributions

The paper presents four primary contributions:

A Measurement-Driven Controller: A lightweight mechanism that reallocates a fixed sample budget based on observed time and energy per sample, preserving the full participation set and synchronization semantics.
Unified Design: A common allocation design applicable to PMA, DDP, and HFL-DDP, with specific constraints for data-ownership boundaries in hierarchical settings.
Energy-Aware Objective: An objective function with tunable weights that exposes the interaction between per-sample energy, training makespan, and synchronized device energy, challenging the assumption that maximizing energy weight always minimizes total energy.
Comprehensive Evaluation: A replicated study across three learning tasks and three synchronization regimes on a heterogeneous testbed, connecting allocation changes to compute imbalance, synchronization exposure, makespan, and measured energy.

4. Results

Experiments on the heterogeneous testbed yielded the following findings:

Reduction in Makespan and Energy: Adaptive partitioning (specifically "Time only" and "Low energy" policies) consistently reduced training makespan and measured device energy compared to Equal allocation. In PMA, "Time only" reduced median makespan by approximately 71–80% and energy by 28–40% across tasks.
Straggler Mitigation: The controller significantly reduced compute imbalance. In DDP, the coefficient of variation for forward-pass time dropped from ~0.8–1.1 (Equal) to ~0.02–0.18 (adaptive), directly correlating with reduced synchronization waiting.
Convergence: Adaptive policies reached low-objective regions (held-out loss/MSE) faster than Equal allocation, often achieving comparable final quality in significantly less time.
Non-Monotonic Energy Response: Increasing the energy weight ( $w_e$ ) did not always reduce total energy. Assigning more work to energy-efficient but slower devices could extend the synchronized execution window, increasing the accumulated energy of all participants. A small energy weight ( $w_e = 0.25$ ) often provided a better trade-off than "Energy only" ( $w_e = 1$ ).
Overhead: The scheduling computation itself was negligible, consuming less than 0.02% of the round duration.

5. Significance and Claims

The paper claims that measurement-based data partitioning provides a practical method for mitigating stragglers across distributed training regimes without altering the learning configuration or synchronization rules.

Practicality: The method requires no prior profiling, learned policies, or changes to the optimizer or model architecture. It operates as a closed-loop controller using runtime telemetry.
Context-Dependent Optimization: The authors emphasize that the optimal balance between time and energy depends on the specific task and execution regime. While "Time only" is the most reliable default for minimizing makespan, a small energy weight can be beneficial in specific contexts (e.g., HFL-DDP on certain tasks).
Distinction from Prior Work: Unlike participant selection methods (e.g., Oort) that exclude devices, this approach retains all provisioned participants. Unlike asynchronous methods, it preserves the synchronization semantics of the target regime.
Limitations: The authors modestly note that the results are specific to the evaluated eleven-device testbed and fixed network conditions. They do not claim universal superiority over all existing methods but demonstrate that a simple, measurement-driven reallocation of a fixed budget can outperform static equal splits and one-time profiling in heterogeneous environments.