1. Parameter Server Architecture

In the Parameter Server architecture, the nodes in the cluster are divided into two categories: parameter server nodes (Parameter Server) and worker server nodes (Worker).

1.1 Parameter Server

The Parameter Server is used to store the model’s parameters.

Each parameter server node is responsible for managing and updating one part of the model’s parameters, while each worker node only processes the subset of data corresponding to it.

1.2 Worker

Worker server nodes are responsible for executing the model’s training tasks. The training data is distributed across multiple worker server nodes. Each worker server node independently performs the forward and backward computations, ultimately producing gradient information.

1.3 Training Process

1. Initialization: all parameter server nodes complete the initialization of the model weights.
2. Weight retrieval: the worker server nodes pull the corresponding weights from all parameter server nodes.
3. Forward and backward computation: the worker server nodes perform the forward and backward computations and generate gradients.
4. Gradient upload: the worker server nodes push the computed gradients to the corresponding parameter server nodes.
5. Weight update: after receiving the gradients from all worker server nodes, the parameter server nodes aggregate them and update the model weights.
6. Repeated iteration: continue executing steps 2 through 5 until the convergence condition is reached or training ends.

1.4 Maintaining the Structure

Handling a Parameter Server failure

• Parameter reassignment: the parameters of the failed parameter server node are reassigned to other available servers to preserve the integrity of the model.
• Dynamic load balancing: the system automatically adjusts the load to ensure parameters are evenly distributed across servers and to prevent overload.
• Continued training: the workers reconnect to the new parameter servers and continue model training.

Handling a Worker failure

• Task reassignment: the tasks of the disconnected Worker are reassigned to other available Workers to avoid wasting compute resources.
• Progressive convergence strategy: the remaining Workers can continue training on a small amount of data to avoid a long interruption.
• Fault recovery: if the failed Worker recovers, the system rejoins it to training and continues the unfinished tasks.

A new Worker joining

• Data redistribution: after the new Worker joins, the system distributes part of the training data to ensure the load is even.
• Gradient synchronization: the new Worker obtains the latest weights from the parameter servers and synchronizes gradients with the other Workers.
• Dynamic scaling: the new Worker can seamlessly join the training flow, increasing the compute capacity of the whole system.

1.5 Framework Support

• TensorFlow natively supports the Parameter Server architecture, providing tf.distribute.Strategy to implement distributed training.

• PyTorch can use the torch.distributed package or other libraries (such as Ray and Horovod) to implement Parameter Server training.

1.6 Applicable Scenarios

• Large-scale recommendation systems

In recommendation systems, the model parameters may include user preferences and item features, and the number of these parameters can be enormous. Using a Parameter Server allows these parameters to be managed and updated effectively.

• Natural language processing (NLP)

When training large language models such as BERT or GPT, a Parameter Server can help store and update the model’s word embeddings and inter-layer weights in a distributed way.

• Image recognition

When training deep neural networks for image recognition, such as convolutional neural networks (CNNs), a Parameter Server can process large numbers of filter weights in a distributed way.

• Large-scale linear regression

When dealing with linear regression problems that have millions of features, a Parameter Server can store and update the weight matrix in a distributed way.

• Real-time big data analytics

In scenarios that need real-time updates of model parameters to respond to rapidly changing data patterns, the Parameter Server architecture can provide the required flexibility and scalability.

2. AllReduce Architecture

2.1 Establishing the Communication Ring

1. Initialize the Rank

Each Worker is assigned a unique Rank (usually an integer, representing the Worker’s identity and order).

2. Establish the initial communication ring

Based on the Workers’ Ranks, the system pairs each Worker with its left and right neighboring Workers, forming a ring-shaped communication topology. Assuming there are N Workers, each Worker i communicates with i-1 and i+1 (where the left neighbor of i=0 is N-1 and the right neighbor of i=N-1 is 0), thus forming a closed communication ring.

2.2 Training Process

1. Gradient splitting

Each Worker splits the gradients it computed into multiple chunks (usually equal to the number of Workers).

2. Gradient exchange

Each Worker sends the first gradient chunk it holds to the right-neighbor Worker, while receiving the first gradient chunk passed from the left-neighbor Worker. This process runs for multiple rounds of communication, and in each round the Worker sends the next gradient chunk and receives a new gradient chunk.

3. Gradient accumulation

In each round of communication, the Worker not only receives a new gradient chunk but also adds it to the previously accumulated gradient chunk, thereby gradually aggregating the gradients across all Workers.

4. Broadcast the final result

Once all gradient chunks have been exchanged and accumulated, the final accumulated result is broadcast back to each Worker by continuing the ring communication. At this point, every Worker holds the complete and synchronized global gradient information.

2.3 Maintaining the Structure

Handling a Worker failure

• Detect disconnection: the remaining Workers automatically detect the disconnected Worker and trigger the fault-handling mechanism.
• Reassign Ranks: the system automatically reassigns the Rank of the failed Worker and updates the communication topology.
• Rebuild the communication ring: the left and right neighbors of the disconnected Worker re-establish connections to form a new ring topology.

A new Worker joining

• Insert into the communication ring: after the new Worker joins, the system inserts it into the existing communication ring according to its Rank.
• Establish connections: the new Worker establishes communication connections with its left and right neighbors, thereby extending the ring structure.

Synchronization and recovery of the communication ring

• State synchronization: after the new communication ring is established, all Workers resynchronize the previous gradient accumulation state.
• Continue training: once synchronization is complete, the system continues with the subsequent training steps to ensure the stability and consistency of training.

2.4 Framework Support

• PyTorch natively supports the AllReduce operation, providing the torch.distributed module.
• TensorFlow natively supports AllReduce, using the MirroredStrategy policy through the tf.distribute.Strategy module to perform AllReduce operations with NVIDIA NCCL.

2.5 Applicable Scenarios

• Deep learning model training

When training large neural networks with deep learning frameworks such as TensorFlow or PyTorch, the AllReduce algorithm can efficiently synchronize gradients across multiple GPUs or TPUs.

• Distributed optimization algorithms

When implementing a distributed version of stochastic gradient descent (SGD) or other optimization algorithms, AllReduce is used to ensure that all compute nodes use the same global gradient estimate in every iteration.

• Multi-task learning

In multi-task learning scenarios, different tasks may be trained on different compute nodes, and AllReduce can ensure that all tasks share the same model parameter updates.

• Reinforcement learning

In distributed reinforcement learning, multiple agents may need to synchronize their experience or policy updates, and the AllReduce algorithm can play a role in this process.

• Large-scale graph computing

When processing large-scale graph data, such as social network analysis or network traffic analysis, AllReduce can be used to synchronize graph embeddings or node features across multiple compute nodes.

• Scientific computing and simulation

In scientific fields that require large-scale parallel computing, such as climate simulation and physics simulation, AllReduce can be used to synchronize simulation state across multiple compute nodes.