Parameter Server

The parameter server acts as the single source of truth for the model's weights and biases. It maintains a distributed, in-memory key-value store where keys represent parameter names and values are their numerical tensors. This centralized logical view simplifies debugging and checkpointing, as the entire model state can be snapshotted from one location. Sharded parameter servers partition the global key-space across multiple nodes to handle models that exceed the memory capacity of a single machine, using consistent hashing to map keys to server shards.

Workers operate independently without barrier synchronization. The core interaction is a non-blocking loop:

• Push: A worker computes a gradient on its local data shard and sends the update to the parameter server.
• Pull: A worker requests the latest global parameters before starting the next iteration. This asynchronous parallelism maximizes hardware utilization by preventing fast workers from waiting for stragglers. However, it introduces the stale gradient problem, where a worker computes gradients using an outdated version of the parameters, potentially slowing convergence.

The trade-off between convergence speed and system throughput is controlled by the consistency model:

• Sequential Consistency: All workers see updates in the same order. Equivalent to single-threaded execution but negates the benefits of asynchrony.
• Eventual Consistency: Workers pull whatever the latest state is. Maximizes throughput but may cause significant staleness.
• Bounded Staleness: A practical middle ground where the server restricts the lag of any worker to a maximum number of missed updates (the staleness bound), ensuring theoretical convergence guarantees while maintaining high parallelism.

The architecture is designed for commodity hardware where node failures are expected:

• Worker Fault Tolerance: If a worker fails, its incomplete task is simply re-scheduled on another node. No state recovery is needed since workers are stateless compute units.
• Server Fault Tolerance: The global state is replicated across multiple server replicas using chain replication or a consensus protocol like Raft. If the primary shard fails, a replica takes over seamlessly.
• Elastic Scaling: New workers can join a running training job dynamically to accelerate progress, and new server shards can be added to expand capacity without stopping the training process.

The programming model abstracts away network communication through a shared key-value interface. Developers define a scheduler that assigns data partitions to workers and a training function that runs on each worker. The function uses:

• ps.push(gradient_vector, learning_rate) to send updates.
• ps.pull(model_keys) to fetch the latest weights. This abstraction allows the same training code to run on a single machine or a cluster of thousands without modification. Range-based push/pull operations allow workers to synchronize only specific layers, enabling techniques like layer-wise asynchronous updates where different parts of the network update at different frequencies.

Network bandwidth is the primary bottleneck. The parameter server employs several strategies to reduce communication overhead:

• Gradient Compression: Updates are quantized to low-precision integers (e.g., 8-bit) or sparsified by transmitting only the top-k largest gradient values.
• Server-side Aggregation: Instead of applying each worker's push immediately, the server batches multiple gradient pushes and applies the averaged update, reducing write amplification.
• Pipeline Parallelism: The pull-compute-push phases are overlapped so that a worker begins computing the next batch while the current gradient is still being transmitted, hiding network latency behind computation.