Add 86M llama-style checkpoint (epoch 1, step 71k), config, model code, and model card

README.md

@@ -1,3 +1,182 @@
 ---
 license: mit
+language:
+- en
+tags:
+- text-generation
+- pytorch
+- from-scratch
+- decoder-only
+- llama-style
+pipeline_tag: text-generation
 ---
+
+# pre-train-llama
+
+A **Llama-style decoder-only transformer** (RoPE, grouped-query attention,
+SwiGLU MLP, RMSNorm) trained **from scratch** on English public-domain books.
+
+This is a prototype / architecture testbed built ahead of the khmer-asr
+project's own model work. It is trained on English classic literature — **not
+on Khmer speech or text** — and is not usable for ASR as-is.
+
+## Architecture
+
+| Hyperparameter | Value |
+|---|---|
+| Layers | 8 |
+| Attention heads | 8 |
+| KV heads (GQA) | 4 |
+| Head dim | 96 |
+| Hidden dim | 768 |
+| FFN dim (SwiGLU) | 3072 |
+| Max sequence length | 512 |
+| Vocab size | 10,000 |
+| Dropout | 0.1 |
+| Parameters | 86.2M |
+| Weights dtype | float32 |
+
+Positional encoding is RoPE (theta 10,000), attention is grouped-query
+attention with repeat-interleaved KV heads, the MLP is SwiGLU, and
+normalization is RMSNorm applied **pre-block** (before attention and before the
+MLP), with a final RMSNorm before the output projection. Input embeddings and
+the output head are **not** tied. Architecturally this mirrors Llama; it is
+trained from scratch, not initialized from Meta's weights.
+
+## Training data
+
+Tokenizer and model were trained on eleven English-language books from Project
+Gutenberg:
+
+- Moby Dick
+- Frankenstein
+- Dracula
+- Little Women
+- Pride and Prejudice
+- Alice's Adventures in Wonderland
+- Crime and Punishment
+- The Adventures of Tom Sawyer
+- A Tale of Two Cities
+- The Adventures of Sherlock Holmes
+- War and Peace
+
+Gutenberg header/footer boilerplate is stripped, blank lines removed, and the
+books concatenated into a single stream. A byte-level BPE tokenizer
+(`vocab_size=10000`, specials `[pad]`, `[eos]`) was trained on that corpus,
+producing ~2.72M tokens. Training examples are **stride-1 sliding windows** of
+512 tokens, so consecutive examples overlap by 511 tokens.
+
+## Training procedure
+
+| Setting | Value |
+|---|---|
+| Objective | Next-token prediction, cross-entropy, `[pad]` ignored |
+| Optimizer | AdamW, peak LR 5e-4 |
+| Schedule | Linear warmup 2,000 steps (0.01 → 1.0), then cosine decay to 0 |
+| Gradient clipping | Global norm 6.0 |
+| Batch | 8 × 4 gradient-accumulation steps = effective 32 |
+| Precision | fp32 (bf16 matmuls internally on TPU) |
+| Hardware | TPU via `torch_xla`, single core |
+
+### State of this checkpoint
+
+Training is **incomplete** — this is a mid-run checkpoint, not a finished model.
+
+| | |
+|---|---|
+| Checkpoint saved | 2026-07-24 03:21:00 |
+| Epoch | 1 of 2 (in progress) |
+| Micro-batch | 284,000 of ~340,500 |
+| Optimizer steps | 71,000 |
+| Learning rate at save | 3.20e-4 |
+| Last training loss | 0.0513 |
+| Best epoch loss | not yet recorded (no epoch has completed) |
+
+Roughly 1.16B tokens have been processed, but only ~2.72M of them are distinct
+— every token is seen ~512 times across overlapping windows within a single
+epoch.
+
+## Usage
+
+This is not a `transformers` model class. Load `model.safetensors` into the
+`TextGenerationModel` defined in `modeling_llama_custom.py`:
+
+```python
+import json
+
+import torch
+import torch.nn.functional as F
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+from tokenizers import Tokenizer
+
+path = snapshot_download("Panhapich/pre-train-llama")
+
+import sys; sys.path.insert(0, path)
+from modeling_llama_custom import TextGenerationModel, create_causal_mask
+
+config = json.load(open(f"{path}/config.json"))
+model = TextGenerationModel(**config["model_config"])
+model.load_state_dict(load_file(f"{path}/model.safetensors"))
+model.eval()
+
+tokenizer = Tokenizer.from_file(f"{path}/tokenizer.json")
+
+
+@torch.no_grad()
+def generate(prompt, max_new_tokens=40, temperature=0.8):
+    ids = torch.tensor(tokenizer.encode(prompt).ids).unsqueeze(0)
+    for _ in range(max_new_tokens):
+        logits = model(ids)[:, -1, :] / temperature
+        next_id = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
+        ids = torch.cat([ids, next_id], dim=1)
+        if next_id.item() == tokenizer.token_to_id("[eos]"):
+            break
+    return tokenizer.decode(ids[0].tolist())
+
+
+print(generate("Once upon a time,"))
+```
+
+Sequences longer than 512 tokens are not supported — the RoPE tables are
+precomputed to `max_seq_len` and indexing past them will fail. There is no KV
+cache, so generation recomputes the full context each step.
+
+`tokenizer.json` must be the tokenizer these weights were trained with. A
+freshly retrained BPE would assign different ids to the same text and the model
+would emit nonsense without erroring.
+
+## Files
+
+| File | What it is |
+|---|---|
+| `model.safetensors` | Model weights (86.2M params, fp32, ~345 MB) |
+| `config.json` | `model_config` hyperparameters for reconstruction |
+| `modeling_llama_custom.py` | `nn.Module` definitions the weights load into |
+| `tokenizer.json` | The BPE tokenizer the weights were trained against |
+
+## Limitations and biases
+
+- **The 0.0513 training loss is not a generalization result.** There is no
+  held-out validation split, and stride-1 windows mean the model sees each
+  passage hundreds of times per epoch. A loss that low on a 10k vocab indicates
+  the corpus has largely been memorized. Expect the model to reproduce long
+  verbatim spans of the source books, and expect much worse performance on any
+  text outside them.
+- **No evaluation has been run** — no perplexity on held-out data, no
+  benchmarks. The only quality check performed is qualitative sampling.
+- **Training is unfinished** (mid-epoch 1 of 2), so the cosine schedule has not
+  annealed and weights are not at a converged point.
+- Trained on 19th-century literature, so output reflects the vocabulary,
+  style, and social attitudes of that corpus, including period-typical racist
+  and sexist content present in the source texts.
+- **English only.** Despite the surrounding khmer-asr project, this model has
+  no Khmer training data and no speech or audio capability.
+- Small (86M) and trained on ~2.7M unique tokens — orders of magnitude below
+  what general-purpose language models see. Treat output as a demonstration
+  that the architecture and training loop work, not as a useful generator.
+
+## Training code
+
+Trained with `decoder_only_transformer_tpu.ipynb` from the khmer-asr project,
+which runs on TPU, CUDA, MPS, or CPU and resumes from checkpoints in this repo.

config.json

@@ -0,0 +1,13 @@
+{
+  "architecture": "llama_style_decoder_only_custom",
+  "model_config": {
+    "num_layers": 8,
+    "num_heads": 8,
+    "num_kv_heads": 4,
+    "hidden_dim": 768,
+    "max_seq_len": 512,
+    "vocab_size": 10000,
+    "dropout": 0.1
+  },
+  "num_parameters": 86235664
+}

model.safetensors

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

modeling_llama_custom.py

@@ -0,0 +1,139 @@
+# Model definition for Panhapich/pre-train-llama.
+#
+# Llama-style decoder-only transformer (RoPE, grouped-query attention, SwiGLU,
+# RMSNorm), trained from scratch. This is not a `transformers`-library model
+# class -- load model.safetensors into TextGenerationModel directly:
+#
+#     import json
+#     from safetensors.torch import load_file
+#     from modeling_llama_custom import TextGenerationModel
+#
+#     config = json.load(open("config.json"))
+#     model = TextGenerationModel(**config["model_config"])
+#     model.load_state_dict(load_file("model.safetensors"))
+#     model.eval()
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+if not hasattr(nn, "RMSNorm"):
+    class _RMSNormFallback(nn.Module):
+        def __init__(self, dim, eps=1e-6):
+            super().__init__()
+            self.eps = eps
+            self.weight = nn.Parameter(torch.ones(dim))
+
+        def forward(self, x):
+            rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
+            return x * rms * self.weight
+
+    nn.RMSNorm = _RMSNormFallback
+
+
+class RotaryPositionalEncoding(nn.Module):
+    def __init__(self, head_dim, max_seq_len, theta=10000.0):
+        super().__init__()
+        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
+        t = torch.arange(max_seq_len).float()
+        freqs = torch.outer(t, inv_freq)
+        self.register_buffer("cos", torch.cos(freqs), persistent=False)
+        self.register_buffer("sin", torch.sin(freqs), persistent=False)
+
+    def rotate(self, x):
+        T = x.shape[-2]
+        cos = self.cos[:T].unsqueeze(0).unsqueeze(0)
+        sin = self.sin[:T].unsqueeze(0).unsqueeze(0)
+        x1, x2 = x[..., 0::2], x[..., 1::2]
+        rotated = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+        return rotated.flatten(-2)
+
+
+class GQA(nn.Module):
+    def __init__(self, hidden_dim, num_heads, num_kv_heads, dropout=0.1):
+        super().__init__()
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.n_rep = num_heads // num_kv_heads
+        self.head_dim = hidden_dim // num_heads
+
+        self.q_proj = nn.Linear(hidden_dim, num_heads * self.head_dim)
+        self.k_proj = nn.Linear(hidden_dim, num_kv_heads * self.head_dim)
+        self.v_proj = nn.Linear(hidden_dim, num_kv_heads * self.head_dim)
+        self.out_proj = nn.Linear(num_heads * self.head_dim, hidden_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, q, k, v, mask=None, rope=None):
+        B, T, _ = q.shape
+        q = self.q_proj(q).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(k).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(v).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        if rope is not None:
+            q = rope.rotate(q)
+            k = rope.rotate(k)
+        if self.n_rep > 1:
+            k = k.repeat_interleave(self.n_rep, dim=1)
+            v = v.repeat_interleave(self.n_rep, dim=1)
+
+        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        if mask is not None:
+            scores = scores.masked_fill(~mask.unsqueeze(1).bool(), float('-inf'))
+        attn = F.softmax(scores, dim=-1)
+        attn = self.dropout(attn)
+        out = attn @ v
+        out = out.transpose(1, 2).reshape(B, T, -1)
+        return self.out_proj(out)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, hidden_dim, ff_dim):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_dim, ff_dim)
+        self.up_proj = nn.Linear(hidden_dim, ff_dim)
+        self.down_proj = nn.Linear(ff_dim, hidden_dim)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, hidden_dim, num_heads, num_kv_heads, dropout=0.1):
+        super().__init__()
+        self.self_attn = GQA(hidden_dim, num_heads, num_kv_heads, dropout)
+        self.mlp = SwiGLU(hidden_dim, 4 * hidden_dim)
+        self.norm1 = nn.RMSNorm(hidden_dim)
+        self.norm2 = nn.RMSNorm(hidden_dim)
+
+    def forward(self, x, mask=None, rope=None):
+        out = self.norm1(x)
+        out = self.self_attn(out, out, out, mask, rope)
+        x = out + x
+        out = self.norm2(x)
+        out = self.mlp(out)
+        return out + x
+
+class TextGenerationModel(nn.Module):
+    def __init__(self, num_layers, num_heads, num_kv_heads, hidden_dim,
+                 max_seq_len, vocab_size, dropout=0.1):
+        super().__init__()
+        self.rope = RotaryPositionalEncoding(hidden_dim // num_heads, max_seq_len)
+        self.embedding = nn.Embedding(vocab_size, hidden_dim)
+        self.decoders = nn.ModuleList([
+            DecoderLayer(hidden_dim, num_heads, num_kv_heads, dropout)
+            for _ in range(num_layers)
+        ])
+        self.norm = nn.RMSNorm(hidden_dim)
+        self.out = nn.Linear(hidden_dim, vocab_size)
+
+    def forward(self, ids, mask=None):
+        x = self.embedding(ids)
+        for decoder in self.decoders:
+            x = decoder(x, mask, self.rope)
+        x = self.norm(x)
+        return self.out(x)
+
+def create_causal_mask(seq_len, device):
+    return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device))