From e5a28e8e7754d2922b6e1fd8f98531bdddeb045d Mon Sep 17 00:00:00 2001 From: iomgaa Date: Sun, 19 Jul 2026 05:02:24 -0400 Subject: [PATCH] =?UTF-8?q?=E5=B1=822/U3:=20SFTCollator=20=E6=94=BE?= =?UTF-8?q?=E5=BC=80=20prompt-only=20=E6=A8=A1=E5=BC=8F=EF=BC=88=E4=BE=9B?= =?UTF-8?q?=20on-policy=20=E7=94=9F=E6=88=90=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit data.py（对应 docs/03 §5 U3）: - 加 prompt_only 开关：True 时输出 prompts/prompt_attention_mask（不产 labels， 由 U4 生成后重建）；False 时 SFT 双预算路径逐字不变 - max_length 改可选：prompt-only 无总预算；SFT 模式缺它构造即报错 - 兑现 T3 为 on-policy 生成预留的口子；生成用左 padding（右边界对齐） test_data.py: - 新增 prompt_only 模式：返回 prompt 张量/不报错、左 padding、截断、剥末轮 assistant - 回归守卫：SFT 模式仍拒绝 prompt-only 行（"SFT 路径行为不变"） Co-Authored-By: Claude Opus 4.8 (1M context) --- ars_opd/data.py | 89 ++++++++++++++++++++++++++++++++++++++++------ tests/test_data.py | 60 ++++++++++++++++++++++++++++++- 2 files changed, 138 insertions(+), 11 deletions(-) diff --git a/ars_opd/data.py b/ars_opd/data.py index a38bc30..2e861ee 100644 --- a/ars_opd/data.py +++ b/ars_opd/data.py @@ -206,32 +206,48 @@ def _load_raw(dataset_path: str, split: str) -> Dataset: class SFTCollator: - """把一个 batch 的 messages 变成 input_ids/attention_mask/labels。 + """把一个 batch 的 messages 变成训练/生成所需的定长张量，两种模式二选一。 + prompt_only=False（层 1 SFT，默认）——输出 input_ids/attention_mask/labels： 核心设计（继承参考实现的双预算方案，docs/02 §2.3）：prompt 与 completion 各自独立预算——prompt 用 max_prompt_length 截断，completion 上限是 max_length - len(截断后 prompt)。若只用一个总预算从右截断，超长解答会把 - prompt 挤空，模型在"没有题目"的样本上学解答。 + prompt 挤空，模型在"没有题目"的样本上学解答。要求每行末轮是 assistant， + 否则报错（prompt-only 行在纯 SFT 下只产生零 loss = 静默空训练）。 - 与参考实现的差异： - - 不支持 prompt-only 行（直接报错）。参考实现支持是为 on-policy 生成留口， - 纯 SFT 下 prompt-only 行只会静默产生零 loss；层 2 接 on-policy 时再放开。 - - 不返回 prompts/prompt_attention_mask（参考实现留给 vLLM 生成用，层 1 用不到）。 - - 空 的一次性诊断打印改为单元测试断言（契约进测试，不进运行时日志）。 + prompt_only=True（层 2 white-box OPD，docs/03 §5 U3）——只渲染 prompt、 + 输出 prompts/prompt_attention_mask 供 model.generate 做 on-policy 生成； + completion 由生成产生、labels 由 U4 的 DistillTrainer 在生成后重建，故此模式 + 不产 labels、也不吃 max_length。这兑现了参考实现为 on-policy 生成留的口子 + （层 1 曾故意关掉，见此前 git 历史）。 + + 与参考实现的其余差异：空 的一次性诊断打印改为单元测试断言（契约进 + 测试，不进运行时日志）。 """ def __init__( self, tokenizer: "Any", - max_length: int, max_prompt_length: int, + max_length: int | None = None, enable_thinking: bool = False, + prompt_only: bool = False, ) -> None: - """tokenizer 需实现 HF 接口：apply_chat_template / __call__ / pad_token_id。""" + """tokenizer 需实现 HF 接口：apply_chat_template / __call__ / pad_token_id。 + + max_length 仅 SFT 模式需要（completion 预算依赖它）；prompt_only 模式下 + completion 是生成的、无总预算，故 max_length 可为 None。 + """ + if not prompt_only and max_length is None: + raise ValueError( + "SFT 模式（prompt_only=False）必须提供 max_length——completion " + "预算 = max_length - len(prompt)，缺它无法确定解答截断点。" + ) self.tokenizer = tokenizer self.max_length = max_length self.max_prompt_length = max_prompt_length self.enable_thinking = enable_thinking + self.prompt_only = prompt_only # pad→eos 回退：左 padding 位置的 attention_mask 恒为 0，pad 值不参与 # 任何计算，只需要一个合法 token id 占位，借用 eos 即可 if tokenizer.pad_token_id is not None: @@ -242,7 +258,60 @@ class SFTCollator: raise ValueError("tokenizer 既无 pad_token 也无 eos_token，无法 padding") def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: - """batch 的 messages → 定长张量。 + """按模式分派：prompt_only 走生成用 prompt 张量，否则走 SFT 双预算。""" + if self.prompt_only: + return self._collate_prompt_only(examples) + return self._collate_sft(examples) + + def _collate_prompt_only( + self, examples: list[dict[str, Any]] + ) -> dict[str, torch.Tensor]: + """层 2：只渲染 prompt 供 on-policy 生成，不产 completion/labels。 + + 返回（B = batch 大小，P = batch 内最长 prompt 长度）： + - prompts: (B, P) 左 padding + - prompt_attention_mask: (B, P) padding 位置为 0 + + 非显然约束：生成必须左 padding——所有 prompt 右对齐到同一右边界， + model.generate 从该边界统一续写；右 padding 会让短 prompt 的生成从 pad + 中间开始，全乱。这也是层 1 SFT 就选左 padding 的原因（全项目一种约定）。 + """ + all_prompt_ids: list[list[int]] = [] + for example in examples: + messages = example["messages"] + # prompt-only 数据末轮是 user；若末轮已是 assistant 则剥掉，取生成前上下文 + prompt_msgs = ( + messages[:-1] if messages[-1]["role"] == "assistant" else messages + ) + if not prompt_msgs: + raise ValueError( + "prompt_only collator 收到空 prompt（无可生成的上下文）" + ) + # 与 SFT 模式同样带生成引导符渲染（add_generation_prompt=True）： + # prompt 末尾就是 "<|im_start|>assistant\n..."，生成从此续写 + formatted_prompt = self.tokenizer.apply_chat_template( + prompt_msgs, + tokenize=False, + add_generation_prompt=True, + enable_thinking=self.enable_thinking, + ) + prompt_ids: list[int] = self.tokenizer( + formatted_prompt, + truncation=True, + max_length=self.max_prompt_length, + add_special_tokens=False, + )["input_ids"] + all_prompt_ids.append(prompt_ids) + + return { + "prompts": _left_pad(all_prompt_ids, self.pad_token_id), # (B, P) + "prompt_attention_mask": _left_pad( + [[1] * len(ids) for ids in all_prompt_ids], 0 + ), # (B, P) + } + + def _collate_sft(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: + """层 1 SFT：messages（末轮 assistant）→ 定长张量。 返回（B = batch 大小，T = batch 内最长序列长度）： - input_ids: (B, T) 左 padding diff --git a/tests/test_data.py b/tests/test_data.py index e22289c..23e7783 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -233,11 +233,20 @@ def test_nothink模板注入空思考块(): assert text.endswith("") -def test_prompt_only行显式报错(): +def test_sft模式仍拒绝prompt_only行(): + # 回归守卫（docs/03 §5 U3："SFT 路径行为不变"）：SFT 模式下 prompt-only 行 + # 仍是静默空训练闸门，必须报错——放开只发生在显式 prompt_only=True 模式 with pytest.raises(ValueError, match="prompt-only"): make_collator()([row("没有答案的题")]) +def test_sft模式缺max_length构造即报错(): + with pytest.raises(ValueError, match="max_length"): + SFTCollator( + ToyTokenizer(), max_prompt_length=50 + ) # 非 prompt_only 却无 max_length + + def test_左padding对齐(): collator = make_collator() batch = collator([row("ab", "cd"), row("a", "c")]) @@ -258,3 +267,52 @@ def test_pad回退到eos(): max_prompt_length=50, ) assert collator.pad_token_id == 7 + + +# --------------------------------------------------------------------------- +# SFTCollator：prompt_only 模式（层 2 / U3） +# --------------------------------------------------------------------------- + + +def make_prompt_collator(**kw): + defaults = dict(max_prompt_length=100, enable_thinking=False, prompt_only=True) + defaults.update(kw) + return SFTCollator(ToyTokenizer(), **defaults) + + +def test_prompt_only模式返回prompt张量且不报错(): + # 层 2：prompt-only 行是常态，不再报错；渲染带生成引导符供 generate 续写 + collator = make_prompt_collator() + batch = collator([row("ab")]) + expected = "[user]ab[assistant]" + + assert set(batch.keys()) == {"prompts", "prompt_attention_mask"} + assert batch["prompts"][0].tolist() == ids_of(expected) + assert batch["prompt_attention_mask"][0].tolist() == [1] * len(expected) + + +def test_prompt_only模式左padding对齐(): + # 生成要求左 padding：短 prompt 在左侧补 pad，右边界对齐 + collator = make_prompt_collator() + batch = collator([row("abc"), row("a")]) + t = batch["prompts"].shape[1] + short_mask = batch["prompt_attention_mask"][1].tolist() + n_pad = t - short_mask.count(1) + + assert n_pad > 0 + assert short_mask[:n_pad] == [0] * n_pad # padding 在左 + assert batch["prompts"][1].tolist()[:n_pad] == [0] * n_pad # pad_token_id=0 + + +def test_prompt_only模式截断到prompt预算(): + collator = make_prompt_collator(max_prompt_length=8) + batch = collator([row("很长的题目" * 20)]) + assert batch["prompts"].shape[1] == 8 # 恰截到 max_prompt_length + + +def test_prompt_only模式剥掉末轮assistant(): + # 若数据碰巧带了 assistant 轮，取生成前上下文（剥掉它再加生成引导符） + collator = make_prompt_collator() + batch = collator([row("q", "已有答案")]) + expected = "[user]q[assistant]" # 不含"已有答案" + assert batch["prompts"][0].tolist() == ids_of(expected)