CANN矩阵乘与AllReduce融合算子

aclnnMatmulAllReduceV3

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer

📄 查看源码

产品支持情况

说明：使用该接口时，请确保驱动固件包和CANN包都为配套的8.0.RC2版本或者配套的更高版本，否则将会引发报错，比如Bus Error等。

功能说明

• 接口功能：完成MatMul计算与AllReduce通信融合，兼容aclnnMatmulAllReduceV2支持的功能，新增通信引擎参数commMode，支持自选通信引擎。

• 计算公式：

  $$ output = AllReduce(x1 @ x2 + bias + x3) $$

函数原型

每个算子分为两段式接口，必须先调用aclnnMatmulAllReduceV3GetWorkspaceSize接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用aclnnMatmulAllReduceV3接口执行计算。

aclnnStatus aclnnMatmulAllReduceV3GetWorkspaceSize( const aclTensor *x1, const aclTensor *x2, const aclTensor *bias, const aclTensor *x3, const char *group, const char *reduceOp, const char *commMode, int64_t commTurn, int64_t streamMode, const aclTensor *output, uint64_t *workspaceSize, aclOpExecutor **executor)

aclnnStatus aclnnMatmulAllReduceV3( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)

aclnnMatmulAllReduceV3GetWorkspaceSize

• 参数说明

  参数名	输入/输出	描述	使用说明	数据类型	数据格式	维度(shape)	非连续Tensor
  x1	输入	MatMul计算的左矩阵，即计算公式中的x1。	当前版本仅支持二维或者三维输入。 支持不转置场景。	FLOAT16、BFLOAT16	ND	2-3	×
  x2	输入	MatMul计算的右矩阵，即计算公式中的x2。	当前版本仅支持两维输入。 支持转置/不转置场景。 ND格式下支持最后两轴转置情况下的非连续的tensor，其他非连续tensor不支持	FLOAT16、BFLOAT16	ND	2	×
  x3	输入	MatMul计算后的add计算，即计算公式中的x3。	shape与MatMul计算后的shape一致。	FLOAT16、BFLOAT16	ND	2	√
  bias	输入	对应计算公式中的bias偏移。	当前版本仅支持一维输入。	FLOAT16、BFLOAT16	ND	0-1	√
  group	输入	通信域名称。	通过Hccl提供的接口“extern HcclResult HcclGetCommName(HcclComm comm, char* commName);”获取，其中commName即为group。	String	-	-	-
  reduceOp	输入	reduce操作类型。	当前版本仅支持输入"sum"。	String	-	-	-
  commMode	输入	通信引擎参数。	用于指定通信引擎，根据芯片型号有不同输入限制。通信引擎简单介绍和输入约束详见约束说明。	String	-	-	-
  commTurn	输入	通信数据切分数，即总数据量/单次通信量。	当前版本仅支持输入0。	INT64	-	-	-
  streamMode	输入	AscendCL流模式的枚举。	当前版本仅支持枚举值1。	INT64	-	-	-
  output	输出	MatMul计算与AllReduce通信的结果，即计算公式中的output。	output的维数与x1一致。	FLOAT16、BFLOAT16	ND	2-3	√
  workspaceSize	输出	返回需要在Device侧申请的workspace大小。	-	-	-	-	-
  executor	输出	返回op执行器，包含了算子计算流程。	-	-	-	-	-

• 返回值

  返回aclnnStatus状态码，具体参见aclnn返回码。

  第一阶段接口完成入参校验，出现以下场景报错：

  返回值	错误码	描述
  ACLNN_ERR_PARAM_NULLPTR	161001	传入的x1、x2或output是空指针。
  ACLNN_ERR_PARAM_INVALID	161002	x1、x2、bias、x3或output的数据类型不在支持的范围之内。
  		reduceOp、streamMode不在合法范围内。
  		x1、x2、bias、x3或output的shape不符合约束要求。

aclnnMatmulAllReduceV3

• 参数说明

  参数名	输入/输出	描述
  workspace	输入	在Device侧申请的workspace内存地址。
  workspaceSize	输入	在Device侧申请的workspace大小，由第一段接口aclnnMatmulAllReduceV3GetWorkspaceSize获取。
  executor	输入	op执行器，包含了算子计算流程。
  stream	输入	指定执行任务的Stream。

• 返回值

  返回aclnnStatus状态码，具体参见aclnn返回码。

约束说明

• 通信引擎commMode支持度：
  • Atlas A2 训练系列产品/Atlas A2 推理系列产品 ：目前不支持指定通信引擎，commMode仅允许输入为空字符串，默认使用AICPU。
  • Ascend 950PR/Ascend 950DT ：目前通信引擎支持AICPU和CCU，commMode允许输入为"ai_cpu"、"ccu"和空字符串，前两者直接指定AICPU和CCU，空字符串表示自适应模式，优先选择最佳的通信引擎，当前版本默认使用CCU作为通信引擎。
  • 同一条通信链路内，只能选择同一种通信引擎。
  • AICPU和CCU通信引擎简单介绍：
    • AICPU：不占用计算核，通信效率高，但通信静态开销较大，对小数据量通信场景不友好。适用于大数据高带宽场景。
    • CCU：能够减少访存带宽与计算核占用，但受限于片上资源，支持的通信域数量有限。适用于高带宽、低时延的通信场景。
    • 更详细的通信引擎介绍，请参考通信引擎。
• 确定性计算：
  • Atlas A2 训练系列产品/Atlas A2 推理系列产品 ：aclnnMatmulAllReduceV3默认非确定性实现，支持通过配置HCCL_DETERMINISTIC环境变量为true开启确定性计算。
  • Ascend 950PR/Ascend 950DT ：aclnnMatmulAllReduceV3默认确定性实现。
• 增量场景不使能MC2，全量场景使能MC2。
• 输入x1可为二维或者三维，其shape为(b, s, k)或者(m, k)。x2必须是二维，其shape为(k, n)，轴满足mm算子入参要求，k轴相等。bias若非空，其shape为(n)。
• b*s、m、k、n的值均不得超过2147483647(INT32_MAX)。
• 当输入x1的shape为(b, s, k)时，输出output的shape为(b, s, n)；当输入x1的shape为(m, k)时，输出output的shape为(m, n)。
• x1、x2、bias计算输入的数据类型要和output计算输出的数据类型一致，传入的x1、x2、x3或者output不为空指针。
• 仅支持HCCS链路all mesh组网。
  • Atlas A2 训练系列产品/Atlas A2 推理系列产品 ：支持1、2、4、8卡。
  • Ascend 950PR/Ascend 950DT ：支持1、2、4、8、16、32、64卡。
• Atlas A2 训练系列产品/Atlas A2 推理系列产品 ：一个模型中的通算融合MC2算子，仅支持相同通信域。
• 空tensor支持度：
  • 支持k为0的场景，输出为bias + x3。
  • 支持bs/m/n为0，此时传入的输出也应该是空tensor，此场景不进入kernel计算，直接返回。

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考编译与运行样例。

说明：本示例代码调用了部分HCCL集合通信库接口：HcclGetCommName、HcclCommInitAll、HcclCommDestroy, 请参考<<HCCL API (C)>>。

• Atlas A2 训练系列产品/Atlas A2 推理系列产品 、 Ascend 950PR/Ascend 950DT ：

  #include <iostream> #include <vector> #include <thread> #include "hccl/hccl.h" #include "aclnn/opdev/fp16_t.h" #include "aclnnop/aclnn_matmul_all_reduce_v2.h" int ndev = 2; #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vector<int64_t> &shape) { int64_t shapeSize = 1; for (auto i: shape) { shapeSize *= i; } return shapeSize; } template<typename T> int CreateAclTensor(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr, aclDataType dataType, aclTensor **tensor) { auto size = GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret); // 计算连续tensor的strides std::vector<int64_t> strides(shape.size(), 1); for (int64_t i = shape.size() - 2; i >= 0; i--) { strides[i] = shape[i + 1] * strides[i + 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } struct Args { uint32_t rankId; HcclComm hcclComm; aclrtStream stream; aclrtContext context; }; int launchOneThreadMatmulAllReduce(Args &args) { int ret; ret = aclrtSetCurrentContext(args.context); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetCurrentContext failed. ERROR: %d\n", ret); return ret); char hcom_name[128]; ret = HcclGetCommName(args.hcclComm, hcom_name); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclGetCommName failed. ret = %d \n", ret); return -1); LOG_PRINT("[INFO] rank %d hcom: %s stream: %p, context : %p\n", args.rankId, hcom_name, args.stream, args.context); std::vector<int64_t> x1Shape = {32, 64}; std::vector<int64_t> x2Shape = {64, 128}; std::vector<int64_t> biasShape = {128}; std::vector<int64_t> x3Shape = {32, 128}; std::vector<int64_t> outShape = {32, 128}; void *x1DeviceAddr = nullptr; void *x2DeviceAddr = nullptr; void *biasDeviceAddr = nullptr; void *x3DeviceAddr = nullptr; void *outDeviceAddr = nullptr; aclTensor *x1 = nullptr; aclTensor *x2 = nullptr; aclTensor *bias = nullptr; aclTensor *x3 = nullptr; aclTensor *out = nullptr; int64_t commTurn = 0; int64_t streamMode = 1; uint64_t workspaceSize = 0; aclOpExecutor *executor; void *workspaceAddr = nullptr; long long x1ShapeSize = GetShapeSize(x1Shape); long long x2ShapeSize = GetShapeSize(x2Shape); long long biasShapeSize = GetShapeSize(biasShape); long long x3ShapeSize = GetShapeSize(x3Shape); long long outShapeSize = GetShapeSize(outShape); std::vector<op::fp16_t> x1HostData(x1ShapeSize, 1); std::vector<op::fp16_t> x2HostData(x2ShapeSize, 1); std::vector<op::fp16_t> biasHostData(biasShapeSize, 1); std::vector<op::fp16_t> x3HostData(x3ShapeSize, 1); std::vector<op::fp16_t> outHostData(outShapeSize, 0); // 创建 tensor ret = CreateAclTensor(x1HostData, x1Shape, &x1DeviceAddr, aclDataType::ACL_FLOAT16, &x1); CHECK_RET(ret == ACL_SUCCESS, return ret); ret = CreateAclTensor(x2HostData, x2Shape, &x2DeviceAddr, aclDataType::ACL_FLOAT16, &x2); CHECK_RET(ret == ACL_SUCCESS, return ret); ret = CreateAclTensor(biasHostData, biasShape, &biasDeviceAddr, aclDataType::ACL_FLOAT16, &bias); CHECK_RET(ret == ACL_SUCCESS, return ret); ret = CreateAclTensor(x3HostData, x3Shape, &x3DeviceAddr, aclDataType::ACL_FLOAT16, &x3); CHECK_RET(ret == ACL_SUCCESS, return ret); ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT16, &out); CHECK_RET(ret == ACL_SUCCESS, return ret); // 调用第一段接口, x3位置传入out ret = aclnnMatmulAllReduceV3GetWorkspaceSize(x1, x2, bias, x3, hcom_name, "sum", "", commTurn, streamMode, out, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMatmulAllReduceV3GetWorkspaceSize failed. ERROR: %d\n", ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 if (workspaceSize > 0) { ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret); } // 调用第二段接口 ret = aclnnMatmulAllReduceV3(workspaceAddr, workspaceSize, executor, args.stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMatmulAllReduceV3 failed. ERROR: %d\n", ret); return ret); //（固定写法）同步等待任务执行结束 ret = aclrtSynchronizeStreamWithTimeout(args.stream, 10000); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret); LOG_PRINT("device%d aclnnMatmulAllReduceV3 execute success \n", args.rankId); // 释放device资源，需要根据具体API的接口定义修改 if (x1 != nullptr) { aclDestroyTensor(x1); } if (x2 != nullptr) { aclDestroyTensor(x2); } if (bias != nullptr) { aclDestroyTensor(bias); } if (x3 != nullptr) { aclDestroyTensor(x3); } if (out != nullptr) { aclDestroyTensor(out); } if (x1DeviceAddr != nullptr) { aclrtFree(x1DeviceAddr); } if (x2DeviceAddr != nullptr) { aclrtFree(x2DeviceAddr); } if (biasDeviceAddr != nullptr) { aclrtFree(biasDeviceAddr); } if (x3DeviceAddr != nullptr) { aclrtFree(x3DeviceAddr); } if (outDeviceAddr != nullptr) { aclrtFree(outDeviceAddr); } if (workspaceSize > 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(args.stream); HcclCommDestroy(args.hcclComm); aclrtDestroyContext(args.context); aclrtResetDevice(args.rankId); return 0; } int main(int argc, char *argv[]) { int ret; int32_t devices[ndev]; for (int i = 0; i < ndev; i++) { devices[i] = i; } HcclComm comms[128]; ret = aclInit(nullptr); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret); // 初始化集合通信域 for (int i = 0; i < ndev; i++) { ret = aclrtSetDevice(devices[i]); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret); } ret = HcclCommInitAll(ndev, devices, comms); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("HcclCommInitAll failed. ERROR: %d\n", ret); return ret); Args args[ndev]; aclrtStream stream[ndev]; aclrtContext context[ndev]; for (uint32_t rankId = 0; rankId < ndev; rankId++) { ret = aclrtSetDevice(rankId); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret); ret = aclrtCreateContext(&context[rankId], rankId); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateContext failed. ERROR: %d\n", ret); return ret); ret = aclrtCreateStream(&stream[rankId]); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret); } // 启动多线程 std::vector<std::unique_ptr<std::thread>> threads(ndev); for (uint32_t rankId = 0; rankId < ndev; rankId++) { args[rankId].rankId = rankId; args[rankId].hcclComm = comms[rankId]; args[rankId].stream = stream[rankId]; args[rankId].context = context[rankId]; threads[rankId].reset(new(std::nothrow) std::thread(&launchOneThreadMatmulAllReduce, std::ref(args[rankId]))); } for (uint32_t rankId = 0; rankId < ndev; rankId++) { threads[rankId]->join(); } aclFinalize(); return 0; }

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer