Documentation
Download
文档中心
aclnnGemmaRmsNorm

aclnnGemmaRmsNorm

📄 View Source Code

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×

Function

  • Description: The GemmaRmsNorm operator is a normalization operation commonly used in large models. Compared with the RmsNorm operator, it performs a +1 operation on gamma during computation.

  • Formula:

GemmaRmsNorm⁡(xi)=1Rms⁡(x)∗xi∗(1+gammai), where Rms⁡(x)=1n∑i=1nxi2+epsilon\operatorname{GemmaRmsNorm}(x_i)=\frac{1}{\operatorname{Rms}(\mathbf{x})} * x_i * (1 + gamma_i), \quad \text { where } \operatorname{Rms}(\mathbf{x})=\sqrt{\frac{1}{n} \sum_{i=1}^n x_i^2+epsilon}

Function Prototype

Each operator uses a two-phase API. You must first call the aclnnGemmaRmsNormGetWorkspaceSize API to obtain the input parameters and calculate the required workspace size based on the computation process, and then call the aclnnGemmaRmsNorm API to perform the computation.

aclnnStatus aclnnGemmaRmsNormGetWorkspaceSize(
  const aclTensor *x,
  const aclTensor *gamma,
  double           epsilon,
  const aclTensor *yOut,
  const aclTensor *rstdOut,
  uint64_t        *workspaceSize,
  aclOpExecutor  **executor)
aclnnStatus aclnnGemmaRmsNorm(
  void          *workspace,
  uint64_t       workspaceSize,
  aclOpExecutor *executor,
  aclrtStream    stream)

aclnnGemmaRmsNormGetWorkspaceSize

  • Parameters

    Parameter Input/Output Description Instruction Data Type Data Format Dimension (Shape) Non-Contiguous Tensor
    x (aclTensor*) Input Indicates the input tensor to be normalized, corresponding to `x` in the formula. Supports null tensor. FLOAT32, FLOAT16, BFLOAT16 ND 1-8 √
    gamma (aclTensor*) Input Represents the data scaling factor, corresponding to `gamma` in the formula.
    • Supports null tensor.
    • The data type must be consistent with the data type of `x`.
    • The shape must satisfy gamma_shape = x_shape[n:], where n < x_shape.dims().
    		 FLOAT32, FLOAT16, BFLOAT16 ND 1-8 √
    epsilon (double) Input Indicates the value added to the variance to avoid division by zero. Corresponds to `epsilon` in the formula. The recommended value is 1e-6. - - - -
    yOut (aclTensor*) Output Represents the normalized output data, corresponding to `GemmaRmsNorm(x)` in the formula.
    • Supports null Tensor.
    • The shape and data type are consistent with those of the input `x`.
    		 FLOAT32, FLOAT16, BFLOAT16 ND 1-8 √
    rstdOut (aclTensor*) Output Represents the reciprocal of the standard deviation of x, used for normalization. Corresponds to the reciprocal of `Rms(x)` in the formula.
    • Supports null tensor.
    • The shape is consistent with the first several dimensions of the input `x`, where the first several dimensions refer to the dimensions of `x` minus the dimensions of `gamma`, indicating the dimensions that do not require normalization.
    		 FLOAT32 ND 1-8 √
    workspaceSize (uint64_t*) Output Returns the workspace size that the user needs to apply for 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 Code Error Code Description
    ACLNN_ERR_PARAM_NULLPTR 161001 If the input parameter is a mandatory input, output, or mandatory attribute and is a null pointer, returns 161001.
    ACLNN_ERR_PARAM_INVALID 161002 The data type of the input or output is not within the supported range.
    ACLNN_ERR_INNER_TILING_ERROR 561002 The parameter does not meet the requirements in Parameters.

aclnnGemmaRmsNorm

  • 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, obtained from the first-phase API aclnnGemmaRmsNormGetWorkspaceSize.
    executor Input Operator executor that contains the operator computation process.
    stream Input Specifies the stream for task execution.

  • Return Value

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

Constraints

  • Functional dimensions:

    • Data type support:

  • x, gamma, and yOut support FLOAT32, FLOAT16, and BFLOAT16.

  • rstdOut supports FLOAT32.

  • Data format support: ND.

  • Unsupported type description:

    DOUBLE: The instruction does not support DOUBLE.

  • Boundary value scenario description:

    • When the input is Inf, the output is Inf.

    • When the input is NaN, the output is NaN.

  • Supported data types by platform:

    x Data Typegamma Data TypeyOut Data TyperstdOut Data Type
    FLOAT16FLOAT16FLOAT16FLOAT32
    BFLOAT16BFLOAT16BFLOAT16FLOAT32
    FLOAT32FLOAT32FLOAT32FLOAT32

  • Deterministic computation:

    • aclnnGemmaRmsNorm defaults to a deterministic implementation.

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 "acl/acl.h"
#include "aclnnop/aclnn_gemma_rms_norm.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)

int64_t GetShapeSize(const std::vector<int64_t>& shape)
{
    int64_t shape_size = 1;
    for (auto i : shape) {
        shape_size *= i;
    }
    return shape_size;
}

int Init(int32_t deviceId, aclrtStream* stream)
{
    // Boilerplate: initialize resources.
    auto ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
    ret = aclrtSetDevice(deviceId);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
    ret = aclrtCreateStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);

    return 0;
}

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 the host-side data to the 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);

    // Calculate the strides of the 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 aclCreateTensor.
    *tensor = aclCreateTensor(
        shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(),
        *deviceAddr);
    return 0;
}

int main()
{
    // 1. (Boilerplate) Initialize the device/stream. See the ACL API manual.
    // Set deviceId based on the actual device.
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = Init(deviceId, &stream);
    // Handle the check as needed.
    CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
    // 2. Construct the input and output. Customize the construction based on the API interface.
    std::vector<int64_t> xShape = {2, 16};
    std::vector<int64_t> gammaShape = {16};
    std::vector<int64_t> yShape = {2, 16};
    std::vector<int64_t> rstdShape = {2, 1};
    void* xDeviceAddr = nullptr;
    void* gammaDeviceAddr = nullptr;
    void* yDeviceAddr = nullptr;
    void* rstdDeviceAddr = nullptr;
    aclTensor* x = nullptr;
    aclTensor* gamma = nullptr;
    aclTensor* y = nullptr;
    aclTensor* rstd = nullptr;
    std::vector<float> xHostData = {0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7,
                                    0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7};
    std::vector<float> gammaHostData = {0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7};
    std::vector<float> yHostData = {0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7,
                                    0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7};
    std::vector<float> rstdHostData = {1, 2};
    float epsilon = 1e-6;

    // Create the x aclTensor.
    ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT, &x);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the gamma aclTensor.
    ret = CreateAclTensor(gammaHostData, gammaShape, &gammaDeviceAddr, aclDataType::ACL_FLOAT, &gamma);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the y aclTensor.
    ret = CreateAclTensor(yHostData, yShape, &yDeviceAddr, aclDataType::ACL_FLOAT, &y);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the rstd aclTensor.
    ret = CreateAclTensor(rstdHostData, rstdShape, &rstdDeviceAddr, aclDataType::ACL_FLOAT, &rstd);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    // 3. Call the CANN operator library API. Replace it with the specific API.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;
    // Call the first-phase API of aclnnGemmaRmsNorm.
    ret = aclnnGemmaRmsNormGetWorkspaceSize(x, gamma, epsilon, y, rstd, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnGemmaRmsNormGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
    // Apply for device memory based on the workspaceSize calculated by the first-phase API.
    void* workspaceAddr = nullptr;
    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 of aclnnGemmaRmsNorm.
    ret = aclnnGemmaRmsNorm(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnGemmaRmsNorm failed. ERROR: %d\n", ret); return ret);
    // 4. (Boilerplate) Synchronously wait for task execution to finish.
    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
    // 5. Obtain the output value and copy the result from the device side memory to the host side. Modify based on the specific API definition.
    auto size = GetShapeSize(yShape);
    std::vector<float> resultData(size, 0);
    ret = aclrtMemcpy(
        resultData.data(), resultData.size() * sizeof(resultData[0]), yDeviceAddr, size * sizeof(float),
        ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
    for (int64_t i = 0; i < size; i++) {
        LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]);
    }

    // 6. Release aclTensor and aclScalar. Modify based on the specific API definition.
    aclDestroyTensor(x);
    aclDestroyTensor(gamma);
    aclDestroyTensor(y);
    aclDestroyTensor(rstd);

    // 7. Release device resources. Modify based on the specific API definition.
    aclrtFree(xDeviceAddr);
    aclrtFree(gammaDeviceAddr);
    aclrtFree(yDeviceAddr);
    aclrtFree(rstdDeviceAddr);
    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    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