CANN/catccos AllGather反量化算子

AllGather矩阵乘法反量化算子设计文档【免费下载链接】catccosCATCCOS昇腾计算-通信融合算子模板库是一个聚焦于提供高性能计算通信融合类算子基础模板的代码库。项目地址: https://gitcode.com/cann/catccos1. 算子概述1.1 功能描述AllGather矩阵乘法反量化算子AllGatherMatmulDequant是一个支持INT8量化和偏置Bias的分布式矩阵乘法算子它结合AllGather通信模式用于推理场景。算子的执行是一个流水线过程首先AIV核执行AllGather操作将各个Rank的INT8输入矩阵A收集到一块共享工作空间中。接着AIC核基于这块共享空间中的A矩阵和本地的B矩阵执行INT8矩阵乘法加上偏置并将INT32结果写回。最后AIV核对INT32结果执行反量化得到最终的FP16输出。1.2 算子签名void AllGatherDequantMatmul( uint64_t fftsAddr, GM_ADDR aDevice, // 输入矩阵A: [M, K], int8 GM_ADDR bDevice, // 输入矩阵B: [K, N], int8 GM_ADDR cDevice, // INT32累加结果矩阵: [M*rankSize, N], int32 GM_ADDR bias, // 偏置向量: [N], int32 GM_ADDR symmetricPtr, // 用于Rank间通信的共享内存工作空间 (workspace) GM_ADDR dDevice, // 输出矩阵: [M*rankSize, N], fp16 GM_ADDR deviceScale, // per-channel 量化缩放因子: [N], float32 uint32_t m, uint32_t n, uint32_t k );1.3 输入输出规格参数形状数据类型描述aDevice[M, K]int8量化后的输入矩阵AbDevice[K, N]int8量化后的输入矩阵BcDevice[M*rankSize, N]int32INT32累加器结果矩阵bias[N]int32加到累加结果上的偏置向量dDevice[M*rankSize, N]fp16输出矩阵deviceScale[N]float32B矩阵的per-channel量化缩放因子symmetricPtr-GM_ADDR用于AllGather通信的共享内存工作区2. 量化算法设计2.1 核心计算流程2.2 量化与反量化公式// 伪代码: AllGather (on AIV) // Gathers matrix A from all ranks into a shared workspace gathered_a_int8 allgather(a_int8); // 伪代码: Matmul Bias (on AIC) // Uses the gathered A to perform the matmul result_int32 matmul(gathered_a_int8, b_int8) bias; // 伪代码: 反量化 (on AIV) // Note: per-token scale is not used result_fp32 result_int32 * scale; output_fp16 cast_to_fp16(result_fp32);3. 核心实现架构3.1 计算与通信分离算子采用计算AIC和通信/后处理AIV分离的设计。AIC (AI Core): 负责执行高密度的INT8 × INT8 → INT32矩阵乘法计算并加上偏置。AIV (AI Vector Core): 负责执行AllGather通信以及后续的Dequantize反量化处理。3.2 主要模块BlockMmad:catlass库提供的矩阵乘法模块通过MmadAtlasA2Pingpong调度策略执行分块的INT8矩阵乘法。CommBlockEpilogue:catccos库提供的通信Epilogue用于执行AllGather操作。它将所有Rank的INT8矩阵A收集到每个Rank。BlockEpilogueDequant:catlass库提供的后处理Epilogue用于在AllGather完成后对INT32结果进行反量化并转换为FP16。4. 内存布局设计4.1 全局内存 (Global Memory)输入布局:aDevice,bDevice,bias,deviceScale均存储在GM中。中间结果布局: 每个Rank计算出的INT32累加器结果存储在各自的GM空间中形状为[M*rankSize, N]。输出布局: 最终的FP16输出也存储在GM中形状为[M*rankSize, N]。4.2 共享内存 (Symmetric Memory)用途:symmetricPtr指向的共享内存区域被用作AllGather操作的临时工作空间Workspace。工作方式: 在AllGather过程中每个Rank需要将自己的INT8矩阵A写入到共享内存区域然后所有Rank从这个区域读取完整的矩阵A。布局:Catlass::layout::RowMajor layoutSymmetric{ WORKSPACE_STAGES * rankSize * commSizeM, K, K };5. 通信模式适配5.1 流水线计算与通信算子的核心是一个计算与通信流水线而不是分立的阶段。通信 (AIV): AIV核启动AllGather操作将一个分片(chunk)的A矩阵从所有Rank收集到共享内存工作区。计算 (AIC): 一旦某个分片的A矩阵数据准备就绪AIC核立即对该分片执行Matmul(A_chunk, B) Bias计算并将结果累加到C矩阵。后处理 (AIV): 在所有计算完成后AIV核对最终的C矩阵INT32执行反量化生成最终的FP16输出。这个过程通过多级缓冲WORKSPACE_STAGES进行流水线处理以重叠通信和计算隐藏延迟。5.2 流程图6. 工作空间管理6.1 多阶段流水线算子使用WORKSPACE_STAGES 2的多阶段流水线设计通过commInterval 3控制通信间隔实现计算与通信的重叠。6.2 内存复用策略共享内存复用: 使用共享内存作为AllGather操作的临时缓冲区避免重复分配。流水线缓冲: 通过多阶段设计实现计算与通信的并行执行。7. 总结该量化算子通过将高成本的通信操作AllGather在INT8数据上完成避免了在FP16或FP32上进行通信从而优化了性能。同时通过精确的量化参数处理确保了在多Rank环境下的计算精度。共享内存Symmetric Memory在此过程中扮演了关键的临时数据交换区的角色。8. 使用指南8.1 编译cd examples/allgather_matmul_dequant/ bash scripts/build.sh8.2 运行# 在2个设备上运行设备0和1 bash scripts/run.sh 0,1 # 在4个设备上运行设备1, 3, 5, 7 bash scripts/run.sh 1,3,5,78.3 测试形状测试形状定义在scripts/test_shapes.csv中M,K,N 16384,27392,4096 131072,8192,3072 64,16384,71688.4 数据文件在output/目录中生成以下数据文件x1_gm.bin: 量化输入矩阵A (int8)x2_gm.bin: 量化输入矩阵B (int8)bias_gm.bin: 偏置向量 (int32)c_gm.bin: 中间累加结果矩阵 (int32)d_gm.bin: 输出矩阵 (fp16)scale_x2_gm.bin: 矩阵B的per-channel量化缩放因子 (float32)golden.bin: 验证用的期望输出 (fp16)output.bin: 算子的实际输出 (fp16)8.5 验证脚本会自动验证输出结果与黄金参考的差异使用现有的verify_result.py进行精度验证。该验证脚本支持 fp16 数据类型的精度检查。8.6 调试模式设置环境变量启用调试模式export debug1 bash scripts/run.sh 0,1调试模式下会使用全1矩阵和固定缩放因子便于问题排查。8.7 环境要求Ascend Toolkit 已正确安装SHMEM 环境已配置PyTorch 支持用于数据生成和验证支持 fp16 数据类型的硬件环境【免费下载链接】catccosCATCCOS昇腾计算-通信融合算子模板库是一个聚焦于提供高性能计算通信融合类算子基础模板的代码库。项目地址: https://gitcode.com/cann/catccos创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考