nanoVLM

---

[!TIP] 我们撰写了一篇 nanoVLM 教程，将引导你了解整个仓库并帮助你快速上手。

---

[!NOTE] 我们在 2025 年 9 月 9 日推送了一些重大变更（breaking changes）。这些更新包括使用图像分割（image splitting）以及在多节点上训练。这些功能用于 FineVision 发布的消融实验（ablations）。代码库中一些支持脚本（例如 notebook 或内存评估）可能已经无法正常工作。同样，旧版本的 nanoVLM 训练模型也可能受到影响（类似于下方的说明）。如果你发现某些功能无法正常工作，请在 Issues 中告知我们或提交 PR！

---

[!NOTE] 我们在 2025 年 6 月 4 日对仓库推送了一些重大变更（breaking changes）。为了实现更智能的打包（smarter packing），我们重构了图像和文本嵌入（embeddings）的组合方式。为了保持一切尽可能顺畅，我们使用新的流程训练了一个新的 nanoVLM-450M 模型，同时保留旧的 nanoVLM-222M 与旧流程兼容。如果你现在克隆此仓库或将更新拉取到本地，默认将是新的 450M 模型。如果你希望获得更简单的理解和更简洁的代码库，可以使用 v0.1 版本。该版本与旧的 222M 模型开箱即用。

---

nanoVLM 是用于训练/微调小型视觉语言模型（Vision-Language Model, VLM）的最简仓库，采用纯 PyTorch 轻量级实现。代码本身非常易读且易于理解，模型由视觉骨干网络（Vision Backbone，models/vision_transformer.py 约 150 行）、语言解码器（Language Decoder，models/language_model.py 约 250 行）、模态投影（Modality Projection，models/modality_projection.py 约 50 行）、VLM 本身（models/vision_language_model.py 约 100 行）以及简单的训练循环（train.py 约 200 行）组成。

类似于 Andrej Karpathy 的 nanoGPT，我们希望为社区提供一个非常简单的视觉语言模型实现和训练脚本。我们并不声称这是一个新的 SOTA（State-of-the-Art，最先进）模型，而是一项教育性工作——如果你拥有合适的硬件，它能发挥相当大的威力！你应该能够在短时间内调整和尝试代码。

nanoVLM 能做什么？

本仓库的模型定义和训练逻辑约 750 行，外加一些样板化的日志记录和参数加载代码。 使用 SigLIP-B/16-224-85M 和 HuggingFaceTB/SmolLM2-135M 作为骨干网络，可得到一个 222M 的 nanoVLM。在单张 H100 GPU 上训练约 6 小时，使用约 170 万条 the cauldron 样本，可在 MMStar 上达到 35.3% 的准确率。

因此，这是一个简单但强大的视觉语言模型入门平台。非常适合尝试不同的设置和配置，探索小型视觉语言模型的能力和效率！

快速开始

你可以克隆仓库、设置环境并使用脚本开始，或者直接 在 Colab 中打开。你也可以使用 交互式 notebook 来上手！

环境设置

我们非常喜欢 uv，推荐将其作为包管理器。但你可以随意使用自己喜欢的工具。

首先克隆仓库：

git clone https://github.com/huggingface/nanoVLM.git
cd nanoVLM

如果你想使用 uv：

uv init --bare --python 3.12
uv sync --python 3.12
source .venv/bin/activate
uv add torch numpy torchvision pillow datasets huggingface-hub transformers wandb
# 可选：如需 lmms-eval 集成，需要从源码安装，参见"使用 lmms-eval 评估"部分

如果你更喜欢其他环境管理器，只需安装这些包：

pip install torch numpy torchvision pillow datasets huggingface-hub transformers wandb
# 可选：如需 lmms-eval 集成，需要从源码安装，参见"使用 lmms-eval 评估"部分

依赖项：

• torch <3
• numpy <3
• torchvision 用于图像处理器（image processors）
• pillow 用于图像加载
• datasets 用于训练数据集
• huggingface-hub 和 transformers 用于加载预训练骨干网络
• wandb 用于日志记录

训练

要训练 nanoVLM，你可以直接使用提供的训练脚本。训练完成后，你的模型将被上传到 Hub！

wandb login --relogin
huggingface-cli login
python train.py

这将使用默认的 models/config.py。

生成

要尝试 训练好的模型，你可以直接使用提供的生成脚本

python generate.py

或者，要使用你自己训练的模型，你可以直接运行：

python generate.py --checkpoint /your/path/to/trained_models

如果我们将 assets/image.png 中的示例图像与问题一起输入模型，我们会得到以下输出。即使经过短期训练，模型也能识别出图片中的猫。

输入：
图像 + 'What is this?'

输出：
生成 1:  This is a cat sitting on the ground. I think this is a cat sitting on the ground.
生成 2:  This picture is clicked outside. In the center there is a brown color cat seems to be sitting on
生成 3:  This is a cat sitting on the ground, which is of white and brown in color. This cat
生成 4:  This is a cat sitting on the ground. I think this is a cat sitting on the ground.
生成 5:  This is a cat sitting on the ground, which is covered with a mat. I think this is

使用 lmms-eval 评估

nanoVLM 现在支持使用全面的 lmms-eval 工具包进行评估：

# 安装 lmms-eval（必须从源码安装）
uv pip install git+https://github.com/EvolvingLMMs-Lab/lmms-eval.git

# 确保环境变量设置正确，并且已登录 HF
export HF_HOME="<Path to HF cache>"
huggingface-cli login

# 在多个基准测试上评估训练好的模型
python evaluation.py --model lusxvr/nanoVLM-450M --tasks mmstar,mme

# 如果你想在训练期间使用它，只需导入模块并像命令行一样调用它。
# 你可以传递所有能在命令行中传递的参数。

训练过程中的评估在完整的 DDP（Distributed Data Parallel，分布式数据并行）设置下运行。

from evaluation import cli_evaluate args = argparse.Namespace( model='lusxvr/nanoVLM-450M', # 这可以是检查点路径或模型本身 tasks='mmstar,mmmu,ocrbench', batch_size=128 # 根据你的 GPU 进行调整，需要传递此参数以避免 OOM（Out of Memory，内存不足）错误 ) results = cli_evaluate(args)

## Hub 集成

**nanoVLM** 提供了便捷的方法，用于从 Hugging Face Hub 加载和保存模型。

### 预训练权重

以下是从 Hugging Face Hub 上的仓库加载的方法。这是开始使用预训练权重的推荐方式。

```python
# 从 Hub 加载预训练权重
from models.vision_language_model import VisionLanguageModel

model = VisionLanguageModel.from_pretrained("lusxvr/nanoVLM-450M")

推送到 Hub

当你训练好一个 nanoVLM 模型后，你可能希望将其分享到 Hugging Face Hub。你可以通过以下方式轻松实现：

... # 加载并训练你的模型

# 推送到 `username/my-awesome-nanovlm-model` 仓库
model.push_to_hub("my-awesome-nanovlm-model")

模型将以配置文件 config.json 和权重文件 model.safetensors 的形式保存在 Hub 上。同时还会为你生成一个模型卡片 README.md，其中包含一些高级信息。你可以手动更新它以说明你的工作。

如果仓库不存在，系统会为你自动创建。默认情况下仓库是公开的。如果你不想公开分享，可以传递 private=True 参数。

本地保存/加载

如果你不想将模型托管在 Hugging Face Hub 上，仍然可以将其保存在本地：

... # 加载并训练你的模型

# 保存到本地文件夹
model.save_pretrained("path/to/local/model")

然后你可以从本地路径重新加载：

# 从本地路径加载预训练权重
from models.vision_language_model import VisionLanguageModel

model = VisionLanguageModel.from_pretrained("path/to/local/model")

VRAM 使用量

了解训练时的 VRAM（Video Random Access Memory，显存）需求对于选择合适的硬件和 batch size（批量大小）至关重要。我们在单张 NVIDIA H100 GPU 上对默认的 nanoVLM 模型（222M 参数）进行了基准测试。以下是训练过程中不同 batch size 下观察到的峰值 VRAM 使用量摘要（包括模型、梯度和优化器状态）：

以下是近似峰值 VRAM 使用量的详细分解：

模型加载到设备后的 VRAM 分配: 871.44 MB
--- VRAM 使用量摘要 ---
Batch Size 1: 4448.58 MB
Batch Size 2: 4465.39 MB
Batch Size 4: 4532.29 MB
Batch Size 8: 5373.46 MB
Batch Size 16: 7604.36 MB
Batch Size 32: 12074.31 MB
Batch Size 64: 20995.06 MB
Batch Size 128: 38834.19 MB
Batch Size 256: 74561.08 MB
Batch Size 512: OOM (OOM 前峰值: 80247.67 MB)

请注意，VRAM 测量是在一个小型设置上进行的，使用的是 'SmolLM2-135M'，最大输入序列长度为 128 个 token。这可能与项目中当前的默认配置有所不同。

关键要点：

• 即使 batch size 为 1，你也需要至少约 4.5 GB 的 VRAM 来训练默认模型。
• 拥有约 8 GB VRAM 时，你应该能够使用最大为 16 的 batch size 进行训练。

为你的设置测量：

上述数值是针对默认模型配置的。如果你修改了模型架构（例如，更改 backbone（主干网络）、hidden size（隐藏层大小））或使用不同的序列长度，你的 VRAM 需求将会发生变化。

我们提供了一个脚本 measure_vram.py，允许你在特定机器上测试 VRAM 需求，以及针对你选择的模型配置和 batch size。

使用方法：

1. 确保你拥有支持 CUDA 的 GPU 并安装了 PyTorch。
2. 使用你期望的 batch size 运行脚本。你也可以指定一个模型检查点（如果有的话），或者让它基于默认的 VLMConfig 初始化一个新模型。

# 示例：使用新的默认模型测试 batch size 1, 2, 4, 8
python measure_vram.py --batch_sizes "1 2 4 8"

# 示例：使用特定检查点和不同的 batch size 进行测试
python measure_vram.py --vlm_checkpoint_path path/to/your/model.pth --batch_sizes "16 32 64"

该脚本将输出每个测试 batch size 的峰值 VRAM 分配量，帮助你确定适合你硬件的可行训练配置。

贡献

我们欢迎对 nanoVLM 的贡献！然而，为了保持仓库对简洁性和纯 PyTorch 的关注，我们有以下指导原则：

• 纯 PyTorch： 我们的目标是将 nanoVLM 保持为轻量级的纯 PyTorch 实现。引入 transformers.Trainer、accelerate 或 deepspeed 等依赖的贡献将不被接受。
• 新功能： 如果你有新功能的想法，请先开启一个 issue 讨论范围和实现细节。这有助于确保你的贡献与项目目标保持一致。
• Bug 修复： 欢迎提交 bug 修复的 pull request。

路线图

以下是我们近期计划开展的一些领域。欢迎在这些领域做出贡献：

• 评估： 实现更多评估方法或改进我们的 MMStar 实现（高度重视）
• 数据打包： 实现一种从输入数据创建给定大小 pack（数据包）的方法，以优化训练。
• 多 GPU 训练： 在多个 GPU 上进行训练
• 多图像支持： 使用多张图像进行训练
• 图像分割： 通过图像分割实现更高分辨率，如同 SmolVLM 中所做的那样。
• VLMEvalKit： 集成到 VLMEvalKit 以启用更多基准测试

引用

如果你喜欢这个项目并希望在某些地方使用它，请使用以下引用：

@misc{wiedmann2025nanovlm,
  author = {Luis Wiedmann and Aritra Roy Gosthipaty and Andrés Marafioti},
  title = {nanoVLM},
  year = {2025},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/huggingface/nanoVLM}}
}

nanoVLM 快速上手指南

环境准备

系统要求

• Python: 3.12 推荐
• GPU: NVIDIA GPU，支持 CUDA（训练至少需要 ~4.5GB VRAM）
• 操作系统: Linux/macOS/Windows（WSL）

前置依赖

• CUDA 环境（用于 GPU 训练）
• Git

安装步骤

1. 克隆仓库

git clone https://github.com/huggingface/nanoVLM.git
cd nanoVLM

2. 创建环境（推荐 uv）

# 安装 uv（如未安装）
curl -LsSf https://astral.sh/uv/install.sh | sh

# 初始化环境
uv init --bare --python 3.12
uv sync --python 3.12
source .venv/bin/activate

# 安装依赖
uv add torch numpy torchvision pillow datasets huggingface-hub transformers wandb

3. 或使用 pip

pip install torch numpy torchvision pillow datasets huggingface-hub transformers wandb

4. 登录必要服务

# 登录 Weights & Biases（训练日志）
wandb login --relogin

# 登录 Hugging Face（模型下载与上传）
huggingface-cli login

国内加速: 使用 Hugging Face 镜像

export HF_ENDPOINT=https://hf-mirror.com

基本使用

快速体验（Colab）

直接打开 nanoVLM.ipynb 无需本地配置。

加载预训练模型

from models.vision_language_model import VisionLanguageModel

# 从 Hugging Face Hub 加载
model = VisionLanguageModel.from_pretrained("lusxvr/nanoVLM-450M")

运行推理

# 使用默认预训练模型
python generate.py

# 使用自己的训练模型
python generate.py --checkpoint /path/to/your/model

训练模型

python train.py

• 默认配置位于 models/config.py
• 训练完成后自动上传至 Hugging Face Hub

模型保存与加载

# 保存到本地
model.save_pretrained("path/to/local/model")

# 从本地加载
model = VisionLanguageModel.from_pretrained("path/to/local/model")

# 推送到 Hub
model.push_to_hub("my-awesome-nanovlm-model")

评估模型

# 安装 lmms-eval（需从源码安装）
uv pip install git+https://github.com/EvolvingLMMs-Lab/lmms-eval.git

# 运行评估
python evaluation.py --model lusxvr/nanoVLM-450M --tasks mmstar,mme

检测显存需求

# 测试不同 batch size 的显存占用
python measure_vram.py --batch_sizes "1 2 4 8 16"

---

参考资源

• 详细教程：Hugging Face Blog - nanoVLM
• 预训练模型：lusxvr/nanoVLM-450M