273 lines
6.9 KiB
Markdown
273 lines
6.9 KiB
Markdown
# swarm
|
|
|
|
Docker Swarm distributed training for heterogeneous CPU clusters.
|
|
|
|
**Directory:** `swarm/`
|
|
**Status:** Complete with 14 passing tests
|
|
|
|
---
|
|
|
|
## Overview
|
|
|
|
The swarm module provides distributed training across heterogeneous hardware using a parameter server architecture:
|
|
|
|
- **Coordinator:** Holds master weights, aggregates gradients, distributes updates
|
|
- **Workers:** Compute forward/backward passes, send compressed gradients
|
|
|
|
Key features:
|
|
- Async gradient updates (workers don't block each other)
|
|
- TopK gradient compression (90% bandwidth reduction)
|
|
- FP16 compression option
|
|
- Hardware profile auto-detection
|
|
- Fault tolerance (workers can drop and rejoin)
|
|
- Docker Swarm deployment support
|
|
|
|
---
|
|
|
|
## Architecture
|
|
|
|
```
|
|
┌─────────────────┐
|
|
│ Coordinator │
|
|
│ (holds weights)│
|
|
└────────┬────────┘
|
|
│
|
|
┌───────────────────┼───────────────────┐
|
|
│ │ │
|
|
┌────▼────┐ ┌────▼────┐ ┌────▼────┐
|
|
│ Worker │ │ Worker │ │ Worker │
|
|
│ (Pi 4) │ │ (Laptop)│ │ (Server)│
|
|
└─────────┘ └─────────┘ └─────────┘
|
|
```
|
|
|
|
---
|
|
|
|
## Quick Start
|
|
|
|
### Running Locally
|
|
|
|
```bash
|
|
# Terminal 1: Start coordinator
|
|
./bin/swarm_coordinator --port 7100
|
|
|
|
# Terminal 2: Start worker
|
|
./bin/swarm_worker --host localhost --port 7100
|
|
```
|
|
|
|
### Docker Swarm Deployment
|
|
|
|
```bash
|
|
cd swarm/docker
|
|
./deploy.sh
|
|
```
|
|
|
|
---
|
|
|
|
## Files
|
|
|
|
| File | Description |
|
|
|------|-------------|
|
|
| `protocol.h/c` | Wire protocol, compression, message handling |
|
|
| `network.h/c` | TCP socket abstraction |
|
|
| `config.h/c` | TOML-like config file parsing |
|
|
| `coordinator.h/c` | Parameter server |
|
|
| `worker.h/c` | Training worker |
|
|
| `coordinator_main.c` | Coordinator CLI entry point |
|
|
| `worker_main.c` | Worker CLI entry point |
|
|
|
|
---
|
|
|
|
## Protocol
|
|
|
|
### Message Types
|
|
|
|
| Type | Direction | Description |
|
|
|------|-----------|-------------|
|
|
| `REGISTER` | Worker → Coord | Worker registration |
|
|
| `REGISTER_ACK` | Coord → Worker | Registration confirmed |
|
|
| `WEIGHTS` | Coord → Worker | Full weight tensor |
|
|
| `GRADIENTS` | Worker → Coord | Compressed gradients |
|
|
| `UPDATE` | Coord → Worker | Weight update delta |
|
|
| `START` | Coord → Worker | Start training |
|
|
| `PAUSE/RESUME` | Coord → Worker | Training control |
|
|
| `SHUTDOWN` | Coord → Worker | Graceful shutdown |
|
|
| `HEARTBEAT` | Both | Connection keepalive |
|
|
| `CHECKPOINT` | Coord → Worker | Save checkpoint |
|
|
|
|
### Compression
|
|
|
|
| Type | Ratio | Description |
|
|
|------|-------|-------------|
|
|
| `NONE` | 1:1 | No compression |
|
|
| `FP16` | 2:1 | Float16 conversion |
|
|
| `TOPK` | 10:1 | Top-K by magnitude |
|
|
| `QSGD` | Variable | Quantized SGD |
|
|
|
|
---
|
|
|
|
## API Reference
|
|
|
|
### Coordinator
|
|
|
|
```c
|
|
typedef struct {
|
|
uint16_t port;
|
|
float learning_rate;
|
|
float beta1, beta2;
|
|
float epsilon;
|
|
int min_workers;
|
|
int checkpoint_interval;
|
|
char checkpoint_dir[SW_MAX_PATH];
|
|
sw_compress_type compression;
|
|
float topk_ratio;
|
|
} sw_coordinator_config;
|
|
|
|
sw_coordinator* sw_coordinator_new(tm_tensor** weights, int num_weights,
|
|
sw_coordinator_config* config);
|
|
void sw_coordinator_free(sw_coordinator* coord);
|
|
|
|
bool sw_coordinator_start(sw_coordinator* coord);
|
|
void sw_coordinator_stop(sw_coordinator* coord);
|
|
void sw_coordinator_wait(sw_coordinator* coord);
|
|
|
|
int sw_coordinator_num_workers(sw_coordinator* coord);
|
|
int sw_coordinator_get_step(sw_coordinator* coord);
|
|
float sw_coordinator_get_loss(sw_coordinator* coord);
|
|
```
|
|
|
|
### Worker
|
|
|
|
```c
|
|
typedef struct {
|
|
char coordinator_host[256];
|
|
uint16_t coordinator_port;
|
|
int reconnect_attempts;
|
|
int reconnect_delay_sec;
|
|
int num_threads;
|
|
int batch_size;
|
|
char data_path[SW_MAX_PATH];
|
|
int shard_index;
|
|
sw_compress_type compression;
|
|
float topk_ratio;
|
|
char worker_id[SW_MAX_WORKER_ID];
|
|
} sw_worker_config;
|
|
|
|
sw_worker* sw_worker_new(sw_worker_config* config);
|
|
void sw_worker_free(sw_worker* w);
|
|
|
|
bool sw_worker_connect(sw_worker* w);
|
|
void sw_worker_disconnect(sw_worker* w);
|
|
|
|
void sw_worker_set_weights(sw_worker* w, tm_tensor** weights, int num_weights);
|
|
void sw_worker_set_forward(sw_worker* w, sw_forward_fn fn, void* ctx);
|
|
void sw_worker_set_backward(sw_worker* w, sw_backward_fn fn, void* ctx);
|
|
void sw_worker_set_data_loader(sw_worker* w, sw_next_batch_fn next,
|
|
sw_reset_data_fn reset, sw_free_batch_fn free_batch,
|
|
void* ctx);
|
|
|
|
bool sw_worker_fetch_weights(sw_worker* w);
|
|
void sw_worker_train(sw_worker* w); // Blocking
|
|
void sw_worker_stop(sw_worker* w);
|
|
```
|
|
|
|
### Training Callbacks
|
|
|
|
```c
|
|
// Forward: compute loss from batch
|
|
typedef float (*sw_forward_fn)(tm_tensor** weights, int num_weights,
|
|
void* batch_data, void* ctx);
|
|
|
|
// Backward: compute gradients
|
|
typedef void (*sw_backward_fn)(tm_tensor** weights, int num_weights,
|
|
void* batch_data, void* ctx);
|
|
|
|
// Data loading
|
|
typedef void* (*sw_next_batch_fn)(void* ctx);
|
|
typedef void (*sw_reset_data_fn)(void* ctx);
|
|
typedef void (*sw_free_batch_fn)(void* batch, void* ctx);
|
|
```
|
|
|
|
---
|
|
|
|
## Configuration Files
|
|
|
|
### coordinator.toml
|
|
```toml
|
|
[network]
|
|
port = 7100
|
|
|
|
[training]
|
|
learning_rate = 1e-4
|
|
beta1 = 0.9
|
|
beta2 = 0.999
|
|
min_workers = 2
|
|
|
|
[checkpointing]
|
|
interval = 1000
|
|
directory = "/checkpoints"
|
|
|
|
[compression]
|
|
type = "topk"
|
|
topk_ratio = 0.1
|
|
```
|
|
|
|
### worker.toml
|
|
```toml
|
|
[network]
|
|
coordinator_host = "coordinator"
|
|
coordinator_port = 7100
|
|
reconnect_attempts = 10
|
|
reconnect_delay_sec = 5
|
|
|
|
[hardware]
|
|
num_threads = 0 # auto
|
|
max_memory_mb = 0 # auto
|
|
|
|
[training]
|
|
batch_size = 0 # auto from profile
|
|
accumulation_steps = 0 # auto
|
|
|
|
[data]
|
|
path = "/data/train.jsonl"
|
|
shard_index = -1 # auto-assign
|
|
|
|
[compression]
|
|
type = "topk"
|
|
topk_ratio = 0.1
|
|
```
|
|
|
|
---
|
|
|
|
## Hardware Profiles
|
|
|
|
Auto-detected profiles with recommended batch sizes:
|
|
|
|
| Profile | Memory | Batch Size | Description |
|
|
|---------|--------|------------|-------------|
|
|
| `RPI3` | <1GB | 2 | Raspberry Pi 3 |
|
|
| `RPI4` | 1-4GB | 4-8 | Raspberry Pi 4 |
|
|
| `RPI5` | 4-8GB | 8-16 | Raspberry Pi 5 |
|
|
| `LAPTOP` | 8-16GB | 16-32 | Standard laptop |
|
|
| `DESKTOP` | 16-32GB | 32-64 | Desktop PC |
|
|
| `WORKSTATION` | 32-64GB | 64-128 | Workstation |
|
|
| `SERVER` | >64GB | 128-256 | Server |
|
|
|
|
---
|
|
|
|
## Docker Deployment
|
|
|
|
See `swarm/docker/` for:
|
|
- `Dockerfile.coordinator`
|
|
- `Dockerfile.worker`
|
|
- `docker-compose.yml` (local testing)
|
|
- `docker-stack.yml` (Swarm deployment)
|
|
- `deploy.sh` (deployment script)
|
|
|
|
---
|
|
|
|
## See Also
|
|
|
|
- [distributed](distributed.md) - Local distributed primitives
|
|
- [tensor](tensor.md) - Tensor operations
|
|
- [optim](optim.md) - Optimizers
|