| 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) |
