[FEATURE]: BF16 support for LLM

[binmakeswell]

Describe the feature

Using BF16 Optimizer rather than FP16 is necessary for LLM training, which has been verified by BLOOM, OPT, and Megatron-Turing.

[kurisusnowdeng]

BF16 is much easier to use than FP16 since its range is the same as FP32, so that usually we only need to use autocast to wrap computation. Example usage of BF16 via torch AMP could be found in my experimental scripts (train_fsdp.py)

[binmakeswell]

Hi @kurisusnowdeng Thanks for your insight.

But what we need is to provide this function to users through colossalai rather than a personal experimental repo, which is actually invisible and unavailable for users.

Can you add it officially via PRs? Thanks.

[kurisusnowdeng]

@binmakeswell Our mixed precision API is still under refactoring. I would give some suggestions based on my experience to the developers.