2D RoPE + CLIP updates

[RdoubleA]

Context

Two-dimensional rotary positional embeddings have been added to vision transformers to improve performance. This was explored in papers such as Eva-02-CLIP, which found that 2D RoPE improved performance and had more stable training as opposed to 1D RoPE for images. Another novel architecture FiT (Flexible Vision Transformer for Diffusion Model) similarly employs 2D RoPE for image resolution generalization. Pixtral, a multimodal LLM, also uses a similar 2D RoPE mechanism as well, as seen in Hugging Face. A full survey of 2D RoPE can be found in this paper.

Here, we add VisionRotaryPositionalEmbeddings as a general component. The forward is identical to RotaryPositionalEmbeddings; the major difference is in the build_rope_cache, where we need to take into account the x and y positions of the patches in the image grid, as defined by patch_size and tile_size. It simply applies 1D RoPE with half the usual dim on x and y independently and concatenates them together.

This is exposed as use_rope in the clip_vision_encoder builder, which will enable 2D RoPE. I also include some various fixes for CLIP and additional configurability

Changelog

What are the changes made in this PR?

• Add VisionRotaryPositionalEmbeddings with tests
• Expose ability to configure RoPE on clip encoder
• intermediate_act is not used in the clip_vision_encoder. This is superceded by activation, so I just removed it
• apply_lora_to_output is not used in lora_clip_vision_encoder. There is a TODO to add this to the CLSProjection. I removed it so users don't accidentally configure it and get unexpected behavior
• CLIP's CLSEmbedding adds the CLS token to the beginning of the input sequence. I expose this as a configurable parameter to allow the option of adding it to the end of the sequence instead, which some variations of CLIP may do.
• attn_bias was not being propagated in lora_clip_vision_encoder, so I fixed this and changed default to False to match current behavior.
• Typing fixes for llama3 rope parameters
• Move Phi3 rope tests to its own file

Test plan

• Recreated clip_vision_encoder with append_cls_token=True and manually verified that cls token embedding is in correct position, added test
• Tests for VisionRotaryPositionalEmbeddings and parity check against an internal reference implementation

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1973

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [DomainsOnly] Jobs fail with GLIBC version not found

✅ No Failures

As of commit 77ecb82 with merge base bca5899 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[joecummings]

buttttt tests that run in CI?

[joecummings]

This looks good to me, but should probably get @pbontrager 's eyes on this as well.

[acisseJZhong]

a single tile with shape (400, 400)? if max_num_tiles=1, then tile size should be equal to patch_size?

[RdoubleA]

ooh ya good catch, let me update this and the docstring where I copied this from

[felipemello1]

I think that "your image as a single tile. with shape (40, 40) each." ---> "your image as a single tile with shape (400, 400)."

I didnt read the code, just sharing what i know from llama 3.2 vision. This explains it visually:

torchtune/torchtune/modules/vision_transformer.py

Line 61 in bd89bb8

Image: shape (8x8)

A image is broken into tiles. E.g, if max_num_tiles=4 and tile_size=400, a 1600x1600 image is broken into 4 tiles of 400x400, e.g. (num images, num tiles, x, y) = (1,4,400x400)

Then each tile is broken down into patches.

If max_num_tiles=1, then it means that 1 image = 1 tile. The tile dimension doesnt go away, e.g. (num images, num tiles, x, y) = (1,1,x, y)