aclnnWeightQuantMatmulAllReduce

📄 View Source Code

Applicable Products

NOTE: When using this API, ensure that the driver firmware package and the CANN package are both of the matching 8.0.RC2 version or a matching later version. Otherwise, errors such as BUS ERROR will occur.

Function

• Description: Performs pseudo-quantization computation on the input x2, and then completes Matmul and AllReduce computation. Supports pertensor, perchannel, and pergroup quantization modes.

• Formula:

  $$output = AllReduce(x1 @ ((x2 + antiquantOffset) * antiquantScale) + bias + x3)$$

Function Prototype

Each operator consists of a two-phase API. You must first call the aclnnWeightQuantMatmulAllReduceGetWorkspaceSize API to obtain the workspace size required for computation and the executor that contains the operator computation process, and then call the aclnnWeightQuantMatmulAllReduce API to perform the computation.

aclnnStatus aclnnWeightQuantMatmulAllReduceGetWorkspaceSize(
    const aclTensor  *x1,
    const aclTensor  *x2,
    const aclTensor  *bias,
    const aclTensor  *antiquantScale,
    const aclTensor  *antiquantOffset,
    const aclTensor  *x3,
    const char       *group,
    const char       *reduceOp,
    int64_t          commTurn,
    int64_t          streamMode,
    int64_t          antiquantGroupSize,
    const aclTensor *output,
    uint64_t        *workspaceSize,
    aclOpExecutor **executor)

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

aclnnWeightQuantMatmulAllReduceGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-Contiguous Tensor
  x1	Input	Left matrix of the MatMul computation, i.e., x1 in the formula.	Only two-dimensional or three-dimensional input is supported in the current version. Supports the non-transpose scenario.	BFLOAT16, FLOAT16	See Constraints.	2-3	×
  x2	Input	The right matrix of the MatMul computation, i.e., x2 in the formula.	Currently, only two-dimensional input is supported. Transpose and non-transpose scenarios are supported. In ND format, non-contiguous tensors with the last two axes transposed are supported; other non-contiguous tensors are not supported.	See Constraints.	ND, FRACTAL_NZ	2	×
  bias	Input	Corresponds to the bias offset in the formula, i.e., the bias in the formula.	Supports passing a null pointer. When non-null, only one-dimensional input is supported in the current version.	See Constraints.	ND	1	√
  antiquantScale	Input	That is, the antiquantScale in the formula.	In the pertensor scenario, the shape is (1). In the perchannel scenario, the shape is (n)/(1,n), where n is the size of the last dimension of x2. In the pergroup scenario, the shape is (ceil(k,antiquantGroupSize),n).	BFLOAT16, FLOAT16	ND	1-2	√
  antiquantOffset	Input	The offset parameter for pseudo-quantization computation on x2, i.e., antiquantOffset in the formula.	Supports passing a null pointer. When non-null, the shape is consistent with that of antiquantScale. When the data format of x2 is FLOAT8_E4M3FN or HIFLOAT8, this parameter is not supported. Pass a null pointer.	BFLOAT16, FLOAT16	ND	1-2	√
  x3	Input	Add operation after the MatMul computation, i.e., x3 in the formula.	Supports passing a null pointer. When non-null, the shape is the same as the shape after the mm computation.	See Constraints.	ND	2-3	√
  group	Input	Name of the communication domain.	Obtained through the Hccl API "extern HcclResult HcclGetCommName(HcclComm comm, char* commName);", where commName is the group.	String	-	-	-
  reduceOp	Input	Reduce operation type.	The current version supports only "sum".	String	-	-	-
  commTurn	Input	Number of communication data splits, i.e., total data volume divided by single communication volume.	Currently, only 0 is supported.	INT64	-	-	-
  streamMode	Input	Enumeration of the stream mode.	Currently, only enumeration value 1 is supported.	INT64	-	-	-
  antiquantGroupSize	Input	Input of groupSize for dequantizing x2 in pseudo-quantization per-group mode.	This parameter must be passed in per-group quantization scenarios. The value range is [32, min(k-1, INT_MAX)] and must be a multiple of 32. The value range of k is consistent with that of the mm API, which is [1, 65535]. Only 0 is supported in non-per-group quantization scenarios.	INT64	-	-	-
  output	Output	Result of the MatMul computation and AllReduce communication, i.e., output in the formula.	The dimensions of output are consistent with those of x1.	-	ND	2-3	√
  workspaceSize	Output	Returns the workspace size to be allocated on the Device side.	-	-	-	-	-
  executor	Output	Returns the operator 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:

aclnnWeightQuantMatmulAllReduce

• Parameters

  Parameter	Input/Output	Description
  workspace	Input	Specifies the workspace memory address applied for on the Device side.
  workspaceSize	Input	Specifies the workspace size applied for on the Device side, obtained through the first-phase API aclnnWeightQuantMatmulAllReduceGetWorkspaceSize.
  executor	Input	Operator executor that 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

• Deterministic computation:

  • Atlas A2 training products/Atlas A2 inference products: aclnnWeightQuantMatmulAllReduce defaults to a non-deterministic implementation. Deterministic computation can be enabled by setting the HCCL_DETERMINISTIC environment variable to true.

  • Ascend 950PR/Ascend 950DT: aclnnWeightQuantMatmulAllReduce defaults to a deterministic implementation.

• MC2 is disabled in incremental scenarios and enabled in full scenarios.

• The input x1 can be two-dimensional or three-dimensional, with a shape of (b, s, k) or (m, k).

• x2 must be two-dimensional. Its shape is (k, n), where the k-axis meets the input requirements of the mm operator, the k-axes are equal, m ranges from [1, 2147483647], and k and n range from [1, 65535].

• The passed x1, x2, antiquantScale, or output is not a null pointer.

• When the shape of input x1 is (b, s, k), the shape of x3 (non-null scenario) and output is (b, s, n); when the shape of input x1 is (m, k), the shape of x3 (non-null scenario) and output is (m, n).

• If bias is non-null, its shape size equals the size of the last dimension of output. The shape of antiquantScale is (1) in the pertensor scenario, (1, n)/(n) in the perchannel scenario, and (ceil(k, antiquantGroupSize), n) in the pergroup scenario. If antiquantOffset is non-null, its shape is the same as that of antiquantScale.

• The data types and data formats of x1, x2, x3 (Non-Null Scenario), antiquantScale, antiquantOffset (Non-Null Scenario), output, and bias (Non-Null Scenario) must fall within the supported ranges.

• x1, antiquantScale, antiquantOffset (Non-Null Scenario), x3 (Non-Null Scenario), bias (Non-Null Scenario), and output have the same data type. The value of antiquantGroupSize must fall within the valid range and be a multiple of 32.

• In the pergroup scenario, when x2 is transposed, antiquantScale and antiquantOffset must be transposed together to maintain continuity.

• In long sequence scenarios, as b/s or m increases, OOM or computation timeout may occur.

• Only HCCS all-mesh networking is supported.

  • Atlas A2 training products/Atlas A2 inference products: Supports 1, 2, 4, or 8 cards.

  • Ascend 950PR/Ascend 950DT: Supports 1, 2, 4, 8, 16, 32, or 64 cards.

• Atlas A2 training products/Atlas A2 inference products:

  • The general-purpose computation fused MC2 operators in a model support only the same communication domain.

  • The data format of input x2 supports ND (only two-dimensional input is supported in the current version) and FRACTAL_NZ format (only four-dimensional input is supported in the current version). When the data format of x2 is FRACTAL_NZ, use aclnnCalculateMatmulWeightSizeV2 and aclnnTransMatmulWeight to convert the input from ND to NZ. Non-contiguous tensors support only the transpose scenario.

• Ascend 950PR/Ascend 950DT:

  • The data format of input x2 supports ND (2D input only). In the current version, when the data type is INT8, N and K must be 32-aligned; when the data type is INT4, N and K must be 64-aligned.

• Null tensor support:

  • Only the scenario where k is 0 is supported, with the output being bias + x3. Null tensor inputs where bs, m, or n is 0 are not supported.

The input and output support the following data type combinations:

• Atlas A2 training products/Atlas A2 inference products:

  x1	x2	bias	antiquantScale	antiquantOffset	x3	output	Constraints
  BFLOAT16	INT8, INT4	null, BFLOAT16	BFLOAT16	null, BFLOAT16	null, BFLOAT16	BFLOAT16	-
  FLOAT16	INT8, INT4	null, FLOAT16	FLOAT16	null, FLOAT16	null, FLOAT16	FLOAT16	-

• Ascend 950PR/Ascend 950DT

  x1	x2	bias	antiquantScale	antiquantOffset	x3	output	Limitation
  BFLOAT16	INT8, INT4	null, BFLOAT16	BFLOAT16	null, BFLOAT16	null, BFLOAT16	BFLOAT16	Supports pertensor, perchannel, and pergroup quantization scenarios.
  BFLOAT16	FLOAT8_E4M3FN, HIFLOAT8	null, BFLOAT16	BFLOAT16	null, BFLOAT16	null, BFLOAT16	BFLOAT16	Only supports the per-channel quantization scenario.
  FLOAT16	INT8, INT4	null, FLOAT16	FLOAT16	null, FLOAT16	null, FLOAT16	FLOAT16	Supports pertensor, perchannel, and pergroup quantization scenarios.
  FLOAT16	FLOAT8_E4M3FN, HIFLOAT8	null, FLOAT16	FLOAT16	null, FLOAT16	null, FLOAT16	FLOAT16	Only supports the per-channel quantization scenario.

Example

The following provides example code for reference only. For details about compilation and running, see Compile and Run Samples.

Note: This sample code calls some HCCL collective communication library APIs: HcclGetCommName, HcclCommInitAll, and HcclCommDestroy. See HCCL API (C).

• Atlas A2 training products/Atlas A2 inference products, Ascend 950PR/Ascend 950DT:

  #include <iostream>
  #include <vector>
  #include <thread>
  #include <string.h>
  #include "hccl/hccl.h"
  #include "aclnn/opdev/fp16_t.h"
  #include "aclnnop/aclnn_weight_quant_matmul_all_reduce.h"
  
  int ndev = 2;
  
  #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;
  }
  
  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 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;
  }
  
  struct Args {
      uint32_t rankId;
      HcclComm hcclComm;
      aclrtStream stream;
      aclrtContext context;
  };
  
  int launchOneThreadweightQuantmatmulAllReduce(Args &args) {
      int ret;
      ret = aclrtSetCurrentContext(args.context);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetCurrentContext failed. ERROR: %d\n", ret); return ret);
      char hcom_name[128];
      ret = HcclGetCommName(args.hcclComm, hcom_name);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclGetCommName failed. ret = %d \n", ret); return -1);
      LOG_PRINT("[INFO] rank %d hcom: %s stream: %p, context : %p\n", args.rankId, hcom_name, args.stream,
              args.context);
  
      std::vector<int64_t> x1Shape = {32, 64};
      std::vector<int64_t> x2Shape = {64, 128};
      std::vector<int64_t> biasShape = {128};
      std::vector<int64_t> antiquantScaleShape = {128};
      std::vector<int64_t> antiquantOffsetShape = {128};
      std::vector<int64_t> x3Shape = {32, 128};
      std::vector<int64_t> outShape = {32, 128};
      void *x1DeviceAddr = nullptr;
      void *x2DeviceAddr = nullptr;
      void *biasDeviceAddr = nullptr;
      void *antiquantScaleDeviceAddr = nullptr;
      void *antiquantOffsetDeviceAddr = nullptr;
      void *x3DeviceAddr = nullptr;
      void *outDeviceAddr = nullptr;
      aclTensor *x1 = nullptr;
      aclTensor *x2 = nullptr;
      aclTensor *bias = nullptr;
      aclTensor *antiquantScale = nullptr;
      aclTensor *antiquantOffset = nullptr;
      aclTensor *x3 = nullptr;
      aclTensor *out = nullptr;
  
      int64_t commTurn = 0;
      int64_t streamMode = 1;
      int64_t antiquantGroupSize = 0;
      uint64_t workspaceSize = 0;
      aclOpExecutor *executor;
      void *workspaceAddr = nullptr;
  
      long long x1ShapeSize = GetShapeSize(x1Shape);
      long long x2ShapeSize = GetShapeSize(x2Shape);
      long long biasShapeSize = GetShapeSize(biasShape);
      long long antiquantScaleShapeSize = GetShapeSize(antiquantScaleShape);
      long long antiquantOffsetShapeSize = GetShapeSize(antiquantOffsetShape);
      long long x3ShapeSize = GetShapeSize(x3Shape);
      long long outShapeSize = GetShapeSize(outShape);
      std::vector<op::fp16_t> x1HostData(x1ShapeSize, 1);
      std::vector<int8_t> x2HostData(x2ShapeSize, 1);
      std::vector<op::fp16_t> biasHostData(biasShapeSize, 1);
      std::vector<op::fp16_t> antiquantScaleHostData(antiquantScaleShapeSize, 1);
      std::vector<op::fp16_t> antiquantOffsetHostData(antiquantOffsetShapeSize, 1);
      std::vector<op::fp16_t> x3HostData(x3ShapeSize, 1);
      std::vector<op::fp16_t> outHostData(outShapeSize, 0);
      // Create a tensor.
      ret = CreateAclTensor(x1HostData, x1Shape, &x1DeviceAddr, aclDataType::ACL_FLOAT16, &x1);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(x2HostData, x2Shape, &x2DeviceAddr, aclDataType::ACL_INT8, &x2);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(biasHostData, biasShape, &biasDeviceAddr, aclDataType::ACL_FLOAT16, &bias);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(antiquantScaleHostData, antiquantScaleShape, &antiquantScaleDeviceAddr,
                          aclDataType::ACL_FLOAT16, &antiquantScale);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(antiquantOffsetHostData, antiquantOffsetShape, &antiquantOffsetDeviceAddr,
                          aclDataType::ACL_FLOAT16, &antiquantOffset);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(x3HostData, x3Shape, &x3DeviceAddr, aclDataType::ACL_FLOAT16, &x3);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT16, &out);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      // Call the first-phase API.
      ret = aclnnWeightQuantMatmulAllReduceGetWorkspaceSize(x1, x2, bias, antiquantScale, antiquantOffset, x3, hcom_name,
                                                          "sum", commTurn, streamMode, antiquantGroupSize, out,
                                                          &workspaceSize, &executor);
      CHECK_RET(ret == ACL_SUCCESS,
              LOG_PRINT("aclnnWeightQuantMatmulAllReduceGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
      // Apply for device memory based on the workspaceSize computed by the first-phase API.
      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.
      ret = aclnnWeightQuantMatmulAllReduce(workspaceAddr, workspaceSize, executor, args.stream);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnWeightQuantMatmulAllReduce failed. ERROR: %d\n", ret); return     ret);
      //(Boilerplate) Synchronously wait for task execution to finish.
      ret = aclrtSynchronizeStreamWithTimeout(args.stream, 10000);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
      LOG_PRINT("device%d aclnnWeightQuantMatmulAllReduce execute success \n", args.rankId);
      // Release device resources. Modify based on the specific API definition.
      if (x1 != nullptr) {
          aclDestroyTensor(x1);
      }
      if (x2 != nullptr) {
          aclDestroyTensor(x2);
      }
      if (bias != nullptr) {
          aclDestroyTensor(bias);
      }
      if (antiquantScale != nullptr) {
          aclDestroyTensor(antiquantScale);
      }
      if (antiquantOffset != nullptr) {
          aclDestroyTensor(antiquantOffset);
      }
      if (x3 != nullptr) {
          aclDestroyTensor(x3);
      }
      if (out != nullptr) {
          aclDestroyTensor(out);
      }
      if (x1DeviceAddr != nullptr) {
          aclrtFree(x1DeviceAddr);
      }
      if (x2DeviceAddr != nullptr) {
          aclrtFree(x2DeviceAddr);
      }
      if (biasDeviceAddr != nullptr) {
          aclrtFree(biasDeviceAddr);
      }
      if (antiquantScaleDeviceAddr != nullptr) {
          aclrtFree(antiquantScaleDeviceAddr);
      }
      if (antiquantOffsetDeviceAddr != nullptr) {
          aclrtFree(antiquantOffsetDeviceAddr);
      }
      if (x3DeviceAddr != nullptr) {
          aclrtFree(x3DeviceAddr);
      }
      if (outDeviceAddr != nullptr) {
          aclrtFree(outDeviceAddr);
      }
      if (workspaceSize > 0) {
          aclrtFree(workspaceAddr);
      }
      aclrtDestroyStream(args.stream);
      HcclCommDestroy(args.hcclComm);
      aclrtDestroyContext(args.context);
      aclrtResetDevice(args.rankId);
      return 0;
  }
  
  int main(int argc, char *argv[]) {
      int ret;
      int32_t devices[ndev];
      for (int i = 0; i < ndev; i++) {
          devices[i] = i;
      }
      HcclComm comms[128];
      ret = aclInit(nullptr);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
      // Initialize the communication domain.
      for (int i = 0; i < ndev; i++) {
          ret = aclrtSetDevice(devices[i]);
          CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
      }
      ret = HcclCommInitAll(ndev, devices, comms);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("HcclCommInitAll failed. ERROR: %d\n", ret); return ret);
      Args args[ndev];
      aclrtStream stream[ndev];
      aclrtContext context[ndev];
      for (uint32_t rankId = 0; rankId < ndev; rankId++) {
          ret = aclrtSetDevice(rankId);
          CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
          ret = aclrtCreateContext(&context[rankId], rankId);
          CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateContext failed. ERROR: %d\n", ret); return ret);
          ret = aclrtCreateStream(&stream[rankId]);
          CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
      }
      // Start multiple threads.
      std::vector<std::unique_ptr<std::thread>> threads(ndev);
      for (uint32_t rankId = 0; rankId < ndev; rankId++) {
          args[rankId].rankId = rankId;
          args[rankId].hcclComm = comms[rankId];
          args[rankId].stream = stream[rankId];
          args[rankId].context = context[rankId];
          threads[rankId].reset(
                  new(std::nothrow) std::thread(&launchOneThreadweightQuantmatmulAllReduce, std::ref(args[rankId])));
      }
      for (uint32_t rankId = 0; rankId < ndev; rankId++) {
          threads[rankId]->join();
      }
      aclFinalize();
      return 0;
  }