Update README.md

README.md

@@ -1,3 +1,183 @@
 ---
 license: apache-2.0
+language:
+- en
+tags:
+- text-generation
+- gpt
+- nanogpt
+- pytorch
+- from-scratch
+- small-language-model
+- slm
+- pretraining
+- muon-optimizer
+- single-gpu
+- base-model
+- research
+pipeline_tag: text-generation
 ---
+
+# GTM-v2-base
+
+A ~120M parameter, decoder-only GPT-style **base** language model, trained
+from scratch on a single RTX Pro 6000, on roughly 4 billion tokens streamed
+from a mix of open web/educational/math corpora.
+
+**This is a base (pretrained) model, not an instruction-tuned or chat
+model.** It completes/continues text; it does not reliably follow
+instructions or answer questions directly. It also has **no fine-tuning
+for factual accuracy** — it can and does confidently generate fluent,
+plausible-sounding but factually incorrect content. Treat it as a small,
+from-scratch research/hobby model, not a production system.
+
+A separate instruction-tuned variant (SFT on UltraChat) is planned as a
+follow-up release; this repo is the base-only checkpoint.
+
+## Model details
+
+- **Architecture**: nanoGPT-style decoder-only transformer
+  - 14 layers, 8 attention heads, 704 embedding dim
+  - ~119.4M parameters (weight-tied embeddings/output head)
+  - Context length: 1024 tokens
+  - Uses PyTorch's fused `scaled_dot_product_attention` (flash-attention kernel)
+- **Tokenizer**: `tiktoken` GPT-2 BPE encoding (`tiktoken.get_encoding("gpt2")`),
+  vocab size 50,257. No custom tokenizer was trained.
+- **Optimizer**: Muon (for 2D weight matrices) + AdamW (for embeddings,
+  layernorms, biases) — a hybrid setup, following the approach popularized
+  in recent efficient-pretraining recipes.
+- **Precision**: trained with bf16 autocast; released weights are fp32.
+- **Training tokens seen**: ~4.16B tokens (40,690 steps x effective batch
+  100 x 1024 context)
+- **Training data**: streamed via HuggingFace `datasets`, mixed and
+  tokenized on the fly (no fixed local copy of the source datasets):
+  - FineWeb-Edu (`HuggingFaceFW/fineweb-edu`, `sample-10BT`) — 45%
+  - Cosmopedia-v2 (`HuggingFaceTB/cosmopedia-v2`) — 30%
+  - FineMath (`HuggingFaceTB/finemath`, `finemath-4plus`) — 15%
+  - FineWeb (`HuggingFaceFW/fineweb`, `sample-10BT`) — 10%
+
+  Note: code data (The Stack v2, StarCoderData, the-stack-smol) was
+  intentionally left out of this mix — every BigCode-hosted code corpus we
+  checked is gated behind a terms-of-use click-through on HuggingFace, so
+  none of them are pulled in without the user separately accepting those
+  terms and authenticating.
+
+This is the second model in this series. An earlier ~100M parameter
+version (GTM-v1-base, trained on ~3B tokens with the same data mix and
+optimizer) informed several of this model's design choices, including the
+Muon/AdamW optimizer split and the corpus weighting above.
+
+## Benchmarks
+
+Evaluated via likelihood-scoring (comparing the model's per-token loss
+across candidate answers, no generation/sampling involved), 200 examples
+per benchmark, final checkpoint at step 40,690 (~4.16B tokens seen):
+
+| Benchmark      | GTM-v2-base | GTM-v1-base | GPT-2 (124M) | Random baseline |
+|----------------|:-----------:|:-----------:|:------------:|:----------------:|
+| HellaSwag      | 33.0%       | 32.0%       | ~28-29%      | 25%              |
+| ARC-Easy       | 43.0%       | 42.5%       | ~39.2%       | ~25%             |
+| ARC-Challenge  | 24.0%       | 26.0%       | ~22.5%       | ~25%             |
+
+Corpus perplexity (per-source, lower is better) on freshly streamed,
+non-training-shard text:
+
+| Source         | Perplexity |
+|----------------|:----------:|
+| Cosmopedia-v2  | 11.80      |
+| FineMath       | 13.04      |
+| FineWeb-Edu    | 27.68      |
+| FineWeb        | 47.34      |
+| **Overall**    | **21.19** (vs. GTM-v1-base's 22.87) |
+
+GTM-v2-base improves on GTM-v1-base (an earlier ~100M param, ~3B token
+version of this same series) on HellaSwag, ARC-Easy, and overall
+perplexity, but is *not* a clean sweep — it's slightly behind v1 on
+ARC-Challenge (24.0% vs. 26.0%), though both remain above GPT-2's ~22.5%
+reference. ARC-Challenge is a small (200-example), noisy evaluation at
+this scale, so a few-point gap in either direction shouldn't be read as a
+strong signal on its own.
+
+GPT-2 (124M) reference numbers are from a from-scratch reproduction
+verified to match official GPT-2 evaluation results. GTM-v2-base was
+trained on roughly 1/10th of GPT-2's training tokens (~4B vs. ~40B), on a
+single consumer/workstation GPU rather than a large multi-GPU cluster.
+
+**These numbers should be read as "beats a 2019 baseline on a few narrow
+multiple-choice benchmarks," not "as good as or better than GPT-2 in
+general."** GPT-2 was trained and evaluated far more broadly; free-form
+generation quality, factual grounding, and behavior outside these specific
+benchmark formats have not been rigorously compared.
+
+## Known limitations
+
+- **No reliable factual recall.** E.g. prompted with "The capital of France
+  is", the model does not consistently produce "Paris" — it may produce
+  fluent but factually invented content instead. This is expected: the
+  training corpus is not dense in discrete factual content, and this
+  model's parameter/token budget is small for reliably memorizing specific
+  facts.
+- **No code capability.** No code-specific training data was included (see
+  above). The model can produce code-*shaped* text (recognizing "write
+  code" prompts and using plausible syntax) but not functionally correct
+  code.
+- **Not instruction-tuned.** Does not follow instructions or answer
+  questions in a chat-like way — it continues text. (A separate SFT/chat
+  release is planned; this is the base-only checkpoint.)
+- **Repetition tendency.** Base generation (greedy or low-temperature
+  sampling) can fall into repetition loops or lock onto structural
+  templates (e.g. numbered-list/fill-in-the-blank formatting) when it
+  lacks a confident continuation. A repetition penalty (see usage example)
+  substantially reduces this.
+
+## Usage
+
+Requires `model.py` (included in this repo) alongside the checkpoint —
+this is a plain PyTorch model, not a `transformers` `AutoModel`.
+
+```bash
+pip install torch safetensors tiktoken
+```
+
+```python
+import json
+import torch
+from safetensors.torch import load_file
+from model import GPT, GPTConfig
+
+with open("config.json") as f:
+    cfg_dict = json.load(f)
+config = GPTConfig(
+    vocab_size=cfg_dict["vocab_size"], block_size=cfg_dict["block_size"],
+    n_layer=cfg_dict["n_layer"], n_head=cfg_dict["n_head"],
+    n_embd=cfg_dict["n_embd"], dropout=cfg_dict["dropout"], bias=cfg_dict["bias"],
+)
+model = GPT(config)
+state_dict = load_file("model.safetensors")
+model.load_state_dict(state_dict)
+model.eval()
+
+import tiktoken
+enc = tiktoken.get_encoding("gpt2")
+
+prompt = "Once upon a time,"
+ids = enc.encode_ordinary(prompt)
+x = torch.tensor([ids], dtype=torch.long)
+
+with torch.no_grad():
+    out = model.generate(
+        x, max_new_tokens=128, temperature=0.8, top_k=50,
+        eot_token=enc.eot_token, repetition_penalty=1.3,
+    )
+
+print(enc.decode(out[0].tolist()))
+```
+
+## License
+
+Apache 2.0 for this repo's contents (model weights, `model.py`, and this
+README). The underlying training data retains its own licenses regardless
+of the license on this trained model — FineWeb, FineWeb-Edu, Cosmopedia-v2,
+and FineMath are each ODC-BY-1.0 (see their respective HF dataset cards for
+full terms and attribution requirements). This repo distributes model
+weights, not the training data itself.