Consolidate the naming of named_parameter and state_dict for CheckpointWrapper

test/distributed/fsdp/test_checkpoint_wrapper.py

@@ -176,6 +176,12 @@ def check_fn(l):
                     self.assertTrue(param.requires_grad)
                     self.assertFalse(param.grad is None)
 
+    def test_fqn(self):
+        lin = nn.Linear(10, 10, bias=False)
+        lin = checkpoint_wrapper(lin)
+        state_dict = lin.state_dict()
+        for fqn, _ in lin.named_parameters():
+            self.assertTrue(fqn in state_dict, msg=f"{fqn} not in state_dict.")
 
 if __name__ == "__main__":
     run_tests()

torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py

@@ -7,7 +7,7 @@
 from torch.distributed.utils import _replace_by_prefix
 from torch.distributed.fsdp.wrap import _recursive_wrap, lambda_auto_wrap_policy
 import torch.nn as nn
-from typing import Dict, Any
+from typing import Any, Dict, Iterator, Tuple
 from functools import partial
 
 _CHECKPOINT_PREFIX = "_checkpoint_wrapped_module"
@@ -61,6 +61,18 @@ def forward(self, *args, **kwargs):
                 **kwargs,
             )
 
+    def named_parameters(
+        self,
+        *args,
+        **kwargs,
+    ) -> Iterator[Tuple[str, torch.nn.Parameter]]:
+        """
+        Overrides :meth:`named_parameters()` to intercept parameter names and
+        remove all occurrences of _CHECKPOINT_PREFIX.
+        """
+        for param_name, param in super().named_parameters(*args, **kwargs):
+            yield param_name.replace(f"{_CHECKPOINT_PREFIX}.", ""), param
+
     @staticmethod
     def _post_state_dict_hook(
         module: nn.Module,