nanoMFM: Educational Multimodal Foundation Models

Learn to build multimodal generative "foundation" models from scratch. This educational repository walks you through implementing minimal versions of GPT, MaskGIT, 4M, Flow Matching, and vision-language models step by step. Versions of these modules are taught in VILAB's CS-503 Visual Intelligence and COM-304 Intelligent Systems: Communications & AI courses at EPFL.

Overview

The repository currently contains five modules:

1. nanoGPT - Autoregressive transformer for language and image generation, in the style of GPT-2, inspired by nanoGPT
2. nanoMaskGIT - Image (and language) generation through unimodal masked modeling, inspired by MaskGIT
3. nano4M - Multimodal any-to-any model (4M-style any-to-any generation through multimodal masked modeling)
4. nanoFlow - Continuous-time generative modeling with Rectified Flow / Flow Matching and a DiT-style denoiser
5. nanoVLM - Vision-language model similar to LLaVA-style models for image-text understanding and visual question answering, based on nanoVLM

In addition, we provide a 4M Tutorial (notebooks/nanoMFM_part3_extension/) that builds upon a pretrained 4M model for multimodal generation and retrieval applications. This serves as an extension to nano4M (module 3).

nanoMFM is designed with two approaches in mind:

• Exercises and notebooks: We provide Jupyter notebooks (in notebooks/) for each module that guide you through the implementation, training, and inference steps. Students are asked to fill in the gaps in the code (in nanomfm/_exercises/), train models, and visualize model predictions in the Jupyter notebooks. If you wish to follow the exercises and implement the modules yourself, please see the instructions in the Jupyter notebooks after following the setup instructions below.
• Get started right away: For those who want to tinker with a ready codebase or check solutions for the exercises, we provide our reference implementation (in nanomfm/_solutions/).

Note

• nanoMFM is designed to be as minimal as possible and makes several trade-offs for educational purposes. The goal of this repository is to teach the modeling fundamentals for training a tiny language model, a generative model, a multimodal any-to-any model, and a vision-language model.
• Solving the exercises requires a basic understanding of deep learning and PyTorch, as well as familiarity with the Transformer architecture in theory.
• To be as accessible as possible, we use relatively small datasets (TinyStories, MNIST, MultimodalCLEVR, The Cauldron) and train small models with modest compute requirements (tested with 1-4 V100 GPUs). If you have more recent hardware, feel free to adapt the configs (e.g. adjusting batch size, number of GPUs, mixed precision, ...).
• We do not include the most recent architecture and training techniques, but we may extend the repository at a later date. The modeling fundamentals you learn here will translate when scaled to somewhat larger settings, but we do not go into infrastructure and training specifics that become necessary when scaling.

Repository Structure

nanoMFM/
├── nanomfm/                     # Main Python package
│   ├── __init__.py
│   ├── data/                    # Data loading (shared, includes data/vlm for nanoVLM)
│   ├── utils/                   # Utilities (shared)
│   ├── modeling/                # Public interface (switches based on mode)
│   ├── models/                  # Public interface (switches based on mode)
│   ├── training/                # Training loop, scheduler, logging (nanoVLM)
│   ├── _solutions/              # Complete implementations
│   │   ├── data/
│   │   ├── modeling/
│   │   └── models/
│   └── _exercises/              # Exercise stubs with TODOs
│       ├── data/
│       ├── modeling/
│       └── models/
├── cfgs/                        # Training configurations
│   ├── nanoGPT/
│   ├── nanoMaskGIT/
│   ├── nano4M/
│   └── nanoVLM/
├── scripts/                     # Training scripts
│   ├── run_training.py
│   ├── train_vlm.py
│   ├── download_datasets.sh
│   └── setup_env.sh
├── notebooks/                   # Jupyter notebooks for learning
│   ├── nanoMFM_part1_nanoGPT.ipynb
│   ├── nanoMFM_part2_nanoMaskGIT.ipynb
│   ├── nanoMFM_part3_nano4M.ipynb
│   ├── nanoMFM_part3_extension/      # 4M Tutorial: using a pretrained 4M (generation + retrieval)
│   ├── nanoMFM_part4_nanoFlow.ipynb
│   └── nanoMFM_part5_nanoVLM.ipynb
└── pyproject.toml

Installation

git clone https://github.com/EPFL-VILAB/nanoMFM.git
cd nanoMFM
bash scripts/setup_env.sh
conda activate nanomfm

This will create a conda environment named nanomfm, install all dependencies, and set up the Jupyter kernel.

For the 4M Tutorial Extension (notebooks/nanoMFM_part3_extension/), install the additional dependencies after the base setup:

pip install -r notebooks/nanoMFM_part3_extension/requirements.txt

Data

Modules 1 & 2: nanoGPT and nanoMaskGIT

The first two modules use MNIST and TinyStories datasets, which are downloaded automatically when you run training for the first time. No manual setup is required.

nanoGPT generates TinyStories text and MNIST digits autoregressively, token by token:

Text generation (TinyStories)	Image generation (MNIST)

nanoMaskGIT generates images and text by iteratively unmasking tokens:

Text generation (TinyStories)	Image generation (MNIST)

Module 3: nano4M and 4M Tutorial Extension

The third module (nano4M) requires MultimodalCLEVR, and the 4M Tutorial Extension requires a subset of COCO validation images. Both are downloaded with a single script:

./scripts/download_datasets.sh ~/datasets

This downloads and extracts both datasets to ~/datasets/ and cleans up the archive files.

nano4M generates any subset of modalities (RGB, depth, normals, text) from any other.

Module 4: nanoFlow

The fourth module (nanoFlow) uses MNIST and CIFAR-10, both downloaded automatically via torchvision when you run the notebook for the first time. No manual setup is required. Unlike the first three modules, nanoFlow training is run entirely inside its notebook (notebooks/nanoMFM_part4_nanoFlow.ipynb) rather than through scripts/run_training.py.

MNIST	CIFAR-10

Module 5: nanoVLM

The fifth module (nanoVLM) uses The Cauldron and MMStar datasets. They are downloaded automatically from Hugging Face when you run training for the first time.

nanoVLM answers questions about an input image.

Exercise and Training Modes

As mentioned above, we offer two ways to use this codebase: either implement the modules one by one using the Jupyter notebooks as guides, or get started right away with our reference implementation.

Mode Switching

We use an environment variable to switch between modes:

# Default (NANOMFM_EXERCISES not set or set to "0"): uses complete solutions
from nanomfm.models import GPT  # → nanomfm._solutions.models.gpt.GPT

# With NANOMFM_EXERCISES=1: uses exercise stubs
import os
os.environ["NANOMFM_EXERCISES"] = "1"
from nanomfm.models import GPT  # → nanomfm._exercises.models.gpt.GPT

For Students

To work through the exercises:

1. Set exercise mode:

   export NANOMFM_EXERCISES=1

   In exercise mode, imports resolve to stub implementations with TODOs for you to complete. Without this environment variable set, the reference implementation will be used instead.

2. Open the notebooks in notebooks/ using JupyterLab

3. Edit files in nanomfm/_exercises/ to complete the TODOs

4. Reference nanomfm/_solutions/ if you get stuck

For Running Complete Models

Simply use the package normally. Solutions (in nanomfm/_solutions/) are the default, so there is no need to set the NANOMFM_EXERCISES environment variable:

from nanomfm.models import GPT, MaskGIT, FourM, RectifiedFlow
from nanomfm.modeling import TransformerTrunk, Attention, DiT_Llama

Training

Don't forget to login with Weights & Biases first (wandb login <KEY> or export WANDB_API_KEY=<KEY>) and set the wandb_entity in the configs!

Module 1: nanoGPT

TinyStories (text generation — ~2h on 2× V100, target val loss ≈ 1.25):

OMP_NUM_THREADS=1 torchrun --nproc_per_node=2 scripts/run_training.py --config cfgs/nanoGPT/tinystories_d8w512.yaml

MNIST (image generation — a few minutes on 1 GPU, target val loss < 0.45):

OMP_NUM_THREADS=1 torchrun --nproc_per_node=1 scripts/run_training.py --config cfgs/nanoGPT/mnist_d8w512.yaml

Module 2: nanoMaskGIT

MNIST (a few minutes on 1 GPU, target val loss < 0.57):

OMP_NUM_THREADS=1 torchrun --nproc_per_node=1 scripts/run_training.py --config cfgs/nanoMaskGIT/mnist_d8w512.yaml

TinyStories (~2h on 2× V100, target val loss ≈ 1.93):

OMP_NUM_THREADS=1 torchrun --nproc_per_node=2 scripts/run_training.py --config cfgs/nanoMaskGIT/tinystories_d8w512.yaml

Module 3: nano4M

MultimodalCLEVR — 5B tokens (recommended starting point, ~5h on 4× V100, target val loss ≈ 3.4):

OMP_NUM_THREADS=1 torchrun --nproc_per_node=4 scripts/run_training.py --config cfgs/nano4M/multiclevr_d6-6w512_5B.yaml

MultimodalCLEVR — 25B tokens (longer run for better results, ~4h on 4× H100):

OMP_NUM_THREADS=1 torchrun --nproc_per_node=4 scripts/run_training.py --config cfgs/nano4M/multiclevr_d6-6w512_25B.yaml

Module 4: nanoFlow

nanoFlow training runs entirely inside notebooks/nanoMFM_part4_nanoFlow.ipynb — no external training script is needed. MNIST trains in ~5–10 minutes and CIFAR-10 in ~1–2 hours on a single GPU.

Module 5: nanoVLM

nanoVLM downloads pretrained backbones from Hugging Face, so a Hugging Face token is required:

export HUGGINGFACE_HUB_TOKEN=<your_hf_token>
export WANDB_API_KEY=<your_wandb_api_key>

OMP_NUM_THREADS=1 torchrun --nproc_per_node=2 scripts/train_vlm.py --config cfgs/nanoVLM/default.yaml

Multi-node training with SLURM

Click to expand SLURM configuration example

For multi-node distributed training on a SLURM cluster, create an sbatch script with the following configuration. This example uses 2 nodes with 1 GPU per node (2 GPUs total):

#!/bin/bash
#SBATCH --job-name=nano-gpt-run
#SBATCH --nodes=2                  # total number of nodes
#SBATCH --ntasks-per-node=1        # total number of tasks per node
#SBATCH --gpus-per-node=1
#SBATCH --time=01:00:00
#SBATCH --output=logs/%x_%j.log    # stdout location
#SBATCH --error=logs/%x_%j.err     # stderr location
#SBATCH --account=<slurm-account-name>

# Set environment variables
export WANDB_API_KEY=<wandb-api-key>
export MASTER_PORT=25678
export MASTER_ADDR=$(hostname)

# Run training across all nodes
srun -ul bash -c "
  # Change to nanoMFM directory
  cd /path/to/nanoMFM-dev
  pip install wandb[media]

  TORCHRUN_ARGS=\"
    --node-rank=\${SLURM_PROCID} \
    --master-addr=\${MASTER_ADDR} \
    --master-port=\${MASTER_PORT} \
    --nnodes=\${SLURM_NNODES} \
    --nproc-per-node=\${SLURM_GPUS_PER_NODE} \
  \"

  torchrun \${TORCHRUN_ARGS} scripts/run_training.py --config cfgs/nanoGPT/tinystories_d8w512.yaml
"

Make sure to:

• Replace <slurm-account-name> with your SLURM account
• Replace <wandb-api-key> with your Weights & Biases API key
• Update /path/to/nanoMFM-dev with the actual path to your nanoMFM directory
• Create a logs/ directory before running: mkdir -p logs
• Adjust --nodes, --gpus-per-node, and --time as needed

Submit the job with sbatch your_script.sh.

License

Apache 2.0 - See LICENSE for details.

Acknowledgments

• 4M (Apple) - Original 4M implementation
• nanoGPT (Karpathy) - Inspiration for nano implementations
• Hugging Face nanoVLM - Main implementation for our nanoVLM exercise

Citation

If you find this repository useful, please consider citing:

@misc{bachmann2026nanomfm,
  author    = {Bachmann, Roman and Gao, Zhitong and Khattak, Muhammad Uzair and Ye, Mingqiao and Zamir, Amir},
  title     = {nanoMFM: Educational Multimodal Foundation Models},
  publisher = {GitHub},
  year      = {2026},
  url       = {https://github.com/EPFL-VILAB/nanoMFM}
}