xai-org/x-algorithm · critical · RuntimeError

non-finite loss/grad at step {step}: loss={lv} grad_norm={gn

Error message

non-finite loss/grad at step {step}: loss={lv} grad_norm={gn} — stopping for investigation (never zero-and-continue)

What it means

Every 500 training steps, train_head checks that loss and gradient norm are finite. NaN/Inf in either halts training with RuntimeError by design: non-finite optimization is treated as a bug to investigate, never something to zero out and continue past.

Source

Thrown at bdsm/training/train_head.py:251

            order = rng.permutation(n)
            cursor = 0
        idx = jnp.asarray(order[cursor : cursor + args.batch_size])
        cursor += args.batch_size

        params, opt_state, loss_val, grad_norm, stats = train_step(
            params, opt_state, x_all[idx], y_all[idx], mask_all[idx]
        )

        if step == 0 and args.focal_gamma > 0:
            assert "focal_mean_weight" in stats, (
                "focal loss not wired: focal_mean_weight missing from stats"
            )
            log.info(f"step 0: focal_mean_weight={float(stats['focal_mean_weight']):.4f}")

        if step % 500 == 0:
            lv, gn = float(loss_val), float(grad_norm)
            if not (np.isfinite(lv) and np.isfinite(gn)):
                raise RuntimeError(
                    f"non-finite loss/grad at step {step}: loss={lv} grad_norm={gn} — "
                    "stopping for investigation (never zero-and-continue)"
                )
            log.info(
                f"[{step:6d}/{args.steps}] loss={lv:.4f} grad={gn:.3f} "
                f"lr={float(schedule(step)):.2e}"
            )

        if step > 0 and step % args.eval_every == 0:
            params_np = {k: np.asarray(v) for k, v in params.items()}
            m = evaluate(params_np, holdout)
            log.info(
                f"eval@{step}: "
                + json.dumps(
                    {k: (round(v, 4) if isinstance(v, float) else v) for k, v in m.items()}
                )
            )
            save_checkpoint(run_dir, step, params_np, run_config)

View on GitHub (pinned to 24c60942c5)

Solutions

  1. Inspect the step's logged loss/grad to see which is non-finite and since when
  2. Lower the learning rate or add LR warmup; if fp16, switch to fp32 or add GradScaler
  3. Audit the cached parquet for NaN/Inf cls rows and re-run the nonfinite-row drop in load_cached
  4. Check focal-loss hyperparameters and class-weight counts for division by zero

Example fix

# before
optimizer = torch.optim.SGD(head.parameters(), lr=1e0)
# after
optimizer = torch.optim.SGD(head.parameters(), lr=1e-3)
scheduler = warmup_cosine(optimizer, warmup=500, total=args.steps)
Defensive patterns

Strategy: try-catch

Validate before calling

assert all(np.isfinite(x).all() for x in batch.values()), 'non-finite inputs'

Try / catch

try:
    train(args)
except RuntimeError as e:
    if 'non-finite loss/grad' in str(e):
        dump_batch_for_postmortem(); halve_lr_and_restart()
    else:
        raise

Prevention

When it happens

Trigger: Loss overflow with too-high learning rate or fp16 instability; bad data (NaN embeddings/labels from the cached parquet); focal-loss gamma/alpha misconfiguration; division by zero in class weighting producing Inf grads.

Common situations: Learning-rate warmup missing after a schedule change; a corrupted parquet shard with NaN cls rows; class-imbalance weighting hitting zero-count classes; mixed-precision underflow.

Related errors


AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/c0bdfc5f608f533b. Report an issue: GitHub.