Spaces:
Running on Zero
Running on Zero
File size: 5,161 Bytes
87608ea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | """Runtime wrappers for DINOv2 checkpointing and optimized attention.
The vendored DINOv2 source is kept close to upstream. These two helpers apply
PXDepth-specific runtime behavior without editing every upstream block.
"""
from typing import Any
import torch
import torch.nn as nn
import torch.nn.functional as F
try:
from flash_attn.flash_attn_interface import flash_attn_func
FLASH_ATTN_AVAILABLE = True
except Exception:
flash_attn_func = None
FLASH_ATTN_AVAILABLE = False
def wrap_module_with_gradient_checkpointing(module: nn.Module) -> nn.Module:
"""Recompute a module's forward pass during backward to save memory.
A dynamic subclass replaces only the supplied module instance's class. Its
original ``forward`` remains the computation being checkpointed, so
parameters and state-dict names are unchanged.
Args:
module: DINO transformer block to modify in place. Its forward inputs
and output retain the upstream block contract, normally token
tensors ``[B, N, C]`` plus optional attention metadata.
Returns:
The same module instance with a checkpointed ``forward`` method.
"""
from torch.utils.checkpoint import checkpoint
class _CheckpointingWrapper(module.__class__):
"""Per-instance subclass that checkpoints the inherited forward call.
It introduces no parameters or buffers, preserving the wrapped DINO
block's state-dict schema and tensor interface.
Instances are created by mutating the class of one existing block.
"""
def forward(self, *args: Any, **kwargs: Any) -> Any:
"""Run the inherited module under non-reentrant checkpointing.
Args:
*args: Positional arguments accepted by the wrapped block,
typically a token tensor ``[B,N,C]``.
**kwargs: Keyword arguments accepted by the wrapped block.
Returns:
The wrapped block's original output, normally ``[B,N,C]``.
"""
return checkpoint(super().forward, *args, use_reentrant=False, **kwargs)
module.__class__ = _CheckpointingWrapper
return module
def wrap_dinov2_attention_with_sdpa(module: nn.Module) -> nn.Module:
"""Replace one DINOv2 attention forward path with Flash-Attention or SDPA.
Flash-Attention is selected only for CUDA FP16/BF16 tensors without an
additive attention bias. Other inputs use PyTorch scaled dot-product
attention. The wrapper reuses the upstream QKV and output projections, so
checkpoint keys and numerical semantics remain compatible.
Args:
module: DINOv2 attention module with ``qkv``, ``num_heads``, ``scale``,
``proj``, and ``proj_drop`` attributes.
Returns:
The same attention module instance with an optimized forward method.
"""
if torch.__version__ < "2.0":
raise RuntimeError("SDPA requires PyTorch 2.0 or later")
class _AttentionWrapper(module.__class__):
"""Per-instance attention subclass dispatching to optimized kernels.
It reuses all upstream projections and introduces no new checkpoint
state, so wrapping does not affect serialization compatibility.
Instances are created by mutating one existing DINO attention module.
"""
def forward(self, x: torch.Tensor, attn_bias: Any = None) -> torch.Tensor:
"""Apply multi-head self-attention to DINO patch and special tokens.
Args:
x: Input token tensor ``[B, N, C]``.
attn_bias: Optional bias broadcastable to ``[B,H,N,N]``. A
non-``None`` value disables the external Flash-Attention path.
Returns:
Projected token tensor ``[B, N, C]``.
"""
batch, length, channels = x.shape
qkv = self.qkv(x).reshape(
batch,
length,
3,
self.num_heads,
channels // self.num_heads,
).permute(2, 0, 3, 1, 4)
query, key, value = torch.unbind(qkv, 0)
use_flash = (
FLASH_ATTN_AVAILABLE
and attn_bias is None
and query.is_cuda
and query.dtype in (torch.float16, torch.bfloat16)
)
if use_flash:
query = query.permute(0, 2, 1, 3).contiguous()
key = key.permute(0, 2, 1, 3).contiguous()
value = value.permute(0, 2, 1, 3).contiguous()
out = flash_attn_func(
query,
key,
value,
dropout_p=0.0,
softmax_scale=self.scale,
causal=False,
).reshape(batch, length, channels)
else:
out = F.scaled_dot_product_attention(query, key, value, attn_bias)
out = out.permute(0, 2, 1, 3).reshape(batch, length, channels)
return self.proj_drop(self.proj(out))
module.__class__ = _AttentionWrapper
return module
|