损失 | Twinkle

InfoNCE 损失

Mon, 01 Jan 0001 00:00:00 +0000

InfonceLoss 实现带批内负样本和可选跨 rank 聚合的对比学习损失，用于嵌入/检索模型训练。

使用方法

from twinkle.loss import InfonceLoss

loss_fn = InfonceLoss(
 temperature=0.1,
 use_batch=True, # 启用批内负样本
 hard_negatives=7, # 固定每样本负样本数
 mask_fake_negative=True, # 遮蔽假负样本
 fake_neg_margin=0.1, # 假负样本检测阈值
)

model.set_loss(loss_fn)

输入格式

每个样本按 锚点(1) + 正样本(1) + 负样本(n) 排列。inputs['labels'] 是一维掩码，1 标记每组的起始位置。

embeddings: [a0, p0, n0_1, n0_2, a1, p1, n1_1, n1_2, ...]
labels: [ 1, 0, 0, 0, 1, 0, 0, 0, ...]

参数

参数	类型	默认值	说明
`temperature`	float	0.1	相似度缩放因子
`use_batch`	bool	True	使用跨样本批内负样本
`hard_negatives`	int	None	固定每样本负样本数（截断/上采样）
`mask_fake_negative`	bool	False	遮蔽高于 positive + margin 的 logit
`fake_neg_margin`	float	0.1	假负样本遮蔽阈值
`include_qq`	bool	False	添加 query-query 相似度块
`include_dd`	bool	False	添加 doc-doc 相似度块

跨 Rank 聚合

当 use_batch=True 且分布式训练激活时，嵌入会从所有 DP rank 聚合以最大化批内负样本多样性。仅本地分片保留梯度。

相似度块

该损失支持三种相似度块，提供全面的对比学习信号：

Q→D（默认）：Query 到所有 Document — 主要对比信号
Q→Q（include_qq=True）：Query 到其他所有 Query — 防止 query 坍缩
D→D（include_dd=True）：Document 到其他所有 Document — Qwen3-Embedding 风格

示例：Embedding 训练

from twinkle.loss import InfonceLoss
from twinkle.metric import EmbeddingMetric

# 配置 Embedding 模型
model.set_loss(InfonceLoss(temperature=0.05, use_batch=True, include_qq=True))
model.set_metric(EmbeddingMetric(device_mesh=mesh, process_group=pg))

# 训练循环
for batch in dataloader:
 model.forward_backward(batch)
 model.clip_grad_and_step()

交叉熵

Mon, 01 Jan 0001 00:00:00 +0000

交叉熵是模型SFT和PT训练中最常用的一类损失。用于对labels的精确概率拟合。

class CrossEntropyLoss(Loss):

 def __init__(self, **kwargs):
 self.reduction = kwargs.get('reduction', 'mean')

 def __call__(self, inputs, outputs, **kwargs):
 import torch
 logits = outputs['logits'].view(-1, outputs['logits'].shape[-1])
 labels = inputs['labels'].view(-1)
 return torch.nn.CrossEntropyLoss(reduction=self.reduction)(logits, labels)

构造中可以传入reduction参数，支持sum, mean, none等（和torch.nn.CrossEntropyLoss输入相同）。

在Transformers模型中目前使用sum。目的是在optimizer.step之前统计有效token数量并在grad层面取单token平均。

分块交叉熵

Mon, 01 Jan 0001 00:00:00 +0000

交叉熵损失的内存优化变体，通过在词表维度上分块处理来减少 GPU 峰值内存使用。

from twinkle.loss import ChunkedCrossEntropyLoss

loss_fn = ChunkedCrossEntropyLoss(
 chunk_size=1024, # 词表分块大小
 reduction='mean',
)

model.set_loss(loss_fn)

参数:

chunk_size: 每块处理的词表 token 数量（默认: 1024）
reduction: 归约模式 — sum, mean, 或 none

实现使用自定义 autograd 函数，沿词表维度将 logit 到损失的计算分块进行。这避免了实例化完整的 [batch*seq_len, vocab_size] 概率张量，显著减少了大词表模型的内存占用。

当训练大词表模型时标准交叉熵导致 OOM 错误时非常有用。

DPO 损失

Mon, 01 Jan 0001 00:00:00 +0000

直接偏好优化（DPO）及其变体用于在不需要单独奖励模型的情况下将模型与人类偏好对齐。

DPOLoss

标准 DPO 损失，支持多种损失类型和可选的无参考模式。

from twinkle.loss import DPOLoss

loss_fn = DPOLoss(
 loss_type='sigmoid', # 'sigmoid', 'hinge', 'ipo', 'kto'
 beta=0.1,
 sft_weight=0.0, # 可选的 SFT 正则化权重
 reference_free=False,
)

model.set_loss(loss_fn)

参数:

loss_type: DPO 变体 — sigmoid（默认）, hinge, ipo, 或 kto
beta: 控制偏好强度的温度参数
sft_weight: chosen 响应上额外 SFT 损失的权重
reference_free: 为 True 时跳过参考模型的对数概率

损失函数期望批次中 chosen/rejected 样本交替排列。它计算序列级对数概率，优化策略使其偏好 chosen 而非 rejected 响应。

SimPOLoss

简化偏好优化，通过使用长度归一化的对数概率来消除对参考模型的需求。

from twinkle.loss import SimPOLoss

loss_fn = SimPOLoss(beta=2.0, gamma=1.0)

参数:

beta: logit 差异的缩放因子
gamma: 添加到偏好差距的 margin 项

CPOLoss

对比偏好优化，将偏好学习与行为克隆相结合。

from twinkle.loss import CPOLoss

loss_fn = CPOLoss(beta=0.1, cpo_alpha=1.0)

参数:

beta: 偏好损失的温度
cpo_alpha: chosen 响应上行为克隆（NLL）损失的权重

ORPOLoss

赔率比偏好优化，在单一损失中统一 SFT 和偏好对齐。

from twinkle.loss import ORPOLoss

loss_fn = ORPOLoss(beta=0.1)

该损失将 chosen 响应上的标准 NLL 项与对数赔率比惩罚相结合，推动模型远离 rejected 响应。

所有偏好损失都继承自 PreferenceLossBase 的共享工具方法，包括对数概率计算、chosen/rejected 拆分和序列级聚合。

GKD 损失

Mon, 01 Jan 0001 00:00:00 +0000

广义知识蒸馏（GKD）损失使用 Jensen-Shannon 散度将知识从教师模型蒸馏到学生模型。

from twinkle.loss import GKDLoss

loss_fn = GKDLoss(
 teacher_mode='full', # 'full', 'topk_local', 'topk_remote'
 beta=0.5, # JSD 的插值权重
 temperature=1.0,
)

model.set_loss(loss_fn)

参数:

teacher_mode: 获取教师 logits 的方式
- full: 来自本地教师模型的全词表 logits
- topk_local: 来自本地教师的 top-k logits，使用分块计算以节省内存
- topk_remote: 来自远程 API 教师的 top-k logits
beta: 学生和教师分布在 JSD 中的插值权重（0 = 纯学生，1 = 纯教师）
temperature: 学生和教师分布的 softmax 温度

GKD 损失内部实现了分块计算，以减少处理大词表时的峰值内存使用。

GKD 适用于训练模仿大型教师模型行为的小型学生模型，同时支持本地和远程教师设置。

GRPO 损失

Mon, 01 Jan 0001 00:00:00 +0000

组相对策略优化（GRPO）及其变体实现了带有 PPO 风格裁剪和 KL 正则化的策略梯度损失。

GRPOLoss

标准 GRPO 损失，带有重要性采样、PPO 裁剪和可选的 KL 惩罚。

from twinkle.loss import GRPOLoss

loss_fn = GRPOLoss(
 clip_range=0.2,
 beta=0.01, # KL 惩罚系数
)

model.set_loss(loss_fn)

参数:

clip_range: 重要性权重的 PPO 裁剪范围（默认: 0.2）
beta: KL 散度惩罚系数。设为 0 以禁用 KL 正则化

损失函数同时处理标准批次和打包序列（通过 position_ids 检测）。它计算每个 token 的重要性权重，应用 PPO 裁剪，并可选地添加针对参考策略的 KL 惩罚项。

变体

Twinkle 提供了多种 GRPO 变体:

GSPOLoss

序列级重要性采样变体，在序列级别而非 token 级别计算重要性权重。

from twinkle.loss import GSPOLoss
loss_fn = GSPOLoss(clip_range=0.2, beta=0.01)

SAPOLoss

软门控优势策略优化，在优势值上应用 sigmoid 门控来控制优化方向。

from twinkle.loss import SAPOLoss
loss_fn = SAPOLoss(clip_range=0.2, beta=0.01, tau=1.0)

CISPOLoss

裁剪重要性采样策略优化，在与优势值相乘之前对重要性权重进行显式裁剪。

from twinkle.loss import CISPOLoss
loss_fn = CISPOLoss(clip_range=0.2, beta=0.01)

BNPOLoss

批归一化策略优化，在聚合之前对批次内的每 token 损失进行归一化。

from twinkle.loss import BNPOLoss
loss_fn = BNPOLoss(clip_range=0.2, beta=0.01)

DRGRPOLoss

动态比率 GRPO，使用固定分母进行重要性权重计算。

from twinkle.loss import DRGRPOLoss
loss_fn = DRGRPOLoss(clip_range=0.2, beta=0.01)

所有 GRPO 变体共享相同的打包序列处理、对数概率对齐和 KL 惩罚计算基础流水线。它们的主要区别在于重要性权重和优势值的组合方式。

MSE 损失

Mon, 01 Jan 0001 00:00:00 +0000

均方误差损失，用于回归式训练任务。

from twinkle.loss import MSELoss

loss_fn = MSELoss()
model.set_loss(loss_fn)

MSELoss 计算模型输出 logits 与目标 labels 之间的均方误差。适用于奖励模型训练或价值函数估计等任务。

构建新的 Loss

Mon, 01 Jan 0001 00:00:00 +0000

Twinkle 中的 loss 基类定义为：

class Loss:

 def __call__(self, inputs: InputFeature, outputs: ModelOutput, **kwargs):
 ...

损失的输入为模型的 InputFeature，输出为模型标准 ModelOutput，kwargs 可以在模型的 calculate_loss 中传入。由于它是一个带有 __call__ 方法的类，因此开发者也可以使用 Callable：

def my_loss(inputs: InputFeature, outputs: ModelOutput, extra_data1: int, extra_data2: dict):
 ...
 return loss

在模型中这样使用：

model.set_loss(my_loss)
model.calculate_loss(extra_data1=10, extra_data2={})

你也可以将 Loss 上传到 ModelScope/Hugging Face 的 Hub 中，在使用时动态拉取：

model.set_loss('ms://my_group/my_loss')

具体可以参考插件文档的介绍。