v11 SFT Setup

This folder contains the complete SFT pipeline for the v11 model line.

Directory Structure

v11/
├── mkm_train/
│   ├── model.py          # GPT architecture (245M params, 18 layers)
│   ├── train.py          # Pretraining script
│   ├── sft.py            # SFT script (this folder)
│   └── checkpoints/      # Training checkpoints
├── models/
│   ├── sft/              # Annealed checkpoints (a-f variants)
│   │   ├── a_books_heavy/
│   │   ├── b_plays_x3/
│   │   └── ...
│   └── sft_v2/           # New SFT checkpoints (created by sft.py)
├── data/
│   ├── pretrain_gutenberg.npy
│   ├── pretrain_blbooks.npy
│   └── anneal/           # Annealing data (npy files)
├── json/                 # SFT conversation data
│   ├── ranked_convs_v2.jsonl
│   ├── chat_named.jsonl
│   ├── chat_unattributed.jsonl
│   └── ...
├── tokenizers/
│   └── tokenizer_v11_16k.model
└── ouroboros/            # Chat generation & evaluation

Quick Start

1. Pick an Annealed Checkpoint

We have 6 annealed variants in models/sft/:

2. Pick SFT Data

Available in json/:

• ranked_convs_v2.jsonl - Graded conversations (best quality)
• chat_named.jsonl - Named speaker conversations
• chat_unattributed.jsonl - Unattributed conversations
• play_exchanges.jsonl - Play dialogue
• book_exchanges.jsonl - Book dialogue

3. Run SFT

cd /Users/melissa/Code/tiny-chat

python3 v11/mkm_train/sft.py \
    --init-from v11/models/sft/a_books_heavy/ckpt_best.pt \
    --data v11/json/ranked_convs_v2.jsonl \
    --checkpoint-dir v11/models/sft_v2/a_books_heavy_r3b \
    --max-iters 300 \
    --batch-size 8 \
    --lr 2e-5

4. Evaluate

The script automatically:

• Splits conversations into train/val (by unique conversation)
• Masks loss to assistant tokens only
• Saves best checkpoint based on smoothed val loss
• Records sampling settings (temp, top_k, repetition_penalty)

Use ckpt_best.pt (not the final checkpoint) - SFT overfits fast.

Key Design Decisions

1. Masked Loss: Only assistant tokens are supervised (user turns are context)
2. Conversation-level Split: Prevents val from measuring memorization
3. Smoothed Val Selection: Single-draw val is noisy; 3-point trailing mean picks better checkpoints
4. Self-contained: All code and data in v11/ - no prose/ dependencies

Local vs Cloud

• Local: Use this sft.py (v11/mkm_train/)
• Cloud: Use prose/mkm_train/sft.py (same logic, different imports)

The v11 version uses v11/mkm_train/model.py which has a different architecture than prose/mkm_train/model.py.