MoeDistributeCombine

产品支持情况

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

功能说明

算子功能：当存在TP域通信时，先进行ReduceScatterV通信，再进行AlltoAllV通信，最后将接收的数据整合（乘权重再相加）；当不存在TP域通信时，进行AlltoAllV通信，最后将接收的数据整合（乘权重再相加）。

• 不存在TP域通信时：

  $$\begin{gathered} ataOut=AllToAllV(expandX) \\ xOut=Sum(expertScales\ast ataOut+expertScales\ast sharedExpertX) \end{gathered}$$

• 存在TP域通信时：

  $$\begin{gathered} rsOut=ReduceScatterV(expandX) \\ ataOut=AllToAllV(rsOut) \\ xOut=Sum(expertScales\ast ataOut+expertScales\ast sharedExpertX) \end{gathered}$$

注意该算子必须与MoeDistributeDispatch配套使用，相当于按MoeDistributeDispatch算子收集数据的路径原路返还。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
expandX	输入	根据expertIds进行扩展过的token特征。	FLOAT16、BFLOAT16	ND
expertIds	输入	每个token的topK个专家索引。	INT32	ND
expandIdx	输入	表示同一专家收到的token个数，对应MoeDistributeDispatch中的expandIdx输出。	INT32	ND
epSendCounts	输入	从EP通信域各卡接收的token数，对应MoeDistributeDispatch中的epRecvCounts输出。	INT32	ND
expertScales	输入	每个token的topK个专家的权重。	FLOAT32	ND
tpSendCounts	可选输入	从TP通信域各卡接收的token数，对应MoeDistributeDispatch中的tpRecvCounts输出，若有TP域通信需传参，若无TP域通信传空指针。	INT32	ND
xActiveMask	可选输入	预留参数，当前版本不支持，传空指针即可。	-	ND
activationScale	可选输入	预留参数，当前版本不支持，传空指针即可。	-	ND
weightScale	可选输入	预留参数，当前版本不支持，传空指针即可。	-	ND
groupList	可选输入	预留参数，当前版本不支持，传空指针即可。	-	ND
expandScales	可选输入	表示本卡输入Token的权重，对应MoeDistributeDispatch中的expandScales输出。	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。同一个EP通信域中各卡的tpRankId不重复。 默认值为0。	INT64	ND
expertShardType	可选属性	表示共享专家卡分布类型，当前仅支持传0，表示共享专家卡排在MoE专家卡前面。 默认值为0。	INT64	ND
sharedExpertNum	可选属性	表示共享专家数量（一个共享专家可复制部署到多个卡上）。 默认值为1。	INT64	ND
sharedExpertRankNum	可选属性	表示共享专家卡数量。 默认值为0。	INT64	ND
globalBS	可选属性	EP域全局的batch size大小。 默认值为0。	INT64	ND
outDtype	可选属性	用于指定输出x的数据类型，预留参数，当前版本不支持，传0即可。 默认值为0。	INT64	ND
commQuantMode	可选属性	通信量化类型，取值范围0或2；0表示通信不量化，2表示通信int8量化。 默认值为0。	INT64	ND
groupListType	可选属性	groupList格式，预留参数，当前版本不支持，传0即可。 默认值为0。	INT64	ND
x	输出	表示处理后的token，数据类型、数据格式与expandX保持一致。	FLOAT16、BFLOAT16	ND

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

  • 不支持共享专家场景，不支持expertShardType、sharedExpertNum、sharedExpertRankNum属性。
  • 当前不支持TP域通信，不支持groupTp、tpWorldSize、tpRankId属性，且tpSendCounts为无效内容。
  • 仅设置环境变量HCCL_INTRA_PCIE_ENABLE = 1和HCCL_INTRA_ROCE_ENABLE = 0时，必须传入expandScales。

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

  • 不支持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属性取值所有卡需保持一致，网络中不同层中也需保持一致，且和MoeDistributeDispatch对应参数也保持一致。

• 参数说明里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，其中Align8(x) = ((x + 8 - 1) / 8) * 8。
  • 参数约束：
    • epWorldSize：取值支持16、32、64。
    • moeExpertNum：需满足moeExpertNum / epWorldSize <= 24。
      • 环境变量HCCL_INTRA_PCIE_ENABLE = 1和HCCL_INTRA_ROCE_ENABLE = 0时，无上述约束。
    • globalBS：当每个rank的BS数一致时，globalBS = BS * epWorldSize 或 globalBS = 0；当每个rank的BS数不一致时，globalBS = maxBS * epWorldSize或者globalBS = 256 * epWorldSize，其中maxBS表示表示单rank BS最大值，建议按maxBS * epWorldSize传入，固定按256 * epWorldSize传入在后续版本BS支持大于256的场景下会无法支持。
    • commQuantMode取值范围0或2，0表示通信不量化，2表示通信int8量化（2仅当HCCL_INTRA_PCIE_ENABLE=1、HCCL_INTRA_ROCE_ENABLE=0且驱动版本≥25.0.RC1.1时支持）。
  • 组网约束：多机场景仅支持交换机组网，不支持双机直连组网。

• 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]。
  • 参数约束：
    • epWorldSize：取值支持8、16、32、64、128、144、256、288。
    • groupTp：字符串长度范围为[1, 128)，不能和groupEp相同。
    • tpWorldSize：取值范围[0, 2]，0和1表示无tp域通信，有tp域通信时仅支持2。
    • tpRankId：取值范围[0, 1]，同一个TP通信域中各卡的tpRankId不重复。无TP域通信时，传0即可。
    • sharedExpertRankNum：当前取值范围[0, epWorldSize)，不为0时需满足epWorldSize % sharedExpertRankNum = 0。
    • globalBS：当每个rank的BS数一致时，globalBS = BS * epWorldSize 或 globalBS = 0；当每个rank的BS数不一致时，globalBS = maxBS * epWorldSize，其中maxBS表示单卡BS最大值。
  • HCCL_BUFFSIZE：调用本算子前需检查HCCL_BUFFSIZE环境变量取值是否合理，该环境变量表示单个通信域占用内存大小，单位MB，不配置时默认为200MB，要求 >= 2且满足1024 ^ 2 * (HCCL_BUFFSIZE - 2) / 2 >= BS * 2 * (H + 128) * (epWorldSize * localExpertNum + K + 1)，localExpertNum需使用MoE专家卡的本卡专家数。

• Ascend 950PR/Ascend 950DT：

  • 参数约束：
    • epWorldSize：取值支持2、4、8、16、32、64、128、144、256、288。
    • sharedExpertRankNum：当前取值范围[0, epWorldSize)，不为0时需满足epWorldSize % sharedExpertRankNum = 0。
    • globalBS：当每个rank的BS数一致时，globalBS = BS * epWorldSize 或 globalBS = 0；当每个rank的BS数不一致时，globalBS = maxBS * epWorldSize，其中maxBS表示单卡BS最大值。
    • commQuantMode取值范围0或2，0表示通信不量化，2表示通信int8量化。
  • HCCL_BUFFSIZE：调用本算子前需检查HCCL_BUFFSIZE环境变量取值是否合理，该环境变量表示单个通信域占用内存大小，单位MB，不配置时默认为200MB，要求 >= 2且满足1024 ^ 2 * (HCCL_BUFFSIZE - 2) / 2 >= BS * 2 * (H + 128) * (epWorldSize * localExpertNum + K + 1)，localExpertNum需使用MoE专家卡的本卡专家数。

调用说明

调用方式	样例代码	说明
aclnn接口	test_aclnn_moe_distribute_combine.cpp	通过aclnnMoeDistributeCombine接口方式调用moe_distribute_combine算子。