Upload chess-mamba-vs-xformer/mamba_lm.py with huggingface_hub

chess-mamba-vs-xformer/mamba_lm.py

@@ -5,7 +5,10 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from mamba import Mamba, MambaConfig, RMSNorm
+#from mamba import Mamba, MambaConfig, RMSNorm
+from mamba_ssm import MambaLMHeadModel
+from mamba_ssm.models.config_mamba import MambaConfig
+from mamba_ssm.ops.triton.layernorm import RMSNorm
 
 """
 
@@ -22,15 +25,18 @@ class MambaLMConfig(MambaConfig):
     pad_vocab_size_multiple: int = 8
 
     def __post_init__(self):
-        super().__post_init__()
+        pass
+        #super().__post_init__()
 
         #if self.vocab_size % self.pad_vocab_size_multiple != 0:
         #    self.vocab_size += (self.pad_vocab_size_multiple - self.vocab_size % self.pad_vocab_size_multiple)
 
     def to_mamba_config(self) -> MambaConfig:
-        mamba_config_fields = {field.name for field in fields(MambaConfig)}
-        filtered_dict = {k: v for k, v in asdict(self).items() if k in mamba_config_fields}
-        return MambaConfig(**filtered_dict)
+        #mamba_config_fields = {field.name for field in fields(MambaConfig)}
+        #print(mamba_config_fields)
+        #filtered_dict = {k: v for k, v in asdict(self).items() if k in mamba_config_fields}
+        #return MambaConfig(**filtered_dict)
+        return MambaConfig(d_model=self.d_model, n_layer=self.n_layer, vocab_size=self.vocab_size, ssm_cfg=self.ssm_cfg)
 
 # adapted from https://github.com/johnma2006/mamba-minimal
 def from_pretrained(name: str):
@@ -65,7 +71,8 @@ def from_pretrained(name: str):
     config_data = load_config_hf(name)
     config = MambaLMConfig(d_model=config_data['d_model'], n_layers=config_data['n_layer'], vocab_size=config_data['vocab_size'])
 
-    model = MambaLM(config)
+    #model = MambaLM(config)
+    model = MambaLMHeadModel(config)
 
     # copy weights
     state_dict = load_state_dict_hf(name)
@@ -90,7 +97,7 @@ class MambaLM(nn.Module):
         self.config = lm_config.to_mamba_config()
 
         self.embedding = nn.Embedding(self.lm_config.vocab_size, self.config.d_model)
-        self.mamba = Mamba(self.config)
+        self.mamba = Mamba(**self.config.__dict__)
         self.norm_f = RMSNorm(self.config.d_model)
 
         self.lm_head = nn.Linear(self.config.d_model, self.lm_config.vocab_size, bias=False)