[RFC] Adding overrides for max cache seq length

[SalmanMohammadi]

Context

What is the purpose of this PR? Is it to

• add a new feature
• fix a bug
• update tests and/or documentation
• other (please add here)

#1364

Changelog

This PR:

• Adds support for overriding the maximum sequence length used when setting up KV-caches.
• Adds support for correctly setting up caches for self-attention, cross-attention, and fusion layers by exposing encoder and decoder max_seq_len args. These arguments are exposed to the top-level transformer class (i.e. TransformerDecoder, DeepFusionModel, so that the API for all models remains the same model.setup_caches(bsz, dtype, encoder_max_seq_len, decoder_max_seq_len)
• We also remove the use of input_pos to update and retrieve from the KV-cache. Instead, the KV cache tracks its own position.

Test plan

Please make sure to do each of the following if applicable to your PR. (If you're not sure 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.
Example of docstring:

torchtune/torchtune/modules/vision_transformer.py

Line 285 in 6a7951f

Examples:

Example in our docs: https://pytorch.org/torchtune/main/tutorials/qat_finetune.html#applying-qat-to-llama3-models

• 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/1449

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 3fc1135 with merge base bc2c013 ():
💚 Looks good so far! There are no failures yet. 💚

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

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 81.40704% with 37 lines in your changes missing coverage. Please review.

Project coverage is 72.90%. Comparing base (726abb0) to head (3fc1135).
Report is 7 commits behind head on main.

Files with missing lines	Patch %	Lines
torchtune/modules/transformer.py	55.88%	30 Missing ⚠️
torchtune/models/gemma/transformer.py	0.00%	6 Missing ⚠️
recipes/eleuther_eval.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1449      +/-   ##
==========================================
+ Coverage   70.72%   72.90%   +2.18%     
==========================================
  Files         288      289       +1     
  Lines       14213    14336     +123     
==========================================
+ Hits        10052    10452     +400     
+ Misses       4161     3884     -277     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[SalmanMohammadi]

What's up with the eleuther eval tests? They're passing locally for me.

[ebsmothers]

Overall it looks reasonable to me, will leave it to @pbontrager or @joecummings for the final sign-off

[SalmanMohammadi]

Thanks so much for the reviews @ebsmothers. Will address once #1424 lands and things are merged.

[joecummings]

No concerns, but pls address @ebsmothers comments.