Improve communication overlapping in FP8 distributed optimizer

[timmoon10]

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:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• 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:

• 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 contains specific people who can review PRs to various areas.

Additional Information

• This PR was implemented together with Support scaled optimizer state in distributed Adam optimizer apex#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 Use distributed optimizer support for multiple dtypes #7359
• Distributed optimizer support for FP8 models was added in Add distopt support for FP8 params and BF16 optimizer state #7909

[timmoon10]

jenkins

[timmoon10]

jenkins

[timmoon10]

jenkins

[timmoon10]

jenkins

[dimapihtar]

LGTM. Thank you!