cuSPARSELt开发NVIDIA Ampere结构化稀疏性

cuSPARSELt开发NVIDIA Ampere结构化稀疏性

深度神经网络在各种领域（例如计算机视觉，语音识别和自然语言处理）中均具有出色的性能。处理这些神经网络所需的计算能力正在迅速提高，因此有效的模型和计算至关重要。神经网络剪枝（删除不必要的模型参数以生成稀疏网络）是一种在保持准确性的同时降低模型复杂性的有用方法。

为了利用细粒度的网络剪枝，NVIDIA Ampere GPU架构引入了细粒度的结构稀疏性的概念。在NVIDIA A100 GPU上，结构显示为2：4模式：每四个元素中至少有两个必须为零。通过使用新的NVIDIA Sparse Tensor Core跳过零值的计算，这可以将一个矩阵乘法（也称为GEMM）操作数的数据占用空间和带宽减少2倍，并使吞吐量翻倍。

cuSPARSELt：用于稀疏矩阵-密集矩阵乘法的高性能CUDA库

为了简化NVIDIA Ampere架构稀疏功能的使用，NVIDIA引入了cuSPARSELt ，这是一种高性能CUDA库，专用于常规矩阵操作，其中至少一个操作数是稀疏矩阵。cuSPARSELt库可以使用NVIDIA第三代Tensor Core稀疏矩阵乘累加（SpMMA）操作，而无需进行底层编程。该库还提供用于剪枝和压缩矩阵的辅助函数。

cuSPARSELt的主要功能包括：

· NVIDIA Sparse Tensor Core支持

· 混合精度支持：

o FP16输入/输出，FP32张量核心累积

o BFLOAT16输入/输出，FP32张量核心累积

o INT8输入/输出，INT32张量核心累积

· Row-major and column-major memory layouts的内存布局

· 矩阵剪枝和压缩实用程序

· 自动调整功能

• NVIDIA Sparse Tensor Core support
• Mixed-precision support:
• FP16 inputs/output, FP32 Tensor Core accumulation
• BFLOAT16 inputs/output, FP32 Tensor Core accumulation
• INT8 inputs/output, INT32 Tensor Core accumulation
• Row-major and column-major memory layouts
• Matrix pruning and compression utilities
• Auto-tuning functionality

定制工作流程

cuSPARSELt库遵循等效方法，并采用与cuBLASLt和cuTENSOR类似的概念。库编程模型要求以某种方式组织计算，以使相同的设置可以重复用于不同的输入。

该模型尤其依赖于以下高层阶段：

· 问题定义：指定矩阵形状，数据类型，操作等。

· 用户偏好和约束：提供算法选择或限制可行实现（候选）的搜索空间。

· 计划：收集执行的描述符，并在需要时“找到”最佳实施。

· 执行：执行实际计算。

通用工作流程包括以下步骤：

1. 初始化库句柄：cusparseLtInit。

2. 指定输入/输出矩阵特征：cusparseLtDenseDescriptorInit， cusparseLtStructuredDescriptorInit。

3. 初始化矩阵乘法描述符和它的属性（例如操作，计算类型等）： cusparseLtMatmulDescriptorInit。

4. 初始化算法选择描述符：cusparseLtMatmulAlgSelectionInit。

5. 初始化矩阵乘法计划：cusparseLtMatmulPlanInit。

6. 剪枝A矩阵：cusparseLtSpMMAPrune。如果用户提供已经满足2：4结构化稀疏性约束的矩阵，例如由ASP库生成的权重矩阵，则不需要此步骤。

7. 压缩剪枝后的矩阵：cusparseLtSpMMACompress。

8. 执行矩阵乘法：cusparseLtMatmul。可以使用不同的输入多次重复此步骤。

9. 销毁矩阵乘法计划和库句柄：cusparseLtMatmulPlanDestroy，cusparseLtDestroy。

稀疏的GEMM性能

与密集矩阵乘法一样，稀疏矩阵乘法的性能随GEMM尺寸，布局和数据类型而变化。这是当前软件与稀疏GEMM相对性能的快照。

下表显示了cuSPARSELt和cuBLAS在以下操作中的性能：

D = alpha * op（A）* op（B）+ beta * C

在该操作中，A，B，和 D=C分别是尺寸的密集矩阵MXK，KXN，和M×N个。矩阵的布局A和B与Ñ为列主顺序（OP是非转置）和Ť为行优先顺序（OP调换）。

为了展示使用cuSPARSELt可以针对实际工作负载实现的性能，下表显示了带有主要列TN FP16内核的剪枝后的BERT-Large模型（seqlen = 128，BS = 128）使用的一些常见GEMM大小。通常，工作量越大，稀疏性可以提供的帮助越多。

表1. BERT-Large模型和不同层的cuSPARSELt性能。

结构化稀疏矩阵-矩阵乘法代码示例

已经看到了可用的性能，下面是一个示例，该示例使用NVIDIA A100或GA100 GPU中的稀疏Tensor内核在cuSPARSELt库中执行具有结构稀疏性的矩阵乘法。有关更多信息，请参见NVIDIA / CUDALibrarySamples / tree / master / cuSPARSELt / spmma GitHub存储库。

首先，包括cuSPARSELt标头，设置一些设备指针和数据结构，并初始化cuSPARSELt句柄。

#include <cusparseLt.h> // cusparseLt header // Device pointers and coefficient definitions float alpha = 1.0f; float beta = 0.0f; __half* dA = ... __half* dB = ... __half* dC = ... // cusparseLt data structures and handle initialization cusparseLtHandle_t handle; cusparseLtMatDescriptor_t matA, matB, matC; cusparseLtMatmulDescriptor_t matmul; cusparseLtMatmulAlgSelection_t alg_sel; cusparseLtMatmulPlan_t plan; cudaStream_t stream = nullptr; cusparseLtInit(&handle);

接下来，初始化结构化的稀疏输入矩阵（matrix A），密集输入矩阵（matrix B）和密集输出矩阵（matrix C）描述符。

cusparseLtStructuredDescriptorInit(&handle, &matA, num_A_rows, num_A_cols, lda, alignment, type, order, CUSPARSELT_SPARSITY_50_PERCENT); cusparseLtDenseDescriptorInit(&handle, &matB, num_B_rows, num_B_cols, ldb, alignment, type, order); cusparseLtDenseDescriptorInit(&handle, &matC, num_C_rows, num_C_cols, ldc, alignment, type, order);

准备好描述符后，可以准备矩阵乘法运算的描述符，选择用于执行matmul运算的算法，并初始化matmul计划。

cusparseLtMatmulDescriptorInit(&handle, &matmul, opA, opB, &matA, &matB, &matC, &matC, compute_type); cusparseLtMatmulAlgSelectionInit(&handle, &alg_sel, &matmul, CUSPARSELT_MATMUL_ALG_DEFAULT); int alg = 0; // set algorithm ID cusparseLtMatmulAlgSetAttribute(&handle, &alg_sel, CUSPARSELT_MATMUL_ALG_CONFIG_ID, &alg, sizeof(alg)); size_t workspace_size, compressed_size; cusparseLtMatmulGetWorkspace(&handle, &alg_sel, &workspace_size); cusparseLtMatmulPlanInit(&handle, &plan, &matmul, &alg_sel, workspace_size);

如果稀疏矩阵尚未被其他进程剪枝，则可以在此时进行。不要忘记检查稀疏模式的有效性，以确保可以使用稀疏张量核心来加速它。

cusparseLtSpMMAPrune(&handle, &matmul, dA, dA, CUSPARSELT_PRUNE_SPMMA_TILE, stream); // checking the correctness int is_valid = 0; cusparseLtSpMMAPruneCheck(&handle, &matmul, dA, &is_valid, stream); if (is_valid != 0) { std::printf("!!!! The matrix does not conform to the SpMMA sparsity pattern. " "cusparseLtMatmul does not provide correct results\n"); return EXIT_FAILURE; }

现在已将矩阵A剪枝为2：4稀疏度，可以将其压缩到大约原始大小的一半。与实际的矩阵乘法（小于5％）相比，该步骤的执行时间可以忽略不计。

cusparseLtSpMMACompressedSize(&handle, &plan, &compressed_size); cudaMalloc((void**) &dA_compressed, compressed_size); cusparseLtSpMMACompress(&handle, &plan, dA, dA_compressed, stream);

设置完成后，执行matmul操作。cusparseLtMatmul使用不同的B矩阵可以多次重复调用 。只需设置一次稀疏矩阵。对于A矩阵值更改的用例，cusparseLtSpMMACompress必须再次调用该例程以设置稀疏矩阵的数据结构。

void* d_workspace = nullptr; int num_streams = 0; cudaStream_t* streams = nullptr; cusparseLtMatmul(&handle, &plan, &alpha, dA_compressed, dB, &beta, dC, dD, d_workspace, streams, num_streams) )

最后，通过破坏matmul计划和cuSPARSELt句柄来清理已使用的内存。

cusparseLtMatmulPlanDestroy（＆plan）; cusparseLtDestroy（＆handle）;

cuSPARSELt

通过cuSPARSELt库，可以轻松利用NVIDIA Sparse Tensor Core运算，从而在不降低网络准确性的情况下，显着提高了用于深度学习应用程序的矩阵矩阵乘法的性能。该库还提供了用于矩阵压缩，剪枝和性能自动调整的实用程序。简而言之，与普通的密集数学方法相比，cuSPARSELt减少了计算，功耗，执行时间和内存存储。