Asynchronous Model Average

In synchronous communication algorithms such as Gradient AllReduce, every worker needs to be in the same iteration in a lock-step style. When there is no straggler in the system, such synchronous algorithms are reasonably efficient, and gives deterministic training results that are easier to reason about. However, when there are stragglers in the system, with synchronous algorithms faster workers have to wait for the slowest worker in each iteration, which can dramatically harm the performance of the whole system¹. To deal with stragglers, we can use asynchronous algorithms where workers are not required to be synchronized. The Asynchronous Model Average algorithm provided by Bagua is one of such algorithms.

Algorithm

The Asynchronous Model Average algorithm can be described as follows:

Every worker maintains a local model $\mathbf{x}$. The $i$-th worker maintains ${\mathbf{x}}^{(i)}$. Every worker runs two threads in parallel. The first thread does gradient computation (called computation thread) and the other one does communication (called communication thread). For each worker $i$, a lock $m_i$ controls the access to its model.

The computation thread on the $i$-th worker repeats the following steps:

1. Acquire lock $m_i$.
2. Calculate a local gradient $\nabla F({\mathbf{x}}^{(i)})$ on a batch of input data.
3. Release lock $m_i$.
4. Update the model with local gradient, ${\mathbf{x}}^{(i)}={\mathbf{x}}^{(i)}-\gamma \nabla F({\mathbf{x}}^{(i)})$.

The communication thread on the $i$-th worker repeats the following steps:

1. Acquire lock $m_i$.
2. Average local model ${\mathbf{x}}^{(i)}$ with all other workers' models: ${\mathbf{x}}^{(i)}=\frac{1}{n}{\sum }_{j=1}^n{\mathbf{x}}^{(j)}$.
3. Release lock $m_i$.

Every worker run the two threads independently and concurrently.

Example usage

First initialize the Bagua algorithm (see API documentation for more options):

from bagua.torch_api.algorithms import async_model_average
algorithm = async_model_average.AsyncModelAverageAlgorithm()

Then use the algorithm for the model

model = model.with_bagua([optimizer], algorithm)

Unlike running synchronous algorithms, you need to stop the communication thread explicitly when the training process is done (for example when you want to run test):

model.bagua_algorithm.abort(model)

To resume the communication thread when you start training again, do:

model.bagua_algorithm.resume(model)

A complete example of running the Asynchronous Model Average algorithm can be found in Bagua examples with --algorithm async command line argument.