Building torchao from source installs unnecessary torch and nvidia packages every time

[gau-nernst]

To reproduce

conda create -n test_ao python=3.10
conda activate test_ao
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
pip install git+https://github.com/pytorch/ao -v

Outputs (the packages are already cached in this case)

Using pip 24.0 from /home/ubuntu/miniconda3/envs/test_ao/lib/python3.10/site-packages/pip (python 3.10)
Collecting git+https://github.com/pytorch/ao
  Cloning https://github.com/pytorch/ao to /tmp/pip-req-build-bcuh0mqg
  Running command git version
  git version 2.34.1
  Running command git clone --filter=blob:none https://github.com/pytorch/ao /tmp/pip-req-build-bcuh0mqg
  Cloning into '/tmp/pip-req-build-bcuh0mqg'...
  Running command git rev-parse HEAD
  b91b6be24afd1220331790ff0866f5b091165cd5
  Resolved https://github.com/pytorch/ao to commit b91b6be24afd1220331790ff0866f5b091165cd5
  Running command git rev-parse HEAD
  b91b6be24afd1220331790ff0866f5b091165cd5
  Running command pip subprocess to install build dependencies
  Collecting setuptools
    Using cached setuptools-69.5.1-py3-none-any.whl.metadata (6.2 kB)
  Collecting wheel
    Using cached wheel-0.43.0-py3-none-any.whl.metadata (2.2 kB)
  Collecting ninja
    Using cached ninja-1.11.1.1-py2.py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl.metadata (5.3 kB)
  Collecting torch
    Using cached torch-2.3.0-cp310-cp310-manylinux1_x86_64.whl.metadata (26 kB)
  Collecting filelock (from torch)
    Using cached filelock-3.14.0-py3-none-any.whl.metadata (2.8 kB)
  Collecting typing-extensions>=4.8.0 (from torch)
    Using cached typing_extensions-4.11.0-py3-none-any.whl.metadata (3.0 kB)
  Collecting sympy (from torch)
    Using cached sympy-1.12-py3-none-any.whl.metadata (12 kB)
  Collecting networkx (from torch)
    Using cached networkx-3.3-py3-none-any.whl.metadata (5.1 kB)
  Collecting jinja2 (from torch)
    Using cached jinja2-3.1.4-py3-none-any.whl.metadata (2.6 kB)
  Collecting fsspec (from torch)
    Using cached fsspec-2024.3.1-py3-none-any.whl.metadata (6.8 kB)
  Collecting nvidia-cuda-nvrtc-cu12==12.1.105 (from torch)
    Using cached nvidia_cuda_nvrtc_cu12-12.1.105-py3-none-manylinux1_x86_64.whl.metadata (1.5 kB)
  Collecting nvidia-cuda-runtime-cu12==12.1.105 (from torch)
    Using cached nvidia_cuda_runtime_cu12-12.1.105-py3-none-manylinux1_x86_64.whl.metadata (1.5 kB)
  Collecting nvidia-cuda-cupti-cu12==12.1.105 (from torch)
    Using cached nvidia_cuda_cupti_cu12-12.1.105-py3-none-manylinux1_x86_64.whl.metadata (1.6 kB)
  Collecting nvidia-cudnn-cu12==8.9.2.26 (from torch)
    Using cached nvidia_cudnn_cu12-8.9.2.26-py3-none-manylinux1_x86_64.whl.metadata (1.6 kB)
  Collecting nvidia-cublas-cu12==12.1.3.1 (from torch)
    Using cached nvidia_cublas_cu12-12.1.3.1-py3-none-manylinux1_x86_64.whl.metadata (1.5 kB)
  Collecting nvidia-cufft-cu12==11.0.2.54 (from torch)
    Using cached nvidia_cufft_cu12-11.0.2.54-py3-none-manylinux1_x86_64.whl.metadata (1.5 kB)
  Collecting nvidia-curand-cu12==10.3.2.106 (from torch)
    Using cached nvidia_curand_cu12-10.3.2.106-py3-none-manylinux1_x86_64.whl.metadata (1.5 kB)
  Collecting nvidia-cusolver-cu12==11.4.5.107 (from torch)
    Using cached nvidia_cusolver_cu12-11.4.5.107-py3-none-manylinux1_x86_64.whl.metadata (1.6 kB)
  Collecting nvidia-cusparse-cu12==12.1.0.106 (from torch)
    Using cached nvidia_cusparse_cu12-12.1.0.106-py3-none-manylinux1_x86_64.whl.metadata (1.6 kB)
  Collecting nvidia-nccl-cu12==2.20.5 (from torch)
    Using cached nvidia_nccl_cu12-2.20.5-py3-none-manylinux2014_x86_64.whl.metadata (1.8 kB)
  Collecting nvidia-nvtx-cu12==12.1.105 (from torch)
    Using cached nvidia_nvtx_cu12-12.1.105-py3-none-manylinux1_x86_64.whl.metadata (1.7 kB)
  Collecting triton==2.3.0 (from torch)
    Using cached triton-2.3.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (1.4 kB)
  Collecting nvidia-nvjitlink-cu12 (from nvidia-cusolver-cu12==11.4.5.107->torch)
    Using cached nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)
  Collecting MarkupSafe>=2.0 (from jinja2->torch)
    Using cached MarkupSafe-2.1.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (3.0 kB)
  Collecting mpmath>=0.19 (from sympy->torch)
    Using cached mpmath-1.3.0-py3-none-any.whl.metadata (8.6 kB)
  Using cached setuptools-69.5.1-py3-none-any.whl (894 kB)
  Using cached wheel-0.43.0-py3-none-any.whl (65 kB)
  Using cached ninja-1.11.1.1-py2.py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl (307 kB)
  Using cached torch-2.3.0-cp310-cp310-manylinux1_x86_64.whl (779.1 MB)
  Using cached nvidia_cublas_cu12-12.1.3.1-py3-none-manylinux1_x86_64.whl (410.6 MB)
  Using cached nvidia_cuda_cupti_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (14.1 MB)
  Using cached nvidia_cuda_nvrtc_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (23.7 MB)
  Using cached nvidia_cuda_runtime_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (823 kB)
  Using cached nvidia_cudnn_cu12-8.9.2.26-py3-none-manylinux1_x86_64.whl (731.7 MB)
  Using cached nvidia_cufft_cu12-11.0.2.54-py3-none-manylinux1_x86_64.whl (121.6 MB)
  Using cached nvidia_curand_cu12-10.3.2.106-py3-none-manylinux1_x86_64.whl (56.5 MB)
  Using cached nvidia_cusolver_cu12-11.4.5.107-py3-none-manylinux1_x86_64.whl (124.2 MB)
  Using cached nvidia_cusparse_cu12-12.1.0.106-py3-none-manylinux1_x86_64.whl (196.0 MB)
  Using cached nvidia_nccl_cu12-2.20.5-py3-none-manylinux2014_x86_64.whl (176.2 MB)
  Using cached nvidia_nvtx_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (99 kB)
  Using cached triton-2.3.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (168.1 MB)
  Using cached typing_extensions-4.11.0-py3-none-any.whl (34 kB)
  Using cached filelock-3.14.0-py3-none-any.whl (12 kB)
  Using cached fsspec-2024.3.1-py3-none-any.whl (171 kB)
  Using cached jinja2-3.1.4-py3-none-any.whl (133 kB)
  Using cached networkx-3.3-py3-none-any.whl (1.7 MB)
  Using cached sympy-1.12-py3-none-any.whl (5.7 MB)
  Using cached MarkupSafe-2.1.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (25 kB)
  Using cached mpmath-1.3.0-py3-none-any.whl (536 kB)
  Using cached nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (21.1 MB)
  Installing collected packages: ninja, mpmath, wheel, typing-extensions, sympy, setuptools, nvidia-nvtx-cu12, nvidia-nvjitlink-cu12, nvidia-nccl-cu12, nvidia-curand-cu12, nvidia-cufft-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, networkx, MarkupSafe, fsspec, filelock, triton, nvidia-cusparse-cu12, nvidia-cudnn-cu12, jinja2, nvidia-cusolver-cu12, torch
  Successfully installed MarkupSafe-2.1.5 filelock-3.14.0 fsspec-2024.3.1 jinja2-3.1.4 mpmath-1.3.0 networkx-3.3 ninja-1.11.1.1 nvidia-cublas-cu12-12.1.3.1 nvidia-cuda-cupti-cu12-12.1.105 nvidia-cuda-nvrtc-cu12-12.1.105 nvidia-cuda-runtime-cu12-12.1.105 nvidia-cudnn-cu12-8.9.2.26 nvidia-cufft-cu12-11.0.2.54 nvidia-curand-cu12-10.3.2.106 nvidia-cusolver-cu12-11.4.5.107 nvidia-cusparse-cu12-12.1.0.106 nvidia-nccl-cu12-2.20.5 nvidia-nvjitlink-cu12-12.4.127 nvidia-nvtx-cu12-12.1.105 setuptools-69.5.1 sympy-1.12 torch-2.3.0 triton-2.3.0 typing-extensions-4.11.0 wheel-0.43.0
  Installing build dependencies ... done

On my machine, it takes 30s just to install the cached packages. Note that this is done every time I re-install torchao. During development, it's quite annoying to have this extra 30s every time I need to re-compile CUDA/C++ code (pip install -e . only works for python code).
If this is the first time installing torchao from source, there will be extra time downloading the packages (which are huge).

The culprit seems to be torch being a build-system.requires in pyproject.toml. Perhaps this is a limitation of pip not being able to recognize torch from conda? During the build process, it's also not clear if it is using existing torch (from conda) or pip-torch (may cause issues if the two versions mismatch? I'm using the latest version so issues may not arise).

[msaroufim]

Indeed this limitation was added because for whatever reason in CI machines only but not locally installing ao was crashing because the setup.py needed torch as a dependency. I think it's due to some new hermetic installation stuff pip install . but I haven't looked at it in much detail.