FedAvg, FedProx, FedAdam, SCAFFOLD: what actually breaks in production federated learning

Federated learning papers describe algorithms in their idealized form: synchronous rounds, balanced clients, IID data, no dropouts. Production federated learning runs nothing like that. Clients arrive late, leave mid-round, have wildly different data distributions, and lie about their compute capacity.

SolanaLM ships four federated learning algorithms: FedAvg, FedProx, FedAdam, and SCAFFOLD. Each one was added because the previous ones broke on a workload we cared about. This post walks through what each algorithm is actually for, where it breaks, and how the runtime picks between them.

FedAvg: the default that is wrong more often than you think

FedAvg (McMahan et al., 2017) is the foundational algorithm: clients train locally for some number of epochs, then the server averages their weights. It is also the algorithm everyone reaches for first because it is simple and most tutorials describe it.

It works when:

• Client data is approximately IID.
• Clients have similar compute and bandwidth.
• The number of clients per round is large.

It breaks when:

• Clients have non-IID data — say, each client has a skewed distribution. The averaged update drifts toward an objective none of the clients actually want.
• Some clients have much more data than others. Naive averaging weighs them equally; weighted averaging weighs them too much.
• Clients drop out mid-round. The average over the survivors is biased toward who finished, not toward the population.

SolanaLM ships FedAvg as the default for simple, low-stakes coordination — fine-tuning a small adapter, for example, where the model is robust to update noise. The runtime handles weighted averaging by client data size and survival weighting for partial completion. The aggregation step looks like:

def fedavg_aggregate(
    client_updates: dict[str, ClientUpdate],
    weights: dict[str, float],
) -> ModelState:
    total_weight = sum(weights.values())
    aggregated = {
        name: torch.zeros_like(param)
        for name, param in next(iter(client_updates.values())).state.items()
    }
    for client_id, update in client_updates.items():
        w = weights[client_id] / total_weight
        for name, param in update.state.items():
            aggregated[name] += w * param
    return ModelState(state=aggregated)

That code looks innocent. The bugs hide in weights. Get the weighting wrong on non-IID data and the global model degrades each round.

FedProx: when clients are heterogeneous

FedProx (Li et al., 2018) adds a proximal term to the local objective. Each client trains against a regularized loss that penalizes drift from the global model. The server then averages as in FedAvg.

This helps when clients are heterogeneous in compute. A slow client that only completes one local epoch produces a less-drifted update than one that completes ten epochs. The proximal term encourages updates that stay near the global model, so the gap between a slow client and a fast client is smaller.

In SolanaLM, the proximal coefficient (commonly written mu) is configurable per round. Higher mu = more pull toward the global model = more robust to heterogeneity but slower convergence. We default to mu=0.01 and recommend tuning between 0.001 and 0.1 based on observed client variance.

def fedprox_local_loss(
    local_loss: torch.Tensor,
    local_params: dict[str, torch.Tensor],
    global_params: dict[str, torch.Tensor],
    mu: float,
) -> torch.Tensor:
    proximal_term = 0
    for name in local_params:
        proximal_term += ((local_params[name] - global_params[name]) ** 2).sum()
    return local_loss + (mu / 2) * proximal_term

Where FedProx breaks: if the data distribution shift between clients is large, the proximal term holds everyone too close to a global model that is wrong for any of them. You converge slowly to a bad answer instead of quickly to a bad answer.

FedAdam: adaptive server optimization

FedAvg averages updates. FedAdam (Reddi et al., 2020) treats the averaged update as a gradient and applies Adam on the server side. This decouples client-side optimization from server-side aggregation.

The benefit is meaningful: if individual clients are bouncing around their local minima, the server-side Adam smooths the trajectory of the global model. We see consistently better convergence on heterogeneous fine-tuning tasks.

The cost is operational. The server now has state (Adam’s m and v moments) that has to persist across rounds. If your registry restarts, you either lose that state or you need to checkpoint it carefully. SolanaLM persists optimizer state in Postgres alongside the model checkpoint.

class FedAdam:
    def __init__(self, lr=0.01, beta1=0.9, beta2=0.99, eps=1e-3):
        self.lr, self.b1, self.b2, self.eps = lr, beta1, beta2, eps
        self.m, self.v, self.t = {}, {}, 0

    def step(self, global_params, aggregated_update):
        self.t += 1
        new_params = {}
        for name, param in global_params.items():
            g = aggregated_update[name] - param
            self.m[name] = self.b1 * self.m.get(name, 0) + (1 - self.b1) * g
            self.v[name] = self.b2 * self.v.get(name, 0) + (1 - self.b2) * (g ** 2)
            m_hat = self.m[name] / (1 - self.b1 ** self.t)
            v_hat = self.v[name] / (1 - self.b2 ** self.t)
            new_params[name] = param - self.lr * m_hat / (torch.sqrt(v_hat) + self.eps)
        return new_params

SCAFFOLD: when non-IID is bad enough that nothing else converges

SCAFFOLD (Karimireddy et al., 2020) addresses “client drift” directly. Each client maintains a control variate that approximates the difference between its local gradient and the global gradient. Both client-side updates and server-side aggregation use these control variates to correct for drift.

This is the algorithm to reach for when client data is so non-IID that FedAvg and FedProx visibly diverge. You pay for it in communication cost (control variates double the upload size) and in implementation complexity (the client needs to track its own state across rounds).

SolanaLM uses SCAFFOLD as the heavyweight option when round metrics indicate severe drift. The runtime monitors the angle between local updates and the global average; when that angle exceeds a threshold, we promote the next round to SCAFFOLD automatically and notify the operator.

How we pick

The runtime exposes algorithm selection per-round, but it also provides an auto-select mode that runs heuristics on the first few rounds and picks an algorithm based on observed behavior:

1. Start with FedAvg.
2. After three rounds, compute per-client gradient variance. If clients are heterogeneous in compute (round latency variance > 2x), upgrade to FedProx.
3. If global loss is not decreasing monotonically after five rounds, upgrade to FedAdam.
4. If per-client update angles diverge (cosine similarity below 0.3), upgrade to SCAFFOLD.

The auto-select logic is opinionated; we expose every parameter so you can override it.

What is not in the box

We do not ship Federated Averaging with Momentum (FedAvgM), MIME, or FedYogi. They are useful but our user research did not surface workloads that hit their sweet spot. If you need one of them, the algorithm interface is small — about 80 lines — and contributions are welcome.

We also do not yet support partial model federation (only training some layers per client). That is on the roadmap; it matters most for very large models where bandwidth is the bottleneck.

Privacy is orthogonal

All four algorithms support our differential privacy and secure aggregation layers. You can run SCAFFOLD with (epsilon=1.0, delta=1e-5) and secure_agg=True without changing the algorithm code. The privacy layer wraps the local update with noise and the aggregation layer with cryptographic masking.

A common mistake is to assume differential privacy on its own gives you secure aggregation. It does not. Differential privacy keeps any individual client’s contribution unidentifiable in the aggregate; secure aggregation keeps the server from seeing individual updates at all. Use both for sensitive workloads.

When federated learning is the wrong answer

We end every FL conversation with this caveat: if your data can legally and operationally be centralized, centralize it. FL is for cases where the data genuinely cannot leave its owner’s environment. The complexity overhead is real. Run a centralized training job if you can.