Add reproducible single-GPU sub-one-pass optimizer benchmark

docs/sub_one_pass/代码审查.md

@@ -0,0 +1,64 @@
+## 2026-07-28 14:37
+
+## 代码审查报告
+
+### 审查目标
+
+- 类型：本地改动
+- 范围：sub-one-pass token budget、数据读取、20-optimizer runner、结果汇总与实验文档
+- 审查时间：2026-07-28 14:37
+
+### 审查输入证据
+
+- Diff 范围：当前工作树相对 `fdecd5d`
+- 单元测试：`python3 -m unittest discover -s tests -v`，14/14 通过
+- 语法检查：`python3 -m compileall -q src scripts/sub_one_pass tests`，通过
+- 新增文件格式：Black 与 isort，均通过
+- Diff whitespace：`git diff --check`，通过
+- dry-run 数量：smoke 20、tune 60、main 20
+- GPU 检查：尚未执行，卡已关闭
+
+### 已发现并修复的问题
+
+#### 🔴 validation/train 文档边界泄漏
+
+- 位置：`scripts/sub_one_pass/prepare_fineweb_stream.py`
+- 问题：旧实现可能把填满 validation 后的同一源文档剩余 tokens 写入 training。
+- 风险：训练集和验证集共享同一文档，validation loss 偏乐观。
+- 修复：填满 validation 后丢弃当前文档剩余 tokens，从下一篇源文档开始 training；manifest 记录丢弃数量和 split unit。
+- 回归测试：`test_fineweb_stream.py` 验证边界文档不会跨 split。
+
+#### 🟡 单卡队列中断后 manifest 信息丢失
+
+- 位置：`scripts/sub_one_pass/run_all.py`
+- 问题：旧 runner 只在整个队列结束后写 manifest，且再次启动会覆盖同阶段历史。
+- 修复：每次状态变化原子更新；以 experiment directory upsert；保留多次 invocation 和扩展 LR 网格结果；有效 `summary.json` 默认跳过。
+
+#### 🟡 任意单 LR 不足以公平比较
+
+- 位置：`scripts/sub_one_pass/run_all.py`
+- 问题：20 个 optimizer 各使用一个未经同预算筛选的 LR，排名会混合算法差异和调参质量。
+- 修复：增加 tune stage，统一 `0.3×/1×/3×` 网格；增加 `select_lr.py`；边界最优时返回非零状态，阻止直接进入 main。
+
+### 剩余问题
+
+#### 🟡 GPU 兼容性门禁待执行
+
+- `soap`、`sophiag`、`d-muon` 和 `muon-pytorch` 必须在目标 GPU 环境完成 smoke。
+- `muon-pytorch` 需要所有正式实验共用的 PyTorch 环境提供 `torch.optim.Muon`。
+- 在完成 20-way smoke 前不得启动 S1 或报告预计账单时间。
+
+#### 🟢 全仓格式门禁
+
+- 全仓 Black/isort 检查会命中多个未由本实验修改的原始文件。
+- 本轮只验证并格式化新增实验脚本与测试，避免产生无关算法 diff。
+
+### 结论
+
+- 结论：`Approved`（本地设计与实现）
+- 严重问题：0 个未解决
+- 建议修改：1 个外部 GPU 验证门禁
+- 可选优化：1 个全仓格式治理
+- 下一步建议：重新开卡后选择已预装兼容 PyTorch 的镜像，只运行 S0；根据 20 个 optimizer 的实测吞吐量再确认 S1/S2 总时长。
+
+---

docs/sub_one_pass/方案设计.md

@@ -0,0 +1,117 @@
+## 2026-07-28 14:37
+
+## Material Passport
+
+- Origin Skill: experiment-agent
+- Origin Mode: plan
+- Origin Date: 2026-07-28
+- Verification Status: UNVERIFIED
+- Version Label: code_plan_v1
+
+## Experiment Overview
+
+- Title: Fixed-corpus sub-one-pass LLM optimizer benchmark
+- Objective: 比较仓库 20 个 optimizer 在相同训练 token 暴露下的早期样本效率。
+- Hypothesis: 不预设某个 optimizer 必然领先；检验排名是否随 token budget 与 seed 改变。
+- Type: training
+- Working Directory: `/root/autodl-tmp/llm-optimizer-benchmark`
+- Entry Command: `python scripts/sub_one_pass/run_all.py --stage smoke ...`
+- Environment: 单张 NVIDIA A800 80GB；所有正式运行共用同一 PyTorch/CUDA 用户态软件栈。
+
+## 技术方案概述
+
+使用一个固定、预 tokenized 的大语料建立确定性无放回 sequence stream。所有 optimizer 共享模型初始化、数据顺序、batch、gradient accumulation、precision、warmup 和 validation batches。主实验是一条持续到 256M tokens 的轨迹，并在嵌套 token checkpoints 评估。单卡成本通过分阶段漏斗控制，而不是直接运行 20 个完整多种子实验。
+
+## 研究问题与变量
+
+- 研究问题：在固定大语料的 sub-one-pass 区间内，不同 optimizer 达到相同 validation quality 所需的训练 token 是否不同？
+- 自变量：optimizer。
+- 主因变量：固定 token checkpoints 的 validation cross-entropy。
+- 次因变量：log-token learning-curve AUC、tokens-to-threshold、wall-clock、tokens/s、峰值显存、失败率。
+- 控制变量：模型、初始化 seed、训练数据及顺序、validation slice、tokenizer、batch、sequence length、precision、gradient clipping、weight decay、warmup policy。
+- 潜在混杂：各 optimizer 的学习率敏感性、框架 kernel 差异、optimizer step 开销、随机 seed。
+
+## 四阶段实验
+
+| 阶段 | 覆盖 | 默认预算 | 用途 | 放行条件 |
+|---|---:|---:|---|---|
+| S0 smoke | 20 optimizers × 1 seed | 2M tokens | API、OOM、NaN、checkpoint 检查 | 20 个均完成或留下可解释失败 |
+| S1 LR screening | 20 × 3 LR 倍率 × 1 seed | 16M tokens | 为每个 optimizer 选学习率区间 | 每个 optimizer 至少一个稳定配置 |
+| S2 main curve | 20 × 1 seed | 16M–256M nested | 完整早期学习曲线 | 结果齐全且无数据重复 |
+| S3 confirmation | AdamW + 领先组 × seeds 1,2 | 16M–256M nested | 检查排名稳定性 | paired-seed 结果齐全 |
+
+S1 默认倍率为 `0.3, 1.0, 3.0`，倍率同时作用于该 optimizer 的主学习率与配套学习率。若三个候选中最优值落在边界，先向该方向补一个候选，不能直接进入 S2。
+
+## 固定数据协议
+
+- 推荐 train corpus：1,073,741,824 tokens；validation：8,388,608 tokens。
+- 最大训练预算：268,435,456 tokens，对应约 25% 数据暴露。
+- sequence length：512；micro batch：16；gradient accumulation：8。
+- 每个 optimizer step：65,536 target tokens。
+- checkpoints：16,777,216、33,554,432、67,108,864、134,217,728、268,435,456 tokens。
+- 数据顺序由相同 `data_seed` 的确定性仿射排列生成；固定 sequence boundaries；one pass 前不重复 sequence。
+- 每次 validation 从 step 0 开始，使用完全相同的 64 个 batches。
+
+## 模型与调度
+
+- 模型：Llama，8 layers，6 heads，embedding width 384，约 30M 参数。
+- precision：A800 上优先 bfloat16；允许 TF32。
+- scheduler：16-step linear warmup 后 constant LR。该 schedule 不依赖最终 horizon，保证 16M checkpoint 在 256M 轨迹中有明确含义。
+- optimizer 内部 horizon：AdEMAMix 的 `alpha/beta3` warmup 在 S0/S1/S2/S3 均固定为 256M-token horizon，短阶段只截取主轨迹前缀。
+- gradient clipping：0.5；weight decay：0.1，除非算法定义要求单独处理。
+
+## 公平性边界
+
+固定 token 数衡量的是样本效率，不是固定 FLOPs 或固定 wall-clock。SOAP、Muon、Sophia 等 optimizer 的 step 成本不同，因此不能只用 validation loss 排名后声称“计算效率更高”。正式报告必须同时给出：
+
+1. validation loss vs training tokens；
+2. validation loss vs wall-clock；
+3. tokens/s 与峰值显存；
+4. 不稳定或失败配置数量。
+
+## 统计策略
+
+- S2 是全算法 exploratory comparison。
+- S3 对 seed 配对：同一 seed 下各 optimizer 共享完全相同的模型初始化和数据顺序。
+- 报告每个 checkpoint 的 paired difference，而不是只报告各自均值。
+- 三个 seed 只能提供稳定性证据，不足以支撑精细正态近似；优先报告原始 seed 点、median/range，并对领先组补到 5 seeds（若计算预算允许）。
+- 主要结论预先指定为 256M-token validation loss；学习曲线 AUC 为次要结论，避免从五个 checkpoint 中事后挑最有利点。
+
+## 监控配置
+
+- 硬超时：只在用户确认后按阶段设置；除硬超时外不自动终止。
+- 进程监控：PID、GPU utilization、显存、`run.log` 最新更新时间。
+- 结果监控：`manifest_<stage>.json`、每个实验目录的 `summary.json`。
+- 异常策略：NaN、OOM、非零退出立即记录并通知；不静默重试。
+- 成功条件：S0 的 20 个配置均生成有效 summary，或留下可复现且经用户决定如何处理的失败。
+
+## 模块划分
+
+| 模块 | 职责 |
+|---|---|
+| `src/experiment_budget.py` | token budget、steps、data exposure 的精确约束 |
+| `src/data/reader.py` | 固定边界、无放回、常数内存排列 |
+| `scripts/sub_one_pass/run_all.py` | 分阶段单卡队列、LR 倍率、恢复与 manifest |
+| `scripts/sub_one_pass/summarize.py` | 汇总 checkpoint 指标 |
+| `scripts/sub_one_pass/select_lr.py` | 从 S1 结果生成每个 optimizer 的 LR 倍率表 |
+| `scripts/sub_one_pass/estimate_runtime.py` | 用 calibration 吞吐量估算后续 GPU 时间 |
+
+## 关键设计决策
+
+| 决策点 | 选择 | 理由 |
+|---|---|---|
+| pass vs token budget | 固定唯一语料，使用 exposure < 1 | 大语料上 one pass 已超出单卡预算 |
+| 多预算实现 | 一条 nested trajectory | 避免重复训练，且 constant schedule 不依赖 horizon |
+| 调参 | 等预算 LR screening | 单一任意 LR 会把超参数质量误当作算法质量 |
+| 全算法复现 | S2 全 20；S3 只扩种子 | 保留覆盖面，同时控制单卡成本 |
+| 环境 | 所有算法使用同一 PyTorch 版本 | 避免 kernel/数值栈成为 optimizer 混杂因素 |
+
+## 实现计划
+
+1. 审计并测试 20 个 optimizer 的命令生成。
+2. 为 runner 增加 stage、LR scale、multi-seed、已完成跳过、增量 manifest。
+3. 增加 LR 选择与时间估算脚本。
+4. 完成单元测试和 CPU smoke；GPU 开卡后只做 S0 calibration。
+5. calibration 后用实测吞吐量确认 S1/S2 队列，再开始正式训练。
+
+---

docs/sub_one_pass/需求理解.md

@@ -0,0 +1,67 @@
+## 2026-07-28 14:37
+
+## 需求概述
+
+在一张 GPU 上，将仓库当前暴露的 20 个 optimizer 接入统一、可恢复、可审计的 sub-one-pass 语言模型训练协议。研究对象是固定大语料上的早期样本效率，而不是 Datablations，也不是固定总计算量实验。
+
+## 功能需求
+
+### P0 - 必须实现
+
+- [ ] 固定模型结构、tokenizer、训练/验证语料、数据顺序和 token checkpoints。
+- [ ] 训练样本在达到 one pass 前不重复，且精确记录 `tokens_seen / unique_train_tokens`。
+- [ ] 覆盖 `adamw`、`sgd`、`muon`、`soap`、`ademamix`、`lion`、`sf-adamw`、`sf-sgd`、`signsgd`、`signum`、`prodigy`、`sophiag`、`adopt`、`mars`、`adafactor`、`lamb`、`scion`、`scion-light`、`d-muon`、`muon-pytorch`。
+- [ ] 单卡串行执行，单个 optimizer 失败时记录失败并保留已完成结果。
+- [ ] 输出每个 token checkpoint 的 validation loss、perplexity、accuracy、吞吐量、训练时间和峰值显存。
+- [ ] 支持 smoke、学习率筛选、主实验和多种子确认四个阶段。
+
+### P1 - 重要功能
+
+- [ ] 已完成实验默认跳过，允许卡被关闭后继续队列。
+- [ ] manifest 在每个子实验结束后立即原子更新。
+- [ ] 学习率筛选对每个 optimizer 使用相同 token 预算和相同候选倍率。
+- [ ] 主实验能够读取每个 optimizer 选出的学习率倍率。
+- [ ] 结果目录和实验名包含阶段、optimizer、seed、学习率倍率和最大 token budget。
+
+### P2 - 可选功能
+
+- [ ] 根据 GPU calibration 结果自动估算剩余队列时间。
+- [ ] 对 paired-seed 结果输出置信区间和曲线汇总。
+
+## 非功能需求
+
+- 性能：默认约 30M 参数 Llama；每次只运行一个训练进程；数据排列为常数内存。
+- 可恢复性：中断后不覆盖完整结果；失败不静默重试。
+- 可复现性：保存完整命令、随机种子、数据 manifest、软件版本和 git commit。
+- 公平性：主指标按相同训练 token 数比较；wall-clock 与显存作为独立效率指标报告。
+
+## 边界与约束
+
+- 不包含：Datablations、数据量消融、跨模型规模结论、多卡 scaling。
+- 硬件限制：一张 GPU，正式运行前必须先 calibration。
+- 数据限制：最大训练 token budget 必须小于或等于无放回数据流的一次遍历容量。
+- 环境限制：20 个 optimizer 必须在同一 PyTorch 软件栈下比较；不允许仅为单个 optimizer 更换框架版本后直接混合排名。
+
+## 验收标准
+
+### 固定 token 预算
+
+- Given 每步 token 数为 65,536，When 指定 16M、32M、64M、128M、256M checkpoints，Then 精确映射为 256、512、1024、2048、4096 optimizer steps。
+
+### 无放回数据
+
+- Given 固定 data seed 和 fixed boundaries，When 消耗小于一次遍历的数据，Then 每个训练 target sequence 最多出现一次，且不同 optimizer 看到相同顺序。
+
+### 全 optimizer 覆盖
+
+- Given 仓库 CLI 的 optimizer choices，When 执行 runner dry-run，Then 恰好生成 20 个唯一 optimizer 配置，无遗漏、无额外算法。
+
+### 可恢复队列
+
+- Given 某个实验目录已存在有效 `summary.json`，When 重启相同阶段，Then 默认跳过该实验并在 manifest 中标记 `skipped-completed`。
+
+### 公平调参
+
+- Given LR screening 阶段，When 比较所有 optimizer，Then 每个 optimizer 使用相同训练 token budget、相同 seed 集合和相同 LR 倍率网格。
+
+---

exps/sub_one_pass/formal/confirm_selection.json

@@ -0,0 +1,116 @@
+{
+  "selection_rule": "AdamW baseline plus the three non-AdamW optimizers with the lowest S2 final validation loss at 268435456 tokens.",
+  "selected_optimizers": [
+    "adamw",
+    "ademamix",
+    "soap",
+    "muon"
+  ],
+  "seeds": [
+    1,
+    2
+  ],
+  "ranking": [
+    {
+      "optimizer": "ademamix",
+      "final_val_loss": 4.01993989944458,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_ademamix_seed0_lrscale1_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "soap",
+      "final_val_loss": 4.038823127746582,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_soap_seed0_lrscale1_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "muon",
+      "final_val_loss": 4.043658256530762,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_muon_seed0_lrscale3_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "d-muon",
+      "final_val_loss": 4.057147979736328,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_d-muon_seed0_lrscale3_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "adamw",
+      "final_val_loss": 4.083443641662598,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_adamw_seed0_lrscale1_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "prodigy",
+      "final_val_loss": 4.091464042663574,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_prodigy_seed0_lrscale0p3_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "mars",
+      "final_val_loss": 4.1218061447143555,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_mars_seed0_lrscale0p3_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "lion",
+      "final_val_loss": 4.166975021362305,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_lion_seed0_lrscale0p3_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "adopt",
+      "final_val_loss": 4.170770168304443,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_adopt_seed0_lrscale0p3_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "sophiag",
+      "final_val_loss": 4.209933280944824,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_sophiag_seed0_lrscale0p3_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "lamb",
+      "final_val_loss": 4.210540771484375,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_lamb_seed0_lrscale27_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "signum",
+      "final_val_loss": 4.239973068237305,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_signum_seed0_lrscale1_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "sf-adamw",
+      "final_val_loss": 4.304592132568359,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_sf-adamw_seed0_lrscale1_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "adafactor",
+      "final_val_loss": 4.321353912353516,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_adafactor_seed0_lrscale0p3_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "muon-pytorch",
+      "final_val_loss": 4.4571852684021,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_muon-pytorch_seed0_lrscale3_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "signsgd",
+      "final_val_loss": 5.291903495788574,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_signsgd_seed0_lrscale0p3_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "scion",
+      "final_val_loss": 6.072164535522461,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_scion_seed0_lrscale0p3_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "scion-light",
+      "final_val_loss": 6.240045070648193,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_scion-light_seed0_lrscale0p3_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "sgd",
+      "final_val_loss": 6.9730939865112305,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_sgd_seed0_lrscale9_tokens268435456/summary.json"
+    },
+    {
+      "optimizer": "sf-sgd",
+      "final_val_loss": 6.992140293121338,
+      "summary_path": "/root/autodl-tmp/llm-optimizer-results/formal/sub1_main_token-bin_sf-sgd_seed0_lrscale729_tokens268435456/summary.json"
+    }
+  ],
+  "created_at_unix": 1785263889.9600036
+}

exps/sub_one_pass/formal/data_manifests/fineweb-512m-manifest.json

@@ -0,0 +1,17 @@
+{
+  "dataset": "HuggingFaceFW/fineweb",
+  "config": "sample-10BT",
+  "sources": [
+    {
+      "path": "/root/autodl-tmp/llm-optimizer-data/fineweb-source/000_00000.parquet",
+      "size_bytes": 2147292183,
+      "sha256": "6b552ea48424648dc86d00df276f93fdfc55e9ad342ce3e4affc23a3a370792b"
+    }
+  ],
+  "tokenizer": "gpt2",
+  "train_tokens": 536870912,
+  "val_tokens": 8388608,
+  "documents_consumed": 788777,
+  "split_unit": "source_document",
+  "validation_boundary_discarded_tokens": 475
+}

exps/sub_one_pass/formal/data_manifests/preflight-manifest.json

@@ -0,0 +1,9 @@
+{
+  "purpose": "optimizer-compatibility-preflight-only",
+  "synthetic": true,
+  "seed": 20260728,
+  "vocab_size": 50304,
+  "train_tokens": 4194304,
+  "val_tokens": 524288,
+  "eligible_for_scientific_comparison": false
+}