re-introduce init_optim_state (#946)

Summary: Pull Request resolved: #946

Reviewed By: galrotem

Differential Revision: D65922636

fbshipit-source-id: 0965aa03ce000d1c1d544635238e3465f0a5dc5e

Author: JKSenthil

torchtnt/framework/_test_utils.py

@@ -146,7 +146,9 @@ def __init__(self, input_dim: int) -> None:
         ]
         self.applied_optims: List[torch.optim.Optimizer] = []
         for module, optim in zip(self.modules, self.optims):
-            self.applied_optims.append(optim(module.parameters(), lr=0.1))
+            o = optim(module.parameters(), lr=0.1)
+            self.applied_optims.append(o)
+            setattr(self, f"optimizer_{optim.__name__}", o)
 
     def train_step(
         self, state: State, data: Batch

torchtnt/framework/callbacks/dcp_saver.py

@@ -11,6 +11,8 @@
 from concurrent.futures import Future
 from typing import Any, cast, Dict, Iterable, List, Optional, Union
 
+import torch
+
 import torch.distributed as dist
 from pyre_extensions import none_throws
 from torch.distributed import checkpoint as dcp
@@ -24,6 +26,16 @@
     DefaultSavePlanner,
 )
 from torch.distributed.checkpoint.planner import LoadPlanner, SavePlanner
+
+try:
+    from torch.distributed.checkpoint.state_dict import _init_optim_state
+except ImportError:
+
+    def noop(_: Any) -> None:
+        return None
+
+    _init_optim_state = noop
+
 from torch.distributed.checkpoint.storage import StorageReader, StorageWriter
 
 from torchtnt.framework.callbacks._checkpoint_utils import (
@@ -351,6 +363,15 @@ def restore_with_id(
             predict_dataloader,
         )
 
+        # necessary for loading optimizers since states are initialized lazy
+        for obj in app_state.values():
+            # sometimes optimizers are actually held in a wrapper which handles calling
+            # state_dict and load_state_dict, sa is the case for
+            # `torchtnt.utils.prepare_module.FSDPOptimizerWrapper`, this handles that case.
+            optimizer = getattr(obj, "optimizer", obj)
+            if isinstance(optimizer, torch.optim.Optimizer):
+                _init_optim_state(optimizer)
+
         with get_or_create_gloo_pg(candidate_pg=process_group) as pg:
             dcp.load(
                 {"app_state": MultiStateful(app_state)},