0
代码
代码
Issues
Pull Requests
讨论
Wiki
分析
0
cann-robotcann-robotmc2 docs fix
ebf8e05e创建于 4月14日历史提交

MoeDistributeDispatchV3

产品支持情况

产品 是否支持
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：如果quant_mode=0（非量化场景）：

allToAllXOut = AllToAllV(X)\\ expand_x_out = \begin{cases} AllToAllV(X), & 无TP通信域 \\ AllGatherV(allToAllXOut), & 有TP通信域 \\ \end{cases}

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

xFp32 = CastToFp32(X) \times scales \\ dynamicScales = dstTypeMax/Max(Abs(xFp32)) \\ quantOut = CastToInt8(xFp32 \times dynamicScales) \\ allToAllXOut = AllToAllV(quantOut) \\ allToAllDynamicScalesOut = AllToAllV(1.0/dynamicScales) \\ expand_x_out = \begin{cases} AllToAllV(quantOut), & 无TP通信域 \\ AllGatherV(allToAllXOut), & 有TP通信域 \\ \end{cases} \\ dynamic_scales_out = \begin{cases} AllGatherV(allToAllDynamicScalesOut), & 无TP通信域 \\ allToAllDynamicScalesOut, & 有TP通信域 \\ \end{cases}

其中，emaxemax表示该类型最大正规数对应的指数部分的值。

  • Atlas A3 训练系列产品/Atlas A3 推理系列产品：该算子必须与MoeDistributeCombineV3一起使用。

相较于MoeDistributeDispatchV3算子，该算子变更如下：

  • 新增context入参，存入通信域相关信息；
  • 新增ccl_buffer_size入参，指定当前通信域大小；
  • 减少group_ep以及group_tp通信域名称入参;

详细说明请参考以下参数说明。

参数说明

参数名 输入/输出/属性 描述 数据类型 数据格式
context 输入 本卡通信域信息数据。 FLOAT16、BFLOAT16 ND
x 输入 本卡发送的token数据。 FLOAT16、BFLOAT16 ND
expert_ids 输入 每个token的topK个专家索引。 INT32 ND
scales_optional 可选输入 每个专家的量化平滑参数，非量化场景传空指针，动态量化可传有效数据或空指针。 FLOAT32 ND
x_active_mask_optional 可选输入 表示token是否参与通信，可传有效数据或空指针；1D时true需排在false前（例：{true, false, true}非法），2D时token对应K个值全为false则不参与通信；默认所有token参与通信；各卡BS不一致时所有token需有效。 BOOL ND
expert_scales_optional 可选输入 每个token的topK个专家权重。 FLOAT32 ND
elastic_info_optional 可选输入 EP通信域动态缩容信息。 FLOAT32 ND
performance_info_optional 可选输入 表示本卡等待各卡数据的通信时间，单位为us（微秒）。单次算子调用各卡通信耗时会累加到该Tensor上，算子内部不进行自动清零，因此用户每次启用此Tensor开始记录耗时前需对Tensor清零。 INT64 ND
ep_world_size 属性 EP通信域大小。 INT64 ND
ep_rank_id 属性 EP域本卡Id，取值范围[0, ep_world_size)，同一个EP通信域中各卡的ep_rank_id不重复。 INT64 ND
moe_expert_num 属性 MoE专家数量，满足moe_expert_num % (ep_world_size - shared_expert_num) = 0。 INT64 ND
ccl_buffer_size 属性
  • 当前通信域Buffer大小。
  • 默认值为""。
    		•  STRING ND
    tp_world_size 可选属性
  • TP通信域大小，取值范围[0, 2]，0和1表示无TP域通信，有TP域通信时仅支持2。
  • 默认值为0。
    		•  INT64 ND
    tp_rank_id 可选属性
  • TP域本卡Id，取值范围[0, 1]，同一个TP通信域中各卡的tp_rank_id不重复；无TP域通信时传0即可。
  • 默认值为0。
    		•  INT64 ND
    expert_shard_type 可选属性
  • 表示共享专家卡分布类型，当前仅支持传0，表示共享专家卡排在MoE专家卡前面。
  • 默认值为0。
    		•  INT64 ND
    shared_expert_num 可选属性
  • 表示共享专家数量（一个共享专家可复制部署到多个卡上）。
  • 默认值为1。
    		•  INT64 ND
    shared_expert_rank_num 可选属性
  • 表示共享专家卡数量，取值范围[0, ep_world_size)；为0时需满足shared_expert_num为0或1，不为0时需满足shared_expert_rank_num % shared_expert_num = 0。
  • 默认值为0。
    		•  INT64 ND
    quant_mode 可选属性
  • 表示量化模式，支持0：非量化，2：动态量化。
  • 默认值为0。
    		•  INT64 ND
    global_bs 可选属性
  • EP域全局的batch size大小；各rank BS一致时，global_bs = BS * ep_world_size 或 0；各rank BS不一致时，global_bs = max_bs * ep_world_size（max_bs为单卡BS最大值）。
  • 默认值为0。
    		•  INT64 ND
    expert_token_nums_type 可选属性
  • 输出expert_token_nums中值的语义类型，支持0：expert_token_nums中的输出为每个专家处理的token数的前缀和，1：expert_token_nums中的输出为每个专家处理的token数量。
  • 默认值为1。
    		•  INT64 ND
    comm_alg 可选属性
  • 表示通信亲和内存布局算法。
  • 默认值为""。
    		•  STRING ND
    zero_expert_num 可选属性
  • 零专家数量。
  • 默认值为0。
    		•  INT64 ND
    copy_expert_num 可选属性
  • copy专家数量。
  • 默认值为0。
    		•  INT64 ND
    const_expert_num 可选属性
  • 常量专家数量。
  • 默认值为0。
    		•  INT64 ND
    expand_x_out 输出 根据expert_ids进行扩展过的token特征。 FLOAT16、BFLOAT16、INT8 ND
    dynamic_scales_out 输出 量化场景下，表示本卡输出Token的量化系数，仅quant_mode=2时有该输出。 FLOAT32 ND
    assist_info_for_combine_out 输出 表示给同一专家发送的token个数。 INT32 ND
    expert_token_nums_out 输出 表示每个专家收到的token个数。 INT64 ND
    ep_recv_count_out 输出 从EP通信域各卡接收的token数。 INT32 ND
    tp_recv_count_out 输出 从TP通信域各卡接收的token数，有TP域通信则有该输出，无TP域通信则无该输出。 INT32 ND
    expand_scales_out 输出 表示本卡输出token的权重。 FLOAT32 ND
    • Atlas A3 训练系列产品/Atlas A3 推理系列产品：
      • 不支持expand_scales_out。

    约束说明

    • MoeDistributeDispatchV3与CombineV3系列算子必须配套使用，具体参考调用示例。

    • 在不同产品型号、不同通信算法或不同版本中，MoeDistributeDispatchV3的Tensor输出assist_info_for_combine_out、ep_recv_count_out、tp_recv_count_out、expand_scales_out中的元素值可能不同，使用时直接将上述Tensor传给CombineV3系列算子对应参数即可，模型其他业务逻辑不应对其存在依赖。

    • 调用算子过程中使用的ep_world_size、moe_expert_num、ccl_buffer_size、tp_world_size、expert_shard_type、shared_expert_num、shared_expert_rank_num、global_bs、comm_alg参数取值所有卡需保持一致，网络中不同层中也需保持一致，且和CombineV3系列算子对应参数也保持一致。

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

      • A：表示本卡可能接收的最大token数量，取值范围如下：
        • 对于共享专家，要满足A = BS * ep_world_size * shared_expert_num / shared_expert_rank_num。
        • 对于MoE专家，当global_bs为0时，要满足A >= BS * ep_world_size * min(local_expert_num, K)；当global_bs非0时，要满足A >= global_bs * min(local_expert_num, K)。
      • K：表示选取topK个专家，取值范围为0 < K ≤ 16同时满足0 < K ≤ moe_expert_num + zero_expert_num + copy_expert_num + const_expert_num。
      • local_expert_num：表示本卡专家数量。
        • 对于共享专家卡，local_expert_num = 1
        • 对于MoE专家卡，local_expert_num = moe_expert_num / (ep_world_size - shared_expert_rank_num)，local_expert_num > 1时，不支持TP域通信。

    • 属性约束：

      • zero_expert_num：取值范围：[0, MAX_INT32)，MAX_INT32 = 2^31 - 1, 合法的零专家的ID的值是[moe_expert_num, moe_expert_num + zero_expert_num)。
      • copy_expert_num：取值范围：[0, MAX_INT32)，MAX_INT32 = 2^31 - 1, 合法的copy专家的ID的值是[moe_expert_num + zero_expert_num, moe_expert_num + zero_expert_num + copy_expert_num)。
      • const_expert_num：取值范围：[0, MAX_INT32)，MAX_INT32 = 2^31 - 1, 合法的常量专家的ID的值是[moe_expert_num + zero_expert_num + copy_expert_num, moe_expert_num + zero_expert_num + copy_expert_num + const_expert_num)。

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

    • 通信域使用约束：

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

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

      • 该场景下单卡包含双DIE（简称为“晶粒”或“裸片”），因此参数说明里的“本卡”均表示单DIE。
      • 参数约束：
        • elastic_info_optional：当前版本不支持，传空指针即可。
        • ep_world_size：取值范围[2, 768]。
        • moe_expert_num：取值范围(0, 1024]。
        • shared_expert_num：取值支持[0, 4]。
        • comm_alg：当前版本仅支持""，"fullmesh_v1"，"fullmesh_v2"三种输入方式。
          • ""：默认值，使能fullmesh_v1模板。
          • "fullmesh_v1"：使能fullmesh_v1模板。
          • "fullmesh_v2"：使能fullmesh_v2模板，其中comm_alg仅在tp_world_size取值为1时生效，且不支持在各卡BS不一致、输入xActiveMask和特殊专家场景下使能。
        • ep_recv_count_out：要求shape为 (ep_world_size * max(tp_world_size, 1) * local_expert_num, )。
        • performance_Info_optional：预留参数，当前版本不支持，传空指针即可。
        • ccl_buffer_size：调用get_low_latency_ccl_buffer_size接口(../../torch_extension/npu_ops_transformer/ops/deep_ep.py)。
      • 参数说明里shape格式说明：
        • H：表示hidden size隐藏层大小，取值范围[1024, 8192]。
        • BS：表示batch sequence size，即本卡最终输出的token数量，取值范围为[1, 512]。

    调用说明

    调用方式 样例代码 说明
    npu_low_latency_dispatch接口 deepep.py 通过npu_low_latency_dispatch接口方式调用moe_distribute_dispatch_v3算子。