AIFrontierLab/TorchUMM

A unified multimodal model toolkit

TorchUMM – Unified Multimodal Model Toolkit

What it is – TorchUMM is a Python library that lets you run, evaluate, and fine‑tune a wide range of modern multimodal models (text‑to‑image, image‑understanding, image‑editing, etc.) through a single, config‑driven interface. It bundles adapters for 14 different back‑bones, provides ready‑made scripts for more than ten standard benchmarks, and supports post‑training methods such as supervised fine‑tuning (SFT) and instruction‑following reinforcement (IRG). The code can be run locally, on AMD‑ROC‑m clusters, or in the cloud via Modal.

Why it matters – Multimodal research is fragmented: each new model often ships with its own inference script, data loaders, and evaluation code. TorchUMM removes that friction by:

  • Standardising the API – a single InferencePipeline class works for any supported model.
  • Ensuring reproducibility – all experiments are driven by YAML configs; swapping a model or benchmark only requires editing the config, not the code.
  • Facilitating fair comparison – the same data preprocessing, metric implementations, and evaluation scripts are used for every model.
  • Scaling effortlessly – built‑in Modal support creates container images with the right CUDA/Flash‑Attention wheels, so you can launch large‑scale runs on cloud GPUs without manual setup.

Key components

  • src/umm/backbones/ – thin wrappers that translate a model’s native API into TorchUMM’s unified interface.
  • src/umm/cli/ – command‑line entry points (infer, eval, train).
  • src/umm/post_training/ – implementations of SFT, IRG, recA, UniCot, etc.
  • configs/ – YAML files describing which model, which benchmark, and which hyper‑parameters to use. Separate folders for inference, evaluation, and post‑training.
  • modal/ – Docker‑style definitions for Modal, letting you run the same code on a managed GPU cluster.
  • eval/ – scripts that orchestrate benchmark runs (e.g., DPG‑Bench, MME, MMMU, WISE) and aggregate scores.

Supported models (adapters are provided for each; see the linked guides for model‑specific dependencies):

  • Bagel, DeepGen, OmniGen2, Emu3 / Emu3.5, MMaDA, Janus family, Show‑o family, BLIP3‑o, TokenFlow, Ovis‑U1.
  • Most models require Flash‑Attention; Emu3.5 can also run on vLLM for faster inference.

Benchmarks covered – generation (DPG‑Bench, GenEval, WISE), understanding (MME, MMMU, MMBench, MM‑Vet, MathVista), and editing (GEdit‑Bench, ImgEdit‑Bench). Results for each model are reproduced in the repo and can be re‑generated with a single CLI command.

Typical workflow

  1. Install the package (pip install -e .) and the model‑specific requirements.txt you need.
  2. Prepare data – most benchmark datasets are auto‑downloaded; a few (MME, MMBench, etc.) have simple wget scripts in the README.
  3. Run inference – e.g., python -m umm.cli.main infer --config configs/inference/modal_bagel_generation.yaml.
  4. Evaluate – pick a benchmark config, e.g. python -m umm.cli.main eval --config configs/eval/dpg_bench/dpg_bench_bagel.yaml.
  5. Post‑train – fine‑tune with python -m umm.cli.main train --config configs/posttrain/bagel_sft.yaml.
  6. Scale – replace the config prefix (modal_, amd_, or none) to run on Modal or an AMD HPC cluster.

Example Python usage

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"},
)

# Text‑to‑image generation
gen = pipeline.run(InferenceRequest(
    backbone="bagel",
    task="generation",
    prompt="A cat sitting on a rainbow",
    params={"num_timesteps": 50},
))

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

Reproducibility notes – The repo’s tables list scores that were reproduced with the same evaluation scripts and the same scoring VLMs (e.g., Qwen2.5‑VL‑72B‑Instruct for WISE). The README warns where differences from original papers arise (different judge models, slightly different pipelines).

Who should use it – Researchers comparing multimodal models, engineers building evaluation pipelines, and anyone needing a plug‑and‑play way to run many state‑of‑the‑art vision‑language models without writing custom wrappers.

Citation – The README provides a BibTeX entry (not reproduced here) for the associated paper (arXiv:2604.10784).

Related

  • Project
    meta-pytorch/torchtune

    torchtune is a PyTorch‑based library that provides ready‑to‑run recipes and configs for fine‑tuning, RL‑HF, distillation, quantization‑aware training, and inference of many modern LLMs (Llama, Gemma, Qwen, Phi, etc.). It abstracts distributed training, memory‑saving tricks, and device support behind a simple CLI (`tune`) and YAML files, integrating with Hugging Face, torch‑ao, FSDP2, and logging tools. Though development stopped in 2025, it remains a practical, open‑source toolkit for LLM experimentation.

  • Project
    AISBench/benchmark

    A model evaluation tool based on OpenCompass that provides accuracy and performance benchmarking for both local and service-based AI models across text and multimodal tasks.

  • Project
    EvolvingLMMs-Lab/lmms-eval

    lmms‑eval is a Python toolkit that provides a single, reproducible pipeline for benchmarking multimodal LLMs (vision, video, audio). It ships with 100+ tasks, 30+ model wrappers, statistical reporting, async serving, an HTTP evaluation server, and an optional web UI, all installable via `uv` or pip.

  • Project
    Ma-Zhuang/OmniNWM

    OmniNWM is a research codebase that builds a panoramic, multi‑modal world model for autonomous‑driving simulation. It jointly generates RGB, semantic, depth, and 3‑D occupancy maps from a vehicle trajectory, uses normalized Plücker ray‑maps for precise control, and provides occupancy‑based dense rewards for closed‑loop policy evaluation. The repo includes installation steps (including a manual patch to `transformers`), data preparation instructions for nuScenes, pretrained checkpoint download links, and scripts for inference, out‑of‑distribution testing, and staged training.

  • Project
    jmaczan/tiny-vllm

    tiny‑vllm is a C++/CUDA LLM inference engine (a smaller, educational version of vLLM) that loads a Llama 3.2 1B Instruct model from safetensors, runs the full forward pass on GPU (prefill + decode), implements KV‑cache, static/continuous batching, FlashAttention‑style online softmax and PagedAttention, and includes a step‑by‑step tutorial for learning or teaching GPU‑based LLM serving.