Fix FP16 gradient scaler crash: ensure FP32 master weights

from_pretrained() may load weights in FP16 (safetensors dtype), but
FP16 mixed precision requires FP32 parameters for gradient unscaling.
Added model.float() in create_optimizer() before constructing param groups.

Also: force pip reinstall in setup cell, verify eps=1.0 and float() fixes
are present via source inspection, add optimizer logging.

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

Author: Sid Mohan

pii-ner-v1/notebooks/full_training.ipynb

@@ -31,7 +31,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "import os, sys\n\n# Clone or force-update to latest\nif not os.path.exists(\"/content/datafog-labs\"):\n    !git clone https://github.com/DataFog/datafog-labs.git /content/datafog-labs\nelse:\n    !cd /content/datafog-labs && git fetch origin && git reset --hard origin/main\n\n!pip install -e \"/content/datafog-labs/pii-ner-v1[dev]\" -q\n\nsys.path.insert(0, \"/content/datafog-labs/pii-ner-v1/src\")\n\nimport datafog_pii_ner\nprint(f\"datafog_pii_ner loaded from: {datafog_pii_ner.__file__}\")\n!cd /content/datafog-labs && git log -1 --oneline"
+   "source": "import os, sys\n\n# Clone or force-update to latest\nif not os.path.exists(\"/content/datafog-labs\"):\n    !git clone https://github.com/DataFog/datafog-labs.git /content/datafog-labs\nelse:\n    !cd /content/datafog-labs && git fetch origin && git reset --hard origin/main\n\n# Force reinstall to ensure latest code (pip may cache stale editable installs)\n!pip install -e \"/content/datafog-labs/pii-ner-v1[dev]\" --force-reinstall --no-deps -q\n!pip install -e \"/content/datafog-labs/pii-ner-v1[dev]\" -q  # install deps if missing\n\nsys.path.insert(0, \"/content/datafog-labs/pii-ner-v1/src\")\n\n# Reload to pick up fresh install\nimport importlib\nimport datafog_pii_ner\nimportlib.reload(datafog_pii_ner)\nprint(f\"datafog_pii_ner loaded from: {datafog_pii_ner.__file__}\")\n!cd /content/datafog-labs && git log -1 --oneline\n\n# Verify critical fixes are present\nimport inspect\nfrom datafog_pii_ner.training.train import PiiTrainer\nsource = inspect.getsource(PiiTrainer.create_optimizer)\nassert \"eps\" in source, \"MISSING: eps=1.0 fix for AdamW NaN\"\nassert \"model.float()\" in source or \"self.model.float()\" in source, \"MISSING: FP32 master weights fix\"\nprint(\"Verified: eps=1.0 and FP32 master weights fixes are present\")"
   },
   {
    "cell_type": "code",

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

@@ -35,18 +35,35 @@ def __init__(self, *args, lr_backbone: float | None = None, lr_head: float | Non
     def create_optimizer(self):
         """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.
+        Key details:
+        - model.float() ensures FP32 master weights. from_pretrained() may load
+          in the saved dtype (FP16 safetensors), but FP16 mixed precision requires
+          FP32 parameters so the gradient scaler can unscale FP32 gradients.
+        - Backbone eps=1.0 dampens AdamW's adaptive per-parameter scaling.
+          Without this, bias-corrected updates at step 1 apply ~lr per weight
+          regardless of gradient magnitude, producing NaN in DeBERTa on PyTorch 2.9+.
         """
         if self.lr_backbone is not None and self.lr_head is not None:
+            import torch
             from torch.optim import AdamW
 
+            # Ensure FP32 master weights — required for FP16 gradient scaler.
+            # from_pretrained() may load in FP16 if safetensors are stored that way.
+            self.model.float()
+
             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]
+
+            logger.info(
+                f"Optimizer: backbone={len(backbone_params)} tensors (lr={self.lr_backbone}, eps=1.0), "
+                f"head={len(head_params)} tensors (lr={self.lr_head})"
+            )
+
+            # Verify all params are FP32
+            non_fp32 = [(n, p.dtype) for n, p in self.model.named_parameters() if p.dtype != torch.float32]
+            if non_fp32:
+                logger.warning(f"Non-FP32 params after .float(): {non_fp32[:5]}")
+
             self.optimizer = AdamW(
                 [
                     {"params": backbone_params, "lr": self.lr_backbone, "eps": 1.0},