File size: 5,115 Bytes
a8f07a3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """Find the optimal training config: batch, seq, and fused CE.
The bottleneck is the logits tensor: [B, T, 50257] = 1.6 GB at batch=16, seq=1024.
The cross-entropy + backward through this is expensive.
Optimizations to test:
1. Larger batch + shorter seq (amortize weight reads)
2. Longer seq + smaller batch (more tokens per forward)
3. Fused cross-entropy (avoid materializing full logits)
4. Gradient checkpointing (trade compute for memory, fit bigger batches)
"""
import time
import torch
import torch.nn.functional as F
from config import Config
from model import build_model
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
cfg = Config.v5_500m()
model = build_model(cfg, "cuda")
model.train()
# Freeze medusa
for name, param in model.named_parameters():
if name.startswith("medusa_"):
param.requires_grad = False
opt = torch.optim.AdamW(
[p for p in model.parameters() if p.requires_grad],
lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1, fused=True)
print(f"Model: {sum(p.numel() for p in model.parameters())/1e6:.1f}M params")
print(f"Trainable: {sum(p.numel() for p in model.parameters() if p.requires_grad)/1e6:.1f}M")
print()
# --- Standard CE vs Fused CE ---
def standard_ce(model, ids, cfg):
h = model(ids)
logits = model.lm_head(h)
loss = F.cross_entropy(
logits[:, :-1].reshape(-1, cfg.vocab_size),
ids[:, 1:].reshape(-1))
return loss
def fused_ce(model, ids, cfg):
"""Fused CE: use F.linear + F.cross_entropy without materializing full logits.
Actually PyTorch's F.cross_entropy already fuses the softmax + NLL.
The issue is the [B*T, V] logits tensor. We can't avoid it with standard CE.
But we CAN use the 'chunked' approach: compute CE in vocab chunks.
This trades compute for memory, letting us use bigger batches.
"""
h = model(ids) # [B, T, d]
# Tied embeddings: lm_head weight = tok_emb weight = [V, d]
weight = model.tok_emb.weight # [V, d]
# Reshape for CE: h_flat [B*T, d], targets [B*T]
h_flat = h[:, :-1].reshape(-1, cfg.d_model)
targets = ids[:, 1:].reshape(-1)
# F.cross_entropy with weight = compute logits on the fly
# Actually, F.cross_entropy(input, target) where input is [N, C]
# still materializes [N, C] logits.
# The real fused version: use torch's internal fused kernel
logits = F.linear(h_flat, weight) # [B*T, V]
loss = F.cross_entropy(logits, targets)
return loss
def chunked_ce(model, ids, cfg, chunk_size=8192):
"""Chunked CE: compute loss in chunks to save memory.
Allows much larger batch sizes.
"""
h = model(ids) # [B, T, d]
weight = model.tok_emb.weight # [V, d]
h_flat = h[:, :-1].reshape(-1, cfg.d_model) # [B*T, d]
targets = ids[:, 1:].reshape(-1) # [B*T]
N = h_flat.shape[0]
total_loss = 0.0
for i in range(0, N, chunk_size):
chunk_h = h_flat[i:i+chunk_size]
chunk_t = targets[i:i+chunk_size]
chunk_logits = F.linear(chunk_h, weight)
total_loss = total_loss + F.cross_entropy(chunk_logits, chunk_t) * chunk_h.shape[0]
return total_loss / N
# --- Test configs ---
configs = [
# (batch, seq, ce_fn, label)
(16, 1024, standard_ce, "standard CE"),
(16, 1024, fused_ce, "fused CE (same speed)"),
(16, 1024, chunked_ce, "chunked CE (8k)"),
(32, 1024, standard_ce, "batch=32"),
(32, 1024, chunked_ce, "batch=32 chunked"),
(64, 1024, chunked_ce, "batch=64 chunked"),
(16, 2048, standard_ce, "seq=2048"),
(16, 2048, chunked_ce, "seq=2048 chunked"),
(32, 2048, chunked_ce, "batch=32 seq=2048 chunked"),
(8, 4096, chunked_ce, "seq=4096 chunked"),
(16, 4096, chunked_ce, "batch=16 seq=4096 chunked"),
]
print(f"{'config':>30} {'tokens':>8} {'step(ms)':>9} {'tok/s':>10} {'VRAM(GB)':>9}")
print("-" * 70)
for batch, seq, ce_fn, label in configs:
try:
ids = torch.randint(0, cfg.vocab_size, (batch, seq), device="cuda")
# Warmup
for _ in range(3):
opt.zero_grad(set_to_none=True)
loss = ce_fn(model, ids, cfg)
loss.backward()
opt.step()
torch.cuda.synchronize()
# Measure
torch.cuda.reset_peak_memory_stats()
times = []
for _ in range(5):
torch.cuda.synchronize()
t0 = time.perf_counter()
opt.zero_grad(set_to_none=True)
loss = ce_fn(model, ids, cfg)
loss.backward()
opt.step()
torch.cuda.synchronize()
times.append(time.perf_counter() - t0)
times.sort()
step_ms = times[2] * 1000
tokens = batch * seq
tok_s = tokens / (step_ms / 1000)
vram = torch.cuda.max_memory_allocated() / 1e9
print(f"{label:>30} {tokens:>8} {step_ms:>9.1f} {tok_s:>10,.0f} {vram:>9.2f}")
except RuntimeError as e:
if "out of memory" in str(e).lower():
print(f"{label:>30} {batch*seq:>8} {'OOM':>9}")
torch.cuda.empty_cache()
else:
raise
|