MoeDistributeDispatch

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

功能说明

• 算子功能：对Token数据进行量化（可选），当存在TP域通信时，先进行EP（Expert Parallelism）域的AllToAllV通信，再进行TP（Tensor Parallelism）域的AllGatherV通信；当不存在TP域通信时，进行EP（Expert Parallelism）域的AllToAllV通信。

• 计算公式：

  • 情形1：如果quaneMode=0（非量化场景）：

  $$\begin{gathered} allToAllXOut=AllToAllV(X) \\ expandXOut=\{\begin{array}{ll}AllToAllV(X),& 无TP通信域\\ AllGatherV(allToAllXOut),& 有TP通信域\end{array} \end{gathered}$$

  • 情形2：如果quaneMode=2（pertoken动态量化场景）：

  $$\begin{gathered} xFp32=CastToFp32(X)\times scales \\ dynamicScales=dstTypeMax/Max(Abs(xFp32)) \\ quantOut=CastToInt8(xFp32\times dynamicScalesValue) \\ allToAllXOut=AllToAllV(quantOut) \\ allToAllDynamicScalesOut=AllToAllV(1.0/dynamicScales) \\ expandXOut=\{\begin{array}{ll}AllToAllV(X),& 无TP通信域\\ AllGatherV(allToAllXOut),& 有TP通信域\end{array} \\ dynamicScalesOut=\{\begin{array}{ll}AllGatherV(allToAllDynamicScalesOut),& 无TP通信域\\ allToAllDynamicScalesOut,& 有TP通信域\end{array} \end{gathered}$$

  注意该算子必须与MoeDistributeCombine配套使用。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
x	输入	本卡发送的token数据。	FLOAT16、BFLOAT16	ND
expertIds	输入	每个token的topK个专家索引。	INT32	ND
scales	可选输入	每个专家的平滑权重、融合量化平滑权重的量化系数或量化系数。	FLOAT32	ND
xActiveMask	可选输入	预留参数，当前版本不支持，传空指针即可。	-	ND
expertScales	可选输入	每个Token的topK个专家权重。	FLOAT32	ND
groupEp	属性	EP通信域名称（专家并行通信域），字符串长度范围为[1, 128)，不能和groupTp相同。	STRING	ND
epWorldSize	属性	EP通信域大小。	INT64	ND
epRankId	属性	EP域本卡ID，取值范围[0, epWorldSize)，同一个EP通信域中各卡的epRankId不重复。	INT64	ND
moeExpertNum	属性	MoE专家数量，取值范围(0, 512]，且满足moeExpertNum % (epWorldSize - sharedExpertRankNum) = 0。	INT64	ND
groupTp	可选属性	TP通信域名称（数据并行通信域）。 默认值为""。	STRING	ND
tpWorldSize	可选属性	TP通信域大小。 默认值为0。	INT64	ND
tpRankId	可选属性	TP域本卡ID。 默认值为0。	INT64	ND
expertShardType	可选属性	表示共享专家卡分布类型，当前仅支持传0，表示共享专家卡排在MoE专家卡前面，默认值为0。 默认值为0。	INT64	ND
sharedExpertNum	可选属性	表示共享专家数量（一个共享专家可复制部署到多个卡上），当前取值范围[0, 1]，0表示无共享专家，1表示一个共享专家，当前版本仅支持传1。 默认值为1。	INT64	ND
sharedExpertRankNum	可选属性	表示共享专家卡数量，当前取值范围[0, epWorldSize)，不为0时需满足epWorldSize % sharedExpertRankNum = 0。 默认值为0。	INT64	ND
quantMode	可选属性	表示量化模式，支持0：非量化，1：静态量化，2：pertoken动态量化，3：pergroup动态量化，4：mx量化。 默认值为0。	INT64	ND
globalBS	可选属性	EP域全局的batch size大小，各rank BS一致时，globalBS = BS * epWorldSize 或 0；各rank BS不一致时，globalBS = maxBS * epWorldSize（maxBS为单rank BS最大值）。 默认值为0。	INT64	ND
expertTokenNumsType	可选属性	输出expertTokenNums中值的语义类型，支持0：expertTokenNums中的输出为每个专家处理的token数的前缀和，1：expertTokenNums中的输出为每个专家处理的token数量。 默认值为0。	INT64	ND
expandX	输出	根据expertIds进行扩展过的token特征。	FLOAT16、BFLOAT16、INT8、FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8	ND
dynamicScales	输出	量化场景下，表示本卡输出Token的量化系数。	FLOAT32、FLOAT8_E8M0	ND
expandIdx	输出	表示给同一专家发送的token个数（对应算子MoeDistributeCombine中的expandIdx）。	INT32	ND
expertTokenNums	输出	表示每个专家收到的token个数。	INT64	ND
epRecvCounts	输出	从EP通信域各卡接收的token数（对应MoeDistributeCombine中的epSendCounts）。	INT32	ND
tpRecvCounts	输出	从TP通信域各卡接收的token数（对应MoeDistributeCombine中的tpSendCounts）；若有TP域通信则有该输出，若无TP域通信则无该输出。	INT32	ND
expandScales	输出	表示本卡输出Token的权重（对应MoeDistributeCombine中的expandScales）。	FLOAT32	ND

• Atlas A2 训练系列产品/Atlas A2 推理系列产品：

  • 支持FLOAT16、BFLOAT16数据类型。
  • quantMode属性仅支持0和2。
  • 不支持共享专家场景，不支持expertShardType、sharedExpertNum、sharedExpertRankNum属性。
  • 仅支持EP域，无TP域，不支持groupTp、tpWorldSize、tpRankId属性，tpRecvCounts为无效内容。
  • 仅设置环境变量HCCL_INTRA_PCIE_ENABLE = 1和HCCL_INTRA_ROCE_ENABLE = 0时，expandScales内容有效。

• Atlas A3 训练系列产品/Atlas A3 推理系列产品：

  • 不支持FLOAT8_E4M3FN、FLOAT8_E5M2、HIFLOAT8、FLOAT8_E8M0数据类型。
  • quantMode属性仅支持0和2。
  • 不支持expandScales。

• Ascend 950PR/Ascend 950DT：

  • 不支持expandScales。
  • 当前不支持TP域通信，不支持groupTp、tpWorldSize、tpRankId属性，且tpSendCounts为无效内容。

约束说明

• MoeDistributeDispatch算子与MoeDistributeCombine算子必须配套使用，具体参考调用示例。

• 算子通信域各节点的驱动版本应当相同。

• 在不同产品型号、不同通信算法或不同版本中，MoeDistributeDispatch的Tensor输出expandIdx、epRecvCounts、tpRecvCounts、expandScales中的元素值可能不同，使用时直接将上述Tensor传给MoeDistributeCombine对应参数即可，模型其他业务逻辑不应对其存在依赖。

• 调用算子过程中使用的groupEp、epWorldSize、moeExpertNum、groupTp、tpWorldSize、expertShardType、sharedExpertNum、sharedExpertRankNum、globalBS属性取值所有卡需保持一致，网络中不同层中也需保持一致，且和MoeDistributeCombine对应参数也保持一致。

• 参数说明里shape格式说明：

  • A：表示本卡可能接收的最大token数量，取值范围如下：
    • 对于共享专家，要满足A = BS * epWorldSize * sharedExpertNum / sharedExpertRankNum。
    • 对于MoE专家，当globalBS为0时，要满足A >= BS * epWorldSize * min(localExpertNum, K)；当globalBS非0时，要满足A >= globalBS * min(localExpertNum, K)。
  • localExpertNum：表示本卡专家数量。
    • 对于共享专家卡，localExpertNum = 1
    • 对于MoE专家卡，localExpertNum = moeExpertNum / (epWorldSize - sharedExpertRankNum)，localExpertNum > 1时，不支持TP域通信。

• 本文公式中的"/"表示整除。

• 通信域使用约束：

  • 一个模型中的MoeDistributeCombine和MoeDistributeDispatch仅支持相同EP通信域，且该通信域中不允许有其他算子。
  • 一个模型中的MoeDistributeCombine和MoeDistributeDispatch仅支持相同TP通信域或都不支持TP通信域，有TP通信域时该通信域中不允许有其他算子。

• Atlas A2 训练系列产品/Atlas A2 推理系列产品：

  • 参数说明里shape格式说明：
    • H：表示hidden size隐藏层大小，取值范围(0, 7168]，且保证是32的整数倍。
    • BS：表示batch sequence size，即本卡最终输出的token数量，取值范围为[1, 256]。
    • K：表示选取topK个专家，需满足0 < K ≤ moeExpertNum，取值范围为[1, 16]。
  • HCCL_BUFFSIZE：调用本算子前需检查HCCL_BUFFSIZE环境变量取值是否合理，该环境变量表示单个通信域占用内存大小，单位MB，不配置时默认为200MB，要求 >= (BS * epWorldSize * min(localExpertNum, K) * H * 4B + 4MB)。
  • HCCL_INTRA_PCIE_ENABLE和HCCL_INTRA_ROCE_ENABLE：设置环境变量HCCL_INTRA_PCIE_ENABLE = 1和HCCL_INTRA_ROCE_ENABLE = 0可以减少跨机通信数据量，可能提升算子性能。此时，要求HCCL_BUFFSIZE >= moeExpertNum * BS * (H * 2 + 16 * Align8(K))B + 104MB。并且，对于入参moeExpertNum，只要求moeExpertNum % epWorldSize = 0，不要求moeExpertNum / epWorldSize <= 24，但不支持scales特性，其中Align8(x) = ((x + 8 - 1) / 8) * 8。
  • epWorldSize：取值支持16、32、64。
  • quantMode相关约束：
    • quantMode取值为2时，表示pertoken动态量化场景，expandX的数据类型支持INT8。
      • 输入scales可传入空指针。
      • 若输入scales传入有效数据时，其shape为 (moeExpertNum, H)。
  • 组网约束：多机场景仅支持交换机组网，不支持双机直连组网。

• Atlas A3 训练系列产品/Atlas A3 推理系列产品：

  • 该场景下单卡包含双DIE（简称为“晶粒”或“裸片”），因此参数说明里的“本卡”均表示单DIE。
  • 参数说明里shape格式说明：
    • H：表示hidden size隐藏层大小，取值为7168。
    • BS：表示batch sequence size，即本卡最终输出的token数量，取值范围为[1, 512]。
    • K：表示选取topK个专家，需满足0 < K ≤ moeExpertNum，取值范围为[1, 8]。
  • HCCL_BUFFSIZE：调用本算子前需检查HCCL_BUFFSIZE环境变量取值是否合理，该环境变量表示单个通信域占用内存大小，单位MB，不配置时默认为200MB，要求 >= 2且满足1024 ^ 2 * (HCCL_BUFFSIZE - 2) / 2 >= BS * 2 * (H + 128) * (epWorldSize * localExpertNum + K + 1)，localExpertNum需使用MoE专家卡的本卡专家数。
  • epWorldSize：取值支持8、16、32、64、128、144、256、288。
  • quantMode相关约束：
    • quantMode取值为2时，表示pertoken动态量化场景，expandX的数据类型支持INT8。
      • 输入scales可传入空指针。
      • 若输入scales传入有效数据且存在共享专家卡时，其shape为 (sharedExpertNum + moeExpertNum, H)。
      • 若输入scales传入有效数据且不存在共享专家卡时，其shape为 (moeExpertNum, H)。

• Ascend 950PR/Ascend 950DT：

  • 参数说明里shape格式说明：
    • H：表示hidden size隐藏层大小，取值为7168。
    • BS：表示batch sequence size，即本卡最终输出的token数量，取值范围为[1, 512]。
    • K：表示选取topK个专家，需满足0 < K ≤ moeExpertNum，取值范围为[1, 8]。
  • epWorldSize：取值支持2、4、8、16、32、64、128、144、256、288。
  • HCCL_BUFFSIZE：调用本算子前需检查HCCL_BUFFSIZE环境变量取值是否合理，该环境变量表示单个通信域占用内存大小，单位MB，不配置时默认为200MB，要求 >= 2且满足1024 ^ 2 * (HCCL_BUFFSIZE - 2) / 2 >= BS * 2 * (H + 128) * (epWorldSize * localExpertNum + K + 1)，localExpertNum需使用MoE专家卡的本卡专家数。
  • quantMode相关约束：
    • quantMode取值为2时，表示pertoken动态量化场景，expandX的数据类型支持INT8。
      • 输入scales可传入空指针。
      • 若输入scales传入有效数据且存在共享专家卡时，其shape为 (sharedExpertNum + moeExpertNum, H)。
      • 若输入scales传入有效数据且不存在共享专家卡时，其shape为 (moeExpertNum, H)。

调用说明

调用方式	样例代码	说明
aclnn接口	test_aclnn_moe_distribute_dispatch.cpp	通过aclnnMoeDistributeDispatch接口方式调用moe_distribute_dispatch算子。