Skip to content

[RELEASE ONLY CHANGES] Use a minimum TorchAO requirement for 1.6 #24318

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
JacobSzwejbka merged 1 commit into from
Oct 30, 2026
Merged

[RELEASE ONLY CHANGES] Use a minimum TorchAO requirement for 1.6 #24318

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 12 additions & 11 deletions .ci/scripts/tests/test_cu134_dependencies.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ def test_all_install_steps_preserve_exact_cu134_selection(self):
"torch==2.15.0+cu134",
"torchvision==0.29.2+cu134",
"torchaudio==2.11.1+cu134",
f"torchao=={self.installer.CU134_TORCHAO_NIGHTLY_VERSION}+{ao_variant}",
f"torchao>={self.installer.CU134_TORCHAO_NIGHTLY_VERSION}",
}
for index, command in enumerate(commands):
required = (
Expand All @@ -78,7 +78,7 @@ def test_all_install_steps_preserve_exact_cu134_selection(self):
else {
requirement
for requirement in expected
if requirement.startswith(("torch==", "torchao=="))
if requirement.startswith(("torch==", "torchao>="))
}
)
self.assertTrue(required.issubset(command), command)
Expand Down Expand Up @@ -108,7 +108,7 @@ def test_other_cuda_trains_keep_existing_pins(self):
)
self.assertIn("torch==2.15.0", core)
self.assertIn(
f"torchao=={self.installer.TORCHAO_NIGHTLY_VERSION}", core
f"torchao>={self.installer.TORCHAO_NIGHTLY_VERSION}", core
)
self.assertIn("torchvision==0.29.2", domains)
self.assertIn("torchaudio==2.11.1", domains)
Expand All @@ -127,7 +127,7 @@ def test_source_pinned_torch_is_not_replaced(self):
def test_no_cuda_keeps_default_pins(self):
core, _, domains, _ = self.install_commands(None)
self.assertIn("torch==2.15.0", core)
self.assertIn(f"torchao=={self.installer.TORCHAO_NIGHTLY_VERSION}", core)
self.assertIn(f"torchao>={self.installer.TORCHAO_NIGHTLY_VERSION}", core)
self.assertIn("torchvision==0.29.2", domains)
self.assertIn("https://download.pytorch.org/whl/cpu", core)

Expand Down Expand Up @@ -263,7 +263,7 @@ def test_package_install_preserves_source_pinned_torchao(self):
self.assertIn(".", commands[-1])
core = commands[0]
torchao = Requirement(
next(arg for arg in core if arg.startswith("torchao=="))
next(arg for arg in core if arg.startswith("torchao>="))
)
version = next(iter(torchao.specifier)).version
self.assertIn(version, metadata.specifier)
Expand Down Expand Up @@ -300,9 +300,7 @@ def test_cu134_keeps_explicit_torchao_source_build(self):
self.assertIn("third-party/ao", commands[1])
self.assertIn(".", commands[2])
for command in commands:
self.assertFalse(
any(arg.startswith("torchao==") for arg in command)
)
self.assertFalse(any(arg.startswith("torchao") for arg in command))
self.assertIn("torch==2.15.0+cu134", commands[-1])
source_version = self.installer.TORCHAO_NIGHTLY_VERSION.partition(
".dev"
Expand All @@ -318,15 +316,14 @@ def test_cu134_keeps_explicit_torchao_source_build(self):
f"{source_version}+git{source_commit[:7]}", metadata.specifier
)

def test_wheel_torchao_bound_matches_selected_train(self):
def test_wheel_torchao_minimum_matches_installer(self):
for cuda in ((13, 4), (13, 2), None):
version = (
self.installer.CU134_TORCHAO_NIGHTLY_VERSION
if cuda == (13, 4)
else self.installer.TORCHAO_NIGHTLY_VERSION
)
major, minor = (int(part) for part in version.split(".")[:2])
expected = f"torchao>={version},<{major}.{minor + 1}"
expected = f"torchao>={version}"
self.utils.determine_torch_url.cache_clear()
with (
patch.object(
Expand All @@ -338,6 +335,10 @@ def test_wheel_torchao_bound_matches_selected_train(self):
patch.object(self.installer.platform, "system", return_value="Linux"),
):
self.assertEqual(self.torchao_requirement(), expected)
requirement = Requirement(self.torchao_requirement())
self.assertNotIn("0.18.0", requirement.specifier)
for newer in ("0.19.0.dev20260908", "0.19.0", "0.20.0", "1.0.0"):
self.assertIn(newer, requirement.specifier)


if __name__ == "__main__":
Expand Down
22 changes: 7 additions & 15 deletions install_requirements.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,10 +43,7 @@ def cu134_requirements(torch_url, include_domains=False):
CU134_TORCH_PACKAGES if include_domains else CU134_TORCH_PACKAGES[:1]
)
if not torchao_from_source():
torchao_variant = (
"cpu" if platform.machine().lower() in ("aarch64", "arm64") else "cu134"
)
packages.append(f"torchao=={CU134_TORCHAO_NIGHTLY_VERSION}+{torchao_variant}")
packages.append(f"torchao>={CU134_TORCHAO_NIGHTLY_VERSION}")
return packages


Expand All @@ -71,17 +68,12 @@ def install_requirements(use_pytorch_nightly):
cu134_packages = cu134_requirements(torch_url)
if cu134_packages and not use_pytorch_nightly:
cu134_packages[0] = "torch"
# torchao's CUDA channel publishes x86_64 only, so asking for a CUDA build makes the pin
# unsatisfiable on aarch64. Only that case is special-cased: falling back everywhere would
# change which torchao a CPU x86_64 install resolves, and the CUDA build is genuinely wanted
# where it exists. This nightly is what a development checkout is tested against, and the
# wheel's own torchao lower bound is this same version so that installing the package
# afterwards leaves this pin in place rather than replacing it.
# TorchAO's CUDA channel publishes x86_64 only, so use the CPU channel on
# aarch64. The installer and wheel metadata share a minimum version, allowing
# newer TorchAO packages without downgrading a preinstalled compatible build.
if platform.machine().lower() in ("aarch64", "arm64"):
# The cpu channel specifically, not the index root. The root carries every variant, and a
# pin without a local segment admits all of them while ordering a local segment highest,
# so the xpu channel's pure python wheel would win on version before pip compares wheel
# tags, silently replacing the compiled aarch64 build.
# The index root carries every variant; a local segment sorts above the
# base version, so an XPU wheel could replace the compiled aarch64 build.
torchao_url = f"{TORCHAO_URL_BASE}/cpu"
else:
torchao_url = determine_torch_url(TORCHAO_URL_BASE)
Expand All @@ -92,7 +84,7 @@ def install_requirements(use_pytorch_nightly):
# that we don't need to set any version number there because they have already
# been installed on CI before this step, so pip won't reinstall them
("torch==2.15.0" if use_pytorch_nightly else "torch"),
f"torchao=={TORCHAO_NIGHTLY_VERSION}",
f"torchao>={TORCHAO_NIGHTLY_VERSION}",
]

# Install the requirements for core ExecuTorch package.
Expand Down
10 changes: 2 additions & 8 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1116,12 +1116,7 @@ def _package_relative_depth(library: Path) -> int:


def _torchao_requirement() -> str:
"""The torchao dependency, pinned to the series install_requirements.py installs.

Derived from that module rather than written out, so a nightly bump cannot move the
pin without moving this bound with it. A bump into the next series would otherwise
silently stop satisfying the lower bound, and installing this package over a
development checkout would replace the torchao that was just installed.
"""The TorchAO minimum version, shared with install_requirements.py.

Loaded by path, the way install_utils is above, because setuptools executes this
file without the project directory on sys.path, so a plain import does not resolve.
Expand All @@ -1142,8 +1137,7 @@ def _torchao_requirement() -> str:
and not module.torchao_from_source()
):
version = module.CU134_TORCHAO_NIGHTLY_VERSION
major, minor = (int(part) for part in version.split(".")[:2])
return f"torchao>={version},<{major}.{minor + 1}"
return f"torchao>={version}"


def _base_dependencies() -> List[str]:
Expand Down
Loading