层1/T4: trainer.py 掩码 SFT 损失 + 最小 HF Trainer 子类（docs/02 §2.4）

- 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 <noreply@anthropic.com>

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)