RopeGradOperation

功能

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

算子功能实现描述（Python）

训练算子，实现对于旋转位置编码（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;
};

参数列表

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

输入

输出

规格约束

• 当前只支持Atlas 800I A2 推理产品/Atlas A2 训练系列产品。
• nTokens和qSeqLen需要满足以下关系：