NVIDIA / NVIDIA/Megatron-LM

[QUESTION] Can we avoid creation of separate master weights in BF16?

Open
#1,821 0 comments 0 reactions 0 assignees View on GitHub
community-request module: transformer engine question
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.