diff --git a/docs/00-roadmap.md b/docs/00-roadmap.md index 0fad931..c9daa68 100644 --- a/docs/00-roadmap.md +++ b/docs/00-roadmap.md @@ -11,7 +11,8 @@ - **当前层**: 层 1(SFT 基线),待开工 - **层 0**: ✅ 已关账(2026-07-18)。两端 pytest 4/4 全绿;本地 env `ars-opd`(torch 2.13 cu130,4070Ti 可做小规模 GPU 调试);远程 env `/data/zym/envs/ars-opd`(torch 2.10 cu128,8 卡可见);gitea 双端打通(SSH 222) - **已完成学习**: 第一章全部精讲(式 1-8、§3.1-3.6、detach 命门专题);第二章已写好待读 -- **下一步**: Claude 编写 T1→T3→T4→T2→T5(工作模式已改:Claude 写码、用户精读提问,见 CLAUDE.md URGENT.1);默认参数已通过(teacher=MiniMax-M3 经自建网关 / enable_thinking=False / max_length=4096 / 1k 子集);T1/T3/T4 已完成入库 +- **层 1 代码全部完成**(T1 configs / T3 data / T4 trainer / T2 teacher / T5 train_sft 脚本,46+ 单测全绿;teacher=MiniMax-M3 经自建网关,2026-07-18 改定) +- **下一步(运行阶段)**: ① 本地跑完 teacher 生成 1k(scripts/generate_teacher_completions.py,断点续传)→ ② scp parquet+缓存到远程 data/ → ③ 远程 `bash scripts/train_sft.sh sanity`(看首样本自检与 loss 下降)→ ④ `bash scripts/train_sft.sh` 正式 1 epoch → 关账判据见 docs/02 §5 - **层 1 讨论已完成的**: docs/02 全部难点已精讲(collator 五步流水线与坑二、FSDP 决策与 DDP 触发点、删除/替代清单逐项理由、hash 不稳定演示) - **未精讲的文档账**: docs/01 的 §3.7(KL 锚三处实现差异)、§3.8(论文外稳定器)、§4(训练步流程走读) - **未精讲的文档账**: docs/01 的 §3.7(KL 锚三处实现差异)、§3.8(论文外稳定器)、§4(训练步流程走读) diff --git a/scripts/train_sft.py b/scripts/train_sft.py new file mode 100644 index 0000000..ea23ace --- /dev/null +++ b/scripts/train_sft.py @@ -0,0 +1,148 @@ +"""层 1:SFT 基线训练入口(由 train_sft.sh 经 torchrun 启动,勿直接 python 运行)。 + +自包含实验脚本:全部参数写死在下方 FULL 配置里,零参数复现;sanity 模式只是 +对 FULL 的两处显式覆盖(50 步 + 独立输出目录)。 +""" + +# ---- FSDP 前置块(必须在一切 transformers/accelerate import 之前)---- +# 非显然约束:FSDP 的激活检查点开关是 accelerate 在 import 时读取的环境变量 +# (参考实现 train_distillation.py:15-21 的著名坑);写在 import 后会静默无效。 +# DDP 下本变量是无害 no-op——现在就位是为了未来换 4B 学生/FSDP 时只改此处一行, +# 且改完必须 nvidia-smi 实测显存验证生效(docs/02 §2.6)。 +import os + +os.environ.setdefault("FSDP_ACTIVATION_CHECKPOINTING", "false") + +import dataclasses +import sys + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments + +from ars_opd.configs import SFTConfig +from ars_opd.data import IGNORE_INDEX, SFTCollator, load_sft_dataset +from ars_opd.trainer import SFTTrainer + +STUDENT_MODEL = "Qwen/Qwen3-0.6B" + +FULL = SFTConfig( + dataset_path="data/dapo-math-17k-unique.parquet", + output_dir="/data/zym/outputs/sft_qwen3-0.6b_dapo1k", + teacher_completions_path="data/teacher_completions_dapo1k_minimax-m3.jsonl", + subset_size=1000, + seed=42, # 非显然约束:与 generate_teacher_completions.py 一致,否则缓存大面积 miss + max_length=4096, + max_prompt_length=1024, + enable_thinking=False, + learning_rate=2e-5, + per_device_train_batch_size=8, + gradient_accumulation_steps=2, # 全局 batch = 8 × 4 卡 × 2 = 64 + num_train_epochs=1, + max_steps=-1, + lr_scheduler_type="linear", + warmup_ratio=0.0, + gradient_checkpointing=False, + bf16=True, + logging_steps=1, + save_steps=100, + save_total_limit=2, + report_to="none", # 层 1 先靠 tmux 实时日志;W&B 触发条件见 appendix C 表 +) + + +def build_config() -> SFTConfig: + """按命令行模式产出配置。frozen dataclass 的换参方式:replace 构造新实例。""" + mode = sys.argv[1] if len(sys.argv) > 1 else "full" + if mode == "full": + return FULL + if mode == "sanity": + return dataclasses.replace( + FULL, max_steps=50, output_dir=FULL.output_dir + "-sanity" + ) + raise ValueError(f"未知模式 {mode!r},只接受 full / sanity") + + +def smoke_check_first_batch(dataset, collator, tokenizer) -> None: + """训练前解码第一个 batch 供肉眼核对(只在 rank0 打印一次)。 + + 单测用玩具 tokenizer 钉死了预算/边界的算法(tests/test_data.py),但真 + tokenizer 的模板渲染只能在这里肉眼验证:掩码边界是否落在 assistant 起点、 + no-think 时空 块是否在 prompt 侧。这是参考实现"一次性诊断打印" + 的合理化版本(docs/02 §2.3)。 + """ + batch = collator([dataset[0]]) + ids, labels = batch["input_ids"][0], batch["labels"][0] + masked = labels == IGNORE_INDEX + prompt_text = tokenizer.decode(ids[masked], skip_special_tokens=False) + completion_text = tokenizer.decode(ids[~masked], skip_special_tokens=False) + print( + "=" * 30 + + " 首样本自检(人工核对掩码边界)" + + "=" * 30 + + f"\n[prompt 段 | {int(masked.sum())} tok | 不产生 loss]\n" + + f"…{prompt_text[-300:]}\n" + + f"\n[completion 段 | {int((~masked).sum())} tok | 监督目标]\n" + + f"{completion_text[:300]}…\n" + + "=" * 80, + flush=True, + ) + + +def main() -> None: + cfg = build_config() + rank0 = int(os.environ.get("RANK", "0")) == 0 + + # 加载顺序刻意 fail-fast:数据(毫秒级,最易配错)→ tokenizer(几 MB)→ + # 模型(GB 级下载)。teacher 缓存缺失要在下模型之前炸出来 + dataset = load_sft_dataset(cfg) + tokenizer = AutoTokenizer.from_pretrained(STUDENT_MODEL) + collator = SFTCollator( + tokenizer, + max_length=cfg.max_length, + max_prompt_length=cfg.max_prompt_length, + enable_thinking=cfg.enable_thinking, + ) + if rank0: + smoke_check_first_batch(dataset, collator, tokenizer) + + model = AutoModelForCausalLM.from_pretrained(STUDENT_MODEL, dtype=torch.float32) + + args = TrainingArguments( + output_dir=cfg.output_dir, + # 非显然约束:必须关掉列裁剪。HF Trainer 默认删除模型 forward 签名里 + # 没有的数据列——"messages" 会被整列删光,collator 收到空字典且不报错 + remove_unused_columns=False, + learning_rate=cfg.learning_rate, + per_device_train_batch_size=cfg.per_device_train_batch_size, + gradient_accumulation_steps=cfg.gradient_accumulation_steps, + num_train_epochs=cfg.num_train_epochs, + max_steps=cfg.max_steps, + lr_scheduler_type=cfg.lr_scheduler_type, + warmup_ratio=cfg.warmup_ratio, + gradient_checkpointing=cfg.gradient_checkpointing, + bf16=cfg.bf16, + seed=cfg.seed, + logging_steps=cfg.logging_steps, + logging_first_step=True, + save_strategy="steps", + save_steps=cfg.save_steps, + save_total_limit=cfg.save_total_limit, + report_to=cfg.report_to, + ddp_find_unused_parameters=False, # 全参训练无闲置参数,省一次全模型扫描 + dataloader_num_workers=2, # collator 逐 batch 分词在 CPU,双 worker 与 GPU 重叠 + ) + trainer = SFTTrainer( + model=model, + args=args, + train_dataset=dataset, + data_collator=collator, + ) + trainer.train() + trainer.save_model() # 终态模型(save_pretrained 格式,含 config) + if rank0: + tokenizer.save_pretrained(cfg.output_dir) + print(f"训练完成,模型已存至 {cfg.output_dir}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/train_sft.sh b/scripts/train_sft.sh new file mode 100644 index 0000000..9757c3d --- /dev/null +++ b/scripts/train_sft.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# 层 1:SFT 基线训练(远程 gpu-a800-060 专用;本地不跑训练)。 +# +# 用法(tmux 内执行,日志实时可查): +# bash scripts/train_sft.sh sanity # 50 步冒烟:看首样本自检 + loss 是否从 ~2-3 下降 +# bash scripts/train_sft.sh # 正式:1k 子集 1 epoch +# +# 前置检查清单: +# 1. nvidia-smi 确认下方 GPUS 四张卡空闲(只许用 8 卡中的 4 张,严禁自动选卡) +# 2. data/ 下已有两个文件(gitignore 不随 git 走,本地 scp 上来): +# scp data/dapo-math-17k-unique.parquet data/teacher_completions_dapo1k_minimax-m3.jsonl \ +# <远程>:/data/zym/ars-opd-rebuild/data/ +# 3. 代码是最新:git -C /data/zym/ars-opd-rebuild pull +set -euo pipefail +cd "$(dirname "$0")/.." # 锚定仓库根:py 内 data/... 相对路径以此为基准 + +GPUS=0,1,2,3 # ⚠️ 改这里前先 nvidia-smi +MODE=${1:-full} + +export CUDA_VISIBLE_DEVICES=$GPUS +export PYTHONUNBUFFERED=1 # 禁止日志缓存(CLAUDE.md §5) +export HF_ENDPOINT=${HF_ENDPOINT:-https://hf-mirror.com} +export HF_HOME=${HF_HOME:-/data/zym/hf_cache} # 模型缓存落 /data,根分区已满 + +torchrun --nproc_per_node=4 --master_port=29571 scripts/train_sft.py "$MODE"