[DCP][AC] Add test for apply AC with FSDP1 (pytorch#126935)

Adding test for this cherry pick. pytorch#126559

Pull Request resolved: pytorch#126935
Approved by: https://github.com/fegin

test/distributed/checkpoint/test_state_dict.py

@@ -548,7 +548,7 @@ def test_cpu_offload_full_state_dict(self) -> None:
 
     @with_comms
     @skip_if_lt_x_gpu(1)
-    def test_activation_ckpt_fqns(self) -> None:
+    def test_activation_ckpt_fqns_ddp(self) -> None:
         """Tests that activation checkpointing prefixes are removed from module names"""
         model = CompositeParamModel(device=torch.device("cuda"))
         original_keys = get_model_state_dict(model).keys()
@@ -559,6 +559,25 @@ def test_activation_ckpt_fqns(self) -> None:
 
         self.assertEqual(original_keys, new_keys)
 
+    @with_comms
+    @skip_if_lt_x_gpu(1)
+    def test_activation_ckpt_fqns_fsdp1(self) -> None:
+        self.run_subtests(
+            {"use_orig_params": [True, False]},
+            self._test_activation_ckpt_fqns_fsdp1,
+        )
+
+    def _test_activation_ckpt_fqns_fsdp1(self, use_orig_params: bool) -> None:
+        """Tests that activation checkpointing prefixes are removed from module names"""
+        model = CompositeParamModel(device=torch.device("cuda"))
+        original_keys = get_model_state_dict(model).keys()
+
+        apply_activation_checkpointing(model)
+        model = FSDP(model, use_orig_params=use_orig_params)
+        new_keys = get_model_state_dict(model).keys()
+
+        self.assertEqual(original_keys, new_keys)
+
     @with_comms
     @skip_if_lt_x_gpu(1)
     def test_extra_state(self) -> None: