version 1.24.19 documentationParallelization of Lazylinops using pmatmat#
First of all, you need to install Lazylinop if not done already, see the installation guide. If you’re not yet introduced to the basic use of Lazylinop, we invite you to jump into the library by following this quick start guide.
In this notebook we will put the focus on LazyLinOp parallelization and precisely on the pmatmat function which is a function to automatically parallelize a matmat implementation. Indeed, the multiplication is not defaultly parallelized in LazyLinOp.
The pmatmat function is fully documented in the Lazylinop’s API documentation. However the parallelization path is full of traps and for someone not totally aware it can rapidly become counterproductive. Hence in this notebook we propose some hints about when and how to parallelize a LazyLinOp and when not to.
1. The basic principle and use of pmatmat parallelization#
First thing first, let us explain what exactly is parallelized in a LazyLinOp when using pmatmat. As the function name suggests, the operation which is parallelized is the multiplication of a LazyLinOp by an array (also known as the matmat function). Hence the letter p in pmatmat stands for parallelization.
So, in principle, how does pmatmat parallelize the multiplication L @ a ?#
A bit of notation before leading to the response: L is any LazyLinOp and a is a \(m \times n\) 2d-array with \(n > 1\). Note that it can in fact be any N-dimensional numpy.ndarray, but in this notebook we are limiting to the 2d case. The N-dimensional case is a sequential repetition of the 2d-case which is itself subject to parallelization.
Let us give an insight of the parallelization principle behind pmatmat. You have multiple workers available, for example CPU cores on your workstation, each one of them or a subgroup only can be assigned a number of columns of the array a for which it is responsible to compute the multiplication. More precisely, if you have a number of \(w\) workers available, pmatmat will assign the multiplication of \(s = \lfloor n / w \rfloor\) columns of \(a\) to each worker (the
nonzero remainder \(r = n - w s\) can additionally be distributed to the first workers because \(r < w\)). Hence, if \(r=0\) the worker \(0 \le i \lt w\) will compute the multiplication L @ a[:, s * i: s * (i+1)]. Because ideally all workers do the job in parallel a speedup is expected. At the end, all the block of columns are grouped together to form the whole product L @ a.
The question then is: what speedup should we expect?#
Well, it depends! In the simple and most efficient case you can expect an acceleration of a factor up to the number of workers you provided (supposedly here the number of CPU cores). But that is not as simple as that as stated in Amdhal’s law. Even though the multiplication of a bunch of columns is independent to one another, the configuration used might lower the speedup to something a way smaller than the number of cores effectively used. So we will give some hints to avoid the pitfalls and plainly take advantage of the hardware you have access to.
What is the simplest code to use pmatmat?#
pmatmat is written to be used very naturally, as follows:
Again, L is the LazyLinOp you want a parallelized version of.
pL = pmatmat(L)
pL is the parallelized LazyLinOp. It does all the same as L but pL @ a is computed in parallel.
However there are several possible configurations of the function. In the next we will present the specificities of each one of them, starting with a concrete and useful example of parallelization.
2. Thread-based parallelization of a LazyLinOp combining sparse matrices#
We shall present the basic use of pmatmat on a concrete case which is a LazyLinOp that results of the combination of SciPy sparse matrices using basic operators provided in the Lazylinop library (namely vstack, hstack, diag,
scalar multiplication and block_diag). We shall see that it is very easy to parallelize a complex LazyLinOp combining many operations.
Let’s first build the sequential LazyLinOp:
[1]:
import lazylinop as lz
import scipy.sparse as sp
import numpy as np
n = 768
# basic sparse matrices
a, b, c = [sp.random(n, n, .2) for _ in range(3)]
# intermediate building-block LazyLinOp-s
Lv = lz.vstack((a, b, c))
Lh = lz.hstack((Lv, 2 * Lv))
v = np.arange(Lh.shape[0])
Ld = lz.diag(v)
Ldh = Ld @ Lh
# final sequential LazyLinOp
L = lz.block_diag(Ldh, 2 * Ldh, 3 * Ldh, 4 * Ldh)
Now let’s build the parallel version with a simple pmatmat function call.
[2]:
from lazylinop.wip.parallel import pmatmat
pL = pmatmat(L)
No need to wait further to compare the computation times of the multiplication with the sequential L and parallelized pL.
[3]:
d = np.random.rand(L.shape[1], n * 4)
print("sequential L time:")
%timeit -n 1 -r 1 L @ d
sequential L time:
17 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)
[4]:
d = np.random.rand(L.shape[1], n * 4)
print("parallel pL time:")
%timeit -n 1 -r 1 pL @ d
parallel pL time:
8.76 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)
Not that bad! But let’s build the SciPy equivalent L matrix from scratch to verify if pL is still faster.
[5]:
spV = sp.vstack((a, b, c))
spH = sp.hstack((spV, 2 * spV))
spD = sp.spdiags(v, [0], len(v), len(v))
spHD = spD @ spH # costs something
spL = sp.block_diag((spHD, 2 * spHD, 3 * spHD, 4 * spHD))
# np.allclose(spL.toarray(), pL.toarray())
[6]:
print("Pure SciPy spL time:")
%timeit -n 1 -r 1 spL @ d
Pure SciPy spL time:
32.7 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)
Nice! We really won something using a parallel LazyLinOp.
pL @ d, that is the default method used by pmatmat (the two others are presented in next sections). Threads execute concurrently or preferably in parallel. They share the same memory space together in the same program (here Python), so they have the advantage to access/share the same data directly without any copy or transfer. Find more about multithreading on
Wikipedia.pmatmat uses as many threads that your configuration provides, generally 2 per CPU cores when hyperthreading is available. But in certain situations (for example if you are running something else while the computation is going on) it might be useful to use less threads, the pmatmat nworkers argument is here in that purpose.3. Process-based parallelization: workaround the Python GIL issue#
In section 2. we presented the multithread parallelization. Another one also proposed by pmatmat is the process-based parallelization.
So what is the difference between a thread and a process?
We usually say that a thread is a lightweight process. It is in fact practically very similar. Indeed, several processes can, as threads, run concurrently or in parallel. However they are considered heavy at the operating system level because they all have their own separate memory space. To be more accurate, at the time you create several processes from the parent process (by what we call a fork), they all share physically the same memory space even if virtually each one has its own. Until
the first write in the process virtual memory no data transfer is needed (a kind of OS lazyness!) but once you wrote data an extra computational cost is likely to happen compared to threads.
The Python GIL problem
LazyLinOp matmat function whose implementation is arbitrary, it is totally possible to be in a GIL situation.The example below shows such a matmat function. This is a toy example that computes manually a matrix chain product. In a real world case we would rely on a BLAS efficient implementation as in NumPy but the point here is to show what can happen about the parallelization of a pure python matmat. We run it sequentially (with L) and then parallelize it with threads (the default method used for ptL) and finally with method='process' (for ppL) in order to compare the
computation times.
Note: the next cell should not take more than 3 minutes of computation.
[7]:
from timeit import timeit
F = [np.random.rand(32, 32) for _ in range(50)]
F.reverse()
def matmat(x):
"""
"""
for f in F:
out = np.zeros((f.shape[0], x.shape[1]))
for i in range(f.shape[0]):
for j in range(f.shape[1]):
for k in range(x.shape[1]):
out[i, k] += f[i, j] * x[j, k]
x = out
return out
# the tranpose case use NumPy, not our concern here
L = lz.LazyLinOp(F[-1].shape, matmat=matmat, rmatmat=lambda y: a.T @ y)
# multithreaded L
ptL = pmatmat(L, method='thread')
# multiprocessed L
ppL = pmatmat(L, method='process')
# computational times
m = np.random.rand(F[-1].shape[1], 1024)
print("sequential time:", timeit(lambda: L @ m, number=1))
print("multithread time:", timeit(lambda: ptL @ m, number=1))
print("multiprocessed time:", timeit(lambda: ppL @ m, number=1))
sequential time: 40.43935673599481
multithread time: 71.55681976102642
multiprocessed time: 23.066745530988555
The results clearly demonstrate that Python multithreading is prone to the discussed GIL issue which leads to a counterproductive parallelization. It is also clear that the multiprocessing method is a solution to the GIL issue and provides here a significant speedup.
4. MPI-Based parallelization#
pmatmat. This is again a process-based parallelization but this one is based on Message Passing Interface (MPI).pmatmat(method='process') can only use the CPUs available on one system, but with MPI it is possible to use an arbitrary number of nodes of a computer center network (a grid).To use pmatmat(method='mpi') you will need to install mpi4py which is an optional dependency of Lazylinop. So contrary to multithreading and multiprocessing shown above it is not defaultly available.
[8]:
#!pip install mpi4py # uncomment to install
Once mpi4py is installed, you can create your own script based on the previous code given in 3. Just replace the end of the script as shown below:
test_mpi_pmatmat.py:
# ... insert the code in 3. here ...
mpiL = pmatmat(L, method='mpi')
# computational time
m = np.random.rand(F[-1].shape[1], 1024)
print("MPI time:", timeit(lambda: mpiL @ m, number=1))
It will allow to run only the MPI-based LazyLinOp.
Then to run it as a MPI parallel program (here only on a single computer) you might use mpiexec upon the Python interpreter. The command would be:
module load mpi # to load MPI for example on Linux using Lmod package
mpiexec python test_mpi_pmatmat.py
Defaultly it will use as many CPU cores as available on your computer. The running performance should be equivalent or even better than the one obtained with pmatmat(method='process').
5. When not to use pmatmat: oversubscribing#
If the LazyLinOp at work is based on an underlying library that handles pretty well the parallelization of the multiplication by itself you must absolutely avoid to parallelize furthermore (except if MPI is used because, as previously explained, in that case more CPUs can be used). Too much parallelization is called oversubscribing. The idea is that too many work is assigned to each core which are in return overwhelmed. It results obviously in a performance falling and makes pmatmat
counterproductive.
We show a very simple example with a basic NumPy array taken as a LazyLinOp. As you can guess NumPy already provides parallelization, so the corresponding pmatmat can hardly do better. Let’s see:
[9]:
m = np.random.rand(1024, 1024)
Lm = lz.aslazylinop(m)
pLm = pmatmat(Lm)
%timeit Lm @ m.T
%timeit pLm @ m.T
38 ms ± 7.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
57.8 ms ± 4.93 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
There is some ways to disable the parallelization used directly in NumPy (which can be parallelized with OpenMP, OpenBLAS or MKL) but doing your own threads parallelization using LazyLinOp won’t have any advantage relatively to NumPy parallelization. You can try as an exercise, numthreads is a way to disable
NumPy multithreading.
6. Benchmark the different methods#
Now that the pmatmat different methods have been introduced, we propose a benchmark to get a precise idea of their potential interest. So below is a benchmark of the main basic operators provided by Lazylinop and parallelized using pmatmat. It consists to measuring computation times of the parallelized LazyLinOp-array product varying the dimensions from \(2^3\) to \(2^{15}\). The cumulative time of 30 multiplications for each dimension and LazyLinOp are recorded and shown
in the figure below. pylops operators were also benchmarked using their default configuration (that is without pmatmat parallelization). We see clearly in the figure that multithread and MPI show a significant advantage.
Note that sequential, ‘thread’, ‘process’ and ‘pylops’ were all executed on the same CBP machine apollo4air3 using 16 workers (threads or processes when parallel). appolo4air3 CPU is: Intel(R) Xeon(R) Gold 5218 CPU @ 2.30GHz with 32 cores (and 64 threads for hyperthreading).
It is noteworthy that however ‘mpi’ was tested apart at PSMN computer center, Lake-flix partition. Configuration of nodes are: (CPU) Gold 6242 @ 2.8GHz, 32 cores, 384 GiB of RAM (12 GiB/core), Infiniband 56 GiB/s. A number of 14 nodes with 16 cores per node, that is a total of 224 cores were used to compute L @ A and outperform ‘thread’ method run on appolo4air3.
Scripts to reproduce the results/figure
sequential: pmatmat_baseline_benchmark.py
thread, process and pylops: pmatmat_multithread_multiproc_benchmark.py,
mpi: pmatmat_mpi_benchmark.py with the SLURM script to run the benchmark at PSMN sbatch_pmatmat_mpi_benchmark-lake-flix.sh.
figure: pmatmat_benchmark_plot.py
utility module: pmatmat_benchmark_utils.py
We acknowledge both of these computer centers for their work in running these nodes, please consult their websites linked above for more information.
Conclusion: the pmatmat recipe#
pmatmat is a function for the automatic parallelization of a LazyLinOp multiplication, proceeding by batch of columns, each one assigned to a CPU worker.a thread (
pmatmat(method='thread')),a process (
pmatmat(method='process')),a MPI process (
pmatmat(method='mpi')).
pmatmat is to verify that the underlying implementation of the LazyLinOp is not already efficiently parallelized. Indeed, a parallelization in that kind of situation might easily lead to oversubscribing (see 5.). Otherwise, if it is worth it, the obvious method to use on a single computer is multithreading. With the exception of
LazyLinOp-s whose matmat functions are mostly in Python, because a GIL issue might occur (see 3.). In this last case, the method='process' can be a workaround for an efficient parallelization.