%config InlineBackend.figure_formats = ['svg']
# Page-only settings: show every worker's log lines (Ray deduplicates
# identical lines by default) and silence Ray's accelerator tip.
import os
os.environ["RAY_DEDUP_LOGS"] = "0"
os.environ["RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO"] = "0"
"""
Utility functions and classes for Jupyter Notebooks lessons.
"""
from collections import OrderedDict
import logging
from logging import INFO
from flwr.client import Client, ClientApp, NumPyClient
from flwr.common.logger import (
ConsoleHandler,
console_handler,
FLOWER_LOGGER,
LOG_COLORS,
)
from logging import LogRecord
from typing import Dict, List, Optional, Tuple, Union
from flwr.server import ServerAppComponents
from flwr.client.mod import arrays_size_mod
from flwr.common import (
Context,
EvaluateRes,
ndarrays_to_parameters,
MessageType,
Parameters,
Scalar,
parameters_to_ndarrays,
)
from flwr.common.logger import (
log,
update_console_handler,
)
from flwr.server import ClientManager, ServerApp, ServerConfig
from flwr.server.strategy import FedAvg
from flwr.simulation import run_simulation
import torch
from transformers import AutoModelForCausalLM, GPTNeoXForCausalLM
# Customize logging for the course.
class InfoFilter(logging.Filter):
def filter(self, record):
return record.levelno == INFO
FLOWER_LOGGER.removeHandler(console_handler)
# To filter logging coming from the Simulation Engine
# so it is more readable in notebooks
from logging import ERROR
backend_setup = {"init_args": {"logging_level": ERROR, "log_to_driver": True}}
class ConsoleHandlerV2(ConsoleHandler):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def format(self, record: LogRecord) -> str:
"""Format function that adds colors to log level."""
if self.json:
log_fmt = "{lvl='%(levelname)s', time='%(asctime)s', msg='%(message)s'}"
else:
log_fmt = (
f"{LOG_COLORS[record.levelname] if self.colored else ''}"
f"%(levelname)s {'%(asctime)s' if self.timestamps else ''}"
f"{LOG_COLORS['RESET'] if self.colored else ''}"
f": %(message)s"
)
formatter = logging.Formatter(log_fmt)
return formatter.format(record)
console_handlerv2 = ConsoleHandlerV2(
timestamps=False,
json=False,
colored=True,
)
console_handlerv2.setLevel(INFO)
console_handlerv2.addFilter(InfoFilter())
FLOWER_LOGGER.addHandler(console_handlerv2)
# Route Flower logs to stdout so they render on this page,
# and silence library FutureWarnings as the other helper files do.
import sys
import warnings
console_handlerv2.setStream(sys.stdout)
warnings.filterwarnings("ignore")
def get_weights(net):
ndarrays = [
val.cpu().numpy() for _, val in net.state_dict().items()
]
return ndarrays
def set_weights(net, parameters):
params_dict = zip(net.state_dict().keys(), parameters)
state_dict = OrderedDict(
{k: torch.tensor(v) for k, v in params_dict}
)
net.load_state_dict(state_dict, strict=True)Bandwidth in Federated Learning
machine-learning
federated-learning
flower
A formula for estimating the bandwidth of a federated run, a worked example with a 14M parameter model, and techniques for reducing communication.
In every round of federated learning, models travel back and forth between the server and the clients. In the context of LLMs, where models get larger and larger, it is important to understand the bandwidth implications. This page builds a formula for reasoning about bandwidth usage in theory, then measures actual consumption with Flower in practice, and closes with the main techniques for reducing it.
Estimating Bandwidth
Build the estimate one factor at a time.
Start with the size of the model sent out to one individual client. To that, add the size of the model update received back from that client. The two sizes are not always the same. Sometimes the server sends out full model parameters, but the client returns compressed gradients, so the update coming back is smaller than the model that went out. The sum gives the bandwidth for serving one client for one round.
Multiply by the cohort size, the total number of clients in the system. With five clients, that is a factor of five. But not all clients are selected in every round, so multiply by the fraction selected per round. With 100 clients of which 20 percent participate, that factor is 0.2. Everything so far is the bandwidth of a single round, so finally multiply by the number of rounds.
\[ \text{bandwidth} \;=\; (\text{size}_{\text{out}} + \text{size}_{\text{in}}) \times \text{cohort size} \times \text{fraction selected} \times \text{rounds} \]
When the outgoing model and the incoming update have the same size, the first factor simplifies to twice the model size.
Worked Example
Take the EleutherAI Pythia 14M model, a language model with 14 million parameters and a size of 53 MB. Following the simplified formula, multiply by two for one client’s round trip, giving 106 MB. Train across two clients and both are selected each round, so the cohort factor is two and the fraction is 1.0. With a cohort of 100 it would be reasonable to select only 50 per round, a fraction of 0.5, but not here. Run a single round. The total is
\[ 53 \times 2 \times 2 \times 1.0 \times 1 = 212 \text{ MB} \]
for one round of federated learning with a small 14M parameter model. Even a single round can quickly eat up a lot of bandwidth.
Review Questions
1. Write the bandwidth formula and name its four factors.
TipAnswer
Bandwidth equals (size of model sent out + size of update received back) times cohort size times fraction of clients selected per round times number of rounds. The first factor covers one client’s round trip, the cohort and fraction factors count how many clients actually participate, and the rounds factor extends it over the whole run.
1. Why can the outgoing size and the incoming size differ?
TipAnswer
The server typically sends full model parameters, but the client may return something smaller, such as compressed gradients. In that case the incoming update is smaller than the outgoing model, and the two terms must be kept separate rather than using twice the model size.
1. A 53 MB model is trained across 10 clients with half selected each round, for 20 rounds, with symmetric updates. Estimate the bandwidth.
TipAnswer
\(106 \times 10 \times 0.5 \times 20 = 10{,}600\) MB, so roughly 10.6 GB. The per-client round trip is \(53 \times 2 = 106\) MB, five clients participate per round on average, and the run has 20 rounds.
Lab: Measuring Bandwidth in Practice
The 212 MB estimate can be checked against a real run. Two measurement points cover the two directions. On the client side, a built-in Flower mod tracks the size of the arrays each client transmits. On the server side, a custom strategy logs the size of every model sent and received. The Flower client itself skips actual training and evaluation, since only the bandwidth is being measured.
NoteLab Files Download
Everything this lab needs, next to your notebook.
- utils5.py (3 KB), the helper file with the compatibility fix below already applied
- requirements.txt (1 KB), the package versions the course shipped
The Pythia 14M model weights (about 53.7 MiB) are downloaded automatically from Hugging Face by the cell that loads the model, not by the first cell, so there is nothing else to fetch by hand.
ImportantFour changes from the course notebook
The course pinned Flower 1.10, and this page runs on Flower 1.33 with current transformers and Torch. Four things changed (updated 2026-08-31).
- The client mod was renamed. What the course imported as
parameters_size_modisarrays_size_modin current Flower. - The checkpoint is loaded as
float32explicitly. Newertransformersreleases load this model in its stored half precision by default, which would halve the measured size. - The byte accounting was corrected. The course truncates each transfer to whole mebibytes before summing, which discards about 2.7 MiB across four legs and yields 212. This page accumulates exact bytes and converts once, giving 225,083,392 bytes or 214.66 MiB.
- The units are named honestly. Everything computed here divides by \(1024^2\), so the figures are mebibytes, and the page says so rather than labeling them MB.
The Pythia 14M model, the tracking strategy, the simulation setup and the single-round design are the course’s own.
The helper file is small. It customizes Flower’s logging so the notebook output stays readable, and pulls in the framework pieces the lab uses.
NoteHelper functions from utils5.py
Loading the Model
The client-side measurement hook is the built-in arrays_size_mod, imported from Flower’s mod collection.
from flwr.client.mod import arrays_size_modThe model is the EleutherAI Pythia 14M language model from the bandwidth formula example. AutoModelForCausalLM.from_pretrained downloads the weights on first use and caches them in the pythia-14m/cache directory, and the dtype argument requests full 32-bit precision, matching the course setup.
model = AutoModelForCausalLM.from_pretrained(
"EleutherAI/pythia-14m",
cache_dir="./pythia-14m/cache",
dtype=torch.float32,
)Summing the element sizes across all parameter tensors in the model’s state_dict confirms the number used in the calculation. Each of the roughly 14 million parameters takes 4 bytes. Dividing the byte total by \(1024^2\) gives mebibytes (MiB), the binary unit, which is what people usually mean when they say “MB” about a file size. Keep the exact byte count as well, so the totals later do not lose anything to rounding.
vals = model.state_dict().values()
total_size_bytes = sum(p.element_size() * p.numel() for p in vals)
total_size_mib = total_size_bytes / (1024**2)
log(INFO, "Model size is: {:,} bytes ({:.2f} MiB)".format(total_size_bytes, total_size_mib))INFO : Model size is: 56,270,848 bytes (53.66 MiB)
Client That Only Moves Parameters
The FlowerClient looks like the ones from the earlier labs, except that fit and evaluate do no actual work. They load the received parameters and hand them straight back.
Leaving the training out is safe here, and the reason is worth being explicit about. Training changes the values inside the parameters, but it does not change how many parameters there are, and the bytes on the wire depend only on the count. A trained update and an untrained one are the same size. Skipping the training therefore gives exactly the same measurement in a small fraction of the time, and it keeps the lab focused on the one quantity being studied.
The ClientApp attaches arrays_size_mod, which logs the size of the arrays each client transmits.
class FlowerClient(NumPyClient):
def __init__(self, net):
self.net = net
def fit(self, parameters, config):
set_weights(self.net, parameters)
# No actual training here
return get_weights(self.net), 1, {}
def evaluate(self, parameters, config):
set_weights(self.net, parameters)
# No actual evaluation here
return float(0), int(1), {"accuracy": 0}
def client_fn(context: Context) -> FlowerClient:
return FlowerClient(model).to_client()
client = ClientApp(
client_fn,
mods=[arrays_size_mod],
)Server That Tracks Sizes
On the server side, a custom strategy called BandwidthTrackingFedAvg extends the usual federated averaging strategy. In configure_fit it calculates and logs the size of the model about to be sent to each client, and in aggregate_fit it calculates and logs the size of each received model update. Every size is appended to the bandwidth_sizes list, and the total bandwidth is the sum of that list at the end. Both methods call super() so FedAvg still does the real configuration and aggregation work.
The server is the right place to count from, because every byte in the system passes through it, so nothing has to be collected from the clients and added up afterwards. The two directions are measured separately rather than one being doubled, and that is deliberate. As the formula noted, what goes out and what comes back are not always the same size, since a client may return something compressed. Measuring each direction on its own keeps the code honest when that happens, and it also lets you see the doubling actually occur instead of assuming it.
bandwidth_sizes = []
class BandwidthTrackingFedAvg(FedAvg):
def aggregate_fit(self, server_round, results, failures):
if not results:
return None, {}
# Track sizes of models received
for _, res in results:
ndas = parameters_to_ndarrays(res.parameters)
# Keep raw bytes in the list. Rounding each leg to whole MiB before
# summing would quietly lose a couple of MiB across four transfers.
nbytes = sum(n.nbytes for n in ndas)
log(INFO, f"Server receiving model size: {nbytes / (1024**2):.2f} MiB")
bandwidth_sizes.append(nbytes)
# Call FedAvg for actual aggregation
return super().aggregate_fit(server_round, results, failures)
def configure_fit(self, server_round, parameters, client_manager):
# Call FedAvg for actual configuration
instructions = super().configure_fit(
server_round, parameters, client_manager
)
# Track sizes of models to be sent
for _, ins in instructions:
ndas = parameters_to_ndarrays(ins.parameters)
nbytes = sum(n.nbytes for n in ndas)
log(INFO, f"Server sending model size: {nbytes / (1024**2):.2f} MiB")
bandwidth_sizes.append(nbytes)
return instructionsThe strategy sets fraction_evaluate to zero to disable client side evaluation, and the run is a single round. One round is enough, because in this setup the bandwidth requirements do not change over consecutive rounds.
params = ndarrays_to_parameters(get_weights(model))
def server_fn(context: Context):
strategy = BandwidthTrackingFedAvg(
fraction_evaluate=0.0,
initial_parameters=params,
)
config = ServerConfig(num_rounds=1)
return ServerAppComponents(
strategy=strategy,
config=config,
)
server = ServerApp(server_fn=server_fn)Running It
The simulation runs with two clients, matching the worked example.
run_simulation(server_app=server,
client_app=client,
num_supernodes=2,
backend_config=backend_setup
)INFO : Starting Flower ServerApp, config: num_rounds=1, no round_timeout INFO : INFO : [INIT] INFO : Using initial global parameters provided by strategy INFO : Starting evaluation of initial global parameters INFO : Evaluation returned no results (`None`) INFO : INFO : [ROUND 1] INFO : Server sending model size: 53.66 MiB INFO : Server sending model size: 53.66 MiB INFO : configure_fit: strategy sampled 2 clients (out of 2) INFO : aggregate_fit: received 2 results and 0 failures INFO : Server receiving model size: 53.66 MiB INFO : Server receiving model size: 53.66 MiB INFO : configure_evaluate: no clients selected, skipping evaluation INFO : INFO : [SUMMARY] INFO : Run finished 1 round(s) in 15.32s INFO :
(ClientAppActor pid=58270) INFO : Incoming `ArrayRecord` size statistics: (ClientAppActor pid=58270) INFO : {'fitins.parameters': {'elements': 14067712, 'bytes': 56280718}} (ClientAppActor pid=58270) INFO : Total array elements received: 56280718 bytes (ClientAppActor pid=58270) INFO : Outgoing `ArrayRecord` size statistics: (ClientAppActor pid=58270) INFO : {'fitres.parameters': {'elements': 14067712, 'bytes': 56280718}} (ClientAppActor pid=58270) INFO : Total array elements sent: 56280718 bytes (ClientAppActor pid=58262) INFO : Incoming `ArrayRecord` size statistics: (ClientAppActor pid=58262) INFO : {'fitins.parameters': {'elements': 14067712, 'bytes': 56280718}} (ClientAppActor pid=58262) INFO : Total array elements received: 56280718 bytes (ClientAppActor pid=58262) INFO : Outgoing `ArrayRecord` size statistics: (ClientAppActor pid=58262) INFO : {'fitres.parameters': {'elements': 14067712, 'bytes': 56280718}} (ClientAppActor pid=58262) INFO : Total array elements sent: 56280718 bytes
The logs show the server sending a 53.66 MiB model to each of the two clients and receiving a 53.66 MiB update back from each. Summing the tracked byte counts gives the total.
total_bytes = sum(bandwidth_sizes)
log(INFO, "Total bandwidth used: {:,} bytes ({:.2f} MiB)".format(
total_bytes, total_bytes / (1024**2)))INFO : Total bandwidth used: 225,083,392 bytes (214.66 MiB)
The total comes out at 225,083,392 bytes, or about 214.66 MiB across the four transfers. That is exactly \(4 \times 53.6640625\) MiB, matching what the formula predicted from the model size. (Multiply the rounded 53.66 instead and you get 214.64, which is the rounding showing, not a real difference.) The rougher 212 MiB figure comes from truncating each leg to a whole 53 MiB before adding, which quietly discards about 2.7 MiB.
Note that every size on this page is in mebibytes, powers of 1024, because that is what the code computes. Network bandwidth is often quoted in decimal megabytes instead, and 214.66 MiB is 225.08 MB in those units.
One caveat about what this number is. The strategy counts only the bytes of the decoded parameter arrays. It does not include the configuration dictionaries, the returned metrics, the message envelopes, or any gRPC and TLS framing, so real on-the-wire traffic is somewhat higher. As an estimate of the dominant cost, the model weights, it is the right number to reason about.
Review Questions
1. Why was measuring a single round sufficient?
TipAnswer
Because in this setup every round moves the same amount of data, the same model size to the same number of selected clients and back. The per-round cost is constant, so the total is just the single-round measurement times the number of rounds.
1. Where do the two measurement hooks live, and what does each observe?
TipAnswer
On the client, the arrays_size_mod tracks the size of the arrays the client transmits. On the server, a custom strategy extending federated averaging logs the outgoing model size in configure_fit and each incoming update size in aggregate_fit. Together they cover both directions of the round trip.
1. Why do fit and evaluate in this lab’s client do no work?
TipAnswer
The lab measures communication, not learning. The client still receives parameters and sends parameters back, which is the traffic being measured, and skipping the training and evaluation keeps the run fast without changing a single byte of what travels.
Reducing Bandwidth
There are many ways to reduce bandwidth usage in federated learning, and they fall into two categories, reducing the size of an individual update, and simply communicating less.
To reduce the size of an update, one option is sparsification, for example top-\(k\) sparsification. If gradient values to be communicated are below a certain threshold, they are treated as zero and their transmission can be skipped, which saves communication. This is especially likely to help toward the end of training, when more elements of the gradients are small in magnitude. Another option is quantization, which in its many forms reduces the number of bits used to represent scalars, and in turn the size of the updates exchanged between client and server.
To communicate less, one approach is to leverage pre-trained models. In many settings it is realistic to assume a pre-trained model exists that is useful for the application, and federated learning then continues the training. In such cases not every layer needs training, and only the layers modified by the federated training need to be communicated. Another approach is to simply train longer locally before exchanging updates with the server, for example five local epochs instead of one. Be aware that this can also prevent convergence. When local models train too many epochs, they diverge more and more from each other, which can cause the aggregated model to become worse instead of better.
Review Questions
1. Name the two categories of bandwidth reduction and one technique in each.
TipAnswer
Reducing the size of an individual update, for example through top-\(k\) sparsification or quantization, and communicating less often or less content, for example starting from a pre-trained model and communicating only the layers that federated training modifies, or training more local epochs between exchanges.
1. Why does sparsification tend to save more toward the end of training?
TipAnswer
As training converges, more gradient elements become small in magnitude, so more of them fall below the threshold and can be skipped, transmitted implicitly as zero. Early in training, gradients are large and dense, so there is less to drop.
1. Training many local epochs between exchanges reduces communication. What is the risk?
TipAnswer
Divergence. The longer each client trains locally, the further the local models drift apart, and aggregating heavily diverged models can make the global model worse instead of better. Local epochs trade bandwidth against the stability of convergence.