resolve CPU OOM with FSDP checkpointer

Summary: Pull Request resolved: https://github.com/facebookresearch/d2go/pull/632

Reviewed By: yzhao30

Differential Revision: D50663689

fbshipit-source-id: 5c4c1dd2e5d2087be5aec268672bb5e7fc329df9

d2go/checkpoint/fsdp_checkpoint.py

@@ -220,12 +220,15 @@ class FSDPCheckpointer(QATCheckpointer):
         self.logger.info("Finished saving checkpoint to {}".format(filename))
 
     def _load_file(self, f: str):
+        if isinstance(self.model, FSDPWrapper):
             with (
                 interleave_by_rank(concurrency_limit=self._concurrency_limit_fetcher())
-            if isinstance(self.model, FSDPWrapper)
-            and self.model.state_dict_type != StateDictType.FULL_STATE_DICT
+                if self.model.state_dict_type != StateDictType.FULL_STATE_DICT
                 else nullcontext()  # FULL_STATE_DICT doesn't need interleaving
             ):
+                # use mmap for FSDP checkpoints
+                return torch.load(f, map_location=torch.device("cpu"), mmap=True)
+        else:
             return super()._load_file(f)
 
     def _save_metadata(self, path):