| """Small LLM1: ~5M parameter Llama-style model trained from scratch on TinyStories.""" |
| import math, os, time |
| import torch |
| from datasets import load_dataset |
| from tokenizers import Tokenizer, models, pre_tokenizers, decoders, trainers, processors |
| from transformers import LlamaConfig, LlamaForCausalLM, PreTrainedTokenizerFast |
|
|
| OUT = "small-llm1" |
| N_STORIES = 400_000 |
| VOCAB = 4096 |
| SEQ = 256 |
| BATCH = 64 |
| STEPS = 6000 |
| LR = 2e-3 |
| torch.manual_seed(0) |
|
|
| |
| print("loading data...") |
| ds = load_dataset("roneneldan/TinyStories", split="train", streaming=True) |
| texts = [ex["text"].strip() for _, ex in zip(range(N_STORIES), ds) if ex["text"].strip()] |
| print(f"{len(texts)} stories") |
|
|
| |
| tok = Tokenizer(models.BPE()) |
| tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) |
| tok.decoder = decoders.ByteLevel() |
| trainer = trainers.BpeTrainer(vocab_size=VOCAB, special_tokens=["<|endoftext|>"], |
| initial_alphabet=pre_tokenizers.ByteLevel.alphabet()) |
| tok.train_from_iterator(texts[:100_000], trainer) |
| tok.post_processor = processors.ByteLevel(trim_offsets=False) |
| hf_tok = PreTrainedTokenizerFast(tokenizer_object=tok, bos_token="<|endoftext|>", |
| eos_token="<|endoftext|>", unk_token="<|endoftext|>", |
| pad_token="<|endoftext|>") |
| eos = hf_tok.eos_token_id |
|
|
| ids = [] |
| for i in range(0, len(texts), 10_000): |
| for seq in hf_tok(texts[i:i + 10_000])["input_ids"]: |
| ids.extend(seq + [eos]) |
| data = torch.tensor(ids, dtype=torch.long) |
| n_val = SEQ * 2000 |
| train_data, val_data = data[:-n_val], data[-n_val:] |
| print(f"{len(data)/1e6:.1f}M tokens") |
|
|
| |
| config = LlamaConfig(vocab_size=VOCAB, hidden_size=256, intermediate_size=512, |
| num_hidden_layers=6, num_attention_heads=8, num_key_value_heads=8, |
| max_position_embeddings=SEQ, tie_word_embeddings=True, |
| bos_token_id=eos, eos_token_id=eos, pad_token_id=eos) |
| dev = "cuda" |
| model = LlamaForCausalLM(config).to(dev) |
| n_params = sum(p.numel() for p in model.parameters()) |
| print(f"params: {n_params/1e6:.2f}M") |
|
|
| opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.1, betas=(0.9, 0.95), fused=True) |
| warmup = 200 |
| sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1, (s + 1) / warmup) * |
| (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(s, STEPS) / STEPS)))) |
|
|
| def batch(src): |
| ix = torch.randint(len(src) - SEQ - 1, (BATCH,)) |
| x = torch.stack([src[i:i + SEQ] for i in ix]).to(dev, non_blocking=True) |
| return x |
|
|
| @torch.no_grad() |
| def evaluate(): |
| model.eval() |
| losses = [] |
| for i in range(0, len(val_data) - SEQ, SEQ * BATCH): |
| x = torch.stack([val_data[j:j + SEQ] for j in range(i, min(i + SEQ * BATCH, len(val_data) - SEQ), SEQ)]).to(dev) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| losses.append(model(input_ids=x, labels=x).loss.item()) |
| model.train() |
| return sum(losses) / len(losses) |
|
|
| t0 = time.time() |
| model.train() |
| for step in range(STEPS): |
| x = batch(train_data) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| loss = model(input_ids=x, labels=x).loss |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| opt.step(); sched.step(); opt.zero_grad(set_to_none=True) |
| if step % 500 == 0 or step == STEPS - 1: |
| print(f"step {step} loss {loss.item():.3f} lr {sched.get_last_lr()[0]:.2e} {time.time()-t0:.0f}s", flush=True) |
|
|
| val_loss = evaluate() |
| print(f"val loss {val_loss:.3f} ppl {math.exp(val_loss):.2f} time {time.time()-t0:.0f}s") |
|
|
| model.generation_config.pad_token_id = eos |
| model.save_pretrained(OUT) |
| hf_tok.save_pretrained(OUT) |
| with open(os.path.join(OUT, "train_stats.txt"), "w") as f: |
| f.write(f"params={n_params}\ntokens={len(data)}\nsteps={STEPS}\nval_loss={val_loss:.4f}\n") |
| print("saved to", OUT) |
|
|
