[Low-bit optim] Improve compile time + Fix PyTorch 2.3 support for 4-bit optim

[gau-nernst]

Static-shape compile optim step for single parameter + disable cache size limit.

• For a given model, the number of different argument combinations to single_param_adam() is fixed -> safe to disable cache limit without the risk of always re-compiling.

Benefits

• Improve compile time for 8-bit and FP8 optim (since we don't compile optim step for all parameters at once anymore) -> no noticeable compile time now 🤯
• Improve speed for 4-bit optim (thanks to static shape) -> on par with 8-bit optim now
• Fix PyTorch 2.3 support for 4-bit optim (introduced by Move more utils to TorchAOBaseTensor #784 (comment))
• (Unintended) Fix unusual memory usage of FP8 optim -> same memory footprint as 8-bit optim now

Others

• Re-enable FSDP test for low-bit optim with torch nightly FSDP 2 low bit optim broken on pytorch nightlies #652
• Remove notes on slow compile time

TODO:

• Update benchmarks with ViT-H (benchmarks/benchmark_low_bit_adam.py)
• Update benchmarks with Llama2-7B ([Low-bit optim] Add Llama2-7B finetune benchmarks #746)

Llama2-7B benchmarks

Fine-tune Llama2-7B on Alpaca dataset. PyTorch 2.4, full BF16, 1 epoch, A100, fixed random seed. Benchmark is done with torchtune 52d1b838.

AdamW impl	Peak memory allocated (GB)	toks/s	truthfulqa_mc2 acc
Not fine-tuned	-	-	38.95
PyTorch (fused)	51.6	3200	42.61
bnb 8-bit	39.3	3000	42.75
ao 8-bit	39.1	2900	41.50
ao 4-bit	33.2	2900	42.27

NOTE: lpmm's 4-bit AdamW does not support BF16 weights.

[pytorch-bot]

🔗 Helpful Links

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

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

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