Fix NaN loss: revert MONA to fp32, fix loss logging

- MONA fp16 causes NaN (gamma=1e-6 underflows in fp16 min subnormal ~6e-8)
- Revert MONA forward to fp32 with autocast(enabled=False), cast output back
- Fix loss CSV: save raw_loss before backward() (tensor consumed after backward)
- Verified: loss=3.78, no NaN, bs=48 peak=21.4 GB

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

src/models/adapters.py

@@ -68,7 +68,7 @@ class MonaAdapter(nn.Module):
         self.gammax = nn.Parameter(torch.ones(dim))
 
     def forward(self, x: torch.Tensor, hw: tuple[int, int]) -> torch.Tensor:
-        """Apply MONA adapter in AMP-native fp16.
+        """Apply MONA adapter in fp32 (gamma=1e-6 and small bottleneck underflow in fp16).
 
         Args:
             x: Token features [B, N, D] where N includes CLS + register + patch tokens.
@@ -77,6 +77,9 @@ class MonaAdapter(nn.Module):
         Returns:
             Adapted features [B, N, D] (same shape, residual connection).
         """
+        orig_dtype = x.dtype
+        with torch.amp.autocast("cuda", enabled=False):
+            x = x.float()
             identity = x
             # Scaled LayerNorm.
             x = self.norm(x) * self.gamma + x * self.gammax
@@ -103,7 +106,7 @@ class MonaAdapter(nn.Module):
             x = self.dropout(x)
             x = self.up(x)  # [B, N, D]
 
-        return identity + x
+            return (identity + x).to(orig_dtype)
 
 
 # ---------------------------------------------------------------------------

src/training/train_gtauav.py

@@ -540,6 +540,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
             )
 
             # Scale loss by accumulation steps so gradients average correctly.
+            raw_loss = float(loss_dict["total"].item())  # save before backward
             total_loss = loss_dict["total"] / accum
             scaler.scale(total_loss).backward()
 
@@ -568,7 +569,6 @@ def train(cfg: TrainConfigGTAUAV) -> None:
                 global_step += 1
 
             # --- Per-batch tracking (log unscaled loss) ---
-            raw_loss = total_loss.item() * accum  # undo /accum for logging
             step_metrics = {
                 "loss": raw_loss,
                 "temperature": float(loss_dict["temperature"].item()),
@@ -585,9 +585,9 @@ def train(cfg: TrainConfigGTAUAV) -> None:
 
             pbar.set_postfix(
                 loss=f"{raw_loss:.3f}",
-                tau=f"{loss_dict['temperature'].item():.4f}",
-                gq=f"{loss_dict['gate_q'].item():.3f}",
-                gg=f"{loss_dict['gate_g'].item():.3f}",
+                tau=f"{step_metrics['temperature']:.4f}",
+                gq=f"{step_metrics['gate_q']:.3f}",
+                gg=f"{step_metrics['gate_g']:.3f}",
             )
 
             # --- Profiler step ---