VLM: Onboarding native resolution, native aspect ratio, interleaved VLM training (#1615)

First PR to onboarding modern VLM training to torchtitan.

## Features:
- Native Aspect Ratio: not limited to square crops.
- Native Resolution: images in a batch can have different sizes, no more
image tiles and thumbnails.
- Native Interleaved data: training samples can have variable number of
images, interleaved with text at different position. You can train more
than just a captioning model.

## Design

Distributed training usually does not play nice with input of varying
shapes. To handle a varying number of images and image sizes, we
requires two additional hyperparameters, number of images per batch `N`
and max image patches length `L`, then we pad the actual image patches
to this fixed size.

<img width="1398" height="840" alt="Screenshot 2025-08-21 at 16 21 57"
src="https://github.com/user-attachments/assets/63fcbbc1-c587-4a63-8246-411cb72f5789"
/>

- After `tok_embedding`, we obtain tokens of shape `BxS`.
- After `encoder`, we obtain visual tokens of shape `NxL`.
- We extract the valid visual tokens only 
- Then scatter those tokens to their actual positions in the LLM input
tokens.

This requires the dataloader to handle the following aspect:
- Interleave the correct precise numbers of image tokens in the inputs
token based on encoder's patch size and input images' size
- Convert images/videos to 1D sequence of patchs:
- `rearrange(pixels, 'n (t pt) (h ph) (w pw) c -> n (t h w) (pt p pw
c)', pt=temporal_ps, ph=patch_size, pw=patch_size)`
- Pad all image patches sequence to a fixed length and return
`pixel_values.shape == [N, L, D]`
- Return a `grid_thw.shape == [N, L, 3]` to keep track of the location
indicies of each patches in the images. Padding image can be tracked in
the same tensors with values `-1`.
  
This result in a very simple and general interface to train modern VLM
with interleaved data and native resolution & aspect ratio:
- Depending on data mixtures, we can set dataloader's hyperparameters
`N, L` to have minimal empty image padding (in batch dimension).
- Use modern pytorch features (Flex Attention, compile etc) for
efficient handling of different attention mask per (padding in sequence
dimension).
- Interface nicely with TP, PP, etc

## In this PR
- Minimal interleaved Obelics dataloader with native resolution and
aspect ratio.
- The dataloader is currently very slow, as it need to download images
from internet everytime you run. (Same thing for the current imp in the
`multimodal` experiment).
- Siglip2 model code, mostly based on HF.
- VLM model code called `Llama3Siglip2` connecting the two vision
encoder and language decoder.
- Minimal infra code for debug model to run

<img width="1672" height="1303" alt="Screenshot 2025-08-21 at 15 25 25"
src="https://github.com/user-attachments/assets/c5c70ae2-04cf-4459-90d7-77d045291b88"
/>


## Todo:
- [x] Add support for captioning HF dataset that has images stored
inside the dataset (CC12M like Flux exp?) so it's not super slow to load
- Flex Attention for encoder.
- Modify Llama3 tokenizer to add special tokens.
- Script to combine Siglip2 + Llama3 weights and load. 
- Test Siglip2 encoder correctness.
- Multimodal CE loss to correct for image token bias
- All the parallelisms DP, CP, TP, PP.

---------

Co-authored-by: Jiani Wang <40016222+wwwjn@users.noreply.github.com>
Co-authored-by: Ankit Singh <ankitsingh135@gmail.com>
Co-authored-by: tianyu-l <150487191+tianyu-l@users.noreply.github.com>

Author: 4 people

tests/assets/tokenizer/tokenizer.json

@@ -2029,7 +2029,11 @@
       "land": 1994,
       "?\n": 1995,
       " respect": 1996,
-      "ances": 1997
+      "ances": 1997,
+      "<|image|>": 1998,
+      "<|begin_of_image|>": 1999,
+      "<|end_of_image|>": 2000,
+      "<|pad|>": 2001
     },
     "merges": [
     ]

tests/assets/tokenizer/tokenizer_config.json

@@ -15,11 +15,47 @@
       "rstrip": false,
       "single_word": false,
       "special": true
+    },
+    "1998": {
+      "content": "<|image|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1999": {
+      "content": "<|begin_of_image|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2000": {
+      "content": "<|end_of_image|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2001": {
+      "content": "<|pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
     }
   },
   "bos_token": "<|begin_of_text|>",
   "clean_up_tokenization_spaces": true,
   "eos_token": "<|end_of_text|>",
+  "img_token": "<|image|>",
+  "boi_token": "<|begin_of_image|>",
+  "eoi_token": "<|end_of_image|>",
+  "pad_token": "<|pad|>",
   "model_input_names": [
     "input_ids",
     "attention_mask"

torchtitan/experiments/__init__.py

@@ -7,3 +7,4 @@
 import torchtitan.experiments.llama4  # noqa: F401
 import torchtitan.experiments.qwen3
 import torchtitan.experiments.simple_fsdp  # noqa: F401
+import torchtitan.experiments.vlm  # noqa: F401

torchtitan/experiments/llama4/__init__.py

@@ -30,7 +30,7 @@
         dim=256,
         n_layers=6,
         n_heads=16,
-        vocab_size=2000,
+        vocab_size=2048,
         rope_theta=500000,
     ),
     "17bx16e": TransformerModelArgs(
@@ -59,7 +59,7 @@
         dim=256,
         n_layers=6,
         n_heads=16,
-        vocab_size=2000,
+        vocab_size=2048,
         rope_theta=500000,
         every_n_layers_nope=4,
         fixed_attn_block_size=256,

torchtitan/experiments/vlm/README.md

@@ -0,0 +1,57 @@
+# Vision Language Model training in `torchtitan`
+
+**under active development**
+
+This folder showcases how to train modern Vision Language Model (vlm) in torchtitan.
+
+
+## Features:
+- Native Aspect Ratio: not limited to square crops.
+- Native Resolution: images in a batch can have different sizes, no more image tiles and thumbnails.
+- Native Interleaved data: training samples can have variable number of images, interleaved with text at different position. You can train more than just a captioning model.
+
+
+## Design
+Distributed training usually does not play nice with input of varying shapes. To handle a varying number of images and image sizes, we requires two hyperparameters, image batch size `N` and image length `L` (in patches), and pad the actual image patches to this fixed size.
+Then we scatter the patch embeddings to their actual positions in the LLM input tokens.
+
+<img width="1398" height="840" alt="Screenshot 2025-08-21 at 16 21 57" src="https://github.com/user-attachments/assets/63fcbbc1-c587-4a63-8246-411cb72f5789" />
+
+- After `tok_embedding`, we obtain tokens of shape `BxS`.
+- After `encoder`, we obtain visual tokens of shape `NxL`.
+- We extract the valid visual tokens only
+- Then scatter those tokens to their actual positions in the LLM input tokens.
+
+
+This results in a very simple and general interface to train modern VLM with interleaved data and native resolution & aspect ratio:
+- Depending on data mixtures, we can set dataloader's hyperparameters `N, L` to have minimal empty image padding (in batch dimension).
+- We use modern Pytorch features like FlexAttention and torch.compile to efficient efficiently handle variable sequence length.
+- Interface nicely with TP, PP, etc
+
+
+## Implementation
+
+### Dataloader
+This approach requires the dataloader to handle the following aspect:
+- [x] Interleave the correct precise numbers of image tokens in the inputs token based on encoder's patch size and input images' size
+- [x] Convert images/videos to 1D sequence of patchs:
+  - `rearrange(pixels, 'n (t pt) (h ph) (w pw) c -> n (t h w) (pt p pw c)', pt=temporal_ps, ph=patch_size, pw=patch_size)`
+  - Pad all image patches sequence to a fixed length and return `pixel_values.shape == [N, L, D]`
+- [x] Return a `grid_thw.shape == [N, L, 3]` to keep track of the location indicies of each patches in the images. Padding image can be tracked in the same tensors with values `-1`.
+- [x] LLM Sample / Document Packing.
+- [x] Captioning dataset: CC12M
+- [x] Interleaved dataset: Obelics
+
+
+
+### Model
+We also need a pretrained vision encoder with support for native resolution and aspect ratio. There is relatively few Vision Encoder that have this capability up until recently, including Siglip2, AimV2, and most recently DINOv3.
+- [ ] Currently we support Siglip2 encoder using Positional Embedding interpolation approach.
+    - [x] Base modelling code.
+    - [ ] Weights conversion and loading from HF.
+- [x] FSDP for both Encoder and Decoder
+- [x] Context Parallel for LLM only, since we will use FlexAttention for Encoder.
+- [ ] FlexAttention for with different seq len per image.
+- [ ] Compile for Encoder + Deocoder
+- [ ] Tensor Parallel
+- [ ] Pipeline Parallel

torchtitan/experiments/vlm/__init__.py

@@ -0,0 +1,57 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from dataclasses import asdict, replace
+
+from torchtitan.components.loss import build_cross_entropy_loss
+from torchtitan.components.lr_scheduler import build_lr_schedulers
+from torchtitan.components.optimizer import build_optimizers
+from torchtitan.components.tokenizer import build_hf_tokenizer
+from torchtitan.components.validate import build_validator
+from torchtitan.models.llama3 import llama3_configs
+from torchtitan.protocols.train_spec import register_train_spec, TrainSpec
+
+from .datasets.mm_datasets import build_mm_dataloader
+from .infra.parallelize import parallelize_vlm
+from .model.args import Llama3Siglip2ModelArgs, Siglip2ModelArgs
+from .model.model import Llama3Siglip2Transformer
+
+__all__ = [
+    "parallelize_vlm",
+    "Llama3Siglip2ModelArgs",
+    "Llama3Siglip2Transformer",
+    "llama3_siglip2_configs",
+]
+
+
+llama3_siglip2_configs = {
+    "debugmodel": Llama3Siglip2ModelArgs(
+        **asdict(replace(llama3_configs["debugmodel"], vocab_size=2048)),
+        encoder=Siglip2ModelArgs(
+            dim=128,
+            ffn_dim=256,
+            n_layers=4,
+            n_heads=2,
+        ),
+    ),
+}
+
+
+register_train_spec(
+    TrainSpec(
+        name="llama3-siglip2",
+        model_cls=Llama3Siglip2Transformer,
+        model_args=llama3_siglip2_configs,
+        parallelize_fn=parallelize_vlm,
+        pipelining_fn=None,
+        build_optimizers_fn=build_optimizers,
+        build_lr_schedulers_fn=build_lr_schedulers,
+        build_dataloader_fn=build_mm_dataloader,
+        build_tokenizer_fn=build_hf_tokenizer,
+        build_loss_fn=build_cross_entropy_loss,
+        build_validator_fn=build_validator,
+    )
+)

torchtitan/experiments/vlm/assets/job_config.py

@@ -0,0 +1,34 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class Data:
+    max_images_per_batch: int = 10
+    """Vision encoder batch size (N)"""
+    max_patches_per_image: int = 256
+    """Vision encoder sequence length (L)"""
+    patch_size: int = 16
+    """ Patch size of the vision encoder.
+    For example, image size 256x256, patch size 16
+        Number of visual tokens is: (256/16)**2=256
+    """
+    spatial_merge_size: int = 1
+    """ Spatially merge visual tokens after encoder.  Default 1 means no merging.
+    For example: image size 256x256, patch size 16, spaitl merge size is 2
+        Number of visual tokens for the LLM: (256/16/2)**2 = 8
+    """
+    packing_buffer_size: int = 0
+    """ Set to a value >0 to enable sample packing.
+    This control the buffer uses to store training samples avaliable for packing.
+    """
+
+
+@dataclass
+class JobConfig:
+    data: Data = field(default_factory=Data)