README.md

---
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.