Fix: Ensure ViT includes the final hidden state in its output

[tkdcjf159]

Context

What is the purpose of this PR? Is it to

• fix a bug

When using the out_indices parameter to select specific hidden states from the Vision Transformer (ViT) based on their layer index, the implementation incorrectly fails to return the hidden state from the final transformer layer. This occurs even if the index corresponding to the final layer is explicitly included in the out_indices list.

Changelog

What are the changes made in this PR?
*

Test plan

Please make sure to do each of the following if applicable to your PR. If you're unsure about any one of these just ask and we will happily help. We also have a contributing page for some guidance on contributing.

• run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
• add unit tests for any new functionality
• update docstrings for any new or updated methods or classes
• run unit tests via pytest tests
• run recipe tests via pytest tests -m integration_test
• manually run any new or modified recipes with sufficient proof of correctness
• include relevant commands and any other artifacts in this summary (pastes of loss curves, eval results, etc.)

UX

If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Here is a docstring example
and a tutorial example

• I did not change any public API
• I have added an example to docs or docstrings

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2610

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [Infra] Jobs got frequently cancelled, sometimes mid-checkout

❌ 2 New Failures, 4 Cancelled Jobs

As of commit 3117b82 with merge base 1075d9c ():

NEW FAILURES - The following jobs have failed:

• GPU tests / gpu_test (3.11, stable) (gh)
  tests/torchtune/modules/test_vision_transformer.py::TestVisionTransformer::test_vision_transformer_return_hidden_layers
• Unit Test / unit_tests (3.11) (gh)
  tests/torchtune/modules/test_vision_transformer.py::TestVisionTransformer::test_vision_transformer_return_hidden_layers

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• GPU tests / gpu_test (3.10, stable) (gh)
  tests/torchtune/modules/test_vision_transformer.py::TestVisionTransformer::test_vision_transformer_return_hidden_layers
• GPU tests / gpu_test (3.9, stable) (gh)
  tests/torchtune/modules/test_vision_transformer.py::TestVisionTransformer::test_vision_transformer_return_hidden_layers
• Unit Test / unit_tests (3.10) (gh)
  ##[error]The operation was canceled.
• Unit Test / unit_tests (3.9) (gh)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[facebook-github-bot]

Hi @tkdcjf159!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at [email protected]. Thanks!

[facebook-github-bot]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[pbontrager]

Thanks for adding this. It was assumed that you'd get the last hidden layer from the x output, but it makes sense that you might want it before the layer norm and positional embeddings. Please put the enumerate back to keep things a bit more compact but otherwise it's good to go.

[pbontrager]

You can still keep the enumerate, the layer_idx variable is still accessible after the for loop

[tkdcjf159]

@pbontrager Fix it!
Initially, I misunderstood that the range of out_indices was 1 ~ num_layer instead of 0 ~ num_layer-1, so I made modifications based on that misconception. If the range of out_indices is 0~num_layer-1, then after the for loop, I should just check if layer_idx is included in self.out_indices and add it if it is. I've made the corrections according to your advice.

[pbontrager]

You need to fix the test and this should be good to land.