Spaces:
Sleeping
Sleeping
| """Generic fine-tuning loop: AdamW with split encoder/head LRs, linear warmup, | |
| bf16/fp16 autocast on CUDA, grad accumulation, early stopping on a dev metric.""" | |
| from __future__ import annotations | |
| import json | |
| import logging | |
| import math | |
| import time | |
| from pathlib import Path | |
| from typing import Callable | |
| import torch | |
| from torch import nn | |
| from torch.utils.data import DataLoader | |
| from transformers import get_linear_schedule_with_warmup | |
| from ..utils import autocast | |
| log = logging.getLogger(__name__) | |
| NO_DECAY = ("bias", "LayerNorm.weight", "layernorm.weight", "layer_norm.weight", "norm.weight") | |
| def param_groups(model: nn.Module, lr: float, head_lr: float, weight_decay: float) -> list[dict]: | |
| groups = {(e, d): [] for e in (True, False) for d in (True, False)} | |
| for name, p in model.named_parameters(): | |
| if not p.requires_grad: | |
| continue | |
| is_encoder = name.startswith("encoder.") | |
| decay = not any(k in name for k in NO_DECAY) | |
| groups[(is_encoder, decay)].append(p) | |
| out = [] | |
| for (is_encoder, decay), params in groups.items(): | |
| if params: | |
| out.append({"params": params, "lr": lr if is_encoder else head_lr, | |
| "weight_decay": weight_decay if decay else 0.0}) | |
| return out | |
| def _nonfinite_report(model: nn.Module) -> str: | |
| bad = [n for n, p in model.named_parameters() if not torch.isfinite(p).all()] | |
| return f"{len(bad)} weight tensors non-finite, e.g. {bad[:3]}" if bad else "all weights finite" | |
| def fit(model: nn.Module, dataset, collate: Callable, hp, device: torch.device, | |
| amp_dtype: torch.dtype | None, evaluate: Callable[[], dict], metric: str, | |
| save: Callable[[dict], None], log_path: Path, seed: int, log_every: int = 50) -> dict: | |
| """Train; after each epoch call `evaluate()`, keep the best by `metric` via `save()`.""" | |
| generator = torch.Generator().manual_seed(seed) | |
| loader = DataLoader(dataset, batch_size=hp.batch_size, shuffle=True, collate_fn=collate, | |
| generator=generator, num_workers=0) # 0 workers: safest on Windows | |
| optimizer = torch.optim.AdamW(param_groups(model, hp.lr, hp.head_lr, hp.weight_decay)) | |
| steps_per_epoch = math.ceil(len(loader) / hp.grad_accum) | |
| total = steps_per_epoch * hp.epochs | |
| scheduler = get_linear_schedule_with_warmup(optimizer, int(hp.warmup_ratio * total), total) | |
| scaler = torch.amp.GradScaler("cuda", enabled=amp_dtype == torch.float16) | |
| best, best_metrics, bad_epochs, skipped = -math.inf, {}, 0, 0 | |
| log_path.parent.mkdir(parents=True, exist_ok=True) | |
| with open(log_path, "a", encoding="utf-8", newline="\n") as logf: | |
| for epoch in range(1, hp.epochs + 1): | |
| model.train() | |
| t0, running, n = time.time(), 0.0, 0 | |
| optimizer.zero_grad(set_to_none=True) | |
| for step, batch in enumerate(loader, 1): | |
| batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()} | |
| with autocast(device, amp_dtype): | |
| loss = model(**batch)["loss"] | |
| if not torch.isfinite(loss): | |
| raise FloatingPointError( | |
| f"Non-finite loss at epoch {epoch} step {step} " | |
| f"({_nonfinite_report(model)}). If weights are fine, the forward pass " | |
| f"overflowed: try --set runtime.precision=fp32.") | |
| scaler.scale(loss / hp.grad_accum).backward() | |
| running += loss.item() | |
| n += 1 | |
| if step % hp.grad_accum == 0 or step == len(loader): | |
| scaler.unscale_(optimizer) | |
| norm = nn.utils.clip_grad_norm_(model.parameters(), hp.max_grad_norm) | |
| if not torch.isfinite(norm) and amp_dtype != torch.float16: | |
| # fp16's GradScaler skips these itself; for bf16/fp32 we must, or the | |
| # NaN gets written into the weights and surfaces a step later. | |
| skipped += 1 | |
| bad = [name for name, p in model.named_parameters() | |
| if p.grad is not None and not torch.isfinite(p.grad).all()] | |
| log.warning("non-finite gradients at epoch %d step %d, update skipped " | |
| "(%d params, e.g. %s)", epoch, step, len(bad), bad[:3]) | |
| optimizer.zero_grad(set_to_none=True) | |
| if skipped > hp.max_skipped_steps: | |
| raise FloatingPointError( | |
| f"{skipped} updates skipped for non-finite gradients; first bad " | |
| f"params: {bad[:5]}. Try --set runtime.precision=fp32.") | |
| continue | |
| scaler.step(optimizer) | |
| scaler.update() | |
| scheduler.step() | |
| optimizer.zero_grad(set_to_none=True) | |
| if step % log_every == 0: | |
| log.info("epoch %d step %d/%d loss %.4f", epoch, step, len(loader), running / n) | |
| metrics = evaluate() | |
| metrics.update(epoch=epoch, train_loss=running / max(n, 1), epoch_seconds=time.time() - t0) | |
| logf.write(json.dumps(metrics) + "\n") | |
| logf.flush() | |
| improved = metrics[metric] > best | |
| log.info("epoch %d %s=%.4f%s (%.0fs)", epoch, metric, metrics[metric], | |
| " *best*" if improved else "", metrics["epoch_seconds"]) | |
| if improved: | |
| best, best_metrics, bad_epochs = metrics[metric], dict(metrics), 0 | |
| save(best_metrics) | |
| else: | |
| bad_epochs += 1 | |
| if bad_epochs >= hp.patience: | |
| log.info("early stop: no %s improvement for %d epochs", metric, hp.patience) | |
| break | |
| return best_metrics |