VENOM: A vectorized N:M format for unleashing the power of sparse tensor cores

VENOM: A Vectorized N:M Format for Unleashing the Power of Sparse Tensor Cores Roberto L. Castro∗ [email protected] CITIC Universidade da Coruña A Coruña, Spain Andrei Ivanov [email protected] Department of Computer Science ETH Zürich Zürich, Switzerland Diego Andrade [email protected] CITIC Universidade da Coruña A Coruña, Spain Tal Ben-Nun [email protected] Department of Computer Science ETH Zürich Zürich, Switzerland Basilio B. Fraguela [email protected] CITIC Universidade da Coruña A Coruña, Spain Torsten Hoefler [email protected] Department of Computer Science ETH Zürich Zürich, Switzerland ABSTRACT The increasing success and scaling of Deep Learning models demands higher computational efficiency and power. Sparsification can lead to both smaller models as well as higher compute efficiency, and accelerated hardware is becoming available. However, exploiting it efficiently requires kernel implementations, pruning algorithms, and storage formats, to utilize hardware support of specialized sparse vector units. An example of those are the NVIDIA’s Sparse Tensor Cores (SPTCs), which promise a 2 × speedup. However, SPTCs only support the 2:4 format, limiting achievable sparsity ratios to 50%. We present the V:N:M format, which enables the execution of arbitrary N:M ratios on SPTCs. To efficiently exploit the resulting format, we propose Spatha, a high-performance sparselibrary for DL routines. We show that Spatha achieves up to 37 × speedup over cuBLAS. We also demonstrate a second-order pruning technique that enables sparsification to high sparsity ratios with V:N:M and little to no loss in accuracy in modern transformers. KEYWORDS Neural Networks, Pruning, GPGPU, CUDA, Sparse Tensor Cores ACM Reference Format: Roberto L. Castro, Andrei Ivanov, Diego Andrade, Tal Ben-Nun, Basilio B. Fraguela, and Torsten Hoefler. 2023. VENOM: A Vectorized N:M Format for Unleashing the Power of Sparse Tensor Cores. In The International Conference for High Performance Computing, Networking, Storage and Analysis (SC ’23), November 12–17, 2023, Denver, CO, USA. ACM, New York, NY, USA, 12 pages. https://doi.org/10.1145/3581784.3607087 1 INTRODUCTION The rapid progress of Deep Learning (DL) is revolutionizing Artificial Intelligence (AI) in areas such as Natural Language Processing ∗ Corresponding author: Roberto L. Castro ([email protected]), Universidade da Coruña, CITIC, Computer Architecture Group, 15071 A Coruña, Spain Publication rights licensed to ACM. ACM acknowledges that this contribution was authored or co-authored by an employee, contractor or affiliate of a national government. As such, the Government retains a nonexclusive, royalty-free right to publish or reproduce this article, or to allow others to do so, for Government purposes only. SC ’23, November 12–17, 2023, Denver, CO, USA ©2023 Copyright held by the owner/author(s). Publication rights licensed to ACM. ACM ISBN 979-8-4007-0109-2/23/11...$15.00 https://doi.org/10.1145/3581784.3607087 (NLP). Large Language Models (LLMs) are at the forefront of modern NLP systems [ 5 , 34 ]; however, their massive growth has led to unprecedented computational requirements [ 1 , 12 , 16 , 30 ]. As a result, training transformers has become a dominant task in DL, with costs reaching millions of dollars and significant energy and carbon emissions [ 32 ]. Therefore, it is critical to improve their inference and training performance. One of the most widely used techniques for this purpose is network pruning [ 13 ], which removes the less significant weights to produce simpler and compressed, yet accurate models. There is a plethora of pruning algorithms and sparse formats focused on accelerating tensor operations such as matrix-matrix multiplications (MMMs) by means of specialized hardware like Tensor Core Units (TCUs) [ 37 ]. While these algorithms and formats reduce the number of arithmetic operations and memory usage compared to their dense counterparts, achieving significant speedup on these accelerators while maintaining model accuracy is challenging [ 15 ]. Semi-structured pruning can yield practical speedups at moderate sparsity levels (e.g., 80 −90%) [2, 3, 23]. However, the irregularity of the sparse input matrices still limits performance and makes difficult to reach the theoretical peak considering the reduction of the number of arithmetic operations [9]. Last generations of NVIDIA GPUs include Sparse Tensor Cores (SPTCs) that are specifically designed for sparse computation [ 25 ]. SPTCs promise to accelerate math operations by up to 2 × at 50% sparsity. The data layout proposed to use SPTCs imposes strict constraints (i.e., 2:4 format, where every consecutive 4 elements have 2 nonzero values), but it reduces the irregularity of the sparse input w.r.t. other performance-aware sparse formats (e.g., vector-wise, block-wise). This makes the N:M format very suitable to execute on GPUs since it favors key aspects of the execution of tensor operations such as interand intra-warp load balance. However, there is an important limitation related to the usage of SPTCs and the 2:4 format: recent models like LLMs commonly have hundreds of millions to trillions of parameters, making it feasible to prune them to higher sparsity ratios with little or no loss in accuracy [ 19 ]. Unfortunately, there is currently no hardware support for executing arbitrary N:M formats with higher compression ratios, which limits the total achievable speedup. SC ’23, November 12–17, 2023, Denver, CO, USA Castro R.L., et al. Recent research has explored the N:M format [ 4 , 6 ]. However, these investigations have been limited to a theoretical perspective, such as network pruning, or have relied on CPU implementations due to a lack of hardware support for alternative N:M patterns on GPUs. To address these limitations, we propose the Vectorized N:M format, which we refer to as V:N:M 1 . This format introduces an abstraction layer over SPTCs, enabling the execution of alternative N:M formats and arbitrary sparsity ratios. The vectorization aspect is derived from the selection of vertical vectors of elements that are stacked together to provide the row-wise N:M pattern. This approach enables the conversion from generic N:M formats to the 2:4 that is accepted by SPTCs. To efficiently exploit the benefits of the V:N:M format, we propose Spatha 2 , a template-based library dedicated to general matrix-matrix multiplication on half precision where one of the operands is sparse (SpMM). Spatha serves as an open-source alternative to cuSparseLt [ 26 ] and removes its 2:4 restriction. The main contributions of this paper are: • A new sparse matrix format V:N:M which enables arbitrary N:M patterns on SPTCs. • Highly optimized SpMM kernels to efficiently exploit the V:N:M format. Specifically, we propose a template-based implementation that can be tuned depending on the input dynamics, such as GEMM size or the V:N:M format configuration. • A second-order pruning technique tailored for the V:N:M format and scalable to the dimensionality of LLMs. This technique allows the sparsification to high sparsity ratios with little to no loss in accuracy (e.g., ∼ 2% drop in BERT F1 score on the SQuAD dataset with 2:16 sparsity), which is required for the full exploitation of the V:N:M format. • Spatha achieves unprecedented speedups w.r.t. its dense counterpart versions (e.g., cuBLAS) yielding up to 37 × faster MMMs on matrices extracted from real-world DL models. Furthermore, Spatha implementation provides speedups of up to 1 . 38 × over the vendor library for 2:4 sparsity, cuSparseLt. • For end-to-end sparse LLMs inference, Spatha shows a GEMM time reduction of 11 × at 2:32 sparsity on real-world models such as GPT-3. 2 BACKGROUND This section presents the technical background of the paper, covering network pruning techniques and the Sparse Tensor Cores of NVIDIA GPUs. 2.1 Network pruning In DL, pruning is a technique used to reduce memory usage, which can also reduce the computational load when combined with compressed storage formats and efficient sparse kernels. Pruning techniques can be categorized based on various criteria, such as the pruning strategy employed, or the granularity of the pruning. Pruning schemes are often based on weight saliency metrics, which directly correlate with the expected impact on accuracy when those weights are removed from the network. Various methods exist 1Pronounced “venom” 2 SParse linear Algebra rouTines for High-performance Applications. The name is motivated by the analogy with the Cutlass library, with the accent on sparse computation - a sharp and efficient tool to cut through the complexity of sparse routines to select the candidate weights for removal, including magnitude pruning [ 18 ], which selects weights with lower absolute values, and gradient-based methods that use the gradient applied to each weight to identify those that are trending towards to zero faster. Within the gradient-based methods, we can find first-order techniques based on the first-derivative information [ 31 , 38 ], and second-order ones [ 7 , 19 , 21 ], which pursue to find the set of weights whose removal will generate a minimum loss increase in the network. Secondorder methods have proven to be effective in pruning convolutional networks in the past, but they have recently been optimized for Large Language Models (LLMs) [19]. As for the granularity of the pruning, unstructured methods [ 11 ] remove weights individually, with gradual magnitude pruning (GMP) being the most commonly used variant [ 8 ]. On the other end of the granularity spectrum, structured methods [ 24 , 35 ] prune complete components like layers, or heads, in the case of transformers networks [ 36 ]. In between, semi-structured methods prune groups of weights. These latter methods aim to balance performance and accuracy by defining specific formats that promote the exploitation of the underlying hardware more efficiently. These methods often imply the usage of tailored compressed storage formats and custom kernels [ 9 , 20 ]. The N:M format, which enables the use of Sparse Tensor Cores (SPTCs) in NVIDIA GPUs, can be classified in this last group. 2.2 Sparse Tensor Cores of NVIDIA GPUs The CUDA programming model organizes GPU kernels into three granularity levels: thread-blocks, warps, and threads. A thread block is composed of a set of warps, with warps being the basic scheduling unit in CUDA. Each warp consists of 32 threads. NVIDIA GPUs consist of an array of Streaming Multiprocessors (SMs), with all SMs sharing the L2 cache, and a DRAM memory, also called Global Memory (GMEM). Each SM is divided in processing blocks, each one having a Register File (RF), a warp scheduler, and an L0 instruction cache. All the processing blocks within an SM share a L1 cache, which is partially used as Shared Memory (SMEM). Each processing block is also equipped with four types of units: Floating-Point Units (FPU), Tensor-Core Units (TCU), Int Units (ALU) and Special Function Units (SFU). 0 1 2 3 4 5 6 7 0 3 1 2 1 2 0 1 0 1 2 3 4 5 6 7 x 0 1 2 3 4 5 6 7 0 1 2 3 0 3 1 2 Select x 0 3 5 6 0 1 2 3 4 5 6 7 Sparse Tensor Core 2:4 Mux 2:4 Mux A (Sparse) B (Dense) R K K/2 Non-zero values metadata indices K/2 2-bit indices C K result Compress Figure 1: The 2:4 format and its mapping to SPTCs Last generations of NVIDIA GPUs have extended their TCUs to also handle row-wise 2:4 sparsity. These updated TCUs include hardware support for sparse computation, and are referred to as Sparse Tensor Cores (SPTCs). To exploit SPTCs, the first argument in tensor operations must be stored in NVIDIA’s N:M sparse format, where 𝑁 represents the maximum number of non-zero elements in VENOM: A Vectorized N:M Format for Unleashing the Power of Sparse Tensor Cores SC ’23, November 12–17, 2023, Denver, CO, USA a block of 𝑀 values. Figure 1 illustrates this format. The left side of the figure shows an uncompressed sparse matrix following the row-wise 2:4 pattern. The compression of that 𝑅×𝐾 matrix requires two structures: (1) a 𝑅×𝐾/ 2matrix representing the values of the non-zero elements, and (2) a metadata structure which contains the position of each nonzero value within each group of 4values. Finally, Figure 1, right side, illustrates the mapping of a 2:4 sparse operation onto SPTCs. Notice that the metadata structure is also used by the hardware to select the corresponding elements in the dense matrix 𝐵 and perform the Matrix Multiply-Accumulate (MMA) operation. Precision Format Supported shapes fp32 1:2 𝑘8,𝑘16 half (fp16) 2:4 k32, k16 uint8 2:4 𝑘32,𝑘64 uint4 2:4 𝑘64,𝑘128 Table 1: Matrix Shapes for mma.sp on SPTCs. 𝑀 and 𝑁 dimensions are fixed to 16 and 8, respectively (𝑚16𝑛8) SPTCs can be accessed in CUDA using the NVPTX API which includes the mma.sp instruction. SPTCs support various shapes of this instruction depending on the data precision (Table 1). This instruction multiplies a 𝑚×𝑘 matrix by a 𝑘×𝑛 matrix, where 𝑚= 16, 𝑛= 8are fixed dimensions, and 𝑘 represents the sparsified dimension which can vary in size. This paper focuses on half precision kernels. Instruction shapes define the sizes of the left-hand-side (LHS) and the right-hand-side (RHS) operands as inputs to TCUs. For example, 𝑘= 32 implies that the LHS operand has a shape of 𝑚×𝑘= 16 × 32 while the RHS is 𝑘×𝑛= 32 × 8. It is important to note that the LHS is 50% sparse, meaning that its real size will be 16 ×16( 32 / 2 ) . NVIDIA’s notation for this instruction is 𝑚 16 𝑛 8 𝑘 32. 3 THE V:N:M FORMAT This section presents the new V:N:M format, which enables pruning to arbitrary N:M ratios retaining the use of SPTCs, which are designed to support only 2:4 patterns natively. 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 2:6 1 4 7 10 13 2 5 8 11 14 3 6 9 12 15 16 19 22 17 20 23 18 21 24 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 V Vector-wise N:M V:N:M 1 0 2 0 0 5 0 6 9 0 10 0 2:4 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 V Block-wise V 123 123 + + + 4 Figure 2: The V:N:M pruning procedure Sparse compression formats are of great significance in many HPC areas other than DL. However, the characteristics of the sparse matrices in DL workloads differ from those in other areas in several aspects [ 9 ]: (1) the sparsity level is generally much lower, (2) the number of non-zeros per row is higher and (3) the load imbalance is more pronounced. To address these challenges, ad-hoc solutions for DL workloads have been developed in two different planes: compression formats and pruning techniques, often interlinked. They seek the efficient exploitation of the hardware during the execution of tensor operations in DL workloads. A new area of research is focused on enhancing control over the distribution of non-zero elements in sparse matrices. This involves, for example, selecting 2D dense groups with size 𝑣×𝑣 (Figure 2, 1 ) or 1D groups of length 𝑣 , either row-wise or column-wise 2 . The aim is to create sparse matrices that are more regular, making them more suitable for efficient execution on GPUs. Block-based pruning techniques ( 1 and 2 ) are particularly useful on improving data reuse on L1 cache or registers during the multiplication of sparse matrices. Furthermore, optimized sparse formats, which compress their data, can be designed to facilitate traversal for the access patterns that arise during matrix multiplication [23, 28]. On the one hand, 1 can be overly aggressive in dropping blocks of elements, leading to a significant reduction in accuracy as the sparsity level increases. On the other hand, 2 offers more flexibility and enables higher sparsification ratios. However, using small vector lengths is a limiting factor to prevent accuracy loss (e.g., 𝑣≤ 8). Furthermore, in these approaches, the different number of elements per row can generate load imbalance and inherent negative effects such as thread divergence, inefficient memory transactions and low occupancy ratios. The N:M format 3 provides an alternative that overcomes most of the weaknesses of other performance-aware methods. Moreover, NVIDIA GPUs recently included hardware support for this format, but it is limited to 2:4. This paper introduces the new V:N:M format 4 which combines block-wise storage, and vector-wise and N:M pruning to enable the exploitation of SPTCs for arbitrary N:M patterns, leveraging higher compression ratios and reducing further the number of arithmetic operations required in MMMs. 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 V=4 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 0 2 0 1 1 3 1 2 0 2 2 3 1 2 1 2 1 2 0 2 0 1 1 3 2 3 0 3 0 1 0 2 0 2 5 7 1 2 5 7 0 3 6 7 1 2 3 6 Non-zero values m-indices 2bits/elem column-loc Sparse Matrix N:M=2:8 R KK/M*2 K/M*2 K/M*4 R/V Figure 3: The V:N:M compression format Figure 2, illustrates how this approach starts by partitioning the original dense matrix in blocks of 𝑉×𝑀 elements (block-wise). Then, the four most significant columns of each block are selected (vector-wise pruning), and for each row of four elements in a block, the two most meaningful weights are kept (2:4 pruning). These two levels of pruning (vector-wise and N:M) enable the exploitation of SPTCs for matrices with arbitrary levels of sparsity, as the vectorwise pruning stage diversifies the sparsity level, and N:M pruning imposes the restrictions required later by SPTCs. That is, in 4 , SC ’23, November 12–17, 2023, Denver, CO, USA Castro R.L., et al. the SPTC vector is 2:4, but it belongs to a 6-columns row, where 2 columns were fully pruned. It is actually an implementation of a 2:6 sparsity pattern that it is mapped onto SPTCs as the required 2:4. Finally, the data is represented using a new block-wise compression format shown in Figure 3. As for the NVIDIA 2:4 layout (Figure 1), the format requires an array with the non-zero values, and a 2-bit metadata index per non-zero (m-indices). Notice that now, each 2-bit metadata index refers to one of the 4columns that we have selected in each block and not to each column of the original dense input matrix (see 4 in Figure 2). Furthermore, the size of these two structures depends on the M value, more specifically their shape now is 𝑅×𝐾/𝑀× 2. This format requires a third structure column-loc of size 𝑅/𝑉×𝐾/𝑀× 4, that indicates which 4 columns (out of 𝑀 ) of each block were selected in the vector-wise pruning stage. Thus, while CSR storage overhead can reach up to 200% [ 25 ], V:N:M can be represented with 2-bit m-indices metadata per value, and up to 8-bit column-loc metadata per value every 𝑉 rows. For 𝑉=128, the overhead is ∼13.28%. 4 SPATHA: A HIGH-PERFORMANCE SPARSE LIBRARY FOR SPARSE MMM This section provides an in-depth description of the sparse kernel implementation associated to the V:N:M format, Spatha. The Sparse Matrix-Matrix multiplication (SpMM) is an important workload in DL that serves as the sparse counterpart to Matrix-Matrix Multiplication (MMM). This routine is widely used in various components of modern DL models. For instance, in the forward pass of a pruned model, the sparse weight matrix is multiplied by a dense activation matrix. Similarly, in transformers, the self-attention operation is performed by multiplying a sparse attention weight matrix by a dense one. Thus, optimizing this routine is crucial to improve the efficiency and the performance of our models. K v R Compress 4:2:8 sparsity (2 non-zeros out of 8) Select x ... Select Sparse Tensor Core Non-zero values m-indices column-loc K C A (Sparse) B (Dense) C (Dense) result Select matching K/M*4 elements out of K Figure 4: Mapping a 4:2:8 format onto Sparse Tensor Core (only native support to 2:4 format) Figure 4 shows an example of how the new V:N:M format (4:2:8 in the figure) is mapped onto SPTCs, which natively only support the 2:4 format. It shows how the SPTC is fed with the appropriate values from a row of the sparse matrix and a column of the dense matrix. The LHS operand is a 𝑅×𝐾/ 4dense matrix after having been pruned with sparsity of 75% (2:8). This pruning reduces the required multiply-and-add operations by 4(from 16 to 4), but also halves the rows loaded from the dense matrix B (selected by the values contained in column-loc). 4.1 Kernel design The design of an efficient CUDA kernel mostly depends on three main stages: (1) the efficient loading of the data to the top levels of the memory hierarchy (i.e., GMEM->SMEM->RF), (2) the computation, and (3) the storage of the results (i.e. RF->SMEM- >GMEM). Figure 5 covers stage 1, particularly the data movement from GMEM to RF, which is divided into 3 steps ( 11 - 13 ). Figure 6 focuses on stage 2, and shows how the data in the RF is mapped onto SPTCs in three steps ( 21 - 23 ). Finally, Figure 8 illustrates how stage 3 is performed (steps 31-32). ... 0 2 5 7 column-loc V V row0 row 2 row 5 row 7 < BSr=V < BSk < BSc Warp 0 V < WSr < WSk < BSk < R < K < C < WSc WSk Non-zero values m-indices (Global Memory) (Shared Memory) (Register File) batchSize 11 Pipelining (async) 1213 Blocked SpMM Thread Block Tile Warp Tile A (Sparse) B (Dense) C (Dense) Figure 5: Thread-Block Tile and Warp tile view (stage 1) Spatha is designed as a template-based library, where several parameters can be tuned depending on the input properties. Considering a 𝑅×𝐾×𝐶 GEMM problem, these parameters are: the threadblock tile size ( 𝐵𝑆𝑟×𝐵𝑆𝑘×𝐵𝑆𝑐 ), the warp tile size ( 𝑊 𝑆𝑟×𝑊 𝑆𝑘×𝑊 𝑆𝑐 ), the mma instruction shape ( 𝑀𝑀𝐴𝑟×𝑀𝑀𝐴𝑘×𝑀𝑀𝐴𝑐 ) and the level of memory pipelining (batchSize). 4.1.1 Stage 1-Data loading. Figure 5 shows the Spatha procedure to load the operands from GMEM onto RF. There are two dimensions to be taken into account: the data location (i.e., GMEM, SMEM, and RF), and the scope of this data from the NVIDIA programming model perspective (i.e., thread-block, and warp). Step 11 loads the column-loc structure from GMEM to SMEM with a two-level prefetching strategy. Note that the column-loc information is used to select the rows of B to be loaded from GMEM (Figure 5, left side) to SMEM (step 12 ). Pre-fetching this information breaks the data dependency with the activation matrix. Furthermore, column-loc is small, so it is convenient to load the information of multiple tiles together to maximize memory bandwidth. Next, step 12 loads the corresponding A and B tiles from GMEM to SMEM. Each threadblock is responsible for an output block of size 𝐵𝑆𝑟×𝐵𝑆𝑐 . More specifically, 𝐵𝑆𝑟=𝑉 , so each thread-block will load only the rows of B selected by the column-loc structure. In order to avoid memory stalls due to data dependencies with the next steps, we pipelined step 12 with step 13 and stage 2 (computation) taking advantage of CUDA asynchronous copies. The pipelining degree depends on the batchSize variable previously mentioned. Finally, in 13 , each warp is responsible for an output block of size 𝑊 𝑆𝑟×𝑊 𝑆𝑐 , so VENOM: A Vectorized N:M Format for Unleashing the Power of Sparse Tensor Cores SC ’23, November 12–17, 2023, Denver, CO, USA the corresponding tiles are loaded from SMEM to RF. Emphasize that all the previously mentioned memory transactions have been optimized to use 128-bit instructions. At this point, we also load directly to the RF the m-indices information. 4.1.2 Stage 2-Computation. When all the data is loaded in the RF, stage 2 starts, which performs the Matrix Multiply-Accumulate (mma.sp) on this data using SPTCs. Figure 6 shows a detailed view of stage 2 , depicting how the data in the RF is mapped onto SPTCs to be executed. Each warp has to break down the warp tile into instruction tiles, which depends on the instruction shapes available on SPTCs, in this example m16n8k32. The first step 21 , selects 𝑀𝑀𝐴𝑘= 16 elements from the warp tile and maps this data to SPTCs following step 22 layout. This layout represents the LHS fragment to the mma.sp instruction. That means that, if 𝑊 𝑆𝑟= 32, we will need to iterate twice over the rows of A’s warp tile. Similarly, the next step maps the B’s warp tile information into SPTCs following step 23 layout, which represents the RHS fragment to the mma.sp instruction. At this point, the mma.sp instruction is executed. b0 b0 b0 b0 b1 b1 b1 b1 ... ... ... ... b0 b0 b0 b0 b1 b1 b1 b1 ... ... ... ... ... WSr MMAk=16 (32/2) WSk WSc 0 1 2 3 4 5 6 7 8 a0 a1 a0 a1 a2 a3 a2 a3 a4 a5 ... a4 a5 a6 a7 a6 a7 T0T4 0 1 ... 6 7 T24 T28 0 1 .. 6 7 Tx = Thread x 0 1 ... 6 7 T0T3 a0 a1 a0 a1 a2 a3 a2 a3 a4 a5 ... a4 a5 a6 a7 a6 a7 0 1 ... 6 7 T16 T19 0 1 2 3 4 5 6 7 8 ... 15 T0{a2,a3} T3{a2,a3} ... ... ... T28{a2,a3} T31{a2,a3} 8 ... 31 T0{c0,c1} T3{c0,c1} T4{c0,c1} T7{c0,c1} T8{c0,c1} ... T11{c0,c1} T12{c0,c1} T15{c0,c1} T16{c0,c1} T19{c0,c1} T20{c0,c1} ... T23{c0,c1} T24{c0,c1} T27{c0,c1} T28{c0,c1} T31{c0,c1} 0 1 2 3 4 5 6 7 0-1 ... 6-7 8 ... 15 T0{c2,c3} T3{c2,c3} ... ... ... T28{c2,c3} T31{c2,c3} WSr=32 16 a8 a9 a8 a9 a10 a11 a10 a11 a12 a13 ... a12 a13 a14 a15 a14 a15 8 9 ... 15 16 T0T3 a8 a9 a8 a9 a10 a11 a10 a11 a12 a13 ... a12 a13 a14 a15 a14 a15 8 9 ... 15 16 T16 T19 T0{a8,a9} T3{a8,a9} T4{a8,a9} T7{a8,a9} T8{a8,a9} T11{a8,a9} T12{a8,a9} ... T15{a8,a9} T16{a8,a9} T19{a8,a9} T20{a8,a9} T23{a8,a9} T24{a8,a9} ... T27{a8,a9} T28{a8,a9} T31{a8,a9} 8-9 ... 14-15 T0{a8,a9} T3{a8,a9} ... ... ... T28{a8,a9} T31{a8,a9 x4 times MMAr=16 MMAc=8 WSkMMAk=32 T1,T2 T1,T2 T17,T18 T17,T18 T0{a0,a1} T3{a0,a1} T4{a0,a1} T7{a0,a1} T8{a0,a1} T11{a0,a1} T12{a0,a1} ... T15{a0,a1} T16{a0,a1} T19{a0,a1} T20{a0,a1} T23{a0,a1} T24{a0,a1} ... T27{a0,a1} T28{a0,a1} T31{a0,a1} 0-1 ... 6-7 21 22 23 A B C T3T7T27 T31 Figure 6: SPTCs view Storage order. Related to stage 1 and 2 , we propose a specific order to store the non-zero values and the m-indices structure of the V:N:M format, which merges, once again, the block-wise and the N:M principles. This order is represented in Figure 7, and it seeks to optimize the data traversal during the data loading and computation. In this representation, half of the non-zero structure shows the access pattern followed to store the data, while the other half shows how the second half-warp is mapped into this structure. This storage order enables 128-bit memory transactions, ensures memory coalescence, and can dispense with the ldmatrix instruction, which is known to cause bank conflicts and can require more SMEM transactions to sequentially serve the memory access [33]. 4.1.3 Stage 3-Result storage. Once the product is calculated, we have to write the output tiles back to GMEM (stage 3 ). This requires storing the intermediate partial results in SMEM. On NVIDIA GPUs, shared memory is partitioned into banks, each one of 32 bits. Each bank can only address one position at a time, so if a quarter-warp (128-bit instructions) tries to access the same bank, the instruction will be serialized. This effect is known as bank conflict. An example of thread mapping to SMEM with 𝐵𝑆𝑐= 64 is shown in Figure 8. T16 T17 T18 T19 T16 T17 T18 T19 T16 T17 T18 T19 T16 T17 T18 T19 T20 T21 T22 T23 T20 T21 T22 T23 T20 T21 T22 T23 T20 T21 T22 T23 T24 T25 T26 T27 T24 T25 T26 T27 T24 T25 T26 T27 T24 T25 T26 T27 T28 T29 T30 T31 T28 T29 T30 T31 T28 T29 T30 T31 T28 T29 T30 T31 MMAm=16 MMAk=32 Non-zero values m-indices < WSm=32 < WSk 32 bits 128 bits 16 bits Figure 7: Storage order The left side of the figure shows how the threads in a warp are mapped to SMEM banks during the storage of their partial results (step 31 ). These stores are performed with 128-bit instructions. Padding elements have been added to avoid bank conflicts. In this specific example, each thread has accumulated 8partial results ( 𝐵𝑆𝑐/𝑀𝑀𝐴𝑐= 64 / 8 ) , so the thread mapping is repeated 8times, meaning that each thread needs 8iterations to store its partial results. Each color represents a quarter-warp, so we can see that each group of 8consecutive threads accesses a different memory bank in the same iteration. 0 1 4 5 8 9 12 13 16 17 20 21 24 25 28 29 PAD PAD PAD PAD 2 3 6 7 10 11 14 15 18 19 22 23 26 27 30 31 PAD 0 1 4 5 8 9 12 13 16 17 20 21 24 25 28 29 PAD PAD PAD PAD 2 3 6 7 10 11 14 15 18 19 22 23 26 27 30 31 PAD ... ... ... ... ... ... ... ... 0 1 4 5 8 9 12 13 16 17 20 21 24 25 28 29 PAD PAD PAD PAD 2 3 6 7 10 11 14 15 18 19 22 23 26 27 30 31 PAD x8 times 0 0 8 8 16 16 24 24 0 0 8 8 16 16 24 24 PAD PAD PAD PAD 0 0 8 8 16 16 24 24 0 0 8 8 16 16 24 24 PAD 1 1 9 9 17 17 25 25 1 1 9 9 17 17 25 25 PAD PAD PAD PAD 1 1 9 9 17 17 25 25 1 1 9 9 17 17 25 25 PAD ... ... ... ... ... ... ... ... 7 7 15 15 23 23 31 31 7 7 15 15 23 23 31 31 PAD PAD PAD PAD 7 7 15 15 23 23 31 31 7 7 15 15 23 23 31 31 PAD 128 bytes 128 bytes Bank 0~3 3132 Figure 8: Conflict-free accesses for output tiles on SMEM The right side of Figure 8 shows step 32 , that is, the SMEM thread mapping designed to read the previously stored intermediate results, and finally, write them back to GMEM. The loads from SMEM and the stores to GMEM are performed with 128-bit instructions. Once again, each thread will need to access SMEM 8times to read all the data. We have colored the accesses related to the first quarter-warp, what depicts a conflict-free layout. Ablation study - Spatha performance and column-loc overhead.In Figure 9, we present the results of a microbenchmark study on matrices of fixed outer dimensions (corresponding to the size of one BERT 𝑙𝑎𝑟𝑔𝑒 weight linear layer), but varying the inner (sparsified) one, 𝐾 (1024 ×𝐾× 4096). The study was conducted using different sparsity levels, specified by different N:M combinations (from 2:10 to 2:100), while the vector size 𝑉 was kept constant at 128. Furthermore, to measure the effect of using the columnloc mechanism, we tested the performance with and without this structure. In the latter we used fixed indexes to simulate an ideal situation with no memory accesses. These experiments are performed on an NVIDIA RTX 3090 GPU, equipped with SPTCs. The SC ’23, November 12–17, 2023, Denver, CO, USA Castro R.L., et al. 768 1536 2304 3072 3840 4608 5376 6144 6912 7680 8448 9216 9984 10752 11520 12288 K 80% [128:2:10] 0 1 2 3 4 768 1536 2304 3072 3840 4608 5376 6144 6912 7680 8448 9216 9984 10752 11520 12288 K 90% [128:2:20] 0 2 4 6 8 768 1536 2304 3072 3840 4608 5376 6144 6912 7680 8448 9216 9984 10752 11520 12288 K 95% [128:2:40] 0 5 10 15 w/ column-loc w/o column-loc 768 1536 2304 3072 3840 4608 5376 6144 6912 7680 8448 9216 9984 10752 11520 12288 K 98% [128:2:100] 0 10 20 30 40 SpeedUp w.r.t. cuBLAS Sparsity [%] (V:N:M) Figure 9: Ablation study of column-loc with different sizes of the inner 𝐾dimension and different V:N:M formats (𝐵𝐸𝑅𝑇𝑙𝑎𝑟𝑔𝑒 ) 32 64 128 V 71% [V:2:7] 0.0 2.5 5.0 7.5 10.0 12.5 15.0 17.5 20.0 32 64 128 V 75% [V:2:8] 32 64 128 V 80% [V:2:10] 32 64 128 V 90% [V:2:20] 32 64 128 V 95% [V:2:40] 32 64 128 V 98% [V:2:100] 32-bit 128-bit Sparsity [%] (V:N:M) Speedup w.r.t. cuBLAS Figure 10: Scaling study of wide shared memory stores for different V:N:M configurations results show that Spatha achieves speedups for sparse computation, approaching theoretical peak performance for a given sparsity level considering the operation count reduction w.r.t. the dense counterpart version. This effect becomes more pronounced as the GEMM problem size increases, as it tend to have higher arithmetic intensity. For instance, at a sparsity level of 80% (2:10 format), the speedup is approximately 4 . 5 × , where 5 × is the ideal scenario. Then, the speedups reported are 8 . 5 × ,17 . 5 × , and 37 × for sparsity levels of 90% (2:20), 95% (2:40) and 98% (2:100), whose theoretical caps are 10 × ,20 × and 50 × , respectively. It can be observed that, for every sparsity ratio, the column-loc structure’s overhead has a negligible effect on the overall time, despite being a software approach to support arbitrary N:M ratios. However, the impact of column-loc becomes slightly more noticeable when dealing with 2:100 sparsity, which is not practical for DL applications in real-world scenarios. Scaling study - Impact of V and output layout format.The 𝑉 variable in our V:N:M format can be used to define trade-offs between performance and accuracy in the same way that the blocksize in block-wise pruning, for example. To study this, we performed a second ablation study on one matrix from 𝐵𝐸𝑅𝑇𝑙𝑎𝑟𝑔𝑒 (size 1024 × 4096 × 4096). Figure 10 shows the performance results of Spatha on this matrix using three different vector lengths: 32,64 and 128. This test is conducted for different sparsity levels, in practice, the test explores different configurations of the V:N:M values. Furthermore, in order to study the impact of the previously proposed layout for writing back results (Figure 8), it is compared the effect of using 50% (2:4) 60% (2:5) 75% (2:8) 80% (2:10) 90% (2:20) 95% (2:40) Sparsity [%] (N:M) 0.0 0.2 0.4 0.6 0.8 Energy BERT-base encoder.layer.8.attention.self.query.weight: 768x768 ideal 1:N:M 16:N:M 32:N:M 64:N:M 128:N:M vw_4 vw_8 vw_16 vw_32 Figure 11: Energy evaluation study on the V:N:M format such layout, enabling 128-bit SMEM stores instead of 32-bit ones. As we can see in Figure 10, the difference in terms of speedups between the three selected vector lengths is noticeable, the value of 𝑉 being conditioned by the accuracy loss. The effect of using 128-bit stores instead of 32-bit ones is noticeable in this problem size, bringing up to a 2 × difference in the final speedup. We performed a similar ablation test for a matrix of a GPT-3 model (size 36864 × 12288 × 4096) and the effect of using 128-bit stores was attenuated, as the weight of the output phase in the total execution time is smaller. 5 ENERGY EVALUATION OF V:N:M DL pruning techniques aim to achieve the highest possible sparsity levels in the pruned models while ensuring little to no loss in accuracy. This becomes especially challenging when the target sparse format requires a specific pruning scheme, and when high sparsity levels are targeted. In these scenarios, the percentage of non-zero values is low, and their location is heavily influenced by the format. Therefore, it is crucial to demonstrate the effectiveness of new sparse formats, to ensure its applicability with minimal or no impact on accuracy. The energy evaluation metric (magnitude preservation) [ 22 , 29 ] is employed to measure the flexibility of a format by comparing the total magnitude of the model (sum of the individual weights) before and after pruning to a specific format. Let us assume a welloptimized dense model 𝑤∗∈R𝑑 , where 𝑑 is the total number of weights. We wish to prune 𝑤∗ to a target sparsity 𝑠∈ ( 0 , 1 ] by VENOM: A Vectorized N:M Format for Unleashing the Power of Sparse Tensor Cores SC ’23, November 12–17, 2023, Denver, CO, USA zeroing out 𝑠×𝑑 weights. The result is a sparse model 𝑤∈R𝑠×𝑑 . The energy metric is defined as follows: 𝑒𝑛𝑒𝑟𝑔𝑦 = Í𝑠×𝑑 𝑖=0|𝑤𝑖| Í𝑑 𝑖=0|𝑤∗ 𝑖| This metric yields a normalized score between 0 ∼ 1, the higher the better. Figure 11 presents the energy evaluation study for a weight tensor extracted from an encoder layer of BERT 𝑏𝑎𝑠𝑒 . This figure compares three weight selection policies: unstructured (ideal), V:N:M with different 𝑉 values, and vector-wise pruning with several vector lengths 𝑙 ( 𝑣𝑤_𝑙 ). The evaluation is done for different sparsity levels, whose value in the V:N:M format is controlled by the N:M ratio. Unstructured magnitude-pruning represents the ideal non-zero selection policy, as it does not impose any restrictions on the location of non-zeros. Vector-wise pruning can accelerate sparse routines on GPUs. However, if the vector length is > 8, it can significantly reduce the accuracy [ 2 , 3 , 23 ]. The results demonstrate that the V:N:M format occupies an intermediate position between unstructured and vector-wise pruning. Moreover, it is highly robust to changes in the vector length, allowing the usage of 𝑉= 128 while consistently preserving more energy than 𝑣𝑤_8and 𝑣𝑤_4. Additionally to the previous conclusions, independently of the selected pruning method, we can also see the tremendous impact on the energy of magnitude-based weight selection policies. At 50% of sparsity, unstructured pruning already lost 20% of the original dense matrix energy. At the other side, at 95% only 20% of the original energy remain in the pruned dense matrix. Thus, we can conclude that, in order to achieve moderate to high sparsity ratios in models with the dimensionality of BERT, more sophisticated pruning methods must be used. Second-order pruning offers an alternative to these problems. 6 SECOND-ORDER PRUNING Magnitude-based pruning techniques provide a straightforward approach to reducing the size of our models without requiring model evaluation for weight selection. However, while magnitude pruning can be effective at moderate sparsity levels, it becomes more challenging to select the "least significant" weights to remove when aiming for high sparsity ratios, and this can significantly impact network accuracy. In contrast, second-order pruning methods offer a more sophisticated approach to select weight candidates for removal, by considering the difference in loss relative to the current model. Hence, they target to find the set of weights whose removal will generate a minimum loss increase. In this context, the Hessian matrix is a key component of second-order pruning methods which represents the matrix of second-order derivatives of the loss function w.r.t. the weights, mathematically expressed as 𝐻=∇2 𝑤𝐿 , for a twice-differentiable loss 𝐿 . The Fisher matrix is very similar to the Hessian matrix but in the probabilistic setting, used to estimate the curvature of the loss function around the current value. As a result, this approximation allows to identify the weight parameters that have less impact in the loss function, and therefore are candidates to be pruned [13]. 6.1 The V:N:M format in 2nd order methods This section introduces a new second-order pruning method based on [ 19 ] and tailored for the V:N:M format. This type of approach yields state-of-the-art results in LLMs for unstructured and semistructured (block) compression. Let us assume we have a well-optimized dense model 𝑤∗∈R𝑑 , where 𝑑 is the total number of weights. Our target is to identify a set of weights 𝑄that we can prune with a minimum loss increase. The following saliency score function is defined to rank groups of weights [19]: 𝜌𝑄= 1 2(𝐸𝑄𝑤∗)𝑇(𝐸𝑄b 𝐹−1(𝑤∗)𝐸𝑇 𝑄)−1𝐸𝑄𝑤∗ where, •b 𝐹−1(𝑤) ∈ R𝑑×𝑑is the Fisher matrix. •𝐸𝑄∈R|𝑄|×𝑑 is a matrix composed of the corresponding canonical basis vectors for a set of 𝑄weights. Thus, the set of canonical basis vectors 𝐸𝑄 depends on the specific sparse format we are using. For instance, in 2:4 sparsity, the canonical vectors are: 𝐸𝑄=[[1,1,0,0],[1,0,1,0],[1,0,0,1],[0,1,1,0],[0,1,0,1],[0,0,1,1]] As observed, 𝐸𝑄 encompasses all possible correlations between 2weights, in a set of 4elements. In general, for an N:M format, this approach requires evaluating 𝑀 𝑁 combinations to determine the best one, which can turn into an intractable combinatorial problem. Furthermore, in the V:N:M format, the addition of a new dimension 𝑉 amplifies the complexity as it requires finding the optimal set of 𝑉×𝑁weights, leading to a combinatorial explosion. To address these challenges, we adopt a similar approach as [ 19 ] between sets of 𝑄 elements, which involves disregarding correlations between rows within 𝑉×𝑀 blocks. This simplification drops the number of combinations to evaluate. Additionally, to mitigate combinatorial issues that may still arise within 1 ×𝑀 groups, we propose a pair-wise approach where correlations are calculated between pairs of elements, that is: 𝐸𝑄=[[1,0],[0,1],[1,1]] The overhead of this pair-wise approach represents < 1% of the sparsification process. Then, depending on the 𝑁 and 𝑀 values, we can modulate the complexity of the problem to be solved by dynamically selecting the m-combinatorial or the pair-wise approach. 6.1.1 Gradual pruning definition. The N:M format prunes a model to a target sparsity 𝑠∈ ( 0 , 1 ] . Typically, the 𝑠×𝑑 weights are removed in one step (one-shot pruning). For 50% (2:4) sparsity, this approach can be applied in most cases and the models still recover the original accuracy. However, for higher sparsity ratios, one-shot pruning reduces severely the model performance and makes hard to recover the original accuracy using additional finetuning steps. This negative effect on accuracy also happens in second-order methods, where one-shot pruning can result in worse Taylor approximations of the function. We propose a structure decay scheduler for the V:N:M format, which performs N:M pruning across different 𝛽 steps, for increasing sparsity levels. This scheduler starts with a high initial value of 𝑁0>> 𝑁𝛽 (lower sparsity), where 𝑁𝛽 is our target 𝑁 value, and gradually decreases 𝑁 (conversely SC ’23, November 12–17, 2023, Denver, CO, USA Castro R.L., et al. increasing sparsity) until it reaches the 𝑁 target value. This gradual pruning approach mitigates the adverse effects on network accuracy and improves accuracy recovery in subsequent fine-tuning stages. 7 EVALUATION We evaluate the performance on an NVIDIA RTX 3090 GPU of the Ampere architecture equipped with SPTCs. We compare the performance of Spatha with different sparse libraries (cuSparseLt, cuSparse, CLASP, and Sputnik) and also with a dense counterpart version (cuBLAS). We build our benchmarks on matrices from the DLMC dataset [ 10 ], and real-world LLMs. Additionally to these micro benchmarks, we also conduct a case study on real-world applications. At this point, we demonstrate the proposed second-order pruning technique, and we benchmark the end-to-end performance of Spatha on different LLM models (BERT, GPT-2, and GPT-3). 7.1 Comparison with existing dense and sparse libraries Firstly, we evaluate our baseline implementation for 1:2:4 sparsity (50%). Since higher N:M ratios will depend on this baseline’s performance, it is crucial to have good speedup results in this configuration. We selected cuBLAS GEMM as our dense counterpart, and for exploiting the 2:4 format on SPTCs, we used the cuSparseLt SpMM implementation, which represents the reference library on this format. Our experiments involve varying sizes of a 𝑅×𝐾×𝐶 GEMM problem, where 𝑅 and 𝐶 are predetermined values from two BERT’s weight linear layers (768 and 4096 for BERT 𝑏𝑎𝑠𝑒 ,1024 and 4096 for BERT 𝑙𝑎𝑟𝑔𝑒 ). The inner dimension 𝐾 of the product, which is the sparsified one, is variable in these experiments. Note that the inner dimension is usually scaled up to enhance the network accuracy. For instance, GPT-3 uses a hidden size of 12288 [ 1 ]. Figure 12 reports the performance of the three contending implementations (cuBLAS, cuSparseLt and Spatha) and the speedups of the selected sparse libraries w.r.t. cuBLAS. The results show that the performance of the sparse implementation improves with the GEMM size, as larger GEMMs tend to have larger arithmetic intensity. In these microbenchmarks, BERT 𝑙𝑎𝑟𝑔𝑒 matrices (right side) increase the computation intensity w.r.t. BERT 𝑏𝑎𝑠𝑒 (left). Notably, for larger GEMM sizes, the performance of cuSparseLt and Spatha is similar and approaches the peak for 2:4 sparsity, 2 × . However, our implementation shows better performance on smaller sizes, which constitutes an interesting feature, since Spatha can probably cover a more variety of network architectures. Overall, Spatha achieves up to 1.38×speedup over the vendor library, cuSparseLt. 768 1536 2304 3072 3840 4608 5376 6144 6912 7680 8448 9216 9984 10752 11520 12288 GEMM-K 0 20 40 60 80 TFLOPS/s GEMM-R=768, GEMM-C=4096 (BERT-base) 768 1536 2304 3072 3840 4608 5376 6144 6912 7680 8448 9216 9984 10752 11520 12288 GEMM-K 0 20 40 60 80 TFLOPS/s GEMM-R=1024, GEMM-C=4096 (BERT-large) Spatha cuBLAS cuSparseLt 1.0 1.2 1.4 1.6 1.8 2.0 SpeedUp w.r.t. cuBLAS 1.0 1.2 1.4 1.6 1.8 2.0 SpeedUp w.r.t. cuBLAS Spatha cuSparseLt Figure 12: Baseline performance at 50% sparsity (2:4 format) Figures 13 and 14 compare the performance of Spatha to other dense and sparse libraries for higher sparsity levels ranging from 50 ∼ 98%. In this context, cuSparseLt SpMM is the reference library to exploit the 2:4 format on SPTCs. However, cuSparseLt’s results are strictly limited to 50% sparsity and cannot be executed for the entire sparsity range because it only supports 2:4 sparse matrices. Since there are no SpMM GPU implementations for arbitrary N:M sparsity levels, we have considered the following third-party libraries for half-precision: Sputnik [ 9 ], cuSparse [ 27 ], and CLASP [ 2 ] which extends vectorSparse [ 3 ] to the latest generations of NVIDIA GPU architectures. Concerning semi-structured sparse matrices, [ 2 ] focuses on the column-vector format, which supports vector lengths 𝑙= 2 , 4and 8. Additionally, cuSparse introduces the Blocked-ELL format to improve the performance of block-wise sparse matrices, for which we considered the same block lengths (2 , 4 , 8) as [ 2 , 3 ]. For non-structured sparsity, [ 9 ] represents the SOTA implementation over cuSparse+CSR, excluded here for simplicity. Figure 13 benchmarks are built using the DLMC dataset [ 10 ]. Spatha, CLASP, and cuSparse configurations are referenced in the rows with the notation V:N:M, 𝑣𝑤_𝑙 , and 𝑏𝑤_𝑙 , respectively. In the absence of sparse matrices compatible with the N:M format, we have generated the requisite matrices by pruning the original dense models [ 8 ]. Most of the dense matrices in this dataset are too small (e.g., 64 × 64, from ResNet-50) to fill the GPU resources. This translates into low arithmetic intensities, hindering peak performance. As Figure 13 shows, when we increase the arithmetic intensity by means of 𝐶 (batch size), cuSparseLt approaches the 2 × peak [ 25 ] (still 1 . 05 × on average for 𝐶= 2048). However, Sputnik, cuSparse and CLASP performance shows important limitations. Existing SpMM kernels are mostly designed considering small models, where the LHS operand can be a tiny matrix. That influences the SpMM design, since data can be loaded directly into registers, but these design decisions represent a serious scalability problem. However, Spatha is able to scale with the problem size, reaching speedups of up to ∼ 25 × (98% sparsity). The scalability property is critical with the increasing models size, especially on LLMs. Figure 14 benchmarks are built using sparse matrices from weightpruned linear layers extracted from BERT. Spatha, cuSparse, and CLASP configurations have been referenced in the columns of Figure 14 with the same notation as Figure 13. The first row of Figure 14 shows the speedup results on sparse matrices extracted from BERT 𝑏𝑎𝑠𝑒 while the second one reports that performance on BERT 𝑙𝑎𝑟𝑔𝑒 . The y-axis is represented in a logarithmic scale to make the results more readable. As we saw in Figure 13, existing implementations outperform the dense counterpart version at sparsity levels above 80% when the matrices are small. However, when we evaluate these implementations on medium or big matrices extracted from larger models (e.g., LLMs), the performance is even worse, and they only outperform cuBLAS at sparsity levels above 90%. Thus, their kernel design presents scalability issues when the problem size grows. The fact that Spatha reaches 2 × speedup at 50% sparsity enables the achievement of high speedups as the sparsity increases, yielding up to 27 × in BERT-like matrices. We can also appreciate that the best performance in our implementation is reached with the augmentation of arithmetic intensity, peaking for BERT _𝑙𝑎𝑟𝑔𝑒 with batch size 16. A discussion of how these speedups affect the global LLMs latency will be presented next. VENOM: A Vectorized N:M Format for Unleashing the Power of Sparse Tensor Cores SC ’23, November 12–17, 2023, Denver, CO, USA 0.1 0.2 0.3 0.5 1.0 2.0 4.0 6.0 C=256 0.1 0.2 0.3 0.5 1.0 2.0 4.0 8.0 C=512 0.1 0.3 0.5 1.0 2.0 3.0 6.0 10.0 15.0 C=1024 0.1 0.3 0.5 1.0 2.0 3.0 6.0 10.0 15.0 27.0 C=2048 50 70 80 90 95 98 0.1 0.3 0.5 1.0 2.0 4.0 7.0 50 70 80 90 95 98 0.1 0.3 0.5 1.0 2.0 4.0 8.0 50 70 80 90 95 98 0.1 0.3 0.5 1.0 2.0 3.0 6.0 10.0 15.0 50 70 80 90 95 98 0.1 0.3 0.5 1.0 2.0 4.0 8.0 15.0 27.0 Sparsity [%] SpeedUp (log-scale) 64:N:M vw_4,bw_4 SpeedUp (log-scale) 128:N:M vw_4,bw_8 cuBLAS CLASP cuSparse cuSparseLt Spatha Sputnik Figure 13: Speedup results on sparse matrices from the neural networks on the DLMC dataset. The problem size is 𝐴𝑅×𝐾×𝐵𝐾×𝐶 where 𝐴is the sparse matrix. The 𝑅and 𝐾sizes are given in the DLMC dataset, while 𝐶is selected from {256,512,1024,2048} 0.1 0.3 0.5 1.0 2.0 4.0 6.0 10.0 15.0 25.0 batch size=8 64:N:M , vw_4 , bw_4 batch size=16 64:N:M , vw_4 , bw_4 batch size=8 128:N:M , vw_8 , bw_8 batch size=16 128:N:M , vw_8 , bw_8 50 70 75 80 90 95 98 0.1 0.3 0.5 1.0 2.0 4.0 6.0 10.0 15.0 28.0 50 70 75 80 90 95 98 50 70 75 80 90 95 98 50 70 75 80 90 95 98 Sparsity [%] SpeedUp (log scale) BERT-base SpeedUp (log scale) BERT-large cuBLAS Spatha cuSparseLt Sputnik CLASP cuSparse Figure 14: Speedup results on BERT 𝑏𝑎𝑠𝑒 and BERT 𝑙𝑎𝑟𝑔𝑒 with sequence length=512. The notation 𝑉 :N:M represents the vector length 𝑉 used on Spatha, while vw _𝑙 and bw _𝑙 represents the vector length 𝑙 used on CLASP and cuSparse, respectively. The N:M pattern related to each of the considered sparsity levels are in ascending order of sparsity: 2:4, 2:7, 2:8, 2:10, 2:20, 2:40 and 2:100 7.2 Case study: sparse LLMs LLMs have revolutioned the NLP field with their unrivaled performance in various domains. Nowadays, these models are widely used in everyday technologies, such as ChatGPT. Transformer LLMs typically consist of multiple transformer layers with self-attention. There are two major sub-components inside a transformer architecture: the multi head attention (MHA), and the fully connected feed forward network (FFN). At a higher level, the model size is determined by different configurable components, such as the head dimension, the number of heads and the number of layers, depending on the specific architecture used. This case study focuses on weight pruning, and explores the on computational speedups achievable with Spatha. In LLMs weight tensors are present in Linear Layers, which can be found in both the MHA and the MLP sub-components. Figure 15 illustrates a pruned MHA where four GEMM instructions are converted to SpMMs by sparsifying the corresponding weight tensors. In this study we demonstrate the efficiency of Spatha on different LLMs. However, it Castro, et al. python3 benchmark/energy.py (7) Since reproducing results on Table 2 can take a significant amount of time, we provide three different scripts to alleviate this process. conda activate sparseml_artf cd sparseml # Script that contains a subset of the experiments with the most aggressive configurations using the pair-wise version of the sparsifier # about 4 days ./sparseml_SS1.sh # Script that contains all the sparsity-format configurations but relaxed with pair-wise version of the sparsifier # about 10 days ./sparseml_SS2.sh # Script that contains all the sparsity-format configurations and performs the exhaustive search process # about 25 days ./sparseml_SS3.sh Note: each script in integrations/huggingface-transformers/scripts has two execution possibilities. Please, uncomment the first line if you want to use a single-GPU, or the second one with the total number of GPUs available for multiple-GPU execution. # single-GPU CUDA_VISIBLE_DEVICES=0 python3.10 src/sparseml/transformers/question_answering.py # multi-GPU (3 in this example) python3.10 -m torch.distributed.launch –nproc_per_node=3 src/sparseml/transformers/question_answering.py Step 3: check plots cd /projects/venom/result scp *.pdf username@hostmachine:/host/path/target