Concrete near-term tasks (monorepo fork prep + W1 audio smoke) on top of mochi's research feasibility doc.
This commit is contained in:
Fam Zheng
@@ -0,0 +1,207 @@
|
||||
# NanoChat-Omni 多模态训练可行性调研
|
||||
|
||||
> Issue: fam/nanochat-omni#2
|
||||
> 作者: @mochi
|
||||
> 日期: 2026-05-05
|
||||
|
||||
---
|
||||
|
||||
## 0. 背景与定位
|
||||
|
||||
NanoChat（karpathy/nanochat，submodule 锁定在 `dc54a1a`）是单节点可跑通的极简
|
||||
LLM 训练框架。其核心模型 `nanochat/gpt.py` 是一个标准 decoder-only Transformer：
|
||||
|
||||
- RoPE、QK-norm、relu² MLP、GQA、Flash-Attention 3
|
||||
- 复杂度由单一参数 `--depth` 决定，宽度/头数/lr/horizon 等全部自动派生
|
||||
- d12 ≈ GPT-1 量级（约 100M 参数），d20 ≈ 350M，d26 ≈ GPT-2 1.6B
|
||||
|
||||
**Omni 目标**：在保留 NanoChat 训练栈极简风格的前提下，给它接上**多模态输入端**
|
||||
（语音为主，图像兼顾），做出"质感感知语音输入"的对话模型。
|
||||
非目标：多模态输出（TTS / 图像生成不在本期范围）。
|
||||
|
||||
---
|
||||
|
||||
## 1. 总体技术方案
|
||||
|
||||
### 1.1 架构选型
|
||||
|
||||
主流可选路径有三类，复杂度与效果递增：
|
||||
|
||||
| 路径 | 代表工作 | 训练难度 | 显存压力 | 推荐度 |
|
||||
|---|---|---|---|---|
|
||||
| **A. 线性 / MLP Projection** | LLaVA-1.0、MiniGPT-4 | 低 | 低 | ★★★★★（MVP 首选） |
|
||||
| **B. Q-Former / Perceiver** | BLIP-2、Flamingo | 中 | 中 | ★★★（token 数固定，长音频友好） |
|
||||
| **C. Cross-Attention 注入** | Flamingo、IDEFICS | 高 | 高 | ★（架构改动大，与 nanochat 单一 dial 哲学冲突） |
|
||||
|
||||
**结论**：MVP 走 A，量到一定程度再评估 B。C 直接淘汰——一旦动到 attention 层
|
||||
就破坏了 nanochat "改 depth 一切自洽" 的简洁性，回归收益不值得。
|
||||
|
||||
### 1.2 模块图（语音通路）
|
||||
|
||||
```
|
||||
┌─────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐
|
||||
│ 16 kHz wav │ -> │ Speech Enc. │ -> │ Projection │ -> │ NanoChat GPT │
|
||||
│ (variable) │ │ (frozen) │ │ (trainable) │ │ (frozen+LoRA)│
|
||||
└─────────────┘ └──────────────┘ └──────────────┘ └──────────────┘
|
||||
Whisper-small/ 2-layer MLP + d20~d26
|
||||
HuBERT-base down-sampler + LoRA r=16
|
||||
768/1024 dim → n_embd
|
||||
```
|
||||
|
||||
关键设计点：
|
||||
|
||||
1. **下采样**：Whisper 输出 50 Hz，30s 音频 = 1500 token，过长。Projection
|
||||
前先做 stride-2 一维卷积或 mean-pool，把序列长度压到 ~250。
|
||||
2. **特殊 token**：复用 nanochat tokenizer 的 BPE 表，新增 `<|audio_start|>`
|
||||
`<|audio_end|>` 两个 reserved token（tokenizer 已有 256 个保留位）。
|
||||
3. **占位策略**：训练样本里把音频特征拼到 `<|audio_start|>` 之后、文本 prompt
|
||||
之前；attention mask 让文本部分照常自回归，音频部分作为前缀 context。
|
||||
4. **图像通路（可选）**：同构，把 Speech Encoder 换成 SigLIP-So400m 或
|
||||
CLIP-ViT-L/14，Projection 单独训练。两路 Projection 互不干扰，可分阶段开。
|
||||
|
||||
### 1.3 训练策略（三阶段）
|
||||
|
||||
| 阶段 | 可训参数 | 数据 | 目标 | 时长（d20 / 单卡 4090 估算） |
|
||||
|---|---|---|---|---|
|
||||
| S1 Pretrain Projection | 仅 Projection（~10M） | 0.3–1M 弱对齐音频-字幕对 | 让 Projection 把 audio 特征对齐到 LM embedding 空间 | 8–24 h |
|
||||
| S2 Instruction Tune (LoRA) | Projection + LoRA(r=16) on LM | 50–200k 高质量音频指令对 | 让模型听懂指令、按要求回答 | 12–36 h |
|
||||
| S3 (可选) RLHF / DPO | 同 S2 | 1–5k 偏好对 | 风格、安全、拒答 | 6–12 h |
|
||||
|
||||
**重要**：底座 LM 用 NanoChat 已经预训完的 d20 / d26 checkpoint（Karpathy 公开），
|
||||
**不重训 base**。重训 base 在单卡上不现实（参考 nanochat 自己的速度表，d26 在
|
||||
8×H100 上还要近 2 小时——单卡 4090 大致 ×30 = 60 小时，没意义）。
|
||||
|
||||
---
|
||||
|
||||
## 2. 硬件需求调研
|
||||
|
||||
下表中**显存估算**基于 d20 + LoRA(r=16) + Projection + Whisper-small encoder
|
||||
（frozen, fp16）+ batch=4、seq=1024 的训练配置。bf16 权重 + fp32 优化器状态。
|
||||
|
||||
| GPU | VRAM | 内存带宽 | BF16 算力 | FP4/FP8 | 训练适用性 | 预估 step/s（d20） | 备注 |
|
||||
|---|---|---|---|---|---|---|---|
|
||||
| **RTX 3090** | 24 GB | 936 GB/s | ~71 TF | ❌ | S1+S2 勉强（bs≤2，需 grad-accum） | ~0.6 | Ampere，无 FP8；二手便宜，做 S1 验证够用 |
|
||||
| **RTX 4090** | 24 GB | 1008 GB/s | ~330 TF | FP8 ✓ | **S1+S2 推荐**（bs=4 顺畅） | ~2.2 | 性价比之王；MVP 主力卡 |
|
||||
| **RTX 5090** | 32 GB | 1.79 TB/s | ~836 TF | FP8/FP4 ✓ | S1+S2+S3 全流程舒适 | ~4.5 | Blackwell；多 8GB 显存让 d26+LoRA、长音频上下文都能塞下 |
|
||||
| **GB10 (DGX Spark)** | 128 GB 统一内存 | ~250 GB/s | 中等（≈ 4090 量级） | FP4 ✓ | 开发友好；显存大，吞吐受带宽掣肘 | ~1.5 | LPDDR5X 带宽是短板，但 128 GB 让 d26 全量微调成为可能；适合"试菜" |
|
||||
| **B40** | 待官方确认（业内传 48 GB Blackwell workstation） | —— | —— | —— | 若上市，定位介于 4090 与 RTX PRO 之间 | —— | 截至 2026-05 NVIDIA 未发布正式 spec，本文不做精确估算；落地前需重做评估 |
|
||||
|
||||
### 2.1 推荐组合
|
||||
|
||||
- **MVP 阶段（S1+S2 跑通）**：单张 RTX 4090，二手价格友好，bf16 + FP8 权重激活
|
||||
量化（FP8 仅推理 / 评估阶段开）能把 d20+LoRA 塞得很舒服。
|
||||
- **质量冲刺阶段（S2 长上下文 + S3）**：单张 RTX 5090；32 GB 显存让 d26+LoRA、
|
||||
60 s 音频上下文不再是 OOM 风险。
|
||||
- **数据预处理 + 长跑实验**：GB10 适合放在桌边长跑（128 GB 让 d26 全量微调可行），
|
||||
但训练吞吐别期待超过 4090。
|
||||
- **3090**：仅推荐做 Projection-only 的 S1 验证或离线 eval，主训不建议。
|
||||
- **B40**：等 NV 实际发布参数再纳入，本期先按 4090/5090 排期。
|
||||
|
||||
### 2.2 显存压力红线
|
||||
|
||||
经验公式（bf16 + AdamW + grad-checkpointing on）：
|
||||
|
||||
```
|
||||
VRAM ≈ 2·P_total + 4·P_train + activations(seq_len, batch)
|
||||
其中 P_total 是模型总参数，P_train 是可训参数（LoRA + Projection ≈ 0.3% × P_total）
|
||||
```
|
||||
|
||||
- d20（~350M）+ LoRA(r=16) ≈ 1.4 GB 权重 + 几百 MB 优化器 + activations
|
||||
→ bs=4、seq=1024 在 24 GB 上**留出 8 GB 给音频特征 / 评估**，是安全的。
|
||||
- d26（~1.6B）+ LoRA → bs=2、seq=2048 时 24 GB 紧张，5090 的 32 GB 是甜点。
|
||||
|
||||
---
|
||||
|
||||
## 3. 数据集规划
|
||||
|
||||
### 3.1 阶段对应数据
|
||||
|
||||
| 阶段 | 类型 | 候选数据集 | 规模 | 获取 | 许可 |
|
||||
|---|---|---|---|---|---|
|
||||
| S1 弱对齐 | 音频-文本（ASR/字幕） | LibriSpeech 960h、GigaSpeech S/M、AISHELL-3、Common Voice 17 zh | 1k–10k 小时 | HuggingFace Datasets、官方下载 | CC-BY / CC0 / 商用友好混合 |
|
||||
| S1 弱对齐（图像） | 图文对 | LAION-COCO-NLLB（subset）、CC3M、ShareGPT4V | 0.5–3M 对 | HF | 注意 LAION 部分授权 |
|
||||
| S2 指令 | 音频指令 | AudioBench、Dynamic-SUPERB、自合成（用现有 LLM 生成 prompt + TTS） | 50–200k | HF + 自建 | 自合成部分许可清晰 |
|
||||
| S2 指令（图像） | 视觉指令 | LLaVA-Instruct-150k、ShareGPT4V-7M（采样） | 50–150k | HF | 商用须审 |
|
||||
| S3 偏好 | 人工偏好对 | 自标 / 用 GPT-4o 当 judge 蒸馏 | 1k–5k | 自建 | 自有 |
|
||||
|
||||
### 3.2 中文 / "质感感知" 部分
|
||||
|
||||
repo 描述强调"质感感知语音输入"，纯靠开源数据吃不到这个词所暗示的**音色 /
|
||||
情感 / 韵律**信号。建议：
|
||||
|
||||
1. **基础**：CommonVoice-zh + AISHELL-3 + WenetSpeech（subset）覆盖通用 ASR 对齐。
|
||||
2. **质感层**：抓 ESD（Emotional Speech Dataset）+ RAVDESS（英文情感）做 S2 指令
|
||||
"请描述这段语音的语气 / 情绪 / 说话人风格"。
|
||||
3. **自合成**：用 CosyVoice / GPT-SoVITS 生成带不同情感标签的同文本变体，
|
||||
构造对比指令，~10k 样本就能让模型对"音色/语气"产生区分度。
|
||||
|
||||
### 3.3 存储与预处理
|
||||
|
||||
- 预提取 Whisper encoder 特征落盘（fp16 npy 或 webdataset tar），训练时直接读，
|
||||
绕过音频解码开销。1000h × 50Hz × 768dim × 2byte ≈ **270 GB**。
|
||||
- 准备 2 TB NVMe 是宽裕的；GB10 可直接当数据节点+训练节点合一。
|
||||
|
||||
---
|
||||
|
||||
## 4. 项目执行计划
|
||||
|
||||
### 4.1 整体时间盘
|
||||
|
||||
以单卡 RTX 4090（必要时跳到 5090）+ 一名全职工程师为基准：
|
||||
|
||||
| 周 | 里程碑 | 交付 |
|
||||
|---|---|---|
|
||||
| W1 | 接 Whisper encoder + 写 Projection；冻结 LM 端到端跑通 forward | `omni/encoder.py`, `omni/projector.py`, smoke test |
|
||||
| W2 | S1 数据 pipeline（LibriSpeech 100h subset）+ 跑 S1 训练 | Loss 曲线、特征对齐 PCA 可视化 |
|
||||
| W3 | S2 指令数据准备（5w 条 mix）+ LoRA 接入 + 跑 S2 | 第一个 demo：能听音频回答"在说什么" |
|
||||
| W4 | **MVP 验收**：开放 web UI（复用 nanochat `chat_web.py`），加录音上传 | demo video + bench 报告 |
|
||||
| W5–W6 | 扩规模（数据 → 1000h，模型 → d26）、对比实验 | 内部 eval：ASR WER、AudioBench 子集 |
|
||||
| W7 | 引入"质感"指令数据 + S3 偏好 | 风格 / 情感问答 demo |
|
||||
| W8 | 文档、blog、checkpoint 发布；图像通路启动设计 | 公开 release |
|
||||
|
||||
**总计：约 8 周做到可发布版本，4 周做到 MVP。**
|
||||
|
||||
### 4.2 MVP 定义（W4 末）
|
||||
|
||||
满足全部以下条件即视为 MVP 达成：
|
||||
|
||||
1. 输入 ≤30s 16 kHz 单声道 wav，模型输出合理中文/英文回答；
|
||||
2. AudioBench-mini（自建 200 题）正确率 ≥ 40%（baseline 是纯 ASR + LM，~25%）；
|
||||
3. 单卡 4090 可在本地启动 web demo，端到端首 token 延迟 < 2 s；
|
||||
4. 训练 / 推理脚本 ≤ 500 行新增代码，遵循 nanochat 的极简风格。
|
||||
|
||||
### 4.3 风险与应对
|
||||
|
||||
| 风险 | 概率 | 影响 | 缓解 |
|
||||
|---|---|---|---|
|
||||
| Projection 学不动（loss 平台） | 中 | 高 | 先用极小 d12 + 小数据 sanity check；改 MLP 深度；先 align 后 instruct |
|
||||
| 单卡显存不够（升 d26 时） | 中 | 中 | 上 5090 / 加 grad checkpointing / 降 seq |
|
||||
| 中文情感数据稀缺 | 高 | 中 | 自合成补齐 + 借用英文情感模式 |
|
||||
| nanochat 上游变动破坏 submodule pin | 低 | 低 | 锁版本，必要时 fork |
|
||||
| B40 / GB10 实测带宽不及预期 | 中 | 中 | MVP 不依赖这两张，仅作扩展评估 |
|
||||
|
||||
---
|
||||
|
||||
## 5. 结论
|
||||
|
||||
**结论：可行，且建议立即启动。**
|
||||
|
||||
理由：
|
||||
|
||||
1. **架构上**最小侵入路径（Encoder + Projection + LoRA）成熟、可复现、与 nanochat
|
||||
"改 depth 一切自洽"的哲学不冲突；新增代码量预计 < 500 行。
|
||||
2. **硬件上**单张 4090 即可跑通 MVP，5090 是质量冲刺甜点；不依赖多卡。
|
||||
3. **数据上**开源量足够 MVP，"质感感知"通过情感数据集 + 自合成可补齐。
|
||||
4. **时间上**4 周 MVP、8 周 release 是基于 nanochat 既有训练栈和成熟 LoRA / 投影
|
||||
范式的现实估计，不是乐观估计。
|
||||
|
||||
**建议的下一步**（不等本 issue 关闭）：
|
||||
|
||||
- 开 issue #3 跟踪 S1 pipeline 实现；
|
||||
- 在 `nanochat-omni/omni/` 下起子目录，与上游 nanochat 严格隔离；
|
||||
- W1 结束前给出 forward smoke test 的复现脚本。
|
||||
|
||||
---
|
||||
|
||||
*完。如果对 hardware 段或数据段有补充信息（比如 B40 实测、私有数据集），
|
||||
随时迭代本文档。*
|
||||
+62
@@ -0,0 +1,62 @@
|
||||
# nanochat-omni TODO
|
||||
|
||||
定位：**质感感知语音输入**（audio-first，输出仅 text，vision 排后期）。
|
||||
参考：[research_feasibility.md](research_feasibility.md)（mochi, 2026-05-05）的 W1–W8 时间盘。
|
||||
|
||||
---
|
||||
|
||||
## 近期 — 仓库结构 / 工程基建
|
||||
|
||||
- [ ] **submodule 展平为 monorepo fork**
|
||||
- `git remote add upstream https://github.com/karpathy/nanochat.git`
|
||||
- 重写 `main`：`git reset --hard upstream/main` + cherry-pick 我们的 7 个 commits（CI / smoke / wandb / gitignore / README）
|
||||
- force push（repo 没人 fork，安全），删 `.gitmodules` + `upstream/nanochat/`
|
||||
- 拉 upstream 更新就 `git fetch upstream && git merge upstream/main`
|
||||
- [ ] CN mirror patch 直接落到 `pyproject.toml` / `nanochat/dataset.py`（fork 后不用 sed）
|
||||
- [ ] CI smoke 跟着 fork 重路径化（`upstream/nanochat/` → 根目录）
|
||||
|
||||
## W1 — Whisper encoder + Projector forward smoke
|
||||
|
||||
参考 research §1.2 模块图。
|
||||
|
||||
- [ ] `nanochat/audio.py`：WhisperEncoder wrapper（冻结，从 HF mirror 拉权重）+ Projector（MLP，输出维度对齐 nanochat `model_dim`）
|
||||
- [ ] `nanochat/gpt.py` `GPT.forward()` 加可选 `audio_features` 参数，作为 soft tokens prepend 到 text embedding 前面
|
||||
- [ ] mini dataset：1–10 段 5s wav + 字幕，落 `data/audio_smoke/`（git 内不存音频，仅清单 + 下载脚本）
|
||||
- [ ] `scripts/audio_align_smoke.py`：50 步、d6 nanochat base、loss 下降即过
|
||||
- [ ] CI 加 audio smoke job（ailab runner 装 ffmpeg；whisper 走 transformers 即可）
|
||||
|
||||
## W2 — S1 弱对齐训练
|
||||
|
||||
- [ ] 拉 LibriSpeech 100h（HF mirror），预提 Whisper-base encoder 特征落盘 webdataset
|
||||
- [ ] `scripts/audio_align_train.py`：冻结 LM + Whisper，只训 Projector
|
||||
- [ ] PCA 可视化对齐效果（特征→文本嵌入空间是否聚类）
|
||||
- [ ] wandb 项目：`nanochat-omni-audio`（跟 nanochat 文本 base 的 `nanochat` 分开）
|
||||
|
||||
## W3 — S2 指令 + LoRA
|
||||
|
||||
- [ ] LoRA 接入 nanochat `Linear`（rank=16，仅 attention/MLP）
|
||||
- [ ] 5w 条音频指令数据 mix（AudioBench + 自合成）
|
||||
- [ ] eval：自建 200 题 AudioBench-mini
|
||||
|
||||
## W4 — MVP demo
|
||||
|
||||
- [ ] 复用 `scripts/chat_web.py`，加录音上传
|
||||
- [ ] AudioBench-mini 准确率 ≥40%（baseline 25%）
|
||||
- [ ] 4090 端到端首 token <2s
|
||||
|
||||
## W5+ — 扩规模 / 质感数据 / vision
|
||||
|
||||
参考 research §4.1，留到 W5–W8 展开。
|
||||
|
||||
## 决定事项
|
||||
|
||||
- **backbone**：nanochat 自训 d12 → d20 → d26（不借现成 gemma/qwen，保持 hackable 灵魂）
|
||||
- **顺序**：audio 先，vision 排 W7+，多模态输出（TTS/imagegen）不做
|
||||
- **infra**：训练 + smoke CI 都跑在 ailab（5090, 32G），CN mirror 走 sjtu/aliyun/hf-mirror
|
||||
- **monorepo fork pattern**：上游 nanochat 的代码就是我们的代码，omni 改动直接进 `nanochat/` 包
|
||||
|
||||
## 暂搁 / 待定
|
||||
|
||||
- [ ] vision 通路：W7+ 启动，参考 LLaVA recipe，跟 audio 复用 Projector 抽象
|
||||
- [ ] 质感数据自合成：用 ailab CosyVoice 或 IndexTTS 生情感变体（s1/i7 上有现成 server，跨机数据生产链待定）
|
||||
- [ ] B40 / GB10 实测：MVP 不依赖
|
||||
Reference in New Issue
Block a user