Revert to 7ed7003 and squash history

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,72 @@
+
+---
+language:
+  - en
+tags:
+  - text-generation
+  - causal-lm
+  - custom-architecture
+  - slm
+  - small-language-model
+license: mit
+---
+
+# Flare (130M Parameters)
+
+Flare is a highly efficient small language model developed by SurjoLabs, demonstrating that this custom XSA recurrent architecture is viable and can yield exceptional reasoning capabilities in tiny architectures.
+
+## Methodology
+
+Flare is built to prove that unique recurrent architectures can match or beat standard transformers even at small scales.
+
+* **Custom Architecture:** Built on a Llama-variant framework but features **XSA Attention** (value-subtraction projection) and **Recurrent Layers**. By reusing weights across recurrent passes, the model achieves extended effective depth while carrying a reduced set of unique weights.
+* **Optimized Tokenizer:** Uses a custom 32,768-vocabulary (2^15) English tokenizer. This power-of-two vocabulary optimizes tensor parallelism and compute alignment while providing strong compression rates across text, code, and math.
+* **Extreme Overtraining:** Trained on **26.21 Billion tokens** (25,000 steps at 1,048,576 tokens per step), resulting in a ~1:200 parameter-to-token ratio. This forces the architecture to memorize syntax and knowledge far beyond Chinchilla-optimal limits.
+* **Data Mixture:** 
+  * 60% Finephrase (Synthetic FineWeb-Edu)
+  * 20% DCLM (Web text)
+  * 10% FineMath (Mathematics)
+  * 10% CornStack (Code)
+* **Optimizer:** Utilizes a hybrid Muon (for 2D weight matrices) and AdamW (for embeddings/norms) optimizer setup for stable, rapid convergence.
+
+## Benchmark Results (Open SLM Leaderboard)
+
+Evaluated at checkpoint-24999.
+
+| Benchmark | Score (acc_norm) |
+| :--- | :--- |
+| HellaSwag | 39.43% |
+| ARC-Easy | 51.14% |
+| ARC-Challenge | 26.88% |
+| PIQA | 65.94% |
+| ArithMark-3 | 39.20% |
+| **Intelligence Index** | **22.50** |
+
+We used the Language Model Evaluation Harness for HellaSwag, ARC-Easy, ARC-Challenge, and PIQA (0-shot).  
+We used the provided script for ArithMark-3.
+
+## Usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+model_id = "SurjoLabs/Flare"
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id, 
+    trust_remote_code=True, 
+    torch_dtype=torch.bfloat16
+).cuda()
+
+prompt = "The capital of France is"
+inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+outputs = model.generate(**inputs, max_new_tokens=20)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+```
+
+## Acknowledgement
+We would like to thank [AxiomicLabs](https://huggingface.co/AxiomicLabs) for proving that XSA architecture is excellent for token efficiency.
+
+## Limitations
+This is an early test of the larger Surjo Project. The code is not very stable and we do not recommend using the modeling file without doing edits for training your own model.

chat_template.jinja

@@ -0,0 +1,4 @@
+{{ bos_token }}{% for message in messages %}{{ '<|im_start|>' + message['role'] + '
+' + message['content'] + '<|im_end|>
+' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
+' }}{% endif %}

checkpoint-1000/chat_template.jinja

@@ -0,0 +1,4 @@
+{{ bos_token }}{% for message in messages %}{{ '<|im_start|>' + message['role'] + '
+' + message['content'] + '<|im_end|>
+' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
+' }}{% endif %}

checkpoint-1000/config.json

@@ -0,0 +1,43 @@
+{
+  "architectures": [
+    "SurjoExpForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_surjo_exp.SurjoExpConfig",
+    "AutoModelForCausalLM": "modeling_surjo_exp.SurjoExpForCausalLM"
+  },
+  "bos_token_id": 0,
+  "coda_layers": 16,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "gradient_checkpointing": false,
+  "head_dim": 72,
+  "hidden_act": "silu",
+  "hidden_size": 576,
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "max_position_embeddings": 2048,
+  "mlp_bias": false,
+  "model_type": "surjo_exp",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 2,
+  "pad_token_id": 1,
+  "prelude_layers": 4,
+  "pretraining_tp": 1,
+  "recurrent_layers": 12,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.15.0",
+  "use_cache": false,
+  "use_flash_attn": false,
+  "vocab_size": 32768,
+  "xsa_projection": true
+}

checkpoint-1000/configuration_surjo_exp.py

@@ -0,0 +1,21 @@
+
+from transformers import LlamaConfig
+
+class SurjoExpConfig(LlamaConfig):
+    model_type = "surjo_exp"
+    def __init__(self, *args, xsa_projection=True, rope_theta=10000.0, attention_bias=False,
+                 prelude_layers=4, recurrent_layers=12, coda_layers=16,
+                 gradient_checkpointing=True, use_flash_attn=True, **kwargs):
+        kwargs["num_hidden_layers"] = prelude_layers + recurrent_layers + coda_layers
+        kwargs.setdefault("use_cache", False)
+        super().__init__(*args, rope_theta=rope_theta, attention_bias=attention_bias, **kwargs)
+        self.xsa_projection = xsa_projection
+        self.rope_theta = rope_theta
+        self.attention_bias = attention_bias
+        self.prelude_layers = prelude_layers
+        self.recurrent_layers = recurrent_layers
+        self.coda_layers = coda_layers
+        self.gradient_checkpointing = gradient_checkpointing
+        self.use_flash_attn = use_flash_attn
+        if not hasattr(self, 'rope_parameters') or self.rope_parameters is None:
+            self.rope_parameters = {"rope_type": "default", "factor": 1.0, "rope_theta": rope_theta}

checkpoint-1000/generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 2,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 1,
+  "transformers_version": "5.15.0",
+  "use_cache": false
+}

checkpoint-1000/model.safetensors

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

checkpoint-1000/modeling_surjo_exp.py

@@ -0,0 +1,417 @@
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from transformers import LlamaConfig, LlamaModel, LlamaForCausalLM
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+from transformers.models.llama.modeling_llama import LlamaMLP
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding, apply_rotary_pos_emb
+from transformers.cache_utils import DynamicCache
+
+try:
+    from .configuration_surjo_exp import SurjoExpConfig
+except ImportError:
+    from configuration_surjo_exp import SurjoExpConfig
+
+try:
+    from flash_attn import flash_attn_varlen_func
+    FLASH_ATTN_AVAILABLE = True
+except ImportError:
+    FLASH_ATTN_AVAILABLE = False
+
+@torch._dynamo.disable()
+def _flash_varlen(q, k, v, cu_seqlens, max_seqlen, dropout_p):
+    ms = int(max_seqlen.item()) if torch.is_tensor(max_seqlen) else int(max_seqlen)
+    return flash_attn_varlen_func(
+        q, k, v, cu_seqlens, cu_seqlens, ms, ms,
+        dropout_p=dropout_p, causal=True,
+    )
+
+class ClampedLlamaMLP(LlamaMLP):
+    def forward(self, x):
+        gate = F.silu(self.gate_proj(x).clamp(-15.0, 15.0))
+        up = self.up_proj(x)
+        return self.down_proj(gate * up)
+
+class XSAAttention(nn.Module):
+    def __init__(self, config, layer_idx=None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.recurrent_cache_idx = None
+        self._use_recurrent_slot = False
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
+        self.attention_bias = getattr(config, "attention_bias", False)
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=self.attention_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.attention_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.attention_bias)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=self.attention_bias)
+
+        self.q_norm = LlamaRMSNorm(self.head_dim, eps=1e-6)
+        self.k_norm = LlamaRMSNorm(self.head_dim, eps=1e-6)
+
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None,
+                output_attentions=False, use_cache=False, cache_position=None, position_embeddings=None,
+                expected_batch_size=None, cu_seqlens=None, max_seqlen=None, **kwargs):
+        past_kv = past_key_value if past_key_value is not None else kwargs.get("past_key_values", None)
+
+        if hidden_states.ndim == 2:
+            if expected_batch_size is None:
+                raise RuntimeError(
+                    f"XSAAttention received 2D hidden_states {hidden_states.shape} "
+                    f"without an expected_batch_size to safely restore the batch dim."
+                )
+            hidden_states = hidden_states.reshape(expected_batch_size, -1, self.hidden_size)
+
+        bsz, q_len, _ = hidden_states.size()
+
+        if expected_batch_size is not None and bsz != expected_batch_size:
+            raise RuntimeError(
+                f"XSAAttention: hidden_states batch size {bsz} does not match "
+                f"expected_batch_size {expected_batch_size}."
+            )
+
+        query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim)
+        key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+        value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+
+        query_states = self.q_norm(query_states)
+        key_states = self.k_norm(key_states)
+
+        cos, sin = position_embeddings
+
+        use_flash = (
+            cu_seqlens is not None
+            and past_kv is None
+            and getattr(self.config, "use_flash_attn", False)
+            and FLASH_ATTN_AVAILABLE
+        )
+
+        if use_flash:
+            total = bsz * q_len
+            q = query_states.reshape(total, self.num_heads, self.head_dim)
+            k = key_states.reshape(total, self.num_key_value_heads, self.head_dim)
+            v = value_states.reshape(total, self.num_key_value_heads, self.head_dim)
+
+            q = q.to(torch.bfloat16)
+            k = k.to(torch.bfloat16)
+            v = v.to(torch.bfloat16)
+
+            cos_f = cos.reshape(-1, cos.shape[-1]).to(torch.bfloat16)
+            sin_f = sin.reshape(-1, sin.shape[-1]).to(torch.bfloat16)
+            q, k = apply_rotary_pos_emb(q, k, cos_f, sin_f, unsqueeze_dim=1)
+
+            attn_output = _flash_varlen(
+                q, k, v, cu_seqlens, max_seqlen,
+                self.config.attention_dropout if self.training else 0.0,
+            )
+
+            if getattr(self.config, 'xsa_projection', True):
+                y = attn_output.view(total, self.num_key_value_heads, self.num_key_value_groups, self.head_dim)
+                v_grouped = v.unsqueeze(2)
+                dot_yv = (y * v_grouped).sum(dim=-1, keepdim=True).float()
+                dot_vv = v_grouped.pow(2).sum(dim=-1, keepdim=True).clamp_min(1e-4).float()
+                scale = (dot_yv / dot_vv).to(y.dtype)
+                attn_output = (y - scale * v_grouped).reshape(total, self.num_heads, self.head_dim)
+
+            attn_output = self.o_proj(attn_output.reshape(bsz, q_len, self.hidden_size))
+            return (attn_output, None)
+
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        current_v = value_states
+
+        target_idx = self.layer_idx
+        if self._use_recurrent_slot and self.recurrent_cache_idx is not None:
+            target_idx = self.recurrent_cache_idx
+
+        if past_kv is not None:
+            while len(past_kv) <= target_idx:
+                past_kv.update(
+                    torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=key_states.dtype, device=key_states.device),
+                    torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=value_states.dtype, device=value_states.device),
+                    len(past_kv)
+                )
+            key_states, value_states = past_kv.update(key_states, value_states, target_idx)
+
+        key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
+        value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)
+
+        kv_len = key_states.shape[-2]
+
+        if attention_mask is not None:
+            if attention_mask.ndim == 2:
+                if attention_mask.shape[-1] < kv_len:
+                    attention_mask = F.pad(attention_mask, (0, kv_len - attention_mask.shape[-1]), value=1)
+                elif attention_mask.shape[-1] > kv_len:
+                    attention_mask = attention_mask[:, -kv_len:]
+
+                pad_mask = (1.0 - attention_mask[:, None, None, :].to(query_states.dtype)) * torch.finfo(query_states.dtype).min
+
+                if q_len > 1:
+                    if cache_position is None:
+                        cache_position = torch.arange(kv_len - q_len, kv_len, device=query_states.device)
+                    kv_positions = torch.arange(kv_len, device=query_states.device)
+
+                    neg_inf = torch.finfo(query_states.dtype).min
+                    causal_mask = torch.zeros((q_len, kv_len), dtype=query_states.dtype, device=query_states.device)
+                    causal_mask = causal_mask.masked_fill(kv_positions[None, :] > cache_position[:, None], neg_inf)
+                    attn_mask = causal_mask[None, None, :, :] + pad_mask
+
+                    diag_idx = torch.arange(q_len, device=attn_mask.device)
+                    start_idx = attn_mask.shape[-1] - q_len
+                    attn_mask[:, :, diag_idx, start_idx + diag_idx] = 0.0
+                else:
+                    attn_mask = pad_mask
+            else:
+                if attention_mask.shape[0] != bsz:
+                    raise RuntimeError(
+                        f"attention_mask batch size {attention_mask.shape[0]} does not "
+                        f"match hidden_states batch size {bsz}."
+                    )
+                attn_mask = attention_mask.to(dtype=query_states.dtype)
+            is_causal = False
+        else:
+            is_causal = True
+            attn_mask = None
+
+        attn_output = F.scaled_dot_product_attention(
+            query_states, key_states, value_states, attn_mask=attn_mask,
+            dropout_p=0.0 if not self.training else self.config.attention_dropout, is_causal=is_causal
+        )
+
+        if getattr(self.config, 'xsa_projection', True):
+            y = attn_output.reshape(bsz, self.num_key_value_heads, self.num_key_value_groups, q_len, self.head_dim)
+            v_grouped = current_v.unsqueeze(2)
+            dot_yv = (y * v_grouped).sum(dim=-1, keepdim=True).float()
+            dot_vv = v_grouped.pow(2).sum(dim=-1, keepdim=True).clamp_min(1e-4).float()
+            scale = (dot_yv / dot_vv).to(y.dtype)
+            attn_output = (y - scale * v_grouped).reshape(bsz, self.num_heads, q_len, self.head_dim)
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+        attn_output = self.o_proj(attn_output)
+
+        return (attn_output, None)
+
+@torch._dynamo.disable()
+def _checkpointed_layer_forward(layer, hidden_states, attention_mask, position_ids,
+                                 cache_position, cos, sin, expected_batch_size, cu_seqlens, max_seqlen):
+    out = layer(
+        hidden_states, attention_mask=attention_mask, position_ids=position_ids,
+        past_key_value=None, use_cache=False,
+        cache_position=cache_position, position_embeddings=(cos, sin),
+        expected_batch_size=expected_batch_size,
+        cu_seqlens=cu_seqlens, max_seqlen=max_seqlen,
+    )
+    hs_out = out[0] if isinstance(out, tuple) else out
+    if hs_out.ndim != 3 or hs_out.shape[0] != expected_batch_size:
+        raise RuntimeError(
+            f"Layer output shape {tuple(hs_out.shape)} does not match expected "
+            f"batch size {expected_batch_size}."
+        )
+    return hs_out
+
+class SurjoExpModel(LlamaModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        assert config.prelude_layers + config.recurrent_layers + config.coda_layers == config.num_hidden_layers, \
+            "prelude_layers + recurrent_layers + coda_layers must equal num_hidden_layers"
+
+        if getattr(config, "use_flash_attn", False) and not FLASH_ATTN_AVAILABLE:
+            raise ImportError(
+                "config.use_flash_attn=True but flash_attn is not importable. "
+                "Install the FA2 wheel or set use_flash_attn=False."
+            )
+
+        p1 = config.prelude_layers
+        r1 = p1 + config.recurrent_layers
+
+        for i, layer in enumerate(self.layers):
+            layer.self_attn = XSAAttention(config, layer_idx=i)
+            layer.mlp = ClampedLlamaMLP(config)
+
+        for i, layer in enumerate(self.layers[p1:r1]):
+            layer.self_attn.recurrent_cache_idx = config.num_hidden_layers + p1 + i
+
+        self.gradient_checkpointing = getattr(config, "gradient_checkpointing", True)
+
+    def gradient_checkpointing_enable(self):
+        self.gradient_checkpointing = True
+
+    def gradient_checkpointing_disable(self):
+        self.gradient_checkpointing = False
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None,
+                past_key_values=None, use_cache=None, output_attentions=False, output_hidden_states=False,
+                cache_position=None, return_dict=True, cu_seqlens=None, max_seqlen=None, **kwargs):
+        if use_cache is None:
+            use_cache = False
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        bsz, seq_len = inputs_embeds.shape[0], inputs_embeds.shape[1]
+
+        if cache_position is None:
+            past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(past_seen, past_seen + seq_len, dtype=torch.long, device=inputs_embeds.device)
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0).expand(bsz, -1)
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        cos, sin = position_embeddings
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        p1 = self.config.prelude_layers
+        r1 = p1 + self.config.recurrent_layers
+        c1 = r1 + self.config.coda_layers
+
+        prelude = self.layers[:p1]
+        recurrent = self.layers[p1:r1]
+        coda = self.layers[r1:c1]
+
+        use_ckpt = self.training and self.gradient_checkpointing and not use_cache
+
+        def run_layer(layer, hs):
+            if cu_seqlens is not None:
+                torch._dynamo.mark_dynamic(cu_seqlens, 0)
+
+            out = layer(
+                hs, attention_mask=attention_mask, position_ids=position_ids,
+                past_key_value=past_key_values if use_cache else None, use_cache=use_cache,
+                cache_position=cache_position, position_embeddings=position_embeddings,
+                expected_batch_size=bsz, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen,
+            )
+            hs_out = out[0] if isinstance(out, tuple) else out
+            if hs_out.ndim != 3 or hs_out.shape[0] != bsz:
+                raise RuntimeError(
+                    f"Layer output shape {tuple(hs_out.shape)} does not match expected "
+                    f"batch size {bsz}."
+                )
+            return hs_out
+
+        def run_layer_maybe_ckpt(layer, hs):
+            if use_ckpt:
+                return torch.utils.checkpoint.checkpoint(
+                    _checkpointed_layer_forward,
+                    layer, hs, attention_mask, position_ids, cache_position, cos, sin, bsz,
+                    cu_seqlens, max_seqlen,
+                    use_reentrant=False,
+                )
+            return run_layer(layer, hs)
+
+        for layer in prelude:
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        if self.training:
+            hidden_states = hidden_states + torch.randn_like(hidden_states) * 0.02
+
+        for layer in recurrent:
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        if self.training:
+            hidden_states = hidden_states + torch.randn_like(hidden_states) * 0.02
+
+        for layer in recurrent:
+            layer.self_attn._use_recurrent_slot = True
+            try:
+                hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+            finally:
+                layer.self_attn._use_recurrent_slot = False
+
+        for layer in coda:
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=past_key_values)
+
+class SurjoExpForCausalLM(LlamaForCausalLM):
+    config_class = SurjoExpConfig
+    def __init__(self, config):
+        super(LlamaForCausalLM, self).__init__(config)
+        self.model = SurjoExpModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def gradient_checkpointing_enable(self, **kwargs):
+        self.model.gradient_checkpointing_enable()
+
+    def gradient_checkpointing_disable(self):
+        self.model.gradient_checkpointing_disable()
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, inputs_embeds=None,
+                use_cache=None, num_logits_to_keep=0, position_ids=None, past_key_values=None,
+                cache_position=None, cu_seqlens=None, max_seqlen=None, **kwargs):
+        if use_cache is None:
+            use_cache = False if (self.training or labels is not None) else True
+
+        if num_logits_to_keep == 0 and "logits_to_keep" in kwargs:
+            num_logits_to_keep = kwargs["logits_to_keep"]
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
+        hidden_states = outputs[0]
+
+        expected_bsz = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]
+        if hidden_states.ndim != 3 or hidden_states.shape[0] != expected_bsz:
+            raise RuntimeError(
+                f"SurjoExpModel returned hidden_states with shape {tuple(hidden_states.shape)}, "
+                f"expected batch size {expected_bsz}."
+            )
+
+        loss = None
+        logits = None
+
+        if labels is not None:
+            shift_hidden = hidden_states[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            num_chunks = 8
+            h_chunks = shift_hidden.chunk(num_chunks, dim=0)
+            l_chunks = shift_labels.chunk(num_chunks, dim=0)
+
+            total_loss = hidden_states.new_zeros((), dtype=torch.float32)
+            total_tokens = 0
+            for h_c, l_c in zip(h_chunks, l_chunks):
+                logits_c = self.lm_head(h_c)
+                chunk_loss = F.cross_entropy(
+                    logits_c.view(-1, logits_c.size(-1)).float(),
+                    l_c.view(-1),
+                    reduction="sum",
+                )
+                total_loss = total_loss + chunk_loss
+                total_tokens += l_c.numel()
+            loss = (total_loss / total_tokens).to(hidden_states.dtype)
+        else:
+            slice_hidden = hidden_states if num_logits_to_keep == 0 else hidden_states[:, -num_logits_to_keep:, :]
+            logits = self.lm_head(slice_hidden)
+
+        return CausalLMOutputWithPast(
+            loss=loss, logits=logits, past_key_values=outputs.past_key_values
+        )

checkpoint-1000/optimizer.bin

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

checkpoint-1000/optimizer_1.bin

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

checkpoint-1000/random_states_0.pkl

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

checkpoint-1000/tokenizer_config.json

@@ -0,0 +1,85 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|bos|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|eos|>",
+  "extra_special_tokens": [
+    "<|unk|>",
+    "<|pad|>",
+    "<|bos|>",
+    "<|eos|>",
+    "<|mask|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|system|>",
+    "<|user|>",
+    "<|assistant|>",
+    "<think>",
+    "</think>",
+    "<|begin_of_thought|>",
+    "<|end_of_thought|>",
+    "<answer>",
+    "</answer>",
+    "<|step|>",
+    "<|/step|>",
+    "<context>",
+    "</context>",
+    "<|doc_start|>",
+    "<|doc_end|>",
+    "<|search|>",
+    "<|search_results|>",
+    "<|tool_list_start|>",
+    "<|tool_list_end|>",
+    "<tools>",
+    "</tools>",
+    "<|tool_call_start|>",
+    "<|tool_call_end|>",
+    "<|tool_call|>",
+    "<|/tool_call|>",
+    "<|tool_response_start|>",
+    "<|tool_response_end|>",
+    "<|tool_response|>",
+    "<|/tool_response|>",
+    "<|image|>",
+    "<|image_pad|>",
+    "<|image_placeholder|>",
+    "<|audio|>",
+    "<|audio_pad|>",
+    "<|audio_placeholder|>",
+    "<|video|>",
+    "<|video_pad|>",
+    "<|fim_prefix|>",
+    "<|fim_suffix|>",
+    "<|fim_middle|>",
+    "<|repo_name|>",
+    "<|file_separator|>",
+    "<|reward|>",
+    "<|reserved_0|>",
+    "<|reserved_1|>",
+    "<|reserved_2|>",
+    "<|reserved_3|>",
+    "<|reserved_4|>",
+    "<|reserved_5|>",
+    "<|reserved_6|>",
+    "<|reserved_7|>",
+    "<|reserved_8|>",
+    "<|reserved_9|>",
+    "<|reserved_10|>",
+    "<|reserved_11|>",
+    "<|reserved_12|>",
+    "<|reserved_13|>",
+    "<|reserved_14|>",
+    "<|reserved_15|>",
+    "<|reserved_16|>",
+    "<|reserved_17|>",
+    "<|reserved_18|>",
+    "<|reserved_19|>"
+  ],
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "<|mask|>",
+  "model_max_length": 10000000,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|unk|>"
+}

checkpoint-10000/chat_template.jinja

@@ -0,0 +1,4 @@
+{{ bos_token }}{% for message in messages %}{{ '<|im_start|>' + message['role'] + '
+' + message['content'] + '<|im_end|>
+' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
+' }}{% endif %}

checkpoint-10000/config.json

@@ -0,0 +1,43 @@
+{
+  "architectures": [
+    "SurjoExpForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_surjo_exp.SurjoExpConfig",
+    "AutoModelForCausalLM": "modeling_surjo_exp.SurjoExpForCausalLM"
+  },
+  "bos_token_id": 0,
+  "coda_layers": 16,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "gradient_checkpointing": false,
+  "head_dim": 72,
+  "hidden_act": "silu",
+  "hidden_size": 576,
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "max_position_embeddings": 2048,
+  "mlp_bias": false,
+  "model_type": "surjo_exp",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 2,
+  "pad_token_id": 1,
+  "prelude_layers": 4,
+  "pretraining_tp": 1,
+  "recurrent_layers": 12,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.15.0",
+  "use_cache": false,
+  "use_flash_attn": false,
+  "vocab_size": 32768,
+  "xsa_projection": true
+}

checkpoint-10000/configuration_surjo_exp.py

@@ -0,0 +1,21 @@
+
+from transformers import LlamaConfig
+
+class SurjoExpConfig(LlamaConfig):
+    model_type = "surjo_exp"
+    def __init__(self, *args, xsa_projection=True, rope_theta=10000.0, attention_bias=False,
+                 prelude_layers=4, recurrent_layers=12, coda_layers=16,
+                 gradient_checkpointing=True, use_flash_attn=True, **kwargs):
+        kwargs["num_hidden_layers"] = prelude_layers + recurrent_layers + coda_layers
+        kwargs.setdefault("use_cache", False)
+        super().__init__(*args, rope_theta=rope_theta, attention_bias=attention_bias, **kwargs)
+        self.xsa_projection = xsa_projection
+        self.rope_theta = rope_theta
+        self.attention_bias = attention_bias
+        self.prelude_layers = prelude_layers
+        self.recurrent_layers = recurrent_layers
+        self.coda_layers = coda_layers
+        self.gradient_checkpointing = gradient_checkpointing
+        self.use_flash_attn = use_flash_attn
+        if not hasattr(self, 'rope_parameters') or self.rope_parameters is None:
+            self.rope_parameters = {"rope_type": "default", "factor": 1.0, "rope_theta": rope_theta}

checkpoint-10000/generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 2,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 1,
+  "transformers_version": "5.15.0",
+  "use_cache": false
+}

checkpoint-10000/model.safetensors

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

checkpoint-10000/modeling_surjo_exp.py

@@ -0,0 +1,417 @@
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from transformers import LlamaConfig, LlamaModel, LlamaForCausalLM
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+from transformers.models.llama.modeling_llama import LlamaMLP
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding, apply_rotary_pos_emb
+from transformers.cache_utils import DynamicCache
+
+try:
+    from .configuration_surjo_exp import SurjoExpConfig
+except ImportError:
+    from configuration_surjo_exp import SurjoExpConfig
+
+try:
+    from flash_attn import flash_attn_varlen_func
+    FLASH_ATTN_AVAILABLE = True
+except ImportError:
+    FLASH_ATTN_AVAILABLE = False
+
+@torch._dynamo.disable()
+def _flash_varlen(q, k, v, cu_seqlens, max_seqlen, dropout_p):
+    ms = int(max_seqlen.item()) if torch.is_tensor(max_seqlen) else int(max_seqlen)
+    return flash_attn_varlen_func(
+        q, k, v, cu_seqlens, cu_seqlens, ms, ms,
+        dropout_p=dropout_p, causal=True,
+    )
+
+class ClampedLlamaMLP(LlamaMLP):
+    def forward(self, x):
+        gate = F.silu(self.gate_proj(x).clamp(-15.0, 15.0))
+        up = self.up_proj(x)
+        return self.down_proj(gate * up)
+
+class XSAAttention(nn.Module):
+    def __init__(self, config, layer_idx=None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.recurrent_cache_idx = None
+        self._use_recurrent_slot = False
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
+        self.attention_bias = getattr(config, "attention_bias", False)
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=self.attention_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.attention_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.attention_bias)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=self.attention_bias)
+
+        self.q_norm = LlamaRMSNorm(self.head_dim, eps=1e-6)
+        self.k_norm = LlamaRMSNorm(self.head_dim, eps=1e-6)
+
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None,
+                output_attentions=False, use_cache=False, cache_position=None, position_embeddings=None,
+                expected_batch_size=None, cu_seqlens=None, max_seqlen=None, **kwargs):
+        past_kv = past_key_value if past_key_value is not None else kwargs.get("past_key_values", None)
+
+        if hidden_states.ndim == 2:
+            if expected_batch_size is None:
+                raise RuntimeError(
+                    f"XSAAttention received 2D hidden_states {hidden_states.shape} "
+                    f"without an expected_batch_size to safely restore the batch dim."
+                )
+            hidden_states = hidden_states.reshape(expected_batch_size, -1, self.hidden_size)
+
+        bsz, q_len, _ = hidden_states.size()
+
+        if expected_batch_size is not None and bsz != expected_batch_size:
+            raise RuntimeError(
+                f"XSAAttention: hidden_states batch size {bsz} does not match "
+                f"expected_batch_size {expected_batch_size}."
+            )
+
+        query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim)
+        key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+        value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+
+        query_states = self.q_norm(query_states)
+        key_states = self.k_norm(key_states)
+
+        cos, sin = position_embeddings
+
+        use_flash = (
+            cu_seqlens is not None
+            and past_kv is None
+            and getattr(self.config, "use_flash_attn", False)
+            and FLASH_ATTN_AVAILABLE
+        )
+
+        if use_flash:
+            total = bsz * q_len
+            q = query_states.reshape(total, self.num_heads, self.head_dim)
+            k = key_states.reshape(total, self.num_key_value_heads, self.head_dim)
+            v = value_states.reshape(total, self.num_key_value_heads, self.head_dim)
+
+            q = q.to(torch.bfloat16)
+            k = k.to(torch.bfloat16)
+            v = v.to(torch.bfloat16)
+
+            cos_f = cos.reshape(-1, cos.shape[-1]).to(torch.bfloat16)
+            sin_f = sin.reshape(-1, sin.shape[-1]).to(torch.bfloat16)
+            q, k = apply_rotary_pos_emb(q, k, cos_f, sin_f, unsqueeze_dim=1)
+
+            attn_output = _flash_varlen(
+                q, k, v, cu_seqlens, max_seqlen,
+                self.config.attention_dropout if self.training else 0.0,
+            )
+
+            if getattr(self.config, 'xsa_projection', True):
+                y = attn_output.view(total, self.num_key_value_heads, self.num_key_value_groups, self.head_dim)
+                v_grouped = v.unsqueeze(2)
+                dot_yv = (y * v_grouped).sum(dim=-1, keepdim=True).float()
+                dot_vv = v_grouped.pow(2).sum(dim=-1, keepdim=True).clamp_min(1e-4).float()
+                scale = (dot_yv / dot_vv).to(y.dtype)
+                attn_output = (y - scale * v_grouped).reshape(total, self.num_heads, self.head_dim)
+
+            attn_output = self.o_proj(attn_output.reshape(bsz, q_len, self.hidden_size))
+            return (attn_output, None)
+
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        current_v = value_states
+
+        target_idx = self.layer_idx
+        if self._use_recurrent_slot and self.recurrent_cache_idx is not None:
+            target_idx = self.recurrent_cache_idx
+
+        if past_kv is not None:
+            while len(past_kv) <= target_idx:
+                past_kv.update(
+                    torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=key_states.dtype, device=key_states.device),
+                    torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=value_states.dtype, device=value_states.device),
+                    len(past_kv)
+                )
+            key_states, value_states = past_kv.update(key_states, value_states, target_idx)
+
+        key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
+        value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)
+
+        kv_len = key_states.shape[-2]
+
+        if attention_mask is not None:
+            if attention_mask.ndim == 2:
+                if attention_mask.shape[-1] < kv_len:
+                    attention_mask = F.pad(attention_mask, (0, kv_len - attention_mask.shape[-1]), value=1)
+                elif attention_mask.shape[-1] > kv_len:
+                    attention_mask = attention_mask[:, -kv_len:]
+
+                pad_mask = (1.0 - attention_mask[:, None, None, :].to(query_states.dtype)) * torch.finfo(query_states.dtype).min
+
+                if q_len > 1:
+                    if cache_position is None:
+                        cache_position = torch.arange(kv_len - q_len, kv_len, device=query_states.device)
+                    kv_positions = torch.arange(kv_len, device=query_states.device)
+
+                    neg_inf = torch.finfo(query_states.dtype).min
+                    causal_mask = torch.zeros((q_len, kv_len), dtype=query_states.dtype, device=query_states.device)
+                    causal_mask = causal_mask.masked_fill(kv_positions[None, :] > cache_position[:, None], neg_inf)
+                    attn_mask = causal_mask[None, None, :, :] + pad_mask
+
+                    diag_idx = torch.arange(q_len, device=attn_mask.device)
+                    start_idx = attn_mask.shape[-1] - q_len
+                    attn_mask[:, :, diag_idx, start_idx + diag_idx] = 0.0
+                else:
+                    attn_mask = pad_mask
+            else:
+                if attention_mask.shape[0] != bsz:
+                    raise RuntimeError(
+                        f"attention_mask batch size {attention_mask.shape[0]} does not "
+                        f"match hidden_states batch size {bsz}."
+                    )
+                attn_mask = attention_mask.to(dtype=query_states.dtype)
+            is_causal = False
+        else:
+            is_causal = True
+            attn_mask = None
+
+        attn_output = F.scaled_dot_product_attention(
+            query_states, key_states, value_states, attn_mask=attn_mask,
+            dropout_p=0.0 if not self.training else self.config.attention_dropout, is_causal=is_causal
+        )
+
+        if getattr(self.config, 'xsa_projection', True):
+            y = attn_output.reshape(bsz, self.num_key_value_heads, self.num_key_value_groups, q_len, self.head_dim)
+            v_grouped = current_v.unsqueeze(2)
+            dot_yv = (y * v_grouped).sum(dim=-1, keepdim=True).float()
+            dot_vv = v_grouped.pow(2).sum(dim=-1, keepdim=True).clamp_min(1e-4).float()
+            scale = (dot_yv / dot_vv).to(y.dtype)
+            attn_output = (y - scale * v_grouped).reshape(bsz, self.num_heads, q_len, self.head_dim)
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+        attn_output = self.o_proj(attn_output)
+
+        return (attn_output, None)
+
+@torch._dynamo.disable()
+def _checkpointed_layer_forward(layer, hidden_states, attention_mask, position_ids,
+                                 cache_position, cos, sin, expected_batch_size, cu_seqlens, max_seqlen):
+    out = layer(
+        hidden_states, attention_mask=attention_mask, position_ids=position_ids,
+        past_key_value=None, use_cache=False,
+        cache_position=cache_position, position_embeddings=(cos, sin),
+        expected_batch_size=expected_batch_size,
+        cu_seqlens=cu_seqlens, max_seqlen=max_seqlen,
+    )
+    hs_out = out[0] if isinstance(out, tuple) else out
+    if hs_out.ndim != 3 or hs_out.shape[0] != expected_batch_size:
+        raise RuntimeError(
+            f"Layer output shape {tuple(hs_out.shape)} does not match expected "
+            f"batch size {expected_batch_size}."
+        )
+    return hs_out
+
+class SurjoExpModel(LlamaModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        assert config.prelude_layers + config.recurrent_layers + config.coda_layers == config.num_hidden_layers, \
+            "prelude_layers + recurrent_layers + coda_layers must equal num_hidden_layers"
+
+        if getattr(config, "use_flash_attn", False) and not FLASH_ATTN_AVAILABLE:
+            raise ImportError(
+                "config.use_flash_attn=True but flash_attn is not importable. "
+                "Install the FA2 wheel or set use_flash_attn=False."
+            )
+
+        p1 = config.prelude_layers
+        r1 = p1 + config.recurrent_layers
+
+        for i, layer in enumerate(self.layers):
+            layer.self_attn = XSAAttention(config, layer_idx=i)
+            layer.mlp = ClampedLlamaMLP(config)
+
+        for i, layer in enumerate(self.layers[p1:r1]):
+            layer.self_attn.recurrent_cache_idx = config.num_hidden_layers + p1 + i
+
+        self.gradient_checkpointing = getattr(config, "gradient_checkpointing", True)
+
+    def gradient_checkpointing_enable(self):
+        self.gradient_checkpointing = True
+
+    def gradient_checkpointing_disable(self):
+        self.gradient_checkpointing = False
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None,
+                past_key_values=None, use_cache=None, output_attentions=False, output_hidden_states=False,
+                cache_position=None, return_dict=True, cu_seqlens=None, max_seqlen=None, **kwargs):
+        if use_cache is None:
+            use_cache = False
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        bsz, seq_len = inputs_embeds.shape[0], inputs_embeds.shape[1]
+
+        if cache_position is None:
+            past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(past_seen, past_seen + seq_len, dtype=torch.long, device=inputs_embeds.device)
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0).expand(bsz, -1)
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        cos, sin = position_embeddings
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        p1 = self.config.prelude_layers
+        r1 = p1 + self.config.recurrent_layers
+        c1 = r1 + self.config.coda_layers
+
+        prelude = self.layers[:p1]
+        recurrent = self.layers[p1:r1]
+        coda = self.layers[r1:c1]
+
+        use_ckpt = self.training and self.gradient_checkpointing and not use_cache
+
+        def run_layer(layer, hs):
+            if cu_seqlens is not None:
+                torch._dynamo.mark_dynamic(cu_seqlens, 0)
+
+            out = layer(
+                hs, attention_mask=attention_mask, position_ids=position_ids,
+                past_key_value=past_key_values if use_cache else None, use_cache=use_cache,
+                cache_position=cache_position, position_embeddings=position_embeddings,
+                expected_batch_size=bsz, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen,
+            )
+            hs_out = out[0] if isinstance(out, tuple) else out
+            if hs_out.ndim != 3 or hs_out.shape[0] != bsz:
+                raise RuntimeError(
+                    f"Layer output shape {tuple(hs_out.shape)} does not match expected "
+                    f"batch size {bsz}."
+                )
+            return hs_out
+
+        def run_layer_maybe_ckpt(layer, hs):
+            if use_ckpt:
+                return torch.utils.checkpoint.checkpoint(
+                    _checkpointed_layer_forward,
+                    layer, hs, attention_mask, position_ids, cache_position, cos, sin, bsz,
+                    cu_seqlens, max_seqlen,
+                    use_reentrant=False,
+                )
+            return run_layer(layer, hs)
+
+        for layer in prelude:
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        if self.training:
+            hidden_states = hidden_states + torch.randn_like(hidden_states) * 0.02
+
+        for layer in recurrent:
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        if self.training:
+            hidden_states = hidden_states + torch.randn_like(hidden_states) * 0.02
+
+        for layer in recurrent:
+            layer.self_attn._use_recurrent_slot = True
+            try:
+                hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+            finally:
+                layer.self_attn._use_recurrent_slot = False
+
+        for layer in coda:
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=past_key_values)
+
+class SurjoExpForCausalLM(LlamaForCausalLM):
+    config_class = SurjoExpConfig
+    def __init__(self, config):
+        super(LlamaForCausalLM, self).__init__(config)
+        self.model = SurjoExpModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def gradient_checkpointing_enable(self, **kwargs):
+        self.model.gradient_checkpointing_enable()
+
+    def gradient_checkpointing_disable(self):
+        self.model.gradient_checkpointing_disable()
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, inputs_embeds=None,
+                use_cache=None, num_logits_to_keep=0, position_ids=None, past_key_values=None,
+                cache_position=None, cu_seqlens=None, max_seqlen=None, **kwargs):
+        if use_cache is None:
+            use_cache = False if (self.training or labels is not None) else True
+
+        if num_logits_to_keep == 0 and "logits_to_keep" in kwargs:
+            num_logits_to_keep = kwargs["logits_to_keep"]
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
+        hidden_states = outputs[0]
+
+        expected_bsz = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]
+        if hidden_states.ndim != 3 or hidden_states.shape[0] != expected_bsz:
+            raise RuntimeError(
+                f"SurjoExpModel returned hidden_states with shape {tuple(hidden_states.shape)}, "
+                f"expected batch size {expected_bsz}."
+            )
+
+        loss = None
+        logits = None
+
+        if labels is not None:
+            shift_hidden = hidden_states[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            num_chunks = 8
+            h_chunks = shift_hidden.chunk(num_chunks, dim=0)
+            l_chunks = shift_labels.chunk(num_chunks, dim=0)
+
+            total_loss = hidden_states.new_zeros((), dtype=torch.float32)
+            total_tokens = 0
+            for h_c, l_c in zip(h_chunks, l_chunks):
+                logits_c = self.lm_head(h_c)
+                chunk_loss = F.cross_entropy(
+                    logits_c.view(-1, logits_c.size(-1)).float(),
+                    l_c.view(-1),
+                    reduction="sum",
+                )
+                total_loss = total_loss + chunk_loss
+                total_tokens += l_c.numel()
+            loss = (total_loss / total_tokens).to(hidden_states.dtype)
+        else:
+            slice_hidden = hidden_states if num_logits_to_keep == 0 else hidden_states[:, -num_logits_to_keep:, :]
+            logits = self.lm_head(slice_hidden)
+
+        return CausalLMOutputWithPast(
+            loss=loss, logits=logits, past_key_values=outputs.past_key_values
+        )

checkpoint-10000/optimizer.bin

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

checkpoint-10000/optimizer_1.bin

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

checkpoint-10000/random_states_0.pkl

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

checkpoint-10000/tokenizer_config.json

@@ -0,0 +1,85 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|bos|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|eos|>",
+  "extra_special_tokens": [
+    "<|unk|>",
+    "<|pad|>",
+    "<|bos|>",
+    "<|eos|>",
+    "<|mask|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|system|>",
+    "<|user|>",
+    "<|assistant|>",
+    "<think>",
+    "</think>",
+    "<|begin_of_thought|>",
+    "<|end_of_thought|>",
+    "<answer>",
+    "</answer>",
+    "<|step|>",
+    "<|/step|>",
+    "<context>",
+    "</context>",
+    "<|doc_start|>",
+    "<|doc_end|>",
+    "<|search|>",
+    "<|search_results|>",
+    "<|tool_list_start|>",
+    "<|tool_list_end|>",
+    "<tools>",
+    "</tools>",
+    "<|tool_call_start|>",
+    "<|tool_call_end|>",
+    "<|tool_call|>",
+    "<|/tool_call|>",
+    "<|tool_response_start|>",
+    "<|tool_response_end|>",
+    "<|tool_response|>",
+    "<|/tool_response|>",
+    "<|image|>",
+    "<|image_pad|>",
+    "<|image_placeholder|>",
+    "<|audio|>",
+    "<|audio_pad|>",
+    "<|audio_placeholder|>",
+    "<|video|>",
+    "<|video_pad|>",
+    "<|fim_prefix|>",
+    "<|fim_suffix|>",
+    "<|fim_middle|>",
+    "<|repo_name|>",
+    "<|file_separator|>",
+    "<|reward|>",
+    "<|reserved_0|>",
+    "<|reserved_1|>",
+    "<|reserved_2|>",
+    "<|reserved_3|>",
+    "<|reserved_4|>",
+    "<|reserved_5|>",
+    "<|reserved_6|>",
+    "<|reserved_7|>",
+    "<|reserved_8|>",
+    "<|reserved_9|>",
+    "<|reserved_10|>",
+    "<|reserved_11|>",
+    "<|reserved_12|>",
+    "<|reserved_13|>",
+    "<|reserved_14|>",
+    "<|reserved_15|>",
+    "<|reserved_16|>",
+    "<|reserved_17|>",
+    "<|reserved_18|>",
+    "<|reserved_19|>"
+  ],
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "<|mask|>",
+  "model_max_length": 10000000,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|unk|>"
+}

checkpoint-11000/chat_template.jinja

@@ -0,0 +1,4 @@
+{{ bos_token }}{% for message in messages %}{{ '<|im_start|>' + message['role'] + '
+' + message['content'] + '<|im_end|>
+' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
+' }}{% endif %}

checkpoint-11000/config.json

@@ -0,0 +1,43 @@
+{
+  "architectures": [
+    "SurjoExpForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_surjo_exp.SurjoExpConfig",
+    "AutoModelForCausalLM": "modeling_surjo_exp.SurjoExpForCausalLM"
+  },
+  "bos_token_id": 0,
+  "coda_layers": 16,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "gradient_checkpointing": false,
+  "head_dim": 72,
+  "hidden_act": "silu",
+  "hidden_size": 576,
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "max_position_embeddings": 2048,
+  "mlp_bias": false,
+  "model_type": "surjo_exp",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 2,
+  "pad_token_id": 1,
+  "prelude_layers": 4,
+  "pretraining_tp": 1,
+  "recurrent_layers": 12,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.15.0",
+  "use_cache": false,
+  "use_flash_attn": false,
+  "vocab_size": 32768,
+  "xsa_projection": true
+}

checkpoint-11000/configuration_surjo_exp.py

@@ -0,0 +1,21 @@
+
+from transformers import LlamaConfig
+
+class SurjoExpConfig(LlamaConfig):
+    model_type = "surjo_exp"
+    def __init__(self, *args, xsa_projection=True, rope_theta=10000.0, attention_bias=False,
+                 prelude_layers=4, recurrent_layers=12, coda_layers=16,
+                 gradient_checkpointing=True, use_flash_attn=True, **kwargs):
+        kwargs["num_hidden_layers"] = prelude_layers + recurrent_layers + coda_layers
+        kwargs.setdefault("use_cache", False)
+        super().__init__(*args, rope_theta=rope_theta, attention_bias=attention_bias, **kwargs)
+        self.xsa_projection = xsa_projection
+        self.rope_theta = rope_theta
+        self.attention_bias = attention_bias
+        self.prelude_layers = prelude_layers
+        self.recurrent_layers = recurrent_layers
+        self.coda_layers = coda_layers
+        self.gradient_checkpointing = gradient_checkpointing
+        self.use_flash_attn = use_flash_attn
+        if not hasattr(self, 'rope_parameters') or self.rope_parameters is None:
+            self.rope_parameters = {"rope_type": "default", "factor": 1.0, "rope_theta": rope_theta}

checkpoint-11000/generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 2,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 1,
+  "transformers_version": "5.15.0",
+  "use_cache": false
+}

checkpoint-11000/model.safetensors

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

checkpoint-11000/modeling_surjo_exp.py

@@ -0,0 +1,417 @@
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from transformers import LlamaConfig, LlamaModel, LlamaForCausalLM
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+from transformers.models.llama.modeling_llama import LlamaMLP
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding, apply_rotary_pos_emb
+from transformers.cache_utils import DynamicCache
+
+try:
+    from .configuration_surjo_exp import SurjoExpConfig
+except ImportError:
+    from configuration_surjo_exp import SurjoExpConfig
+
+try:
+    from flash_attn import flash_attn_varlen_func
+    FLASH_ATTN_AVAILABLE = True
+except ImportError:
+    FLASH_ATTN_AVAILABLE = False
+
+@torch._dynamo.disable()
+def _flash_varlen(q, k, v, cu_seqlens, max_seqlen, dropout_p):
+    ms = int(max_seqlen.item()) if torch.is_tensor(max_seqlen) else int(max_seqlen)
+    return flash_attn_varlen_func(
+        q, k, v, cu_seqlens, cu_seqlens, ms, ms,
+        dropout_p=dropout_p, causal=True,
+    )
+
+class ClampedLlamaMLP(LlamaMLP):
+    def forward(self, x):
+        gate = F.silu(self.gate_proj(x).clamp(-15.0, 15.0))
+        up = self.up_proj(x)
+        return self.down_proj(gate * up)
+
+class XSAAttention(nn.Module):
+    def __init__(self, config, layer_idx=None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.recurrent_cache_idx = None
+        self._use_recurrent_slot = False
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
+        self.attention_bias = getattr(config, "attention_bias", False)
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=self.attention_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.attention_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.attention_bias)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=self.attention_bias)
+
+        self.q_norm = LlamaRMSNorm(self.head_dim, eps=1e-6)
+        self.k_norm = LlamaRMSNorm(self.head_dim, eps=1e-6)
+
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None,
+                output_attentions=False, use_cache=False, cache_position=None, position_embeddings=None,
+                expected_batch_size=None, cu_seqlens=None, max_seqlen=None, **kwargs):
+        past_kv = past_key_value if past_key_value is not None else kwargs.get("past_key_values", None)
+
+        if hidden_states.ndim == 2:
+            if expected_batch_size is None:
+                raise RuntimeError(
+                    f"XSAAttention received 2D hidden_states {hidden_states.shape} "
+                    f"without an expected_batch_size to safely restore the batch dim."
+                )
+            hidden_states = hidden_states.reshape(expected_batch_size, -1, self.hidden_size)
+
+        bsz, q_len, _ = hidden_states.size()
+
+        if expected_batch_size is not None and bsz != expected_batch_size:
+            raise RuntimeError(
+                f"XSAAttention: hidden_states batch size {bsz} does not match "
+                f"expected_batch_size {expected_batch_size}."
+            )
+
+        query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim)
+        key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+        value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+
+        query_states = self.q_norm(query_states)
+        key_states = self.k_norm(key_states)
+
+        cos, sin = position_embeddings
+
+        use_flash = (
+            cu_seqlens is not None
+            and past_kv is None
+            and getattr(self.config, "use_flash_attn", False)
+            and FLASH_ATTN_AVAILABLE
+        )
+
+        if use_flash:
+            total = bsz * q_len
+            q = query_states.reshape(total, self.num_heads, self.head_dim)
+            k = key_states.reshape(total, self.num_key_value_heads, self.head_dim)
+            v = value_states.reshape(total, self.num_key_value_heads, self.head_dim)
+
+            q = q.to(torch.bfloat16)
+            k = k.to(torch.bfloat16)
+            v = v.to(torch.bfloat16)
+
+            cos_f = cos.reshape(-1, cos.shape[-1]).to(torch.bfloat16)
+            sin_f = sin.reshape(-1, sin.shape[-1]).to(torch.bfloat16)
+            q, k = apply_rotary_pos_emb(q, k, cos_f, sin_f, unsqueeze_dim=1)
+
+            attn_output = _flash_varlen(
+                q, k, v, cu_seqlens, max_seqlen,
+                self.config.attention_dropout if self.training else 0.0,
+            )
+
+            if getattr(self.config, 'xsa_projection', True):
+                y = attn_output.view(total, self.num_key_value_heads, self.num_key_value_groups, self.head_dim)
+                v_grouped = v.unsqueeze(2)
+                dot_yv = (y * v_grouped).sum(dim=-1, keepdim=True).float()
+                dot_vv = v_grouped.pow(2).sum(dim=-1, keepdim=True).clamp_min(1e-4).float()
+                scale = (dot_yv / dot_vv).to(y.dtype)
+                attn_output = (y - scale * v_grouped).reshape(total, self.num_heads, self.head_dim)
+
+            attn_output = self.o_proj(attn_output.reshape(bsz, q_len, self.hidden_size))
+            return (attn_output, None)
+
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        current_v = value_states
+
+        target_idx = self.layer_idx
+        if self._use_recurrent_slot and self.recurrent_cache_idx is not None:
+            target_idx = self.recurrent_cache_idx
+
+        if past_kv is not None:
+            while len(past_kv) <= target_idx:
+                past_kv.update(
+                    torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=key_states.dtype, device=key_states.device),
+                    torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=value_states.dtype, device=value_states.device),
+                    len(past_kv)
+                )
+            key_states, value_states = past_kv.update(key_states, value_states, target_idx)
+
+        key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
+        value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)
+
+        kv_len = key_states.shape[-2]
+
+        if attention_mask is not None:
+            if attention_mask.ndim == 2:
+                if attention_mask.shape[-1] < kv_len:
+                    attention_mask = F.pad(attention_mask, (0, kv_len - attention_mask.shape[-1]), value=1)
+                elif attention_mask.shape[-1] > kv_len:
+                    attention_mask = attention_mask[:, -kv_len:]
+
+                pad_mask = (1.0 - attention_mask[:, None, None, :].to(query_states.dtype)) * torch.finfo(query_states.dtype).min
+
+                if q_len > 1:
+                    if cache_position is None:
+                        cache_position = torch.arange(kv_len - q_len, kv_len, device=query_states.device)
+                    kv_positions = torch.arange(kv_len, device=query_states.device)
+
+                    neg_inf = torch.finfo(query_states.dtype).min
+                    causal_mask = torch.zeros((q_len, kv_len), dtype=query_states.dtype, device=query_states.device)
+                    causal_mask = causal_mask.masked_fill(kv_positions[None, :] > cache_position[:, None], neg_inf)
+                    attn_mask = causal_mask[None, None, :, :] + pad_mask
+
+                    diag_idx = torch.arange(q_len, device=attn_mask.device)
+                    start_idx = attn_mask.shape[-1] - q_len
+                    attn_mask[:, :, diag_idx, start_idx + diag_idx] = 0.0
+                else:
+                    attn_mask = pad_mask
+            else:
+                if attention_mask.shape[0] != bsz:
+                    raise RuntimeError(
+                        f"attention_mask batch size {attention_mask.shape[0]} does not "
+                        f"match hidden_states batch size {bsz}."
+                    )
+                attn_mask = attention_mask.to(dtype=query_states.dtype)
+            is_causal = False
+        else:
+            is_causal = True
+            attn_mask = None
+
+        attn_output = F.scaled_dot_product_attention(
+            query_states, key_states, value_states, attn_mask=attn_mask,
+            dropout_p=0.0 if not self.training else self.config.attention_dropout, is_causal=is_causal
+        )
+
+        if getattr(self.config, 'xsa_projection', True):
+            y = attn_output.reshape(bsz, self.num_key_value_heads, self.num_key_value_groups, q_len, self.head_dim)
+            v_grouped = current_v.unsqueeze(2)
+            dot_yv = (y * v_grouped).sum(dim=-1, keepdim=True).float()
+            dot_vv = v_grouped.pow(2).sum(dim=-1, keepdim=True).clamp_min(1e-4).float()
+            scale = (dot_yv / dot_vv).to(y.dtype)
+            attn_output = (y - scale * v_grouped).reshape(bsz, self.num_heads, q_len, self.head_dim)
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+        attn_output = self.o_proj(attn_output)
+
+        return (attn_output, None)
+
+@torch._dynamo.disable()
+def _checkpointed_layer_forward(layer, hidden_states, attention_mask, position_ids,
+                                 cache_position, cos, sin, expected_batch_size, cu_seqlens, max_seqlen):
+    out = layer(
+        hidden_states, attention_mask=attention_mask, position_ids=position_ids,
+        past_key_value=None, use_cache=False,
+        cache_position=cache_position, position_embeddings=(cos, sin),
+        expected_batch_size=expected_batch_size,
+        cu_seqlens=cu_seqlens, max_seqlen=max_seqlen,
+    )
+    hs_out = out[0] if isinstance(out, tuple) else out
+    if hs_out.ndim != 3 or hs_out.shape[0] != expected_batch_size:
+        raise RuntimeError(
+            f"Layer output shape {tuple(hs_out.shape)} does not match expected "
+            f"batch size {expected_batch_size}."
+        )
+    return hs_out
+
+class SurjoExpModel(LlamaModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        assert config.prelude_layers + config.recurrent_layers + config.coda_layers == config.num_hidden_layers, \
+            "prelude_layers + recurrent_layers + coda_layers must equal num_hidden_layers"
+
+        if getattr(config, "use_flash_attn", False) and not FLASH_ATTN_AVAILABLE:
+            raise ImportError(
+                "config.use_flash_attn=True but flash_attn is not importable. "
+                "Install the FA2 wheel or set use_flash_attn=False."
+            )
+
+        p1 = config.prelude_layers
+        r1 = p1 + config.recurrent_layers
+
+        for i, layer in enumerate(self.layers):
+            layer.self_attn = XSAAttention(config, layer_idx=i)
+            layer.mlp = ClampedLlamaMLP(config)
+
+        for i, layer in enumerate(self.layers[p1:r1]):
+            layer.self_attn.recurrent_cache_idx = config.num_hidden_layers + p1 + i
+
+        self.gradient_checkpointing = getattr(config, "gradient_checkpointing", True)
+
+    def gradient_checkpointing_enable(self):
+        self.gradient_checkpointing = True
+
+    def gradient_checkpointing_disable(self):
+        self.gradient_checkpointing = False
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None,
+                past_key_values=None, use_cache=None, output_attentions=False, output_hidden_states=False,
+                cache_position=None, return_dict=True, cu_seqlens=None, max_seqlen=None, **kwargs):
+        if use_cache is None:
+            use_cache = False
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        bsz, seq_len = inputs_embeds.shape[0], inputs_embeds.shape[1]
+
+        if cache_position is None:
+            past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(past_seen, past_seen + seq_len, dtype=torch.long, device=inputs_embeds.device)
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0).expand(bsz, -1)
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        cos, sin = position_embeddings
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        p1 = self.config.prelude_layers
+        r1 = p1 + self.config.recurrent_layers
+        c1 = r1 + self.config.coda_layers
+
+        prelude = self.layers[:p1]
+        recurrent = self.layers[p1:r1]
+        coda = self.layers[r1:c1]
+
+        use_ckpt = self.training and self.gradient_checkpointing and not use_cache
+
+        def run_layer(layer, hs):
+            if cu_seqlens is not None:
+                torch._dynamo.mark_dynamic(cu_seqlens, 0)
+
+            out = layer(
+                hs, attention_mask=attention_mask, position_ids=position_ids,
+                past_key_value=past_key_values if use_cache else None, use_cache=use_cache,
+                cache_position=cache_position, position_embeddings=position_embeddings,
+                expected_batch_size=bsz, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen,
+            )
+            hs_out = out[0] if isinstance(out, tuple) else out
+            if hs_out.ndim != 3 or hs_out.shape[0] != bsz:
+                raise RuntimeError(
+                    f"Layer output shape {tuple(hs_out.shape)} does not match expected "
+                    f"batch size {bsz}."
+                )
+            return hs_out
+
+        def run_layer_maybe_ckpt(layer, hs):
+            if use_ckpt:
+                return torch.utils.checkpoint.checkpoint(
+                    _checkpointed_layer_forward,
+                    layer, hs, attention_mask, position_ids, cache_position, cos, sin, bsz,
+                    cu_seqlens, max_seqlen,
+                    use_reentrant=False,
+                )
+            return run_layer(layer, hs)
+
+        for layer in prelude:
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        if self.training:
+            hidden_states = hidden_states + torch.randn_like(hidden_states) * 0.02
+
+        for layer in recurrent:
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        if self.training:
+            hidden_states = hidden_states + torch.randn_like(hidden_states) * 0.02
+
+        for layer in recurrent:
+            layer.self_attn._use_recurrent_slot = True
+            try:
+                hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+            finally:
+                layer.self_attn._use_recurrent_slot = False
+
+        for layer in coda:
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=past_key_values)
+
+class SurjoExpForCausalLM(LlamaForCausalLM):
+    config_class = SurjoExpConfig
+    def __init__(self, config):
+        super(LlamaForCausalLM, self).__init__(config)
+        self.model = SurjoExpModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def gradient_checkpointing_enable(self, **kwargs):
+        self.model.gradient_checkpointing_enable()
+
+    def gradient_checkpointing_disable(self):
+        self.model.gradient_checkpointing_disable()
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, inputs_embeds=None,
+                use_cache=None, num_logits_to_keep=0, position_ids=None, past_key_values=None,
+                cache_position=None, cu_seqlens=None, max_seqlen=None, **kwargs):
+        if use_cache is None:
+            use_cache = False if (self.training or labels is not None) else True
+
+        if num_logits_to_keep == 0 and "logits_to_keep" in kwargs:
+            num_logits_to_keep = kwargs["logits_to_keep"]
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
+        hidden_states = outputs[0]
+
+        expected_bsz = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]
+        if hidden_states.ndim != 3 or hidden_states.shape[0] != expected_bsz:
+            raise RuntimeError(
+                f"SurjoExpModel returned hidden_states with shape {tuple(hidden_states.shape)}, "
+                f"expected batch size {expected_bsz}."
+            )
+
+        loss = None
+        logits = None
+
+        if labels is not None:
+            shift_hidden = hidden_states[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            num_chunks = 8
+            h_chunks = shift_hidden.chunk(num_chunks, dim=0)
+            l_chunks = shift_labels.chunk(num_chunks, dim=0)
+
+            total_loss = hidden_states.new_zeros((), dtype=torch.float32)
+            total_tokens = 0
+            for h_c, l_c in zip(h_chunks, l_chunks):
+                logits_c = self.lm_head(h_c)
+                chunk_loss = F.cross_entropy(
+                    logits_c.view(-1, logits_c.size(-1)).float(),
+                    l_c.view(-1),
+                    reduction="sum",
+                )
+                total_loss = total_loss + chunk_loss
+                total_tokens += l_c.numel()
+            loss = (total_loss / total_tokens).to(hidden_states.dtype)
+        else:
+            slice_hidden = hidden_states if num_logits_to_keep == 0 else hidden_states[:, -num_logits_to_keep:, :]
+            logits = self.lm_head(slice_hidden)
+
+        return CausalLMOutputWithPast(
+            loss=loss, logits=logits, past_key_values=outputs.past_key_values
+        )

checkpoint-11000/optimizer.bin

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

checkpoint-11000/optimizer_1.bin

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

checkpoint-11000/random_states_0.pkl

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

checkpoint-11000/tokenizer_config.json

@@ -0,0 +1,85 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|bos|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|eos|>",
+  "extra_special_tokens": [
+    "<|unk|>",
+    "<|pad|>",
+    "<|bos|>",
+    "<|eos|>",
+    "<|mask|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|system|>",
+    "<|user|>",
+    "<|assistant|>",
+    "<think>",
+    "</think>",
+    "<|begin_of_thought|>",
+    "<|end_of_thought|>",
+    "<answer>",
+    "</answer>",
+    "<|step|>",
+    "<|/step|>",
+    "<context>",
+    "</context>",
+    "<|doc_start|>",
+    "<|doc_end|>",
+    "<|search|>",
+    "<|search_results|>",
+    "<|tool_list_start|>",
+    "<|tool_list_end|>",
+    "<tools>",
+    "</tools>",
+    "<|tool_call_start|>",
+    "<|tool_call_end|>",
+    "<|tool_call|>",
+    "<|/tool_call|>",
+    "<|tool_response_start|>",
+    "<|tool_response_end|>",
+    "<|tool_response|>",
+    "<|/tool_response|>",
+    "<|image|>",
+    "<|image_pad|>",
+    "<|image_placeholder|>",
+    "<|audio|>",
+    "<|audio_pad|>",
+    "<|audio_placeholder|>",
+    "<|video|>",
+    "<|video_pad|>",
+    "<|fim_prefix|>",
+    "<|fim_suffix|>",
+    "<|fim_middle|>",
+    "<|repo_name|>",
+    "<|file_separator|>",
+    "<|reward|>",
+    "<|reserved_0|>",
+    "<|reserved_1|>",
+    "<|reserved_2|>",
+    "<|reserved_3|>",
+    "<|reserved_4|>",
+    "<|reserved_5|>",
+    "<|reserved_6|>",
+    "<|reserved_7|>",
+    "<|reserved_8|>",
+    "<|reserved_9|>",
+    "<|reserved_10|>",
+    "<|reserved_11|>",
+    "<|reserved_12|>",
+    "<|reserved_13|>",
+    "<|reserved_14|>",
+    "<|reserved_15|>",
+    "<|reserved_16|>",
+    "<|reserved_17|>",
+    "<|reserved_18|>",
+    "<|reserved_19|>"
+  ],
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "<|mask|>",
+  "model_max_length": 10000000,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|unk|>"
+}

checkpoint-12000/chat_template.jinja

@@ -0,0 +1,4 @@
+{{ bos_token }}{% for message in messages %}{{ '<|im_start|>' + message['role'] + '
+' + message['content'] + '<|im_end|>
+' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
+' }}{% endif %}

checkpoint-12000/config.json

@@ -0,0 +1,43 @@
+{
+  "architectures": [
+    "SurjoExpForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_surjo_exp.SurjoExpConfig",
+    "AutoModelForCausalLM": "modeling_surjo_exp.SurjoExpForCausalLM"
+  },
+  "bos_token_id": 0,
+  "coda_layers": 16,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "gradient_checkpointing": false,
+  "head_dim": 72,
+  "hidden_act": "silu",
+  "hidden_size": 576,
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "max_position_embeddings": 2048,
+  "mlp_bias": false,
+  "model_type": "surjo_exp",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 2,
+  "pad_token_id": 1,
+  "prelude_layers": 4,
+  "pretraining_tp": 1,
+  "recurrent_layers": 12,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.15.0",
+  "use_cache": false,
+  "use_flash_attn": false,
+  "vocab_size": 32768,
+  "xsa_projection": true
+}

checkpoint-12000/configuration_surjo_exp.py

@@ -0,0 +1,21 @@
+
+from transformers import LlamaConfig
+
+class SurjoExpConfig(LlamaConfig):
+    model_type = "surjo_exp"
+    def __init__(self, *args, xsa_projection=True, rope_theta=10000.0, attention_bias=False,
+                 prelude_layers=4, recurrent_layers=12, coda_layers=16,
+                 gradient_checkpointing=True, use_flash_attn=True, **kwargs):
+        kwargs["num_hidden_layers"] = prelude_layers + recurrent_layers + coda_layers
+        kwargs.setdefault("use_cache", False)
+        super().__init__(*args, rope_theta=rope_theta, attention_bias=attention_bias, **kwargs)
+        self.xsa_projection = xsa_projection
+        self.rope_theta = rope_theta
+        self.attention_bias = attention_bias
+        self.prelude_layers = prelude_layers
+        self.recurrent_layers = recurrent_layers
+        self.coda_layers = coda_layers
+        self.gradient_checkpointing = gradient_checkpointing
+        self.use_flash_attn = use_flash_attn
+        if not hasattr(self, 'rope_parameters') or self.rope_parameters is None:
+            self.rope_parameters = {"rope_type": "default", "factor": 1.0, "rope_theta": rope_theta}

checkpoint-12000/generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 2,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 1,
+  "transformers_version": "5.15.0",
+  "use_cache": false
+}

checkpoint-12000/model.safetensors

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

checkpoint-12000/modeling_surjo_exp.py

@@ -0,0 +1,417 @@
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from transformers import LlamaConfig, LlamaModel, LlamaForCausalLM
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+from transformers.models.llama.modeling_llama import LlamaMLP
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding, apply_rotary_pos_emb
+from transformers.cache_utils import DynamicCache
+
+try:
+    from .configuration_surjo_exp import SurjoExpConfig
+except ImportError:
+    from configuration_surjo_exp import SurjoExpConfig
+
+try:
+    from flash_attn import flash_attn_varlen_func
+    FLASH_ATTN_AVAILABLE = True
+except ImportError:
+    FLASH_ATTN_AVAILABLE = False
+
+@torch._dynamo.disable()
+def _flash_varlen(q, k, v, cu_seqlens, max_seqlen, dropout_p):
+    ms = int(max_seqlen.item()) if torch.is_tensor(max_seqlen) else int(max_seqlen)
+    return flash_attn_varlen_func(
+        q, k, v, cu_seqlens, cu_seqlens, ms, ms,
+        dropout_p=dropout_p, causal=True,
+    )
+
+class ClampedLlamaMLP(LlamaMLP):
+    def forward(self, x):
+        gate = F.silu(self.gate_proj(x).clamp(-15.0, 15.0))
+        up = self.up_proj(x)
+        return self.down_proj(gate * up)
+
+class XSAAttention(nn.Module):
+    def __init__(self, config, layer_idx=None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.recurrent_cache_idx = None
+        self._use_recurrent_slot = False
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
+        self.attention_bias = getattr(config, "attention_bias", False)
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=self.attention_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.attention_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.attention_bias)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=self.attention_bias)
+
+        self.q_norm = LlamaRMSNorm(self.head_dim, eps=1e-6)
+        self.k_norm = LlamaRMSNorm(self.head_dim, eps=1e-6)
+
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None,
+                output_attentions=False, use_cache=False, cache_position=None, position_embeddings=None,
+                expected_batch_size=None, cu_seqlens=None, max_seqlen=None, **kwargs):
+        past_kv = past_key_value if past_key_value is not None else kwargs.get("past_key_values", None)
+
+        if hidden_states.ndim == 2:
+            if expected_batch_size is None:
+                raise RuntimeError(
+                    f"XSAAttention received 2D hidden_states {hidden_states.shape} "
+                    f"without an expected_batch_size to safely restore the batch dim."
+                )
+            hidden_states = hidden_states.reshape(expected_batch_size, -1, self.hidden_size)
+
+        bsz, q_len, _ = hidden_states.size()
+
+        if expected_batch_size is not None and bsz != expected_batch_size:
+            raise RuntimeError(
+                f"XSAAttention: hidden_states batch size {bsz} does not match "
+                f"expected_batch_size {expected_batch_size}."
+            )
+
+        query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim)
+        key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+        value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+
+        query_states = self.q_norm(query_states)
+        key_states = self.k_norm(key_states)
+
+        cos, sin = position_embeddings
+
+        use_flash = (
+            cu_seqlens is not None
+            and past_kv is None
+            and getattr(self.config, "use_flash_attn", False)
+            and FLASH_ATTN_AVAILABLE
+        )
+
+        if use_flash:
+            total = bsz * q_len
+            q = query_states.reshape(total, self.num_heads, self.head_dim)
+            k = key_states.reshape(total, self.num_key_value_heads, self.head_dim)
+            v = value_states.reshape(total, self.num_key_value_heads, self.head_dim)
+
+            q = q.to(torch.bfloat16)
+            k = k.to(torch.bfloat16)
+            v = v.to(torch.bfloat16)
+
+            cos_f = cos.reshape(-1, cos.shape[-1]).to(torch.bfloat16)
+            sin_f = sin.reshape(-1, sin.shape[-1]).to(torch.bfloat16)
+            q, k = apply_rotary_pos_emb(q, k, cos_f, sin_f, unsqueeze_dim=1)
+
+            attn_output = _flash_varlen(
+                q, k, v, cu_seqlens, max_seqlen,
+                self.config.attention_dropout if self.training else 0.0,
+            )
+
+            if getattr(self.config, 'xsa_projection', True):
+                y = attn_output.view(total, self.num_key_value_heads, self.num_key_value_groups, self.head_dim)
+                v_grouped = v.unsqueeze(2)
+                dot_yv = (y * v_grouped).sum(dim=-1, keepdim=True).float()
+                dot_vv = v_grouped.pow(2).sum(dim=-1, keepdim=True).clamp_min(1e-4).float()
+                scale = (dot_yv / dot_vv).to(y.dtype)
+                attn_output = (y - scale * v_grouped).reshape(total, self.num_heads, self.head_dim)
+
+            attn_output = self.o_proj(attn_output.reshape(bsz, q_len, self.hidden_size))
+            return (attn_output, None)
+
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        current_v = value_states
+
+        target_idx = self.layer_idx
+        if self._use_recurrent_slot and self.recurrent_cache_idx is not None:
+            target_idx = self.recurrent_cache_idx
+
+        if past_kv is not None:
+            while len(past_kv) <= target_idx:
+                past_kv.update(
+                    torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=key_states.dtype, device=key_states.device),
+                    torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=value_states.dtype, device=value_states.device),
+                    len(past_kv)
+                )
+            key_states, value_states = past_kv.update(key_states, value_states, target_idx)
+
+        key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
+        value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)
+
+        kv_len = key_states.shape[-2]
+
+        if attention_mask is not None:
+            if attention_mask.ndim == 2:
+                if attention_mask.shape[-1] < kv_len:
+                    attention_mask = F.pad(attention_mask, (0, kv_len - attention_mask.shape[-1]), value=1)
+                elif attention_mask.shape[-1] > kv_len:
+                    attention_mask = attention_mask[:, -kv_len:]
+
+                pad_mask = (1.0 - attention_mask[:, None, None, :].to(query_states.dtype)) * torch.finfo(query_states.dtype).min
+
+                if q_len > 1:
+                    if cache_position is None:
+                        cache_position = torch.arange(kv_len - q_len, kv_len, device=query_states.device)
+                    kv_positions = torch.arange(kv_len, device=query_states.device)
+
+                    neg_inf = torch.finfo(query_states.dtype).min
+                    causal_mask = torch.zeros((q_len, kv_len), dtype=query_states.dtype, device=query_states.device)
+                    causal_mask = causal_mask.masked_fill(kv_positions[None, :] > cache_position[:, None], neg_inf)
+                    attn_mask = causal_mask[None, None, :, :] + pad_mask
+
+                    diag_idx = torch.arange(q_len, device=attn_mask.device)
+                    start_idx = attn_mask.shape[-1] - q_len
+                    attn_mask[:, :, diag_idx, start_idx + diag_idx] = 0.0
+                else:
+                    attn_mask = pad_mask
+            else:
+                if attention_mask.shape[0] != bsz:
+                    raise RuntimeError(
+                        f"attention_mask batch size {attention_mask.shape[0]} does not "
+                        f"match hidden_states batch size {bsz}."
+                    )
+                attn_mask = attention_mask.to(dtype=query_states.dtype)
+            is_causal = False
+        else:
+            is_causal = True
+            attn_mask = None
+
+        attn_output = F.scaled_dot_product_attention(
+            query_states, key_states, value_states, attn_mask=attn_mask,
+            dropout_p=0.0 if not self.training else self.config.attention_dropout, is_causal=is_causal
+        )
+
+        if getattr(self.config, 'xsa_projection', True):
+            y = attn_output.reshape(bsz, self.num_key_value_heads, self.num_key_value_groups, q_len, self.head_dim)
+            v_grouped = current_v.unsqueeze(2)
+            dot_yv = (y * v_grouped).sum(dim=-1, keepdim=True).float()
+            dot_vv = v_grouped.pow(2).sum(dim=-1, keepdim=True).clamp_min(1e-4).float()
+            scale = (dot_yv / dot_vv).to(y.dtype)
+            attn_output = (y - scale * v_grouped).reshape(bsz, self.num_heads, q_len, self.head_dim)
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+        attn_output = self.o_proj(attn_output)
+
+        return (attn_output, None)
+
+@torch._dynamo.disable()
+def _checkpointed_layer_forward(layer, hidden_states, attention_mask, position_ids,
+                                 cache_position, cos, sin, expected_batch_size, cu_seqlens, max_seqlen):
+    out = layer(
+        hidden_states, attention_mask=attention_mask, position_ids=position_ids,
+        past_key_value=None, use_cache=False,
+        cache_position=cache_position, position_embeddings=(cos, sin),
+        expected_batch_size=expected_batch_size,
+        cu_seqlens=cu_seqlens, max_seqlen=max_seqlen,
+    )
+    hs_out = out[0] if isinstance(out, tuple) else out
+    if hs_out.ndim != 3 or hs_out.shape[0] != expected_batch_size:
+        raise RuntimeError(
+            f"Layer output shape {tuple(hs_out.shape)} does not match expected "
+            f"batch size {expected_batch_size}."
+        )
+    return hs_out
+
+class SurjoExpModel(LlamaModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        assert config.prelude_layers + config.recurrent_layers + config.coda_layers == config.num_hidden_layers, \
+            "prelude_layers + recurrent_layers + coda_layers must equal num_hidden_layers"
+
+        if getattr(config, "use_flash_attn", False) and not FLASH_ATTN_AVAILABLE:
+            raise ImportError(
+                "config.use_flash_attn=True but flash_attn is not importable. "
+                "Install the FA2 wheel or set use_flash_attn=False."
+            )
+
+        p1 = config.prelude_layers
+        r1 = p1 + config.recurrent_layers
+
+        for i, layer in enumerate(self.layers):
+            layer.self_attn = XSAAttention(config, layer_idx=i)
+            layer.mlp = ClampedLlamaMLP(config)
+
+        for i, layer in enumerate(self.layers[p1:r1]):
+            layer.self_attn.recurrent_cache_idx = config.num_hidden_layers + p1 + i
+
+        self.gradient_checkpointing = getattr(config, "gradient_checkpointing", True)
+
+    def gradient_checkpointing_enable(self):
+        self.gradient_checkpointing = True
+
+    def gradient_checkpointing_disable(self):
+        self.gradient_checkpointing = False
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None,
+                past_key_values=None, use_cache=None, output_attentions=False, output_hidden_states=False,
+                cache_position=None, return_dict=True, cu_seqlens=None, max_seqlen=None, **kwargs):
+        if use_cache is None:
+            use_cache = False
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        bsz, seq_len = inputs_embeds.shape[0], inputs_embeds.shape[1]
+
+        if cache_position is None:
+            past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(past_seen, past_seen + seq_len, dtype=torch.long, device=inputs_embeds.device)
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0).expand(bsz, -1)
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        cos, sin = position_embeddings
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        p1 = self.config.prelude_layers
+        r1 = p1 + self.config.recurrent_layers
+        c1 = r1 + self.config.coda_layers
+
+        prelude = self.layers[:p1]
+        recurrent = self.layers[p1:r1]
+        coda = self.layers[r1:c1]
+
+        use_ckpt = self.training and self.gradient_checkpointing and not use_cache
+
+        def run_layer(layer, hs):
+            if cu_seqlens is not None:
+                torch._dynamo.mark_dynamic(cu_seqlens, 0)
+
+            out = layer(
+                hs, attention_mask=attention_mask, position_ids=position_ids,
+                past_key_value=past_key_values if use_cache else None, use_cache=use_cache,
+                cache_position=cache_position, position_embeddings=position_embeddings,
+                expected_batch_size=bsz, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen,
+            )
+            hs_out = out[0] if isinstance(out, tuple) else out
+            if hs_out.ndim != 3 or hs_out.shape[0] != bsz:
+                raise RuntimeError(
+                    f"Layer output shape {tuple(hs_out.shape)} does not match expected "
+                    f"batch size {bsz}."
+                )
+            return hs_out
+
+        def run_layer_maybe_ckpt(layer, hs):
+            if use_ckpt:
+                return torch.utils.checkpoint.checkpoint(
+                    _checkpointed_layer_forward,
+                    layer, hs, attention_mask, position_ids, cache_position, cos, sin, bsz,
+                    cu_seqlens, max_seqlen,
+                    use_reentrant=False,
+                )
+            return run_layer(layer, hs)
+
+        for layer in prelude:
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        if self.training:
+            hidden_states = hidden_states + torch.randn_like(hidden_states) * 0.02
+
+        for layer in recurrent:
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        if self.training:
+            hidden_states = hidden_states + torch.randn_like(hidden_states) * 0.02
+
+        for layer in recurrent:
+            layer.self_attn._use_recurrent_slot = True
+            try:
+                hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+            finally:
+                layer.self_attn._use_recurrent_slot = False
+
+        for layer in coda:
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=past_key_values)
+
+class SurjoExpForCausalLM(LlamaForCausalLM):
+    config_class = SurjoExpConfig
+    def __init__(self, config):
+        super(LlamaForCausalLM, self).__init__(config)
+        self.model = SurjoExpModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def gradient_checkpointing_enable(self, **kwargs):
+        self.model.gradient_checkpointing_enable()
+
+    def gradient_checkpointing_disable(self):
+        self.model.gradient_checkpointing_disable()
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, inputs_embeds=None,
+                use_cache=None, num_logits_to_keep=0, position_ids=None, past_key_values=None,
+                cache_position=None, cu_seqlens=None, max_seqlen=None, **kwargs):
+        if use_cache is None:
+            use_cache = False if (self.training or labels is not None) else True
+
+        if num_logits_to_keep == 0 and "logits_to_keep" in kwargs:
+            num_logits_to_keep = kwargs["logits_to_keep"]
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
+        hidden_states = outputs[0]
+
+        expected_bsz = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]
+        if hidden_states.ndim != 3 or hidden_states.shape[0] != expected_bsz:
+            raise RuntimeError(
+                f"SurjoExpModel returned hidden_states with shape {tuple(hidden_states.shape)}, "
+                f"expected batch size {expected_bsz}."
+            )
+
+        loss = None
+        logits = None
+
+        if labels is not None:
+            shift_hidden = hidden_states[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            num_chunks = 8
+            h_chunks = shift_hidden.chunk(num_chunks, dim=0)
+            l_chunks = shift_labels.chunk(num_chunks, dim=0)
+
+            total_loss = hidden_states.new_zeros((), dtype=torch.float32)
+            total_tokens = 0
+            for h_c, l_c in zip(h_chunks, l_chunks):
+                logits_c = self.lm_head(h_c)
+                chunk_loss = F.cross_entropy(
+                    logits_c.view(-1, logits_c.size(-1)).float(),
+                    l_c.view(-1),
+                    reduction="sum",
+                )
+                total_loss = total_loss + chunk_loss
+                total_tokens += l_c.numel()
+            loss = (total_loss / total_tokens).to(hidden_states.dtype)
+        else:
+            slice_hidden = hidden_states if num_logits_to_keep == 0 else hidden_states[:, -num_logits_to_keep:, :]
+            logits = self.lm_head(slice_hidden)
+
+        return CausalLMOutputWithPast(
+            loss=loss, logits=logits, past_key_values=outputs.past_key_values
+        )

checkpoint-12000/optimizer.bin

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

checkpoint-12000/optimizer_1.bin

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

checkpoint-12000/random_states_0.pkl

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

checkpoint-12000/tokenizer_config.json

@@ -0,0 +1,85 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|bos|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|eos|>",
+  "extra_special_tokens": [
+    "<|unk|>",
+    "<|pad|>",
+    "<|bos|>",
+    "<|eos|>",
+    "<|mask|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|system|>",
+    "<|user|>",
+    "<|assistant|>",
+    "<think>",
+    "</think>",
+    "<|begin_of_thought|>",
+    "<|end_of_thought|>",
+    "<answer>",
+    "</answer>",
+    "<|step|>",
+    "<|/step|>",
+    "<context>",
+    "</context>",
+    "<|doc_start|>",
+    "<|doc_end|>",
+    "<|search|>",
+    "<|search_results|>",
+    "<|tool_list_start|>",
+    "<|tool_list_end|>",
+    "<tools>",
+    "</tools>",
+    "<|tool_call_start|>",
+    "<|tool_call_end|>",
+    "<|tool_call|>",
+    "<|/tool_call|>",
+    "<|tool_response_start|>",
+    "<|tool_response_end|>",
+    "<|tool_response|>",
+    "<|/tool_response|>",
+    "<|image|>",
+    "<|image_pad|>",
+    "<|image_placeholder|>",
+    "<|audio|>",
+    "<|audio_pad|>",
+    "<|audio_placeholder|>",
+    "<|video|>",
+    "<|video_pad|>",
+    "<|fim_prefix|>",
+    "<|fim_suffix|>",
+    "<|fim_middle|>",
+    "<|repo_name|>",
+    "<|file_separator|>",
+    "<|reward|>",
+    "<|reserved_0|>",
+    "<|reserved_1|>",
+    "<|reserved_2|>",
+    "<|reserved_3|>",
+    "<|reserved_4|>",
+    "<|reserved_5|>",
+    "<|reserved_6|>",
+    "<|reserved_7|>",
+    "<|reserved_8|>",
+    "<|reserved_9|>",
+    "<|reserved_10|>",
+    "<|reserved_11|>",
+    "<|reserved_12|>",
+    "<|reserved_13|>",
+    "<|reserved_14|>",
+    "<|reserved_15|>",
+    "<|reserved_16|>",
+    "<|reserved_17|>",
+    "<|reserved_18|>",
+    "<|reserved_19|>"
+  ],
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "<|mask|>",
+  "model_max_length": 10000000,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|unk|>"
+}

checkpoint-13000/chat_template.jinja

@@ -0,0 +1,4 @@
+{{ bos_token }}{% for message in messages %}{{ '<|im_start|>' + message['role'] + '
+' + message['content'] + '<|im_end|>
+' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
+' }}{% endif %}

checkpoint-13000/config.json

@@ -0,0 +1,43 @@
+{
+  "architectures": [
+    "SurjoExpForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_surjo_exp.SurjoExpConfig",
+    "AutoModelForCausalLM": "modeling_surjo_exp.SurjoExpForCausalLM"
+  },
+  "bos_token_id": 0,
+  "coda_layers": 16,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "gradient_checkpointing": false,
+  "head_dim": 72,
+  "hidden_act": "silu",
+  "hidden_size": 576,
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "max_position_embeddings": 2048,
+  "mlp_bias": false,
+  "model_type": "surjo_exp",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 2,
+  "pad_token_id": 1,
+  "prelude_layers": 4,
+  "pretraining_tp": 1,
+  "recurrent_layers": 12,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.15.0",
+  "use_cache": false,
+  "use_flash_attn": false,
+  "vocab_size": 32768,
+  "xsa_projection": true
+}

checkpoint-13000/configuration_surjo_exp.py

@@ -0,0 +1,21 @@
+
+from transformers import LlamaConfig
+
+class SurjoExpConfig(LlamaConfig):
+    model_type = "surjo_exp"
+    def __init__(self, *args, xsa_projection=True, rope_theta=10000.0, attention_bias=False,
+                 prelude_layers=4, recurrent_layers=12, coda_layers=16,
+                 gradient_checkpointing=True, use_flash_attn=True, **kwargs):
+        kwargs["num_hidden_layers"] = prelude_layers + recurrent_layers + coda_layers
+        kwargs.setdefault("use_cache", False)
+        super().__init__(*args, rope_theta=rope_theta, attention_bias=attention_bias, **kwargs)
+        self.xsa_projection = xsa_projection
+        self.rope_theta = rope_theta
+        self.attention_bias = attention_bias
+        self.prelude_layers = prelude_layers
+        self.recurrent_layers = recurrent_layers
+        self.coda_layers = coda_layers
+        self.gradient_checkpointing = gradient_checkpointing
+        self.use_flash_attn = use_flash_attn
+        if not hasattr(self, 'rope_parameters') or self.rope_parameters is None:
+            self.rope_parameters = {"rope_type": "default", "factor": 1.0, "rope_theta": rope_theta}