Improve communication overlapping in FP8 distributed optimizer
jenkins jenkins jenkins jenkins
# What does this PR do ?

When training GPT, the Apex distributed Adam optimizer overlaps its first parameter all-gather with the optimizer step. This optimization has been applied to both FP8 and non-FP8 models.

**Collection**: NLP

# Changelog 
- Wait until all distopt buckets finish optimizer step before updating FP8 scaling factors
- Support dist opt buckets with multiple dtype configs
- Put GPT "leftover params" (i.e. embeddings, layer norm params, biases) in same distopt bucket as first layer params

# Usage

Run GPT, e.g. with the config at https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/language_modeling/conf/megatron_gpt_config.yaml.

Enable FP8 support with `model.fp8=True`, FP8 parameters with `model.fp8_params=True`, the distributed optimizer with `model.optim.name=distributed_fused_adam`, and overlapped param all-gathers with `model.optim.overlap_param_sync=True`

# Jenkins CI
To run Jenkins, a NeMo User with write access must comment `jenkins` on the PR.

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [x] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [x] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* This PR was implemented together with https://github.com/NVIDIA/apex/pull/1771. Each of the two PRs should be able to run without the other, but there may be performance degradations when training GPT with FP8 params. With the Apex PR and without the NeMo PR, there will be redundant amax reductions. Without the Apex PR and with the NeMo PR, the first bucket all-gather will not be overlapped with the optimizer step.
* This communication optimization used to be present for non-FP8 models, but was removed in https://github.com/NVIDIA/NeMo/pull/7359
* Distributed optimizer support for FP8 models was added in https://github.com/NVIDIA/NeMo/pull/7909