fixes

Author: sayakpaul

tests/models/test_modeling_common.py

@@ -837,12 +837,13 @@ def test_effective_gradient_checkpointing(self, loss_tolerance=1e-5, param_grad_
                 continue
             self.assertTrue(torch_all_close(param.grad.data, named_params_2[name].grad.data, atol=param_grad_tol))
 
+    @unittest.skipIf(torch_device == "mps", "This test is not supported for MPS devices.")
     def test_gradient_checkpointing_is_applied(
         self, expected_set=None, attention_head_dim=None, num_attention_heads=None, block_out_channels=None
     ):
         if not self.model_class._supports_gradient_checkpointing:
             return  # Skip test if model does not support gradient checkpointing
-        if torch_device == "mps" and self.model_class.__name__ in [
+        if self.model_class.__name__ in [
             "UNetSpatioTemporalConditionModel",
             "AutoencoderKLTemporalDecoder",
         ]:
@@ -876,6 +877,8 @@ def _set_gradient_checkpointing_new(self, module, value=False):
         model = model_class_copy(**init_dict)
         model.enable_gradient_checkpointing()
 
+        print(f"{set(modules_with_gc_enabled.keys())=}, {expected_set=}")
+
         assert set(modules_with_gc_enabled.keys()) == expected_set
         assert all(modules_with_gc_enabled.values()), "All modules should be enabled"
 

tests/models/transformers/test_models_transformer_allegro.py

@@ -77,3 +77,7 @@ def prepare_init_args_and_inputs_for_common(self):
         }
         inputs_dict = self.dummy_input
         return init_dict, inputs_dict
+
+    def test_gradient_checkpointing_is_applied(self):
+        expected_set = {"AllegroTransformer3DModel"}
+        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)

tests/models/transformers/test_models_transformer_cogview3plus.py

@@ -83,3 +83,7 @@ def prepare_init_args_and_inputs_for_common(self):
         }
         inputs_dict = self.dummy_input
         return init_dict, inputs_dict
+
+    def test_gradient_checkpointing_is_applied(self):
+        expected_set = {"CogView3PlusTransformer2DModel"}
+        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)

tests/models/transformers/test_models_transformer_sd3.py

@@ -84,6 +84,10 @@ def prepare_init_args_and_inputs_for_common(self):
     def test_set_attn_processor_for_determinism(self):
         pass
 
+    def test_gradient_checkpointing_is_applied(self):
+        expected_set = {"SD3Transformer2DModel"}
+        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+
 
 class SD35TransformerTests(ModelTesterMixin, unittest.TestCase):
     model_class = SD3Transformer2DModel