Fix AdamW NaN: use eps=1.0 for backbone to dampen adaptive scaling

AdamW's bias-corrected update at step 1 applies ~lr per weight regardless
of gradient magnitude (update ≈ lr × sign(g)). On PyTorch 2.9+ with
DeBERTa v3, this produces NaN weights after a single optimizer step.

Setting eps=1.0 for the backbone param group makes updates scale with
gradient magnitude like SGD (update ≈ lr × g), which is numerically safe.
The head retains standard eps=1e-8 for adaptive scaling.

Diagnosis confirmed via isolation tests:
- Default AdamW (eps=1e-8): NaN at step 2
- Reference AdamW (fused=False): NaN at step 2
- Manual SGD: 3 steps OK
- AdamW with eps=1.0: 10 steps OK, loss 480→149

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: Sid Mohan

pii-ner-v1/docs/smoke_test_walkthrough.md

@@ -195,3 +195,65 @@ Updated thresholds:
 | Verdict | Complete failure | PASS |
 
 **The lesson:** When combining pretrained and randomly initialized components, differential learning rates aren't optional — they're essential. The pretrained backbone needs preservation; the new head needs speed.
+
+---
+
+## AdamW NaN on Full Training (PyTorch 2.9+)
+
+When scaling from 100 smoke-test examples to 135K+ real examples, training immediately produced NaN loss. Every metric was zero from epoch 1. This section documents the diagnosis and fix.
+
+### Symptoms
+
+- Training loss displayed as `0.000000` (HF Trainer v5.0 renders NaN as 0)
+- Validation loss: `nan`
+- All F1 scores: `0.000`
+- Loss was valid at step 1 (~490–920), then NaN from step 2 onward
+
+### Root Cause: AdamW Bias-Correction Amplification
+
+AdamW's adaptive per-parameter scaling makes the effective update magnitude **independent of gradient magnitude**:
+
+```
+At step 1 with bias correction:
+  m_hat = g                    (first moment, bias-corrected)
+  v_hat = g²                   (second moment, bias-corrected)
+  update = lr * g / (|g| + ε)  ≈ lr × sign(g)
+```
+
+With standard `ε=1e-8`, every backbone weight changes by ≈ ±lr (≈ ±2e-5) regardless of gradient clipping. While ±2e-5 seems small, it was sufficient to push DeBERTa v3's internal computations into NaN territory on PyTorch 2.9+ with CUDA.
+
+Manual SGD (`p -= lr * clipped_grad`) produces updates ~8000× smaller because it preserves the clipped gradient magnitude: ≈ lr × |g_clipped| ≈ 2e-5 × 1.2e-4 ≈ 2.4e-9 per weight.
+
+### Diagnostic Evidence
+
+| Test | Result | Conclusion |
+|------|--------|------------|
+| Default AdamW (eps=1e-8) | NaN at step 2 | Baseline fails |
+| Reference AdamW (fused=False, foreach=False) | NaN at step 2 | Not a CUDA kernel bug |
+| Manual SGD (p -= lr * grad) | 3 steps OK | Raw weight updates are safe |
+| Frozen DeBERTa (head only) | 5 steps OK, loss 496→291 | Head learns without backbone |
+| **AdamW with eps=1.0 for backbone** | **10 steps OK, loss 480→149** | **Fix confirmed** |
+
+### The Fix
+
+Set `eps=1.0` for the backbone parameter group in AdamW. This dampens the adaptive scaling so updates scale with gradient magnitude (like SGD) rather than being a fixed ±lr:
+
+```python
+# In PiiTrainer.create_optimizer():
+self.optimizer = AdamW([
+    {"params": backbone_params, "lr": 2e-5, "eps": 1.0},   # SGD-like
+    {"params": head_params,     "lr": 1e-3},                # standard AdamW
+], weight_decay=0.01)
+```
+
+With `eps=1.0` and clipped gradients (`|g| < 1.0`):
+```
+denom = |g| + 1.0 ≈ 1.0
+update ≈ lr × g       (same as SGD — scales with gradient magnitude)
+```
+
+The head retains standard AdamW (`eps=1e-8`) because it's randomly initialized and benefits from adaptive scaling. The backbone gets SGD-like updates because pretrained weights need minimal, gradient-proportional adjustments.
+
+### Why the Smoke Test Passed
+
+The smoke test used 100 examples with a different Colab environment (older PyTorch/Transformers). The NaN manifests specifically on real data at scale with PyTorch 2.9+ and Transformers 5.0. The CRF loss on real data (~490) is comparable to the smoke test's initial loss (~362), so the scale of loss is not the differentiator — the PyTorch version and its optimizer internals are.

pii-ner-v1/notebooks/full_training.ipynb

@@ -168,7 +168,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "from transformers import TrainingArguments\nfrom datafog_pii_ner.data.collator import PiiDataCollator\nfrom datafog_pii_ner.training.metrics import compute_metrics\nfrom datafog_pii_ner.training.train import PiiTrainer\n\ncollator = PiiDataCollator(tokenizer=tokenizer, max_char_len=CONFIG[\"max_char_len\"])\n\ntraining_args = TrainingArguments(\n    output_dir=CONFIG[\"output_dir\"],\n    num_train_epochs=CONFIG[\"epochs\"],\n    per_device_train_batch_size=CONFIG[\"batch_size\"],\n    per_device_eval_batch_size=CONFIG[\"batch_size\"],\n    gradient_accumulation_steps=CONFIG[\"gradient_accumulation_steps\"],\n    learning_rate=CONFIG[\"lr_backbone\"],\n    warmup_ratio=CONFIG[\"warmup_ratio\"],\n    weight_decay=CONFIG[\"weight_decay\"],\n    fp16=CONFIG.get(\"fp16\", False),\n    bf16=CONFIG.get(\"bf16\", False),\n    eval_strategy=\"epoch\",\n    save_strategy=\"epoch\",\n    metric_for_best_model=\"overall_f1\",\n    load_best_model_at_end=True,\n    report_to=\"wandb\",\n    run_name=CONFIG[\"run_name\"],\n    logging_steps=50,\n    remove_unused_columns=False,\n    dataloader_num_workers=2,\n    save_total_limit=3,\n)\n\nprint(f\"Backbone LR: {CONFIG['lr_backbone']}, Head LR: {CONFIG['lr_head']}\")\nprint(f\"Mixed precision: {'bf16' if CONFIG.get('bf16') else 'fp16' if CONFIG.get('fp16') else 'none'}\")\n\n# PiiTrainer handles differential learning rates internally via create_optimizer().\n# This avoids passing a custom optimizer through optimizers=(), which bypasses\n# Accelerate's optimizer wrapping and causes FP16/BF16 gradient scaler issues.\ntrainer = PiiTrainer(\n    model=model,\n    args=training_args,\n    train_dataset=datasets[\"train\"],\n    eval_dataset=datasets[\"validation\"],\n    data_collator=collator,\n    compute_metrics=compute_metrics,\n    lr_backbone=CONFIG[\"lr_backbone\"],\n    lr_head=CONFIG[\"lr_head\"],\n)\n\nprint(f\"\\nStarting training for {CONFIG['epochs']} epochs...\")\ntrain_result = trainer.train()\nprint(f\"\\nTraining complete. Final loss: {train_result.training_loss:.4f}\")"
+   "source": "from transformers import TrainingArguments\nfrom datafog_pii_ner.data.collator import PiiDataCollator\nfrom datafog_pii_ner.training.metrics import compute_metrics\nfrom datafog_pii_ner.training.train import PiiTrainer\n\ncollator = PiiDataCollator(tokenizer=tokenizer, max_char_len=CONFIG[\"max_char_len\"])\n\ntraining_args = TrainingArguments(\n    output_dir=CONFIG[\"output_dir\"],\n    num_train_epochs=CONFIG[\"epochs\"],\n    per_device_train_batch_size=CONFIG[\"batch_size\"],\n    per_device_eval_batch_size=CONFIG[\"batch_size\"],\n    gradient_accumulation_steps=CONFIG[\"gradient_accumulation_steps\"],\n    learning_rate=CONFIG[\"lr_backbone\"],\n    warmup_ratio=CONFIG[\"warmup_ratio\"],\n    weight_decay=CONFIG[\"weight_decay\"],\n    fp16=CONFIG.get(\"fp16\", False),\n    bf16=CONFIG.get(\"bf16\", False),\n    eval_strategy=\"epoch\",\n    save_strategy=\"epoch\",\n    metric_for_best_model=\"overall_f1\",\n    load_best_model_at_end=True,\n    report_to=\"wandb\",\n    run_name=CONFIG[\"run_name\"],\n    logging_steps=50,\n    remove_unused_columns=False,\n    dataloader_num_workers=2,\n    save_total_limit=3,\n)\n\nprint(f\"Backbone LR: {CONFIG['lr_backbone']}, Head LR: {CONFIG['lr_head']}\")\nprint(f\"Mixed precision: {'bf16' if CONFIG.get('bf16') else 'fp16' if CONFIG.get('fp16') else 'none'}\")\n\n# PiiTrainer handles differential learning rates internally via create_optimizer().\n# The backbone param group uses eps=1.0 to dampen AdamW's adaptive scaling,\n# preventing NaN weight updates caused by bias-correction amplification at step 1\n# (see smoke_test_walkthrough.md \"AdamW NaN on Full Training\" section).\ntrainer = PiiTrainer(\n    model=model,\n    args=training_args,\n    train_dataset=datasets[\"train\"],\n    eval_dataset=datasets[\"validation\"],\n    data_collator=collator,\n    compute_metrics=compute_metrics,\n    lr_backbone=CONFIG[\"lr_backbone\"],\n    lr_head=CONFIG[\"lr_head\"],\n)\n\nprint(f\"\\nStarting training for {CONFIG['epochs']} epochs...\")\ntrain_result = trainer.train()\nprint(f\"\\nTraining complete. Final loss: {train_result.training_loss:.4f}\")"
   },
   {
    "cell_type": "markdown",

pii-ner-v1/src/datafog_pii_ner/training/train.py

@@ -33,15 +33,23 @@ def __init__(self, *args, lr_backbone: float | None = None, lr_head: float | Non
         super().__init__(*args, **kwargs)
 
     def create_optimizer(self):
-        """Create AdamW with differential learning rates for backbone vs head."""
+        """Create AdamW with differential learning rates for backbone vs head.
+
+        The backbone group uses eps=1.0 to dampen AdamW's adaptive per-parameter
+        scaling. Without this, AdamW's bias-corrected update at step 1 applies
+        ~lr per weight regardless of gradient magnitude, which produces NaN in
+        DeBERTa on PyTorch 2.9+. Setting eps=1.0 makes backbone updates behave
+        like SGD (update ≈ lr * grad), which is numerically safe and standard
+        for fine-tuning pretrained transformers.
+        """
         if self.lr_backbone is not None and self.lr_head is not None:
             from torch.optim import AdamW
 
             backbone_params = [p for n, p in self.model.named_parameters() if "deberta" in n and p.requires_grad]
             head_params = [p for n, p in self.model.named_parameters() if "deberta" not in n and p.requires_grad]
             self.optimizer = AdamW(
                 [
-                    {"params": backbone_params, "lr": self.lr_backbone},
+                    {"params": backbone_params, "lr": self.lr_backbone, "eps": 1.0},
                     {"params": head_params, "lr": self.lr_head},
                 ],
                 weight_decay=self.args.weight_decay,