aclnnFusedLinearCrossEntropyLossGrad

产品支持情况

产品	是否支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]	√
[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]	√
[object Object]Atlas 200I/500 A2 推理产品[object Object]	×
[object Object]Atlas 推理系列产品 [object Object]	×
[object Object]Atlas 训练系列产品[object Object]	×

功能说明

• 算子功能：本算子是词汇表并行场景下交叉熵损失计算模块中的一部分，解决超大规模词汇表下的显存和计算效率问题，当前部分为梯度计算实现，用于计算叶子节点input和weight的梯度。 需要获得aclnnFusedLinearOnlineMaxSum、aclnnFusedCrossEntropyLossWithMaxSum的相关输出，以及logits相关的全局通信结果作为本接口输入。
• 计算公式：

  高性能模式，softmaxOptional非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}$$ [object Object]$$\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}$$

函数原型

每个算子分为，必须先调用aclnnFusedLinearCrossEntropyLossGradGetWorkspaceSize接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用aclnnFusedLinearCrossEntropyLossGrad接口执行计算。

• 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

• 参数说明：

  • grad(aclTensor*, 计算输入)：当前节点的梯度，公式中的grad，Device侧的aclTensor。数据类型支持FLOAT32，shape支持1维。支持，支持空tensor，支持ND。
  • input(aclTensor*, 计算输入)：矩阵乘的输入矩阵，公式中的input，Device侧的aclTensor。数据类型支持FLOAT16、BFLOAT16，shape支持2维，其中第1维长度与输入grad的长度相同。支持，支持空tensor，支持ND。
  • weight(aclTensor*, 计算输入)：矩阵乘的权重矩阵，公式中的weight，Device侧的aclTensor。数据类型支持FLOAT16、BFLOAT16，数据类型与输入input相同，shape支持2维，其中第1维长度不支持小于128，第2维长度与输入input的第2维长度相同。支持，支持空tensor，支持ND。
  • targetMask(aclTensor*, 计算输入)：中间变量，代表对应词ID是否在目标范围内，公式中的target_mask，Device侧的aclTensor。数据类型支持UINT8，其中每1bit数据代表1个布尔值，0代表false，1代表true。shape支持1维，长度乘以8须不小于输入grad的长度。支持，支持空tensor，支持ND。
  • maskedTarget(aclTensor*, 计算输入)：中间变量，代表对应词ID映射到当前设备词汇表分片的局部索引，无效目标会被掩码targetMask处理，公式中的masked_target，Device侧的aclTensor。数据类型支持INT64、INT32，shape支持1维，长度与输入grad相同。支持，支持空tensor，支持ND。
  • labelSmoothing(float, 计算输入)：标签平滑系数，用于缓解过拟合，目前仅支持取值为0，属性参数。
  • logitsMaxOptional(aclTensor*, 可选输入)：中间变量，全局logits的最大值，公式中的logits_max，Device侧的aclTensor。可选输入，支持输入nullptr，输入nullptr的场景需要提供有效的softmaxOptional输入。数据类型支持FLOAT32，shape支持1维，长度与输入grad相同。支持，支持空tensor，支持ND。
  • sumExpLogitsOptional(aclTensor*, 可选输入)：中间变量，处理后的logits，公式中的sum_exp_logits，Device侧的aclTensor。可选输入，支持输入nullptr，输入nullptr的场景需要提供有效的softmaxOptional输入。数据类型支持FLOAT32，shape支持1维，长度与输入grad相同。支持，支持空tensor，支持ND。
  • softmaxOptional(aclTensor*, 计算输入)：中间变量，矩阵乘的结果，公式中的softmax，Device侧的aclTensor。可选输入，支持输入nullptr，输入nullptr时须提供有效的logitsMaxOptional、sumExpLogitsOptional输入；输入非nullptr时，logitsMaxOptional、sumExpLogitsOptional输入无效。数据类型支持FLOAT32，shape支持2维，其中第1维长度与输入grad长度相同，第二维长度与输入weight的第1维长度相同。支持，支持空tensor，支持ND。
  • inputGradOut(aclTensor*, 计算输出)：对应叶子节点input的梯度，公式中的grad_input，Device侧的aclTensor。数据类型支持FLOAT16、BFLOAT16，数据类型与输入input相同，shape支持2维，第1维长度与输入grad长度相同，第2维长度与输入weight的第2维长度相同。支持ND。
  • weightGradOut(aclTensor*, 计算输出)：对应叶子节点weight的梯度，公式中的grad_weight，Device侧的aclTensor。数据类型支持FLOAT16、BFLOAT16，数据类型与输入input相同，shape支持2维，第1维长度与输入weight的第1维长度相同，第2维长度与输入grad长度相同。支持ND。
  • workspaceSize(uint64_t*, 出参)：返回需要在Device侧申请的workspace大小。
  • executor(aclOpExecutor**, 出参)：返回op执行器，包含了算子计算流程。

• 返回值：

  aclnnStatus：返回状态码，具体参见。

  [object Object]

aclnnFusedLinearCrossEntropyLossGrad

• 参数说明：

  • workspace(void*, 入参)：在Device侧申请的workspace内存地址。
  • workspaceSize(uint64_t, 入参)：在Device侧申请的workspace大小，由第一段接口aclnnFusedLinearCrossEntropyLossGradGetWorkspaceSize获取。
  • executor(aclOpExecutor*, 入参)：op执行器，包含了算子计算流程。
  • stream(aclrtStream, 入参)：指定执行任务的Stream。

• 返回值：

  aclnnStatus：返回状态码，具体参见。

约束说明

无。

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考。