HeyMSK
Chapter 7 of 10

Teach a new format

Our model learned from examples like 23+19=42. Now I want it to understand a different format, Q:23+19;A:42. Q marks the question and A marks the answer.

I load the trained model and continue training with examples in this format. This is called fine-tuning. It starts with what the model has already learned instead of starting with random weights.

More training can also change how it answers the old questions. I'll check both formats afterwards.

Can it learn the new format?

Here's an example in the new format:

Q:23+19;A:42

Before the extra training, I checked how the base model answered validation questions in this format.

Before fine-tuning: 0 of 1,006

These saved results count answers that were correct and stopped in the right place.

Did it still understand the old format?

After choosing the fine-tuned model using validation questions, I tested it on a separate set. It got 486 of 502 new-format answers right. On the original format, it got 501 of 502 right.

The training examples included the old format too, to help the model retain it. The final test questions weren't used to choose the model.

Try before and after fine-tuning

Run Q:23+19;A: with Improved model, then Fine-tuned Q and A. Compare the answers, then try the old 23+19= format with both models.

Runs on your device. First load is about 14 MB.

Exact input: "Q:23+19;A:"

Model output
No output yet
Input + output
Q:23+19;A:

Run to see a prediction.

Model and input details

Printable ASCII is passed verbatim. BOS is added once. All 98 tokens remain available. Context: 32 tokens. Run budget: 16 output tokens, capped by remaining context. EOS ends generation. Unexpected PAD or BOS tokens are visible.

Improved model. Epoch 140. 1.74 MB of weights. The WASM runtime is downloaded on first use.

Optional code

Code and commands for this step

Fine-tuning starts from saved weights. I change the example format, then repeat weight updates. The target mask selects only answer positions and EOS.

# Each row already holds the correct answer.
prompt = f"Q:{row['a']}+{row['b']};A:"
text = prompt + row["answer"]

# logits: [batch, positions, 98]; y: target token IDs
starts = torch.tensor([len(r["prompt"]) for r in values])[:, None]
mask = (torch.arange(y.shape[1])[None, :] >= starts) & (y != 0)
loss = F.cross_entropy(logits[mask], y[mask])

The original-format examples are included too. This helps the model retain its previous behavior. The fine-tuning script supplies the batches, device handling and training loop.

Run the fine-tuning experiment after the base model is trained.

Read the complete experiments.py
"""Predeclared fine-tuning, frozen evaluations, calibration and candidate generation."""
import argparse
import hashlib
import json
import math
import random
import sys
import time
from collections import defaultdict
from pathlib import Path
import torch
from torch.nn import functional as F

parser = argparse.ArgumentParser()
parser.add_argument('--source', type=Path, required=True)
parser.add_argument('--action', choices=['finetune', 'evaluate', 'evaluate-rl', 'candidates'], required=True)
parser.add_argument('--device', default='mps')
args = parser.parse_args()
sys.path.insert(0, str(args.source))
from tiny_gpt.checkpoint import load_checkpoint, save_checkpoint
from tiny_gpt.tokenizer import CharacterTokenizer
from tiny_gpt.examples import encode_record, make_batch
from tiny_gpt.evaluation import evaluate, better_validation
from tiny_gpt.generation import greedy_generate_batch, generate_answer

torch.set_num_threads(4)
tokenizer = CharacterTokenizer()
out = Path('public/tiny-llm')
run = Path('scripts/tiny-llm/runs/finetune')
run.mkdir(parents=True, exist_ok=True)
base_path = args.source / 'runs/pretrain-v2/best.pt'
def sha(path): return hashlib.sha256(path.read_bytes()).hexdigest()
def write(path, content): path.write_text(json.dumps(content, indent=2, allow_nan=False) + '\n')
def rows(split, formatted=False):
    values = [json.loads(line) for line in (args.source / f'data/addition-v1/{split}.jsonl').read_text().splitlines()]
    if formatted:
        values = [{**r, 'prompt': f"Q:{r['a']}+{r['b']};A:", 'text': f"Q:{r['a']}+{r['b']};A:{r['answer']}"} for r in values]
    return values
def data(values): return {'rows': values, 'batch': make_batch([encode_record(r, tokenizer) for r in values])}
def completion_mask(batch, values, device):
    y = batch['target_ids'].to(device)
    starts = torch.tensor([len(r['prompt']) for r in values], device=device)[:, None]
    return (torch.arange(y.shape[1], device=device)[None, :] >= starts) & (y != 0)

if args.action == 'finetune':
    if (run / 'best.pt').exists(): raise RuntimeError('Fine-tuning checkpoint already exists. Preserve it; use a new run for another experiment.')
    torch.manual_seed(73)
    model, _ = load_checkpoint(base_path, device=args.device)
    values = rows('train') + rows('train', True)
    batch = data(values)['batch']
    x, y = batch['input_ids'].to(args.device), batch['target_ids'].to(args.device)
    mask = completion_mask(batch, values, args.device)
    validation = data(rows('validation', True))
    canonical = data(rows('validation'))
    initial = evaluate(model, validation)
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.0002, betas=(0.9, 0.95), weight_decay=0.01)
    generator = torch.Generator().manual_seed(73)
    best, history, step = None, [], 0
    started = time.perf_counter()
    for epoch in range(1, 41):
        model.train()
        order = torch.randperm(len(values), generator=generator).to(args.device)
        summed, count = 0, 0
        for start in range(0, len(values), 256):
            chosen = order[start:start + 256]
            optimizer.zero_grad(set_to_none=True)
            logits = model(x[chosen])
            loss = F.cross_entropy(logits[mask[chosen]], y[chosen][mask[chosen]])
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1, error_if_nonfinite=True)
            optimizer.step(); step += 1
            n = int(mask[chosen].sum()); summed += float(loss.detach()) * n; count += n
        record = {'epoch': epoch, 'steps': step, 'completion_loss': summed / count}
        if epoch == 1 or epoch % 5 == 0:
            measured = evaluate(model, validation)
            regression = evaluate(model, canonical)
            record.update(validation=measured, canonical_validation=regression)
            if better_validation(measured, best):
                best = measured
                save_checkpoint(model, run / 'best.pt', metadata={'stage': 'finetuned-qa', 'epoch': epoch, 'steps': step, 'seed': 73, 'validation': measured, 'base_sha256': sha(base_path)})
            print(f"epoch {epoch}: formatted {measured['correct']}/1006, canonical {regression['correct']}/1006", flush=True)
        history.append(record)
        write(run / 'history.json', history)
    write(out / 'finetune.json', {'protocol': '07-experiment-protocol.md', 'seed': 73, 'source_sha256': sha(base_path), 'selected_sha256': sha(run / 'best.pt'), 'initial_formatted_validation': initial, 'selected_validation': best, 'history': history, 'elapsed_seconds': time.perf_counter() - started, 'training_rows': len(values), 'maximum_epochs': 40, 'test_used_for_selection': False})

def generated_metrics(model, values):
    grouped = defaultdict(list)
    for r in values: grouped[len(r['prompt'])].append(r)
    outcomes = []
    for group in grouped.values():
        for start in range(0, len(group), 256):
            part = group[start:start + 256]
            answers = greedy_generate_batch(model, [r['prompt'] for r in part], max_new_tokens=4)
            for r, answer in zip(part, answers):
                outcomes.append({'id': r['id'], 'prompt': r['prompt'], 'expected': r['answer'], 'generated': answer.text, 'token_ids': list(answer.token_ids), 'ended': answer.ended, 'correct': answer.matches(r['answer']), 'carry': (r['a'] % 10 + r['b'] % 10 >= 10) or ((r['a'] // 10) % 10 + (r['b'] // 10) % 10 + (r['a'] % 10 + r['b'] % 10 >= 10) >= 10), 'digits': len(r['answer'])})
    def summary(subset):
        return {'examples': len(subset), 'correct': sum(r['correct'] for r in subset), 'terminated': sum(r['ended'] for r in subset)}
    return {**summary(outcomes), 'slices': {**{f'carry_{flag}': summary([r for r in outcomes if r['carry'] == flag]) for flag in (True, False)}, **{f'digits_{n}': summary([r for r in outcomes if r['digits'] == n]) for n in (1, 2, 3)}}, 'mistakes': [r for r in outcomes if not r['correct']][:12]}

@torch.inference_mode()
def collected_logits(model, values):
    batch = data(values)['batch']
    selected_logits, selected_targets = [], []
    device = next(model.parameters()).device
    for start in range(0, len(values), 256):
        part = values[start:start + 256]
        b = {k: v[start:start + 256] for k, v in batch.items() if isinstance(v, torch.Tensor)}
        mask = completion_mask(b, part, device)
        logits = model(b['input_ids'].to(device))
        selected_logits.append(logits[mask].cpu().double())
        selected_targets.append(b['target_ids'].to(device)[mask].cpu())
    return torch.cat(selected_logits), torch.cat(selected_targets)

def confidence_metrics(logits, targets, temperature):
    p = (logits / temperature).softmax(-1)
    confidence, choice = p.max(-1)
    correct = choice == targets
    brier = (p.square().sum(-1) - 2 * p.gather(-1, targets[:, None]).squeeze(-1) + 1).mean()
    bins = []
    for i in range(10):
        selected = (confidence >= i / 10) & ((confidence < (i + 1) / 10) if i < 9 else (confidence <= 1))
        count = int(selected.sum())
        bins.append({'lower': i / 10, 'upper': (i + 1) / 10, 'count': count, 'confidence': float(confidence[selected].mean()) if count else None, 'accuracy': float(correct[selected].double().mean()) if count else None})
    return {'tokens': len(targets), 'nll': float(F.cross_entropy(logits / temperature, targets)), 'brier': float(brier), 'bins': bins}

if args.action in ['evaluate', 'evaluate-rl']:
    if args.action == 'evaluate':
        if (out / 'evaluation.json').exists(): raise RuntimeError('Frozen evaluation already exists. Do not overwrite it.')
        policies = [('base', base_path, False), ('first', args.source / 'runs/pretrain-v1/best.pt', False), ('finetuned', run / 'best.pt', True)]
        reports, fits = {}, {}
    else:
        policies = [('rl', Path('scripts/tiny-llm/runs/rlhf/policy.pt'), True)]
        reports = json.loads((out / 'evaluation.json').read_text())['policies']
        fits = json.loads((out / 'calibration.json').read_text())
        if 'rl' in reports: raise RuntimeError('RL frozen evaluation already exists. Do not overwrite it.')
    for stage, checkpoint, formatted in policies:
        model, _ = load_checkpoint(checkpoint, device=args.device)
        calibrated_logits, calibrated_targets = collected_logits(model, rows('calibration', formatted))
        calibrated_logits = calibrated_logits.clone()
        calibrated_targets = calibrated_targets.clone()
        log_t = torch.zeros((), dtype=torch.float64, requires_grad=True)
        optimizer = torch.optim.Adam([log_t], lr=0.03)
        best_nll, best_t = math.inf, 1.0
        for iteration in range(200):
            optimizer.zero_grad()
            loss = F.cross_entropy(calibrated_logits / log_t.exp(), calibrated_targets)
            if float(loss.detach()) < best_nll: best_nll, best_t = float(loss.detach()), float(log_t.detach().exp())
            loss.backward(); optimizer.step()
            with torch.no_grad(): log_t.clamp_(math.log(0.05), math.log(10))
        test_logits, test_targets = collected_logits(model, rows('test', formatted))
        fits[stage] = {'temperature': best_t, 'checkpoint_sha256': sha(checkpoint), 'fit_split': 'calibration', 'fit_examples': 496, 'fit_target': 'completion tokens and EOS', 'fit_nll': best_nll, 'held_out_split': 'test', 'raw': confidence_metrics(test_logits, test_targets, 1), 'fitted': confidence_metrics(test_logits, test_targets, best_t)}
        reports[stage] = {'checkpoint_sha256': sha(checkpoint), 'formatted': formatted, 'test': generated_metrics(model, rows('test', formatted)), 'ood': generated_metrics(model, rows('ood', formatted))}
        if formatted: reports[stage]['canonical_test_regression'] = generated_metrics(model, rows('test'))
        print(stage, 'test', reports[stage]['test']['correct'], 'ood', reports[stage]['ood']['correct'], 'temperature', best_t, flush=True)
    write(out / 'evaluation.json', {'protocol': 'four output tokens including EOS; frozen selected policies', 'policies': reports, 'data_sha256': {split: sha(args.source / f'data/addition-v1/{split}.jsonl') for split in ['train', 'validation', 'calibration', 'test', 'ood']}})
    write(out / 'calibration.json', fits)

if args.action == 'candidates':
    model, _ = load_checkpoint(run / 'best.pt')
    torch.manual_seed(83)
    rng = random.Random(83)
    groups = {r['group_id']: r for r in rows('train', True)}
    values = list(groups.values()); rng.shuffle(values)
    pairs = []
    @torch.inference_mode()
    def sample(prompt, temperature):
        ids = tokenizer.encode(prompt, add_bos=True); generated = []
        for _ in range(min(16, 32 - len(ids))):
            p = (model(torch.tensor([ids]))[0, -1] / temperature).softmax(-1)
            selected = int(torch.multinomial(p, 1)); ids.append(selected); generated.append(selected)
            if selected == 2: break
        return generated
    def candidate(ids):
        ended = bool(ids and ids[-1] == 2)
        return {'token_ids': ids, 'text': tokenizer.decode(ids[:-1] if ended else ids), 'ended': ended}
    for r in values:
        a = list(generate_answer(model, r['prompt'], max_new_tokens=min(16, 31 - len(r['prompt']))).token_ids)
        b = sample(r['prompt'], 2)
        temperature, attempts = 2, 1
        while a == b and attempts <= 20:
            b = sample(r['prompt'], 3); temperature = 3; attempts += 1
        if a == b: continue
        split_number = int(hashlib.sha256(('preference-v1' + r['group_id']).encode()).hexdigest()[:8], 16) % 100
        split = 'train' if split_number < 70 else 'validation' if split_number < 85 else 'test'
        if rng.random() < 0.5: a, b = b, a; greedy_side = 'b'
        else: greedy_side = 'a'
        pairs.append({'id': 'preference-' + str(len(pairs) + 1).zfill(3), 'group_id': r['group_id'], 'split': split, 'prompt': r['prompt'], 'a': candidate(a), 'b': candidate(b), 'greedy_side': greedy_side, 'sample_temperature': temperature, 'attempts': attempts})
        if len(pairs) >= 120: break
    write(out / 'preferences.json', {'checkpoint_sha256': sha(run / 'best.pt'), 'seed': 83, 'policy': 'actual fine-tuned model; greedy and all-token sampling', 'pairs': pairs, 'labels_collected': 0})
    print('Generated', len(pairs), 'actual preference pairs', flush=True)
Download experiments.py
Technical details

Supervised fine-tuning uses completion-only cross entropy, seed 73 and a fixed forty-epoch budget. Selection uses formatted validation exact count, then answer NLL. Original-format validation is logged as a regression measurement.

Code and training steps