[Cherry-pick][DCP][AC] Add test for apply AC with FSDP1 (#126935) (#126992)

wz337 · web-flow · commit cd033a128cf5 · 2024-05-24T10:35:55.000-04:00
[DCP][AC] Add test for apply AC with FSDP1 (#126935) Adding test for this cherry pick. #126559 Pull Request resolved: #126935 Approved by: https://github.com/fegin
diff --git a/test/distributed/checkpoint/test_state_dict.py b/test/distributed/checkpoint/test_state_dict.py
@@ -506,7 +506,7 @@ def is_cpu(v):
 
     @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()
@@ -517,6 +517,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)
+
 
 if __name__ == "__main__":
     run_tests()
