diff --git a/megatron/core/models/mimo/comm/__init__.py b/megatron/core/models/mimo/comm/__init__.py
new file mode 100644
index 00000000000..26496bfed70
--- /dev/null
+++ b/megatron/core/models/mimo/comm/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/models/mimo/comm/colocated_communicator.py b/megatron/core/models/mimo/comm/colocated_communicator.py
new file mode 100644
index 00000000000..4c43dcdf3cd
--- /dev/null
+++ b/megatron/core/models/mimo/comm/colocated_communicator.py
@@ -0,0 +1,325 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+from dataclasses import dataclass
+from enum import Enum
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.distributed as dist
+
+from megatron.core.hyper_comm_grid import HyperCommGrid
+
+
+@dataclass
+class SliceInfo:
+    """Batch dimension slice information for a rank's data partition."""
+
+    start: int
+    size: int
+
+
+class BridgeDirection(str, Enum):
+    """Which side of the bridge scales up, if any.
+
+    ``FAN_IN`` — src has more DP replicas than dest; forward all-gathers
+    src outputs along the batch dim, backward narrows the sibling dest
+    gradient down to this src rank's slot.
+
+    ``FAN_OUT`` — dest has more DP replicas; forward narrows, backward
+    all-gathers across the sibling dest DP ranks (the adjoint of narrow
+    is not zero-pad-and-scatter because every dest rank consumes a
+    different slice of the same src activation).
+
+    ``EQUAL`` — matching DP; the bridge is a pure passthrough.
+    """
+
+    FAN_IN = "fan_in"
+    FAN_OUT = "fan_out"
+    EQUAL = "equal"
+
+
+class ColocatedBridgeCommunicator:
+    """Bridges tensors between colocated modules with different TP/DP layouts.
+
+    Default ``dim_mapping`` assumes 3D ``(b, s, h)``. Callers bridging
+    ``MimoModel``'s pre-flattened ``(s*b, h)`` encoder output should pass
+    ``dim_mapping={'b': 0, 'h': 1}``; this relies on a uniform token count per
+    sample so dim 0 divides evenly by the DP scale.
+
+    Precondition: the input must be TP-replicated across the src TP group —
+    i.e. all TP ranks inside a src DP replica hold the same tensor on the
+    batch dim. The bridge never gathers along TP; violating this silently
+    produces wrong results.
+    """
+
+    def __init__(
+        self,
+        src_grid: HyperCommGrid,
+        dest_grid: HyperCommGrid,
+        src_module_name: str = "src",
+        dest_module_name: str = "dest",
+        dim_mapping: Optional[Dict[str, int]] = None,
+    ):
+        self.src_grid = src_grid
+        self.dest_grid = dest_grid
+        self.src_module_name = src_module_name
+        self.dest_module_name = dest_module_name
+        self.dim_mapping = dim_mapping or {'b': 0, 's': 1, 'h': 2}
+        self.current_rank = dist.get_rank()
+
+        self._validate_grids()
+        self._extract_parallelism_info()
+        self._build_rank_mappings()
+
+        # At most one direction is active; fan-in and fan-out are mutually
+        # exclusive (one of ``src_dp / dest_dp`` is >1, the other is 1).
+        # Equal DP uses no collective at all. Unify behind a single
+        # ``gather_pg`` + ``direction`` + ``scale`` rather than a fan-in
+        # and fan-out pair of attributes.
+        self.gather_pg: Optional[dist.ProcessGroup] = None
+        self.gather_group_ranks: List[List[int]] = []
+
+        if self.src_dp_size > self.dest_dp_size:
+            self.direction = BridgeDirection.FAN_IN
+            self.scale = self.src_dp_size // self.dest_dp_size
+            self.gather_group_ranks = self._build_gather_groups(
+                iter_size=self.dest_dp_size,
+                sibling_tp_size=self.src_tp_size,
+                scale=self.scale,
+                rank_to_pos=self.rank_to_src_pos,
+            )
+            self.gather_pg, _ = dist.new_subgroups_by_enumeration(
+                self.gather_group_ranks, backend='nccl'
+            )
+        elif self.dest_dp_size > self.src_dp_size:
+            self.direction = BridgeDirection.FAN_OUT
+            self.scale = self.dest_dp_size // self.src_dp_size
+            self.gather_group_ranks = self._build_gather_groups(
+                iter_size=self.src_dp_size,
+                sibling_tp_size=self.dest_tp_size,
+                scale=self.scale,
+                rank_to_pos=self.rank_to_dest_pos,
+            )
+            self.gather_pg, _ = dist.new_subgroups_by_enumeration(
+                self.gather_group_ranks, backend='nccl'
+            )
+        else:
+            self.direction = BridgeDirection.EQUAL
+            self.scale = 1
+
+        logging.info(
+            f"[Rank {self.current_rank}] ColocatedBridgeCommunicator: "
+            f"{src_module_name}({self.src_tp_size}TP/{self.src_dp_size}DP) -> "
+            f"{dest_module_name}({self.dest_tp_size}TP/{self.dest_dp_size}DP), "
+            f"direction={self.direction.value}, scale={self.scale}"
+        )
+
+    def _validate_grids(self):
+        if self.src_grid.size != self.dest_grid.size:
+            raise ValueError(
+                f"Grids must span same number of ranks: "
+                f"src={self.src_grid.size}, dest={self.dest_grid.size}"
+            )
+
+        if self.src_grid.rank_offset != self.dest_grid.rank_offset:
+            raise ValueError(
+                f"Grids must have same rank offset: "
+                f"src={self.src_grid.rank_offset}, dest={self.dest_grid.rank_offset}"
+            )
+
+        # Per-grid dim checks: tp/dp required; pp and cp (if present) must be 1.
+        # CP>1 also corrupts dp_idx when iterating get_rank_enum(['tp']) groups.
+        for name, grid in [("src", self.src_grid), ("dest", self.dest_grid)]:
+            for required in ('tp', 'dp'):
+                if required not in grid.dim_names:
+                    raise ValueError(
+                        f"{name} grid must have '{required}' dimension, "
+                        f"got dim_names={grid.dim_names}"
+                    )
+            for singleton in ('pp', 'cp'):
+                if singleton in grid.dim_names:
+                    size = grid.shape[grid.dim_names.index(singleton)]
+                    if size != 1:
+                        raise ValueError(
+                            f"{name} {singleton.upper()} must be 1 for "
+                            f"ColocatedBridgeCommunicator, got {size}"
+                        )
+
+        src_dp = self.src_grid.shape[self.src_grid.dim_names.index('dp')]
+        dest_dp = self.dest_grid.shape[self.dest_grid.dim_names.index('dp')]
+        if src_dp % dest_dp != 0 and dest_dp % src_dp != 0:
+            raise ValueError(
+                f"DP sizes must be evenly divisible: src_dp={src_dp}, dest_dp={dest_dp}"
+            )
+
+    def _extract_parallelism_info(self):
+        self.src_tp_size = self.src_grid.shape[self.src_grid.dim_names.index('tp')]
+        self.src_dp_size = self.src_grid.shape[self.src_grid.dim_names.index('dp')]
+        self.dest_tp_size = self.dest_grid.shape[self.dest_grid.dim_names.index('tp')]
+        self.dest_dp_size = self.dest_grid.shape[self.dest_grid.dim_names.index('dp')]
+
+    def _build_rank_mappings(self):
+        self.rank_to_src_pos: Dict[int, Tuple[int, int]] = {}
+        self.rank_to_dest_pos: Dict[int, Tuple[int, int]] = {}
+
+        src_tp_groups = self.src_grid.get_rank_enum(['tp'])
+        for dp_idx, tp_group in enumerate(src_tp_groups):
+            for tp_idx, rank in enumerate(tp_group):
+                self.rank_to_src_pos[rank] = (dp_idx, tp_idx)
+
+        dest_tp_groups = self.dest_grid.get_rank_enum(['tp'])
+        for dp_idx, tp_group in enumerate(dest_tp_groups):
+            for tp_idx, rank in enumerate(tp_group):
+                self.rank_to_dest_pos[rank] = (dp_idx, tp_idx)
+
+    @staticmethod
+    def _build_gather_groups(
+        iter_size: int, sibling_tp_size: int, scale: int, rank_to_pos: Dict[int, Tuple[int, int]]
+    ) -> List[List[int]]:
+        """Build ``iter_size * sibling_tp_size`` gather groups of ``scale`` ranks.
+
+        For each slot on the "iterating" side and each TP shard on the
+        sibling side, collect the ``scale`` sibling ranks whose DP indices
+        map into that slot. Append order equals group-local-rank order,
+        which ``all_gather_into_tensor`` uses to concatenate outputs — do
+        not sort.
+        """
+        groups: List[List[int]] = []
+        for iter_idx in range(iter_size):
+            sibling_dp_indices = range(iter_idx * scale, (iter_idx + 1) * scale)
+            for sibling_tp_idx in range(sibling_tp_size):
+                group_ranks = []
+                for sibling_dp_idx in sibling_dp_indices:
+                    for rank, (dp, tp) in rank_to_pos.items():
+                        if dp == sibling_dp_idx and tp == sibling_tp_idx:
+                            group_ranks.append(rank)
+                            break
+                groups.append(group_ranks)
+        return groups
+
+    def is_fan_in(self) -> bool:
+        """True if src DP > dest DP (forward all-gathers)."""
+        return self.direction is BridgeDirection.FAN_IN
+
+    def is_fan_out(self) -> bool:
+        """True if src DP < dest DP (forward narrows)."""
+        return self.direction is BridgeDirection.FAN_OUT
+
+    def get_slice_info(self, batch_size: int) -> SliceInfo:
+        """Compute this rank's slice of ``batch_size`` on the narrowing side.
+
+        For FAN_OUT this is the forward narrow; for FAN_IN it is the
+        backward narrow against the post-gather batch. EQUAL returns the
+        identity slice.
+
+        Raises ``ValueError`` if ``batch_size`` is not divisible by ``scale``.
+        """
+        if self.direction is BridgeDirection.EQUAL:
+            return SliceInfo(start=0, size=batch_size)
+        self._check_divisible(batch_size)
+        if self.direction is BridgeDirection.FAN_OUT:
+            dp_idx = self.rank_to_dest_pos[self.current_rank][0]
+        else:  # FAN_IN
+            dp_idx = self.rank_to_src_pos[self.current_rank][0]
+        slot = dp_idx % self.scale
+        slice_size = batch_size // self.scale
+        return SliceInfo(start=slot * slice_size, size=slice_size)
+
+    def _check_divisible(self, batch_size: int) -> None:
+        if batch_size % self.scale != 0:
+            raise ValueError(
+                f"ColocatedBridgeCommunicator: batch dim size {batch_size} is "
+                f"not divisible by {self.direction.value} scale={self.scale}."
+            )
+
+    def communicate(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Transform ``tensor`` from src TP/DP layout to dest TP/DP layout.
+
+        Raises ``ValueError`` when FAN_OUT and the batch dim is not
+        divisible by ``scale``; FAN_IN only slices on the backward pass
+        and re-checks via ``get_slice_info`` there.
+        """
+        if self.direction is BridgeDirection.FAN_OUT:
+            self._check_divisible(tensor.shape[self.dim_mapping['b']])
+        return _ColocatedCommunicate.apply(tensor, self)
+
+    def destroy(self) -> None:
+        """Release the NCCL subgroup created by this communicator.
+
+        NCCL caps concurrent communicators; long-lived or repeated
+        construction leaks PGs without this call.
+        """
+        if self.gather_pg is not None:
+            dist.destroy_process_group(self.gather_pg)
+            self.gather_pg = None
+
+
+class _ColocatedCommunicate(torch.autograd.Function):
+    """Autograd function for colocated communication with correct backward pass."""
+
+    @staticmethod
+    def forward(ctx, tensor: torch.Tensor, comm: ColocatedBridgeCommunicator) -> torch.Tensor:
+        """Reshape the batch dim across the bridge: narrow on fan-out, all-gather on fan-in."""
+        ctx.comm = comm
+        ctx.batch_dim = comm.dim_mapping['b']
+
+        if comm.direction is BridgeDirection.FAN_OUT:
+            # Narrow this rank's slice out of the full src batch.
+            slice_info = comm.get_slice_info(tensor.shape[ctx.batch_dim])
+            return tensor.narrow(ctx.batch_dim, slice_info.start, slice_info.size).contiguous()
+
+        if comm.direction is BridgeDirection.FAN_IN:
+            # All-gather sibling src outputs into a single full-batch tensor.
+            return _all_gather_along_batch_dim(tensor, comm.gather_pg, ctx.batch_dim)
+
+        # EQUAL: pure passthrough.
+        return tensor.contiguous()
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]:
+        """Adjoint of forward: narrow for fan-in, all-gather for fan-out.
+
+        Fan-out's forward is ``narrow``, whose naive adjoint is zero-pad.
+        That would leave each src rank with only its own dest rank's
+        slice of the gradient, missing the contributions from every
+        other dest rank that consumed a different slice of the same src
+        activation. Instead we all-gather across the fan-out sibling
+        group, reconstructing the full src-batch gradient (symmetric
+        with the fan-in forward's all-gather).
+        """
+        comm = ctx.comm
+        batch_dim = ctx.batch_dim
+
+        if comm.direction is BridgeDirection.FAN_OUT:
+            return _all_gather_along_batch_dim(grad_output, comm.gather_pg, batch_dim), None
+
+        if comm.direction is BridgeDirection.FAN_IN:
+            slice_info = comm.get_slice_info(grad_output.shape[batch_dim])
+            return (
+                grad_output.narrow(batch_dim, slice_info.start, slice_info.size).contiguous(),
+                None,
+            )
+
+        return grad_output.contiguous(), None
+
+
+def _all_gather_along_batch_dim(
+    tensor: torch.Tensor, group: dist.ProcessGroup, batch_dim: int
+) -> torch.Tensor:
+    """All-gather ``tensor`` along an arbitrary batch dim into a single tensor.
+
+    ``all_gather_into_tensor`` concatenates along dim 0, so when the
+    batch dim is not 0 we move it, gather, then restore.
+    """
+    world_size = dist.get_world_size(group)
+    src = tensor.contiguous()
+    if batch_dim != 0:
+        src = src.movedim(batch_dim, 0).contiguous()
+    out_shape = list(src.shape)
+    out_shape[0] *= world_size
+    out = torch.empty(out_shape, dtype=tensor.dtype, device=tensor.device)
+    dist.all_gather_into_tensor(out, src, group=group)
+    if batch_dim != 0:
+        out = out.movedim(0, batch_dim).contiguous()
+    return out
diff --git a/megatron/core/models/mimo/config/base_configs.py b/megatron/core/models/mimo/config/base_configs.py
index a92484a5a48..0eda09465e0 100644
--- a/megatron/core/models/mimo/config/base_configs.py
+++ b/megatron/core/models/mimo/config/base_configs.py
@@ -23,9 +23,11 @@ class MimoModelConfig:
             in the input_ids to insert the modality embeddings at the correct positions.
         module_to_grid_map (Optional[Dict[str, HyperCommGrid]]):
             Dictionary mapping module keys (e.g., "vision", "language") to their
-            corresponding HyperCommGrid configurations for non-colocated pipeline
-            parallelism. The language model must use the key MIMO_LANGUAGE_MODULE_KEY.
-            When None, all modules are assumed to be colocated on the same ranks.
+            corresponding HyperCommGrid configurations. The language model must use
+            the key MIMO_LANGUAGE_MODULE_KEY.
+            When grids span the same ranks → colocated (same or different TP/DP).
+            When grids span disjoint ranks → non-colocated (pipeline parallel).
+            When None → colocated with legacy global parallel_state.
         kv_format (str):
             Key-value format for attention: "sbhd" (seq-batch-head-dim) or "thd" (total-head-dim).
             Default is "sbhd".
@@ -43,3 +45,18 @@ class MimoModelConfig:
     special_token_ids: Dict[str, int] = field(default_factory=dict)
     module_to_grid_map: Optional[Dict[str, HyperCommGrid]] = None
     kv_format: str = "sbhd"
+
+    def __post_init__(self):
+        if not self.module_to_grid_map:
+            return
+        # Local import avoids circular imports at dataclass-module import time.
+        from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY
+
+        expected_keys = set(self.modality_submodules_spec.keys()) | {MIMO_LANGUAGE_MODULE_KEY}
+        grid_keys = set(self.module_to_grid_map.keys())
+        if grid_keys != expected_keys:
+            raise ValueError(
+                f"module_to_grid_map keys must match modality module names + "
+                f"'{MIMO_LANGUAGE_MODULE_KEY}'. Missing: {expected_keys - grid_keys}, "
+                f"Extra: {grid_keys - expected_keys}"
+            )
diff --git a/megatron/core/models/mimo/config/role.py b/megatron/core/models/mimo/config/role.py
index 77c2512e8e6..411791f1e5c 100644
--- a/megatron/core/models/mimo/config/role.py
+++ b/megatron/core/models/mimo/config/role.py
@@ -5,7 +5,7 @@
 import logging
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import torch.distributed as dist
 
@@ -24,22 +24,17 @@ class ModuleLayout(Enum):
     Determines how modules are distributed across ranks and which
     forward path is used.
 
-    UNIFIED: No module_to_grid_map. All modules share same ranks and
-        parallelism. Uses the unified forward path (_forward_all_modules).
+    COLOCATED: All modules share the same ranks. Covers both legacy
+        (no grid map, global parallel_state) and heterogeneous TP/DP
+        (grid map with overlapping ranks). Uses _forward_all_modules.
 
     NON_COLOCATED: module_to_grid_map is set with non-overlapping rank
         ranges. Each rank runs EITHER encoder(s) OR the language model.
         Uses role-based dispatch with separate forward paths.
-
-    COLOCATED: (future) module_to_grid_map is set with overlapping rank
-        ranges. Encoder(s) and language model share ranks but have
-        different parallelism configs. Uses role-based dispatch but
-        allows both module types on the same rank.
     """
 
-    UNIFIED = "unified"
-    NON_COLOCATED = "non_colocated"
     COLOCATED = "colocated"
+    NON_COLOCATED = "non_colocated"
 
 
 @dataclass
@@ -70,50 +65,50 @@ class RankRole:
     """
 
     modules: Dict[str, ModuleStageInfo] = field(default_factory=dict)
-    mode: ModuleLayout = ModuleLayout.UNIFIED
+    mode: ModuleLayout = ModuleLayout.COLOCATED
+
+    @classmethod
+    def build(
+        cls,
+        modality_module_names: List[str],
+        module_to_grid_map: Optional[Dict[str, 'HyperCommGrid']] = None,
+    ) -> 'RankRole':
+        """Build a RankRole, dispatching by whether grids share ranks.
+
+        No grid map or all grids span the same ranks → COLOCATED.
+        Grids differ → NON_COLOCATED with PP-stage info per module.
+        """
+        if module_to_grid_map is None or cls._all_grids_colocated(module_to_grid_map):
+            return cls._colocated(modality_module_names)
+        return cls._from_grid_map(module_to_grid_map)
+
+    @staticmethod
+    def _all_grids_colocated(module_to_grid_map: Dict[str, 'HyperCommGrid']) -> bool:
+        grids = list(module_to_grid_map.values())
+        first = grids[0]
+        return all(g.rank_offset == first.rank_offset and g.size == first.size for g in grids[1:])
 
     @classmethod
-    def unified(cls, module_names: List[str]) -> 'RankRole':
-        """Create a role for the unified case: every module, first+last stage."""
+    def _colocated(cls, modality_module_names: List[str]) -> 'RankRole':
+        """Colocated layout: every module on every rank, PP=1."""
+        all_module_names = list(modality_module_names) + [MIMO_LANGUAGE_MODULE_KEY]
         return cls(
             modules={
                 name: ModuleStageInfo(is_first_stage=True, is_last_stage=True)
-                for name in module_names
+                for name in all_module_names
             },
-            mode=ModuleLayout.UNIFIED,
+            mode=ModuleLayout.COLOCATED,
         )
 
     @classmethod
-    def from_grid_map(
-        cls, module_to_grid_map: Dict[str, HyperCommGrid], modality_module_names: List[str]
-    ) -> 'RankRole':
-        """Create a role from a module-to-grid mapping for non-colocated PP.
-
-        Determines which modules the current rank participates in and its
-        pipeline stage position within each module.
+    def _from_grid_map(cls, module_to_grid_map: Dict[str, HyperCommGrid]) -> 'RankRole':
+        """Non-colocated role for this rank from a module-to-grid mapping.
 
-        Args:
-            module_to_grid_map: Dict mapping module names to HyperCommGrid objects.
-                Must contain keys matching modality_module_names + MIMO_LANGUAGE_MODULE_KEY.
-            modality_module_names: List of modality module names (e.g., ["images", "audio"]).
-
-        Returns:
-            RankRole for the current rank.
+        Grid map keys are validated by ``MimoModelConfig.__post_init__``.
 
         Raises:
-            ValueError: If grid map keys don't match expected module names.
             RuntimeError: If current rank is not in any module grid.
         """
-        # Validate keys
-        expected_keys = set(modality_module_names) | {MIMO_LANGUAGE_MODULE_KEY}
-        grid_keys = set(module_to_grid_map.keys())
-        if grid_keys != expected_keys:
-            raise ValueError(
-                f"module_to_grid_map keys must match modality module names + "
-                f"'{MIMO_LANGUAGE_MODULE_KEY}'. Missing: {expected_keys - grid_keys}, "
-                f"Extra: {grid_keys - expected_keys}"
-            )
-
         current_rank = dist.get_rank()
         modules = {}
 
@@ -131,7 +126,7 @@ def from_grid_map(
             is_first = pp_rank == 0
             is_last = pp_rank == pp_size - 1
             logger.info(
-                f"[RankRole.from_grid_map] Rank {current_rank}: module={module_name}, "
+                f"[RankRole._from_grid_map] Rank {current_rank}: module={module_name}, "
                 f"pp_rank={pp_rank}/{pp_size}, is_first_stage={is_first}, is_last_stage={is_last}"
             )
             modules[module_name] = ModuleStageInfo(is_first_stage=is_first, is_last_stage=is_last)
diff --git a/megatron/core/models/mimo/model/base.py b/megatron/core/models/mimo/model/base.py
index b1c12f521c3..bdfe4289dd0 100644
--- a/megatron/core/models/mimo/model/base.py
+++ b/megatron/core/models/mimo/model/base.py
@@ -7,6 +7,7 @@
 import torch
 
 from megatron.core.distributed import DistributedDataParallel
+from megatron.core.models.mimo.comm.colocated_communicator import ColocatedBridgeCommunicator
 from megatron.core.models.mimo.config import MimoModelConfig
 from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY, ModuleLayout, RankRole
 from megatron.core.models.mimo.partition.utils import PartitionAdapter, PartitionConfig
@@ -59,10 +60,12 @@ def __init__(self, mimo_config: MimoModelConfig, cp_group=None, tp_group=None) -
 
         self.mimo_config = mimo_config
         modality_names = list(mimo_config.modality_submodules_spec.keys())
-        if mimo_config.module_to_grid_map:
-            self.role = RankRole.from_grid_map(mimo_config.module_to_grid_map, modality_names)
-        else:
-            self.role = RankRole.unified(modality_names + [MIMO_LANGUAGE_MODULE_KEY])
+        self.colocated_comms = {}
+        self.role = RankRole.build(modality_names, mimo_config.module_to_grid_map)
+        if self.role.mode is ModuleLayout.COLOCATED and mimo_config.module_to_grid_map:
+            # Per-encoder bridge needed iff modules share ranks but may differ
+            # in TP/DP within those ranks.
+            self._build_colocated_communicators()
 
         # Use special token IDs from the config
         self.special_token_ids = (
@@ -358,7 +361,7 @@ def forward(
         # Get any tensors passed via set_input_tensor
         input_tensors = getattr(self, 'input_tensors', None)
 
-        if self.role.mode == ModuleLayout.UNIFIED:
+        if self.role.mode == ModuleLayout.COLOCATED:
             return self._forward_all_modules(
                 input_ids,
                 position_ids,
@@ -491,6 +494,47 @@ def _forward_language_module(
 
         return lm_output
 
+    def _build_colocated_communicators(self):
+        grid_map = self.mimo_config.module_to_grid_map
+        if any(
+            'tp' not in grid.dim_names or 'dp' not in grid.dim_names for grid in grid_map.values()
+        ):
+            logger.info(
+                "Skipping colocated communicator setup because module_to_grid_map "
+                "does not define TP/DP topology for every module."
+            )
+            return
+
+        lang_key = MIMO_LANGUAGE_MODULE_KEY
+        lang_grid = grid_map[lang_key]
+        for mod_name in self.mimo_config.modality_submodules_spec:
+            if mod_name == lang_key:
+                continue
+            self.colocated_comms[(mod_name, lang_key)] = ColocatedBridgeCommunicator(
+                src_grid=grid_map[mod_name],
+                dest_grid=lang_grid,
+                src_module_name=mod_name,
+                dest_module_name=lang_key,
+                dim_mapping={'b': 0, 'h': 1},
+            )
+
+    def destroy(self) -> None:
+        """Release process groups owned by this MimoModel."""
+        for comm in self.colocated_comms.values():
+            comm.destroy()
+        self.colocated_comms.clear()
+
+    def _apply_colocated_comms(self, modality_embeddings):
+        """Transform encoder embeddings from encoder TP/DP to LLM TP/DP layout."""
+        lang_key = MIMO_LANGUAGE_MODULE_KEY
+        for modality_name in list(modality_embeddings.keys()):
+            comm = self.colocated_comms.get((modality_name, lang_key))
+            if comm is not None:
+                modality_embeddings[modality_name] = comm.communicate(
+                    modality_embeddings[modality_name]
+                )
+        return modality_embeddings
+
     def _forward_all_modules(
         self,
         input_ids: torch.Tensor,
@@ -533,6 +577,10 @@ def _forward_all_modules(
                         f"Generated embeddings for {modality_name} with shape {embeddings.shape}"
                     )
 
+        # Apply colocated communication if configured (no-op when colocated_comms is empty)
+        if self.colocated_comms:
+            modality_embeddings = self._apply_colocated_comms(modality_embeddings)
+
         # Get text embeddings
         text_embeddings = self.get_text_embeddings(input_ids, position_ids, self.special_token_ids)
         logger.debug(f"Generated text embeddings with shape {text_embeddings.shape}")
diff --git a/tests/unit_tests/models/test_mimo_1f1b_schedule.py b/tests/unit_tests/models/test_mimo_1f1b_schedule.py
index 44be0c7911e..836382b21cc 100644
--- a/tests/unit_tests/models/test_mimo_1f1b_schedule.py
+++ b/tests/unit_tests/models/test_mimo_1f1b_schedule.py
@@ -60,6 +60,27 @@
 _embedding_pg_cache: dict = {}
 
 
+def build_no_sync_func(mimo_model):
+    """Build a no_sync_func that stacks DDP no_sync over each sub-module.
+
+    Shared by 1F1B pipeline tests and colocated-correctness tests — both need
+    DDP's gradient sync disabled during microbatches and resumed via the
+    schedule's finalize_grads_func.
+    """
+
+    @contextmanager
+    def no_sync_func():
+        with ExitStack() as stack:
+            if mimo_model.language_model is not None:
+                stack.enter_context(mimo_model.language_model.no_sync())
+            for submodule in mimo_model.modality_submodules.values():
+                if submodule is not None:
+                    stack.enter_context(submodule.no_sync())
+            yield
+
+    return no_sync_func
+
+
 def create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1):
     """Create a HyperCommGrid with specified parallelism."""
     grid = HyperCommGrid(
@@ -183,13 +204,40 @@ def is_rank_in_grid(grid):
 
 
 def get_language_model_spec(
-    num_layers, hidden_size, num_attention_heads, vocab_size, seq_len, pg_collection
+    num_layers,
+    hidden_size,
+    num_attention_heads,
+    vocab_size,
+    seq_len,
+    pg_collection,
+    bf16=True,
+    bias=True,
+    dropout=True,
+    per_token_loss=False,
 ):
-    """Get the language model spec."""
+    """Get the language model spec.
+
+    ``bf16=False`` switches pipeline dtype and autocast to fp32. Correctness
+    tests also pass ``bias=False, dropout=False`` to remove bias-update and
+    stochastic noise from the cross-config diff signal.
+
+    ``per_token_loss=True`` sets ``calculate_per_token_loss=True`` on the
+    TransformerConfig, which pins DDP's gradient_scaling_factor to 1.0
+    (pure SUM reduction). Callers that flip this must supply a 3-tuple
+    loss_func and drive the external divide in their finalize hook.
+    """
     pp_rank = dist.get_rank(pg_collection.pp)
     pp_size = dist.get_world_size(pg_collection.pp)
     tp_size = pg_collection.tp.size() if pg_collection.tp is not None else 1
 
+    pipeline_dtype = torch.bfloat16 if bf16 else torch.float32
+    extra_kwargs = {}
+    if not bias:
+        extra_kwargs['add_bias_linear'] = False
+    if not dropout:
+        extra_kwargs['attention_dropout'] = 0.0
+        extra_kwargs['hidden_dropout'] = 0.0
+
     lm_config = TransformerConfig(
         num_layers=num_layers,
         hidden_size=hidden_size,
@@ -199,10 +247,12 @@ def get_language_model_spec(
         moe_token_dispatcher_type='alltoall',
         tensor_model_parallel_size=tp_size,
         pipeline_model_parallel_size=pp_size,
-        pipeline_dtype=torch.bfloat16,
-        bf16=True,
+        pipeline_dtype=pipeline_dtype,
+        bf16=bf16,
         cross_entropy_loss_fusion=True,
         cross_entropy_fusion_impl='te',
+        calculate_per_token_loss=per_token_loss,
+        **extra_kwargs,
     )
     return ModuleSpec(
         module=GPTModel,
@@ -218,12 +268,12 @@ def get_language_model_spec(
     )
 
 
-def get_projection_config(hidden_size):
+def get_projection_config(hidden_size, bias=True):
     """Return a TransformerConfig for the vision projection MLP."""
     cfg = TransformerConfig(num_layers=1, hidden_size=hidden_size, num_attention_heads=1)
     cfg.ffn_hidden_size = hidden_size
-    cfg.bias_activation_fusion = True
-    cfg.add_bias_linear = True
+    cfg.bias_activation_fusion = bool(bias)
+    cfg.add_bias_linear = bool(bias)
     cfg.activation_func = torch.nn.functional.gelu
     return cfg
 
@@ -239,15 +289,38 @@ def get_projection_layer_spec():
 
 
 def get_vision_submodules_spec(
-    num_layers, hidden_size, num_attention_heads, language_hidden_size, pg_collection
+    num_layers,
+    hidden_size,
+    num_attention_heads,
+    language_hidden_size,
+    pg_collection,
+    bf16=True,
+    bias=True,
+    dropout=True,
+    per_token_loss=False,
 ):
-    """Get the submodule spec for the vision modality."""
+    """Get the submodule spec for the vision modality.
+
+    ``bias=False`` / ``dropout=False`` mirror the LM-spec kwargs for
+    correctness tests. ``per_token_loss=True`` sets
+    ``calculate_per_token_loss=True`` on the encoder's TransformerConfig so
+    the encoder DDP also pure-SUMs across DP (needed for the heterogeneous-DP
+    colocated path).
+    """
     from megatron.core.transformer.transformer_block import TransformerBlock
 
     tp_size = pg_collection.tp.size() if pg_collection.tp is not None else 1
     pp_size = pg_collection.pp.size() if pg_collection.pp is not None else 1
     pp_rank = dist.get_rank(pg_collection.pp)
 
+    pipeline_dtype = torch.bfloat16 if bf16 else torch.float32
+    extra_kwargs = {}
+    if not bias:
+        extra_kwargs['add_bias_linear'] = False
+    if not dropout:
+        extra_kwargs['attention_dropout'] = 0.0
+        extra_kwargs['hidden_dropout'] = 0.0
+
     vision_config = TransformerConfig(
         num_layers=num_layers,
         hidden_size=hidden_size,
@@ -257,8 +330,10 @@ def get_vision_submodules_spec(
         moe_token_dispatcher_type='alltoall',
         tensor_model_parallel_size=tp_size,
         pipeline_model_parallel_size=pp_size,
-        pipeline_dtype=torch.bfloat16,
-        bf16=True,
+        pipeline_dtype=pipeline_dtype,
+        bf16=bf16,
+        calculate_per_token_loss=per_token_loss,
+        **extra_kwargs,
     )
     vision_encoder_spec = ModuleSpec(
         module=TransformerBlock,
@@ -274,7 +349,7 @@ def get_vision_submodules_spec(
     vision_projection_spec = ModuleSpec(
         module=MultimodalProjector,
         params={
-            "config": get_projection_config(hidden_size=language_hidden_size),
+            "config": get_projection_config(hidden_size=language_hidden_size, bias=bias),
             "submodules": get_projection_layer_spec().submodules,
             "projector_type": "mlp",
             "input_size": vision_config.hidden_size,
@@ -293,9 +368,38 @@ def get_vision_submodules_spec(
 
 
 def get_mimo_model(
-    encoder_name, encoder_grid, llm_grid, hidden_size, num_layers, vocab_size, seq_len
+    encoder_name,
+    encoder_grid,
+    llm_grid,
+    hidden_size,
+    num_layers,
+    vocab_size,
+    seq_len,
+    ddp_config=None,
+    bf16=True,
+    bias=True,
+    dropout=True,
+    per_token_loss=False,
 ):
-    """Create MIMO model with TransformerBlock encoder and GPTModel LLM."""
+    """Create MIMO model with TransformerBlock encoder and GPTModel LLM.
+
+    Args:
+        ddp_config: Optional override for the Megatron DDP config. Default
+            matches the 1F1B schedule tests' config.
+        bf16: If True (default) build the model in bf16; if False build in
+            fp32 end-to-end for deterministic numerics in correctness tests.
+        bias: If False, disable ``add_bias_linear`` in LM/vision configs and
+            the projection MLP — removes bias-update noise from diffs.
+        dropout: If False, force attention/hidden dropout to 0.0.
+        per_token_loss: If True, set ``calculate_per_token_loss=True`` on
+            both sub-model configs. This pins the encoder and LLM DDP
+            gradient_scaling_factor to 1.0 (pure SUM across DP). The caller
+            MUST supply a 3-tuple loss_func ``(sum_loss, num_tokens,
+            log_dict)`` and a custom ``finalize_model_grads_func`` that
+            divides grads by the correct global divisor on both sides;
+            hetero-DP callers use this to land ``1/B_full`` on both encoder
+            and LLM without relying on the per-DDP built-in scaling.
+    """
     language_pg = get_pg_collection_with_embedding_groups(llm_grid, is_language_model=True)
     vision_pg = get_pg_collection_with_embedding_groups(encoder_grid, is_language_model=False)
 
@@ -306,6 +410,10 @@ def get_mimo_model(
         vocab_size=vocab_size,
         seq_len=seq_len,
         pg_collection=language_pg,
+        bf16=bf16,
+        bias=bias,
+        dropout=dropout,
+        per_token_loss=per_token_loss,
     )
     vision_submodule_spec = get_vision_submodules_spec(
         num_layers=num_layers,
@@ -313,6 +421,10 @@ def get_mimo_model(
         num_attention_heads=8,
         language_hidden_size=hidden_size,
         pg_collection=vision_pg,
+        bf16=bf16,
+        bias=bias,
+        dropout=dropout,
+        per_token_loss=per_token_loss,
     )
 
     module_to_grid_map = {encoder_name: encoder_grid, MIMO_LANGUAGE_MODULE_KEY: llm_grid}
@@ -326,12 +438,15 @@ def get_mimo_model(
     )
 
     mimo_model = MimoModel(mimo_config)
-    mimo_model.to(torch.device("cuda")).to(torch.bfloat16)
-
-    # Wrap with DDP
-    ddp_config = DistributedDataParallelConfig(
-        overlap_grad_reduce=True, bucket_size=10000, use_distributed_optimizer=True
-    )
+    mimo_model.to(torch.device("cuda"))
+    if bf16:
+        mimo_model.to(torch.bfloat16)
+
+    # Wrap with DDP (caller may override e.g. for heterogeneous-DP scaling).
+    if ddp_config is None:
+        ddp_config = DistributedDataParallelConfig(
+            overlap_grad_reduce=True, bucket_size=10000, use_distributed_optimizer=True
+        )
 
     if mimo_model.language_model is not None:
         mimo_model.language_model = DistributedDataParallel(
@@ -485,16 +600,7 @@ def run_mimo_1f1b_test(
         seq_len=seq_length,
     )
 
-    # Build schedule functions using pre-created pg_collections (no leaks)
-    @contextmanager
-    def no_sync_func():
-        with ExitStack() as stack:
-            if mimo_model.language_model is not None:
-                stack.enter_context(mimo_model.language_model.no_sync())
-            for submodule in mimo_model.modality_submodules.values():
-                if submodule is not None:
-                    stack.enter_context(submodule.no_sync())
-            yield
+    no_sync_func = build_no_sync_func(mimo_model)
 
     def finalize_grads_func(*args, **kwargs):
         if mimo_model.language_model is not None:
@@ -595,30 +701,35 @@ def loss_func(loss_mask, output_tensor):
 
     optimizer.zero_grad()
 
-    losses = schedule.forward_backward_pipelining_without_interleaving(
-        forward_step_func=step_func,
-        data_iterator=data_iterator,
-        model=[mimo_model],
-        num_microbatches=num_microbatches,
-        seq_length=seq_length,
-        micro_batch_size=micro_batch_size,
-        forward_only=False,
-        p2p_communicator=communicator,
-        pg_collection=pg_collection,
-    )
-
-    # Optimizer step with global gradient clipping
-    success, grad_norm, num_zeros = optimizer.step()
-    assert success, "Optimizer step failed"
-    assert grad_norm is not None and grad_norm > 0, f"Expected positive grad norm, got {grad_norm}"
-
-    # Verify results on last LLM stage
-    if is_rank_in_grid(llm_grid) and is_pp_last_stage(llm_grid.get_pg("pp")):
-        assert len(losses) > 0, "Expected losses on last LLM stage"
-        for loss_dict in losses:
-            assert 'loss_reduced' in loss_dict
+    try:
+        losses = schedule.forward_backward_pipelining_without_interleaving(
+            forward_step_func=step_func,
+            data_iterator=data_iterator,
+            model=[mimo_model],
+            num_microbatches=num_microbatches,
+            seq_length=seq_length,
+            micro_batch_size=micro_batch_size,
+            forward_only=False,
+            p2p_communicator=communicator,
+            pg_collection=pg_collection,
+        )
 
-    return losses
+        # Optimizer step with global gradient clipping
+        success, grad_norm, num_zeros = optimizer.step()
+        assert success, "Optimizer step failed"
+        assert (
+            grad_norm is not None and grad_norm > 0
+        ), f"Expected positive grad norm, got {grad_norm}"
+
+        # Verify results on last LLM stage
+        if is_rank_in_grid(llm_grid) and is_pp_last_stage(llm_grid.get_pg("pp")):
+            assert len(losses) > 0, "Expected losses on last LLM stage"
+            for loss_dict in losses:
+                assert 'loss_reduced' in loss_dict
+
+        return losses
+    finally:
+        mimo_model.destroy()
 
 
 # ============================================================================
diff --git a/tests/unit_tests/models/test_mimo_colocated_communicator.py b/tests/unit_tests/models/test_mimo_colocated_communicator.py
new file mode 100644
index 00000000000..67cee551a0f
--- /dev/null
+++ b/tests/unit_tests/models/test_mimo_colocated_communicator.py
@@ -0,0 +1,543 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+import logging
+import os
+import sys
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from megatron.core.hyper_comm_grid import HyperCommGrid
+from megatron.core.models.mimo.comm.colocated_communicator import ColocatedBridgeCommunicator
+
+logging.basicConfig(level=logging.DEBUG, stream=sys.stderr)
+
+_active_grids: list = []
+_active_comms: list = []
+
+
+def create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1):
+    grid = HyperCommGrid(
+        shape=[tp, cp, pp, dp],
+        dim_names=["tp", "cp", "pp", "dp"],
+        rank_offset=offset,
+        backend="nccl",
+    )
+    grid.create_pg(["tp"])
+    grid.create_pg(["cp"])
+    grid.create_pg(["pp"])
+    grid.create_pg(["dp"])
+    _active_grids.append(grid)
+    return grid
+
+
+def make_comm(*args, **kwargs):
+    comm = ColocatedBridgeCommunicator(*args, **kwargs)
+    _active_comms.append(comm)
+    return comm
+
+
+def destroy_all_grids():
+    # Destroy communicators first so their NCCL subgroups are freed before we
+    # tear down the parent grids. NCCL caps concurrent communicators at ~500;
+    # leaked PGs from per-test fixtures blow that budget quickly.
+    for comm in _active_comms:
+        comm.destroy()
+    _active_comms.clear()
+    for grid in _active_grids:
+        grid.destroy()
+    _active_grids.clear()
+
+
+# ── Test 1: Rank mappings ──────────────────────────────────────────────────────
+
+
+class TestRankMappings:
+
+    @classmethod
+    def setup_class(cls):
+        if not dist.is_initialized():
+            dist.init_process_group(backend="nccl")
+        if torch.cuda.is_available():
+            torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
+
+    def teardown_method(self):
+        destroy_all_grids()
+
+    @pytest.mark.parametrize(
+        "src_tp, src_dp, dest_tp, dest_dp, expected_src_pos, expected_dest_pos",
+        [
+            # Fan-in: TP2/DP4 → TP4/DP2
+            (
+                2,
+                4,
+                4,
+                2,
+                {
+                    0: (0, 0),
+                    1: (0, 1),
+                    2: (1, 0),
+                    3: (1, 1),
+                    4: (2, 0),
+                    5: (2, 1),
+                    6: (3, 0),
+                    7: (3, 1),
+                },
+                {
+                    0: (0, 0),
+                    1: (0, 1),
+                    2: (0, 2),
+                    3: (0, 3),
+                    4: (1, 0),
+                    5: (1, 1),
+                    6: (1, 2),
+                    7: (1, 3),
+                },
+            ),
+            # Fan-out: TP4/DP2 → TP2/DP4
+            (
+                4,
+                2,
+                2,
+                4,
+                {
+                    0: (0, 0),
+                    1: (0, 1),
+                    2: (0, 2),
+                    3: (0, 3),
+                    4: (1, 0),
+                    5: (1, 1),
+                    6: (1, 2),
+                    7: (1, 3),
+                },
+                {
+                    0: (0, 0),
+                    1: (0, 1),
+                    2: (1, 0),
+                    3: (1, 1),
+                    4: (2, 0),
+                    5: (2, 1),
+                    6: (3, 0),
+                    7: (3, 1),
+                },
+            ),
+        ],
+        ids=["fan_in", "fan_out"],
+    )
+    def test_rank_mappings(
+        self, src_tp, src_dp, dest_tp, dest_dp, expected_src_pos, expected_dest_pos
+    ):
+        src_grid = create_hypercomm_grid(tp=src_tp, dp=src_dp)
+        dest_grid = create_hypercomm_grid(tp=dest_tp, dp=dest_dp)
+        comm = make_comm(src_grid, dest_grid)
+
+        assert comm.rank_to_src_pos == expected_src_pos
+        assert comm.rank_to_dest_pos == expected_dest_pos
+
+    def test_rank_mappings_with_rank_offset(self):
+        # 4-rank grids at offset=4 (covering ranks 4-7). Exercises the
+        # rank_offset propagation that previously only ran with offset=0.
+        if dist.get_world_size() < 8:
+            pytest.skip("requires at least 8 ranks")
+        src_grid = create_hypercomm_grid(offset=4, tp=2, dp=2)
+        dest_grid = create_hypercomm_grid(offset=4, tp=1, dp=4)
+        comm = make_comm(src_grid, dest_grid)
+
+        assert comm.rank_to_src_pos == {4: (0, 0), 5: (0, 1), 6: (1, 0), 7: (1, 1)}
+        assert comm.rank_to_dest_pos == {4: (0, 0), 5: (1, 0), 6: (2, 0), 7: (3, 0)}
+
+
+# ── Test 2: All-gather groups ──────────────────────────────────────────────────
+
+
+class TestAllGatherGroups:
+
+    @classmethod
+    def setup_class(cls):
+        if not dist.is_initialized():
+            dist.init_process_group(backend="nccl")
+        if torch.cuda.is_available():
+            torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
+
+    def teardown_method(self):
+        destroy_all_grids()
+
+    def test_fan_in_all_gather_groups(self):
+        # Fan-in TP2/DP4 → TP4/DP2. Groups are keyed (dest_dp_idx, src_tp_idx)
+        # and members must appear in src_dp_idx order so all_gather_into_tensor
+        # concatenates in slot order on the backward path.
+        src_grid = create_hypercomm_grid(tp=2, dp=4)
+        dest_grid = create_hypercomm_grid(tp=4, dp=2)
+        comm = make_comm(src_grid, dest_grid)
+
+        assert comm.gather_group_ranks == [[0, 2], [1, 3], [4, 6], [5, 7]]
+        assert comm.gather_pg is not None
+
+    def test_fan_out_gather_groups(self):
+        # Fan-out TP4/DP2 → TP2/DP4. Groups are keyed (src_dp_idx, dest_tp_idx);
+        # membership order must track dest_dp_idx so the backward all-gather
+        # reconstructs the full-batch gradient in the correct layout.
+        src_grid = create_hypercomm_grid(tp=4, dp=2)
+        dest_grid = create_hypercomm_grid(tp=2, dp=4)
+        comm = make_comm(src_grid, dest_grid)
+
+        assert comm.gather_group_ranks == [[0, 2], [1, 3], [4, 6], [5, 7]]
+        assert comm.gather_pg is not None
+
+
+# ── Test 3b: _validate_grids negative tests ───────────────────────────────────
+
+
+class TestValidateGrids:
+    """One negative test per raise path in ColocatedBridgeCommunicator._validate_grids.
+
+    Each case builds a pair of grids that violates exactly one invariant and
+    asserts that the constructor raises ValueError.
+    """
+
+    @classmethod
+    def setup_class(cls):
+        if not dist.is_initialized():
+            dist.init_process_group(backend="nccl")
+        if torch.cuda.is_available():
+            torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
+
+    def teardown_method(self):
+        destroy_all_grids()
+
+    def _grid_missing_tp(self, offset=0, dp=1):
+        # Build a grid without a 'tp' dim to exercise the "missing 'tp'" raise.
+        grid = HyperCommGrid(shape=[dp], dim_names=["dp"], rank_offset=offset, backend="nccl")
+        grid.create_pg(["dp"])
+        _active_grids.append(grid)
+        return grid
+
+    def test_missing_tp_dim(self):
+        src_grid = self._grid_missing_tp(dp=8)
+        dest_grid = create_hypercomm_grid(tp=4, dp=2)
+        with pytest.raises(ValueError, match="must have 'tp' dimension"):
+            make_comm(src_grid, dest_grid)
+
+    def test_size_mismatch(self):
+        src_grid = create_hypercomm_grid(tp=2, dp=4)  # 8 ranks
+        dest_grid = create_hypercomm_grid(offset=4, tp=2, dp=2)  # 4 ranks
+        with pytest.raises(ValueError, match="span same number of ranks"):
+            make_comm(src_grid, dest_grid)
+
+    def test_rank_offset_mismatch(self):
+        src_grid = create_hypercomm_grid(offset=0, tp=2, dp=2)
+        dest_grid = create_hypercomm_grid(offset=4, tp=2, dp=2)
+        with pytest.raises(ValueError, match="same rank offset"):
+            make_comm(src_grid, dest_grid)
+
+    @pytest.mark.parametrize(
+        "side,dim,expected",
+        [
+            ("src", "pp", "src PP must be 1"),
+            ("dest", "pp", "dest PP must be 1"),
+            ("src", "cp", "CP must be 1"),
+        ],
+    )
+    def test_pp_or_cp_gt_one_rejected(self, side, dim, expected):
+        bad = {dim: 2, "tp": 2, "dp": 2}
+        good = {"tp": 4, "dp": 2}
+        if side == "src":
+            src_grid = create_hypercomm_grid(**bad)
+            dest_grid = create_hypercomm_grid(**good)
+        else:
+            src_grid = create_hypercomm_grid(**good)
+            dest_grid = create_hypercomm_grid(**bad)
+        with pytest.raises(ValueError, match=expected):
+            make_comm(src_grid, dest_grid)
+
+    def test_dp_not_divisible(self):
+        # 6-rank grids with DP sizes (3 vs 2) that neither divides the other.
+        # Fits inside an 8-rank world (HyperCommGrid enforces size <= world - offset).
+        if dist.get_world_size() < 6:
+            pytest.skip("requires at least 6 ranks")
+        src_grid = HyperCommGrid(
+            shape=[2, 1, 1, 3], dim_names=["tp", "cp", "pp", "dp"], backend="nccl"
+        )
+        dest_grid = HyperCommGrid(
+            shape=[3, 1, 1, 2], dim_names=["tp", "cp", "pp", "dp"], backend="nccl"
+        )
+        for g in (src_grid, dest_grid):
+            _active_grids.append(g)
+        with pytest.raises(ValueError, match="evenly divisible"):
+            make_comm(src_grid, dest_grid)
+
+
+# ── Test 3c: communicate() runtime preconditions ──────────────────────────────
+
+
+class TestCommunicatePreconditions:
+    """Runtime-input checks enforced by ``communicate()``."""
+
+    @classmethod
+    def setup_class(cls):
+        if not dist.is_initialized():
+            dist.init_process_group(backend="nccl")
+        if torch.cuda.is_available():
+            torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
+
+    def teardown_method(self):
+        destroy_all_grids()
+
+    def test_non_divisible_batch_raises_fan_out(self):
+        # Fan-out: dest_dp=4, src_dp=2 → scale=2. Pass a batch dim of size 3
+        # so 3 % 2 != 0 and the forward communicate() raises before slicing.
+        src_grid = create_hypercomm_grid(tp=4, dp=2)
+        dest_grid = create_hypercomm_grid(tp=2, dp=4)
+        comm = make_comm(src_grid, dest_grid, dim_mapping={'b': 0, 'h': 1})
+        tensor = torch.zeros(3, 8, device='cuda')
+        with pytest.raises(ValueError, match="not divisible by fan_out"):
+            comm.communicate(tensor)
+
+    def test_non_divisible_batch_raises_fan_in_backward_narrow(self):
+        # Fan-in forward all-gathers (no slice), so the forward path never
+        # divides. The backward path narrows the post-gather output via
+        # get_slice_info, which raises on a non-divisible size. Call
+        # get_slice_info directly with an odd size to exercise that path.
+        src_grid = create_hypercomm_grid(tp=2, dp=4)
+        dest_grid = create_hypercomm_grid(tp=4, dp=2)
+        comm = make_comm(src_grid, dest_grid)
+        with pytest.raises(ValueError, match="not divisible by fan_in"):
+            comm.get_slice_info(batch_size=3)
+
+
+# ── Test 3d: destroy() releases PGs ──────────────────────────────────────────
+
+
+class TestDestroy:
+    """``destroy()`` must null out both PG attributes."""
+
+    @classmethod
+    def setup_class(cls):
+        if not dist.is_initialized():
+            dist.init_process_group(backend="nccl")
+        if torch.cuda.is_available():
+            torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
+
+    def teardown_method(self):
+        destroy_all_grids()
+
+    def test_destroy_releases_fan_in_pg(self):
+        src_grid = create_hypercomm_grid(tp=2, dp=4)
+        dest_grid = create_hypercomm_grid(tp=4, dp=2)
+        # Don't track via make_comm — destroy() is exactly what we're testing.
+        comm = ColocatedBridgeCommunicator(src_grid, dest_grid)
+        assert comm.gather_pg is not None
+        comm.destroy()
+        assert comm.gather_pg is None
+
+    def test_destroy_releases_fan_out_pg(self):
+        src_grid = create_hypercomm_grid(tp=4, dp=2)
+        dest_grid = create_hypercomm_grid(tp=2, dp=4)
+        comm = ColocatedBridgeCommunicator(src_grid, dest_grid)
+        assert comm.gather_pg is not None
+        comm.destroy()
+        assert comm.gather_pg is None
+
+    def test_destroy_is_idempotent(self):
+        # Calling destroy twice must not raise — leftover test fixtures often
+        # double-destroy during exception cleanup.
+        src_grid = create_hypercomm_grid(tp=2, dp=4)
+        dest_grid = create_hypercomm_grid(tp=4, dp=2)
+        comm = ColocatedBridgeCommunicator(src_grid, dest_grid)
+        comm.destroy()
+        comm.destroy()
+
+
+# ── Test 3e: Bridge gradient correctness (bitwise exact) ─────────────────────
+
+
+def _shape_for_dim_mapping(dim_mapping, B, S, H):
+    s = [0, 0, 0]
+    s[dim_mapping['b']] = B
+    s[dim_mapping['s']] = S
+    s[dim_mapping['h']] = H
+    return s
+
+
+# Parametrize dim_mapping for the fan-in tests (tests 1 & 2 per AXIOM spec).
+_DIM_MAPPINGS = [{'s': 0, 'b': 1, 'h': 2}, {'b': 0, 's': 1, 'h': 2}]
+_DIM_MAPPING_IDS = ["sbh", "bsh"]
+
+
+class TestBridgeGradients:
+    """Bitwise-exact gradient tests for ``ColocatedBridgeCommunicator``.
+
+    This class is **intentionally distinct** from the model-level correctness
+    tests in ``test_mimo_colocated_correctness.py`` (see PR review comment 19).
+    The bridge forward and backward are pure data
+    movement (``narrow`` / ``all_gather_into_tensor``) with no FP compute, so
+    the mathematical adjoint relationship can — and should — be asserted at
+    ``rtol=0, atol=0``:
+
+        * fan-in forward == ``torch.cat`` of sibling inputs in slot order
+        * fan-in backward == ``grad_output.narrow`` at this rank's slot
+        * fan-out forward == ``input.narrow`` at this rank's slot
+        * fan-out backward == ``cat`` of every sibling's grad (catches
+          zero-pad-without-gather, wrong slot order, double-counting,
+          missing siblings — the four failure modes of the adjoint)
+        * equal-DP is a pure identity (forward + backward)
+
+    The MimoModel-level tests validate the full training stack including GEMM
+    reduction order and DDP scaling, and can only assert approximate FP32
+    closeness. These tests localise the bridge's own invariants and fail
+    first when one of them regresses.
+    """
+
+    S = 8
+    B_PER_RANK = 2
+    H = 128
+
+    @classmethod
+    def setup_class(cls):
+        if not dist.is_initialized():
+            dist.init_process_group(backend="nccl")
+        if torch.cuda.is_available():
+            torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
+
+    def teardown_method(self):
+        destroy_all_grids()
+
+    # ── Test 1: fan-in forward = torch.cat of sibling inputs ─────────────────
+    @pytest.mark.parametrize("src_tp,src_dp,dest_tp,dest_dp", [(2, 4, 4, 2)], ids=["2x_fan_in"])
+    @pytest.mark.parametrize("dim_mapping", _DIM_MAPPINGS, ids=_DIM_MAPPING_IDS)
+    def test_fan_in_forward_equals_torch_cat(self, src_tp, src_dp, dest_tp, dest_dp, dim_mapping):
+        src_grid = create_hypercomm_grid(tp=src_tp, dp=src_dp)
+        dest_grid = create_hypercomm_grid(tp=dest_tp, dp=dest_dp)
+        comm = make_comm(src_grid, dest_grid, dim_mapping=dim_mapping)
+
+        rank = dist.get_rank()
+        shape = _shape_for_dim_mapping(dim_mapping, self.B_PER_RANK, self.S, self.H)
+
+        # Distinct inputs per rank so the cat reveals ordering bugs.
+        torch.manual_seed(1000 + rank)
+        local_input = torch.randn(*shape, device='cuda')
+
+        actual = comm.communicate(local_input)
+
+        # Expected: manual all_gather over the communicator's fan-in group,
+        # then cat along batch_dim. all_gather preserves group-local-rank
+        # order, which is the same order the communicator uses.
+        group = comm.gather_pg
+        gathered = [torch.empty_like(local_input) for _ in range(dist.get_world_size(group))]
+        dist.all_gather(gathered, local_input, group=group)
+        expected = torch.cat(gathered, dim=dim_mapping['b'])
+
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+    # ── Test 2: fan-in backward = grad_output.narrow for this rank's slot ────
+    @pytest.mark.parametrize("src_tp,src_dp,dest_tp,dest_dp", [(2, 4, 4, 2)], ids=["2x_fan_in"])
+    @pytest.mark.parametrize("dim_mapping", _DIM_MAPPINGS, ids=_DIM_MAPPING_IDS)
+    def test_fan_in_backward_equals_narrow(self, src_tp, src_dp, dest_tp, dest_dp, dim_mapping):
+        src_grid = create_hypercomm_grid(tp=src_tp, dp=src_dp)
+        dest_grid = create_hypercomm_grid(tp=dest_tp, dp=dest_dp)
+        comm = make_comm(src_grid, dest_grid, dim_mapping=dim_mapping)
+
+        rank = dist.get_rank()
+        batch_dim = dim_mapping['b']
+        b_local = self.B_PER_RANK
+        shape = _shape_for_dim_mapping(dim_mapping, b_local, self.S, self.H)
+
+        torch.manual_seed(1000 + rank)
+        local_input = torch.randn(*shape, device='cuda', requires_grad=True)
+        out = comm.communicate(local_input)
+
+        # grad_output is TP-replicated within the dest DP group: seed the same
+        # on every rank so every rank in the fan-in group backward-narrows the
+        # same upstream gradient. out shape is identical across group members,
+        # so seeded randn produces the same tensor on each.
+        torch.manual_seed(42)
+        grad_output = torch.randn_like(out)
+        out.backward(grad_output)
+
+        slot = comm.rank_to_src_pos[rank][0] % comm.scale
+        expected = grad_output.narrow(batch_dim, slot * b_local, b_local).contiguous()
+        torch.testing.assert_close(local_input.grad, expected, rtol=0, atol=0)
+
+    # ── Test 3: fan-out forward = input.narrow for this rank's slot ─────────
+    @pytest.mark.parametrize("src_tp,src_dp,dest_tp,dest_dp", [(4, 2, 2, 4)], ids=["2x_fan_out"])
+    def test_fan_out_forward_equals_narrow(self, src_tp, src_dp, dest_tp, dest_dp):
+        dim_mapping = {'b': 0, 's': 1, 'h': 2}
+        src_grid = create_hypercomm_grid(tp=src_tp, dp=src_dp)
+        dest_grid = create_hypercomm_grid(tp=dest_tp, dp=dest_dp)
+        comm = make_comm(src_grid, dest_grid, dim_mapping=dim_mapping)
+
+        rank = dist.get_rank()
+        batch_dim = dim_mapping['b']
+        b_per_dest = self.B_PER_RANK
+        b_full = b_per_dest * comm.scale
+        shape = _shape_for_dim_mapping(dim_mapping, b_full, self.S, self.H)
+
+        # Input is TP-replicated on the batch dim (bridge contract). Seed
+        # identically across all ranks to satisfy it.
+        torch.manual_seed(42)
+        input_tensor = torch.randn(*shape, device='cuda')
+
+        actual = comm.communicate(input_tensor)
+
+        slot = comm.rank_to_dest_pos[rank][0] % comm.scale
+        expected = input_tensor.narrow(batch_dim, slot * b_per_dest, b_per_dest).contiguous()
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+    # ── Test 4 (CRITICAL): fan-out backward = concat of all sibling grads ──
+    @pytest.mark.parametrize("src_tp,src_dp,dest_tp,dest_dp", [(4, 2, 2, 4)], ids=["2x_fan_out"])
+    def test_fan_out_backward_equals_concat_of_sibling_grads(
+        self, src_tp, src_dp, dest_tp, dest_dp
+    ):
+        """Fan-out backward must all-gather sibling grads in slot order.
+
+        Catches four distinct regressions with a single assertion:
+          * zero-pad-without-gather (other slots would be zero),
+          * wrong slot order (values would be scrambled),
+          * double-counting (values would be multiplied),
+          * missing siblings (shape or zeros would diverge).
+        """
+        dim_mapping = {'b': 0, 's': 1, 'h': 2}
+        src_grid = create_hypercomm_grid(tp=src_tp, dp=src_dp)
+        dest_grid = create_hypercomm_grid(tp=dest_tp, dp=dest_dp)
+        comm = make_comm(src_grid, dest_grid, dim_mapping=dim_mapping)
+
+        rank = dist.get_rank()
+        batch_dim = dim_mapping['b']
+        scale = comm.scale
+        b_per_dest = self.B_PER_RANK
+        b_full = b_per_dest * scale
+        shape = _shape_for_dim_mapping(dim_mapping, b_full, self.S, self.H)
+
+        torch.manual_seed(42)  # identical input across ranks (TP-replicated)
+        input_tensor = torch.randn(*shape, device='cuda', requires_grad=True)
+        out = comm.communicate(input_tensor)  # narrowed to (b_per_dest, S, H)
+
+        # Distinct grad per slot so the cat reveals both membership and order.
+        slot = comm.rank_to_dest_pos[rank][0] % scale
+        grad_output = (slot + 1) * torch.ones_like(out)
+        out.backward(grad_output)
+
+        slot_shape = _shape_for_dim_mapping(dim_mapping, b_per_dest, self.S, self.H)
+        expected = torch.cat(
+            [(i + 1) * torch.ones(*slot_shape, device='cuda') for i in range(scale)], dim=batch_dim
+        )
+        torch.testing.assert_close(input_tensor.grad, expected, rtol=0, atol=0)
+
+    # ── Test 5: equal DP is a pure identity forward and backward ────────────
+    @pytest.mark.parametrize("src_tp,src_dp,dest_tp,dest_dp", [(4, 2, 4, 2)], ids=["tp4_dp2"])
+    def test_equal_dp_is_bitwise_identity_fwd_and_bwd(self, src_tp, src_dp, dest_tp, dest_dp):
+        dim_mapping = {'b': 0, 's': 1, 'h': 2}
+        src_grid = create_hypercomm_grid(tp=src_tp, dp=src_dp)
+        dest_grid = create_hypercomm_grid(tp=dest_tp, dp=dest_dp)
+        comm = make_comm(src_grid, dest_grid, dim_mapping=dim_mapping)
+
+        shape = _shape_for_dim_mapping(dim_mapping, self.B_PER_RANK, self.S, self.H)
+        torch.manual_seed(1000 + dist.get_rank())
+        x = torch.randn(*shape, device='cuda', requires_grad=True)
+
+        out = comm.communicate(x)
+        torch.testing.assert_close(out, x, rtol=0, atol=0)
+
+        grad_output = torch.randn_like(x)
+        out.backward(grad_output)
+        torch.testing.assert_close(x.grad, grad_output, rtol=0, atol=0)
diff --git a/tests/unit_tests/models/test_mimo_colocated_correctness.py b/tests/unit_tests/models/test_mimo_colocated_correctness.py
new file mode 100644
index 00000000000..e2d91bdf83e
--- /dev/null
+++ b/tests/unit_tests/models/test_mimo_colocated_correctness.py
@@ -0,0 +1,1183 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+"""Gradient-scaling correctness for colocated MimoModel under heterogeneous DP.
+
+Verifies that a heterogeneous-DP MimoModel produces the same post-step
+encoder weights as an **equal-DP** reference built on the SAME encoder
+TP/DP layout as the dist model (so the bridge is the identity
+passthrough — ``BridgeDirection.EQUAL`` in
+``ColocatedBridgeCommunicator``). Under correct grad scaling, both
+configs yield the DP=1 gradient on every encoder shard, so the Adam
+update lands on identical values and the sharded post-step weights
+compare directly.
+
+Why an equal-DP reference is the right oracle:
+  * Encoder sharding matches exactly — ref and dist both use
+    ``enc_tp=dist_enc_tp, enc_dp=dist_enc_dp``. Shards line up 1:1,
+    so there is no gather-and-slice in the weight comparison and no
+    TP=1-vs-TP>1 accumulation-order drift to contend with.
+  * ``enc_dp == llm_dp`` on the ref side → the bridge is identity and
+    every encoder rank feeds its colocated LLM rank with no
+    redistribution collective.
+  * Both sides set ``calculate_per_token_loss=True`` on their
+    TransformerConfigs, which pins DDP's ``gradient_scaling_factor=1.0``
+    — pure SUM across DP. The custom
+    ``finalize_grads_func`` in ``_wire_training_hooks`` all-reduces
+    ``total_num_tokens`` over the LLM DP group, then calls
+    ``scale_gradients(1/N_global)`` on both encoder and LLM. This lands
+    the true global per-token mean on every shard without touching
+    ``DistributedDataParallel``.
+
+LLM TP differs between ref (``llm_tp=dist_enc_tp``) and dist
+(``llm_tp=dist_llm_tp``), so ref's LLM weights are copied into dist via
+all-gather-across-ref-TP + slice-for-dist-TP. The LLM forward then
+diverges numerically by fp32 TP accumulation order, but the aggregate
+gradient that flows back into the encoder remains the DP=1 gradient in
+both models, which is what the post-step encoder weight oracle checks.
+The test runs in fp32 with ``add_bias_linear=False`` and dropout
+disabled to minimize non-bridge numerical noise — this keeps the
+post-bridge hidden states bit-exact and surfaces only TP-shape drift
+in the logits oracle.
+
+If the heterogeneous-DP scaling is wrong (e.g. dividing by encoder_dp
+when it should be 1, or letting either DDP apply its default ``1/dp_size``
+on top of the per-token mean already delivered by the finalize hook),
+the dist encoder's post-step weights diverge from the ref encoder's
+weights — a single Adam step is enough to detect.
+
+Run with::
+
+    uv run python -m torch.distributed.run --nproc_per_node=8 \\
+        -m pytest tests/unit_tests/models/test_mimo_colocated_correctness.py -v -s
+"""
+
+import os
+from functools import partial
+
+import pytest
+import torch
+import torch.distributed as dist
+from packaging import version
+
+import megatron.core.pipeline_parallel.schedules as schedule
+from megatron.core.distributed import DistributedDataParallelConfig
+from megatron.core.distributed.finalize_model_grads import finalize_model_grads
+from megatron.core.models.mimo.optimizer import get_mimo_optimizer
+from megatron.core.optimizer.optimizer_config import OptimizerConfig
+from megatron.core.transformer.enums import ModelType
+from megatron.core.utils import unwrap_model
+from tests.unit_tests.models.test_mimo_1f1b_schedule import (
+    build_no_sync_func,
+    create_all_embedding_groups,
+    create_hypercomm_grid,
+    destroy_all_grids,
+    get_mimo_model,
+)
+from tests.unit_tests.test_utilities import Utils
+
+
+def loss_func(loss_mask, output_tensor):
+    """Per-token-loss 3-tuple: raw local sum + local valid-token count.
+
+    Returns ``(local_sum, local_num_tokens, log_dict)`` — the contract the
+    schedule expects when ``calculate_per_token_loss=True`` is set on the
+    TransformerConfig. No ``1/num_tokens`` or ``1/num_microbatches``
+    division is applied here; the schedule skips the per-microbatch
+    division (see ``schedules.py:270-274``) and aggregates ``num_tokens``
+    across microbatches for the finalize step.
+
+    Paired with ``mimo_finalize_grads_func`` below, which all-reduces
+    ``total_num_tokens`` over the LLM DP group to obtain ``N_global`` and
+    then divides both encoder and LLM grads by ``1/N_global`` directly via
+    ``scale_gradients`` — landing the true global per-token mean on every
+    shard without touching DDP.
+
+    ``output_tensor`` is per-token CE from
+    ``GPTModel.compute_language_model_loss`` with shape ``[b, s]``.
+    """
+    if output_tensor is None:
+        zero_loss = torch.tensor(0.0, device='cuda', requires_grad=True)
+        zero_count = torch.tensor(0, device='cuda', dtype=torch.int)
+        return zero_loss, zero_count, {'loss_reduced': 0.0}
+
+    masked = output_tensor.float() * loss_mask.float()
+    local_sum = masked.sum()
+    local_num_tokens = loss_mask.float().sum().to(torch.int)
+    return local_sum, local_num_tokens, {'loss_reduced': local_sum.detach().item()}
+
+
+def forward_step(data_iterator, model, encoder_grid, llm_grid, encoder_name):
+    """Forward step with per-rank data slicing for heterogeneous DP."""
+    batch = next(data_iterator) if data_iterator is not None else {'input_ids': None}
+
+    if batch.get('input_ids') is None:
+        output_tensor, loss_mask = model(**batch)
+        return output_tensor, partial(loss_func, loss_mask)
+
+    encoder_dp = encoder_grid.get_pg("dp").size()
+    llm_dp = llm_grid.get_pg("dp").size()
+
+    if encoder_dp > llm_dp:
+        # Fan-in: input was pre-sliced to LLM-DP (larger per-rank batch).
+        # Narrow modality_inputs to the encoder's smaller per-rank slice.
+        scale = encoder_dp // llm_dp
+        encoder_dp_idx = encoder_grid.get_pg("dp").rank()
+        slot = encoder_dp_idx % scale
+
+        if 'modality_inputs' in batch and batch['modality_inputs'] is not None:
+            for mod_name, mod_data in batch['modality_inputs'].items():
+                for enc_name, enc_data in mod_data.items():
+                    for key, tensor in enc_data.items():
+                        if tensor is not None and isinstance(tensor, torch.Tensor):
+                            batch_size = tensor.shape[1]  # [seq, batch, hidden]
+                            slice_size = batch_size // scale
+                            start = slot * slice_size
+                            enc_data[key] = tensor[:, start : start + slice_size, :].contiguous()
+
+    elif llm_dp > encoder_dp:
+        # Fan-out: input was pre-sliced to encoder-DP (larger per-rank batch).
+        # Narrow the LLM-side tensors to this LLM-DP rank's slice.
+        scale = llm_dp // encoder_dp
+        llm_dp_idx = llm_grid.get_pg("dp").rank()
+        slot = llm_dp_idx % scale
+
+        batch_size = batch['input_ids'].shape[0]
+        slice_size = batch_size // scale
+        start = slot * slice_size
+
+        for key in ['input_ids', 'labels', 'loss_mask', 'position_ids']:
+            if key in batch and batch[key] is not None:
+                batch[key] = batch[key][start : start + slice_size].contiguous()
+
+    output_tensor, loss_mask = model(**batch)
+    return output_tensor, partial(loss_func, loss_mask)
+
+
+def _set_deterministic_env():
+    for k, v in {
+        "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0",
+        "CUDA_DEVICE_MAX_CONNECTIONS": "1",
+        "CUBLAS_WORKSPACE_CONFIG": ":4096:8",
+    }.items():
+        os.environ[k] = v
+    os.environ.pop('NVTE_FLASH_ATTN', None)
+    os.environ.pop('NVTE_FUSED_ATTN', None)
+    os.environ.pop('NVTE_UNFUSED_ATTN', None)
+
+
+def _wire_training_hooks(mimo_model, language_pg, vision_pg):
+    """Attach no_sync / finalize_grads / grad_scale hooks to a MimoModel.
+
+    The finalize hook implements the heterogeneous-DP grad-scaling story
+    without touching ``DistributedDataParallel``. Both sub-model configs
+    set ``calculate_per_token_loss=True``, so both DDPs pure-SUM across
+    their own DP group (``gradient_scaling_factor=1.0``). After backward
+    and DDP reduce, every rank's ``main_grad`` holds the un-normalized
+    full-batch sum of per-token gradients.
+
+    This hook then:
+      1. all-reduces the schedule's ``total_num_tokens`` across the LLM
+         DP group to obtain ``N_global`` (total valid tokens in the global
+         batch). Since ranks are colocated, every rank now knows
+         ``N_global``.
+      2. Calls ``finalize_model_grads(num_tokens=None)`` per side — runs
+         the usual DDP grad finish + layernorm/embedding AR work without
+         letting the built-in divisor path fire.
+      3. Calls ``scale_gradients(1/N_global)`` on each side — lands the
+         true global per-token mean uniformly on encoder and LLM grads.
+
+    Note: encoder has no loss_func (so nothing emits a per-encoder-DP
+    ``num_tokens`` to feed ``finalize_model_grads``' internal all-reduce).
+    Doing the all-reduce once ourselves and calling ``scale_gradients``
+    directly avoids engineering a fictitious per-encoder-rank count whose
+    sum happens to equal ``N_global``.
+    """
+
+    no_sync_func = build_no_sync_func(mimo_model)
+
+    def finalize_grads_func(model_list, num_tokens, force_all_reduce=False, **kwargs):
+        # Schedule passes the per-rank sum-across-microbatches of what the
+        # loss_func returned. Because loss_func runs only on the LLM side,
+        # this is the LLM-local token count.
+        assert num_tokens is not None, (
+            "finalize_grads_func expects calculate_per_token_loss=True on the "
+            "TransformerConfig so the schedule forwards total_num_tokens; got None."
+        )
+
+        # Phase 1: lift the all-reduce. After this, every rank (including
+        # encoder-only replicas) has N_global = total non-padded tokens in
+        # the global batch.
+        llm_dp_pg = language_pg.dp_cp if language_pg.dp_cp is not None else language_pg.dp
+        dist.all_reduce(num_tokens, group=llm_dp_pg, op=dist.ReduceOp.SUM)
+        n_global = num_tokens.item()
+
+        # Phase 2: per-side DDP finish without built-in num_tokens scaling.
+        # Forward ``force_all_reduce`` so PP grad-sync semantics (if ever
+        # exercised here) aren't silently dropped.
+        if mimo_model.language_model is not None:
+            finalize_model_grads(
+                [mimo_model.language_model],
+                num_tokens=None,
+                pg_collection=language_pg,
+                force_all_reduce=force_all_reduce,
+            )
+        for submodule in mimo_model.modality_submodules.values():
+            if submodule is not None:
+                finalize_model_grads(
+                    [submodule],
+                    num_tokens=None,
+                    pg_collection=vision_pg,
+                    force_all_reduce=force_all_reduce,
+                )
+
+        # Phase 3: uniform divide by N_global. Guard div-by-zero for the
+        # degenerate fully-masked batch.
+        if n_global > 0:
+            inv = 1.0 / n_global
+            if mimo_model.language_model is not None:
+                mimo_model.language_model.scale_gradients(inv)
+            for submodule in mimo_model.modality_submodules.values():
+                if submodule is not None:
+                    submodule.scale_gradients(inv)
+
+    mimo_model.config.no_sync_func = no_sync_func
+    mimo_model.config.finalize_model_grads_func = finalize_grads_func
+    mimo_model.config.grad_scale_func = lambda loss: (
+        torch.tensor(loss, dtype=torch.float32, device='cuda', requires_grad=True)
+        if isinstance(loss, (int, float))
+        else loss
+    )
+
+
+def _generate_and_broadcast_global_batches(
+    global_mbs,
+    seq_length,
+    hidden_size,
+    vocab_size,
+    encoder_name,
+    num_batches,
+    image_token_id=50257,
+    mask_pattern="uniform",
+):
+    """Generate global batches on rank 0 and broadcast so every rank sees
+    identical data. Dist pre-slices per rank; ref consumes the full batch.
+
+    ``mask_pattern``:
+      * ``"uniform"`` — every sample has the same valid-token count (image
+        tokens masked, text tokens all valid). Local/global denominators
+        coincide up to DP-rank partitioning.
+      * ``"asymmetric"`` — each sample zeros out an additional sample-
+        dependent number of trailing text tokens, so different samples
+        (and therefore different DP-rank slices) carry different valid-
+        token counts. This exercises the num+den global-mean CE path
+        where the old local-mean recipe would be only approximately
+        correct.
+    """
+    if mask_pattern not in ("uniform", "asymmetric"):
+        raise ValueError(f"Unknown mask_pattern: {mask_pattern!r}")
+
+    rank = dist.get_rank()
+    image_seq_length = seq_length // 2
+    batches = []
+
+    for batch_idx in range(num_batches):
+        if rank == 0:
+            encoder_hidden_states = torch.randn(
+                image_seq_length, global_mbs, hidden_size, device='cuda', dtype=torch.float32
+            )
+            image_tokens = torch.full(
+                (global_mbs, image_seq_length), image_token_id, dtype=torch.long, device='cuda'
+            )
+            text_tokens = torch.randint(
+                1, vocab_size, (global_mbs, seq_length - image_seq_length), device='cuda'
+            )
+            input_ids = torch.cat([image_tokens, text_tokens], dim=1)
+        else:
+            encoder_hidden_states = torch.empty(
+                image_seq_length, global_mbs, hidden_size, device='cuda', dtype=torch.float32
+            )
+            input_ids = torch.empty(global_mbs, seq_length, dtype=torch.long, device='cuda')
+
+        dist.broadcast(encoder_hidden_states, src=0)
+        dist.broadcast(input_ids, src=0)
+
+        labels = input_ids.clone()
+        labels[input_ids == image_token_id] = -100
+        loss_mask = torch.ones(global_mbs, seq_length, device='cuda', dtype=torch.float32)
+        loss_mask[input_ids == image_token_id] = 0.0
+
+        if mask_pattern == "asymmetric":
+            # Zero out a sample-dependent trailing run of text tokens so
+            # each sample ends up with a different valid-token count.
+            # Counts are deterministic given (batch_idx, sample_idx) so the
+            # broadcast-on-rank-0 pattern is reproducible on every rank.
+            text_len = seq_length - image_seq_length
+            for sample_idx in range(global_mbs):
+                n_drop = ((batch_idx * 7 + sample_idx * 3) % (text_len - 1)) + 1
+                loss_mask[sample_idx, seq_length - n_drop :] = 0.0
+                labels[sample_idx, seq_length - n_drop :] = -100
+        position_ids = (
+            torch.arange(seq_length, device='cuda').unsqueeze(0).expand(global_mbs, -1).clone()
+        )
+
+        batches.append(
+            {
+                "input_ids": input_ids,
+                "labels": labels,
+                "loss_mask": loss_mask,
+                "position_ids": position_ids,
+                "modality_inputs": {
+                    encoder_name: {
+                        "clip_encoder": {
+                            'hidden_states': encoder_hidden_states,
+                            'attention_mask': None,
+                        }
+                    }
+                },
+            }
+        )
+
+    return batches
+
+
+def _slice_batch(global_batch, split_dp, split_rank):
+    """Return the ``split_rank``-th of ``split_dp`` slices along the batch dim."""
+    batch_dim = global_batch['input_ids'].shape[0]
+    slice_size = batch_dim // split_dp
+    start = split_rank * slice_size
+    end = start + slice_size
+
+    per_rank = {}
+    for key in ['input_ids', 'labels', 'loss_mask', 'position_ids']:
+        per_rank[key] = global_batch[key][start:end].contiguous()
+
+    mod_inputs_new = {}
+    for mod_name, mod_data in global_batch['modality_inputs'].items():
+        mod_inputs_new[mod_name] = {}
+        for enc_name, enc_data in mod_data.items():
+            mod_inputs_new[mod_name][enc_name] = {}
+            for key, tensor in enc_data.items():
+                if tensor is not None and isinstance(tensor, torch.Tensor):
+                    # modality hidden_states is [seq, batch, hidden] — slice dim 1
+                    mod_inputs_new[mod_name][enc_name][key] = tensor[:, start:end, :].contiguous()
+                else:
+                    mod_inputs_new[mod_name][enc_name][key] = tensor
+    per_rank['modality_inputs'] = mod_inputs_new
+    return per_rank
+
+
+def _slice_global_batch_for_dist(global_batch, encoder_grid, llm_grid):
+    """Pre-slice a global batch to the per-rank batch that ``forward_step`` expects.
+
+    ``forward_step`` assumes each rank already has its LLM-DP slice
+    (fan-in) or encoder-DP slice (fan-out); this helper performs that
+    slicing so both models can consume the same underlying global batch.
+    When ``enc_dp == llm_dp`` there is no fan-in/fan-out to pre-slice for
+    (``forward_step`` also skips slicing), and the full batch is returned.
+    """
+    enc_dp = encoder_grid.get_pg("dp").size()
+    llm_dp = llm_grid.get_pg("dp").size()
+
+    if enc_dp > llm_dp:
+        return _slice_batch(global_batch, llm_dp, llm_grid.get_pg("dp").rank())
+    if llm_dp > enc_dp:
+        return _slice_batch(global_batch, enc_dp, encoder_grid.get_pg("dp").rank())
+    return global_batch
+
+
+def _slice_global_batch_by_dp(global_batch, dp_pg):
+    """Slice a global batch along the batch dim by ``dp_pg`` rank.
+
+    For the equal-DP reference (``enc_dp == llm_dp``, bridge is identity),
+    each rank consumes 1/``dp_size`` of the global batch directly.
+    ``_slice_global_batch_for_dist`` returns the full batch in that case,
+    so this helper does the DP-rank split explicitly.
+    """
+    dp_size = dist.get_world_size(dp_pg)
+    if dp_size <= 1:
+        return global_batch
+    return _slice_batch(global_batch, dp_size, dist.get_rank(dp_pg))
+
+
+def _copy_ref_params_to_dist(ref_module, dist_module, ref_tp_group, dist_tp_group):
+    """Copy ref params into dist, handling differing TP shardings.
+
+    When ref and dist params have the same shape (same TP size and layout
+    at offset=0), shards align 1:1 and we copy directly. When shapes differ
+    (different TP sizes), we all-gather ref's shards across ``ref_tp_group``
+    to reconstruct the full weight, then slice by the dist ``partition_dim``
+    for this rank's dist TP shard.
+
+    Must be called **before** constructing the distributed optimizer, which
+    clones current param data into fp32 master weights at __init__.
+    """
+    ref_tp_size = dist.get_world_size(ref_tp_group)
+    dist_tp_rank = dist.get_rank(dist_tp_group)
+    dist_tp_size = dist.get_world_size(dist_tp_group)
+    ref_params = dict(ref_module.named_parameters())
+
+    with torch.no_grad():
+        for name, dist_param in dist_module.named_parameters():
+            assert name in ref_params, f"Param '{name}' in dist but not in ref"
+            ref_param = ref_params[name]
+            partition_dim = getattr(dist_param, 'partition_dim', -1)
+
+            if ref_param.shape == dist_param.shape:
+                # Same shard size (same TP layout or both replicated).
+                dist_param.data.copy_(ref_param.data.to(dist_param.dtype))
+                continue
+
+            assert partition_dim >= 0, (
+                f"Param '{name}': shapes differ "
+                f"(ref={tuple(ref_param.shape)}, dist={tuple(dist_param.shape)}) "
+                f"but partition_dim<0 — cannot reshard a replicated param."
+            )
+
+            # Different TP sizes: gather ref shards, then slice for dist.
+            shards = [torch.empty_like(ref_param.data) for _ in range(ref_tp_size)]
+            dist.all_gather(shards, ref_param.data.contiguous(), group=ref_tp_group)
+            full_weight = torch.cat(shards, dim=partition_dim)
+            dist_slice = torch.tensor_split(full_weight, dist_tp_size, dim=partition_dim)[
+                dist_tp_rank
+            ]
+
+            assert dist_slice.shape == dist_param.shape, (
+                f"Param '{name}': sliced.shape={tuple(dist_slice.shape)} != "
+                f"dist.shape={tuple(dist_param.shape)} "
+                f"(ref_tp={ref_tp_size}, dist_tp={dist_tp_size}, "
+                f"partition_dim={partition_dim})"
+            )
+            dist_param.data.copy_(dist_slice.to(dist_param.dtype))
+
+
+def _global_abs_diff_stats(a, b, pg=None):
+    """Absolute-diff stats plus reference-tensor magnitude stats, across ``pg``.
+
+    Reports both the abs-diff distribution AND the magnitude of ``b`` (the
+    reference tensor) so the caller can judge scale: a max abs-diff of 1.0
+    is catastrophic for values of O(1), but fine for values of O(100). The
+    relative-diff column (``rel_max = max_diff / ref_max``) gives a quick
+    percentage read.
+
+    Useful when the per-rank tensors cover different shards — all-reducing
+    MAX/MIN (and MAX of per-rank p95/p99 as a conservative worst-case) lets
+    rank 0 print a global view of drift across every shard in ``pg``. Mean
+    is SUM/world_size, which is the true global mean when every rank holds
+    the same number of elements (true here — shards have the same shape).
+    """
+    diff = (a.float() - b.float()).abs().flatten()
+    ref = b.float().abs().flatten()
+    n = diff.numel()
+
+    if n == 0:
+        zero = torch.tensor(0.0, device='cuda')
+        local_min = local_max = local_mean = local_p50 = local_p95 = local_p99 = zero
+        local_ref_max = local_ref_p95 = local_ref_mean = zero
+    else:
+        local_min = diff.min()
+        local_max = diff.max()
+        local_mean = diff.mean()
+        local_p50 = diff.quantile(0.50)
+        local_p95 = diff.quantile(0.95)
+        local_p99 = diff.quantile(0.99)
+        local_ref_max = ref.max()
+        local_ref_p95 = ref.quantile(0.95)
+        local_ref_mean = ref.mean()
+
+    world = dist.get_world_size(pg) if dist.is_initialized() else 1
+    if world > 1:
+        g_min = local_min.clone()
+        g_max = local_max.clone()
+        g_mean = local_mean.clone()
+        g_p50 = local_p50.clone()
+        g_p95 = local_p95.clone()
+        g_p99 = local_p99.clone()
+        g_ref_max = local_ref_max.clone()
+        g_ref_p95 = local_ref_p95.clone()
+        g_ref_mean = local_ref_mean.clone()
+        dist.all_reduce(g_min, op=dist.ReduceOp.MIN, group=pg)
+        dist.all_reduce(g_max, op=dist.ReduceOp.MAX, group=pg)
+        dist.all_reduce(g_mean, op=dist.ReduceOp.SUM, group=pg)
+        dist.all_reduce(g_p50, op=dist.ReduceOp.MAX, group=pg)
+        dist.all_reduce(g_p95, op=dist.ReduceOp.MAX, group=pg)
+        dist.all_reduce(g_p99, op=dist.ReduceOp.MAX, group=pg)
+        dist.all_reduce(g_ref_max, op=dist.ReduceOp.MAX, group=pg)
+        dist.all_reduce(g_ref_p95, op=dist.ReduceOp.MAX, group=pg)
+        dist.all_reduce(g_ref_mean, op=dist.ReduceOp.SUM, group=pg)
+        g_mean = g_mean / world
+        g_ref_mean = g_ref_mean / world
+        return {
+            'min': g_min.item(),
+            'max': g_max.item(),
+            'mean': g_mean.item(),
+            'p50_worst': g_p50.item(),
+            'p95_worst': g_p95.item(),
+            'p99_worst': g_p99.item(),
+            'ref_max': g_ref_max.item(),
+            'ref_p95': g_ref_p95.item(),
+            'ref_mean': g_ref_mean.item(),
+            'numel_per_rank': n,
+            'ranks': world,
+        }
+    return {
+        'min': local_min.item(),
+        'max': local_max.item(),
+        'mean': local_mean.item(),
+        'p50_worst': local_p50.item(),
+        'p95_worst': local_p95.item(),
+        'p99_worst': local_p99.item(),
+        'ref_max': local_ref_max.item(),
+        'ref_p95': local_ref_p95.item(),
+        'ref_mean': local_ref_mean.item(),
+        'numel_per_rank': n,
+        'ranks': 1,
+    }
+
+
+def _fmt_diff_stats(s):
+    ref_max = s.get('ref_max', 0.0)
+    rel_max = (s['max'] / ref_max) if ref_max > 0 else float('inf')
+    return (
+        f"min={s['min']:.2e} p50={s['p50_worst']:.2e} mean={s['mean']:.2e} "
+        f"p95={s['p95_worst']:.2e} p99={s['p99_worst']:.2e} "
+        f"max={s['max']:.2e} | ref_max={ref_max:.2e} ref_p95={s.get('ref_p95', 0.0):.2e} "
+        f"ref_mean={s.get('ref_mean', 0.0):.2e} rel_max={rel_max:.1%} "
+        f"(n_per_rank={s['numel_per_rank']}, ranks={s['ranks']})"
+    )
+
+
+def _print_from_rank0(msg):
+    if not dist.is_initialized() or dist.get_rank() == 0:
+        print(msg, flush=True)
+
+
+def _register_logits_capture(mimo_model):
+    """Forward hook on the LLM ``output_layer``; captures per-microbatch logits.
+
+    The hook runs on every microbatch forward. ``output`` from
+    ``ColumnParallelLinear`` is ``(logits, bias)`` with logits shape
+    ``[s, b, v/tp]`` — this rank's per-DP-slot, per-TP-vocab-shard slice
+    of the global logits tensor. Cloning so backward doesn't mutate.
+
+    Returns ``(captures, handle)``; caller must ``handle.remove()`` after
+    the schedule completes.
+    """
+    gpt = unwrap_model(mimo_model.language_model)
+    captures = []
+
+    def hook(_module, _inputs, output):
+        logits = output[0] if isinstance(output, tuple) else output
+        captures.append(logits.detach().clone())
+
+    handle = gpt.output_layer.register_forward_hook(hook)
+    return captures, handle
+
+
+def _register_llm_input_capture(mimo_model):
+    """Forward pre-hook on the GPT ``decoder``; captures post-bridge hidden states.
+
+    This is the activation entering the transformer block AFTER embedding
+    (skipped when MIMO passes ``decoder_input``) AND after the bridge has
+    moved the encoder output into the LLM's TP/DP layout. Shape is
+    ``[s, b_local, h_full]`` — hidden dim is NOT TP-sharded at this point.
+
+    Comparing dist vs ref at this capture isolates "does the bridge deliver
+    mathematically equivalent inputs to the LLM?" from downstream LLM TP
+    forward drift. If this oracle passes but ``llm_logits`` fails, the
+    divergence is inside the LLM TP forward; if this fails, the bridge
+    (fan_in/fan_out vs equal) is not equivalent.
+    """
+    gpt = unwrap_model(mimo_model.language_model)
+    captures = []
+
+    def pre_hook(_module, args, kwargs):
+        hidden = kwargs.get('hidden_states', None)
+        if hidden is None and args:
+            hidden = args[0]
+        if hidden is not None:
+            captures.append(hidden.detach().clone())
+
+    handle = gpt.decoder.register_forward_pre_hook(pre_hook, with_kwargs=True)
+    return captures, handle
+
+
+def _gather_bs_dp(local_tensor, llm_dp_pg):
+    """All-gather ``[s, b, h]`` across LLM DP along the batch dim."""
+    dp_size = dist.get_world_size(llm_dp_pg)
+    if dp_size <= 1:
+        return local_tensor.contiguous()
+    contig = local_tensor.contiguous()
+    shards = [torch.empty_like(contig) for _ in range(dp_size)]
+    dist.all_gather(shards, contig, group=llm_dp_pg)
+    return torch.cat(shards, dim=1)
+
+
+def _assert_llm_input_match(
+    ref_captures, dist_captures, ref_llm_grid, dist_llm_grid, rtol=1e-3, atol=1e-3
+):
+    """Post-bridge oracle: hidden states entering the LLM decoder match.
+
+    Hidden dim is not TP-sharded at the decoder input, so only DP-gather
+    across the LLM DP group is needed to reconstruct the full-batch tensor.
+    """
+    assert len(ref_captures) == len(dist_captures), (
+        f"Microbatch count mismatch: ref={len(ref_captures)}, " f"dist={len(dist_captures)}"
+    )
+    ref_dp_pg = ref_llm_grid.get_pg("dp")
+    dist_dp_pg = dist_llm_grid.get_pg("dp")
+
+    mismatches = []
+    for mbs_idx, (ref_local, dist_local) in enumerate(zip(ref_captures, dist_captures)):
+        ref_full = _gather_bs_dp(ref_local, ref_dp_pg)
+        dist_full = _gather_bs_dp(dist_local, dist_dp_pg)
+        assert ref_full.shape == dist_full.shape, (
+            f"mbs[{mbs_idx}]: gathered llm-input shape mismatch — "
+            f"ref={tuple(ref_full.shape)}, dist={tuple(dist_full.shape)}"
+        )
+        stats = _global_abs_diff_stats(dist_full, ref_full, pg=dist.group.WORLD)
+        _print_from_rank0(
+            f"[llm-input-diff] mbs[{mbs_idx}] shape={tuple(ref_full.shape)} "
+            f"{_fmt_diff_stats(stats)}"
+        )
+        try:
+            torch.testing.assert_close(dist_full, ref_full, rtol=rtol, atol=atol)
+        except AssertionError as e:
+            mismatches.append((mbs_idx, str(e)))
+
+    if mismatches:
+        rank = dist.get_rank()
+        details = "\n".join(f"  mbs[{i}]: {msg}" for i, msg in mismatches)
+        raise AssertionError(
+            f"Rank {rank}: llm-input diverged on {len(mismatches)} microbatch(es):\n" f"{details}"
+        )
+
+
+def _gather_logits_full_batch(local_logits, llm_tp_pg, llm_dp_pg):
+    """All-gather ``[s, b, v/tp]`` across LLM TP (vocab dim) then DP (batch dim).
+
+    Returns ``[s, b * dp_size, v]`` — the full global-batch logits,
+    identical on every rank of the LLM grid. Used to compare dist vs ref
+    on the same global slots regardless of how TP/DP slices them.
+    """
+    tp_size = dist.get_world_size(llm_tp_pg)
+    dp_size = dist.get_world_size(llm_dp_pg)
+
+    vocab_full = local_logits.contiguous()
+    if tp_size > 1:
+        shards = [torch.empty_like(vocab_full) for _ in range(tp_size)]
+        dist.all_gather(shards, vocab_full, group=llm_tp_pg)
+        vocab_full = torch.cat(shards, dim=-1)
+
+    batch_full = vocab_full.contiguous()
+    if dp_size > 1:
+        shards = [torch.empty_like(batch_full) for _ in range(dp_size)]
+        dist.all_gather(shards, batch_full, group=llm_dp_pg)
+        batch_full = torch.cat(shards, dim=1)
+
+    return batch_full
+
+
+def _assert_llm_logits_match(
+    ref_captures, dist_captures, ref_llm_grid, dist_llm_grid, rtol=1e-2, atol=1e-2
+):
+    """Logits oracle: TP+DP-gathered full-batch logits match microbatch-by-microbatch.
+
+    Dist and ref share the same global batch on every rank (broadcast from
+    rank 0), and with the HyperCommGrid layout both reconstruct global
+    batch rows 0..N in the same order after TP+DP all-gather (see
+    ``_slice_global_batch_*`` helpers for how the slicing lines up).
+    The only numerical difference between the two gathered logits is
+    fp32 accumulation order across a different LLM TP shape — hence the
+    loose ``rtol=atol=1e-2`` default.
+    """
+    assert len(ref_captures) == len(dist_captures), (
+        f"Microbatch count mismatch: ref={len(ref_captures)}, " f"dist={len(dist_captures)}"
+    )
+    ref_tp_pg = ref_llm_grid.get_pg("tp")
+    ref_dp_pg = ref_llm_grid.get_pg("dp")
+    dist_tp_pg = dist_llm_grid.get_pg("tp")
+    dist_dp_pg = dist_llm_grid.get_pg("dp")
+
+    mismatches = []
+    for mbs_idx, (ref_local, dist_local) in enumerate(zip(ref_captures, dist_captures)):
+        ref_full = _gather_logits_full_batch(ref_local, ref_tp_pg, ref_dp_pg)
+        dist_full = _gather_logits_full_batch(dist_local, dist_tp_pg, dist_dp_pg)
+        assert ref_full.shape == dist_full.shape, (
+            f"mbs[{mbs_idx}]: gathered logits shape mismatch — "
+            f"ref={tuple(ref_full.shape)}, dist={tuple(dist_full.shape)}"
+        )
+        # Gathered full-batch logits are identical on every LLM-grid rank,
+        # so stats at rank 0 represent the tensor globally — no reduction
+        # needed across other ranks.
+        stats = _global_abs_diff_stats(dist_full, ref_full, pg=dist.group.WORLD)
+        _print_from_rank0(
+            f"[logits-diff] mbs[{mbs_idx}] shape={tuple(ref_full.shape)} "
+            f"{_fmt_diff_stats(stats)}"
+        )
+        try:
+            torch.testing.assert_close(dist_full, ref_full, rtol=rtol, atol=atol)
+        except AssertionError as e:
+            mismatches.append((mbs_idx, str(e)))
+
+    if mismatches:
+        rank = dist.get_rank()
+        details = "\n".join(f"  mbs[{i}]: {msg}" for i, msg in mismatches)
+        raise AssertionError(
+            f"Rank {rank}: logits diverged on {len(mismatches)} microbatch(es):\n" f"{details}"
+        )
+
+
+def _snapshot_first_layer_encoder_grads(mimo_model, encoder_name):
+    """Clone ``param.main_grad`` for every ``.layers.0.`` encoder param.
+
+    ``main_grad`` holds the post-DDP-reduction gradient (reduced across
+    encoder DP), populated by the backward pass and consumed by
+    ``optimizer.step()``. Snapshot between backward and step so the values
+    aren't yet zeroed.
+    """
+    encoder = mimo_model.modality_submodules[encoder_name].module
+    snap = {}
+    for name, param in encoder.named_parameters():
+        if '.layers.0.' not in name:
+            continue
+        grad = getattr(param, 'main_grad', None)
+        if grad is None:
+            continue
+        snap[name] = grad.detach().clone()
+    return snap
+
+
+def _assert_first_layer_grads_match(ref_snap, dist_snap, rtol=1e-3, atol=1e-3):
+    """First-layer encoder grad oracle: shard-to-shard match between ref and dist.
+
+    Ref and dist use identical encoder TP/DP layout, so for every
+    ``layers.0.*`` encoder parameter their local shards line up 1:1.
+    Under correct grad scaling both main_grads equal the DP=1 gradient,
+    so the per-shard values must match within fp32 precision. Tighter
+    tolerances than the logits oracle are possible because the encoder
+    forward is identical on both sides — only the LLM TP layout differs,
+    and that noise enters via the gradient flowing back into the encoder.
+    """
+    assert set(ref_snap.keys()) == set(dist_snap.keys()), (
+        f"First-layer param name mismatch — "
+        f"ref-only: {set(ref_snap) - set(dist_snap)}, "
+        f"dist-only: {set(dist_snap) - set(ref_snap)}"
+    )
+    mismatches = []
+    for name in sorted(ref_snap):
+        ref_g = ref_snap[name]
+        dist_g = dist_snap[name]
+        assert ref_g.shape == dist_g.shape, (
+            f"Param '{name}': grad shape {tuple(ref_g.shape)} != "
+            f"{tuple(dist_g.shape)} — caller must match encoder TP."
+        )
+        # Every rank holds its own TP shard of this param; all-reduce
+        # across the full world so rank 0 prints the worst-case drift
+        # across all shards.
+        stats = _global_abs_diff_stats(dist_g, ref_g, pg=dist.group.WORLD)
+        _print_from_rank0(
+            f"[grad-diff] {name} shape={tuple(ref_g.shape)} " f"{_fmt_diff_stats(stats)}"
+        )
+        try:
+            torch.testing.assert_close(dist_g, ref_g, rtol=rtol, atol=atol)
+        except AssertionError as e:
+            mismatches.append((name, str(e)))
+
+    if mismatches:
+        rank = dist.get_rank()
+        details = "\n".join(f"  {n}: {msg}" for n, msg in mismatches)
+        raise AssertionError(
+            f"Rank {rank}: {len(mismatches)} first-layer encoder grad(s) "
+            f"diverged between dist and ref:\n{details}"
+        )
+
+
+def _assert_encoder_weights_match(ref_module, dist_module, rtol=1e-3, atol=1e-3):
+    """Assert every dist encoder shard matches the ref encoder shard.
+
+    Caller is responsible for ensuring ref and dist have the same encoder TP
+    layout (same ``enc_tp`` and ``enc_dp``), so each rank's shards line up
+    1:1 and can be compared directly. Under correct grad scaling and
+    identical initial state, one Adam step yields shard-wise equal post-step
+    weights — modulo fp32 TP accumulation-order drift from the LLM TP
+    layout differing between the two models.
+    """
+    ref_params = dict(ref_module.named_parameters())
+
+    mismatches = []
+    for name, dist_param in dist_module.named_parameters():
+        ref_param = ref_params[name]
+        assert ref_param.shape == dist_param.shape, (
+            f"Param '{name}': ref.shape={tuple(ref_param.shape)} != "
+            f"dist.shape={tuple(dist_param.shape)} — caller must match encoder TP."
+        )
+        stats = _global_abs_diff_stats(dist_param.data, ref_param.data, pg=dist.group.WORLD)
+        _print_from_rank0(
+            f"[weight-diff] {name} shape={tuple(ref_param.shape)} " f"{_fmt_diff_stats(stats)}"
+        )
+        try:
+            torch.testing.assert_close(dist_param.data, ref_param.data, rtol=rtol, atol=atol)
+        except AssertionError as e:
+            mismatches.append((name, str(e)))
+
+    if mismatches:
+        rank = dist.get_rank()
+        details = "\n".join(f"  {n}: {msg}" for n, msg in mismatches)
+        raise AssertionError(
+            f"Rank {rank}: {len(mismatches)} encoder param(s) diverged between "
+            f"heterogeneous-DP dist model and equal-DP reference:\n{details}"
+        )
+
+
+class _BatchIterator:
+    """Minimal iterator over a pre-generated list of batches."""
+
+    def __init__(self, batches):
+        self.batches = batches
+        self.idx = 0
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self.idx >= len(self.batches):
+            raise StopIteration
+        b = self.batches[self.idx]
+        self.idx += 1
+        return b
+
+
+def _run_forward_backward(
+    mimo_model,
+    batches,
+    enc_grid,
+    llm_grid,
+    encoder_name,
+    language_pg,
+    micro_batch_size,
+    seq_length,
+    num_microbatches,
+):
+    """One forward/backward pass through the mimo schedule."""
+    return schedule.forward_backward_no_pipelining(
+        forward_step_func=partial(
+            forward_step, encoder_grid=enc_grid, llm_grid=llm_grid, encoder_name=encoder_name
+        ),
+        data_iterator=_BatchIterator(batches),
+        model=[mimo_model],
+        num_microbatches=num_microbatches,
+        seq_length=seq_length,
+        micro_batch_size=micro_batch_size,
+        forward_only=False,
+        pg_collection=language_pg,
+    )
+
+
+class TestColocatedGradientScalingCorrectness:
+    """Verify heterogeneous-DP encoder grad scaling against an equal-DP reference.
+
+    The critical invariant: with ``calculate_per_token_loss=True`` on both
+    sub-model configs, DDP's ``gradient_scaling_factor`` is pinned to
+    1.0 and each side's DDP reduction is a pure SUM. The custom
+    ``finalize_grads_func`` then divides both encoder and LLM grads by
+    ``1/N_global`` (true global valid-token count), so the aggregate
+    gradient on every encoder shard equals the DP=1 per-token-mean
+    gradient. The reference uses the same encoder TP/DP as dist but with
+    ``enc_tp == llm_tp`` and ``enc_dp == llm_dp`` (identity bridge), so
+    after one Adam step the dist model's sharded weights match the ref
+    model's sharded weights within fp32 precision.
+
+    If the scaling were wrong (e.g., if either DDP applied its default
+    ``1/dp_size`` on top of the per-token mean, or if the custom finalize
+    used the encoder DP group's sum-of-local-counts instead of the
+    globally lifted ``N_global``), the encoder's reduced grad would be
+    skewed and post-step weights would diverge — a single optimizer step
+    is sufficient to detect.
+    """
+
+    @classmethod
+    def setup_class(cls):
+        Utils.initialize_distributed()
+        cls.world_size = dist.get_world_size()
+
+    @classmethod
+    def teardown_class(cls):
+        Utils.destroy_model_parallel()
+
+    def setup_method(self):
+        # Track MimoModels built by the test so teardown can release any
+        # ColocatedBridgeCommunicator subgroups before destroy_all_grids.
+        self._mimo_models = []
+
+    def teardown_method(self):
+        torch.use_deterministic_algorithms(False)
+        for model in self._mimo_models:
+            model.destroy()
+        self._mimo_models.clear()
+        destroy_all_grids()
+
+    @pytest.mark.skipif(
+        version.parse(torch.__version__) < version.parse("2.3.0"), reason="Requires PyTorch 2.3+"
+    )
+    @pytest.mark.parametrize(
+        "enc_tp,enc_dp,llm_tp,llm_dp", [(2, 4, 4, 2), (4, 2, 2, 4)], ids=["fan_in", "fan_out"]
+    )
+    @pytest.mark.parametrize(
+        "mask_pattern", ["uniform", "asymmetric"], ids=["uniform", "asymmetric"]
+    )
+    @pytest.mark.parametrize("num_microbatches", [1, 4], ids=["mbs1", "mbs4"])
+    def test_dist_matches_dp1_reference_post_step_weights(
+        self, enc_tp, enc_dp, llm_tp, llm_dp, mask_pattern, num_microbatches
+    ):
+        """Heterogeneous-DP dist post-step encoder weights match equal-DP reference.
+
+        Builds two MimoModels on every rank:
+
+        * Dist: the heterogeneous TP/DP config under test, with
+          ``calculate_per_token_loss=True`` + custom finalize hook that
+          pure-SUMs DDP and externally divides by ``N_global``.
+        * Ref: equal-DP uniform with ``enc_tp=dist_enc_tp``,
+          ``enc_dp=dist_enc_dp``, ``llm_tp=dist_enc_tp``,
+          ``llm_dp=dist_enc_dp`` — bridge is
+          ``BridgeDirection.EQUAL`` (identity passthrough), and the
+          encoder TP sharding matches dist's exactly so shards line up
+          1:1 for comparison.
+
+        Both models run the same finalize wiring; both DDPs pure-SUM
+        across their own DP group, then divide uniformly by ``N_global``.
+        LLM TP differs between the two models, which introduces fp32 TP
+        accumulation-order drift in the gradient flowing back to the
+        encoder but does not change the per-token-mean invariant that the
+        post-step encoder oracle checks.
+
+        Reference weights are copied into the distributed model so both
+        start from identical state. One Adam step later, the dist shards
+        should match the ref shards within fp32 precision.
+        """
+        if self.world_size != 8:
+            pytest.skip(f"Requires 8 GPUs, got {self.world_size}")
+
+        _set_deterministic_env()
+        torch.use_deterministic_algorithms(True)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+        encoder_name = "images"
+        hidden_size, seq_length, vocab_size = 256, 64, 1000
+        micro_batch_size = 2
+
+        # Global batch spans the larger DP side; dist pre-slices per rank
+        # before forward_step (which further slices encoder/LLM side).
+        global_batch_size = micro_batch_size * max(enc_dp, llm_dp)
+
+        # Grids: dist is heterogeneous; ref is equal-DP uniform matching
+        # dist's encoder so the bridge is identity and encoder shards
+        # align 1:1 for direct comparison.
+        dist_enc_grid = create_hypercomm_grid(offset=0, tp=enc_tp, cp=1, pp=1, dp=enc_dp)
+        dist_llm_grid = create_hypercomm_grid(offset=0, tp=llm_tp, cp=1, pp=1, dp=llm_dp)
+        ref_enc_grid = create_hypercomm_grid(offset=0, tp=enc_tp, cp=1, pp=1, dp=enc_dp)
+        ref_llm_grid = create_hypercomm_grid(offset=0, tp=enc_tp, cp=1, pp=1, dp=enc_dp)
+        create_all_embedding_groups([dist_enc_grid, dist_llm_grid, ref_enc_grid, ref_llm_grid])
+
+        # Both sub-model TransformerConfigs set calculate_per_token_loss=True
+        # (via per_token_loss=True on get_mimo_model), which pins DDP's
+        # gradient_scaling_factor to 1.0 — pure SUM across DP on both sides.
+        # Under the 3-tuple loss_func + custom finalize_grads_func in
+        # _wire_training_hooks, grads are divided uniformly by N_global,
+        # which is the true global per-token mean on every shard.
+        ddp_config = DistributedDataParallelConfig(
+            overlap_grad_reduce=True, bucket_size=10000, use_distributed_optimizer=True
+        )
+
+        # Build dist first (heterogeneous TP/DP).
+        torch.manual_seed(12345)
+        dist_mimo, _, _, dist_language_pg, dist_vision_pg = get_mimo_model(
+            encoder_name=encoder_name,
+            encoder_grid=dist_enc_grid,
+            llm_grid=dist_llm_grid,
+            hidden_size=hidden_size,
+            num_layers=2,
+            vocab_size=vocab_size,
+            seq_len=seq_length,
+            ddp_config=ddp_config,
+            bf16=False,
+            bias=False,
+            dropout=False,
+            per_token_loss=True,
+        )
+        dist_mimo.model_type = ModelType.encoder_or_decoder
+        self._mimo_models.append(dist_mimo)
+
+        # Reference with equal-DP uniform (enc_tp == llm_tp, enc_dp == llm_dp).
+        torch.manual_seed(12345)
+        ref_mimo, _, _, ref_language_pg, ref_vision_pg = get_mimo_model(
+            encoder_name=encoder_name,
+            encoder_grid=ref_enc_grid,
+            llm_grid=ref_llm_grid,
+            hidden_size=hidden_size,
+            num_layers=2,
+            vocab_size=vocab_size,
+            seq_len=seq_length,
+            ddp_config=ddp_config,
+            bf16=False,
+            bias=False,
+            dropout=False,
+            per_token_loss=True,
+        )
+        ref_mimo.model_type = ModelType.encoder_or_decoder
+        self._mimo_models.append(ref_mimo)
+
+        # Force identical initial state: encoder shards already match
+        # (same TP layout), so the helper copies shard-to-shard. LLM
+        # shards don't match (ref_llm_tp=enc_tp, dist_llm_tp=llm_tp), so
+        # the helper all-gathers ref's shards across ref's TP group and
+        # re-slices for dist's TP group.
+        _copy_ref_params_to_dist(
+            ref_mimo.modality_submodules[encoder_name].module,
+            dist_mimo.modality_submodules[encoder_name].module,
+            ref_enc_grid.get_pg("tp"),
+            dist_enc_grid.get_pg("tp"),
+        )
+        _copy_ref_params_to_dist(
+            ref_mimo.language_model.module,
+            dist_mimo.language_model.module,
+            ref_llm_grid.get_pg("tp"),
+            dist_llm_grid.get_pg("tp"),
+        )
+
+        _wire_training_hooks(dist_mimo, dist_language_pg, dist_vision_pg)
+        _wire_training_hooks(ref_mimo, ref_language_pg, ref_vision_pg)
+
+        # Distributed optimizers snapshot current param.data into fp32 master
+        # weights at __init__, so both must be built AFTER the ref-to-dist
+        # param copy above.
+        opt_config = OptimizerConfig(
+            optimizer='adam',
+            lr=1e-4,
+            weight_decay=0.01,
+            clip_grad=1.0,
+            bf16=False,
+            use_distributed_optimizer=True,
+        )
+        dist_optimizer = get_mimo_optimizer(dist_mimo, opt_config)
+        ref_optimizer = get_mimo_optimizer(ref_mimo, opt_config)
+
+        # Data: one deterministic global batch, identical on every rank.
+        torch.manual_seed(99999)
+        global_batches = _generate_and_broadcast_global_batches(
+            global_mbs=global_batch_size,
+            seq_length=seq_length,
+            hidden_size=hidden_size,
+            vocab_size=vocab_size,
+            encoder_name=encoder_name,
+            num_batches=num_microbatches,
+            mask_pattern=mask_pattern,
+        )
+        dist_batches = [
+            _slice_global_batch_for_dist(b, dist_enc_grid, dist_llm_grid) for b in global_batches
+        ]
+        # Ref is uniform (enc_dp == llm_dp), so _slice_global_batch_for_dist
+        # returns the full batch; slice explicitly by enc_dp so each rank
+        # sees the same per-rank batch size as dist's encoder does.
+        ref_batches = [
+            _slice_global_batch_by_dp(b, ref_enc_grid.get_pg("dp")) for b in global_batches
+        ]
+        ref_per_rank_batch_size = global_batch_size // enc_dp
+
+        # Logits capture: hook fires on every microbatch forward.
+        # Registered before forward/backward, removed right after so the
+        # hook doesn't leak across the second model's run.
+        dist_logits, dist_logits_hook = _register_logits_capture(dist_mimo)
+        ref_logits, ref_logits_hook = _register_logits_capture(ref_mimo)
+        dist_llm_input, dist_input_hook = _register_llm_input_capture(dist_mimo)
+        ref_llm_input, ref_input_hook = _register_llm_input_capture(ref_mimo)
+
+        try:
+            # One optimizer step on dist (heterogeneous forward_step slicing).
+            dist_optimizer.zero_grad()
+            _run_forward_backward(
+                mimo_model=dist_mimo,
+                batches=dist_batches,
+                enc_grid=dist_enc_grid,
+                llm_grid=dist_llm_grid,
+                encoder_name=encoder_name,
+                language_pg=dist_language_pg,
+                micro_batch_size=micro_batch_size,
+                seq_length=seq_length,
+                num_microbatches=num_microbatches,
+            )
+            # Snapshot encoder first-layer grads AFTER backward and BEFORE
+            # optimizer.step() consumes/zeros the grad buffer.
+            dist_first_layer_grads = _snapshot_first_layer_encoder_grads(dist_mimo, encoder_name)
+            dist_success, dist_grad_norm, _ = dist_optimizer.step()
+            assert dist_success, "Dist optimizer step failed"
+            assert dist_grad_norm is not None and dist_grad_norm > 0, (
+                f"Dist grad_norm={dist_grad_norm} — encoder grads may have been "
+                "silently zeroed by wrong scaling"
+            )
+
+            # One optimizer step on ref (enc_dp == llm_dp → forward_step skips slicing).
+            ref_optimizer.zero_grad()
+            _run_forward_backward(
+                mimo_model=ref_mimo,
+                batches=ref_batches,
+                enc_grid=ref_enc_grid,
+                llm_grid=ref_llm_grid,
+                encoder_name=encoder_name,
+                language_pg=ref_language_pg,
+                micro_batch_size=ref_per_rank_batch_size,
+                seq_length=seq_length,
+                num_microbatches=num_microbatches,
+            )
+            ref_first_layer_grads = _snapshot_first_layer_encoder_grads(ref_mimo, encoder_name)
+            ref_success, ref_grad_norm, _ = ref_optimizer.step()
+            assert ref_success, "Ref optimizer step failed"
+            assert ref_grad_norm is not None and ref_grad_norm > 0, f"Ref grad_norm={ref_grad_norm}"
+        finally:
+            dist_logits_hook.remove()
+            ref_logits_hook.remove()
+            dist_input_hook.remove()
+            ref_input_hook.remove()
+
+        # Run all three oracles regardless of individual failures so the
+        # diff-stats print covers every layer. Order: encoder weights /
+        # first-layer grads first (tightest — same encoder TP/DP layout
+        # → shards align 1:1), then LLM logits last (loosest — different
+        # LLM TP layout drives fp32 accumulation drift). Each oracle
+        # printed its own min/mean/p95/p99/max before its assertion ran,
+        # so the user sees the full drift distribution for every test.
+        failures = []
+
+        try:
+            _assert_encoder_weights_match(
+                ref_mimo.modality_submodules[encoder_name].module,
+                dist_mimo.modality_submodules[encoder_name].module,
+                rtol=1e-3,
+                atol=1e-3,
+            )
+        except AssertionError as e:
+            failures.append(('encoder_weights', str(e)))
+
+        try:
+            _assert_first_layer_grads_match(
+                ref_first_layer_grads, dist_first_layer_grads, rtol=1e-3, atol=1e-3
+            )
+        except AssertionError as e:
+            failures.append(('first_layer_grads', str(e)))
+
+        try:
+            _assert_llm_input_match(
+                ref_llm_input, dist_llm_input, ref_llm_grid, dist_llm_grid, rtol=1e-3, atol=1e-3
+            )
+        except AssertionError as e:
+            failures.append(('llm_input', str(e)))
+
+        try:
+            _assert_llm_logits_match(
+                ref_logits, dist_logits, ref_llm_grid, dist_llm_grid, rtol=1e-2, atol=1e-2
+            )
+        except AssertionError as e:
+            failures.append(('llm_logits', str(e)))
+
+        if failures:
+            summary = "\n\n".join(f"== {oracle} ==\n{msg}" for oracle, msg in failures)
+            raise AssertionError(f"{len(failures)} oracle(s) failed:\n{summary}")
diff --git a/tests/unit_tests/models/test_mimo_model.py b/tests/unit_tests/models/test_mimo_model.py
index e1c4b6e89bf..0ef62ff570f 100644
--- a/tests/unit_tests/models/test_mimo_model.py
+++ b/tests/unit_tests/models/test_mimo_model.py
@@ -528,15 +528,13 @@ def test_grid_validation_rejects_mismatched_keys(self):
             self.hidden_size, self.img_h, self.img_w, self.patch_dim
         )
 
-        mimo_config = MimoModelConfig(
-            language_model_spec=language_model_spec,
-            modality_submodules_spec={"images": vision_submodule_spec},
-            special_token_ids={"images": 50257},
-            module_to_grid_map={MIMO_LANGUAGE_MODULE_KEY: MockGrid()},
-        )
-
         with pytest.raises(ValueError, match="module_to_grid_map keys must match"):
-            MimoModel(mimo_config)
+            MimoModelConfig(
+                language_model_spec=language_model_spec,
+                modality_submodules_spec={"images": vision_submodule_spec},
+                special_token_ids={"images": 50257},
+                module_to_grid_map={MIMO_LANGUAGE_MODULE_KEY: MockGrid()},
+            )
 
     def test_role_determination(self):
         """Test role correctly identifies modules and stage positions."""
@@ -550,7 +548,7 @@ def test_role_determination(self):
             self.patch_dim,
             {"images": 50257},
         )
-        assert model_no_grid.role.mode == ModuleLayout.UNIFIED
+        assert model_no_grid.role.mode == ModuleLayout.COLOCATED
         assert model_no_grid.role.has_language_module is True
         assert model_no_grid.role.has_modality_modules is True
 
@@ -564,12 +562,15 @@ def test_role_determination(self):
         assert model_language.role.has_modality_modules is False
         assert model_language.role.has_language_module is True
 
-        # Stage info with PP
+        # Stage info with PP. language_in_grid=False so encoder and language
+        # grids have distinct rank_offsets and role.build dispatches to
+        # _from_grid_map (rather than collapsing to the COLOCATED path).
         model_pp = MimoModel(
-            self._make_config(encoder_in_grid=True, language_in_grid=True, pp_rank=1, pp_size=3)
+            self._make_config(encoder_in_grid=True, language_in_grid=False, pp_rank=1, pp_size=3)
         )
         assert model_pp.role.is_first_stage("images") is False
         assert model_pp.role.is_last_stage("images") is False
+        assert model_pp.colocated_comms == {}
 
     def test_selective_init_encoder_only(self):
         """Test encoder-only rank initializes encoder but not language model."""