给定 n 个专家,平均分布在 m 张 NPU 上(每张卡上一组,组内专家编号连续)。算法分三步:
按组取代表 组大小为 g=n/m。对每一组,找到组内最大概率以及对应的专家编号,作为该组代表值。
选路由目标 NPU(选 p 个组)
将所有组按“代表概率”从大到小排序,取前 p 个组(对应的 NPU)。若 p>m,直接输出 error
。
MOE 模型训练时,token 根据概率发送到 topk 个不同的专家进行计算。这些专家分布在多个 NPU 卡上。Device−Limitedr outing 算法将 token 的路由目标限制在 P 个 NPU 上,可以有效降低通信成本。具体的:
把 n 个专家平均分配在 m 个 NPU 上,每个 NPU 上的专家为一个组;设 n 个专家的编号为 N=[0,1,2,…,n−1] ,同一个专家组上的专家编号是连续的;
每个专家对应一个概率,表示被路由到的可能性;用每个组中的最大概率作为本组代表,从所有组中选择概率最大的 p 个组,其所在的 NPU 即为路由目标限制 NPU ;
再从上述 p 个 NPU 对应的所有专家概率中选择 k 个最大的概率对应的专家编号作为最终路由目标。
试着编写一段程序,实现以上路由算法。
第一行有 4 个处于区间 [1,10000] 之内的整数,第 1 个表示专家的个数 n ,第 2 个表示 NPU 个数 m ,第 3 个表示路由目标限制 NPU 个数 p ,第 4 个表示目标路由专家个数 k ;
第二行有 n 个处于区间 (0,1) 之内的浮点数,表示每个专家对应的概率值,这 n 个数对应的专家的编号为 [0,1,2,...,n−1] ;
如果,n 不能被 m 整除或者获取不到 k 个专家编号,输出 error ;
否则,按照从小到大的顺序,输出 k 个专家编号,任意相邻两数之间有空格,最后一个数字(行尾没有空格)
输入
8 4 4 2
0.5 0.01 0.09 0.023 0.027 0.05 0.1 0.2
输出
0 7
说明
将专家分成 4 组,分别为:(1)0.5 0.01 (2)0.09 0.023 (3)0.027 0.05 (4)0.1 0.2
限定专家为 4 ,则 4 组都被选定,从选定的 4 组中,选择 2 个专家,分别是 0.5 和 0.2 对应的专家,对应的编号分别是 0 和 7 ,按照升序排个列,输出: 0 7
输入
8 4 5 2
0.3 0.04 0.06 0.45 0.05 0.01 0.03 0.06
输出
error
说明
NPU 一共只有 4 个,需要限定 5 个 NPU ,不满足条件,输出 error