aclnnMoeTokenUnpermuteWithRoutingMapGrad

📄 View Source Code

Applicable Products

Function

• Description: Backward propagation of aclnnMoeTokenUnpermuteWithRoutingMap.

• Formula:

  (1) probs is not None:

  $$permutedTokensGrad[outIndex[i]] = unpermutedTokensGrad[permuteTokenId[i]]$$ $$permutedProbsGrad = permutedTokensGrad * permutedTokensOptional$$ $$probsGradExpertOrder = \sum_{j=0}^{hidden\_size}(permutedProbsGrad_{i,j})$$
  • When dropAndPad is false:
  $$probsGradOut = masked\_scatter(routingMapOptional^T,probsGradExpertOrder)$$ $$permutedProbs = probsOptional^T.masked\_select(routingMapOptional^T)$$ $$permutedTokensGradOut = permutedProbs.unsqueeze(-1) * permutedTokensGrad$$
  • When dropAndPad is true:
  $$probsGradOut[permuteTokenId[i], outIndex[i]/capacity] = probsGradExpertOrder[outIndex[i]]$$ $$permutedProbs[outIndex[i]] = probsOptional.view(1)[i]$$ $$permutedTokensGradOut = permutedProbs * permutedTokensGrad$$

  (2) probs is None:

  $$permutedTokensGradOut[outIndex[i]] = unpermutedTokensGrad[permuteTokenId[i]]$$

  1. hidden_size refers to the size of the first dimension of unpermutedTokensGrad.

  2. When dropAndPad is true, each expert can process a fixed number of capacity tokens. The first dimension of the input routingMapOptional is experts_num, i.e., the number of experts, and the 0th dimension of the input outIndex is experts_num * capacity. capacity can be derived from these two dimensions.

  3. When dropAndPad is false, each token is processed by a fixed number of topK_num experts. The 0th dimension of the input unpermutedTokensGrad is tokens_num, i.e., the number of tokens, and the 0th dimension of the input outIndex is tokens_num * capacity. topK_num can be derived from these two dimensions.

Function Prototype

Each operator is divided into a two-phase API. You must first call the "aclnnMoeTokenUnpermuteWithRoutingMapGradGetWorkspaceSize" API to obtain the workspace size required for computation and the executor that contains the operator computation flow, and then call the "aclnnMoeTokenUnpermuteWithRoutingMapGrad" API to perform the computation.

aclnnStatus aclnnMoeTokenUnpermuteWithRoutingMapGradGetWorkspaceSize(
    const aclTensor*   unpermutedTokensGrad,
    const aclTensor*   outIndex,
    const aclTensor*   permuteTokenId,
    const aclTensor*   routingMapOptional,
    const aclTensor*   permutedTokensOptional,
    const aclTensor*   probsOptional,
    bool               dropAndPad,
    const aclIntArray* restoreShapeOptional,
    const aclTensor*   permutedTokensGradOut,
    const aclTensor*   probsGradOutOptional,
    uint64_t*          workspaceSize,
    aclOpExecutor**    executor)

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

aclnnMoeTokenUnpermuteWithRoutingMapGradGetWorkspaceSize

• Parameters:

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
  unpermutedTokensGrad (aclTensor*)	Input	unpermutedTokensGrad in the formula, representing the gradient of the forward output unpermutedTokens.	-	BFLOAT16, FLOAT16, FLOAT	ND	(tokens_num, hidden_size).	√
  outIndex (aclTensor*)	Input	outIndex in the formula, representing the output position index.	When dropAndPad is false, the value range is [0, tokens_num * topK_num - 1]. When dropAndPad is true, the value range is [0, experts_num * capacity - 1].	INT32	ND	When dropAndPad is false, the shape is (tokens_num * topK_num). When dropAndPad is true, the shape is (experts_num * capacity).	√
  permuteTokenId (aclTensor*)	Input	permuteTokenId in the formula, representing the token index corresponding to each position of the input permutedTokens.	Value range: [0, tokens_num-1].	INT32	ND	Same as outIndex.	√
  routingMapOptional (aclTensor*)	Optional input	When the input probsOptional is a null pointer, this input is not required and a null pointer should be passed. routingMapOptional in the formula indicates whether the token at the corresponding position is processed by the corresponding expert.	When the data type is INT8, the value supports 0 and 1. When the data type is BOOL, the value supports true and false.	INT8, BOOL	ND	(tokens_num, experts_num).	√
  permutedTokensOptional (aclTensor*)	Optional input.	When the input probsOptional is a null pointer, this input is not required and a null pointer should be passed in.	The data type is the same as that of unpermutedTokensGrad.	BFLOAT16, FLOAT16, FLOAT	ND	When dropAndPad is false, the shape is (tokens_num*topK_num, hidden_size). When dropAndPad is true, the shape is (experts_num*capacity, hidden_size).	√
  probsOptional (aclTensor*)	Optional Input	Null pointer when not needed.	The data type is the same as unpermutedTokensGrad, or when unpermutedTokensGrad is BFLOAT16, probsOptional supports FLOAT.	BFLOAT16, FLOAT16, FLOAT	ND	Same as routingMapOptional.	√
  dropAndPad (bool)	Attribute	true indicates that dropAndPad is enabled, and false indicates that dropAndPad is disabled.	-	BOOL	-	-	-
  restoreShapeOptional (aclIntArray*)	Attribute	An aclIntArray of the INT64 type. When dropAndPad is true, represents the shape of unpermutedTokensGrad.	-	INT64	-	-	-
  permutedTokensGradOut (aclTensor*)	Output	permutedTokensGradOut in the formula represents the gradient of the input permutedTokens.	The data type is the same as unpermutedTokensGrad.	BFLOAT16, FLOAT16, FLOAT	ND	When dropAndPad is false, the shape is (tokens_num*topK_num, hidden_size). When dropAndPad is true, the shape is (experts_num*capacity, hidden_size).	×
  probsGradOutOptional (aclTensor*)	Optional output	Null pointer when probsOptional is not input. Inputs the gradient of probs.	The data type is the same as probsOptional.	BFLOAT16, FLOAT16, FLOAT	ND	Same as routingMapOptional.	×
  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 Value	Return Code	Description
  ACLNN_ERR_PARAM_NULLPTR	161001	The passed mandatory input, mandatory output, or mandatory attribute is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The data type and data format of the input and output are not within the supported range.
  ACLNN_ERR_INNER_TILING_ERROR	561002	When the input probsOptional is non-empty and dropAndPad is false, topK_num > 512.
  		When the input probsOptional is non-empty and dropAndPad is false, topK_num is greater than experts_num.
  		When the input probsOptional is non-empty and dropAndPad is false, (ubSize - (probTypeLen + 1) * numExpertAlign - (tokenTypeLen + 8) * 256) / (6 * tokenTypeLen + 12) < 1.
  		When the input probsOptional is non-empty and dropAndPad is true, capacity is greater than tokens_num.
  		When the input probsOptional is non-empty and dropAndPad is true, hidden_size > 256 * (ubSize - 2080) / (8 + tokenTypeLen).
  		When the input probsOptional is non-empty, the input routingMapOptional or permutedTokensOptional is empty.
  		When the input probsOptional is non-empty, the data type of probsOptional differs from that of unpermutedTokensGrad, and unpermutedTokensGrad is not BFLOAT16.
  		The shape of the input or output does not meet the requirements.

aclnnMoeTokenUnpermuteWithRoutingMapGrad

• Parameters:

  Parameter	Input/Output	Description
  workspace	Input	Address of the workspace memory applied on the Device side.
  workspaceSize	Input	Size of the workspace applied on the Device side, obtained from the first-phase API aclnnMoeTokenUnpermuteWithRoutingMapGradGetWorkspaceSize.
  executor	Input	Operator executor, which contains the operator computation flow.
  stream	Input	Specifies the stream on which the task is executed.

• Return value

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

Constraints

• Deterministic computation:

  • aclnnMoeTokenUnpermuteWithRoutingMapGrad defaults to a deterministic implementation.

• tokens_num indicates the number of input tokens, hidden_size indicates the word vector dimension, and experts_num indicates the number of experts.

• The following two modes are distinguished by dropAndPad: when dropAndPad is true, each expert can process a fixed number of capacity tokens. When dropAndPad is false, each token is processed by a fixed number of topK_num experts.

• When the input probsOptional is non-empty and dropAndPad is false:

  • topK_num must be ≤ 512 and topK_num ≤ experts_num.

  • experts_num must satisfy (ubSize - (probTypeLen + 1) numExpertAlign - (tokenTypeLen + 8) 256) / (6 * tokenTypeLen + 12) >= 1, where ubSize is the size of the chip's UB space, probTypeLen is the number of bytes corresponding to the data type of the input probsOptional, tokenTypeLen is the number of bytes corresponding to the data type of the input unpermutedTokensGrad, and numExpertAlign is the result of aligning experts_num upward to 32.

• When the input probsOptional is non-empty and dropAndPad is true:

  • capacity must be ≤ tokens_num.

  • Requires hidden_size <= 256 * (ubSize - 2080) / (8 + tokenTypeLen), where ubSize is the size of the chip's UB space, and tokenTypeLen is the number of bytes of the data type of the input unpermutedTokensGrad.

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_moe_token_unpermute_with_routing_map_grad.h"
#include <iostream>

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

void PrintOutResult(std::vector<int64_t> &shape, void **deviceAddr) {
  auto size = GetShapeSize(shape);
  std::vector<float> resultData(size, 0);
  auto ret = aclrtMemcpy(
      resultData.data(), resultData.size() * sizeof(resultData[0]), *deviceAddr,
      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 );
  for (int64_t i = 0; i < size; i++) {
    LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]);
  }
}

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 the 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 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;
}

int main() {
  // 1. (boilerplate) Initialize device/stream. See the ACL API manual.
  // Fill in 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. Customize the construction based on the API.
  bool dropAndPad = false;
  int32_t tokenNum = 1;
  int32_t hiddenSize = 2;
  int32_t expertNum = 2;
  int32_t topK = 2;
  int32_t outTokenNum = tokenNum * topK;
  std::vector<int64_t> permutedTokensShape = {outTokenNum, hiddenSize};
  std::vector<int64_t> unpermutedTokensGradShape = {tokenNum, hiddenSize};
  std::vector<int64_t> probsShape = {tokenNum, expertNum};
  std::vector<int64_t> outIndexShape = {outTokenNum};
  std::vector<int64_t> permuteTokenIdShape = {outTokenNum};
  std::vector<int64_t> routingMapShape = {tokenNum, expertNum};
  std::vector<int64_t> permutedTokensGradShape = {outTokenNum, hiddenSize};
  std::vector<int64_t> probsGradShape = {tokenNum, expertNum};
  void* permutedTokensDeviceAddr = nullptr;
  void* unpermutedTokensGradDeviceAddr = nullptr;
  void* probsDeviceAddr = nullptr;
  void* outIndexDeviceAddr = nullptr;
  void* permuteTokenIdDeviceAddr = nullptr;
  void* routingMapDeviceAddr = nullptr;
  void* permutedTokensGradDeviceAddr = nullptr;
  void* probsGradDeviceAddr = nullptr;

  aclTensor* permutedTokens = nullptr;
  aclTensor* unpermutedTokensGrad = nullptr;
  aclTensor* probs = nullptr;
  aclTensor* outIndex = nullptr;
  aclTensor* permuteTokenId = nullptr;
  aclTensor* routingMap = nullptr;
  aclTensor *permutedTokensGrad = nullptr;
  aclTensor *probsGrad = nullptr;

  std::vector<float> permutedTokensHostData = {1, 1, 1, 1};
  std::vector<float> unpermutedTokensGradHostData = {1, 1};
  std::vector<float> probsHostData = {1, 1};
  std::vector<int> outIndexHostData = {0, 1};
  std::vector<int> permuteTokenIdHostData = {0, 0};
  std::vector<int8_t> routingMapHostData = {1, 1};
  std::vector<float> permutedTokensGradHostData = {0, 0, 0, 0};
  std::vector<float> probsGradHostData = {0, 0};

  ret = CreateAclTensor(unpermutedTokensGradHostData, unpermutedTokensGradShape, &unpermutedTokensGradDeviceAddr, aclDataType::ACL_FLOAT, &unpermutedTokensGrad);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(outIndexHostData, outIndexShape, &outIndexDeviceAddr, aclDataType::ACL_INT32, &outIndex);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(permuteTokenIdHostData, permuteTokenIdShape, &permuteTokenIdDeviceAddr, aclDataType::ACL_INT32, &permuteTokenId);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(routingMapHostData, routingMapShape, &routingMapDeviceAddr, aclDataType::ACL_BOOL, &routingMap);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(permutedTokensHostData, permutedTokensShape, &permutedTokensDeviceAddr, aclDataType::ACL_FLOAT, &permutedTokens);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(probsHostData, probsShape, &probsDeviceAddr, aclDataType::ACL_FLOAT, &probs);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(permutedTokensGradHostData, permutedTokensGradShape, &permutedTokensGradDeviceAddr, aclDataType::ACL_FLOAT, &permutedTokensGrad);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(probsGradHostData, probsGradShape, &probsGradDeviceAddr, aclDataType::ACL_FLOAT, &probsGrad);
  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 aclnnMoeTokenUnpermuteWithRoutingMapGrad.
  ret = aclnnMoeTokenUnpermuteWithRoutingMapGradGetWorkspaceSize(unpermutedTokensGrad, outIndex, permuteTokenId, routingMap, permutedTokens, probs, dropAndPad, nullptr, permutedTokensGrad, probsGrad, &workspaceSize, &executor);
  CHECK_RET(
      ret == ACL_SUCCESS,
      LOG_PRINT("aclnnMoeTokenUnpermuteWithRoutingMapGradGetWorkspaceSize 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 aclnnMoeTokenUnpermuteWithRoutingMapGrad.
  ret = aclnnMoeTokenUnpermuteWithRoutingMapGrad(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS,
            LOG_PRINT("aclnnMoeTokenUnpermuteWithRoutingMapGrad 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 values and copy the results from the device memory to the host memory. Modify this based on the specific API definition.
  LOG_PRINT("permutedTokensGrad \n");
  PrintOutResult(permutedTokensGradShape, &permutedTokensGradDeviceAddr);
  LOG_PRINT("probsGrad \n");
  PrintOutResult(probsGradShape, &probsGradDeviceAddr);

  // 6. Release aclTensor and aclScalar. Modify based on the specific API definition.
  aclDestroyTensor(permutedTokens);
  aclDestroyTensor(unpermutedTokensGrad);
  aclDestroyTensor(outIndex);
  aclDestroyTensor(permuteTokenId);
  aclDestroyTensor(routingMap);
  aclDestroyTensor(probs);
  aclDestroyTensor(permutedTokensGrad);
  aclDestroyTensor(probsGrad);

  // 7. Release device resources.
  aclrtFree(permutedTokensDeviceAddr);
  aclrtFree(unpermutedTokensGradDeviceAddr);
  aclrtFree(probsDeviceAddr);
  aclrtFree(outIndexDeviceAddr);
  aclrtFree(permuteTokenIdDeviceAddr);
  aclrtFree(routingMapDeviceAddr);
  aclrtFree(permutedTokensGradDeviceAddr);
  aclrtFree(probsGradDeviceAddr);

  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();

  return 0;
}