Update on "[llama-mm] Onboard Llama3.2 mm vision encoder"

Summary: Add llama3.2 mm vision encoder to examples/models.

We need to do a module swapping for TilePositionEmbedding to make sure
vision encoder is exportable.

Test Plan: Unit tests.

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]

Author: larryliu0820

.ci/scripts/gather_test_models.py

@@ -24,6 +24,7 @@
         "ic4": "linux.12xlarge",
         "resnet50": "linux.12xlarge",
         "llava": "linux.12xlarge",
+        "llama3_2_vision_encoder": "linux.12xlarge",
         # This one causes timeout on smaller runner, the root cause is unclear (T161064121)
         "dl3": "linux.12xlarge",
         "emformer_join": "linux.12xlarge",

examples/models/llama3_2_vision/vision_encoder/model.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 from dataclasses import dataclass, field
+from typing import Optional
 
 import torch
 
@@ -41,8 +42,10 @@ class VisionEncoderConfig:
 
 
 class FlamingoVisionEncoderModel(EagerModelBase):
-    def __init__(self, config: VisionEncoderConfig = VisionEncoderConfig()):
+    def __init__(self, config: Optional[VisionEncoderConfig] = None):
         super().__init__()
+        if config is None:
+            config = VisionEncoderConfig()
         self.config = config
         self.model = flamingo_vision_encoder(
             patch_size=config.patch_size,
@@ -56,6 +59,7 @@ def __init__(self, config: VisionEncoderConfig = VisionEncoderConfig()):
             max_num_tiles=config.max_num_tiles,
             in_channels=config.in_channels,
         )
+        self.model = replace_tile_positional_embedding(self.model)
         self.image = torch.randn(
             1, 1, 4, 3, self.config.tile_size, self.config.tile_size
         )
@@ -66,7 +70,6 @@ def __init__(self, config: VisionEncoderConfig = VisionEncoderConfig()):
         )
 
     def get_eager_model(self, **kwargs):
-        self.model = replace_tile_positional_embedding(self.model)
         return self.model
 
     def get_example_inputs(self):

examples/models/llama3_2_vision/vision_encoder/test/test_vision_encoder.py

@@ -14,17 +14,16 @@
 
 from executorch.examples.models.llama3_2_vision.vision_encoder import (
     FlamingoVisionEncoderModel,
-    VisionEncoderConfig,
 )
-from torch._inductor.package import load_package, package_aoti
+from torch._inductor.package import package_aoti
 
 
 class FlamingoVisionEncoderTest(unittest.TestCase):
     def setUp(self) -> None:
         super().setUp()
 
     def test_flamingo_vision_encoder(self) -> None:
-        model = FlamingoVisionEncoderModel(VisionEncoderConfig())
+        model = FlamingoVisionEncoderModel()
         encoder = model.model
         eager_res = encoder.forward(*model.get_example_inputs())
 
@@ -38,7 +37,7 @@ def test_flamingo_vision_encoder(self) -> None:
         with tempfile.TemporaryDirectory() as tmpdir:
             path = package_aoti(os.path.join(tmpdir, "vision_encoder.pt2"), so)
             print(path)
-            encoder_aoti = load_package(path)
+            encoder_aoti = torch._inductor.aoti_load_package(path)
 
             y = encoder_aoti(*model.get_example_inputs())
 

extension/llm/modules/_position_embeddings.py

@@ -188,10 +188,13 @@ def forward(self, x: torch.Tensor, aspect_ratio: torch.Tensor) -> torch.Tensor:
             torch._check(n_tiles_w >= 1)
             torch._check(n_tiles_h <= self.max_num_tiles)
             torch._check(n_tiles_w <= self.max_num_tiles)
+            # TODO: Remove this once pytorch/pytorch#120288 is fixed
             padded_embedding = F.pad(self.embedding, (0, 0, 0, 0, 0, 1, 0, 1))
             pos_embed = padded_embedding[:n_tiles_h, :n_tiles_w, :, :]
 
-            # Add pos encoding to the non padded tiles.
+            # We need to do a clone here in order to make this model export
+            # friendly as the reshape is collapsing dim 0 and dim 1 into a
+            # single dim.
             pos_embed = pos_embed.clone()
             pos_embed = pos_embed.reshape(n_non_padded_tiles, 1, self.embed_dim)