Publish measured results, curves, and pinned training source

70a25ae verified about 1 month ago

19.1 kB

"""Instruction tuning with assistant-only loss and the same chat format at inference.
    python sft.py --prepare
    python sft.py --check
    python sft.py --checkpoint checkpoints/main/ckpt.pt
    python sft.py --checkpoint checkpoints/sft/ckpt.pt --evaluate
"""

import argparse
from collections import Counter
import hashlib
import json
import math
import os
import re
import subprocess
import time

import numpy as np
import torch

import tokenizer
from data import stories
from eval import load
from inference import FIELDS, TEMPLATE, instruction, prompt
from model import Sprout
from train import (BETAS, DEVICE, GRAD_CLIP, SEED, WEIGHT_DECAY, autocast, fingerprint, lr_at,
                   optimizer, split, token_bytes, val_metrics)

REPO = "roneneldan/TinyStoriesInstruct"
REVISION = "ee050ed1f8720795be342921335e821856a2b42e"
FIELD = re.compile(r"^(Summary|Words|Random sentence|Features):\s*(.*)$")


def digest(path):
    with open(path, "rb") as f:
        return hashlib.file_digest(f, "sha256").hexdigest()


def parse(text):
    """Headers may occur before Story: or after the story's final paragraph."""
    fields, body, in_story = {}, [], False
    for line in text.splitlines():
        match = FIELD.fullmatch(line)
        if match:
            key, value = match.groups()
            if key in fields:
                raise ValueError(f"duplicate field: {key}")
            fields[key] = value.strip()
        elif line.startswith("Story:"):
            if in_story:
                raise ValueError("duplicate Story header")
            in_story = True
            body.append(line[len("Story:"):].lstrip())
        elif in_story:
            body.append(line)
        elif line.strip():
            raise ValueError("text before Story header")
    story = "\n".join(body).strip()
    if not story or not any(fields.values()):
        raise ValueError("missing story or instruction")
    if any(s in story or any(s in v for v in fields.values()) for s in tokenizer.SPECIALS):
        raise ValueError("reserved token in content")
    return fields, story


def encode(tok, fields, story):
    prefix = prompt(tok, fields)
    ids = prefix + tok.encode(story) + [tok.special["<|end|>"]]
    assert tok.decode(ids) == TEMPLATE.format(instruction=instruction(fields)) + story + "<|end|>"
    return ids, len(prefix)


def story_hash(text):
    # Collapse whitespace only for deduplication, never for training text.
    return hashlib.sha256(" ".join(text.split()).encode()).digest()


def format_hash(data):
    # The binary index carries its encoding contract, so cached tokens cannot
    # silently be used with another tokenizer or a changed conversation format.
    contract = [digest(f"{data}/tokenizer.model"), TEMPLATE, FIELDS, REVISION]
    return hashlib.sha256(json.dumps(contract).encode()).digest()


def prepare(args):
    from huggingface_hub import hf_hub_download
    os.makedirs(args.instruct, exist_ok=True)
    tok = tokenizer.load(f"{args.data}/tokenizer.model")
    excluded = set()
    with open(f"{args.data}/val.jsonl") as f:
        for line in f:
            excluded.add(story_hash(json.loads(line)["text"]))
    # Process validation first so its stories cannot also enter training.
    seen = set()
    for name in ("val", "train"):
        if args.prepare_split and name != args.prepare_split:
            continue
        source = "valid" if name == "val" else "train"
        raw = hf_hub_download(REPO, f"TinyStories-Instruct-{source}.txt",
                               repo_type="dataset", revision=REVISION,
                               local_dir=args.instruct)
        counts, offset = Counter(), 0
        path = f"{args.instruct}/{name}"
        with open(path + ".bin.tmp", "wb") as tokens, open(path + ".idx.tmp", "wb") as index:
            index.write(format_hash(args.data))
            for text in stories(raw):
                counts["source"] += 1
                try:
                    fields, story = parse(text)
                except ValueError:
                    counts["malformed"] += 1
                    continue
                key = story_hash(story)
                if key in excluded or key in seen:
                    counts["duplicate_or_base_val"] += 1
                    continue
                seen.add(key)
                ids, prefix = encode(tok, fields, story)
                if len(ids) - 1 > args.ctx:
                    counts["over_context"] += 1
                    continue
                np.asarray(ids, dtype=np.uint16).tofile(tokens)
                np.asarray([offset, len(ids), prefix], dtype=np.uint64).tofile(index)
                offset += len(ids)
                counts["kept"] += 1
                counts["assistant_tokens"] += len(ids) - prefix
                if counts["source"] % 10000 == 0:
                    print(name, dict(counts), flush=True)
        assert counts["kept"], counts
        os.replace(path + ".bin.tmp", path + ".bin")
        os.replace(path + ".idx.tmp", path + ".idx")
        print(name, dict(counts), flush=True)


def dataset(root, name, data="data"):
    with open(f"{root}/{name}.idx", "rb") as f:
        assert f.read(32) == format_hash(data), "prepared data encoding differs; run --prepare"
    ids = np.memmap(f"{root}/{name}.bin", mode="r", dtype=np.uint16)
    index = np.memmap(f"{root}/{name}.idx", mode="r", dtype=np.uint64, offset=32).reshape(-1, 3)
    assert len(index) and int(index[-1, 0] + index[-1, 1]) == len(ids)
    return ids, index


def collate(data, indices, pad, device):
    ids, rows = data
    selected = rows[indices]
    width = int(selected[:, 1].max()) - 1
    x = torch.full((len(indices), width), pad, dtype=torch.long)
    y = torch.full_like(x, -100)
    for i, (offset, length, prefix) in enumerate(selected.tolist()):
        assert 0 < prefix < length
        seq = torch.from_numpy(ids[offset:offset + length].astype(np.int64))
        x[i, :length - 1] = seq[:-1]
        # Logit at prefix-1 predicts the FIRST assistant token. The last
        # non-padding label is the assistant end token.
        y[i, prefix - 1:length - 1] = seq[prefix:]
    return x.to(device), y.to(device)


@torch.no_grad()
def validation(model, data, args, pad):
    was_training = model.training
    model.eval()
    nats = targets = 0
    for start in range(0, len(data[1]), args.micro_batch):
        indices = np.arange(start, min(start + args.micro_batch, len(data[1])))
        x, y = collate(data, indices, pad, args.device)
        n = int((y != -100).sum())
        with autocast(args.device):
            loss = model(x, y)[1]
        nats += loss.item() * n
        targets += n
    model.train(was_training)
    return nats / targets


def save(path, model, opt, cfg, fp, step, best, run_id):
    state = dict(model=model.state_dict(), optimizer=opt.state_dict(), config=cfg["model"],
                 fingerprint=fp, step=step, sft=cfg, best_val=best, wandb_id=run_id,
                 torch_rng=torch.get_rng_state(),
                 cuda_rng=torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
                 mps_rng=torch.mps.get_rng_state() if next(model.parameters()).device.type == "mps" else None)
    torch.save(state, path + ".tmp")
    os.replace(path + ".tmp", path)


def upload(path, hub):
    if hub:
        from huggingface_hub import HfApi
        HfApi().upload_file(path_or_fileobj=path, path_in_repo="sft/" + os.path.basename(path),
                           repo_id=hub, commit_message="Instruction tuning checkpoint")


def train(args):
    assert args.epochs in (1, 2, 3) and 0 < args.lr <= 5e-5
    assert args.micro_batch > 0 and args.accum > 0 and args.eval_every > 0
    torch.manual_seed(SEED)
    if args.device == "mps":
        torch.mps.manual_seed(SEED)
    model, tok, model_cfg, _ = load(args.checkpoint, args.data, args.device)
    train_data = dataset(args.instruct, "train", args.data)
    val_data = dataset(args.instruct, "val", args.data)
    assert max(train_data[1][:, 1].max(), val_data[1][:, 1].max()) - 1 <= model.ctx
    pad = tok.special["<|endoftext|>"]
    size = args.micro_batch * args.accum
    per_epoch = math.ceil(len(train_data[1]) / size)
    steps = per_epoch * args.epochs
    cfg = dict(model=model_cfg, template=TEMPLATE, fields=FIELDS, revision=REVISION,
               checkpoint_sha256=digest(args.checkpoint), seed=SEED, epochs=args.epochs,
               micro_batch=args.micro_batch, accum=args.accum, lr=args.lr,
               weight_decay=WEIGHT_DECAY, betas=BETAS, grad_clip=GRAD_CLIP,
               eval_every=args.eval_every, device=args.device,
               warmup=min(100, max(1, steps // 20)), steps=steps,
               data={f"{name}.{ext}": digest(f"{args.instruct}/{name}.{ext}")
                     for name in ("train", "val") for ext in ("bin", "idx")})
    fp = fingerprint(args.data)
    opt = optimizer(model, args.device)
    os.makedirs(args.out, exist_ok=True)
    path = f"{args.out}/ckpt.pt"
    step, best, run_id = 0, float("inf"), None
    if args.resume:
        state = torch.load(path, map_location="cpu", weights_only=False)
        assert state["sft"] == cfg and state["fingerprint"] == fp
        model.load_state_dict(state["model"])
        opt.load_state_dict(state["optimizer"])
        step, best, run_id = state["step"], state["best_val"], state["wandb_id"]
        torch.set_rng_state(state["torch_rng"])
        if state["cuda_rng"] is not None:
            assert args.device == "cuda", "resume on the original device"
            torch.cuda.set_rng_state_all(state["cuda_rng"])
        if state["mps_rng"] is not None:
            assert args.device == "mps", "resume on the original device"
            torch.mps.set_rng_state(state["mps_rng"])
    wandb = None
    if not args.no_wandb:
        import wandb
        commit = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip()
        run = wandb.init(project="sprout", name=os.path.basename(args.out), id=run_id,
                         resume="must" if run_id else None, config=cfg | {"commit": commit})
        run_id = run.id
    print(json.dumps(cfg), flush=True)
    model.train()
    order, epoch = None, -1
    stop = min(steps, args.stop_after or steps)
    with open(f"{args.out}/losses.txt", "a" if args.resume else "w") as log:
        if step == 0:
            initial = validation(model, val_data, args, pad)
            print(f"step 0 sft_val_loss {initial:.6f}", flush=True)
            print(json.dumps(dict(step=0, sft_val_loss=initial)), file=log, flush=True)
            if wandb:
                wandb.log({"sft_val_loss": initial}, step=0)
        while step < stop:
            current_epoch, batch_id = divmod(step, per_epoch)
            if current_epoch != epoch:
                epoch = current_epoch
                order = torch.randperm(len(train_data[1]), generator=torch.Generator().manual_seed(SEED + epoch)).numpy()
            rows = order[batch_id * size:(batch_id + 1) * size]
            # Weight accumulation by supervised tokens, not padded batch size.
            target_count = int((train_data[1][rows, 1] - train_data[1][rows, 2]).sum())
            lr = lr_at(step, steps, cfg["warmup"], args.lr)
            for group in opt.param_groups:
                group["lr"] = lr
            t0, total = time.time(), 0.0
            opt.zero_grad(set_to_none=True)
            for start in range(0, len(rows), args.micro_batch):
                x, y = collate(train_data, rows[start:start + args.micro_batch], pad, args.device)
                weight = int((y != -100).sum()) / target_count
                with autocast(args.device):
                    loss = model(x, y)[1]
                (loss * weight).backward()
                total += loss.item() * weight
            grad = torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP, error_if_nonfinite=True)
            opt.step()
            step += 1
            metrics = dict(step=step, loss=total, lr=lr, grad_norm=float(grad),
                           assistant_tokens_per_sec=target_count / (time.time() - t0))
            print(json.dumps(metrics), flush=True)
            print(json.dumps(metrics), file=log, flush=True)
            if wandb:
                wandb.log(metrics, step=step)
            if step % args.eval_every == 0 or step % per_epoch == 0 or step == stop:
                val = validation(model, val_data, args, pad)
                print(json.dumps(dict(step=step, sft_val_loss=val)), file=log, flush=True)
                print(f"step {step} sft_val_loss {val:.6f}", flush=True)
                if wandb:
                    wandb.log({"sft_val_loss": val}, step=step)
                if val < best:
                    best = val
                    save(f"{args.out}/best.pt", model, opt, cfg, fp, step, best, run_id)
                    upload(f"{args.out}/best.pt", args.hub)
                save(path, model, opt, cfg, fp, step, best, run_id)
                upload(path, args.hub)
    if wandb:
        wandb.finish()


@torch.inference_mode()
def evaluate(args):
    model, tok, cfg, step = load(args.checkpoint, args.data, args.device)
    args.ctx = cfg["ctx"]
    loss, bpb = val_metrics(model, split(args.data, "val"), token_bytes(tok), args)
    print(json.dumps(dict(checkpoint=args.checkpoint, sha256=digest(args.checkpoint),
                          step=step, device=args.device, torch=torch.__version__,
                          dtype="bfloat16" if args.device == "cuda" else "float32",
                          template=TEMPLATE, revision=REVISION, micro_batch=args.micro_batch,
                          base_val_loss=loss, base_val_bpb=bpb)), flush=True)
    raw = f"{args.instruct}/TinyStories-Instruct-valid.txt"
    records = []
    for text in stories(raw):
        try:
            fields, story = parse(text)
        except ValueError:
            continue
        if len(encode(tok, fields, story)[0]) - 1 <= model.ctx and fields.get("Words") and fields.get("Random sentence"):
            records.append(fields)
        if len(records) == 20:
            break
    assert len(records) == 20, "need 20 held-out word-and-sentence prompts"
    totals = Counter()
    for i, fields in enumerate(records):
        torch.manual_seed(SEED + i)
        if args.device == "mps":
            torch.mps.manual_seed(SEED + i)
        prefix = prompt(tok, fields)
        # Stay inside training context so the instruction remains visible.
        limit = model.ctx - len(prefix) + 1
        end = tok.special["<|end|>"]
        x = torch.tensor([prefix], device=args.device)
        with autocast(args.device):
            ids = model.generate(x, limit, temperature=0.8, stop=end, top_k=50)[0, len(prefix):].tolist()
        ended = bool(ids and ids[-1] == end)
        result = tok.decode(ids[:-1] if ended else ids)
        words = [w.strip() for w in fields["Words"].split(",")]
        included = [bool(re.search(r"(?<!\w)" + re.escape(w) + r"(?!\w)", result, re.I)) for w in words]
        sentence = fields["Random sentence"] in result
        totals.update(dict(prompts=1, all_words=int(all(included)), sentence=int(sentence), ended=int(ended)))
        print(json.dumps(dict(prompt_id=i + 1, fields=fields, continuation=result,
                              token_ids=ids, seed=SEED + i, temperature=0.8, top_k=50,
                              max_tokens=limit, use_cache=True, words_included=included,
                              sentence_included=sentence, ended=ended), ensure_ascii=False), flush=True)
    print(json.dumps(dict(behavior=dict(totals))), flush=True)


def check(args):
    torch.manual_seed(SEED)
    torch.set_num_threads(2)
    tok = tokenizer.load(f"{args.data}/tokenizer.model")
    fields = {"Summary": "A café cat smiles.", "Words": "cat, happy", "Random sentence": "The cat is happy."}
    pairs = [(fields, "The cat is happy."), ({"Words": "hi"}, "Hi!")]
    flat, rows = [], []
    for f, story in pairs:
        ids, prefix = encode(tok, f, story)
        assert ids[:prefix] == prompt(tok, f)
        rows.append([len(flat), len(ids), prefix])
        flat.extend(ids)
    data = np.array(flat, dtype=np.uint16), np.array(rows, dtype=np.uint64)
    x, y = collate(data, np.arange(2), tok.special["<|endoftext|>"], "cpu")
    for i, (offset, length, prefix) in enumerate(rows):
        assert (y[i, :prefix - 1] == -100).all()
        assert y[i, prefix - 1].item() == flat[offset + prefix]
        assert y[i, length - 2].item() == tok.special["<|end|>"]
        assert (y[i, length - 1:] == -100).all()
        assert int((y[i] != -100).sum()) == length - prefix
    model = Sprout(n_layer=1, dim=32, n_head=4, ctx=x.shape[1], vocab=len(tok.vocab))
    logits, loss = model(x, y)
    logits.retain_grad()
    manual = torch.nn.functional.cross_entropy(logits[y != -100], y[y != -100])
    assert torch.allclose(loss, manual)
    manual.backward()
    assert (logits.grad[y == -100] == 0).all()
    full_grad = [p.grad.clone() for p in model.parameters()]
    model.zero_grad(set_to_none=True)
    for i in range(2):
        xx, yy = collate(data, [i], tok.special["<|endoftext|>"], "cpu")
        (model(xx, yy)[1] * (yy != -100).sum() / (y != -100).sum()).backward()
    error = max((p.grad - g).abs().max().item() for p, g in zip(model.parameters(), full_grad))
    assert error < 1e-6, error
    parsed, story = parse("Words: cat, happy\nStory: \nThe cat is happy.\nSummary: A happy cat.\n")
    assert story == "The cat is happy." and parsed["Summary"] == "A happy cat."
    print(f"template round-trip, first assistant label, end label, padding and prompt gradients passed; accumulation max error {error:.2e}")
    print(f"supervised tokens: {(y != -100).sum().item()}; ignored labels: {(y == -100).sum().item()}")


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--data", default="data")
    p.add_argument("--instruct", default="data/instruct")
    p.add_argument("--checkpoint", default="checkpoints/main/ckpt.pt")
    p.add_argument("--out", default="checkpoints/sft")
    p.add_argument("--device", choices=("cpu", "mps", "cuda"), default=DEVICE)
    p.add_argument("--ctx", type=int, default=512)
    p.add_argument("--epochs", type=int, default=1)
    p.add_argument("--lr", type=float, default=3e-5)
    p.add_argument("--micro-batch", type=int, default=16)
    p.add_argument("--accum", type=int, default=8)
    p.add_argument("--eval-every", type=int, default=500)
    p.add_argument("--stop-after", type=int)
    p.add_argument("--hub", default="")
    p.add_argument("--no-wandb", action="store_true")
    p.add_argument("--resume", action="store_true")
    p.add_argument("--prepare", action="store_true")
    p.add_argument("--prepare-split", choices=("val",))
    p.add_argument("--check", action="store_true")
    p.add_argument("--evaluate", action="store_true")
    args = p.parse_args()
    if args.prepare:
        prepare(args)
    elif args.check:
        check(args)
    elif args.evaluate:
        evaluate(args)
    else:
        train(args)


if __name__ == "__main__":
    main()

	"""Instruction tuning with assistant-only loss and the same chat format at inference.

	python sft.py --prepare
	python sft.py --check
	python sft.py --checkpoint checkpoints/main/ckpt.pt
	python sft.py --checkpoint checkpoints/sft/ckpt.pt --evaluate
	"""

	import argparse
	from collections import Counter
	import hashlib
	import json
	import math
	import os
	import re
	import subprocess
	import time

	import numpy as np
	import torch

	import tokenizer
	from data import stories
	from eval import load
	from inference import FIELDS, TEMPLATE, instruction, prompt
	from model import Sprout
	from train import (BETAS, DEVICE, GRAD_CLIP, SEED, WEIGHT_DECAY, autocast, fingerprint, lr_at,
	optimizer, split, token_bytes, val_metrics)

	REPO = "roneneldan/TinyStoriesInstruct"
	REVISION = "ee050ed1f8720795be342921335e821856a2b42e"
	FIELD = re.compile(r"^(Summary\|Words\|Random sentence\|Features):\s(.)$")


	def digest(path):
	with open(path, "rb") as f:
	return hashlib.file_digest(f, "sha256").hexdigest()


	def parse(text):
	"""Headers may occur before Story: or after the story's final paragraph."""
	fields, body, in_story = {}, [], False
	for line in text.splitlines():
	match = FIELD.fullmatch(line)
	if match:
	key, value = match.groups()
	if key in fields:
	raise ValueError(f"duplicate field: {key}")
	fields[key] = value.strip()
	elif line.startswith("Story:"):
	if in_story:
	raise ValueError("duplicate Story header")
	in_story = True
	body.append(line[len("Story:"):].lstrip())
	elif in_story:
	body.append(line)
	elif line.strip():
	raise ValueError("text before Story header")
	story = "\n".join(body).strip()
	if not story or not any(fields.values()):
	raise ValueError("missing story or instruction")
	if any(s in story or any(s in v for v in fields.values()) for s in tokenizer.SPECIALS):
	raise ValueError("reserved token in content")
	return fields, story


	def encode(tok, fields, story):
	prefix = prompt(tok, fields)
	ids = prefix + tok.encode(story) + [tok.special["<\|end\|>"]]
	assert tok.decode(ids) == TEMPLATE.format(instruction=instruction(fields)) + story + "<\|end\|>"
	return ids, len(prefix)


	def story_hash(text):
	# Collapse whitespace only for deduplication, never for training text.
	return hashlib.sha256(" ".join(text.split()).encode()).digest()


	def format_hash(data):
	# The binary index carries its encoding contract, so cached tokens cannot
	# silently be used with another tokenizer or a changed conversation format.
	contract = [digest(f"{data}/tokenizer.model"), TEMPLATE, FIELDS, REVISION]
	return hashlib.sha256(json.dumps(contract).encode()).digest()


	def prepare(args):
	from huggingface_hub import hf_hub_download
	os.makedirs(args.instruct, exist_ok=True)
	tok = tokenizer.load(f"{args.data}/tokenizer.model")
	excluded = set()
	with open(f"{args.data}/val.jsonl") as f:
	for line in f:
	excluded.add(story_hash(json.loads(line)["text"]))
	# Process validation first so its stories cannot also enter training.
	seen = set()
	for name in ("val", "train"):
	if args.prepare_split and name != args.prepare_split:
	continue
	source = "valid" if name == "val" else "train"
	raw = hf_hub_download(REPO, f"TinyStories-Instruct-{source}.txt",
	repo_type="dataset", revision=REVISION,
	local_dir=args.instruct)
	counts, offset = Counter(), 0
	path = f"{args.instruct}/{name}"
	with open(path + ".bin.tmp", "wb") as tokens, open(path + ".idx.tmp", "wb") as index:
	index.write(format_hash(args.data))
	for text in stories(raw):
	counts["source"] += 1
	try:
	fields, story = parse(text)
	except ValueError:
	counts["malformed"] += 1
	continue
	key = story_hash(story)
	if key in excluded or key in seen:
	counts["duplicate_or_base_val"] += 1
	continue
	seen.add(key)
	ids, prefix = encode(tok, fields, story)
	if len(ids) - 1 > args.ctx:
	counts["over_context"] += 1
	continue
	np.asarray(ids, dtype=np.uint16).tofile(tokens)
	np.asarray([offset, len(ids), prefix], dtype=np.uint64).tofile(index)
	offset += len(ids)
	counts["kept"] += 1
	counts["assistant_tokens"] += len(ids) - prefix
	if counts["source"] % 10000 == 0:
	print(name, dict(counts), flush=True)
	assert counts["kept"], counts
	os.replace(path + ".bin.tmp", path + ".bin")
	os.replace(path + ".idx.tmp", path + ".idx")
	print(name, dict(counts), flush=True)


	def dataset(root, name, data="data"):
	with open(f"{root}/{name}.idx", "rb") as f:
	assert f.read(32) == format_hash(data), "prepared data encoding differs; run --prepare"
	ids = np.memmap(f"{root}/{name}.bin", mode="r", dtype=np.uint16)
	index = np.memmap(f"{root}/{name}.idx", mode="r", dtype=np.uint64, offset=32).reshape(-1, 3)
	assert len(index) and int(index[-1, 0] + index[-1, 1]) == len(ids)
	return ids, index


	def collate(data, indices, pad, device):
	ids, rows = data
	selected = rows[indices]
	width = int(selected[:, 1].max()) - 1
	x = torch.full((len(indices), width), pad, dtype=torch.long)
	y = torch.full_like(x, -100)
	for i, (offset, length, prefix) in enumerate(selected.tolist()):
	assert 0 < prefix < length
	seq = torch.from_numpy(ids[offset:offset + length].astype(np.int64))
	x[i, :length - 1] = seq[:-1]
	# Logit at prefix-1 predicts the FIRST assistant token. The last
	# non-padding label is the assistant end token.
	y[i, prefix - 1:length - 1] = seq[prefix:]
	return x.to(device), y.to(device)


	@torch.no_grad()
	def validation(model, data, args, pad):
	was_training = model.training
	model.eval()
	nats = targets = 0
	for start in range(0, len(data[1]), args.micro_batch):
	indices = np.arange(start, min(start + args.micro_batch, len(data[1])))
	x, y = collate(data, indices, pad, args.device)
	n = int((y != -100).sum())
	with autocast(args.device):
	loss = model(x, y)[1]
	nats += loss.item() * n
	targets += n
	model.train(was_training)
	return nats / targets


	def save(path, model, opt, cfg, fp, step, best, run_id):
	state = dict(model=model.state_dict(), optimizer=opt.state_dict(), config=cfg["model"],
	fingerprint=fp, step=step, sft=cfg, best_val=best, wandb_id=run_id,
	torch_rng=torch.get_rng_state(),
	cuda_rng=torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
	mps_rng=torch.mps.get_rng_state() if next(model.parameters()).device.type == "mps" else None)
	torch.save(state, path + ".tmp")
	os.replace(path + ".tmp", path)


	def upload(path, hub):
	if hub:
	from huggingface_hub import HfApi
	HfApi().upload_file(path_or_fileobj=path, path_in_repo="sft/" + os.path.basename(path),
	repo_id=hub, commit_message="Instruction tuning checkpoint")


	def train(args):
	assert args.epochs in (1, 2, 3) and 0 < args.lr <= 5e-5
	assert args.micro_batch > 0 and args.accum > 0 and args.eval_every > 0
	torch.manual_seed(SEED)
	if args.device == "mps":
	torch.mps.manual_seed(SEED)
	model, tok, model_cfg, _ = load(args.checkpoint, args.data, args.device)
	train_data = dataset(args.instruct, "train", args.data)
	val_data = dataset(args.instruct, "val", args.data)
	assert max(train_data[1][:, 1].max(), val_data[1][:, 1].max()) - 1 <= model.ctx
	pad = tok.special["<\|endoftext\|>"]
	size = args.micro_batch * args.accum
	per_epoch = math.ceil(len(train_data[1]) / size)
	steps = per_epoch * args.epochs
	cfg = dict(model=model_cfg, template=TEMPLATE, fields=FIELDS, revision=REVISION,
	checkpoint_sha256=digest(args.checkpoint), seed=SEED, epochs=args.epochs,
	micro_batch=args.micro_batch, accum=args.accum, lr=args.lr,
	weight_decay=WEIGHT_DECAY, betas=BETAS, grad_clip=GRAD_CLIP,
	eval_every=args.eval_every, device=args.device,
	warmup=min(100, max(1, steps // 20)), steps=steps,
	data={f"{name}.{ext}": digest(f"{args.instruct}/{name}.{ext}")
	for name in ("train", "val") for ext in ("bin", "idx")})
	fp = fingerprint(args.data)
	opt = optimizer(model, args.device)
	os.makedirs(args.out, exist_ok=True)
	path = f"{args.out}/ckpt.pt"
	step, best, run_id = 0, float("inf"), None
	if args.resume:
	state = torch.load(path, map_location="cpu", weights_only=False)
	assert state["sft"] == cfg and state["fingerprint"] == fp
	model.load_state_dict(state["model"])
	opt.load_state_dict(state["optimizer"])
	step, best, run_id = state["step"], state["best_val"], state["wandb_id"]
	torch.set_rng_state(state["torch_rng"])
	if state["cuda_rng"] is not None:
	assert args.device == "cuda", "resume on the original device"
	torch.cuda.set_rng_state_all(state["cuda_rng"])
	if state["mps_rng"] is not None:
	assert args.device == "mps", "resume on the original device"
	torch.mps.set_rng_state(state["mps_rng"])
	wandb = None
	if not args.no_wandb:
	import wandb
	commit = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip()
	run = wandb.init(project="sprout", name=os.path.basename(args.out), id=run_id,
	resume="must" if run_id else None, config=cfg \| {"commit": commit})
	run_id = run.id
	print(json.dumps(cfg), flush=True)
	model.train()
	order, epoch = None, -1
	stop = min(steps, args.stop_after or steps)
	with open(f"{args.out}/losses.txt", "a" if args.resume else "w") as log:
	if step == 0:
	initial = validation(model, val_data, args, pad)
	print(f"step 0 sft_val_loss {initial:.6f}", flush=True)
	print(json.dumps(dict(step=0, sft_val_loss=initial)), file=log, flush=True)
	if wandb:
	wandb.log({"sft_val_loss": initial}, step=0)
	while step < stop:
	current_epoch, batch_id = divmod(step, per_epoch)
	if current_epoch != epoch:
	epoch = current_epoch
	order = torch.randperm(len(train_data[1]), generator=torch.Generator().manual_seed(SEED + epoch)).numpy()
	rows = order[batch_id * size:(batch_id + 1) * size]
	# Weight accumulation by supervised tokens, not padded batch size.
	target_count = int((train_data[1][rows, 1] - train_data[1][rows, 2]).sum())
	lr = lr_at(step, steps, cfg["warmup"], args.lr)
	for group in opt.param_groups:
	group["lr"] = lr
	t0, total = time.time(), 0.0
	opt.zero_grad(set_to_none=True)
	for start in range(0, len(rows), args.micro_batch):
	x, y = collate(train_data, rows[start:start + args.micro_batch], pad, args.device)
	weight = int((y != -100).sum()) / target_count
	with autocast(args.device):
	loss = model(x, y)[1]
	(loss * weight).backward()
	total += loss.item() * weight
	grad = torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP, error_if_nonfinite=True)
	opt.step()
	step += 1
	metrics = dict(step=step, loss=total, lr=lr, grad_norm=float(grad),
	assistant_tokens_per_sec=target_count / (time.time() - t0))
	print(json.dumps(metrics), flush=True)
	print(json.dumps(metrics), file=log, flush=True)
	if wandb:
	wandb.log(metrics, step=step)
	if step % args.eval_every == 0 or step % per_epoch == 0 or step == stop:
	val = validation(model, val_data, args, pad)
	print(json.dumps(dict(step=step, sft_val_loss=val)), file=log, flush=True)
	print(f"step {step} sft_val_loss {val:.6f}", flush=True)
	if wandb:
	wandb.log({"sft_val_loss": val}, step=step)
	if val < best:
	best = val
	save(f"{args.out}/best.pt", model, opt, cfg, fp, step, best, run_id)
	upload(f"{args.out}/best.pt", args.hub)
	save(path, model, opt, cfg, fp, step, best, run_id)
	upload(path, args.hub)
	if wandb:
	wandb.finish()


	@torch.inference_mode()
	def evaluate(args):
	model, tok, cfg, step = load(args.checkpoint, args.data, args.device)
	args.ctx = cfg["ctx"]
	loss, bpb = val_metrics(model, split(args.data, "val"), token_bytes(tok), args)
	print(json.dumps(dict(checkpoint=args.checkpoint, sha256=digest(args.checkpoint),
	step=step, device=args.device, torch=torch.__version__,
	dtype="bfloat16" if args.device == "cuda" else "float32",
	template=TEMPLATE, revision=REVISION, micro_batch=args.micro_batch,
	base_val_loss=loss, base_val_bpb=bpb)), flush=True)
	raw = f"{args.instruct}/TinyStories-Instruct-valid.txt"
	records = []
	for text in stories(raw):
	try:
	fields, story = parse(text)
	except ValueError:
	continue
	if len(encode(tok, fields, story)[0]) - 1 <= model.ctx and fields.get("Words") and fields.get("Random sentence"):
	records.append(fields)
	if len(records) == 20:
	break
	assert len(records) == 20, "need 20 held-out word-and-sentence prompts"
	totals = Counter()
	for i, fields in enumerate(records):
	torch.manual_seed(SEED + i)
	if args.device == "mps":
	torch.mps.manual_seed(SEED + i)
	prefix = prompt(tok, fields)
	# Stay inside training context so the instruction remains visible.
	limit = model.ctx - len(prefix) + 1
	end = tok.special["<\|end\|>"]
	x = torch.tensor([prefix], device=args.device)
	with autocast(args.device):
	ids = model.generate(x, limit, temperature=0.8, stop=end, top_k=50)[0, len(prefix):].tolist()
	ended = bool(ids and ids[-1] == end)
	result = tok.decode(ids[:-1] if ended else ids)
	words = [w.strip() for w in fields["Words"].split(",")]
	included = [bool(re.search(r"(?<!\w)" + re.escape(w) + r"(?!\w)", result, re.I)) for w in words]
	sentence = fields["Random sentence"] in result
	totals.update(dict(prompts=1, all_words=int(all(included)), sentence=int(sentence), ended=int(ended)))
	print(json.dumps(dict(prompt_id=i + 1, fields=fields, continuation=result,
	token_ids=ids, seed=SEED + i, temperature=0.8, top_k=50,
	max_tokens=limit, use_cache=True, words_included=included,
	sentence_included=sentence, ended=ended), ensure_ascii=False), flush=True)
	print(json.dumps(dict(behavior=dict(totals))), flush=True)


	def check(args):
	torch.manual_seed(SEED)
	torch.set_num_threads(2)
	tok = tokenizer.load(f"{args.data}/tokenizer.model")
	fields = {"Summary": "A café cat smiles.", "Words": "cat, happy", "Random sentence": "The cat is happy."}
	pairs = [(fields, "The cat is happy."), ({"Words": "hi"}, "Hi!")]
	flat, rows = [], []
	for f, story in pairs:
	ids, prefix = encode(tok, f, story)
	assert ids[:prefix] == prompt(tok, f)
	rows.append([len(flat), len(ids), prefix])
	flat.extend(ids)
	data = np.array(flat, dtype=np.uint16), np.array(rows, dtype=np.uint64)
	x, y = collate(data, np.arange(2), tok.special["<\|endoftext\|>"], "cpu")
	for i, (offset, length, prefix) in enumerate(rows):
	assert (y[i, :prefix - 1] == -100).all()
	assert y[i, prefix - 1].item() == flat[offset + prefix]
	assert y[i, length - 2].item() == tok.special["<\|end\|>"]
	assert (y[i, length - 1:] == -100).all()
	assert int((y[i] != -100).sum()) == length - prefix
	model = Sprout(n_layer=1, dim=32, n_head=4, ctx=x.shape[1], vocab=len(tok.vocab))
	logits, loss = model(x, y)
	logits.retain_grad()
	manual = torch.nn.functional.cross_entropy(logits[y != -100], y[y != -100])
	assert torch.allclose(loss, manual)
	manual.backward()
	assert (logits.grad[y == -100] == 0).all()
	full_grad = [p.grad.clone() for p in model.parameters()]
	model.zero_grad(set_to_none=True)
	for i in range(2):
	xx, yy = collate(data, [i], tok.special["<\|endoftext\|>"], "cpu")
	(model(xx, yy)[1] * (yy != -100).sum() / (y != -100).sum()).backward()
	error = max((p.grad - g).abs().max().item() for p, g in zip(model.parameters(), full_grad))
	assert error < 1e-6, error
	parsed, story = parse("Words: cat, happy\nStory: \nThe cat is happy.\nSummary: A happy cat.\n")
	assert story == "The cat is happy." and parsed["Summary"] == "A happy cat."
	print(f"template round-trip, first assistant label, end label, padding and prompt gradients passed; accumulation max error {error:.2e}")
	print(f"supervised tokens: {(y != -100).sum().item()}; ignored labels: {(y == -100).sum().item()}")


	def main():
	p = argparse.ArgumentParser(description=__doc__)
	p.add_argument("--data", default="data")
	p.add_argument("--instruct", default="data/instruct")
	p.add_argument("--checkpoint", default="checkpoints/main/ckpt.pt")
	p.add_argument("--out", default="checkpoints/sft")
	p.add_argument("--device", choices=("cpu", "mps", "cuda"), default=DEVICE)
	p.add_argument("--ctx", type=int, default=512)
	p.add_argument("--epochs", type=int, default=1)
	p.add_argument("--lr", type=float, default=3e-5)
	p.add_argument("--micro-batch", type=int, default=16)
	p.add_argument("--accum", type=int, default=8)
	p.add_argument("--eval-every", type=int, default=500)
	p.add_argument("--stop-after", type=int)
	p.add_argument("--hub", default="")
	p.add_argument("--no-wandb", action="store_true")
	p.add_argument("--resume", action="store_true")
	p.add_argument("--prepare", action="store_true")
	p.add_argument("--prepare-split", choices=("val",))
	p.add_argument("--check", action="store_true")
	p.add_argument("--evaluate", action="store_true")
	args = p.parse_args()
	if args.prepare:
	prepare(args)
	elif args.check:
	check(args)
	elif args.evaluate:
	evaluate(args)
	else:
	train(args)


	if __name__ == "__main__":
	main()