File size: 4,590 Bytes
c21e887 e520ea7 c21e887 e520ea7 e378e01 e520ea7 c21e887 e378e01 c21e887 e520ea7 c21e887 e520ea7 c21e887 e520ea7 c21e887 e520ea7 c21e887 e520ea7 c21e887 e378e01 | 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 | import torch
import torch.nn.functional as F
import logging
import time
from dataclasses import asdict
import numpy as np
import wandb
from torch.cuda.amp import autocast, GradScaler
from tqdm import tqdm
from src.engine.generate import generate_greedy
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
def calc_cross_entropy_loss(logits, targets):
B, T, V = logits.shape
return F.cross_entropy(logits.view(B * T, V), targets.view(B * T))
def calc_perplexity(loss: float) -> float:
return float(np.exp(loss))
def train(
model, optimizer, train_loader, val_loader, device, cfg, sample_prompt, tokenizer
):
run = wandb.init(
entity="pjawale-student",
project="gpt2-pytorch",
config=asdict(cfg),
)
scaler = GradScaler()
history = []
tokens_per_batch = cfg.batch_size * cfg.context_window_size
global_step = 0
try:
for epoch in range(1, cfg.num_epochs + 1):
model.train()
train_loss = 0.0
epoch_start = time.perf_counter()
window_start = epoch_start
window_batches = 0
pbar = tqdm(
train_loader,
desc=f"Epoch {epoch}/{cfg.num_epochs}",
unit="batch",
leave=True,
)
for batch_idx, (inputs, targets) in enumerate(pbar, 1):
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
with autocast(dtype=torch.bfloat16):
loss = calc_cross_entropy_loss(model(inputs), targets)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
train_loss += loss.item()
window_batches += 1
global_step += 1
if batch_idx % 10 == 0 or batch_idx == len(train_loader):
now = time.perf_counter()
elapsed = now - window_start
tok_s = (window_batches * tokens_per_batch) / elapsed if elapsed > 0 else 0.0
window_start = now
window_batches = 0
pbar.set_postfix(
loss=f"{loss.item():.4f}",
tok_s=f"{tok_s:,.0f}",
)
wandb.log(
{
"train/batch_loss": loss.item(),
"train/batch_perplexity": calc_perplexity(loss.item()),
"train/tok_s": tok_s,
"epoch": epoch,
},
step=global_step,
)
epoch_elapsed = time.perf_counter() - epoch_start
epoch_tokens = len(train_loader) * tokens_per_batch
epoch_tok_s = epoch_tokens / epoch_elapsed if epoch_elapsed > 0 else 0.0
model.eval()
val_loss = 0.0
with torch.no_grad():
for inputs, targets in tqdm(val_loader, desc="Validating", leave=False):
inputs, targets = inputs.to(device), targets.to(device)
with autocast(dtype=torch.bfloat16):
val_loss += calc_cross_entropy_loss(model(inputs), targets).item()
avg_train = train_loss / len(train_loader)
avg_val = val_loss / len(val_loader)
avg_train_perplexity = calc_perplexity(avg_train)
avg_val_perplexity = calc_perplexity(avg_val)
sample = generate_greedy(model, tokenizer, device, sample_prompt)
logging.info(
f"Epoch {epoch:2d}/{cfg.num_epochs} | train={avg_train:.4f} | val={avg_val:.4f} | "
f"train_perplexity={avg_train_perplexity:.4f} | val_perplexity={avg_val_perplexity:.4f} | "
f"{epoch_tok_s:,.0f} tok/s (epoch avg)"
)
wandb.log(
{
"train/loss": avg_train,
"val/loss": avg_val,
"train/perplexity": avg_train_perplexity,
"val/perplexity": avg_val_perplexity,
"train/tok_s_epoch": epoch_tok_s,
"sample": sample,
"epoch": epoch,
},
step=global_step,
)
history.append((avg_train, avg_val))
logging.info(f" Sample: {sample}\n")
finally:
wandb.finish()
return np.array(history) |