diff --git a/doc/w1-audio-smoke.md b/doc/w1-audio-smoke.md new file mode 100644 index 0000000..e80bbdb --- /dev/null +++ b/doc/w1-audio-smoke.md @@ -0,0 +1,153 @@ +# W1 — 音频通路 forward smoke 总结 + +> 阶段: W1(参考 [`doc/todo.md`](todo.md) / [`doc/research_feasibility.md`](research_feasibility.md) §1.2) +> 作者: @mochi +> 日期: 2026-05-05 +> 状态: ✅ 通路打通,CI 接入留作 W2 同步 + +--- + +## 0. 目标 + +W1 的唯一目标是 **proof of plumbing**: + +``` +wav → WhisperEncoder (frozen) → Projector → prepend 到 text embedding + → 随机初始化 d6 GPT → CE loss on text only +``` + +跑通这条链,并验证梯度能从 loss 反传到 Projector。**不**追求 alignment 质量、 +不追求 transcribe 能力——这些是 W2/W3 的事。 + +Pass criterion 故意做得宽:50 步训练后 `loss[-1] < loss[0] - 0.5`。 + +--- + +## 1. 实现切片 + +W1 的全部代码改动落在三个 commit 里,逻辑互不耦合: + +| Commit | 范围 | 核心改动 | +|---|---|---| +| [`d760915`](../../../commit/d760915) | `nanochat/gpt.py` | `forward()` 加 `audio_features` 关键字参数 | +| [`7cc94cf`](../../../commit/7cc94cf) | `nanochat/audio.py`(新) | `WhisperEncoder` + `Projector` 模块 | +| [`3c1cc33`](../../../commit/3c1cc33) | `scripts/`、`data/` | 合成数据集 + 50 步 smoke 脚本 | + +### 1.1 GPT.forward 的 audio prepend + +LLaVA-style:把 projector 输出的 soft tokens 拼在 text embedding **前面**, +其余照旧。改动 18 行,要点: + +- **prepend 时机**:smear 之后、transformer trunk 之前。smear 的 prev-token + 语义对 soft tokens 没定义,所以让它仍然是严格的 "text-only" 操作。 +- **rotary 位置**:原来按 text 长度切片 cos/sin;audio 在场时改切 `[0, T_a + T_text)`。 + rotary 缓存在 `__init__` 时已经按 `sequence_len * 10` 过度分配,覆盖得过来。 +- **value embedding**:`self.value_embeds[i](idx)` 需要 token id;audio 位置用 0 + 填充。`ve_gate` 是 input-dependent 的,会自己学到压制这些假行——W1 不操心。 +- **targets 对齐**:传入 targets 时自动在 audio 位置 prepend `-1`(`ignore_index`), + loss 只统计文本位置。 +- **不支持的路径**:`kv_cache is not None` 时直接 `assert`。KV-cache 是 prefill + + decode 的协议,给它写 audio 语义需要重新设计,W1 只跑 train-style forward, + 现在无所谓。 + +### 1.2 nanochat/audio.py + +两个类,零隐式状态。 + +**`WhisperEncoder`** + +- 权重加载顺序:`WHISPER_MS_ID` 设了就先 ModelScope(CN 镜像政策,详见 + `doc/todo.md` 决定事项),失败/没设就 HF(`WHISPER_HF_ID`,默认 + `openai/whisper-base`)。HF 路径自动 honor `HF_ENDPOINT=hf-mirror.com`, + 跟 `scripts/smoke.sh` 现有 env 兼容。 +- `__init__` 即 freeze(`requires_grad = False` + `eval()`),调用方不需要 + 记得"冻结一下"——少一种忘记踩坑的方式。 +- `preprocess(audios)` 走 transformers' `WhisperFeatureExtractor`, + `forward(input_features)` 走 encoder.last_hidden_state。 +- **设计妥协**:Whisper 把每段音频 pad 到 30 s → encoder 输出 1500 帧不变。 + 5 s 的样本浪费 6× 算力,W1 不优化;W2 可以换 streaming chunking。 + +**`Projector`** + +- 两层 MLP:`in_dim → out_dim → out_dim`,GELU 激活,bias 全无。 +- 用 `nanochat.gpt.Linear`,master 权重 fp32、forward 时按输入 dtype(bf16) + cast,对齐 nanochat 主模型风格。 +- **`fc2` 零初始化**:模型在第 0 步对音频"完全无视"——从一个干净的 baseline + 起步,audio 路径是 opt-out by default、训练后才 opt-in。这对 debug 很友好: + forward 走通了但 loss 不动?立刻就能定位到 projector 没在学。 + +### 1.3 W1 smoke + +**合成数据**(`scripts/audio_smoke_data.py`):5 段 5 s 正弦(220/330/440/660/ +880 Hz,加二次谐波 + 1% 高斯噪声防止纯音 log-mel 退化),文字标签依次是 +"low / mid low / middle / mid high / high tone"。文件在 `data/audio_smoke/wavs/` +(gitignored),`manifest.jsonl` 入 git。stdlib `wave` 写 PCM16,零额外依赖。 + +为什么是合成而不是 LibriSpeech?W1 是 forward proof,不需要数据真实——网络 +依赖反而会让 smoke 不稳定。W2 上真数据。 + +**对齐脚本**(`scripts/audio_align_smoke.py`): +- **字节级 tokenizer**:UTF-8 字节 + 单独的 ``,vocab=257。绕开 nanochat + BPE 的 `tok_train` 前置依赖,让 W1 完全 standalone。W2 切回真 BPE。 +- 流程:load 5 个 wav → 一次性预提 Whisper 特征(encoder 冻结,每步重算就 + 是浪费)→ 50 步 AdamW,projector + LM 一起练。 +- pass 判据:`losses[0] - losses[-1] >= 0.5`,宽容到挡不住任何真实失败。 + +--- + +## 2. 实测结果 + +cpc-i7(RTX 4090 24G,bf16,CUDA 12.8): + +``` +input_features: (5, 80, 3000) # batch × n_mels × T_mel +whisper features: (5, 1500, 512) # batch × T_a × d_audio (whisper-base) +GPT: depth=6 n_embd=384 n_head=6 +text idx: (5, 13) # max(transcript) - 1 + +step 000 | loss 5.5533 +step 005 | loss 1.7214 +step 010 | loss 0.9479 +... +step 049 | loss 0.1658 + +Done 50 steps in 0.9s | start=5.5533 end=0.1658 drop=5.3875 +PASS +``` + +- 50 步训练 ~0.9 s(不含 Whisper 首次下载和编码) +- `tests/` 13 passed / 10 skipped — `forward()` 改动没破坏既有路径 +- 显存峰值未测;whisper-base + d6 + B=5 应该 < 2 GB + +--- + +## 3. 已知限制 / 留给后续 + +按"留意度"降序: + +1. **loss 下降并不能证明对齐学会了**。LM 也在训,5 段短字符串完全可以靠 + LM 死记。projector 是不是真的在传递音频信息,得 W2 freeze LM 才能验证 + ——那时候唯一能改 loss 的路径就是 projector → audio 通路。 +2. **`last_hidden_state` 是 Whisper 最偏文本语义的层**。"质感感知" 这个项目 + 定位要求 timbre / prosody / emotion 等非文本信号能传到 LM;W3+/W5+ 时 + 应当切到中间层、多层 weighted sum,或者干脆换 wav2vec2 / w2v-bert + (`facebook/w2v-bert-2.0` 在 cpc-i7 cache 里就有)。 +3. **30 s 强 padding** 让短音频浪费 6× 算力。短期是 W2 数据准备的常数代价; + 长期需要 streaming-style chunking 或者直接换非-Whisper backbone。 +4. **CI smoke job 暂未接入**。W1 在 4090 本地跑通,按计划在 W2 同步把 audio + smoke 加进 `scripts/smoke.sh` + `.gitea/workflows/smoke.yml`,统一在 ailab + runner 上跑。 + +--- + +## 4. 衔接 W2 + +[`doc/todo.md`](todo.md) W2 段已经写好,关键交接点: + +- **数据**:LibriSpeech 100h(HF mirror),预提 Whisper-base 特征落 webdataset +- **冻结策略**:Whisper + LM **都**冻结,只训 Projector —— 这是真正的 + 弱对齐,也是验证 W1 第 1 条限制的实验 +- **可视化**:projector 输出对 LM 嵌入空间做 PCA,看不同样本是否在文本嵌入 + 空间里聚类 +- **wandb**:项目名分到 `nanochat-omni-audio`,跟 nanochat 文本 base 的 + `nanochat` 互不污染