from typing import List 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.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": 8, # "moe_loss_weight": config.moe_loss_weight, # "device": config.init_device, # Handled by FSDP "bf16": False, "fp16": False, "bias": False, "return_bias": False, } return MoEArgs(**kwargs) class MolmoeConfig(PretrainedConfig): model_type = "molmoe" keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, vocab_size=50304, embedding_size=50304, hidden_size=4096, intermediate_size=11008, num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=None, max_position_embeddings=2048, initializer_range=0.02, use_cache=True, layer_norm_eps: float = 1e-5, rope_theta=10000.0, clip_qkv=None, qkv_bias: bool = False, weight_tying: bool = False, use_position_ids: bool=True, tie_word_embeddings: bool=True, moe_num_experts: int = 64, moe_top_k: int = 8, **kwargs, ): self.vocab_size = vocab_size self.embedding_size = embedding_size self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.layer_norm_eps = layer_norm_eps self.weight_tying = weight_tying self.use_position_ids = use_position_ids # for backward compatibility if num_key_value_heads is None: num_key_value_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.initializer_range = initializer_range self.use_cache = use_cache self.rope_theta = rope_theta self.clip_qkv = clip_qkv self.qkv_bias = qkv_bias self.tie_word_embeddings = tie_word_embeddings super().__init__( tie_word_embeddings=tie_word_embeddings, **kwargs, ) MolmoeConfig.register_for_auto_class()