Fam Zheng
8c7210abeb
smoke / nanochat-smoke (push) Successful in 33s
Walkthrough of the three W1 commits, the 4090 result (50 steps in ~1s, loss 5.55 → 0.17), and the limitations to keep in mind before reading into the loss-down (LM is also random + tiny vocab, so the drop is mostly memorisation, not Whisper-Projector alignment — W2 freezes the LM specifically to test that). Includes the W2 hand-off checklist.
154 lines
6.7 KiB
Markdown
154 lines
6.7 KiB
Markdown
# 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 字节 + 单独的 `<BOS>`，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` 互不污染
|