feat(mimo): phase 2 (model provider, DDP wrapping, process groups)

Signed-off-by: Ali Roshan Ghias <[email protected]>

Author: Ali Roshan Ghias

src/megatron/bridge/models/__init__.py

@@ -210,6 +210,10 @@
     EncoderTransformerConfig,
     GenericVisionEncoderProvider,
 )
+from megatron.bridge.models.mimo_provider import (
+    MIMOModelProvider,
+    MIMOModelProviderResult,
+)
 from megatron.bridge.models.t5_provider import T5ModelProvider
 
 
@@ -253,6 +257,8 @@
     "EncoderProvider",
     "EncoderTransformerConfig",
     "GenericVisionEncoderProvider",
+    "MIMOModelProvider",
+    "MIMOModelProviderResult",
     "LlamaModelProvider",
     "Llama2ModelProvider7B",
     "Llama2ModelProvider13B",

src/megatron/bridge/models/encoder_provider.py

@@ -2,45 +2,98 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Optional
+from typing import Any, Optional
 
 from megatron.core.transformer.spec_utils import ModuleSpec
 
 
 @dataclass
 class EncoderTransformerConfig:
-    """Lightweight base config for encoder providers."""
+    """Lightweight base config for encoder providers.
+
+    Attributes:
+        num_layers: Number of transformer layers in the encoder.
+        hidden_size: Hidden dimension size of the encoder.
+        num_attention_heads: Number of attention heads.
+        seq_length: Sequence length for the encoder.
+        projector_type: Type of projector (e.g., "mlp", "linear", "qformer").
+            None means no projection is needed.
+        projector_input_size: Input size for projector. Defaults to hidden_size.
+        projector_output_size: Output size for projector (e.g., LLM hidden size).
+            Required if projector_type is set.
+        projector_config: Optional TransformerConfig for the projector module.
+    """
 
     num_layers: int
     hidden_size: int
     num_attention_heads: int
     seq_length: int
 
+    # Projector support for VLM setups
+    projector_type: Optional[str] = None
+    projector_input_size: Optional[int] = None
+    projector_output_size: Optional[int] = None
+    projector_config: Optional[Any] = None
+
+    def __post_init__(self) -> None:
+        """Set default projector_input_size to hidden_size if not specified."""
+        if self.projector_input_size is None:
+            self.projector_input_size = self.hidden_size
+
 
 class EncoderProvider(ABC):
-    """Interface for encoder providers used in MIMO setups."""
+    """Interface for encoder providers used in MIMO setups.
+
+    Subclasses must set the `config` attribute to an EncoderTransformerConfig.
+    """
+
+    config: EncoderTransformerConfig
 
     @abstractmethod
     def provide_model(self, pg_collection) -> object:
         """Create the encoder module (unwrapped)."""
 
     @abstractmethod
     def get_transformer_layer_spec(self) -> ModuleSpec:
-        """Return the ModuleSpec for the encoder stack."""
+        """Return the ModuleSpec for the encoder transformer layers."""
 
     @abstractmethod
     def get_projection_spec(self) -> Optional[ModuleSpec]:
-        """Optional projection ModuleSpec for encoder outputs."""
+        """Optional projection ModuleSpec for encoder outputs.
+
+        Returns None if no projection is needed.
+        """
+
+    def has_projector(self) -> bool:
+        """Check if this encoder requires a projector."""
+        return self.config.projector_type is not None
+
+    def validate_projector_config(self) -> None:
+        """Validate projector configuration consistency.
+
+        Raises:
+            ValueError: If projector_type is set but required fields are missing,
+                or if projector_type is set but get_projection_spec() returns None.
+        """
+        if self.config.projector_type is not None:
+            if self.config.projector_output_size is None:
+                raise ValueError(
+                    f"projector_output_size must be set when projector_type='{self.config.projector_type}'"
+                )
+            if self.get_projection_spec() is None:
+                raise ValueError(
+                    f"get_projection_spec() must return a ModuleSpec when "
+                    f"projector_type='{self.config.projector_type}'"
+                )
 
 
 class GenericVisionEncoderProvider(EncoderProvider):
-    """Minimal stub encoder provider for Phase 1 wiring."""
+    """Minimal stub encoder provider for Phase 1/2 wiring."""
 
     def __init__(self, config: EncoderTransformerConfig) -> None:
         self.config = config
 
     def provide_model(self, pg_collection) -> object:
-        # Stub: actual encoder creation will be implemented in Phase 2.
         raise NotImplementedError("GenericVisionEncoderProvider.provide_model not implemented.")
 
     def get_transformer_layer_spec(self) -> ModuleSpec: