ScoreNormalize
产品支持情况
| 产品 | 是否支持 |
|---|---|
| Atlas A2 训练系列产品 | 是 |
功能说明
- 算子功能:MoE路由专家得分进行归一化。
- 计算公式:
x=x∑i=1kxi×2.5x = \frac{x}{\sum_{i=1}^{k} x_i} \times 2.5
- 这里xx是一个二维张量,表示路由专家的分数,shape为[rowsrows, kk],每行表示一个token选择的topk个专家的分数。计算公式对每个token选择的所有专家分数进行归一化,除以所有分数的总和再乘2.5。
参数说明
| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 |
|---|---|---|---|---|
| x | 输入/输出 | 公式中的输入/输出张量x,shape为 (rows, k) | BFLOAT16 | ND |
| rows | 输入 | 输入/输出x的第一维 | INT | - |
| k | 输入 | 输入/输出x的第二维 | INT | - |
约束说明
- x 仅支持BFLOAT16类型。
- k <= 16。
调用说明
torch.ops.npu_ops_transformer_ext.score_normalize(x, rows, k)