AIFrontierLab/TorchUMM

TorchUMM – 统一的多模态模型工具包

是什么 – TorchUMM 是一个 Python 库，通过单一、配置驱动的接口，让您能够运行、评估和微调一系列现代多模态模型（文本到图像、图像理解、图像编辑等）。它捆绑了 14 种不同骨干模型的适配器，提供了超过十个标准基准的现成脚本，并支持 SFT（监督微调）和 IRG（指令跟随强化学习）等后训练方法。代码可在本地、AMD-ROC-m 集群或通过 Modal 云平台运行。

为何重要 – 多模态研究目前是碎片化的：每个新模型通常都自带自己的推理脚本、数据加载器和评估代码。TorchUMM 消除了这种摩擦：

标准化 API – 一个 InferencePipeline 类可适用于任何支持的模型。
确保可复现性 – 所有实验均由 YAML 配置驱动；只需修改配置文件即可切换模型或基准，无需更改代码。
促进公平比较 – 所有模型均使用相同的数据预处理、指标实现和评估脚本。
轻松扩展 – 内置 Modal 支持可自动生成包含正确 CUDA/Flash-Attention 轮子的容器镜像，让您无需手动设置即可在云 GPU 上启动大规模运行。

核心组件

src/umm/backbones/ – 将模型原生 API 转换为 TorchUMM 统一接口的轻量级包装器。
src/umm/cli/ – 命令行入口点（infer、eval、train）。
src/umm/post_training/ – SFT、IRG、recA、UniCot 等方法的实现。
configs/ – YAML 文件，描述使用哪个模型、哪个基准和哪些超参数。分别用于推理、评估和后训练的独立文件夹。
modal/ – Modal 用的 Docker 风格定义，可让您在托管 GPU 集群上运行相同代码。
eval/ – 协调基准运行（如 DPG-Bench、MME、MMMU、WISE）并聚合分数的脚本。

支持的模型（为每种模型提供适配器；详见链接指南了解模型特定依赖）：

Bagel、DeepGen、OmniGen2、Emu3 / Emu3.5、MMaDA、Janus 系列、Show-o 系列、BLIP3-o、TokenFlow、Ovis-U1。
多数模型需要 Flash-Attention；Emu3.5 也可在 vLLM 上运行以实现更快推理。

覆盖的基准 – 生成（DPG-Bench、GenEval、WISE）、理解（MME、MMMU、MMBench、MM-Vet、MathVista）和编辑（GEdit-Bench、ImgEdit-Bench）。每个模型的结果已在仓库中复现，可通过单个 CLI 命令重新生成。

典型工作流

安装 – 安装包（pip install -e .）并安装所需的模型特定 requirements.txt。
准备数据 – 多数基准数据集会自动下载；少数（MME、MMBench 等）在 README 中提供简单的 wget 脚本。
运行推理 – 例如：python -m umm.cli.main infer --config configs/inference/modal_bagel_generation.yaml。
评估 – 选择一个基准配置，例如：python -m umm.cli.main eval --config configs/eval/dpg_bench/dpg_bench_bagel.yaml。
后训练 – 使用 python -m umm.cli.main train --config configs/posttrain/bagel_sft.yaml 进行微调。
扩展 – 将配置前缀（modal_、amd_ 或无）替换为在 Modal 或 AMD HPC 集群上运行。

Python 使用示例

from umm.inference.pipeline import InferencePipeline
from umm.inference.multimodal_inputs import InferenceRequest

pipeline = InferencePipeline(
    backbone_name="bagel",
    backbone_cfg={"model_path": "/path/to/BAGEL-7B-MoT", "max_mem_per_gpu": "80GiB"},
)

# 文本到图像生成
gen = pipeline.run(InferenceRequest(
    backbone="bagel",
    task="generation",
    prompt="A cat sitting on a rainbow",
    params={"num_timesteps": 50},
))

# 图像理解
understand = pipeline.run(InferenceRequest(
    backbone="bagel",
    task="understanding",
    prompt="Describe this image in detail.",
    images=["/tmp/cat.jpg"],
))

可复现性说明 – 仓库中的表格列出了使用相同评估脚本和相同评分 VLM（例如 WISE 使用 Qwen2.5-VL-72B-Instruct）复现的分数。README 中警告了与原始论文的差异（不同的裁判模型、略有不同的管道）。

谁应该使用它 – 比较多模态模型的研究人员、构建评估流水线的工程师，以及希望无需编写自定义包装器即可运行大量前沿视觉-语言模型的任何人。

引用 – README 提供了相关论文（arXiv:2604.10784）的 BibTeX 条目（此处未重复）。