From 51543064532345478b5b9b6b66f7f7d891481e23 Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Thu, 29 Oct 2026 19:59:21 -0700 Subject: [PATCH] [RELEASE ONLY CHANGES] Use a minimum TorchAO requirement for 1.6 --- .ci/scripts/tests/test_cu134_dependencies.py | 23 ++++++++++---------- install_requirements.py | 22 ++++++------------- setup.py | 10 ++------- 3 files changed, 21 insertions(+), 34 deletions(-) diff --git a/.ci/scripts/tests/test_cu134_dependencies.py b/.ci/scripts/tests/test_cu134_dependencies.py index 0dd5aff9f92..3f5da1f882c 100644 --- a/.ci/scripts/tests/test_cu134_dependencies.py +++ b/.ci/scripts/tests/test_cu134_dependencies.py @@ -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 = ( @@ -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) @@ -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) @@ -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) @@ -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) @@ -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" @@ -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( @@ -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__": diff --git a/install_requirements.py b/install_requirements.py index 1d8570f9f46..123f971edb8 100644 --- a/install_requirements.py +++ b/install_requirements.py @@ -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 @@ -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) @@ -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. diff --git a/setup.py b/setup.py index 95af786ad82..076b2324bb3 100644 --- a/setup.py +++ b/setup.py @@ -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. @@ -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]: