Upload all training checkpoint weights

chck_100M/config.json

@@ -0,0 +1,17 @@
+{
+  "auto_map": {
+    "AutoConfig": "modeling_msit.MSITGPTBERTHFConfig",
+    "AutoModel": "modeling_msit.MSITGPTBERTModelWrapper",
+    "AutoModelForCausalLM": "modeling_msit.MSITGPTBERTForCausalLM"
+  },
+  "vocab_size": 16384,
+  "block_size": 512,
+  "d_model": 384,
+  "hidden_size": 384,
+  "d_thin": 192,
+  "num_layers": 6,
+  "num_blocks": 6,
+  "capacity_factor": 2.0,
+  "dropout": 0.1,
+  "model_type": "msit_gptbert"
+}

chck_100M/configuration_msit.py

@@ -0,0 +1,35 @@
+from transformers import PretrainedConfig
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)

chck_100M/modeling_msit.py

@@ -0,0 +1,436 @@
+import os
+import math
+import random
+import inspect
+from typing import Optional, Tuple, Dict, Any
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin, AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+# ─────────────────────────────────────────────────────────────
+# Configuration Classes
+# ─────────────────────────────────────────────────────────────
+
+@dataclass
+class MSITGPTBERTConfig:
+    vocab_size: int = 16384
+    block_size: int = 512
+    d_model: int = 384
+    d_thin: int = 192
+    num_layers: int = 6
+    num_blocks: int = 6
+    capacity_factor: float = 2.0
+    dropout: float = 0.1
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)
+
+# ─────────────────────────────────────────────────────────────
+# ROPE HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def _precompute_rope_freqs(head_dim: int, seq_len: int, device: torch.device, theta: float = 10000.0):
+    assert head_dim % 2 == 0, "head_dim must be divisible by 2 for RoPE"
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos(), emb.sin()
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    L = x.size(2)
+    cos = cos[:L, :].unsqueeze(0).unsqueeze(1) 
+    sin = sin[:L, :].unsqueeze(0).unsqueeze(1) 
+    
+    half_dim = x.size(-1) // 2
+    x1 = x[..., :half_dim]
+    x2 = x[..., half_dim:]
+    rotated_x = torch.cat((-x2, x1), dim=-1)
+    
+    return (x * cos) + (rotated_x * sin)
+
+# ─────────────────────────────────────────────────────────────
+# 1.  SLIDING WINDOW ATTENTION
+# ─────────────────────────────────────────────────────────────
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size=None):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.num_heads  = num_heads
+        self.window_size = window_size
+        self.head_dim   = dim // num_heads
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        B, L, D = x.size()
+
+        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            past_len = past_k.size(2)
+            q_cos, q_sin = _precompute_rope_freqs(self.head_dim, past_len + L, x.device)
+            q = _apply_rope(q, q_cos[past_len:, :], q_sin[past_len:, :])
+            k = _apply_rope(k, q_cos[past_len:, :], q_sin[past_len:, :])
+            
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        else:
+            cos, sin = _precompute_rope_freqs(self.head_dim, L, x.device)
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+
+        L_kv = k.size(2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        if bidirectional:
+            if self.window_size is not None:
+                past_len = L_kv - L
+                pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+                pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+                dist  = torch.abs(pos_i - pos_j)
+                win_mask = dist < self.window_size
+                scores = scores.masked_fill(
+                    ~win_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+                )
+        else:
+            past_len = L_kv - L
+            pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+            pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+            dist  = pos_i - pos_j
+
+            causal_mask = dist >= 0
+            if self.window_size is not None:
+                causal_mask = causal_mask & (dist < self.window_size)
+
+            scores = scores.masked_fill(
+                ~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+            )
+
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn, v)                                
+        out = out.transpose(1, 2).contiguous().view(B, L, D)
+        out = self.o_proj(out)
+
+        if use_cache:
+            if self.window_size is not None:
+                present_kv = (
+                    k[:, :, -self.window_size:, :],
+                    v[:, :, -self.window_size:, :]
+                )
+            else:
+                present_kv = (k, v)
+        else:
+            present_kv = None
+
+        return out, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 2.  MSIT BRANCH BLOCK (Pre-Norm)
+# ─────────────────────────────────────────────────────────────
+
+class MSITBranchBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size):
+        super().__init__()
+        self.ln1  = nn.LayerNorm(dim)
+        self.attn = SlidingWindowAttention(dim, num_heads, window_size)
+        self.ln2  = nn.LayerNorm(dim)
+        self.ffn  = nn.Sequential(
+            nn.Linear(dim, dim * 4, bias=False),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim, bias=False),
+        )
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        attn_out, present_kv = self.attn(
+            self.ln1(x), past_kv, use_cache, bidirectional
+        )
+        x = x + attn_out
+        x = x + self.ffn(self.ln2(x))
+        return x, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 3.  MoEP-MSIT ARCHITECTURE BLOCK  (Expert Choice routing)
+# ─────────────────────────────────────────────────────────────
+
+class MoEPMSITBlock(nn.Module):
+    def __init__(self, d_model: int = 512, d_thin: int = 192, num_blocks: int = 14,
+                 capacity_factor: float = 2.0):
+        super().__init__()
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+
+        # 1. Global Block (Dense, d_model)
+        num_heads_global = max(1, d_model // 64)
+        self.global_block = MSITBranchBlock(d_model, num_heads_global, window_size=None)
+
+        # 2. Router
+        self.router_ln = nn.LayerNorm(d_model)
+        self.w_router = nn.Linear(d_model, num_blocks, bias=False)
+
+        # 3. Shrink Projection
+        self.w_down = nn.Linear(d_model, d_thin, bias=False)
+
+        # 4. Thin Parallel Blocks (d_thin)
+        self.windows = [None, 16, 8, 4, 2] + [None] * (num_blocks - 5)
+        self.heads = [max(1, d_thin // 64)] * num_blocks  
+        
+        self.thin_blocks = nn.ModuleList([
+            MSITBranchBlock(d_thin, self.heads[i], self.windows[i])
+            for i in range(num_blocks)
+        ])
+
+        # 6. Grow Projection
+        self.w_up = nn.Linear(d_thin, d_model, bias=False)
+        self.last_topk_indices = None
+
+    def forward(self, x_0: torch.Tensor, past_kvs=None, use_cache: bool = False, bidirectional: bool = False):
+        B, T, D = x_0.size()
+        n_tokens = B * T
+
+        # Step 1: Global Block
+        pkv_g = past_kvs[0] if past_kvs else None
+        x_1, nkv_g = self.global_block(x_0, pkv_g, use_cache, bidirectional)
+
+        # Step 2: Gated input stream
+        x_2 = x_1 + x_0
+
+        # Router scores
+        r_logits = self.w_router(self.router_ln(x_2)).view(n_tokens, self.num_blocks)
+        r_probs = F.softmax(r_logits, dim=-1)
+
+        # Per-expert capacity: k = (n * c) / e
+        k_capacity = max(1, int(round(n_tokens * self.capacity_factor / self.num_blocks)))
+        k_capacity = min(k_capacity, n_tokens)
+
+        # Expert Choice routing: topk over the token axis for each expert
+        expert_token_scores = r_probs.transpose(0, 1) # (num_blocks, n_tokens)
+        topk_scores, topk_token_idx = torch.topk(expert_token_scores, k_capacity, dim=-1)
+        self.last_topk_indices = topk_token_idx
+
+        # Load balancing is guaranteed by construction in Expert Choice
+        layer_aux_loss = x_2.new_zeros(())
+
+        # Shrink projection
+        x_2_thin = self.w_down(x_2) # (B, T, d_thin)
+        x_2_thin_flat = x_2_thin.view(n_tokens, self.d_thin)
+
+        # Expert computations
+        new_kvs = [nkv_g]
+        expert_outputs_flat = torch.zeros(n_tokens, self.d_model, device=x_2.device, dtype=x_2.dtype)
+
+        for i, block in enumerate(self.thin_blocks):
+            sel_idx = topk_token_idx[i]
+            bucket_in = x_2_thin_flat[sel_idx].unsqueeze(0) # (1, k_capacity, d_thin)
+
+            pkv_i = past_kvs[i + 1] if past_kvs else None
+            bucket_out, nkv_i = block(bucket_in, pkv_i, use_cache, bidirectional)
+            if use_cache:
+                new_kvs.append(nkv_i)
+
+            bucket_out = bucket_out.squeeze(0) # (k_capacity, d_thin)
+            bucket_out_full = self.w_up(bucket_out) # (k_capacity, d_model)
+
+            gate = topk_scores[i].unsqueeze(-1) # (k_capacity, 1)
+            expert_outputs_flat.index_add_(0, sel_idx, bucket_out_full * gate)
+
+        x_3_full = expert_outputs_flat.view(B, T, self.d_model)
+        out = x_2 + x_3_full
+
+        present_kvs = tuple(new_kvs) if use_cache else None
+        return out, present_kvs, layer_aux_loss
+
+# ─────────────────────────────────────────────────────────────
+# 4.  RAW MSIT-GPT-BERT MODEL
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.wte      = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop_emb = nn.Dropout(cfg.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(
+                d_model=cfg.d_model, 
+                d_thin=cfg.d_thin, 
+                num_blocks=cfg.num_blocks, 
+                capacity_factor=cfg.capacity_factor
+            )
+            for _ in range(cfg.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(cfg.d_model)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.wte.weight = self.lm_head.weight
+
+    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None, bidirectional: bool = False):
+        x = self.drop_emb(self.wte(input_ids))
+        total_aux_loss = 0.0
+        
+        for block in self.blocks:
+            x, _, layer_aux = block(x, past_kvs=None, use_cache=False, bidirectional=bidirectional)
+            total_aux_loss += layer_aux
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            ce_loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), targets.view(-1), ignore_index=-100)
+            avg_aux_loss = total_aux_loss / self.cfg.num_layers
+            loss = ce_loss + (0.01 * avg_aux_loss)
+
+        return logits, loss
+
+# ─────────────────────────────────────────────────────────────
+# 5.  HUGGING FACE MODEL WRAPPERS
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModelWrapper(PreTrainedModel):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.wte      = nn.Embedding(config.vocab_size, config.d_model)
+        self.drop_emb = nn.Dropout(config.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(config.d_model, config.d_thin, config.num_blocks, config.capacity_factor)
+            for _ in range(config.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(config.d_model)
+        self.post_init()
+
+    def forward(self, input_ids, **kwargs):
+        x = self.drop_emb(self.wte(input_ids))
+        for block in self.blocks:
+            x, _, _ = block(x, past_kvs=None, use_cache=False, bidirectional=False)
+        x = self.ln_f(x)
+        return BaseModelOutputWithPast(last_hidden_state=x)
+
+
+class MSITGPTBERTForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+    _no_split_modules = ["MoEPMSITBlock"]
+    _tied_weights_keys = {"transformer.lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.transformer = MSITGPTBERTModelWrapper(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.post_init()
+        
+        # State-dict pre-hook for backwards compatibility with checkpoint key naming
+        def _prefix_cleaner(state_dict, prefix, local_metadata, Moore, missing_keys, unexpected_keys, error_msgs):
+            keys = list(state_dict.keys())
+            for k in keys:
+                if k.startswith("transformer."):
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                elif f"{prefix}transformer." in k:
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                    
+        self._register_load_state_dict_pre_hook(_prefix_cleaner)
+
+    def tie_weights(self, **kwargs):
+        if hasattr(self, "transformer") and hasattr(self.transformer, "wte") and hasattr(self.transformer, "lm_head"):
+            self.transformer.wte.weight = self.lm_head.weight
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(self,
+                input_ids: Optional[torch.LongTensor] = None,
+                attention_mask: Optional[torch.FloatTensor] = None,
+                labels: Optional[torch.LongTensor] = None,
+                **kwargs) -> CausalLMOutputWithPast:
+        outputs = self.transformer(input_ids)
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# Register with auto-mapping
+AutoConfig.register("msit_gptbert", MSITGPTBERTHFConfig)
+AutoModel.register(MSITGPTBERTHFConfig, MSITGPTBERTModelWrapper)
+AutoModelForCausalLM.register(MSITGPTBERTHFConfig, MSITGPTBERTForCausalLM)

chck_100M/pytorch_model.bin

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

chck_100M/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "[UNK]"
+}

chck_10M/config.json

@@ -0,0 +1,17 @@
+{
+  "auto_map": {
+    "AutoConfig": "modeling_msit.MSITGPTBERTHFConfig",
+    "AutoModel": "modeling_msit.MSITGPTBERTModelWrapper",
+    "AutoModelForCausalLM": "modeling_msit.MSITGPTBERTForCausalLM"
+  },
+  "vocab_size": 16384,
+  "block_size": 512,
+  "d_model": 384,
+  "hidden_size": 384,
+  "d_thin": 192,
+  "num_layers": 6,
+  "num_blocks": 6,
+  "capacity_factor": 2.0,
+  "dropout": 0.1,
+  "model_type": "msit_gptbert"
+}

chck_10M/configuration_msit.py

@@ -0,0 +1,35 @@
+from transformers import PretrainedConfig
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)

chck_10M/modeling_msit.py

@@ -0,0 +1,436 @@
+import os
+import math
+import random
+import inspect
+from typing import Optional, Tuple, Dict, Any
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin, AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+# ─────────────────────────────────────────────────────────────
+# Configuration Classes
+# ─────────────────────────────────────────────────────────────
+
+@dataclass
+class MSITGPTBERTConfig:
+    vocab_size: int = 16384
+    block_size: int = 512
+    d_model: int = 384
+    d_thin: int = 192
+    num_layers: int = 6
+    num_blocks: int = 6
+    capacity_factor: float = 2.0
+    dropout: float = 0.1
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)
+
+# ─────────────────────────────────────────────────────────────
+# ROPE HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def _precompute_rope_freqs(head_dim: int, seq_len: int, device: torch.device, theta: float = 10000.0):
+    assert head_dim % 2 == 0, "head_dim must be divisible by 2 for RoPE"
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos(), emb.sin()
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    L = x.size(2)
+    cos = cos[:L, :].unsqueeze(0).unsqueeze(1) 
+    sin = sin[:L, :].unsqueeze(0).unsqueeze(1) 
+    
+    half_dim = x.size(-1) // 2
+    x1 = x[..., :half_dim]
+    x2 = x[..., half_dim:]
+    rotated_x = torch.cat((-x2, x1), dim=-1)
+    
+    return (x * cos) + (rotated_x * sin)
+
+# ─────────────────────────────────────────────────────────────
+# 1.  SLIDING WINDOW ATTENTION
+# ─────────────────────────────────────────────────────────────
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size=None):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.num_heads  = num_heads
+        self.window_size = window_size
+        self.head_dim   = dim // num_heads
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        B, L, D = x.size()
+
+        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            past_len = past_k.size(2)
+            q_cos, q_sin = _precompute_rope_freqs(self.head_dim, past_len + L, x.device)
+            q = _apply_rope(q, q_cos[past_len:, :], q_sin[past_len:, :])
+            k = _apply_rope(k, q_cos[past_len:, :], q_sin[past_len:, :])
+            
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        else:
+            cos, sin = _precompute_rope_freqs(self.head_dim, L, x.device)
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+
+        L_kv = k.size(2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        if bidirectional:
+            if self.window_size is not None:
+                past_len = L_kv - L
+                pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+                pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+                dist  = torch.abs(pos_i - pos_j)
+                win_mask = dist < self.window_size
+                scores = scores.masked_fill(
+                    ~win_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+                )
+        else:
+            past_len = L_kv - L
+            pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+            pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+            dist  = pos_i - pos_j
+
+            causal_mask = dist >= 0
+            if self.window_size is not None:
+                causal_mask = causal_mask & (dist < self.window_size)
+
+            scores = scores.masked_fill(
+                ~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+            )
+
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn, v)                                
+        out = out.transpose(1, 2).contiguous().view(B, L, D)
+        out = self.o_proj(out)
+
+        if use_cache:
+            if self.window_size is not None:
+                present_kv = (
+                    k[:, :, -self.window_size:, :],
+                    v[:, :, -self.window_size:, :]
+                )
+            else:
+                present_kv = (k, v)
+        else:
+            present_kv = None
+
+        return out, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 2.  MSIT BRANCH BLOCK (Pre-Norm)
+# ─────────────────────────────────────────────────────────────
+
+class MSITBranchBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size):
+        super().__init__()
+        self.ln1  = nn.LayerNorm(dim)
+        self.attn = SlidingWindowAttention(dim, num_heads, window_size)
+        self.ln2  = nn.LayerNorm(dim)
+        self.ffn  = nn.Sequential(
+            nn.Linear(dim, dim * 4, bias=False),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim, bias=False),
+        )
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        attn_out, present_kv = self.attn(
+            self.ln1(x), past_kv, use_cache, bidirectional
+        )
+        x = x + attn_out
+        x = x + self.ffn(self.ln2(x))
+        return x, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 3.  MoEP-MSIT ARCHITECTURE BLOCK  (Expert Choice routing)
+# ─────────────────────────────────────────────────────────────
+
+class MoEPMSITBlock(nn.Module):
+    def __init__(self, d_model: int = 512, d_thin: int = 192, num_blocks: int = 14,
+                 capacity_factor: float = 2.0):
+        super().__init__()
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+
+        # 1. Global Block (Dense, d_model)
+        num_heads_global = max(1, d_model // 64)
+        self.global_block = MSITBranchBlock(d_model, num_heads_global, window_size=None)
+
+        # 2. Router
+        self.router_ln = nn.LayerNorm(d_model)
+        self.w_router = nn.Linear(d_model, num_blocks, bias=False)
+
+        # 3. Shrink Projection
+        self.w_down = nn.Linear(d_model, d_thin, bias=False)
+
+        # 4. Thin Parallel Blocks (d_thin)
+        self.windows = [None, 16, 8, 4, 2] + [None] * (num_blocks - 5)
+        self.heads = [max(1, d_thin // 64)] * num_blocks  
+        
+        self.thin_blocks = nn.ModuleList([
+            MSITBranchBlock(d_thin, self.heads[i], self.windows[i])
+            for i in range(num_blocks)
+        ])
+
+        # 6. Grow Projection
+        self.w_up = nn.Linear(d_thin, d_model, bias=False)
+        self.last_topk_indices = None
+
+    def forward(self, x_0: torch.Tensor, past_kvs=None, use_cache: bool = False, bidirectional: bool = False):
+        B, T, D = x_0.size()
+        n_tokens = B * T
+
+        # Step 1: Global Block
+        pkv_g = past_kvs[0] if past_kvs else None
+        x_1, nkv_g = self.global_block(x_0, pkv_g, use_cache, bidirectional)
+
+        # Step 2: Gated input stream
+        x_2 = x_1 + x_0
+
+        # Router scores
+        r_logits = self.w_router(self.router_ln(x_2)).view(n_tokens, self.num_blocks)
+        r_probs = F.softmax(r_logits, dim=-1)
+
+        # Per-expert capacity: k = (n * c) / e
+        k_capacity = max(1, int(round(n_tokens * self.capacity_factor / self.num_blocks)))
+        k_capacity = min(k_capacity, n_tokens)
+
+        # Expert Choice routing: topk over the token axis for each expert
+        expert_token_scores = r_probs.transpose(0, 1) # (num_blocks, n_tokens)
+        topk_scores, topk_token_idx = torch.topk(expert_token_scores, k_capacity, dim=-1)
+        self.last_topk_indices = topk_token_idx
+
+        # Load balancing is guaranteed by construction in Expert Choice
+        layer_aux_loss = x_2.new_zeros(())
+
+        # Shrink projection
+        x_2_thin = self.w_down(x_2) # (B, T, d_thin)
+        x_2_thin_flat = x_2_thin.view(n_tokens, self.d_thin)
+
+        # Expert computations
+        new_kvs = [nkv_g]
+        expert_outputs_flat = torch.zeros(n_tokens, self.d_model, device=x_2.device, dtype=x_2.dtype)
+
+        for i, block in enumerate(self.thin_blocks):
+            sel_idx = topk_token_idx[i]
+            bucket_in = x_2_thin_flat[sel_idx].unsqueeze(0) # (1, k_capacity, d_thin)
+
+            pkv_i = past_kvs[i + 1] if past_kvs else None
+            bucket_out, nkv_i = block(bucket_in, pkv_i, use_cache, bidirectional)
+            if use_cache:
+                new_kvs.append(nkv_i)
+
+            bucket_out = bucket_out.squeeze(0) # (k_capacity, d_thin)
+            bucket_out_full = self.w_up(bucket_out) # (k_capacity, d_model)
+
+            gate = topk_scores[i].unsqueeze(-1) # (k_capacity, 1)
+            expert_outputs_flat.index_add_(0, sel_idx, bucket_out_full * gate)
+
+        x_3_full = expert_outputs_flat.view(B, T, self.d_model)
+        out = x_2 + x_3_full
+
+        present_kvs = tuple(new_kvs) if use_cache else None
+        return out, present_kvs, layer_aux_loss
+
+# ─────────────────────────────────────────────────────────────
+# 4.  RAW MSIT-GPT-BERT MODEL
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.wte      = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop_emb = nn.Dropout(cfg.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(
+                d_model=cfg.d_model, 
+                d_thin=cfg.d_thin, 
+                num_blocks=cfg.num_blocks, 
+                capacity_factor=cfg.capacity_factor
+            )
+            for _ in range(cfg.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(cfg.d_model)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.wte.weight = self.lm_head.weight
+
+    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None, bidirectional: bool = False):
+        x = self.drop_emb(self.wte(input_ids))
+        total_aux_loss = 0.0
+        
+        for block in self.blocks:
+            x, _, layer_aux = block(x, past_kvs=None, use_cache=False, bidirectional=bidirectional)
+            total_aux_loss += layer_aux
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            ce_loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), targets.view(-1), ignore_index=-100)
+            avg_aux_loss = total_aux_loss / self.cfg.num_layers
+            loss = ce_loss + (0.01 * avg_aux_loss)
+
+        return logits, loss
+
+# ─────────────────────────────────────────────────────────────
+# 5.  HUGGING FACE MODEL WRAPPERS
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModelWrapper(PreTrainedModel):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.wte      = nn.Embedding(config.vocab_size, config.d_model)
+        self.drop_emb = nn.Dropout(config.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(config.d_model, config.d_thin, config.num_blocks, config.capacity_factor)
+            for _ in range(config.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(config.d_model)
+        self.post_init()
+
+    def forward(self, input_ids, **kwargs):
+        x = self.drop_emb(self.wte(input_ids))
+        for block in self.blocks:
+            x, _, _ = block(x, past_kvs=None, use_cache=False, bidirectional=False)
+        x = self.ln_f(x)
+        return BaseModelOutputWithPast(last_hidden_state=x)
+
+
+class MSITGPTBERTForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+    _no_split_modules = ["MoEPMSITBlock"]
+    _tied_weights_keys = {"transformer.lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.transformer = MSITGPTBERTModelWrapper(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.post_init()
+        
+        # State-dict pre-hook for backwards compatibility with checkpoint key naming
+        def _prefix_cleaner(state_dict, prefix, local_metadata, Moore, missing_keys, unexpected_keys, error_msgs):
+            keys = list(state_dict.keys())
+            for k in keys:
+                if k.startswith("transformer."):
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                elif f"{prefix}transformer." in k:
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                    
+        self._register_load_state_dict_pre_hook(_prefix_cleaner)
+
+    def tie_weights(self, **kwargs):
+        if hasattr(self, "transformer") and hasattr(self.transformer, "wte") and hasattr(self.transformer, "lm_head"):
+            self.transformer.wte.weight = self.lm_head.weight
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(self,
+                input_ids: Optional[torch.LongTensor] = None,
+                attention_mask: Optional[torch.FloatTensor] = None,
+                labels: Optional[torch.LongTensor] = None,
+                **kwargs) -> CausalLMOutputWithPast:
+        outputs = self.transformer(input_ids)
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# Register with auto-mapping
+AutoConfig.register("msit_gptbert", MSITGPTBERTHFConfig)
+AutoModel.register(MSITGPTBERTHFConfig, MSITGPTBERTModelWrapper)
+AutoModelForCausalLM.register(MSITGPTBERTHFConfig, MSITGPTBERTForCausalLM)

chck_10M/pytorch_model.bin

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

chck_10M/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "[UNK]"
+}

chck_1M/config.json

@@ -0,0 +1,17 @@
+{
+  "auto_map": {
+    "AutoConfig": "modeling_msit.MSITGPTBERTHFConfig",
+    "AutoModel": "modeling_msit.MSITGPTBERTModelWrapper",
+    "AutoModelForCausalLM": "modeling_msit.MSITGPTBERTForCausalLM"
+  },
+  "vocab_size": 16384,
+  "block_size": 512,
+  "d_model": 384,
+  "hidden_size": 384,
+  "d_thin": 192,
+  "num_layers": 6,
+  "num_blocks": 6,
+  "capacity_factor": 2.0,
+  "dropout": 0.1,
+  "model_type": "msit_gptbert"
+}

chck_1M/configuration_msit.py

@@ -0,0 +1,35 @@
+from transformers import PretrainedConfig
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)

chck_1M/modeling_msit.py

@@ -0,0 +1,436 @@
+import os
+import math
+import random
+import inspect
+from typing import Optional, Tuple, Dict, Any
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin, AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+# ─────────────────────────────────────────────────────────────
+# Configuration Classes
+# ─────────────────────────────────────────────────────────────
+
+@dataclass
+class MSITGPTBERTConfig:
+    vocab_size: int = 16384
+    block_size: int = 512
+    d_model: int = 384
+    d_thin: int = 192
+    num_layers: int = 6
+    num_blocks: int = 6
+    capacity_factor: float = 2.0
+    dropout: float = 0.1
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)
+
+# ─────────────────────────────────────────────────────────────
+# ROPE HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def _precompute_rope_freqs(head_dim: int, seq_len: int, device: torch.device, theta: float = 10000.0):
+    assert head_dim % 2 == 0, "head_dim must be divisible by 2 for RoPE"
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos(), emb.sin()
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    L = x.size(2)
+    cos = cos[:L, :].unsqueeze(0).unsqueeze(1) 
+    sin = sin[:L, :].unsqueeze(0).unsqueeze(1) 
+    
+    half_dim = x.size(-1) // 2
+    x1 = x[..., :half_dim]
+    x2 = x[..., half_dim:]
+    rotated_x = torch.cat((-x2, x1), dim=-1)
+    
+    return (x * cos) + (rotated_x * sin)
+
+# ─────────────────────────────────────────────────────────────
+# 1.  SLIDING WINDOW ATTENTION
+# ─────────────────────────────────────────────────────────────
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size=None):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.num_heads  = num_heads
+        self.window_size = window_size
+        self.head_dim   = dim // num_heads
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        B, L, D = x.size()
+
+        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            past_len = past_k.size(2)
+            q_cos, q_sin = _precompute_rope_freqs(self.head_dim, past_len + L, x.device)
+            q = _apply_rope(q, q_cos[past_len:, :], q_sin[past_len:, :])
+            k = _apply_rope(k, q_cos[past_len:, :], q_sin[past_len:, :])
+            
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        else:
+            cos, sin = _precompute_rope_freqs(self.head_dim, L, x.device)
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+
+        L_kv = k.size(2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        if bidirectional:
+            if self.window_size is not None:
+                past_len = L_kv - L
+                pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+                pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+                dist  = torch.abs(pos_i - pos_j)
+                win_mask = dist < self.window_size
+                scores = scores.masked_fill(
+                    ~win_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+                )
+        else:
+            past_len = L_kv - L
+            pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+            pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+            dist  = pos_i - pos_j
+
+            causal_mask = dist >= 0
+            if self.window_size is not None:
+                causal_mask = causal_mask & (dist < self.window_size)
+
+            scores = scores.masked_fill(
+                ~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+            )
+
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn, v)                                
+        out = out.transpose(1, 2).contiguous().view(B, L, D)
+        out = self.o_proj(out)
+
+        if use_cache:
+            if self.window_size is not None:
+                present_kv = (
+                    k[:, :, -self.window_size:, :],
+                    v[:, :, -self.window_size:, :]
+                )
+            else:
+                present_kv = (k, v)
+        else:
+            present_kv = None
+
+        return out, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 2.  MSIT BRANCH BLOCK (Pre-Norm)
+# ─────────────────────────────────────────────────────────────
+
+class MSITBranchBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size):
+        super().__init__()
+        self.ln1  = nn.LayerNorm(dim)
+        self.attn = SlidingWindowAttention(dim, num_heads, window_size)
+        self.ln2  = nn.LayerNorm(dim)
+        self.ffn  = nn.Sequential(
+            nn.Linear(dim, dim * 4, bias=False),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim, bias=False),
+        )
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        attn_out, present_kv = self.attn(
+            self.ln1(x), past_kv, use_cache, bidirectional
+        )
+        x = x + attn_out
+        x = x + self.ffn(self.ln2(x))
+        return x, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 3.  MoEP-MSIT ARCHITECTURE BLOCK  (Expert Choice routing)
+# ─────────────────────────────────────────────────────────────
+
+class MoEPMSITBlock(nn.Module):
+    def __init__(self, d_model: int = 512, d_thin: int = 192, num_blocks: int = 14,
+                 capacity_factor: float = 2.0):
+        super().__init__()
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+
+        # 1. Global Block (Dense, d_model)
+        num_heads_global = max(1, d_model // 64)
+        self.global_block = MSITBranchBlock(d_model, num_heads_global, window_size=None)
+
+        # 2. Router
+        self.router_ln = nn.LayerNorm(d_model)
+        self.w_router = nn.Linear(d_model, num_blocks, bias=False)
+
+        # 3. Shrink Projection
+        self.w_down = nn.Linear(d_model, d_thin, bias=False)
+
+        # 4. Thin Parallel Blocks (d_thin)
+        self.windows = [None, 16, 8, 4, 2] + [None] * (num_blocks - 5)
+        self.heads = [max(1, d_thin // 64)] * num_blocks  
+        
+        self.thin_blocks = nn.ModuleList([
+            MSITBranchBlock(d_thin, self.heads[i], self.windows[i])
+            for i in range(num_blocks)
+        ])
+
+        # 6. Grow Projection
+        self.w_up = nn.Linear(d_thin, d_model, bias=False)
+        self.last_topk_indices = None
+
+    def forward(self, x_0: torch.Tensor, past_kvs=None, use_cache: bool = False, bidirectional: bool = False):
+        B, T, D = x_0.size()
+        n_tokens = B * T
+
+        # Step 1: Global Block
+        pkv_g = past_kvs[0] if past_kvs else None
+        x_1, nkv_g = self.global_block(x_0, pkv_g, use_cache, bidirectional)
+
+        # Step 2: Gated input stream
+        x_2 = x_1 + x_0
+
+        # Router scores
+        r_logits = self.w_router(self.router_ln(x_2)).view(n_tokens, self.num_blocks)
+        r_probs = F.softmax(r_logits, dim=-1)
+
+        # Per-expert capacity: k = (n * c) / e
+        k_capacity = max(1, int(round(n_tokens * self.capacity_factor / self.num_blocks)))
+        k_capacity = min(k_capacity, n_tokens)
+
+        # Expert Choice routing: topk over the token axis for each expert
+        expert_token_scores = r_probs.transpose(0, 1) # (num_blocks, n_tokens)
+        topk_scores, topk_token_idx = torch.topk(expert_token_scores, k_capacity, dim=-1)
+        self.last_topk_indices = topk_token_idx
+
+        # Load balancing is guaranteed by construction in Expert Choice
+        layer_aux_loss = x_2.new_zeros(())
+
+        # Shrink projection
+        x_2_thin = self.w_down(x_2) # (B, T, d_thin)
+        x_2_thin_flat = x_2_thin.view(n_tokens, self.d_thin)
+
+        # Expert computations
+        new_kvs = [nkv_g]
+        expert_outputs_flat = torch.zeros(n_tokens, self.d_model, device=x_2.device, dtype=x_2.dtype)
+
+        for i, block in enumerate(self.thin_blocks):
+            sel_idx = topk_token_idx[i]
+            bucket_in = x_2_thin_flat[sel_idx].unsqueeze(0) # (1, k_capacity, d_thin)
+
+            pkv_i = past_kvs[i + 1] if past_kvs else None
+            bucket_out, nkv_i = block(bucket_in, pkv_i, use_cache, bidirectional)
+            if use_cache:
+                new_kvs.append(nkv_i)
+
+            bucket_out = bucket_out.squeeze(0) # (k_capacity, d_thin)
+            bucket_out_full = self.w_up(bucket_out) # (k_capacity, d_model)
+
+            gate = topk_scores[i].unsqueeze(-1) # (k_capacity, 1)
+            expert_outputs_flat.index_add_(0, sel_idx, bucket_out_full * gate)
+
+        x_3_full = expert_outputs_flat.view(B, T, self.d_model)
+        out = x_2 + x_3_full
+
+        present_kvs = tuple(new_kvs) if use_cache else None
+        return out, present_kvs, layer_aux_loss
+
+# ─────────────────────────────────────────────────────────────
+# 4.  RAW MSIT-GPT-BERT MODEL
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.wte      = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop_emb = nn.Dropout(cfg.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(
+                d_model=cfg.d_model, 
+                d_thin=cfg.d_thin, 
+                num_blocks=cfg.num_blocks, 
+                capacity_factor=cfg.capacity_factor
+            )
+            for _ in range(cfg.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(cfg.d_model)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.wte.weight = self.lm_head.weight
+
+    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None, bidirectional: bool = False):
+        x = self.drop_emb(self.wte(input_ids))
+        total_aux_loss = 0.0
+        
+        for block in self.blocks:
+            x, _, layer_aux = block(x, past_kvs=None, use_cache=False, bidirectional=bidirectional)
+            total_aux_loss += layer_aux
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            ce_loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), targets.view(-1), ignore_index=-100)
+            avg_aux_loss = total_aux_loss / self.cfg.num_layers
+            loss = ce_loss + (0.01 * avg_aux_loss)
+
+        return logits, loss
+
+# ─────────────────────────────────────────────────────────────
+# 5.  HUGGING FACE MODEL WRAPPERS
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModelWrapper(PreTrainedModel):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.wte      = nn.Embedding(config.vocab_size, config.d_model)
+        self.drop_emb = nn.Dropout(config.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(config.d_model, config.d_thin, config.num_blocks, config.capacity_factor)
+            for _ in range(config.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(config.d_model)
+        self.post_init()
+
+    def forward(self, input_ids, **kwargs):
+        x = self.drop_emb(self.wte(input_ids))
+        for block in self.blocks:
+            x, _, _ = block(x, past_kvs=None, use_cache=False, bidirectional=False)
+        x = self.ln_f(x)
+        return BaseModelOutputWithPast(last_hidden_state=x)
+
+
+class MSITGPTBERTForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+    _no_split_modules = ["MoEPMSITBlock"]
+    _tied_weights_keys = {"transformer.lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.transformer = MSITGPTBERTModelWrapper(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.post_init()
+        
+        # State-dict pre-hook for backwards compatibility with checkpoint key naming
+        def _prefix_cleaner(state_dict, prefix, local_metadata, Moore, missing_keys, unexpected_keys, error_msgs):
+            keys = list(state_dict.keys())
+            for k in keys:
+                if k.startswith("transformer."):
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                elif f"{prefix}transformer." in k:
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                    
+        self._register_load_state_dict_pre_hook(_prefix_cleaner)
+
+    def tie_weights(self, **kwargs):
+        if hasattr(self, "transformer") and hasattr(self.transformer, "wte") and hasattr(self.transformer, "lm_head"):
+            self.transformer.wte.weight = self.lm_head.weight
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(self,
+                input_ids: Optional[torch.LongTensor] = None,
+                attention_mask: Optional[torch.FloatTensor] = None,
+                labels: Optional[torch.LongTensor] = None,
+                **kwargs) -> CausalLMOutputWithPast:
+        outputs = self.transformer(input_ids)
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# Register with auto-mapping
+AutoConfig.register("msit_gptbert", MSITGPTBERTHFConfig)
+AutoModel.register(MSITGPTBERTHFConfig, MSITGPTBERTModelWrapper)
+AutoModelForCausalLM.register(MSITGPTBERTHFConfig, MSITGPTBERTForCausalLM)

chck_1M/pytorch_model.bin

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

chck_1M/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "[UNK]"
+}

chck_20M/config.json

@@ -0,0 +1,17 @@
+{
+  "auto_map": {
+    "AutoConfig": "modeling_msit.MSITGPTBERTHFConfig",
+    "AutoModel": "modeling_msit.MSITGPTBERTModelWrapper",
+    "AutoModelForCausalLM": "modeling_msit.MSITGPTBERTForCausalLM"
+  },
+  "vocab_size": 16384,
+  "block_size": 512,
+  "d_model": 384,
+  "hidden_size": 384,
+  "d_thin": 192,
+  "num_layers": 6,
+  "num_blocks": 6,
+  "capacity_factor": 2.0,
+  "dropout": 0.1,
+  "model_type": "msit_gptbert"
+}

chck_20M/configuration_msit.py

@@ -0,0 +1,35 @@
+from transformers import PretrainedConfig
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)

chck_20M/modeling_msit.py

@@ -0,0 +1,436 @@
+import os
+import math
+import random
+import inspect
+from typing import Optional, Tuple, Dict, Any
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin, AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+# ─────────────────────────────────────────────────────────────
+# Configuration Classes
+# ─────────────────────────────────────────────────────────────
+
+@dataclass
+class MSITGPTBERTConfig:
+    vocab_size: int = 16384
+    block_size: int = 512
+    d_model: int = 384
+    d_thin: int = 192
+    num_layers: int = 6
+    num_blocks: int = 6
+    capacity_factor: float = 2.0
+    dropout: float = 0.1
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)
+
+# ─────────────────────────────────────────────────────────────
+# ROPE HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def _precompute_rope_freqs(head_dim: int, seq_len: int, device: torch.device, theta: float = 10000.0):
+    assert head_dim % 2 == 0, "head_dim must be divisible by 2 for RoPE"
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos(), emb.sin()
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    L = x.size(2)
+    cos = cos[:L, :].unsqueeze(0).unsqueeze(1) 
+    sin = sin[:L, :].unsqueeze(0).unsqueeze(1) 
+    
+    half_dim = x.size(-1) // 2
+    x1 = x[..., :half_dim]
+    x2 = x[..., half_dim:]
+    rotated_x = torch.cat((-x2, x1), dim=-1)
+    
+    return (x * cos) + (rotated_x * sin)
+
+# ─────────────────────────────────────────────────────────────
+# 1.  SLIDING WINDOW ATTENTION
+# ─────────────────────────────────────────────────────────────
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size=None):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.num_heads  = num_heads
+        self.window_size = window_size
+        self.head_dim   = dim // num_heads
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        B, L, D = x.size()
+
+        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            past_len = past_k.size(2)
+            q_cos, q_sin = _precompute_rope_freqs(self.head_dim, past_len + L, x.device)
+            q = _apply_rope(q, q_cos[past_len:, :], q_sin[past_len:, :])
+            k = _apply_rope(k, q_cos[past_len:, :], q_sin[past_len:, :])
+            
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        else:
+            cos, sin = _precompute_rope_freqs(self.head_dim, L, x.device)
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+
+        L_kv = k.size(2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        if bidirectional:
+            if self.window_size is not None:
+                past_len = L_kv - L
+                pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+                pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+                dist  = torch.abs(pos_i - pos_j)
+                win_mask = dist < self.window_size
+                scores = scores.masked_fill(
+                    ~win_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+                )
+        else:
+            past_len = L_kv - L
+            pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+            pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+            dist  = pos_i - pos_j
+
+            causal_mask = dist >= 0
+            if self.window_size is not None:
+                causal_mask = causal_mask & (dist < self.window_size)
+
+            scores = scores.masked_fill(
+                ~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+            )
+
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn, v)                                
+        out = out.transpose(1, 2).contiguous().view(B, L, D)
+        out = self.o_proj(out)
+
+        if use_cache:
+            if self.window_size is not None:
+                present_kv = (
+                    k[:, :, -self.window_size:, :],
+                    v[:, :, -self.window_size:, :]
+                )
+            else:
+                present_kv = (k, v)
+        else:
+            present_kv = None
+
+        return out, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 2.  MSIT BRANCH BLOCK (Pre-Norm)
+# ─────────────────────────────────────────────────────────────
+
+class MSITBranchBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size):
+        super().__init__()
+        self.ln1  = nn.LayerNorm(dim)
+        self.attn = SlidingWindowAttention(dim, num_heads, window_size)
+        self.ln2  = nn.LayerNorm(dim)
+        self.ffn  = nn.Sequential(
+            nn.Linear(dim, dim * 4, bias=False),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim, bias=False),
+        )
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        attn_out, present_kv = self.attn(
+            self.ln1(x), past_kv, use_cache, bidirectional
+        )
+        x = x + attn_out
+        x = x + self.ffn(self.ln2(x))
+        return x, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 3.  MoEP-MSIT ARCHITECTURE BLOCK  (Expert Choice routing)
+# ─────────────────────────────────────────────────────────────
+
+class MoEPMSITBlock(nn.Module):
+    def __init__(self, d_model: int = 512, d_thin: int = 192, num_blocks: int = 14,
+                 capacity_factor: float = 2.0):
+        super().__init__()
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+
+        # 1. Global Block (Dense, d_model)
+        num_heads_global = max(1, d_model // 64)
+        self.global_block = MSITBranchBlock(d_model, num_heads_global, window_size=None)
+
+        # 2. Router
+        self.router_ln = nn.LayerNorm(d_model)
+        self.w_router = nn.Linear(d_model, num_blocks, bias=False)
+
+        # 3. Shrink Projection
+        self.w_down = nn.Linear(d_model, d_thin, bias=False)
+
+        # 4. Thin Parallel Blocks (d_thin)
+        self.windows = [None, 16, 8, 4, 2] + [None] * (num_blocks - 5)
+        self.heads = [max(1, d_thin // 64)] * num_blocks  
+        
+        self.thin_blocks = nn.ModuleList([
+            MSITBranchBlock(d_thin, self.heads[i], self.windows[i])
+            for i in range(num_blocks)
+        ])
+
+        # 6. Grow Projection
+        self.w_up = nn.Linear(d_thin, d_model, bias=False)
+        self.last_topk_indices = None
+
+    def forward(self, x_0: torch.Tensor, past_kvs=None, use_cache: bool = False, bidirectional: bool = False):
+        B, T, D = x_0.size()
+        n_tokens = B * T
+
+        # Step 1: Global Block
+        pkv_g = past_kvs[0] if past_kvs else None
+        x_1, nkv_g = self.global_block(x_0, pkv_g, use_cache, bidirectional)
+
+        # Step 2: Gated input stream
+        x_2 = x_1 + x_0
+
+        # Router scores
+        r_logits = self.w_router(self.router_ln(x_2)).view(n_tokens, self.num_blocks)
+        r_probs = F.softmax(r_logits, dim=-1)
+
+        # Per-expert capacity: k = (n * c) / e
+        k_capacity = max(1, int(round(n_tokens * self.capacity_factor / self.num_blocks)))
+        k_capacity = min(k_capacity, n_tokens)
+
+        # Expert Choice routing: topk over the token axis for each expert
+        expert_token_scores = r_probs.transpose(0, 1) # (num_blocks, n_tokens)
+        topk_scores, topk_token_idx = torch.topk(expert_token_scores, k_capacity, dim=-1)
+        self.last_topk_indices = topk_token_idx
+
+        # Load balancing is guaranteed by construction in Expert Choice
+        layer_aux_loss = x_2.new_zeros(())
+
+        # Shrink projection
+        x_2_thin = self.w_down(x_2) # (B, T, d_thin)
+        x_2_thin_flat = x_2_thin.view(n_tokens, self.d_thin)
+
+        # Expert computations
+        new_kvs = [nkv_g]
+        expert_outputs_flat = torch.zeros(n_tokens, self.d_model, device=x_2.device, dtype=x_2.dtype)
+
+        for i, block in enumerate(self.thin_blocks):
+            sel_idx = topk_token_idx[i]
+            bucket_in = x_2_thin_flat[sel_idx].unsqueeze(0) # (1, k_capacity, d_thin)
+
+            pkv_i = past_kvs[i + 1] if past_kvs else None
+            bucket_out, nkv_i = block(bucket_in, pkv_i, use_cache, bidirectional)
+            if use_cache:
+                new_kvs.append(nkv_i)
+
+            bucket_out = bucket_out.squeeze(0) # (k_capacity, d_thin)
+            bucket_out_full = self.w_up(bucket_out) # (k_capacity, d_model)
+
+            gate = topk_scores[i].unsqueeze(-1) # (k_capacity, 1)
+            expert_outputs_flat.index_add_(0, sel_idx, bucket_out_full * gate)
+
+        x_3_full = expert_outputs_flat.view(B, T, self.d_model)
+        out = x_2 + x_3_full
+
+        present_kvs = tuple(new_kvs) if use_cache else None
+        return out, present_kvs, layer_aux_loss
+
+# ─────────────────────────────────────────────────────────────
+# 4.  RAW MSIT-GPT-BERT MODEL
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.wte      = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop_emb = nn.Dropout(cfg.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(
+                d_model=cfg.d_model, 
+                d_thin=cfg.d_thin, 
+                num_blocks=cfg.num_blocks, 
+                capacity_factor=cfg.capacity_factor
+            )
+            for _ in range(cfg.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(cfg.d_model)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.wte.weight = self.lm_head.weight
+
+    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None, bidirectional: bool = False):
+        x = self.drop_emb(self.wte(input_ids))
+        total_aux_loss = 0.0
+        
+        for block in self.blocks:
+            x, _, layer_aux = block(x, past_kvs=None, use_cache=False, bidirectional=bidirectional)
+            total_aux_loss += layer_aux
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            ce_loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), targets.view(-1), ignore_index=-100)
+            avg_aux_loss = total_aux_loss / self.cfg.num_layers
+            loss = ce_loss + (0.01 * avg_aux_loss)
+
+        return logits, loss
+
+# ─────────────────────────────────────────────────────────────
+# 5.  HUGGING FACE MODEL WRAPPERS
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModelWrapper(PreTrainedModel):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.wte      = nn.Embedding(config.vocab_size, config.d_model)
+        self.drop_emb = nn.Dropout(config.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(config.d_model, config.d_thin, config.num_blocks, config.capacity_factor)
+            for _ in range(config.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(config.d_model)
+        self.post_init()
+
+    def forward(self, input_ids, **kwargs):
+        x = self.drop_emb(self.wte(input_ids))
+        for block in self.blocks:
+            x, _, _ = block(x, past_kvs=None, use_cache=False, bidirectional=False)
+        x = self.ln_f(x)
+        return BaseModelOutputWithPast(last_hidden_state=x)
+
+
+class MSITGPTBERTForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+    _no_split_modules = ["MoEPMSITBlock"]
+    _tied_weights_keys = {"transformer.lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.transformer = MSITGPTBERTModelWrapper(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.post_init()
+        
+        # State-dict pre-hook for backwards compatibility with checkpoint key naming
+        def _prefix_cleaner(state_dict, prefix, local_metadata, Moore, missing_keys, unexpected_keys, error_msgs):
+            keys = list(state_dict.keys())
+            for k in keys:
+                if k.startswith("transformer."):
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                elif f"{prefix}transformer." in k:
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                    
+        self._register_load_state_dict_pre_hook(_prefix_cleaner)
+
+    def tie_weights(self, **kwargs):
+        if hasattr(self, "transformer") and hasattr(self.transformer, "wte") and hasattr(self.transformer, "lm_head"):
+            self.transformer.wte.weight = self.lm_head.weight
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(self,
+                input_ids: Optional[torch.LongTensor] = None,
+                attention_mask: Optional[torch.FloatTensor] = None,
+                labels: Optional[torch.LongTensor] = None,
+                **kwargs) -> CausalLMOutputWithPast:
+        outputs = self.transformer(input_ids)
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# Register with auto-mapping
+AutoConfig.register("msit_gptbert", MSITGPTBERTHFConfig)
+AutoModel.register(MSITGPTBERTHFConfig, MSITGPTBERTModelWrapper)
+AutoModelForCausalLM.register(MSITGPTBERTHFConfig, MSITGPTBERTForCausalLM)

chck_20M/pytorch_model.bin

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

chck_20M/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "[UNK]"
+}

chck_2M/config.json

@@ -0,0 +1,17 @@
+{
+  "auto_map": {
+    "AutoConfig": "modeling_msit.MSITGPTBERTHFConfig",
+    "AutoModel": "modeling_msit.MSITGPTBERTModelWrapper",
+    "AutoModelForCausalLM": "modeling_msit.MSITGPTBERTForCausalLM"
+  },
+  "vocab_size": 16384,
+  "block_size": 512,
+  "d_model": 384,
+  "hidden_size": 384,
+  "d_thin": 192,
+  "num_layers": 6,
+  "num_blocks": 6,
+  "capacity_factor": 2.0,
+  "dropout": 0.1,
+  "model_type": "msit_gptbert"
+}

chck_2M/configuration_msit.py

@@ -0,0 +1,35 @@
+from transformers import PretrainedConfig
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)

chck_2M/modeling_msit.py

@@ -0,0 +1,436 @@
+import os
+import math
+import random
+import inspect
+from typing import Optional, Tuple, Dict, Any
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin, AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+# ─────────────────────────────────────────────────────────────
+# Configuration Classes
+# ─────────────────────────────────────────────────────────────
+
+@dataclass
+class MSITGPTBERTConfig:
+    vocab_size: int = 16384
+    block_size: int = 512
+    d_model: int = 384
+    d_thin: int = 192
+    num_layers: int = 6
+    num_blocks: int = 6
+    capacity_factor: float = 2.0
+    dropout: float = 0.1
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)
+
+# ─────────────────────────────────────────────────────────────
+# ROPE HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def _precompute_rope_freqs(head_dim: int, seq_len: int, device: torch.device, theta: float = 10000.0):
+    assert head_dim % 2 == 0, "head_dim must be divisible by 2 for RoPE"
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos(), emb.sin()
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    L = x.size(2)
+    cos = cos[:L, :].unsqueeze(0).unsqueeze(1) 
+    sin = sin[:L, :].unsqueeze(0).unsqueeze(1) 
+    
+    half_dim = x.size(-1) // 2
+    x1 = x[..., :half_dim]
+    x2 = x[..., half_dim:]
+    rotated_x = torch.cat((-x2, x1), dim=-1)
+    
+    return (x * cos) + (rotated_x * sin)
+
+# ─────────────────────────────────────────────────────────────
+# 1.  SLIDING WINDOW ATTENTION
+# ─────────────────────────────────────────────────────────────
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size=None):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.num_heads  = num_heads
+        self.window_size = window_size
+        self.head_dim   = dim // num_heads
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        B, L, D = x.size()
+
+        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            past_len = past_k.size(2)
+            q_cos, q_sin = _precompute_rope_freqs(self.head_dim, past_len + L, x.device)
+            q = _apply_rope(q, q_cos[past_len:, :], q_sin[past_len:, :])
+            k = _apply_rope(k, q_cos[past_len:, :], q_sin[past_len:, :])
+            
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        else:
+            cos, sin = _precompute_rope_freqs(self.head_dim, L, x.device)
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+
+        L_kv = k.size(2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        if bidirectional:
+            if self.window_size is not None:
+                past_len = L_kv - L
+                pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+                pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+                dist  = torch.abs(pos_i - pos_j)
+                win_mask = dist < self.window_size
+                scores = scores.masked_fill(
+                    ~win_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+                )
+        else:
+            past_len = L_kv - L
+            pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+            pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+            dist  = pos_i - pos_j
+
+            causal_mask = dist >= 0
+            if self.window_size is not None:
+                causal_mask = causal_mask & (dist < self.window_size)
+
+            scores = scores.masked_fill(
+                ~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+            )
+
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn, v)                                
+        out = out.transpose(1, 2).contiguous().view(B, L, D)
+        out = self.o_proj(out)
+
+        if use_cache:
+            if self.window_size is not None:
+                present_kv = (
+                    k[:, :, -self.window_size:, :],
+                    v[:, :, -self.window_size:, :]
+                )
+            else:
+                present_kv = (k, v)
+        else:
+            present_kv = None
+
+        return out, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 2.  MSIT BRANCH BLOCK (Pre-Norm)
+# ─────────────────────────────────────────────────────────────
+
+class MSITBranchBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size):
+        super().__init__()
+        self.ln1  = nn.LayerNorm(dim)
+        self.attn = SlidingWindowAttention(dim, num_heads, window_size)
+        self.ln2  = nn.LayerNorm(dim)
+        self.ffn  = nn.Sequential(
+            nn.Linear(dim, dim * 4, bias=False),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim, bias=False),
+        )
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        attn_out, present_kv = self.attn(
+            self.ln1(x), past_kv, use_cache, bidirectional
+        )
+        x = x + attn_out
+        x = x + self.ffn(self.ln2(x))
+        return x, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 3.  MoEP-MSIT ARCHITECTURE BLOCK  (Expert Choice routing)
+# ─────────────────────────────────────────────────────────────
+
+class MoEPMSITBlock(nn.Module):
+    def __init__(self, d_model: int = 512, d_thin: int = 192, num_blocks: int = 14,
+                 capacity_factor: float = 2.0):
+        super().__init__()
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+
+        # 1. Global Block (Dense, d_model)
+        num_heads_global = max(1, d_model // 64)
+        self.global_block = MSITBranchBlock(d_model, num_heads_global, window_size=None)
+
+        # 2. Router
+        self.router_ln = nn.LayerNorm(d_model)
+        self.w_router = nn.Linear(d_model, num_blocks, bias=False)
+
+        # 3. Shrink Projection
+        self.w_down = nn.Linear(d_model, d_thin, bias=False)
+
+        # 4. Thin Parallel Blocks (d_thin)
+        self.windows = [None, 16, 8, 4, 2] + [None] * (num_blocks - 5)
+        self.heads = [max(1, d_thin // 64)] * num_blocks  
+        
+        self.thin_blocks = nn.ModuleList([
+            MSITBranchBlock(d_thin, self.heads[i], self.windows[i])
+            for i in range(num_blocks)
+        ])
+
+        # 6. Grow Projection
+        self.w_up = nn.Linear(d_thin, d_model, bias=False)
+        self.last_topk_indices = None
+
+    def forward(self, x_0: torch.Tensor, past_kvs=None, use_cache: bool = False, bidirectional: bool = False):
+        B, T, D = x_0.size()
+        n_tokens = B * T
+
+        # Step 1: Global Block
+        pkv_g = past_kvs[0] if past_kvs else None
+        x_1, nkv_g = self.global_block(x_0, pkv_g, use_cache, bidirectional)
+
+        # Step 2: Gated input stream
+        x_2 = x_1 + x_0
+
+        # Router scores
+        r_logits = self.w_router(self.router_ln(x_2)).view(n_tokens, self.num_blocks)
+        r_probs = F.softmax(r_logits, dim=-1)
+
+        # Per-expert capacity: k = (n * c) / e
+        k_capacity = max(1, int(round(n_tokens * self.capacity_factor / self.num_blocks)))
+        k_capacity = min(k_capacity, n_tokens)
+
+        # Expert Choice routing: topk over the token axis for each expert
+        expert_token_scores = r_probs.transpose(0, 1) # (num_blocks, n_tokens)
+        topk_scores, topk_token_idx = torch.topk(expert_token_scores, k_capacity, dim=-1)
+        self.last_topk_indices = topk_token_idx
+
+        # Load balancing is guaranteed by construction in Expert Choice
+        layer_aux_loss = x_2.new_zeros(())
+
+        # Shrink projection
+        x_2_thin = self.w_down(x_2) # (B, T, d_thin)
+        x_2_thin_flat = x_2_thin.view(n_tokens, self.d_thin)
+
+        # Expert computations
+        new_kvs = [nkv_g]
+        expert_outputs_flat = torch.zeros(n_tokens, self.d_model, device=x_2.device, dtype=x_2.dtype)
+
+        for i, block in enumerate(self.thin_blocks):
+            sel_idx = topk_token_idx[i]
+            bucket_in = x_2_thin_flat[sel_idx].unsqueeze(0) # (1, k_capacity, d_thin)
+
+            pkv_i = past_kvs[i + 1] if past_kvs else None
+            bucket_out, nkv_i = block(bucket_in, pkv_i, use_cache, bidirectional)
+            if use_cache:
+                new_kvs.append(nkv_i)
+
+            bucket_out = bucket_out.squeeze(0) # (k_capacity, d_thin)
+            bucket_out_full = self.w_up(bucket_out) # (k_capacity, d_model)
+
+            gate = topk_scores[i].unsqueeze(-1) # (k_capacity, 1)
+            expert_outputs_flat.index_add_(0, sel_idx, bucket_out_full * gate)
+
+        x_3_full = expert_outputs_flat.view(B, T, self.d_model)
+        out = x_2 + x_3_full
+
+        present_kvs = tuple(new_kvs) if use_cache else None
+        return out, present_kvs, layer_aux_loss
+
+# ─────────────────────────────────────────────────────────────
+# 4.  RAW MSIT-GPT-BERT MODEL
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.wte      = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop_emb = nn.Dropout(cfg.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(
+                d_model=cfg.d_model, 
+                d_thin=cfg.d_thin, 
+                num_blocks=cfg.num_blocks, 
+                capacity_factor=cfg.capacity_factor
+            )
+            for _ in range(cfg.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(cfg.d_model)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.wte.weight = self.lm_head.weight
+
+    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None, bidirectional: bool = False):
+        x = self.drop_emb(self.wte(input_ids))
+        total_aux_loss = 0.0
+        
+        for block in self.blocks:
+            x, _, layer_aux = block(x, past_kvs=None, use_cache=False, bidirectional=bidirectional)
+            total_aux_loss += layer_aux
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            ce_loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), targets.view(-1), ignore_index=-100)
+            avg_aux_loss = total_aux_loss / self.cfg.num_layers
+            loss = ce_loss + (0.01 * avg_aux_loss)
+
+        return logits, loss
+
+# ─────────────────────────────────────────────────────────────
+# 5.  HUGGING FACE MODEL WRAPPERS
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModelWrapper(PreTrainedModel):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.wte      = nn.Embedding(config.vocab_size, config.d_model)
+        self.drop_emb = nn.Dropout(config.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(config.d_model, config.d_thin, config.num_blocks, config.capacity_factor)
+            for _ in range(config.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(config.d_model)
+        self.post_init()
+
+    def forward(self, input_ids, **kwargs):
+        x = self.drop_emb(self.wte(input_ids))
+        for block in self.blocks:
+            x, _, _ = block(x, past_kvs=None, use_cache=False, bidirectional=False)
+        x = self.ln_f(x)
+        return BaseModelOutputWithPast(last_hidden_state=x)
+
+
+class MSITGPTBERTForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+    _no_split_modules = ["MoEPMSITBlock"]
+    _tied_weights_keys = {"transformer.lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.transformer = MSITGPTBERTModelWrapper(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.post_init()
+        
+        # State-dict pre-hook for backwards compatibility with checkpoint key naming
+        def _prefix_cleaner(state_dict, prefix, local_metadata, Moore, missing_keys, unexpected_keys, error_msgs):
+            keys = list(state_dict.keys())
+            for k in keys:
+                if k.startswith("transformer."):
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                elif f"{prefix}transformer." in k:
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                    
+        self._register_load_state_dict_pre_hook(_prefix_cleaner)
+
+    def tie_weights(self, **kwargs):
+        if hasattr(self, "transformer") and hasattr(self.transformer, "wte") and hasattr(self.transformer, "lm_head"):
+            self.transformer.wte.weight = self.lm_head.weight
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(self,
+                input_ids: Optional[torch.LongTensor] = None,
+                attention_mask: Optional[torch.FloatTensor] = None,
+                labels: Optional[torch.LongTensor] = None,
+                **kwargs) -> CausalLMOutputWithPast:
+        outputs = self.transformer(input_ids)
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# Register with auto-mapping
+AutoConfig.register("msit_gptbert", MSITGPTBERTHFConfig)
+AutoModel.register(MSITGPTBERTHFConfig, MSITGPTBERTModelWrapper)
+AutoModelForCausalLM.register(MSITGPTBERTHFConfig, MSITGPTBERTForCausalLM)

chck_2M/pytorch_model.bin

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

chck_2M/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "[UNK]"
+}

chck_30M/config.json

@@ -0,0 +1,17 @@
+{
+  "auto_map": {
+    "AutoConfig": "modeling_msit.MSITGPTBERTHFConfig",
+    "AutoModel": "modeling_msit.MSITGPTBERTModelWrapper",
+    "AutoModelForCausalLM": "modeling_msit.MSITGPTBERTForCausalLM"
+  },
+  "vocab_size": 16384,
+  "block_size": 512,
+  "d_model": 384,
+  "hidden_size": 384,
+  "d_thin": 192,
+  "num_layers": 6,
+  "num_blocks": 6,
+  "capacity_factor": 2.0,
+  "dropout": 0.1,
+  "model_type": "msit_gptbert"
+}

chck_30M/configuration_msit.py

@@ -0,0 +1,35 @@
+from transformers import PretrainedConfig
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)

chck_30M/modeling_msit.py

@@ -0,0 +1,436 @@
+import os
+import math
+import random
+import inspect
+from typing import Optional, Tuple, Dict, Any
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin, AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+# ─────────────────────────────────────────────────────────────
+# Configuration Classes
+# ─────────────────────────────────────────────────────────────
+
+@dataclass
+class MSITGPTBERTConfig:
+    vocab_size: int = 16384
+    block_size: int = 512
+    d_model: int = 384
+    d_thin: int = 192
+    num_layers: int = 6
+    num_blocks: int = 6
+    capacity_factor: float = 2.0
+    dropout: float = 0.1
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)
+
+# ─────────────────────────────────────────────────────────────
+# ROPE HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def _precompute_rope_freqs(head_dim: int, seq_len: int, device: torch.device, theta: float = 10000.0):
+    assert head_dim % 2 == 0, "head_dim must be divisible by 2 for RoPE"
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos(), emb.sin()
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    L = x.size(2)
+    cos = cos[:L, :].unsqueeze(0).unsqueeze(1) 
+    sin = sin[:L, :].unsqueeze(0).unsqueeze(1) 
+    
+    half_dim = x.size(-1) // 2
+    x1 = x[..., :half_dim]
+    x2 = x[..., half_dim:]
+    rotated_x = torch.cat((-x2, x1), dim=-1)
+    
+    return (x * cos) + (rotated_x * sin)
+
+# ─────────────────────────────────────────────────────────────
+# 1.  SLIDING WINDOW ATTENTION
+# ─────────────────────────────────────────────────────────────
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size=None):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.num_heads  = num_heads
+        self.window_size = window_size
+        self.head_dim   = dim // num_heads
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        B, L, D = x.size()
+
+        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            past_len = past_k.size(2)
+            q_cos, q_sin = _precompute_rope_freqs(self.head_dim, past_len + L, x.device)
+            q = _apply_rope(q, q_cos[past_len:, :], q_sin[past_len:, :])
+            k = _apply_rope(k, q_cos[past_len:, :], q_sin[past_len:, :])
+            
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        else:
+            cos, sin = _precompute_rope_freqs(self.head_dim, L, x.device)
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+
+        L_kv = k.size(2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        if bidirectional:
+            if self.window_size is not None:
+                past_len = L_kv - L
+                pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+                pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+                dist  = torch.abs(pos_i - pos_j)
+                win_mask = dist < self.window_size
+                scores = scores.masked_fill(
+                    ~win_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+                )
+        else:
+            past_len = L_kv - L
+            pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+            pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+            dist  = pos_i - pos_j
+
+            causal_mask = dist >= 0
+            if self.window_size is not None:
+                causal_mask = causal_mask & (dist < self.window_size)
+
+            scores = scores.masked_fill(
+                ~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+            )
+
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn, v)                                
+        out = out.transpose(1, 2).contiguous().view(B, L, D)
+        out = self.o_proj(out)
+
+        if use_cache:
+            if self.window_size is not None:
+                present_kv = (
+                    k[:, :, -self.window_size:, :],
+                    v[:, :, -self.window_size:, :]
+                )
+            else:
+                present_kv = (k, v)
+        else:
+            present_kv = None
+
+        return out, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 2.  MSIT BRANCH BLOCK (Pre-Norm)
+# ─────────────────────────────────────────────────────────────
+
+class MSITBranchBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size):
+        super().__init__()
+        self.ln1  = nn.LayerNorm(dim)
+        self.attn = SlidingWindowAttention(dim, num_heads, window_size)
+        self.ln2  = nn.LayerNorm(dim)
+        self.ffn  = nn.Sequential(
+            nn.Linear(dim, dim * 4, bias=False),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim, bias=False),
+        )
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        attn_out, present_kv = self.attn(
+            self.ln1(x), past_kv, use_cache, bidirectional
+        )
+        x = x + attn_out
+        x = x + self.ffn(self.ln2(x))
+        return x, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 3.  MoEP-MSIT ARCHITECTURE BLOCK  (Expert Choice routing)
+# ─────────────────────────────────────────────────────────────
+
+class MoEPMSITBlock(nn.Module):
+    def __init__(self, d_model: int = 512, d_thin: int = 192, num_blocks: int = 14,
+                 capacity_factor: float = 2.0):
+        super().__init__()
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+
+        # 1. Global Block (Dense, d_model)
+        num_heads_global = max(1, d_model // 64)
+        self.global_block = MSITBranchBlock(d_model, num_heads_global, window_size=None)
+
+        # 2. Router
+        self.router_ln = nn.LayerNorm(d_model)
+        self.w_router = nn.Linear(d_model, num_blocks, bias=False)
+
+        # 3. Shrink Projection
+        self.w_down = nn.Linear(d_model, d_thin, bias=False)
+
+        # 4. Thin Parallel Blocks (d_thin)
+        self.windows = [None, 16, 8, 4, 2] + [None] * (num_blocks - 5)
+        self.heads = [max(1, d_thin // 64)] * num_blocks  
+        
+        self.thin_blocks = nn.ModuleList([
+            MSITBranchBlock(d_thin, self.heads[i], self.windows[i])
+            for i in range(num_blocks)
+        ])
+
+        # 6. Grow Projection
+        self.w_up = nn.Linear(d_thin, d_model, bias=False)
+        self.last_topk_indices = None
+
+    def forward(self, x_0: torch.Tensor, past_kvs=None, use_cache: bool = False, bidirectional: bool = False):
+        B, T, D = x_0.size()
+        n_tokens = B * T
+
+        # Step 1: Global Block
+        pkv_g = past_kvs[0] if past_kvs else None
+        x_1, nkv_g = self.global_block(x_0, pkv_g, use_cache, bidirectional)
+
+        # Step 2: Gated input stream
+        x_2 = x_1 + x_0
+
+        # Router scores
+        r_logits = self.w_router(self.router_ln(x_2)).view(n_tokens, self.num_blocks)
+        r_probs = F.softmax(r_logits, dim=-1)
+
+        # Per-expert capacity: k = (n * c) / e
+        k_capacity = max(1, int(round(n_tokens * self.capacity_factor / self.num_blocks)))
+        k_capacity = min(k_capacity, n_tokens)
+
+        # Expert Choice routing: topk over the token axis for each expert
+        expert_token_scores = r_probs.transpose(0, 1) # (num_blocks, n_tokens)
+        topk_scores, topk_token_idx = torch.topk(expert_token_scores, k_capacity, dim=-1)
+        self.last_topk_indices = topk_token_idx
+
+        # Load balancing is guaranteed by construction in Expert Choice
+        layer_aux_loss = x_2.new_zeros(())
+
+        # Shrink projection
+        x_2_thin = self.w_down(x_2) # (B, T, d_thin)
+        x_2_thin_flat = x_2_thin.view(n_tokens, self.d_thin)
+
+        # Expert computations
+        new_kvs = [nkv_g]
+        expert_outputs_flat = torch.zeros(n_tokens, self.d_model, device=x_2.device, dtype=x_2.dtype)
+
+        for i, block in enumerate(self.thin_blocks):
+            sel_idx = topk_token_idx[i]
+            bucket_in = x_2_thin_flat[sel_idx].unsqueeze(0) # (1, k_capacity, d_thin)
+
+            pkv_i = past_kvs[i + 1] if past_kvs else None
+            bucket_out, nkv_i = block(bucket_in, pkv_i, use_cache, bidirectional)
+            if use_cache:
+                new_kvs.append(nkv_i)
+
+            bucket_out = bucket_out.squeeze(0) # (k_capacity, d_thin)
+            bucket_out_full = self.w_up(bucket_out) # (k_capacity, d_model)
+
+            gate = topk_scores[i].unsqueeze(-1) # (k_capacity, 1)
+            expert_outputs_flat.index_add_(0, sel_idx, bucket_out_full * gate)
+
+        x_3_full = expert_outputs_flat.view(B, T, self.d_model)
+        out = x_2 + x_3_full
+
+        present_kvs = tuple(new_kvs) if use_cache else None
+        return out, present_kvs, layer_aux_loss
+
+# ─────────────────────────────────────────────────────────────
+# 4.  RAW MSIT-GPT-BERT MODEL
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.wte      = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop_emb = nn.Dropout(cfg.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(
+                d_model=cfg.d_model, 
+                d_thin=cfg.d_thin, 
+                num_blocks=cfg.num_blocks, 
+                capacity_factor=cfg.capacity_factor
+            )
+            for _ in range(cfg.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(cfg.d_model)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.wte.weight = self.lm_head.weight
+
+    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None, bidirectional: bool = False):
+        x = self.drop_emb(self.wte(input_ids))
+        total_aux_loss = 0.0
+        
+        for block in self.blocks:
+            x, _, layer_aux = block(x, past_kvs=None, use_cache=False, bidirectional=bidirectional)
+            total_aux_loss += layer_aux
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            ce_loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), targets.view(-1), ignore_index=-100)
+            avg_aux_loss = total_aux_loss / self.cfg.num_layers
+            loss = ce_loss + (0.01 * avg_aux_loss)
+
+        return logits, loss
+
+# ─────────────────────────────────────────────────────────────
+# 5.  HUGGING FACE MODEL WRAPPERS
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModelWrapper(PreTrainedModel):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.wte      = nn.Embedding(config.vocab_size, config.d_model)
+        self.drop_emb = nn.Dropout(config.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(config.d_model, config.d_thin, config.num_blocks, config.capacity_factor)
+            for _ in range(config.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(config.d_model)
+        self.post_init()
+
+    def forward(self, input_ids, **kwargs):
+        x = self.drop_emb(self.wte(input_ids))
+        for block in self.blocks:
+            x, _, _ = block(x, past_kvs=None, use_cache=False, bidirectional=False)
+        x = self.ln_f(x)
+        return BaseModelOutputWithPast(last_hidden_state=x)
+
+
+class MSITGPTBERTForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+    _no_split_modules = ["MoEPMSITBlock"]
+    _tied_weights_keys = {"transformer.lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.transformer = MSITGPTBERTModelWrapper(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.post_init()
+        
+        # State-dict pre-hook for backwards compatibility with checkpoint key naming
+        def _prefix_cleaner(state_dict, prefix, local_metadata, Moore, missing_keys, unexpected_keys, error_msgs):
+            keys = list(state_dict.keys())
+            for k in keys:
+                if k.startswith("transformer."):
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                elif f"{prefix}transformer." in k:
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                    
+        self._register_load_state_dict_pre_hook(_prefix_cleaner)
+
+    def tie_weights(self, **kwargs):
+        if hasattr(self, "transformer") and hasattr(self.transformer, "wte") and hasattr(self.transformer, "lm_head"):
+            self.transformer.wte.weight = self.lm_head.weight
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(self,
+                input_ids: Optional[torch.LongTensor] = None,
+                attention_mask: Optional[torch.FloatTensor] = None,
+                labels: Optional[torch.LongTensor] = None,
+                **kwargs) -> CausalLMOutputWithPast:
+        outputs = self.transformer(input_ids)
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# Register with auto-mapping
+AutoConfig.register("msit_gptbert", MSITGPTBERTHFConfig)
+AutoModel.register(MSITGPTBERTHFConfig, MSITGPTBERTModelWrapper)
+AutoModelForCausalLM.register(MSITGPTBERTHFConfig, MSITGPTBERTForCausalLM)

chck_30M/pytorch_model.bin

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

chck_30M/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "[UNK]"
+}

chck_3M/config.json

@@ -0,0 +1,17 @@
+{
+  "auto_map": {
+    "AutoConfig": "modeling_msit.MSITGPTBERTHFConfig",
+    "AutoModel": "modeling_msit.MSITGPTBERTModelWrapper",
+    "AutoModelForCausalLM": "modeling_msit.MSITGPTBERTForCausalLM"
+  },
+  "vocab_size": 16384,
+  "block_size": 512,
+  "d_model": 384,
+  "hidden_size": 384,
+  "d_thin": 192,
+  "num_layers": 6,
+  "num_blocks": 6,
+  "capacity_factor": 2.0,
+  "dropout": 0.1,
+  "model_type": "msit_gptbert"
+}

chck_3M/configuration_msit.py

@@ -0,0 +1,35 @@
+from transformers import PretrainedConfig
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)

chck_3M/modeling_msit.py

@@ -0,0 +1,436 @@
+import os
+import math
+import random
+import inspect
+from typing import Optional, Tuple, Dict, Any
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin, AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+# ─────────────────────────────────────────────────────────────
+# Configuration Classes
+# ─────────────────────────────────────────────────────────────
+
+@dataclass
+class MSITGPTBERTConfig:
+    vocab_size: int = 16384
+    block_size: int = 512
+    d_model: int = 384
+    d_thin: int = 192
+    num_layers: int = 6
+    num_blocks: int = 6
+    capacity_factor: float = 2.0
+    dropout: float = 0.1
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)
+
+# ─────────────────────────────────────────────────────────────
+# ROPE HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def _precompute_rope_freqs(head_dim: int, seq_len: int, device: torch.device, theta: float = 10000.0):
+    assert head_dim % 2 == 0, "head_dim must be divisible by 2 for RoPE"
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos(), emb.sin()
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    L = x.size(2)
+    cos = cos[:L, :].unsqueeze(0).unsqueeze(1) 
+    sin = sin[:L, :].unsqueeze(0).unsqueeze(1) 
+    
+    half_dim = x.size(-1) // 2
+    x1 = x[..., :half_dim]
+    x2 = x[..., half_dim:]
+    rotated_x = torch.cat((-x2, x1), dim=-1)
+    
+    return (x * cos) + (rotated_x * sin)
+
+# ─────────────────────────────────────────────────────────────
+# 1.  SLIDING WINDOW ATTENTION
+# ─────────────────────────────────────────────────────────────
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size=None):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.num_heads  = num_heads
+        self.window_size = window_size
+        self.head_dim   = dim // num_heads
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        B, L, D = x.size()
+
+        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            past_len = past_k.size(2)
+            q_cos, q_sin = _precompute_rope_freqs(self.head_dim, past_len + L, x.device)
+            q = _apply_rope(q, q_cos[past_len:, :], q_sin[past_len:, :])
+            k = _apply_rope(k, q_cos[past_len:, :], q_sin[past_len:, :])
+            
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        else:
+            cos, sin = _precompute_rope_freqs(self.head_dim, L, x.device)
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+
+        L_kv = k.size(2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        if bidirectional:
+            if self.window_size is not None:
+                past_len = L_kv - L
+                pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+                pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+                dist  = torch.abs(pos_i - pos_j)
+                win_mask = dist < self.window_size
+                scores = scores.masked_fill(
+                    ~win_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+                )
+        else:
+            past_len = L_kv - L
+            pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+            pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+            dist  = pos_i - pos_j
+
+            causal_mask = dist >= 0
+            if self.window_size is not None:
+                causal_mask = causal_mask & (dist < self.window_size)
+
+            scores = scores.masked_fill(
+                ~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+            )
+
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn, v)                                
+        out = out.transpose(1, 2).contiguous().view(B, L, D)
+        out = self.o_proj(out)
+
+        if use_cache:
+            if self.window_size is not None:
+                present_kv = (
+                    k[:, :, -self.window_size:, :],
+                    v[:, :, -self.window_size:, :]
+                )
+            else:
+                present_kv = (k, v)
+        else:
+            present_kv = None
+
+        return out, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 2.  MSIT BRANCH BLOCK (Pre-Norm)
+# ─────────────────────────────────────────────────────────────
+
+class MSITBranchBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size):
+        super().__init__()
+        self.ln1  = nn.LayerNorm(dim)
+        self.attn = SlidingWindowAttention(dim, num_heads, window_size)
+        self.ln2  = nn.LayerNorm(dim)
+        self.ffn  = nn.Sequential(
+            nn.Linear(dim, dim * 4, bias=False),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim, bias=False),
+        )
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        attn_out, present_kv = self.attn(
+            self.ln1(x), past_kv, use_cache, bidirectional
+        )
+        x = x + attn_out
+        x = x + self.ffn(self.ln2(x))
+        return x, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 3.  MoEP-MSIT ARCHITECTURE BLOCK  (Expert Choice routing)
+# ─────────────────────────────────────────────────────────────
+
+class MoEPMSITBlock(nn.Module):
+    def __init__(self, d_model: int = 512, d_thin: int = 192, num_blocks: int = 14,
+                 capacity_factor: float = 2.0):
+        super().__init__()
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+
+        # 1. Global Block (Dense, d_model)
+        num_heads_global = max(1, d_model // 64)
+        self.global_block = MSITBranchBlock(d_model, num_heads_global, window_size=None)
+
+        # 2. Router
+        self.router_ln = nn.LayerNorm(d_model)
+        self.w_router = nn.Linear(d_model, num_blocks, bias=False)
+
+        # 3. Shrink Projection
+        self.w_down = nn.Linear(d_model, d_thin, bias=False)
+
+        # 4. Thin Parallel Blocks (d_thin)
+        self.windows = [None, 16, 8, 4, 2] + [None] * (num_blocks - 5)
+        self.heads = [max(1, d_thin // 64)] * num_blocks  
+        
+        self.thin_blocks = nn.ModuleList([
+            MSITBranchBlock(d_thin, self.heads[i], self.windows[i])
+            for i in range(num_blocks)
+        ])
+
+        # 6. Grow Projection
+        self.w_up = nn.Linear(d_thin, d_model, bias=False)
+        self.last_topk_indices = None
+
+    def forward(self, x_0: torch.Tensor, past_kvs=None, use_cache: bool = False, bidirectional: bool = False):
+        B, T, D = x_0.size()
+        n_tokens = B * T
+
+        # Step 1: Global Block
+        pkv_g = past_kvs[0] if past_kvs else None
+        x_1, nkv_g = self.global_block(x_0, pkv_g, use_cache, bidirectional)
+
+        # Step 2: Gated input stream
+        x_2 = x_1 + x_0
+
+        # Router scores
+        r_logits = self.w_router(self.router_ln(x_2)).view(n_tokens, self.num_blocks)
+        r_probs = F.softmax(r_logits, dim=-1)
+
+        # Per-expert capacity: k = (n * c) / e
+        k_capacity = max(1, int(round(n_tokens * self.capacity_factor / self.num_blocks)))
+        k_capacity = min(k_capacity, n_tokens)
+
+        # Expert Choice routing: topk over the token axis for each expert
+        expert_token_scores = r_probs.transpose(0, 1) # (num_blocks, n_tokens)
+        topk_scores, topk_token_idx = torch.topk(expert_token_scores, k_capacity, dim=-1)
+        self.last_topk_indices = topk_token_idx
+
+        # Load balancing is guaranteed by construction in Expert Choice
+        layer_aux_loss = x_2.new_zeros(())
+
+        # Shrink projection
+        x_2_thin = self.w_down(x_2) # (B, T, d_thin)
+        x_2_thin_flat = x_2_thin.view(n_tokens, self.d_thin)
+
+        # Expert computations
+        new_kvs = [nkv_g]
+        expert_outputs_flat = torch.zeros(n_tokens, self.d_model, device=x_2.device, dtype=x_2.dtype)
+
+        for i, block in enumerate(self.thin_blocks):
+            sel_idx = topk_token_idx[i]
+            bucket_in = x_2_thin_flat[sel_idx].unsqueeze(0) # (1, k_capacity, d_thin)
+
+            pkv_i = past_kvs[i + 1] if past_kvs else None
+            bucket_out, nkv_i = block(bucket_in, pkv_i, use_cache, bidirectional)
+            if use_cache:
+                new_kvs.append(nkv_i)
+
+            bucket_out = bucket_out.squeeze(0) # (k_capacity, d_thin)
+            bucket_out_full = self.w_up(bucket_out) # (k_capacity, d_model)
+
+            gate = topk_scores[i].unsqueeze(-1) # (k_capacity, 1)
+            expert_outputs_flat.index_add_(0, sel_idx, bucket_out_full * gate)
+
+        x_3_full = expert_outputs_flat.view(B, T, self.d_model)
+        out = x_2 + x_3_full
+
+        present_kvs = tuple(new_kvs) if use_cache else None
+        return out, present_kvs, layer_aux_loss
+
+# ─────────────────────────────────────────────────────────────
+# 4.  RAW MSIT-GPT-BERT MODEL
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.wte      = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop_emb = nn.Dropout(cfg.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(
+                d_model=cfg.d_model, 
+                d_thin=cfg.d_thin, 
+                num_blocks=cfg.num_blocks, 
+                capacity_factor=cfg.capacity_factor
+            )
+            for _ in range(cfg.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(cfg.d_model)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.wte.weight = self.lm_head.weight
+
+    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None, bidirectional: bool = False):
+        x = self.drop_emb(self.wte(input_ids))
+        total_aux_loss = 0.0
+        
+        for block in self.blocks:
+            x, _, layer_aux = block(x, past_kvs=None, use_cache=False, bidirectional=bidirectional)
+            total_aux_loss += layer_aux
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            ce_loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), targets.view(-1), ignore_index=-100)
+            avg_aux_loss = total_aux_loss / self.cfg.num_layers
+            loss = ce_loss + (0.01 * avg_aux_loss)
+
+        return logits, loss
+
+# ─────────────────────────────────────────────────────────────
+# 5.  HUGGING FACE MODEL WRAPPERS
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModelWrapper(PreTrainedModel):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.wte      = nn.Embedding(config.vocab_size, config.d_model)
+        self.drop_emb = nn.Dropout(config.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(config.d_model, config.d_thin, config.num_blocks, config.capacity_factor)
+            for _ in range(config.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(config.d_model)
+        self.post_init()
+
+    def forward(self, input_ids, **kwargs):
+        x = self.drop_emb(self.wte(input_ids))
+        for block in self.blocks:
+            x, _, _ = block(x, past_kvs=None, use_cache=False, bidirectional=False)
+        x = self.ln_f(x)
+        return BaseModelOutputWithPast(last_hidden_state=x)
+
+
+class MSITGPTBERTForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+    _no_split_modules = ["MoEPMSITBlock"]
+    _tied_weights_keys = {"transformer.lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.transformer = MSITGPTBERTModelWrapper(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.post_init()
+        
+        # State-dict pre-hook for backwards compatibility with checkpoint key naming
+        def _prefix_cleaner(state_dict, prefix, local_metadata, Moore, missing_keys, unexpected_keys, error_msgs):
+            keys = list(state_dict.keys())
+            for k in keys:
+                if k.startswith("transformer."):
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                elif f"{prefix}transformer." in k:
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                    
+        self._register_load_state_dict_pre_hook(_prefix_cleaner)
+
+    def tie_weights(self, **kwargs):
+        if hasattr(self, "transformer") and hasattr(self.transformer, "wte") and hasattr(self.transformer, "lm_head"):
+            self.transformer.wte.weight = self.lm_head.weight
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(self,
+                input_ids: Optional[torch.LongTensor] = None,
+                attention_mask: Optional[torch.FloatTensor] = None,
+                labels: Optional[torch.LongTensor] = None,
+                **kwargs) -> CausalLMOutputWithPast:
+        outputs = self.transformer(input_ids)
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# Register with auto-mapping
+AutoConfig.register("msit_gptbert", MSITGPTBERTHFConfig)
+AutoModel.register(MSITGPTBERTHFConfig, MSITGPTBERTModelWrapper)
+AutoModelForCausalLM.register(MSITGPTBERTHFConfig, MSITGPTBERTForCausalLM)

chck_3M/pytorch_model.bin

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

chck_3M/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "[UNK]"
+}

chck_40M/config.json

@@ -0,0 +1,17 @@
+{
+  "auto_map": {
+    "AutoConfig": "modeling_msit.MSITGPTBERTHFConfig",
+    "AutoModel": "modeling_msit.MSITGPTBERTModelWrapper",
+    "AutoModelForCausalLM": "modeling_msit.MSITGPTBERTForCausalLM"
+  },
+  "vocab_size": 16384,
+  "block_size": 512,
+  "d_model": 384,
+  "hidden_size": 384,
+  "d_thin": 192,
+  "num_layers": 6,
+  "num_blocks": 6,
+  "capacity_factor": 2.0,
+  "dropout": 0.1,
+  "model_type": "msit_gptbert"
+}

chck_40M/configuration_msit.py

@@ -0,0 +1,35 @@
+from transformers import PretrainedConfig
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)

chck_40M/modeling_msit.py

@@ -0,0 +1,436 @@
+import os
+import math
+import random
+import inspect
+from typing import Optional, Tuple, Dict, Any
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin, AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+# ─────────────────────────────────────────────────────────────
+# Configuration Classes
+# ─────────────────────────────────────────────────────────────
+
+@dataclass
+class MSITGPTBERTConfig:
+    vocab_size: int = 16384
+    block_size: int = 512
+    d_model: int = 384
+    d_thin: int = 192
+    num_layers: int = 6
+    num_blocks: int = 6
+    capacity_factor: float = 2.0
+    dropout: float = 0.1
+
+class MSITGPTBERTHFConfig(PretrainedConfig):
+    model_type = "msit_gptbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 384,
+        d_thin: int = 192,
+        num_layers: int = 6,
+        num_blocks: int = 6,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+
+        super().__init__(**kwargs)
+
+# ─────────────────────────────────────────────────────────────
+# ROPE HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def _precompute_rope_freqs(head_dim: int, seq_len: int, device: torch.device, theta: float = 10000.0):
+    assert head_dim % 2 == 0, "head_dim must be divisible by 2 for RoPE"
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos(), emb.sin()
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    L = x.size(2)
+    cos = cos[:L, :].unsqueeze(0).unsqueeze(1) 
+    sin = sin[:L, :].unsqueeze(0).unsqueeze(1) 
+    
+    half_dim = x.size(-1) // 2
+    x1 = x[..., :half_dim]
+    x2 = x[..., half_dim:]
+    rotated_x = torch.cat((-x2, x1), dim=-1)
+    
+    return (x * cos) + (rotated_x * sin)
+
+# ─────────────────────────────────────────────────────────────
+# 1.  SLIDING WINDOW ATTENTION
+# ─────────────────────────────────────────────────────────────
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size=None):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.num_heads  = num_heads
+        self.window_size = window_size
+        self.head_dim   = dim // num_heads
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        B, L, D = x.size()
+
+        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            past_len = past_k.size(2)
+            q_cos, q_sin = _precompute_rope_freqs(self.head_dim, past_len + L, x.device)
+            q = _apply_rope(q, q_cos[past_len:, :], q_sin[past_len:, :])
+            k = _apply_rope(k, q_cos[past_len:, :], q_sin[past_len:, :])
+            
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        else:
+            cos, sin = _precompute_rope_freqs(self.head_dim, L, x.device)
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+
+        L_kv = k.size(2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        if bidirectional:
+            if self.window_size is not None:
+                past_len = L_kv - L
+                pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+                pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+                dist  = torch.abs(pos_i - pos_j)
+                win_mask = dist < self.window_size
+                scores = scores.masked_fill(
+                    ~win_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+                )
+        else:
+            past_len = L_kv - L
+            pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+            pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+            dist  = pos_i - pos_j
+
+            causal_mask = dist >= 0
+            if self.window_size is not None:
+                causal_mask = causal_mask & (dist < self.window_size)
+
+            scores = scores.masked_fill(
+                ~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+            )
+
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn, v)                                
+        out = out.transpose(1, 2).contiguous().view(B, L, D)
+        out = self.o_proj(out)
+
+        if use_cache:
+            if self.window_size is not None:
+                present_kv = (
+                    k[:, :, -self.window_size:, :],
+                    v[:, :, -self.window_size:, :]
+                )
+            else:
+                present_kv = (k, v)
+        else:
+            present_kv = None
+
+        return out, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 2.  MSIT BRANCH BLOCK (Pre-Norm)
+# ─────────────────────────────────────────────────────────────
+
+class MSITBranchBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size):
+        super().__init__()
+        self.ln1  = nn.LayerNorm(dim)
+        self.attn = SlidingWindowAttention(dim, num_heads, window_size)
+        self.ln2  = nn.LayerNorm(dim)
+        self.ffn  = nn.Sequential(
+            nn.Linear(dim, dim * 4, bias=False),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim, bias=False),
+        )
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        attn_out, present_kv = self.attn(
+            self.ln1(x), past_kv, use_cache, bidirectional
+        )
+        x = x + attn_out
+        x = x + self.ffn(self.ln2(x))
+        return x, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 3.  MoEP-MSIT ARCHITECTURE BLOCK  (Expert Choice routing)
+# ─────────────────────────────────────────────────────────────
+
+class MoEPMSITBlock(nn.Module):
+    def __init__(self, d_model: int = 512, d_thin: int = 192, num_blocks: int = 14,
+                 capacity_factor: float = 2.0):
+        super().__init__()
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+
+        # 1. Global Block (Dense, d_model)
+        num_heads_global = max(1, d_model // 64)
+        self.global_block = MSITBranchBlock(d_model, num_heads_global, window_size=None)
+
+        # 2. Router
+        self.router_ln = nn.LayerNorm(d_model)
+        self.w_router = nn.Linear(d_model, num_blocks, bias=False)
+
+        # 3. Shrink Projection
+        self.w_down = nn.Linear(d_model, d_thin, bias=False)
+
+        # 4. Thin Parallel Blocks (d_thin)
+        self.windows = [None, 16, 8, 4, 2] + [None] * (num_blocks - 5)
+        self.heads = [max(1, d_thin // 64)] * num_blocks  
+        
+        self.thin_blocks = nn.ModuleList([
+            MSITBranchBlock(d_thin, self.heads[i], self.windows[i])
+            for i in range(num_blocks)
+        ])
+
+        # 6. Grow Projection
+        self.w_up = nn.Linear(d_thin, d_model, bias=False)
+        self.last_topk_indices = None
+
+    def forward(self, x_0: torch.Tensor, past_kvs=None, use_cache: bool = False, bidirectional: bool = False):
+        B, T, D = x_0.size()
+        n_tokens = B * T
+
+        # Step 1: Global Block
+        pkv_g = past_kvs[0] if past_kvs else None
+        x_1, nkv_g = self.global_block(x_0, pkv_g, use_cache, bidirectional)
+
+        # Step 2: Gated input stream
+        x_2 = x_1 + x_0
+
+        # Router scores
+        r_logits = self.w_router(self.router_ln(x_2)).view(n_tokens, self.num_blocks)
+        r_probs = F.softmax(r_logits, dim=-1)
+
+        # Per-expert capacity: k = (n * c) / e
+        k_capacity = max(1, int(round(n_tokens * self.capacity_factor / self.num_blocks)))
+        k_capacity = min(k_capacity, n_tokens)
+
+        # Expert Choice routing: topk over the token axis for each expert
+        expert_token_scores = r_probs.transpose(0, 1) # (num_blocks, n_tokens)
+        topk_scores, topk_token_idx = torch.topk(expert_token_scores, k_capacity, dim=-1)
+        self.last_topk_indices = topk_token_idx
+
+        # Load balancing is guaranteed by construction in Expert Choice
+        layer_aux_loss = x_2.new_zeros(())
+
+        # Shrink projection
+        x_2_thin = self.w_down(x_2) # (B, T, d_thin)
+        x_2_thin_flat = x_2_thin.view(n_tokens, self.d_thin)
+
+        # Expert computations
+        new_kvs = [nkv_g]
+        expert_outputs_flat = torch.zeros(n_tokens, self.d_model, device=x_2.device, dtype=x_2.dtype)
+
+        for i, block in enumerate(self.thin_blocks):
+            sel_idx = topk_token_idx[i]
+            bucket_in = x_2_thin_flat[sel_idx].unsqueeze(0) # (1, k_capacity, d_thin)
+
+            pkv_i = past_kvs[i + 1] if past_kvs else None
+            bucket_out, nkv_i = block(bucket_in, pkv_i, use_cache, bidirectional)
+            if use_cache:
+                new_kvs.append(nkv_i)
+
+            bucket_out = bucket_out.squeeze(0) # (k_capacity, d_thin)
+            bucket_out_full = self.w_up(bucket_out) # (k_capacity, d_model)
+
+            gate = topk_scores[i].unsqueeze(-1) # (k_capacity, 1)
+            expert_outputs_flat.index_add_(0, sel_idx, bucket_out_full * gate)
+
+        x_3_full = expert_outputs_flat.view(B, T, self.d_model)
+        out = x_2 + x_3_full
+
+        present_kvs = tuple(new_kvs) if use_cache else None
+        return out, present_kvs, layer_aux_loss
+
+# ─────────────────────────────────────────────────────────────
+# 4.  RAW MSIT-GPT-BERT MODEL
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.wte      = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop_emb = nn.Dropout(cfg.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(
+                d_model=cfg.d_model, 
+                d_thin=cfg.d_thin, 
+                num_blocks=cfg.num_blocks, 
+                capacity_factor=cfg.capacity_factor
+            )
+            for _ in range(cfg.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(cfg.d_model)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.wte.weight = self.lm_head.weight
+
+    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None, bidirectional: bool = False):
+        x = self.drop_emb(self.wte(input_ids))
+        total_aux_loss = 0.0
+        
+        for block in self.blocks:
+            x, _, layer_aux = block(x, past_kvs=None, use_cache=False, bidirectional=bidirectional)
+            total_aux_loss += layer_aux
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            ce_loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), targets.view(-1), ignore_index=-100)
+            avg_aux_loss = total_aux_loss / self.cfg.num_layers
+            loss = ce_loss + (0.01 * avg_aux_loss)
+
+        return logits, loss
+
+# ─────────────────────────────────────────────────────────────
+# 5.  HUGGING FACE MODEL WRAPPERS
+# ─────────────────────────────────────────────────────────────
+
+class MSITGPTBERTModelWrapper(PreTrainedModel):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.wte      = nn.Embedding(config.vocab_size, config.d_model)
+        self.drop_emb = nn.Dropout(config.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(config.d_model, config.d_thin, config.num_blocks, config.capacity_factor)
+            for _ in range(config.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(config.d_model)
+        self.post_init()
+
+    def forward(self, input_ids, **kwargs):
+        x = self.drop_emb(self.wte(input_ids))
+        for block in self.blocks:
+            x, _, _ = block(x, past_kvs=None, use_cache=False, bidirectional=False)
+        x = self.ln_f(x)
+        return BaseModelOutputWithPast(last_hidden_state=x)
+
+
+class MSITGPTBERTForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = MSITGPTBERTHFConfig
+    base_model_prefix = "transformer"
+    _no_split_modules = ["MoEPMSITBlock"]
+    _tied_weights_keys = {"transformer.lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config: MSITGPTBERTHFConfig):
+        super().__init__(config)
+        self.transformer = MSITGPTBERTModelWrapper(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.post_init()
+        
+        # State-dict pre-hook for backwards compatibility with checkpoint key naming
+        def _prefix_cleaner(state_dict, prefix, local_metadata, Moore, missing_keys, unexpected_keys, error_msgs):
+            keys = list(state_dict.keys())
+            for k in keys:
+                if k.startswith("transformer."):
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                elif f"{prefix}transformer." in k:
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                    
+        self._register_load_state_dict_pre_hook(_prefix_cleaner)
+
+    def tie_weights(self, **kwargs):
+        if hasattr(self, "transformer") and hasattr(self.transformer, "wte") and hasattr(self.transformer, "lm_head"):
+            self.transformer.wte.weight = self.lm_head.weight
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(self,
+                input_ids: Optional[torch.LongTensor] = None,
+                attention_mask: Optional[torch.FloatTensor] = None,
+                labels: Optional[torch.LongTensor] = None,
+                **kwargs) -> CausalLMOutputWithPast:
+        outputs = self.transformer(input_ids)
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# Register with auto-mapping
+AutoConfig.register("msit_gptbert", MSITGPTBERTHFConfig)
+AutoModel.register(MSITGPTBERTHFConfig, MSITGPTBERTModelWrapper)
+AutoModelForCausalLM.register(MSITGPTBERTHFConfig, MSITGPTBERTForCausalLM)

chck_40M/pytorch_model.bin

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

chck_40M/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "[UNK]"
+}

chck_4M/config.json

@@ -0,0 +1,17 @@
+{
+  "auto_map": {
+    "AutoConfig": "modeling_msit.MSITGPTBERTHFConfig",
+    "AutoModel": "modeling_msit.MSITGPTBERTModelWrapper",
+    "AutoModelForCausalLM": "modeling_msit.MSITGPTBERTForCausalLM"
+  },
+  "vocab_size": 16384,
+  "block_size": 512,
+  "d_model": 384,
+  "hidden_size": 384,
+  "d_thin": 192,
+  "num_layers": 6,
+  "num_blocks": 6,
+  "capacity_factor": 2.0,
+  "dropout": 0.1,
+  "model_type": "msit_gptbert"
+}