Build with nightly Torch instead of the latest (#1681)

Summary:
Pull Request resolved: #1681

Changes in the PyTorch C++ layer affect Monarch's tensor engine builds. Without this change, the CI builds will pick up the latest PyTorch through build-requirements.txt. This diff removes this dependency.

We see this affecting #1680

Reviewed By: amirafzali

Differential Revision: D85670693

fbshipit-source-id: 1de2188ac5c84271c41bf253f9a12fd2e2cd66e6

Author: allenwang28

.github/workflows/build-cpu.yml

@@ -29,5 +29,9 @@ jobs:
         # Setup build environment (conda + system deps + rust + build deps)
         setup_build_environment
 
+        # Install torch nightly (CPU version)
+        pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
+        pip install -r build-requirements.txt
+
         # Build monarch (No tensor engine, CPU version)
         USE_TENSOR_ENGINE=0 python setup.py bdist_wheel

.github/workflows/build-cuda.yml

@@ -34,6 +34,10 @@ jobs:
         # Setup build environment (conda + system deps + rust + build deps)
         setup_build_environment
 
+        # Install torch nightly
+        pip install ${{ matrix.torch-spec }}
+        pip install -r build-requirements.txt
+
         # Setup Tensor Engine
         setup_tensor_engine
 

.github/workflows/test-cpu-python.yml

@@ -31,6 +31,9 @@ jobs:
         # Disable tensor engine
         export USE_TENSOR_ENGINE=0
 
+        # Install PyTorch nightly
+        pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
+
         # Install the built wheel from artifact
         install_wheel_from_artifact
 

.github/workflows/test-gpu-python.yml

@@ -45,6 +45,10 @@ jobs:
         # Setup CUDA environment and library paths
         setup_cuda_environment
 
+        # Install torch nightly before installing the wheel,
+        # so that we can test the wheel against the latest nightly
+        pip install ${{ matrix.torch-spec }}
+
         # Install the built wheel from artifact
         install_wheel_from_artifact
 

README.md

@@ -70,6 +70,9 @@ torchmonarch-nightly is built with torch nightly.
 
 ### Build and Install from Source
 
+If you're building Monarch from source, you should be building it with the nightly PyTorch as well for ABI compatibility.
+
+
 #### On Fedora distributions
 
 ```sh
@@ -87,7 +90,7 @@ rustup default nightly
 conda install libunwind -y
 
 # Install the correct cuda and cuda-toolkit versions for your machine
-sudo dnf install cuda-toolkit-12-0 cuda-12-0
+sudo dnf install cuda-toolkit-12-8 cuda-12-8
 
 # Install clang-dev and nccl-dev
 sudo dnf install clang-devel libnccl-devel
@@ -99,6 +102,7 @@ conda update -n monarchenv --all -c conda-forge -y
 sudo dnf install -y libibverbs rdma-core libmlx5 libibverbs-devel rdma-core-devel
 
 # Install build dependencies
+pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128
 pip install -r build-requirements.txt
 # Install test dependencies
 pip install -r python/tests/requirements.txt
@@ -134,7 +138,11 @@ sudo apt install -y clang
 export CC=clang
 export CXX=clang++
 
+# Install the correct cuda and cuda-toolkit versions for your machine
+sudo apt install -y cuda-toolkit-12-8 cuda-12-8
+
 # Install build dependencies
+pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128
 pip install -r build-requirements.txt
 # Install test dependencies
 pip install -r python/tests/requirements.txt
@@ -153,9 +161,9 @@ pip install --no-build-isolation -e .
 pip list | grep monarch
 ```
 
-#### On MacOS
+#### On non-CUDA machines
 
-You can also build Monarch to run locally on a MacOS system.
+You can also build Monarch to run on non-CUDA machines, e.g. locally on a MacOS system.
 
 Note that this does not support tensor engine, which is tied to CUDA and RDMA (via ibverbs).
 
@@ -172,6 +180,7 @@ rustup toolchain install nightly
 rustup default nightly
 
 # Install build dependencies
+pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
 pip install -r build-requirements.txt
 # Install test dependencies
 pip install -r python/tests/requirements.txt

scripts/common-setup.sh

@@ -40,11 +40,6 @@ setup_rust_toolchain() {
     cargo install cargo-nextest --locked
 }
 
-install_build_dependencies() {
-    echo "Installing build dependencies..."
-    pip install -r build-requirements.txt ${1:-}
-}
-
 # Install Python test dependencies
 install_python_test_dependencies() {
     echo "Installing test dependencies..."
@@ -115,7 +110,6 @@ setup_build_environment() {
     setup_conda_environment "${python_version}"
     install_system_dependencies
     setup_rust_toolchain
-    install_build_dependencies "${install_args}"
 }
 
 # Detect and configure CUDA environment for linking