[RELEASE ONLY CHANGES] Retain existing TorchAO nightly pins for 1.6

Author: JacobSzwejbka

.ci/scripts/test-rocm-aoti.sh

@@ -10,7 +10,7 @@ set -euo pipefail
 ROCM_VERSION="${ROCM_VERSION:-7.2}"
 ROCM_PATH="${ROCM_PATH:-/opt/rocm}"
 PYTORCH_ROCM_INDEX="${PYTORCH_ROCM_INDEX:-https://download.pytorch.org/whl/rocm${ROCM_VERSION}}"
-TORCHAO_ROCM_WHEEL_BASE="${TORCHAO_ROCM_WHEEL_BASE:-https://download.pytorch.org/whl/rocm${ROCM_VERSION}}"
+TORCHAO_ROCM_WHEEL_BASE="${TORCHAO_ROCM_WHEEL_BASE:-https://download.pytorch.org/whl/nightly/rocm${ROCM_VERSION}}"
 ROCM_CI_TMP_ROOT="${RUNNER_TEMP:-/tmp}"
 mkdir -p "${ROCM_CI_TMP_ROOT}" 2>/dev/null || ROCM_CI_TMP_ROOT=/tmp
 ROCM_CI_TMPDIR="$(mktemp -d "${ROCM_CI_TMP_ROOT}/executorch-rocm-ci.XXXXXX")"

.ci/scripts/test-rocm-voxtral.sh

@@ -12,7 +12,7 @@ ROCM_PATH="${ROCM_PATH:-/opt/rocm}"
 EXPECTED_ROCM_ARCH="${EXPECTED_ROCM_ARCH:-gfx950}"
 EXPECTED_WARP_SIZE="${EXPECTED_WARP_SIZE:-64}"
 PYTORCH_ROCM_INDEX="${PYTORCH_ROCM_INDEX:-https://download.pytorch.org/whl/rocm${ROCM_VERSION}}"
-TORCHAO_ROCM_WHEEL_BASE="${TORCHAO_ROCM_WHEEL_BASE:-https://download.pytorch.org/whl/rocm${ROCM_VERSION}}"
+TORCHAO_ROCM_WHEEL_BASE="${TORCHAO_ROCM_WHEEL_BASE:-https://download.pytorch.org/whl/nightly/rocm${ROCM_VERSION}}"
 VOXTRAL_CI_TMP_ROOT="${RUNNER_TEMP:-/tmp}"
 if ! mkdir -p "${VOXTRAL_CI_TMP_ROOT}" 2>/dev/null ||
   [[ ! -w "${VOXTRAL_CI_TMP_ROOT}" ]]; then

.ci/scripts/tests/test_cu134_dependencies.py

@@ -90,12 +90,14 @@ def test_all_install_steps_preserve_exact_cu134_selection(self):
                             )
                         )
                     self.assertIn("https://download.pytorch.org/whl/cu134", command)
+                    self.assertIn(
+                        f"https://download.pytorch.org/whl/nightly/{ao_variant}",
+                        command,
+                    )
                     self.assertNotIn(
                         "https://download.pytorch.org/whl/test/cu134", command
                     )
                     self.assertNotIn("--no-deps", command)
-                    if machine == "aarch64":
-                        self.assertIn("https://download.pytorch.org/whl/cpu", command)
 
     def test_other_cuda_trains_keep_existing_pins(self):
         for cuda in ((12, 6), (13, 0), (13, 2)):

install_requirements.py

@@ -16,10 +16,12 @@
 # The pip repository that hosts released torch packages.
 # This will be dynamically set based on CUDA availability and CUDA backend enabled/disabled.
 TORCH_URL_BASE = "https://download.pytorch.org/whl"
-TORCHAO_URL_BASE = "https://download.pytorch.org/whl"
-TORCHAO_NIGHTLY_VERSION = "0.19.0"
-CU134_TORCHAO_NIGHTLY_VERSION = "0.19.0"
-ROCM_TORCHAO_NIGHTLY_VERSION = "0.19.0"
+TORCHAO_URL_BASE = "https://download.pytorch.org/whl/nightly"
+TORCHAO_NIGHTLY_VERSION = "0.19.0.dev20260907"
+CU134_TORCHAO_NIGHTLY_VERSION = "0.19.0.dev20260907"
+# Newest rocm7.2 TorchAO nightly with compiled kernels (cp310-abi3); later ROCm nightlies
+# are py3-none-any only, and none are published after 0.19.0.dev20260831.
+ROCM_TORCHAO_NIGHTLY_VERSION = "0.19.0.dev20260805"
 CU134_TORCH_PACKAGES = [
     "torch==2.15.0+cu134",
     "torchvision==0.29.2+cu134",
@@ -67,14 +69,12 @@ def install_requirements(use_pytorch_nightly):
     # Determine the appropriate PyTorch URL based on CUDA delegate status
     torch_url = determine_torch_url(TORCH_URL_BASE)
     cu134_packages = cu134_requirements(torch_url)
-    if cu134_packages:
-        torch_url = determine_torch_url(TORCHAO_URL_BASE)
-        if not use_pytorch_nightly:
-            cu134_packages[0] = "torch"
+    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 release is what the checkout is tested against, and the
+    # 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.
     if platform.machine().lower() in ("aarch64", "arm64"):
@@ -183,11 +183,14 @@ def install_optional_example_requirements(use_pytorch_nightly):
         if use_pytorch_nightly
         else []
     )
-    if cu134_packages:
-        torch_url = determine_torch_url(TORCHAO_URL_BASE)
+    torchao_url = (
+        f"{TORCHAO_URL_BASE}/cpu"
+        if platform.machine().lower() in ("aarch64", "arm64")
+        else determine_torch_url(TORCHAO_URL_BASE)
+    )
     torchao_index = (
-        ["--extra-index-url", f"{TORCHAO_URL_BASE}/cpu"]
-        if cu134_packages and platform.machine().lower() in ("aarch64", "arm64")
+        ["--extra-index-url", torchao_url]
+        if cu134_packages and not torchao_from_source()
         else []
     )
 

setup.py

@@ -1118,7 +1118,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 pin bump cannot move the
+    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.
@@ -1170,14 +1170,21 @@ def _base_dependencies() -> List[str]:
         "py-cpuinfo",
         "requests",
         "pytorch-tokenizers>=1.6.0",
+        "torch>=2.15.0,<2.16",
         # Shipped code imports torchao at module scope in many places, so a plain install cannot
         # lower a model without it. Among others: the XNNPACK utilities the partitioner uses
         # (backends/xnnpack/utils/utils.py), the Core ML quantizer, and executorch.export itself.
         # The MLX backend needs it too, though indirectly: it registers a torchao operator that
         # only exists once torchao has been imported.
         #
-        # Keep dependency bounds in the release series used by the installer.
-        "torch>=2.15.0,<2.16",
+        # The lower bound is the nightly install_requirements.py pins, so that installing this
+        # package over a development checkout leaves that pin in place. A bound at the stable
+        # release instead would evict it, because a dev release sorts below its own final.
+        #
+        # The upper bound is what makes naming a pre-release safe. A specifier that names one
+        # accepts pre-releases for this requirement, so without the bound pip would resolve a
+        # pre-release of the next series wherever one is published, such as an index carrying
+        # nightlies. Keep both bounds in the same series.
         _torchao_requirement(),
         "pyyaml",
         "ruamel.yaml",