MoeDistributeCombineAddRmsNorm
产品支持情况
| 产品 | 是否支持 |
|---|---|
| Ascend 950PR/Ascend 950DT | √ |
| Atlas A3 训练系列产品/Atlas A3 推理系列产品 | √ |
| Atlas A2 训练系列产品/Atlas A2 推理系列产品 | × |
| Atlas 200I/500 A2 推理产品 | × |
| Atlas 推理系列产品 | × |
| Atlas 训练系列产品 | × |
功能说明
**算子功能：当存在TP域通信时，先进行ReduceScatterV通信，再进行AllToAllV通信，最后将接收的数据整合（乘权重再相加）；当不存在TP域通信时，进行AllToAllV通信，最后将接收的数据整合（乘权重再相加），之后完成Add + RmsNorm融合。
计算公式：
rsOut=ReduceScatterV(expandX)ataOut=AllToAllV(rsOut)combineOut=Sum(expertScales∗ataOut+expertScales∗sharedExpertX)x=combineOut+residualXy=xRMS(x)∗gamma,whereRMS(x)=1H∑i=1Hxi2+normEpsrsOut = ReduceScatterV(expandX)\\ ataOut = AllToAllV(rsOut)\\ combineOut = Sum(expertScales * ataOut + expertScales * sharedExpertX)\\ x = combineOut + residualX\\ y = \frac{x}{RMS(x)} * gamma,\quad\text{where}RMS(x) = \sqrt{\frac{1}{H}\sum_{i=1}^{H}x_{i}^{2}+normEps}
注意该接口必须与aclnnMoeDistributeDispatchV2配套使用，相当于按MoeDistributeDispatchV2算子收集数据的路径原路返还。
参数说明
| 参数名 | 输入/输出 | 描述 | 使用说明 | 数据类型 | 数据格式 | 维度 | 非连续 |
|---|---|---|---|---|---|---|---|
| expandX | 输入 | 根据expertIds进行扩展过的token特征。 | 要求为2D Tensor，shape为 (max(tpWorldSize, 1) * A , H)。 | BFLOAT16 | ND | 2 | √ |
| expertIds | 输入 | 每个token的topK个专家索引。 | 要求为2D Tensor，shape为 (Bs, K)。 | INT32 | ND | 2 | √ |
| assistInfoForCombine | 输入 | 对应aclnnMoeDistributeDispatchV2中的assistInfoForCombineOut输出。 | 要求为1D Tensor，shape为 (A * 128, )。 | INT32 | ND | 1 | √ |
| epSendCounts | 输入 | 对应aclnnMoeDistributeDispatchV2中的epRecvCounts输出。 | 要求为1D Tensor，shape为 (epWorldSize * max(tpWorldSize, 1) * localExpertNum, )。 | INT32 | ND | 1 | √ |
| expertScales | 输入 | 每个token的topK个专家的权重。 | 要求为2D Tensor，shape为 (Bs, K)。 | FLOAT32 | ND | 2 | √ |
| residualX | 输入 | AddRmsNorm中Add的右矩阵。 | 要求为3D Tensor，shape为 (Bs，1，H)。 | BFLOAT16 | ND | 3 | √ |
| gamma | 输入 | RmsNorm中的gamma输入。 | 要求为1D Tensor，shape为 (H, )。 | BFLOAT16 | ND | 1 | √ |
| tpSendCountsOptional | 输入 | 对应aclnnMoeDistributeDispatchV2中的tpRecvCounts输出。 | 有TP域通信需传参，无TP域通信传空指针；有TP域通信时为1D Tensor，shape为 (tpWorldSize, )。 | INT32 | ND | - | √ |
| xActiveMaskOptional | 输入 | 表示token是否参与通信。 |
|
BOOL | ND | - | √ |
| activationScaleOptional | 输入 | 预留参数。 | 当前版本不支持，传空指针即可。 | - | ND | - | - |
| weightScaleOptional | 输入 | 预留参数。 | 当前版本不支持，传空指针即可。 | - | ND | - | - |
| groupListOptional | 输入 | 预留参数。 | 当前版本不支持，传空指针即可。 | - | ND | - | - |
| expandScalesOptional | 输入 | 对应aclnnMoeDistributeDispatchV2中的expandScales输出；预留参数。 | 当前版本不支持，传空指针即可。 | - | ND | - | - |
| sharedExpertXOptional | 输入 | 表示共享专家计算后的Token。 | 可传有效数据或空指针，2D时shape为(Bs, H)，3D时shape为(Bs, 1, H)） | BFLOAT16 | ND | - | √ |
| groupEp | 输入 | EP通信域名称（专家并行通信域）。 | 字符串长度范围为[1, 128)，不能和groupTp相同。 | STRING | ND | - | - |
| epWorldSize | 输入 | EP通信域大小。 | - | INT64 | ND | - | - |
| epRankId | 输入 | EP域本卡Id。 | 取值范围[0, epWorldSize)，同一个EP通信域中各卡的epRankId不重复。 | INT64 | ND | - | - |
| moeExpertNum | 输入 | MoE专家数量。 | 满足moeExpertNum % (epWorldSize - sharedExpertRankNum) = 0。 | INT64 | ND | - | - |
| groupTp | 输入 | TP通信域名称（数据并行通信域）。 | 不能和groupEp相同。 | STRING | ND | - | - |
| tpWorldSize | 输入 | TP通信域大小。 | 取值范围[0, 2]，0和1表示无TP域通信，有TP域通信时仅支持2。 | INT64 | ND | - | - |
| tpRankId | 输入 | TP域本卡Id。 | 取值范围[0, 1]，同一个TP通信域中各卡的tpRankId不重复；无TP域通信时传0即可。 | INT64 | ND | - | - |
| expertShardType | 输入 | 表示共享专家卡分布类型。 | 当前仅支持传0，表示共享专家卡排在MoE专家卡前面 | INT64 | ND | - | - |
| sharedExpertNum | 输入 | 表示共享专家数量。 | 当前版本不支持，传0即可。 | INT64 | ND | - | - |
| sharedExpertRankNum | 输入 | 表示共享专家卡数量。 | 当前版本不支持，传0即可。 | INT64 | ND | - | - |
| globalBS | 输入 | EP域全局的batch size大小。 |
|
INT64 | ND | - | - |
| outDtype | 输入 | 用于指定输出x的数据类型，预留参数 | 当前版本不支持，传0即可。 | INT64 | ND | - | - |
| commQuantMode | 输入 | 通信量化类型。 | 当前版本不支持，传0即可。 | INT64 | ND | - | - |
| groupListType | 输入 | group List格式，预留参数。 | 当前版本不支持，传0即可。 | INT64 | ND | - | - |
| commAlg | 输入 | 表示通信亲和内存布局算法。 | 预留字段，当前版本不支持，传入空指针即可。 | STRING | ND | - | - |
| normEps | 输入 | 用于防止AddRmsNorm除0错误。 | 可取值为1e-6。 | FLOAT | ND | - | - |
| yOut | 输出 | RmsNorm后的输出结果。 | 要求为3D Tensor。 | BFLOAT16 | ND | 3 | - |
| rstdOut | 输出 | RmsNorm后的输出结果。 | 要求为3D Tensor，shape为（Bs，1，1）。 | FLOAT32 | ND | 3 | √ |
| xOut | 输出 | Add后的输出结果。 | 要求为3D Tensor，shape为 (Bs, 1，H)。 | BFLOAT16 | ND | 3 | - |
| workspaceSize | 输出 | 返回需要在Device侧申请的workspace大小。 | - | - | - | - | - |
| executor | 输出 | 返回op执行器，包含了算子的计算流程。 | - | - | - | - | - |
约束说明
-
aclnnMoeDistributeDispatchV2接口与aclnnMoeDistributeCombineAddRmsNorm接口必须配套使用，具体参考调用示例。
-
调用接口过程中使用的groupEp、epWorldSize、moeExpertNum、groupTp、tpWorldSize、expertShardType、sharedExpertNum、sharedExpertRankNum、globalBs参数取值所有卡需保持一致，网络中不同层中也需保持一致，且和aclnnMoeDistributeDispatchV2对应参数也保持一致。
-
Atlas A3 训练系列产品/Atlas A3 推理系列产品：该场景下单卡包含双DIE（简称为“晶粒”或“裸片”），因此参数说明里的“本卡”均表示单DIE。
-
参数说明里shape格式说明：
- A：表示本卡需要分发的最大token数量，取值范围如下：
- 对于共享专家，要满足A = Bs * epWorldSize * sharedExpertNum / sharedExpertRankNum。
- 对于MoE专家，当globalBs为0时，要满足A >= Bs * epWorldSize * min(localExpertNum, K)；当globalBs非0时，要满足A >= globalBs * min(localExpertNum, K)。
- H：表示hidden size隐藏层大小，取值范围为[1024, 8192]。
- Bs：表示batch sequence size，即本卡最终输出的token数量，取值范围为0 < Bs ≤ 512。
- K：表示选取topK个专家，取值范围为0 < K ≤ 16同时满足0 < K ≤ moeExpertNum。
- localExpertNum：表示本卡专家数量。
- 对于共享专家卡，localExpertNum = 1
- 对于MoE专家卡，localExpertNum = moeExpertNum / (epWorldSize - sharedExpertRankNum)，localExpertNum > 1时，不支持TP域通信。
- A：表示本卡需要分发的最大token数量，取值范围如下：
-
HCCL_BUFFSIZE： 调用本接口前需检查HCCL_BUFFSIZE环境变量取值是否合理，该环境变量表示单个通信域占用内存大小，单位MB，不配置时默认为200MB。要求 >= 2且满足1024 ^ 2 * (HCCL_BUFFSIZE - 2) / 2 >= (BS * 2 * (H + 128) * (epWorldSize * localExpertNum + K + 1))，localExpertNum需使用MoE专家卡的本卡专家数。
-
通信域使用约束：
- 一个模型中的aclnnMoeDistributeCombineAddRmsNorm和aclnnMoeDistributeDispatchV2仅支持相同EP通信域，且该通信域中不允许有其他算子。
- 一个模型中的aclnnMoeDistributeCombineAddRmsNorm和aclnnMoeDistributeDispatchV2仅支持相同TP通信域或都不支持TP通信域，有TP通信域时该通信域中不允许有其他算子。
调用说明
| 调用方式 | 样例代码 | 说明 |
|---|---|---|
| aclnn接口 | test_aclnn_moe_distribute_combine_add_rms_norm.cpp | 通过aclnnMoeDistributeCombineAddRmsNorm接口方式调用moe_distribute_combine_add_rms_norm算子。 |