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

@@ -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(