[QUESTION] Can we avoid creation of separate master weights in BF16?
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Background
As far as I can tell, the Megatron optimizer wrappers have a strict requirement that master weights (and a `master_param`) attribute are present. The torch optimizers do not appear to do this, and TE FusedAdam has a `master_weights` param which can be set to `False`, although this is not passed through to the constructor in Megatron-LM.
I recognise the benefit to precision and convergence of high precision masters, but the power to choose as a user and test this would be appreciated, and currently it's quite a costly difference between Megatron-LM and other trainers. (Differences may also be reduced by methods like [stochastic rounding](https://arxiv.org/pdf/2010.06192) but this is not my area of expertise).
## Questions:
1. Is there any way to disable using a separate master parameters to save on memory usage when doing lower-than-fp32 precision training?
2. If the answer to (1) is no, can this be added as a toggle? (Can make a feature request)
3. In its current state, is the most memory-efficient way to perform bf16 training to use the precision-aware optimizer, and to set `store_param_remainders` so that active and master params consume a total 4 bytes, or is there a more efficient setup?
Contributor guide
Assessment
This issue has not been assessed yet.