Files

T

smoke / nanochat-smoke (push) Successful in 33s

Details

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.

2026-05-05 22:40:57 +01:00

6.7 KiB

Raw Blame History

W1 — 音频通路 forward smoke 总结

阶段: W1（参考 doc/todo.md / doc/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`	`nanochat/gpt.py`	`forward()` 加 `audio_features` 关键字参数
`7cc94cf`	`nanochat/audio.py`（新）	`WhisperEncoder` + `Projector` 模块
`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. 已知限制 / 留给后续

按"留意度"降序：

loss 下降并不能证明对齐学会了。LM 也在训，5 段短字符串完全可以靠 LM 死记。projector 是不是真的在传递音频信息，得 W2 freeze LM 才能验证 ——那时候唯一能改 loss 的路径就是 projector → audio 通路。
last_hidden_state 是 Whisper 最偏文本语义的层。"质感感知" 这个项目定位要求 timbre / prosody / emotion 等非文本信号能传到 LM；W3+/W5+ 时应当切到中间层、多层 weighted sum，或者干脆换 wav2vec2 / w2v-bert （facebook/w2v-bert-2.0 在 cpc-i7 cache 里就有）。
30 s 强 padding 让短音频浪费 6× 算力。短期是 W2 数据准备的常数代价；长期需要 streaming-style chunking 或者直接换非-Whisper backbone。
CI smoke job 暂未接入。W1 在 4090 本地跑通，按计划在 W2 同步把 audio smoke 加进 scripts/smoke.sh + .gitea/workflows/smoke.yml，统一在 ailab runner 上跑。

4. 衔接 W2

doc/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 互不污染

6.7 KiB Raw Blame History Unescape Escape