Auto-select precision: try BF16, fallback to FP32, skip FP16

FP16 is incompatible with custom optimizer param groups on Accelerate 1.x+
(gradient scaler crashes on unscale). BF16 has no scaler and is preferred
on A100/H100. FP32 is the guaranteed fallback.

Preflight check now auto-selects: BF16 first → FP32 if NaN detected.
Catches crashes via try/except so fallback runs automatically.

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

Author: Sid Mohan

pii-ner-v1/notebooks/full_training.ipynb

@@ -77,16 +77,16 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "# Training configuration\n# Adjust batch sizes based on your GPU memory:\n#   A100 (40GB): batch_size=32, grad_accum=1 -> effective=32\n#   T4 (16GB):   batch_size=8,  grad_accum=4 -> effective=32\n\nCONFIG = {\n    # Model\n    \"backbone\": \"microsoft/deberta-v3-xsmall\",\n    \"max_seq_len\": 256,\n    \"max_char_len\": 20,\n    \"dropout\": 0.1,\n    \n    # Training\n    \"epochs\": 10,\n    \"batch_size\": 32,\n    \"gradient_accumulation_steps\": 1,\n    \"lr_backbone\": 2e-5,\n    \"lr_head\": 1e-3,\n    \"warmup_ratio\": 0.1,\n    \"weight_decay\": 0.01,\n    \n    # Mixed precision: always FP16 (never BF16).\n    # DeBERTa v3's disentangled attention produces NaN under BF16 (7-bit mantissa)\n    # at scale with real data, even though synthetic validation passes.\n    # FP16 (10-bit mantissa) is safe when combined with:\n    #   - eps=1.0 for backbone AdamW (dampens adaptive scaling)\n    #   - CRF head forced to FP32 via fused.float() + autocast(enabled=False)\n    \"fp16\": True,\n    \"bf16\": False,\n    \n    # Data\n    \"val_ratio\": 0.1,\n    \"test_ratio\": 0.1,\n    \"seed\": 42,\n    \n    # Output\n    \"output_dir\": \"/content/pii_ner_v1_output\",\n    \"run_name\": \"pii-ner-v1-full\",\n}\n\n# Auto-adjust batch size for GPU memory\nif torch.cuda.is_available():\n    gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9\n    gpu_name = torch.cuda.get_device_name(0)\n    if gpu_mem < 20:\n        print(f\"GPU: {gpu_name} ({gpu_mem:.0f}GB) — adjusting batch size for T4\")\n        CONFIG[\"batch_size\"] = 8\n        CONFIG[\"gradient_accumulation_steps\"] = 4\n    else:\n        print(f\"GPU: {gpu_name} ({gpu_mem:.0f}GB) — using batch_size=32\")\n\neffective_batch = CONFIG[\"batch_size\"] * CONFIG[\"gradient_accumulation_steps\"]\nprint(f\"Effective batch size: {effective_batch}\")\nprint(f\"Mixed precision: FP16 (CRF head runs in FP32)\")"
+   "source": "# Training configuration\n# Adjust batch sizes based on your GPU memory:\n#   A100 (40GB): batch_size=32, grad_accum=1 -> effective=32\n#   T4 (16GB):   batch_size=8,  grad_accum=4 -> effective=32\n\nCONFIG = {\n    # Model\n    \"backbone\": \"microsoft/deberta-v3-xsmall\",\n    \"max_seq_len\": 256,\n    \"max_char_len\": 20,\n    \"dropout\": 0.1,\n    \n    # Training\n    \"epochs\": 10,\n    \"batch_size\": 32,\n    \"gradient_accumulation_steps\": 1,\n    \"lr_backbone\": 2e-5,\n    \"lr_head\": 1e-3,\n    \"warmup_ratio\": 0.1,\n    \"weight_decay\": 0.01,\n    \n    # Mixed precision strategy (set by preflight check):\n    # - FP16 is NOT usable: Accelerate's gradient scaler crashes with custom\n    #   optimizer param groups (ValueError: Attempting to unscale FP16 gradients)\n    # - BF16 is preferred on A100/H100: no gradient scaler, same exponent range as FP32\n    # - FP32 is the guaranteed fallback: slower but always correct\n    # The preflight check will try BF16 first, fall back to FP32 if NaN detected.\n    \"fp16\": False,\n    \"bf16\": False,  # Set by preflight check below\n    \n    # Data\n    \"val_ratio\": 0.1,\n    \"test_ratio\": 0.1,\n    \"seed\": 42,\n    \n    # Output\n    \"output_dir\": \"/content/pii_ner_v1_output\",\n    \"run_name\": \"pii-ner-v1-full\",\n}\n\n# Auto-adjust batch size for GPU memory\nif torch.cuda.is_available():\n    gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9\n    gpu_name = torch.cuda.get_device_name(0)\n    cc = torch.cuda.get_device_properties(0).major\n    if gpu_mem < 20:\n        print(f\"GPU: {gpu_name} ({gpu_mem:.0f}GB) — adjusting batch size\")\n        CONFIG[\"batch_size\"] = 8\n        CONFIG[\"gradient_accumulation_steps\"] = 4\n    else:\n        print(f\"GPU: {gpu_name} ({gpu_mem:.0f}GB) — using batch_size=32\")\n    \n    # BF16 requires compute capability >= 8.0 (A100, H100, etc.)\n    if cc >= 8:\n        print(f\"Compute capability {cc}.x — BF16 available (preflight will verify)\")\n        CONFIG[\"_try_bf16\"] = True\n    else:\n        print(f\"Compute capability {cc}.x — BF16 not available, using FP32\")\n        CONFIG[\"_try_bf16\"] = False\n\neffective_batch = CONFIG[\"batch_size\"] * CONFIG[\"gradient_accumulation_steps\"]\nprint(f\"Effective batch size: {effective_batch}\")"
   },
   {
    "cell_type": "markdown",
-   "source": "## 2.5 Preflight Check\n\nRun training steps on synthetic data at **full batch size** to verify the optimizer + mixed precision pipeline is numerically stable. Explicitly checks for NaN in loss and model weights (HF Trainer v5.0 silently renders NaN as 0.000000).",
+   "source": "## 2.5 Preflight Check (auto-selects precision)\n\nTries BF16 first (faster, no gradient scaler). If NaN detected, falls back to FP32 (slower but guaranteed correct). FP16 is **not** attempted — incompatible with custom optimizer param groups on Accelerate 1.x+.",
    "metadata": {}
   },
   {
    "cell_type": "code",
-   "source": "import math\nimport numpy as np\nfrom datasets import Dataset\nfrom transformers import AutoTokenizer, TrainingArguments\nfrom datafog_pii_ner.data.collator import PiiDataCollator\nfrom datafog_pii_ner.data.label_schema import NUM_LABELS\nfrom datafog_pii_ner.model.pii_model import PiiNerConfig, PiiNerModel\nfrom datafog_pii_ner.training.train import PiiTrainer\n\nprint(\"=== PREFLIGHT CHECK ===\")\nprint(f\"Testing: batch_size={CONFIG['batch_size']}, fp16={CONFIG['fp16']}, bf16={CONFIG['bf16']}\")\n\n_tokenizer = AutoTokenizer.from_pretrained(CONFIG[\"backbone\"])\n_seq_len = CONFIG[\"max_seq_len\"]  # Use real sequence length, not short\n_n_samples = CONFIG[\"batch_size\"] * 2  # Enough for 2 full batches\n\n# Synthetic data at full sequence length and batch size\n_fake_data = {\n    \"input_ids\": np.random.randint(1, 1000, (_n_samples, _seq_len)).tolist(),\n    \"attention_mask\": np.ones((_n_samples, _seq_len), dtype=int).tolist(),\n    \"labels\": np.random.randint(0, NUM_LABELS, (_n_samples, _seq_len)).tolist(),\n    \"char_ids\": np.random.randint(0, 100, (_n_samples, _seq_len, CONFIG[\"max_char_len\"])).tolist(),\n}\n_ds = Dataset.from_dict(_fake_data)\n_collator = PiiDataCollator(tokenizer=_tokenizer, max_char_len=CONFIG[\"max_char_len\"])\n\n_config = PiiNerConfig(backbone=CONFIG[\"backbone\"], num_labels=NUM_LABELS)\n_model = PiiNerModel(_config)\n\n_args = TrainingArguments(\n    output_dir=\"/tmp/preflight_check\",\n    num_train_epochs=1,\n    per_device_train_batch_size=CONFIG[\"batch_size\"],\n    learning_rate=CONFIG[\"lr_backbone\"],\n    fp16=CONFIG[\"fp16\"],\n    bf16=CONFIG[\"bf16\"],\n    report_to=\"none\",\n    logging_steps=1,\n    max_steps=5,  # 5 steps to catch delayed NaN\n    remove_unused_columns=False,\n    save_strategy=\"no\",\n)\n\n_trainer = PiiTrainer(\n    model=_model,\n    args=_args,\n    train_dataset=_ds,\n    data_collator=_collator,\n    lr_backbone=CONFIG[\"lr_backbone\"],\n    lr_head=CONFIG[\"lr_head\"],\n)\n\n_result = _trainer.train()\n_loss = _result.training_loss\n\n# === NaN checks ===\nchecks_passed = True\n\n# Check 1: Training loss is a real number\nif math.isnan(_loss) or math.isinf(_loss):\n    print(f\"FAIL: Training loss is {_loss}\")\n    checks_passed = False\nelse:\n    print(f\"PASS: Training loss = {_loss:.4f}\")\n\n# Check 2: Model weights contain no NaN\n_nan_params = []\nfor name, param in _model.named_parameters():\n    if torch.isnan(param).any():\n        _nan_params.append(name)\nif _nan_params:\n    print(f\"FAIL: NaN in {len(_nan_params)} parameter tensors:\")\n    for p in _nan_params[:5]:\n        print(f\"  - {p}\")\n    checks_passed = False\nelse:\n    print(f\"PASS: All {sum(1 for _ in _model.parameters())} parameter tensors are finite\")\n\n# Check 3: Loss decreased (model is actually learning, not stuck)\n_log_history = _trainer.state.log_history\n_losses = [h[\"loss\"] for h in _log_history if \"loss\" in h]\nif len(_losses) >= 2:\n    if any(math.isnan(l) for l in _losses):\n        print(f\"FAIL: NaN in step losses: {_losses}\")\n        checks_passed = False\n    elif _losses[-1] < _losses[0]:\n        print(f\"PASS: Loss decreasing ({_losses[0]:.1f} → {_losses[-1]:.1f})\")\n    else:\n        print(f\"WARN: Loss not decreasing ({_losses[0]:.1f} → {_losses[-1]:.1f}) — may be OK for 5 steps\")\n\nif checks_passed:\n    print(\"\\n=== PREFLIGHT PASSED — safe to proceed ===\")\nelse:\n    raise RuntimeError(\n        \"PREFLIGHT FAILED — do not proceed with training. \"\n        \"Check mixed precision settings and optimizer config.\"\n    )\n\ndel _model, _trainer, _ds, _collator, _args, _result\nimport gc; gc.collect()\nif torch.cuda.is_available():\n    torch.cuda.empty_cache()",
+   "source": "import math\nimport gc\nimport numpy as np\nfrom datasets import Dataset\nfrom transformers import AutoTokenizer, TrainingArguments\nfrom datafog_pii_ner.data.collator import PiiDataCollator\nfrom datafog_pii_ner.data.label_schema import NUM_LABELS\nfrom datafog_pii_ner.model.pii_model import PiiNerConfig, PiiNerModel\nfrom datafog_pii_ner.training.train import PiiTrainer\n\n\ndef run_preflight(batch_size, seq_len, max_char_len, backbone, lr_backbone, lr_head, fp16, bf16):\n    \"\"\"Run 5 training steps and return (passed: bool, loss: float).\"\"\"\n    label = \"BF16\" if bf16 else \"FP16\" if fp16 else \"FP32\"\n    print(f\"\\n--- Testing {label} (batch_size={batch_size}) ---\")\n    \n    tokenizer = AutoTokenizer.from_pretrained(backbone)\n    n_samples = batch_size * 2\n\n    fake_data = {\n        \"input_ids\": np.random.randint(1, 1000, (n_samples, seq_len)).tolist(),\n        \"attention_mask\": np.ones((n_samples, seq_len), dtype=int).tolist(),\n        \"labels\": np.random.randint(0, NUM_LABELS, (n_samples, seq_len)).tolist(),\n        \"char_ids\": np.random.randint(0, 100, (n_samples, seq_len, max_char_len)).tolist(),\n    }\n    ds = Dataset.from_dict(fake_data)\n    collator = PiiDataCollator(tokenizer=tokenizer, max_char_len=max_char_len)\n    config = PiiNerConfig(backbone=backbone, num_labels=NUM_LABELS)\n    model = PiiNerModel(config)\n\n    args = TrainingArguments(\n        output_dir=\"/tmp/preflight_check\",\n        num_train_epochs=1,\n        per_device_train_batch_size=batch_size,\n        learning_rate=lr_backbone,\n        fp16=fp16,\n        bf16=bf16,\n        report_to=\"none\",\n        logging_steps=1,\n        max_steps=5,\n        remove_unused_columns=False,\n        save_strategy=\"no\",\n    )\n\n    trainer = PiiTrainer(\n        model=model, args=args, train_dataset=ds, data_collator=collator,\n        lr_backbone=lr_backbone, lr_head=lr_head,\n    )\n\n    try:\n        result = trainer.train()\n        loss = result.training_loss\n    except Exception as e:\n        print(f\"  FAIL: {type(e).__name__}: {e}\")\n        del model, trainer, ds, collator, args\n        gc.collect(); torch.cuda.empty_cache()\n        return False, float(\"nan\")\n\n    # NaN checks\n    passed = True\n\n    if math.isnan(loss) or math.isinf(loss):\n        print(f\"  FAIL: Training loss is {loss}\")\n        passed = False\n    else:\n        print(f\"  PASS: Training loss = {loss:.4f}\")\n\n    nan_params = [n for n, p in model.named_parameters() if torch.isnan(p).any()]\n    if nan_params:\n        print(f\"  FAIL: NaN in {len(nan_params)} parameter tensors\")\n        passed = False\n    else:\n        print(f\"  PASS: All parameter tensors are finite\")\n\n    log_losses = [h[\"loss\"] for h in trainer.state.log_history if \"loss\" in h]\n    if log_losses and any(math.isnan(l) for l in log_losses):\n        print(f\"  FAIL: NaN in step losses: {log_losses}\")\n        passed = False\n    elif len(log_losses) >= 2 and log_losses[-1] < log_losses[0]:\n        print(f\"  PASS: Loss decreasing ({log_losses[0]:.1f} → {log_losses[-1]:.1f})\")\n\n    del model, trainer, ds, collator, args, result\n    gc.collect(); torch.cuda.empty_cache()\n    return passed, loss\n\n\n# === Run preflight with auto-fallback ===\nprint(\"=== PREFLIGHT CHECK ===\")\n\npreflight_ok = False\npreflight_args = dict(\n    batch_size=CONFIG[\"batch_size\"],\n    seq_len=CONFIG[\"max_seq_len\"],\n    max_char_len=CONFIG[\"max_char_len\"],\n    backbone=CONFIG[\"backbone\"],\n    lr_backbone=CONFIG[\"lr_backbone\"],\n    lr_head=CONFIG[\"lr_head\"],\n)\n\n# Try 1: BF16 (if GPU supports it)\nif CONFIG.get(\"_try_bf16\", False):\n    ok, loss = run_preflight(**preflight_args, fp16=False, bf16=True)\n    if ok:\n        CONFIG[\"fp16\"] = False\n        CONFIG[\"bf16\"] = True\n        preflight_ok = True\n        print(f\"\\n=== BF16 PREFLIGHT PASSED (loss={loss:.4f}) ===\")\n    else:\n        print(f\"\\nBF16 failed — falling back to FP32\")\n\n# Try 2: FP32 (guaranteed fallback)\nif not preflight_ok:\n    ok, loss = run_preflight(**preflight_args, fp16=False, bf16=False)\n    if ok:\n        CONFIG[\"fp16\"] = False\n        CONFIG[\"bf16\"] = False\n        preflight_ok = True\n        print(f\"\\n=== FP32 PREFLIGHT PASSED (loss={loss:.4f}) ===\")\n    else:\n        raise RuntimeError(\n            \"PREFLIGHT FAILED on FP32 — something is fundamentally broken. \"\n            \"Check model architecture and data pipeline.\"\n        )\n\nprecision = \"BF16\" if CONFIG[\"bf16\"] else \"FP16\" if CONFIG[\"fp16\"] else \"FP32\"\nprint(f\"\\nSelected precision: {precision}\")\nprint(f\"Training will proceed with: fp16={CONFIG['fp16']}, bf16={CONFIG['bf16']}\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []