diff --git a/src/models/adapters.py b/src/models/adapters.py index c2bf6b0..4727cf3 100644 --- a/src/models/adapters.py +++ b/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 - # Scaled LayerNorm. - x = self.norm(x) * self.gamma + x * self.gammax + 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 - x = self.down(x) # [B, N, bottleneck] + x = self.down(x) # [B, N, bottleneck] - B, N, C = x.shape - H, W = hw - n_special = N - H * W # CLS + register tokens + B, N, C = x.shape + H, W = hw + n_special = N - H * W # CLS + register tokens - # Separate special tokens (CLS, registers) from patch tokens. - special = x[:, :n_special] # [B, n_special, C] - patches = x[:, n_special:] # [B, H*W, C] + # Separate special tokens (CLS, registers) from patch tokens. + special = x[:, :n_special] # [B, n_special, C] + patches = x[:, n_special:] # [B, H*W, C] - # Reshape patches to 2D for convolutions. - patches = patches.reshape(B, H, W, C).permute(0, 3, 1, 2) # [B, C, H, W] - patches = self.mona_op(patches) - patches = patches.permute(0, 2, 3, 1).reshape(B, H * W, C) # [B, H*W, C] + # Reshape patches to 2D for convolutions. + patches = patches.reshape(B, H, W, C).permute(0, 3, 1, 2) # [B, C, H, W] + patches = self.mona_op(patches) + patches = patches.permute(0, 2, 3, 1).reshape(B, H * W, C) # [B, H*W, C] - # Recombine. - x = torch.cat([special, patches], dim=1) # [B, N, C] + # Recombine. + x = torch.cat([special, patches], dim=1) # [B, N, C] - x = F.gelu(x) - x = self.dropout(x) - x = self.up(x) # [B, N, D] + x = F.gelu(x) + x = self.dropout(x) + x = self.up(x) # [B, N, D] - return identity + x + return (identity + x).to(orig_dtype) # --------------------------------------------------------------------------- diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 581706e..da04071 100644 --- a/src/training/train_gtauav.py +++ b/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 ---