Pikaliov/caption-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
- 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>
This commit is contained in:
pikaliov
2026-04-21 22:09:49 +03:00
parent 916854f124
commit 9200772bea
2 changed files with 28 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
src/models/adapters.py
View File

@@ -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)
# ---------------------------------------------------------------------------