[hushenwei2000]

本 PR 主要解决两个问题：

对 None 结果的处理

新增了对 None 结果的处理。动机来源于 fused_rotary_position_embedding 可接受 Q、K、V 三个变量，后两个可以为 None。但无论是否有 None，Padlle 都将输出一个三个元素的 Tuple；如果某个变量传入为 None，那么该位置的输出为 Tensor(Not Initialized)。但目前框架未对 None 变量输出进行处理，导致在以下处抛出异常：

# 截取自 tester/base.py : gen_torch_output_and_output_grad()
elif isinstance(outputs, tuple):
            for output in outputs:
                if isinstance(output, torch.Tensor):
                    result_outputs.append(output)
                else:
                    raise ValueError("outputs format not support")

。而在外层会接收该异常并不再继续进行 Torch 的 grad 操作，并只检查前向结果。这导致了原有测试有大量只通过了前向即输出 Pass 的测试。

因此本 PR 增加了对 None 结果的处理：

• 在 accuracy.py 前向结果比较时， None 参与比较。
• 在 base.py 计算梯度时，None 结果不参与梯度计算。

fused_rotary_position_embedding API 精度问题

统计发现所有测试用例的前向结果均能通过测试，而反向结果有错。经过源码分析发现 Paddle 实现正确，而 Torch 的反向结果有误。评论区将详细阐述。

本 PR 进行tester/api_config/5_accuracy/accuracy_gpu_error_uncertain.txt 中的 60 个测试：

• 20 个完全通过
• 20 个 Accuracy 错误（都是反向错，实际上是 Torch 计算错误）
• 20 个 Torch 错误

[paddle-bot]

Thanks for your contribution!

[hushenwei2000]

以下给出了一个最小可复现代码，展示了 Paddle 和 Torch 在计算旋转位置编码反向时的不同。

q 为输入 Tensor。

正向计算：out = q * cos + rotate(q) * sin 。Paddle 和 Torch 一致。

反向计算的 Paddle 版本和正向计算一致，除了正负号相反。而 Torch 反向计算逻辑为 out = q * cos + rotate(q * sin)。从数学推导上来说 Paddle 应该为正确版本。

import torch
import paddle
import numpy
import torch.nn.functional as F
import os
from typing import Optional

device = torch.device("cuda:0")
torch.set_default_device(device)

def init_input(numpy_tensor):
    paddle_x = paddle.to_tensor(numpy_tensor)
    torch_x = torch.tensor(numpy_tensor, requires_grad = True if numpy_tensor.dtype in [numpy.float32, numpy.float64] else False)
    if numpy_tensor.dtype in [numpy.float32, numpy.float64]:
        paddle_x.stop_gradient = False

    numpy.testing.assert_allclose(
        paddle_x.numpy(),
        torch_x.cpu().detach().numpy(),
        1e-10,
        1e-10,
        err_msg='intput diff'
    )
    return paddle_x, torch_x


def fused_rotary_position_embedding(
    q: torch.Tensor,
    k: Optional[torch.Tensor] = None,
    v: Optional[torch.Tensor] = None,
    sin: Optional[torch.Tensor] = None,
    cos: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.Tensor] = None,
    use_neox_rotary_style: bool = True,
    time_major: bool = False,
    rotary_emb_base: float = 10000.0,
) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
    def _deal_qkv_pytorch(init_value: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
        if init_value is None:
            return None
        return init_value.permute(0, 2, 1, 3)

    def _mult_qkv_pytorch(value: Optional[torch.Tensor], cos_tensor: torch.Tensor, sin_tensor: torch.Tensor) -> Optional[torch.Tensor]:
        if value is None:
            return None
        rotate_half_q = torch.stack([-value[..., 1::2], value[..., 0::2]], dim=-1).reshape(value.shape)
        query = value * cos_tensor + rotate_half_q * sin_tensor
        return query

    def _mult_qkv_rotate_half_pytorch(value: Optional[torch.Tensor], cos_tensor: torch.Tensor, sin_tensor: torch.Tensor) -> Optional[torch.Tensor]:
        if value is None:
            return None
        half_dim = value.shape[-1] // 2

        rotate_half_q = torch.cat([-value[..., half_dim:], value[..., :half_dim]], dim=-1)
        query = value * cos_tensor + rotate_half_q * sin_tensor
        return query

    def _get_sin_cos_tensor_pytorch(seq_len: int, head_dim: int, sign: int = 1, rotate_half: bool = False):
        pos_seq = torch.arange(0, seq_len, 1, dtype=torch.float32)
        indices = torch.arange(0, head_dim, 2, dtype=torch.float32)
        indices = 1 / 10000 ** (indices / head_dim)
        sinusoid_inp = pos_seq.unsqueeze(1) * indices.unsqueeze(0)
        sin_sin = np.empty((seq_len * head_dim), dtype=np.float32)
        cos_cos = np.empty((seq_len * head_dim), dtype=np.float32)
        numpy_array = sinusoid_inp.numpy()
        iter_array = np.nditer(numpy_array)
        i = 0
        if rotate_half:
            stride = head_dim // 2
            for value in iter_array:
                sin_sin[i] = sign * np.sin(value)
                cos_cos[i] = np.cos(value)
                sin_sin[i + stride] = np.sin(value)
                cos_cos[i + stride] = np.cos(value)
                i += 1
                if i % head_dim == stride:
                    i += stride
        else:
            for value in iter_array:
                sin_sin[i * 2] = sign * np.sin(value)
                cos_cos[i * 2 + 0] = np.cos(value)
                sin_sin[i * 2 + 1] = np.sin(value)
                cos_cos[i * 2 + 1] = np.cos(value)
                i += 1
        tensor_sin = torch.from_numpy(sin_sin).reshape([1, seq_len, 1, head_dim])
        tensor_cos = torch.from_numpy(cos_cos).reshape([1, seq_len, 1, head_dim])
        return tensor_sin, tensor_cos
    
    init_q, init_k, init_v = q, k, v
    if time_major:
        init_q = init_q.permute(1, 0, 2, 3)
        if init_k is not None:
            init_k = init_k.permute(1, 0, 2, 3)
        if init_v is not None:
            init_v = init_v.permute(1, 0, 2, 3)

    head_dim = init_q.shape[3]
    seq_len = init_q.shape[1]
    
    sin_tensor, cos_tensor = sin, cos
    if sin_tensor is None or cos_tensor is None:
        sin_tensor, cos_tensor = _get_sin_cos_tensor_pytorch(seq_len, head_dim, rotate_half=not use_neox_rotary_style)
        sin_tensor = sin_tensor.to(q.device)
        cos_tensor = cos_tensor.to(q.device)
    
    q_rope = init_q
    k_rope = init_k
    v_rope = init_v

    if position_ids is not None:
        sin_tensor = sin_tensor.squeeze((0, 2))[position_ids].unsqueeze(2)
        cos_tensor = cos_tensor.squeeze((0, 2))[position_ids].unsqueeze(2)
    
    if use_neox_rotary_style:
        query = _mult_qkv_pytorch(q_rope, cos_tensor, sin_tensor)
        value = _mult_qkv_pytorch(v_rope, cos_tensor, sin_tensor)
        key = _mult_qkv_pytorch(k_rope, cos_tensor, sin_tensor)
    else:
        query = _mult_qkv_rotate_half_pytorch(q_rope, cos_tensor, sin_tensor)
        value = _mult_qkv_rotate_half_pytorch(v_rope, cos_tensor, sin_tensor)
        key = _mult_qkv_rotate_half_pytorch(k_rope, cos_tensor, sin_tensor)

    r_query = query
    r_key = key
    r_value = value

    if time_major:
        r_query = r_query.permute(1, 0, 2, 3)
        if r_key is not None:
            r_key = r_key.permute(1, 0, 2, 3)
        if r_value is not None:
            r_value = r_value.permute(1, 0, 2, 3)
            
    return r_query, r_key, r_value


q_shape = [1, 4, 2, 4]
sin_cos_shape = [1, 4, 1, 4]
dtype = "float32"

numpy_tensor_q =  (numpy.random.random(q_shape) * 2 - 1).astype(dtype)
paddle_q, torch_q = init_input(numpy_tensor_q)

numpy_tensor_sin = (numpy.random.random(sin_cos_shape) * 2 - 1).astype(dtype)
paddle_sin, torch_sin = init_input(numpy_tensor_sin)

numpy_tensor_cos = (numpy.random.random(sin_cos_shape) * 2 - 1).astype(dtype)
paddle_cos, torch_cos = init_input(numpy_tensor_cos)


# API test

torch_out = fused_rotary_position_embedding(torch_q, None, None, torch_sin, torch_cos, None, False, False)

paddle_out = paddle.incubate.nn.functional.fused_rotary_position_embedding(paddle_q, None, None, paddle_sin, paddle_cos, position_ids=None, use_neox_rotary_style=False, time_major=False)

rtol=1e-2
atol=1e-2

paddle_out_q, paddle_out_k, paddle_out_v = paddle_out
torch_out_q, torch_out_k, torch_out_v = torch_out

numpy.testing.assert_allclose(
    paddle_out_q.numpy(),
    torch_out_q.cpu().detach().numpy(),
    rtol,
    atol,
    equal_nan = True
)

print("q forward pass")


numpy_tensor = (numpy.random.random(q_shape) * 2 - 1).astype("float32")
paddle_grad_q, torch_grad_q = init_input(numpy_tensor)

numpy_tensor = (numpy.random.random(sin_cos_shape) * 2 - 1).astype("float32")
paddle_grad_sin, torch_grad_sin = init_input(numpy_tensor)

numpy_tensor = (numpy.random.random(sin_cos_shape) * 2 - 1).astype("float32")
paddle_grad_cos, torch_grad_cos = init_input(numpy_tensor)

# Two calculation methods

# First: the same calculation method as forward. OUT = q * cos + rotate(q) * sin
half_dim = torch_grad_q.shape[-1] // 2
rotate_torch_grad_q = torch.cat([torch_grad_q[..., half_dim:], -torch_grad_q[..., :half_dim]], dim=-1)
paddle_style_out_grad_q = torch_grad_q * torch_cos + rotate_torch_grad_q * torch_sin
print("paddle_style_out_grad_q : ", paddle_style_out_grad_q)

# Second: OUT = q * cos + rotate(q * sin)
sin_v = torch_grad_q * torch_sin
rotated_sin_v = torch.cat([sin_v[..., half_dim:], -sin_v[..., :half_dim]], dim=-1)
cos_v = torch_grad_q * torch_cos
torch_style_out_grad_q = cos_v + rotated_sin_v
print("torch_style_out_grad_q : ", torch_style_out_grad_q)

# Run real grad
torch_grad_q_out = torch.autograd.grad([torch_out_q], [torch_q, torch_sin, torch_cos], grad_outputs=[torch_grad_q], allow_unused=True, retain_graph=True)

paddle_grad_q_out = paddle.grad([paddle_out_q], [paddle_q, paddle_sin, paddle_cos], grad_outputs=[paddle_grad_q], allow_unused=True, retain_graph=True)

torch_grad_q_out = torch_grad_q_out[0]
paddle_grad_q_out = paddle_grad_q_out[0]

if numpy.allclose(
    paddle_grad_q_out.numpy(),
    torch_grad_q_out.cpu().detach().numpy(),
    rtol,
    atol,
):
    print("q backward pass")
else:
    print("q backward pass failed!")
    print("paddle_grad_q", paddle_grad_q_out.numpy())
    print("------")
    print("torch_grad_q", torch_grad_q_out.cpu().detach().numpy())

print("")

if numpy.allclose(
    torch_style_out_grad_q.cpu().detach().numpy(),
    torch_grad_q_out.cpu().detach().numpy(),
    rtol,
    atol,
):
    print("torch_style_out_grad_q and torch_grad_q_out are the same")

if numpy.allclose(
    paddle_style_out_grad_q.cpu().detach().numpy(),
    paddle_grad_q_out.numpy(),
    rtol,
    atol,
):
    print("paddle_style_out_grad_q and paddle_grad_q_out are the same")

[cangtianhuang]

@LCStayingdullCircuit 同学也对 fused_rotary_position_embedding 的 Rule 进行了修改工作，麻烦邹同学也看看这个 pr～

相关工作：

• #416

此外，这份代码中对 np 对使用不正确，我正在进行修改😉：

• #458

[cangtianhuang]

能否修改 gen_torch_output_and_output_grad , 允许 None 值进入后续比较，这样做的可能好处有：

• 支持列表长度的比较，比如 paddle 与 torch 的输出都有三个，其中后两个是 None；
• 不同顺序的比较可以报错，比如 [1, None, 2] 与 [1, 2, None] 的比较

[hushenwei2000]

尝试过如果在 gen_torch_output_and_output_grad 函数里允许 None 值进入后续处理，有多处例如 "xx.shape"、"xx.dtype" 需要一并修改以处理 None，修改会比较多。

[hushenwei2000]

更新了对 None 结果的处理方式。在 accuracy.py 前向结果比较时， None 参与比较。在 base.py 计算梯度时，None 结果不参与梯度计算。

[cangtianhuang]

Suggested change

	for output in outputs:
	if (output is not None) and (not isinstance(output, paddle.Tensor) or output.size != 0):
	result_outputs.append(output)
	result_outputs = [output for output in outputs
	if output is not None and isinstance(output, paddle.Tensor) and output.size != 0]

[cangtianhuang]

为什么需要检查 paddle_out.size == 0 呢？

[cangtianhuang]

这个是否会影响 0size 的测试？

[cangtianhuang]

此处应该是冗余的

[cangtianhuang]

同冗余

[cangtianhuang]

Suggested change

	for output in outputs:
	if output is not None:
	result_outputs.append(output)
	result_outputs = [output for output in outputs if output is not None]

[cangtianhuang]

同冗余