Building ComputeShare: A Lightweight Multi‑Machine Distributed Training System with PyTorch

Quick Overview of ComputeShare

ComputeShare enables multiple machines to cooperate over a LAN or the public Internet to train a neural network. Workers automatically select the best available hardware: NVIDIA GPUs via CUDA, Apple Silicon via MPS, or CPUs when no accelerator is present.

Each worker trains on a partition of the dataset, computes gradients, and sends them to a central parameter server, which aggregates the gradients by averaging and updates the global model with SGD.

System Architecture

Parameter Server – Stores the global model, receives gradients from workers through FastAPI endpoints, aggregates them by averaging, and updates the model with SGD.

Workers – Download the latest model, train on a local data shard, compute gradients, and push the results back. Workers can run on GPUs, Apple Silicon, CPUs, or edge devices such as Raspberry Pi or Android phones.

Linear Scaling Rule

Adding workers increases the effective batch size; keeping the learning rate unchanged slows convergence. ComputeShare scales the learning rate linearly with the number of workers:

scaled_lr = 0.01 * BUFFER_SIZE
optimizer = torch.optim.SGD(
    global_model.parameters(),
    lr=scaled_lr,
    momentum=0.9
)

With three workers the effective learning rate becomes 0.03.

Modular Design

The project ships with a lightweight CNN called SimpleNet (≈21 k parameters) optimized for MNIST, but the architecture is model‑agnostic and can accommodate ResNet, EfficientNet, or Transformer‑based models.

Engineering Challenges

Stale Gradients (Synchronous Solution)

When workers finish at different speeds, a slow worker may submit gradients computed on an outdated model version. Each gradient payload includes the worker’s model version; the server rejects mismatched versions:

if worker_version != model_version:
    raise HTTPException(
        status_code=409,
        detail="Stale gradients rejected."
    )

Asynchronous SGD (Experimental Async Branch)

An async branch ( feature/async) lets workers push gradients without waiting. The server applies gradients immediately, removing the “slowest worker” bottleneck but introducing instability because stale gradients may be applied to newer weights, causing accuracy fluctuations.

Bandwidth Compression

Gradients are serialized with PyTorch’s native binary format, compressed with gzip, and transmitted as binary payloads, reducing sustained traffic to under 50 KB/s for SimpleNet:

buffer = io.BytesIO()
torch.save(cpu_grads, buffer)
compressed_payload = gzip.compress(
    buffer.getvalue()
)

Unstable Connections

For internet‑scale training, Cloudflare Tunnels ( cloudflared) replace LocalTunnel/ngrok to provide more reliable TCP connections, longer timeouts, and retry logic that tolerates occasional packet loss.

Authentication

A simple PIN‑based header ( X-Auth-Pin) protects the server from unauthorized gradient submissions.

Universal Dataset Factory

A unified factory in utils.py normalizes dataset creation across torchvision datasets (e.g., handling train=True vs split="train"), automatically adjusts model output dimensions, and installs missing dependencies when needed. It supports MNIST, Fashion‑MNIST, CIFAR‑10/100, and others without manual reconfiguration.

Comparison with PyTorch DDP and FSDP

ComputeShare uses a centralized parameter‑server model, whereas DDP employs decentralized AllReduce communication that scales better in high‑performance GPU clusters. FSDP shards model parameters, gradients, and optimizer states for memory efficiency; ComputeShare replicates the full model on each worker, simplifying heterogeneous hardware participation.

The system emphasizes ease of configuration for heterogeneous, low‑cost devices rather than raw scalability.

Summary

Centralized parameter server

Hardware‑agnostic workers (CUDA, MPS, CPU)

Compressed binary gradient transfer

Automatic dataset handling via a universal factory

Stale‑gradient protection

Experimental asynchronous training support

Project repository: https://github.com/alphastar-avi/computeShare