ForwardKLLoss

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

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

参数:

ignore_index (int) – 指定要忽略且不计入输入梯度的目标值。损失值除以未被忽略的目标数量。默认值：-100。

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

参数:

• 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) – 是否按未遮掩的元素数量归一化损失值。

返回值:

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

返回类型:

torch.Tensor