[Llama] Make torchao's Llama trainable

[gau-nernst]

Fixes #674

To make Llama trainable, I changed the following:

• Do not initialize KV cache if training=True is passed to setup_caches()
• When input_pos is not passed to the model, handle it accordingly and use F.sdpa(is_causal=True)

Other minor changes:

• Ignore .safetensors weights in scripts/download.py to save bandwidth/speed up download time.
• Use torchao's Llama as an example for benchmarks/quantized_training

I manually ran

python torchao/_models/llama/generate.py --checkpoint_path checkpoints/meta-llama/Llama-2-7b-chat-hf/model.pth --prompt "Hello, my name is" --compile --precision float16

to confirm that the generated outputs are identical.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/728

• 📄 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 806721c with merge base 99644e9 ():
💚 Looks good so far! There are no failures yet. 💚

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