05_train.py

"""
Phase 4: Train on TinyShakespeare
-----------------------------------
Downloads Shakespeare text, builds a character-level vocab,
and trains the dLLM to denoise masked sequences.

  python 05_train.py

Checkpoints saved to checkpoints/dllm.pt every 500 steps.
Training ~1-2 hrs on RTX 3050 for 5000 steps.
"""

import os
import glob
import math
import time
import json
import urllib.request
import torch
import torch.nn as nn
import torch.nn.functional as F


# ── Download data ─────────────────────────────────────────────────────────────
DATA_URL  = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
DATA_PATH = "data/shakespeare.txt"

def download_data():
    os.makedirs("data", exist_ok=True)
    if not os.path.exists(DATA_PATH):
        print("Downloading TinyShakespeare...")
        urllib.request.urlretrieve(DATA_URL, DATA_PATH)
        print("Done.")
    with open(DATA_PATH) as f:
        return f.read()


# ── Character-level tokenizer ─────────────────────────────────────────────────
class CharTokenizer:
    def __init__(self, text):
        chars = sorted(set(text))
        self.vocab_size = len(chars)
        self.c2i = {c: i for i, c in enumerate(chars)}
        self.i2c = {i: c for i, c in enumerate(chars)}

    def encode(self, text):
        return [self.c2i[c] for c in text]

    def decode(self, ids):
        return ''.join(self.i2c.get(i, '?') for i in ids)

    def save(self, path):
        json.dump(self.c2i, open(path, 'w'))

    @classmethod
    def load(cls, path):
        obj = cls.__new__(cls)
        obj.c2i = json.load(open(path))
        obj.i2c = {int(v): k for k, v in obj.c2i.items()}
        obj.vocab_size = len(obj.c2i)
        return obj


# ── Model (same as 04) ────────────────────────────────────────────────────────
class SelfAttention(nn.Module):
    def __init__(self, hidden, n_heads):
        super().__init__()
        self.n_heads  = n_heads
        self.head_dim = hidden // n_heads
        self.qkv = nn.Linear(hidden, 3 * hidden, bias=False)
        self.out  = nn.Linear(hidden, hidden, bias=False)
    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        def split(t): return t.view(B,T,self.n_heads,self.head_dim).transpose(1,2)
        q, k, v = split(q), split(k), split(v)
        scores = (q @ k.transpose(-2,-1)) / math.sqrt(self.head_dim)
        out = (F.softmax(scores,-1) @ v).transpose(1,2).contiguous().view(B,T,C)
        return self.out(out)

class FeedForward(nn.Module):
    def __init__(self, hidden):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(hidden, 4*hidden), nn.GELU(), nn.Linear(4*hidden, hidden))
    def forward(self, x): return self.net(x)

class TransformerBlock(nn.Module):
    def __init__(self, hidden, n_heads):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden)
        self.attn  = SelfAttention(hidden, n_heads)
        self.norm2 = nn.LayerNorm(hidden)
        self.ff    = FeedForward(hidden)
    def forward(self, x):
        return x + self.ff(self.norm2(x + self.attn(self.norm1(x))))

class TinyDLLM(nn.Module):
    def __init__(self, vocab_size, hidden=256, n_layers=4, n_heads=4, max_seq=128):
        super().__init__()
        self.mask_token_id = vocab_size
        full_vocab = vocab_size + 1
        self.token_emb = nn.Embedding(full_vocab, hidden)
        self.pos_emb   = nn.Embedding(max_seq, hidden)
        self.blocks    = nn.Sequential(*[TransformerBlock(hidden, n_heads) for _ in range(n_layers)])
        self.norm = nn.LayerNorm(hidden)
        self.head = nn.Linear(hidden, vocab_size, bias=False)
        self.head.weight = nn.Parameter(self.token_emb.weight[:vocab_size])
        for m in self.modules():
            if isinstance(m, (nn.Linear, nn.Embedding)):
                nn.init.normal_(m.weight, std=0.02)
    def forward(self, token_ids):
        B, T = token_ids.shape
        pos = torch.arange(T, device=token_ids.device)
        x = self.token_emb(token_ids) + self.pos_emb(pos)
        return self.head(self.norm(self.blocks(x)))


# ── Diffusion ─────────────────────────────────────────────────────────────────
class MaskedDiffusion:
    def __init__(self, mask_token_id):
        self.mask_id = mask_token_id

    def loss(self, model, tokens):
        B, T = tokens.shape
        t = torch.rand(B, device=tokens.device)
        mask_prob = t.unsqueeze(1).expand(B, T)
        mask = torch.bernoulli(mask_prob).bool()
        noisy = tokens.clone()
        noisy[mask] = self.mask_id
        logits = model(noisy)
        logits_masked = logits[mask]
        targets = tokens[mask]
        if logits_masked.numel() == 0:
            return torch.tensor(0.0, device=tokens.device)
        return F.cross_entropy(logits_masked, targets)

    @torch.no_grad()
    def sample(self, model, seq_len, n_steps=20, device='cpu'):
        model.eval()
        tokens = torch.full((1, seq_len), self.mask_id, dtype=torch.long, device=device)
        for step in range(n_steps):
            frac = (step + 1) / n_steps
            target = int(frac * seq_len)
            logits = model(tokens)
            probs  = F.softmax(logits, dim=-1)
            predicted = torch.multinomial(probs.view(seq_len, -1), 1).view(1, seq_len)
            confidence, _ = probs.max(dim=-1)
            still_masked = (tokens == self.mask_id)
            confidence[~still_masked] = -1.0
            already = (~still_masked).sum().item()
            to_unmask = max(0, target - already)
            if to_unmask > 0 and still_masked.any():
                _, idx = confidence.view(-1).topk(min(to_unmask, still_masked.sum().item()))
                flat = tokens.view(-1)
                flat[idx] = predicted.view(-1)[idx]
                tokens = flat.view(1, seq_len)
        model.train()
        return tokens


# ── Data loader ───────────────────────────────────────────────────────────────
def get_batch(data, seq_len, batch_size, device):
    ix = torch.randint(len(data) - seq_len, (batch_size,))
    x  = torch.stack([data[i:i+seq_len] for i in ix])
    return x.to(device)


# ── Training ──────────────────────────────────────────────────────────────────
def train():
    # config
    HIDDEN    = 384
    N_LAYERS  = 6
    N_HEADS   = 6
    SEQ_LEN   = 128
    BATCH     = 32
    LR        = 3e-4
    STEPS     = 50000
    EVAL_INT  = 100
    SAVE_INT  = 1000
    WARMUP    = 400

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Device: {device}")

    # data
    text      = download_data()
    tokenizer = CharTokenizer(text)
    data      = torch.tensor(tokenizer.encode(text), dtype=torch.long)
    split     = int(0.9 * len(data))
    train_data, val_data = data[:split], data[split:]

    print(f"Vocab size:  {tokenizer.vocab_size}")
    print(f"Train tokens: {len(train_data):,}")

    os.makedirs("checkpoints", exist_ok=True)
    tokenizer.save("checkpoints/vocab.json")

    # model
    model     = TinyDLLM(tokenizer.vocab_size, HIDDEN, N_LAYERS, N_HEADS, SEQ_LEN).to(device)
    diffusion = MaskedDiffusion(model.mask_token_id)
    optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.1)

    # resume from latest checkpoint if available
    ckpts = sorted(glob.glob('checkpoints/dllm_step*.pt'), key=lambda p: int(p.split('step')[1].split('.')[0]))
    start_step = 0
    if ckpts:
        ckpt = torch.load(ckpts[-1], map_location=device, weights_only=True)
        model.load_state_dict(ckpt['model'])
        optimizer.load_state_dict(ckpt['optimizer'])
        start_step = ckpt['step']
        print(f"Resuming from step {start_step}")

    total_params = sum(p.numel() for p in model.parameters())
    print(f"Parameters: {total_params:,}")

    # linear warmup + cosine decay schedule
    def lr_schedule(step):
        if step < WARMUP:
            return step / WARMUP
        progress = (step - WARMUP) / (STEPS - WARMUP)
        return 0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * progress))

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_schedule)

    model.train()
    t0 = time.time()

    for step in range(start_step + 1, STEPS + 1):
        batch = get_batch(train_data, SEQ_LEN, BATCH, device)
        loss  = diffusion.loss(model, batch)

        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()

        if step % EVAL_INT == 0:
            # val loss
            with torch.no_grad():
                val_batch = get_batch(val_data, SEQ_LEN, BATCH, device)
                val_loss  = diffusion.loss(model, val_batch)

            elapsed = time.time() - t0
            lr_now  = scheduler.get_last_lr()[0] * LR
            print(f"step {step:5d} | train {loss.item():.4f} | val {val_loss.item():.4f} "
                  f"| lr {lr_now:.2e} | {elapsed:.0f}s")

            # sample
            ids  = diffusion.sample(model, seq_len=80, n_steps=20, device=device)
            text_out = tokenizer.decode(ids[0].tolist())
            print(f"  sample: {repr(text_out[:80])}")
            t0 = time.time()

        if step % SAVE_INT == 0:
            torch.save({
                'step': step,
                'model': model.state_dict(),
                'optimizer': optimizer.state_dict(),
                'vocab_size': tokenizer.vocab_size,
                'hidden': HIDDEN, 'n_layers': N_LAYERS,
                'n_heads': N_HEADS, 'seq_len': SEQ_LEN,
            }, f"checkpoints/dllm_step{step}.pt")
            print(f"  ✅ Saved checkpoint at step {step}")


if __name__ == '__main__':
    train()