train llm from scratch

The Problem

Training a large language model (LLM) from raw text to a usable assistant normally requires a stack of heavyweight libraries (transformers, trl, peft). This repo eliminates those dependencies and shows how to build the entire pipeline—tokenisation, transformer, loss, fine‑tuning, reward modelling, RL, and chat—in plain PyTorch. The goal is a single‑GPU trainable LLM that can be aligned for reasoning tasks without external tooling.

What This Does

Model definition lives in src/models/*.py (e.g., transformer.py, attention.py).
Data pipelines are in data_loader/*.py and the generic data_loader.py loads raw text, tokenises it, and builds PromptDataset, SFTDataset, or PreferenceDataset.
Training loops are in scripts/*.py (e.g., pretrain_base.py, train_sft.py, train_dpo.py).
Post‑training (SFT, reward model, PPO/DPO/GRPO) is implemented in src/post_training/*.py.
UI is a Streamlit app (ui/app.py) that orchestrates training, evaluation, and chat.

All configuration is JSON (configs/*.json) and loaded via config/loader.py.

How It Is Wired

Entry point: src/post_training/cli.py.

CLI parses arguments and calls train(), evaluate(), or inference() from src/post_training/*.
train() imports the chosen trainer (e.g., train_sft() from src/post_training/sft.py) which builds a dataset via data_loader.* and a model via src/models/transformer.py.
The model is trained in a PyTorch loop defined in src/post_training/optim.py and src/post_training/rollout.py.
After training, the checkpoint is saved to checkpoints/ and can be resumed (tests/test_checkpoint_resume.py).
Post‑training: src/post_training/reward_train.py trains a Bradley‑Terry reward model, src/post_training/dpo.py runs DPO, src/post_training/ppo.py runs PPO, and src/post_training/grpo.py runs GRPO.
Evaluation (src/post_training/evaluation.py) and inference (src/post_training/inference.py) use the trained checkpoints.

UI wiring: ui/app.py imports ui/stage_page.py and ui/pages/*. Each page imports the corresponding post‑training module and calls the same functions used by the CLI, providing a web interface for training and chat.

External effects:

Reads/writes checkpoints in checkpoints/.
Generates tokenised datasets in data/.
Writes logs to logs/ via src/post_training/logging_utils.py.
Serves a web UI on localhost:8501 via Streamlit.

How To Use It

Setup ``bash git clone https://github.com/moses-y/train-llm-from-scratch cd train-llm-from-scratch pip install -r requirements.txt ``
Configuration Edit the JSON files under configs/ (e.g., configs/pretrain.json) to set dataset paths, model hyper‑parameters, and GPU settings.
Training ``bash python scripts/pretrain_base.py --config configs/pretrain.json python scripts/train_sft.py --config configs/sft.json python scripts/train_reward.py --config configs/reward.json python scripts/train_dpo.py --config configs/dpo.json ``
Evaluation / Inference ``bash python scripts/eval_post_training.py --config configs/eval.json python scripts/generate_text.py --prompt "Explain quantum computing." ``
UI ``bash streamlit run ui/app.py ``

Real‑World Use

A research lab can use this repo to prototype a reasoning‑aligned assistant. Train a 13‑M parameter model on a private dataset, fine‑tune with DPO, evaluate with the built‑in tool, and expose the model via the Streamlit UI for quick demos.

Code Health & Issues

Low – Missing lockfile (pyproject.toml declares dependencies but no lockfile; reproducibility depends on pip).
Low – No Dockerfile or container build; deployment requires manual environment setup.
Low – CI limited to documentation (.github/workflows/docs.yml); no automated test runs.
Low – License is MIT, clear for reuse.

The repository contains 7 test files (tests/*.py) that cover checkpoint resume and smoke tests, but deeper integration tests for all training stages are absent.

The Bottom Line

This repo is an excellent educational platform for building an LLM from scratch in pure PyTorch. It demonstrates each training stage, from tokenisation to reinforcement learning, and offers a ready‑made UI. Production use would require adding a lockfile, containerisation, and a robust CI pipeline, but for research or teaching it delivers a clear, single‑GPU‑ready workflow.

What the analyser found

Deployment readiness

5/7

Container image
CI pipeline
Lockfile committed
Test suite
README
License
No committed secrets

Composition

186 files

Python77
Markdown24
JSON16
Shell3
YAML2
TOML2

Flask

Architecture

Top-level areas of the codebase, sized by module count. Arrows show how many imports cross from one area into another.

graph TD g0["src 26 modules"] g1["ui 18 modules"] g2["scripts 16 modules"] g3["tests 6 modules"] g4["data_loader 5 modules"] g5["config 4 modules"] g6["images 2 modules"] g2 -->|55| g0 g3 -->|22| g0 g2 -->|8| g5 g2 -->|7| g4 g1 -->|3| g5 g1 -->|3| g0 g4 -->|1| g0 g0 -->|1| g5 g3 -->|1| g2 g3 -->|1| g5

Module dependencies

The 10 most depended-upon modules of 78, from static import analysis. Red outline marks a module in an import cycle.

graph LR n0_src_models_transformer["transformer"] n1_src_post_training_utils["utils"] n2_config_post_training_config["post_training_config"] n3__post_training_chat_template["chat_template"] n4_src_post_training_rollout["rollout"] n5_ui_stages["stages"] n6_ui_stage_page["stage_page"] n7_st_training_rewards___init__["__init__"] n8_src_post_training_cli["cli"] n9_rc_post_training_distributed["distributed"] n4_src_post_training_rollout --> n3__post_training_chat_template n1_src_post_training_utils --> n0_src_models_transformer n6_ui_stage_page --> n5_ui_stages n5_ui_stages --> n2_config_post_training_config

Dependencies

9 declared dependencies, none a major version behind. Resolved 9 of 9 against the registry.

Findings

LowDependencies declared without a lockfile - non-reproducible builds pyproject.toml

What would stop me shipping this

Ranked by severity × confidence × production reach. Reach is the honest discriminator across a collection that is mostly other people's code: the same finding matters more in something that ships.

high2

medium4

low1

Commit a lockfile beside the manifesthigh

pyproject.toml

manifest with no lockfile in scope

An unlocked range means the artifact you tested and the artifact you ship can contain different transitive code, so a malicious patch release reaches production with no diff.

Fix: Run the package manager once and commit the generated lockfile.

Make CI invoke the test suite it hashigh7 occurrences

.github/workflows

7 test files, no test command in any workflow

A green check that never executed an assertion is worse than no check, because reviewers trust it.

Fix: Add a test step to the existing workflow rather than a new workflow.

Enable Dependabot or Renovatemedium

3 manifest(s), no update bot configured

Without a bot a published advisory sits unpatched until someone audits by hand, which across 1,322 repositories means never.

Fix: Commit .github/dependabot.yml covering the repo ecosystems plus github-actions.

Gate pull requests on a dependency vulnerability scanmedium

.github/workflows

no dependency scan in CI

This is the one gate that would catch a known-vulnerable package before it reaches a build, and no repository in the sample had it.

Fix: Add dependency-review-action on pull_request, or osv-scanner on push and a schedule.

Split the oversized notebook cell into stepsmedium

sft_rlhf_guide.ipynb

largest cell 7438 characters of 76615 (1 notebooks sampled)

A failure anywhere in a 22,000-character cell means rerunning the whole pipeline from load to model, so nothing can be debugged in isolation.

Fix: Break at the load, clean, feature, model and evaluate boundaries.

Set persist-credentials: false on checkoutmedium

.github/workflows/docs.yml

checkout keeps the token, then dependencies are installed

The token stays in .git/config for every later step, so a malicious postinstall script reads a pushable credential without one ever being passed to it.

Fix: Add with: persist-credentials: false, and pass an explicit token only to the step that pushes.

Set timeout-minutes on the workflow jobslow

.github/workflows/docs.yml

1 workflow(s) declare no job timeout

A wedged step runs to the six-hour platform default, which on a two-hourly schedule means three runs overlap behind it.

Fix: Add timeout-minutes with a realistic bound to each job.

Checked deterministically against the repository tree and a bounded set of its files: committed credentials, unpinned actions and base images, missing lockfiles and update bots, workflows that discard failures, published advisories against the declared dependencies, runtime configuration, licensing and notebook reproducibility. No language model is involved in this section.