facebookresearch / facebookresearch/vjepa2
Native encoder ~500x slower than HF port: fp16 SDPA dtype crash + un-fused attention (frozen-feature extraction)
- Dominant language
- Python
- Stars
- 4.6k
- Forks
- 573
- PR merge metrics
- No merged PRs in 30d
Description
**Two issues in the native V-JEPA 2 / 2.1 encoder path that make frozen-feature extraction ~500× slower than the HuggingFace port.**
Setup: RTX 3070 Ti 8 GB, torch 2.6.0+cu124, extracting frozen features from `vjepa2_1_vit_large_384` and `vjepa2_vit_large` via `torch.hub.load("facebookresearch/vjepa2", ...)`.
**1. fp16 forward crashes (mixed dtype in SDPA).** Casting the encoder + input to `.half()` throws inside attention:
```
RuntimeError: Expected query, key, and value to have the same dtype,
but got query.dtype: float key.dtype: float and value.dtype: c10::Half instead.
```
Some tensor (a norm/pos-embed buffer or a projection) stays fp32 while value is fp16, so `scaled_dot_product_attention` rejects the mix. `.half()` on the whole module isn't enough.
**2. Attention is not memory-efficient / fused at the token counts these models actually use.** At the native `vjepa2_1_vit_large_384` setting (384px, 64 frames → ~18k tokens) a **single-clip forward exceeds 250 s** on this GPU; even at 256px/16 frames (~2k tokens) it is far slower than expected. By contrast the HF port `facebook/vjepa2-vitl-fpc64-256` with `attn_implementation="sdpa"` extracts the same clips at **~0.8 s/clip**. The gap is consistent with the native path not routing attention through a fused/flash kernel.
**Workaround we used:** run the native models in **fp32** (avoids #1) at **256px / 16 frames** (makes #2 tractable, ~0.5 s/clip) — but that forces a resolution/frame-budget compromise on 384-native checkpoints.
**Suggestions:** (a) make the encoder fp16-safe (cast pos-embed/norm buffers, or document that fp16 is unsupported); (b) route attention through `F.scaled_dot_product_attention` so the native path gets flash/mem-efficient kernels like the HF port. Happy to share a minimal repro script.
Repo @ main, 2026-07-23. (Related: #174, the committed `localhost:8300` base-URL that also blocks hub loads.)
Contributor guide
Research direction
Start from the native encoder path reached by torch.hub.load("facebookresearch/vjepa2", ...), then reproduce the fp16 failure and the 256px/16-frame and 384px/64-frame timing cases described here. Trace the attention inputs and compare the native path with the HuggingFace port's SDPA configuration. Done means the reported dtype crash is resolved and native attention no longer has the documented performance gap at the tested settings.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100