src/megatron/bridge/models/mimo/mimo_config.py (1)

126-131: Consider validating data_parallel before finalize() in heterogeneous mode.

In heterogeneous mode, parallelism.finalize(None) is called (Line 131), which will fail with "world_size must be provided" if data_parallel is None. While _validate_heterogeneous() later checks for this, the error message from finalize() may be confusing. Consider moving the data_parallel is None check earlier:

♻️ Suggested improvement

         if self.deployment_mode in ("colocated", "homogeneous"):
             for parallelism in self.module_parallelisms.values():
                 parallelism.finalize(world_size)
         else:
+            # Heterogeneous requires data_parallel to be pre-set
+            for name, parallelism in self.module_parallelisms.items():
+                if parallelism.data_parallel is None:
+                    raise ValueError(
+                        f"data_parallel must be set for module '{name}' in heterogeneous deployment."
+                    )
             for parallelism in self.module_parallelisms.values():
                 parallelism.finalize(None)

src/megatron/bridge/models/mimo/mimo_builder.py (2)

3-3: Remove unused import Optional.

Optional is imported but never used in this file.

🧹 Suggested fix

-from typing import Dict, List, Optional
+from typing import Dict, List

---

62-62: Simplify redundant dict comprehension.

The dict comprehension {name: grid for name, grid in grids.items()} creates an identical copy of grids. If a copy is intended, use grids.copy(); otherwise, pass grids directly.

♻️ Suggested fix

-    module_to_grid_map = {name: grid for name, grid in grids.items()}
     topology = _default_topology(mimo_parallelism_config)
     return ColocatedCommConfig(
-        module_to_grid_map=module_to_grid_map,
+        module_to_grid_map=grids,
         topology=topology,
         dim_mapping={"b": 0, "s": 1, "h": 2},
     )

src/megatron/bridge/models/mimo/mimo_provider.py (1)

201-204: Rename unused loop variable encoder_name to _encoder_name.

The loop variable encoder_name is not used within the loop body. Following Python conventions, prefix it with underscore to indicate it's intentionally unused.

🧹 Suggested fix

-            for encoder_name, encoder_spec in spec.submodules["encoders"].items():
+            for _encoder_name, encoder_spec in spec.submodules["encoders"].items():

tests/unit_tests/training/mimo/test_mimo_config.py (1)

14-17: Use raw string for regex pattern with metacharacters.

The pattern "world_size .* not divisible" contains regex metacharacters (.*). Use a raw string prefix for clarity and to avoid potential escaping issues.

🧹 Suggested fix

 def test_module_parallelism_finalize_invalid_world_size():
     parallelism = ModuleParallelismConfig(tensor_parallel=3, pipeline_parallel=2)
-    with pytest.raises(ValueError, match="world_size .* not divisible"):
+    with pytest.raises(ValueError, match=r"world_size .* not divisible"):
         parallelism.finalize(world_size=10)

src/megatron/bridge/models/mimo/llava_provider.py (1)

130-137: Encoder key mismatch: "clip_encoder" vs modality key "images".

The encoder is registered under "clip_encoder" in the submodules dict (line 134), but the modality itself is keyed as "images" (line 77-78). This is likely intentional for semantic clarity, but ensure downstream code correctly maps between modality names and encoder names when applying per-encoder parallelism.

tests/unit_tests/models/mimo/test_mimo_provider.py (1)

4-11: Unused imports: pytest and torch.

Static analysis correctly identifies that pytest and torch are imported but not used. While pytest is typically needed for fixtures and markers, neither is used in this file. Consider removing them or adding a # noqa: F401 comment if they're kept for future use.

Suggested fix

-from unittest.mock import MagicMock, Mock, patch
-
-import pytest
-import torch
+from unittest.mock import MagicMock, Mock, patch
 
 from megatron.core.transformer.spec_utils import ModuleSpec