层1/T3: data.py 数据管线——加载、messages 归一、teacher 缓存挂接、双预算 collator（docs/02 §2.2-2.3）

- to_messages 三分支归一，except:pass 改显式报错（差异标注在注释）
- prompt_key: sha256 内容寻址，作为与 teacher.py 的缓存契约单点定义
- attach_teacher_completions: 缺键一次性报全，绝不静默跳过
- SFTCollator: 双预算截断 + 未截断长度定边界（坑二）+ -100 掩码 + 左 padding；
  prompt-only 行显式报错（层 2 接 on-policy 再放开）
- tests/test_data.py: 19 个单测，玩具字符级 tokenizer 覆盖超长解答/超长题目/
  enable_thinking 双取值/左 padding/缓存契约

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

ars_opd/data.py

@@ -0,0 +1,328 @@
+"""数据管线（IO 边缘，无论文锚点）：加载 → messages 归一 → 挂接 teacher 解答 → collator。
+
+层 1 的数据流（对应 docs/02 §1 的基线定义：SFT = 在 teacher rollout 上的离线蒸馏）：
+
+    DAPO parquet（prompt-only）
+      → to_messages 归一成 [{"role","content"}] 列表
+      → 按 seed 抽子集
+      → attach_teacher_completions 从 JSONL 缓存挂上 teacher 解答（assistant 轮）
+      → SFTCollator 分词、双预算截断、-100 掩码、左 padding
+
+本模块与 teacher.py 的缓存契约由 `prompt_key` 单点定义：teacher.py 生成缓存、
+本模块消费缓存，双方必须用同一个函数算键。
+"""
+
+from __future__ import annotations
+
+import ast
+import hashlib
+import json
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import torch
+from datasets import Dataset, load_dataset
+
+if TYPE_CHECKING:
+    from ars_opd.configs import SFTConfig
+
+# F.cross_entropy 的 ignore_index 默认值；标了它的位置不产生 loss
+IGNORE_INDEX = -100
+
+
+# ---------------------------------------------------------------------------
+# messages 归一
+# ---------------------------------------------------------------------------
+
+
+def _parse_stringified_list(value: str, column: str) -> list:
+    """parquet 有时把 list 存成其字符串形态，用 ast 还原。
+
+    差异标注：参考实现(train_distillation.py:299-305)在这里 `except: pass` 静默吞错，
+    坏行会以原始字符串流进 collator，在 apply_chat_template 处以难懂的方式炸；
+    我们显式报错，错误信息直接指向坏数据本身。
+    """
+    try:
+        parsed = ast.literal_eval(value)
+    except (ValueError, SyntaxError) as e:
+        raise ValueError(
+            f"列 {column!r} 是字符串但无法解析为 Python 字面量（坏数据行）："
+            f"{value[:200]!r}"
+        ) from e
+    if not isinstance(parsed, (list, tuple)):
+        raise ValueError(f"列 {column!r} 解析结果不是列表：{type(parsed).__name__}")
+    return list(parsed)
+
+
+def to_messages(example: dict[str, Any]) -> dict[str, list[dict[str, str]]]:
+    """把三种来源格式归一成 messages 列：[{"role": ..., "content": ...}, ...]。
+
+    支持（与参考实现 train_distillation.py:295-321 相同的三分支）：
+    - ``messages`` 列：直取；
+    - ``prompt`` 列（DAPO parquet，列名不副实——装的是完整 chat 列表）：改名；
+    - ``question`` 列（gsm8k 风格纯文本）：包成单 user 轮。
+
+    差异标注：参考实现对不认识的行 `return x` 静默放行，我们显式报错。
+    """
+    if "messages" in example:
+        msgs = example["messages"]
+        column = "messages"
+    elif "prompt" in example:
+        msgs = example["prompt"]
+        column = "prompt"
+    elif "question" in example:
+        return {"messages": [{"role": "user", "content": example["question"]}]}
+    else:
+        raise ValueError(
+            f"无法识别的数据行：既无 messages/prompt 也无 question 列，"
+            f"实有列 {sorted(example.keys())}"
+        )
+
+    if isinstance(msgs, str):
+        msgs = _parse_stringified_list(msgs, column)
+    msgs = list(msgs)
+
+    if not msgs:
+        raise ValueError(f"列 {column!r} 是空列表（坏数据行）")
+    for m in msgs:
+        if not isinstance(m, dict) or "role" not in m or "content" not in m:
+            raise ValueError(
+                f"列 {column!r} 中存在非 {{role, content}} 结构的元素：{m!r}"
+            )
+
+    return {"messages": [{"role": m["role"], "content": m["content"]} for m in msgs]}
+
+
+# ---------------------------------------------------------------------------
+# teacher 解答缓存（与 teacher.py 的契约）
+# ---------------------------------------------------------------------------
+
+
+def prompt_key(messages: list[dict[str, str]]) -> str:
+    """teacher 缓存的键：对 messages 的规范化 JSON 取 sha256。
+
+    差异标注：参考实现(distillation_trainer.py:984)用 `str(hash(prompt))`——
+    Python 对 str 的 hash 默认加盐，跨进程/跨次运行不稳定，缓存必然失效重生成。
+    sha256 内容寻址：同一道题永远同一个键。
+
+    只取 role/content 两个字段参与哈希：DAPO 行里其余元数据（data_source 等）
+    变了不应导致缓存失效。
+    """
+    canon = [{"role": m["role"], "content": m["content"]} for m in messages]
+    return hashlib.sha256(json.dumps(canon, ensure_ascii=False).encode()).hexdigest()
+
+
+def attach_teacher_completions(dataset: Dataset, jsonl_path: str) -> Dataset:
+    """把 teacher 解答缓存（JSONL，每行 {"key", "completion"}）挂到数据集上。
+
+    - 末轮已是 assistant 的行保持原样（数据自带解答，不覆盖）；
+    - 任何 prompt-only 行在缓存中查不到键 → 收集齐所有缺失后一次性报错，
+      提示先运行 teacher 生成——绝不静默跳过（跳过 = 悄悄改变训练集组成）。
+    """
+    path = Path(jsonl_path)
+    if not path.exists():
+        raise FileNotFoundError(
+            f"teacher 解答缓存不存在：{jsonl_path}。先运行 teacher.py 的批量生成。"
+        )
+
+    cache: dict[str, str] = {}
+    with open(path, encoding="utf-8") as f:
+        for line_no, line in enumerate(f, 1):
+            if not line.strip():
+                continue
+            rec = json.loads(line)  # 坏行直接炸，带行号
+            if "key" not in rec or "completion" not in rec:
+                raise ValueError(f"{jsonl_path}:{line_no} 缺少 key/completion 字段")
+            cache[rec["key"]] = rec["completion"]
+
+    # 先整体扫描缺失，一次性报全——比在 .map 里炸第一条更省来回
+    missing = [
+        i
+        for i, ex in enumerate(dataset)
+        if ex["messages"][-1]["role"] != "assistant"
+        and prompt_key(ex["messages"]) not in cache
+    ]
+    if missing:
+        raise KeyError(
+            f"{len(missing)}/{len(dataset)} 行在 teacher 缓存中查不到解答"
+            f"（首个缺失行 index={missing[0]}）。检查：teacher 生成是否用了同一"
+            f"子集与同一 seed？（子集抽取在 load_sft_dataset 中先于挂接发生，"
+            f"两侧 seed 不同则键集合不同）"
+        )
+
+    def _attach(ex: dict[str, Any]) -> dict[str, Any]:
+        msgs = ex["messages"]
+        if msgs[-1]["role"] == "assistant":
+            return ex
+        completion = cache[prompt_key(msgs)]
+        return {"messages": msgs + [{"role": "assistant", "content": completion}]}
+
+    return dataset.map(_attach)
+
+
+# ---------------------------------------------------------------------------
+# 数据集加载（入口）
+# ---------------------------------------------------------------------------
+
+
+def load_sft_dataset(cfg: "SFTConfig") -> Dataset:
+    """层 1 数据管线入口：加载 → 归一 → 抽子集 → 挂 teacher 解答。
+
+    返回只含 ``messages`` 一列的 Dataset，每行末轮是 assistant（可直接喂 SFTCollator）。
+    """
+    ds = _load_raw(cfg.dataset_path, cfg.dataset_split)
+    ds = ds.map(
+        to_messages,
+        remove_columns=[c for c in ds.column_names if c != "messages"],
+    )
+    if cfg.subset_size is not None and cfg.subset_size < len(ds):
+        # 非显然约束：抽子集必须在挂接 teacher 解答之前、且由 seed 完全确定——
+        # teacher.py 生成缓存时会走完全相同的"加载→归一→抽子集"路径，两侧 seed
+        # 一致才能得到同一批题；否则 attach 处大面积缓存 miss 报错。
+        ds = ds.shuffle(seed=cfg.seed).select(range(cfg.subset_size))
+    if cfg.teacher_completions_path is not None:
+        ds = attach_teacher_completions(ds, cfg.teacher_completions_path)
+    return ds
+
+
+def _load_raw(dataset_path: str, split: str) -> Dataset:
+    """三分支加载：parquet 目录 / 单 parquet 文件 / HF Hub 数据集名。
+
+    差异标注：参考实现(train_distillation.py:292)对 Hub 分支硬编码 config 名
+    "main"（gsm8k 专用）；我们不硬编码——需要特定 config 的数据集请下载成
+    parquet 本地加载。
+    """
+    p = Path(dataset_path)
+    if p.is_dir():
+        return load_dataset("parquet", data_dir=dataset_path, split=split)
+    if dataset_path.endswith(".parquet"):
+        return load_dataset("parquet", data_files=dataset_path, split=split)
+    return load_dataset(dataset_path, split=split)
+
+
+# ---------------------------------------------------------------------------
+# Collator：本层最核心的一段（对拍 distillation_trainer.py:210-343）
+# ---------------------------------------------------------------------------
+
+
+class SFTCollator:
+    """把一个 batch 的 messages 变成 input_ids/attention_mask/labels。
+
+    核心设计（继承参考实现的双预算方案，docs/02 §2.3）：prompt 与 completion
+    各自独立预算——prompt 用 max_prompt_length 截断，completion 上限是
+    max_length - len(截断后 prompt)。若只用一个总预算从右截断，超长解答会把
+    prompt 挤空，模型在"没有题目"的样本上学解答。
+
+    与参考实现的差异：
+    - 不支持 prompt-only 行（直接报错）。参考实现支持是为 on-policy 生成留口，
+      纯 SFT 下 prompt-only 行只会静默产生零 loss；层 2 接 on-policy 时再放开。
+    - 不返回 prompts/prompt_attention_mask（参考实现留给 vLLM 生成用，层 1 用不到）。
+    - 空 <think> 的一次性诊断打印改为单元测试断言（契约进测试，不进运行时日志）。
+    """
+
+    def __init__(
+        self,
+        tokenizer: "Any",
+        max_length: int,
+        max_prompt_length: int,
+        enable_thinking: bool = False,
+    ) -> None:
+        """tokenizer 需实现 HF 接口：apply_chat_template / __call__ / pad_token_id。"""
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.max_prompt_length = max_prompt_length
+        self.enable_thinking = enable_thinking
+        # pad→eos 回退：左 padding 位置的 attention_mask 恒为 0，pad 值不参与
+        # 任何计算，只需要一个合法 token id 占位，借用 eos 即可
+        if tokenizer.pad_token_id is not None:
+            self.pad_token_id: int = tokenizer.pad_token_id
+        elif tokenizer.eos_token_id is not None:
+            self.pad_token_id = tokenizer.eos_token_id
+        else:
+            raise ValueError("tokenizer 既无 pad_token 也无 eos_token，无法 padding")
+
+    def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
+        """batch 的 messages → 定长张量。
+
+        返回（B = batch 大小，T = batch 内最长序列长度）：
+        - input_ids: (B, T)  左 padding
+        - attention_mask: (B, T)  padding 位置为 0
+        - labels: (B, T)  padding 与 prompt 位置为 -100，completion 位置为 token id
+        """
+        all_input_ids: list[list[int]] = []
+        all_labels: list[list[int]] = []
+
+        for example in examples:
+            messages = example["messages"]
+            if len(messages) < 2 or messages[-1]["role"] != "assistant":
+                raise ValueError(
+                    "SFTCollator 收到 prompt-only 行（末轮不是 assistant）。纯 SFT 下"
+                    "它只会产生全 -100 的零 loss 样本——静默空训练。检查 teacher "
+                    "解答是否挂接成功。"
+                )
+
+            # prompt = 末轮 assistant 之前的全部轮次，渲染时带生成引导符
+            # （"<|im_start|>assistant\n..."），这样 completion 是纯解答文本的分词
+            formatted_prompt = self.tokenizer.apply_chat_template(
+                messages[:-1],
+                tokenize=False,
+                add_generation_prompt=True,
+                enable_thinking=self.enable_thinking,
+            )
+            # prompt 自己的预算内截断。沿用 tokenizer 默认右截断（与参考实现一致）：
+            # 超预算的题目被截掉尾部（含生成引导符）——1024 预算下 DAPO 极少触发，
+            # 触发时该样本退化但不会污染边界（边界用未截断长度算，见下）
+            prompt_ids: list[int] = self.tokenizer(
+                formatted_prompt,
+                truncation=True,
+                max_length=self.max_prompt_length,
+                add_special_tokens=False,
+            )["input_ids"]
+
+            # 非显然约束（docs/02 坑一/坑二）：completion 边界必须用"未截断 prompt
+            # 的分词长度"从整段渲染中切出。BPE 分词不满足拼接稳定性，分开渲染
+            # prompt 和 completion 再拼接 ≠ 整段渲染后分词；而若用截断后长度当切分
+            # 点，会把 prompt 尾部的 token 误标成 completion——静默的语义错误。
+            formatted_full = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=False,
+                enable_thinking=self.enable_thinking,
+            )
+            full_ids: list[int] = self.tokenizer(
+                formatted_full, truncation=False, add_special_tokens=False
+            )["input_ids"]
+            untruncated_prompt_len = len(
+                self.tokenizer(
+                    formatted_prompt, truncation=False, add_special_tokens=False
+                )["input_ids"]
+            )
+            completion_ids = full_ids[untruncated_prompt_len:]
+
+            # completion 预算 = 总预算 - 截断后 prompt 实长。配置校验
+            # （max_prompt_length < max_length）保证它恒 > 0
+            completion_budget = self.max_length - len(prompt_ids)
+            completion_ids = completion_ids[:completion_budget]
+
+            all_input_ids.append(prompt_ids + completion_ids)
+            # prompt 位置标 -100：题目不产生 loss，只学解答
+            all_labels.append([IGNORE_INDEX] * len(prompt_ids) + completion_ids)
+
+        # 左 padding：batch 内所有序列右对齐。纯 SFT 用右 padding 也行，但左 padding
+        # 让 trainer 能用一个标量 prompt_length 切 batch（docs/02 §2.4），且与
+        # 层 2+ 的生成场景（生成必须左 padding）统一，全项目只有一种 padding 约定
+        return {
+            "input_ids": _left_pad(all_input_ids, self.pad_token_id),  # (B, T)
+            "attention_mask": _left_pad(
+                [[1] * len(ids) for ids in all_input_ids], 0
+            ),  # (B, T)
+            "labels": _left_pad(all_labels, IGNORE_INDEX),  # (B, T)
+        }
+
+
+def _left_pad(seqs: list[list[int]], pad_value: int) -> torch.Tensor:
+    """把变长序列在左侧补齐成 (B, T) 张量，T = batch 内最大长度。"""
+    t_max = max(len(s) for s in seqs)
+    return torch.tensor(
+        [[pad_value] * (t_max - len(s)) + s for s in seqs], dtype=torch.long
+    )

tests/test_data.py

@@ -0,0 +1,260 @@
+"""层 1 / T3：数据管线单测（docs/02 §5.1 规定的验证项）。
+
+用字符级玩具 tokenizer 在 CPU 上对拍 collator 行为，不依赖网络下载真模型。
+玩具模板刻意模仿 Qwen3 的关键结构：生成引导符 + no-think 时注入空思考块。
+"""
+
+import json
+
+import pytest
+from datasets import Dataset
+
+from ars_opd.data import (
+    IGNORE_INDEX,
+    SFTCollator,
+    attach_teacher_completions,
+    prompt_key,
+    to_messages,
+)
+
+
+class ToyTokenizer:
+    """字符级 tokenizer：一个字符一个 token（id = 码点）。
+
+    模板契约与真 chat 模板同构：
+    - 每轮渲染成 "[role]content"；
+    - assistant 轮（或生成引导符后）在 no-think 模式下注入 "<T></T>"（模仿
+      Qwen3 的空 <think>\\n\\n</think>）；
+    - 完整渲染 == prompt 渲染 + 解答文本，保证边界可精确断言。
+    """
+
+    def __init__(self, pad_token_id=0, eos_token_id=1):
+        self.pad_token_id = pad_token_id
+        self.eos_token_id = eos_token_id
+
+    def apply_chat_template(
+        self,
+        messages,
+        tokenize=False,
+        add_generation_prompt=False,
+        enable_thinking=False,
+    ):
+        think = "" if enable_thinking else "<T></T>"
+        parts = []
+        for m in messages:
+            prefix = think if m["role"] == "assistant" else ""
+            parts.append(f"[{m['role']}]{prefix}{m['content']}")
+        text = "".join(parts)
+        if add_generation_prompt:
+            text += f"[assistant]{think}"
+        return text
+
+    def __call__(
+        self,
+        text,
+        truncation=False,
+        max_length=None,
+        padding=False,
+        add_special_tokens=False,
+    ):
+        ids = [ord(c) for c in text]
+        if truncation and max_length is not None:
+            ids = ids[:max_length]
+        return {"input_ids": ids}
+
+
+def ids_of(text):
+    return [ord(c) for c in text]
+
+
+def row(question, answer=None):
+    msgs = [{"role": "user", "content": question}]
+    if answer is not None:
+        msgs.append({"role": "assistant", "content": answer})
+    return {"messages": msgs}
+
+
+# ---------------------------------------------------------------------------
+# to_messages
+# ---------------------------------------------------------------------------
+
+
+def test_dapo_prompt列直接归一():
+    ex = {"prompt": [{"role": "user", "content": "1+1=?"}], "data_source": "dapo"}
+    assert to_messages(ex) == {"messages": [{"role": "user", "content": "1+1=?"}]}
+
+
+def test_字符串化的列表被还原():
+    ex = {"prompt": "[{'role': 'user', 'content': 'hi'}]"}
+    assert to_messages(ex)["messages"] == [{"role": "user", "content": "hi"}]
+
+
+def test_坏字符串显式报错而非静默放行():
+    # 参考实现 except:pass 会让这行以字符串形态流进 collator
+    with pytest.raises(ValueError, match="无法解析"):
+        to_messages({"prompt": "[{'role': broken"})
+
+
+def test_question列包成单user轮():
+    assert to_messages({"question": "2+2=?"}) == {
+        "messages": [{"role": "user", "content": "2+2=?"}]
+    }
+
+
+def test_无法识别的行报错():
+    with pytest.raises(ValueError, match="无法识别"):
+        to_messages({"foo": 1})
+
+
+# ---------------------------------------------------------------------------
+# prompt_key（与 teacher.py 的缓存契约）
+# ---------------------------------------------------------------------------
+
+
+def test_同题同键_不同题不同键():
+    m1 = [{"role": "user", "content": "q"}]
+    m2 = [{"role": "user", "content": "q'"}]
+    assert prompt_key(m1) == prompt_key(m1)
+    assert prompt_key(m1) != prompt_key(m2)
+
+
+def test_额外元数据字段不影响键():
+    plain = [{"role": "user", "content": "q"}]
+    noisy = [{"role": "user", "content": "q", "source": "dapo"}]
+    assert prompt_key(plain) == prompt_key(noisy)
+
+
+# ---------------------------------------------------------------------------
+# attach_teacher_completions
+# ---------------------------------------------------------------------------
+
+
+def write_cache(path, entries):
+    with open(path, "w", encoding="utf-8") as f:
+        for msgs, completion in entries:
+            f.write(
+                json.dumps({"key": prompt_key(msgs), "completion": completion}) + "\n"
+            )
+
+
+def test_挂接teacher解答(tmp_path):
+    q = [{"role": "user", "content": "1+1=?"}]
+    cache = tmp_path / "cache.jsonl"
+    write_cache(cache, [(q, "答案是 2")])
+    ds = Dataset.from_list([{"messages": q}])
+    out = attach_teacher_completions(ds, str(cache))
+    assert out[0]["messages"][-1] == {"role": "assistant", "content": "答案是 2"}
+
+
+def test_缓存缺键一次性报全部缺失(tmp_path):
+    cache = tmp_path / "cache.jsonl"
+    write_cache(cache, [])
+    ds = Dataset.from_list([row("q1"), row("q2")])
+    with pytest.raises(KeyError, match="2/2"):
+        attach_teacher_completions(ds, str(cache))
+
+
+def test_自带解答的行不被覆盖(tmp_path):
+    cache = tmp_path / "cache.jsonl"
+    write_cache(cache, [])
+    ds = Dataset.from_list([row("q", "人写的答案")])
+    out = attach_teacher_completions(ds, str(cache))
+    assert out[0]["messages"][-1]["content"] == "人写的答案"
+
+
+def test_缓存文件不存在报错():
+    ds = Dataset.from_list([row("q")])
+    with pytest.raises(FileNotFoundError):
+        attach_teacher_completions(ds, "/不存在/cache.jsonl")
+
+
+# ---------------------------------------------------------------------------
+# SFTCollator
+# ---------------------------------------------------------------------------
+
+
+def make_collator(**kw):
+    defaults = dict(max_length=1000, max_prompt_length=100, enable_thinking=False)
+    defaults.update(kw)
+    return SFTCollator(ToyTokenizer(), **defaults)
+
+
+def test_基本形态_掩码与边界():
+    collator = make_collator()
+    batch = collator([row("ab", "cd")])
+    prompt_text = "[user]ab[assistant]<T></T>"
+
+    labels = batch["labels"][0].tolist()
+    # prompt 全 -100，completion 位置是解答的 token
+    assert labels[: len(prompt_text)] == [IGNORE_INDEX] * len(prompt_text)
+    assert labels[len(prompt_text) :] == ids_of("cd")
+    assert batch["input_ids"][0].tolist() == ids_of(prompt_text + "cd")
+    assert batch["attention_mask"][0].tolist() == [1] * (len(prompt_text) + 2)
+
+
+def test_超长解答不挤占prompt():
+    # 头号正确性卖点：completion 被截，prompt 一个 token 不少
+    prompt_text = "[user]ab[assistant]<T></T>"
+    collator = make_collator(max_length=len(prompt_text) + 3)
+    batch = collator([row("ab", "x" * 50)])
+
+    input_ids = batch["input_ids"][0].tolist()
+    assert input_ids[: len(prompt_text)] == ids_of(prompt_text)  # prompt 完整
+    assert len(input_ids) == len(prompt_text) + 3  # completion 只剩预算内 3 个
+
+
+def test_超长题目截断但边界不错位():
+    # 坑二场景：prompt 超预算被截断，completion 的 token 必须仍然精确
+    # （切分点用未截断长度，而非截断后长度）
+    collator = make_collator(max_prompt_length=10)
+    batch = collator([row("很长的题目" * 20, "答案")])
+
+    labels = batch["labels"][0].tolist()
+    non_masked = [t for t in labels if t != IGNORE_INDEX]
+    assert non_masked == ids_of("答案")  # 解答 token 一个不错
+    assert sum(t == IGNORE_INDEX for t in labels) == 10  # prompt 恰被截到预算
+
+
+def test_enable_thinking两种取值边界都正确():
+    for thinking in (False, True):
+        collator = make_collator(enable_thinking=thinking)
+        batch = collator([row("q", "ans")])
+        non_masked = [t for t in batch["labels"][0].tolist() if t != IGNORE_INDEX]
+        assert non_masked == ids_of("ans"), f"enable_thinking={thinking} 时边界错位"
+
+
+def test_nothink模板注入空思考块():
+    # 参考实现的一次性诊断打印，在这里变成永久契约
+    text = ToyTokenizer().apply_chat_template(
+        [{"role": "user", "content": "q"}],
+        add_generation_prompt=True,
+        enable_thinking=False,
+    )
+    assert text.endswith("<T></T>")
+
+
+def test_prompt_only行显式报错():
+    with pytest.raises(ValueError, match="prompt-only"):
+        make_collator()([row("没有答案的题")])
+
+
+def test_左padding对齐():
+    collator = make_collator()
+    batch = collator([row("ab", "cd"), row("a", "c")])
+    t = batch["input_ids"].shape[1]
+    short_mask = batch["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["labels"][1].tolist()[:n_pad] == [IGNORE_INDEX] * n_pad
+    assert batch["input_ids"][1].tolist()[:n_pad] == [0] * n_pad  # pad_token_id=0
+
+
+def test_pad回退到eos():
+    collator = SFTCollator(
+        ToyTokenizer(pad_token_id=None, eos_token_id=7),
+        max_length=100,
+        max_prompt_length=50,
+    )
+    assert collator.pad_token_id == 7