层2/U4: DistillTrainer 编排（on-policy 生成→teacher no_grad 前向→反向KL）

trainer.py（对应 docs/03 §5 U4）:
- build_generated_batch: 纯函数，generate 输出重建 ids/attention/labels；
  "首个 eos 及之前有效"用 cumsum-self==0 实现，稳健对付 pad==eos / pad!=eos / 无eos
- DistillTrainer.compute_loss 四步：生成(no_grad,unwrap)→重建→双前向→移位+divergence
- 损失几何复用 compute_prompt_length（与 sft_loss 同款）；梯度只经 student
- 构造时校验 teacher/student 同 tokenizer（白盒前提，比 DT:2876 更早）
- teacher eval+冻结、设备迁移推迟到 compute_loss；同层1退出梯度累积新式契约

test_distill.py（对应 docs/03 §2.5）:
- build_generated_batch 6 例：eos居中屏蔽/无eos全监督/batch混长/pad==eos/立即eos/prompt段-100

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

tests/test_distill.py

@@ -0,0 +1,91 @@
+"""层 2 / U4：生成重建纯逻辑单测（docs/03 §2.5）。
+
+只测 build_generated_batch——把 model.generate 输出重建成 input_ids/attention/labels
+的张量逻辑。DistillTrainer 的编排（生成→双前向→divergence）需真模型，由远程冒烟
+验证。重点盯"首个 eos 之后一律屏蔽"这条最易 off-by-one 的规则。
+"""
+
+import torch
+
+from ars_opd.data import IGNORE_INDEX
+from ars_opd.trainer import build_generated_batch
+
+EOS = 99
+PAD = 0
+
+
+def test_eos在中间_其后全屏蔽():
+    # 行内：prompt=[pad,u,u]，生成=[a, EOS, pad]；a 与 eos 有效，eos 后的 pad 无效
+    prompts = torch.tensor([[PAD, 5, 5]])
+    prompt_mask = torch.tensor([[0, 1, 1]])
+    gen_output = torch.tensor([[PAD, 5, 5, 7, EOS, PAD]])  # (1, P+G)=(1,6)
+
+    ids, attn, labels = build_generated_batch(prompts, prompt_mask, gen_output, EOS)
+
+    assert ids.tolist() == gen_output.tolist()  # input_ids 即生成全序列
+    # prompt 段全 -100；生成段 [7, EOS, -100]
+    assert labels[0].tolist() == [IGNORE_INDEX] * 3 + [7, EOS, IGNORE_INDEX]
+    # attention：prompt 左 pad=0，生成段 eos 及之前=1、其后=0
+    assert attn[0].tolist() == [0, 1, 1, 1, 1, 0]
+
+
+def test_无eos撞max时整段生成有效():
+    prompts = torch.tensor([[5, 5, 5]])
+    prompt_mask = torch.tensor([[1, 1, 1]])
+    gen_output = torch.tensor([[5, 5, 5, 8, 9, 10]])  # 生成三 token，无 eos
+
+    _, attn, labels = build_generated_batch(prompts, prompt_mask, gen_output, EOS)
+
+    assert labels[0].tolist() == [IGNORE_INDEX] * 3 + [8, 9, 10]  # 全监督
+    assert attn[0].tolist() == [1, 1, 1, 1, 1, 1]
+
+
+def test_batch内不同生成长度_各自正确对齐():
+    # 行0 提前 eos（右侧被补 pad 到 batch 宽度）；行1 撞 max。二者共用同一 (B,P+G)
+    prompts = torch.tensor([[PAD, 5, 5], [5, 5, 5]])
+    prompt_mask = torch.tensor([[0, 1, 1], [1, 1, 1]])
+    gen_output = torch.tensor(
+        [
+            [PAD, 5, 5, 7, EOS, PAD],  # 行0：生成 [7, EOS]，末位 pad 补齐
+            [5, 5, 5, 8, 9, 10],  # 行1：生成 [8, 9, 10]
+        ]
+    )
+
+    _, attn, labels = build_generated_batch(prompts, prompt_mask, gen_output, EOS)
+
+    assert labels[0].tolist() == [IGNORE_INDEX] * 3 + [7, EOS, IGNORE_INDEX]
+    assert labels[1].tolist() == [IGNORE_INDEX] * 3 + [8, 9, 10]
+    assert attn[0].tolist() == [0, 1, 1, 1, 1, 0]
+    assert attn[1].tolist() == [1, 1, 1, 1, 1, 1]
+
+
+def test_pad等于eos也不误判():
+    # 关键 corner：pad_token == eos_token。首个 eos 有效、其后补位的 eos 全屏蔽
+    prompts = torch.tensor([[5, 5, 5]])
+    prompt_mask = torch.tensor([[1, 1, 1]])
+    gen_output = torch.tensor([[5, 5, 5, 7, EOS, EOS]])  # 末位补的 pad 恰好==eos
+
+    _, attn, labels = build_generated_batch(prompts, prompt_mask, gen_output, EOS)
+
+    assert labels[0].tolist() == [IGNORE_INDEX] * 3 + [7, EOS, IGNORE_INDEX]
+    assert attn[0].tolist() == [1, 1, 1, 1, 1, 0]  # 第二个 eos 被当补位屏蔽
+
+
+def test_立即eos_只留一个token():
+    prompts = torch.tensor([[5, 5, 5]])
+    prompt_mask = torch.tensor([[1, 1, 1]])
+    gen_output = torch.tensor([[5, 5, 5, EOS, PAD, PAD]])  # 第一步就 eos
+
+    _, attn, labels = build_generated_batch(prompts, prompt_mask, gen_output, EOS)
+
+    assert labels[0].tolist() == [IGNORE_INDEX] * 3 + [EOS, IGNORE_INDEX, IGNORE_INDEX]
+    assert attn[0].tolist() == [1, 1, 1, 1, 0, 0]
+
+
+def test_prompt段恒为负100():
+    prompts = torch.tensor([[PAD, PAD, 5, 5]])
+    prompt_mask = torch.tensor([[0, 0, 1, 1]])
+    gen_output = torch.tensor([[PAD, PAD, 5, 5, 8, 9]])
+
+    _, _, labels = build_generated_batch(prompts, prompt_mask, gen_output, EOS)
+    assert labels[0, :4].tolist() == [IGNORE_INDEX] * 4  # prompt 段（含左 pad）全 -100