RopeGradOperation

Description

Implements backward rotary position embedding.

Operator Function Implementation (Python)

Train the operator to implement gradient backward computation for RoPE. The computation process is as follows:

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]

Definition

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

Parameters

Unpad scenario: When the input sequence length (seq length) is dynamic, if all input sequences are calculated based on the maximum length, a large number of redundant calculations exist. In the unpad solution, the sequence length is calculated based on the actual length instead of the maximum length in the decoder process, reducing the calculation workload.

Input

Output

Restrictions

• Currently, only the Atlas A2 inference products is supported.
•