PR #45497: Add V-JEPA 2.1 inference support

• Repository: huggingface/transformers
• Author: davevanveen
• State: open | merged: False
• Link: https://github.com/huggingface/transformers/pull/45497

Description (problem / solution / changelog)

What does this PR do?

Adds inference support and checkpoint conversion for Meta V-JEPA 2.1 pretrained backbones within the existing vjepa2 model family.

V-JEPA 2.1 was released by Meta on 2026-03-16 with four pretrained encoders at 384 resolution: ViT-B (80M), ViT-L (300M), ViT-g (1B), and ViT-G (2B). This PR extends the existing VJEPA2Model with backward-compatible config fields and modeling changes to support loading these checkpoints.

Changes

Config (configuration_vjepa2.py): 8 new fields with backward-compatible defaults that preserve V-JEPA 2.0 behavior:

• use_rope_interleave, use_modality_embeddings, interpolate_rope, return_all_tokens, img_temporal_dim_size
• Predictor-only: teacher_embed_dim, n_output_distillation, hierarchical_layers

Modeling (modeling_vjepa2.py):

• Corrected RoPE implementation (repeat_interleave vs repeat) with config toggle
• Learnable modality embeddings (image/video) for both encoder and predictor, with top-level modality routing
• Separate image patch embedding (tubelet_size=1) when img_temporal_dim_size is set
• Hierarchical feature extraction with per-layer norms_block in encoder
• Smart encoder output: concatenated features when predictor needs them (n_output_distillation > 1), single-norm output for get_vision_features() / skip_predictor=True
• RoPE position interpolation for flexible resolution (encoder only, matches Meta's behavior)
• Multi-layer predictor embed (Linear+GELU+Linear) for n_output_distillation > 1
• Predictor context token projection (proj_context) for return_all_tokens
• Teacher embedding dimension-based output projection sizing
• VJEPA2ForVideoClassification guard against unsupported hierarchical configs

Converter (convert_vjepa2_to_hf.py):

• Four new 2.1 model variants with correct architecture parameters
• Key remappings for all new layers (modality embeddings, norms_block, patch_embed_img, proj_context)
• Checkpoint key handling (ema_encoder for distilled models, target_encoder for self-supervised)
• Updated test function to handle 2.1 tuple predictor returns

Tests (test_modeling_vjepa2.py):

• Config defaults test (verifies 2.0 backward compatibility)
• Fast forward pass with shape assertions for n_output_distillation=1 (distilled path)
• Fast forward pass with shape assertions for n_output_distillation=4 (multi-layer path)
• All 90 existing tests pass unchanged

Docs (vjepa2.md): V-JEPA 2.1 section with checkpoint table and architecture notes.

Verification

Verified end-to-end against Meta's reference implementation on two checkpoints covering both architecture paths:

ViT-B/384 (80M, n_output_distillation=1, distilled):

• All weight keys match (strict load)
• Encoder: max diff 0.0001
• Predictor: max diff 0.008

ViT-g/384 (1B, n_output_distillation=4, self-supervised):

• All weight keys match (strict load)
• Encoder: max diff 0.004
• Predictor target: max diff 0.0002
• Predictor context: max diff 0.002

All diffs are within SDPA floating-point precision.

2.1 checkpoint summary

Fixes https://github.com/huggingface/transformers/issues/45496

Code Agent Policy

• I confirm that this is not a pure code agent PR.

AI assistance (Claude Code) was used for implementation. I reviewed and verified all changes, ran tests, and validated checkpoint conversion end-to-end on both architecture paths (ViT-B and ViT-g).

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a Github issue or the forum? https://github.com/huggingface/transformers/issues/45496
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

Who can review?

@yonigozlan — vision models reviewer and original V-JEPA 2 contributor @molbap — vision models reviewer

Changed files

• docs/source/en/model_doc/vjepa2.md (modified, +16/-1)
• src/transformers/models/vjepa2/configuration_vjepa2.py (modified, +24/-0)
• src/transformers/models/vjepa2/convert_vjepa2_to_hf.py (modified, +143/-7)
• src/transformers/models/vjepa2/modeling_vjepa2.py (modified, +146/-39)
• tests/models/vjepa2/test_modeling_vjepa2.py (modified, +107/-0)