TT Lab
Get started
Learn Learning paths Courses

MiniMind — Train a Small Language Model Yourself, End to End

Pretrain for 300 steps with MiniMind's training loop and read the loss curve

Continue in TT Lab

Goal

Write yourself a training loop that shrinks the skeleton of MiniMind's train_pretrain.py (AdamW, MiniMind cosine learning rate, gradient clipping 1.0) to CPU size, fix the seed and check that it reproduces, and then train 300 steps. You read the loss curve in numbers and see what breaks when the learning rate is too large.

Why it matters

The core of a training loop is two lines, but what decides the result is the decisions around them — the learning rate and its schedule, gradient clipping, the seed, and how often to measure the validation loss. If these decisions are wrong, you find out only after running for hours. That is why you need the habit of first checking the shape of the curve with a training of a few minutes. The loss curve also shows what the model learns first. The first few dozen steps learn token frequency, and after that it learns rules that depend on context. If you compare it with the loss of a model that knows only frequency without context (the unigram entropy), you can tell which part of the curve is the stretch where it "began to use context".

Steps

  1. Write /root/mm/pre/train.py (arguments --steps, --lr, --seed, --eval-every, --out, and --log), run 20 steps twice with the same seed, and leave /root/mm/pre/log_a.csv and /root/mm/pre/log_b.csv.
  2. Train 300 steps (batch 8, length 128, lr 3e-3) and leave the weights in /root/mm/pre/ckpt.pth and the log in /root/mm/pre/log.csv. The validation loss must come down below 1.6.
  3. Extract the shape of the curve from log.csv and write it to /root/mm/pre/curve.json.
  4. Write the unigram entropy of the training array to /root/mm/pre/unigram.json.
  5. Run 50 steps at lr 0.05 and leave /root/mm/pre/log_high.csv.
  6. Using ckpt.pth, continue the text after "Garam village's" for 16 tokens with greedy decoding and save it to /root/mm/pre/sample.txt.
  7. In /root/mm/pre/report.md, write the three sections ## 손실 곡선, ## 학습률, and ## 이어 쓰기 (the Korean headings mean "Loss curve", "Learning rate", and "Continuation"), and include the final validation loss and the unigram entropy as numbers.

Notes

Same seed, same loss

Write /root/mm/pre/train.py. It builds the model from /opt/mm/ref/config.json, trains by drawing 8 batches of 128 tokens at random positions from /opt/mm/ref/train.npy, and writes step,train_loss,val_loss,lr at every step to a CSV at the --log path. Run it twice with --steps 20 to make /root/mm/pre/log_a.csv and /root/mm/pre/log_b.csv.

To reproduce, you must fix both random sources — weight initialization (mmkit.seed_all(시드) before creating the model, where the placeholder is the seed) and batch drawing (passing torch.Generator().manual_seed(시드) to randint, where the placeholder is the seed). The first-step loss must be near ln 1024 ≈ 6.93.

Pretrain for 300 steps

With train.py, train 300 steps (batch 8, length 128, lr 3e-3, seed 42, validation every 50 steps) and leave the weights in /root/mm/pre/ckpt.pth (model.state_dict()) and the log in /root/mm/pre/log.csv.

The learning rate must follow MiniMind's cosine (mmkit.minimind_lr), being 3e-3 at the first step and 3e-4 at the last. The grader loads ckpt.pth, remeasures the validation loss, and checks that it is below 1.6 and equal to the last val_loss in the log.

The shape of the curve in numbers

From log.csv, write the first loss (first_loss), the step at which the training loss first went below 3 (step_below_3), the final training loss (final_train), and the final and lowest validation losses (final_val and best_val) to /root/mm/pre/curve.json.

See at which step the steeply falling stretch ends. After that, it takes many more steps to bring it down by the same amount.

The loss of a model that knows only frequency

Compute the unigram entropy (−Σ p·ln p, in nats) from the token frequencies of /opt/mm/ref/train.npy and write it to /root/mm/pre/unigram.json as unigram_entropy.

This value is the cross entropy of a model that looks at no context and always outputs the same distribution (the token frequencies). If the curve goes below this value, it means the model has begun to use the preceding token to narrow down the next. numpy.bincount is fast.

When the learning rate is too large

Run train.py with --lr 0.05 --steps 50 and leave /root/mm/pre/log_high.csv. Since the seed is the same, you can compare with steps 41–50 of the step 2 log.

A large learning rate seems to drop fast for the first few steps, then spikes once and sticks at some value and does not come down. Gradient clipping (1.0) cannot limit the size of the step either — the size of the step is decided by lr.

Continue a text with what it learned

Load ckpt.pth and continue after [bos] + '가람 마을의' for 16 tokens with greedy decoding (do_sample=False, top_k=0, top_p=1.0, eos_token_id=2), and save the whole text without bos to /root/mm/pre/sample.txt. The Korean phrase in the code means "of Garam village".

Load with mmkit.load_model(경로) (the placeholder is the path) and call MiniMind's model.generate. Greedy generation always produces the same text for the same weights, so the grader regenerates with your ckpt.pth and compares. A model that has only been pretrained does not answer questions but continues like the corpus.

A record of reading the curve

In /root/mm/pre/report.md, write the three sections ## 손실 곡선, ## 학습률, and ## 이어 쓰기 (the Korean headings mean "Loss curve", "Learning rate", and "Continuation"), and include the final_val from step 3 and the unigram_entropy from step 4 as numbers.

Write one line, using the unigram entropy as the reference, on which stretch of the curve learned frequency and which learned context. In the learning rate section, write what you saw in step 5.