generalize torch compatibility check

[liangel-02]

Summary

per #2919, torchao .so files built using PyTorch 2.8.0 are not ABI compatible with PyTorch 2.9+. this PR modifies the check added in #2908 for future versions of torchao. we keep the checks for current stable versions for clarity and add a check for most recent versions (torch 2.9.0.dev and torchao 0.14.0dev). this check should be updated with every release to correctly reflect compatibility.

Testing

bash script ./scripts/test_torch_version_torchao_version_compatibility.sh. this script was ran locally to test and should be updated and run before every release/version bump.

[pytorch-bot]

🔗 Helpful Links

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

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

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

[andrewor14]

Looks good to me, thanks! Just left a few minor questions

[jerryzh168]

why is this only torch 2.9.0 nightly and torchao == 0.14.0 instead of something like torch >= 2.8 and torchao >= 0.13.0?

[liangel-02]

we want to do stricter checks like == instead of >=
if we did torch >= x and torchao >= y we can have false positives if torchao version is y+2 but torch is x+1 when y+2 is built against torch x+2 (so not guaranteed compatible with torch x+1)

we also only need to check the latest torchao version against the latest torch version because this code will never run for a previous torchao version (it is redundant to consider the case when torchao == 0.13.0 because if that were the case, this code would not exist)

[jerryzh168]

if we did torch >= x and torchao >= y we can have false positives if torchao version is y+2 but torch is x+1 when y+2 is built against torch x+2 (so not guaranteed compatible with torch x+1)

I think changing to >= for both would skip all future inclusions of this? e.g. torchao >= 0.14.0.dev and torch >= 2.9.0.dev (this will be satisfied for all current and future torch and torchao versions)

explanation for 0.13.0 and below makes sense

[jerryzh168]

also to clarify, current version for torchao should be 0.14.0.dev I think, not sure if you intended to check 0.14.0 or not

[vkuzo]

I think you should spell out all the versions here:

if torch_version == X and torchao_version == Y:
  ...
elif torch_version == X+1 and torchao_version == Y+1:
  ...
else:
  ...

this removes all confusion about what is supported or not. It's ok that if torch_version == X and torchao_version == Y will become dead code, IMO more important to make it crystal clear what's going on. IMO having the entire support matrix here is the best way to do that.

If you really don't want dead code, I'd still have a comment very clearly explaining why only the latest thing is here and what the behavior was before.

[jerryzh168]

should we test more combinations? e.g.

torchao = (0.13, 0.14.dev)
pytorch = (2.8, 2.9.dev, 2.9 (future))

[liangel-02]

adding more test, but don't need to test previous versions i think (see #3042 (comment), #3042 (comment))

[vkuzo]

comments inline

[vkuzo]

thank you for improving this!