Optimize MONA: fp16, remove conv7x7, bottleneck 64→32

- Remove forced fp32 cast in MONA forward (runs in AMP fp16 now)
- Remove conv7x7 from MonaOp (keep 3x3 + 5x5 only)
- Reduce default bottleneck from 64 to 32
- MONA params: 3.5M (was 7.0M, -50%)
- Total trainable: 7.0M (was 10.5M)
- Peak VRAM at bs=24: 18.6 GB (was 20.3 GB before fp16)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-21 21:40:23 +03:00
parent 6b7bcae198
commit d279a6e745
2 changed files with 6 additions and 11 deletions

View File

@@ -29,18 +29,17 @@ coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(lev
# ---------------------------------------------------------------------------
class MonaOp(nn.Module):
"""Multi-cognitive visual filter: parallel depthwise convs (3×3, 5×5, 7×7)."""
"""Multi-cognitive visual filter: parallel depthwise convs (3×3, 5×5)."""
def __init__(self, channels: int) -> None:
super().__init__()
self.conv3 = nn.Conv2d(channels, channels, kernel_size=3, padding=1, groups=channels)
self.conv5 = nn.Conv2d(channels, channels, kernel_size=5, padding=2, groups=channels)
self.conv7 = nn.Conv2d(channels, channels, kernel_size=7, padding=3, groups=channels)
self.projector = nn.Conv2d(channels, channels, kernel_size=1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
identity = x
x = (self.conv3(x) + self.conv5(x) + self.conv7(x)) / 3.0 + identity
x = (self.conv3(x) + self.conv5(x)) / 2.0 + identity
return x + self.projector(x)
@@ -56,7 +55,7 @@ class MonaAdapter(nn.Module):
dropout: Dropout rate.
"""
def __init__(self, dim: int = 1024, bottleneck: int = 64, dropout: float = 0.1) -> None:
def __init__(self, dim: int = 1024, bottleneck: int = 32, dropout: float = 0.1) -> None:
super().__init__()
self.down = nn.Linear(dim, bottleneck)
self.up = nn.Linear(bottleneck, dim)
@@ -68,7 +67,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 (runs in fp32 to avoid AMP overflow).
"""Apply MONA adapter in AMP-native fp16.
Args:
x: Token features [B, N, D] where N includes CLS + register + patch tokens.
@@ -77,10 +76,6 @@ class MonaAdapter(nn.Module):
Returns:
Adapted features [B, N, D] (same shape, residual connection).
"""
with torch.amp.autocast("cuda", enabled=False):
return self._forward_fp32(x.float(), hw)
def _forward_fp32(self, x: torch.Tensor, hw: tuple[int, int]) -> torch.Tensor:
identity = x
# Scaled LayerNorm.
x = self.norm(x) * self.gamma + x * self.gammax
@@ -151,7 +146,7 @@ class LoRALinear(nn.Module):
def inject_mona_into_dinov3(
model: nn.Module,
bottleneck: int = 64,
bottleneck: int = 32,
dropout: float = 0.1,
) -> int:
"""Inject MONA adapters into a frozen DINOv3ViT model.

View File

@@ -298,7 +298,7 @@ class AsymmetricEncoder(nn.Module):
init_gate: float = 0.7,
baseline_mode: bool = False,
shared_encoder: bool = True,
mona_bottleneck: int = 64,
mona_bottleneck: int = 32,
lora_rank: int = 4,
device: str = "cuda",
) -> None: