Add BERT support for overlapping forward compute with distopt communication

[timmoon10]

What does this PR do ?

This PR adds experimental support in BERT for the distributed Adam optimizer to overlap its parameter all-gather with the forward compute.

Collection: NLP

Changelog

• Add BERT support for overlapping forward compute with distopt communication

Usage

Set the optimizer to distributed_fused_adam in the config file:

NeMo/examples/nlp/language_modeling/conf/megatron_bert_config.yaml

Line 144 in a7c2a04

name: fused_adam

Configure the optimizer with overlap_param_sync: True.

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.

Pinging @shanmugamr1992.

Additional Information

• Follow-up to Distributed Adam optimizer overlaps param all-gather with forward compute #5684.

[ericharper]

LGTM. Thanks!