Documentation
Download

aclnnQuantReduceScatter

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 fused quant + reduceScatter computation.

  • Formula:

    output=Reduce(AllToAllScales∗AllToAllData)output=Reduce(AllToAllScales * AllToAllData) AllToAllData=AllToAll(x)AllToAllData=AllToAll(x) AllToAllScales=AllToAll(scales)AllToAllScales=AllToAll(scales)

    The Reduce computation here performs a reduce operation on data from different ranks.

Function Prototype

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

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

aclnnQuantReduceScatterGetWorkspaceSize

  • Parameters

    Parameter Input/Output Description Instruction Data Type Data Format Dimension (Shape) Non-contiguous Tensor
    x Input Input x in the formula.
    • Null Not Supported.
    • Supported shapes: (BS, H) or (B, S, H). B is batch size, S is sequence length, and H is hidden size. In the current version, H of input x supports any 128-aligned generalization from 1024 to 8192.
    		 INT8, HIFLOAT8, FLOAT8_E4M3FN, FLOAT8_E5M2 ND 2-3 √
    scales Input Input scales in the formula.
    • Null tensors are 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 Input Communication domain identifier. Communication domain identifier. String - - -
    reduceOp Input Type of the reduce operation in the formula. Currently only "sum" is supported. string - - -
    output Output Output in the formula.
    • Null tensor is not supported.
    • When the shape of x is (BS, H), the shape of output must be (BS/rankNum, H); when the shape of x is (B, S, H), the shape of output must be (B*S/rankNum, H). rankNum indicates the communication domain size.
    		 FLOAT, FLOAT16, BFLOAT16 ND 2 √
    workspaceSize Output Returns the Device Side-Applied workspace size. - - - - -
    executor 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 Return 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 type and shape of x and scales do not match.
    The dimensions of x and scales are not within the supported range.

aclnnQuantReduceScatter

  • Parameters

    Parameter Input/Output Description
    workspace Input Side-applied workspace memory address on the Device.
    workspaceSize Input Side-applied workspace size on the Device, obtained from the first-phase API aclnnQuantReduceScatterGetWorkspaceSize.
    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

  • 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).

  • Only enabled on Ascend 950 series platforms.

  • Null tensor input is not supported.

  • The communicator size supports 2, 4, and 8.

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

  • 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 communicator, in MB. It defaults to 200 MB when not configured. The requirement is HCCL_BUFFSIZE >= 2 * (xDataSize + scalesDataSize + 1). Here, xDataSize is the data size of the input x, with the formula: xDataSize = BS * H * 1 (Byte). scalesDataSize is the data size of scales. When the quantization mode is pertoken-pergroup quantization, the formula is: scalesDataSize = BS * H / 128 * 4 (Byte). When the quantization mode is mx quantization, the formula is: scalesDataSize = BS * H / 32 * 1 (Byte).

  • The range of H 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.

  • Ascend 950PR/Ascend 950DT:

    #include <thread>
#include <iostream>
#include <vector>
#include <string>
#include <cstring>
#include "hccl/hccl.h"
#include "aclnnop/aclnn_quant_reduce_scatter.h"

#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;

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 {
    int rankId;
    HcclComm hcclComm;
    aclrtStream stream;
    aclrtContext context;
};

int LaunchOneThreadQtReduceScatter(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};
    std::vector<int64_t> scalesShape = {1024, 40};
    std::vector<int64_t> outputShape = {1024 / DEV_NUM, 5120};
    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 the 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_FLOAT, &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 = aclnnQuantReduceScatterGetWorkspaceSize(
        x, scales, hcomName, "sum", output, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS,
        LOG_PRINT("[ERROR] aclnnQuantReduceScatterGetWorkspaceSize 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 = aclnnQuantReduceScatter(workspaceAddr, workspaceSize, executor, args.stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnQuantReduceScatter failed. ret = %d \n", ret);
            return ret);
    // (Boilerplate) Synchronously wait for 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 aclnnQuantReduceScatter 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 = aclrtDestroyStream(args.stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtDestroyStream failed. ret = %d \n", ret);
            return ret);

    ret = HcclCommDestroy(args.hcclComm);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclCommDestroy 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);

    ret = aclrtResetDevice(args.rankId);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtResetDevice 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 multiple threads.
    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(&LaunchOneThreadQtReduceScatter, std::ref(args[rankId])));
    }
    for (uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
        threads[rankId]->join();
    }
    aclFinalize();
    return 0;
}
This site uses cookies to improve your browsing experience. By continuing to browse this site, you agree to our use of cookies.View Details