ForwardKLLoss

class torchtune.modules.loss.ForwardKLLoss(ignore_index: int = - 100)

有效索引的Kullback-Leibler散度损失。 实现于 https://github.com/jongwooko/distillm/blob/17c0f98bc263b1861a02d5df578c84aea652ee65/distillm/losses.py

Parameters:

ignore_index (int) – 指定一个被忽略的目标值，该值不会对输入梯度产生影响。 损失会在非忽略的目标上进行平均。 默认值：-100。

forward(student_logits: Tensor, teacher_logits: Tensor, labels: Tensor, normalize: bool = True) → Tensor

Parameters:

• student_logits (torch.Tensor) – 来自学生模型的logits，形状为 (batch_size*num_tokens, vocab_size)。

• teacher_logits (torch.Tensor) – 来自教师模型的logits，形状为 (batch_size*num_tokens, vocab_size)。

• labels (torch.Tensor) – 形状为 (batch_size, vocab_size) 的真实标签。

• normalize (bool) – 是否通过未屏蔽元素的数量来归一化损失。

Returns:

形状为 (1,) 的KL散度损失。

Return type:

torch.Tensor