人人可懂的 AI 第1篇 — 用15分钟从零训练一个1600万参数的语言模型

from transformers import AutoTokenizer
from datasets import load_dataset

tok = AutoTokenizer.from_pretrained("openai-community/gpt2")
ds = load_dataset("roneneldan/TinyStories", split="train[:40000]")

ids = []
for ex in ds:
    ids.extend(tok(ex["text"]).input_ids + [tok.eos_token_id])
    if len(ids) >= 4_000_000:
        break
data = torch.tensor(ids[:4_000_000], dtype=torch.long)

CTX, DIM, LAYERS, HEADS = 128, 256, 4, 8

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(DIM), nn.LayerNorm(DIM)
        self.attn = nn.MultiheadAttention(DIM, HEADS, batch_first=True)
        self.mlp = nn.Sequential(
            nn.Linear(DIM, 4 * DIM), nn.GELU(), nn.Linear(4 * DIM, DIM)
        )

    def forward(self, x, mask):
        h = self.ln1(x)
        x = x + self.attn(h, h, h, attn_mask=mask, need_weights=False)[0]
        return x + self.mlp(self.ln2(x))


class TinyGPT(nn.Module):
    def __init__(self, vocab):
        super().__init__()
        self.emb = nn.Embedding(vocab, DIM)
        self.pos = nn.Embedding(CTX, DIM)
        self.blocks = nn.ModuleList(Block() for _ in range(LAYERS))
        self.ln = nn.LayerNorm(DIM)
        mask = torch.triu(torch.full((CTX, CTX), float("-inf")), 1)
        self.register_buffer("mask", mask)

    def forward(self, x):
        T = x.shape[1]
        h = self.emb(x) + self.pos(torch.arange(T, device=x.device))
        for b in self.blocks:
            h = b(h, self.mask[:T, :T])
        return self.ln(h) @ self.emb.weight.T  # weight tying

model = TinyGPT(len(tok)).to("cuda")
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)

step = 0
while not run.over_budget():          # 15分钟预算
    ix = torch.randint(0, len(data) - CTX - 1, (64,))
    x = torch.stack([data[i:i + CTX] for i in ix]).to("cuda")
    y = torch.stack([data[i + 1:i + CTX + 1] for i in ix]).to("cuda")
    with torch.autocast("cuda", torch.bfloat16):
        loss = F.cross_entropy(model(x).flatten(0, 1), y.flatten())
    opt.zero_grad(); loss.backward(); opt.step()
    step += 1

[    0.6s] step      0 loss=168.4468 ppl=...
[   38.2s] step   1000 loss=5.2213  ppl=185.13
[  114.7s] step   3000 loss=3.6841  ppl=39.81
[  305.9s] step   6700 loss=3.0350  ppl=20.80
[  662.6s] step  16100 loss=2.2882  ppl=9.86
[  901.2s] step  22358 loss=2.0948  ppl=8.12

PROMPT: Once upon a time
Once upon a time, there was a little boy named Timmy. Timmy liked to play
with his toy cars instead. One day, Timmy wanted to play with the red cars,
but he couldn't find it.

As Timmy was leaving a fight on the slide, he saw a big tree. Timmy thought
it looked like fun, so he started to tremble. He mixed the leaves and

PROMPT: Lily went to the
Lily went to the park with her mom. They saw a big tree, a hole and a hole.
They were curious and wanted to see what was inside.

They heard a hole in the hole. It was a zipper, but it was in the hole.