Add V-JEPA 2.1 inference support

PULL REQUEST #45497: Add V-JEPA 2.1 inference support
State: open  |  Merged: False
Author: davevanveen  |  Base: main
Labels:
Created: 2026-04-17T22:05:24Z

# 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

| Model | Params | Distilled | `n_output_distillation` | `teacher_embed_dim` |
|-------|--------|-----------|------------------------|---------------------|
| ViT-B/16, 384 | 80M | Yes (ViT-G) | 1 | 1664 |
| ViT-L/16, 384 | 300M | Yes (ViT-G) | 1 | 1664 |
| ViT-g/16, 384 | 1B | No | 4 | — |
| ViT-G/16, 384 | 2B | No | 4 | — |

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

## Code Agent Policy

- [x] 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).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? https://github.com/huggingface/transformers/issues/45496
- [x] Did you make sure to update the documentation with your changes?
- [x] Did you write any new necessary tests?

## Who can review?

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

--- Comment by yonigozlan at 2026-05-01T14:45:59Z ---
Hello @davevanveen ! Thanks a lot for opening this PR. Sorry for the delay in reviewing this, we have been very busy with lots of new models to review and recent core changes. I'll try to get to this PR next week!

--- Comment by davevanveen at 2026-05-01T23:05:43Z ---
Appreciate the note, @yonigozlan! No rush from me; I imagine y'all are pretty slammed. Thanks for all the work you do. Cheers.

--- Comment by davevanveen at 2026-05-16T19:51:00Z ---
Thank you for the feedback, @yonigozlan! I appreciate your attention to detail. Happy to iterate further.

Pushed a refactor addressing the API simplification comments. Included inline replies above for changes worth flagging. Summary below:

- replaced `img_temporal_dim_size` with `use_image_patch_embedder` (boolean) and simplified the single-frame check
- added `tubelet_size` arg on `VJEPA2PatchEmbeddings3D` instead of the config-copy workaround
- replaced predictor-specific `is_predictor` with an explicit `interpolate_rope` flag
- renamed `n_output_distillation` to `num_distillation_outputs` (default 1) and `return_all_tokens` to `use_context_projection`
- added `VJEPA2EncoderOutput.hierarchical_hidden_state` and dropped the `return_hierarchical` argument
- added `VJEPA2PredictorOutput.context_predictions` and stopped overloading `last_hidden_state`
- extended `VJEPA2WithMaskedInputPredictorOutput` and `VJEPA2WithMaskedInputModelOutput` with the new fields, top-level forward chooses predictor input from hierarchical state when present
- simplified distillation naming (`distillation_layers` / `distillation_norms`) and projection-dim logic
- removed the classification-head guard
- updated converter validation to use the new public output fields

Validated the converter on `vit_base_2_1_384` and `vit_giant_2_1_384` — `allclose` passes against Meta's checkpoints for both encoder and predictor.

--- Comment by github-actions[bot] at 2026-05-21T21:12:55Z ---
**[For maintainers]** Suggested jobs to run (before merge)

run-slow: vjepa2

--- Comment by yonigozlan at 2026-05-08T20:54:03Z ---
Let's add a `tubelet_size` arg to VJEPA2PatchEmbeddings3D instead of this workaround, and take the config value when the arg is not set.

--- Comment by yonigozlan at 2026-05-08T20:56:44Z ---
Is `config.img_temporal_dim_size` ever something other than 1? otherwise let's change `img_temporal_dim_size` to a boolean flag in the config like `use_image_patch_embedder` and simplify this logic

--- Comment by yonigozlan at 2026-05-08T21:02:06Z ---
Looks like `is_predictor` is only used to interpolate rope or not, so let's use an `interpolate_rope` flag instead (which would override the config one if specified)

--- Comment by yonigozlan at 2026-05-08T21:04:44Z ---
We would only need this then:
```suggestion
        if self.interpolate_rope and self.grid_size > 1:
```

--- Comment by yonigozlan at 2026-05-08T21:11:07Z ---
Let's call this `num_distillation_outputs`, and set it to 1 by default, so that we don't need the "> 0" check in the modeling code (no need to have a redundant condition when hierarchical_layers is set)

--- Comment by yonigozlan at 2026-05-08T21:14:13Z ---
We can slightly simplify and use more explicit naming:

```suggestion
        if config.hierarchical_layers is not None:
            self.distillation_layers = config.hierarchical_layers[-config.num_distillation_outputs :]
            self.distillation_norms = nn.ModuleList(
                [nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) for _ in self.distillation_layers]
            )
            self.layernorm = None
        else:
            self.distillation_norms = None
            self.distillation_layers = None
            self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
```

--- Comment by yonigozlan at 2026-05-08T21:20:02Z ---
No need for checking both norms_block and _extraction_layers as they are defined together, + no need for nested if:

```suggestion
        for i, layer_module in enumerate(self.layer):
            hidden_states = layer_module(hidden_states, None, **kwargs)[0]

            if self.distillation_norms is not None and i in self.distillation_layers:
                norm_idx = self.distillation_layers.index(i)
                hierarchical_outputs.append(self.distillation_norms[norm_idx](hidden_states))
```

--- Comment by yonigozlan at 2026-05-08T21:29:20Z ---
This is getting complex, I think we could avoid a lot of added logic by making custom model outputs for VJEPA2Encoder (VJEPA2EncoderOutput) and VJEPA2Predictor.

Here, `class VJEPA2EncoderOutput(BaseModelOutput)` could add a `hierarchical_hidden_state` field that is populated when `hierarchical_outputs` is not empty, otherwise set to None, so that we avoid the `return_hierarchical` altogether

--- Comment by yonigozlan at 2026-05-08T21:30:21Z ---
```suggestion
        if config.num_distillation_outputs> 1:
            self.predictor_embeddings = nn.Sequential(
                nn.Linear(config.hidden_size * config.num_distillation_outputs, config.hidden_size, bias=True),
```

--- Comment by yonigozlan at 2026-05-08T21:32:37Z ---
Seems like this can be simplified to this:

```suggestion
        proj_output_dim = (
            config.teacher_embed_dim
            if config.teacher_embed_dim is not None
            else config.num_distillation_outputs * config.hidden_size
        )
```

--- Comment by yonigozlan at 2026-05-08T21:42:49Z ---
Same as VJEPA2EncoderOutput, we could have a custom `class VJEPA2PredictorOutput(BaseModelOutput)` adding a `context_predictions` field, so that we don't have `last_hidden_state` representing two different things depending on `return_all_tokens`, and we can simplify the logic to something like:

```suggestion
        target_predictions = self.proj(hidden_states[:, N_ctxt:])

        context_predictions = (
            self.proj_context(hidden_states[:, :N_ctxt]) if self.proj_context is not None else None
        )

        return VJEPA2PredictorOutput(
            last_hidden_state=target_predictions,
            context_predictions=context_predictions,
```

--- Comment by yonigozlan at 2026-05-08T21:44:17Z ---
We can then simplify to this:

```suggestion
        is_image = self.config.use_image_patch_embedder and pixel_values_videos.shape[1] == 1

        encoder_outputs: VJEPA2EncoderOutput = self.encoder(
            pixel_values_videos=pixel_values_videos,
            **kwargs,
        )
```

--- Comment by yonigozlan at 2026-05-08T21:59:28Z ---
We'll also need to get the correct predictor input here (hierarchical_hidden_state if not None, otherwise sequence_output)

--- Comment by yonigozlan at 2026-05-09T18:29:40Z ---
This is not the ideal API imo, We could instead change `VJEPA2WithMaskedInputPredictorOutput` in the modeling code to add a `context_predictions` field so that, we'll have:

```suggestion
            assert torch.allclose(predictor_outputs.last_hidden_state, og_target, atol=1e-2)
            assert torch.allclose(predictor_outputs.context_predictions, og_context, atol=1e-2)
```

--- Comment by yonigozlan at 2026-05-11T14:18:55Z ---
Same here:

```suggestion
            assert torch.allclose(predictor_outputs.last_hidden_state, og_target, atol=1e-2)
            assert torch.allclose(predictor_outputs.context_predictions, og_context, atol=1e-2)
```

--- Comment by yonigozlan at 2026-05-11T14:24:48Z ---
We'll also need to add a `hierarchical_hidden_state` field here as well since it wouldn't be in `last_hidden_state` anymore

--- Comment by yonigozlan at 2026-05-11T14:26:05Z ---
Same here (hierarchical_hidden_state if not None, otherwise sequence_output)

--- Comment by yonigozlan at 2026-05-11T14:26:52Z ---
Why have this guard here? I think we can get rid of it fully and let users experiment

--- Comment by davevanveen at 2026-05-16T19:51:43Z ---
Confirmed it's always 1. Renamed to `use_image_patch_embedder: bool` (default `False`) and simplified to `is_image = self.config.use_image_patch_embedder and num_frames == 1`.

--- Comment by davevanveen at 2026-05-16T19:53:10Z ---
Done. Renamed to `num_distillation_outputs` with default `1`. Also took the opportunity to rename `return_all_tokens` to `use_context_projection`, since after the predictor refactor the old name described deleted behavior (the concatenation into `last_hidden_state`). Happy to revert if you'd rather keep `return_all_tokens`.

--- Comment by davevanveen at 2026-05-16T19:54:32Z ---
Done. Added `VJEPA2EncoderOutput(BaseModelOutput)` with a `hierarchical_hidden_state` field, populated only when `config.hierarchical_layers` is set. Encoder forward no longer takes `return_hierarchical`. `last_hidden_state` is now consistently the post-loop hidden state across both V-JEPA 2 and 2.1 paths.

--- Comment by davevanveen at 2026-05-16T19:56:38Z ---
Done. Applied the simpler form.

Flagging one thing: this assumes `num_distillation_outputs == len(hierarchical_layers)`, which holds for the released V-JEPA 2.1 checkpoints but isn't enforced by the config. If a future config sets `num_distillation_outputs < len(hierarchical_layers)` (i.e. more taps instantiated than emitted), the projection size will silently differ from the previous `len(hierarchical_layers) * out_embed_dim` expression. Validated against `vit_base_2_1_384` and `vit_giant_2_1_384` (the only checkpoint where this branch is actually exercised, since base/large set `teacher_embed_dim`).

Happy to add an assertion in `VJEPA2Config.__post_init__` or leave the invariant implicit. Your call.

--- Comment by davevanveen at 2026-05-16T19:57:30Z ---
Done. Added `VJEPA2PredictorOutput(BaseModelOutput)` with `context_predictions`. `last_hidden_state` now consistently represents target predictions; context goes to `context_predictions` when `use_context_projection=True`, else `None`. The concat in the forward is gone.

--- Comment by davevanveen at 2026-05-16T19:58:39Z ---
Done at both sites. Validated end-to-end against Meta's `vit_base_2_1_384` and `vit_giant_2_1_384` checkpoints. `allclose` passes for full-mask and partial-mask cases.

--- Comment by davevanveen at 2026-05-16T19:59:51Z ---
Done. Added `hierarchical_hidden_state` to `VJEPA2WithMaskedInputPredictorOutput`. Also added it to `VJEPA2WithMaskedInputModelOutput` so users calling `VJEPA2Model(..., skip_predictor=True)` for feature extraction can still reach hierarchical features when `predictor_output` is `None`.

--- Comment by davevanveen at 2026-05-16T20:02:07Z ---
Guard removed.

Added a small smoke test (`test_classification_head_with_hierarchical_distillation`) that instantiates `VJEPA2ForVideoClassification` with a 2.1-style config (`num_distillation_outputs=4`, `hierarchical_layers` set) and runs a forward pass. The pooler consumes `last_hidden_state` (post-loop, dim `hidden_size`), not `hierarchical_hidden_state`, so shapes work out for any `num_distillation_outputs`.

Documented in the classification head's docstring that for V-JEPA 2.1 the pooler sees unnormalized features. Users wanting hierarchical features should use `VJEPA2Model` for feature extraction.


PULL REQUEST #45495: revert sha commit pointing to main for transformers_amd_ci_  workflows
State: closed  |  Merged: True
Author: paulinebm  |  Base: main
Labels:
Created: 2026-04-17T14:18:37Z

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes # (issue)

## Code Agent Policy

The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. We are currently bottlenecked by our ability to review and respond to them. As a result,
**we ask that new users do not submit pure code agent PRs** at this time.
You may use code agents in drafting or to help you diagnose issues. We'd also ask autonomous "OpenClaw"-like agents
not to open any PRs or issues for the moment.

PRs that appear to be fully agent-written will probably be closed without review, and we may block users who do this
repeatedly or maliciously.

This is a rapidly-evolving situation that's causing significant shockwaves in the open-source community. As a result,
this policy is likely to be updated regularly in the near future. For more information, please read [`CONTRIBUTING.md`](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md).

- [ ] I confirm that this is not a pure code agent PR.

## 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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?

## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc
- kernels: @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->

--- Comment by HuggingFaceDocBuilderDev at 2026-04-17T14:27:26Z ---
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_45495). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

--- Comment by Copilot at 2026-04-17T14:22:17Z ---
Using a mutable ref (`@main`) for a reusable workflow means this scheduled CI can change behavior without any change in this repo, which hurts reproducibility and can unexpectedly break scheduled runs. Since this workflow also inherits secrets on self-hosted runners, pinning to a specific commit SHA (or a version tag) is safer and makes changes auditable.

--- Comment by Copilot at 2026-04-17T14:22:17Z ---
Using a mutable ref (`@main`) for a reusable workflow means this scheduled CI can change behavior without any change in this repo, which hurts reproducibility and can unexpectedly break scheduled runs. Since this workflow also inherits secrets on self-hosted runners, pinning to a specific commit SHA (or a version tag) is safer and makes changes auditable.

--- Comment by Copilot at 2026-04-17T14:22:18Z ---
Using a mutable ref (`@main`) for a reusable workflow means this scheduled CI can change behavior without any change in this repo, which hurts reproducibility and can unexpectedly break scheduled runs. Since this workflow also inherits secrets on self-hosted runners, pinning to a specific commit SHA (or a version tag) is safer and makes changes auditable.

--- Comment by Copilot at 2026-04-17T14:22:18Z ---
Using a mutable ref (`@main`) for a reusable workflow means this scheduled CI can change behavior without any change in this repo, which hurts reproducibility and can unexpectedly break scheduled runs. Since this workflow also inherits secrets on self-hosted runners, pinning to a specific commit SHA (or a version tag) is safer and makes changes auditable.

--- Comment by Copilot at 2026-04-17T14:22:18Z ---
Using a mutable ref (`@main`) for a reusable workflow means this scheduled CI can change behavior without any change in this repo, which hurts reproducibility and can unexpectedly break scheduled runs. Since this workflow also inherits secrets on self-hosted runners, pinning to a specific commit SHA (or a version tag) is safer and makes changes auditable.

--- Comment by Copilot at 2026-04-17T14:22:19Z ---
Using a mutable ref (`@main`) for a reusable workflow means this scheduled CI can change behavior without any change in this repo, which hurts reproducibility and can unexpectedly break scheduled runs. Since this workflow also inherits secrets on self-hosted runners, pinning to a specific commit SHA (or a version tag) is safer and makes changes auditable.

--- Comment by Copilot at 2026-04-17T14:22:19Z ---
Using a mutable ref (`@main`) for a reusable workflow means this scheduled CI can change behavior without any change in this repo, which hurts reproducibility and can unexpectedly break scheduled runs. Since this workflow also inherits secrets on self-hosted runners, pinning to a specific commit SHA (or a version tag) is safer and makes changes auditable.

--- Comment by Copilot at 2026-04-17T14:22:19Z ---
Using a mutable ref (`@main`) for a reusable workflow means this scheduled CI can change behavior without any change in this repo, which hurts reproducibility and can unexpectedly break scheduled runs. Since this workflow also inherits secrets on self-hosted runners, pinning to a specific commit SHA (or a version tag) is safer and makes changes auditable.