Training checkpoint / run progress

The difference from “fine-tuning completed”: there is one final push, here they are intermediate. Training your own model runs for hours, and there are two moments when you want to know now rather than at the end:

• milestone — an epoch passed, a checkpoint was saved, val-loss improved;
• failure — NaN/Inf in the loss, early-stopping, the loss started climbing. Speed matters here: the sooner you stop a stuck run, the fewer GPU-hours you burn.

Hook this straight into the training loop via a callback.

PyTorch: a callback on epoch-end and on anomaly

import os, math, requests

NOTIFLY_URL   = os.environ["NOTIFLY_URL"]
NOTIFLY_TOKEN = os.environ["NOTIFLY_TOKEN"]

def notify(title, msg, prio):
    requests.post(f"{NOTIFLY_URL}/message",
                  params={"token": NOTIFLY_TOKEN},
                  json={"title": title, "message": msg, "priority": prio},
                  timeout=5)

class NotifyCallback:
    def __init__(self, run_name, milestone_every=1):
        self.run = run_name
        self.every = milestone_every          # send a milestone every N epochs
        self.best_val = math.inf

    def on_epoch_end(self, epoch, train_loss, val_loss):
        # 1) failures — loud and immediate
        if not math.isfinite(train_loss) or not math.isfinite(val_loss):
            notify("🔥 NaN/Inf in the loss",
                   f"{self.run}: epoch {epoch}, train={train_loss}, val={val_loss}. "
                   f"The run is most likely ruined — stop it and lower the LR.",
                   prio=10)
            return

        # 2) milestone — quiet-normal
        if epoch % self.every == 0:
            tag = "🟢 new best" if val_loss < self.best_val else ""
            notify(f"📊 Epoch {epoch} {tag}".strip(),
                   f"{self.run}\ntrain_loss={train_loss:.4f}\n"
                   f"val_loss={val_loss:.4f} (best={min(val_loss, self.best_val):.4f})",
                   prio=4)

        self.best_val = min(self.best_val, val_loss)

    def on_early_stop(self, epoch, patience):
        notify("🛑 Early stopping triggered",
               f"{self.run}: stopped at epoch {epoch}, "
               f"val loss hasn't improved for {patience} epochs. Grab the best checkpoint.",
               prio=6)

The priorities are chosen by importance: a milestone is 4 (a normal push, you can glance at it over coffee), early-stopping is 6 (a routine finish but it needs action), NaN is 10 (the run is on fire, put it out).

Wiring it into a plain loop:

cb = NotifyCallback("llama-lora-run7", milestone_every=1)
for epoch in range(EPOCHS):
    tr = train_one_epoch(model, train_loader)
    vl = validate(model, val_loader)
    torch.save(model.state_dict(), f"ckpt/epoch{epoch}.pt")
    cb.on_epoch_end(epoch, tr, vl)

HuggingFace Trainer: the same idea via TrainerCallback

If you train through transformers.Trainer, you don’t need to touch the loop — there is a ready extension point:

from transformers import TrainerCallback

class NotiflyTrainerCallback(TrainerCallback):
    def on_evaluate(self, args, state, control, metrics=None, **kw):
        loss = (metrics or {}).get("eval_loss")
        if loss is not None and not math.isfinite(loss):
            notify("🔥 NaN eval_loss", f"step {state.global_step}", 10)
        elif loss is not None:
            notify("📊 Eval checkpoint",
                   f"step {state.global_step}, eval_loss={loss:.4f}", 4)

Heartbeat: “is the training even alive”

A separate signal from progress — that the process is not hung. Let the loop ping a heartbeat after each step; if the pings stop (the GPU host froze, an OOM, someone pulled the power), Notifly itself sends an alert on timeout, even when the training code can no longer send anything.

requests.get(f"{NOTIFLY_URL}/heartbeat/ping/H<token>", timeout=3)  # once per step/minute

What to put in the alert text

• the run name and the epoch/step number (for deduplication);
• train/val loss and the current best;
• the learning rate and batch size — to reason about the cause of an anomaly;
• the path to the latest checkpoint in S3, to fetch/roll back.