Support multimodule communication - #2031

Merged

yaoyu-33 merged 23 commits into

NVIDIA:mainfrom

yaoyu-33:shifang/multi_module_comm

Jan 24, 2026

Commits on Oct 5, 2025

Support multimodule communication
shifangx
committed
5aec371

Commits on Oct 6, 2025

add test
shifangx
committed
01f2b4f
add test for MultiModulePipelineCommunicator init
shifangx
committed
3425fac
add test for compute_total_pipeline_stages
shifangx
committed
cc2b6ee
add test_send_forward_recv_forward
shifangx
committed
e940fcb
add test_send_backward_recv_backward
shifangx
committed
b689f16
add test_send_forward_recv_backward_send_backward_recv_forward
shifangx
committed
9565a98
skip if torch version is low
shifangx
committed
381986a
delete duplicated dist.broadcast in BridgeCommunicator
shifangx
committed
f064f05

Commits on Oct 7, 2025

add test_send_forward_recv_forward_with_transformer_blocks
shifangx
committed
11ae642
add test_send_forward_recv_forward_with_transformer_blocks_and_different_parallelisms
shifangx
committed
698e454
chore: Format files
Mcore Bot
committed
babbb54

Commits on Oct 14, 2025

update comment and logging
shifangx
committed
2ac3216
rename from pp_stage to pp_rank, and import some unit function from test_bridge_communicator
shifangx
committed
e9d03ba
add test for the case when pp(image_encoder) != pp(audio_encoder)
shifangx
committed
7c10cb0
chore: Format files
Mcore Bot
committed
b28287f

Commits on Oct 22, 2025

add test with llm pp size larger than 2, such as llm pp=4
shifangx
committed
fbeb6e9

Commits on Nov 13, 2025

move comments to docstrings
yaoyu-33
committed
3b52ec4

Commits on Dec 10, 2025

Merge branch 'main' into shifang/multi_module_comm
yaoyu-33
authored
aa19652

Commits on Jan 6, 2026

add copyright
shifangx
committed
10bf0c4
Merge branch 'main' into shifang/multi_module_comm
shifangx
authored
19f0ecd

Commits on Jan 16, 2026

Merge branch 'main' into shifang/multi_module_comm
yaoyu-33
authored
2fe2eeb

Commits on Jan 24, 2026

Merge branch 'main' into shifang/multi_module_comm
shifangx
authored
762ec65