Upload 7 files

config.json

@@ -0,0 +1,39 @@
+{
+  "architectures": [
+    "BertForSequenceClassification"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoModelForSequenceClassification": "modeling_monarch_bert.MonarchBertForSequenceClassification"
+  },
+  "classifier_dropout": null,
+  "finetuning_task": "mnli",
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1",
+    "2": "LABEL_2"
+  },
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1,
+    "LABEL_2": 2
+  },
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "monarch_groups": 16,
+  "monarch_start_layer": 0,
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "transformers_version": "4.56.2",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 30522
+}

modeling_monarch_bert.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+from transformers import BertModel, BertPreTrainedModel
+from transformers.models.bert.modeling_bert import BertEncoder, BertLayer
+
+# --- 1. Monarch Low-Level Operations ---
+class MonarchUp(nn.Module):
+    def __init__(self, d_model=768, hidden_dim=3072, n_blocks=16):
+        super().__init__()
+        self.n_blocks = n_blocks
+        self.in_block = d_model // n_blocks
+        self.out_block_exp = hidden_dim // self.in_block
+        
+        self.b1 = nn.Parameter(torch.randn(n_blocks, self.in_block, self.in_block) * 0.02)
+        self.b2 = nn.Parameter(torch.randn(self.in_block, self.out_block_exp, n_blocks) * 0.02)
+        self.bias = nn.Parameter(torch.zeros(hidden_dim))
+
+    def forward(self, x):
+        B, S, D = x.shape
+        x = x.view(B, S, self.n_blocks, self.in_block)
+        x = torch.einsum('bsni,noi->bsno', x, self.b1)
+        x = x.transpose(-1, -2) # Implicit fusion friendly for Triton later
+        x = torch.einsum('bsni,noi->bsno', x, self.b2)
+        return x.reshape(B, S, -1) + self.bias
+
+class MonarchDown(nn.Module):
+    def __init__(self, d_model=768, hidden_dim=3072, n_blocks=16):
+        super().__init__()
+        self.n_blocks = n_blocks
+        self.out_block = d_model // n_blocks
+        self.in_block_exp = hidden_dim // self.out_block
+        
+        self.b1 = nn.Parameter(torch.randn(self.out_block, n_blocks, self.in_block_exp) * 0.02)
+        self.b2 = nn.Parameter(torch.randn(n_blocks, self.out_block, self.out_block) * 0.02)
+        self.bias = nn.Parameter(torch.zeros(d_model))
+
+    def forward(self, x):
+        B, S, D = x.shape
+        x = x.view(B, S, self.out_block, self.in_block_exp)
+        x = torch.einsum('bsni,noi->bsno', x, self.b1)
+        x = x.transpose(-1, -2)
+        x = torch.einsum('bsni,noi->bsno', x, self.b2)
+        return x.reshape(B, S, -1) + self.bias
+
+class MonarchFFN(nn.Module):
+    def __init__(self, d_model, hidden_dim, groups, act_fn):
+        super().__init__()
+        self.monarch_up = MonarchUp(d_model, hidden_dim, n_blocks=groups)
+        self.act = act_fn
+        self.monarch_down = MonarchDown(d_model, hidden_dim, n_blocks=groups)
+
+    def forward(self, x):
+        x = self.monarch_up(x)
+        x = self.act(x)
+        x = self.monarch_down(x)
+        return x
+
+class FFNWrapper(nn.Module):
+    def __init__(self, new_ffn_module):
+        super().__init__()
+        self.ffn = new_ffn_module
+    def forward(self, hidden_states):
+        return self.ffn(hidden_states)
+
+# --- 2. The Model Architecture ---
+class MonarchBertForSequenceClassification(BertPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        
+        # 1. Load Standard BERT
+        self.bert = BertModel(config)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        
+        # 2. Inject Monarch Layers based on Config
+        # This reconstructs the architecture exactly as it was during distillation
+        monarch_start_layer = getattr(config, "monarch_start_layer", 0) 
+        n_groups = getattr(config, "monarch_groups", 16)
+        
+        bert_layers = self.bert.encoder.layer
+        # Backward replacement logic (11 -> start_layer)
+        for i in range(11, monarch_start_layer - 1, -1):
+            d_model = config.hidden_size
+            h_dim = config.intermediate_size
+            
+            # Create Monarch Module
+            monarch_ffn = MonarchFFN(d_model, h_dim, n_groups, nn.GELU())
+            
+            # Surgery
+            bert_layers[i].intermediate = FFNWrapper(monarch_ffn)
+            bert_layers[i].output.dense = nn.Identity()
+
+        self.init_weights()
+        
+    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, labels=None):
+        outputs = self.bert(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
+        pooled_output = outputs[1]
+        logits = self.classifier(pooled_output)
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            
+        return type('SequenceClassifierOutput', (object,), {'loss': loss, 'logits': logits})()

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df08af71ced68de7c27a397d9fe6324915c9d23f44d9f846bd01b7009b3ccfef
+size 219794023

special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

tokenizer_config.json

@@ -0,0 +1,58 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "101": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "102": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "103": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "do_basic_tokenize": true,
+  "do_lower_case": true,
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "model_max_length": 512,
+  "never_split": null,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]"
+}