Is this page helpful?

core.models.mimo.model.base#

Module Contents#

Classes#

MimoModel

Multimodal In/Out Model supporting arbitrary combinations of modalities.

Data#

logger

_LANGUAGE_INPUT_PROJECTIONS_ATTR

API#

core.models.mimo.model.base.logger#

‘getLogger(…)’

core.models.mimo.model.base._LANGUAGE_INPUT_PROJECTIONS_ATTR#

‘mimo_input_projections’

class core.models.mimo.model.base.MimoModel(
mimo_config: megatron.core.models.mimo.config.MimoModelConfig,
cp_group=None,
tp_group=None,
)#

Bases: megatron.core.transformer.MegatronModule

Multimodal In/Out Model supporting arbitrary combinations of modalities.

.. warning:: EXPERIMENTAL: This class is experimental, still under active development, and the API is subject to change without notice. Use at your own risk.

.. note:: This implementation is in development and may undergo API changes.

This model processes multiple modalities (e.g., vision, audio) alongside text, combining their embeddings before passing them through a language model.

Parameters:

mimo_config (MimoModelConfig) – Configuration for the model, including language model and modality submodules

Initialization

Initialize the multimodal model.

.. rubric:: Example

# Create a model with default configuration
model = MimoModel(mimo_config)
sharded_state_dict(prefix='', sharded_offsets=(), metadata=None)#

Build sharded state dict, bypassing parallel_state global fallbacks.

Iterates modality_submodules manually (ModuleDict lacks sharded_state_dict) and injects dp_cp_group from each module’s pg_collection.

static _validate_precomputed_token_indices(
modality_embeddings: Dict[str, torch.Tensor],
modality_token_indices: Dict[str, torch.Tensor],
num_positions: int,
) → None#

Validate the metadata-only contract for precomputed token positions.

align_embeddings_by_token_positions(
modality_embeddings: Dict[str, torch.Tensor],
input_ids: torch.Tensor,
special_token_ids: Dict[str, int],
modality_token_indices: Optional[Dict[str, torch.Tensor]] = None,
) → torch.Tensor#

Align embeddings from different modalities based on special token positions in input_ids.

Parameters:

  • modality_embeddings – Dictionary mapping modality names to their embeddings. For all modalities: tensor of shape (N, H). Shape: (num_tokens_for_modality, hidden_dim)

  • input_ids – Input token IDs. Shape: (B, S) or (S,) Contains special tokens that mark where each modality’s embeddings should go. The number of special tokens for each modality should exactly match the number of embeddings for that modality.

  • special_token_ids – Dictionary mapping modality names to their special token IDs

  • modality_token_indices – Optional complete mapping from modality names to flat, batch-major torch.long token indices on the embedding device. Each tensor must contain the same positions, in the same order, that the corresponding special_token_ids mask would select; text contains the complement of all special-token positions. When provided, the keys must exactly match modality_embeddings and their lengths must sum to B * S. Index values are a trusted producer contract and are not revalidated because doing so would add device reductions and host synchronization to the forward hot path. Pass None to use the mask-based path for every modality.

Returns:

(S, B, H)

Return type:

Combined embeddings tensor. Shape

_initialize_submodules() → None#

Initialize modality submodules from the ModuleSpec configurations.

When role is set, only initializes submodules this rank participates in. Stage info is passed to from_spec() to conditionally skip projection.

_finish_init_quantization() → None#

Apply per-module quantization recipes to initialized projections and encoders.

_initialize_language_input_projections() → None#

Install modality input projections on the first language stage.

property language_model_input_projections: Optional[torch.nn.ModuleDict]#

Return the input projections installed on the language model.

_initialize_language_model() → None#

Initialize the language model.

When role is set, only initializes if this rank participates in language module.

set_input_tensor(input_tensor)#

Set input tensor for pipeline parallelism.

This method is required by Megatron’s pipeline parallel mechanism. It passes the output tensor from the previous stage as input to this stage.

Parameters:

input_tensor –

Either:

  • Dict[str, Tensor]: Maps module names to their input tensors (for multi-module PP)

  • Tensor or List[Tensor]: Single tensor for language model (backward compat)

Returns:

None

_active_submodules()#

Yield this rank’s present submodules.

_active_ddp_modules()#

Yield this rank’s active DDP-wrapped submodules.

property remove_forward_pre_hook_handles: Dict[torch.nn.Module, Any]#

Expose the active inner DDP parameter-gather hooks to the stock train loop.

no_sync()#

Disable grad-ready registration on overlapped inner DDP modules.

enable_forward_pre_hook()#

Enable parameter-gather pre-hooks on overlapped inner DDP modules.

disable_forward_pre_hook(param_sync: bool = True)#

Disable parameter-gather pre-hooks on overlapped inner DDP modules.

start_param_sync(
*unused,
force_sync: bool = False,
force_dispatch: bool = False,
)#

Start parameter synchronization on overlapped inner DDP modules.

start_grad_sync(*unused)#

Start gradient synchronization on overlapped inner DDP modules.

free_overlap_buffers()#

Release parameter-gather buffers owned by overlapped inner DDP modules.

zero_grad_buffer()#

Zero each active submodule’s DDP grad buffer.

get_text_embeddings(
input_ids: torch.Tensor,
position_ids: torch.Tensor,
special_token_ids: Dict[str, int],
text_token_indices: Optional[torch.Tensor] = None,
) → torch.Tensor#

Get embeddings for text tokens in the input.

Parameters:

  • input_ids – Input token IDs. Shape: (B, S) Contains text tokens and potentially special tokens for other modalities.

  • position_ids – Position IDs corresponding to input tokens, used for positional encoding. Shape: (B, S)

  • special_token_ids – Dictionary mapping modality names to their special token IDs. Used to identify non-text tokens in the input_ids.

  • text_token_indices – Optional flat, logical row-major torch.long indices of text tokens in the [B, S] input grid.

Returns:

Embeddings for text tokens. Shape: (N, H), where N is the number of text tokens.

Return type:

torch.Tensor

forward(
input_ids: torch.Tensor,
position_ids: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
loss_mask: Optional[torch.Tensor] = None,
labels: Optional[torch.Tensor] = None,
modality_inputs: Optional[Dict[str, Dict[str, Any]]] = None,
packing_kwargs: Optional[dict] = None,
modality_token_indices: Optional[Dict[str, torch.Tensor]] = None,
)#

Forward pass through the multimodal model.

Parameters:

  • input_ids – Input token IDs. Shape: (B, S)

  • position_ids – Position IDs. Shape: (B, S)

  • attention_mask – Attention mask. Shape: (B, S)

  • loss_mask – Loss mask. Shape: (B, S)

  • labels – Labels for training. Shape: (B, S)

  • modality_inputs – Dictionary mapping modality names to encoder inputs. For example: { “images”: { “clip_encoder”: {“pixel_values”: clip_images}, “vit_encoder”: {“images”: vit_images} }, “audio”: { “whisper_encoder”: {“input_features”: whisper_features} } }

  • packing_kwargs – Optional dictionary of kwargs to construct PackedSeqParams if packed_seq_params is not provided. For example: { “cu_seqlens_q”: cu_seqlens, “cu_seqlens_kv”: cu_seqlens, “cu_seqlens_q_padded”: cu_seqlens_padded, “cu_seqlens_kv_padded”: cu_seqlens_padded, “max_seqlen_q”: torch.tensor( max(seqlens_padded), dtype=torch.int32 ), “max_seqlen_kv”: torch.tensor( max(seqlens_padded), dtype=torch.int32 ), }

  • modality_token_indices – Optional complete mapping from every active modality name, including text, to flat logical row-major torch.long indices in the [B, S] input grid. Encoder-only ranks ignore this argument. Pass None to derive every position from input_ids. See align_embeddings_by_token_positions for the trusted-producer contract.

Returns:

(output, loss_mask) where output semantics depend on role: - Encoder-only ranks: Dict[str, Tensor] of encoder outputs - Language module ranks: language model output (logits or loss) - No role (all modules colocated): language model output

Return type:

tuple

_forward_encoders(
input_ids: Optional[torch.Tensor],
modality_inputs: Optional[Dict[str, Dict[str, Any]]],
input_tensors: Optional[Dict[str, torch.Tensor]],
) → Dict[str, torch.Tensor]#

Forward pass for encoder modules on this rank.

Parameters:

  • modality_inputs – Raw inputs for each modality (images, audio, etc.)

  • input_tensors – Hidden states from previous pipeline stages

Returns:

Dict mapping encoder names to their output tensors

_attach_modality_split_sizes(
output: torch.Tensor,
input_ids: Optional[torch.Tensor],
encoder_name: str,
) → None#

Annotate flat modality outputs with per-sample split sizes for bridge fan-out.

Only attaches when per-sample token counts are non-uniform. Uniform counts give equal splits, which the bridge’s torch.tensor_split fallback already produces, so the metadata would be a no-op.

TODO(mimo): non-uniform per-sample counts in fan-in (encoder DP > LM DP) are not supported. Multiple encoder ranks contribute slices to a single LM peer, and the receiver-side torch.cat path in BridgeCommunicator has no metadata channel today, so per-sample boundaries are lost on the LM rank. Lift this by routing per-sample sizes through the bridge alongside the activations and adding a sample-aligned concat path.

_has_encoder_tokens(
input_ids: Optional[torch.Tensor],
encoder_name: str,
) → bool#

Return whether the batch contains tokens for an encoder module.

_empty_encoder_output(encoder_name: str) → torch.Tensor#

Return the bridge payload for text-only non-colocated batches.

_build_packed_seq_params(
packing_kwargs: Optional[dict],
) → Optional[megatron.core.packed_seq_params.PackedSeqParams]#

Build THD PackedSeqParams from packing_kwargs (None if not packing).

_shard_language_inputs(
embeddings: Optional[torch.Tensor],
labels: Optional[torch.Tensor],
loss_mask: Optional[torch.Tensor],
packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None,
) → Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[megatron.core.packed_seq_params.PackedSeqParams]]#

Apply CP/SP sharding via the partition adapter, or pass through if inactive.

embeddings are sequence-first (S, B, H) (None on non-first PP stages) and come back in (S/(cp*tp), B, H); labels/loss_mask are (B, S).

_language_model_owns_mtp() → bool#

Return whether this rank executes the language model’s MTP block.

static _materialize_mtp_input_mask(
input_ids: torch.Tensor,
special_token_ids: Dict[str, int],
text_token_indices: Optional[torch.Tensor] = None,
) → torch.Tensor#

Return positions backed by the language model’s token embedding table.

_prepare_mtp_inputs(
input_ids: Optional[torch.Tensor],
position_ids: Optional[torch.Tensor],
packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams],
owns_mtp: bool,
text_token_indices: Optional[torch.Tensor] = None,
) → Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]#

Prepare CP-local position IDs and optional MTP token metadata.

MTP consumes token IDs only on the stage that owns its prediction block. Under context parallelism, those IDs, their validity mask, and position IDs must all use the same local sequence partition as the language-model hidden states.

_forward_language_module(
input_ids: torch.Tensor,
position_ids: Optional[torch.Tensor],
attention_mask: Optional[torch.Tensor],
loss_mask: Optional[torch.Tensor],
labels: Optional[torch.Tensor],
input_tensors: Optional[Dict[str, torch.Tensor]],
packing_kwargs: Optional[dict] = None,
modality_token_indices: Optional[Dict[str, torch.Tensor]] = None,
) → Tuple[Any, Optional[torch.Tensor]]#

Forward pass for language module on this rank.

Parameters:

  • input_ids – Token IDs

  • position_ids – Position IDs

  • attention_mask – Attention mask. Must be None under context parallelism (CP-local hidden states cannot line up with a dense mask); mask via a causal attn_mask_type or packed_seq_params instead.

  • loss_mask – Loss mask for per-token loss normalization

  • labels – Labels for loss computation

  • input_tensors – Hidden states or embeddings from previous stage

  • packing_kwargs – Optional kwargs to construct packed (THD) sequence params.

  • modality_token_indices – Optional complete mapping of trusted flat logical row-major token indices. See align_embeddings_by_token_positions.

Returns:

Tuple of (language model output, possibly CP-sharded loss mask). The output is hidden states, logits, or loss depending on the stage.

_build_colocated_communicators()#
destroy() → None#

Release process groups owned by this MimoModel.

_apply_colocated_comms(modality_embeddings)#

Transform encoder embeddings from encoder TP/DP to LLM TP/DP layout.

_forward_all_modules(
input_ids: torch.Tensor,
position_ids: Optional[torch.Tensor],
attention_mask: Optional[torch.Tensor],
loss_mask: Optional[torch.Tensor],
labels: Optional[torch.Tensor],
modality_inputs: Optional[Dict[str, Dict[str, Any]]],
packing_kwargs: Optional[dict] = None,
modality_token_indices: Optional[Dict[str, torch.Tensor]] = None,
)#

Forward pass when all modules are on all ranks (no multi-module PP).

This is the original behavior, preserved for backward compatibility.