LLMs with PyTorch

General

A GPT-style LLM is a decoder-only Transformer trained on next-token prediction: given tokens t₁…tₙ, predict tₙ₊₁. The loss is cross-entropy between the logits and the inputs shifted by one. Finetuning reuses the same loss on curated data.

Params ≈ 12·L·d² + V·d. Shapes flow as [B, T] → [B, T, d] → [B, T, V].

Building blocks

Tokenize

Text → sub-word ids with BPE. Always use the tokenizer and chat template that match the model.

from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
ids = tok("Hello world", return_tensors="pt").input_ids  # [1,T]
msgs = [{"role": "user", "content": "Explain LoRA"}]
prompt = tok.apply_chat_template(msgs, tokenize=False,
                                 add_generation_prompt=True)
tok.pad_token = tok.pad_token or tok.eos_token  # for batching

Small open models (Hugging Face IDs)

Generate

inputs = tok(prompt, return_tensors="pt").to(dev)
out = model.generate(**inputs, max_new_tokens=200,
                     do_sample=True, temperature=0.7, top_p=0.9)
print(tok.decode(out[0], skip_special_tokens=True))

Build a tiny GPT

import torch, torch.nn as nn, torch.nn.functional as F

class Block(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.h, self.n1, self.n2 = h, nn.RMSNorm(d), nn.RMSNorm(d)
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.proj = nn.Linear(d, d, bias=False)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(),
                                 nn.Linear(4 * d, d))
    def forward(self, x):                        # x: [B, T, d]
        B, T, C = x.shape
        q, k, v = self.qkv(self.n1(x)).split(C, dim=-1)
        q, k, v = (t.view(B, T, self.h, -1).transpose(1, 2)
                   for t in (q, k, v))           # [B, h, T, d/h]
        a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        x = x + self.proj(a.transpose(1, 2).reshape(B, T, C))
        return x + self.mlp(self.n2(x))
class TinyGPT(nn.Module):
    def __init__(self, V, d=384, h=6, L=6, T=256):
        super().__init__()
        self.tok, self.pos = nn.Embedding(V, d), nn.Embedding(T, d)
        self.blocks = nn.Sequential(*(Block(d, h) for _ in range(L)))
        self.norm = nn.RMSNorm(d)
        self.head = nn.Linear(d, V, bias=False)
        self.head.weight = self.tok.weight       # weight tying
    def forward(self, idx, targets=None):        # idx: [B, T]
        pos = torch.arange(idx.size(1), device=idx.device)
        x = self.blocks(self.tok(idx) + self.pos(pos))
        logits = self.head(self.norm(x))         # [B, T, V]
        if targets is None: return logits, None
        loss = F.cross_entropy(logits.flatten(0, 1),
                               targets.flatten(), ignore_index=-100)
        return logits, loss

Pretraining loop

data = torch.tensor(tok.encode(text))        # 1-D token stream
def get_batch(B=32, T=256):
    i = torch.randint(len(data) - T - 1, (B,))
    x = torch.stack([data[j:j + T] for j in i])
    y = torch.stack([data[j + 1:j + T + 1] for j in i])
    return x.to(dev), y.to(dev)
model = torch.compile(TinyGPT(V=len(tok)).to(dev))
opt = torch.optim.AdamW(model.parameters(), lr=3e-4,
                        betas=(0.9, 0.95), weight_decay=0.1)
for step in range(5000):
    x, y = get_batch()
    with torch.autocast("cuda", dtype=torch.bfloat16):
        _, loss = model(x, y)
    loss.backward()
    nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step(); sched.step(); opt.zero_grad(set_to_none=True)

y is x shifted by one token. Use warmup + cosine decay for sched. Track perplexity = exp(val loss); random-init loss ≈ ln V.

Efficiency & memory tricks

Finetuning

GPU memory per parameter (training)

LoRA

h = W·x + (α/r)·B·A·x, A: r×k, B: d×r, r ≪ d. B starts at 0, so step 0 equals the base model.

class LoRALinear(nn.Module):
    def __init__(self, base, r=8, alpha=16):
        super().__init__()
        self.base = base.requires_grad_(False)   # freeze W
        self.A = nn.Linear(base.in_features, r, bias=False)
        self.B = nn.Linear(r, base.out_features, bias=False)
        nn.init.zeros_(self.B.weight)            # dW = 0 at start
        self.s = alpha / r
    def forward(self, x):
        return self.base(x) + self.s * self.B(self.A(x))
from peft import LoraConfig, get_peft_model   # with PEFT
cfg = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05,
        target_modules="all-linear", task_type="CAUSAL_LM")
model = get_peft_model(model, cfg)
model.print_trainable_parameters()

Typical: r 8–64, α = r or 2r, lr 1e-4–2e-4 (≈10× full-FT lr).

QLoRA

from transformers import AutoModelForCausalLM
from transformers import BitsAndBytesConfig as BnB
from peft import prepare_model_for_kbit_training
bnb = BnB(load_in_4bit=True,
    bnb_4bit_quant_type="nf4",          # NormalFloat4
    bnb_4bit_use_double_quant=True,     # quantize scales too
    bnb_4bit_compute_dtype=torch.bfloat16)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B",
    quantization_config=bnb, device_map="auto")
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, cfg)      # same LoRA config

SFT, save & merge