nanochat-omni/docs/research_feasibility.md

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 实测、私有数据集）， 随时迭代本文档。