Shiyin Dong

From-scratch model training

Mary LLM

A small Chinese language model, trained from tokenizer to preference alignment.

Code
Left: text is tokenized by a BPE tokenizer trained on the corpus, and a decoder-only Transformer predicts the next token — every part is trained from scratch. Right: the four training stages.
Left: text is tokenized by a BPE tokenizer trained on the corpus, and a decoder-only Transformer predicts the next token — every part is trained from scratch. Right: the four training stages.
29M / 104M
small / base parameters
2048
max sequence length
4
training stages

Two model sizes, one pipeline, run end to end on 4 × RTX 3090.

Overview

Mary LLM is a study project: instead of fine-tuning an existing checkpoint, every stage of the language model pipeline is written out and run — training a tokenizer on the corpus, defining the model architecture, pre-training it on next-token prediction, then aligning it with supervised fine-tuning and preference optimization.

The point is not a competitive model. It is to keep each stage small enough to read and cheap enough to re-run, so the effect of a stage is visible. The final checkpoint is served through a Gradio interface for quick manual testing.

Techniques used

The four stages, and what each one actually does.

Tokenizer
Text has to become integers before a model can read it. A tokenizer is trained on the corpus itself so that common Chinese sequences get a single token, which keeps the vocabulary compact and the sequences short.
Pre-training
Next-token prediction on raw text. The model only learns to continue text here — it does not yet know how to answer a question.
SFT (supervised fine-tuning)
Training on instruction–response pairs, which is what turns a text continuer into something that replies to a prompt.
DPO (direct preference optimization)
Alignment from pairs of preferred / rejected answers. It optimizes the same objective RLHF targets, but directly on the pairs, so no separate reward model or RL loop is needed.

Pipeline

  1. 01

    Train the vocabulary

    Build a compact tokenizer for the training corpus.

  2. 02

    Pre-train the base

    Next-token prediction on the 29M small and 104M base configurations.

  3. 03

    Supervised fine-tuning

    Adapt the base model to instruction–response data.

  4. 04

    Preference optimization + demo

    Run DPO on preference pairs, then expose the result through a Gradio interface.

Repository map

Each script corresponds to one decision in the pipeline.

  • model.pyModel architecture and configuration.
  • pre_train.pyPre-training loop.
  • sft.pySupervised fine-tuning.
  • dpo.pyDirect preference optimization.
  • gradio_llm.pyInteractive interface for the trained model.

Training environment

Hardware
RTX 3090 × 4
Stack
PyTorch 2.1.2 + CUDA 12.1 · Python 3.10