diff --git a/megatron/core/pipeline_parallel/bridge_communicator.py b/megatron/core/pipeline_parallel/bridge_communicator.py
index a67ded6bf08..f1e74a2f16d 100644
--- a/megatron/core/pipeline_parallel/bridge_communicator.py
+++ b/megatron/core/pipeline_parallel/bridge_communicator.py
@@ -628,9 +628,6 @@ def send_forward_recv_backward(
                 dist.broadcast(
                     shape_tensor, src=self.current_rank, group=self.src_grid_broadcast_pg
                 )
-                dist.broadcast(
-                    shape_tensor, src=self.current_rank, group=self.src_grid_broadcast_pg
-                )
 
                 # Broadcast the tensors to all ranks in the group
                 dist.broadcast(
diff --git a/megatron/core/pipeline_parallel/multimodule_communicator.py b/megatron/core/pipeline_parallel/multimodule_communicator.py
new file mode 100644
index 00000000000..1e8da3468e2
--- /dev/null
+++ b/megatron/core/pipeline_parallel/multimodule_communicator.py
@@ -0,0 +1,531 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.distributed as dist
+
+from megatron.core.hyper_comm_grid import HyperCommGrid
+from megatron.core.model_parallel_config import ModelParallelConfig
+from megatron.core.pipeline_parallel.bridge_communicator import BridgeCommunicator
+from megatron.core.pipeline_parallel.p2p_communication import P2PCommunicator
+
+# Types
+Shape = Union[List[int], torch.Size]
+
+
+@dataclass
+class RankModuleInfo:
+    """Information about a rank in a module.
+
+    Attributes:
+        pp_rank: The stage index of the current rank within the module's pipeline.
+        pp_size: The total number of pipeline stages (ranks) in the module.
+        p2p_communicator: Intra-module point-to-point communicator.
+        bridge_comms_as_src_module: Bridge communicators for outgoing connections
+            from this module to downstream modules. One module may have multiple
+            bridge communicators if it has multiple outgoing connections.
+        bridge_comms_as_dest_module: Bridge communicators for incoming connections
+            to this module from upstream modules. One module may have multiple
+            bridge communicators if it has multiple incoming connections.
+        is_source_stage: True if this rank is at the absolute first stage in the
+            overall model (no incoming connections).
+        is_terminal_stage: True if this rank is at the absolute last stage in the
+            overall model (no outgoing connections).
+    """
+
+    pp_rank: int
+    pp_size: int
+    p2p_communicator: Optional[P2PCommunicator]
+    bridge_comms_as_src_module: Optional[List[BridgeCommunicator]]
+    bridge_comms_as_dest_module: Optional[List[BridgeCommunicator]]
+    is_source_stage: Optional[bool] = True
+    is_terminal_stage: Optional[bool] = True
+
+
+class MultiModulePipelineCommunicator:
+    """Communicator for a multi-module pipeline."""
+
+    def __init__(
+        self,
+        module_to_grid_map: Dict[str, HyperCommGrid],
+        topology: Dict[str, List[str]],
+        config: ModelParallelConfig,
+        dim_mapping: Dict[str, List[int]] = None,
+    ):
+        """
+        Initialize the MultiModulePipelineCommunicator.
+
+        Args:
+            module_to_grid_map (dict): A dictionary mapping module names to HyperCommGrids.
+                Example:
+                    module_to_grid_map = {
+                        'image_encoder': image_encoder_grid,
+                        'audio_encoder': audio_encoder_grid,
+                        'llm': llm_grid,
+                        'generator': generator_grid
+                    }
+            topology (dict): A dictionary mapping module names to lists of outgoing modules.
+                Example:
+                    topology = {
+                        'image_encoder': ['llm'],
+                        'audio_encoder': ['llm'],
+                        'llm': ['generator'],
+                        'generator': []
+                    }
+            config (ModelParallelConfig): A ModelParallelConfig object.
+            dim_mapping (Dict[str, List[int]]): Dimension mapping for sequence, batch, hidden.
+                Example:
+                    dim_mapping = {'s': 0, 'h': 2, 'b': 1}
+                Default: None
+        """
+        self.module_to_grid_map = module_to_grid_map
+        self.topology = topology
+        self.config = config
+        self.dim_mapping = dim_mapping
+        self.current_rank = dist.get_rank()
+
+        # Build bridge communicators for all modules
+        self.bridge_comms = []
+        self._build_bridge_comms()
+
+        self.rank_module_map = {}
+        self._build_rank_module_info_map()
+
+    def _build_bridge_comms(self):
+        """Construct and store BridgeCommunicator objects that describe the outgoing
+        communication relationships for all of the modules.
+        """
+        for src_module_name, src_grid in self.module_to_grid_map.items():
+            for dest_module_name in self.topology[src_module_name]:
+                dest_grid = self.module_to_grid_map[dest_module_name]
+                bridge_comm = BridgeCommunicator(
+                    src_grid=src_grid,
+                    dest_grid=dest_grid,
+                    dim_mapping=self.dim_mapping,
+                    comm_dtype=self.config.pipeline_dtype,
+                    src_module_name=src_module_name,
+                    dest_module_name=dest_module_name,
+                )
+                self.bridge_comms.append(bridge_comm)
+
+    @property
+    def is_pp_first_stage(self):
+        """Return True if the current rank has the absolute first stage in the overall model.
+
+        The absolute first stage is defined as:
+        1. The current rank must be in the first PP stage (pp_rank == 0) of some module
+        2. That module must be a source module (no incoming connections in topology)
+        """
+        for module_name, rank_module_info in self.rank_module_map.items():
+            # Check if this rank is at the first PP stage of this module
+            if rank_module_info.pp_rank == 0:
+                # Check if this module is a source module (no incoming connections)
+                if self._is_source_module(module_name):
+                    return True
+        return False
+
+    @property
+    def is_pp_last_stage(self):
+        """Return True if the current rank has the absolute last stage in the overall model.
+
+        The absolute last stage is defined as:
+        1. The current rank must be in the last PP stage of some module
+        2. That module must be a sink module (no outgoing connections in topology)
+        """
+        for module_name, rank_module_info in self.rank_module_map.items():
+            # Check if this rank is at the last PP stage of this module
+            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
+                # Check if this module is a sink module (no outgoing connections)
+                if self._is_sink_module(module_name):
+                    return True
+        return False
+
+    def _is_source_module(self, module_name: str) -> bool:
+        """Check if a module is a source module (has no incoming connections)."""
+        # A module is a source if no other module lists it as a destination
+        for src_module, dest_modules in self.topology.items():
+            if module_name in dest_modules:
+                return False
+        return True
+
+    def _is_sink_module(self, module_name: str) -> bool:
+        """Check if a module is a sink module (has no outgoing connections)."""
+        return len(self.topology.get(module_name, [])) == 0
+
+    def is_current_rank_in_grid(self, grid: HyperCommGrid) -> bool:
+        """Check if the current rank is in the grid."""
+        return grid.rank_offset <= self.current_rank < grid.rank_offset + grid.size
+
+    @property
+    def num_warmup_microbatches(self):
+        """Calculate the number of warmup microbatches for the current rank.
+
+        Uses the same simple logic as P2PCommunicator:
+        total_pipeline_stages - current_rank_stage - 1
+
+        Returns:
+            int: Number of warmup microbatches for this rank
+        """
+        # Get total pipeline depth across all modules
+        total_stages = self.compute_total_pipeline_stages(self.topology, self.module_to_grid_map)
+
+        # Get current rank's position in the overall pipeline (0-indexed)
+        # Use compute_total_pipeline_stages with current rank to get cumulative position
+        if self.rank_module_map:
+            # Take the first module this rank belongs to
+            # TODO: ykarnati - improve this logic.
+            module_name = next(iter(self.rank_module_map.keys()))
+            current_stage = (
+                self.compute_total_pipeline_stages(
+                    self.topology,
+                    self.module_to_grid_map,
+                    rank=self.current_rank,
+                    module_name=module_name,
+                )
+                - 1
+            )  # Convert from 1-indexed to 0-indexed
+        else:
+            current_stage = 0
+
+        assert (
+            current_stage <= total_stages
+        ), f"current_stage: {current_stage} is greater than total_stages: {total_stages}"
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"current_stage: {current_stage} total_stages: {total_stages} "
+            f"num_warmup_microbatches: {total_stages - current_stage - 1}"
+        )
+        return total_stages - current_stage - 1
+
+    def _build_rank_module_info_map(self):
+        """For each module in the current rank, initialize the P2P communicator
+        and build the bridge communicator info for the module.
+        Each rank may hold multiple modules when colocated.
+        """
+        for module_name, module_grid in self.module_to_grid_map.items():
+            if self.is_current_rank_in_grid(module_grid):
+                # Initialize P2P communicator
+                pp_group = module_grid.get_pg('pp')
+                p2p_comm = P2PCommunicator(pp_group, self.config)
+                pp_size = dist.get_world_size(pp_group)
+                rank_in_pp_group = dist.get_group_rank(pp_group, self.current_rank)
+                pp_rank = rank_in_pp_group % pp_size
+
+                bridge_comms_as_dest_module = []
+                bridge_comms_as_src_module = []
+                # If first stage, check if the module has any incoming modules
+                # If so, initialize bridge communicator
+                if pp_rank == 0:
+                    for bridge_comm in self.bridge_comms:
+                        if (
+                            bridge_comm.is_current_rank_in_grid(bridge_comm.dest_grid)
+                            and bridge_comm.dest_module_name == module_name
+                        ):
+                            bridge_comms_as_dest_module.append(bridge_comm)
+                # If last stage, check if the module has any outgoing modules
+                # If so, initialize bridge communicator
+                if pp_rank == pp_size - 1:
+                    for bridge_comm in self.bridge_comms:
+                        if (
+                            bridge_comm.is_current_rank_in_grid(bridge_comm.src_grid)
+                            and bridge_comm.src_module_name == module_name
+                        ):
+                            bridge_comms_as_src_module.append(bridge_comm)
+                # Build RankModuleInfo for the module
+                rank_module_info = RankModuleInfo(
+                    pp_rank=pp_rank,
+                    pp_size=pp_size,
+                    p2p_communicator=p2p_comm,
+                    bridge_comms_as_dest_module=bridge_comms_as_dest_module,
+                    bridge_comms_as_src_module=bridge_comms_as_src_module,
+                )
+                self.rank_module_map[module_name] = rank_module_info
+
+    def recv_forward(
+        self, tensor_shape: Optional[Shape] = None, is_first_stage: bool = False
+    ) -> Dict[str, torch.Tensor]:
+        """Receive forward activation tensor.
+
+        Args:
+            tensor_shape: Expected activation tensor shape
+
+        Returns:
+            A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[receive_forward] tensors_shape: {tensor_shape}, is_first_stage: {is_first_stage}"
+        )
+        input_dict = {}
+        for module_name, rank_module_info in self.rank_module_map.items():
+
+            if rank_module_info.pp_rank == 0:
+                # If first stage, and has incoming modules, receive forward activation
+                # from incoming modules.
+                for bridge_comm in rank_module_info.bridge_comms_as_dest_module:
+                    input_dict[bridge_comm.src_module_name] = bridge_comm.recv_forward()
+            else:
+                # If not first stage, receive forward activation tensor from P2P communicator.
+                input_dict[module_name] = rank_module_info.p2p_communicator.recv_forward(
+                    tensor_shapes=tensor_shape, is_first_stage=False
+                )
+        return input_dict
+
+    def send_forward(self, output_dict: Dict[str, torch.Tensor], is_last_stage: bool = False):
+        """Send forward activation tensor.
+
+        Args:
+            output_dict: A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[send_forward] output_dict keys: {output_dict.keys()}, is_last_stage: {is_last_stage}"
+        )
+        for module_name, rank_module_info in self.rank_module_map.items():
+            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
+                # If last stage, and has outgoing modules, send forward activation
+                # by using bridge communicator.
+                for bridge_comm in rank_module_info.bridge_comms_as_src_module:
+                    bridge_comm.send_forward(output_dict[module_name])
+            else:
+                # If not last stage, send forward activation by using P2P communicator.
+                rank_module_info.p2p_communicator.send_forward(
+                    output_dict[module_name], is_last_stage=False
+                )
+
+    def send_forward_recv_backward(
+        self,
+        output_dict: Dict[str, torch.Tensor],
+        tensor_shape: Optional[Shape] = None,
+        is_last_stage: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        """Send forward activation tensor and receive backward activation tensor.
+
+        Args:
+            output_dict: A dictionary mapping module names to tensors.
+            tensor_shape: Expected gradient tensor shape
+
+        Returns:
+            A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[send_forward_recv_backward] output_dict keys: {output_dict.keys()}, "
+            f"tensor_shape: {tensor_shape}, is_last_stage: {is_last_stage}"
+        )
+        grad_dict = {}
+        for module_name, rank_module_info in self.rank_module_map.items():
+            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
+                # If last stage, and has outgoing modules, send forward activation and
+                # receive backward gradient by using bridge communicator.
+                for bridge_comm in rank_module_info.bridge_comms_as_src_module:
+                    grad_dict[bridge_comm.src_module_name] = bridge_comm.send_forward_recv_backward(
+                        output_dict[module_name]
+                    )
+            else:
+                # If not last stage, send forward activation and receive backward gradient
+                # by using P2P communicator.
+                grad_dict[module_name] = (
+                    rank_module_info.p2p_communicator.send_forward_recv_backward(
+                        output_dict[module_name], tensor_shapes=tensor_shape, is_last_stage=False
+                    )
+                )
+        return grad_dict
+
+    def send_backward_recv_forward(
+        self,
+        grad_dict: Dict[str, torch.Tensor],
+        tensor_shape: Optional[Shape] = None,
+        is_first_stage: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        """Send backward activation tensor and receive forward activation tensor.
+
+        Args:
+            grad_dict: A dictionary mapping module names to tensors.
+            tensor_shape: Expected gradient tensor shape
+
+        Returns:
+            A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[send_backward_recv_forward] grad_dict keys: {grad_dict.keys()}, "
+            f"tensor_shape: {tensor_shape}, is_first_stage: {is_first_stage}"
+        )
+        input_dict = {}
+        for module_name, rank_module_info in self.rank_module_map.items():
+            if rank_module_info.pp_rank == 0:
+                for bridge_comm in rank_module_info.bridge_comms_as_dest_module:
+                    # If first stage, and has incoming modules, send backward gradient and
+                    # receive forward activation by using bridge communicator.
+                    input_dict[bridge_comm.src_module_name] = (
+                        bridge_comm.send_backward_recv_forward(
+                            grad_dict[bridge_comm.src_module_name]
+                        )
+                    )
+            else:
+                # If not first stage, send backward gradient and receive forward activation
+                # by using P2P communicator.
+                input_dict[module_name] = (
+                    rank_module_info.p2p_communicator.send_backward_recv_forward(
+                        grad_dict[module_name], tensor_shapes=tensor_shape, is_first_stage=False
+                    )
+                )
+        return input_dict
+
+    def recv_backward(
+        self, tensor_shape: Optional[Shape] = None, is_last_stage: bool = False
+    ) -> Dict[str, torch.Tensor]:
+        """Receive backward activation tensor.
+
+        Args:
+            tensor_shape: Expected gradient tensor shape
+
+        Returns:
+            A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[recv_backward] tensor_shape: {tensor_shape}, is_last_stage: {is_last_stage}"
+        )
+        grad_dict = {}
+        for module_name, rank_module_info in self.rank_module_map.items():
+            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
+                # If last stage, and has incoming modules, receive backward gradient
+                # by using bridge communicator.
+                for bridge_comm in rank_module_info.bridge_comms_as_src_module:
+                    grad_dict[bridge_comm.src_module_name] = bridge_comm.recv_backward()
+            else:
+                # If not last stage, receive backward gradient by using P2P communicator.
+                grad_dict[module_name] = rank_module_info.p2p_communicator.recv_backward(
+                    tensor_shapes=tensor_shape, is_last_stage=False
+                )
+        return grad_dict
+
+    def send_backward(self, grad_dict: Dict[str, torch.Tensor], is_first_stage: bool = False):
+        """Send backward activation tensor.
+
+        Args:
+            grad_dict: A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[send_backward] grad_dict keys: {grad_dict.keys()}, is_first_stage: {is_first_stage}"
+        )
+        for module_name, rank_module_info in self.rank_module_map.items():
+            if rank_module_info.pp_rank == 0:
+                # If first stage, and has incoming modules, send backward activation
+                # by using bridge communicator.
+                for bridge_comm in rank_module_info.bridge_comms_as_dest_module:
+                    bridge_comm.send_backward(grad_dict[bridge_comm.src_module_name])
+            else:
+                # If not first stage, send backward activation by using P2P communicator.
+                rank_module_info.p2p_communicator.send_backward(
+                    grad_dict[module_name], is_first_stage=False
+                )
+
+    @staticmethod
+    def compute_total_pipeline_stages(
+        topology: Dict[str, List[str]],
+        module_to_grid_map: Dict[str, HyperCommGrid],
+        rank: Optional[int] = None,
+        module_name: Optional[str] = None,
+    ) -> int:
+        """Compute the total number of pipeline stages across a multi-module chain.
+
+        Interprets ``topology`` as a directed acyclic graph (DAG) where nodes are modules
+        and edges indicate forward data flow from source to destination modules. Each node
+        is assigned a weight equal to its pipeline parallel size (number of PP stages).
+
+        The total number of stages is defined as the length of the longest path in this DAG
+        under node weights.
+
+        If ``rank`` is None (default), returns the maximum over all terminal (sink) modules of
+        the sum of PP sizes along a path ending at that terminal. For example, given:
+
+            image_encoder ->\
+                              -> llm -> generator
+            audio_encoder  ->/
+
+        the total is: max(pp(image_encoder), pp(audio_encoder)) + pp(llm) + pp(generator).
+
+        If ``rank`` is provided, the result is the total number of pipeline stages up to (and
+        including) the PP stage that ``rank`` occupies inside its module. In this case, the
+        weight of the target module equals (pp_rank_index(rank) + 1) instead of the module's
+        full PP size; other modules still contribute their full PP sizes. If the rank belongs to
+        multiple modules (colocation), pass ``module_name`` to disambiguate; otherwise the
+        maximum across all candidate modules containing the rank is returned.
+
+        Args:
+            topology: Mapping from a module to its list of outgoing modules.
+            module_to_grid_map: Mapping from module name to its ``HyperCommGrid``.
+
+        Returns:
+            The total number of pipeline stages along the longest path given the constraints.
+
+        Raises:
+            ValueError: If the topology contains cycles; or has no terminal nodes when
+                ``rank`` is None
+        """
+        nodes = set(module_to_grid_map.keys())
+        # Build adjacency and reverse-adjacency (predecessors).
+        adj: Dict[str, List[str]] = {node: list(topology.get(node, [])) for node in nodes}
+        preds: Dict[str, List[str]] = {node: [] for node in nodes}
+        for src, outs in adj.items():
+            for dst in outs:
+                preds[dst].append(src)
+
+        # Identify terminal nodes (no outgoing edges) for the rank=None case.
+        sinks = [node for node, outs in adj.items() if not outs]
+        if rank is None and not sinks:
+            raise ValueError(
+                "Topology must be a DAG with at least one terminal (no outgoing) module."
+            )
+
+        def pp_size(name: str) -> int:
+            grid = module_to_grid_map[name]
+            pp_dim_index = grid.dim_names.index('pp')
+            return grid.shape[pp_dim_index]
+
+        def partial_weight_for_target(target: str) -> Optional[int]:
+            if rank is None:
+                return None
+            grid = module_to_grid_map.get(target)
+            rank_groups = grid._gen_rank_enum(['pp'])
+            stage_index: Optional[int] = None
+            for group in rank_groups:
+                if rank in group:
+                    stage_index = group.index(rank)
+                    break
+            return stage_index + 1
+
+        def longest_path_to(target: str) -> int:
+            visiting = set()
+            partial = partial_weight_for_target(target)
+
+            def weight(name: str) -> int:
+                if partial is not None and name == target:
+                    return partial
+                return pp_size(name)
+
+            def dfs(node: str) -> int:
+                if node in visiting:
+                    raise ValueError("Topology contains cycles; expected a DAG.")
+                visiting.add(node)
+                best = 0
+                for p in preds.get(node, []):
+                    val = dfs(p)
+                    if val > best:
+                        best = val
+                visiting.remove(node)
+                return weight(node) + best
+
+            return dfs(target)
+
+        if rank is None:
+            return max(longest_path_to(sink) for sink in sinks)
+
+        return longest_path_to(module_name)
diff --git a/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py b/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py
new file mode 100644
index 00000000000..22f790cc0a9
--- /dev/null
+++ b/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py
@@ -0,0 +1,782 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+import os
+import sys
+
+import pytest
+import torch
+import torch.distributed as dist
+from packaging import version
+
+from megatron.core import parallel_state
+from megatron.core.hyper_comm_grid import HyperCommGrid
+from megatron.core.model_parallel_config import ModelParallelConfig
+from megatron.core.pipeline_parallel.multimodule_communicator import MultiModulePipelineCommunicator
+from tests.unit_tests.pipeline_parallel.test_bridge_communicator import (
+    _avg_params,
+    _create_transformer_block,
+    _get_pg_collection_from_grid,
+    create_hypercomm_grid,
+    get_transformer_block_and_grid,
+)
+from tests.unit_tests.test_utilities import Utils
+
+
+class TestMultiModulePipelineCommunicator:
+
+    @classmethod
+    def setup_class(cls):
+        """Set up distributed environment for the entire test class."""
+        if not dist.is_initialized():
+            dist.init_process_group(backend="nccl")
+        if torch.cuda.is_available():
+            torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+
+        world_size = dist.get_world_size()
+        if world_size != 8:
+            pytest.skip(
+                f"These tests require 8 GPUs, but only {world_size} are available.",
+                allow_module_level=True,
+            )
+
+    def teardown_class(cls):
+        Utils.destroy_model_parallel()
+
+    def test_multimodule_communicator_init(self):
+        """Test MultiModulePipelineCommunicator initialization."""
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
+        generator_grid = create_hypercomm_grid(offset=6, tp=2, cp=1, pp=1, dp=1)
+
+        # Define module-grid mapping
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        # Define module computation topology
+        topology = {
+            'image_encoder': ['llm'],
+            'audio_encoder': ['llm'],
+            'llm': ['generator'],
+            'generator': [],
+        }
+        config = ModelParallelConfig(bf16=True)
+        # Initialize communicator
+        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
+        # Test attributes match expectations
+        assert mllm_comm.module_to_grid_map == module_to_grid_map
+        assert mllm_comm.topology == topology
+        assert mllm_comm.config == config
+        assert mllm_comm.current_rank == dist.get_rank()
+
+    def test_compute_total_pipeline_stages(self):
+        """Test compute_total_pipeline_stages for overall chain and until specific ranks."""
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
+        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
+
+        # Define module-grid mapping and topology
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        topology = {
+            'image_encoder': ['llm'],
+            'audio_encoder': ['llm'],
+            'llm': ['generator'],
+            'generator': [],
+        }
+
+        # Overall total pipeline stages: max(1,1) + 2 + 1 = 4
+        total = MultiModulePipelineCommunicator.compute_total_pipeline_stages(
+            topology, module_to_grid_map
+        )
+        assert total == 4
+
+        llm_pp_rank = MultiModulePipelineCommunicator.compute_total_pipeline_stages(
+            topology, module_to_grid_map, rank=2, module_name='llm'
+        )
+        assert llm_pp_rank == 2
+
+    def test_send_forward_recv_forward(self):
+        """Test send_forward and recv_forward operations."""
+        if not dist.is_initialized():
+            pytest.skip("Distributed not initialized")
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
+        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
+
+        # Set up module-grid mapping and topology
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        topology = {
+            'image_encoder': ['llm'],
+            'audio_encoder': ['llm'],
+            'llm': ['generator'],
+            'generator': [],
+        }
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
+
+        # Simulate forward communication for each module
+        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
+            # Image encoder sends output forward
+            output_dict = {'image_encoder': torch.randn(2, 8, 128).cuda()}
+            mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
+            # Audio encoder sends output forward
+            output_dict = {'audio_encoder': torch.randn(2, 16, 128).cuda()}
+            mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+            if dist.get_rank() == 2 or dist.get_rank() == 3:
+                # LLM stage receives both image and audio outputs
+                input_dict = mllm_comm.recv_forward()
+                assert input_dict['image_encoder'].shape == (2, 8, 128)
+                assert input_dict['audio_encoder'].shape == (2, 16, 128)
+                mllm_comm.send_forward(output_dict)
+            else:
+                # LLM stage receives concatenated LLM outputs
+                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 32, 128))
+                assert input_dict['llm'].shape == (2, 32, 128)
+                mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(generator_grid):
+            # Generator module receives final LLM output
+            input_dict = mllm_comm.recv_forward()
+            assert input_dict['llm'].shape == (1, 32, 128)
+
+    def test_send_forward_recv_forward_with_different_pp_size(self):
+        """Test for the case when pp(image_encoder) != pp(audio_encoder)."""
+        if not dist.is_initialized():
+            pytest.skip("Distributed not initialized")
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=2, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=4, tp=1, cp=1, pp=4, dp=1)
+
+        # Set up module-grid mapping and topology
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+        }
+        topology = {'image_encoder': ['llm'], 'audio_encoder': ['llm'], 'llm': []}
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
+
+        # Simulate forward communication for each module
+        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
+            output_dict = {'image_encoder': torch.randn(2, 8, 128).cuda()}
+            if dist.get_rank() == 0:
+                # Image encoder sends output forward
+                mllm_comm.send_forward(output_dict)
+            else:
+                # Image stage receives image outputs
+                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 8, 128))
+                assert input_dict['image_encoder'].shape == (2, 8, 128)
+                mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
+            # Audio encoder sends output forward
+            output_dict = {'audio_encoder': torch.randn(2, 16, 128).cuda()}
+            mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+            if dist.get_rank() == 4:
+                # LLM stage receives both image and audio outputs
+                input_dict = mllm_comm.recv_forward()
+                assert input_dict['image_encoder'].shape == (2, 8, 128)
+                assert input_dict['audio_encoder'].shape == (2, 16, 128)
+                mllm_comm.send_forward(output_dict)
+            elif dist.get_rank() == 5 or dist.get_rank() == 6:
+                # LLM stage receives concatenated LLM outputs
+                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 32, 128))
+                assert input_dict['llm'].shape == (2, 32, 128)
+                mllm_comm.send_forward(output_dict)
+            elif dist.get_rank() == 7:
+                # LLM stage receives concatenated LLM outputs
+                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 32, 128))
+                assert input_dict['llm'].shape == (2, 32, 128)
+
+    def test_send_backward_recv_backward(self):
+        """Test send_backward and recv_backward operations."""
+        if not dist.is_initialized():
+            pytest.skip("Distributed not initialized")
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
+        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
+
+        # Set up module-grid mapping and topology
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        topology = {
+            'image_encoder': ['llm'],
+            'audio_encoder': ['llm'],
+            'llm': ['generator'],
+            'generator': [],
+        }
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
+
+        # Simulate backward communication for each module
+        if mllm_comm.is_current_rank_in_grid(generator_grid):
+            # Generator sends gradient backward
+            grad_dict = {'llm': torch.randn(1, 32, 128).cuda()}
+            mllm_comm.send_backward(grad_dict)
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            if dist.get_rank() == 4 or dist.get_rank() == 5:
+                # LLM receives expanded gradient and sends backward
+                received_grad = mllm_comm.recv_backward()
+                assert received_grad['llm'].shape == (2, 32, 128)
+                grad_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+                mllm_comm.send_backward(grad_dict)
+            else:
+                # LLM receives gradient and sends backward to both image/audio encoders
+                received_grad = mllm_comm.recv_backward(tensor_shape=(2, 32, 128))
+                assert received_grad['llm'].shape == (2, 32, 128)
+                grad_dict = {
+                    'image_encoder': torch.randn(2, 8, 128).cuda(),
+                    'audio_encoder': torch.randn(2, 16, 128).cuda(),
+                }
+                mllm_comm.send_backward(grad_dict)
+        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
+            # Image encoder receives its gradient
+            received_grad = mllm_comm.recv_backward()
+            assert received_grad['image_encoder'].shape == (2, 8, 128)
+        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
+            # Audio encoder receives its gradient
+            received_grad = mllm_comm.recv_backward()
+            assert received_grad['audio_encoder'].shape == (2, 16, 128)
+
+    @pytest.mark.skipif(
+        version.parse(torch.__version__) < version.parse('2.3.0'),
+        reason="Feature requires PyTorch 2.3 or later",
+    )
+    def test_send_forward_recv_backward_send_backward_recv_forward(self):
+        """Test send_forward_recv_backward and send_backward_recv_forward operations."""
+        if not dist.is_initialized():
+            pytest.skip("Distributed not initialized")
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
+        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
+
+        # Set up module-grid mapping and topology
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        topology = {
+            'image_encoder': ['llm'],
+            'audio_encoder': ['llm'],
+            'llm': ['generator'],
+            'generator': [],
+        }
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
+
+        # Simulate bidirectional send/recv for forward and backward in pipeline
+
+        # Encoder stages send forward to the first stage of LLM, and receive backward from the first stage of LLM
+        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
+            output_dict = {'image_encoder': torch.randn(2, 8, 128).cuda()}
+            received_grad = mllm_comm.send_forward_recv_backward(output_dict)
+            assert received_grad['image_encoder'].shape == (2, 8, 128)
+        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
+            output_dict = {'audio_encoder': torch.randn(2, 16, 128).cuda()}
+            received_grad = mllm_comm.send_forward_recv_backward(output_dict)
+            assert received_grad['audio_encoder'].shape == (2, 16, 128)
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            if dist.get_rank() == 2 or dist.get_rank() == 3:
+                grad_dict = {
+                    'image_encoder': torch.randn(2, 8, 128).cuda(),
+                    'audio_encoder': torch.randn(2, 16, 128).cuda(),
+                }
+                input_dict = mllm_comm.send_backward_recv_forward(grad_dict)
+                assert input_dict['image_encoder'].shape == (2, 8, 128)
+                assert input_dict['audio_encoder'].shape == (2, 16, 128)
+
+        # First stage of LLM sends forward to the second stage of LLM, and receive backward from the second stage of LLM
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            if dist.get_rank() == 2 or dist.get_rank() == 3:
+                output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+                received_grad = mllm_comm.send_forward_recv_backward(
+                    output_dict, tensor_shape=(2, 32, 128)
+                )
+                assert received_grad['llm'].shape == (2, 32, 128)
+            if dist.get_rank() == 4 or dist.get_rank() == 5:
+                grad_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+                input_dict = mllm_comm.send_backward_recv_forward(
+                    grad_dict, tensor_shape=(2, 32, 128)
+                )
+                assert input_dict['llm'].shape == (2, 32, 128)
+
+        # Second stage of LLM sends forward to generator, and receive backward from generator
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            if dist.get_rank() == 4 or dist.get_rank() == 5:
+                output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+                received_grad = mllm_comm.send_forward_recv_backward(output_dict)
+                assert received_grad['llm'].shape == (2, 32, 128)
+        if mllm_comm.is_current_rank_in_grid(generator_grid):
+            grad_dict = {'llm': torch.randn(1, 32, 128).cuda()}
+            input_dict = mllm_comm.send_backward_recv_forward(grad_dict)
+            assert input_dict['llm'].shape == (1, 32, 128)
+
+    @pytest.mark.skipif(
+        version.parse(torch.__version__) < version.parse('2.3.0'),
+        reason="Feature requires PyTorch 2.3 or later",
+    )
+    def test_send_forward_recv_forward_with_transformer_blocks(self):
+        """Test send_forward and recv_forward operations."""
+
+        # Set model/test dimensions for easier debugging and output comparison
+        hidden_size = 16
+        sequence_length = 2
+        micro_batch_size = 2
+
+        # For reproducibility, set a fixed seed
+        torch.manual_seed(12345)
+        dtype = torch.float32
+
+        # Create random input hidden states tensor
+        hidden_states = torch.randn(
+            (sequence_length, micro_batch_size, hidden_size), device="cuda"
+        ).to(dtype)
+        current_rank = dist.get_rank()
+
+        # ========== Initialize tensor model-parallel environment ==========
+        parallel_state_tp = 2
+        Utils.initialize_model_parallel(tensor_model_parallel_size=2)
+
+        # ========== Build reference 1D grid and transformer block for weight sharing ==========
+        ref_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=8)
+        ref_pg_collection = _get_pg_collection_from_grid(ref_grid)
+        ref_block = _create_transformer_block(
+            dtype=dtype, hidden_size=hidden_size, pg_collection=ref_pg_collection
+        )
+        _avg_params(
+            ref_block, ref_grid.get_pg("dp")
+        )  # Ensure parameters are averaged across data parallel (DP)
+
+        # ========== Create different transformer blocks for each model stage ==========
+        # Image encoder
+        image_encoder_block, image_encoder_grid = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=1,
+            cp_size=1,
+            pp_size=1,
+            dp_size=1,
+            grid_offset=0,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        # Audio encoder
+        audio_encoder_block, audio_encoder_grid = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=1,
+            cp_size=1,
+            pp_size=1,
+            dp_size=1,
+            grid_offset=1,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        # LLM (Large Language Model) block with tensor & pipeline parallelism
+        llm_block, llm_grid = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=2,
+            cp_size=1,
+            pp_size=2,
+            dp_size=1,
+            grid_offset=2,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        # Generator block (final stage) with DP=2
+        generator_block, generator_grid = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=1,
+            cp_size=1,
+            pp_size=1,
+            dp_size=2,
+            grid_offset=6,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+
+        # ========== Define module-to-grid correspondence and pipeline topology ==========
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        topology = {
+            'image_encoder': ['llm'],  # image_encoder sends output to llm
+            'audio_encoder': ['llm'],  # audio_encoder sends output to llm
+            'llm': ['generator'],  # llm sends output to generator
+            'generator': [],  # generator is the final module
+        }
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        # Define dimension mapping for sequence, batch, hidden
+        dim_mapping = {'s': 0, 'h': 2, 'b': 1}
+        seq_dim = dim_mapping['s']
+
+        # Communication handler for multi-module pipeline (send/recv abstraction)
+        mllm_comm = MultiModulePipelineCommunicator(
+            module_to_grid_map, topology, config, dim_mapping=dim_mapping
+        )
+
+        # ========== Run actual distributed pipeline blocks (per process, depending on role) ==========
+        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
+            # Image encoder rank: run forward and send output
+            image_encoder_output = image_encoder_block(
+                hidden_states=hidden_states, attention_mask=None
+            )
+            output_dict = {'image_encoder': image_encoder_output}
+            mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
+            # Audio encoder rank: run forward and send output
+            audio_encoder_output = audio_encoder_block(
+                hidden_states=hidden_states, attention_mask=None
+            )
+            output_dict = {'audio_encoder': audio_encoder_output}
+            mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            if dist.get_rank() == 2 or dist.get_rank() == 3:
+                # LLM stage 0 (receives both image and audio, concatenates along seq_dim)
+                input_dict = mllm_comm.recv_forward()
+                llm_output = llm_block(
+                    hidden_states=torch.cat(
+                        [input_dict['image_encoder'], input_dict['audio_encoder']], dim=seq_dim
+                    ),
+                    attention_mask=None,
+                )
+                output_dict = {'llm': llm_output}
+                mllm_comm.send_forward(output_dict)
+            else:
+                # LLM stage 1 (receives output of previous LLM stage)
+                input_dict = mllm_comm.recv_forward(
+                    tensor_shape=(sequence_length * 2, micro_batch_size, hidden_size)
+                )
+                llm_output = llm_block(hidden_states=input_dict['llm'], attention_mask=None)
+                output_dict = {'llm': llm_output}
+                mllm_comm.send_forward(output_dict)
+
+        if mllm_comm.is_current_rank_in_grid(generator_grid):
+            # Generator block: only receives from llm and runs forward
+            input_dict = mllm_comm.recv_forward()
+            generator_output = generator_block(hidden_states=input_dict['llm'], attention_mask=None)
+
+        # ========== Build a reference (serial/global) pipeline for correctness checking ==========
+        global_image_encoder_block, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        global_audio_encoder_block, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        global_llm_block_pp_rank_0, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        global_llm_block_pp_rank_1, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        global_generator_block, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+
+        # Run each stage sequentially as a global pipeline (for truth)
+        global_image_encoder_output = global_image_encoder_block(
+            hidden_states=hidden_states, attention_mask=None
+        )
+        global_audio_encoder_output = global_audio_encoder_block(
+            hidden_states=hidden_states, attention_mask=None
+        )
+        # Compare output between global and distributed blocks for image/audio stage
+        if current_rank == 0:
+            torch.testing.assert_close(
+                global_image_encoder_output, image_encoder_output, rtol=1e-3, atol=1e-3
+            )
+        if current_rank == 1:
+            torch.testing.assert_close(
+                global_audio_encoder_output, audio_encoder_output, rtol=1e-3, atol=1e-3
+            )
+
+        # Feed outputs to LLM stages (emulate pipeline cut with concatenation)
+        global_llm_input = torch.cat(
+            [global_image_encoder_output, global_audio_encoder_output], dim=seq_dim
+        )
+        global_llm_pp_rank_0_output = global_llm_block_pp_rank_0(
+            hidden_states=global_llm_input, attention_mask=None
+        )
+        if current_rank == 2 or current_rank == 3:
+            torch.testing.assert_close(
+                global_llm_pp_rank_0_output, llm_output, rtol=1e-3, atol=1e-3
+            )
+        global_llm_pp_rank_1_output = global_llm_block_pp_rank_1(
+            hidden_states=global_llm_pp_rank_0_output, attention_mask=None
+        )
+        if current_rank == 4 or current_rank == 5:
+            torch.testing.assert_close(
+                global_llm_pp_rank_1_output, llm_output, rtol=1e-3, atol=1e-3
+            )
+
+        # Generator output and comparison to distributed output (for each DP chunk)
+        global_generator_block_output = global_generator_block(
+            hidden_states=global_llm_pp_rank_1_output, attention_mask=None
+        )
+        global_generator_block_chunks = torch.split(
+            global_generator_block_output, global_generator_block_output.shape[1] // 2, dim=1
+        )
+        if current_rank == 6:
+            torch.testing.assert_close(
+                global_generator_block_chunks[0], generator_output, rtol=1e-3, atol=1e-3
+            )
+        if current_rank == 7:
+            torch.testing.assert_close(
+                global_generator_block_chunks[1], generator_output, rtol=1e-3, atol=1e-3
+            )
+
+    @pytest.mark.skipif(
+        version.parse(torch.__version__) < version.parse('2.3.0'),
+        reason="Feature requires PyTorch 2.3 or later",
+    )
+    @pytest.mark.parametrize(
+        "grid1_tp, grid1_pp, grid1_dp, grid2_tp, grid2_pp, grid2_dp, parallel_state_tp",
+        [
+            (2, 1, 1, 2, 1, 1, 2),  # TP2PP1DP1 to TP2PP1DP1
+            (2, 1, 1, 2, 2, 1, 2),  # TP2PP1DP1 to TP2PP2DP1
+            (2, 2, 1, 2, 2, 1, 2),  # TP2PP2DP1 to TP2PP2DP1
+            (4, 1, 1, 4, 1, 1, 4),  # TP4DP1 to TP4DP1
+            (2, 1, 2, 4, 1, 1, 2),  # TP2DP2 to TP4DP1
+            (4, 1, 1, 2, 1, 2, 2),  # TP4DP1 to TP2DP2
+            (2, 1, 2, 1, 1, 4, 2),  # TP2DP2 to TP1DP4
+        ],
+    )
+    def test_send_forward_recv_forward_with_transformer_blocks_and_different_parallelisms(
+        self, grid1_tp, grid1_pp, grid1_dp, grid2_tp, grid2_pp, grid2_dp, parallel_state_tp
+    ):
+        """Test bridge communicator with two transformer blocks having different process group configurations."""
+        # Model and input configuration
+        hidden_size = 16
+        sequence_length = 2
+        micro_batch_size = 8
+        torch.manual_seed(12345)
+        dtype = torch.float32
+
+        # Create random input tensor on CUDA
+        hidden_states = torch.randn(
+            (sequence_length, micro_batch_size, hidden_size), device="cuda"
+        ).to(dtype)
+        hidden_states_ref = hidden_states.clone()
+        current_rank = dist.get_rank()
+
+        # Initialize model parallel with desired TP
+        Utils.initialize_model_parallel(tensor_model_parallel_size=parallel_state_tp)
+
+        # Build a reference grid and block for parameter sharing & DP averaging
+        ref_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=8)
+        ref_pg_collection = _get_pg_collection_from_grid(ref_grid)
+        ref_block = _create_transformer_block(
+            dtype=dtype, hidden_size=hidden_size, pg_collection=ref_pg_collection
+        )
+        _avg_params(
+            ref_block, ref_grid.get_pg("dp")
+        )  # Synchronize parameters across DP for reproducibility
+
+        # ====== Create two transformer block+grid pairs with different TP/DP settings ======
+        block_grid_1, grid_1 = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=grid1_tp,
+            pp_size=grid1_pp,
+            dp_size=grid1_dp,
+            grid_offset=0,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+
+        block_grid_2, grid_2 = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=grid2_tp,
+            pp_size=grid2_pp,
+            dp_size=grid2_dp,
+            grid_offset=grid_1.size,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+
+        dist.barrier()  # Synchronize ranks before communication
+
+        # Module-grid map and pipeline communication topology
+        module_to_grid_map = {'image_encoder': grid_1, 'llm': grid_2}
+        topology = {
+            'image_encoder': ['llm'],  # image_encoder sends forward results to llm
+            'llm': [],  # llm is the last stage here
+        }
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        mllm_comm = MultiModulePipelineCommunicator(
+            module_to_grid_map, topology, config, dim_mapping={'s': 0, 'h': 2, 'b': 1}
+        )
+
+        output_grid_2 = None
+        # If current rank is in the first grid, run first block and send output
+        if grid_1 is not None and mllm_comm.is_current_rank_in_grid(grid_1):
+            rank_module_info = mllm_comm.rank_module_map['image_encoder']
+            if rank_module_info.pp_rank == 0:
+                hidden_states = block_grid_1(hidden_states=hidden_states, attention_mask=None)
+                mllm_comm.send_forward({'image_encoder': hidden_states})
+            else:
+                input_dict = mllm_comm.recv_forward(
+                    tensor_shape=(sequence_length, micro_batch_size, hidden_size)
+                )
+                hidden_states = input_dict['image_encoder']
+                hidden_states = block_grid_1(hidden_states=hidden_states, attention_mask=None)
+                mllm_comm.send_forward({'image_encoder': hidden_states})
+
+        # If current rank is in second grid, receive and run the second block
+        if grid_2 is not None and mllm_comm.is_current_rank_in_grid(grid_2):
+            rank_module_info = mllm_comm.rank_module_map['llm']
+            if rank_module_info.pp_rank == 0:
+                input_dict = mllm_comm.recv_forward()
+                hidden_states = input_dict['image_encoder']
+                hidden_states = block_grid_2(hidden_states=hidden_states, attention_mask=None)
+                if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
+                    output_grid_2 = hidden_states
+                else:
+                    mllm_comm.send_forward({'llm': hidden_states})
+            elif rank_module_info.pp_rank < rank_module_info.pp_size - 1:
+                input_dict = mllm_comm.recv_forward(
+                    tensor_shape=(
+                        sequence_length,
+                        (grid1_dp * micro_batch_size) // grid2_dp,
+                        hidden_size,
+                    )
+                )
+                hidden_states = input_dict['llm']
+                hidden_states = block_grid_2(hidden_states=hidden_states, attention_mask=None)
+                mllm_comm.send_forward({'llm': hidden_states})
+            else:
+                input_dict = mllm_comm.recv_forward(
+                    tensor_shape=(
+                        sequence_length,
+                        (grid1_dp * micro_batch_size) // grid2_dp,
+                        hidden_size,
+                    )
+                )
+                hidden_states = input_dict['llm']
+                output_grid_2 = block_grid_2(hidden_states=hidden_states, attention_mask=None)
+
+                # Compute expected output shape based on change in DP size (chunk/expand batch dimension appropriately)
+                factor = max(grid1_dp, grid2_dp) // min(grid1_dp, grid2_dp)
+                expected_output_shape = (
+                    sequence_length,
+                    (
+                        micro_batch_size * factor
+                        if grid1_dp > grid2_dp
+                        else micro_batch_size // factor
+                    ),
+                    hidden_size,
+                )
+                assert (
+                    output_grid_2.shape == expected_output_shape
+                ), f"Output2 shape mismatch: {output_grid_2.shape}"
+
+        # ====== Reference: global (replicated) pipeline forward for correctness checking ======
+        global_block_1, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        global_block_2, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+
+        for i in range(grid1_pp):
+            hidden_states_ref = global_block_1(hidden_states=hidden_states_ref, attention_mask=None)
+
+        for i in range(grid2_pp):
+            hidden_states_ref = global_block_2(hidden_states=hidden_states_ref, attention_mask=None)
+
+        # Output comparison under different DP compositions between grids
+        if (
+            grid_2 is not None
+            and mllm_comm.is_current_rank_in_grid(grid_2)
+            and rank_module_info.pp_rank == rank_module_info.pp_size - 1
+        ):
+            if grid1_dp == grid2_dp:
+                # DP size matches: all outputs directly compared
+                torch.testing.assert_close(hidden_states_ref, output_grid_2, rtol=1e-3, atol=1e-3)
+            elif grid1_dp < grid2_dp:
+                # If grid2 expands DP: each output_grid_2 chunk corresponds to a split of the reference output
+                grid2_dp_ranks = grid_2._gen_rank_enum([x for x in grid_2.dim_names if x != "dp"])
+                global_block_2_chunks = torch.split(
+                    hidden_states_ref, hidden_states_ref.shape[1] // (grid2_dp // grid1_dp), dim=1
+                )
+                relevant_chunk = None
+                for i, dp_ranks in enumerate(grid2_dp_ranks):
+                    if current_rank in dp_ranks:
+                        relevant_chunk = global_block_2_chunks[i % len(global_block_2_chunks)]
+                torch.testing.assert_close(relevant_chunk, output_grid_2, rtol=1e-3, atol=1e-3)
+            else:
+                # If DP shrinks (grid1_dp > grid2_dp): just compare the relevant first chunk
+                output_grid_2_first_chunk = torch.chunk(output_grid_2, grid1_dp // grid2_dp, dim=1)[
+                    0
+                ]
+                torch.testing.assert_close(
+                    hidden_states_ref, output_grid_2_first_chunk, rtol=1e-3, atol=1e-3
+                )