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,33 +77,36 @@ class MonaAdapter(nn.Module):
         Returns:
             Adapted features [B, N, D] (same shape, residual connection).
         """
-        identity = x
+        orig_dtype = x.dtype
-        # Scaled LayerNorm.
+        with torch.amp.autocast("cuda", enabled=False):
-        x = self.norm(x) * self.gamma + x * self.gammax
+            x = x.float()
+            identity = x
+            # Scaled LayerNorm.
+            x = self.norm(x) * self.gamma + x * self.gammax
 
-        x = self.down(x)  # [B, N, bottleneck]
+            x = self.down(x)  # [B, N, bottleneck]
 
-        B, N, C = x.shape
+            B, N, C = x.shape
-        H, W = hw
+            H, W = hw
-        n_special = N - H * W  # CLS + register tokens
+            n_special = N - H * W  # CLS + register tokens
 
-        # Separate special tokens (CLS, registers) from patch tokens.
+            # Separate special tokens (CLS, registers) from patch tokens.
-        special = x[:, :n_special]  # [B, n_special, C]
+            special = x[:, :n_special]  # [B, n_special, C]
-        patches = x[:, n_special:]  # [B, H*W, C]
+            patches = x[:, n_special:]  # [B, H*W, C]
 
-        # Reshape patches to 2D for convolutions.
+            # Reshape patches to 2D for convolutions.
-        patches = patches.reshape(B, H, W, C).permute(0, 3, 1, 2)  # [B, C, H, W]
+            patches = patches.reshape(B, H, W, C).permute(0, 3, 1, 2)  # [B, C, H, W]
-        patches = self.mona_op(patches)
+            patches = self.mona_op(patches)
-        patches = patches.permute(0, 2, 3, 1).reshape(B, H * W, C)  # [B, H*W, C]
+            patches = patches.permute(0, 2, 3, 1).reshape(B, H * W, C)  # [B, H*W, C]
 
-        # Recombine.
+            # Recombine.
-        x = torch.cat([special, patches], dim=1)  # [B, N, C]
+            x = torch.cat([special, patches], dim=1)  # [B, N, C]
 
-        x = F.gelu(x)
+            x = F.gelu(x)
-        x = self.dropout(x)
+            x = self.dropout(x)
-        x = self.up(x)  # [B, N, D]
+            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}",
+                tau=f"{step_metrics['temperature']:.4f}",
-                gq=f"{loss_dict['gate_q'].item():.3f}",
+                gq=f"{step_metrics['gate_q']:.3f}",
-                gg=f"{loss_dict['gate_g'].item():.3f}",
+                gg=f"{step_metrics['gate_g']:.3f}",
             )
 
             # --- Profiler step ---