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 項目（此處未重複）。