download.py
Скачивает сырой корпус с Hugging Face.
36 строк
· Разбираем в главе 9
· Открыть как текст
"""Fetch the raw corpus from Hugging Face into raw/ (about 6.5 GB).
python download.py --out raw
"""
import argparse
import os
import urllib.request
HF = 'https://huggingface.co/datasets'
FILES = {
'TinyStoriesV2-GPT4-train.txt': f'{HF}/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-train.txt',
'TinyStoriesV2-GPT4-valid.txt': f'{HF}/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-valid.txt',
'soda-train.parquet': f'{HF}/allenai/soda/resolve/main/train.parquet',
'everyday-train.parquet': f'{HF}/HuggingFaceTB/everyday-conversations-llama3.1-2k/resolve/main/data/train_sft-00000-of-00001.parquet',
**{f'simplestories-{i}.parquet': f'{HF}/SimpleStories/SimpleStories/resolve/main/data/train-0000{i}-of-00007.parquet'
for i in range(7)},
}
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--out', default='raw')
args = ap.parse_args()
os.makedirs(args.out, exist_ok=True)
for name, url in FILES.items():
path = os.path.join(args.out, name)
if os.path.exists(path):
print('have', name)
continue
print('get ', name, flush=True)
urllib.request.urlretrieve(url, path + '.part')
os.rename(path + '.part', path)
if __name__ == '__main__':
main()
tokenizer.py
Байтовый BPE: обучение, кодирование и декодирование.
132 строк
· Разбираем в главе 6
· Открыть как текст
"""Byte-level BPE: turn text into a list of integers and back.
Training starts from the 256 possible bytes and repeatedly glues together the
pair of neighbouring tokens that occurs most often. Every glued pair becomes a
new token. Encoding replays the same merges in the same order.
"""
import json
from collections import Counter, defaultdict
import regex as re
# Split text into chunks first, so a merge never crosses a word boundary:
# contractions, words with their leading space, single digits, punctuation, spaces.
SPLIT = re.compile(r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
SPECIALS = ['<|endoftext|>', '<|user|>', '<|assistant|>', '<|end|>']
class Tokenizer:
def __init__(self, merges=(), specials=SPECIALS):
self.merges = {tuple(pair): 256 + i for i, pair in enumerate(merges)}
self.vocab = {i: bytes([i]) for i in range(256)}
for (a, b), idx in self.merges.items():
self.vocab[idx] = self.vocab[a] + self.vocab[b]
first_special = 256 + len(self.merges)
self.specials = {s: first_special + i for i, s in enumerate(specials)}
for s, idx in self.specials.items():
self.vocab[idx] = s.encode()
self.special_split = re.compile('(' + '|'.join(re.escape(s) for s in specials) + ')')
self.cache = {}
@property
def vocab_size(self):
return len(self.vocab)
def __getitem__(self, special):
return self.specials[special]
# ---- training -------------------------------------------------------
@classmethod
def train(cls, texts, vocab_size, specials=SPECIALS, verbose=False):
n_merges = vocab_size - 256 - len(specials)
words = Counter()
for text in texts:
for chunk in SPLIT.findall(text):
words[chunk] += 1
# every distinct chunk is a list of token ids, counted once with its frequency
seqs = [list(w.encode()) for w in words]
freqs = list(words.values())
pairs = Counter()
where = defaultdict(set) # pair -> indices of chunks that contain it
for i, seq in enumerate(seqs):
for pair in zip(seq, seq[1:]):
pairs[pair] += freqs[i]
where[pair].add(i)
merges = []
for step in range(n_merges):
if not pairs:
break
best = max(pairs, key=pairs.get)
new_id = 256 + step
merges.append(best)
for i in list(where[best]):
seq, f = seqs[i], freqs[i]
for pair in zip(seq, seq[1:]): # forget this chunk's old pairs
pairs[pair] -= f
if pairs[pair] <= 0:
del pairs[pair]
where[pair].discard(i)
seq = merge(seq, best, new_id)
seqs[i] = seq
for pair in zip(seq, seq[1:]): # and count the new ones
pairs[pair] += f
where[pair].add(i)
if verbose and step % 500 == 0:
tok = cls(merges).vocab[new_id]
print(f'merge {step:5d}: {tok!r}')
return cls(merges, specials)
# ---- encoding -------------------------------------------------------
def encode_chunk(self, chunk):
ids = self.cache.get(chunk)
if ids is not None:
return ids
ids = list(chunk.encode())
while len(ids) > 1:
# the pair that was learned earliest gets merged first
pair = min(zip(ids, ids[1:]), key=lambda p: self.merges.get(p, float('inf')))
if pair not in self.merges:
break
ids = merge(ids, pair, self.merges[pair])
if len(self.cache) < 500_000:
self.cache[chunk] = ids
return ids
def encode(self, text, allow_special=True):
out = []
parts = self.special_split.split(text) if allow_special else [text]
for part in parts:
if part in self.specials and allow_special:
out.append(self.specials[part])
continue
for chunk in SPLIT.findall(part):
out.extend(self.encode_chunk(chunk))
return out
def decode(self, ids):
return b''.join(self.vocab[i] for i in ids).decode('utf-8', errors='replace')
# ---- saving ---------------------------------------------------------
def save(self, path):
with open(path, 'w') as f:
json.dump({'merges': [list(p) for p in self.merges], 'specials': list(self.specials)}, f)
@classmethod
def load(cls, path):
with open(path) as f:
data = json.load(f)
return cls(data['merges'], data['specials'])
def merge(ids, pair, new_id):
"""Replace every occurrence of `pair` in `ids` with `new_id`."""
out, i = [], 0
while i < len(ids):
if i < len(ids) - 1 and ids[i] == pair[0] and ids[i + 1] == pair[1]:
out.append(new_id)
i += 2
else:
out.append(ids[i])
i += 1
return out
prepare.py
Чистит и дедуплицирует тексты, обучает токенизатор и превращает корпус в файлы с номерами токенов.
149 строк
· Разбираем в главе 9
· Открыть как текст
"""Build the pre-training corpus: clean, deduplicate, train the tokenizer, encode.
Sources (all English, all openly licensed):
- TinyStories V2 (GPT-4 part) - short stories in a 3-4 year old's vocabulary
- SimpleStories - simple stories with many topics, themes and styles
- SODA - everyday two-person dialogues with a short narrative
Output: data/train.bin and data/val.bin (uint16 token ids, every document starts
with <|endoftext|>), data/tokenizer.json and data/corpus_stats.json.
"""
import argparse
import glob
import hashlib
import json
import os
import random
import time
from multiprocessing import Pool
import numpy as np
import pyarrow.parquet as pq
from tokenizer import Tokenizer
REPLACE = str.maketrans({'‘': "'", '’': "'", '“': '"', '”': '"', '–': '-',
'—': ' - ', '…': '...', ' ': ' ', '\r': ''})
def clean(text):
"""Normalise quotes and dashes, drop stray spaces; None if the text should go."""
text = text.translate(REPLACE).strip()
lines = [' '.join(line.split()) for line in text.split('\n')]
text = '\n'.join(lines).strip()
while '\n\n\n' in text:
text = text.replace('\n\n\n', '\n\n')
if len(text) < 80:
return None
if sum(ord(c) > 126 for c in text) > 0.002 * len(text): # other scripts, mojibake
return None
return text
def tinystories(raw):
with open(f'{raw}/TinyStoriesV2-GPT4-train.txt', encoding='utf-8') as f:
for story in f.read().split('<|endoftext|>'):
yield story
def simplestories(raw):
for path in sorted(glob.glob(f'{raw}/simplestories-[0-9].parquet')):
yield from pq.read_table(path, columns=['story']).column('story').to_pylist()
def soda(raw):
table = pq.read_table(f'{raw}/soda-train.parquet', columns=['narrative', 'dialogue', 'speakers'])
for narrative, lines, speakers in zip(*(table.column(c).to_pylist() for c in table.column_names)):
if len(speakers) != len(lines):
continue
dialogue = '\n'.join(f'{who}: {line}' for who, line in zip(speakers, lines))
yield f'{narrative}\n\n{dialogue}'
SOURCES = {'tinystories': tinystories, 'simplestories': simplestories, 'soda': soda}
_tok = None
def _init(path):
global _tok
_tok = Tokenizer.load(path)
def _encode(texts):
eot = _tok['<|endoftext|>']
out = []
for t in texts:
out.append(eot)
out.extend(_tok.encode(t, allow_special=False))
return np.array(out, dtype=np.uint16)
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--raw', default='raw')
ap.add_argument('--out', default='data')
ap.add_argument('--vocab', type=int, default=8192)
ap.add_argument('--tok-docs', type=int, default=40_000, help='documents per source for BPE training')
ap.add_argument('--val-frac', type=float, default=0.005)
args = ap.parse_args()
os.makedirs(args.out, exist_ok=True)
rng = random.Random(1337)
# 1. clean + exact deduplication
docs, stats = {}, {}
for name, read in SOURCES.items():
seen, kept, raw_count = set(), [], 0
for text in read(args.raw):
raw_count += 1
text = clean(text)
if text is None:
continue
h = hashlib.md5(text.lower().encode()).digest()
if h in seen:
continue
seen.add(h)
kept.append(text)
rng.shuffle(kept)
docs[name] = kept
stats[name] = {'raw': raw_count, 'kept': len(kept), 'chars': sum(map(len, kept))}
print(f'{name}: {raw_count} -> {len(kept)} documents', flush=True)
# 2. tokenizer on a balanced sample
tok_path = f'{args.out}/tokenizer.json'
if not os.path.exists(tok_path):
t0 = time.time()
sample = [d for name in docs for d in docs[name][:args.tok_docs]]
tok = Tokenizer.train(sample, args.vocab, verbose=True)
tok.save(tok_path)
print(f'tokenizer: {tok.vocab_size} tokens in {time.time() - t0:.0f}s', flush=True)
# 3. encode in parallel; a small slice of every source goes to validation
tok = Tokenizer.load(tok_path)
eot = tok['<|endoftext|>']
train, val = [], []
with Pool(max(1, os.cpu_count() - 2), initializer=_init, initargs=(tok_path,)) as pool:
for name, texts in docs.items():
n_val = int(len(texts) * args.val_frac)
for split, part in ((val, texts[:n_val]), (train, texts[n_val:])):
chunks = [part[i:i + 2000] for i in range(0, len(part), 2000)]
split.append(np.concatenate(pool.map(_encode, chunks)))
stats[name]['tokens'] = int(len(train[-1]) + len(val[-1]))
print(f'{name}: {stats[name]["tokens"] / 1e6:.1f}M tokens', flush=True)
# shuffle whole runs of documents between sources, so any stretch of the file is a fair mix
blocks = []
for arr in train:
starts = np.flatnonzero(arr == eot)
cuts = starts[np.searchsorted(starts, np.arange(1_000_000, len(arr), 1_000_000))]
blocks.extend(np.split(arr, np.unique(cuts)))
rng.shuffle(blocks)
np.concatenate(blocks).tofile(f'{args.out}/train.bin')
np.concatenate(val).tofile(f'{args.out}/val.bin')
stats['train_tokens'] = int(sum(len(b) for b in blocks))
stats['val_tokens'] = int(sum(len(v) for v in val))
with open(f'{args.out}/corpus_stats.json', 'w') as f:
json.dump(stats, f, indent=1)
if __name__ == '__main__':
main()
model.py
Сама модель: RMSNorm, RoPE, внимание с QK-norm, SwiGLU, общие эмбеддинги.
174 строк
· Разбираем в главе 8
· Открыть как текст
"""Sprout: a tiny GPT, all of it in one file.
A decoder-only transformer in the style of today's open models:
- pre-norm blocks with RMSNorm,
- rotary position embeddings (RoPE),
- multi-head causal self-attention with QK-norm,
- a SwiGLU feed-forward layer,
- input and output embeddings shared (weight tying),
- no biases anywhere.
"""
import math
from dataclasses import dataclass, asdict
import torch
import torch.nn as nn
import torch.nn.functional as F
@dataclass
class Config:
vocab_size: int = 8192 # how many different tokens the model knows
context: int = 512 # how many tokens it can look back at
n_layer: int = 8 # transformer blocks stacked on top of each other
n_head: int = 8 # attention heads per block
d_model: int = 512 # width of the residual stream
d_ff: int = 1408 # hidden width of the feed-forward layer
rope_base: float = 10000.0
def to_dict(self):
return asdict(self)
class RMSNorm(nn.Module):
"""Rescale a vector to unit root-mean-square, then apply a learned gain."""
def __init__(self, dim, eps=1e-6):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def forward(self, x):
return F.rms_norm(x.float(), (x.size(-1),), self.weight.float(), self.eps).type_as(x)
def rope_tables(head_dim, context, base):
"""cos/sin for every position and every frequency, shape (context, head_dim/2)."""
inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
angles = torch.outer(torch.arange(context).float(), inv_freq)
return angles.cos(), angles.sin()
def apply_rope(x, cos, sin):
"""Rotate pairs (x[i], x[i + d/2]) by a position-dependent angle.
x: (batch, heads, time, head_dim); cos/sin: (time, head_dim/2).
"""
d = x.size(-1) // 2
x1, x2 = x[..., :d], x[..., d:]
return torch.cat((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1)
class Attention(nn.Module):
def __init__(self, cfg):
super().__init__()
self.n_head = cfg.n_head
self.head_dim = cfg.d_model // cfg.n_head
self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
self.q_norm = RMSNorm(self.head_dim)
self.k_norm = RMSNorm(self.head_dim)
def forward(self, x, cos, sin):
B, T, C = x.shape
q, k, v = self.qkv(x).split(C, dim=-1)
# (B, T, C) -> (B, heads, T, head_dim)
q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
q, k = self.q_norm(q), self.k_norm(k)
q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
# softmax(q·k / sqrt(d)) · v, every token sees only itself and the past
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
y = y.transpose(1, 2).contiguous().view(B, T, C)
return self.proj(y)
class FeedForward(nn.Module):
"""SwiGLU: silu(x·W1) ⊙ (x·W3), projected back with W2."""
def __init__(self, cfg):
super().__init__()
self.w1 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
self.w3 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
self.w2 = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)
def forward(self, x):
return self.w2(F.silu(self.w1(x)) * self.w3(x))
class Block(nn.Module):
def __init__(self, cfg):
super().__init__()
self.norm1 = RMSNorm(cfg.d_model)
self.attn = Attention(cfg)
self.norm2 = RMSNorm(cfg.d_model)
self.ffn = FeedForward(cfg)
def forward(self, x, cos, sin):
x = x + self.attn(self.norm1(x), cos, sin) # tokens exchange information
x = x + self.ffn(self.norm2(x)) # each token thinks on its own
return x
class GPT(nn.Module):
def __init__(self, cfg):
super().__init__()
self.cfg = cfg
self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layer))
self.norm = RMSNorm(cfg.d_model)
self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
self.head.weight = self.embed.weight # weight tying
cos, sin = rope_tables(cfg.d_model // cfg.n_head, cfg.context, cfg.rope_base)
self.register_buffer('cos', cos, persistent=False)
self.register_buffer('sin', sin, persistent=False)
self.apply(self._init)
# residual projections start small so every block begins close to "do nothing"
for name, p in self.named_parameters():
if name.endswith('proj.weight') or name.endswith('w2.weight'):
nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer))
@staticmethod
def _init(m):
if isinstance(m, (nn.Linear, nn.Embedding)):
nn.init.normal_(m.weight, mean=0.0, std=0.02)
def num_params(self):
return sum(p.numel() for p in self.parameters())
def forward(self, idx, targets=None, loss_mask=None):
T = idx.size(1)
x = self.embed(idx)
cos, sin = self.cos[:T], self.sin[:T]
for block in self.blocks:
x = block(x, cos, sin)
logits = self.head(self.norm(x))
if targets is None:
return logits, None
loss = F.cross_entropy(logits.float().view(-1, logits.size(-1)), targets.reshape(-1), reduction='none')
if loss_mask is None:
return logits, loss.mean()
mask = loss_mask.reshape(-1).float()
return logits, (loss * mask).sum() / mask.sum().clamp(min=1)
@torch.no_grad()
def generate(self, idx, max_new, temperature=0.8, top_k=None, top_p=0.95, stop=None):
for _ in range(max_new):
logits, _ = self(idx[:, -self.cfg.context:])
logits = logits[:, -1, :].float() / max(temperature, 1e-5)
if top_k:
kth = torch.topk(logits, top_k).values[:, -1, None]
logits[logits < kth] = -float('inf')
probs = F.softmax(logits, dim=-1)
if top_p and top_p < 1.0:
sorted_p, order = probs.sort(descending=True)
drop = sorted_p.cumsum(-1) - sorted_p > top_p
sorted_p[drop] = 0
probs = torch.zeros_like(probs).scatter(-1, order, sorted_p)
probs /= probs.sum(-1, keepdim=True)
nxt = torch.multinomial(probs, 1)
idx = torch.cat((idx, nxt), dim=1)
if stop is not None and nxt.item() in stop:
break
return idx
muon.py
Оптимизатор Muon для скрытых матриц.
63 строк
· Разбираем в главе 10
· Открыть как текст
"""Muon: momentum, then orthogonalise the update of every weight matrix.
For a matrix W, plain SGD adds -lr * G. Muon replaces the momentum-averaged G
by the closest orthogonal matrix U·Vᵀ (from G = U·S·Vᵀ): every direction gets
the same step size, so rare but useful directions are not drowned out by a
few dominant ones. Five Newton–Schulz iterations approximate U·Vᵀ cheaply.
Only for 2-D hidden weights. Embeddings and norm gains stay with AdamW.
After Keller Jordan, https://kellerjordan.github.io/posts/muon/
"""
from collections import defaultdict
import torch
def orthogonalize(G, steps=5):
"""G: (..., rows, cols) -> the nearest semi-orthogonal matrices, batched."""
a, b, c = 3.4445, -4.7750, 2.0315 # tuned for fast convergence
X = G.bfloat16()
tall = X.size(-2) > X.size(-1)
if tall:
X = X.mT
X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
for _ in range(steps):
A = X @ X.mT
X = a * X + (b * A + c * A @ A) @ X
if tall:
X = X.mT
return X.to(G.dtype)
class Muon(torch.optim.Optimizer):
def __init__(self, params, lr=0.02, momentum=0.95, weight_decay=0.0, split=None):
# split: {param: n} treats a fused weight (like q, k, v) as n separate matrices
self.split = split or {}
super().__init__(params, dict(lr=lr, momentum=momentum, weight_decay=weight_decay))
@torch.no_grad()
def step(self):
for group in self.param_groups:
# matrices of the same shape are orthogonalised together, in one batch
by_shape = defaultdict(list)
for p in group['params']:
if p.grad is None:
continue
state = self.state[p]
if 'momentum' not in state:
state['momentum'] = torch.zeros_like(p)
buf = state['momentum']
buf.lerp_(p.grad, 1 - group['momentum'])
g = p.grad.lerp(buf, group['momentum']) # Nesterov look-ahead
n = self.split.get(p, 1)
by_shape[(p.size(0) // n, p.size(1))].append((p, g.view(n, -1, p.size(1))))
for (rows, cols), items in by_shape.items():
updates = orthogonalize(torch.cat([g for _, g in items]))
scale = max(1.0, rows / cols) ** 0.5
i = 0
for p, g in items:
u = updates[i:i + g.size(0)].reshape_as(p)
i += g.size(0)
if group['weight_decay']:
p.mul_(1 - group['lr'] * group['weight_decay'])
p.add_(u, alpha=-group['lr'] * scale)
train.py
Цикл предобучения: батчи, bf16, torch.compile, расписание шага, образцы по ходу.
211 строк
· Разбираем в главе 10
· Открыть как текст
"""Pre-train Sprout: show it text, ask for the next token, nudge the weights.
python train.py --out runs/base --tokens 400e6
Writes runs/<name>/log.jsonl (loss, learning rate, samples as it learns)
and checkpoints ckpt_<step>.pt + ckpt_final.pt.
"""
import argparse
import json
import math
import os
import time
import numpy as np
import torch
from model import GPT, Config
from muon import Muon
from tokenizer import Tokenizer
PROMPTS = ['Once upon a time', 'Lily looked at the sky and said', 'Tom: Hi! How are you today?\nAnna:',
'The best thing about summer is']
def get_device():
if torch.cuda.is_available():
return 'cuda'
if torch.backends.mps.is_available():
return 'mps'
return 'cpu'
class Batches:
"""Random windows of context+1 tokens from a flat file of token ids."""
def __init__(self, path, device, seed=0):
self.data = np.memmap(path, dtype=np.uint16, mode='r')
self.device = device
self.rng = np.random.default_rng(seed)
def get(self, batch, context):
starts = self.rng.integers(0, len(self.data) - context - 1, batch)
rows = np.stack([self.data[s:s + context + 1] for s in starts]).astype(np.int64)
rows = torch.from_numpy(rows).to(self.device, non_blocking=True)
return rows[:, :-1], rows[:, 1:]
def lr_factor(step, total, warmup, cooldown):
"""Warm up, hold, then decay linearly to zero over the last `cooldown` share."""
if step < warmup:
return (step + 1) / warmup
decay_start = total * (1 - cooldown)
if step < decay_start:
return 1.0
return max(0.0, (total - step) / (total - decay_start))
def context_at(step, total, context, warm):
"""Sequence-length warm-up: short windows first, full length later."""
if not warm:
return context
if step < 0.1 * total:
return max(64, context // 4)
if step < 0.3 * total:
return context // 2
return context
def make_optimizers(model, args):
matrices = [p for n, p in model.named_parameters() if p.ndim == 2 and 'embed' not in n]
others = [p for n, p in model.named_parameters() if not (p.ndim == 2 and 'embed' not in n)]
embed = [p for p in others if p.ndim == 2]
gains = [p for p in others if p.ndim < 2]
adam_groups = [dict(params=embed, lr=args.lr_embed, weight_decay=0.0),
dict(params=gains, lr=args.lr_embed, weight_decay=0.0)]
if args.optim == 'muon':
split = {blk.attn.qkv.weight: 3 for blk in model.blocks}
muon = Muon(matrices, lr=args.lr, weight_decay=args.wd, split=split)
adam = torch.optim.AdamW(adam_groups, betas=(0.9, 0.95), eps=1e-10)
return [muon, adam]
adam_groups.append(dict(params=matrices, lr=args.lr, weight_decay=args.wd))
return [torch.optim.AdamW(adam_groups, betas=(0.9, 0.95), eps=1e-10)]
@torch.no_grad()
def evaluate(model, val, args, ctx):
model.eval()
rng_state = val.rng
val.rng = np.random.default_rng(1234) # the same batches every time
losses = []
for _ in range(args.eval_batches):
x, y = val.get(args.batch_tokens // args.context, args.context)
with ctx:
_, loss = model(x, y)
losses.append(loss.item())
val.rng = rng_state
model.train()
return sum(losses) / len(losses)
@torch.no_grad()
def samples(model, tok, device, n_tokens=80):
model.eval()
out = []
g = torch.Generator(device='cpu').manual_seed(7)
for prompt in PROMPTS:
torch.manual_seed(int(torch.randint(0, 2**31, (1,), generator=g)))
idx = torch.tensor([[tok['<|endoftext|>']] + tok.encode(prompt)], device=device)
ids = model.generate(idx, n_tokens, temperature=0.8, top_p=0.95, stop={tok['<|endoftext|>']})
out.append(tok.decode([i for i in ids[0, 1:].tolist() if i != tok['<|endoftext|>']]))
model.train()
return out
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--data', default='data')
ap.add_argument('--out', default='runs/base')
ap.add_argument('--d-model', type=int, default=384)
ap.add_argument('--n-layer', type=int, default=8)
ap.add_argument('--n-head', type=int, default=6)
ap.add_argument('--d-ff', type=int, default=1024)
ap.add_argument('--context', type=int, default=512)
ap.add_argument('--batch-tokens', type=int, default=32768)
ap.add_argument('--tokens', type=float, default=400e6, help='how many tokens to train on')
ap.add_argument('--optim', choices=['muon', 'adamw'], default='muon')
ap.add_argument('--lr', type=float, default=0.02, help='Muon lr (or AdamW lr for matrices)')
ap.add_argument('--lr-embed', type=float, default=0.006)
ap.add_argument('--wd', type=float, default=0.0)
ap.add_argument('--warmup', type=int, default=200)
ap.add_argument('--cooldown', type=float, default=0.3)
ap.add_argument('--seq-warmup', action='store_true')
ap.add_argument('--eval-every', type=int, default=250)
ap.add_argument('--eval-batches', type=int, default=20)
ap.add_argument('--sample-every', type=int, default=500)
ap.add_argument('--sample-at', default='', help='extra steps with samples, e.g. 25,50,100')
ap.add_argument('--save-at', default='', help='extra checkpoint steps, e.g. 100,1000,5000')
ap.add_argument('--no-compile', action='store_true')
ap.add_argument('--seed', type=int, default=0)
ap.add_argument('--device', default='', help='cuda / mps / cpu (default: the best available)')
ap.add_argument('--threads', type=int, default=0, help='CPU threads (0 = default)')
args = ap.parse_args()
torch.manual_seed(args.seed)
if args.threads:
torch.set_num_threads(args.threads)
device = args.device or get_device()
os.makedirs(args.out, exist_ok=True)
tok = Tokenizer.load(f'{args.data}/tokenizer.json')
cfg = Config(vocab_size=tok.vocab_size, context=args.context, n_layer=args.n_layer,
n_head=args.n_head, d_model=args.d_model, d_ff=args.d_ff)
model = GPT(cfg).to(device)
print(f'{model.num_params() / 1e6:.2f}M parameters on {device}')
train_model = model if args.no_compile else torch.compile(model)
opts = make_optimizers(model, args)
for opt in opts:
for g in opt.param_groups:
g['base_lr'] = g['lr']
ctx = torch.autocast(device_type=device, dtype=torch.bfloat16) if device != 'cpu' else torch.autocast('cpu', enabled=False)
train = Batches(f'{args.data}/train.bin', device, seed=args.seed)
val = Batches(f'{args.data}/val.bin', device)
steps = int(args.tokens // args.batch_tokens)
save_at = {int(s) for s in args.save_at.split(',') if s}
sample_at = {int(s) for s in args.sample_at.split(',') if s}
log = open(f'{args.out}/log.jsonl', 'a')
json.dump({'config': cfg.to_dict(), 'args': vars(args), 'steps': steps, 'params': model.num_params()}, log)
log.write('\n')
t0, seen = time.time(), 0
for step in range(steps + 1):
last = step == steps
if step % args.eval_every == 0 or step in sample_at or last:
vl = evaluate(train_model, val, args, ctx)
rec = {'step': step, 'tokens': seen, 'val': round(vl, 4), 'time': round(time.time() - t0, 1)}
if step % args.sample_every == 0 or step in sample_at or last:
rec['samples'] = samples(model, tok, device)
print(json.dumps(rec), flush=True)
log.write(json.dumps(rec) + '\n')
log.flush()
if step in save_at or last:
name = 'final' if last else step
torch.save({'config': cfg.to_dict(), 'model': model.state_dict(), 'step': step, 'tokens': seen},
f'{args.out}/ckpt_{name}.pt')
if last:
break
T = context_at(step, steps, args.context, args.seq_warmup)
x, y = train.get(args.batch_tokens // T, T)
with ctx:
_, loss = train_model(x, y)
loss.backward()
norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
f = lr_factor(step, steps, args.warmup, args.cooldown)
for opt in opts:
for g in opt.param_groups:
g['lr'] = g['base_lr'] * f
opt.step()
opt.zero_grad(set_to_none=True)
seen += x.numel()
if step % 20 == 0:
rec = {'step': step, 'loss': round(loss.item(), 4), 'lr': round(f, 4), 'norm': round(norm.item(), 3),
'tokens': seen, 'tok_s': round(seen / (time.time() - t0)), 'T': T}
log.write(json.dumps(rec) + '\n')
if step % 100 == 0:
print(json.dumps(rec), flush=True)
log.close()
if __name__ == '__main__':
main()
build_chats.py
Собирает диалоги для обучения разговору.
122 строк
· Разбираем в главе 12
· Открыть как текст
"""Collect the conversations Sprout learns to chat from -> data/chats.jsonl
- sft/*.jsonl conversations written for Sprout's level
- everyday-conversations (HF) short everyday chats with an assistant
- SimpleStories "tell me a story about ..." requests, built from its labels
- TinyStories "continue this story: ..." requests
"""
import argparse
import glob
import json
import random
import re
import pyarrow.parquet as pq
from prepare import REPLACE
def tidy(text):
text = text.translate(REPLACE).strip()
return '\n'.join(' '.join(line.split()) for line in text.split('\n'))
def conversation(messages, source):
msgs = [{'role': m['role'], 'content': tidy(m['content'])} for m in messages if m['role'] in ('user', 'assistant')]
ok = msgs and msgs[0]['role'] == 'user' and msgs[-1]['role'] == 'assistant' and all(m['content'] for m in msgs)
ok = ok and all(a['role'] != b['role'] for a, b in zip(msgs, msgs[1:]))
return {'messages': msgs, 'source': source} if ok else None
STORY_ASKS = [
'Tell me a story about {subject}.', 'Can you write a story about {subject}?', 'story about {subject} pls',
'I want a short story about {subject}.', 'Tell me a bedtime story about {subject}.',
'Write a story where the hero is {subject}.', 'Could you tell me a story about {subject}, please?',
'Make up a story about {subject}!',
]
# "Once upon a time, there was a little fox named ..." -> "a little fox"
SUBJECT = re.compile(r"\bthere (?:was|lived) (a|an) ([a-z]+(?: [a-z]+){0,3}?)(?= named| who| that| in | with| called|,|\.)")
def story_requests(raw, n, rng, per_noun=120):
"""Ask for a story about exactly the hero the story is about, so request and answer agree.
(A first version asked for SimpleStories' abstract topic labels like "bygone eras"; the stories
only loosely matched them, and the model learned that a story need not fit the request.)
"""
with open(f'{raw}/TinyStoriesV2-GPT4-train.txt', encoding='utf-8') as f:
text = f.read(150_000_000)
per = {}
out = []
stories = text.split('<|endoftext|>')[1:-1]
rng.shuffle(stories)
for s in stories:
s = tidy(s)
m = SUBJECT.search(s[:220])
if not m or len(s.split()) > 170 or not 1 <= len(m.group(2).split()) <= 3:
continue
noun = m.group(2).split()[-1]
if per.get(noun, 0) >= per_noun:
continue
per[noun] = per.get(noun, 0) + 1
ask = rng.choice(STORY_ASKS).format(subject=f'{m.group(1)} {m.group(2)}')
out.append(conversation([{'role': 'user', 'content': ask}, {'role': 'assistant', 'content': s}], 'story'))
if len(out) >= n:
break
return out
def continuations(raw, n, rng):
with open(f'{raw}/TinyStoriesV2-GPT4-valid.txt', encoding='utf-8') as f:
stories = [s.strip() for s in f.read().split('<|endoftext|>') if 300 < len(s.strip()) < 1200]
out = []
for s in rng.sample(stories, min(n, len(stories))):
sentences = re.split(r'(?<=[.!?])\s+', tidy(s))
if len(sentences) < 4:
continue
k = rng.choice([1, 2])
start, rest = ' '.join(sentences[:k]), ' '.join(sentences[k:])
ask = rng.choice(['Continue this story: {s}', 'Can you finish this story? "{s}"', 'What happens next? {s}',
'Please continue: {s}'])
out.append(conversation([{'role': 'user', 'content': ask.format(s=start)},
{'role': 'assistant', 'content': rest}], 'continue'))
return out
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--raw', default='raw')
ap.add_argument('--written', default='sft')
ap.add_argument('--out', default='data/chats.jsonl')
ap.add_argument('--stories', type=int, default=2500)
ap.add_argument('--continues', type=int, default=800)
ap.add_argument('--everyday', type=int, default=1200)
ap.add_argument('--written-repeat', type=int, default=2)
ap.add_argument('--self-repeat', type=int, default=4, help='conversations about Sprout itself (files 08-*)')
args = ap.parse_args()
rng = random.Random(7)
written = []
for path in sorted(glob.glob(f'{args.written}/*.jsonl')):
repeat = args.self_repeat if path.split('/')[-1].startswith('08-') else args.written_repeat
for line in open(path):
if line.strip():
written += [conversation(json.loads(line)['messages'], 'written')] * repeat
everyday = [conversation(m, 'everyday')
for m in pq.read_table(f'{args.raw}/everyday-train.parquet').column('messages').to_pylist()]
parts = {
'written': [c for c in written if c],
'everyday': rng.sample([c for c in everyday if c], args.everyday),
'story': [c for c in story_requests(args.raw, args.stories, rng) if c],
'continue': [c for c in continuations(args.raw, args.continues, rng) if c],
}
chats = [c for part in parts.values() for c in part]
rng.shuffle(chats)
with open(args.out, 'w') as f:
for c in chats:
f.write(json.dumps(c) + '\n')
print({k: len(v) for k, v in parts.items()}, '->', len(chats))
if __name__ == '__main__':
main()
sft.py
Дообучение для чата: ошибка считается только на словах ассистента.
154 строк
· Разбираем в главе 12
· Открыть как текст
"""Teach the base model to chat: supervised fine-tuning on conversations.
python sft.py --base runs/base/ckpt_final.pt --chats data/chats.jsonl --out runs/chat
A conversation becomes one line of tokens:
<|endoftext|><|user|>hi!<|end|><|assistant|>Hello! How are you?<|end|><|user|>...
The loss is counted only on the assistant's words and its closing <|end|>, so the
model learns to answer, not to imitate the user.
"""
import argparse
import json
import os
import random
import time
import torch
from model import GPT, Config
from muon import Muon
from tokenizer import Tokenizer
from train import get_device, lr_factor
def render(tok, messages):
"""Token ids and a 0/1 mask saying which positions the model should learn to predict."""
ids, mask = [tok['<|endoftext|>']], [0]
for m in messages:
role = tok['<|user|>'] if m['role'] == 'user' else tok['<|assistant|>']
body = tok.encode(m['content'].strip(), allow_special=False) + [tok['<|end|>']]
ids += [role] + body
learn = 1 if m['role'] == 'assistant' else 0
mask += [0] + [learn] * len(body)
return ids, mask
def pack(examples, context):
"""Glue conversations into rows of exactly context+1 tokens (the rest is padding)."""
rows, cur_ids, cur_mask = [], [], []
for ids, mask in examples:
if len(ids) > context + 1:
continue
if len(cur_ids) + len(ids) > context + 1:
rows.append((cur_ids, cur_mask))
cur_ids, cur_mask = [], []
cur_ids += ids
cur_mask += mask
if cur_ids:
rows.append((cur_ids, cur_mask))
x = torch.zeros(len(rows), context + 1, dtype=torch.long)
m = torch.zeros(len(rows), context + 1, dtype=torch.long)
for i, (ids, mask) in enumerate(rows):
x[i, :len(ids)] = torch.tensor(ids)
m[i, :len(mask)] = torch.tensor(mask)
return x, m
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--base', required=True)
ap.add_argument('--chats', default='data/chats.jsonl')
ap.add_argument('--data', default='data')
ap.add_argument('--out', default='runs/chat')
ap.add_argument('--epochs', type=float, default=3)
ap.add_argument('--batch', type=int, default=32)
ap.add_argument('--lr', type=float, default=0.004)
ap.add_argument('--lr-embed', type=float, default=0.0008)
ap.add_argument('--val-frac', type=float, default=0.03)
ap.add_argument('--seed', type=int, default=0)
args = ap.parse_args()
torch.manual_seed(args.seed)
device = get_device()
os.makedirs(args.out, exist_ok=True)
tok = Tokenizer.load(f'{args.data}/tokenizer.json')
ck = torch.load(args.base, map_location='cpu')
cfg = Config(**ck['config'])
model = GPT(cfg)
model.load_state_dict(ck['model'])
model.to(device)
chats = [json.loads(line) for line in open(args.chats)]
# some conversations appear several times on purpose; hold out whole conversations,
# every copy of them, so validation text is really never trained on
key = lambda c: json.dumps(c['messages'], sort_keys=True)
unique = sorted({key(c) for c in chats})
random.Random(args.seed).shuffle(unique)
held_out = set(unique[:int(len(unique) * args.val_frac)])
val = list({key(c): c for c in chats if key(c) in held_out}.values())
train = [c for c in chats if key(c) not in held_out]
random.Random(args.seed).shuffle(train)
val_x, val_m = pack([render(tok, c['messages']) for c in val], cfg.context)
train_x, train_m = pack([render(tok, c['messages']) for c in train], cfg.context)
print(f'{len(chats)} conversations -> {len(train_x)} train rows, {len(val_x)} val rows')
matrices = [p for n, p in model.named_parameters() if p.ndim == 2 and 'embed' not in n]
rest = [p for n, p in model.named_parameters() if not (p.ndim == 2 and 'embed' not in n)]
muon = Muon(matrices, lr=args.lr, split={b.attn.qkv.weight: 3 for b in model.blocks})
adam = torch.optim.AdamW(rest, lr=args.lr_embed, betas=(0.9, 0.95), weight_decay=0.0)
opts = [muon, adam]
for opt in opts:
for g in opt.param_groups:
g['base_lr'] = g['lr']
ctx = torch.autocast(device_type=device, dtype=torch.bfloat16)
@torch.no_grad()
def evaluate():
model.eval()
losses = []
for i in range(0, len(val_x), args.batch):
x, m = val_x[i:i + args.batch].to(device), val_m[i:i + args.batch].to(device)
with ctx:
_, loss = model(x[:, :-1], x[:, 1:], loss_mask=m[:, 1:])
losses.append(loss.item())
model.train()
return sum(losses) / len(losses)
steps = int(args.epochs * len(train_x) / args.batch)
log = open(f'{args.out}/log.jsonl', 'w')
log.write(json.dumps({'config': cfg.to_dict(), 'args': vars(args), 'params': model.num_params(), 'steps': steps,
'conversations': len(chats), 'train_rows': len(train_x)}) + '\n')
t0 = time.time()
order = torch.randperm(len(train_x))
pos = 0
for step in range(steps + 1):
if step % 50 == 0 or step == steps:
rec = {'step': step, 'val': round(evaluate(), 4), 'time': round(time.time() - t0, 1)}
print(json.dumps(rec), flush=True)
log.write(json.dumps(rec) + '\n')
log.flush()
if step == steps:
break
if pos + args.batch > len(order):
order, pos = torch.randperm(len(train_x)), 0
idx = order[pos:pos + args.batch]
pos += args.batch
x, m = train_x[idx].to(device), train_m[idx].to(device)
with ctx:
_, loss = model(x[:, :-1], x[:, 1:], loss_mask=m[:, 1:])
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
f = lr_factor(step, steps, warmup=20, cooldown=0.5)
for opt in opts:
for g in opt.param_groups:
g['lr'] = g['base_lr'] * f
opt.step()
opt.zero_grad(set_to_none=True)
if step % 10 == 0:
log.write(json.dumps({'step': step, 'loss': round(loss.item(), 4), 'lr': round(f, 4)}) + '\n')
torch.save({'config': cfg.to_dict(), 'model': model.state_dict(), 'step': steps}, f'{args.out}/ckpt_final.pt')
log.close()
if __name__ == '__main__':
main()
export.py
Упаковывает веса в int8 для браузерного движка.
83 строк
· Разбираем в главе 11
· Открыть как текст
"""Pack a checkpoint for the browser: int8 weights with one scale per row.
python export.py runs/chat/ckpt_final.pt ../../public/llm/models/sprout-chat.bin
File layout: b'SPRT', uint32 version, uint32 header length, JSON header,
then the tensors, each aligned to 16 bytes. A matrix W (rows x cols) is stored
as int8 q plus float32 scales s so that W[r] ≈ q[r] * s[r].
"""
import argparse
import json
import struct
import numpy as np
import torch
MAGIC, VERSION = b'SPRT', 1
def quantize_rows(w):
scale = np.abs(w).max(axis=1) / 127.0
scale[scale == 0] = 1.0
q = np.clip(np.round(w / scale[:, None]), -127, 127).astype(np.int8)
return q, scale.astype(np.float32)
def pack(state, config, meta, path, quantize=True):
tensors, blobs, offset = [], [], 0
def add(name, arr, dtype):
nonlocal offset
data = arr.astype(dtype).tobytes()
pad = (-len(data)) % 16
tensors.append({'name': name, 'dtype': np.dtype(dtype).name, 'shape': list(arr.shape), 'offset': offset})
blobs.append(data + b'\0' * pad)
offset += len(data) + pad
for name, t in state.items():
if name == 'head.weight': # tied to embed.weight
continue
w = t.detach().float().cpu().numpy()
if quantize and w.ndim == 2:
q, s = quantize_rows(w)
add(name, q, np.int8)
add(name + '.scale', s, np.float32)
else:
add(name, w, np.float32)
header = json.dumps({'config': config, 'meta': meta, 'tensors': tensors}).encode()
header += b' ' * ((-(12 + len(header))) % 16)
with open(path, 'wb') as f:
f.write(MAGIC + struct.pack('<II', VERSION, len(header)) + header)
for b in blobs:
f.write(b)
return 12 + len(header) + offset
def dequantized(state):
"""The weights exactly as the browser will see them, to measure the damage."""
out = {}
for name, t in state.items():
w = t.detach().float().cpu()
if w.ndim == 2 and name != 'head.weight':
q, s = quantize_rows(w.numpy())
w = torch.from_numpy(q.astype(np.float32) * s[:, None])
out[name] = w
out['head.weight'] = out['embed.weight']
return out
def main():
ap = argparse.ArgumentParser()
ap.add_argument('ckpt')
ap.add_argument('out')
ap.add_argument('--name', default='')
ap.add_argument('--float', action='store_true', help='keep float32 (tiny models)')
args = ap.parse_args()
ck = torch.load(args.ckpt, map_location='cpu')
meta = {'name': args.name, 'step': ck.get('step'), 'tokens': ck.get('tokens'), 'kind': 'gpt'}
size = pack(ck['model'], ck['config'], meta, args.out, quantize=not args.float)
print(f'{args.out}: {size / 1e6:.1f} MB')
if __name__ == '__main__':
main()
chat.py
Разговор с моделью в терминале.
54 строк
· Разбираем в главе 12
· Открыть как текст
"""Talk to Sprout in the terminal.
python chat.py runs/chat/ckpt_final.pt # the chat model
python chat.py runs/base/ckpt_final.pt --base # the base model just continues your text
"""
import argparse
import torch
from model import GPT, Config
from tokenizer import Tokenizer
from train import get_device
def main():
ap = argparse.ArgumentParser()
ap.add_argument('ckpt')
ap.add_argument('--data', default='data')
ap.add_argument('--base', action='store_true', help='plain continuation, no chat template')
ap.add_argument('--temperature', type=float, default=0.7)
ap.add_argument('--top-p', type=float, default=0.9)
ap.add_argument('--max', type=int, default=200)
args = ap.parse_args()
device = get_device()
tok = Tokenizer.load(f'{args.data}/tokenizer.json')
ck = torch.load(args.ckpt, map_location='cpu')
model = GPT(Config(**ck['config']))
model.load_state_dict(ck['model'])
model.to(device).eval()
eot, user, assistant, end = (tok[s] for s in ('<|endoftext|>', '<|user|>', '<|assistant|>', '<|end|>'))
history = [eot]
print('Sprout is listening. Empty line to quit.')
while True:
text = input('you> ').strip()
if not text:
break
if args.base:
ids = [eot] + tok.encode(text, allow_special=False)
else:
history += [user] + tok.encode(text, allow_special=False) + [end, assistant]
ids = history[-(model.cfg.context - args.max):]
out = model.generate(torch.tensor([ids], device=device), args.max, temperature=args.temperature,
top_p=args.top_p, stop={end, eot})
new = out[0, len(ids):].tolist()
reply = [i for i in new if i not in (end, eot)]
print('sprout>', (text + ' ' if args.base else '') + tok.decode(reply).strip())
if not args.base:
history += reply + [end]
if __name__ == '__main__':
main()
lora.py
Стилевой адаптер LoRA, который потом вливается в веса.
141 строк
· Разбираем в главе 13
· Открыть как текст
"""LoRA: teach the chat model a new style by training two thin matrices per layer.
python lora.py --base runs/chat/ckpt_final.pt --chats sft/style/pirate.jsonl --out runs/pirate
A frozen weight W gets a detour: y = W x + (alpha / r) · B (A x), A: r × in, B: out × r.
B starts at zero, so at step 0 the model is exactly the chat model. Only A and B learn —
with rank 8 that is about 2.4 % of Sprout's parameters. At the end B·A is merged into W,
and the result is an ordinary Sprout checkpoint again.
"""
import argparse
import json
import math
import os
import random
import time
import torch
import torch.nn as nn
from model import GPT, Config
from sft import pack, render
from tokenizer import Tokenizer
from train import get_device, lr_factor
class LoRALinear(nn.Module):
def __init__(self, base, rank=8, alpha=16):
super().__init__()
self.base = base
self.A = nn.Parameter(torch.randn(rank, base.in_features) / math.sqrt(base.in_features))
self.B = nn.Parameter(torch.zeros(base.out_features, rank))
self.scale = alpha / rank
def forward(self, x):
return self.base(x) + (x @ self.A.T @ self.B.T) * self.scale
def merge(self):
self.base.weight.data += (self.B @ self.A).to(self.base.weight.dtype) * self.scale
return self.base
TARGETS = [('attn', 'qkv'), ('attn', 'proj'), ('ffn', 'w1'), ('ffn', 'w2'), ('ffn', 'w3')]
def add_lora(model, rank, alpha):
"""Freeze everything, then wrap every projection in the blocks with a LoRA detour."""
for p in model.parameters():
p.requires_grad_(False)
for block in model.blocks:
for part, name in TARGETS:
parent = getattr(block, part)
setattr(parent, name, LoRALinear(getattr(parent, name), rank, alpha))
return [p for p in model.parameters() if p.requires_grad]
def merge_lora(model):
for block in model.blocks:
for part, name in TARGETS:
parent = getattr(block, part)
layer = getattr(parent, name)
if isinstance(layer, LoRALinear):
setattr(parent, name, layer.merge())
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--base', required=True)
ap.add_argument('--chats', required=True)
ap.add_argument('--data', default='data')
ap.add_argument('--out', required=True)
ap.add_argument('--rank', type=int, default=8)
ap.add_argument('--alpha', type=float, default=16)
ap.add_argument('--epochs', type=float, default=9)
ap.add_argument('--batch', type=int, default=16)
ap.add_argument('--lr', type=float, default=3e-3)
ap.add_argument('--seed', type=int, default=0)
args = ap.parse_args()
torch.manual_seed(args.seed)
device = get_device()
os.makedirs(args.out, exist_ok=True)
tok = Tokenizer.load(f'{args.data}/tokenizer.json')
ck = torch.load(args.base, map_location='cpu')
cfg = Config(**ck['config'])
model = GPT(cfg)
model.load_state_dict(ck['model'])
trainable = add_lora(model, args.rank, args.alpha)
model.to(device)
n_train = sum(p.numel() for p in trainable)
print(f'LoRA rank {args.rank}: {n_train / 1e3:.0f}k trainable of {model.num_params() / 1e6:.2f}M')
chats = [json.loads(line) for line in open(args.chats)]
random.Random(args.seed).shuffle(chats)
n_val = max(8, len(chats) // 20)
val_x, val_m = pack([render(tok, c['messages']) for c in chats[:n_val]], cfg.context)
train_x, train_m = pack([render(tok, c['messages']) for c in chats[n_val:]], cfg.context)
opt = torch.optim.AdamW(trainable, lr=args.lr, weight_decay=0.0)
ctx = torch.autocast(device_type=device, dtype=torch.bfloat16)
steps = max(1, int(args.epochs * len(train_x) / args.batch))
log = open(f'{args.out}/log.jsonl', 'w')
log.write(json.dumps({'config': cfg.to_dict(), 'args': vars(args), 'params': model.num_params() - n_train,
'trainable': n_train, 'steps': steps, 'conversations': len(chats)}) + '\n')
@torch.no_grad()
def evaluate():
model.eval()
x, m = val_x.to(device), val_m.to(device)
with ctx:
_, loss = model(x[:, :-1], x[:, 1:], loss_mask=m[:, 1:])
model.train()
return loss.item()
t0 = time.time()
for step in range(steps + 1):
if step % 10 == 0 or step == steps:
rec = {'step': step, 'val': round(evaluate(), 4), 'time': round(time.time() - t0, 1)}
print(json.dumps(rec), flush=True)
log.write(json.dumps(rec) + '\n')
if step == steps:
break
idx = torch.randint(0, len(train_x), (args.batch,))
x, m = train_x[idx].to(device), train_m[idx].to(device)
with ctx:
_, loss = model(x[:, :-1], x[:, 1:], loss_mask=m[:, 1:])
loss.backward()
for g in opt.param_groups:
g['lr'] = args.lr * lr_factor(step, steps, warmup=10, cooldown=0.5)
opt.step()
opt.zero_grad(set_to_none=True)
log.write(json.dumps({'step': step, 'loss': round(loss.item(), 4)}) + '\n')
log.close()
adapter = {n: p.detach().cpu() for n, p in model.named_parameters() if p.requires_grad}
torch.save({'rank': args.rank, 'alpha': args.alpha, 'adapter': adapter}, f'{args.out}/adapter.pt')
merge_lora(model)
torch.save({'config': cfg.to_dict(), 'model': model.state_dict(), 'step': steps}, f'{args.out}/ckpt_final.pt')
if __name__ == '__main__':
main()
snapshots.py
Маленькие модели из первых глав: биграммы, нейросети на буквах и токенах.
148 строк
· Разбираем в главе 1
· Открыть как текст
"""The small models the course shows along the way (the "Sprout now" cards).
python snapshots.py bigram # letter bigram: counts -> probabilities
python snapshots.py mlp3 # 3 letters -> next letter (embeddings + hidden layer)
python snapshots.py mlp8 # 8 letters -> next letter (two hidden layers)
python snapshots.py emb2d # the same idea with 2-D embeddings, to draw them
python snapshots.py mlptok # 8 BPE tokens -> next token
Writes browser files into --out (weights in the same SPRT format as export.py).
"""
import argparse
import json
import math
import os
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from export import pack
# id 0 = start/end of a text, then '\n' and the 95 printable ASCII characters (as chars.js)
ALPHABET = '\n' + ''.join(chr(i) for i in range(32, 127))
STOI = {c: i + 1 for i, c in enumerate(ALPHABET)}
V = len(ALPHABET) + 1
def char_corpus(data, n_chars):
"""A slice of the pre-training text, as letters (decoded from the token file)."""
from tokenizer import Tokenizer
tok = Tokenizer.load(f'{data}/tokenizer.json')
ids = np.memmap(f'{data}/train.bin', dtype=np.uint16, mode='r')[: n_chars // 3]
eot = tok['<|endoftext|>']
text = tok.decode([int(i) for i in ids]).replace('<|endoftext|>', '\0')
out = np.array([0 if c == '\0' else STOI.get(c, -1) for c in text], dtype=np.int64)
return out[out >= 0]
def bigram(args):
ids = char_corpus(args.data, args.chars)
counts = np.zeros((V, V), dtype=np.int64)
np.add.at(counts, (ids[:-1], ids[1:]), 1)
probs = (counts + 1) / (counts + 1).sum(1, keepdims=True)
logits = np.log(probs).astype(np.float32)
state = {'embed.weight': torch.from_numpy(logits)}
cfg = {'vocab_size': V, 'context': 1}
pack(state, cfg, {'kind': 'mlp', 'name': 'bigram-char', 'chars': int(len(ids))}, f'{args.out}/models/bigram-char.bin', quantize=False)
with open(f'{args.out}/data/bigram-counts.json', 'w') as f:
json.dump({'alphabet': ALPHABET, 'counts': counts.tolist(), 'chars': int(len(ids))}, f, separators=(',', ':'))
nll = -np.log(probs[ids[:-1], ids[1:]]).mean()
print(f'bigram: {len(ids)} letters, loss {nll:.4f} nats/letter')
class MLP(nn.Module):
def __init__(self, vocab, context, d_embed, hidden):
super().__init__()
self.context = context
self.embed = nn.Embedding(vocab, d_embed)
dims = [context * d_embed] + hidden
self.hidden = nn.ModuleList(nn.Linear(a, b) for a, b in zip(dims, dims[1:]))
self.out = nn.Linear(dims[-1], vocab)
def forward(self, x): # x: (batch, context)
h = self.embed(x).flatten(1)
for layer in self.hidden:
h = torch.tanh(layer(h))
return self.out(h)
def windows(ids, context, pad=0):
"""(context letters, next letter) pairs; the window restarts at every text start."""
return torch.from_numpy(np.lib.stride_tricks.sliding_window_view(np.concatenate([np.full(context, pad), ids]), context + 1).copy())
def train_mlp(name, ids, vocab, context, d_embed, hidden, args, lr=3e-3, steps=None, batch=512):
torch.manual_seed(0)
data = windows(ids, context)
n_val = len(data) // 50
val, tr = data[:n_val], data[n_val:]
model = MLP(vocab, context, d_embed, hidden)
opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
steps = steps or args.steps
t0 = time.time()
log = []
for step in range(steps + 1):
if step % 500 == 0 or step == steps:
with torch.no_grad():
idx = torch.randint(0, len(val), (8192,), generator=torch.Generator().manual_seed(0))
vb = val[idx]
vl = F.cross_entropy(model(vb[:, :-1]), vb[:, -1]).item()
log.append({'step': step, 'val': round(vl, 4)})
print(f'{name} step {step} val {vl:.4f} ({time.time() - t0:.0f}s)', flush=True)
if step == steps:
break
b = tr[torch.randint(0, len(tr), (batch,))]
loss = F.cross_entropy(model(b[:, :-1]), b[:, -1])
opt.zero_grad(set_to_none=True)
loss.backward()
opt.step()
for g in opt.param_groups:
g['lr'] = lr * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * step / steps)))
return model, log
def export_mlp(model, name, cfg, args, log):
state = {k: v for k, v in model.state_dict().items()}
size = pack(state, cfg, {'kind': 'mlp', 'name': name, 'log': log}, f'{args.out}/models/{name}.bin', quantize=True)
print(f'{name}: {sum(p.numel() for p in model.parameters()) / 1e3:.0f}k params, {size / 1e3:.0f} KB')
def main():
ap = argparse.ArgumentParser()
ap.add_argument('what', choices=['bigram', 'mlp3', 'mlp8', 'emb2d', 'mlptok'])
ap.add_argument('--data', default='data')
ap.add_argument('--out', default='../../public/llm')
ap.add_argument('--chars', type=int, default=30_000_000)
ap.add_argument('--steps', type=int, default=20000)
ap.add_argument('--threads', type=int, default=4)
args = ap.parse_args()
torch.set_num_threads(args.threads)
os.makedirs(f'{args.out}/models', exist_ok=True)
os.makedirs(f'{args.out}/data', exist_ok=True)
if args.what == 'bigram':
bigram(args)
elif args.what in ('mlp3', 'mlp8', 'emb2d'):
ids = char_corpus(args.data, args.chars)
spec = {'mlp3': (3, 16, [256]), 'mlp8': (8, 24, [512, 512]), 'emb2d': (3, 2, [128])}[args.what]
context, d_embed, hidden = spec
name = {'mlp3': 'mlp-char-3', 'mlp8': 'mlp-char-8', 'emb2d': 'emb2d'}[args.what]
model, log = train_mlp(name, ids, V, context, d_embed, hidden, args)
if args.what == 'emb2d':
emb = model.embed.weight.detach().numpy()
with open(f'{args.out}/data/emb2d.json', 'w') as f:
json.dump({'alphabet': ALPHABET, 'xy': emb.round(4).tolist(), 'log': log}, f, separators=(',', ':'))
cfg = {'vocab_size': V, 'context': context, 'd_embed': d_embed, 'hidden': hidden}
export_mlp(model, name, cfg, args, log)
elif args.what == 'mlptok':
ids = np.memmap(f'{args.data}/train.bin', dtype=np.uint16, mode='r')[:40_000_000].astype(np.int64)
model, log = train_mlp('mlp-token', ids, 8192, 8, 48, [384], args, lr=2e-3, batch=512)
cfg = {'vocab_size': 8192, 'context': 8, 'd_embed': 48, 'hidden': [384]}
export_mlp(model, 'mlp-token', cfg, args, log)
if __name__ == '__main__':
main()
course_data.py
JSON-файлы, из которых рисуют виджеты учебника.
110 строк
· Открыть как текст
"""Small JSON files the course widgets draw from (public/llm/data/*.json).
python course_data.py --work ~/claude-projects/tinyllm-work --out ../../public/llm/data
"""
import argparse
import collections
import glob
import json
import os
import random
import numpy as np
from prepare import SOURCES, clean
from tokenizer import Tokenizer
def corpus(work, out, tok):
stats = json.load(open(f'{work}/data/corpus_stats.json'))
rng = random.Random(3)
samples = {}
lengths = {}
for name, read in SOURCES.items():
docs, n = [], 0
for text in read(f'{work}/raw'):
n += 1
if n % 997 == 0:
t = clean(text)
if t:
docs.append(t)
if len(docs) >= 400:
break
lens = [len(tok.encode(d, allow_special=False)) for d in docs]
lengths[name] = np.histogram(lens, bins=20, range=(0, 800))[0].tolist()
samples[name] = [d[:1400] for d in rng.sample(docs, 24)]
ids = np.memmap(f'{work}/data/val.bin', dtype=np.uint16, mode='r')
counts = collections.Counter(ids.tolist())
top = [{'id': int(i), 'text': tok.decode([int(i)]), 'n': int(c)} for i, c in counts.most_common(400)]
ranks = sorted(counts.values(), reverse=True)
zipf = [[r + 1, int(ranks[r])] for r in sorted({int(x) for x in np.unique(np.geomspace(1, len(ranks), 60).astype(int) - 1)})]
json.dump({'stats': stats, 'samples': samples, 'lengths': lengths, 'length_bins': 40, 'top': top, 'zipf': zipf,
'val_tokens': int(len(ids))}, open(f'{out}/corpus.json', 'w'), separators=(',', ':'))
def train_log(work, out, run='base', name='train-log'):
path = f'{work}/runs/{run}/log.jsonl'
if not os.path.exists(path):
return
lines = [json.loads(l) for l in open(path) if l.strip()]
head = lines[0]
train = [{k: r[k] for k in ('step', 'loss', 'lr', 'tokens', 'T', 'norm') if k in r} for r in lines[1:] if 'loss' in r]
evals = [{k: r[k] for k in ('step', 'val', 'tokens', 'time', 'samples') if k in r} for r in lines[1:] if 'val' in r]
extra = {k: head[k] for k in ('train_rows', 'conversations', 'trainable') if k in head}
json.dump({'config': head.get('config'), 'args': head.get('args'), 'params': head.get('params'), 'steps': head.get('steps'),
**extra, 'train': train, 'evals': evals}, open(f'{out}/{name}.json', 'w'), separators=(',', ':'))
def lora_logs(work, out):
res = {}
for style in ('pirate', 'poet'):
path = f'{work}/runs/{style}/log.jsonl'
if not os.path.exists(path):
continue
lines = [json.loads(l) for l in open(path) if l.strip()]
head = lines[0]
evals = [r for r in lines[1:] if 'val' in r]
res[style] = {'rank': head['args']['rank'], 'alpha': head['args']['alpha'], 'trainable': head['trainable'],
'params': head['params'], 'steps': head['steps'], 'conversations': head.get('conversations'),
'minutes': round((evals[-1]['time'] if evals else 0) / 60, 1),
'train': [{'step': r['step'], 'loss': r['loss']} for r in lines[1:] if 'loss' in r],
'evals': [{'step': r['step'], 'val': r['val']} for r in evals]}
if res:
json.dump(res, open(f'{out}/train-log-lora.json', 'w'), separators=(',', ':'))
def chats(work, out):
path = f'{work}/data/chats.jsonl'
if not os.path.exists(path):
return
rows = [json.loads(l) for l in open(path)]
rng = random.Random(5)
by = collections.defaultdict(list)
for r in rows:
by[r['source']].append(r)
sample = {k: rng.sample(v, min(12, len(v))) for k, v in by.items()}
json.dump({'counts': {k: len(v) for k, v in by.items()}, 'samples': sample}, open(f'{out}/chats.json', 'w'), separators=(',', ':'))
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--work', default='.')
ap.add_argument('--out', default='../../public/llm/data')
ap.add_argument('--only', default='')
args = ap.parse_args()
os.makedirs(args.out, exist_ok=True)
tok = Tokenizer.load(f'{args.work}/data/tokenizer.json')
todo = args.only.split(',') if args.only else ['corpus', 'log', 'chats']
if 'corpus' in todo:
corpus(args.work, args.out, tok)
if 'log' in todo:
train_log(args.work, args.out)
for run in ('gpt1', 'gpt4', 'chat'):
train_log(args.work, args.out, run, f'train-log-{run}')
lora_logs(args.work, args.out)
if 'chats' in todo:
chats(args.work, args.out)
if __name__ == '__main__':
main()