Skip to content

Support multimodule communication - #2031

Merged
yaoyu-33 merged 23 commits into
NVIDIA:mainfrom
yaoyu-33:shifang/multi_module_comm
Jan 24, 2026
Merged

Support multimodule communication#2031
yaoyu-33 merged 23 commits into
NVIDIA:mainfrom
yaoyu-33:shifang/multi_module_comm

Conversation

@yaoyu-33

@yaoyu-33 yaoyu-33 commented Oct 29, 2025

Copy link
Copy Markdown
Contributor

This is a ported gitlab PR 4069. Expert review is done, needs final review.

shifangx and others added 17 commits October 5, 2025 16:59
@shifangx
Support multimodule communication
5aec371
@shifangx
add test
01f2b4f
@shifangx
add test for MultiModulePipelineCommunicator init
3425fac
@shifangx
add test for compute_total_pipeline_stages
cc2b6ee
@shifangx
add test_send_forward_recv_forward
e940fcb
@shifangx
add test_send_backward_recv_backward
b689f16
@shifangx
add test_send_forward_recv_backward_send_backward_recv_forward
9565a98
@shifangx
skip if torch version is low
381986a
@shifangx
delete duplicated dist.broadcast in BridgeCommunicator
f064f05
@shifangx
add test_send_forward_recv_forward_with_transformer_blocks
11ae642
@shifangx
add test_send_forward_recv_forward_with_transformer_blocks_and_differ…
698e454 
…ent_parallelisms
Mcore Bot
chore: Format files
babbb54
@shifangx
update comment and logging
2ac3216
@shifangx
rename from pp_stage to pp_rank, and import some unit function from t…
e9d03ba 
…est_bridge_communicator
@shifangx
add test for the case when pp(image_encoder) != pp(audio_encoder)
7c10cb0
Mcore Bot
chore: Format files
b28287f
@shifangx
add test with llm pp size larger than 2, such as llm pp=4
fbeb6e9
@yaoyu-33
yaoyu-33 requested review from a team as code owners October 29, 2025 18:27
@copy-pr-bot

copy-pr-bot Bot commented Oct 29, 2025

Copy link
Copy Markdown

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.

jaredcasper
jaredcasper reviewed Oct 29, 2025
View reviewed changes
Comment thread megatron/core/pipeline_parallel/multimodule_communicator.py Outdated
@yashaswikarnati yashaswikarnati mentioned this pull request Nov 13, 2025
Closed
6 tasks
@yaoyu-33
yaoyu-33 requested review from a team as code owners November 13, 2025 20:23
@yanring yanring added the Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. label Nov 19, 2025
@yanring yanring added this to the Core 0.16 milestone Nov 19, 2025
@yaoyu-33

yaoyu-33 commented Jan 16, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 2fe2eeb

@shifangx

shifangx commented Jan 24, 2026

Copy link
Copy Markdown
Contributor

/ok to test 762ec65

@yaoyu-33
yaoyu-33 added this pull request to the merge queue Jan 24, 2026
Merged via the queue into NVIDIA:main with commit 485ed18 Jan 24, 2026
52 of 54 checks passed
@yaoyu-33
yaoyu-33 deleted the shifang/multi_module_comm branch January 24, 2026 04:05
ko3n1g added a commit to ko3n1g/Megatron-LM that referenced this pull request Jan 24, 2026
@ko3n1g
Revert "Support multimodule communication (NVIDIA#2031)"
913178e 
This reverts commit 485ed18.
ko3n1g added a commit that referenced this pull request Jan 24, 2026
@ko3n1g
Revert "Support multimodule communication (#2031)" (#3068)
0972f02
ko3n1g added a commit that referenced this pull request Jan 27, 2026
@ko3n1g
Reapply "Support multimodule communication (#2031)" (#3068)
35e85a6 
This reverts commit 0972f02.
daiyaanarfeen pushed a commit to daiyaanarfeen/Megatron-LM that referenced this pull request Feb 23, 2026
@yaoyu-33 @shifangx
Mcore Bot @daiyaanarfeen
Support multimodule communication (NVIDIA#2031)
f91932c 
Co-authored-by: shifangx <shifangx@nvidia.com>
Co-authored-by: Mcore Bot <mcore-bot@nvidia.com>
daiyaanarfeen pushed a commit to daiyaanarfeen/Megatron-LM that referenced this pull request Feb 23, 2026
@ko3n1g @daiyaanarfeen
Revert "Support multimodule communication (NVIDIA#2031)" (NVIDIA#3068)
558fe76
daiyaanarfeen pushed a commit to daiyaanarfeen/Megatron-LM that referenced this pull request Feb 23, 2026
@ko3n1g @daiyaanarfeen
Reapply "Support multimodule communication (NVIDIA#2031)" (NVIDIA#3068)
848dc09 
This reverts commit 0972f02.
@shifangx shifangx mentioned this pull request Apr 15, 2026
Closed
6 tasks
yangbofun pushed a commit to xlm-research/Megatron-LM that referenced this pull request May 22, 2026
@yaoyu-33 @shifangx Mcore Bot
Support multimodule communication (NVIDIA#2031)
1bbc530 
Co-authored-by: shifangx <shifangx@nvidia.com>
Co-authored-by: Mcore Bot <mcore-bot@nvidia.com>
yangbofun pushed a commit to xlm-research/Megatron-LM that referenced this pull request May 22, 2026
@ko3n1g
Revert "Support multimodule communication (NVIDIA#2031)" (NVIDIA#3068)
7ee7b71
yangbofun pushed a commit to xlm-research/Megatron-LM that referenced this pull request May 22, 2026
@ko3n1g
Reapply "Support multimodule communication (NVIDIA#2031)" (NVIDIA#3068)
bbad49a 
This reverts commit 7ee7b71.
terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026
@yaoyu-33 @shifangx
Mcore Bot @terminator123
Support multimodule communication (NVIDIA#2031)
a9f42dd 
Co-authored-by: shifangx <shifangx@nvidia.com>
Co-authored-by: Mcore Bot <mcore-bot@nvidia.com>
terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026
@ko3n1g @terminator123
Revert "Support multimodule communication (NVIDIA#2031)" (NVIDIA#3068)
f6a2e25
terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026
@ko3n1g @terminator123
Reapply "Support multimodule communication (NVIDIA#2031)" (NVIDIA#3068)
4697f7b 
This reverts commit 0972f02.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@jaredcasper jaredcasper jaredcasper approved these changes
@ericharper ericharper ericharper approved these changes
@ko3n1g ko3n1g ko3n1g approved these changes
@yashaswikarnati yashaswikarnati yashaswikarnati approved these changes

Assignees

No one assigned

Labels

complexity: low Run tests

Projects

None yet

Milestone

Core 0.16

Development

Successfully merging this pull request may close these issues.

7 participants

@yaoyu-33 @shifangx @ko3n1g @jaredcasper @ericharper @yashaswikarnati @yanring