From 6c6fb576ee8503bc6dabb505eb4a925bb0260810 Mon Sep 17 00:00:00 2001
From: iomgaa <iomgaaycz@gmail.com>
Date: Wed, 15 Jul 2026 08:22:18 -0400
Subject: [PATCH] =?UTF-8?q?feat(harness):=20checkpoint=20=E5=AD=98=20unit?=
 =?UTF-8?q?=5Fid=20=E5=BA=8F=E5=88=97=EF=BC=8C=E6=96=AD=E7=82=B9=E7=BB=AD?=
 =?UTF-8?q?=E8=B7=91=E5=AD=AA=E7=94=9F=E5=AF=B9=E4=B8=8D=E6=8B=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

核心算法保真 #3（断点续跑）：checkpoint 从逐题 question_id 改为存 unit_id
序列（孪生对折叠为单个 unit_id），恢复时 build_units + 按完整 unit 展开，
续跑后 pair 两成员同进同出、绝不被劈开。

- _batch_unit_ids/_batch_from_ids 对称折叠/展开，保序去重，纯非 AR 下
  unit_id==question_id、与旧逐题序列逐字节一致。
- momentum 采样抽取为 _sample_momentum_candidates 纯函数，docstring 显式
  记录 Phase 1 设计偏差：仅保证纯非 AR byte-identical，混格 momentum 不保证。
- 新增 test_checkpoint_pair（unit_id 落盘往返、pair 不拆）与
  test_non_ar_byte_identical（pools→batching→checkpoint→momentum 端到端黄金）。
---
 app/harness/checkpoint.py                 |   3 +-
 app/harness/runner.py                     |  89 +++++--
 tests/integration/test_checkpoint_pair.py | 254 ++++++++++++++++++
 tests/unit/test_non_ar_byte_identical.py  | 297 ++++++++++++++++++++++
 4 files changed, 627 insertions(+), 16 deletions(-)
 create mode 100644 tests/integration/test_checkpoint_pair.py
 create mode 100644 tests/unit/test_non_ar_byte_identical.py

diff --git a/app/harness/checkpoint.py b/app/harness/checkpoint.py
index d08b0a5..f686d3b 100644
--- a/app/harness/checkpoint.py
+++ b/app/harness/checkpoint.py
@@ -233,7 +233,8 @@ def write_checkpoint(
         global_step: 全局 step 序号。
         total_steps: 全局总 step 数。
         version_snapshot: skills/prompts 版本快照。
-        epoch_batches: 本 epoch 的 batch 划分（question_id 列表的列表）。
+        epoch_batches: 本 epoch 的 batch 划分（unit_id 列表的列表，孪生对折叠为
+            单个 unit_id；纯非 AR 下 unit_id==question_id）。
         config: 训练配置对象，用于计算 config_fingerprint。
 
     关键实现细节:
diff --git a/app/harness/runner.py b/app/harness/runner.py
index 80c6110..4b27a65 100644
--- a/app/harness/runner.py
+++ b/app/harness/runner.py
@@ -199,18 +199,78 @@ def _accumulate_slow_packs(diagnosis: DiagnosisResult, state: _TrainState) -> No
     state.tool_packs.extend(diagnosis.tool_case_packs.values())
 
 
-def _batch_from_ids(pools: Pools, ids: list[str]) -> list[GeneratedQuestion]:
-    """按 question_id 从诊断池重建一个 batch（保持原 epoch 划分）。
+def _batch_unit_ids(batch: list[GeneratedQuestion]) -> list[str]:
+    """把一个 batch 的扁平题目折叠为 unit_id 序列（孪生对成员去重为单个 unit_id）。
+
+    checkpoint 存 unit_id 序列而非逐题 question_id：断点续跑恢复时按完整 unit 展开，
+    保证孪生对整体重建、绝不被劈开（核心算法保真 #3 断点续跑）。
+
+    参数:
+        batch: 一个 mini-batch 的扁平题目列表（pair 两成员相邻）。
+
+    返回:
+        unit_id 列表，按题目在 batch 中的首次出现顺序去重；single 的 unit_id 即
+        question_id，故纯非 AR 输入下与旧逐题 question_id 序列逐字节一致。
+
+    关键实现:
+        用 dict 保序去重（pair 两成员共享 unit_id，仅记一次），无需额外集合。
+    """
+    ordered: dict[str, None] = {}
+    for q in batch:
+        ordered[q.unit_id] = None
+    return list(ordered)
+
+
+def _batch_from_ids(pools: Pools, unit_ids: list[str]) -> list[GeneratedQuestion]:
+    """按 unit_id 序列从诊断池重建一个 batch，按完整 unit 展开成题目列表。
+
+    与 _batch_unit_ids 对称：恢复时以完整 unit 为单位展开（pair 两成员同进同出），
+    断点续跑后孪生对绝不被拆开（核心算法保真 #3）。
 
     参数:
         pools: 三池容器。
-        ids: 一个 batch 的 question_id 列表。
+        unit_ids: 一个 batch 的 unit_id 序列（checkpoint 存的粒度）。
 
     返回:
-        按 ids 顺序取出的 GeneratedQuestion 列表。
+        按 unit_ids 顺序展开的 GeneratedQuestion 列表；每个 unit_id 展开为其全部
+        成员题（single 1 题、pair 2 题），顺序与原 batch 一致。
+
+    关键实现:
+        直接以 units_by_id[uid] 取值，unit_id 缺失触发 KeyError（P5 防静默兜底），
+        强制 checkpoint 与当前诊断池一致；纯非 AR 下 unit_id==question_id、单元即
+        单题，与旧逐题重建逐字节一致。
     """
-    by_id = {q.question_id: q for q in pools.diagnosis}
-    return [by_id[i] for i in ids]
+    units_by_id = {u.unit_id: u for u in build_units(pools.diagnosis)}
+    return [q for uid in unit_ids for q in units_by_id[uid].questions]
+
+
+def _sample_momentum_candidates(
+    pool: list[GeneratedQuestion],
+    allowed_task_types: set[str],
+    momentum_samples: int,
+    epoch: int,
+) -> list[GeneratedQuestion]:
+    """为 momentum 从诊断池按题型过滤后做确定性抽样（逐题粒度）。
+
+    设计偏差（Phase 1 显式记录）：momentum 采样在逐题粒度进行、不折叠 QuestionUnit，
+    故仅保证「纯非 AR 题库」的抽样序列与引入 QuestionUnit 前逐字节一致；混格题库下
+    孪生对可能被半采样、且候选集长度/顺序随 AR 成员增减而漂移，**Phase 1 不保证混格
+    momentum 的 byte-identical**（属可接受偏差，纯非 AR 必须不漂）。
+
+    参数:
+        pool: 诊断池扁平题目列表。
+        allowed_task_types: 允许参与的题型集合。
+        momentum_samples: 目标采样数上限。
+        epoch: 采样种子（同 epoch 可复现）。
+
+    返回:
+        采样到的题目列表；候选不足则全取，候选为空返回空列表。
+    """
+    candidates = [q for q in pool if q.task_type in allowed_task_types]
+    n = min(momentum_samples, len(candidates))
+    if n <= 0:
+        return []
+    return random.Random(epoch).sample(candidates, n)
 
 
 def _snapshot_current_skills(skills_dir: Path) -> dict[str, str]:
@@ -753,7 +813,8 @@ class Runner:
                     correct_ratio=self._config.batch_correct_ratio,
                 )
                 step_from = 0
-            batch_ids = [[q.question_id for q in b] for b in batches]
+            # checkpoint 存 unit_id 序列：断点续跑按完整 unit 展开，孪生对不拆
+            batch_unit_ids = [_batch_unit_ids(b) for b in batches]
             for step in range(step_from, len(batches)):
                 await self._run_step(epoch, step, total_steps, batches[step], pools, state)
                 state.global_step += 1
@@ -766,7 +827,7 @@ class Runner:
                     global_step=state.global_step,
                     total_steps=total_steps,
                     version_snapshot=self._current_version_snapshot(),
-                    epoch_batches=batch_ids,
+                    epoch_batches=batch_unit_ids,
                     config=self._config,
                 )
             await self._slow_update_cycle(epoch, pools, state)
@@ -782,7 +843,7 @@ class Runner:
                 global_step=state.global_step,
                 total_steps=total_steps,
                 version_snapshot=self._current_version_snapshot(),
-                epoch_batches=batch_ids,
+                epoch_batches=batch_unit_ids,
                 config=self._config,
             )
             if _should_early_stop(
@@ -1610,12 +1671,10 @@ class Runner:
         prev_skill = state.epoch_start_skills.get(target_file, skill_content)
         prev_guidance = momentum_inner(skill_content)
 
-        # 采样
-        allowed = set(task_types)
-        candidates = [q for q in pools.diagnosis if q.task_type in allowed]
-        rng = random.Random(epoch)
-        n = min(self._config.momentum_samples, len(candidates))
-        sampled = rng.sample(candidates, n) if n > 0 else []
+        # 采样（逐题粒度，不折叠 unit；混格偏差见 _sample_momentum_candidates docstring）
+        sampled = _sample_momentum_candidates(
+            pools.diagnosis, set(task_types), self._config.momentum_samples, epoch
+        )
 
         if not sampled:
             skill_path.write_text(
diff --git a/tests/integration/test_checkpoint_pair.py b/tests/integration/test_checkpoint_pair.py
new file mode 100644
index 0000000..0b5591a
--- /dev/null
+++ b/tests/integration/test_checkpoint_pair.py
@@ -0,0 +1,254 @@
+"""断点续跑 checkpoint 的 unit_id 粒度契约集成测试（Task 10）。
+
+覆盖三条铁律：
+- checkpoint 存 **unit_id 序列**而非逐题 question_id（孪生对折叠为单个 unit_id）；
+- 续跑恢复以 **完整 unit** 展开（pair 两成员同进同出，绝不被劈开）；
+- 真实经 write_checkpoint → load_checkpoint 落盘往返后，unit_id 序列无损存活、
+  重建的 batch 与原 batch 逐字节一致（core 算法保真 #3 断点续跑）。
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from app.harness.batching import build_batches
+from app.harness.checkpoint import load_checkpoint, write_checkpoint
+from app.harness.pools import Pools
+from app.harness.question_units import build_units, validate_units
+from app.harness.runner import _batch_from_ids, _batch_unit_ids
+from core.types import GeneratedQuestion
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+# ---------------------------------------------------------------------------
+# 辅助构造
+# ---------------------------------------------------------------------------
+
+
+def _make_single(qid: str, task_type: str = "RETRIEVAL", video_id: str = "v1") -> GeneratedQuestion:
+    """构造最小 single 题目。"""
+    return GeneratedQuestion(
+        question_id=qid,
+        video_id=video_id,
+        task_type=task_type,
+        question=f"q_{qid}",
+        options=("A. a", "B. b", "C. c", "D. d"),
+        answer="A",
+        source_nodes=("n1",),
+        difficulty="medium",
+        question_role="single",
+    )
+
+
+def _make_pair(
+    pair_id: str, task_type: str = "AR", video_id: str = "v1"
+) -> tuple[GeneratedQuestion, GeneratedQuestion]:
+    """构造一个孪生对（original + mirror），共享 pair_id / flip_axis。"""
+    common = {
+        "video_id": video_id,
+        "task_type": task_type,
+        "question": "?",
+        "options": ("A. a", "B. b", "C. c", "D. d"),
+        "answer": "A",
+        "source_nodes": ("n1",),
+        "difficulty": "medium",
+        "pair_id": pair_id,
+        "flip_axis": "before_after",
+    }
+    original = GeneratedQuestion(
+        question_id=f"{pair_id}_o", question_role="pair_original", **common
+    )
+    mirror = GeneratedQuestion(question_id=f"{pair_id}_m", question_role="pair_mirror", **common)
+    return original, mirror
+
+
+def _pools_with(diagnosis: list[GeneratedQuestion]) -> Pools:
+    """用给定诊断池构造最小 Pools（其余池置空，_batch_from_ids 只读 diagnosis）。"""
+    return Pools(
+        diagnosis=diagnosis,
+        validation=[],
+        test=[],
+        baseline_run_id="baseline",
+        baseline_val_accuracy=0.0,
+        correctness={q.question_id: False for q in diagnosis},
+    )
+
+
+# ---------------------------------------------------------------------------
+# checkpoint 落盘所需的最小 _TrainState / RunConfig 替身
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class _FakeState:
+    """serialize_state 只读的可持久化字段（本测试用空累加包即可）。"""
+
+    correctness: dict[str, bool] = field(default_factory=dict)
+    eval_prev_acc: float = 0.0
+    eval_prev_run_id: str = "run-0"
+    baseline_skills_version: str = "v1"
+    baseline_prompts_version: str = "v1"
+    steps_since_best_improved: int = 0
+    epoch_start_skills: str = "v1"
+    changed_task_types_this_epoch: set[str] = field(default_factory=set)
+    rejected_buffer: dict = field(default_factory=dict)
+    system_packs: list = field(default_factory=list)
+    tool_packs: list = field(default_factory=list)
+    probations: dict = field(default_factory=dict)
+    gate_cooldown: dict = field(default_factory=dict)
+    gate_epoch_observed: dict = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class _FakeConfig:
+    """compute_fingerprint 读取的结构性 + 决策性字段。"""
+
+    batch_size: int = 4
+    min_class_per_batch: int = 2
+    epochs: int = 3
+    diag_size: int = 30
+    val_size: int = 50
+    batch_correct_ratio: float = 0.0
+    edit_budget_start: int = 6
+    edit_budget_end: int = 3
+    early_stop_patience: int = 3
+    use_slow_momentum: bool = True
+    skill_update_mode: str = "patch"
+    appendix_consolidate_threshold: int = 10
+    momentum_samples: int = 20
+    gate_e_confirm: float = 20.0
+    gate_e_provisional: float = 6.0
+    gate_w_net_min: int = 2
+    gate_delta_min: float = 0.02
+    gate_lambda_dir: float = -3.0
+    gate_e_rollback: float = 10.0
+    gate_block: int = 4
+    gate_n_max: int = 40
+    gate_p_low: float = 0.1
+    gate_p_high: float = 0.9
+    gate_probe_quota: float = 0.2
+    gate_gamma_decay: float = 0.9
+    gate_cooldown_steps: int = 2
+    gate_guard_err: float = 0.3
+
+
+# ---------------------------------------------------------------------------
+# checkpoint 存 unit_id 序列（孪生对折叠为单个 unit_id）
+# ---------------------------------------------------------------------------
+
+
+class TestCheckpointStoresUnitIds:
+    """一个含孪生对的 batch，checkpoint 序列应折叠为 unit_id（pair 只记一次）。"""
+
+    def test_pair_folds_to_single_unit_id(self) -> None:
+        po, pm = _make_pair("pA")
+        s1, s2 = _make_single("s1"), _make_single("s2")
+        batch = [s1, po, pm, s2]  # pair 两成员相邻
+
+        unit_ids = _batch_unit_ids(batch)
+
+        # 4 题折叠为 3 个 unit：s1 / pA(pair) / s2
+        assert unit_ids == ["s1", "pA", "s2"]
+        # 逐题 question_id 序列（旧行为）会含 4 项且把 pair 拆成两条
+        assert [q.question_id for q in batch] == ["s1", "pA_o", "pA_m", "s2"]
+
+    def test_pure_single_unit_ids_equal_question_ids(self) -> None:
+        """纯 single batch 的 unit_id 序列与逐题 question_id 序列逐字节一致。"""
+        batch = [_make_single(f"s{i}") for i in range(5)]
+        assert _batch_unit_ids(batch) == [q.question_id for q in batch]
+
+
+# ---------------------------------------------------------------------------
+# 恢复以完整 unit 展开：pair 不拆
+# ---------------------------------------------------------------------------
+
+
+class TestResumeExpandsFullUnit:
+    """按 unit_id 序列恢复时，孪生对整体展开、绝不被劈开。"""
+
+    def test_pair_not_split_on_resume(self) -> None:
+        po, pm = _make_pair("pA")
+        s1, s2 = _make_single("s1"), _make_single("s2")
+        batch = [s1, po, pm, s2]
+        # 诊断池顺序故意与 batch 不同，验证恢复以池折叠的 unit 为准
+        pools = _pools_with([s2, pm, s1, po])
+
+        unit_ids = _batch_unit_ids(batch)
+        restored = _batch_from_ids(pools, unit_ids)
+
+        # 逐字节一致（顺序保留）
+        assert [q.question_id for q in restored] == ["s1", "pA_o", "pA_m", "s2"]
+        # 恢复后可无损重建为完整孪生对单元
+        units = validate_units(build_units(restored))
+        pair_units = [u for u in units if u.kind == "pair"]
+        assert len(pair_units) == 1
+        assert pair_units[0].unit_id == "pA"
+        assert pair_units[0].size == 2
+
+    def test_roundtrip_preserves_order_mixed(self) -> None:
+        """混格 batch 经 unit_id 折叠 → 展开 round-trip 逐字节还原原顺序。"""
+        po, pm = _make_pair("pA")
+        qo, qm = _make_pair("pB")
+        singles = [_make_single(f"s{i}") for i in range(3)]
+        batch = [singles[0], po, pm, singles[1], qo, qm, singles[2]]
+        pools = _pools_with(list(batch))
+
+        restored = _batch_from_ids(pools, _batch_unit_ids(batch))
+        assert [q.question_id for q in restored] == [q.question_id for q in batch]
+
+    def test_missing_unit_id_raises(self) -> None:
+        """checkpoint 引用了诊断池不存在的 unit_id 时硬失败（防静默兜底）。"""
+        pools = _pools_with([_make_single("s1")])
+        try:
+            _batch_from_ids(pools, ["s1", "ghost"])
+        except KeyError:
+            return
+        raise AssertionError("引用缺失 unit_id 应触发 KeyError")
+
+
+# ---------------------------------------------------------------------------
+# 真实落盘往返：write_checkpoint → load_checkpoint → 重建 batch
+# ---------------------------------------------------------------------------
+
+
+class TestCheckpointPersistenceRoundtrip:
+    """unit_id 序列经真实 checkpoint.json 原子写/读回后无损、重建逐字节一致。"""
+
+    def test_persisted_unit_ids_rebuild_batches(self, tmp_path: Path) -> None:
+        po, pm = _make_pair("pA")
+        singles = [_make_single(f"s{i}", task_type="RETRIEVAL") for i in range(6)]
+        items = [*singles, po, pm]
+        correctness = {q.question_id: False for q in items}
+        batches, _ = build_batches(items, correctness, batch_size=4, min_class_per_batch=2, seed=7)
+        pools = _pools_with(list(items))
+
+        # 存 unit_id 序列（train() 落盘路径）
+        epoch_batches = [_batch_unit_ids(b) for b in batches]
+        write_checkpoint(
+            tmp_path,
+            state=_FakeState(correctness=correctness),
+            epoch=1,
+            step_completed=0,
+            phase="in_epoch",
+            global_step=1,
+            total_steps=len(batches),
+            version_snapshot={"skills": "skills/v1", "prompts": "prompts/v1"},
+            epoch_batches=epoch_batches,
+            config=_FakeConfig(),
+        )
+
+        ckpt = load_checkpoint(tmp_path)
+        assert ckpt is not None
+        rebuilt = [_batch_from_ids(pools, ids) for ids in ckpt["epoch_batches"]]
+
+        # 落盘往返后逐批逐字节还原
+        assert [[q.question_id for q in b] for b in rebuilt] == [
+            [q.question_id for q in b] for b in batches
+        ]
+        # 含孪生对的批次恢复后 pair 仍成对
+        for b in rebuilt:
+            for u in validate_units(build_units(b)):
+                if u.kind == "pair":
+                    assert u.size == 2
diff --git a/tests/unit/test_non_ar_byte_identical.py b/tests/unit/test_non_ar_byte_identical.py
new file mode 100644
index 0000000..0cee871
--- /dev/null
+++ b/tests/unit/test_non_ar_byte_identical.py
@@ -0,0 +1,297 @@
+"""纯非 AR 题库端到端 byte-identical 黄金测试（Task 10）。
+
+Phase 1 验收铁律：引入 QuestionUnit 后，纯 single 题库过 pools 抽样 → batching 分批
+→ checkpoint 折叠/恢复 → momentum 采样，端到端结果与"引入 QuestionUnit 前"的逐题
+旧算法逐字节一致。golden 用固定 seed 的确定性对照（忠实重实现旧逐题逻辑），非空断言。
+
+momentum 设计偏差：momentum 采样在逐题粒度进行、不折叠 unit，故仅保证纯非 AR 的抽样
+序列 byte-identical；Phase 1 不保证混格 momentum byte-identical（见 runner
+_sample_momentum_candidates docstring）。本测试保守证明纯非 AR momentum 不漂。
+"""
+
+from __future__ import annotations
+
+import math
+import random
+
+from app.harness.batching import build_batches
+from app.harness.pools import build_pools
+from app.harness.question_units import build_units
+from app.harness.runner import _batch_from_ids, _batch_unit_ids, _sample_momentum_candidates
+from core.types import GeneratedQuestion
+
+# ---------------------------------------------------------------------------
+# 辅助构造
+# ---------------------------------------------------------------------------
+
+
+def _single(qid: str, task_type: str = "RETRIEVAL", video_id: str = "v1") -> GeneratedQuestion:
+    """构造最小 single 题目。"""
+    return GeneratedQuestion(
+        question_id=qid,
+        video_id=video_id,
+        task_type=task_type,
+        question=f"q_{qid}",
+        options=("A. a", "B. b", "C. c", "D. d"),
+        answer="A",
+        source_nodes=("n1",),
+        difficulty="medium",
+        question_role="single",
+    )
+
+
+def _pair(
+    pair_id: str, task_type: str = "AR", video_id: str = "v1"
+) -> tuple[GeneratedQuestion, GeneratedQuestion]:
+    """构造一个孪生对（original + mirror）。"""
+    common = {
+        "video_id": video_id,
+        "task_type": task_type,
+        "question": "?",
+        "options": ("A. a", "B. b", "C. c", "D. d"),
+        "answer": "A",
+        "source_nodes": ("n1",),
+        "difficulty": "medium",
+        "pair_id": pair_id,
+        "flip_axis": "before_after",
+    }
+    original = GeneratedQuestion(
+        question_id=f"{pair_id}_o", question_role="pair_original", **common
+    )
+    mirror = GeneratedQuestion(question_id=f"{pair_id}_m", question_role="pair_mirror", **common)
+    return original, mirror
+
+
+# ---------------------------------------------------------------------------
+# 引入 QuestionUnit 前的旧逐题 build_batches 忠实副本（单一 rng，无 unit 折叠）
+# ---------------------------------------------------------------------------
+
+
+def _reference_select_mixed(
+    items: list[GeneratedQuestion],
+    correctness: dict[str, bool],
+    correct_ratio: float,
+    rng: random.Random,
+) -> dict[str, list[GeneratedQuestion]]:
+    """旧版 _select_mixed_by_task_type 的忠实副本（逐题、单一 rng）。"""
+    errors_by_type: dict[str, list[GeneratedQuestion]] = {}
+    correct_by_type: dict[str, list[GeneratedQuestion]] = {}
+    for q in items:
+        qid = q.question_id
+        if correctness.get(qid) is False:
+            errors_by_type.setdefault(q.task_type, []).append(q)
+        elif correctness.get(qid, False):
+            correct_by_type.setdefault(q.task_type, []).append(q)
+    if correct_ratio <= 0:
+        return errors_by_type
+    grouped: dict[str, list[GeneratedQuestion]] = {}
+    for task_type in sorted(errors_by_type):
+        errs = errors_by_type[task_type]
+        n_correct = round(len(errs) * correct_ratio / (1 - correct_ratio))
+        available = correct_by_type.get(task_type, [])
+        sampled = (
+            list(available) if len(available) <= n_correct else rng.sample(available, n_correct)
+        )
+        grouped[task_type] = errs + sampled
+    return grouped
+
+
+def _reference_build_batches(
+    items: list[GeneratedQuestion],
+    correctness: dict[str, bool],
+    batch_size: int,
+    min_class_per_batch: int,
+    seed: int,
+    correct_ratio: float = 0.0,
+) -> list[list[GeneratedQuestion]]:
+    """引入 QuestionUnit 前的旧版 build_batches 忠实副本（逐题、单一 rng）。"""
+    rng = random.Random(seed)
+    grouped = _reference_select_mixed(items, correctness, correct_ratio, rng)
+    total = sum(len(g) for g in grouped.values())
+    if total == 0:
+        return []
+    nb = max(1, math.ceil(total / batch_size))
+    batches: list[list[GeneratedQuestion]] = [[] for _ in range(nb)]
+    small = {t: g for t, g in grouped.items() if len(g) <= min_class_per_batch}
+    large = {t: g for t, g in grouped.items() if len(g) > min_class_per_batch}
+    for t in sorted(small, key=lambda t: (-len(small[t]), t)):
+        group = small[t]
+        placed = False
+        for b in batches:
+            if len(b) + len(group) <= batch_size:
+                b.extend(group)
+                placed = True
+                break
+        if not placed:
+            batches.append(list(group))
+    nb_live = len(batches)
+    pointer = 0
+    for task_type in sorted(large):
+        group = list(large[task_type])
+        rng.shuffle(group)
+        for q in group:
+            for offset in range(nb_live):
+                idx = (pointer + offset) % nb_live
+                if len(batches[idx]) < batch_size:
+                    batches[idx].append(q)
+                    pointer = (idx + 1) % nb_live
+                    break
+    return [b for b in batches if b]
+
+
+def _ids(batches: list[list[GeneratedQuestion]]) -> list[list[str]]:
+    """提取 batch 的 question_id 序列，便于逐字节对比。"""
+    return [[q.question_id for q in b] for b in batches]
+
+
+# ---------------------------------------------------------------------------
+# (1) 纯非 AR 走 size=1 unit：unit_id == question_id
+# ---------------------------------------------------------------------------
+
+
+class TestPureSingleAreSizeOneUnits:
+    """纯 single 输入折叠出的每个 unit 都是 size=1、unit_id 等于 question_id。"""
+
+    def test_all_units_size_one(self) -> None:
+        items = [_single(f"s{i}", task_type=f"t{i % 3}") for i in range(15)]
+        units = build_units(items)
+        assert len(units) == len(items)
+        assert all(u.kind == "single" and u.size == 1 for u in units)
+        assert [u.unit_id for u in units] == [q.question_id for q in items]
+
+
+# ---------------------------------------------------------------------------
+# (2) batching 逐字节一致（黄金对照）
+# ---------------------------------------------------------------------------
+
+
+class TestBatchingByteIdentical:
+    """纯 single build_batches 与旧逐题算法逐字节一致。"""
+
+    def _assert(self, items, correctness, batch_size, min_cls, seed, ratio) -> None:
+        got, _ = build_batches(items, correctness, batch_size, min_cls, seed, ratio)
+        ref = _reference_build_batches(items, correctness, batch_size, min_cls, seed, ratio)
+        assert _ids(got) == _ids(ref)
+
+    def test_pure_errors(self) -> None:
+        items = [_single(f"q{i}", task_type=f"t{i % 3}") for i in range(20)]
+        correctness = {f"q{i}": False for i in range(20)}
+        self._assert(items, correctness, 5, 2, 42, 0.0)
+
+    def test_mixed_ratio(self) -> None:
+        items = [_single(f"q{i}", task_type=f"t{i % 4}") for i in range(40)]
+        correctness = {f"q{i}": (i % 3 == 0) for i in range(40)}
+        self._assert(items, correctness, 8, 3, 7, 0.5)
+
+
+# ---------------------------------------------------------------------------
+# (3) 端到端：pools 抽样 → batching → checkpoint 折叠/恢复 逐字节一致
+# ---------------------------------------------------------------------------
+
+
+class TestEndToEndByteIdentical:
+    """pools → batching → checkpoint unit_id round-trip 全链纯非 AR 逐字节稳定。"""
+
+    def _make_pool_questions(self) -> tuple[list[GeneratedQuestion], dict[str, bool]]:
+        items = [_single(f"s{i}", task_type=f"t{i % 3}") for i in range(60)]
+        correctness = {f"s{i}": (i % 2 == 0) for i in range(60)}
+        return items, correctness
+
+    def test_pools_deterministic_and_pure_single(self) -> None:
+        items, correctness = self._make_pool_questions()
+        cfg = {
+            "diag_cfg": {
+                "size": 18,
+                "correct_ratio": 0.5,
+                "task_types": None,
+                "seed": 5,
+                "min_per_class": None,
+            },
+            "val_cfg": {
+                "size": 12,
+                "correct_ratio": 0.5,
+                "task_types": None,
+                "seed": 5,
+                "min_per_class": None,
+            },
+            "test_cfg": {"size": 10, "seed": 5},
+            "baseline_run_id": "baseline",
+        }
+        p1 = build_pools(items, correctness, **cfg)
+        p2 = build_pools(items, correctness, **cfg)
+        # 确定性：同 seed 同输入池划分逐字节一致
+        assert [q.question_id for q in p1.diagnosis] == [q.question_id for q in p2.diagnosis]
+        # 诊断池纯 single：折叠出的 unit 与逐题一一对应
+        units = build_units(p1.diagnosis)
+        assert [u.unit_id for u in units] == [q.question_id for q in p1.diagnosis]
+
+    def test_batching_then_checkpoint_roundtrip(self) -> None:
+        items, correctness = self._make_pool_questions()
+        p = build_pools(
+            items,
+            correctness,
+            diag_cfg={
+                "size": 18,
+                "correct_ratio": 0.5,
+                "task_types": None,
+                "seed": 5,
+                "min_per_class": None,
+            },
+            val_cfg={
+                "size": 12,
+                "correct_ratio": 0.5,
+                "task_types": None,
+                "seed": 5,
+                "min_per_class": None,
+            },
+            test_cfg={"size": 10, "seed": 5},
+            baseline_run_id="baseline",
+        )
+        batches, _ = build_batches(p.diagnosis, p.correctness, 6, 2, seed=3, correct_ratio=0.5)
+
+        # checkpoint 折叠为 unit_id（纯 single 即 question_id）后恢复
+        epoch_batches = [_batch_unit_ids(b) for b in batches]
+        assert epoch_batches == _ids(batches)  # 纯非 AR：unit_id 序列 == question_id 序列
+        rebuilt = [_batch_from_ids(p, ids) for ids in epoch_batches]
+        # 恢复的 batch 与原 batch 逐字节一致（inference 消费的题序不变）
+        assert _ids(rebuilt) == _ids(batches)
+
+
+# ---------------------------------------------------------------------------
+# (4) momentum 纯非 AR byte-identical（混格偏差已在 runner docstring 记录）
+# ---------------------------------------------------------------------------
+
+
+class TestMomentumPureNonARByteIdentical:
+    """momentum 采样纯非 AR 与旧逐题 random.Random(epoch).sample 逐字节一致。"""
+
+    def test_matches_legacy_rng(self) -> None:
+        pool = [_single(f"s{i}", task_type="RETRIEVAL") for i in range(30)]
+        allowed = {"RETRIEVAL"}
+        epoch = 11
+        samples = 8
+        got = _sample_momentum_candidates(pool, allowed, samples, epoch)
+
+        candidates = [q for q in pool if q.task_type in allowed]
+        ref = random.Random(epoch).sample(candidates, samples)
+        assert [q.question_id for q in got] == [q.question_id for q in ref]
+
+    def test_ar_pairs_of_other_type_do_not_shift(self) -> None:
+        """向池中加入其它题型的 AR pair，不改变非 AR momentum 的抽样序列。"""
+        singles = [_single(f"s{i}", task_type="RETRIEVAL") for i in range(30)]
+        allowed = {"RETRIEVAL"}
+        epoch = 11
+        samples = 8
+        base = _sample_momentum_candidates(singles, allowed, samples, epoch)
+
+        mixed = list(singles)
+        for k in range(4):
+            po, pm = _pair(f"p{k}", task_type="AR")
+            mixed.extend([po, pm])
+        after = _sample_momentum_candidates(mixed, allowed, samples, epoch)
+        assert [q.question_id for q in base] == [q.question_id for q in after]
+
+    def test_fewer_candidates_than_samples_returns_all(self) -> None:
+        pool = [_single(f"s{i}", task_type="RETRIEVAL") for i in range(3)]
+        got = _sample_momentum_candidates(pool, {"RETRIEVAL"}, 10, epoch=1)
+        assert {q.question_id for q in got} == {"s0", "s1", "s2"}