9200772bea
- 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>
294 lines
11 KiB
Python
294 lines
11 KiB
Python
from __future__ import annotations
|
||
|
||
"""MONA and LoRA adapters for DINOv3 and DGTRS-CLIP.
|
||
|
||
MONA (Multi-Cognitive One-Shot Nested Adaptation, CVPR 2025):
|
||
Vision-specific adapter with multi-scale depthwise convolutions.
|
||
Applied to DINOv3 ViT-L/16 (after MSA and MLP in each block).
|
||
Original: github.com/Leiyi-Hu/mona (MIT license).
|
||
|
||
LoRA (Low-Rank Adaptation):
|
||
Low-rank matrices on Q/V attention projections.
|
||
Applied to DGTRS-CLIP text encoder (all 12 transformer blocks).
|
||
"""
|
||
|
||
import logging
|
||
import math
|
||
|
||
import coloredlogs
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
LOGGER = logging.getLogger("caption_test.adapters")
|
||
coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# MONA adapter (for DINOv3 ViT — 2D vision tokens)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class MonaOp(nn.Module):
|
||
"""Multi-cognitive visual filter: parallel depthwise convs (3×3, 5×5, 7×7)."""
|
||
|
||
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
|
||
return x + self.projector(x)
|
||
|
||
|
||
class MonaAdapter(nn.Module):
|
||
"""MONA adapter module for ViT blocks.
|
||
|
||
ScaledLayerNorm → Down(dim→bottleneck) → reshape 2D → MonaOp → reshape 1D
|
||
→ GELU → Dropout → Up(bottleneck→dim) + residual
|
||
|
||
Args:
|
||
dim: Input feature dimension (e.g. 1024 for ViT-L).
|
||
bottleneck: Bottleneck dimension for down/up projections.
|
||
dropout: Dropout rate.
|
||
"""
|
||
|
||
def __init__(self, dim: int = 1024, bottleneck: int = 64, dropout: float = 0.1) -> None:
|
||
super().__init__()
|
||
self.down = nn.Linear(dim, bottleneck)
|
||
self.up = nn.Linear(bottleneck, dim)
|
||
self.mona_op = MonaOp(bottleneck)
|
||
self.norm = nn.LayerNorm(dim)
|
||
self.dropout = nn.Dropout(p=dropout)
|
||
# Scaled init: gamma ≈ 0 at start (near-identity), gammax ≈ 1.
|
||
self.gamma = nn.Parameter(torch.ones(dim) * 1e-6)
|
||
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).
|
||
|
||
Args:
|
||
x: Token features [B, N, D] where N includes CLS + register + patch tokens.
|
||
hw: (H, W) spatial shape of patch tokens.
|
||
|
||
Returns:
|
||
Adapted features [B, N, D] (same shape, residual connection).
|
||
"""
|
||
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]
|
||
|
||
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]
|
||
|
||
# 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]
|
||
|
||
x = F.gelu(x)
|
||
x = self.dropout(x)
|
||
x = self.up(x) # [B, N, D]
|
||
|
||
return (identity + x).to(orig_dtype)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# LoRA adapter (for DGTRS-CLIP text encoder — 1D text tokens)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class LoRALinear(nn.Module):
|
||
"""LoRA adapter for a linear layer: output = W·x + (B·A)·x.
|
||
|
||
Args:
|
||
in_features: Input dimension.
|
||
out_features: Output dimension.
|
||
rank: LoRA rank (low-rank decomposition dimension).
|
||
alpha: LoRA scaling factor (effective scale = alpha / rank).
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
in_features: int,
|
||
out_features: int,
|
||
rank: int = 4,
|
||
alpha: float = 1.0,
|
||
) -> None:
|
||
super().__init__()
|
||
self.rank = rank
|
||
self.scale = alpha / rank
|
||
# A: down-projection, B: up-projection. Init: A ~ N(0, 1/rank), B = 0.
|
||
self.A = nn.Parameter(torch.randn(rank, in_features) / math.sqrt(rank))
|
||
self.B = nn.Parameter(torch.zeros(out_features, rank))
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
"""Compute LoRA delta: (B @ A) @ x * scale. Always fp32."""
|
||
with torch.amp.autocast("cuda", enabled=False):
|
||
x = x.float()
|
||
return (x @ self.A.t() @ self.B.t()) * self.scale
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Injection functions — apply adapters to existing frozen models
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def inject_mona_into_dinov3(
|
||
model: nn.Module,
|
||
bottleneck: int = 64,
|
||
dropout: float = 0.1,
|
||
last_n_blocks: int = 12,
|
||
) -> int:
|
||
"""Inject MONA adapters into a frozen DINOv3ViT model.
|
||
|
||
Adds two MonaAdapter modules per block (after attention, after MLP).
|
||
Only adapter parameters are trainable. Early blocks remain pure frozen
|
||
DINOv3 (low-level features don't need spatial adaptation).
|
||
|
||
Args:
|
||
model: DINOv3ViT model (already frozen).
|
||
bottleneck: MONA bottleneck dimension.
|
||
dropout: MONA dropout rate.
|
||
last_n_blocks: Number of last blocks to adapt (default 12 of 24).
|
||
|
||
Returns:
|
||
Number of trainable adapter parameters added.
|
||
"""
|
||
dim = model.embed_dim
|
||
n_adapters = 0
|
||
total_blocks = len(model.layer)
|
||
start_idx = total_blocks - last_n_blocks
|
||
|
||
for i, block in enumerate(model.layer):
|
||
if i < start_idx:
|
||
continue # Skip early blocks — frozen, no adaptation.
|
||
|
||
# Create adapters.
|
||
block.mona_attn = MonaAdapter(dim, bottleneck, dropout)
|
||
block.mona_mlp = MonaAdapter(dim, bottleneck, dropout)
|
||
|
||
# Wrap the original forward to inject adapters.
|
||
original_forward = block.forward
|
||
|
||
def make_mona_forward(blk, orig_fwd):
|
||
def mona_forward(x: torch.Tensor) -> torch.Tensor:
|
||
B, N, C = x.shape
|
||
# DINOv3: [CLS, 4 registers, H*W patches]
|
||
n_special = 1 + 4 # CLS + registers
|
||
n_patches = N - n_special
|
||
H = W = int(math.sqrt(n_patches))
|
||
|
||
# Original block forward (attention + MLP with layer scale).
|
||
x = x + blk.layer_scale1(blk.attention(blk.norm1(x)))
|
||
x = blk.mona_attn(x, (H, W))
|
||
x = x + blk.layer_scale2(blk.mlp(blk.norm2(x)))
|
||
x = blk.mona_mlp(x, (H, W))
|
||
return x
|
||
return mona_forward
|
||
|
||
block.forward = make_mona_forward(block, original_forward)
|
||
n_adapters += 2
|
||
|
||
n_params = sum(p.numel() for n, p in model.named_parameters() if "mona" in n)
|
||
LOGGER.info(
|
||
"🔧 MONA injected: %d adapters (blocks %d-%d of %d), %s trainable params (bottleneck=%d)",
|
||
n_adapters, start_idx, total_blocks - 1, total_blocks, f"{n_params:,}", bottleneck,
|
||
)
|
||
return n_params
|
||
|
||
|
||
def inject_lora_into_dgtrs(
|
||
model: nn.Module,
|
||
rank: int = 4,
|
||
alpha: float = 1.0,
|
||
) -> int:
|
||
"""Inject LoRA adapters into DGTRS-CLIP text encoder attention layers.
|
||
|
||
Adds LoRA to Q and V projections in each ResidualAttentionBlock.
|
||
Original weights stay frozen, LoRA params are trainable.
|
||
|
||
Args:
|
||
model: DGTRSTextEncoder model.
|
||
rank: LoRA rank.
|
||
alpha: LoRA scaling factor.
|
||
|
||
Returns:
|
||
Number of trainable LoRA parameters added.
|
||
"""
|
||
n_adapted = 0
|
||
|
||
for i, block in enumerate(model.transformer.resblocks):
|
||
attn = block.attn
|
||
# nn.MultiheadAttention uses in_proj_weight [3*dim, dim] for Q, K, V.
|
||
embed_dim = attn.embed_dim
|
||
|
||
# LoRA on Q and V (not K — standard practice).
|
||
block.lora_q = LoRALinear(embed_dim, embed_dim, rank, alpha)
|
||
block.lora_v = LoRALinear(embed_dim, embed_dim, rank, alpha)
|
||
|
||
# Wrap attention forward to add LoRA deltas.
|
||
original_attention = block.attention
|
||
|
||
def make_lora_attention(blk):
|
||
def lora_attention(x: torch.Tensor) -> torch.Tensor:
|
||
# x shape: [T, B, D] (sequence-first for nn.MultiheadAttention)
|
||
attn_mod = blk.attn
|
||
attn_mask = blk.attn_mask
|
||
if attn_mask is not None:
|
||
attn_mask = attn_mask.to(dtype=x.dtype, device=x.device)
|
||
|
||
# Compute QKV using original in_proj_weight.
|
||
T, B, D = x.shape
|
||
|
||
# Original QKV.
|
||
qkv = F.linear(x, attn_mod.in_proj_weight, attn_mod.in_proj_bias)
|
||
q, k, v = qkv.chunk(3, dim=-1)
|
||
|
||
# Add LoRA deltas to Q and V.
|
||
# LoRA expects [B, T, D], but x is [T, B, D].
|
||
x_btd = x.permute(1, 0, 2) # [B, T, D]
|
||
q = q + blk.lora_q(x_btd).permute(1, 0, 2)
|
||
v = v + blk.lora_v(x_btd).permute(1, 0, 2)
|
||
|
||
# Multi-head attention.
|
||
num_heads = attn_mod.num_heads
|
||
head_dim = D // num_heads
|
||
q = q.reshape(T, B * num_heads, head_dim).transpose(0, 1)
|
||
k = k.reshape(T, B * num_heads, head_dim).transpose(0, 1)
|
||
v = v.reshape(T, B * num_heads, head_dim).transpose(0, 1)
|
||
|
||
attn_output = F.scaled_dot_product_attention(
|
||
q, k, v, attn_mask=attn_mask,
|
||
)
|
||
attn_output = attn_output.transpose(0, 1).reshape(T, B, D)
|
||
attn_output = attn_mod.out_proj(attn_output)
|
||
return attn_output
|
||
return lora_attention
|
||
|
||
block.attention = make_lora_attention(block)
|
||
n_adapted += 1
|
||
|
||
n_params = sum(
|
||
p.numel() for n, p in model.named_parameters() if "lora" in n
|
||
)
|
||
LOGGER.info(
|
||
"🔧 LoRA injected: %d blocks, rank=%d, %s trainable params",
|
||
n_adapted, rank, f"{n_params:,}",
|
||
)
|
||
return n_params
|