aclnnFusedLinearCrossEntropyLossGrad

📄 View Source Code

Applicable Products

Function

• Description: This operator is part of the cross-entropy loss computation module in vocabulary parallelism scenarios, addressing memory and computational efficiency issues for extremely large vocabularies. The current part implements gradient computation, used to compute the gradients of the leaf nodes input and weight.

  The relevant outputs of aclnnFusedLinearOnlineMaxSum and aclnnFusedCrossEntropyLossWithMaxSum, as well as the global communication results related to logits, are required as inputs of this API.

• Formula:

  High-performance mode, softmaxOptional is not nullptr:

$$\text{softmax} \in \mathbb{R}^{BT \times V}$$ $$\text{arange\_1d} = [0, 1, \dots, BT-1] \in \mathbb{N}^{BT}$$ $$\text{softmax\_update} = \mathbf{1} - \text{target\_mask}.view(-1) \in \mathbb{R}^{BT}$$ $$\text{softmax}[\text{arange\_1d}, \text{masked\_target}] \leftarrow \text{softmax}[\text{arange\_1d}, \text{masked\_target}] - \text{softmax\_update}$$ $$\text{softmax} \leftarrow \text{softmax} \odot \text{grad}.unsqueeze(-1) \in \mathbb{R}^{BT \times V}$$ $$\text{grad\_input} = \text{softmax} \cdot \text{weight}^T \in \mathbb{R}^{BT \times H}$$ $$\text{grad\_weight} = \text{softmax}^T \cdot \text{input} \in \mathbb{R}^{V \times H}$$
  Memory-saving mode, where softmaxOptional is nullptr: $$\text{vocab\_parallel\_logits} = \text{input} \cdot \text{weight}^T \quad \in \mathbb{R}^{BT \times V}$$ $$\text{logits\_sub} = \text{vocab\_parallel\_logits} - \text{logits\_max}.unsqueeze(-1) \quad \in \mathbb{R}^{BT \times V}$$ $$\text{exp\_logits} = \exp(\text{logits\_sub}) \quad \in \mathbb{R}^{BT \times V}$$ $$\text{exp\_logits} \gets \frac{\text{exp\_logits}}{\text{sum\_exp\_logits}.unsqueeze(-1)} \quad \in \mathbb{R}^{BT \times V}$$ $$\text{grad\_logits} = \text{exp\_logits} \quad \in \mathbb{R}^{BT \times V}$$ $$\text{grad\_2d} = \text{grad\_logits}.view(-1, \text{partition\_vocab\_size}) \quad \in \mathbb{R}^{BT \times V}$$ $$\text{arange\_1d} = [0, 1, \dots, BT-1] \quad \in \mathbb{N}^{BT}$$ $$\text{softmax\_update} = 1 - \text{target\_mask}.view(-1) \quad \in \mathbb{R}^{BT}$$ $$\text{grad\_2d}[\text{arange\_1d}, \text{masked\_target\_1d}] \gets \text{grad\_2d}[\text{arange\_1d}, \text{masked\_target\_1d}] - \text{softmax\_update}$$ $$\text{grad\_logits} \gets \text{grad\_logits} \odot \text{grad}.unsqueeze(-1) \quad \in \mathbb{R}^{BT \times V}$$ $$\text{grad\_input} = \text{grad\_logits} \cdot \text{weight} \quad \in \mathbb{R}^{BT \times H}$$ $$\text{grad\_weight} = \text{grad\_logits}^T \cdot \text{input} \quad \in \mathbb{R}^{V \times H}$$

Function Prototype

Each operator is divided into a two-phase API. You must first call the aclnnFusedLinearCrossEntropyLossGradGetWorkspaceSize API to obtain the workspace size required for computation and the executor that includes the operator computation process, and then call the aclnnFusedLinearCrossEntropyLossGrad API to perform the computation.

aclnnStatus aclnnFusedLinearCrossEntropyLossGradGetWorkspaceSize(
  const aclTensor   *grad,
  const aclTensor   *input,
  const aclTensor   *weight,
  const aclTensor   *targetMask,
  const aclTensor   *maskedTarget,
  float              labelSmoothing,
  const aclTensor   *logitsMaxOptional,
  const aclTensor   *sumExpLogitsOptional,
  const aclTensor   *softmaxOptional,
  aclTensor         *inputGradOut,
  aclTensor         *weightGradOut,
  uint64_t          *workspaceSize,
  aclOpExecutor    **executor)

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

aclnnFusedLinearCrossEntropyLossGradGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
  grad	Input	Gradient of the current node, which is the input grad in the formula.	Supports empty tensor.	FLOAT32	ND	1	√
  input	Input	The input matrix for matrix multiplication, which is the input in the formula.	The first dimension length is the same as that of the input grad, and supports empty tensors.	FLOAT16, BFLOAT16	ND	2	√
  weight	Input	Weight matrix for matrix multiplication, i.e., weight in the formula.	The data type is the same as that of input. The first dimension length does not support values less than 128, and the second dimension length is the same as the second dimension length of input.	FLOAT16, BFLOAT16	ND	2	√
  targetMask	Input	Intermediate variable, indicating whether the corresponding word ID is within the target range, i.e., target_mask in the formula.	Each bit of data represents a Boolean value, where 0 indicates false and 1 indicates true. The shape length multiplied by 8 must be no less than the length of the input grad.	UINT8	ND	1	√
  maskedTarget	Input	An intermediate variable that represents the local index mapping the corresponding word ID to the vocabulary shard of the current device. Invalid targets are processed by the targetMask, corresponding to masked_target in the formula.	The shape length is the same as that of the input grad.	INT64, INT32	ND	1	√
  labelSmoothing	Input	Label smoothing coefficient, used to mitigate overfitting.	Currently, only the value 0 is supported.	FLOAT32	-	-	-
  logitsMaxOptional	Optional input	Intermediate variable, the maximum value of global logits, logits_max in the formula.	Supports input of nullptr. When nullptr is input, a valid softmaxOptional input must be provided. The shape has the same length as the input grad.	FLOAT32	ND	1	√
  sumExpLogitsOptional	Optional input	Intermediate variable, the processed logits, i.e., sum_exp_logits in the formula.	Supports input of nullptr. When nullptr is input, a valid softmaxOptional input is required. The shape length is the same as that of the input grad.	FLOAT32	ND	1	√
  softmaxOptional	Optional input.	Intermediate variable, the result of matrix multiplication, softmax in the formula.	Supports input of nullptr. When nullptr is input, valid logitsMaxOptional and sumExpLogitsOptional inputs must be provided. When non-nullptr is input, the logitsMaxOptional and sumExpLogitsOptional inputs are invalid. The first dimension length of shape is the same as the length of input grad, and the second dimension length is the same as the first dimension length of input weight.	FLOAT32	ND	2	√
  inputGradOut	Output	Gradient corresponding to the leaf node input, i.e., grad_input in the formula.	The data type is the same as that of the input. The first dimension length is the same as that of the input grad, and the second dimension length is the same as that of the second dimension of the input weight.	FLOAT16, BFLOAT16	ND	2	√
  weightGradOut	Output	Gradient of the corresponding leaf node weight, i.e., grad_weight in the formula.	The data type is the same as the input. The first dimension length of the shape is the same as the first dimension length of the input weight, and the second dimension length is the same as the input grad length.	FLOAT16, BFLOAT16	ND	2	√
  workspaceSize	Output	Returns the workspace size to be side-applied on the Device.	-	-	-	-	-
  executor	Output	Returns the op executor, which includes 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 passed non-optional input is a null pointer.
  		When the passed softmaxOptional is a null pointer, logitsMaxOptional or sumExpLogitsOptional is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The input labelSmoothing value is not supported.
  		The input data type is not within the supported range.
  		The input data format is not within the supported range.
  		The input shape is not within the supported range and does not meet the length requirement.
  ACLNN_ERR_RUNTIME_ERROR	361001	The current platform is not within the supported platform range.

aclnnFusedLinearCrossEntropyLossGrad

• 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 aclnnFusedLinearCrossEntropyLossGradGetWorkspaceSize.
  executor	Input	Operator executor that includes the operator computation process.
  stream	Input	Specifies the stream for task execution.

• Return Value

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

Constraints

• Deterministic description:

  • aclnnFusedLinearCrossEntropyLossGrad 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_fused_linear_cross_entropy_loss_grad.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 shapeSize = 1;
    for (auto i : shape) {
        shapeSize *= i;
    }
    return shapeSize;
}

int Init(int32_t deviceId, aclrtStream *stream)
{
    // Boilerplate: 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>
std::vector<T> GenZeroVector(const std::vector<int64_t>& shape) {
    // 1. Calculate the total number of elements.
    size_t total = 1;
    for (auto dim : shape) {
        total *= dim;
    }

    // 2. Fill with zeros.
    std::vector<T> vec(total);
    for (auto& elem : vec) {
        elem = 0;
    }
    return vec;
}

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 allocate 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);

    // 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];
    }

    // Call aclCreateTensor 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;
}

template <typename T>
int CreateEmptyAclTensor(const std::vector<int64_t> &shape, void **deviceAddr,
                         aclDataType dataType, aclTensor **tensor)
{
    auto size = GetShapeSize(shape) * sizeof(T);
    // Call aclrtMalloc to allocate 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);

    // 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];
    }

    // 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 list of external ACL APIs.
    // Set deviceId based on your actual device.
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = Init(deviceId, &stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

    // 2. Construct inputs and outputs. Custom construction is required based on the API.
    int64_t BT = 1024;
    int64_t V = 1024;
    int64_t H = 1024;
    std::vector<int64_t> gradShape = {BT};
    std::vector<int64_t> inputShape = {BT, H};
    std::vector<int64_t> weightShape = {V, H};
    std::vector<int64_t> targetMaskShape = {BT};
    std::vector<int64_t> maskedTargetShape = {BT};
    std::vector<int64_t> softmaxOptionalShape = {BT, V};
    std::vector<int64_t> inputGradOutShape = {BT, H};
    std::vector<int64_t> weightGradOutShape = {V, H};
    void *gradDeviceAddr = nullptr;
    void *inputDeviceAddr = nullptr;
    void *weightDeviceAddr = nullptr;
    void *targetMaskDeviceAddr = nullptr;
    void *maskedTargetDeviceAddr = nullptr;
    void *softmaxOptionalDeviceAddr = nullptr;
    void *inputGradOutDeviceAddr = nullptr;
    void *weightGradOutDeviceAddr = nullptr;
    aclTensor *grad = nullptr;
    aclTensor *input = nullptr;
    aclTensor *weight = nullptr;
    aclTensor *targetMask = nullptr;
    aclTensor *maskedTarget = nullptr;
    float labelSmoothing = 0.0;
    aclTensor *logitsMaxOptional = nullptr;
    aclTensor *sumExpLogitsOptional = nullptr;
    aclTensor *softmaxOptional = nullptr;
    aclTensor *inputGradOut = nullptr;
    aclTensor *weightGradOut = nullptr;
    // Create aclTensor.
    auto gradData = GenZeroVector<int32_t>(gradShape);
    ret = CreateAclTensor<int32_t>(gradData, gradShape, &gradDeviceAddr, aclDataType::ACL_FLOAT, &grad);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    auto inputData = GenZeroVector<int16_t>(inputShape);
    ret = CreateAclTensor<int16_t>(inputData, inputShape, &inputDeviceAddr, aclDataType::ACL_BF16, &input);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    auto weightData = GenZeroVector<int16_t>(weightShape);
    ret = CreateAclTensor<int16_t>(weightData, weightShape, &weightDeviceAddr, aclDataType::ACL_BF16, &weight);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    auto targetMaskData = GenZeroVector<int8_t>(targetMaskShape);
    ret = CreateAclTensor<int8_t>(targetMaskData, targetMaskShape, &targetMaskDeviceAddr, aclDataType::ACL_UINT8, &targetMask);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    auto maskedTargetData = GenZeroVector<int32_t>(maskedTargetShape);
    ret = CreateAclTensor<int32_t>(maskedTargetData, maskedTargetShape, &maskedTargetDeviceAddr, aclDataType::ACL_INT32, &maskedTarget);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    auto softmaxOptionalData = GenZeroVector<int32_t>(softmaxOptionalShape);
    ret = CreateAclTensor<int32_t>(softmaxOptionalData, softmaxOptionalShape, &softmaxOptionalDeviceAddr, aclDataType::ACL_FLOAT, &softmaxOptional);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    auto inputGradOutData = GenZeroVector<int16_t>(inputGradOutShape);
    ret = CreateAclTensor<int16_t>(inputGradOutData, inputGradOutShape, &inputGradOutDeviceAddr, aclDataType::ACL_BF16, &inputGradOut);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    auto weightGradOutData = GenZeroVector<int16_t>(weightGradOutShape);
    ret = CreateAclTensor<int16_t>(weightGradOutData, weightGradOutShape, &weightGradOutDeviceAddr, aclDataType::ACL_BF16, &weightGradOut);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    // 3. Call the CANN operator library API. Replace with the specific API name.
    uint64_t workspaceSize = 0;
    aclOpExecutor *executor;
    // Call the first-phase API of aclnnFusedLinearCrossEntropyLossGrad.
    ret = aclnnFusedLinearCrossEntropyLossGradGetWorkspaceSize(grad, input, weight, targetMask, maskedTarget, labelSmoothing, logitsMaxOptional, sumExpLogitsOptional, softmaxOptional, inputGradOut, weightGradOut, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFusedLinearCrossEntropyLossGradGetWorkspaceSize 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 aclnnFusedLinearCrossEntropyLossGrad.
    ret = aclnnFusedLinearCrossEntropyLossGrad(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFusedLinearCrossEntropyLossGrad failed. ERROR: %d\n", ret); return ret);

    // 4. (Boilerplate) Synchronously wait for task execution to complete.
    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 memory to the host. Modify based on the specific API definition.
    // inputGradOut
    auto size = GetShapeSize(inputGradOutShape);
    std::vector<float> resultData(size, 0);
    ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), inputGradOutDeviceAddr, 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 < 16; i++) {
        LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]);
    }

    // 6. Release aclTensor. Modify based on the specific API definition.
    aclDestroyTensor(grad);
    aclDestroyTensor(input);
    aclDestroyTensor(weight);
    aclDestroyTensor(targetMask);
    aclDestroyTensor(maskedTarget);
    aclDestroyTensor(softmaxOptional);
    aclDestroyTensor(inputGradOut);
    aclDestroyTensor(weightGradOut);

    // 7. Release device resources. Modify based on the specific API definition.
    aclrtFree(gradDeviceAddr);
    aclrtFree(inputDeviceAddr);
    aclrtFree(weightDeviceAddr);
    aclrtFree(targetMaskDeviceAddr);
    aclrtFree(maskedTargetDeviceAddr);
    aclrtFree(softmaxOptionalDeviceAddr);
    aclrtFree(inputGradOutDeviceAddr);
    aclrtFree(weightGradOutDeviceAddr);
    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();

    return 0;
}