Continual Learning Pipeline for v11

Overview

Iteratively improve the model by:

1. Training on base data (books/plays) to learn Victorian conversation style
2. Adding new graded data (modern prompts → Victorian responses)
3. Replacing base data with held-out portion in subsequent rounds
4. Eventually discarding base data entirely, just adding new graded data

Data Strategy

Base Data (books/plays)

Split into two halves to avoid memorization:

• Round 1: Uses split1 (first 50% of base data)
• Round 2: Uses split2 (second 50% of base data)
• Round 3+: Discards base data entirely

New Data (graded conversations)

Always included, keeps growing:

• prose/json_sft/ready/cues_remapped.jsonl (900)
• prose/json_sft/ready/graded_convs.jsonl (294)
• prose/json_sft/ready/graded_batch400.jsonl (122 edit)
• prose/json_sft/ready/handpicked.jsonl (97 keep)
• prose/json_sft/ready/danglers_person.jsonl (1,132 pass)
• prose/json_sft/reviewed/cues_batch1.reviewed.jsonl (299)
• Total: ~3,929 high-quality examples

To add more graded data:

• Just drop new JSONL files into prose/json_sft/ready/ or prose/json_sft/reviewed/
• Or add them to the new_data_paths list in build_continual_dataset.py

Data Size Strategy

Original SFT: 37,943 records, 0.11 passes in 300 steps

Continual learning rounds 1-2:

• Sample 15K base records (not full 324K!)
• Add 3.9K new graded records
• Total: ~18.9K records (~0.3-0.4 passes in 300 steps)
• Ratio: 79% base, 21% new

Why sample?

• Full 324K base data would overwhelm the new data
• Model needs to learn from new prompts, not just memorize base
• 15K is enough to reinforce Victorian style without drowning new data

Pipeline

Step 1: Round 1 - Train from anneal on first half of base data

cd /Users/melissa/Code/tiny-chat
python3 v11/mkm_train/build_continual_dataset.py --round 1
python3 v11/mkm_train/sft.py \
    --init-from v11/models/anneal/checkpoints_anneal_e_bl_mixed/ckpt_final.pt \
    --data v11/data/continual_round1.jsonl \
    --checkpoint-dir v11/models/sft_round2/r1 \
    --lr 2e-5 \
    --val-frac 0.15 \
    --max-iters 300

What happens:

• Automatically creates split files (one-time setup)
• Starts from anneal checkpoint (has learned chat shape)
• Trains on sampled split1 (15K) + graded data (3.9K) = 18.9K total
• Best checkpoint saved to v11/models/sft_round2/r1/ckpt_best.pt

Step 2: Round 2 - Train on second half of base data

cd /Users/melissa/Code/tiny-chat
python3 v11/mkm_train/build_continual_dataset.py --round 2
python3 v11/mkm_train/sft.py \
    --init-from v11/models/sft_round2/r1/ckpt_best.pt \
    --data v11/data/continual_round2.jsonl \
    --checkpoint-dir v11/models/sft_round2/r2 \
    --lr 1e-5 \
    --val-frac 0.15 \
    --max-iters 300

What happens:

• Starts from r1 checkpoint
• Trains on sampled split2 (15K) + graded data (3.9K) = 18.9K total
• Different base data than round 1 → no memorization
• Best checkpoint saved to v11/models/sft_round2/r2/ckpt_best.pt

Step 3+: Round N - Train on new graded data only

cd /Users/melissa/Code/tiny-chat
python3 v11/mkm_train/build_continual_dataset.py --round 3 --no-base
python3 v11/mkm_train/sft.py \
    --init-from v11/models/sft_round2/r2/ckpt_best.pt \
    --data v11/data/continual_round3.jsonl \
    --checkpoint-dir v11/models/sft_round2/r3 \
    --lr 5e-6 \
    --val-frac 0.2 \
    --max-iters 300

What happens:

• Starts from r2 checkpoint
• Trains on graded data only (no base JSON)
• Add more graded data as you create it
• Lower LR to avoid overfitting on small dataset

Note: Round 3+ uses --no-base flag to skip base JSON entirely

Testing Your Model

After each round, test with:

cd /Users/melissa/Code/tiny-chat
python3 v11/chat.py --checkpoint v11/models/sft_round2/r1/ckpt_best.pt

Or for a quick test without interactive chat:

cd /Users/melissa/Code/tiny-chat
python3 << 'EOF'
import sys
sys.path.insert(0, 'v11')
from chat import load_checkpoint, tokenize, generate
import sentencepiece as spm
import torch

device = "mps" if torch.backends.mps.is_available() else "cpu"
model = load_checkpoint('v11/models/sft_round2/r1/ckpt_best.pt', device)
sp = spm.SentencePieceProcessor(model_file='v11/tokenizers/tokenizer_v11_16k.model')

prompts = ["What's for dinner?", "I'm going to the movies."]
for prompt in prompts:
    conversation = [{"role": "user", "content": prompt}]
    prompt_tokens, _ = tokenize(conversation, sp, eos=False)
    response = generate(model, sp, prompt_tokens, device, max_new=40, temp=0.8)
    print(f"Q: {prompt}\nA: {response}\n")
EOF

Key Parameters

Notes

• Round 1 is special: starts from anneal (has learned chat shape), uses first half of base data
• Round 2 uses second half of base data (no overlap with round 1)
• Round 3+ uses only new graded data (base data discarded)
• Early stopping via smoothed validation loss (best checkpoint auto-saved)
• Responses may be too long → use max_new=40-50 for generation