From bcb01bcb6dea67088498ce2fa82239f961d2997f Mon Sep 17 00:00:00 2001
From: pikaliov <pikaliov@gmail.com>
Date: Tue, 21 Apr 2026 18:42:29 +0300
Subject: [PATCH] Fix NaN: compute loss in fp32 outside AMP autocast

Root cause: GradScaler scales gradients by ~65536 in fp16, causing
logit_scale.exp() gradient to overflow. The learnable temperature
and similarity logits must stay in fp32.

Fix: model forward runs inside autocast(fp16), but loss computation
(similarity @ temperature + cross_entropy) runs outside in fp32.

Also: clamp logit_scale in logit-space before exp() and force
similarity computation to fp32.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 src/losses/multi_infonce.py  | 15 ++++++++++-----
 src/training/train_gtauav.py | 12 +++++++-----
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/src/losses/multi_infonce.py b/src/losses/multi_infonce.py
index cffcac2..ba70981 100644
--- a/src/losses/multi_infonce.py
+++ b/src/losses/multi_infonce.py
@@ -19,14 +19,15 @@ import torch.nn.functional as F
 def _symmetric_info_nce(
     emb_a: torch.Tensor,
     emb_b: torch.Tensor,
-    temperature: float,
+    temperature: float | torch.Tensor,
     label_smoothing: float,
     weight_a2b: float = 0.5,
     weight_b2a: float = 0.5,
 ) -> torch.Tensor:
     """Weighted symmetric InfoNCE. Positives on the diagonal."""
     batch_size = emb_a.size(0)
-    logits = emb_a @ emb_b.t() / temperature
+    # Compute logits in fp32 to avoid overflow with small temperature.
+    logits = emb_a.float() @ emb_b.float().t() / temperature
     targets = torch.arange(batch_size, device=emb_a.device)
     loss_a2b = F.cross_entropy(logits, targets, label_smoothing=label_smoothing)
     loss_b2a = F.cross_entropy(logits.t(), targets, label_smoothing=label_smoothing)
@@ -118,10 +119,14 @@ class InfoNCELoss(nn.Module):
             Dict with 'total', 'temperature', 'gate'.
         """
         if self.learnable_temperature:
-            # Clamp logit_scale to prevent tau from going out of bounds.
-            logit_scale = self.logit_scale.exp().clamp(
-                min=1.0 / self.tau_max, max=1.0 / self.tau_min,
+            # Clamp logit_scale in logit space first to prevent exp() overflow in fp16.
+            # tau_min=0.01 -> max logit_scale=ln(1/0.01)=4.6
+            # tau_max=0.5  -> min logit_scale=ln(1/0.5)=0.69
+            clamped = self.logit_scale.float().clamp(
+                min=math.log(1.0 / self.tau_max),
+                max=math.log(1.0 / self.tau_min),
             )
+            logit_scale = clamped.exp()
             tau = 1.0 / logit_scale
         else:
             tau = cosine_temperature(
diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py
index 6e5a8fc..cb26238 100644
--- a/src/training/train_gtauav.py
+++ b/src/training/train_gtauav.py
@@ -361,6 +361,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
             drone_img = batch["drone_img"].to(cfg.device, non_blocking=True)
             sat_img = batch["sat_img"].to(cfg.device, non_blocking=True)
 
+            # Model forward in AMP (fp16 for DINOv3/DGTRS encoders).
             with autocast(device_type="cuda", enabled=cfg.use_amp):
                 if cfg.baseline_mode:
                     embeddings = model(drone_img=drone_img, sat_img=sat_img)
@@ -372,11 +373,12 @@ def train(cfg: TrainConfigGTAUAV) -> None:
                         caption_l2=batch["caption_l2"],
                         caption_l3=batch["caption_l3"],
                     )
-                loss_dict = loss_fn(
-                    embeddings=embeddings,
-                    epoch=epoch,
-                    total_epochs=cfg.epochs,
-                )
+            # Loss in fp32 (learnable temperature gradient overflows in fp16).
+            loss_dict = loss_fn(
+                embeddings=embeddings,
+                epoch=epoch,
+                total_epochs=cfg.epochs,
+            )
 
             total_loss = loss_dict["total"]
             scaler.scale(total_loss).backward()