From f5bb852fdefdff77bb039ece3cf4c7ad4d7dcf12 Mon Sep 17 00:00:00 2001 From: iomgaa Date: Sat, 18 Jul 2026 05:02:41 -0400 Subject: [PATCH] =?UTF-8?q?=E5=B1=821/T4:=20trainer.py=20=E6=8E=A9?= =?UTF-8?q?=E7=A0=81=20SFT=20=E6=8D=9F=E5=A4=B1=20+=20=E6=9C=80=E5=B0=8F?= =?UTF-8?q?=20HF=20Trainer=20=E5=AD=90=E7=B1=BB=EF=BC=88docs/02=20=C2=A72.?= =?UTF-8?q?4=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - sft_loss 纯张量函数：batch-min prompt_length、移位切片、labels 重掩码； 空 batch/无 prompt 行显式报错（参考实现静默归零，差异已标注） - SFTTrainer 只重写 compute_loss 与 log；不向模型传 labels（防内部损失 绕过重掩码）；日志并入每步有效 token 数供 sanity 监控 - tests/test_trainer.py: 7 个单测，含手算对拍、移位对齐、窗口外无关性、 漏切重掩码 Co-Authored-By: Claude Fable 5 --- ars_opd/trainer.py | 150 ++++++++++++++++++++++++++++++++++++++++++ tests/test_trainer.py | 118 +++++++++++++++++++++++++++++++++ 2 files changed, 268 insertions(+) create mode 100644 ars_opd/trainer.py create mode 100644 tests/test_trainer.py diff --git a/ars_opd/trainer.py b/ars_opd/trainer.py new file mode 100644 index 0000000..d209156 --- /dev/null +++ b/ars_opd/trainer.py @@ -0,0 +1,150 @@ +"""训练编排（IO 边缘）。层 1 形态：掩码 SFT 损失 + 最小 HF Trainer 子类。 + +论文锚点：§3.1 式(1) 的标准交叉熵 SFT（监督目标是 teacher rollout，见 docs/02 §1）。 +层 5 会在此模块长出式(8) 的 chunk 蒸馏损失与 KL 锚定；届时式(1) 路径保留为基线。 + +结构：损失算法收在纯张量函数 `sft_loss`（可用 toy 张量在 CPU 单测）， +`SFTTrainer` 只做接线——模型前向、调 `sft_loss`、把指标并进日志， +其余一切（优化器、调度、DDP、checkpoint 保存）原样继承 HF Trainer。 +""" + +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn.functional as F +from transformers import Trainer + +from ars_opd.data import IGNORE_INDEX + +# --------------------------------------------------------------------------- +# 损失算法（纯张量函数，对拍 distillation_trainer.py:751-759 + 2776-2874） +# --------------------------------------------------------------------------- + + +def compute_prompt_length(attention_mask: torch.Tensor, labels: torch.Tensor) -> int: + """batch 级 prompt 边界 = batch 内最短的"有效长度 - completion 长度"。 + + 参数（B = batch，T = padding 后长度）： + - attention_mask: (B, T)，左 padding 位置为 0 + - labels: (B, T)，completion 位置为 token id，其余为 -100 + + 返回标量 pl。非显然约束：取 batch **最小值**是为了不切掉任何行的 + completion token——左 padding 下所有序列右对齐，第 r 行的 completion 起点 + 索引是 T - comp_r ≥ prompt_r ≥ min，故切片 [pl:] 必然包含全部 completion； + 代价是长 prompt 行会漏进一些 prompt token，由 sft_loss 用 labels 重掩码兜住。 + """ + full_lengths = attention_mask.sum(dim=1) # (B,) 每行非 padding 的 token 数 + completion_lengths = (labels != IGNORE_INDEX).sum(dim=1) # (B,) + return int((full_lengths - completion_lengths).min().item()) + + +def sft_loss( + logits: torch.Tensor, + input_ids: torch.Tensor, + labels: torch.Tensor, + attention_mask: torch.Tensor, +) -> tuple[torch.Tensor, int]: + """式(1)：completion 位置上的移位交叉熵。 + + 参数（B = batch，T = padding 后长度，V = 词表大小）： + - logits: (B, T, V) 模型对全序列的输出 + - input_ids / labels / attention_mask: (B, T)，SFTCollator 的产物 + + 返回 (标量 loss, 本 batch 有效 completion token 数)。 + + 切片几何（docs/02 §2.4 的四行，此处为权威实现）： + 位置 t-1 的 logit 预测位置 t 的 token，故 logits 取 [pl-1, T-1) 、 + targets 取 [pl, T)，两段长度同为 T-pl，逐位对齐。 + """ + pl = compute_prompt_length(attention_mask, labels) + if pl < 1: + # 需要位置 pl-1 的 logit 存在；pl=0 意味着某行完全没有 prompt token， + # 数据管线出了问题（collator 保证 prompt 至少含模板 token） + raise ValueError(f"prompt_length={pl} < 1，存在无 prompt 的数据行") + + shifted_logits = logits[:, pl - 1 : -1, :] # (B, T, V) -> (B, T-pl, V) + targets = input_ids[:, pl:].clone() # (B, T-pl)；clone: 下面要原地改写 + + # 重掩码：切片里漏进的 prompt token（长 prompt 行）与 padding 全部置 -100。 + # labels 是权威掩码，切片几何只是省算力——正确性完全由这一步保证 + invalid = labels[:, pl:] == IGNORE_INDEX # (B, T-pl) + targets[invalid] = IGNORE_INDEX + + num_valid = int((~invalid).sum().item()) + if num_valid == 0: + # 差异标注：参考实现(trainer:2810-2812)对非有限 loss 返回零梯度标量静默 + # 继续；我们显式报错——collator 已挡下 prompt-only 行，走到这里仍全被 + # 掩码只可能是数据坏了（如 teacher 返回空解答），必须暴露而非跳过 + raise ValueError( + "本 batch 没有任何有效 completion token（全被 -100 掩码）。" + "检查 teacher 解答是否为空、completion 预算是否被截光。" + ) + + vocab = shifted_logits.shape[-1] + loss = F.cross_entropy( + shifted_logits.reshape(-1, vocab), # (B*(T-pl), V) + targets.reshape(-1), # (B*(T-pl),) + ignore_index=IGNORE_INDEX, + ) + return loss, num_valid + + +# --------------------------------------------------------------------------- +# Trainer 接线 +# --------------------------------------------------------------------------- + + +class SFTTrainer(Trainer): + """最小 SFT Trainer：只重写 compute_loss 与 log，其余全部继承 HF Trainer。 + + 用法（见 scripts/ 训练脚本）：与 HF Trainer 完全同参构造， + data_collator 传 SFTCollator，train_dataset 传 load_sft_dataset 的产物。 + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + # 关 dropout（参考实现同款，docs/02 §3 保留清单）：层 2+ 的蒸馏要求 + # student/ref 两次前向可比，前向必须确定；Qwen3 默认无 dropout， + # 此处是无害的统一前置 + for module in self.model.modules(): + if isinstance(module, torch.nn.Dropout): + module.p = 0.0 + self._token_counts: list[int] = [] + + def compute_loss( + self, + model: Any, + inputs: dict[str, torch.Tensor], + return_outputs: bool = False, + num_items_in_batch: int | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, Any]: + # 非显然约束：不把 labels 传给模型前向。HF 模型收到 labels 会自己算 + # "全序列移位 CE"并放进 outputs.loss，那会绕过 batch-min 切片与重掩码； + # 损失的权威实现只能有 sft_loss 一处 + outputs = model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + ) + loss, num_tokens = sft_loss( + outputs.logits, + inputs["input_ids"], + inputs["labels"], + inputs["attention_mask"], + ) + self._token_counts.append(num_tokens) + return (loss, outputs) if return_outputs else loss + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + """在 HF 的常规日志里并入每步有效 token 数均值。 + + sanity run 时盯这个数：若它远小于预期（≈batch 内解答总长），说明掩码 + 把 completion 也吞了——loss 曲线看不出这种错，token 数看得出。 + """ + if self._token_counts: + logs["sft/num_tokens_per_step"] = sum(self._token_counts) / len( + self._token_counts + ) + self._token_counts = [] + super().log(logs, start_time) diff --git a/tests/test_trainer.py b/tests/test_trainer.py new file mode 100644 index 0000000..2074515 --- /dev/null +++ b/tests/test_trainer.py @@ -0,0 +1,118 @@ +"""层 1 / T4：掩码 SFT 损失单测（docs/02 §2.4 的切片几何与重掩码）。 + +只测纯张量函数 sft_loss / compute_prompt_length；SFTTrainer 类是 HF 接线， +由远程 sanity run 验证。张量全部手工构造，期望值可手算。 +""" + +import pytest +import torch +import torch.nn.functional as F + +from ars_opd.data import IGNORE_INDEX +from ars_opd.trainer import compute_prompt_length, sft_loss + +V = 7 # 玩具词表大小 + + +def onehot_logits(target_ids, scale=10.0): + """构造在 target_ids 处放尖峰的 logits。(T,) -> (T, V)""" + t = torch.tensor(target_ids) + return F.one_hot(t, V).float() * scale + + +def batch_of_one(input_ids, labels, attention=None): + """单行 batch 的三件套，logits 另配。""" + ids = torch.tensor([input_ids]) + lab = torch.tensor([labels]) + att = torch.ones_like(ids) if attention is None else torch.tensor([attention]) + return ids, lab, att + + +def test_prompt_length_取batch最小且不数padding(): + # pad p p c c c p p p c c c + labels = torch.tensor( + [ + [IGNORE_INDEX] * 3 + [5, 6, 5], + [IGNORE_INDEX] * 3 + [6, 5, 6], + ] + ) + attention = torch.tensor([[0, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1]]) + # 行 0 有效长 5、completion 3 → prompt 2；行 1 是 6-3=3；batch 最小 = 2 + assert compute_prompt_length(attention, labels) == 2 + + +def test_损失与手算交叉熵一致(): + ids, lab, att = batch_of_one([3, 4, 5, 6], [IGNORE_INDEX, IGNORE_INDEX, 5, 6]) + torch.manual_seed(0) + logits = torch.randn(1, 4, V) + + loss, num = sft_loss(logits, ids, lab, att) + + # pl=2：位置 1、2 的 logit 分别预测位置 2、3 的 token（5 和 6） + expected = F.cross_entropy(logits[0, 1:3], torch.tensor([5, 6])) + assert torch.allclose(loss, expected) + assert num == 2 + + +def test_移位对齐_预测下一个token而非当前(): + ids, lab, att = batch_of_one([3, 4, 5, 6], [IGNORE_INDEX, IGNORE_INDEX, 5, 6]) + + # 位置 t 的尖峰指向位置 t+1 的 token（正确的"预测下一个"）→ loss ≈ 0 + next_logits = onehot_logits([4, 5, 6, 0]).unsqueeze(0) + loss_next, _ = sft_loss(next_logits, ids, lab, att) + # 位置 t 的尖峰指向位置 t 自己的 token（错误的"复读当前"）→ loss 大 + self_logits = onehot_logits([3, 4, 5, 6]).unsqueeze(0) + loss_self, _ = sft_loss(self_logits, ids, lab, att) + + assert loss_next.item() < 0.01 + assert loss_self.item() > 5.0 + + +def test_窗口外的logits不影响损失(): + ids, lab, att = batch_of_one([3, 4, 5, 6], [IGNORE_INDEX, IGNORE_INDEX, 5, 6]) + torch.manual_seed(0) + logits = torch.randn(1, 4, V) + loss_base, _ = sft_loss(logits, ids, lab, att) + + # pl=2 → 用到的窗口是位置 [1, 3)；位置 0（prompt 内部）和 3（末位）不参与 + perturbed = logits.clone() + perturbed[0, 0] += 100.0 + perturbed[0, 3] -= 100.0 + loss_pert, _ = sft_loss(perturbed, ids, lab, att) + assert torch.allclose(loss_base, loss_pert) + + +def test_batch_min切片漏进的prompt_token被重掩码(): + # 行 0：pad1 + prompt2 + comp3；行 1：prompt3 + comp3 → pl = min(2,3) = 2 + ids = torch.tensor([[0, 3, 4, 5, 6, 5], [3, 4, 3, 6, 5, 6]]) + lab = torch.tensor( + [ + [IGNORE_INDEX] * 3 + [5, 6, 5], + [IGNORE_INDEX] * 3 + [6, 5, 6], + ] + ) + att = torch.tensor([[0, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1]]) + torch.manual_seed(1) + logits = torch.randn(2, 6, V) + + loss_base, num = sft_loss(logits, ids, lab, att) + assert num == 6 # 两行各 3 个 completion token，漏进切片的 prompt 位不计数 + + # 位置 1 的 logit 预测位置 2——两行的位置 2 都在切片内但都是 -100（行 0 是 + # prompt 尾、行 1 是漏进来的 prompt token）。改它不该动 loss + perturbed = logits.clone() + perturbed[:, 1] += 100.0 + loss_pert, _ = sft_loss(perturbed, ids, lab, att) + assert torch.allclose(loss_base, loss_pert) + + +def test_全掩码batch显式报错(): + ids, lab, att = batch_of_one([3, 4, 5, 6], [IGNORE_INDEX] * 4) + with pytest.raises(ValueError, match="有效 completion"): + sft_loss(torch.randn(1, 4, V), ids, lab, att) + + +def test_无prompt行显式报错(): + ids, lab, att = batch_of_one([3, 4], [3, 4]) # labels 全有效 → prompt 长 0 + with pytest.raises(ValueError, match="prompt_length"): + sft_loss(torch.randn(1, 2, V), ids, lab, att)