aclnnQuantMatmulAllReduceAddRmsNorm

NOTE

📄 View Source Code

Applicable Products

NOTE When using this API, ensure that both the driver firmware package and the CANN package are of the matching 8.0.RC2 version or a matching later version. Otherwise, errors such as Bus Error may occur.

Function

• Description: Performs mm + all_reduce + add + rms_norm computation.

• Formula:

  $$mm\_out = allReduce(dequantScale * (x1_{int8}@x2_{int8} + bias_{int32}))$$ $$y = mm\_out + residual$$ $$normOut = \frac{y}{RMS(y)} * gamma, RMS(y) = \sqrt{\frac{1}{d} \sum_{i=1}^{d} y_{i}^{2} + epsilon}$$

Function Prototype

• aclnnQuantMatmulAllReduceAddRmsNorm and aclnnInplaceQuantMatmulAllReduceAddRmsNorm implement the same functionality. The differences in usage are as follows. Select the appropriate operator based on your actual scenario.

  • aclnnQuantMatmulAllReduceAddRmsNorm: Two output tensors, normOut and y, need to be created to store the computation results.

  • aclnnInplaceQuantMatmulAllReduceAddRmsNorm: Only one output tensor, normOut, needs to be created. The result that would be stored in the newly created output tensor y in the non-inplace scenario is directly stored in the memory of the input tensor residual.

• Each operator is divided into a two-phase API. You must first call the "aclnnQuantMatmulAllReduceAddRmsNormGetWorkspaceSize" API to obtain the workspace size required for computation and the executor that contains the operator computation process, and then call the "aclnnQuantMatmulAllReduceAddRmsNorm" API to perform the computation.

aclnnStatus aclnnQuantMatmulAllReduceAddRmsNormGetWorkspaceSize(
    const aclTensor *x1,
    const aclTensor *x2,
    const aclTensor *bias,
    const aclTensor *dequantScale,
    const aclTensor *residual,
    const aclTensor *gamma,
    double           epsilon,
    const char      *group,
    const char      *reduceOp,
    int64_t          commTurn,
    int64_t          streamMode,
    const aclTensor *y,
    const aclTensor *normOut,
    uint64_t        *workspaceSize,
    aclOpExecutor  **executor)

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

aclnnQuantMatmulAllReduceAddRmsNormGetWorkspaceSize

• Parameters:

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Shape	Non-contiguous Tensor
  x1 (aclTensor)	Input	Left matrix for the MatMul computation, i.e., x1 in the formula.	Supports empty tensors. Must have the same data type as x2. In the current version, only two-dimensional or three-dimensional inputs are supported.	INT8	ND	2-3	×
  x2 (aclTensor)	Input	The right matrix of the MatMul computation, i.e., x2 in the formula.	Supports empty tensors. Must have the same data type as x1. In the current version, only two-dimensional input is supported, with transpose and non-transpose scenarios supported. Supports non-contiguous tensors in transpose scenarios.	INT8	ND	2	√
  bias (aclTensor)	Input	Refers to bias in the formula.	Supports passing a null pointer. Only one-dimensional input is supported in the current version.	INT32	ND	1	√
  dequantScale (aclTensor)	Input	Full quantization coefficient after MatMul computation, i.e., dequantScale in the formula.	Shape is (1) in the pertensor scenario, and (n) or (1, n) in the perchannel scenario.	UINT64, INT64, BFLOAT16	ND	1-2	×
  residual (aclTensor)	Input	Residual input of the AddRmsNorm fusion operator, i.e., residual in the formula.	In the in-place scenario, residual is used as the output address of y. Currently, only three-dimensional input is supported.	FLOAT16, BFLOAT16	ND	3	×
  gamma (aclTensor)	Input	RmsNorm computation input of the AddRmsNorm fusion operator, i.e., gamma in the formula.	Only one-dimensional input is supported in the current version.	FLOAT16, BFLOAT16	ND	1	×
  epsilon (double)	Input	A double-precision value on the host side, used to prevent division-by-zero errors, i.e., epsilon in the formula.	The epsilon value falls within the range (0, 1).	Double	-	-	-
  group (char)	Input	A string on the host side that identifies the communication domain, i.e., the communication domain name.	Obtained through the HCCL API "extern HcclResult HcclGetCommName(HcclComm comm, char* commName);", where commName is the group.	String	-	-	-
  reduceOp (char)	Input	A string identifying the operation type on the host side, indicating the reduce operation type.	Currently, only "sum" is supported.	String	-	-	-
  commTurn (int64_t)	Input	An integer on the host side, representing the communication data split count, i.e., total data volume divided by single communication volume.	Only the value 0 is supported in the current version.	INT64	-	-	-
  streamMode (int64_t)	Input	Integer on the host side, indicating the stream mode enumeration.	Currently only enumeration value 1 is supported.	INT64	-	-	-
  y (aclTensor)	Output	Result of mm + all_reduce + add, i.e., y in the formula.	Empty tensors are not supported. The data type is the same as the residual input.	FLOAT16, BFLOAT16	ND	3	√
  normOut (aclTensor)	Output	Result of mm + all_reduce + add + rms_norm, i.e., normOut in the formula.	Null tensors are not supported. The data type is the same as that of the residual input.	FLOAT16, BFLOAT16	ND	2	√
  workspaceSize (uint64_t)	Output	Returns the workspace size to be requested on the Device side.	-	-	-	-	-
  executor (aclOpExecutor)	Output	Returns the op executor, which contains the operator computation flow.	-	-	-	-	-

• Return value

  aclnnStatus: return code. For details, see aclnn Return Code.

  The first-phase API performs input parameter validation and returns an error in the following scenarios:

  Return Value	Error Code	Description
  ACLNN_ERR_PARAM_NULLPTR	161001	The input x1, x2, residual, gamma, or dequantScale is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The shapes of x1, x2, and residual do not meet the constraint requirements.
  		The streamMode is not within the valid range.
  		The size of the last dimension of x1, i.e., the k value, is 0.

aclnnQuantMatmulAllReduceAddRmsNorm

• Parameters:

  Parameter	Input/Output	Description
  workspace	Input	Address of the workspace memory applied on the device side.
  workspaceSize	Input	Size of the workspace memory applied on the device side, obtained from the first-phase API aclnnQuantMatmulAllReduceAddRmsNormGetWorkspaceSize.
  executor	Input	Specifies the operator executor, which contains the operator computation flow.
  stream	Input	Specifies the stream for task execution.

• Return value

  aclnnStatus: return code. For details, see aclnn Return Code.

Constraints

• Deterministic computation:

  • aclnnQuantMatmulAllReduceAddRmsNorm defaults to a non-deterministic implementation. Deterministic computation can be enabled by setting the HCCL_DETERMINISTIC environment variable to true.

• The usage scenario is the same as that of the fused operator aclnnQuantMatmulAllReduce: MC2 is disabled in incremental scenarios and enabled in full-scale scenarios.

• The input x1 can be two-dimensional or three-dimensional, with its shape being (b, s, k) or (m, k). x2 must be two-dimensional, with its shape being (k, n). The axes must meet the input requirements of the mm operator, with the k axes being equal. If bias is not empty, bias is one-dimensional with its shape being (n).

• The size of m does not exceed 2147483647. The size of the last dimension of x1 and x2 does not exceed 65535. The last dimension of x1 refers to k, and the last dimension of x2 refers to k when transposed or n when not transposed.

• The input residual must be three-dimensional with the shape (b, s, n). When x1 is two-dimensional, (b*s) of residual equals m of x1. The input gamma has the shape (n).

• The output y and normOut have the same dimensions and data type as residual, with the shape (b, s, n).

• If the output residual type is FLOAT16, the type of dequantScale is UINT64 or INT64. If the output residual type is BFLOAT16, the type of dequantScale is BFLOAT16.

• The data type of x1 and x2 is INT8. The data type of bias is INT32. The data types of residual, gamma, y, and normOut as computation inputs must be consistent.

• Only x2 matrix transposition/non-transposition is supported, and x1 matrix supports the non-transposition scenario.

• 1, 2, 4, and 8 cards are supported, and only HCCS all-mesh networking is supported.

• Empty tensors with (b*s) and n being 0 are supported, but empty tensors with k being 0 are not supported.

• Atlas A2 training products/Atlas A2 inference products: A fused compute MC2 operator in a model supports only the same communication domain.

Example

The following provides example code for reference only. For details about compilation and running, see Compile and Run Samples.

#include <iostream>
#include <vector>
#include <thread>
#include "hccl/hccl.h"
#include "aclnn/opdev/fp16_t.h"
#include "aclnnop/aclnn_quant_matmul_all_reduce_add_rms_norm.h"

int ndev = 8;

#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);
    // Call aclrtMalloc to apply for device-side memory.
    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);
    // Copy host-side data to device memory using aclrtMemcpy.
    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);
    // Calculate the strides of a contiguous tensor.
    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];
    }
    // Create an aclTensor by calling the aclCreateTensor API.
    *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 launchOneThreadQuantMatmulAllReduceAddRmsNorm(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> dequantScaleShape = {128};
    std::vector<int64_t> residualShape = {1, 32, 128};
    std::vector<int64_t> gammaShape = {128};
    std::vector<int64_t> yShape = {1, 32, 128};
    std::vector<int64_t> normOutShape = {1, 32, 128};
    void *x1DeviceAddr = nullptr;
    void *x2DeviceAddr = nullptr;
    void *biasDeviceAddr = nullptr;
    void *dequantScaleDeviceAddr = nullptr;
    void *residualDeviceAddr = nullptr;
    void *gammaDeviceAddr = nullptr;
    void *yDeviceAddr = nullptr;
    void *normOutDeviceAddr = nullptr;
    aclTensor *x1 = nullptr;
    aclTensor *x2 = nullptr;
    aclTensor *bias = nullptr;
    aclTensor *dequantScale = nullptr;
    aclTensor *residual = nullptr;
    aclTensor *gamma = nullptr;
    aclTensor *y = nullptr;
    aclTensor *normOut = nullptr;

    int64_t commTurn = 0;
    int64_t streamMode = 1;
    double  epsilon = 0.000001;
    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 dequantScaleShapeSize = GetShapeSize(dequantScaleShape);
    long long residualShapeSize = GetShapeSize(residualShape);
    long long gammaShapeSize = GetShapeSize(gammaShape);
    long long yShapeSize = GetShapeSize(yShape);
    long long normOutShapeSize = GetShapeSize(normOutShape);

    std::vector<int8_t> x1HostData(x1ShapeSize, 1);
    std::vector<int8_t> x2HostData(x2ShapeSize, 1);
    std::vector<int32_t> biasHostData(biasShapeSize, 1);
    std::vector<uint64_t> dequantScaleHostData(dequantScaleShapeSize, 1);
    std::vector<op::fp16_t> residualHostData(residualShapeSize, 1);
    std::vector<op::fp16_t> gammaHostData(gammaShapeSize, 1);
    std::vector<op::fp16_t> yHostData(yShapeSize, 0);
    std::vector<op::fp16_t> normOutHostData(normOutShapeSize, 0);
    // Create a tensor.
    ret = CreateAclTensor(x1HostData, x1Shape, &x1DeviceAddr, aclDataType::ACL_INT8, &x1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(x2HostData, x2Shape, &x2DeviceAddr, aclDataType::ACL_INT8, &x2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(biasHostData, biasShape, &biasDeviceAddr, aclDataType::ACL_INT32, &bias);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(dequantScaleHostData, dequantScaleShape, &dequantScaleDeviceAddr,
                        aclDataType::ACL_UINT64, &dequantScale);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(residualHostData, residualShape, &residualDeviceAddr, aclDataType::ACL_FLOAT16, &residual);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(gammaHostData, gammaShape, &gammaDeviceAddr, aclDataType::ACL_FLOAT16, &gamma);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(yHostData, yShape, &yDeviceAddr, aclDataType::ACL_FLOAT16, &y);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(normOutHostData, normOutShape, &normOutDeviceAddr, aclDataType::ACL_FLOAT16, &normOut);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Example of calling the aclnnQuantMatmulAllReduceAddRmsNorm API.
    // Call the first-phase API.
    ret = aclnnQuantMatmulAllReduceAddRmsNormGetWorkspaceSize(x1, x2, bias, dequantScale, residual, gamma, epsilon,
                                        hcom_name, "sum", commTurn, streamMode, y, normOut,
                                        &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS,
            LOG_PRINT("aclnnQuantMatmulAllReduceAddRmsNormGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
    // Apply for device memory based on the workspaceSize calculated by the first-phase API.
    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);
    }
    // Call the second-phase API.
    ret = aclnnQuantMatmulAllReduceAddRmsNorm(workspaceAddr, workspaceSize, executor, args.stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnQuantMatmulAllReduceAddRmsNorm failed. ERROR: %d\n", ret); return ret);
    //Synchronously wait for task execution to complete (boilerplate).
    ret = aclrtSynchronizeStreamWithTimeout(args.stream, 10000);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
    LOG_PRINT("device%d aclnnQuantMatmulAllReduceAddRmsNorm execute success \n", args.rankId);
    // Release device resources. Modify based on the specific API definition.

    if (x1 != nullptr) {
        aclDestroyTensor(x1);
    }
    if (x2 != nullptr) {
        aclDestroyTensor(x2);
    }
    if (bias != nullptr) {
        aclDestroyTensor(bias);
    }
    if (dequantScale != nullptr) {
        aclDestroyTensor(dequantScale);
    }
    if (residual != nullptr) {
        aclDestroyTensor(residual);
    }
    if (gamma != nullptr) {
        aclDestroyTensor(gamma);
    }
    if (y != nullptr) {
        aclDestroyTensor(y);
    }
    if (normOut != nullptr) {
        aclDestroyTensor(normOut);
    }
    if (x1DeviceAddr != nullptr) {
        aclrtFree(x1DeviceAddr);
    }
    if (x2DeviceAddr != nullptr) {
        aclrtFree(x2DeviceAddr);
    }
    if (biasDeviceAddr != nullptr) {
        aclrtFree(biasDeviceAddr);
    }
    if (dequantScaleDeviceAddr != nullptr) {
        aclrtFree(dequantScaleDeviceAddr);
    }
    if (residualDeviceAddr != nullptr) {
        aclrtFree(residualDeviceAddr);
    }
    if (gammaDeviceAddr != nullptr) {
        aclrtFree(gammaDeviceAddr);
    }
    if (yDeviceAddr != nullptr) {
        aclrtFree(yDeviceAddr);
    }
    if (normOutDeviceAddr != nullptr) {
        aclrtFree(normOutDeviceAddr);
    }
    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);
    // Initialize the collective communication domain.
    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);
    }
    // Start multithreading.
    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(&launchOneThreadQuantMatmulAllReduceAddRmsNorm, std::ref(args[rankId])));
    }
    for (uint32_t rankId = 0; rankId < ndev; rankId++) {
        threads[rankId]->join();
    }
    aclFinalize();
    return 0;
}