From b585e07dc263db9d6cac35a045d832c23d4eac83 Mon Sep 17 00:00:00 2001 From: Fam Zheng Date: Tue, 5 May 2026 22:21:31 +0100 Subject: [PATCH] omni: CI smoke + docs + README preamble MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - .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 --- .gitea/workflows/smoke.yml | 25 +++++ .gitignore | 3 + README.md | 16 ++- doc/research_feasibility.md | 207 ++++++++++++++++++++++++++++++++++++ doc/todo.md | 62 +++++++++++ scripts/smoke.sh | 55 ++++++++++ 6 files changed, 367 insertions(+), 1 deletion(-) create mode 100644 .gitea/workflows/smoke.yml create mode 100644 doc/research_feasibility.md create mode 100644 doc/todo.md create mode 100755 scripts/smoke.sh diff --git a/.gitea/workflows/smoke.yml b/.gitea/workflows/smoke.yml new file mode 100644 index 0000000..bdfc69e --- /dev/null +++ b/.gitea/workflows/smoke.yml @@ -0,0 +1,25 @@ +name: smoke +on: + push: + workflow_dispatch: + +jobs: + nanochat-smoke: + runs-on: gpu + timeout-minutes: 30 + steps: + - name: checkout + # actions/checkout@v4 mis-resolves server_url on subpath gitea + # (drops /gitea/ prefix), so clone manually. + run: | + git init . + git remote add origin https://famzheng.me/gitea/${{ github.repository }}.git + git fetch --depth=1 origin ${{ github.sha }} + git checkout FETCH_HEAD + - name: nvidia-smi + run: nvidia-smi --query-gpu=name,memory.free,memory.used --format=csv + - name: smoke + env: + WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} + CI_RUN_TAG: smoke-${{ github.run_number }} + run: bash scripts/smoke.sh diff --git a/.gitignore b/.gitignore index 3e92824..81f3f0c 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,6 @@ eval_bundle/ # Local setup CLAUDE.md wandb/ + +# Claude Code runtime +.claude/ diff --git a/README.md b/README.md index 483f3e3..24f7348 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,18 @@ -# nanochat +# nanochat-omni + +> **质感感知语音输入** — multimodal extension on top of [karpathy/nanochat](https://github.com/karpathy/nanochat). +> +> Audio-first (Whisper encoder + Projector → soft tokens, LLaVA-style alignment). +> Vision later. Output stays text. Single GPU (RTX 5090 / 4090). CN mirrors baked +> into `pyproject.toml` / `nanochat/dataset.py` (sjtu pytorch-wheels, hf-mirror). +> +> Roadmap: [`doc/todo.md`](doc/todo.md) · Research: [`doc/research_feasibility.md`](doc/research_feasibility.md) · CI smoke: [`scripts/smoke.sh`](scripts/smoke.sh) +> +> Sync upstream: `git fetch upstream && git merge upstream/master` + +--- + +# nanochat (upstream) ![nanochat logo](dev/nanochat.png) ![scaling laws](dev/scaling_laws_jan26.png) diff --git a/doc/research_feasibility.md b/doc/research_feasibility.md new file mode 100644 index 0000000..0c0cb82 --- /dev/null +++ b/doc/research_feasibility.md @@ -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 实测、私有数据集）， +随时迭代本文档。* diff --git a/doc/todo.md b/doc/todo.md new file mode 100644 index 0000000..16fcba8 --- /dev/null +++ b/doc/todo.md @@ -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 不依赖 diff --git a/scripts/smoke.sh b/scripts/smoke.sh new file mode 100755 index 0000000..e5cf5bf --- /dev/null +++ b/scripts/smoke.sh @@ -0,0 +1,55 @@ +#!/bin/bash +# End-to-end smoke for nanochat-omni: dataset → tokenizer → tiny base_train. +# Runs in-place at the repo root (post-monorepo-fork). +# Idempotent: caches venv + uv-cache + dataset shards under /data/nanochat-smoke/. +# Targets ailab (CN, RTX 5090). CN mirrors (sjtu pytorch, aliyun PyPI, hf-mirror) +# are committed directly into pyproject.toml/uv.lock/nanochat/dataset.py. +set -euo pipefail +export PATH="$HOME/.local/bin:$PATH" + +CACHE_ROOT=${CACHE_ROOT:-/data/nanochat-smoke} +export NANOCHAT_BASE_DIR=$CACHE_ROOT/cache +export UV_CACHE_DIR=$CACHE_ROOT/uv-cache +export UV_DEFAULT_INDEX=${UV_DEFAULT_INDEX:-https://mirrors.aliyun.com/pypi/simple/} +export UV_INDEX_STRATEGY=unsafe-best-match +export HF_ENDPOINT=${HF_ENDPOINT:-https://hf-mirror.com} +export OMP_NUM_THREADS=1 + +mkdir -p "$CACHE_ROOT" "$NANOCHAT_BASE_DIR" + +# wandb: real run if WANDB_API_KEY is set, otherwise fall back to dummy (DummyWandb). +if [ -n "${WANDB_API_KEY:-}" ]; then + RUN_TAG=${CI_RUN_TAG:-smoke-$(date +%Y%m%d-%H%M%S)} +else + RUN_TAG=dummy +fi + +[ -d .venv ] || uv venv +uv sync --extra gpu --index-strategy unsafe-best-match +source .venv/bin/activate + +echo "=== [1/4] download 1 climbmix shard ===" +time python -m nanochat.dataset -n 1 + +echo "=== [2/4] train tokenizer (50M chars) ===" +time python -m scripts.tok_train --max-chars=50000000 + +echo "=== [3/4] tok_eval ===" +time python -m scripts.tok_eval + +echo "=== [4/4] base_train d=6 50 iters ===" +time python -m scripts.base_train \ + --depth=6 \ + --head-dim=64 \ + --window-pattern=L \ + --max-seq-len=512 \ + --device-batch-size=8 \ + --total-batch-size=4096 \ + --eval-every=25 \ + --eval-tokens=131072 \ + --core-metric-every=-1 \ + --sample-every=25 \ + --num-iterations=50 \ + --run="$RUN_TAG" + +echo "=== smoke done ==="