nanoVLM

[!TIP]
我们写了一篇 nanoVLM 教程，它将引导你浏览本仓库，并帮助你快速上手。

[!NOTE]
我们已于 2025 年 9 月 9 日推送了一些重大的破坏性变更。这些更新全部是为了支持图像分割和在多节点上训练，用于 FineVision 版本的消融实验。代码库中关于辅助脚本（例如 notebook 或内存评估）的部分可能已无法使用。旧版本的 nanoVLM 也是如此（参见下面的备注）。如果你发现任何不适用的问题，请在 Issues 中告知我们或提交 PR！

[!NOTE]
我们已于 2025 年 6 月 4 日向仓库推送了一些破坏性变更。为了能进行更智能的打包，我们重构了图像和文本嵌入的组合方式。为了尽可能平滑，我们使用新流程训练了一个新的 nanoVLM-450M 模型，同时保留了与旧流程兼容的旧 nanoVLM-222M 模型。如果你现在克隆此仓库或拉取更新到本地，默认将是新的 450M 模型。如果你想要更简单的理解和更简单的代码库，可以使用 v0.1 版本，它与旧的 222M 模型开箱即用。

nanoVLM 是训练/微调小型视觉语言模型（VLM）的最简单仓库，使用纯 PyTorch 实现，代码轻量。代码本身非常易读且易于上手，模型由视觉骨干（models/vision_transformer.py ~150 行）、语言解码器（models/language_model.py ~250 行）、模态投影（models/modality_projection.py ~50 行）以及 VLM 本身（models/vision_language_model.py ~100 行）组成，还有一个简单的训练循环（train.py ~200 行）。

与 Andrej Karpathy 的 nanoGPT 类似，我们希望为社区提供一个非常简单的 VLM 实现和训练脚本。我们不声称这是新的 SOTA 模型，而更像是一个教育性的工作，如果你有合适的硬件，它也能提供相当不错的性能！你应该很快就能修改和玩弄代码。

nanoVLM 能做什么？

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

因此，它是一个简单而强大的 VLM 入门平台。非常适合尝试不同的设置和配置，探索小型 VLM 的能力和效率！

快速开始

你可以克隆仓库，搭建环境并开始使用脚本，或者直接在 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 用于图像处理器
- 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 中的示例图像并附带一个问题，会得到以下输出。即使只经过短时间训练，模型也能识别出图像中的猫。

输入：
Image + '这是什么？'

输出：
生成 1：这是一只猫坐在地上。我认为这是一只猫坐在地上。
生成 2：这张照片是在室外拍摄的。中心有一只棕色的猫似乎在坐着。
生成 3：这是一只猫坐在地上，颜色是白色和棕色。这只猫
生成 4：这是一只猫坐在地上。我认为这是一只猫坐在地上。
生成 5：这是一只猫坐在地上，上面铺着一张垫子。我认为这是...

使用 lmms-eval 进行评估

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

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

# 确保环境变量设置正确，并且已登录 Hugging Face
export HF_HOME="<HF 缓存路径>"
huggingface-cli login

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

# 如果你希望在训练期间使用它，只需导入模块并像在命令行中那样调用它。
# 你可以传递所有在命令行中也可以传递的参数。
# 训练期间的评估在完整的 DDP 设置下有效。
from evaluation import cli_evaluate
args = argparse.Namespace(
    model='lusxvr/nanoVLM-450M', # 这可以是一个检查点路径或模型本身
    tasks='mmstar,mmmu,ocrbench',
    batch_size=128 # 根据你的 GPU 调整，需要传递以避免内存溢出错误
)
results = cli_evaluate(args)

Hub 集成

nanoVLM 提供了方便的方法来从 Hugging Face Hub 加载和保存模型。

预训练权重

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

# 从 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 需求对于选择正确的硬件和批次大小至关重要。我们在单个 NVIDIA H100 GPU 上对默认的 nanoVLM 模型（222M 参数）进行了基准测试。以下是训练期间不同批次大小（包括模型、梯度和优化器状态）观察到的峰值 VRAM 使用情况总结：

以下是近似峰值 VRAM 使用明细：

将模型加载到设备后分配的 VRAM：871.44 MB
--- VRAM 使用情况总结 ---
批次大小 1：4448.58 MB
批次大小 2：4465.39 MB
批次大小 4：4532.29 MB
批次大小 8：5373.46 MB
批次大小 16：7604.36 MB
批次大小 32：12074.31 MB
批次大小 64：20995.06 MB
批次大小 128：38834.19 MB
批次大小 256：74561.08 MB
批次大小 512：内存溢出（OOM 前峰值：80247.67 MB）

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

关键要点：
- 即使批次大小为 1，训练默认模型也需要至少约 4.5 GB 的 VRAM。
- 使用大约 8 GB 的 VRAM，你应该可以训练到批次大小 16。

为你的设置进行测量：

上述值适用于默认模型配置。如果你修改了模型架构（例如，更改骨干网络、隐藏层大小）或使用了不同的序列长度，你的 VRAM 需求将会改变。

我们提供了一个 measure_vram.py 脚本，允许你在自己的特定机器上，针对你选择的模型配置和批次大小测试 VRAM 需求。

使用方法：
1. 确保你有支持 CUDA 的 GPU 并已安装 PyTorch。
2. 使用你想要的批次大小运行脚本。你还可以指定一个模型检查点（如果有的话），或者让它根据默认的 VLMConfig 初始化一个新模型。

# 示例：使用新的默认模型测试批次大小 1、2、4、8
python measure_vram.py --batch_sizes "1 2 4 8"

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

该脚本将输出每个测试批次大小的峰值 VRAM，帮助你确定适合硬件的可行训练配置。

贡献

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

• 纯 PyTorch： 我们致力于保持 nanoVLM 作为纯 PyTorch 的轻量级实现。引入了类似 transformers.Trainer、accelerate 或 deepspeed 等依赖的贡献将不被接受。
• 新特性： 如果你有新特性的想法，请先开一个 issue 来讨论范围和实现细节。这有助于确保你的贡献符合项目的目标。
• 错误修复： 欢迎随时提交错误修复的 pull request。

路线图

以下是我们未来计划重点关注的领域。特别欢迎在这些领域的贡献：

• 评估： 实现更多评估或改进我们的 MMStar 实现（高度受青睐）
• 数据打包： 实现一种从输入数据创建指定大小的包以优化训练的方法。
• 多 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}}
}