Update on "Add test case to export, quantize and lower vision encoder model for ET"

Differential Revision: [D67878162](https://our.internmc.facebook.com/intern/diff/D67878162)

[ghstack-poisoned]

Author: tarun292

examples/models/llama3_2_vision/vision_encoder/model.py

@@ -14,8 +14,12 @@
     replace_tile_positional_embedding,
     replace_tiled_token_positional_embedding,
 )
+from executorch.extension.llm.modules.attention import (
+    replace_mha_with_inference_mha,
+    replace_sdpa_with_custom_op,
+)
 from torchtune.models.llama3_2_vision._component_builders import llama3_2_vision_encoder
-from executorch.extension.llm.modules.attention import replace_mha_with_inference_mha, replace_sdpa_with_custom_op
+
 
 @dataclass
 class VisionEncoderConfig:
@@ -47,7 +51,11 @@ class VisionEncoderConfig:
 
 
 class FlamingoVisionEncoderModel(EagerModelBase):
-    def __init__(self, config: Optional[VisionEncoderConfig] = None, enable_source_transforms = True):
+    def __init__(
+        self,
+        config: Optional[VisionEncoderConfig] = None,
+        enable_source_transforms=True,
+    ):
         super().__init__()
         if config is None:
             config = demo_config

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

@@ -11,25 +11,25 @@
 import unittest
 
 import torch
+from executorch.backends.transforms.duplicate_dynamic_quant_chain import (
+    DuplicateDynamicQuantChainPass,
+)
+from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
+    XnnpackDynamicallyQuantizedPartitioner,
+)
 
 from executorch.examples.models.llama3_2_vision.vision_encoder import (
     FlamingoVisionEncoderModel,
 )
-from torch.testing import assert_close
-from executorch.exir import to_edge, to_edge_transform_and_lower, EdgeCompileConfig
-from torch._inductor.package import package_aoti
-from torch.nn.attention import SDPBackend
+from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
 from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
 from torch.ao.quantization.quantizer.xnnpack_quantizer import (
     get_symmetric_quantization_config,
     XNNPACKQuantizer,
 )
-from executorch.backends.transforms.duplicate_dynamic_quant_chain import (
-    DuplicateDynamicQuantChainPass
-)
-from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
-                    XnnpackDynamicallyQuantizedPartitioner,
-)
+from torch.nn.attention import SDPBackend
+from torch.testing import assert_close
+
 
 class FlamingoVisionEncoderTest(unittest.TestCase):
     def setUp(self) -> None:
@@ -38,15 +38,30 @@ def setUp(self) -> None:
     def test_flamingo_vision_encoder_et(self) -> None:
         with torch.no_grad():
             vision_model = FlamingoVisionEncoderModel(enable_source_transforms=False)
-            encoder_no_source_transform_outputs = vision_model.model.forward(*vision_model.get_example_inputs())
+            encoder_no_source_transform_outputs = vision_model.model.forward(
+                *vision_model.get_example_inputs()
+            )
             vision_model.source_transofrm()
             encoder = vision_model.model
-            encoder_source_transform_outputs = encoder.forward(*vision_model.get_example_inputs())
-            assert_close(encoder_source_transform_outputs, encoder_no_source_transform_outputs)
+            encoder_source_transform_outputs = encoder.forward(
+                *vision_model.get_example_inputs()
+            )
+            assert_close(
+                encoder_source_transform_outputs, encoder_no_source_transform_outputs
+            )
 
-            with torch.nn.attention.sdpa_kernel([SDPBackend.MATH]), torch.no_grad(), tempfile.TemporaryDirectory() as tmpdir:
-                training_output = torch.export.export_for_training(encoder, vision_model.get_example_inputs(), dynamic_shapes=vision_model.get_dynamic_shapes())
-                assert_close(encoder(*vision_model.get_example_inputs()), training_output.module()(*vision_model.get_example_inputs()))
+            with torch.nn.attention.sdpa_kernel(
+                [SDPBackend.MATH]
+            ), torch.no_grad(), tempfile.TemporaryDirectory() as tmpdir:
+                training_output = torch.export.export_for_training(
+                    encoder,
+                    vision_model.get_example_inputs(),
+                    dynamic_shapes=vision_model.get_dynamic_shapes(),
+                )
+                assert_close(
+                    encoder(*vision_model.get_example_inputs()),
+                    training_output.module()(*vision_model.get_example_inputs()),
+                )
 
                 dynamic_quantizer = XNNPACKQuantizer()
                 operator_config_dynamic = get_symmetric_quantization_config(
@@ -58,11 +73,19 @@ def test_flamingo_vision_encoder_et(self) -> None:
                 convert = convert_pt2e(prepare)
                 DuplicateDynamicQuantChainPass()(convert)
 
-                export_output = torch.export.export(convert, vision_model.get_example_inputs(), dynamic_shapes=vision_model.get_dynamic_shapes())
+                export_output = torch.export.export(
+                    convert,
+                    vision_model.get_example_inputs(),
+                    dynamic_shapes=vision_model.get_dynamic_shapes(),
+                )
 
-                edge = to_edge_transform_and_lower(export_output, partitioner=[
-                    XnnpackDynamicallyQuantizedPartitioner(),
-                ], compile_config=EdgeCompileConfig(_check_ir_validity=False))
+                edge = to_edge_transform_and_lower(
+                    export_output,
+                    partitioner=[
+                        XnnpackDynamicallyQuantizedPartitioner(),
+                    ],
+                    compile_config=EdgeCompileConfig(_check_ir_validity=False),
+                )
                 edge.to_executorch()
 
     def test_flamingo_vision_encoder_aoti(self) -> None:

extension/llm/modules/attention.py

@@ -9,11 +9,11 @@
 
 import torch
 import torchtune.modules.attention as TorchTuneAttention
+from executorch.examples.models.llama.source_transformation.sdpa import SDPACustom
 from executorch.extension.llm.modules.kv_cache import KVCache as InferenceKVCache
 from torch import nn
 from torchtune.modules.attention_utils import _MaskType, _sdpa_or_flex_attention
 from torchtune.modules.kv_cache import KVCache
-from executorch.examples.models.llama.source_transformation.sdpa import SDPACustom
 
 logger = logging.getLogger(__name__)
 
@@ -367,7 +367,6 @@ def forward(
             k = k.unsqueeze(2).expand(expand_shape).flatten(1, 2)
             v = v.unsqueeze(2).expand(expand_shape).flatten(1, 2)
 
-
         output = self._attention_fn(
             q,
             k,
@@ -431,5 +430,6 @@ def _replace_sdpa_with_custom_op(module: torch.nn.Module):
 
 def replace_sdpa_with_custom_op(module: torch.nn.Module) -> torch.nn.Module:
     from executorch.extension.llm.custom_ops import custom_ops
+
     _replace_sdpa_with_custom_op(module)
     return module