Fam Zheng
b585e07dc2
smoke / nanochat-smoke (push) Successful in 2m30s
- .gitea/workflows/smoke.yml: gitea CI on ailab gpu runner (manual git clone since actions/checkout@v4 mis-resolves subpath gitea); injects WANDB_API_KEY + CI_RUN_TAG=smoke-$run_number - scripts/smoke.sh: in-place smoke (uv sync + 1 shard + tokenizer + d=6 50-step base_train); idempotent cache at /data/nanochat-smoke/ - doc/research_feasibility.md: voice-first multimodal feasibility study (mochi) - doc/todo.md: phase-by-phase roadmap (W1 Whisper smoke → W4 MVP) - README.md: omni preamble pointing at upstream nanochat README - .gitignore: exclude .claude/ runtime files
208 lines
11 KiB
Markdown
208 lines
11 KiB
Markdown
# 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 实测、私有数据集），
|
||
随时迭代本文档。*
|