COSTA: Communication-Optimal Shuffle and Transpose Algorithm with Process Relabeling

Kabic, Marko; Pintarelli, Simon; Kozhevnikov, Anton; VandeVondele, Joost

doi:10.1007/978-3-030-78713-4_12

Cited by 4 publications

(3 citation statements)

References 17 publications

Supporting

Mentioning

Contrasting

Order By: Relevance

“…Data distribution. COnf LUX and COnf CHOX provide ScaLA-PACK wrappers by using the highly-optimized COSTA algorithm [38] to transform the matrices between different layouts. In addition, they support the COSTA API for matrix descriptors, which is more general than ScaLAPACK's layout, as it supports matrices distributed in arbitrary grid-like layouts, processor assignments, and local blocks orderings.…”

Section: Methodsmentioning

confidence: 99%

On the parallel I/O optimality of linear algebra kernels

Kwasniewski

Kabic

Ben-Nun

et al. 2021

Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis

Self Cite

View full text Add to dashboard Cite

Matrix factorizations are among the most important building blocks of scientific computing. However, state-of-the-art libraries are not communication-optimal, underutilizing current parallel architectures. We present novel algorithms for Cholesky and LU factorizations that utilize an asymptotically communication-optimal 2.5D decomposition. We first establish a theoretical framework for deriving parallel I/O lower bounds for linear algebra kernels, and then utilize its insights to derive Cholesky and LU schedules, both communicating N 3 /(P √ M) elements per processor, where M is the local memory size. The empirical results match our theoretical analysis: our implementations communicate significantly less than Intel MKL, SLATE, and the asymptotically communication-optimal CANDMC and CAPITAL libraries. Our code outperforms these state-of-the-art libraries in almost all tested scenarios, with matrix sizes ranging from 2,048 to 524,288 on up to 512 CPU nodes of the Piz Daint supercomputer, decreasing the time-to-solution by up to three times. Our code is ScaLAPACK-compatible and available as an open-source library.

show abstract

Section: Methodsmentioning

confidence: 99%

On the parallel I/O optimality of linear algebra kernels

Kwasniewski

Kabic

Ben-Nun

et al. 2021

Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis

Self Cite

View full text Add to dashboard Cite

show abstract

“…Each device computes the attention matrix for all attention heads but only with the samples in their own partitions. Next, we use COSTA [28] to efficiently shuffle partitions so that each device has the attention scores for all samples but with only 1/N heads. This step prepares for the softmax calculation that requires all examples for each embedding dimension, and it can be computed sequentially by each device.…”

Section: Distributed Attentionmentioning

confidence: 99%

A Unified View of Long-Sequence Models towards Modeling Million-Scale Dependencies

Hè¹,

Kabic²

2023

Preprint

View full text Add to dashboard Cite

Nomenclature ΣCovariance matrix G Gram/kernel matrix k(•)Kernel function P(•) Probability density P(•)Token mixing process Re(•) Function that extracts the real component of a complex numberElement at ith position of column vector a A * :jColumn vector in jth row of A A i,jElement in ith row jth column ofmatrix of the embedding dimension F s L×L Vandermonde matrix of the sequence dimension W Weight matix learned with element-wise non-linearity (e.g., ReLU, GELU) W C L×L Weight matix of a single convolution kernel W K D×N Weight matix of attention key (for self-attention, N = M ) W Q D×M Weight matix of attention query W V D×M Weight matix of attention value X Resulting tokens with inductive bias introduced into X X L×D Input sequence of length L and embedding dimension D, where L D * Correspondence to

show abstract

“…scale, transpose or conjugate the received package S i,id if needed 6 Implementation Details COSTA (Algorithm 3) is implemented using the hybrid MPI+OpenMP parallelization model. The code is publicly available under the BSD-3 Clause Licence at [11]. It has the following features: 1) provides the ScaLAPACK wrappers for pxgemr2d and pxtran; 2) supports arbitrary grid-like matrix layouts (not limited to block-cyclic).…”

mentioning

confidence: 99%

COSTA: Communication-Optimal Shuffle and Transpose Algorithm with Process Relabeling

Kabic¹,

Pintarelli²,

Kozhevnikov³

et al. 2021

Preprint

Self Cite

View full text Add to dashboard Cite

Communication-avoiding algorithms for Linear Algebra have become increasingly popular, in particular for distributed memory architectures. In practice, these algorithms assume that the data is already distributed in a specific way, thus making data reshuffling a key to use them. For performance reasons, a straightforward all-to-all exchange must be avoided. Here, we show that process relabeling (i.e. permuting processes in the final layout) can be used to obtain communication optimality for data reshuffling, and that it can be efficiently found by solving a Linear Assignment Problem (Maximum Weight Bipartite Perfect Matching). Based on this, we have developed a Communication-Optimal Shuffle and Transpose Algorithm (COSTA): this highly-optimised algorithm implements A = α • op(B) + β • A, op ∈ {transpose, conjugate-transpose, identity} on distributed systems, where A, B are matrices with potentially different (distributed) layouts and α, β are scalars. COSTA can take advantage of the communication-optimal process relabeling even for heterogeneous network topologies, where latency and bandwidth differ among nodes. Moreover, our algorithm can be easily generalized to even more generic problems, making it suitable for distributed Machine Learning applications. The implementation not only outperforms the best available ScaLAPACK redistribute and transpose routines multiple times, but is also able to deal with more general matrix layouts, in particular it is not limited to block-cyclic layouts. Finally, we use COSTA to integrate a communication-optimal matrix multiplication algorithm into the CP2K quantum chemistry simulation package. This way, we show that COSTA can be used to unlock the full potential of recent Linear Algebra algorithms in applications by facilitating interoperability between algorithms with a wide range of data layouts, in addition to bringing significant redistribution speedups.

show abstract

COSTA: Communication-Optimal Shuffle and Transpose Algorithm with Process Relabeling

Cited by 4 publications

References 17 publications

On the parallel I/O optimality of linear algebra kernels

On the parallel I/O optimality of linear algebra kernels

A Unified View of Long-Sequence Models towards Modeling Million-Scale Dependencies

COSTA: Communication-Optimal Shuffle and Transpose Algorithm with Process Relabeling

Contact Info

Product

Resources

About