support fsdp2 optimizer checkpointing (#977)

Summary: Pull Request resolved: #977

Reviewed By: diego-urgell

Differential Revision: D70337970

fbshipit-source-id: 52f0915f01dba18ee420c9b38282db9f6eb926df

Author: JKSenthil

tests/framework/test_app_state_mixin.py

@@ -19,7 +19,7 @@
 
 from torchtnt.utils.env import init_from_env
 from torchtnt.utils.lr_scheduler import TLRScheduler
-from torchtnt.utils.prepare_module import FSDPOptimizerWrapper
+from torchtnt.utils.prepare_module import FSDP2OptimizerWrapper, FSDPOptimizerWrapper
 from torchtnt.utils.stateful import MultiStateful
 
 
@@ -269,3 +269,15 @@ def test_construct_tracked_optimizers_and_schedulers(self) -> None:
         self.assertIsInstance(result["optimizer"], FSDPOptimizerWrapper)
         self.assertIsInstance(result["optim2"], torch.optim.Optimizer)
         self.assertIsInstance(result["lr_scheduler"], TLRScheduler)
+
+        with patch(
+            "torchtnt.framework.unit._is_fsdp_module", side_effect=lambda m: m == module
+        ), patch(
+            "torchtnt.framework.unit._is_fsdp2_module",
+            side_effect=lambda m: m == module,
+        ):
+            result = auto_unit._construct_tracked_optimizers_and_schedulers()
+
+        self.assertIsInstance(result["optimizer"], FSDP2OptimizerWrapper)
+        self.assertIsInstance(result["optim2"], torch.optim.Optimizer)
+        self.assertIsInstance(result["lr_scheduler"], TLRScheduler)

torchtnt/framework/unit.py

@@ -20,7 +20,12 @@
 
 from torchtnt.framework.state import State
 from torchtnt.utils.lr_scheduler import TLRScheduler
-from torchtnt.utils.prepare_module import _is_fsdp_module, FSDPOptimizerWrapper
+from torchtnt.utils.prepare_module import (
+    _is_fsdp2_module,
+    _is_fsdp_module,
+    FSDP2OptimizerWrapper,
+    FSDPOptimizerWrapper,
+)
 from torchtnt.utils.progress import Progress
 from torchtnt.utils.stateful import MetricStateful, Stateful
 
@@ -199,13 +204,27 @@ def __delattr__(self, name: str) -> None:
 
     def _construct_tracked_optimizers_and_schedulers(
         self,
-    ) -> Dict[str, Union[torch.optim.Optimizer, FSDPOptimizerWrapper, TLRScheduler]]:
+    ) -> Dict[
+        str,
+        Union[
+            torch.optim.Optimizer,
+            FSDPOptimizerWrapper,
+            FSDP2OptimizerWrapper,
+            TLRScheduler,
+        ],
+    ]:
         """
-        Combines tracked optimizers and schedulers. Handles optimizers working on FSDP modules, wrapping them in FSDPOptimizerWrapper.
+        Combines tracked optimizers and schedulers. Handles optimizers working on FSDP modules, wrapping them in FSDPOptimizerWrapper/FSDP2OptimizerWrapper.
         """
         # construct custom tracked optimizers with FSDP optimizers
         tracked_optimizers_and_schedulers: Dict[
-            str, Union[torch.optim.Optimizer, FSDPOptimizerWrapper, TLRScheduler]
+            str,
+            Union[
+                torch.optim.Optimizer,
+                FSDPOptimizerWrapper,
+                FSDP2OptimizerWrapper,
+                TLRScheduler,
+            ],
         ] = {}
         tracked_optimizers_and_schedulers.update(self._construct_tracked_optimizers())
 
@@ -224,25 +243,38 @@ def _construct_tracked_optimizers_and_schedulers(
 
     def _construct_tracked_optimizers(
         self,
-    ) -> Dict[str, Union[torch.optim.Optimizer, FSDPOptimizerWrapper]]:
+    ) -> Dict[
+        str, Union[torch.optim.Optimizer, FSDPOptimizerWrapper, FSDP2OptimizerWrapper]
+    ]:
         """
-        Constructs tracked optimizers. Handles optimizers working on FSDP modules, wrapping them in FSDPOptimizerWrapper.
+        Constructs tracked optimizers. Handles optimizers working on FSDP modules, wrapping them in FSDPOptimizerWrapper/FSDP2OptimizerWrapper.
         """
-        fsdp_tracked_optimizers: Dict[str, FSDPOptimizerWrapper] = {}
+        fsdp_tracked_optimizers: Dict[
+            str, Union[FSDPOptimizerWrapper, FSDP2OptimizerWrapper]
+        ] = {}
         for module in self.tracked_modules().values():
             if _is_fsdp_module(module):
                 # find optimizers for module, if exists
                 optimizer_list = _find_optimizers_for_module(
                     module, self.tracked_optimizers()
                 )
+
+                is_fsdp2 = _is_fsdp2_module(module)
+
                 for optim_name, optimizer in optimizer_list:
-                    fsdp_tracked_optimizers[optim_name] = FSDPOptimizerWrapper(
-                        module, optimizer
-                    )
+                    if is_fsdp2:
+                        fsdp_tracked_optimizers[optim_name] = FSDP2OptimizerWrapper(
+                            module, optimizer
+                        )
+                    else:
+                        fsdp_tracked_optimizers[optim_name] = FSDPOptimizerWrapper(
+                            module, optimizer
+                        )
 
         # construct custom tracked optimizers with FSDP optimizers
         tracked_optimizers: Dict[
-            str, Union[torch.optim.Optimizer, FSDPOptimizerWrapper]
+            str,
+            Union[torch.optim.Optimizer, FSDPOptimizerWrapper, FSDP2OptimizerWrapper],
         ] = {
             key: value
             for key, value in self.tracked_optimizers().items()

torchtnt/utils/prepare_module.py

@@ -33,6 +33,10 @@
     checkpoint_wrapper,
     CheckpointImpl,
 )
+from torch.distributed.checkpoint.state_dict import (
+    get_optimizer_state_dict,
+    set_optimizer_state_dict,
+)
 from torch.distributed.device_mesh import init_device_mesh
 
 try:
@@ -449,6 +453,24 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         self.optimizer.load_state_dict(optim_state_dict)
 
 
+class FSDP2OptimizerWrapper:
+    """
+    Wrapper for FSDP2 optimizer which uses distributed state dict APIs.
+    """
+
+    def __init__(
+        self, module: torch.nn.Module, optimizer: torch.optim.Optimizer
+    ) -> None:
+        self.module = module
+        self.optimizer = optimizer
+
+    def state_dict(self) -> Dict[str, Any]:
+        return get_optimizer_state_dict(self.module, self.optimizer)
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        set_optimizer_state_dict(self.module, self.optimizer, state_dict)
+
+
 def _is_fsdp_module(module: torch.nn.Module) -> bool:
     if isinstance(module, FSDP):
         return True
@@ -461,6 +483,14 @@ def _is_fsdp_module(module: torch.nn.Module) -> bool:
     return False
 
 
+def _is_fsdp2_module(module: torch.nn.Module) -> bool:
+    maybe_composable_state = _get_module_state(module)
+    if maybe_composable_state is not None:
+        return isinstance(maybe_composable_state, FSDPState)
+
+    return False
+
+
 def prepare_module(
     module: torch.nn.Module,
     device: torch.device,

torchtnt/utils/stateful.py

@@ -10,7 +10,7 @@
 
 import torch
 from torchtnt.utils.lr_scheduler import TLRScheduler
-from torchtnt.utils.prepare_module import FSDPOptimizerWrapper
+from torchtnt.utils.prepare_module import FSDP2OptimizerWrapper, FSDPOptimizerWrapper
 from torchtnt.utils.progress import Progress
 
 from typing_extensions import Protocol, runtime_checkable
@@ -28,7 +28,10 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None: ...
 StatefulDict = Dict[str, Stateful]
 ModuleDict = Dict[str, torch.nn.Module]
 OptimizerAndLRSchedulerDict = Dict[
-    str, Union[TLRScheduler, torch.optim.Optimizer, FSDPOptimizerWrapper]
+    str,
+    Union[
+        TLRScheduler, torch.optim.Optimizer, FSDPOptimizerWrapper, FSDP2OptimizerWrapper
+    ],
 ]
 ProgressDict = Dict[str, Progress]