aclnnQuantAllReduce

Applicable Products

ProductSupported
Ascend 950PR/Ascend 950DT√
Atlas A3 training products/Atlas A3 inference products×
Atlas A2 training products/Atlas A2 inference products×
Atlas 200I/500 A2 inference products×
Atlas inference products×
Atlas training 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: Implements AllReduce communication for low-bit data, dequantizes the data during communication, and outputs the communication result.

The specific implementation varies depending on the data size in two cases:

  • Formula:AllGatherData=AllGather(x)AllGatherData = AllGather(x) AllGatherScales=AllGather(scales)AllGatherScales = AllGather(scales) output=Reduce(AllGatherScales∗AllGatherData)output = Reduce(AllGatherScales * AllGatherData) The Reduce computation performs a reduce operation on data from different ranks.

Function Prototype

This operator is divided into a two-phase API. The "aclnnQuantAllReduceGetWorkspaceSize" API must be called first to obtain the workspace size required for computation and the executor that encapsulates the operator computation flow. Then, the "aclnnQuantAllReduce" API is called to perform the computation.

aclnnStatus aclnnQuantAllReduceGetWorkspaceSize(
    const aclTensor*  x,
    const aclTensor*  scales,
    const char*       group,
    const char*       reduceOp,
    aclTensor*        output,
    uint64_t*         workspaceSize,
    aclOpExecutor**   executor)
aclnnStatus aclnnQuantAllReduce(
    void*             workspace,
    uint64_t          workspaceSize,
    aclOpExecutor*    executor,
    const aclrtStream stream)

aclnnQuantAllReduceGetWorkspaceSize

  • Parameters

    Parameter Input/Output Description Instruction Data Type Data Format Dimension (Shape) Non-contiguous Tensor
    x Input Input x in the formula.
    • Null tensor is not supported.
    • Supported shapes: (bs, H) or (b, s, H), where b is batch size, s is sequence length, and H is hidden size.
    INT8, HIFLOAT8, FLOAT8_E4M3FN, FLOAT8_E5M2 ND 2-3 √
    scales Input Input scales in the formula.
    • Null tensor is not supported.
    • When the data type of scales is FLOAT8_E8M0, the data type of x must be FLOAT8_E4M3FN or FLOAT8_E5M2, the shape of x must be (bs, H) or (b, s, H), and the shape of scales must correspond to the shape of x as (bs, H/64, 2) or (b, s, H/64, 2).
    • When the data type of scales is FLOAT, the data type of x must be INT8, HIFLOAT8, FLOAT8_E4M3FN, or FLOAT8_E5M2, the shape of x must be (bs, H) or (b, s, H), and the shape of scales must correspond to the shape of x as (bs, H/128) or (b, s, H/128).
    FLOAT、FLOAT8_E8M0 ND 2-4 √
    group Attribute Communication domain identifier.
    • A string that identifies the communication group on the Host side, representing the communication domain name.
    • Obtained through the Hccl API "extern HcclResult HcclGetCommName(HcclComm comm, char* commName);", where commName is the group.
    Char*, String - - -
    reduceOp Optional attribute. Type of the reduce operation in the formula. Currently, only the "sum" operation is supported. Char*, String - - -
    output Output The output in the formula.
    • Null not supported.
    • Supported shapes are (bs, H) or (b, s, H). The shape of output is consistent with that of x.
    FLOAT, FLOAT16, BFLOAT16 ND 2-3 √
    workspaceSize Output Returns the workspace size to be applied on the device side. - - - - -
    executor Output Returns the operator 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 x, scales, or output is a null pointer.
    ACLNN_ERR_PARAM_INVALID 161002 The data type of x, scales, or output is not within the supported range.
    The data types and shapes of x and scales do not match.
    The dimensions of x and scales are not within the supported range.

aclnnQuantAllReduce

  • Parameters

    Parameter Input/Output Description
    workspace Input Specifies the workspace memory address applied on the Device side.
    workspaceSize Input Specifies the workspace size applied on the Device side, which is obtained by the first-phase API aclnnQuantAllReduceGetWorkspaceSize.
    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:

  • aclnnQuantAllReduce defaults to a non-deterministic implementation. Deterministic computation can be enabled through aclrtCtxSetSysParamOpt.

  • When the data type of x is FLOAT8_E4M3FN or FLOAT8_E5M2 and the data type of scales is FLOAT8_E8M0, the quantization mode of the input data is mx quantization.

  • When the data type of x is INT8, HIFLOAT8, FLOAT8_E4M3FN, or FLOAT8_E5M2 and the data type of scales is FLOAT, the quantization mode of the input data is pertoken-pergroup quantization (groupSize=128).

  • Enabled only on Ascend 950 series platforms.

  • Null tensor input is not supported.

  • The communication domain size supports 2, 4, and 8.

  • Communication domain usage constraint: Only the aclnnQuantAllReduce and aclnnQuantReduceScatter operators are allowed to be executed consecutively within the same communication domain, and no other communication operators are allowed in this communication domain.

  • HCCL_BUFFSIZE: Before calling this operator, check whether the value of the HCCL_BUFFSIZE environment variable is reasonable. This environment variable indicates the memory size occupied by a single communication domain, in MB. The default value is 200 MB if not configured. The requirement is HCCL_BUFFSIZE >= 2 * (xDataSize + scalesDataSize + 1). xDataSize is the data size of the input x, calculated as: xDataSize = b * s * H * 1 (Byte). scalesDataSize is the data size of scales. When the quantization mode is pertoken-pergroup quantization, the formula is: scalesDataSize = b * s * H / 128 * 4 (Byte). When the quantization mode is mx quantization, the formula is: scalesDataSize = b * s * H / 32 * 1 (Byte).

  • The H range supports only [1024, 8192] and must be 128-aligned.

Example

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

NOTE The sample code in this example calls some HCCL collective communication library APIs: HcclCommInitClusterInfoConfig, HcclGetCommName, and HcclCommDestroy. For details, see HCCL API (C).

  • Ascend 950PR/Ascend 950DT:

    #include <thread>
    #include <iostream>
    #include <vector>
    #include <string>
    #include <cstring>
    #include "hccl/hccl.h"
    #include "aclnnop/aclnn_quant_all_reduce.h"
    using namespace std;
    
    #define CHECK_RET(cond, return_expr) \
        do {                             \
            if (!(cond)) {               \
                return_expr;             \
            }                            \
        } while (0)
    
    #define LOG_PRINT(message, ...)         \
        do {                                \
            printf(message, ##__VA_ARGS__); \
        } while (0)
    
    constexpr int DEV_NUM = 2; // Number of devices.
    
    int64_t GetShapeSize(const std::vector<int64_t> &shape)
    {
        int64_t shape_size = 1;
        for (auto i : shape) {
            shape_size *= i;
        }
        return shape_size;
    }
    
    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);
        auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMalloc failed. ret: %d\n", ret);
                return ret);
        ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMemcpy failed. ret: %d\n", ret);
                return ret);
        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];
        }
        *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 LaunchOneThreadQuantAllReduce(Args &args)
    {
        int ret = aclrtSetCurrentContext(args.context);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSetCurrentContext failed. ret = %d\n", ret);
                return ret);
        char hcomName[128] = {0};
        ret = HcclGetCommName(args.hcclComm, hcomName);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclGetCommName failed. ret = %d\n", ret);
                return -1);
        LOG_PRINT("[INFO] rank = %d, hcomName = %s, stream = %p\n", args.rankId, hcomName, args.stream);
        std::vector<int64_t> xShape = {1024, 5120}; // (bs, H)
        std::vector<int64_t> scalesShape = {1024, 80, 2}; // (bs, H/64, 2)
        std::vector<int64_t> outputShape = {1024, 5120}; // (bs, H)
        void *xDeviceAddr = nullptr;
        void *scalesDeviceAddr = nullptr;
        void *outputDeviceAddr = nullptr;
        void *workspaceAddr = nullptr;
    
        aclTensor *x = nullptr;
        aclTensor *scales = nullptr;
        aclTensor *output = nullptr;
        uint64_t workspaceSize = 0;
        aclOpExecutor *executor = nullptr;
    
        long long xShapeSize = GetShapeSize(xShape);
        long long scalesShapeSize = GetShapeSize(scalesShape);
        long long outputShapeSize = GetShapeSize(outputShape);
    
        std::vector<int8_t> xHostData(xShapeSize, 0);
        std::vector<int8_t> scalesHostData(scalesShapeSize, 0);
        std::vector<int16_t> outputHostData(outputShapeSize, 0);
    
        // Create tensor.
        ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT8_E5M2, &x);
        CHECK_RET(ret == ACL_SUCCESS, return ret);
        ret = CreateAclTensor(scalesHostData, scalesShape, &scalesDeviceAddr, aclDataType::ACL_FLOAT8_E8M0, &scales);
        CHECK_RET(ret == ACL_SUCCESS, return ret);
        ret = CreateAclTensor(outputHostData, outputShape, &outputDeviceAddr, aclDataType::ACL_FLOAT16, &output);
        CHECK_RET(ret == ACL_SUCCESS, return ret);
    
        // Call the first-phase API.
        ret = aclnnQuantAllReduceGetWorkspaceSize(
            x, scales, hcomName, "sum", output, &workspaceSize, &executor);
        CHECK_RET(ret == ACL_SUCCESS,
            LOG_PRINT("[ERROR] aclnnQuantAllReduceGetWorkspaceSize failed. ret = %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("[ERROR] aclrtMalloc workspace failed. ret = %d \n", ret);
                    return ret);
        }
        // Call the second-phase API.
        ret = aclnnQuantAllReduce(workspaceAddr, workspaceSize, executor, args.stream);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnQuantAllReduce failed. ret = %d \n", ret); return ret);
        // (Boilerplate) Synchronously wait for the task execution to complete.
        ret = aclrtSynchronizeStreamWithTimeout(args.stream, 10000);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSynchronizeStreamWithTimeout failed. ret = %d \n", ret);
                return ret);
        LOG_PRINT("[INFO] device_%d aclnnQuantAllReduce execute successfully.\n", args.rankId);
        // Release device resources. Modify based on the specific API definition.
        if (x != nullptr) {
            aclDestroyTensor(x);
        }
        if (scales != nullptr) {
            aclDestroyTensor(scales);
        }
        if (output != nullptr) {
            aclDestroyTensor(output);
        }
    
        if (xDeviceAddr != nullptr) {
            aclrtFree(xDeviceAddr);
        }
        if (scalesDeviceAddr != nullptr) {
            aclrtFree(scalesDeviceAddr);
        }
        if (outputDeviceAddr != nullptr) {
            aclrtFree(outputDeviceAddr);
        }
        if (workspaceSize > 0) {
            aclrtFree(workspaceAddr);
        }
        ret = HcclCommDestroy(args.hcclComm);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclCommDestroy failed. ret = %d \n", ret); return ret);
        ret = aclrtDestroyStream(args.stream);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtDestroyStream failed. ret = %d \n", ret); return ret);
        ret = aclrtResetDevice(args.rankId);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtResetDevice failed. ret = %d \n", ret); return ret);
        ret = aclrtDestroyContext(args.context);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtDestroyContext failed. ret = %d \n", ret); return ret);
        return 0;
    }
    
    int main(int argc, char *argv[])
    {
        int ret = aclInit(nullptr);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclInit failed. ret = %d \n", ret);
                return ret);
        aclrtStream stream[DEV_NUM];
        aclrtContext context[DEV_NUM];
        for (uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
            ret = aclrtSetDevice(rankId);
            CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSetDevice failed. ret = %d \n", ret); return ret);
            ret = aclrtCreateContext(&context[rankId], rankId);
            CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateContext failed. ret = %d \n", ret); return ret);
            ret = aclrtCreateStream(&stream[rankId]);
            CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateStream failed. ret = %d \n", ret); return ret);
        }
        int32_t devices[DEV_NUM];
        for (int i = 0; i < DEV_NUM; i++) {
            devices[i] = i;
        }
        // Initialize the collective communication domain.
        HcclComm comms[DEV_NUM];
        ret = HcclCommInitAll(DEV_NUM, devices, comms);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclCommInitAll failed. ret = %d \n", ret); return ret);
    
        Args args[DEV_NUM];
        // Start multithreading.
        std::vector<std::unique_ptr<std::thread>> threads(DEV_NUM);
        for (uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
            args[rankId].rankId = rankId;
            args[rankId].hcclComm = comms[rankId];
            args[rankId].context = context[rankId];
            args[rankId].stream = stream[rankId];
            threads[rankId].reset(new(std::nothrow) std::thread(&LaunchOneThreadQuantAllReduce, std::ref(args[rankId])));
        }
        for (uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
            threads[rankId]->join();
        }
        aclFinalize();
        return 0;
    }