[Reland] Reinitialize ModelWrapper.cache_state_dict because set_model_state_dict has side effects (#977)

Reland #971 due to ghstack
issues.

Author: fegin

torchtitan/components/checkpoint.py

@@ -65,6 +65,11 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
             options=StateDictOptions(strict=False),
         )
         list(map(func, self.model))
+        # `set_model_state_dict()` does change the keys of the input state_dict,
+        # we will need to reinitialize the cache_state_dict.
+        self.cache_state_dict = {
+            k: v for sd in map(get_model_state_dict, self.model) for k, v in sd.items()
+        }
 
 
 class Terminate: