facebookresearch / facebookresearch/fairscale
[feature] Add support for torch.cpu.autocast - fp16
Open
- Dominant language
- Python
- Stars
- 3.4k
- Forks
- 293
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
PyTorch 1.10 provides torch.cpu.autocast support for bf16 right now. Investigate if and how we should support this for fp16 as well so that it works well for users out of the box. Use case for where this would be useful is when using cpu_offload with FSDP.
[Related PL PR](https://github.com/PyTorchLightning/pytorch-lightning/pull/9084)
cc @SeanNaren
Contributor guide
Assessment
This issue has not been assessed yet.