File size: 9,808 Bytes
c855122 | 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 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 | """KernelConfig-mappable drop-in forwards for Hugging Face Transformers ViT encoders.
``NeuronDINOv3ViTModel`` -> ``DINOv3ViTModel`` (all DINOv3 ViT sizes: S/S+/B/L/H+; 7B needs TP)
``NeuronVJEPA2Encoder`` -> ``VJEPA2Encoder`` (V-JEPA2 video encoder, up to 2048 tokens / clip)
Both run the whole transformer stack on the NKI kernels (``encoder.EncoderRunner``) and keep the
model's own patch embedding, RoPE position computation and final LayerNorm. Weights are folded
once, on the first device forward, from the stock module's parameters (checkpoint loading is
untouched) and cached on the module. Anything the kernels do not support (off-device input,
training mode, > 2048 tokens, unsupported head size) falls back to the stock forward.
Usage (graft, works on any Transformers version)::
import types
model.forward = types.MethodType(kernels.NeuronDINOv3ViTModel.forward, model)
model = torch.compile(model, backend="neuron")
``torch.compile(backend="neuron")`` is part of the intended configuration; the classes set
``can_torch_compile = True`` so the ``kernels`` library keeps them under ``Mode.TORCH_COMPILE``.
"""
from __future__ import annotations
import torch
from .encoder import LNC, BlockWeights, EncoderRunner, pack_block, rope_tables
_PACK = "_vit_enc_neuron_packed"
_TABS = "_vit_enc_neuron_tabs"
_MAX_TOKENS = 2048
def _rup(n: int, m: int) -> int:
return ((n + m - 1) // m) * m
def _supported(H: int, heads: int) -> bool:
D = H // heads
return H % 128 == 0 and D in (64, 128) and H == heads * D
def _stack(hs: torch.Tensor, N: int):
B, n, H = hs.shape
x = hs.float()
if N > n:
x = torch.cat([x, x.new_zeros(B, N - n, H)], 1)
return x.reshape(B * N, H).contiguous()
# ------------------------------------------------------------------------------------------------
# DINOv3
# ------------------------------------------------------------------------------------------------
def _dinov3_layers(model):
enc = getattr(model, "model", None) # newer Transformers: model.model.layer
return enc.layer if enc is not None and hasattr(enc, "layer") else model.layer
def _dinov3_block(lyr) -> BlockWeights:
a = lyr.attention
bw = BlockWeights(
ln1_w=lyr.norm1.weight, ln1_b=lyr.norm1.bias, ln1_eps=lyr.norm1.eps,
wq=a.q_proj.weight, bq=a.q_proj.bias, wk=a.k_proj.weight, bk=a.k_proj.bias,
wv=a.v_proj.weight, bv=a.v_proj.bias, wo=a.o_proj.weight, bo=a.o_proj.bias,
ln2_w=lyr.norm2.weight, ln2_b=lyr.norm2.bias, ln2_eps=lyr.norm2.eps,
ls1=lyr.layer_scale1.lambda1, ls2=lyr.layer_scale2.lambda1, scale=float(a.scaling))
m = lyr.mlp
if hasattr(m, "gate_proj"):
bw.gate_w, bw.gate_b = m.gate_proj.weight, m.gate_proj.bias
bw.up_w, bw.up_b = m.up_proj.weight, m.up_proj.bias
bw.down_w, bw.down_b = m.down_proj.weight, m.down_proj.bias
else:
bw.fc1_w, bw.fc1_b = m.up_proj.weight, m.up_proj.bias
bw.fc2_w, bw.fc2_b = m.down_proj.weight, m.down_proj.bias
return bw
def _dinov3_ok(model) -> bool:
c = model.config
if not _supported(c.hidden_size, c.num_attention_heads):
return False
if getattr(c, "use_gated_mlp", False):
return getattr(c, "hidden_act", "silu") == "silu"
return getattr(c, "hidden_act", "gelu") == "gelu"
@torch._dynamo.disable
def _dinov3_freeze(model):
if model.training:
raise RuntimeError("NeuronDINOv3ViTModel: these kernels have no autograd rule; call model.eval() first.")
dev = model.norm.weight.device
heads = model.config.num_attention_heads
packed = [pack_block(_dinov3_block(l), heads, "half", dev) for l in _dinov3_layers(model)]
setattr(model, _PACK, EncoderRunner(packed))
setattr(model, _TABS, {})
@torch._dynamo.disable
def _dinov3_tab(model, pixel_values, n_prefix, N):
key = (tuple(pixel_values.shape[-2:]), N)
tabs = getattr(model, _TABS)
if key not in tabs:
cos, sin = model.rope_embeddings(pixel_values[:1].float()) # [n_patches, D] each
c = model.config
tabs[key] = rope_tables(cos, sin, n_prefix, N, c.num_attention_heads,
c.hidden_size // c.num_attention_heads, "half", model.norm.weight.device)
return tabs[key]
class NeuronDINOv3ViTModel(torch.nn.Module):
"""Drop-in forward for ``transformers.DINOv3ViTModel`` (bf16 NKI encoder, fp32 residual stream).
Returns ``last_hidden_state`` (final-normed, ``[B, 1 + registers + patches, H]``) and
``pooler_output`` (CLS), like the stock model. Any batch size; images up to 2048 tokens
(e.g. 720x720 at patch 16).
"""
can_torch_compile = True
has_backward = False
def forward(self, pixel_values, bool_masked_pos=None, head_mask=None, **kwargs):
from transformers.modeling_outputs import BaseModelOutputWithPooling
pv = pixel_values.to(self.embeddings.patch_embeddings.weight.dtype)
hs = self.embeddings(pv, bool_masked_pos=bool_masked_pos) # [B, n, H]
B, n, H = hs.shape
N = _rup(n, 128 * LNC)
if (hs.device.type != "neuron" or self.training or head_mask is not None or N > _MAX_TOKENS
or not _dinov3_ok(self)):
import transformers.models.dinov3_vit.modeling_dinov3_vit as mv
if head_mask is not None:
kwargs["head_mask"] = head_mask
return mv.DINOv3ViTModel.forward(self, pixel_values, bool_masked_pos=bool_masked_pos, **kwargs)
if not hasattr(self, _PACK):
_dinov3_freeze(self)
n_prefix = 1 + self.config.num_register_tokens
tab = _dinov3_tab(self, pixel_values, n_prefix, N)
x = getattr(self, _PACK)(_stack(hs, N), tab, n, B)
x = x.reshape(B, N, H)[:, :n]
out = self.norm(x.to(self.norm.weight.dtype))
return BaseModelOutputWithPooling(last_hidden_state=out, pooler_output=out[:, 0, :])
# ------------------------------------------------------------------------------------------------
# V-JEPA2
# ------------------------------------------------------------------------------------------------
def _vjepa2_block(lyr) -> BlockWeights:
a = lyr.attention
return BlockWeights(
ln1_w=lyr.norm1.weight, ln1_b=lyr.norm1.bias, ln1_eps=lyr.norm1.eps,
wq=a.query.weight, bq=a.query.bias, wk=a.key.weight, bk=a.key.bias,
wv=a.value.weight, bv=a.value.bias, wo=a.proj.weight, bo=a.proj.bias,
ln2_w=lyr.norm2.weight, ln2_b=lyr.norm2.bias, ln2_eps=lyr.norm2.eps,
fc1_w=lyr.mlp.fc1.weight, fc1_b=lyr.mlp.fc1.bias, fc2_w=lyr.mlp.fc2.weight, fc2_b=lyr.mlp.fc2.bias,
scale=float(a.scaling))
@torch._dynamo.disable
def _vjepa2_tab(encoder, N, n):
"""Signed-sin 3D RoPE tables from the model's own position ids (interleaved pairing)."""
a = encoder.layer[0].attention
H, D = a.all_head_size, a.attention_head_size
dev = a.query.weight.device
pos = a.get_position_ids(torch.zeros(1, n, H, device=dev), masks=None)
to1d = lambda t: (t if t.dim() == 1 else t.reshape(-1, n)[0]).float().cpu()
cos = torch.ones(n, D)
sin = torch.zeros(n, D)
s = 0
for sub, p in zip((a.d_dim, a.h_dim, a.w_dim), (to1d(pos[0]), to1d(pos[1]), to1d(pos[2]))):
half = sub // 2
omega = 1.0 / (10000.0 ** (torch.arange(half, dtype=torch.float32) / float(half)))
f = p[:, None] * omega
c, sn = f.cos(), f.sin()
cos[:, s:s + half] = c; cos[:, s + half:s + sub] = c
sin[:, s:s + half] = sn; sin[:, s + half:s + sub] = sn
s += sub
g = torch.ones(D); g[0::2] = -1.0
return rope_tables(cos, sin * g, 0, N, H // D, D, "interleaved", dev)
@torch._dynamo.disable
def _vjepa2_freeze(encoder):
if encoder.training:
raise RuntimeError("NeuronVJEPA2Encoder: these kernels have no autograd rule; call model.eval() first.")
a = encoder.layer[0].attention
dev = a.query.weight.device
heads = a.all_head_size // a.attention_head_size
setattr(encoder, _PACK, EncoderRunner([pack_block(_vjepa2_block(l), heads, "interleaved", dev)
for l in encoder.layer]))
setattr(encoder, _TABS, {})
class NeuronVJEPA2Encoder(torch.nn.Module):
"""Drop-in forward for ``transformers.VJEPA2Encoder``.
Any batch size; clips up to 2048 tokens (e.g. 16 frames at 256 px, tubelet 2, patch 16).
"""
can_torch_compile = True
has_backward = False
def forward(self, pixel_values_videos=None, head_mask=None, output_attentions=False,
output_hidden_states=False, **kwargs):
from transformers.modeling_outputs import BaseModelOutput
hs = self.embeddings(pixel_values_videos) # [B, n, H]
B, n, H = hs.shape
N = _rup(n, 128 * LNC)
a = self.layer[0].attention
if (hs.device.type != "neuron" or self.training or head_mask is not None or output_attentions
or output_hidden_states or N > _MAX_TOKENS
or not _supported(H, a.all_head_size // a.attention_head_size)):
import transformers.models.vjepa2.modeling_vjepa2 as mv
return mv.VJEPA2Encoder.forward(self, pixel_values_videos, head_mask=head_mask,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states, **kwargs)
if not hasattr(self, _PACK):
_vjepa2_freeze(self)
tabs = getattr(self, _TABS)
if N not in tabs:
tabs[N] = _vjepa2_tab(self, N, n)
x = getattr(self, _PACK)(_stack(hs, N), tabs[N], n, B)
out = self.layernorm(x.reshape(B, N, H)[:, :n])
return BaseModelOutput(last_hidden_state=out)
|