Skip to content

Commit be84494

Browse files
feginfegin
committed
[FSDP][optim_state_dict][10/N] Make optim_state_dict and optim_state_dict_to_load public
Pull Request resolved: #92118 Make optim_state_dict and optim_state_dict_to_load public APIs and consolidate them with state_dict by using the same state_dict_type to decide how to perform the optimizer state_dict save and load. ghstack-source-id: 178869730 Differential Revision: [D42488022](https://our.internmc.facebook.com/intern/diff/D42488022/)
1 parent 4170fbe commit be84494

File tree

7 files changed

+656
-208
lines changed

7 files changed

+656
-208
lines changed

‎test/distributed/_composable/fully_shard/test_fully_shard_optim_checkpoint.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,8 @@ def _test_optim_state_save_load(self, model1, optim1, model2, optim2) -> None:
4646
model(batch).sum().backward()
4747
optim.step()
4848

49-
optim_state_dict1 = FSDP._optim_state_dict(model1, optim1)
50-
optim_state_dict2 = FSDP._optim_state_dict(model2, optim2)
49+
optim_state_dict1 = FSDP.optim_state_dict(model1, optim1)
50+
optim_state_dict2 = FSDP.optim_state_dict(model2, optim2)
5151

5252
self.assertEqual(
5353
len(optim_state_dict1["state"]), len(optim_state_dict2["state"])

0 commit comments

Comments
 (0)