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:
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user