RopeGradOperation

产品支持情况

产品	是否支持
Atlas A3 推理系列产品/Atlas A3 训练系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 训练系列产品	x
Atlas 推理系列产品	x
Atlas 200I/500 A2 推理产品	x

功能说明

旋转位置编码处理的反向。

训练算子，实现对于旋转位置编码（RoPE）的梯度反向计算，计算流程如下：

def RopeGradCalc(self, in_tensors):
    cos_list = [in_tensors[2][:x, :] for x in OP_PARAM['qSeqLen']]
    sin_list = [in_tensors[3][:x, :] for x in OP_PARAM['qSeqLen']]
    cos = torch.cat(cos_list, dim=0)
    sin = torch.cat(sin_list, dim=0)
    sin1 = sin[:,:64]
    sin2 = sin[:,64:]
    rohqgsin = torch.cat((sin2, -sin1), dim=-1)
    q_grad = torch.zeros_like(in_tensors[0])
    bs = int(in_tensors[0].shape[1] / 128)
    for i in range(bs):
        q_grad[:, i * 128:(i + 1) * 128] = in_tensors[0][:, i * 128:(i + 1) * 128] * (cos + rohqgsin)

    k_grad = torch.zeros_like(in_tensors[1])
    for i in range(bs):
        k_grad[:,i * 128:(i + 1) * 128] = in_tensors[1][:, i * 128:(i + 1) * 128] *(cos + rohqgsin)
        return [q_grad, k_grad]

定义

struct RopeGradParam {
    std::vector<int32_t> qSeqLen;
    uint8_t rsv[8] = {0};
};

参数列表

成员名称	类型	默认值	描述
qSeqLen	std::vector< int32_t >	-	存储unpad场景下每个batch实际seqlen的值。size不能为0。qSeqLen的值只能为[2,1,4,4]
rsv[8]	uint8_t	{0}	预留参数。

unpad场景：输入序列长度（seq length）是动态的场景下，如果对所有输入序列都按最大长度计算，存在大量冗余计算。unpad方案是decoder过程序列长度不再按照最大长度计算，而是根据实际的长度进行计算，减少计算量。

输入

参数	维度	数据类型	格式	描述
ropeQ_grad1	[nTokens, hiddenSize]	float16	ND	ropeQ_grad矩阵。
ropeQ_grad2	[nTokens, hiddenSize]	float16	ND	ropeQ_grad矩阵。
cos	[maxSeqLen, head_size]	float16	ND	cos矩阵，maxSeqLen为qSeqLen中的最大元素，head_size为128。
sin	[maxSeqLen, head_size]	float16	ND	sin矩阵。

输出

参数	维度	数据类型	格式	描述
q_grad	[nTokens, hiddenSize]	float16	ND	q_grad矩阵。
k_grad	[nTokens, hiddenSize]	float16	ND	k_grad矩阵。

约束说明

nTokens和qSeqLen需要满足以下关系：

$\text{[math]}$

父主题： atb/train_op_params.h