Add create_all_gather_group configuration option

[jeffnvidia]

Important

The Update branch button must only be pressed in very rare occassions.
An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.

PR related

This PR is directly related to this megatron-lm PR : NVIDIA/Megatron-LM#2663

What does this PR do ?

Adds create_all_gather_group configuration option to enable overlapping reduce-scatter and all-gather.

Collection: Core / Distributed Training

Changelog

• Add create_all_gather_group parameter to ParallelismConfig class in megatron_strategy.py
• Add create_all_gather_group parameter to MegatronStrategy class initialization and setup
• Add create_all_gather_group property and setter to AppState class in app_state.py
• Propagate create_all_gather_group configuration through init_parallel_ranks() and initialize_model_parallel_for_nemo() functions
• Add documentation explaining the purpose of the new parameter (enables separate process group for all-gather operations)

Usage

Users can now enable overlapping of reduce-scatter and all-gather operations by setting the configuration:

from nemo.lightning.pytorch.strategies import MegatronStrategy, ParallelismConfig

Configure parallelism with all-gather group creation

parallel_config = ParallelismConfig(
tensor_model_parallel_size=2,
pipeline_model_parallel_size=2,
create_all_gather_group=True, # Enable separate all-gather process group
)

strategy = MegatronStrategy(
parallel_config=parallel_config,
create_all_gather_group=True,
)# GitHub Actions CI

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

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, P

[jeffnvidia]

The formatting check seems to be failing due to a Git history issue after force-pushing to add sign-offs, not actual formatting problems. I've verified locally that both isort and black checks pass. The workflows are awaiting maintainer approval to run properly.

[jeffnvidia]

Hi @nithinraok, could you review this PR or tell me who I could refer to ? Thanks !

[jeffnvidia]

@ericharper @yaoyu-33

Could you please add the "Run CICD" label to trigger the CI tests when you're ready to review?

Note: The reformat_with_isort_and_black check is currently failing due to a GitHub Actions issue (it can't find an old commit SHA after I rebased to add DCO sign-offs). I've verified locally that all formatting checks pass with python setup.py style.

Thanks!

[jeffnvidia]

@ericharper I changed the manifest, I think it needs to be re-run now (there is an error because of a commit change due to a force push)

[jeffnvidia]

Hey @ericharper @ko3n1g ,

I reverted the changes to the manifest and I added backward compatibility, let me know if I need to do more things now