算子功能:完成路由专家AlltoAllv、量化GroupedMatMul融合并实现与共享专家量化MatMul并行融合,先通信后计算。
计算公式: 假设通信域中的总卡数为epWorldSize,每张卡上通信后路由专家个数为e,每张卡分组矩阵乘只负责本卡专家的计算。对于每张卡的计算公式如下:
本卡共享专家分组矩阵乘计算
[object Object]Alltoallv通信和permute
[object Object]本卡路由专家按专家维度分组矩阵乘计算
[object Object]
每个算子分为两段式接口,必须先调用[object Object]接口获取入参并根据计算流程计算所需workspace大小,再调用[object Object]接口执行计算。
[object Object]
[object Object]
参数说明
[object Object]gmmXQuantMode、gmmWeightQuantMode、mmXQuantMode、mmWeightQuantMode的枚举值跟量化模式关系如下:
- 0: 非量化
- 1: pertensor
- 2: perchannel
- 3: pertoken
- 4: pergroup
- 5: perblock
- 6: mx量化
- 7: pertoken动态量化
返回值
[object Object]
确定性计算:
[object Object]默认确定性实现。
参数说明里shape使用的变量:
- BSK:本卡发送的token数,是sendCounts参数累加之和,取值范围(0, 52428800)。
- H1:表示路由专家hidden size隐藏层大小,取值范围(0, 65536)。
- H2:表示共享专家hidden size隐藏层大小,取值范围(0, 12288]。
- e:表示单卡上专家个数,取值范围(0, 32],e * epWorldSize最大支持256。
- N1:表示路由专家的head_num,取值范围(0, 65536)。
- N2:表示共享专家的head_num,取值范围(0, 65536)。
- BS:batch sequence size。
- K:表示选取TopK个专家,K的范围[2, 8]。
- A:本卡收到的token数,是recvCounts参数累加之和。
- ep通信域内所有卡的 A 参数的累加和等于所有卡上的 BSK 参数的累加和。
- mx量化且gmmX与gmmWeight为FLOAT4_E2M1时,H1和H2必须为偶数且不能为2,同时transGmmWeight和transMmWeight为false情况下,N1和N2必须为偶数。
- groupSize:
- 仅当gmmXScale/gmmWeightScale/mmXScale/mmWeightScale输入都是2维及以上数据时,groupSize取值有效,其他场景需传入0。
- groupSize值支持公式推导:传入的groupSize内部会按如下公式分解得到groupSizeM、groupSizeN、groupSizeK,当其中有1个或多个为0,会根据gmmX/gmmWeight/mmX/mmWeight/gmmXScale/gmmWeightScale/mmXScale/mmWeightScale输入shape重新设置groupSizeM、groupSizeN、groupSizeK用于计算。设置原理:如果groupSizeM=0,表示m方向量化分组值由接口推导,推导公式为groupSizeM = m / scaleM(需保证m能被scaleM整除),其中m与gmmX/mmX shape中的m一致,scaleM与gmmXScale/mmXScale shape中的m一致;如果groupSizeK=0,表示k方向量化分组值由接口推导,推导公式为groupSizeK = k / scaleK(需保证k能被scaleK整除),其中k与gmmX/mmX shape中的k一致,scaleK与gmmXScale/mmXScale shape中的k一致;如果groupSizeN=0,表示n方向量化分组值由接口推导,推导公式为groupSizeN = n / scaleN(需保证n能被scaleN整除),其中n与gmmWeight/mmWeight shape中的n一致,scaleN与gmmWeightScale/mmWeightScale shape中的n一致。
- 如果满足重新设置条件,当gmmXScale/gmmWeightScale/mmXScale/mmWeightScale输入是2维及以上时,且数据类型都为FLOAT8_E8M0时,[groupSizeM,groupSizeN,groupSizeK]取值组合会推导为[1, 1, 32],对应groupSize的值为4295032864。
量化参数约束:
- 当前版本支持pertensor量化、mx量化。
注意:由于量化接口仅支持Atlas 350 加速卡,以下示例基于该系列实现。本示例代码以2卡为例,请根据实际环境卡数修改 [object Object]。
[object Object]