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

57-59: Enforce the documented cfg.model contract at runtime.

Line 57 documents OmniModalProvider, but pretrain_mimo() does not validate it before calling setup_mimo(). Add a fail-fast check so invalid configs raise a clear error instead of later attribute errors.

Suggested patch

 from megatron.bridge.training.config import ConfigContainer, mimo_runtime_config_update
+from megatron.bridge.models.omni_modal.omni_modal_provider import OmniModalProvider
 from megatron.bridge.training.pretrain import _maybe_destroy_process_group
 from megatron.bridge.training.setup_mimo import setup_mimo
@@
 def pretrain_mimo(
@@
 ) -> None:
@@
     logger.info("Starting MIMO pretraining")
+    if not isinstance(cfg.model, OmniModalProvider):
+        raise TypeError(
+            f"pretrain_mimo requires cfg.model to be OmniModalProvider, got {type(cfg.model).__name__}."
+        )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/training/pretrain_mimo.py` around lines 57 - 59,
pretrain_mimo currently documents that cfg.model must be an OmniModalProvider
but doesn't validate this at runtime; add a fail-fast type/attribute check at
the start of pretrain_mimo (before calling setup_mimo) to verify cfg is a
ConfigContainer and cfg.model implements the OmniModalProvider contract (e.g.,
check isinstance(cfg.model, OmniModalProvider) or the presence of required
methods/attributes used by setup_mimo), and raise a clear ValueError/TypeError
with a descriptive message if the check fails so callers get an immediate,
actionable error instead of later attribute errors.

src/megatron/bridge/models/omni_modal/__init__.py (1)

14-20: Sort __all__ to satisfy Ruff RUF022.

__all__ is currently unsorted; aligning ordering avoids lint noise and keeps exports deterministic.

♻️ Suggested change

 __all__ = [
     "LlavaOmniModalProvider",
+    "ModuleParallelismConfig",
     "OmniModalInfra",
+    "OmniModalParallelismConfig",
     "OmniModalProvider",
-    "OmniModalParallelismConfig",
-    "ModuleParallelismConfig",
 ]

As per coding guidelines: `**/*.py` must pass Ruff linting/formatting.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/models/omni_modal/__init__.py` around lines 14 - 20, The
__all__ list in omni_modal/__init__.py is unsorted and causing Ruff RUF022;
reorder the exported names alphabetically (e.g., "LlavaOmniModalProvider",
"ModuleParallelismConfig", "OmniModalInfra", "OmniModalParallelismConfig",
"OmniModalProvider") so the __all__ sequence is deterministic and
lint-compliant; update the __all__ assignment to contain the same symbols sorted
lexicographically.

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

273-277: Add * separator for multiple same-type parameters in validate_data_loader_contract.

Line 273-277 has three int parameters and should enforce keyword-only separation per repo style rule.

♻️ Proposed refactor

 def validate_data_loader_contract(
     infra: OmniModalInfra,
+    *,
     global_batch_size: int,
     micro_batch_size: int,
     num_microbatches: int,
 ):

As per coding guidelines, "**/*.py: Use * separator for functions with multiple same-type parameters".

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/training/mimo_parallel_utils.py` around lines 273 - 277,
The function validate_data_loader_contract has three same-typed int params
(global_batch_size, micro_batch_size, num_microbatches) and needs a keyword-only
separator; update the signature to add a '*' before those params (e.g., def
validate_data_loader_contract(infra: OmniModalInfra, *, global_batch_size: int,
micro_batch_size: int, num_microbatches: int):) and adjust any callers to pass
those three arguments by keyword if necessary.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes