DistStat.jl: Towards Unified Programming for High-Performance Statistical Computing Environments in Julia

The de facto standard for inter-node communication in distributed computing environments is the message passing interface (MPI). The latter defines several ways to communicating between two processes (point-to-point communication) or a group of processes (collective communication). Although MPI is originally defined in C and Fortran, many other high-level languages offer wrappers to this interface: Rmpi (Yu, , 2002) for R, mpi4py (Dalcín et al., , 2008) for Python, and MPI.jl (JuliaParallel Contributors, , 2020) for Julia, etc.

Leveraging on distributed computing environments, there have been several attempts to incorporate array and linear algebra operations through the basic syntax of the base programming language. MATLAB has a distributed array implementation that uses MPI as a backend in the Parallel Computing Toolbox. In Julia, MPIArrays.jl (Janssens, , 2018) defines a matrix-vector multiplication routine that uses MPI as its backend. DistributedArrays.jl (JuliaParallel Contributors, , 2019) is a more general attempt to create a distributed array, allowing various communication modes, including Transmission Control Protocol/Internet Protocol (TCP/IP) and Secure Shell (SSH), and MPI. In R, a package called ddR (Ma et al., , 2016) supports distributed array operations.

For GPU programming, CUDA C for Nvidia GPUs and OpenCL for generic GPUs are by far the most widely used. R package gputools (Buckner et al., , 2010) is one of the earliest efforts to incorporate GPU in R. PyCUDA and PyOpenCL (Klöckner et al., , 2012) for Python and CUDA.jl (previously CUDAnative.jl and CUDAdrv.jl) (Besard et al., , 2018) for Julia allow users to access low-level features of the respective interfaces.

An interface to array and linear algebra operations that works transparently on both CPU and GPU is highly desirable, as it reduces the programming burden tremendously. In MATLAB, the Parallel Computing Toolbox includes the data structure gpuArray. Simply wrapping an ordinary array with the function gpuArray() allows the users to use predefined functions to exploit the single instruction, multiple data (SIMD) parallelism of Nvidia GPUs. In Python, the recent deep learning (LeCun et al., , 2015) boom accelerated development of easy-to-use array interfaces and linear algebra operations along with automatic differentiation for both CPU and GPU. The most popular among them are TensorFlow (Abadi et al., , 2016) and PyTorch (Paszke et al., , 2019). The latter development is worth noting since most statistical computing tasks, not alone neural networks, can benefit from it. For example, the Distributions subpackage of TensorFlow (Dillon et al., , 2017) provides a convenient programming environment for Bayesian computing on GPUs, such as stochastic gradient Monte Carlo Markov chain (Baker et al., , 2018). In Julia, CUDA.jl (previously CuArrays.jl)(Besard et al., , 2019; JuliaGPU Contributors, , 2020) defines many array operations and simple linear algebra routines using the same syntax as the base CPU arrays.

To lessen the burden of programming for different HPC environments, it is desirable to unify these approaches and have a common syntax for within-node CPU/GPU computation and inter-node/inter-GPU communication. In this regard, the authors have previously developed a package dist_stat (Ko et al., , 2020) that implements distributed array and linear algebra operations using PyTorch. Through the examples of nonnegative matrix factorization, multidimensional scaling, and $\ell_{1}$-regularized Cox regression, they have demonstrated the scalability of their approach on both multiple GPUs and multiple CPU nodes on a cloud. This package provides a transparent and easy-to-use interface for distributed arrays on many computing environments, e.g., a single CPU-only node, a single node with a single GPU, a single node with multiple GPUs, multiple nodes with only CPUs, and multiple nodes with multiple GPUs. However, there are a couple of limitations in dist_stat, mainly coming from those of PyTorch. First, all the processes must have data of an equal size, hence the number of processes to be created must divide the total size of the data. This limite the tunability of the computation jobs. Second, it is very difficult to write effective device-specific code for further acceleration without writing code in C or other packages that support just-in-time (JIT) compilation. Hence the authors decided to migrate to Julia in order to avoid these issues and leverage higher flexibility.

In Julia, a function is “an object that maps a tuple of argument values to a return value”. A function can have many different specific implementations, distinguished by the types of input arguments. Each specific implementation is called a method, and a method of a function is selected by the mechanism called multiple dispatch. Many core functions in Julia possess several methods attached to each of them. A user can also define additional methods to existing functions. For example, a method for function f can be defined as follows:

julia> f(x, y) = "foo"
f (generic function with 1 method)

Each argument can be constrained to certain type, for example:

julia> f(x::Float64, y::Float64) = x * y
f (generic function with 2 methods)

julia> f(x::String, y::String) = x * y
f (generic function with 3 methods)

Float64 is the data type for a double-precision (64-bit) floating point number. An asterisk (*) between two String objects means concatenating them. At runtime, the most specific method is used for the given combination of input arguments.

julia> f("Candy", 3.0)
"foo"

julia> f("test", "me")
"testme"

julia> f(2.0, 3.0)
6.0

Methods and types may have parameters, enclosed by a pair of curly braces ({}). A parametric method is defined as follows:

julia> g(x::T, y::T) where {T <: Real} = x * y
g (generic function with 1 method)

The function g() performs multiplication of the two arguments if the two arguments are the same subtype of Real (a type for real numbers, for example, Float64, Int32 (32-bit integer), etc.) and both are of the same type.

julia> g(2.0, 3.0)
6.0

julia> g(2, 3)
6

julia> g(2.0, 3)
ERROR: MethodError: no method matching g(::Float64, ::Int64)
Closest candidates are:
  g(::T<:Real, ::T<:Real) where T<:Real at REPL[17]:1
Stacktrace:
 [1] top-level scope at REPL[28]:1

The third command throws an error, because the two arguments have different types. Such an error can be avoided by defining a more general method:

julia> g(x::Real, y::Real) = x * y
g (generic function with 2 methods)

Here, the exact type of x and y may be different. An example of parametric types, AbstractArray, is discussed in Section 3.2.

An array in Julia is defined as “a collection of objects stored in a multi-dimensional grid”. Each object should be of a specific type for optimized performance, such as Float64, Int32, or String.

The top-level abstract type for a multidimensional array is AbstractArray{T,N}, where parameter T is the type of element (such as Float64, Int32), and N is the number of dimensions. AbstractVector{T} and AbstractMatrix{T} are aliases for AbstractArray{T, 1} and AbstractArray{T, 2}, respectively. Operations for AbstractArrays are provided as fallback methods which would generally work correctly in many cases, but are often slow.

The type DenseArray is a subtype of AbstractArray representing an array stored in contiguous CPU memory. The most frequently used instance of DenseArray is Array, a type for basic CPU array with a grid structure. Vector{T} and Matrix{T} are aliases for Array{T, 1} and Array{T, 2}, respectively. Another subtype of DenseArray is CuArray, defined in CUDA.jl, a contiguous array data type on a CUDA GPU. Many of array operations for CuArray are provided using the same syntax as Arrays.

A Matrix (or an instance of Array{T, 2}) is easily created using a MATLAB-like syntax such as:

julia> A = [1 2; 3 4]
22 Array{Int64,2}:
 1  2
 3  4

An Array can be allocated with undefined values using:

julia> B = Array{Float64}(undef, 2, 3)
23 Array{Float64,2}:
 6.90922e-310  6.90922e-310  6.90922e-310
 6.90922e-310  6.90922e-310  6.90921e-310

There are predefined basic functions for array operations, such as size(A) that returns a tuple of dimensions of A, eltype(A) that returns the type of elements in A, and ndims(A), that shows the number of dimensions of A.

Linear algebra operations in Julia are defined in the basic package LinearAlgebra. The functions in LinearAlgebra can be loaded to the workspace with the keyword using:

julia> using LinearAlgebra

Matrix multiplication in Julia is defined in the function LinearAlgebra.mul!(C, A, B)¹¹1It is a convention in Julia to end the name of a function that changes the value of its arguments with an exclamation mark (!).. This function computes the post-multiplication of matrix B to matrix A, and stores the result in matrix C. The most general definition of LinearAlgebra.mul!() is:

LinearAlgebra.mul!(C::AbstractMatrix, A::AbstractVecOrMat,
                    B::AbstractVecOrMat)

which implements a naive algorithm for matrix multiplication. For a Matrix stored in the CPU memory, the call to LinearAlgebra.mul!() with arguments

LinearAlgebra.mul!(C::Matrix, A::Matrix, B::Matrix)

invokes the gemm (general matrix multiplication) routine of the BLAS, or the basic linear algebra subprograms (Blackford et al., , 2002), e.g.,

julia> A= [1. 2.; 3. 4.]
22 Array{Float64,2}:
 1.0  2.0
 3.0  4.0

julia> B = [3. 4.; 5. 6.]
22 Array{Float64,2}:
 3.0  4.0
 5.0  6.0

julia> C = Array{Float64, 2}(undef, 2, 2)
22 Array{Float64,2}:
 6.90922e-310  6.90921e-310
 6.90922e-310  6.90922e-310

julia> mul!(C, A, B)
22 Array{Float64,2}:
 13.0  16.0
 29.0  36.0

julia> C
22 Array{Float64,2}:
 13.0  16.0
 29.0  36.0

On the other hand, for matrices on GPU, a call to the same function LinearAlgebra.mul!() with arguments

LinearAlgebra.mul!(C::CuMatrix, A::CuMatrix, B::CuMatrix)

results in operations using cuBLAS (NVIDIA, , 2013), a CUDA version of BLAS:

julia> using CUDA

julia> A_d = cu(A)
22 CuArray{Float32,2,Nothing}:
 1.0  2.0
 3.0  4.0

julia> B_d = cu(B)
22 CuArray{Float32,2,Nothing}:
 3.0  4.0
 5.0  6.0

julia> C_d = cu(C)
22 CuArray{Float32,2,Nothing}:
 0.0  0.0
 0.0  0.0

julia> mul!(C_d, A_d, B_d)
22 CuArray{Float32,2,Nothing}:
 13.0  16.0
 29.0  36.0

The function cu() transforms an Array{T, N} into a CuArray{Float32, N}, transferring data from CPU to GPU.

Julia has a special “dot” syntax for vectorization. The dot syntax is invoked by prepending a dot to an operator (e.g., .+) or appending a dot to a function name (e.g., soft_threshold.()). Unlike many other programming languages (e.g., R), vectorization in Julia can be applied to any function without a need to deliberately tailor the corresponding method. Julia’s JIT compiler automatically matches singleton dimensions of array arguments to the dimensions of other array arguments. For example,

julia> a = [1, 2]
2-element Array{Int64,1}:
 1
 2

julia> b = [3 4; 5 6]
22 Array{Int64,2}:
 3  4
 5  6

julia> a .+ b
22 Array{Int64,2}:
 4  5
 7  8

Note that a is a column vector and b is a matrix.

The dot syntax can be extended by defining the method broadcast() for each array interface, allowing its generalization to any underlying hardware architecture. In addition, multiple dots on the same line of code fuse into a single call to broadcast(), translated into a single vectorized loop (for CPU) or a single generated kernel (for GPU) for that line.

While broadcasting is one of the simplest way to specifying elementwise operations, it is often not the fastest option. Broadcasting often allocates excessive memory, thus well-optimized compiled loops without memory allocation may be faster in many cases.

In DistStat.jl, a distributed array data type MPIArray{T,N,AT} is defined. Here, parameter T is the type of the elements of the array, e.g., Float64 or Float32. Parameter N is the dimension of the array, 1 for vector and 2 for matrix, etc. Parameter AT (for “array type”) is the implementation of AbstractArray used for base operations: Array for the basic CPU array, and CuArray for the arrays on Nvidia GPUs (which require CUDA.jl). If there are multiple CUDA devices, a device is assigned to a process automatically by the rank of the process modulo the size. This assignment scheme extends to the setting in which there are multiple GPU devices in multiple CPU nodes. The type MPIArray{T,N,AT} is a subtype of AbstractArray{T, N}. In MPIArray{T,N,AT}, each rank holds a contiguous block of the full data in AT{T,N} split and distributed along the N-th or the last dimension of an MPIArray.

In the special case of two-dimensional arrays, aliased MPIMatrix{T,AT} and created using the function transpose(), the data is column-major ordered and column-split. The transpose of this matrix has type of Transpose{T,MPIMatrix{T,AT}}, which is a lazy wrapper representing row-major ordered and row-split matrix. There also is an alias for one-dimensional array MPIArray{T,1,AT}, which is MPIVector{T,AT}.

a = MPIArray{Float64, 2, Array}(undef, 3, 4)

a.localarray

if DistStat.Rank() == 0
    dat = [1, 2, 3, 4]
else
    dat = Array{Int64}(undef, 0)
end
d = distribute(dat)

fill!(a, 0)

using Random
rand!(a)

function rand!(a::MPIArray{T,N,A}; seed=nothing, common_init=false,
    root=0) where {T,N,A}

The “dot” broadcasting feature of DistStat.jl follows the standard Julia syntax. This syntax provides a convenient way to operate on both multi-node clusters and multi-GPU workstations with the same code. For example, the soft-thresholding operator, which appears commonly in sparse regression, can be defined in the element level:

function soft_threshold(x::T, lambda::T)::T where T <: AbstractFloat
    x > lambda && return (x - lambda)
    x < -lambda && return (x + lambda)
    return zero(T)
end

This function can be applied to each element of an MPIArray using the dot broadcasting, as follows. When the dot operation is used for an MPIArray{T,N,AT}, it is naturally passed to inner array implementation AT. Consider the following arrays filled with random numbers from the standard normal distribution:

a = MPIArray{Float64, 2, Array}(undef, 2, 4) |> randn!
b = MPIArray{Float64, 2, Array}(undef, 2, 4) |> randn!

The function soft_threshold() is applied elementwisely as the following:

a .= soft_threshold.(a .+ 2 .* b, 0.5)

The three dot operations, .=, .+, and .*, are fused into a single loop (in CPU) or a single kernel (in GPU) internally.

A singleton non-last dimension is treated as if the array is repeated along that dimension, just like Array operations. For example,

c = MPIArray{Float64, 2, Array}(undef, 1, 4) |> rand!
a .= soft_threshold.(a .+ 2 .* c, 0.5)

works as if c were a $2\times 4$ array, with its content repeated twice. The terminal dimension is a little bit subtle, as an MPIArray{T,N,AT} is distributed along that dimension. Dot broadcasting along this direction works if the broadcast array has the type AT and holds the same data across the processes. For example,

d = Array{Float64}(undef, 2, 1); fill!(d, -0.1)
a .= soft_threshold.(a .+ 2 .* d, 0.5)

As with any dot operations in Julia, the dot operations for DistStat.jl are convenient but usually not the fastest option. Its implementations can be further optimized by specializing in specific array types. An example of this is given in Section 5.3.

Reduction operations, such as sum(), prod(), maximum(), minimum(), and accumulation operations, such as cumsum(), cumsum!(), cumprod(), cumprod!(), are implemented just like their base counterparts. Example usages of sum() and sum!() are:

sum(a)
sum(abs2, a)
sum(a, dims=1)
sum(a, dims=2)
sum(a, dims=(1,2))
sum!(c, a)
sum!(d, a)

The first line computes the elementwise sum of a. The second line computes the sum of squared absolute values (abs2() is the function that computes the squared absolute values). The third and fourth lines compute the column sums and row sums, respectively. Similar to the dot operations, the third line reduces along the distributed dimensions, and returns a broadcast local Array. The fifth line returns the sum of all elements, but the data type is a $1\times 1$ MPIArray. The syntax sum!(p, q) selects which dimension to reduce based on the shape of p, the first argument. The sixth line computes the columnwise sum and saves it to c, because c is a $1\times 4$ MPIArray. The seventh line computes rowwise sum, because d is a $2\times 1$ local Array.

Given below are examples for cumsum() and cumsum!():

cumsum(a; dims=1)
cumsum(a; dims=2)
cumsum!(b, a; dims=1)
cumsum!(b, a; dims=2)

The first line computes the columnwise cumulative sum, and the second line computes the rowwise cumulative sum. So do the third and fourth lines, but save the results in b, which has the same size as a.

Given a matrix $X\in\mathbb{R}^{m\times n}$ that consists of nonnegative values, NMF approximates $X$ by a product of two rank-$r$ ($r\ll\min\{m,n\}$) nonnegative matrices $V\in\mathbb{R}^{m\times r}$ and $W\in\mathbb{R}^{r\times n}$. This method has been applied in diverse disciplines, such as bioinformatics, recommender systems, astronomy, and image processing (Wang and Zhang, , 2013). In achieving this goal, consider minimizing the cost function

$\displaystyle f(V,W)$

$\displaystyle=\|X-VW\|_{\mathrm{F}}^{2},$

where $\|M\|_{\mathrm{F}}$ is the Frobenius norm of matrix $M$. A famous multiplicative algorithm to compute a local minimum of this cost function is due to Lee and Seung, (1999, 2001):

	$\displaystyle V_{(n+1)}$	$\displaystyle=V_{(n)}\odot XW_{(n)}^{\top}\oslash V_{(n)}W_{(n)}W_{(n)}^{\top}$
	$\displaystyle W_{(n+1)}$	$\displaystyle=W_{(n)}\odot V_{(n+1)}^{\top}X\oslash V_{(n+1)}^{\top}V_{(n+1)}^{\top}W_{(n)},$

where $\odot$ and $\oslash$ are elementwise multiplication and division, respectively. An alternative is the alternating projected gradient (APG) algorithm (Lin, , 2007):

	$\displaystyle V_{(n+1)}$	$\displaystyle=\max\{0,V_{(n)}-\sigma_{n}(V_{(n)}W_{(n)}W_{(n)}^{\top}-XW_{(n)}^{\top})\}$
	$\displaystyle W_{(n+1)}$	$\displaystyle=\max\{0,W_{(n)}-\tau_{n}(V_{(n+1)}^{\top}V_{(n+1)}W_{(n)}-V_{(n+1)}^{\top}X)\},$

where the maximum is applied in an elementwise fashion. The $\sigma_{n}$ and $\tau_{n}$ are the step sizes, for which we can choose $\sigma_{n}=1/2\|W_{(n)}W_{(n)}^{\top}\|_{\mathrm{F}}^{2}$ and $\tau_{n}=1/2\|V_{(n+1)}^{\top}V_{(n+1)}\|_{\mathrm{F}}^{2}$. Such selection is known to promote the algorithm to converge in fewer iterations than the multiplicative algorithm (Ko et al., , 2020).

An implementation of the multiplicative algorithm of NMF in DistStat.jl is

using DistStat, Random, LinearAlgebra
m = 10000; n = 10000; r = 20
T = Float64; A = Array
X  = MPIMatrix{T, A}(undef, m, n); rand!(X)
Vt = MPIMatrix{T, A}(undef, r, m); rand!(Vt)
W  = MPIMatrix{T, A}(undef, r, n); rand!(W)
WXt = MPIMatrix{T, A}(undef, r, m)
WWt = A{T}(undef, r, r)
WWtVt = MPIMatrix{T, A}(undef, r, m)
VtX = MPIMatrix{T, A}(undef, r, n)
VtV = A{T}(undef, r, r)
VtVW = MPIMatrix{T, A}(undef, r, n)
eps = 1e-10
for i in 1:iter
    mul!(WXt, W, transpose(X))
    mul!(WWt, W, transpose(W))
    mul!(WWtVt, WWt, Vt)
    Vt .= Vt .* WXt ./ (WWtVt .+ eps)
    mul!(VtX, Vt, X)
    mul!(VtV, Vt, transpose(Vt))
    mul!(VtVW, VtV, W)
    W .= W .* VtX ./ (VtVW .+ eps)
end

A main loop for the APG algorithm is given by:

for i in 1:iter
    mul!(WXt, W, transpose(X))
    mul!(WWt, W, transpose(W))
    mul!(WWtVt, WWt, Vt)
    sigma = 1.0 / (2 * (sum(WWt.^2)) + eps)
    Vt .= max.(0.0, Vt .- sigma .* (WWtVt .- WXt))
    mul!(VtX, Vt, X)
    mul!(VtV, Vt, transpose(Vt))
    mul!(VtVW, VtV, v.W)
    tau = 1.0 / (2 * (sum(VtV.^2)) + eps)
    W .= max.(0.0, W .- tau .* (VtVW .- VtX))
end

This code runs on a single CPU node, and can be modified to run on a GPU workstation (multiple GPUs are allowed) by importing the package CUDA.jl and changing A = Array to A = CuArray on the third line. This selection is made through a command-line argument for the full implementation (in the examples/ directory), where multi-node clusters are supported. Matrix $W$ (W) and the transpose of $V$ (Vt) are updated for each iteration. The intermediate results, WXt ($WX^{\top}$), WWt ($WW^{\top}$), WWtVt ($WW^{\top}V^{\top}$), VtX ($V^{\top}X$), VtV ($V^{\top}V$), and VtVW ($V^{\top}VW$) are preallocated. A very small number (eps) is added to the denominator for numerical stability.

Thanks to the transparent implementation of distributed arrays and mul!() in Section 4, the main loop is not different from what one would write with native Julia. The same code with slight changes in the first three lines run on various HPC environments including CPU clusters and multi-GPU workstations in a distributed fashion. In our numerical experiments, further optimization for memory efficiency was conducted.

Table 3 compares the performance of the two NMF algorithms on the multi-GPU setting in Table 2 with $10,000\times 10,000$ data for 10,000 iterations. It can be seen that the performance of the two algorithms is comparable, with APG being slightly slower with fixed number of iterations. This is because APG involves slightly more operations. With more than 4 GPUs, the communication burden outweighs the speed-up from using more GPU cores, and the algorithm slows down. The execution time between the dist_stat and DistStat.jl implementations are also largely comparable, with DistStat.jl version being faster in $r=20$ cases. Experiments with 3, 6, or 7 GPUs were not possible with dist_stat, because the size of the data was not divisible by 3, 6, and 7.

Table 4 compares the algorithms and implementations using $200,000\times 200,000$ data on multiple AWS EC2 instances for 1000 iterations. We used two processes per instance to avoid the communication burden. Once again, elapsed time was largely similar between the two algorithms. APG was faster than the multiplicative algorithms in more cases compared to the GPU case, because the multiplicative algorithm on CPU often suffers from slowdown due to creation of denormal numbers (Ko et al., , 2020). Between the two implementations, the DistStat.jl implementation was faster in 24 out of 30 cases.

MDS is a dimensionality reduction method that embeds a high-dimensional data $X=[x_{1},\dotsc,x_{m}]^{\top}\in\mathbb{R}^{m\times n}$ into $\boldsymbol{\theta}=[\theta_{1},\dotsc,\theta_{m}]^{\top}\in\mathbb{R}^{m\times q}$ in a lower dimensional Euclidean space ($q\ll n$) while keeping the distance measures as close as possible. We use the cost function

	$\displaystyle f(\theta)$	$\displaystyle=\sum_{i=1}^{m}\sum_{j\neq i}w_{ij}(y_{ij}-\|\theta_{i}-\theta_{j}\|_{2})^{2}$
		$\displaystyle=\sum_{i=1}^{m}\sum_{j\neq i}\left[-2w_{ij}y_{ij}\|\theta_{i}-\theta_{j}\|_{2}+w_{ij}\|\theta_{i}-\theta_{j}\|_{2}^{2}\right]+\mathrm{const.},$

where $y_{ij}$ is the distance measure between $x_{i}$ and $x_{j}$; and $w_{ij}$ are the weights. Using the majorization-minimization principle (de Leeuw and Heiser, , 1977; de Leeuw, , 1977), the update equation is given by

$\displaystyle\theta_{(n+1)ik}$

$\displaystyle=\left.\left(\sum_{j\neq i}w_{ij}\left[y_{ij}\frac{\theta_{(n)ik}-\theta_{(n)jk}}{\|\theta_{(n)i}-\theta_{(n)j}\|_{2}}+(\theta_{(n)ik}+\theta_{(n)jk})\right]\right)\right/\left(2\sum_{j\neq i}w_{ij}\right).$

for $i=1,\dotsc,m$ and $k=1,\dotsc,q$. The code below is a simple implementation of MDS in DistStat.jl, which runs on the same environment as the example of the previous subsection.

using DistStat, Random, LinearAlgebra
m = 1000; n = 1000; r = 20
T = Float64; A = Array; iter=1000
X = MPIMatrix{T, A}(undef, m, n); rand!(X)
Y = MPIMatrix{T, A}(undef, n, n)
theta = MPIMatrix{T, A}(undef, r, n); rand!(theta)
theta = 2theta .- 1.0 # range (-1, 1)
theta_distances = MPIMatrix{T, A}(undef, n, n)
theta_WmZ = MPIMatrix{T, A}(undef, r, n)
d_dist = MPIMatrix{T, A}(undef, 1, n) # dist. row vector
d_local = A{T}(undef, n) # bcast. col vector
DistStat.euclidean_distance!(Y, X) # compute pairwise dist
W_sums = convert(T, n - 1)
for i in 1:iter
    mul!(theta_distances, transpose(theta), theta)
    diag!(d_dist, theta_distances)
    diag!(d_local, theta_distances)
    theta_distances .= -2theta_distances .+ d_dist .+ d_local
    fill_diag!(theta_distances, Inf)
    Z = Y ./ theta_distances
    Z_sums = sum(Z; dims=1) # Z sums, length n dist. row vector.
    WmZ = 1.0 .- Z
    fill_diag!(WmZ, zero(T)) # weights of diagonal elements are zero
    mul!(theta_WmZ, theta, WmZ)
    theta .= (theta .* (Z_sums .+ W_sums) .+ theta_WmZ) ./ 2W_sums
end

This code can also run with local arrays if diag!() and fill_diag!() are defined for them. The function DistStat.euclidean_distance!(Y, X) computes the pairwise Euclidean distance matrix $Y$ between the input points in $X$, which incorporates the chunking method (Li et al., , 2010).

Table 5 compares the performance of DistStat.jl and dist_stat on multiple GPUs with $10,000\times 10,000$ data matrix $X$. Timing of dist_stat was faster when the number of GPUs employed was small. This is because dist_stat uses highly-optimized precompiled kernels from PyTorch, while DistStat.jl uses just-in-time compiled generated kernels. However, the gap between the two implementations vanishes as more GPUs are utilized.

For the AWS experiments, we used 36 processes per instance, because the step that mainly causes inter-instance communication is the matrix-vector multiplication, and its communication cost is much less than NMF. Note that this setting was not possible with the dist_stat implementation. Table 6 shows the runtime of each experiment for 1000 iterations on $100,000\times 1000$ dataset. It can be easily seen that DistStat.jl implementation is significantly faster.

Another example we consider is the Cox proportional hazards model-based regression (Cox, , 1972) with $\ell_{1}$-regularization. In this regression problem, the covariate matrix $X\in\mathbb{R}^{m\times n}$ is given, as well as the survival time for each sample, which is possibly right-censored. They are denoted by $y=(y_{1},\dotsc,y_{m})$ where $y_{i}=\min\{t_{i},c_{i}\}$, and $t_{i}$ is the time to event and $c_{i}$ is time to right censoring. Let $\delta_{i}=I_{\{t_{i}\leq c_{i}\}}$ indicate whether the event occured before the right censoring of sample $i$. The log-partial likelihood of the Cox proportional hazards model is defined by

$\displaystyle L(\beta)$

$\displaystyle=\sum_{i=1}^{m}\delta_{i}\left[\beta^{\top}x_{i}\log\left(\sum_{j:y_{j}\geq y_{i}}\exp(\beta^{\top}x_{j})\right)\right].$

We maximize $L(\beta)-\lambda\|\beta\|_{1}$. Applying the proximal gradient algorithm, which is equivalent to the iterative soft-thresholding algorithm (Beck and Teboulle, , 2009) to this problem, the iteration is given by:

	$\displaystyle w_{(n+1)i}$	$\displaystyle=\exp(x_{i}^{\top}\beta);\;\;W_{(n+1)j}=\sum_{i:y_{i}\geq y_{j}}w_{(n+1)i}$
	$\displaystyle\pi_{(n+1)ij}$	$\displaystyle=I(t_{i}\geq t_{j})w_{(n+1)i}/W_{(n+1)j}$
	$\displaystyle\Delta_{(n+1)}$	$\displaystyle=X^{\top}(I-P_{(n+1)})\delta,\text{ where $P_{(n+1)}=(\pi_{(n+1)ij})$}$
	$\displaystyle\beta_{(n+1)}$	$\displaystyle=\mathcal{S}_{\lambda}(\beta_{(n)}+\sigma\Delta_{(n+1)}),$

where $\mathcal{S}_{\lambda}(\cdot)$ is the soft-thresholding operator discussed in Section 4.2. If the samples are sorted in nonincreasing order of $y_{i}$, the summation $\sum_{i:y_{i}\geq y_{j}}w_{(n+1)i}$ can be implemented using the function cumsum(). Convergence of the algorithm is guaranteed if we choose the step size $\sigma=1/(2\|X\|_{2}^{2})$ (Ko et al., , 2020). A simple implementation of this algorithm in Julia, assuming no ties in $y_{i}$, can be written as:

using DistStat, LinearAlgebra, Random
m = 1000; n = 1000
T = Float64; A = Array; iter=1000
sigma = 0.00001 # step size
lambda = 0.00001 # penalty param
X = MPIMatrix{T,A}(undef, m, n); randn!(X)
delta = convert(A{T}, rand(m) .> 0.3) # 30% 1, 70% 0.
DistStat.Bcast!(delta) # bcast the delta from rank-0 proc.
y = convert(A{T}, collect(reverse(1:size(X,1))))
# assume decreasing observed time
beta = MPIVector{T,A}(undef, n)
pi_ind = MPIMatrix{T,A}(undef, m, m)
y_dist = distribute(reshape(y, 1, :))
fill!(pi_ind, one(T))
pi_ind .= ((pi_ind .* y_dist) .- y) .<= 0
Xbeta = A{T}(undef, m)
W = A{T}(undef, m)
pi_delta = A{T}(undef, m)
gradient = MPIVector{T, A}(undef, n)
for i in 1:iter
    mul!(Xbeta, X, beta)
    w = exp.(Xbeta)
    cumsum!(W, w)
    W_dist = distribute(reshape(W, 1, :))
    pi = pi_ind .* w ./ W_dist
    mul!(pi_delta, pi, delta)
    dmpd = delta .- pi_delta
    mul!(gradient, transpose(X), dmpd)
    beta .= soft_threshold.(beta .+ sigma .* gradient, lambda)
end

For performance optimization, note that in addition to the memory for $X$, an intermediate storage for two $m\times m$ matrices are needed for pi and pi_ind. This can be avoided by environment-specific implementation. For example, a CPU function to compute $P_{(n+1)}\delta$ can be written using LoopVectorization.jl (Elrod, , 2020) for efficient SIMD parallelization using the Advanced Vector Extensions (AVX, Firasta et al., , 2008). For GPU, a kernel directly computing $P_{(n+1)}\delta$ can be written using CUDA.jl. These environment-specific implementations not only use less memory, but also results in some speed-up. On the local workstation we used, the device-specific CPU implementation with four processes with each process using a single core took roughly half the time compared to the dot broadcasting-based implementation. The GPU implementation with four GPUs was 5-10% faster. Code for accelerating computation of $P_{(n+1)}\delta$ is avialable in Appendix B.

Table 7 demonstrates the scalability of the proximal gradient algorithm for $\ell_{1}$-regularized Cox regression on multiple GPUs. The speed-up with 8 GPUs is clear. While DistStat.jl was slightly slower than dist_stat with fewere GPUs, the gap is in general small. Unfortunately, the underlying algorithm for the cumsum() method in PyTorch is numerically unstable, and could not be used for very small values of $\lambda$. For this reason, Ko et al., (2020) resorted to double precision. On the other hand, the cumsum() function from CUDA.jl is numerically stable for these values of $\lambda$.

For the AWS experiments on DistStat.jl, 36 processes per instance was used once again. Table 8 shows the runtime of the algorithm for 1000 iterations with a simulated $100,000\times 200,000$ dataset. Thanks to the flexibility of the Julia implementation allowing 36 threads per instance, the speed-up of DistStat.jl over dist_stat, which used two threads per instance, is obvious.