Add PyTorch 2.4 tests in CI

[msaroufim]

EDIT: will need to fix this first before merging this #679 (comment)

Before this PR we were testing nightlies which recently moved from PyTorch 2.4 and 2.5 meaning we don't have coverage for 2.4 now

I would have also liked to remove support for PyTorch 2.2 with the policy being "we support latest stable pytorch, the version before that and nightlies"

A few tests needed to be changed

• Autoquant with our nightlies now fail on pytorch 2.4 stable - cc @jerryzh168 @HDCharles
• fp8 tests were moved to nightly only per @vkuzo's request
• Mixed_mm test @msaroufim
• fsdp2 low bit optim tests on 2.4 which means they are completely skipped now @gau-nernst

But will not make this change suddenly, cc some partners in case this is not OK @ebsmothers @joecummings @kimishpatel @Jack-Khuu @awgu @malfet

[pytorch-bot]

🔗 Helpful Links

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

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

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

[msaroufim]

@jerryzh168 there's 12 failing tests for AffineQuantized Tensor with torch 2.4 https://github.com/pytorch/ao/actions/runs/10335226306/job/28609531090?pr=654

This is concerning cause it means our latest nightly release does not work with latest pytorch stable release. I'm worried we have a similar with our official 0.4 release not working with PyTorch 2.4.

If so We might be forced to do a patch release. Can you please sync with @jcaip and figure this out

[vkuzo]

is there a way to move the float8 tests to PT 2.5 in this PR? It would be nice to just support the latest version for now to move fast. It would be just switching

ao/test/float8/test_base.py

Line 21 in 261d0a4

if not TORCH_VERSION_AFTER_2_4:

and comparable lines in test/float8/* to target 2.5 instead.

[msaroufim]

is there a way to move the float8 tests to PT 2.5 in this PR? It would be nice to just support the latest version for now to move fast. It would be just switching

Yeah will land this first since our version guard logic is incorrect #485

[jerryzh168]

@jerryzh168 there's 12 failing tests for AffineQuantized Tensor with torch 2.4 pytorch/ao/actions/runs/10335226306/job/28609531090?pr=654

This is concerning cause it means our latest nightly release does not work with latest pytorch stable release. I'm worried we have a similar with our official 0.4 release not working with PyTorch 2.4.

If so We might be forced to do a patch release. Can you please sync with @jcaip and figure this out

I can repro, opened a PR for export team to take a look here: pytorch/pytorch#133262

[jcaip]

cc @msaroufim @jerryzh168

Testing locally with torch==2.4.0 and torchao==0.4.0 I don't see any failures for test_integration.py (The test that's failing in CI). Is there an easy way to test in CI to make sure?

I'm ~70% sure it's this line that's causing the bug, since it's the only use of get_plain after branch cut.

[msaroufim]

wait huh this is not the same result Jerry is getting, lemme first double check the version guard code, I should have time Wednesday to finally finish this

[jerryzh168]

Testing locally with torch==2.4.0 and torchao==0.4.0 I don't see any failures for test_integration.py (The test that's failing in CI). Is there an easy way to test in CI to make sure?

oh I'm testing with torchao nightly, let me test with 0.4.0 as well

Update: yeah looks like 0.4.0 works, this error only happens in torchao nightly

I think it's related to #609 because previously autoquant does not have multiple levels of tensor subclass

[jerryzh168]

Testing locally with torch==2.4.0 and torchao==0.4.0 I don't see any failures for test_integration.py (The test that's failing in CI). Is there an easy way to test in CI to make sure?

oh I'm testing with torchao nightly, let me test with 0.4.0 as well

Update: yeah looks like 0.4.0 works, this error only happens in torchao nightly

I think it's related to #609 because previously autoquant does not have multiple levels of tensor subclass

we are waiting for pytorch/pytorch#133262 to be resolved, in the meanwhile I feel it's fine to only run these failed autoquant tests on 2.5+ if it works there, but I remember it failed as well (torch nightly + torchao nightly)