Documentation
Download
文档中心
概述

aclnnApplyAdamWQuant

📄 View Source Code

Applicable Products

ProductSupported
Ascend 950PR/Ascend 950DTNo
Atlas A3 training products/Atlas A3 inference productsYes
Atlas A2 training products/Atlas A2 inference productsYes
Atlas 200I/500 A2 inference productsNo
Atlas inference productsNo
Atlas training productsNo

Function

  • Description: Uses the optimizer inputs m and v as indices to look up values in their respective qmaps, multiplies them by the absmax corresponding to each blockSize for dequantization, and then implements the AdamW optimizer functionality. For the updated m and v, one maximum value is taken per blockSize, with each blockSize of m and v corresponding to one absmax. A norm normalization is performed, and binary search is used to find the indices in the qmaps corresponding to m and v as outputs. The absmax also serves as the input for the next round of quantization.

  • Optimizer Formula:

    mt=β1mt−1+(1−β1)gtm_{t}=\beta_{1} m_{t-1}+\left(1-\beta_{1}\right) g_{t} \\ vt=β2vt−1+(1−β2)gt2v_{t}=\beta_{2} v_{t-1}+\left(1-\beta_{2}\right) g_{t}^{2} m^t=mt1−β1t\hat{m}_{t}=\frac{m_{t}}{1-\beta_{1}^{t}} \\ v^t=vt1−β2t\hat{v}_{t}=\frac{v_{t}}{1-\beta_{2}^{t}} \\ θt+1=θt−ηv^t+ϵm^t−η⋅λ⋅θt−1\theta_{t+1}=\theta_{t}-\frac{\eta}{\sqrt{\hat{v}_{t}}+\epsilon} \hat{m}_{t}-\eta \cdot \lambda \cdot \theta_{t-1}

Function Prototype

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

aclnnStatus aclnnApplyAdamWQuantGetWorkspaceSize(
    aclTensor       *varRef,
    const aclTensor *grad,
    aclTensor       *mRef,
    aclTensor       *vRef,
    const aclTensor *qmapM,
    const aclTensor *qmapV,
    aclTensor       *absmaxMRef,
    aclTensor       *absmaxVRef,
    const aclTensor *step,
    double          lr,
    double          beta1,
    double          beta2,
    double          weightDecay,
    double          eps,
    double          gnormScale,
    char *          quantModeOptional,
    int64_t         blockSize,
    uint64_t        *workspaceSize,
    aclOpExecutor  **executor)
aclnnStatus aclnnApplyAdamWQuant(
    void            *workspace,
    uint64_t         workspaceSize,
    aclOpExecutor   *executor,
    aclrtStream      stream)

aclnnApplyAdamWQuantGetWorkspaceSize

  • Parameters:

    Parameter Input/Output Description Instruction Data Type Data Format Dimension (Shape) Non-contiguous Tensor
    varRef (aclTensor*) Input/Output Specifies the weight input to be computed, which also serves as the output, corresponding to θ in the formula. FLOAT16, BFLOAT16, FLOAT32 ND 1-8 √
    grad (aclTensor*) Input Input gradient, g in the formula. The shape must be consistent with the input varRef. Same as varRef. ND 1-8 √
    mRef (aclTensor*) Input/Output Index value of the m parameter in the AdamW optimizer formula before quantization. The specific value in qmapM is derived based on this index. The shape must be consistent with the input varRef. uint8_t ND 1-8 √
    vRef (aclTensor*) Input/Output Index value of the v parameter before quantization in the AdamW optimizer formula. The specific value in qmapV is derived based on this index. The shape must be consistent with the input varRef. uint8_t ND 1-8 √
    qmapM (aclTensor*) Input. Quantization mapping table in ascending order. Each m is selected based on the index value in mRef. The shape must be [256,]. FLOAT32 ND 1 √
    qmapV (aclTensor*) Input The quantization mapping table is sorted in ascending order, and each v is selected based on the index value in vRef. The shape requirement is [256,]. FLOAT32 ND 1 √
    absmaxMRef (aclTensor*) Input/Output Every blockSize (256) vRef elements correspond to one maximum value, which is used to select a value from qmapM for the mRef index and multiply it by the corresponding absmaxMRef for dequantization. Then, a maximum value is selected from every blockSize (256) updated mRef elements as the output of absmaxMRef. The shape requirement is "absmaxMRef.size = mRef.size/blockSize". FLOAT32 ND 1 √
    absmaxVRef (aclTensor*) Input/Output Every blockSize (256) vRef elements correspond to one maximum value. This maximum value is used to select the corresponding value from qmapV for vRef index lookup, and then multiplied by the corresponding absmaxVRef for dequantization. After the update, one maximum value is selected from every blockSize (256) updated vRef elements as the output of absmaxVRef. The shape requirement is "absmaxVRef.size = vRef.size/blockSize". FLOAT32 ND 1 √
    step (aclTensor*) Input t in the formula, representing the number of iterations, greater than 0. The shape must be [1]. INT64 ND 1 √
    lr (double) Input Learning rate, η in the formula. Recommendation: 1e-3, 1e-5, 1e-8. Range: 0 to 1. - DOUBLE - - -
    beta1 (double) Input. Beta1 parameter in the AdamW optimizer formula. The recommended value is 0.9, with a range of 0 to 1. - DOUBLE - - -
    beta2 (double) Input Specifies the beta2 parameter in the AdamW optimizer formula, with a recommended value of 0.999 and a range of 0 to 1. - DOUBLE - - -
    weightDecay (double) Input Weight decay coefficient, the λ parameter in the AdamW optimizer formula. The recommended value is 0.999, and the range is 0 to 1. - DOUBLE - - -
    eps (double) Input Epsilon parameter in the AdamW optimizer formula, added to the denominator to prevent division by zero. The recommended value is 1e-8. - DOUBLE - - -
    gnormScale (double) Input Scaling parameter for the input parameter grad. The recommended value is 0.999, and the range is 0 to 1. - DOUBLE - - -
    blockSize (int64_t) Input Size of each block participating in the computation, fixed to 256. As described in absmax above, one maximum value is selected per blockSize. - INT64 - - -
    quantModeOptional (char*) Input Reserved parameter with no current significance. - - - - -
    workspaceSize (uint64_t*) Output Returns the workspace size to be applied 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 The input and output tensors are null pointers.
    ACLNN_ERR_PARAM_INVALID 161002 The data type or data format of the input or output is not within the supported range.

aclnnApplyAdamWQuant

  • Parameters:

    Parameter Input/Output Description
    workspace Input Memory address of the workspace applied for on the Device side.
    workspaceSize Input Size of the workspace applied for on the Device side, obtained through the first-phase API aclnnApplyAdamWGetWorkspaceSize.
    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

The shape of varRef meets the following constraints:

  • varRef.shape = grad.shape

  • varRef.shape = mRef.shape

    • varRef.shape = vRef.shape

    • varRef.size/blockSize = absmaxMRef.size

    • varRef.size/blockSize = absmaxVRef.size

    Deterministic computation:

    • aclnnApplyAdamWQuant 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_apply_adam_w_quant.h"
#define FAILED 1

#define ERROR_LOG(fmt, args...) fprintf(stderr, "[ERROR] " fmt "\n", ##args)
#define INFO_LOG(fmt, args...) fprintf(stdout, "[INFO] " fmt "\n", ##args)
#define WARN_LOG(fmt, args...) fprintf(stdout, "[WARN]  " fmt "\n", ##args)

#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 for resource initialization.
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);
// Call aclrtMemcpy to copy host-side data to device-side memory.
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];
}

// Call the aclCreateTensor API to create an aclTensor.
*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.
// Fill in the deviceId based on your 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 inputs and outputs based on the API definition.
std::vector<int64_t> VarRefShape = {1,256};
std::vector<int64_t> GradShape = {1,256};
std::vector<int64_t> mRefShape = {1,256};
std::vector<int64_t> vRefShape = {1,256};
std::vector<int64_t> mapMShape = {256,};
std::vector<int64_t> mapVShape = {256,};
std::vector<int64_t> absmaxMRefShape = {1,};
std::vector<int64_t> absmaxVRefShape = {1,};
std::vector<int64_t> stepShape = {1,};

void *VarRefDeviceAddr = nullptr;
void *GradDeviceAddr = nullptr;
void *mRefDeviceAddr = nullptr;
void *vRefDeviceAddr = nullptr;
void *qmapMDeviceAddr = nullptr;
void *qmapVDeviceAddr = nullptr;
void *absmaxMRefDeviceAddr = nullptr;
void *absmaxVRefDeviceAddr = nullptr;
void *stepDeviceAddr = nullptr;

aclTensor *varRef = nullptr;
aclTensor *grad = nullptr;
aclTensor *mRef = nullptr;
aclTensor *vRef = nullptr;
aclTensor *qmapM = nullptr;
aclTensor *qmapV = nullptr;
aclTensor *absmaxMRef = nullptr;
aclTensor *absmaxVRef = nullptr;
aclTensor *step = nullptr;

std::vector<float> inputVarHostData(256);
std::vector<float> inputGradHostData(256);
std::vector<uint8_t> inputMHostData(256);
std::vector<uint8_t> inputVHostData(256);
std::vector<float> inputmapMHostData(256);
std::vector<float> inputmapVHostData(256);
std::vector<float> inputabsmaxMHostData = {5};
std::vector<float> inputabsmaxVHostData = {3};
std::vector<int64_t> inputstepHostData(1);

const float lr = 0.1;
const float beta1 = 0.1;
const float beta2 = 0.1;
const float weightDecay = 0.1;
const float eps = 0.01;
const float gnormScale = 0.1;
const int64_t blockSize = 256;
char* quantModeOptional = "BLOCKWISE";

// Create the gradOutput aclTensor.
ret = CreateAclTensor(inputVarHostData, VarRefShape, &VarRefDeviceAddr, aclDataType::ACL_FLOAT, &varRef);
CHECK_RET(ret == ACL_SUCCESS, return FAILED);

ret = CreateAclTensor(inputGradHostData, GradShape, &GradDeviceAddr, aclDataType::ACL_FLOAT, &grad);
CHECK_RET(ret == ACL_SUCCESS, return FAILED);

ret = CreateAclTensor(inputMHostData, mRefShape, &mRefDeviceAddr, aclDataType::ACL_UINT8, &mRef);
CHECK_RET(ret == ACL_SUCCESS, return FAILED);

ret = CreateAclTensor(inputVHostData, vRefShape, &vRefDeviceAddr, aclDataType::ACL_UINT8, &vRef);
CHECK_RET(ret == ACL_SUCCESS, return FAILED);

ret = CreateAclTensor(inputmapMHostData, mapMShape, &qmapMDeviceAddr, aclDataType::ACL_FLOAT, &qmapM);
CHECK_RET(ret == ACL_SUCCESS, return FAILED);

ret = CreateAclTensor(inputmapVHostData, mapVShape, &qmapVDeviceAddr, aclDataType::ACL_FLOAT, &qmapV);
CHECK_RET(ret == ACL_SUCCESS, return FAILED);

ret = CreateAclTensor(inputabsmaxMHostData, absmaxMRefShape, &absmaxMRefDeviceAddr, aclDataType::ACL_FLOAT, &absmaxMRef);
CHECK_RET(ret == ACL_SUCCESS, return FAILED);

ret = CreateAclTensor(inputabsmaxVHostData, absmaxVRefShape, &absmaxVRefDeviceAddr, aclDataType::ACL_FLOAT, &absmaxVRef);
CHECK_RET(ret == ACL_SUCCESS, return FAILED);

ret = CreateAclTensor(inputstepHostData, stepShape, &stepDeviceAddr, aclDataType::ACL_INT64, &step);
CHECK_RET(ret == ACL_SUCCESS, return FAILED);

// 3. Call the CANN operator library API. Replace this with the specific API.
uint64_t workspaceSize = 0;
aclOpExecutor* executor;
// Call the first-phase API of aclnnApplyAdamWQuantGetWorkspaceSize.
ret = aclnnApplyAdamWQuantGetWorkspaceSize(varRef, grad, mRef, vRef, qmapM, qmapV, absmaxMRef, absmaxVRef, step, lr, beta1, beta2, weightDecay, eps, gnormScale, quantModeOptional, blockSize, &workspaceSize, &executor);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnApplyAdamWQuantGetWorkspaceSize 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 aclnnApplyAdamWQuant.
ret = aclnnApplyAdamWQuant(workspaceAddr, workspaceSize, executor, stream);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnApplyAdamWQuant failed. ERROR: %d\n", ret); return ret);
// 4. Boilerplate: synchronize and wait for task execution to complete.
ret = aclrtSynchronizeStream(stream);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return FAILED);

// 5. Obtain the output value and copy the result from the device-side memory to the host side.
auto size = GetShapeSize(VarRefShape);
std::vector<float> resultData(size, 0);
ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), VarRefDeviceAddr,
                size * sizeof(resultData[0]), 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. Destroy the aclTensor. Modify this based on the specific API definition.
aclDestroyTensor(varRef);
aclDestroyTensor(grad);
aclDestroyTensor(mRef);
aclDestroyTensor(vRef);
aclDestroyTensor(qmapM);
aclDestroyTensor(qmapV);
aclDestroyTensor(absmaxMRef);
aclDestroyTensor(absmaxVRef);
aclDestroyTensor(step);

// 7. Release device resources. Modify this based on the specific API definition.
aclrtFree(VarRefDeviceAddr);
aclrtFree(GradDeviceAddr);
aclrtFree(mRefDeviceAddr);
aclrtFree(vRefDeviceAddr);
aclrtFree(qmapMDeviceAddr);
aclrtFree(qmapVDeviceAddr);
aclrtFree(absmaxMRefDeviceAddr);
aclrtFree(absmaxVRefDeviceAddr);
aclrtFree(stepDeviceAddr);
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