MONA bfloat16: safe low-precision (gamma=1e-6 needs bf16 exponent range)

- Switch MONA from fp32 to bfloat16 (same exponent range as fp32, no underflow)
- fp16 causes NaN: gamma=1e-6 falls into subnormal range (min normal ~6.1e-5)
- bf16 min normal ~1.2e-38, so 1e-6 is safe
- RTX 4090 supports bf16 natively
- Document bf16 vs fp16 vs fp32 comparison in README
- Update model summary: 3.5M MONA (last 12 blocks), 5.6M total trainable

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-21 22:13:35 +03:00
parent 9200772bea
commit e55734425c
2 changed files with 25 additions and 7 deletions

View File

@@ -68,7 +68,11 @@ 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 fp32 (gamma=1e-6 and small bottleneck underflow in fp16).
"""Apply MONA adapter in bfloat16.
bfloat16 has the same exponent range as fp32 (min ~1.2e-38), so gamma=1e-6
is safe. fp16 would underflow (min normal ~6.1e-5). RTX 4090 supports bf16
natively with comparable throughput to fp16.
Args:
x: Token features [B, N, D] where N includes CLS + register + patch tokens.
@@ -78,10 +82,9 @@ class MonaAdapter(nn.Module):
Adapted features [B, N, D] (same shape, residual connection).
"""
orig_dtype = x.dtype
with torch.amp.autocast("cuda", enabled=False):
x = x.float()
with torch.amp.autocast("cuda", dtype=torch.bfloat16):
identity = x
# Scaled LayerNorm.
# Scaled LayerNorm (gamma=1e-6 safe in bf16).
x = self.norm(x) * self.gamma + x * self.gammax
x = self.down(x) # [B, N, bottleneck]