ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter8/cot-distillation/train_student.py

@@ -0,0 +1,341 @@
+#!/usr/bin/env python3
+"""Train the Experiment 8-9 student on verified teacher CoT trajectories.
+
+This is the parameter-update stage missing from the original collection-only
+companion.  It deliberately has no mock training mode: a successful run writes
+a real Hugging Face/PEFT checkpoint plus a provenance manifest.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import importlib.metadata
+import importlib.util
+import json
+import os
+import re
+import platform
+import subprocess
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+
+def sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def load_verified_messages(path: Path) -> list[list[dict[str, str]]]:
+    """Load only complete user/assistant rows with a non-empty final answer."""
+    rows: list[list[dict[str, str]]] = []
+    with path.open(encoding="utf-8") as handle:
+        for line_number, line in enumerate(handle, 1):
+            if not line.strip():
+                continue
+            item = json.loads(line)
+            messages = item.get("messages")
+            if not isinstance(messages, list) or len(messages) != 2:
+                raise ValueError(f"{path}:{line_number}: expected exactly two messages")
+            if [m.get("role") for m in messages] != ["user", "assistant"]:
+                raise ValueError(f"{path}:{line_number}: expected user then assistant")
+            if not all(isinstance(m.get("content"), str) and m["content"].strip() for m in messages):
+                raise ValueError(f"{path}:{line_number}: empty message content")
+            if not re.search(r"Final Answer[:：]", messages[1]["content"], re.IGNORECASE):
+                raise ValueError(f"{path}:{line_number}: assistant lacks verified Final Answer")
+            rows.append(messages)
+    if not rows:
+        raise ValueError(f"{path}: no training samples")
+    return rows
+
+
+@dataclass
+class EncodedExample:
+    input_ids: list[int]
+    labels: list[int]
+
+
+def _chat_template_ids(encoded: Any) -> list[int]:
+    """Normalize Transformers 4.x/5.x chat-template return values.
+
+    Transformers 4.x returned a bare list from ``apply_chat_template`` when
+    ``tokenize=True``.  Transformers 5.x returns a BatchEncoding containing
+    both ``input_ids`` and ``attention_mask``.  Calling ``len`` or slicing the
+    latter operates on mapping keys, which can make every assistant trajectory
+    appear to have only two tokens and defeats the loss-mask safety check.
+    """
+    if isinstance(encoded, dict) or hasattr(encoded, "keys"):
+        encoded = encoded["input_ids"]
+    if hasattr(encoded, "tolist"):
+        encoded = encoded.tolist()
+    if encoded and isinstance(encoded[0], list):
+        if len(encoded) != 1:
+            raise ValueError("expected one chat-template sequence")
+        encoded = encoded[0]
+    if not isinstance(encoded, list) or not all(isinstance(token, int) for token in encoded):
+        raise TypeError("chat template did not return a one-dimensional integer token sequence")
+    return encoded
+
+
+def encode_messages(tokenizer: Any, messages: list[dict[str, str]], max_length: int) -> EncodedExample:
+    """Mask user/prompt tokens and supervise only the teacher assistant trajectory."""
+    prompt_ids = _chat_template_ids(
+        tokenizer.apply_chat_template(messages[:1], tokenize=True, add_generation_prompt=True)
+    )
+    full_ids = _chat_template_ids(
+        tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False)
+    )
+    if len(full_ids) > max_length:
+        full_ids = full_ids[:max_length]
+    prompt_length = min(len(prompt_ids), len(full_ids))
+    labels = [-100] * prompt_length + full_ids[prompt_length:]
+    if not any(label != -100 for label in labels):
+        raise ValueError("max_length truncates the entire assistant response")
+    return EncodedExample(input_ids=full_ids, labels=labels)
+
+
+def _git_commit(root: Path) -> str | None:
+    try:
+        return subprocess.run(
+            ["git", "rev-parse", "HEAD"], cwd=root, check=True,
+            capture_output=True, text=True,
+        ).stdout.strip()
+    except (OSError, subprocess.CalledProcessError):
+        return None
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Experiment 8-9: real student SFT on verified CoT trajectories",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--train-data", type=Path, default=Path("data/sft_cot_distill_aime_kimi_k3.jsonl"))
+    parser.add_argument("--base-model", default="Qwen/Qwen2.5-1.5B-Instruct")
+    parser.add_argument("--output-dir", type=Path, default=Path("checkpoints/cot-student"))
+    parser.add_argument("--max-length", type=int, default=4096)
+    parser.add_argument("--epochs", type=float, default=3.0)
+    parser.add_argument("--learning-rate", type=float, default=2e-5)
+    parser.add_argument("--batch-size", type=int, default=1)
+    parser.add_argument("--gradient-accumulation", type=int, default=16)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--lora-rank", type=int, default=32, help="0 disables LoRA and updates all weights")
+    parser.add_argument("--lora-alpha", type=int, default=64)
+    parser.add_argument("--gradient-checkpointing", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--trust-remote-code", action="store_true")
+    parser.add_argument("--preflight", action="store_true", help="write dependency/GPU readiness evidence without training")
+    parser.add_argument("--preflight-output", type=Path, default=Path("validation/student_sft_preflight.json"))
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if args.max_length <= 0 or args.batch_size <= 0 or args.gradient_accumulation <= 0:
+        raise SystemExit("max-length, batch-size, and gradient-accumulation must be positive")
+    messages = load_verified_messages(args.train_data)
+
+    if args.preflight:
+        dependencies = {
+            name: importlib.util.find_spec(name) is not None
+            for name in ("torch", "transformers", "accelerate", "peft")
+        }
+        dependency_versions = {
+            name: importlib.metadata.version(name) if installed else None
+            for name, installed in dependencies.items()
+        }
+        cuda_available = False
+        gpu_names: list[str] = []
+        torch_version = None
+        trainer_stack_error = None
+        if dependencies["torch"]:
+            import torch
+            torch_version = torch.__version__
+            cuda_available = torch.cuda.is_available()
+            if cuda_available:
+                gpu_names = [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]
+        try:
+            from transformers import Trainer  # noqa: F401
+        except Exception as exc:  # integration errors include incompatible peft/transformers versions
+            trainer_stack_error = f"{type(exc).__name__}: {exc}"
+        trainer_stack_importable = trainer_stack_error is None
+        payload = {
+            "schema_version": 1,
+            "experiment": "8-9",
+            "stage": "student_sft_preflight",
+            "generated_at_utc": datetime.now(timezone.utc).isoformat(),
+            "ready": all(dependencies.values()) and trainer_stack_importable and cuda_available,
+            "training_data": {
+                "path": str(args.train_data.resolve()),
+                "sha256": sha256(args.train_data.resolve()),
+                "samples": len(messages),
+            },
+            "host": {
+                "platform": platform.platform(),
+                "machine": platform.machine(),
+                "torch": torch_version,
+                "cuda_available": cuda_available,
+                "gpu_names": gpu_names,
+            },
+            "dependencies": dependencies,
+            "dependency_versions": dependency_versions,
+            "trainer_stack_importable": trainer_stack_importable,
+            "trainer_stack_error": trainer_stack_error,
+            "blockers": [
+                *[f"missing Python dependency: {name}" for name, ok in dependencies.items() if not ok],
+                *([] if trainer_stack_importable else ["transformers/peft trainer stack is not importable"]),
+                *([] if cuda_available else ["no CUDA device available"]),
+            ],
+        }
+        args.preflight_output.parent.mkdir(parents=True, exist_ok=True)
+        args.preflight_output.write_text(
+            json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
+        )
+        print(json.dumps({"preflight": str(args.preflight_output), "ready": payload["ready"]}, ensure_ascii=False))
+        return
+
+    try:
+        import torch
+    except ImportError as exc:
+        raise SystemExit("PyTorch is missing. Install requirements.txt before training.") from exc
+    if not torch.cuda.is_available():
+        raise SystemExit(
+            "Experiment 8-9 student SFT requires a CUDA host; this runner has no synthetic/CPU success fallback."
+        )
+    try:
+        from torch.utils.data import Dataset
+        from transformers import (
+            AutoModelForCausalLM,
+            AutoTokenizer,
+            Trainer,
+            TrainingArguments,
+            set_seed,
+        )
+    except (ImportError, RuntimeError) as exc:
+        raise SystemExit(
+            f"The transformers/peft training stack is not importable: {type(exc).__name__}: {exc}"
+        ) from exc
+
+    set_seed(args.seed)
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.base_model, trust_remote_code=args.trust_remote_code
+    )
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    encoded = [encode_messages(tokenizer, item, args.max_length) for item in messages]
+
+    class CotDataset(Dataset):
+        def __len__(self) -> int:
+            return len(encoded)
+
+        def __getitem__(self, index: int) -> dict[str, list[int]]:
+            item = encoded[index]
+            return {"input_ids": item.input_ids, "labels": item.labels}
+
+    def collate(batch: list[dict[str, list[int]]]) -> dict[str, Any]:
+        width = max(len(item["input_ids"]) for item in batch)
+        ids, masks, labels = [], [], []
+        for item in batch:
+            padding = width - len(item["input_ids"])
+            ids.append(item["input_ids"] + [tokenizer.pad_token_id] * padding)
+            masks.append([1] * len(item["input_ids"]) + [0] * padding)
+            labels.append(item["labels"] + [-100] * padding)
+        return {
+            "input_ids": torch.tensor(ids, dtype=torch.long),
+            "attention_mask": torch.tensor(masks, dtype=torch.long),
+            "labels": torch.tensor(labels, dtype=torch.long),
+        }
+
+    model = AutoModelForCausalLM.from_pretrained(
+        args.base_model,
+        torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
+        trust_remote_code=args.trust_remote_code,
+    )
+    if args.gradient_checkpointing:
+        model.gradient_checkpointing_enable()
+        model.config.use_cache = False
+    if args.lora_rank:
+        try:
+            from peft import LoraConfig, get_peft_model
+        except ImportError as exc:
+            raise SystemExit("LoRA requested but peft is not installed") from exc
+        model = get_peft_model(model, LoraConfig(
+            r=args.lora_rank,
+            lora_alpha=args.lora_alpha,
+            lora_dropout=0.05,
+            bias="none",
+            task_type="CAUSAL_LM",
+            target_modules="all-linear",
+        ))
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    training_args = TrainingArguments(
+        output_dir=str(args.output_dir),
+        num_train_epochs=args.epochs,
+        learning_rate=args.learning_rate,
+        per_device_train_batch_size=args.batch_size,
+        gradient_accumulation_steps=args.gradient_accumulation,
+        logging_steps=1,
+        save_strategy="epoch",
+        seed=args.seed,
+        bf16=torch.cuda.is_bf16_supported(),
+        fp16=not torch.cuda.is_bf16_supported(),
+        report_to="none",
+        remove_unused_columns=False,
+    )
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=CotDataset(),
+        data_collator=collate,
+    )
+    result = trainer.train()
+    trainer.save_model(str(args.output_dir))
+    tokenizer.save_pretrained(str(args.output_dir))
+
+    root = Path(__file__).resolve().parents[2]
+    manifest = {
+        "schema_version": 1,
+        "experiment": "8-9",
+        "stage": "student_sft",
+        "status": "complete",
+        "generated_at_utc": datetime.now(timezone.utc).isoformat(),
+        "book_git_commit": _git_commit(root),
+        "training_data": {
+            "path": str(args.train_data.resolve()),
+            "sha256": sha256(args.train_data.resolve()),
+            "samples": len(messages),
+        },
+        "base_model": args.base_model,
+        "output_dir": str(args.output_dir.resolve()),
+        "host": {
+            "platform": platform.platform(),
+            "gpu_names": [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())],
+            "torch": torch.__version__,
+        },
+        "dependency_versions": {
+            name: importlib.metadata.version(name)
+            for name in ("torch", "transformers", "accelerate", "peft")
+        },
+        "training": {
+            "epochs": args.epochs,
+            "learning_rate": args.learning_rate,
+            "max_length": args.max_length,
+            "batch_size": args.batch_size,
+            "gradient_accumulation": args.gradient_accumulation,
+            "lora_rank": args.lora_rank,
+            "seed": args.seed,
+            "metrics": result.metrics,
+        },
+    }
+    (args.output_dir / "training_manifest.json").write_text(
+        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
+    )
+    print(json.dumps({"checkpoint": str(args.output_dir), "samples": len(messages)}, ensure_ascii=False))
+
+
+if __name__ == "__main__":
+    main()