CrossEntropy在矩阵比较大的情况下，效果不符合预期

[Yelrose]

bug描述 Describe the Bug

CrossEntropy Loss在输入矩阵比较大的情况下，loss计算有问题。下面是复现的代码，在logits是fp32或者fp16都有问题。fp32需要80G显存复现

import paddle
import numpy as np

paddle.seed(42)

batch_size = 1
vocab_size = 65536
# 有问题
max_seq_len =  65546 #32768
# 没问题
# max_seq_len =  65536 #32768

logits = paddle.rand(shape=[batch_size, max_seq_len, vocab_size], dtype="float16")
labels = paddle.randint(0, vocab_size, shape=[batch_size, max_seq_len, 1], dtype="int64")

loss_fn = paddle.nn.loss.CrossEntropyLoss(reduction="none")

loss = loss_fn(logits, labels)


seqlen = 10
"""计算所有样本的loss，print倒数两个样本loss"""
loss_1 = loss[:, -seqlen:]
print(f"loss_1 {loss_1}")

"""单独计算倒数两个样本的loss"""
loss_2 = loss_fn(logits[:, -seqlen:], labels[:, -seqlen:])
print(f"loss_2 {loss_2}")

其他补充信息 Additional Supplementary Information

No response

[risemeup1]

感谢提出问题，相关RD正在看这个问题，后续会回复你