Fix

config.json

@@ -7,7 +7,7 @@
     "AutoModelForCausalLM": "modeling_molmoe.MOLMoEForCausalLM"
   },
   "clip_qkv": null,
-  "embedding_size": 152064,
+  "embedding_size": 50304,
   "hidden_size": 2048,
   "initializer_range": 0.02,
   "intermediate_size": 1024,

config_molmoe.py

@@ -5,19 +5,22 @@ from transformers import PretrainedConfig, AutoTokenizer
 
 def config_to_moe_args(config):
     from megablocks.layers.arguments import Arguments as MoEArgs
+    import torch.nn.functional as F
+
+    # import pdb; pdb.set_trace()
 
     kwargs = {
         "activation_fn": F.silu,
         "mlp_type": "glu" if "glu" in config.activation_type.lower() else "mlp",
         "mlp_impl": "sparse",
-        "hidden_size": config.hidden_size,
-        "ffn_hidden_size": config.intermediate_size,
-        "moe_num_experts": config.moe_num_experts,
-        "num_layers": config.num_hidden_layers,
+        "hidden_size": config.d_model,
+        "ffn_hidden_size": config.mlp_hidden_size,
+        "moe_num_experts": 64,
+        "num_layers": config.n_layers,
         # Handled by FSDP (https://github.com/databricks/megablocks/issues/57#issuecomment-1854594483)
         "moe_weight_parallelism": False,
         "moe_expert_model_parallelism": False,
-        "moe_top_k": config.moe_top_k,
+        "moe_top_k": 8,
         # "moe_loss_weight": config.moe_loss_weight,
         # "device": config.init_device,
         # Handled by FSDP

modeling_molmoe.py

@@ -235,14 +235,16 @@ class OLMoBlock(nn.Module):
             device=config.init_device
         )
 
-        # Feed-forward output projection.
-        self.ff_out = nn.Linear(
-            int(self.act.output_multiplier * self.hidden_size),
-            config.d_model,
-            bias=config.include_bias,
-            device=config.init_device,
-        )
-        self.ff_out._is_residual = True  # type: ignore
+
+        if self.config.block_type != "moe":
+            # Feed-forward output projection.
+            self.ff_out = nn.Linear(
+                int(self.act.output_multiplier * self.hidden_size),
+                config.d_model,
+                bias=config.include_bias,
+                device=config.init_device,
+            )
+            self.ff_out._is_residual = True  # type: ignore
 
         # Rotary embeddings.
         if self.config.rope:
@@ -423,7 +425,7 @@ class OLMoBlock(nn.Module):
             return OLMoSequentialBlock(layer_id, config, cache)
         elif config.block_type == "llama":
             return OLMoLlamaBlock(layer_id, config, cache)
-        elif config.block_type == BlockType.moe:
+        elif config.block_type == "moe":
             return OLMoEBlock(layer_id, config, cache)
         else:
             raise NotImplementedError(f"Unknown block type: '{config.block_type}'")
@@ -725,7 +727,7 @@ class OLMoEBlock(OLMoBlock):
     (plus another skip connection).
     """
 
-    def __init__(self, layer_id: int, config: ModelConfig, cache: BufferCache):
+    def __init__(self, layer_id: int, config, cache: BufferCache):
         try:
             from megablocks.layers.dmoe import dMoE
             from megablocks.layers.moe import MoE
@@ -733,12 +735,12 @@ class OLMoEBlock(OLMoBlock):
             raise ImportError(
                 "To train MoEs, run `pip install git+https://github.com/Muennighoff/megablocks.git@olmoe`"
             )
-        from .config import config_to_moe_args
+        from .config_molmoe import config_to_moe_args
 
         super().__init__(layer_id, config, cache)
 
         self.moe_args = config_to_moe_args(config)
-        self.ffn = dMoE(self.moe_args) if self.config.moe_dropless else MoE(self.moe_args)
+        self.ffn = dMoE(self.moe_args)
 
         self.attn_norm = LayerNorm.build(config)
         self.ff_norm = LayerNorm.build(config)
@@ -956,12 +958,14 @@ class VisionBackboneConfig:
     image_default_input_size: Tuple[int, int] = (336, 336)
     image_patch_size: int = 14
     image_pos_patch_size: int = 14
+    # image_emb_dim: int = 1024
     image_emb_dim: int = 1024
     image_num_heads: int = 16
     image_num_key_value_heads: int = 16
     image_num_layers: int = 24
     image_head_dim: int = 64
-    image_mlp_dim: int = 4096
+    # image_mlp_dim: int = 4096
+    image_mlp_dim: int = 2048
     image_mlp_activations: str = "gelu"
     image_dropout_rate: float = 0.0
     image_num_pos: int = 577
@@ -990,10 +994,10 @@ class FullMolmoeConfig:
     qkv_bias: bool = False
     clip_qkv: Optional[float] = None
     n_layers: int = 12
-    mlp_ratio: int = 4
+    mlp_ratio: int = 1
     mlp_hidden_size: Optional[int] = None
     activation_type: str = "swiglu"
-    block_type: str = "sequential"
+    block_type: str = "moe"
     block_group_size: int = 1
     alibi: bool = False
     alibi_bias_max: float = 8.0
@@ -1009,7 +1013,7 @@ class FullMolmoeConfig:
     attention_dropout: float = 0.1
     response_attention_dropout: float = 0.0
     multi_query_attention: Optional[bool] = None
-    attention_layer_norm: bool = False
+    attention_layer_norm: bool = True
     residual_dropout: float = 0.1
     response_residual_dropout: float = 0.0
     embedding_dropout: float = 0.1
@@ -1651,6 +1655,9 @@ class OLMoVisionBackbone(nn.Module):
                 [MLP(mlp_config, input_dim), Residual(MLP(config, input_dim))]
             )
         elif config.image_projector == ImageProjectType.mlp:
+            #import pdb; pdb.set_trace()
+            #mlp_config.image_mlp_dim = 2048
+            mlp_config.mlp_hidden_size = 2048
             self.image_projector = MLP(mlp_config, input_dim)
         elif config.image_projector == ImageProjectType.linear:
             self.image_projector = nn.Linear(
@@ -2423,7 +2430,7 @@ class MOLMoEForCausalLM(PreTrainedModel):
     base_model_prefix = "model"
     _no_split_modules = ["OLMoBlock"]
 
-    def __init__(self, config: MolmoeConfig, model: Optional[MOLMo] = None, init_params: bool = False):
+    def __init__(self, config: MolmoeConfig, model: Optional[MOLMoE] = None, init_params: bool = False):
         super().__init__(config)
 
         if not model:
@@ -2447,8 +2454,8 @@ class MOLMoEForCausalLM(PreTrainedModel):
                 additional_vocab_size=128,
                 n_heads=config.num_attention_heads,
                 n_kv_heads=config.num_key_value_heads,
-                rope_theta=1000000.0,
-                layer_norm_eps=1e-6,
+                rope_theta=10000.0,
+                layer_norm_eps=1e-5,
                 layer_norm_type="rms",
                 pad_tokenizer=True,
                 vit_layers=[-2, -9],
@@ -2472,7 +2479,7 @@ class MOLMoEForCausalLM(PreTrainedModel):
                     initializer_range=0.02,
                 )
             )
-            self.model = MOLMo(full_config, init_params=init_params)
+            self.model = MOLMoE(full_config, init_params=init_params)
         else:
             self.model = model