step 10 | ppl=inf

config.json

@@ -0,0 +1,30 @@
+{
+  "vocab_size": 4096,
+  "hidden_size": 256,
+  "tie_embeddings": true,
+  "num_layers": 11,
+  "intermediate_size": 768,
+  "num_heads": 4,
+  "num_kv_heads": 2,
+  "max_seq_len": 1024,
+  "rope_theta": 10000.0,
+  "qk_norm": true,
+  "use_bias": false,
+  "ffn_type": "swiglu",
+  "activation": "silu",
+  "xsa": false,
+  "xsa_window": 128,
+  "dropout": 0.0,
+  "norm_eps": 1e-05,
+  "model_type": "slm",
+  "tie_word_embeddings": true,
+  "auto_map": {
+    "AutoConfig": "model.SLMHFConfig",
+    "AutoModelForCausalLM": "model.SLMHFForCausalLM"
+  },
+  "bos_token_id": 2,
+  "eos_token_id": 3,
+  "pad_token_id": 3,
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.0"
+}

generation_config.json

@@ -0,0 +1,5 @@
+{
+  "bos_token_id": 2,
+  "eos_token_id": 3,
+  "max_length": 1024
+}

model.py

@@ -0,0 +1,575 @@
+"""
+model.py  —  SLM Custom Architecture (v6.5 — fully HF-compatible)
+=================================================================
+Auto-written by the training notebook. Required for HF auto_map.
+
+v6.5 fixes:
+  1. _tied_weights_keys is a LIST — works on transformers 4.x AND 5.x.
+  2. RoPE buffers: persistent=False — NOT saved to safetensors.
+  3. tie_weights() properly delegates so HF re-ties weights on load.
+  4. Full HF forward() signature (return_dict, attention_mask, etc.).
+  5. CausalLMOutputWithPast used — generation helpers (.generate()) work.
+  6. prepare_inputs_for_generation() added.
+  7. Weights saved WITH model. prefix (matches SLMHFForCausalLM layout).
+"""
+from __future__ import annotations
+import math
+from dataclasses import dataclass, asdict
+from typing import Optional, Tuple, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from transformers import PretrainedConfig, PreTrainedModel
+    from transformers.modeling_outputs import CausalLMOutputWithPast
+    _HF_AVAILABLE = True
+except ImportError:
+    _HF_AVAILABLE = False
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# CONFIG DATACLASS
+# ══════════════════════════════════════════════════════════════════════════════
+
+@dataclass
+class SLMConfig:
+    # Vocab / embedding
+    vocab_size:        int   = 32000
+    hidden_size:       int   = 576
+    tie_embeddings:    bool  = True
+
+    # Depth / width
+    num_layers:        int   = 12
+    intermediate_size: int   = 1536
+
+    # Attention
+    num_heads:         int   = 9
+    num_kv_heads:      int   = 3
+    max_seq_len:       int   = 2048
+    rope_theta:        float = 500000.0
+    qk_norm:           bool  = True
+    use_bias:          bool  = False
+
+    # FFN type: 'swiglu' | 'geglu' | 'mlp'
+    ffn_type:          str   = 'swiglu'
+    activation:        str   = 'silu'   # used by 'mlp' only
+
+    # XSA
+    xsa:               bool  = False
+    xsa_window:        int   = 128
+
+    # Regularisation
+    dropout:           float = 0.0
+    norm_eps:          float = 1e-5
+
+    @property
+    def head_dim(self) -> int:
+        return self.hidden_size // self.num_heads
+
+    def validate(self) -> 'SLMConfig':
+        assert self.hidden_size % self.num_heads == 0, \
+            f"hidden_size {self.hidden_size} % num_heads {self.num_heads} != 0"
+        assert self.num_heads % self.num_kv_heads == 0, \
+            f"num_heads {self.num_heads} % num_kv_heads {self.num_kv_heads} != 0"
+        assert self.ffn_type in ('swiglu', 'geglu', 'mlp'), \
+            f"ffn_type must be 'swiglu'|'geglu'|'mlp', got {self.ffn_type!r}"
+        assert self.activation in ('silu', 'gelu', 'relu'), \
+            f"activation must be 'silu'|'gelu'|'relu', got {self.activation!r}"
+        if self.xsa:
+            assert self.xsa_window > 0, "xsa_window must be > 0"
+        return self
+
+    def to_dict(self) -> dict:
+        d = asdict(self)
+        d['model_type'] = 'slm'
+        return d
+
+    @classmethod
+    def from_dict(cls, d: dict) -> 'SLMConfig':
+        valid = {f for f in cls.__dataclass_fields__}
+        return cls(**{k: v for k, v in d.items() if k in valid})
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# PRIMITIVES
+# ══════════════════════════════════════════════════════════════════════════════
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.w   = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.w
+
+
+def precompute_rope(head_dim: int, max_len: int,
+                    theta: float = 10000.0) -> Tuple[torch.Tensor, torch.Tensor]:
+    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
+    t     = torch.arange(max_len, dtype=torch.float32)
+    cos   = torch.cos(torch.outer(t, freqs))
+    sin   = torch.sin(torch.outer(t, freqs))
+    return cos, sin
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor,
+               sin: torch.Tensor) -> torch.Tensor:
+    """Correct split-half RoPE rotation. x: [B,H,T,D]"""
+    T  = x.shape[2]
+    c  = cos[:T].unsqueeze(0).unsqueeze(0)
+    s  = sin[:T].unsqueeze(0).unsqueeze(0)
+    x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
+    return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1)
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# FFN  (SwiGLU / GeGLU / MLP)
+# ══════════════════════════════════════════════════════════════════════════════
+
+class FFN(nn.Module):
+    def __init__(self, cfg: SLMConfig):
+        super().__init__()
+        H, I, b = cfg.hidden_size, cfg.intermediate_size, cfg.use_bias
+        self.ffn_type = cfg.ffn_type
+
+        if cfg.ffn_type in ('swiglu', 'geglu'):
+            self.gate = nn.Linear(H, I, bias=b)
+            self.up   = nn.Linear(H, I, bias=b)
+            self.down = nn.Linear(I, H, bias=b)
+            self._act = F.silu if cfg.ffn_type == 'swiglu' else \
+                        lambda x: F.gelu(x, approximate='tanh')
+        else:  # 'mlp'
+            self.up   = nn.Linear(H, I, bias=b)
+            self.down = nn.Linear(I, H, bias=b)
+            act = cfg.activation
+            self._act = (F.silu if act == 'silu' else
+                         (lambda x: F.gelu(x, approximate='tanh')) if act == 'gelu'
+                         else F.relu)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.ffn_type in ('swiglu', 'geglu'):
+            return self.down(self._act(self.gate(x)) * self.up(x))
+        return self.down(self._act(self.up(x)))
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# ATTENTION  (GQA + RoPE + optional QK-Norm)
+# ══════════════════════════════════════════════════════════════════════════════
+
+class GQAttention(nn.Module):
+    """GQA with optional QK-Norm and two backends: sdpa / chunked."""
+
+    def __init__(self, cfg: SLMConfig):
+        super().__init__()
+        H          = cfg.hidden_size
+        b          = cfg.use_bias
+        self.nh    = cfg.num_heads
+        self.nkv   = cfg.num_kv_heads
+        self.d     = cfg.head_dim
+        self.rep   = self.nh // self.nkv
+        KVD        = self.nkv * self.d
+
+        self.q  = nn.Linear(H, H,   bias=b)
+        self.k  = nn.Linear(H, KVD, bias=b)
+        self.v  = nn.Linear(H, KVD, bias=b)
+        self.o  = nn.Linear(H, H,   bias=b)
+        self.dp = nn.Dropout(cfg.dropout)
+
+        self.qk_norm = cfg.qk_norm
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.d, cfg.norm_eps)
+            self.k_norm = RMSNorm(self.d, cfg.norm_eps)
+
+        cos, sin = precompute_rope(self.d, cfg.max_seq_len, cfg.rope_theta)
+        # FIX: persistent=False so RoPE buffers are NOT saved to safetensors.
+        # They are deterministically recomputed from config at load time.
+        self.register_buffer('cos', cos, persistent=False)
+        self.register_buffer('sin', sin, persistent=False)
+
+        self._backend    = 'chunked'
+        self._chunk_size = 256
+
+    def _qk_normalize(self, q, k):
+        B, nh, T, d  = q.shape
+        q = self.q_norm(q.reshape(-1, d)).reshape(B, nh,       T, d)
+        B, nkv, T, d = k.shape
+        k = self.k_norm(k.reshape(-1, d)).reshape(B, self.nkv, T, d)
+        return q, k
+
+    def _sdpa(self, q, k, v):
+        p = self.dp.p if self.training else 0.0
+        return F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=p)
+
+    def _chunked(self, q, k, v):
+        B, H, T, D = q.shape
+        out   = torch.zeros_like(q)
+        scale = D ** -0.5
+        cs    = self._chunk_size
+        mask  = torch.tril(torch.ones(T, T, device=q.device, dtype=torch.bool))
+        for start in range(0, T, cs):
+            end  = min(start + cs, T)
+            q_c  = q[:, :, start:end, :]
+            k_c  = k[:, :, :end,      :]
+            v_c  = v[:, :, :end,      :]
+            attn = torch.matmul(q_c, k_c.transpose(-2, -1)) * scale
+            attn = attn.masked_fill(
+                ~mask[start:end, :end].unsqueeze(0).unsqueeze(0), float('-inf')
+            )
+            attn = F.softmax(attn.float(), dim=-1).to(q.dtype)
+            if self.training and self.dp.p > 0:
+                attn = self.dp(attn)
+            out[:, :, start:end, :] = torch.matmul(attn, v_c)
+        return out
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, _ = x.shape
+        q = self.q(x).view(B, T, self.nh,  self.d).transpose(1, 2)
+        k = self.k(x).view(B, T, self.nkv, self.d).transpose(1, 2)
+        v = self.v(x).view(B, T, self.nkv, self.d).transpose(1, 2)
+        if self.qk_norm:
+            q, k = self._qk_normalize(q, k)
+        q = apply_rope(q, self.cos, self.sin)
+        k = apply_rope(k, self.cos, self.sin)
+        k = k.repeat_interleave(self.rep, dim=1)
+        v = v.repeat_interleave(self.rep, dim=1)
+        o = self._sdpa(q, k, v) if self._backend == 'sdpa' else self._chunked(q, k, v)
+        return self.o(o.transpose(1, 2).contiguous().view(B, T, self.nh * self.d))
+
+    def set_backend(self, backend: str, chunk_size: int = 256):
+        self._backend    = backend
+        self._chunk_size = chunk_size
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# XSA  (sliding-window cross-sequence attention)
+# ══════════════════════════════════════════════════════════════════════════════
+
+class XSALayer(nn.Module):
+    """
+    Optional sliding-window cross-sequence attention.
+    K/V projections are SHARED with the paired GQAttention (parameter sharing).
+    Stored in a plain Python list to avoid duplicate state_dict keys.
+    """
+
+    def __init__(self, cfg: SLMConfig, paired_attn: GQAttention):
+        super().__init__()
+        H = cfg.hidden_size
+        b = cfg.use_bias
+        self.window  = cfg.xsa_window
+        self.nh      = cfg.num_heads
+        self.nkv     = cfg.num_kv_heads
+        self.d       = cfg.head_dim
+        self.rep     = self.nh // self.nkv
+
+        self.norm = RMSNorm(H, cfg.norm_eps)
+        self.q    = nn.Linear(H, H, bias=b)
+        self.o    = nn.Linear(H, H, bias=b)
+        self.dp   = nn.Dropout(cfg.dropout)
+
+        # Plain list — NOT registered as sub-modules → no duplicate state_dict keys
+        self._shared: List[nn.Linear] = [paired_attn.k, paired_attn.v]
+        self._paired: List[GQAttention] = [paired_attn]
+
+        self.qk_norm = cfg.qk_norm
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.d, cfg.norm_eps)
+            self.k_norm = RMSNorm(self.d, cfg.norm_eps)
+
+    @property
+    def _k(self) -> nn.Linear:  return self._shared[0]
+
+    @property
+    def _v(self) -> nn.Linear:  return self._shared[1]
+
+    @property
+    def cos(self) -> torch.Tensor: return self._paired[0].cos
+
+    @property
+    def sin(self) -> torch.Tensor: return self._paired[0].sin
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, H = x.shape
+        W  = self.window
+        xn = self.norm(x)
+
+        q = self.q(xn).view(B, T, self.nh, self.d).transpose(1, 2)
+        if self.qk_norm:
+            B_, nh, T_, d = q.shape
+            q = self.q_norm(q.reshape(-1, d)).reshape(B_, nh, T_, d)
+        q = apply_rope(q, self.cos, self.sin)
+
+        pad      = torch.zeros(B, W, H, device=x.device, dtype=x.dtype)
+        x_padded = torch.cat([pad, xn], dim=1)
+        ctx      = x_padded.unfold(1, W, 1)
+        ctx      = ctx[:, :T, :, :].permute(0, 1, 3, 2).contiguous()
+        ctx_flat = ctx.reshape(B * T, W, H)
+
+        k_flat = self._k(ctx_flat).view(B*T, W, self.nkv, self.d).transpose(1, 2)
+        v_flat = self._v(ctx_flat).view(B*T, W, self.nkv, self.d).transpose(1, 2)
+        k_flat = k_flat.repeat_interleave(self.rep, dim=1)
+        v_flat = v_flat.repeat_interleave(self.rep, dim=1)
+        if self.qk_norm:
+            BT, nh, W_, d = k_flat.shape
+            k_flat = self.k_norm(k_flat.reshape(-1, d)).reshape(BT, nh, W_, d)
+
+        q_flat = q.permute(0, 2, 1, 3).reshape(B * T, self.nh, 1, self.d)
+        scale  = self.d ** -0.5
+        attn   = torch.matmul(q_flat, k_flat.transpose(-2, -1)) * scale
+        attn   = F.softmax(attn.float(), dim=-1).to(q.dtype)
+        if self.training and self.dp.p > 0:
+            attn = self.dp(attn)
+        out = torch.matmul(attn, v_flat).reshape(B, T, self.nh * self.d)
+        return self.o(out)
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# TRANSFORMER BLOCK
+# ══════════════════════════════════════════════════════════════════════════════
+
+class Block(nn.Module):
+    """Pre-LN transformer block. XSA is inserted between attn and FFN."""
+    def __init__(self, cfg: SLMConfig):
+        super().__init__()
+        self.n1      = RMSNorm(cfg.hidden_size, cfg.norm_eps)
+        self.attn    = GQAttention(cfg)
+        self.n2      = RMSNorm(cfg.hidden_size, cfg.norm_eps)
+        self.ffn     = FFN(cfg)
+        self.use_xsa = cfg.xsa
+        if self.use_xsa:
+            self.xsa = XSALayer(cfg, self.attn)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.n1(x))
+        if self.use_xsa:
+            x = x + self.xsa(x)
+        x = x + self.ffn(self.n2(x))
+        return x
+
+    def set_backend(self, backend: str, chunk_size: int = 256):
+        self.attn.set_backend(backend, chunk_size)
+        if self.use_xsa:
+            self.xsa._paired[0]._backend    = backend
+            self.xsa._paired[0]._chunk_size = chunk_size
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# FULL MODEL
+# ══════════════════════════════════════════════════════════════════════════════
+
+class SLMModel(nn.Module):
+    def __init__(self, cfg: SLMConfig):
+        super().__init__()
+        cfg.validate()
+        self.cfg    = cfg
+        self.embed  = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
+        self.drop   = nn.Dropout(cfg.dropout)
+        self.layers = nn.ModuleList([Block(cfg) for _ in range(cfg.num_layers)])
+        self.norm   = RMSNorm(cfg.hidden_size, cfg.norm_eps)
+        self._init_weights()
+
+    def _init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=0.02)
+                if m.bias is not None: nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Embedding):
+                nn.init.normal_(m.weight, std=0.02)
+
+    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+        x = self.drop(self.embed(ids))
+        for layer in self.layers:
+            x = layer(x)
+        return self.norm(x)
+
+    def set_attention_backend(self, backend: str, chunk_size: int = 256):
+        for blk in self.layers:
+            blk.set_backend(backend, chunk_size)
+
+
+class SLMForCausalLM(nn.Module):
+    """
+    SLM with LM head. Used during training.
+
+    State-dict keys (these get 'model.' prefix added by export_safetensors):
+      backbone.embed.weight
+      backbone.layers.{i}.{n1,n2}.w
+      backbone.layers.{i}.attn.{q,k,v,o}.weight
+      backbone.layers.{i}.attn.{q_norm,k_norm}.w  (if qk_norm)
+      backbone.layers.{i}.ffn.{gate,up,down}.weight
+      backbone.layers.{i}.xsa.{norm,q,o}.*         (if xsa)
+      backbone.norm.w
+      lm_head.weight                                (only if NOT tie_embeddings)
+    RoPE cos/sin: persistent=False → NOT in state_dict.
+    """
+    def __init__(self, cfg: SLMConfig):
+        super().__init__()
+        self.cfg      = cfg
+        self.backbone = SLMModel(cfg)
+        self.lm_head  = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
+        if cfg.tie_embeddings:
+            self.lm_head.weight = self.backbone.embed.weight
+        else:
+            nn.init.normal_(self.lm_head.weight, std=0.02)
+
+    def tie_weights(self) -> None:
+        """Re-apply weight tying — called by HF after load_state_dict."""
+        if self.cfg.tie_embeddings:
+            self.lm_head.weight = self.backbone.embed.weight
+
+    def forward(self, input_ids: torch.Tensor,
+                labels: Optional[torch.Tensor] = None
+                ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        hidden = self.backbone(input_ids)
+        logits = self.lm_head(hidden)
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits.reshape(-1, logits.size(-1)),
+                labels.reshape(-1),
+                ignore_index=-100,
+            )
+        return logits, loss
+
+    def set_attention_backend(self, backend: str, chunk_size: int = 256):
+        self.backbone.set_attention_backend(backend, chunk_size)
+
+    def n_params(self) -> int:
+        return sum(p.numel() for p in self.parameters())
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# HUGGINGFACE WRAPPERS
+# ══════════════════════════════════════════════════════════════════════════════
+
+if _HF_AVAILABLE:
+    class SLMHFConfig(PretrainedConfig):
+        model_type = "slm"
+
+        def __init__(self, **kwargs):
+            valid = set(SLMConfig.__dataclass_fields__)
+            slm_kw = {k: kwargs.pop(k, getattr(SLMConfig(), k)) for k in valid}
+            tie_emb = slm_kw.get('tie_embeddings', True)
+            # Accept canonical HF name coming from a saved config.json
+            if 'tie_word_embeddings' in kwargs:
+                tie_emb = kwargs.pop('tie_word_embeddings')
+            super().__init__(tie_word_embeddings=tie_emb, **kwargs)
+            for k, v in slm_kw.items():
+                setattr(self, k, v)
+            self.tie_embeddings = self.tie_word_embeddings
+
+        def to_slm_config(self) -> SLMConfig:
+            return SLMConfig(**{f: getattr(self, f)
+                                for f in SLMConfig.__dataclass_fields__})
+
+    class SLMHFForCausalLM(PreTrainedModel):
+        config_class                    = SLMHFConfig
+        base_model_prefix               = "model"
+        supports_gradient_checkpointing = False
+
+        # ── FIX: LIST format — compatible with transformers 4.x AND 5.x ────
+        # transformers 4.x: iterates as list of keys to re-tie after loading.
+        # transformers 5.x: _get_tied_weight_keys() calls .keys() — BUT it first
+        # checks isinstance(tied, dict); for a list it falls back to iteration.
+        # A list is the safest format across ALL versions.
+        _tied_weights_keys = ["model.lm_head.weight"]
+
+        def __init__(self, config: SLMHFConfig):
+            super().__init__(config)
+            self.model = SLMForCausalLM(config.to_slm_config())
+            # post_init() sets up all_tied_weights_keys from _tied_weights_keys.
+            # Must be called AFTER all submodules are initialised.
+            self.post_init()
+
+        def forward(
+            self,
+            input_ids=None,
+            attention_mask=None,       # accepted, ignored — we use causal mask
+            position_ids=None,
+            past_key_values=None,
+            inputs_embeds=None,
+            labels=None,
+            use_cache=None,
+            output_attentions=None,
+            output_hidden_states=None,
+            return_dict=None,
+            **kwargs,
+        ):
+            logits, loss = self.model(input_ids, labels)
+            # Always return CausalLMOutputWithPast so .generate() / pipelines work
+            return CausalLMOutputWithPast(loss=loss, logits=logits)
+
+        def get_input_embeddings(self):
+            return self.model.backbone.embed
+
+        def set_input_embeddings(self, v):
+            self.model.backbone.embed = v
+            if self.model.cfg.tie_embeddings:
+                self.model.lm_head.weight = v.weight
+
+        def get_output_embeddings(self):
+            return self.model.lm_head
+
+        def set_output_embeddings(self, v):
+            self.model.lm_head = v
+
+        def tie_weights(self):
+            # ── FIX: Delegate to inner model ─────────────────────────────
+            # HF calls tie_weights() on SLMHFForCausalLM after loading weights.
+            # Without delegation, embed->lm_head link is never restored on load.
+            self.model.tie_weights()
+
+        def prepare_inputs_for_generation(self, input_ids, **kwargs):
+            # ── FIX: Required for model.generate() to work ───────────────
+            return {"input_ids": input_ids}
+
+        def resize_token_embeddings(self, new_num_tokens, pad_to_multiple_of=None):
+            result = super().resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
+            if self.model.cfg.tie_embeddings:
+                self.model.lm_head.weight = self.model.backbone.embed.weight
+            return result
+
+    try:
+        from transformers import AutoConfig, AutoModelForCausalLM
+        SLMHFConfig.register_for_auto_class("AutoConfig")
+        SLMHFForCausalLM.register_for_auto_class("AutoModelForCausalLM")
+    except Exception:
+        pass
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# PUBLIC HELPERS
+# ══════════════════════════════════════════════════════════════════════════════
+
+def build_model(cfg_dict: dict, device: torch.device) -> SLMForCausalLM:
+    cfg = SLMConfig.from_dict(cfg_dict)
+    cfg.validate()
+    return SLMForCausalLM(cfg).to(device)
+
+
+def count_params(model: nn.Module) -> dict:
+    total     = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    return {'total': total, 'trainable': trainable}
+
+
+def build_hf_config_json(cfg: SLMConfig, tokenizer) -> dict:
+    bos = getattr(tokenizer, 'bos_token_id', None) or 0
+    eos = getattr(tokenizer, 'eos_token_id', None) or 2
+    pad = getattr(tokenizer, 'pad_token_id',  None) or eos
+    d   = cfg.to_dict()
+    d['tie_word_embeddings'] = cfg.tie_embeddings
+    d.update({
+        'auto_map': {
+            'AutoConfig':           'model.SLMHFConfig',
+            'AutoModelForCausalLM': 'model.SLMHFForCausalLM',
+        },
+        'bos_token_id': bos,
+        'eos_token_id': eos,
+        'pad_token_id': pad,
+        'torch_dtype':  'float32',
+        'transformers_version': '4.40.0',
+    })
+    return d

model.safetensors

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

resume/ckpt.pt

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

resume/latest_step.txt

@@ -0,0 +1 @@
+10

tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|bos|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|eos|>",
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "<|mask|>",
+  "model_max_length": 4096,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|unk|>"
+}

training_meta.json

@@ -0,0 +1,21 @@
+{
+  "step": 10,
+  "val_loss": Infinity,
+  "val_ppl": Infinity,
+  "params_M": 9.707,
+  "tokens_seen": 10475520,
+  "pushed_at": "2026-06-10T17:35:40.399616",
+  "arch": {
+    "ffn_type": "swiglu",
+    "qk_norm": true,
+    "xsa": false,
+    "xsa_window": 128,
+    "num_layers": 11,
+    "hidden_size": 256,
+    "num_heads": 4,
+    "num_kv_heads": 2,
+    "max_seq_len": 1024
+  },
+  "optimizer": "adamw",
+  "lr_scheduler": "wsd"
+}