层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>

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)