Update architecture and tokenizer

huggingface.py

@@ -120,7 +120,7 @@ class ShramForCausalLM(PreTrainedModel, GenerationMixin):
             num_local_heads=self.config.num_sliding_window_heads,
             local_head_dim=self.config.head_dim,
             num_mosrah_heads=self.config.num_mosrah_heads,
-            mosrah_head_dim=self.config.hidden_size // self.config.num_selected_heads,
+            mosrah_head_dim=self.config.head_dim,
             batch_size=batch_size,
             device=device,
         )