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:

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.