# Optimal parameters for given compute
def chinchilla_optimal_params(compute):
    return 0.074 * compute**0.73  # ~70B for 1.4T tokens

# Optimal tokens
def chinchilla_optimal_tokens(compute):
    return 19.3 * compute**0.27  # ~1.4T for 70B model

Data Parallel (DP): 
  8 GPUs → 8x batch → same model

Tensor Parallel (TP): 
  Layer split across 4 GPUs → W_q on GPU0, W_k on GPU1

Pipeline Parallel (PP): 
  Layers 1–4 on GPU0, 5–8 on GPU1

attn = softmax(Q @ K.T / sqrt(d)) @ V
# → Materialize N×N matrix → O(N²) memory

# Online softmax + tiling
for i in blocks:
    Q_block = Q[i]
    for j in blocks:
        K_block, V_block = K[j], V[j]
        S = Q_block @ K_block.T
        P = softmax(S)
        O += P @ V_block

from flash_attn import flash_attention

attn_output = flash_attention(q, k, v, causal=True)

W = W₀ + ΔW
ΔW = B A    # B: (d, r), A: (r, k) → r << d

class LoRALinear(nn.Module):
    def __init__(self, linear, rank=8):
        super().__init__()
        self.linear = linear
        d = linear.in_features
        self.A = nn.Parameter(torch.randn(rank, d) * 0.01)
        self.B = nn.Parameter(torch.zeros(d, rank))

    def forward(self, x):
        return self.linear(x) + (x @ self.A.T @ self.B.T)

from transformers import AutoModelForCausalLM
import peft
from flash_attn import flash_attention

# Load base model
model = AutoModelForCausalLM.from_pretrained("gpt2")

# Add LoRA
lora_config = peft.LoraConfig(
    r=8, lora_alpha=32, target_modules=["c_attn", "c_proj"], lora_dropout=0.1
)
model = peft.get_peft_model(model, lora_config)

# Use FlashAttention
def forward_with_flash(self, x):
    q, k, v = self.W_q(x), self.W_k(x), self.W_v(x)
    return flash_attention(q, k, v, causal=True)

# Monkey patch
model.transformer.h[0].attn.forward = forward_with_flash

def transformer_flops(batch, seq_len, d_model, layers, vocab):
    # Embedding
    flops = batch * seq_len * vocab * d_model

    # Per layer
    attn = 2 * batch * seq_len**2 * d_model
    ffn = 8 * batch * seq_len * d_model**2
    flops += layers * (attn + ffn)

    # Output
    flops += batch * seq_len * d_model * vocab
    return flops

print(f"GPT-3 175B: {transformer_flops(1, 2048, 12288, 96, 50257):.2e} FLOPs")
# → 3.7e23 FLOPs

# Tensor Parallel (simplified)
class TensorParallelLinear(nn.Module):
    def __init__(self, in_features, out_features, world_size):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_features//world_size, in_features))
        self.world_size = world_size

    def forward(self, x):
        out = x @ self.weight.t()
        # All-gather across GPUs
        return all_gather(out, dim=-1)

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)

# Gradient Checkpointing: trade compute for memory
model = torch.utils.checkpoint.checkpoint_sequential(model, segments=4)

!pip install flash-attn peft transformers

import torch
from transformers import AutoModelForCausalLM
from peft import LoraConfig, get_peft_model
from flash_attn import flash_attention

# Load model
model = AutoModelForCausalLM.from_pretrained("gpt2")

# LoRA
config = LoraConfig(r=8, lora_alpha=32, target_modules=["c_attn", "c_proj"])
model = get_peft_model(model, config)

# Replace attention
def flash_forward(self, x):
    q, k, v = self.c_attn(x).chunk(3, dim=-1)
    return self.c_proj(flash_attention(q, k, v, causal=True))

# Patch first layer
model.transformer.h[0].attn.forward = flash_forward.__get__(model.transformer.h[0].attn)

print(f"Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")