support save&load of fsdp_optim_state #16
Conversation
| We judge whether the optim_state_dict is sharded automatically | ||
|
|
||
| Args: | ||
| optim_state_dict (Dict[str, Any]): The optimizer states to be loaded. |
There was a problem hiding this comment.
Please add some comments to explain what optim_state_dict are passed by other ranks when rank0_only is True.
There was a problem hiding this comment.
There is no explanation here about how the optim_state_dict parameters for other ranks are passed after the optimizer state is loaded in rank 0. We can let users pass None instead of an empty dictionary (since users may not know what the keys are).
There was a problem hiding this comment.
I add the explanation that we will broadcast rank0's optim_state_dict info to other ranks if specified rank0_only, so it's no matter what other rank's pass in and of course other ranks can pass None.
| We judge whether the optim_state_dict is sharded automatically | ||
|
|
||
| Args: | ||
| optim_state_dict (Dict[str, Any]): The optimizer states to be loaded. |
torchacc/dist/fsdp.py
Outdated
| self.model which is sharded. | ||
| """ | ||
| # for sharded optim_state, we return directly | ||
| if 'shard_metadata' in optim_state_dict.keys(): |
There was a problem hiding this comment.
I check the world_size between the stored shard_metadata and current shard_metadata and raise a NotImplementedError(is it suitable?) error.
|
Overall looks good to me. Pass off to @anw90 for final review. |
hanwen-sun
changed the title
| return model, optim | ||
|
|
||
|
|
||
| def _train_step( |
There was a problem hiding this comment.
it’s better to replace _train_step with _train because _train_step might suggest that it only includes a single training step.
| labels = torch.zeros(batch_size, dtype=torch.int64).to(device) | ||
| loss = model(data) | ||
| loss = torch.nn.functional.nll_loss(loss, labels) | ||
| loss.backward() |
There was a problem hiding this comment.
do we need to call loss.backward here?
There was a problem hiding this comment.
No need, but there's no difference whether we do forward only or do forward and backward in the test case.
What this pr do:
TODO: