Federated Learning on AWS

Phase 01

Distributed federated learning in AWS

I provisioned ten EC2 t3.micro instances as independent workers. Each held a private, non-overlapping partition of 6,000 MNIST images (60,000 in total) and was prepared with PyTorch and its data partition through a custom AMI. During every round, a worker downloaded the current global model from S3, trained a LeNet-5 CNN locally for five epochs, and uploaded its updated model to a separate local-model bucket, which triggered the aggregator automatically.

The aggregator was a serverless Lambda function (Python 3.10, 2048 MB) triggered by S3 PUT events on the local-model bucket. Once all ten workers had uploaded their models for a round, it ran Federated Averaging (FedAvg), a weighted average of every local model's weight tensors, evaluated the merged model against a held-out test set, and saved the new global model to models/ and the round metrics to metrics/ in the global bucket. To keep the Lambda package lean, I reimplemented the complete LeNet-5 forward pass (convolution, max-pooling, ReLU and softmax) in NumPy instead of shipping PyTorch.

To avoid costly S3 polling between rounds, I designed an SNS-to-SQS fanout. An S3 event notification on the global bucket published to the fl-global-model-ready SNS topic, which fanned out to ten dedicated SQS queues, one per worker. Each worker long-polled its own queue and woke as soon as the next global model was ready.

Phase 02

Extending the system to the edge

I then converted the workers into managed edge nodes by installing AWS IoT Greengrass Core v2 on all ten EC2 instances and registering each as an IoT core device in a Greengrass Thing Group. Every device received its own IoT certificate, Thing Policy, and Token Exchange Role so it could communicate securely with AWS IoT Core and S3.

The training worker was repackaged as a Greengrass component, com.fl.Worker v1.0.0: a JSON recipe defining its lifecycle, configuration parameters (WORKER_ID, PARTITION_ID) and platform compatibility, alongside a Python artifact. I deployed it locally on each device with the Greengrass CLI.

Instead of SQS polling, round coordination ran entirely over MQTT through AWS IoT Core. Each worker subscribed to fl/<ASU-ID>/next-round. The autograder published the first trigger; after each round, the worker itself published the next one, creating a self-sustaining training loop with no cloud-side orchestration between rounds.

Phase 03

What the project taught me

Implementing FedAvg from scratch made the relationship between local epochs, learning rate, batch size, and global convergence tangible; tuning them took the edge phase past 90% accuracy from a random initialization in a single round. It also showed why federated learning is useful: model quality can improve across distributed data owners without centralizing their raw datasets.

Moving the same workload from cloud-managed queues to edge-managed MQTT changed the coordination model. The workers became autonomous participants reacting to small messages instead of compute nodes driven by continuous polling. MQTT's publish-subscribe model is far cheaper than polling for battery-constrained devices, which maps naturally to healthcare monitoring, autonomous vehicles, industrial automation, and other bandwidth- or privacy-conscious environments.

Triggering aggregation from S3 PUT events instead of polling showed how reactive, serverless pieces compose into a reliable pipeline, and building the SNS-to-SQS fanout end to end made a standard enterprise pattern for broadcasting events to many consumers concrete. Writing LeNet-5 inference in plain NumPy taught me to build lean, dependency-minimal ML for constrained serverless environments.

Coordinating EC2, IAM, S3 notifications, Lambda, SNS, SQS, Greengrass, and IoT Core also reinforced the operational side of distributed systems: consistent naming, least-privilege access, idempotent events, and observable round state matter as much as the model itself.