Transformer

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# 1) Sinusoidal positional encoding (the original Transformer)
class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)

    def forward(self, x):                       # x: (B, T, D)
        return x + self.pe[: x.size(1)]

# 2) Encoder-only classifier (BERT-style)
class TextClassifier(nn.Module):
    def __init__(self, vocab_size, num_classes, d_model=256, nhead=8, num_layers=4, dropout=0.1):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.pos   = PositionalEncoding(d_model)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=4 * d_model,
            dropout=dropout,
            batch_first=True,
            norm_first=True,                    # pre-LN (better gradient flow)
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.classifier = nn.Linear(d_model, num_classes)

    def forward(self, ids, pad_mask=None):      # ids: (B, T)
        x = self.embed(ids) * math.sqrt(self.embed.embedding_dim)
        x = self.pos(x)
        x = self.encoder(x, src_key_padding_mask=pad_mask)
        cls = x[:, 0]                            # use the first token like BERT's [CLS]
        return self.classifier(cls)

model = TextClassifier(vocab_size=30_000, num_classes=4)
print(sum(p.numel() for p in model.parameters()) / 1e6, 'M params')

# 3) Custom self-attention from scratch (for understanding)
class SelfAttention(nn.Module):
    def __init__(self, d_model, nhead):
        super().__init__()
        self.nhead = nhead
        self.d_head = d_model // nhead
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.out = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):            # x: (B, T, D)
        B, T, D = x.shape
        qkv = self.qkv(x).reshape(B, T, 3, self.nhead, self.d_head).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]         # (B, h, T, d_head)

        attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head)
        if mask is not None:
            attn = attn.masked_fill(mask, float('-inf'))
        attn = F.softmax(attn, dim=-1)

        out = (attn @ v).transpose(1, 2).reshape(B, T, D)
        return self.out(out)

# 4) Causal mask — for decoder/GPT-style models
def causal_mask(t):
    return torch.triu(torch.ones(t, t, dtype=torch.bool), diagonal=1)
# Apply: mask = causal_mask(T).to(device)

# 5) Train
optim   = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
sched   = torch.optim.lr_scheduler.CosineAnnealingLR(optim, T_max=epochs)
loss_fn = nn.CrossEntropyLoss(ignore_index=PAD_ID)
scaler  = torch.amp.GradScaler('cuda')

for epoch in range(epochs):
    for ids, labels, mask in train_dl:
        ids, labels, mask = ids.to(device), labels.to(device), mask.to(device)
        optim.zero_grad(set_to_none=True)
        with torch.amp.autocast('cuda', dtype=torch.bfloat16):
            logits = model(ids, pad_mask=mask)
            loss = loss_fn(logits, labels)
        scaler.scale(loss).backward()
        scaler.unscale_(optim)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optim); scaler.update()
    sched.step()

# 6) FlashAttention via SDPA — PyTorch picks the fastest kernel
# scaled_dot_product_attention dispatches to flash / mem-efficient when available
from torch.nn.functional import scaled_dot_product_attention
out = scaled_dot_product_attention(q, k, v, is_causal=True)

# 7) Real-world
# - For LLMs / NLP, use Hugging Face Transformers (`pip install transformers`)
# - Vision: Vision Transformers via `torchvision.models.vit_b_16(weights=...)`
# - Fine-tune with LoRA / QLoRA (peft) for parameter-efficient training

# 8) Tips
# - batch_first=True everywhere — consistency with most other libraries
# - norm_first=True (pre-LN) trains more stably than post-LN
# - Use scaled_dot_product_attention — automatic flash/efficient backends
# - Pad-token masking critical for variable-length batches

import torch.nn as nn
block = nn.TransformerEncoderLayer(d_model=512, nhead=8, dim_feedforward=2048, batch_first=True)
enc   = nn.TransformerEncoder(block, num_layers=6)