add all_gather process-group for overlapping in fsdp disributed training

[jeffnvidia]

What does this PR do ?

This PR intends to separate the all_gather process-group from reduce-scatter and the other operations. The goal is to have overlapping of these 2 collectives which when combined with SHARP have been proven to greatly increase performances (~15%)

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share discuss a design-doc with the team.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Loading

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

1. Attach the Expert Review label when your PR is ready for review.
2. GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

1. Add Final Review label
2. GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[jeffnvidia]

I am having this warning on the PR, any idea why ?

[MCORE][MultiGroupMemPoolAllocator] Failed to deregister mem pool from<torch.distributed.distributed_c10d.ProcessGroup object at 0x400228ee6170>(DATA_PARALLEL_GROUP_WITH_CP_AG) group!!

cursor said it's normal behavior but I'm suspicious

[shjwudp]

LGTM, thanks.

[shjwudp]

Can we have a more accurate description? For example:
"All-gather group used when overlapping all-gather and gradient reduction."

[jeffnvidia]

no problem changed to your description

[youngeunkwon0405]

Logic-wise LGTM. I left some comments regarding variable names and comments.

[youngeunkwon0405]

Could you please write a more specific comment that anyone without the context could understand what you are trying to do?

[jeffnvidia]

sure let me know if it's good now

[shengf-nv]

@jeffnvidia, thanks for creating this PR to add all_gather process group. I have two comments:

1. Only PG for regular data parallelism has been split, not for ET data parallelism. Do you plan to add it in another PR?

2. The PR decides which PG to use for all gather based on the return PG from self.buffer.dist_index.get_fsdp_group(is_expert_parallel=False, all_gather=True). What is your plan when ET data parallelism also exists? My suggestion is: instead of picking PG when all gather is called, we can set the appropriate data parallel process group for DataParallelBuffer of model_weight_buffer in function _init_each_parameter_group_buffers. Each parameter_group has a flag to indicate if it is a regular weight or a ET weight. We can use that flag to choose regular DP or ET DP. Then we do not need to make any change when all gather is being called.

Let me know if it makes sense to you.

Thanks
Sheng

[jeffnvidia]

Hey Sheng thanks for your comments

1. Only PG for regular data parallelism has been split, not for ET data parallelism. Do you plan to add it in another PR?

I can add it in a separate PR, I don't think it would require much work. The reason I didn't do it now was that we never tested it yet as far as I know. Once we move on to testing EP, I can create a new PR.

2. The PR decides which PG to use for all gather based on the return PG from self.buffer.dist_index.get_fsdp_group(is_expert_parallel=False, all_gather=True). What is your plan when ET data parallelism also exists? My suggestion is: instead of picking PG when all gather is called, we can set the appropriate data parallel process group for DataParallelBuffer of model_weight_buffer in function _init_each_parameter_group_buffers. Each parameter_group has a flag to indicate if it is a regular weight or a ET weight. We can use that flag to choose regular DP or ET DP. Then we do not need to make any change when all gather is being called.

I've refactored per your recommendation:
PG selection now happens in _init_each_parameter_group_buffers() based on parameter group flags
All-gather call sites simplified to use the pre-configured group
Much cleaner and easier to extend, thanks !

let me know if it fits what you had in mind

[youngeunkwon0405]

I just realized that if EP is enabled, there will be a different DP group for MoE layers. I think if you run this with EP, it will not show the expected behavior. I think you have to modify the code for this, or at least insert an assertion about this.

@shjwudp, please correct me if I am wrong.

[jeffnvidia]

I just realized that if EP is enabled, there will be a different DP group for MoE layers. I think if you run this with EP, it will not show the expected behavior. I think you have to modify the code for this, or at least insert an assertion about this.

@shjwudp, please correct me if I am wrong.

in the code, I am checking if expert parallelism is on and only doing the changes if not which is the scope of the POC

[shengf-nv]

@jeffnvidia, thanks for the update. The updated code looks good to me.

It is your call if you want to have another PR to support ET. However, more recent models all have ET. I already have Llama4 setup, and plan to try overlapped AG+RS with Llama4 soon.

[youngeunkwon0405]

I just realized that if EP is enabled, there will be a different DP group for MoE layers. I think if you run this with EP, it will not show the expected behavior. I think you have to modify the code for this, or at least insert an assertion about this.
@shjwudp, please correct me if I am wrong.

in the code, I am checking if expert parallelism is on and only doing the changes if not which is the scope of the POC

I don't see the code part that checks if expert_parallel and makes a decision. Can you point out which part of your code is doing that?

[jeffnvidia]

I just realized that if EP is enabled, there will be a different DP group for MoE layers. I think if you run this with EP, it will not show the expected behavior. I think you have to modify the code for this, or at least insert an assertion about this.
@shjwudp, please correct me if I am wrong.

in the code, I am checking if expert parallelism is on and only doing the changes if not which is the scope of the POC

I don't see the code part that checks if expert_parallel and makes a decision. Can you point out which part of your code is doing that?

It's line 1890 : https://github.com/NVIDIA/Megatron-LM/pull/2663/files#diff-da62f73a7a6a4ac7815ed316a147ba348a7915e35a2f4885ceaf1678e5f650fbR1890

if it is expert_parallelism, it wont even try to create the separate group

[jeffnvidia]

I rebased on main to fix the merging error. @youngeunkwon0405, could you re-run the CI ? Thanks a lot

Could you approve the PR @jaredcasper as we said and I'll start working immediately on a PR that puts everything in order. Thanks a lot

[youngeunkwon0405]

/ok to test caf8de5

[jeffnvidia]

/ok to test caf8de5

thanks Youngeun, we're having the same bug again of the CICD which seems to be random, I don't why and if it's a blocker or not

[youngeunkwon0405]

Got it, I will reinitiate the failed runs.

[jeffnvidia]

Hi guys, who do I need to get the approval from to finally merge this PR ?

@ericharper @jaredcasper @NVIDIA/core-adlr ?

[cspades]

Needs an approval from @NVIDIA/core-adlr or else it won't merge. It looks like concerns pertaining to MCore PG management have not been resolved.

[cspades]

M-FSDP is using pg_collection(mcore_fsdp_adapter.py#L193). However, as I understand it, the current MCore pretraining still initializes through the global parallel state (training.py#L1156), right?

And just to chime in here, FSDPDistributedIndex is a one-stop shop to hack in PG groups by argument, so really there shouldn't be any blocker for a fully-argument-based PG management system. And add the same args to fully_shard and you're good to go.