{"record":{"id":"c0bdfc5f608f533b","repo":"xai-org/x-algorithm","slug":"non-finite-loss-grad-at-step-step-loss-lv-gra","errorCode":null,"errorMessage":"non-finite loss/grad at step {step}: loss={lv} grad_norm={gn} — stopping for investigation (never zero-and-continue)","messagePattern":"non-finite loss/grad at step (.+?): loss=(.+?) grad_norm=(.+?) — stopping for investigation \\(never zero-and-continue\\)","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"critical","filePath":"bdsm/training/train_head.py","lineNumber":251,"sourceCode":"            order = rng.permutation(n)\n            cursor = 0\n        idx = jnp.asarray(order[cursor : cursor + args.batch_size])\n        cursor += args.batch_size\n\n        params, opt_state, loss_val, grad_norm, stats = train_step(\n            params, opt_state, x_all[idx], y_all[idx], mask_all[idx]\n        )\n\n        if step == 0 and args.focal_gamma > 0:\n            assert \"focal_mean_weight\" in stats, (\n                \"focal loss not wired: focal_mean_weight missing from stats\"\n            )\n            log.info(f\"step 0: focal_mean_weight={float(stats['focal_mean_weight']):.4f}\")\n\n        if step % 500 == 0:\n            lv, gn = float(loss_val), float(grad_norm)\n            if not (np.isfinite(lv) and np.isfinite(gn)):\n                raise RuntimeError(\n                    f\"non-finite loss/grad at step {step}: loss={lv} grad_norm={gn} — \"\n                    \"stopping for investigation (never zero-and-continue)\"\n                )\n            log.info(\n                f\"[{step:6d}/{args.steps}] loss={lv:.4f} grad={gn:.3f} \"\n                f\"lr={float(schedule(step)):.2e}\"\n            )\n\n        if step > 0 and step % args.eval_every == 0:\n            params_np = {k: np.asarray(v) for k, v in params.items()}\n            m = evaluate(params_np, holdout)\n            log.info(\n                f\"eval@{step}: \"\n                + json.dumps(\n                    {k: (round(v, 4) if isinstance(v, float) else v) for k, v in m.items()}\n                )\n            )\n            save_checkpoint(run_dir, step, params_np, run_config)","sourceCodeStart":233,"sourceCodeEnd":269,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/bdsm/training/train_head.py#L233-L269","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Inspect the step's logged loss/grad to see which is non-finite and since when","Lower the learning rate or add LR warmup; if fp16, switch to fp32 or add GradScaler","Audit the cached parquet for NaN/Inf cls rows and re-run the nonfinite-row drop in load_cached","Check focal-loss hyperparameters and class-weight counts for division by zero"],"exampleFix":"# before\noptimizer = torch.optim.SGD(head.parameters(), lr=1e0)\n# after\noptimizer = torch.optim.SGD(head.parameters(), lr=1e-3)\nscheduler = warmup_cosine(optimizer, warmup=500, total=args.steps)","handlingStrategy":"try-catch","validationCode":"assert all(np.isfinite(x).all() for x in batch.values()), 'non-finite inputs'","typeGuard":null,"tryCatchPattern":"try:\n    train(args)\nexcept RuntimeError as e:\n    if 'non-finite loss/grad' in str(e):\n        dump_batch_for_postmortem(); halve_lr_and_restart()\n    else:\n        raise","preventionTips":["Add LR warmup and gradient clipping from the start","Sanitize NaN/Inf rows in load_cached (the codebase already drops them)","Run short exploratory jobs with finite checks before long training runs"],"tags":["python","training","nan-loss","numerical-stability"],"backgroundTag":"nan-loss-during-training","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}