DGPT v1-base release

.gitattributes

@@ -0,0 +1 @@
+model.npz filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,142 @@
+---
+license: unknown
+language:
+- en
+tags:
+- text-generation
+- from-scratch
+- transformer
+- gpt
+- numpy
+- cupy
+- tinystories
+datasets:
+- roneneldan/TinyStories
+pipeline_tag: text-generation
+---
+
+# DGPT v1-base
+
+**A small, from-scratch base language model. Not an instruction-tuned assistant.**
+
+DGPT v1-base is a 13,049,856-parameter decoder-only Transformer trained from
+scratch (manual forward pass, manual backward pass, manual AdamW — no
+autograd, no PyTorch/JAX/TensorFlow) on the TinyStories dataset. It generates
+short, simple, TinyStories-style children's narratives and nothing more.
+
+**Do not expect:** instruction following, multi-turn conversation, reasoning,
+factual world knowledge, or ChatGPT-comparable capability of any kind. This
+model was never trained or tuned for any of those.
+
+## Model description
+
+- **Model type:** decoder-only Transformer, Pre-LN, GELU (tanh approx), tied
+  token embedding / LM head (no output bias), learned positional embeddings.
+- **Parameters:** 13,049,856
+- **Context length:** 256 tokens
+- **Vocabulary:** 6,000 (locked byte-level BPE, `bpe_6000.json`)
+- **Framework:** none — hand-implemented NumPy/CuPy. Every layer's backward
+  pass was independently verified against finite-difference gradient checks
+  before training.
+
+## Architecture
+
+| param | value |
+|---|---|
+| vocab_size | 6000 |
+| block_size | 256 |
+| d_model | 384 |
+| n_layer | 6 |
+| n_head | 6 |
+| head_dim | 64 |
+| d_ff | 1536 |
+| activation | GELU (tanh approx) |
+| norm | Pre-LN |
+| positions | learned |
+| lm_head | tied to token embedding, no bias |
+
+## Training data
+
+[TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories)
+(`TinyStoriesV2-GPT4-train.txt`): 2,717,495 synthetically generated (GPT-3.5/
+GPT-4) short stories using a deliberately small vocabulary
+([Eldan & Li, 2023](https://arxiv.org/abs/2305.07759)), tokenized to
+371,525,259 tokens. Licensed by its authors under CDLA-Sharing-1.0 — this
+model card does not redistribute the dataset itself.
+
+## Training procedure
+
+- **Optimizer:** manually implemented AdamW (lr, betas, weight decay applied
+  only to matrix params — biases/LayerNorm params excluded from decay).
+- **Stage 2 (validation run):** 50k-story subset, 2000 steps, batch size 64,
+  LR 3e-4 with warmup, used to gate correctness before full training.
+- **Full run:** resumed from the Stage 2 checkpoint, continued on the full
+  371.5M-token corpus at LR 3e-5, to a **final checkpoint at step 5000**.
+- **Tokenizer:** locked, pre-trained externally, never retrained during model
+  training.
+
+## Hardware
+
+- 1x NVIDIA Tesla T4 (Turing, SM75, 16 GB VRAM), Kaggle.
+- CuPy 14.0.1 as the GPU numerical execution backend (no autograd usage).
+- Measured throughput: ~3,300–3,440 tokens/sec at batch size 64 (directly
+  measured, not extrapolated).
+
+## Intended use
+
+- Educational reference for from-scratch Transformer implementation
+  (manual forward/backward/AdamW) at small scale.
+- Generating short, TinyStories-style children's narratives from a prompt.
+- Portfolio / ML-engineering demonstration.
+
+## Out-of-scope use
+
+- Any production or consumer-facing assistant use case.
+- Instruction following, chat, question answering, factual retrieval,
+  reasoning tasks, or code generation.
+- Anything requiring broad world knowledge — the model's effective knowledge
+  is bounded by TinyStories' simplified vocabulary and narrative style.
+- Any use that assumes safety alignment or content filtering — **none was
+  performed.**
+
+## Evaluation
+
+From the training notebook (full-data run, step 5000):
+
+- Train loss ≈ 3.3
+- Val loss ≈ 3.3–3.4
+- Val perplexity ≈ 27–29
+
+No held-out benchmark suite (e.g. downstream NLP tasks) was run — TinyStories
+train/val loss and perplexity are the only reported metrics. Treat any
+numbers as approximate; see the training notebook's step-by-step log for the
+exact source values.
+
+## Known generation issues
+
+- Occasional run-on or abruptly concatenated sentences (short stories
+  sometimes blend into the next without a clean boundary).
+- Repetition of simple phrases/character names across generations.
+- No factual grounding — names, objects, and events are generated freely and
+  are not to be treated as accurate about anything.
+- Context is capped at 256 tokens; longer prompts are truncated from the
+  left before generation.
+
+## How to use
+
+```python
+from src.generate import load_dgpt, generate_text
+from src.tokenizer import BPETokenizer
+
+tok = BPETokenizer("tokenizer/bpe_6000.json")
+model, _ = load_dgpt("model.npz")
+
+print(generate_text(model, tok, "Once upon a time", max_new_tokens=150))
+```
+
+## Licensing
+
+License is marked `unknown` above deliberately. See this repository's main
+`README.md` → "Licensing" for the full breakdown across code, weights,
+tokenizer, and the TinyStories dataset — the weights and tokenizer do not
+have an established license and none is invented here.

config.json

@@ -0,0 +1,30 @@
+{
+  "model_type": "dgpt",
+  "model_name": "DGPT v1-base",
+  "architectures": ["DGPT"],
+  "vocab_size": 6000,
+  "block_size": 256,
+  "d_model": 384,
+  "n_layer": 6,
+  "n_head": 6,
+  "head_dim": 64,
+  "d_ff": 1536,
+  "activation": "gelu_tanh",
+  "norm": "pre_ln",
+  "positional_embedding": "learned",
+  "tie_word_embeddings": true,
+  "lm_head_bias": false,
+  "parameter_count": 13049856,
+  "final_checkpoint_step": 5000,
+  "tokenizer": {
+    "file": "bpe_6000.json",
+    "type": "byte_level_bpe",
+    "vocab_size": 6000,
+    "schema": "bpt_v1"
+  },
+  "generation_defaults": {
+    "temperature": 0.8,
+    "top_k": 40,
+    "max_new_tokens": 200
+  }
+}

model.npz

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

src/generate.py

@@ -0,0 +1,131 @@
+# src/generate.py | 158 lines
+"""
+DGPT v1-base inference: load model.npz + bpe_6000.json, generate text.
+
+Usage:
+    python src/generate.py \
+        --checkpoint checkpoints/model.npz \
+        --tokenizer tokenizer/bpe_6000.json \
+        --prompt "Once upon a time" \
+        --max_new_tokens 200 --temperature 0.8 --top_k 40
+
+Accepts either:
+  - the full training checkpoint (model.npz), which contains
+    `param__*`, `m__*`, `v__*`, and a JSON `__meta__` blob with the model
+    config and optimizer state, OR
+  - an inference-only weights file (model_weights.npz) produced by
+    scripts/extract_weights.py, which contains ONLY `param__*` keys and
+    requires --config to be passed explicitly (defaults to
+    configs/v1-base.json).
+"""
+
+import argparse
+import json
+import os
+import sys
+
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(__file__))
+from model import DGPT, stable_softmax  # noqa: E402
+from tokenizer import BPETokenizer  # noqa: E402
+
+DEFAULT_CONFIG_PATH = os.path.join(
+    os.path.dirname(__file__), "..", "configs", "v1-base.json"
+)
+
+
+def load_config(config_path=None):
+    path = config_path or DEFAULT_CONFIG_PATH
+    with open(path, "r") as f:
+        return json.load(f)
+
+
+def load_dgpt(checkpoint_path, config_path=None):
+    """
+    Loads a DGPT model for inference from either a full training checkpoint
+    or an inference-only weights file. Returns (model, step_or_none).
+    """
+    data = np.load(checkpoint_path, allow_pickle=True)
+    param_keys = [k for k in data.files if k.startswith("param__")]
+    if not param_keys:
+        raise ValueError(f"No 'param__*' arrays found in {checkpoint_path}")
+
+    params = {k[len("param__"):]: data[k] for k in param_keys}
+
+    if "__meta__" in data.files:
+        meta = json.loads(str(data["__meta__"]))
+        config = meta["config"]
+        step = meta.get("step")
+    else:
+        config = load_config(config_path)
+        step = None
+
+    model = DGPT(params, config)
+    return model, step
+
+
+def generate_text(model, tokenizer, prompt, max_new_tokens=200, temperature=0.8, top_k=40,
+                   seed=None):
+    """Autoregressive sampling. NumPy only (CPU inference)."""
+    rng = np.random.default_rng(seed)
+    generated = list(tokenizer.encode(prompt))
+
+    for _ in range(max_new_tokens):
+        context = generated[-model.block_size:]
+        idx = np.asarray([context], dtype=np.int64)
+
+        logits = model.forward(idx)
+        logits_last = logits[0, -1].astype(np.float32, copy=False)
+
+        temperature = max(float(temperature), 1e-6)
+        logits_last = logits_last / temperature
+
+        if top_k is not None and top_k > 0:
+            k = min(int(top_k), logits_last.shape[0])
+            top_idx = np.argpartition(logits_last, -k)[-k:]
+            filtered = np.full_like(logits_last, -1e10)
+            filtered[top_idx] = logits_last[top_idx]
+            logits_last = filtered
+
+        probs = stable_softmax(logits_last, axis=-1)
+        next_token = int(rng.choice(probs.shape[0], p=probs / probs.sum()))
+        generated.append(next_token)
+
+    return tokenizer.decode(generated)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="DGPT v1-base text generation")
+    parser.add_argument("--checkpoint", default="checkpoints/model.npz")
+    parser.add_argument("--tokenizer", default="tokenizer/bpe_6000.json")
+    parser.add_argument("--config", default=None, help="Only needed for weights-only npz files")
+    parser.add_argument("--prompt", default="Once upon a time")
+    parser.add_argument("--max_new_tokens", type=int, default=200)
+    parser.add_argument("--temperature", type=float, default=0.8)
+    parser.add_argument("--top_k", type=int, default=40)
+    parser.add_argument("--seed", type=int, default=None)
+    args = parser.parse_args()
+
+    print(f"Loading tokenizer from {args.tokenizer} ...")
+    tok = BPETokenizer(args.tokenizer)
+    assert tok.vocab_size == 6000, f"Expected vocab_size=6000, got {tok.vocab_size}"
+
+    print(f"Loading model from {args.checkpoint} ...")
+    model, step = load_dgpt(args.checkpoint, args.config)
+    print(f"Model loaded. step={step} vocab={model.vocab_size} block_size={model.block_size}")
+
+    print(f"\nPrompt: {args.prompt!r}\n")
+    output = generate_text(
+        model, tok, args.prompt,
+        max_new_tokens=args.max_new_tokens,
+        temperature=args.temperature,
+        top_k=args.top_k,
+        seed=args.seed,
+    )
+    print("Output:")
+    print(output)
+
+
+if __name__ == "__main__":
+    main()

src/model.py

@@ -0,0 +1,157 @@
+# src/model.py | 158 lines
+"""
+DGPT — from-scratch decoder-only Transformer (inference-only reimplementation).
+
+This is a forward-pass-only extraction of the architecture defined in the
+original training notebook (NanoScratchGPT / DGPT). The original notebook
+also hand-implements the backward pass and an AdamW optimizer for training;
+those are intentionally NOT reproduced here since this file is for loading
+`model.npz` and generating text only.
+
+Locked architecture (see configs/v1-base.json):
+    vocab_size = 6000
+    block_size = 256
+    d_model    = 384
+    n_layer    = 6
+    n_head     = 6
+    head_dim   = 64
+    d_ff       = 1536
+    activation = GELU (tanh approximation, GPT-2 style)
+    norm       = Pre-LN
+    positions  = learned
+    lm_head    = tied to token embedding, no output bias
+
+No PyTorch, no Hugging Face Transformers. NumPy only.
+"""
+
+import numpy as np
+
+
+def stable_softmax(x, axis=-1):
+    x_max = np.max(x, axis=axis, keepdims=True)
+    e = np.exp(x - x_max)
+    return e / np.sum(e, axis=axis, keepdims=True)
+
+
+def gelu(x):
+    c = np.float32((2.0 / np.pi) ** 0.5)
+    inner = c * (x + 0.044715 * x ** 3)
+    return 0.5 * x * (1.0 + np.tanh(inner))
+
+
+class Linear:
+    def __init__(self, W, b=None):
+        self.W = W
+        self.b = b
+
+    def __call__(self, x):
+        out = x @ self.W
+        if self.b is not None:
+            out = out + self.b
+        return out
+
+
+class LayerNorm:
+    def __init__(self, gamma, beta, eps=1e-5):
+        self.gamma = gamma
+        self.beta = beta
+        self.eps = eps
+
+    def __call__(self, x):
+        mu = x.mean(axis=-1, keepdims=True)
+        var = x.var(axis=-1, keepdims=True)
+        x_hat = (x - mu) / np.sqrt(var + self.eps)
+        return self.gamma * x_hat + self.beta
+
+
+class CausalSelfAttention:
+    def __init__(self, params, prefix, n_head):
+        self.Wq = Linear(params[f"{prefix}.Wq.W"], params[f"{prefix}.Wq.b"])
+        self.Wk = Linear(params[f"{prefix}.Wk.W"], params[f"{prefix}.Wk.b"])
+        self.Wv = Linear(params[f"{prefix}.Wv.W"], params[f"{prefix}.Wv.b"])
+        self.Wo = Linear(params[f"{prefix}.Wo.W"], params[f"{prefix}.Wo.b"])
+        self.n_head = n_head
+
+    def __call__(self, x):
+        B, T, C = x.shape
+        H = self.n_head
+        hd = C // H
+
+        Q, K, V = self.Wq(x), self.Wk(x), self.Wv(x)
+
+        def split_heads(t):
+            return t.reshape(B, T, H, hd).transpose(0, 2, 1, 3)
+
+        Qh, Kh, Vh = split_heads(Q), split_heads(K), split_heads(V)
+
+        scale = np.float32(1.0 / (hd ** 0.5))
+        scores = np.matmul(Qh, Kh.transpose(0, 1, 3, 2)) * scale
+
+        mask = np.triu(np.ones((T, T), dtype=bool), k=1)
+        scores = np.where(mask, np.float32(-1e9), scores)
+
+        A = stable_softmax(scores, axis=-1)
+        ctx = np.matmul(A, Vh)
+        ctx_merged = ctx.transpose(0, 2, 1, 3).reshape(B, T, C)
+        return self.Wo(ctx_merged)
+
+
+class FeedForward:
+    def __init__(self, params, prefix):
+        self.fc1 = Linear(params[f"{prefix}.fc1.W"], params[f"{prefix}.fc1.b"])
+        self.fc2 = Linear(params[f"{prefix}.fc2.W"], params[f"{prefix}.fc2.b"])
+
+    def __call__(self, x):
+        return self.fc2(gelu(self.fc1(x)))
+
+
+class TransformerBlock:
+    def __init__(self, params, prefix, n_head):
+        self.ln1 = LayerNorm(params[f"{prefix}.ln1.gamma"], params[f"{prefix}.ln1.beta"])
+        self.attn = CausalSelfAttention(params, f"{prefix}.attn", n_head)
+        self.ln2 = LayerNorm(params[f"{prefix}.ln2.gamma"], params[f"{prefix}.ln2.beta"])
+        self.ffn = FeedForward(params, f"{prefix}.ffn")
+
+    def __call__(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.ffn(self.ln2(x))
+        return x
+
+
+class DGPT:
+    """Inference-only DGPT. Load parameters with `DGPT.from_params(params, config)`."""
+
+    def __init__(self, params, config):
+        self.config = config
+        self.vocab_size = config["vocab_size"]
+        self.block_size = config["block_size"]
+        self.d_model = config["d_model"]
+        self.n_layer = config["n_layer"]
+        self.n_head = config["n_head"]
+
+        self.tok_emb_W = params["tok_emb.W"]
+        self.pos_emb_W = params["pos_emb.W"]
+        self.blocks = [
+            TransformerBlock(params, f"blocks.{i}", self.n_head)
+            for i in range(self.n_layer)
+        ]
+        self.ln_f = LayerNorm(params["ln_f.gamma"], params["ln_f.beta"])
+
+    def num_parameters(self):
+        return sum(v.size for v in self.__dict__.get("_raw_params", {}).values())
+
+    def forward(self, idx):
+        """idx: int array (B, T) with T <= block_size. Returns logits (B, T, vocab_size)."""
+        B, T = idx.shape
+        assert T <= self.block_size, "sequence length exceeds block_size"
+
+        tok = self.tok_emb_W[idx]
+        pos = self.pos_emb_W[:T]
+        x = tok + pos[None, :, :]
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        x = self.ln_f(x)
+        logits = x @ self.tok_emb_W.T  # tied weights, no output bias
+        return logits

src/tokenizer.py

@@ -0,0 +1,109 @@
+# src/tokenizer.py | 96 lines
+"""
+Loader for the LOCKED byte-level BPE tokenizer (`bpe_6000.json`, vocab=6000).
+
+This tokenizer was trained once, outside this repository, and must never be
+retrained or regenerated. This module only deserializes the finished merge
+table and performs encode/decode. There is no training method here by design.
+
+On-disk schema (BPT_V1):
+    {
+        "vocab_size": 6000,
+        "merge_order": [
+            [[a_id, b_id], new_id],
+            ...
+        ]
+    }
+"""
+
+import json
+
+
+class BPETokenizer:
+    def __init__(self, path):
+        self.path = path
+        with open(path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+
+        if "merge_order" not in data:
+            raise KeyError(
+                f"Unrecognized tokenizer file at {path}: expected a 'merge_order' key "
+                f"(BPT_V1 schema). Found top-level keys: {list(data.keys())}. "
+                "Do not attempt to regenerate this tokenizer."
+            )
+
+        merge_order = data["merge_order"]
+
+        # Base vocabulary: raw bytes 0..255
+        self.id_to_bytes = {i: (i,) for i in range(256)}
+        self.bpe_ranks = {}
+
+        for rank, entry in enumerate(merge_order):
+            a_id, b_id, new_id = self._normalize_merge_entry(entry)
+            self.bpe_ranks[(a_id, b_id)] = rank
+            self.id_to_bytes[new_id] = (a_id, b_id)
+
+        self.bytes_to_id = {v: k for k, v in self.id_to_bytes.items()}
+
+        declared_vocab_size = data.get("vocab_size")
+        computed_vocab_size = max(self.id_to_bytes.keys()) + 1
+        if declared_vocab_size is not None and declared_vocab_size != computed_vocab_size:
+            raise ValueError(
+                f"Tokenizer mismatch: declared vocab_size={declared_vocab_size}, "
+                f"computed={computed_vocab_size}. The tokenizer artifact is locked; "
+                "do not regenerate it."
+            )
+        self.vocab_size = computed_vocab_size
+
+    @staticmethod
+    def _normalize_merge_entry(entry):
+        if (
+            isinstance(entry, (list, tuple))
+            and len(entry) == 2
+            and isinstance(entry[0], (list, tuple))
+            and len(entry[0]) == 2
+        ):
+            return int(entry[0][0]), int(entry[0][1]), int(entry[1])
+        raise KeyError(f"Unrecognized merge entry: {entry!r}")
+
+    def _get_pairs(self, seq):
+        return set(zip(seq[:-1], seq[1:]))
+
+    def _bpe_merge(self, seq):
+        seq = list(seq)
+        if len(seq) < 2:
+            return seq
+        while True:
+            pairs = self._get_pairs(seq)
+            ranked = [(self.bpe_ranks[p], p) for p in pairs if p in self.bpe_ranks]
+            if not ranked:
+                break
+            _, best = min(ranked)
+            new_seq, i = [], 0
+            while i < len(seq):
+                if i < len(seq) - 1 and (seq[i], seq[i + 1]) == best:
+                    new_seq.append(self.bytes_to_id[(seq[i], seq[i + 1])])
+                    i += 2
+                else:
+                    new_seq.append(seq[i])
+                    i += 1
+            seq = new_seq
+        return seq
+
+    def encode(self, text):
+        raw_bytes = text.encode("utf-8")
+        return self._bpe_merge(list(raw_bytes))
+
+    def _expand(self, token_id):
+        if token_id < 256:
+            return bytes([token_id])
+        out = bytearray()
+        for part in self.id_to_bytes[token_id]:
+            out.extend(self._expand(part) if part >= 256 else bytes([part]))
+        return bytes(out)
+
+    def decode(self, ids):
+        out = bytearray()
+        for token_id in ids:
+            out.extend(self._expand(int(token_id)))
+        return bytes(out).decode("utf-8", errors="replace")