Cracking the Code of Distributed AI: Moving from FedSGD to FedAvg in Federated Learning

In traditional Machine Learning, data is king, and it all lives in one place—a massive centralized data center or cloud server. But in the modern privacy-first era, centralized training is facing severe bottlenecks due to data privacy regulations (like GDPR) and the sheer logistics of moving terabytes of data from edge devices to the cloud.

Enter Federated Learning (FL): a decentralized machine learning framework where you train models directly on edge devices (like smartphones, IoT devices, or local hospital servers) without ever moving the raw data.
However, FL introduces a massive bottleneck of its own: Network Communication Costs.
In this post, we’ll dive deep into the foundational mechanics of Federated Learning, breaking down how early frameworks evolved from the highly inefficient FedSGD to the industry-standard FedAvg (Federated Averaging), and how you can implement this mindset in your own research.

The Mathematical Objective of Federated Learning
Before looking at the algorithms, let's understand what we are trying to solve. In a standard centralized setup, our goal is to minimize a global loss function f(w) across all available data points (n):

In Federated Learning, this data is shattered across K distinct clients. Each client k holds a local dataset P
k with a size of n k​. Therefore, we rewrite the global objective function as a weighted average of local loss functions
​
This means the central server’s job is to orchestrate the training so that the global model performs optimally across all clients, even though the server never sees the data points (i) directly.
FedSGD: The Chatty Predecessor

The earliest instinct in distributed optimization was FedSGD (Federated Stochastic Gradient Descent).
How FedSGD Works:
The central server broadcasts the current global model weights (w t) to all clients.
Each client k takes a tiny slice of local data and computes the gradient (the direction to reduce error), denoted as The clients immediately upload these gradients to the server.
The server aggregates them and updates the global model:

The Fatal Flaw of FedSGD
FedSGD requires clients to communicate with the server after every single local gradient step. If a neural network has hundreds of millions of parameters, uploading gradients constantly over cellular networks or weak Wi-Fi will instantly crash the network bandwidth. It is highly inefficient, battery-draining, and completely impractical for real-world edge devices.
FedAvg: Shifting the Heavy Lifting to the Edge
To fix this, Google researchers introduced FedAvg (Federated Averaging). The breakthrough realization was simple: What if we let clients do more computation locally before making them report back to the server?
Instead of sending raw gradients after one step, a client takes the global model and runs multiple local epochs over its local dataset using its own processing power (CPU/GPU). The client actually updates its local weights multiple times:
Once the client completes its local training iterations and makes the local model "smarter," it uploads the final model weights back to the server. The server then takes a weighted average of these local models to produce the new global model:

By increasing local computation, FedAvg reduces the required communication rounds by orders of magnitude.
The Three Knobs of FedAvg

The exact amount of local computation vs. communication frequency in FedAvg is controlled by three vital hyperparameters:

C (Fraction of Clients): The percentage of total clients chosen randomly to participate in a given training round. (e.g., C=0.1 means only 10% of devices work per round, saving the battery of the remaining 90%).
E (Local Epochs): The number of times a client loops through its entire local dataset during a single round before uploading its weights.

B (Local Minibatch Size): The batch size used for local client updates. Setting B=∞ treats the full local dataset as a single batch, which collapses FedAvg back into FedSGD if E=1.
The total number of local updates a client performs per round

From Theory to Code: PyTorch Realization
How does this translate to actual lines of code? In a typical research framework like Flower (flwr), you intercept the local training loop on the client side.
While the server handles the global averaging (flwr.server.strategy.FedAvg), your client code simply runs a standard PyTorch loop but loops through the dataset E times locally before returning the state dictionary:
Python
# A typical FedAvg Local Training Loop on the Client Side
def local_train(model, train_loader, epochs, lr):
model.train()
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
criterion = torch.nn.CrossEntropyLoss()

for epoch in range(epochs): # Controlled by Parameter E
for images, labels in train_loader: # Batch size controlled by B
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step() # Local weight update

return model.state_dict() # Uploaded back to Server

What's Next? Navigating Non-IID Data
While FedAvg solves the communication crisis, it introduces a statistical one: Non-IID (Independent and Identically Distributed) data. Because edge devices collect data based on individual user habits, data distributions vary wildly between clients. This causes Client Drift, where local models diverge too far from each other, harming the global model's accuracy.

To fix this, advanced algorithms like FedProx introduce a proximal term to penalize local models that stray too far from the global anchor.

I am currently running local simulations utilizing PyTorch and the Flower Framework to map the trade-offs between FedAvg and FedProx on Non-IID image partitions (CIFAR-10). Mastering this foundational baseline is the mandatory stepping stone before scaling decentralized training up to massive multi-modal architectures like Vision-Language Models (VLMs).
Are you implementing or researching Federated Learning? What strategies are you using to balance the communication-computation trade-off? Let’s talk in the comments below!