src/megatron/bridge/data/mimo/collate.py (1)

56-60: Standard field stacking assumes uniform sequence lengths.

torch.stack will raise a RuntimeError if items have different sequence lengths. Unlike the modality inputs handling (which catches RuntimeError), these standard fields will propagate the exception without a helpful message.

Consider adding a try/except with a more descriptive error, or validating shapes upfront.

♻️ Optional: Add descriptive error handling

     # Stack standard fields
-    input_ids = torch.stack([item["input_ids"] for item in batch])
-    labels = torch.stack([item["labels"] for item in batch])
-    attention_mask = torch.stack([item["attention_mask"] for item in batch])
-    position_ids = torch.stack([item["position_ids"] for item in batch])
+    try:
+        input_ids = torch.stack([item["input_ids"] for item in batch])
+        labels = torch.stack([item["labels"] for item in batch])
+        attention_mask = torch.stack([item["attention_mask"] for item in batch])
+        position_ids = torch.stack([item["position_ids"] for item in batch])
+    except RuntimeError as e:
+        raise RuntimeError(
+            f"Failed to stack batch items. Ensure all items have the same sequence length. "
+            f"Original error: {e}"
+        ) from e

tests/unit_tests/data/mimo/test_collate.py (1)

10-24: Consider explicit Optional type hint.

Per PEP 484, implicit Optional via = None default is discouraged. Static analysis flagged this (RUF013).

♻️ Optional fix

+from typing import Optional
+
 def make_sample(
     seq_length: int = 64,
-    modalities: dict = None,
+    modalities: Optional[dict] = None,
 ) -> dict:

tests/unit_tests/data/mimo/test_dataset.py (1)

371-410: Multi-modality placeholder ordering test relies on dict insertion order.

This test assumes placeholders appear in dict insertion order (vision before audio). While Python 3.7+ guarantees dict ordering, this implicit dependency could be fragile if the ordering contract isn't documented in MimoDataset.

Consider adding a brief comment in the test or verifying the ordering is documented in the dataset implementation.

src/megatron/bridge/data/mimo/loaders.py (3)

6-6: Remove unused import.

Dict is imported but not used in this file.

🔧 Proposed fix

-from typing import TYPE_CHECKING, Dict, Optional, Tuple
+from typing import TYPE_CHECKING, Optional, Tuple

---

20-27: Consider documenting or removing unused train_state parameter.

The train_state parameter is declared but not used in the function body. If this is for API consistency with other data loader builders or planned future use, consider adding a brief comment. Otherwise, it can be removed.

---

66-67: Use TypeError for type validation.

When validating that an argument is of a specific type, TypeError is semantically more appropriate than ValueError.

🔧 Proposed fix

     if not isinstance(cfg.model, MimoModelProvider):
-        raise ValueError("cfg.model must be MimoModelProvider for MIMO data loading.")
+        raise TypeError("cfg.model must be MimoModelProvider for MIMO data loading.")

src/megatron/bridge/training/mimo_ddp.py (1)

78-94: Edge case: empty encoders dictionary.

The condition if hasattr(submodule, 'encoders') and submodule.encoders passes if encoders is a non-empty dict, but next(iter(submodule.encoders.keys())) would fail with StopIteration if called on an empty dict. The current check is correct, but consider adding a more explicit check or comment for clarity.

🔧 Suggested clarification

-            if hasattr(submodule, 'encoders') and submodule.encoders:
+            # Skip submodules without encoders or with empty encoder dict
+            if not (hasattr(submodule, 'encoders') and submodule.encoders):
+                continue
+            
+            if True:  # Encoders present
                 encoder_key = next(iter(submodule.encoders.keys()))

Or simply add a comment to clarify the truthiness check covers the empty dict case.

src/megatron/bridge/data/mimo/hf_provider.py (2)

8-8: Remove unused imports.

List and Union are imported but not used in this file.

🔧 Proposed fix

-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Optional, Tuple

---

122-140: Consider narrowing exception handling.

Catching ValueError broadly may inadvertently suppress errors unrelated to missing splits (e.g., configuration issues). Consider checking the exception message or using a more specific approach.

🔧 Proposed approach

         try:
             dataset = load_dataset(
                 self.hf_dataset_path,
                 name=self.hf_dataset_name,
                 split=split,
                 streaming=self.streaming,
                 trust_remote_code=is_safe_repo(
                     trust_remote_code=self.trust_remote_code,
                     hf_path=self.hf_dataset_path,
                 ),
             )
             return dataset
-        except ValueError:
-            # Split doesn't exist
+        except ValueError as e:
+            # Check if this is specifically a missing split error
+            if "split" in str(e).lower():
+                return None
+            raise
-            return None

src/megatron/bridge/data/mimo/dataset.py (2)

6-6: Remove unused imports.

List and Union are imported but not used in this file.

🔧 Proposed fix

-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, Optional

---

189-212: Potential over-truncation of text.

The text is first tokenized with max_length=self.seq_length (line 193), then truncated again after adding prefix tokens (line 211). This means if prefix tokens take up N positions, the final text could be truncated to seq_length - N tokens, but the initial tokenization may have already truncated content that could have fit.

Consider pre-computing the available text space:

🔧 Proposed optimization

+        # Pre-compute space available for text tokens
+        total_prefix_len = sum(
+            self.encoder_seq_lengths.get(m, 1)
+            for m in modality_inputs.keys()
+            if m in self.special_token_ids
+        )
+        max_text_length = self.seq_length - total_prefix_len
+        
         # Tokenize the text
         encoded = self.tokenizer(
             text,
             truncation=True,
-            max_length=self.seq_length,
+            max_length=max_text_length,
             return_tensors="pt",
         )
         input_ids = encoded["input_ids"].squeeze(0)
         
         # Insert placeholder tokens for each modality at the beginning
-        # The order follows the order of modality_inputs (Python 3.7+ dict ordering)
         prefix_tokens = []
         for modality_name in modality_inputs.keys():
             if modality_name in self.special_token_ids:
                 token_id = self.special_token_ids[modality_name]
                 num_tokens = self.encoder_seq_lengths.get(modality_name, 1)
                 prefix_tokens.extend([token_id] * num_tokens)
         
         if prefix_tokens:
             prefix = torch.tensor(prefix_tokens, dtype=input_ids.dtype)
-            # Truncate text tokens to make room for placeholders
-            max_text_len = self.seq_length - len(prefix_tokens)
-            input_ids = input_ids[:max_text_len]
             input_ids = torch.cat([prefix, input_ids])

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

61-73: Return type annotation uses forward reference but static analysis can't see the lazy import.

The return type "ColocatedCommConfig" on Line 63 is correctly handled via a string annotation (forward reference), and the actual import happens on Line 65 before use. This is a valid pattern for lazy imports.

However, Line 67 creates a redundant copy of the grids dict:

♻️ Minor simplification

 def build_colocated_comm_config(
     mimo_parallelism_config: MimoParallelismConfig, grids: Dict[str, "HyperCommGrid"]
 ) -> "ColocatedCommConfig":
     """Build ColocatedCommConfig with default encoder-to-LLM topology."""
     from megatron.core.models.mimo.config.base_configs import ColocatedCommConfig

-    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},
     )

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

355-376: Unused variable model from static analysis is valid but intentional.

The model variable on Line 372 is assigned but never used because the test focuses on verifying the side effect (parameter freezing) via the mock. While technically unused, the assignment makes the test intent clearer.

Consider prefixing with underscore to silence the warning:

♻️ Silence unused variable warning

-        model = provider.provide()
+        _model = provider.provide()

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

22-22: Unused import DistributedDataParallel flagged by static analysis.

The import at Line 22 is flagged as unused. While DistributedDataParallelConfig is used, DistributedDataParallel appears to not be directly referenced in this file (it's used in mimo_ddp.py instead).

♻️ Remove unused import

-from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
+from megatron.core.distributed import DistributedDataParallelConfig

---

294-299: Unused loop variable encoder_name can be renamed.

The loop at Line 296 iterates over encoder items but only uses encoder_spec, not encoder_name.

♻️ Rename unused loop variable

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