add resolution changes tests to hotswapping test suite.

Author: sayakpaul

tests/models/test_modeling_common.py

@@ -2046,7 +2046,9 @@ def get_linear_module_name_other_than_attn(self, model):
         ]
         return linear_names[0]
 
-    def check_model_hotswap(self, do_compile, rank0, rank1, target_modules0, target_modules1=None):
+    def check_model_hotswap(
+        self, do_compile, rank0, rank1, target_modules0, target_modules1=None, different_resolutions=None
+    ):
         """
         Check that hotswapping works on a small unet.
 
@@ -2056,6 +2058,7 @@ def check_model_hotswap(self, do_compile, rank0, rank1, target_modules0, target_
         - hotswap the second adapter
         - check that the outputs are correct
         - optionally compile the model
+        - optionally check if recompilations happen on different shapes
 
         Note: We set rank == alpha here because save_lora_adapter does not save the alpha scalings, thus the test would
         fail if the values are different. Since rank != alpha does not matter for the purpose of this test, this is
@@ -2110,10 +2113,17 @@ def check_model_hotswap(self, do_compile, rank0, rank1, target_modules0, target_
             model.load_lora_adapter(file_name0, safe_serialization=True, adapter_name="adapter0", prefix=None)
 
             if do_compile:
-                model = torch.compile(model, mode="reduce-overhead")
+                model = torch.compile(model, mode="reduce-overhead", dynamic=different_resolutions is not None)
 
             with torch.inference_mode():
                 output0_after = model(**inputs_dict)["sample"]
+
+                # additionally check if dynamic compilation works.
+                if different_resolutions is not None:
+                    for height, width in self.different_shapes_for_compilation:
+                        new_inputs_dict = self.prepare_dummy_input(height=height, width=width)
+                        _ = model(**new_inputs_dict)
+
             assert torch.allclose(output0_before, output0_after, atol=tol, rtol=tol)
 
             # hotswap the 2nd adapter
@@ -2122,6 +2132,12 @@ def check_model_hotswap(self, do_compile, rank0, rank1, target_modules0, target_
             # we need to call forward to potentially trigger recompilation
             with torch.inference_mode():
                 output1_after = model(**inputs_dict)["sample"]
+
+                if different_resolutions is not None:
+                    for height, width in self.different_shapes_for_compilation:
+                        new_inputs_dict = self.prepare_dummy_input(height=height, width=width)
+                        _ = model(**new_inputs_dict)
+
             assert torch.allclose(output1_before, output1_after, atol=tol, rtol=tol)
 
             # check error when not passing valid adapter name
@@ -2240,3 +2256,21 @@ def test_hotswap_second_adapter_targets_more_layers_raises(self):
                     do_compile=True, rank0=8, rank1=8, target_modules0=target_modules0, target_modules1=target_modules1
                 )
                 assert any("Hotswapping adapter0 was unsuccessful" in log for log in cm.output)
+
+    @parameterized.expand([(11, 11), (7, 13), (13, 7)])
+    @require_torch_version_greater("2.7.1")
+    def test_hotswapping_compile_on_different_shapes(self, rank0, rank1):
+        different_shapes_for_compilation = self.different_shapes_for_compilation
+        if different_shapes_for_compilation is None:
+            pytest.skip(f"Skipping as `different_shapes_for_compilation` is not set for {self.__class__.__name__}.")
+        torch.fx.experimental._config.use_duck_shape = False
+
+        target_modules = ["to_q", "to_k", "to_v", "to_out.0"]
+        with torch._dynamo.config.patch(error_on_recompile=True):
+            self.check_model_hotswap(
+                do_compile=True,
+                rank0=rank0,
+                rank1=rank1,
+                target_modules0=target_modules,
+                different_resolutions=different_shapes_for_compilation,
+            )