Gemma generation hotfix

[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)

Closes #1446

Changelog

Patched in caches_are_enabled so generation isn't broken. Let's address this properly with #1454.

Test plan

(venv) combuter:torchtune salmanmohammadi$ tune run generate --config target/generation_gemma.yaml 
W0830 12:57:49.903000 8424225472 torch/distributed/elastic/multiprocessing/redirects.py:29] NOTE: Redirects are currently not supported in Windows or MacOs.
INFO:torchtune.utils.logging:Running InferenceRecipe with resolved config:

chat_format: null
checkpointer:
  _component_: torchtune.utils.FullModelHFCheckpointer
  checkpoint_dir: ./target/gemma-2b/
  checkpoint_files:
  - model-00001-of-00002.safetensors
  - model-00002-of-00002.safetensors
  model_type: GEMMA
  output_dir: /tmp/gemma
  recipe_checkpoint: null
device: cuda
dtype: bf16
enable_kv_cache: false
instruct_template: null
max_new_tokens: 10
model:
  _component_: torchtune.models.gemma.gemma_2b
prompt: Tell me a joke?
quantizer: null
seed: 1234
temperature: 0.6
tokenizer:
  _component_: torchtune.models.gemma.gemma_tokenizer
  path: ./target/gemma-2b/tokenizer.model
top_k: 300

DEBUG:torchtune.utils.logging:Setting manual seed to local seed 1234. Local seed is seed + rank = 1234 + 0
INFO:torchtune.utils.logging:Model is initialized with precision torch.bfloat16.
INFO:torchtune.utils.logging:Tell me a joke?
My mom is not here to tell the joke

• 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.)

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 c3f5b4a with merge base 0bdd308 ():
💚 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 50.00000% with 1 line in your changes missing coverage. Please review.

Project coverage is 26.93%. Comparing base (0bdd308) to head (c3f5b4a).

Files with missing lines	Patch %	Lines
torchtune/models/gemma/transformer.py	50.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #1462       +/-   ##
===========================================
- Coverage   69.80%   26.93%   -42.87%     
===========================================
  Files         272      272               
  Lines       13053    13055        +2     
===========================================
- Hits         9111     3516     -5595     
- Misses       3942     9539     +5597     

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