bridge.training.pretrain_mimo#

Entry point for MIMO pretraining.

Thin entry point that orchestrates runtime config updates, setup, and training. Mirrors the standard pretrain.py → setup() → train() pattern.

See also:

  • setup_mimo.py: MIMO-specific setup logic (analogous to setup.py)

  • train_mimo.py: MIMO training loop (analogous to train.py)

  • config.py: mimo_runtime_config_update() (analogous to runtime_config_update())

Module Contents#

Functions#

pretrain_mimo

Entry point for MIMO pretraining.

Data#

logger

API#

bridge.training.pretrain_mimo.logger#

‘getLogger(…)’

bridge.training.pretrain_mimo.pretrain_mimo(
cfg: megatron.bridge.training.config.ConfigContainer,
forward_step_func: Callable,
build_data_iterators_fn: Callable,
global_state: Optional[megatron.bridge.training.state.GlobalState] = None,
) → None#

Entry point for MIMO pretraining.

Steps:

  1. Apply MIMO runtime config updates (finalize sub-configs, set data_parallel_size=1)

  2. Call setup_mimo() to get model, optimizer, schedulers, infra, communicators

  3. Call train_mimo() with all components

Parameters:

  • cfg – ConfigContainer with training configuration. cfg.model must be a MimoModelProvider. cfg.optimizer (a BridgeOptimizerConfig) is used to create the MimoOptimizer and per-module LR schedulers.

  • forward_step_func – Forward step function for training.

  • build_data_iterators_fn – Function to build data iterators. Signature: (cfg, mimo_infra) -> (train_iter, valid_iter)

  • global_state – Optional GlobalState. If not provided, creates a new one.

TODO(liding): check if build_data_iterators_fn and global_state are needed (deferred to phase 5 review)