train

下面是更完整版，可直接丢给训练 agent。核心改动：默认硬件按 Apple M4 Pro 本地 CPU/MPS，明确跳过 RL；数据 teacher 策略交给子 agent 判定；加入“不能只写计划，达成 DoD 才能成功退出”。

```text
# 任务：在 InfoLens 上训练一个 Qwen3-0.6B Tiny-NLA（AV + AR）

你是训练执行 agent。你的任务不是写调研报告，而是把一个能跑的 Tiny-NLA 原型训练出来，并交付可复现脚本、checkpoint、评估样例。

请完整读完本任务书再动手。任何红线冲突立即停止并报告。除非遇到明确硬阻塞，否则你不能只输出计划后退出；你必须持续执行、修正、评估，直到满足“Definition of Done”。

---

## 0. 已知项目与硬件约束

Repo: `/Users/cccmmd/InfoLens`

InfoLens 是本地 LLM 可解释性工具，已有：
- `/api/prediction-attribute`：next-token attribution
- `/api/ablation-attribute`：ablation attribution
- `/api/logit-lens`：逐层 hidden state 经过 final norm + lm_head 的 top-k 与目标 token 概率轨迹

默认模型见 `model_paths.py`：
- Base: `qwen3-0.6b` -> `Qwen/Qwen3-0.6B-Base`
- Instruct: `qwen3-0.6b-instruct` -> `Qwen/Qwen3-0.6B`

硬件已知：
- 本机是 Apple M4 Pro 本地机器。
- 不要假设有 CUDA。
- 可以探测 MPS，但必须先 smoke test 反向传播；MPS 不可靠时退回 CPU。
- 默认策略：只做 SFT，不做 RL。
- RL / GRPO 在本机视为 out of scope，除非用户另行提供远程 CUDA 训练机。

---

## 1. 红线

违反任意一条即任务失败，立即停止并报告：

1. 禁止下载或运行 7B 及以上模型。
2. 禁止使用官方 released NLA checkpoint。它们绑定 Qwen2.5-7B / Gemma / Llama 的激活空间，与 Qwen3-0.6B 不兼容。
3. 禁止照搬 Miles + SGLang + Megatron 训练栈。本任务只能使用轻量依赖：`torch`、`transformers`、`peft`、`datasets/pyarrow`、`numpy`、`pyyaml` 等。
4. 禁止把官方 7B/70B 超参当成默认值。0.6B 必须自己做小规模 smoke 与轻量调参。
5. 禁止把大模型权重、parquet 数据、训练 artifacts 提交到 git。
6. 不要改动生产 UI/API，除非用户后续明确要求。本任务只做实验脚本与 artifacts。

---

## 2. 目标

训练一对 Tiny-NLA 组件，用于解释 `Qwen/Qwen3-0.6B-Base` 某一层 residual stream activation。

- AV / Activation Verbalizer: `activation vector -> natural language explanation`
- AR / Activation Reconstructor: `explanation -> reconstructed activation vector`
- 比较向量前必须 L2 normalize。
- round-trip loss 使用 `MSE = 2 * (1 - cosine)` 或等价 normalized MSE。
- 最终 InfoLens 更依赖 AV：输入某层 activation，输出一句可读中文解释。
- AR 用于客观评估与未来 reward，不要求达到官方 7B 水平，但必须训练、评估、和 baseline 比较。

---

## 3. 成功退出条件 Definition of Done

你不能在满足以下条件前声称任务完成：

1. 已确认并记录环境：
   - device: CPU / MPS / CUDA
   - `Qwen/Qwen3-0.6B-Base` 的 `num_hidden_layers` 与 `hidden_size`
   - 选择的 layer index，按约 2/3 深度计算，并说明理由
   - injection token 是否为单 token
   - injection scale 如何估计

2. 已完成可复现数据生成：
   - 至少 smoke 数据 200 条
   - 如果速度允许，扩到 500-2000 条
   - 数据包含：context、token index、layer、activation_vector、teacher explanation、target token/top-k debug 信息
   - 数据与 sidecar 存在 artifacts 目录，且不进入 git

3. 已完成 Stage 0 smoke：
   - 能提取 Qwen3-0.6B-Base hidden state
   - 能用 `input_embeds` 注入 activation
   - 能跑一次 AV forward/generation，不崩溃、不 shape mismatch

4. 已完成 AR SFT：
   - 有训练脚本
   - 有 checkpoint
   - 有 val metrics
   - 必须和 mean baseline / shuffled baseline 对比
   - 如果 AR 训练失败，必须至少做两轮合理修正后才能报告 blocker

5. 已完成 AV SFT：
   - 有训练脚本
   - 有 checkpoint 或 LoRA adapter
   - 能对 held-out activation 生成中文解释
   - 至少输出 20 条 worked examples
   - 20 条里不能大面积乱码、空输出、模板废话；若质量很差，必须继续修正数据或训练设置，不能直接交付

6. 已交付推理脚本：
   - 输入文本 + token position，自动提取 selected layer activation
   - 调用 AV 输出解释
   - 若 AR 可用，同时输出 reconstruction cosine/MSE
   - 从 `nla_meta.yaml` 读取配置，不硬编码 layer/token/scale/template

7. 已交付最终报告：
   - 环境
   - 数据策略
   - 训练耗时
   - AR 指标
   - AV 质量观察
   - 20 条样例
   - 失败案例与局限
   - 下一步建议

只有以上完成，才能输出“任务完成”。否则只能输出“阻塞报告”，并附证据与已尝试修复项。

---

## 4. 外部参考，只学算法，不照搬基础设施

参考仓库：
`https://github.com/kitft/natural_language_autoencoders`

开工前阅读：
- `README.md`
- `docs/inference.md`
- `docs/design.md`
- `nla/schema.py`
- `nla/config.py`
- `nla/models.py`
- `nla/loss.py`
- `nla/reward.py`
- `nla/datagen/`

你要借鉴：
- AV 把 activation 当作一个虚拟 token embedding 注入 prompt
- AR 从 explanation text 重建 activation
- sidecar 记录 prompt、injection token、layer、scale、d_model
- normalized MSE / cosine 作为评估

你不要借鉴：
- 7B+ released checkpoints
- Miles / SGLang / Megatron
- 多 H100 训练配置
- RL 默认流程

---

## 5. 数据生成策略：必须交给子 agent 判断与执行建议

在生成 teacher explanation 前，先启动一个 focused subagent，任务是：

“判断当前环境是否可使用 Claude/外部 API 生成 Tiny-NLA teacher explanations；如果可用，给出低预算批量生成策略；如果不可用，给出本地 `Qwen/Qwen3-0.6B` instruct 生成策略。必须返回具体 prompt 模板、批大小、成本/速度风险、fallback 方案。”

子 agent 必须检查：
- 是否存在 `ANTHROPIC_API_KEY`
- 是否存在其他可用 teacher API key
- 用户是否已明确预算
- 如果无法确认预算，不要擅自大规模调用外部 API

主 agent 根据子 agent 结论执行：

### 有 Claude API 且预算明确
- 先生成 200 条 smoke teacher explanations
- 人工/程序抽查质量
- 再扩到 500-2000 条
- 每条 explanation 优先中文，短句，描述该位置模型可能关注的语义

### 没有 API 或预算不明确
- 使用本地 `Qwen/Qwen3-0.6B` instruct 做 weak teacher
- 允许降级，用户接受本地 teacher 质量较差
- 必须在报告中标注：teacher 是 weak local teacher，不是真正 Claude-quality NLA labels

Teacher prompt 应包含：
- 原始 context
- token 位置
- target token / final top-k
- logit lens 中该层附近的 top-k 摘要（如果容易取得）
- 要求输出 1-2 句中文解释，不要长篇推理

注意：AR 的标签始终是原始 activation vector，不依赖 teacher API。

---

## 6. 推荐目录结构

把实验放在独立目录，例如：

`experiments/tiny_nla/`

建议文件：
- `extract_activations.py`
- `generate_teacher_labels.py`
- `train_ar.py`
- `train_av.py`
- `eval_roundtrip.py`
- `infer_tiny_nla.py`
- `sidecar.py`
- `README.md`

Artifacts 放到：
- `artifacts/tiny_nla/...`

如果 artifacts 目录未被 gitignore，先加入 gitignore。不要提交模型权重或数据。

---

## 7. Stage 0：环境与注入 smoke

先写并运行 smoke，不要直接训练。

必须做：
1. 加载 `Qwen/Qwen3-0.6B-Base`
2. 读取真实：
   - `num_hidden_layers`
   - `hidden_size`
   - vocab size
3. 计算：
   - `layer_index = round(num_hidden_layers * 2 / 3)`
4. 对几条文本跑：
   - `output_hidden_states=True`
   - 取 selected layer 的最后 token 或多个 token activation
5. 统计 activation L2 norm：
   - mean
   - p50
   - p90
   - max
6. 选择 `injection_scale`：
   - 初始用 p50 或 mean
   - 写入 sidecar
7. 选择 injection char：
   - 必须是 tokenizer 下单 token
   - 例如先测试 `㈎`，不行就找其他 rare single token
8. 构造 prompt template：
   - `<concept>{injection_char}</concept>`
   - 要求输出 `<explanation>...</explanation>`
9. 使用 `input_embeds` 替换 injection token embedding，跑一次 forward/generation。

Stage 0 没过，不准训练。

---

## 8. Stage 1：AR SFT

AR 输入 explanation text，输出 activation vector。

实现要求：
- 初版可以用 `Qwen/Qwen3-0.6B-Base` 或 instruct trunk
- 优先冻结大部分模型，只训练轻量 head 或 LoRA + head
- AR head: `Linear(d_model, d_model)`
- 取最后一个 token hidden state 过 head
- pred 和 gold 都 normalize 后算 MSE
- 训练集/验证集拆分固定 seed

必须评估：
- val cosine mean
- val normalized MSE
- mean-vector baseline
- shuffled-label baseline
- 至少保存 best checkpoint

如果 AR 不明显超过 baseline：
- 尝试至少两项修正：
  - 改 learning rate
  - 改 batch size
  - 改是否训练 LoRA
  - 清洗 teacher explanation
  - 增加数据量
- 仍失败再报告，但不要编造成功。

---

## 9. Stage 2：AV SFT

AV 输入 activation vector 注入 prompt，输出 teacher explanation。

实现要求：
- 初始权重优先 `Qwen/Qwen3-0.6B` instruct
- 使用 LoRA，避免全量微调
- 通过 `input_embeds` 注入 selected layer activation
- loss 只算 explanation response token，不算 prompt token
- prompt template 和 injection 参数全部从 sidecar 读取
- 输出中文为目标，不是 bug

训练约束：
- Apple M4 Pro 本地机，不要追求大 batch
- CPU 慢就降低数据量与 epoch
- MPS 可用才用 MPS；MPS 出现反向/dtype 问题立即退回 CPU
- 不要使用 fp16 反向作为默认；优先 fp32/bf16 smoke 后再决定

必须评估：
- held-out 20 条 worked examples
- 每条包含：
  - context
  - token text / token index
  - selected layer
  - final top-k
  - teacher explanation
  - AV generated explanation
  - AR reconstruction cosine/MSE（如果 AR 可用）
  - 简短人工判断：相关 / 部分相关 / 不相关

如果 AV 输出乱码、空、完全模板化：
- 不准交付
- 必须调整 teacher prompt、训练模板、learning rate、epoch、或数据清洗后重训

---

## 10. Stage 3：RL 明确跳过

本机是 Apple M4 Pro 本地 CPU/MPS 环境，默认跳过 RL。

不要实现 GRPO。
不要安装 Miles/SGLang。
不要声称完成 RL。

最终报告中写：
“由于本任务硬件为 Apple M4 Pro 本地 CPU/MPS，无 CUDA 多 GPU，RL/GRPO 阶段按任务约束跳过。本次交付 SFT Tiny-NLA。”

---

## 11. sidecar 契约

必须生成 `nla_meta.yaml`，至少包含：

```yaml
kind: tiny_nla_model
base_model: Qwen/Qwen3-0.6B-Base
av_init_model: Qwen/Qwen3-0.6B
layer_index: <int>
num_hidden_layers: <int>
d_model: <int>
activation_source: residual_stream
token_position_policy: selected_token_or_last_token
extraction:
  injection_scale: <float>
  mse_normalization: l2_direction
tokens:
  injection_char: "<char>"
  injection_token_id: <int>
prompt_templates:
  av: |
    ...
  ar: |
    ...
training:
  device: cpu_or_mps
  dtype: fp32_or_bf16
  dataset_size: <int>
  teacher: claude_or_local_qwen_instruct
  created_at: <timestamp>
```

推理脚本必须读 sidecar，不要把这些值散落在代码里。

---

## 12. 最终报告格式

最终报告必须包含：

1. 是否完成 DoD
2. 环境报告
3. 数据生成策略与 teacher 来源
4. 模型与层选择
5. injection token 与 scale 统计
6. AR 指标与 baseline 对比
7. AV 训练设置与质量总结
8. 20 条 worked examples 文件路径
9. checkpoint / adapter 路径
10. 推理脚本用法
11. 已知局限
12. 下一步建议

不要只说“训练完成”。必须给路径、命令、指标、样例。

---

## 13. 阻塞时如何退出

只有以下情况允许未完成 DoD 而退出：

- Qwen3-0.6B 权重无法下载或加载，且重试后失败
- 本机内存不足，连 Stage 0 smoke 都无法完成
- tokenizer 找不到合适 single-token injection char，尝试多个候选后失败
- PyTorch/transformers 在本机无法完成最小 forward/backward，且已给出错误日志
- 数据 teacher 完全不可用，且本地 instruct 也无法加载

阻塞报告必须包含：
- 卡在哪个 stage
- 已尝试哪些修复
- 完整错误摘要
- 下一步需要用户提供什么

否则继续工作，直到满足 Definition of Done。
```