iomgaa
0ca60ea93f
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>
92 lines
3.7 KiB
Python
92 lines
3.7 KiB
Python
"""层 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
|