स्टेप 6 · ट्रेनिंग लूप जो कन्वर्ज होता है

import math, torch

model = TinyGPT(Config()).to(device)
model = torch.compile(model)          # big speedup, one line

# weight decay on matrices, none on norms and biases
decay = [p for n, p in model.named_parameters() if p.dim() >= 2]
nodecay = [p for n, p in model.named_parameters() if p.dim() < 2]
opt = torch.optim.AdamW(
    [{"params": decay, "weight_decay": 0.1},
     {"params": nodecay, "weight_decay": 0.0}],
    lr=6e-4, betas=(0.9, 0.95), eps=1e-8, fused=True)

max_steps, warmup, min_lr = 20000, 500, 6e-5
def lr_at(step):
    if step < warmup:                                   # linear warmup
        return 6e-4 * (step + 1) / warmup
    r = (step - warmup) / (max_steps - warmup)          # cosine decay
    return min_lr + 0.5 * (6e-4 - min_lr) * (1 + math.cos(math.pi * r))

scaler_dtype = torch.bfloat16
accum = 8                                               # gradient accumulation

for step in range(max_steps):
    for g in opt.param_groups:
        g["lr"] = lr_at(step)

    opt.zero_grad(set_to_none=True)
    for micro in range(accum):                          # simulate a big batch
        x, y = get_batch("train")
        with torch.autocast(device_type="cuda", dtype=scaler_dtype):
            _, loss = model(x, y)
        (loss / accum).backward()

    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()

    if step % 100 == 0:
        print(f"step {step}  loss {loss.item():.3f}  lr {lr_at(step):.2e}  gnorm {norm:.2f}")

def save(step):
    torch.save({
        "step": step,
        "model": model.state_dict(),
        "opt": opt.state_dict(),        # optimiser state is essential
        "config": cfg.__dict__,
        "rng": torch.get_rng_state(),
    }, f"ckpt_{step}.pt")