aclnnWeightQuantMatmulAllReduceAddRmsNorm

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 will occur.

Function

• Description: Completes the computation of mm + AllReduce + add + rms_norm.

• Formula:

  $$mm\_out = allReduce(x1 @ (x2*antiquantScale + antiquantOffset) + bias)$$ $$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

• aclnnWeightQuantMatmulAllReduceAddRmsNorm and aclnnInplaceWeightQuantMatmulAllReduceAddRmsNorm implement the same function. The differences in usage are as follows. Select the appropriate operator based on your actual scenario.

  • aclnnWeightQuantMatmulAllReduceAddRmsNorm: Creates two output tensors, normOut and y, to store the computation results.

  • aclnnInplaceWeightQuantMatmulAllReduceAddRmsNorm: Creates one output tensor normOut. The result that would be stored in the newly created output tensor y in the non-in-place scenario is directly stored in the memory of the input tensor residual.

• Each API uses a two-phase API. You must first call the "aclnnWeightQuantMatmulAllReduceAddRmsNormGetWorkspaceSize" API to obtain the workspace size required for computation and the executor that contains the computation process, and then call the "aclnnWeightQuantMatmulAllReduceAddRmsNorm" API to perform the computation.

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

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

aclnnWeightQuantMatmulAllReduceAddRmsNormGetWorkspaceSize

• Parameters:

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Shape	Non-contiguous Tensor
  x1	Input	Left matrix of the MatMul computation, i.e., x1 in the formula.	Supports empty tensors. Currently only supports two-dimensional or three-dimensional inputs.	BFLOAT16, FLOAT16	ND	2-3	×
  x2	Input	The right matrix of the MatMul computation, i.e., x2 in the formula.	Supports empty tensors. The current version supports only two-dimensional input, with transpose and non-transpose scenarios supported. Supports non-contiguous tensors in transpose scenarios.	INT8, INT4	ND	2	√
  bias	Input	Bias in the formula.	Supports passing a null pointer. Only one-dimensional input is supported in the current version.	BFLOAT16, FLOAT16	ND	1	√
  antiquantScale	Input	That is, the antiquantScale in the formula.	The shape is (1) in the per-tensor scenario, (n) or (1, n) in the per-channel scenario where n is the size of the last dimension of x2, and (ceil(k, antiquantGroupSize), n) in the per-group scenario.	BFLOAT16, FLOAT16	ND	1-2	×
  antiquantOffset	Input	Specifies the offset parameter for pseudo-quantization computation on x2, i.e., the antiquantOffset in the formula.	Optional. Can be null. When non-null, its shape is the same as that of antiquantScale.	BFLOAT16, FLOAT16	ND	1-2	×
  residual	Input	Specifies the residual input of the AddRmsNorm fusion operator, i.e., residual in the formula.	Only three-dimensional input is supported in the current version.	FLOAT16, BFLOAT16	ND	3	×
  gamma	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	Input	A double-precision value on the host side, used to prevent division-by-zero errors, i.e., epsilon in the formula.	The value of epsilon must fall within the range (0, 1).	Double	-	-	-
  group	Input	A string that identifies the communication domain on the host side, 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	Input	Specifies the operation type as a string on the host side, indicating the reduce operation type.	Currently, only "sum" is supported.	String	-	-	-
  commTurn	Input	Host-side integer representing the communication data split count, i.e., total data volume divided by single communication volume.	Currently, only 0 is supported.	INT64	-	-	-
  streamMode	Input	Integer on the host side, representing the stream mode enumeration.	Currently, only the enumeration value 1 is supported.	INT64	-	-	-
  antiquantGroupSize	Input	Specifies the groupSize input for dequantizing x2 in per-group pseudo-quantization mode.	When per-group mode is not supported, pass 0. When supported, the value ranges from [32, min(k-1, INT_MAX)] and must be a multiple of 32. The value range of k is consistent with that of the mm API.	INT64	-	-	-
  y	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	Output	Result of mm + all_reduce + add + rms_norm, i.e., normOut in the formula.	Empty tensor is not supported. The data type is the same as the residual input.	FLOAT16, BFLOAT16	ND	3	√
  workspaceSize	Output	Returns the workspace size to be applied for on the Device side.	-	-	-	-	-
  executor	Output	Returns the operator executor, which contains the operator computation process.	-	-	-	-	-

• 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 passed x1, x2, residual, antiQuantScale, or output is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The shape of x1, x2, residual, or output does not meet the constraint requirements.
  		streamMode is not within the valid range.

aclnnWeightQuantMatmulAllReduceAddRmsNorm

• Parameters:

  Parameter	Input/Output	Description
  workspace	Input	Workspace memory address applied for on the Device side.
  workspaceSize	Input	Workspace size applied for on the Device side, obtained from the first-phase API aclnnWeightQuantMatmulAllReduceAddRmsNormGetWorkspaceSize.
  executor	Input	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:

  • aclnnWeightQuantMatmulAllReduceAddRmsNorm defaults to a non-deterministic implementation. You can enable deterministic computation by setting the HCCL_DETERMINISTIC environment variable to true.

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

• The input x1 can be two-dimensional or three-dimensional, with its shape being (b, s, k) or (s, 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. The range of b*s and s is [1, 2147483647], and the range of k and n is [1, 65535]. If bias is not null, bias is one-dimensional with its shape being (n).

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

• antiquantScale has shape (1) in the per-tensor scenario, shape (1,n)/(n) in the per-channel scenario, and shape (ceil(k,antiquantGroupSize),n) in the per-group scenario. If antiquantOffset is not null, its shape is the same as that of antiquantScale.

• The output y and normOut have the same dimension and data type as residual. If bias is not null, its shape size equals the size of the last dimension of normOut.

• The data type of x2 must be int8 or int4. The data types of x1, bias, residual, gamma, y, and normOut as computation inputs must be consistent.

• Only the transposed/non-transposed scenarios of the x2 matrix are supported. The x1 matrix supports the non-transposed scenario.

• The antiquantGroupSize value falls within the range [32, min(k-1, INT_MAX)] and is a multiple of 32.

• Supports 1, 2, 4, or 8 cards, and only supports HCCS all-mesh networking.

• Supports empty tensors where (b*s) and n are 0, but does not support empty tensors where k is 0.

• Atlas A2 training products/Atlas A2 inference products: A communication-computation fused MC2 operator in a model only supports 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_weight_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 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-side memory by calling 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);
    // Compute 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 launchOneThreadweightQuantmatmulAllReduceAddRmsNorm(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> antiquantScaleShape = {128};
    std::vector<int64_t> antiquantOffsetShape = {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 *antiquantScaleDeviceAddr = nullptr;
    void *antiquantOffsetDeviceAddr = nullptr;
    void *residualDeviceAddr = nullptr;
    void *gammaDeviceAddr = nullptr;
    void *yDeviceAddr = nullptr;
    void *normOutDeviceAddr = nullptr;
    aclTensor *x1 = nullptr;
    aclTensor *x2 = nullptr;
    aclTensor *bias = nullptr;
    aclTensor *antiquantScale = nullptr;
    aclTensor *antiquantOffset = 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;
    int64_t antiquantGroupSize = 0;
    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 antiquantScaleShapeSize = GetShapeSize(antiquantScaleShape);
    long long antiquantOffsetShapeSize = GetShapeSize(antiquantOffsetShape);
    long long residualShapeSize = GetShapeSize(residualShape);
    long long gammaShapeSize = GetShapeSize(gammaShape);
    long long yShapeSize = GetShapeSize(yShape);
    long long normOutShapeSize = GetShapeSize(normOutShape);
    std::vector<op::fp16_t> x1HostData(x1ShapeSize, 1);
    std::vector<int8_t> x2HostData(x2ShapeSize, 1);
    std::vector<op::fp16_t> biasHostData(biasShapeSize, 1);
    std::vector<op::fp16_t> antiquantScaleHostData(antiquantScaleShapeSize, 1);
    std::vector<op::fp16_t> antiquantOffsetHostData(antiquantOffsetShapeSize, 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_FLOAT16, &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_FLOAT16, &bias);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(antiquantScaleHostData, antiquantScaleShape, &antiquantScaleDeviceAddr,
                        aclDataType::ACL_FLOAT16, &antiquantScale);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(antiquantOffsetHostData, antiquantOffsetShape, &antiquantOffsetDeviceAddr,
                        aclDataType::ACL_FLOAT16, &antiquantOffset);
    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 aclnnWeightQuantMatmulAllReduceAddRmsNorm.
    // Call the first-phase API.
    ret = aclnnWeightQuantMatmulAllReduceAddRmsNormGetWorkspaceSize(x1, x2, bias, antiquantScale, antiquantOffset,
                                                        residual, gamma, epsilon, hcom_name,
                                                        "sum", commTurn, streamMode, antiquantGroupSize, y, normOut,
                                                        &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS,
            LOG_PRINT("aclnnWeightQuantMatmulAllReduceAddRmsNormGetWorkspaceSize 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 = aclnnWeightQuantMatmulAllReduceAddRmsNorm(workspaceAddr, workspaceSize, executor, args.stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnWeightQuantMatmulAllReduceAddRmsNorm 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 aclnnWeightQuantMatmulAllReduceAddRmsNorm 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 (antiquantScale != nullptr) {
        aclDestroyTensor(antiquantScale);
    }
    if (antiquantOffset != nullptr) {
        aclDestroyTensor(antiquantOffset);
    }
    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 (antiquantScaleDeviceAddr != nullptr) {
        aclrtFree(antiquantScaleDeviceAddr);
    }
    if (antiquantOffsetDeviceAddr != nullptr) {
        aclrtFree(antiquantOffsetDeviceAddr);
    }
    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(&launchOneThreadweightQuantmatmulAllReduceAddRmsNorm, std::ref(args[rankId])));
    }
    for (uint32_t rankId = 0; rankId < ndev; rankId++) {
        threads[rankId]->join();
    }
    aclFinalize();
    return 0;
}