VLM: Onboarding native resolution, native aspect ratio, interleaved VLM training

[lkhphuc]

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.

• 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

Todo:

• 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.

[wwwjn]

Thanks for making the great PR! I learned a lot from this PR personally. However, I feel like the data preprocessing part in mm_collator_nld.py is a little bit hard to follow and read.

For image preprocessing, it mainly happens inmm_collator_nld.py , and the collator functions contains following steps for images:

1. Patchify
2. Generate Grids with coordinations
3. Padding/ truncate
4. Assemble as batched outputs

And text preprocessing is mainly handled in mm_dataset.py, which also contains several steps, eg padding with <image> tokens, Tokenization, mask out <image> tokens in label.

I was wondering can we future split the image and text preprocessing function into smaller code pieces, adding tensor shape hints, or even adding examples like experiments/multimodal? In this way, we could increase readability and easy to debug

The VLM modeling parts LGTM, it's clear structured!

[wwwjn]

I would suggest we keep a test tokenizer for VLM only, under the experiments/vlm/assets folder. Currently I'm debugging a know issue #1136, adding these vocab will make the test tokenizer's vocab size not dividable by worldsize, which might causing failure of other LLM when resume training

[wwwjn]

n00b question: why we use CLIP mean/std to normalize the dataset? Is it a common practice?

[wwwjn]

Does this means we will drop those images > max_images_per_batch

[wwwjn]

For this path, we can not easily run or test on our GPUs because our cluster don't have Internet access :( I might try finding some alternative dataset rather than HuggingFaceM4/OBELICS

[wwwjn]

By reading the code, T dimension here will always be 1 for image samples?