Improving Distributed Gradient Descent Using Reed-Solomon Codes

Introduction

With the size of today’s datasets, due to high computation and/or memory requirements, it is virtually impossible to run large-scale learning tasks on a single machine; and even if that is possible, the learning process can be extremely slow due to its sequential nature. Therefore, it is highly desirable or, even necessary, to run the tasks in a distributed fashion on multiple machines/cores. For this reason, parallel and distributed computing has attracted a lot of attention in recent years from the machine learning, and other, communities .

When a task is divided among a number of machines, the “computation time” is clearly reduced significantly, since the task is being processed in parallel rather than sequentially. However, the taskmaster has to wait for all the machines in order to be able to recover the exact desired computation. Therefore, in the face of substantial or heterogeneous delays, distributed computing may suffer from being slow, which defeats the purpose of the exercise. Several approaches have been proposed to tackle this problem. One naive yet common way, especially when the task consists of many iterations, is to not wait for all machines, and ignore the straggling machines. One may hope that in this way on average the taskmaster receives enough information from everyone; however, it is clear that the performance of the learning algorithm may be significantly impacted in many cases because of lost updates. An alternative and more appropriate way to resolve this issue, is to introduce some redundancy in the computation of the machines, in order to efficiently trade off computation time for less wait time, and to be able to recover the correct update using only a few machines. But the great challenge here is to design a clever scheme for distributing the task among the machines, such that the computation can be recovered using a few machines, independent of which machines they are.

Over the past few decades, coding theory has been developed to address similar challenges in other domains, and has had enormous success in many applications such as mobile communication, storage, data transmission, and broadcast systems. Despite the existence of a great set of tools developed in coding theory which can be used in many machine learning problems, researchers had not looked at this area until very recently . This work is aimed at bridging the gap between distributed machine learning and coding theory, by introducing a carefully designed coding scheme for efficiently distributing a learning task among a number of machines.

More specifically, we consider gradient-based methods for additively separable cost functions, which are the most common way of training any model in machine learning, and use coding to cleverly distribute each gradient iteration across $n$ machines in an efficient way. To that end, we propose a deterministic construction based on Reed-Solomon codes accompanied with an efficient decoder, which is used to recover the full gradient update from a fixed number of returning machines. Furthermore, we provide a new delay model based on heavy-tail distributions that also incorporates the time required for decoding. We analyze this model theoretically and use it to optimally pick our scheme’s parameters. We compare the performance of our method on the MNIST dataset with other approaches, namely: 1) Ignoring the straggling machines , 2) Waiting for all the machines, and 3) GradientCoding as proposed by Tandon et al. . Our numerical results show that, for the same training time, our scheme achieves better test errors.

As mentioned earlier, coding theory in machine learning is a relatively new area. We summarize the recent related work here. Lee et al. recently employed a coding-theoretic method in two specific distributed tasks, namely matrix multiplication and data shuffling. They showed significant speed-ups are possible in those two tasks by using coding. Dutta et al. proposed a method that speeds up distributed matrix multiplication by sparsifying the inner products computed at each machine. A coded MapReduce framework was introduced by Li et al in which is used to facilitate data shuffling in distributed computing. The closest work to our framework is the work of Tandon et al. , which aims at mitigating the effect of stragglers in distributed gradient descent using Maximum-Distance Separable (MDS) codes. However, no analysis of computation time was provided. Furthermore, in their framework, along with the above-mentioned works, the decoding was assumed to be performed offline which might be impractical in certain settings.

2 Statement of Contributions

In this work, we make the following three main contributions.

We construct a deterministic coding scheme for efficiently distributing gradient descent over a given number of machines. Our scheme is optimal in the sense that it can recover the gradient from the smallest possible number of returning machines, $f$, given a prespecified computational effort per machine.

We provide an efficient online decoder, with time complexity $O(f^{2})$ for recovering the gradient from any $f$ machines, which is faster than the best known method , $O(f^{3})$.

We analyze the total computation time, and provide a method for finding the optimal coding parameters. We consider heavy-tailed delays, which have been widely observed in CPU job runtimes in practice .

The rest of the paper is organized as follow. In Section 2, we describe the problem setup and explain the design objectives in detail. Section 3, provides the construction of our coding scheme, using the idea of balanced Reed-Solomon codes. Our efficient online decoder is presented in Section 4. We then characterize the total computation time, and describe the optimal choice of coding parameters, in Section 5. Finally, we provide our numerical results in Section 6, and conclude in Section 7.

Preliminaries

As it will be explained in detail, $k$ and $w$ are to be chosen in such a way that the total computation time is minimized. For a fixed $k$ and $w$, we want to be able to recover the gradient using the linear combinations received from the fastest $f$ machines at the master (or equivalently tolerate $s:=n-f$ stragglers). Note that we do not assume any prior knowledge about the stragglers, i.e., we shall design a scheme that enables master to recover the gradient from any set of $f$ machines. It is known that for any fixed $k$ and $w$, an upper-bound on the number of stragglers that any scheme can tolerate is:

The scheme proposed in this work achieves this bound. A coding scheme designed to tolerate $s$ stragglers consists of an encoding matrix $\mathbf{B}$, and a collection of decoding vectors $\{\mathbf{a}_{\mathcal{F}}:\mathcal{F}\subset[n],{|\mathcal{F}|}=n-s\}$. The matrix $\mathbf{B}$ should satisfy:

Each row of $\mathbf{B}$ contains exactly $w$ nonzero entries.

The linear space generated by any $f$ rows of $\mathbf{B}$ contains the all-one vector of length $k$, $\mathbf{1}_{1\times k}$.

The values of these nonzero entries prescribe the linear combination sent by $W_{i}$. In other words, the coded partial gradient sent from $W_{i}$ to $M$ is given by

When this holds for any set of indices $\mathcal{F}\subset[n]$ of size $f$, it means that the gradient can be recovered from the set of $f$ machines that return fastest. ,

2 Computational Trade-offs

In a distributed scheme that does not employ redundancy, the taskmaster has to wait for all the workers to finish in order to compute the full gradient. However, in the scheme outlined above, the taskmaster needs to wait for the fastest $f$ machines to recover the full gradient. Clearly, this requires more computation by each machine. Note that in the uncoded setting, the amount of computation that each worker does is $\frac{1}{n}$ of the total work, whereas in the coded setting each machine performs a $\frac{w}{k}$ fraction of the total work. From (2), we know that if a scheme can tolerate $s$ stragglers, the fraction of computation that each worker does is $\frac{w}{k}\geq\frac{s+1}{n}$. Therefore, the computation load of each worker increases by a factor of $(s+1)$. As will be explained further in Section 5, there is a sweet spot for $\frac{w}{k}$ (and consequently $s$) that minimizes the expected total time that the master waits in order to recover the full gradient update.

It is worth noting that it is often assumed that the decoding vectors are precomputed for all possible combinations of returning machines, and the decoding cost is not taken into account in the total computation time. In a practical system, however, it is not very reasonable to compute and store all the decoding vectors, especially as there are ${n\choose f}$ such vectors, which grows quickly with $n$. In this work, we introduce an online algorithm for computing the decoding vectors on the fly, for the indices of the $f$ workers that respond first. The approach is based on the idea of inverting Vandermonde matrices, which can be done very efficiently. In the sequel, we show how to construct an encoding matrix $\mathbf{B}$ for any $w,k$ and $n$, such that the system is resilient to $\left\lfloor\frac{wn}{k}\right\rfloor-1$ stragglers, along with an efficient algorithm for computing the decoding vectors $\{\mathbf{a}_{\mathcal{F}}:\mathcal{F}\subset[n],{|\mathcal{F}|}=f\}$.

Code Construction

The basic building block of our encoding scheme is a matrix $\mathbf{M}\in\{0,1\}^{n\times k}$, where each row is of weight $w$, which serves as a mask for the matrix $\mathbf{B}$, where $w$ is the number of data partitions that is assigned to every machine. Each column of $\mathbf{B}$ will be chosen as a codeword from a suitable Reed–Solomon Code over the complex field, with support dictated by the corresponding column in $\mathbf{M}$. Whereas the authors of choose the rows of $\mathbf{B}$ as codewords from a suitable MDS code, this approach does not immediately work when $k$ is not equal to $n$.

We will utilize techniques from to construct the matrix $\mathbf{M}$ (and then $\mathbf{B}$). For that, we present the following definition.

A matrix $\mathbf{M}\in\{0,1\}^{n\times k}$ is column (row)-balanced if for fixed row (column) weight, the weights of any two columns (rows) differ by at most 1.

While our scheme can handle the case where $\frac{nw}{k}$ is not an integer, in this section we illustrate the case where it is. The general result is described in the appendix.

Ultimately, we are interested in a matrix $\mathbf{M}$ with row weight $w$ that prescribes a mask for the encoding matrix $\mathbf{B}$. As an example, let $n=8$, $k=4$ and $w=3$. Then, $\mathbf{M}$ is given by

where each column is of weight $\frac{nw}{k}=6$. The following algorithm produces a balanced mask matrix. For a fixed column weight $d$, each row has weight either $\left\lfloor\frac{kd}{n}\right\rfloor$ or $\left\lceil\frac{kd}{n}\right\rceil$.

2 Reed–Solomon Codes

It is well-known that any $f$ rows of $\mathbf{G}$ form an invertible matrix, which implies that specifying any $f$ evaluations $\{t(\alpha^{i_{1}}),\ldots,t(\alpha^{i_{f}})\}$ of a polynomial $t(x)$ of degree at most $f-1$ characterizes it. In particular, fixing $f-1$ evaluations of the polynomial to zero characterizes $t(x)$ uniquely up to scaling. This property will give us the ability to construct $\mathbf{B}$ from $\mathbf{M}$.

3 Building the Encoding Matrix from the Mask Matrix

Once a mask matrix $\mathbf{M}$ has been determined using Algorithm 4, the encoding matrix $\mathbf{B}$ can be built by picking appropriate codewords from $\mathsf{RS}[n,f]$. Consider $\mathbf{M}$ in (5) and the following polynomials

The constant $\kappa_{j}$ is chosen such that the constant term of $t_{j}(x)$, i.e. $t_{j}(0)$, is equal to $1$. The evaluations of $t_{j}(x)$ on $\{1,\alpha,\ldots,\alpha^{7}\}$ are collected in the vector $(t_{j}(1),t_{j}(\alpha),\ldots,t_{j}(\alpha^{7}))^{\mathsf{T}}$ which sits as the $j^{\text{th}}$ column of $\mathbf{B}$. The validity of this process can be confirmed using (6), and is generalized in Algorithm 2.

Once the matrix $\mathbf{B}$ is specified, the corresponding decoding vectors required for computing the gradient at the taskmaster have to be characterized.

Efficient Online Decoding

We exploit the fact that $\mathbf{B}$ is constructed using Reed–Solomon codewords and show that each decoding vector $\mathbf{a}_{\mathcal{F}}$ can be computed in $O(f^{2})$ time. Recall that the taskmaster should be able to compute the gradient from any $f$ surviving machines, indexed by $\mathcal{F}\subseteq[n]$, according to (4). The $j^{\text{th}}$ column of $\mathbf{B}$ is determined by a polynomial $t_{j}(x)=\sum_{i=0}^{f-1}t_{j,i}x^{i}$ where $t_{j,0}=1$. We can write $\mathbf{B}$ as $\mathbf{B}=\mathbf{G}\mathbf{T}$, where $\mathbf{T}=\begin{bmatrix}\mathbf{t}_{1}&\cdots&\mathbf{t}_{k}\end{bmatrix}$ and $\mathbf{t}_{j}$ is the vector of coefficients of $t_{j}(x)$, and $\mathbf{G}$ is the matrix given in (6). Now consider $\mathbf{C}_{\mathcal{F}}$, the coded partial gradients received from $\{W_{i}:i\in\mathcal{F}\}$. The rows of $\mathbf{B}$ corresponding to $\mathcal{F}$ are given by

We require a vector $\mathbf{a}_{\mathcal{F}}$ such that $\mathbf{a}^{\mathsf{T}}_{\mathcal{F}}\mathbf{B}_{\mathcal{F}}=\mathbf{1}_{1\times k}$. This is equivalent to finding a vector $\mathbf{a}_{\mathcal{F}}$ such that

Indeed, the matrix $\mathbf{G}_{\mathcal{F}}$ in the above product is a Vandermonde matrix defined by $f$ distinct elements and so it is invertible in $O(f^{2})$ time , which facilitates the online computation of the decoding vectors. This is an improvement compared to previous works where the decoding time is usually $O(f^{3})$. A careful inspection of inverses of Vandermonde matrices built from an $n^{\text{th}}$ root of unity allows us to compute the required decoding vector in a space efficient manner. This is demonstrated in the next subsection.

Note that $\mathbf{a}_{\mathcal{F}}^{\mathsf{T}}$ is nothing but the first row of the inverse of $\mathbf{G}_{\mathcal{F}}$, which can be built from a set of polynomials $\{v_{1}(x),\ldots,v_{f}(x)\}$. Let the $l^{\text{th}}$ column of $\mathbf{G}^{-1}_{\mathcal{F}}$ be $\mathbf{v}_{l}=(v_{l,0},\ldots,v_{l,f-1})^{\mathsf{T}}$ and associate it with $v_{l}(x)=\sum_{i=0}^{f-1}v_{l,i}x^{i}$. The condition $\mathbf{G}_{\mathcal{F}}\mathbf{v}_{l}=\mathbf{e}_{l}$, where $\mathbf{e}_{l}$ is the $l^{\text{th}}$ elementary basis vector of length $f$, implies that $v_{l}(x)$ should vanish on $\{\alpha^{i_{1}},\ldots,\alpha^{i_{f}}\}\setminus\{\alpha^{i_{l}}\}$. Specifically,

The first row of $\mathbf{G}^{-1}_{\mathcal{F}}$ is given by $(v_{1,0},\ldots,v_{f,0})$, where $v_{l,0}$ is the constant term of $v_{l}(x)$. Indeed, we have $v_{l,0}=v_{l}(0)$, which can be computed in closed form according to the following formula,

By choosing $\alpha$ as a primitive $n^{\text{th}}$ root of unity, one is guaranteed that there are only $n-1$ distinct values of $(1-\alpha^{i_{l}-i_{j}})^{-1}$. This observation proposes that the master should precompute and store the set $\{(1-\alpha^{i})^{-1}\}_{i=1}^{n-1}$, and then compute each $v_{l,0}$ by utilizing lookup operations. The following algorithm outlines this procedure.

Analysis of Total Computation time

In this section, we provide a theoretical model which can be used to optimize the choice of parameters that define the encoding scheme. For this purpose, we model the response time of a single computing machine as

where the quantity $t_{0}$ can be thought of the fundamental delay of the machine, i.e. the minimum time required for a machine to return in perfect conditions. Previous works model the return time of a machine as a shifted exponential random variable. We propose using this approach since the heavy-tailed nature of CPU job runtime has been observed in practice .

Let $T_{f}$ denote the expected time of computing the gradient using the first $f$ machines. As a result we have

where $T_{\text{delay}}^{(f)}$ is the $f^{\text{th}}$ ordered statistic of $T_{\text{delay}}$, and $T_{\text{dec}}(f)$ is the time required at the taskmaster for decoding. Here we assume $n$ is large and define $\alpha:=\frac{w}{k}$ as the fraction of the dataset assigned to each machine. For this value of $\alpha$, the number of machines required for successful recovery of the gradient is given by

The expected value of the $f^{\text{th}}$ order statistic of the Pareto distribution with parameter $\xi$ will converge as $n$ grows, i.e.,

Using this result, we can approximate $T_{f}$, for $n\gg 1$,

where we assume that the taskmaster uses Algorithm 3 for decoding. If we assume $c_{m}$ is the time required for one FLOP, the total decoding time is given by $c_{m}(f-1)f\approx c_{m}(1-\alpha)^{2}n^{2}$. Since $\alpha$ is bounded from above by the memory of each machine, one can find the optimal computation time, subject to memory constraints, by minimizing $T_{f}$ with respect to $\alpha$.

In the schemes where the decoding vectors are computed offline, the quantity $T_{\text{dec}}$ does not appear in the total computation time $T_{f}$. Therefore, for large values of $n$, we can write:

This function can be minimized with respect to $\alpha$ by standard calculus to give

Note that this quantity is valid (less than one) if and only if one has $\frac{t_{0}}{c_{g}N\xi}<1$. It has been observed in practice that the parameter $\xi$ is close to one. Therefore, this assumption holds because $N$ is assumed to be large.

Numerical Results

To demonstrate the effectiveness of the scheme, we performed numerical simulations using MATLAB. We train a softmax regression model on a distributed cluster composed of $n=80$ machines to classify 10000 handwritten digits from the MNIST database while artificially introducing delay as a random variable sampled from a Pareto distribution according to (16) with parameters $\xi=1.1$ and $t_{0}=0.001$. Similar to , knowledge of the entire gradient allows us to employ accelerated gradient methods such as the one proposed by Nesterov . Details of the experiment are given in the accompanying description of Figure 2. We compare several schemes by running each of them on the same dataset for a fixed amount of time (in seconds) and then measuring the test error. The results depicted in Figure 2 demonstrate that the scheme proposed in this paper outperforms the four other schemes.

Conclusion

We presented a straggler mitigation scheme that facilitates the implementation of distributed gradient descent in a computing cluster. For a fixed per-machine computational effort, the taskmaster recovers the full gradient from the least number of machines theoretically required, which is done via an algorithm that is efficient in both space and time. Furthermore, we propose a theoretical delay model based on heavy-tailed distributions and incorporates the decoding time, which allows us to minimize the expected running time of the algorithm.

References

Appendix

To lighten notation, we prove correctness for $t=0$. The general case follows immediately.

Let $k,d$ and $n$ be integers where $d<n$. The row weights of matrix $\mathbf{M}\in\{0,1\}^{n\times k}$ produced by Algorithm 1 for $t=0$ are

Proof. The nonzero entries in column $j$ of $\mathbf{M}$ are given by

In case $n\mid kd$, each element in $\mathcal{S}$, after reducing modulo $n$, appears the same number of times. As a result, those indices correspond to columns of equal weight, namely $\frac{kd}{n}$. Hence, the two cases of $w_{i}$ are identical along with their corresponding index sets.

In the case where $n\nmid kd$, each of the first $\left\lfloor\frac{kd}{n}\right\rfloor n$ elements, after reducing modulo $n$, appears the same number of times. As a result, the nonzero entries corresponding to those indices are distributed evenly amongst the $n$ rows, each of which is of weight $\left\lfloor\frac{kd}{n}\right\rfloor$. The remaining indices $\{\left\lfloor\frac{kd}{n}\right\rfloor n,\ldots,kd-1\}_{n}$ contribute an additional nonzero entry to their respective rows, those indexed by $\{0,\ldots,(kd-1)_{n}\}$. Finally, we have that the first $(kd)_{n}$ rows are of weight $\left\lfloor\frac{kd}{n}\right\rfloor+1=\left\lceil\frac{kd}{n}\right\rceil$, while the remaining ones are of weight $\left\lfloor\frac{kd}{n}\right\rfloor$. $\square$

Now consider the case when $t$ is not necessarily equal to zero. This amounts to shifting (cyclically) the entries in each column by $t$ positions downwards. As a result, the rows themselves are shifted by the same amount allowing to conclude us the following.

Let $k,d$ and $n$ be integers where $d<n$. The row weights of matrix $M\in\{0,1\}^{n\times k}$ produced by Algorithm 1 are

2 General construction

The matrices $\mathbf{M}_{h}$ and $\mathbf{M}_{l}$ are constructed using Algorithm 1. Each column of $\mathbf{M}_{h}$ has weight $d_{h}:=\left\lceil\frac{nw}{k}\right\rceil$ and each column of $\mathbf{M}_{l}$ has weight $d_{l}:=\left\lfloor\frac{nw}{k}\right\rfloor$. Note that according to (2), we require $d_{l}\geq 2$ in order to tolerate a positive number of stragglers.

3 Correctness of Algorithm 4

According to the algorithm, the condition $k\mid nw$ implies that $k_{h}=0$ leading to $M=M_{l}$, which is constructed using Algorithm 1.

Moving on to the general case, the matrix $\mathbf{M}$ given by

where each matrix is row-balanced. The particular choice of $t$ in $\mathbf{M}_{l}$ aligns the “heavy” rows of $\mathbf{M}_{h}$ with the “light" rows of $\mathbf{M}_{l}$, and vice-versa. The algorithm works because the choice of parameters equates the number of heavy rows $n_{h}$ of $\mathbf{M}_{l}$ to the number of light rows $n_{l}$ of $\mathbf{M}_{h}$. The following lemma is useful in two ways.

$\left\lfloor\frac{k_{h}d_{h}}{n}\right\rfloor+\left\lceil\frac{k_{l}d_{l}}{n}\right\rceil=\left\lceil\frac{k_{h}d_{h}}{n}\right\rceil+\left\lfloor\frac{k_{l}d_{l}}{n}\right\rfloor=w$.

and conclude that $\left\lceil\frac{k_{h}d_{h}}{n}\right\rceil+\left\lfloor\frac{k_{l}d_{l}}{n}\right\rfloor=w$. $\square$

We have shown the concatenation of a “heavy” row of $\mathbf{M}_{h}$ along with a “light” row of $\mathbf{M}_{l}$ results in one that is of weight $w$. It remains to show that the concatenation of $\mathbf{M}_{h}$ and $\mathbf{M}_{l}$ results of rows of this type only.

We will assume that $n\nmid k_{h}d_{h}$ holds. From Proposition 2, we have $n_{l}=n-(k_{h}d_{h})_{n}$ and $n_{h}=(k_{l}d_{l})_{n}$. We will show that the two quantities are in fact equal. Indeed, we can express $n_{l}$ as

Hence $n_{l}=n_{h}$ and by the choice of $t$, the “light" rows of $\mathbf{M}_{h}$ align with the “heavy" rows of $\mathbf{M}_{l}$, and vice-versa. Furthermore, Lemma 1 guarantees that each row of $\mathbf{M}$ is of weight $\left\lceil\frac{k_{h}d_{h}}{n}\right\rceil+\left\lfloor\frac{k_{l}d_{l}}{n}\right\rfloor=w$. The same holds for the remaining rows, using the fact that $\left\lceil x\right\rceil+\left\lfloor y\right\rfloor=\left\lfloor x\right\rfloor+\left\lceil y\right\rceil$ when both $x$ and $y$ are non-integers.

4 Proof of Proposition 1

From , the expected value of the $f^{\text{th}}$ ordered statistic of the Pareto distribution is:

where $\Gamma(x)$ is the gamma function given by $\Gamma(x)=\int_{0}^{\infty}t^{x-1}e^{-t}dt$. We now assume that $n$ is large and make the standard approximation

Furthermore, (2) implies that the number of machines we wait for is $f=(1-\alpha)n$, for some $\alpha<1$ which leads to

By letting $n\rightarrow\infty$, the first two terms in the product converge to $e^{-}\xi$ and $e^{\xi}$, respectively, which yields

Offline Decoding

For illustrative purposes, we plot the function $T_{f}$ from (21) for a given set of parameters and indicate the optimal point. This plot is given in Figure 3.