[low-bit optim] Fix Adam4bit support on PyTorch 2.3 and 2.4. Update AdamFp8 torch requirement

[gau-nernst]

See #744 (comment)

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures

As of commit d83a1c1 with merge base ba2d3b1 ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/whl/nightl... / linux-job (gh)
  test/dtypes/test_fpx.py::TestFpxTensorCoreAQTLayout::test_to_scaled_tc_fpx_compile_ebits_3_mbits_2_device_cpu
• Run Regression Tests / test (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://download.pytorc... / linux-job (gh)
  test/dtypes/test_fpx.py::TestFpxTensorCoreAQTLayout::test_to_scaled_tc_fpx_compile_ebits_3_mbits_2_device_cpu

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

[gau-nernst]

convert this to draft since I'm also investigating torch version support of FP8 optim. FP8 optim has never run in CI due to sm89 constraint.

[gau-nernst]

Fixed issue with 4-bit Adam. Now 4-bit Adam works with PyTorch 2.3 likes in the past. Hopefully CI is green. The issue seems to be related to this pytorch/pytorch#128649

I kinda feel conflicted about this change, since now the optimizer state is flattened, instead of having the same shape as param. Will try a better solution in the future. I think it has to do with dynamic compile also. 4-bit optim is giving us a lot of headaches 🤣.

stale