facebookresearch / facebookresearch/fairscale
[feat] OSS: Support seamless state consolidation
- Dominant language
- Python
- Stars
- 3.4k
- Forks
- 293
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
Currently `state_dict()` being called on a single rank cannot return the full optimizer state, because this rank only has a shard's worth. Checkpointing the full optimizer state requires calling `consolidate_state()` on all rank prior to pulling the state, which is not intuitive and not pytorch compliant.
Taking inspiration from other projects sharding the model or parameter state, we could use the torch RPC framework to store remote references to other ranks' states, which would allow for a seamless `to_here()` consolidation. In that case, we would need to suppose that both RPC and torch distributed are initialized, which is a feature being tackled.
## Motivation
Pytorch compliance
Smaller API footprint
Ease of use
## Pitch
See the feature description, this is doable and would be a lot more elegant.
## Alternatives
Status quo
## Additional context
Discussed with Pytorch Distributed and MSFT folks
Contributor guide
Assessment
This issue has not been assessed yet.